perf: batch DynamoDB degree queries

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(

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 = {}