restore and improve progress bar updates during adding nodes on archive import

Author: GeigerJ2

src/aiida/tools/archive/imports.py

@@ -109,7 +109,7 @@ def import_archive(
     if not (merge_extras[0] in ['k', 'n'] and merge_extras[1] in ['c', 'n'] and merge_extras[2] in ['l', 'u', 'd']):
         raise ValueError('merge_extras contains invalid values')
     if merge_comments not in ('leave', 'newest', 'overwrite'):
-        raise ValueError(f"merge_comments not in {('leave', 'newest', 'overwrite')!r}")
+        raise ValueError(f'merge_comments not in {("leave", "newest", "overwrite")!r}')
     type_check(group, orm.Group, allow_none=True)
     type_check(test_run, bool)
     backend = backend or get_manager().get_profile_storage()
@@ -232,23 +232,26 @@ def _add_new_entities(
     for (ufield,) in query.distinct().iterall(batch_size=batch_size):
         if ufield not in backend_unique_id:
             ufields.append(ufield)
-    
-    breakpoint()
-    with get_progress_reporter()(desc=f'Adding new {etype.value}(s)', total=total) as progress:
-        rows = [
-            transform(row)
-            for row in QueryBuilder(backend=backend_from)
-            .append(
+
+    with get_progress_reporter()(desc=f'Importing new {etype.value}(s)', total=total) as progress:
+        # For UX: batch large ID lists so queries start returning results faster
+        # Even though the improved IN clause handles any size, query planning for 500k+ IDs can be slow
+        query_batch_size = 50_000
+
+        # Batch the IDs for querying (UX optimization, not a technical requirement)
+        for _, ufields_batch in batch_iter(ufields, query_batch_size):
+            query = QueryBuilder(backend=backend_from).append(
                 entity_type_to_orm[etype],
-                filters={unique_field: {'in': ufields}},
+                filters={unique_field: {'in': ufields_batch}},
                 project=['**'],
                 tag='entity',
             )
-            .dict(batch_size=batch_size)
-        ]
-        new_ids = backend_to.bulk_insert(etype, rows)
-        backend_unique_id.update({row[unique_field]: pk for pk, row in zip(new_ids, rows)})
-        progress.update(len(rows))
+
+            # Batch the results processing for progress updates and memory efficiency
+            for nrows, rows_batch in batch_iter(query.dict(batch_size=batch_size), batch_size, transform):
+                new_ids = backend_to.bulk_insert(etype, rows_batch)
+                backend_unique_id.update({row[unique_field]: pk for pk, row in zip(new_ids, rows_batch)})
+                progress.update(nrows)
 
 
 def _import_users(