Remove 150K data rows limitation

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/dataset.py

@@ -256,10 +256,10 @@ def create_data_rows_sync(self, items) -> None:
             url_param: descriptor_url
         })
 
-    def create_data_rows(
-            self,
-            items,
-            file_upload_thread_count=FILE_UPLOAD_THREAD_COUNT) -> "Task":
+    def create_data_rows(self,
+                         items,
+                         file_upload_thread_count=FILE_UPLOAD_THREAD_COUNT
+                        ) -> "DataUpsertTask":
         """ Asynchronously bulk upload data rows
 
         Use this instead of `Dataset.create_data_rows_sync` uploads for batches that contain more than 1000 data rows.
@@ -572,10 +572,10 @@ def _export(
         is_streamable = res["isStreamable"]
         return Task.get_task(self.client, task_id), is_streamable
 
-    def upsert_data_rows(
-            self,
-            items,
-            file_upload_thread_count=FILE_UPLOAD_THREAD_COUNT) -> "Task":
+    def upsert_data_rows(self,
+                         items,
+                         file_upload_thread_count=FILE_UPLOAD_THREAD_COUNT
+                        ) -> "DataUpsertTask":
         """
         Upserts data rows in this dataset. When "key" is provided, and it references an existing data row,
         an update will be performed. When "key" is not provided a new data row will be created.
@@ -606,18 +606,14 @@ def upsert_data_rows(
         >>>     ])
         >>>     task.wait_till_done()
         """
-        if len(items) > MAX_DATAROW_PER_API_OPERATION:
-            raise MalformedQueryException(
-                f"Cannot upsert more than {MAX_DATAROW_PER_API_OPERATION} DataRows per function call."
-            )
-
         specs = DataRowUpsertItem.build(self.uid, items)
         return self._exec_upsert_data_rows(specs, file_upload_thread_count)
 
     def _exec_upsert_data_rows(
-            self,
-            specs: List[DataRowItemBase],
-            file_upload_thread_count: int = FILE_UPLOAD_THREAD_COUNT) -> "Task":
+        self,
+        specs: List[DataRowItemBase],
+        file_upload_thread_count: int = FILE_UPLOAD_THREAD_COUNT
+    ) -> "DataUpsertTask":
 
         manifest = DataRowUploader.upload_in_chunks(
             client=self.client,

libs/labelbox/tests/data/annotation_import/test_model_run.py

@@ -174,6 +174,7 @@ def test_model_run_split_assignment_by_data_row_ids(
     data_rows = dataset.create_data_rows([{
         "row_data": image_url
     } for _ in range(n_data_rows)])
+    data_rows.wait_till_done()
     data_row_ids = [data_row['id'] for data_row in data_rows.result]
     configured_project_with_one_data_row._wait_until_data_rows_are_processed(
         data_row_ids=data_row_ids)

libs/labelbox/tests/integration/test_task.py

@@ -58,11 +58,12 @@ def test_task_success_json(dataset, image_url, snapshot):
                           'test_task.test_task_success_json.json')
     assert len(task.result)
 
+
 @pytest.mark.export_v1("export_v1 test remove later")
 def test_task_success_label_export(client, configured_project_with_label):
     project, _, _, _ = configured_project_with_label
     # TODO: Move to export_v2
-    res = project.export_labels()
+    project.export_labels()
     user = client.get_user()
     task = None
     for task in user.created_tasks():