paginate mapping endpoints (#1506)

closes lig-4613 - paginate mapping endpoint to allow >400k samples without overloading the API
lightly-ai · Feb 19, 2024 · 51db826 · 51db826
1 parent f779e58
commit 51db826
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/lightly/api/api_workflow_client.py b/lightly/api/api_workflow_client.py
@@ -221,8 +221,13 @@ def get_filenames(self) -> List[str]:
 
         :meta private:  # Skip docstring generation
         """
-        filenames_on_server = self._mappings_api.get_sample_mappings_by_dataset_id(
-            dataset_id=self.dataset_id, field="fileName"
+        filenames_on_server = list(
+            utils.paginate_endpoint(
+                self._mappings_api.get_sample_mappings_by_dataset_id,
+                page_size=25000,
+                dataset_id=self.dataset_id,
+                field="fileName",
+            )
         )
         return filenames_on_server
 

diff --git a/lightly/api/api_workflow_download_dataset.py b/lightly/api/api_workflow_download_dataset.py
@@ -93,8 +93,13 @@ def download_dataset(
             )
 
         # get sample ids
-        sample_ids = self._mappings_api.get_sample_mappings_by_dataset_id(
-            self.dataset_id, field="_id"
+        sample_ids = list(
+            utils.paginate_endpoint(
+                self._mappings_api.get_sample_mappings_by_dataset_id,
+                page_size=25000,
+                dataset_id=self.dataset_id,
+                field="_id",
+            )
         )
 
         indices = BitMask.from_hex(tag.bit_mask_data).to_indices()