Merge pull request #571 from Labelbox/kkim/AL-2219

[AL-2219] Add custom_metadata to input file for create_data_rows()

Author: kkim-labelbox

labelbox/schema/data_row_metadata.py

@@ -307,7 +307,7 @@ def _batch_upsert(
                     data_row_id=m.data_row_id,
                     fields=list(
                         chain.from_iterable(
-                            self.parse_upsert(m) for m in m.fields))).dict(
+                            self._parse_upsert(m) for m in m.fields))).dict(
                                 by_alias=True))
         res = _batch_operations(_batch_upsert, items, self._batch_size)
         return res
@@ -404,7 +404,7 @@ def _bulk_export(_data_row_ids: List[str]) -> List[DataRowMetadata]:
                                  data_row_ids,
                                  batch_size=self._batch_size)
 
-    def parse_upsert(
+    def _parse_upsert(
             self, metadatum: DataRowMetadataField
     ) -> List[_UpsertDataRowMetadataInput]:
         """Format for metadata upserts to GQL"""
@@ -435,6 +435,33 @@ def parse_upsert(
 
         return [_UpsertDataRowMetadataInput(**p) for p in parsed]
 
+    # Convert metadata to DataRowMetadataField objects, parse all fields
+    # and return a dictionary of metadata fields for upsert
+    def parse_upsert_metadata(self, metadata_fields):
+
+        def _convert_metadata_field(metadata_field):
+            if isinstance(metadata_field, DataRowMetadataField):
+                return metadata_field
+            elif isinstance(metadata_field, dict):
+                if not all(key in metadata_field
+                           for key in ("schema_id", "value")):
+                    raise ValueError(
+                        f"Custom metadata field '{metadata_field}' must have 'schema_id' and 'value' keys"
+                    )
+                return DataRowMetadataField(
+                    schema_id=metadata_field["schema_id"],
+                    value=metadata_field["value"])
+            else:
+                raise ValueError(
+                    f"Metadata field '{metadata_field}' is neither 'DataRowMetadataField' type or a dictionary"
+                )
+
+        # Convert all metadata fields to DataRowMetadataField type
+        metadata_fields = [_convert_metadata_field(m) for m in metadata_fields]
+        parsed_metadata = list(
+            chain.from_iterable(self._parse_upsert(m) for m in metadata_fields))
+        return [m.dict(by_alias=True) for m in parsed_metadata]
+
     def _validate_delete(self, delete: DeleteDataRowMetadata):
         if not len(delete.fields):
             raise ValueError(f"No fields specified for {delete.data_row_id}")

labelbox/schema/dataset.py

@@ -5,7 +5,7 @@
 from collections.abc import Iterable
 import time
 import ndjson
-from itertools import islice, chain
+from itertools import islice
 
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import StringIO
@@ -82,13 +82,8 @@ def create_data_row(self, **kwargs) -> "DataRow":
         # Parse metadata fields, if they are provided
         if DataRow.custom_metadata.name in kwargs:
             mdo = self.client.get_data_row_metadata_ontology()
-            metadata_fields = kwargs[DataRow.custom_metadata.name]
-            metadata = list(
-                chain.from_iterable(
-                    mdo.parse_upsert(m) for m in metadata_fields))
-            kwargs[DataRow.custom_metadata.name] = [
-                md.dict(by_alias=True) for md in metadata
-            ]
+            kwargs[DataRow.custom_metadata.name] = mdo.parse_upsert_metadata(
+                kwargs[DataRow.custom_metadata.name])
 
         return self.client._create(DataRow, kwargs)
 
@@ -268,6 +263,13 @@ def validate_attachments(item):
                     )
             return attachments
 
+        def parse_metadata_fields(item):
+            metadata_fields = item.get('custom_metadata')
+            if metadata_fields:
+                mdo = self.client.get_data_row_metadata_ontology()
+                item['custom_metadata'] = mdo.parse_upsert_metadata(
+                    metadata_fields)
+
         def format_row(item):
             # Formats user input into a consistent dict structure
             if isinstance(item, dict):
@@ -308,6 +310,8 @@ def convert_item(item):
             validate_keys(item)
             # Make sure attachments are valid
             validate_attachments(item)
+            # Parse metadata fields if they exist
+            parse_metadata_fields(item)
             # Upload any local file paths
             item = upload_if_necessary(item)
 

tests/integration/test_batch.py

@@ -36,7 +36,7 @@ def test_create_batch(configured_project: Project, big_dataset: Dataset):
 
     data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
     batch = configured_project.create_batch("test-batch", data_rows, 3)
-    assert batch.name == 'test-batch'
+    assert batch.name == "test-batch"
     assert batch.size == len(data_rows)
 
 
@@ -79,4 +79,4 @@ def test_export_data_rows(configured_project: Project, dataset: Dataset):
     exported_data_rows = [dr.uid for dr in result]
 
     assert len(result) == n_data_rows
-    assert set(data_rows) == set(exported_data_rows)
+    assert set(data_rows) == set(exported_data_rows)

tests/integration/test_data_row_metadata.py

@@ -281,4 +281,4 @@ def test_parse_raw_metadata(mdo):
 
     for row in parsed:
         for field in row.fields:
-            assert mdo.parse_upsert(field)
+            assert mdo._parse_upsert(field)

tests/integration/test_data_rows.py

@@ -15,6 +15,10 @@
 EMBEDDING_SCHEMA_ID = "ckpyije740000yxdk81pbgjdc"
 TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
 CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
+EXPECTED_METADATA_SCHEMA_IDS = [
+    SPLIT_SCHEMA_ID, TEST_SPLIT_ID, EMBEDDING_SCHEMA_ID, TEXT_SCHEMA_ID,
+    CAPTURE_DT_SCHEMA_ID
+].sort()
 
 
 def make_metadata_fields():
@@ -31,6 +35,27 @@ def make_metadata_fields():
     return fields
 
 
+def make_metadata_fields_dict():
+    embeddings = [0.0] * 128
+    msg = "A message"
+    time = datetime.utcnow()
+
+    fields = [{
+        "schema_id": SPLIT_SCHEMA_ID,
+        "value": TEST_SPLIT_ID
+    }, {
+        "schema_id": CAPTURE_DT_SCHEMA_ID,
+        "value": time
+    }, {
+        "schema_id": TEXT_SCHEMA_ID,
+        "value": msg
+    }, {
+        "schema_id": EMBEDDING_SCHEMA_ID,
+        "value": embeddings
+    }]
+    return fields
+
+
 def test_get_data_row(datarow, client):
     assert client.get_data_row(datarow.uid)
 
@@ -152,7 +177,7 @@ def test_data_row_single_creation(dataset, rand_gen, image_url):
         assert requests.get(data_row_2.row_data).content == data
 
 
-def test_data_row_single_creation_with_metadata(dataset, rand_gen, image_url):
+def test_create_data_row_with_metadata(dataset, image_url):
     client = dataset.client
     assert len(list(dataset.data_rows())) == 0
 
@@ -167,38 +192,133 @@ def test_data_row_single_creation_with_metadata(dataset, rand_gen, image_url):
         requests.get(data_row.row_data).content
     assert data_row.media_attributes is not None
     assert len(data_row.custom_metadata) == 5
+    assert [m["schemaId"] for m in data_row.custom_metadata
+           ].sort() == EXPECTED_METADATA_SCHEMA_IDS
 
-    with NamedTemporaryFile() as fp:
-        data = rand_gen(str).encode()
-        fp.write(data)
-        fp.flush()
-        data_row_2 = dataset.create_data_row(row_data=fp.name)
-        assert len(list(dataset.data_rows())) == 2
-        assert requests.get(data_row_2.row_data).content == data
 
+def test_create_data_row_with_metadata_dict(dataset, image_url):
+    client = dataset.client
+    assert len(list(dataset.data_rows())) == 0
 
-def test_data_row_single_creation_with_invalid_metadata(dataset, image_url):
+    data_row = dataset.create_data_row(
+        row_data=image_url, custom_metadata=make_metadata_fields_dict())
 
-    def make_invalid_metadata_fields():
-        embeddings = [0.0] * 128
-        msg = "A message"
-        time = datetime.utcnow()
+    assert len(list(dataset.data_rows())) == 1
+    assert data_row.dataset() == dataset
+    assert data_row.created_by() == client.get_user()
+    assert data_row.organization() == client.get_organization()
+    assert requests.get(image_url).content == \
+        requests.get(data_row.row_data).content
+    assert data_row.media_attributes is not None
+    assert len(data_row.custom_metadata) == 5
+    assert [m["schemaId"] for m in data_row.custom_metadata
+           ].sort() == EXPECTED_METADATA_SCHEMA_IDS
 
-        fields = [
-            DataRowMetadataField(schema_id=SPLIT_SCHEMA_ID,
-                                 value=TEST_SPLIT_ID),
-            DataRowMetadataField(schema_id=CAPTURE_DT_SCHEMA_ID, value=time),
-            DataRowMetadataField(schema_id=TEXT_SCHEMA_ID, value=msg),
-            DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID,
-                                 value=embeddings),
-            DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID,
-                                 value=embeddings),
-        ]
-        return fields
+
+def test_create_data_row_with_invalid_metadata(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append(
+        DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID, value=[0.0] * 128))
 
     with pytest.raises(labelbox.exceptions.MalformedQueryException) as excinfo:
