diff --git a/grand/backends/_dynamodb.py b/grand/backends/_dynamodb.py index dcfe5b4..64e06ac 100644 --- a/grand/backends/_dynamodb.py +++ b/grand/backends/_dynamodb.py @@ -502,6 +502,60 @@ def get_edge_count(self) -> int: "ItemCount" ] + def degrees(self, nbunch=None) -> Collection: + if nbunch is not None and not isinstance(nbunch, (list, tuple)): + return self.degree(nbunch) + + requested = None if nbunch is None else list(nbunch) + degrees = {} if requested is None else {node: 0 for node in requested} + requested_by_id = ( + None if requested is None else {str(node): node for node in requested} + ) + for edge in self._scan_table(self._edge_table): + source = edge[self._edge_source_key] + target = edge[self._edge_target_key] + if requested_by_id is None: + degrees[source] = degrees.get(source, 0) + 1 + if self._directed or target != source: + degrees[target] = degrees.get(target, 0) + 1 + continue + if source in requested_by_id: + node = requested_by_id[source] + degrees[node] += 1 + if target in requested_by_id and (self._directed or target != source): + node = requested_by_id[target] + degrees[node] += 1 + return degrees + + def in_degrees(self, nbunch=None) -> Collection: + if not self._directed: + return self.degrees(nbunch) + return self._directed_bulk_degrees(nbunch, self._edge_target_key) + + def out_degrees(self, nbunch=None) -> Collection: + if not self._directed: + return self.degrees(nbunch) + return self._directed_bulk_degrees(nbunch, self._edge_source_key) + + def _directed_bulk_degrees(self, nbunch, endpoint_key): + if nbunch is not None and not isinstance(nbunch, (list, tuple)): + if endpoint_key == self._edge_target_key: + return super().in_degree(nbunch) + return super().out_degree(nbunch) + + requested = None if nbunch is None else list(nbunch) + degrees = {} if requested is None else {node: 0 for node in requested} + requested_by_id = ( + None if requested is None else {str(node): node for node in requested} + ) + for edge in self._scan_table(self._edge_table): + endpoint = edge[endpoint_key] + if requested_by_id is None: + degrees[endpoint] = degrees.get(endpoint, 0) + 1 + elif endpoint in requested_by_id: + degrees[requested_by_id[endpoint]] += 1 + return degrees + # Ingesting def ingest_from_edgelist_dataframe( diff --git a/grand/backends/test_dynamodb.py b/grand/backends/test_dynamodb.py index 2c26aec..031676b 100644 --- a/grand/backends/test_dynamodb.py +++ b/grand/backends/test_dynamodb.py @@ -135,6 +135,69 @@ def test_undirected_neighbors_query_both_indexes_and_deduplicate(backend): assert backend._edge_table.scan.call_args_list == [] +def test_directed_bulk_degrees_scan_edges_once(backend): + backend._edge_table.scan.side_effect = [ + { + "Items": [ + {"ID": "ab", "Source": "A", "Target": "B"}, + {"ID": "ac", "Source": "A", "Target": "C"}, + ], + "LastEvaluatedKey": {"ID": "ac"}, + }, + {"Items": [{"ID": "ca", "Source": "C", "Target": "A"}]}, + ] + + assert backend.degrees(["A", "B", "C", "D"]) == { + "A": 3, + "B": 1, + "C": 2, + "D": 0, + } + first_scan, second_scan = backend._edge_table.scan.call_args_list + assert first_scan.kwargs == {} + assert second_scan.kwargs == {"ExclusiveStartKey": {"ID": "ac"}} + backend._edge_table.query.assert_not_called() + + +def test_directed_bulk_in_and_out_degrees_scan_once(backend): + edges = [ + {"ID": "12", "Source": "1", "Target": "2"}, + {"ID": "13", "Source": "1", "Target": "3"}, + ] + backend._edge_table.scan.return_value = {"Items": edges} + + assert backend.out_degrees([1, 2, 3]) == {1: 2, 2: 0, 3: 0} + backend._edge_table.scan.assert_called_once_with() + backend._edge_table.reset_mock() + backend._edge_table.scan.return_value = {"Items": edges} + assert backend.in_degrees([1, 2, 3]) == {1: 0, 2: 1, 3: 1} + backend._edge_table.scan.assert_called_once_with() + backend._edge_table.query.assert_not_called() + + +def test_directed_all_degrees_scan_once(backend): + backend._edge_table.scan.return_value = { + "Items": [{"ID": "ab", "Source": "A", "Target": "B"}] + } + + assert backend.degrees() == {"A": 1, "B": 1} + backend._edge_table.scan.assert_called_once_with() + + +def test_undirected_bulk_degrees_scan_edges_once_and_deduplicate(backend): + backend._directed = False + backend._edge_table.scan.return_value = { + "Items": [ + {"ID": "ab", "Source": "A", "Target": "B"}, + {"ID": "aa", "Source": "A", "Target": "A"}, + ] + } + + assert backend.degrees(["A", "B", "C"]) == {"A": 2, "B": 1, "C": 0} + backend._edge_table.scan.assert_called_once_with() + backend._edge_table.query.assert_not_called() + + def test_add_edge_uses_collision_safe_bounded_identity(backend): backend._node_table.get_item.return_value = {"Item": {"ID": "exists"}} backend._edge_table.get_item.return_value = {}