Merge pull request #637 from Labelbox/jtso/al-3045

[AL-3045] Max Datarow Upload Per Operation

Author: jtsodapop

labelbox/schema/dataset.py

@@ -15,12 +15,16 @@
 from labelbox.exceptions import InvalidQueryError, LabelboxError, ResourceNotFoundError, InvalidAttributeError
 from labelbox.orm.db_object import DbObject, Updateable, Deletable
 from labelbox.orm.model import Entity, Field, Relationship
+from labelbox.exceptions import MalformedQueryException
 
 if TYPE_CHECKING:
     from labelbox import Task, User, DataRow
 
 logger = logging.getLogger(__name__)
 
+MAX_DATAROW_PER_API_OPERATION = 150000
+MAX_DATAROW_WITH_METADATA = 30000
+
 
 class Dataset(DbObject, Updateable, Deletable):
     """ A Dataset is a collection of DataRows.
@@ -391,14 +395,18 @@ def convert_item(item):
                 f"Must pass an iterable to create_data_rows. Found {type(items)}"
             )
 
+        if len(items) > MAX_DATAROW_PER_API_OPERATION:
+            raise MalformedQueryException(
+                f"Cannot create more than {MAX_DATAROW_PER_API_OPERATION} DataRows per function call."
+            )
+
         # TODO: If any datarows contain metadata, we're limiting max # of datarows
         # until we address performance issues with datarow create with metadata
-        max_datarow_with_metadata = 30_000
-        if (len(items) > max_datarow_with_metadata):
+        if len(items) > MAX_DATAROW_WITH_METADATA:
             for row in items:
                 if 'metadata_fields' in row:
-                    raise ValueError(
-                        f"Cannot create more than {max_datarow_with_metadata} DataRows, if any DataRows contain metadata"
+                    raise MalformedQueryException(
+                        f"Cannot create more than {MAX_DATAROW_WITH_METADATA} DataRows, if any DataRows contain metadata"
                     )
 
         with ThreadPoolExecutor(file_upload_thread_count) as executor:

tests/integration/test_dataset.py

@@ -2,7 +2,8 @@
 import pytest
 import requests
 from labelbox import Dataset
-from labelbox.exceptions import ResourceNotFoundError
+from labelbox.exceptions import ResourceNotFoundError, MalformedQueryException
+from labelbox.schema.dataset import MAX_DATAROW_PER_API_OPERATION
 
 
 def test_dataset(client, rand_gen):
@@ -137,3 +138,12 @@ def test_create_descriptor_file(dataset):
             'content_type': 'application/json',
             'filename': 'json_import.json'
         }
+
+
+def test_max_dataset_datarow_upload(dataset, image_url, rand_gen):
+    external_id = str(rand_gen)
+    items = [dict(row_data=image_url, external_id=external_id)
+            ] * (MAX_DATAROW_PER_API_OPERATION + 1)
+
+    with pytest.raises(MalformedQueryException):
+        dataset.create_data_rows(items)