Batching on MastMissions.get_product_list

Author: snbianco

astroquery/mast/missions.py

@@ -21,6 +21,7 @@
 from astroquery import log
 from astroquery.utils import commons, async_to_sync
 from astroquery.utils.class_or_instance import class_or_instance
+from astropy.utils.console import ProgressBarOrSpinner
 from astroquery.exceptions import InvalidQueryError, MaxResultsWarning, InputWarning, NoResultsWarning
 
 from astroquery.mast import utils
@@ -106,7 +107,13 @@ def _parse_result(self, response, *, verbose=False):  # Used by the async_to_syn
                               MaxResultsWarning)
         elif self.service == self._list_products:
             # Results from post_list_products endpoint need to be handled differently
-            results = Table(response.json()['products'])
+            if isinstance(response, list):  # multiple async responses from batching
+                combined_products = []
+                for resp in response:
+                    combined_products.extend(resp.json().get('products', []))
+                return Table(combined_products)
+
+            results = Table(response.json()['products'])  # single async response
 
         return results
 
@@ -370,13 +377,32 @@ def get_product_list_async(self, datasets):
                             'list of strings, Astropy Row, Astropy Column, or Astropy Table.')
 
         # Filter out empty strings from IDs
-        datasets = [item.strip() for item in datasets if item.strip() != '' and item is not None]
-        if not len(datasets):
+        datasets = [item.strip() for item in datasets if item and item.strip()]
+        if not datasets:
             raise InvalidQueryError("Dataset list is empty, no associated products.")
 
-        # Send async service request
-        params = {'dataset_ids': datasets}
-        return self._service_api_connection.missions_request_async(self.service, params)
+        # Filter out duplicates
+        datasets = list(set(datasets))
+
+        # Batch API calls if number of datasets exceeds maximum
+        max_batch = 1000
+        num_datasets = len(datasets)
+        if num_datasets > max_batch:
+            # Split datasets into chunks
+            dataset_chunks = list(utils.split_list_into_chunks(datasets, max_batch))
+
+            results = []  # list to store responses from each batch
+            with ProgressBarOrSpinner(len(dataset_chunks), f'Fetching products for {num_datasets} unique datasets '
+                                      f'in {len(dataset_chunks)} batches ...') as pb:
+                for i, chunk in enumerate(dataset_chunks):
+                    pb.update(i)
+                    params = {'dataset_ids': chunk}
+                    results.append(self._service_api_connection.missions_request_async(self.service, params))
+            return results
+        else:
+            # Single batch request
+            params = {'dataset_ids': datasets}
+            return self._service_api_connection.missions_request_async(self.service, params)
 
     def get_unique_product_list(self, datasets):
         """

astroquery/mast/tests/test_mast.py

@@ -344,6 +344,11 @@ def test_missions_get_product_list(patch_post):
     result = mast.MastMissions.get_product_list(datasets[0])
     assert isinstance(result, Table)
 
+    # Batching
+    dataset_list = [f'{i}' for i in range(1001)]
+    result = mast.MastMissions.get_product_list(dataset_list)
+    assert isinstance(result, Table)
+
 
 def test_missions_get_unique_product_list(patch_post, caplog):
     unique_products = mast.MastMissions.get_unique_product_list('Z14Z0104T')

astroquery/mast/tests/test_mast_remote.py

@@ -188,31 +188,40 @@ def test_missions_get_product_list_async(self):
             MastMissions.get_product_list_async([' '])
         assert 'Dataset list is empty' in str(err_empty.value)
 
-    def test_missions_get_product_list(self):
+    def test_missions_get_product_list(self, capsys):
         datasets = MastMissions.query_object("M4", radius=0.1)
         test_dataset = datasets[0]['sci_data_set_name']
         multi_dataset = list(datasets[:2]['sci_data_set_name'])
 
         # Compare Row input and string input
-        result1 = MastMissions.get_product_list(test_dataset)
-        result2 = MastMissions.get_product_list(datasets[0])
-        assert isinstance(result1, Table)
-        assert len(result1) == len(result2)
-        assert set(result1['filename']) == set(result2['filename'])
+        result_str = MastMissions.get_product_list(test_dataset)
+        result_row = MastMissions.get_product_list(datasets[0])
+        assert isinstance(result_str, Table)
+        assert len(result_str) == len(result_row)
+        assert set(result_str['filename']) == set(result_row['filename'])
 
         # Compare Table input and list input
-        result1 = MastMissions.get_product_list(multi_dataset)
-        result2 = MastMissions.get_product_list(datasets[:2])
-        assert isinstance(result1, Table)
-        assert len(result1) == len(result2)
-        assert set(result1['filename']) == set(result2['filename'])
+        result_list = MastMissions.get_product_list(multi_dataset)
+        result_table = MastMissions.get_product_list(datasets[:2])
+        assert isinstance(result_list, Table)
+        assert len(result_list) == len(result_table)
+        assert set(result_list['filename']) == set(result_table['filename'])
 
         # Filter datasets based on sci_data_set_name and verify products
         filtered = datasets[datasets['sci_data_set_name'] == 'IBKH03020']
         result = MastMissions.get_product_list(filtered)
         assert isinstance(result, Table)
         assert (result['dataset'] == 'IBKH03020').all()
 
+        # Test batching by creating a list of 1001 different strings
+        # This won't return any results, but will test the batching
+        dataset_list = [f'{i}' for i in range(1001)]
+        result = MastMissions.get_product_list(dataset_list)
+        out, _ = capsys.readouterr()
+        assert isinstance(result, Table)
+        assert len(result) == 0
+        assert 'Fetching products for 1001 unique datasets in 2 batches' in out
+
     def test_missions_get_unique_product_list(self, caplog):
         # Check that no rows are filtered out when all products are unique
         dataset_ids = ['JBTAA8010']

astroquery/mast/utils.py

@@ -159,6 +159,26 @@ def parse_input_location(coordinates=None, objectname=None):
     return obj_coord
 
 
+def split_list_into_chunks(input_list, chunk_size):
+    """
+    Splits a list into chunks of a specified size.
+
+    Parameters
+    ----------
+    input_list : list
+        List to be split into chunks.
+    chunk_size : int
+        Size of each chunk.
+
+    Yields
+    ------
+    chunk : list
+        A chunk of the input list.
+    """
+    for idx in range(0, len(input_list), chunk_size):
+        yield input_list[idx:idx + chunk_size]
+
+
 def mast_relative_path(mast_uri):
     """
     Given one or more MAST dataURI(s), return the associated relative path(s).
@@ -180,7 +200,7 @@ def mast_relative_path(mast_uri):
 
     # Split the list into chunks of 50 URIs; this is necessary
     # to avoid "414 Client Error: Request-URI Too Large".
-    uri_list_chunks = list(_split_list_into_chunks(uri_list, chunk_size=50))
+    uri_list_chunks = list(split_list_into_chunks(uri_list, chunk_size=50))
 
     result = []
     for chunk in uri_list_chunks:
@@ -214,12 +234,6 @@ def mast_relative_path(mast_uri):
     return result
 
 
-def _split_list_into_chunks(input_list, chunk_size):
-    """Helper function for `mast_relative_path`."""
-    for idx in range(0, len(input_list), chunk_size):
-        yield input_list[idx:idx + chunk_size]
-
-
 def remove_duplicate_products(data_products, uri_key):
     """
     Removes duplicate data products that have the same data URI.