reset changes to age_rules.py and graph_traversers.py

Author: GeigerJ2

src/aiida/tools/graph/age_rules.py

@@ -221,30 +221,10 @@ def _load_results(self, target_set, operational_set):
         target_set.empty()
 
         if primkeys:
-            primkeys_list = list(primkeys)
-            batch_size = 10000  # Stay well under 65535 parameter limit
-
-            # If we have fewer keys than the batch size, use the original approach
-            if len(primkeys_list) <= batch_size:
-                self._querybuilder.add_filter(
-                    self._first_tag, {operational_set[self._entity_from].identifier: {'in': primkeys}}
-                )
-                qres = self._querybuilder.dict()
-            else:
-                # Batch the queries for large datasets
-                all_results = []
-                for i in range(0, len(primkeys_list), batch_size):
-                    batch_primkeys = primkeys_list[i:i + batch_size]
-
-                    # Use deepcopy only when we need to batch
-                    batch_qb = deepcopy(self._querybuilder)
-                    batch_qb.add_filter(
-                        self._first_tag, {operational_set[self._entity_from].identifier: {'in': batch_primkeys}}
-                    )
-                    batch_results = batch_qb.dict()
-                    all_results.extend(batch_results)
-
-                qres = all_results
+            self._querybuilder.add_filter(
+                self._first_tag, {operational_set[self._entity_from].identifier: {'in': primkeys}}
+            )
+            qres = self._querybuilder.dict()
 
             # These are the new results returned by the query
             target_set[self._entity_to].add_entities(

src/aiida/tools/graph/graph_traversers.py

@@ -240,17 +240,9 @@ def traverse_graph(
             return {'nodes': set(), 'links': set()}
         return {'nodes': set(), 'links': None}
 
-    # Batch the query to avoid parameter limits
-    existing_pks = set()
-    operational_list = list(operational_set)
-    batch_size = 10000  # Stay well under 65535 parameter limit
-
-    for i in range(0, len(operational_list), batch_size):
-        batch_ids = operational_list[i:i + batch_size]
-        query_nodes = orm.QueryBuilder(backend=backend)
-        query_nodes.append(orm.Node, project=['id'], filters={'id': {'in': batch_ids}})
-        existing_pks.update(query_nodes.all(flat=True))
-
+    query_nodes = orm.QueryBuilder(backend=backend)
+    query_nodes.append(orm.Node, project=['id'], filters={'id': {'in': operational_set}})
+    existing_pks = set(query_nodes.all(flat=True))
     missing_pks = operational_set.difference(existing_pks)
     if missing_pks and missing_callback is None:
         raise exceptions.NotExistent(
@@ -302,4 +294,3 @@ def traverse_graph(
         output['links'] = results['nodes_nodes'].keyset
 
     return output
-