Merge pull request #307 from Labelbox/ms/metadata-tests

metadata updates

Author: msokoloff1

examples/basics/data_row_metadata.ipynb

@@ -169,18 +169,6 @@
     "train_field = mdo.reserved_by_name[\"split\"][\"train\"]"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "uOS2QlHmqAIs",
-   "metadata": {
-    "id": "uOS2QlHmqAIs"
-   },
-   "outputs": [],
-   "source": [
-    "split_field.options"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -567,7 +555,6 @@
     "fields = []\n",
     "# iterate through the fields you want to delete\n",
     "for field in md.fields:\n",
-    "    schema = mdo.field_by_index[field.schema_id]\n",
     "    fields.append(field.schema_id)\n",
     "\n",
     "deletes = DeleteDataRowMetadata(\n",
@@ -649,4 +636,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

labelbox/client.py

@@ -1,6 +1,8 @@
 # type: ignore
 from datetime import datetime, timezone
 import json
+from typing import List, Dict
+from collections import defaultdict
 
 import logging
 import mimetypes
@@ -658,3 +660,29 @@ def create_model(self, name, ontology_id):
             "ontologyId": ontology_id
         })
         return Model(self, result['createModel'])
+
+    def get_data_row_ids_for_external_ids(
+            self, external_ids: List[str]) -> Dict[str, List[str]]:
+        """
+        Returns a list of data row ids for a list of external ids.
+        There is a max of 1500 items returned at a time.
+
+        Args:
+            external_ids: List of external ids to fetch data row ids for
+
+        Returns:
+            A dict of external ids as keys and values as a list of data row ids that correspond to that external id.
+        """
+        query_str = """query externalIdsToDataRowIdsPyApi($externalId_in: [String!]!){
+            externalIdsToDataRowIds(externalId_in: $externalId_in) { dataRowId externalId }
+        }
+        """
+        max_n_per_request = 100
+        result = defaultdict(list)
+        for i in range(0, len(external_ids), max_n_per_request):
+            for row in self.execute(
+                    query_str,
+                {'externalId_in': external_ids[i:i + max_n_per_request]
+                })['externalIdsToDataRowIds']:
+                result[row['externalId']].append(row['dataRowId'])
+        return result

labelbox/schema/data_row_metadata.py

@@ -1,5 +1,5 @@
 # type: ignore
-import datetime
+from datetime import datetime
 import warnings
 from copy import deepcopy
 from enum import Enum
@@ -42,7 +42,7 @@ def id(self):
 
 # Constraints for metadata values
 Embedding: Type[List[float]] = conlist(float, min_items=128, max_items=128)
-DateTime: Type[datetime.datetime] = datetime.datetime  # must be in UTC
+DateTime: Type[datetime] = datetime  # must be in UTC
 String: Type[str] = constr(max_length=500)
 OptionId: Type[SchemaId] = SchemaId  # enum option
 Number: Type[float] = float
@@ -62,7 +62,7 @@ class Config:
 # Metadata base class
 class DataRowMetadataField(_CamelCaseMixin):
     schema_id: SchemaId
-    value: Any
+    value: Union[DataRowMetadataValue, _DataRowMetadataValuePrimitives]
 
 
 class DataRowMetadata(_CamelCaseMixin):
@@ -489,7 +489,6 @@ def _validate_parse_number(
 
 def _validate_parse_datetime(
         field: DataRowMetadataField) -> List[Dict[str, Union[SchemaId, str]]]:
-    # TODO: better validate tzinfo
     return [{
         "schemaId": field.schema_id,
         "value": field.value.isoformat() + "Z",  # needs to be UTC

tests/integration/test_data_row_metadata.py

@@ -7,6 +7,7 @@
 from labelbox.schema.data_row_metadata import DataRowMetadataField, DataRowMetadata, DeleteDataRowMetadata, \
     DataRowMetadataOntology
 
+INVALID_SCHEMA_ID = "1" * 25
 FAKE_SCHEMA_ID = "0" * 25
 FAKE_DATAROW_ID = "D" * 25
 SPLIT_SCHEMA_ID = "cko8sbczn0002h2dkdaxb5kal"
@@ -15,6 +16,7 @@
 EMBEDDING_SCHEMA_ID = "ckpyije740000yxdk81pbgjdc"
 TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
 CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
+PRE_COMPUTED_EMBEDDINGS_ID = 'ckrzang79000008l6hb5s6za1'
 
 FAKE_NUMBER_FIELD = {
     "id": FAKE_SCHEMA_ID,
@@ -40,24 +42,13 @@ def big_dataset(dataset: Dataset, image_url):
             "row_data": image_url,
             "external_id": "my-image"
         },
-    ] * 250)
+    ] * 5)
     task.wait_till_done()
 
     yield dataset
     dataset.delete()
 
 
