temporarily reorganize symmetry query around splitting into separate query chunks

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -659,7 +659,7 @@ def _query_resource(
             use_document_model: if None, will defer to the self.use_document_model attribute
             num_chunks: Maximum number of chunks of data to yield. None will yield all possible.
             chunk_size: Number of data entries per chunk.
-            timeout : Time in seconds to wait until a request timeout error is thrown
+            timeout (float or None): Time in seconds to wait until a request timeout error is thrown
 
         Returns:
             A Resource, a dict with two keys, "data" containing a list of documents, and
@@ -810,9 +810,9 @@ def _submit_requests(
         url: str,
         criteria: dict[str, Any],
         use_document_model: bool,
-        chunk_size: int,
+        chunk_size: int | None,
         num_chunks: int | None = None,
-        timeout: float | None = None,
+        timeout: int | None = None,
         max_batch_size: int = 100,
         norecur: bool = False,
     ) -> dict:
@@ -826,8 +826,8 @@ def _submit_requests(
             criteria (dict of str): dictionary of criteria to filter down
             use_document_model (bool): whether to use the document model
             num_chunks (int or None): Maximum number of chunks of data to yield. None will yield all possible.
-            chunk_size (int): Number of data entries per chunk.
-            timeout (float): Time in seconds to wait until a request timeout error is thrown
+            chunk_size (int or None): Number of data entries per chunk.
+            timeout (int or None): Time in seconds to wait until a request timeout error is thrown
             max_batch_size (int) : Maximum size of a batch when retrieving batches in parallel
             norecur (bool) : Whether to forbid recursive splitting of a query field
                 when a direct query fails
@@ -985,6 +985,12 @@ def _submit_requests(
             if "meta" in data:
                 total_data["meta"] = data["meta"]
 
+        # otherwise, paginate sequentially
+        if chunk_size is None or chunk_size < 1:
+            raise ValueError(
+                "A positive chunk size must be provided to enable pagination"
+            )
+
         # Get max number of response pages
         max_pages = (
             num_chunks if num_chunks is not None else ceil(total_num_docs / chunk_size)
@@ -1024,10 +1030,6 @@ def _submit_requests(
                 pbar.close()
             return new_total_data
 
-        # otherwise, paginate sequentially
-        if chunk_size is None:
-            raise ValueError("A chunk size must be provided to enable pagination")
-
         # Warning to select specific fields only for many results
         if criteria.get("_all_fields", False) and (total_num_docs / chunk_size > 10):
             warnings.warn(

mp_api/client/routes/materials/summary.py

@@ -2,6 +2,7 @@
 
 import warnings
 from collections import defaultdict
+from itertools import chain, product
 
 from emmet.core.summary import HasProps, SummaryDoc
 from emmet.core.symmetry import CrystalSystem
@@ -200,8 +201,9 @@ def search(  # noqa: D417
         mmnd_inv = {v: k for k, v in min_max_name_dict.items() if k != v}
 
         # Set user query params from `locals`
+        _locals = locals()
         user_settings = {
-            k: v for k, v in locals().items() if k in min_max_name_dict and v
+            k: v for k, v in _locals.items() if k in min_max_name_dict and v
         }
 
         # Check to see if user specified _search fields using **kwargs,
@@ -328,10 +330,11 @@ def _csrc(x):
             "spacegroup_number": 230,
             "spacegroup_symbol": 230,
         }
+        batched_symm_query = {}
         for k, cardinality in symm_cardinality.items():
-            if isinstance(symm_vals := locals().get(k), list | tuple | set):
+            if isinstance(symm_vals := _locals.get(k), list | tuple | set):
                 if len(symm_vals) < cardinality // 2:
-                    query_params.update({k: ",".join(str(v) for v in symm_vals)})
+                    batched_symm_query[k] = symm_vals
                 else:
                     raise MPRestError(
                         f"Querying `{k}` by a list of values is only "
@@ -378,6 +381,24 @@ def _csrc(x):
             if query_params[entry] is not None
         }
 
+        if batched_symm_query:
+            ordered_symm_key = sorted(batched_symm_query)
+            return list(
+                chain.from_iterable(
+                    self._search(  # type: ignore[return-value]
+                        num_chunks=num_chunks,
+                        chunk_size=chunk_size,
+                        all_fields=all_fields,
+                        fields=fields,
+                        **query_params,
+                        **{sk: symm_params[i] for i, sk in enumerate(ordered_symm_key)},
+                    )
+                    for symm_params in product(
+                        *[batched_symm_query[k] for k in ordered_symm_key]
+                    )
+                )
+            )
+
         return super()._search(  # type: ignore[return-value]
             num_chunks=num_chunks,
             chunk_size=chunk_size,