-        dataset.create_data_row(row_data=image_url,
-                                custom_metadata=make_invalid_metadata_fields())
+        dataset.create_data_row(row_data=image_url, custom_metadata=fields)
+
+
+def test_create_data_rows_with_metadata(dataset, image_url):
+    client = dataset.client
+    assert len(list(dataset.data_rows())) == 0
+
+    task = dataset.create_data_rows([
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row1",
+            DataRow.custom_metadata: make_metadata_fields()
+        },
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row2",
+            "custom_metadata": make_metadata_fields()
+        },
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row3",
+            DataRow.custom_metadata: make_metadata_fields_dict()
+        },
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row4",
+            "custom_metadata": make_metadata_fields_dict()
+        },
+    ])
+    task.wait_till_done()
+
+    assert len(list(dataset.data_rows())) == 4
+    for r in ["row1", "row2", "row3", "row4"]:
+        row = list(dataset.data_rows(where=DataRow.external_id == r))[0]
+        assert row.dataset() == dataset
+        assert row.created_by() == client.get_user()
+        assert row.organization() == client.get_organization()
+        assert requests.get(image_url).content == \
+            requests.get(row.row_data).content
+        assert row.media_attributes is not None
+        assert len(row.custom_metadata) == 5
+        assert [m["schemaId"] for m in row.custom_metadata
+               ].sort() == EXPECTED_METADATA_SCHEMA_IDS
+
+
+def test_create_data_rows_with_invalid_metadata(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append(
+        DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID, value=[0.0] * 128))
+
+    task = dataset.create_data_rows([{
+        DataRow.row_data: image_url,
+        DataRow.custom_metadata: fields
+    }])
+    task.wait_till_done()
+    assert task.status == "FAILED"
+
+
+def test_create_data_rows_with_metadata_missing_value(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append({"schemaId": "some schema id"})
+
+    with pytest.raises(ValueError) as exc:
+        dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+                DataRow.custom_metadata: fields
+            },
+        ])
+
+
+def test_create_data_rows_with_metadata_missing_schema_id(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append({"value": "some value"})
+
+    with pytest.raises(ValueError) as exc:
+        dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+                DataRow.custom_metadata: fields
+            },
+        ])
+
+
+def test_create_data_rows_with_metadata_wrong_type(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append("Neither DataRowMetadataField or dict")
+
+    with pytest.raises(ValueError) as exc:
+        task = dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+                DataRow.custom_metadata: fields
+            },
+        ])
 
 
 def test_data_row_update(dataset, rand_gen, image_url):