-def wait_for_embeddings_svc(data_row_ids, mdo):
-    for idx in range(5):
-        if all([
-                len(metadata.fields)
-                for metadata in mdo.bulk_export(data_row_ids)
-        ]):
-            return
-        time.sleep((idx + 1)**2)
-    raise Exception("Embedding svc failed to update metadata.")
-
-
 def make_metadata(dr_id) -> DataRowMetadata:
     embeddings = [0.0] * 128
     msg = "A message"
@@ -97,18 +88,20 @@ def test_get_datarow_metadata_ontology(mdo):
 
 
 def test_bulk_upsert_datarow_metadata(datarow, mdo: DataRowMetadataOntology):
-    wait_for_embeddings_svc([datarow.uid], mdo)
     metadata = make_metadata(datarow.uid)
     mdo.bulk_upsert([metadata])
-    assert len(mdo.bulk_export([datarow.uid]))
-    assert len(mdo.bulk_export([datarow.uid])[0].fields) == 5
+    exported = mdo.bulk_export([datarow.uid])
+    assert len(exported)
+    assert len([
+        field for field in exported[0].fields
+        if field.schema_id != PRE_COMPUTED_EMBEDDINGS_ID
+    ]) == 4
 
 
 @pytest.mark.slow
 def test_large_bulk_upsert_datarow_metadata(big_dataset, mdo):
     metadata = []
     data_row_ids = [dr.uid for dr in big_dataset.data_rows()]
-    wait_for_embeddings_svc(data_row_ids, mdo)
     for data_row_id in data_row_ids:
         metadata.append(make_metadata(data_row_id))
     errors = mdo.bulk_upsert(metadata)
@@ -119,14 +112,16 @@ def test_large_bulk_upsert_datarow_metadata(big_dataset, mdo):
         for metadata in mdo.bulk_export(data_row_ids)
     }
     for data_row_id in data_row_ids:
-        assert len(metadata_lookup.get(data_row_id).fields)
+        assert len([
+            f for f in metadata_lookup.get(data_row_id).fields
+            if f.schema_id != PRE_COMPUTED_EMBEDDINGS_ID
+        ]), metadata_lookup.get(data_row_id).fields
 
 
 def test_bulk_delete_datarow_metadata(datarow, mdo):
     """test bulk deletes for all fields"""
     metadata = make_metadata(datarow.uid)
     mdo.bulk_upsert([metadata])
-
     assert len(mdo.bulk_export([datarow.uid])[0].fields)
     upload_ids = [m.schema_id for m in metadata.fields[:-2]]
     mdo.bulk_delete(
@@ -155,7 +150,6 @@ def test_bulk_partial_delete_datarow_metadata(datarow, mdo):
 def test_large_bulk_delete_datarow_metadata(big_dataset, mdo):
     metadata = []
     data_row_ids = [dr.uid for dr in big_dataset.data_rows()]
-    wait_for_embeddings_svc(data_row_ids, mdo)
     for data_row_id in data_row_ids:
         metadata.append(
             DataRowMetadata(data_row_id=data_row_id,
@@ -181,29 +175,33 @@ def test_large_bulk_delete_datarow_metadata(big_dataset, mdo):
     errors = mdo.bulk_delete(deletes)
     assert len(errors) == 0
     for data_row_id in data_row_ids:
-        # 2 remaining because we delete the user provided embedding but text and labelbox generated embeddings still exist
-        fields = mdo.bulk_export([data_row_id])[0].fields
-        assert len(fields) == 2
+        fields = [
+            f for f in mdo.bulk_export([data_row_id])[0].fields
+            if f.schema_id != PRE_COMPUTED_EMBEDDINGS_ID
+        ]
+        assert len(fields) == 1, fields
         assert EMBEDDING_SCHEMA_ID not in [field.schema_id for field in fields]
 
 
 def test_bulk_delete_datarow_enum_metadata(datarow: DataRow, mdo):
     """test bulk deletes for non non fields"""
-    wait_for_embeddings_svc([datarow.uid], mdo)
     metadata = make_metadata(datarow.uid)
     metadata.fields = [
         m for m in metadata.fields if m.schema_id == SPLIT_SCHEMA_ID
     ]
     mdo.bulk_upsert([metadata])
 
-    assert len(mdo.bulk_export([datarow.uid])[0].fields) == len(
+    exported = mdo.bulk_export([datarow.uid])[0].fields
+    assert len(exported) == len(
         set([x.schema_id for x in metadata.fields] +
-            [x.schema_id for x in mdo.bulk_export([datarow.uid])[0].fields]))
+            [x.schema_id for x in exported]))
 
     mdo.bulk_delete([
         DeleteDataRowMetadata(data_row_id=datarow.uid, fields=[SPLIT_SCHEMA_ID])
     ])
-    assert len(mdo.bulk_export([datarow.uid])[0].fields) == 1
+    exported = mdo.bulk_export([datarow.uid])[0].fields
+    assert len(
+        [f for f in exported if f.schema_id != PRE_COMPUTED_EMBEDDINGS_ID]) == 0
 
 
 def test_raise_enum_upsert_schema_error(datarow, mdo):
@@ -223,7 +221,7 @@ def test_upsert_non_existent_schema_id(datarow, mdo):
     metadata = DataRowMetadata(data_row_id=datarow.uid,
                                fields=[
                                    DataRowMetadataField(
-                                       schema_id=FAKE_SCHEMA_ID,
+                                       schema_id=INVALID_SCHEMA_ID,
                                        value="message"),
                                ])
     with pytest.raises(ValueError):

tests/integration/test_data_rows.py

@@ -1,4 +1,6 @@
 from tempfile import NamedTemporaryFile
+import uuid
+import time
 
 import pytest
 import requests
@@ -11,6 +13,36 @@ def test_get_data_row(datarow, client):
     assert client.get_data_row(datarow.uid)
 
 
+def test_lookup_data_rows(client, dataset):
+    uid = str(uuid.uuid4())
+    # 1 external id : 1 uid
+    dr = dataset.create_data_row(row_data="123", external_id=uid)
+    lookup = client.get_data_row_ids_for_external_ids([uid])
+    assert len(lookup) == 1
+    assert lookup[uid][0] == dr.uid
+    # 2 external ids : 1 uid
+    uid2 = str(uuid.uuid4())
+    dr2 = dataset.create_data_row(row_data="123", external_id=uid2)
+    lookup = client.get_data_row_ids_for_external_ids([uid, uid2])
+    assert len(lookup) == 2
+    assert all([len(x) == 1 for x in lookup.values()])
+    assert lookup[uid][0] == dr.uid
+    assert lookup[uid2][0] == dr2.uid
+    #1 external id : 2 uid
+    dr3 = dataset.create_data_row(row_data="123", external_id=uid2)
+    lookup = client.get_data_row_ids_for_external_ids([uid2])
+    assert len(lookup) == 1
+    assert len(lookup[uid2]) == 2
+    assert lookup[uid2][0] == dr2.uid
+    assert lookup[uid2][1] == dr3.uid
+    # Empty args
+    lookup = client.get_data_row_ids_for_external_ids([])
+    assert len(lookup) == 0
+    # Non matching
+    lookup = client.get_data_row_ids_for_external_ids([str(uuid.uuid4())])
+    assert len(lookup) == 0
+
+
 def test_data_row_bulk_creation(dataset, rand_gen, image_url):
     client = dataset.client
     assert len(list(dataset.data_rows())) == 0