[AL-4256] Create datarow metadata by name

whistler · whistler · commit 5838ea61013c · 2022-12-22T05:25:35.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,8 @@
 # Version 3.34.0 (...)
 ### Added
 * Added `get_by_name()` method to MetadataOntology object to access both custom and reserved metadata by name.
+* Added support for adding metadata by name when creating datarows using `DataRowMetadataOntology.bulk_upsert()`.
+* Added support for adding metadata by name when creating datarows using `Dataset.create_data_rows()`, `Dataset.create_data_rows_sync()`, and `Dataset.create_data_row()`.
 
 ### Changed
 * `Dataset.create_data_rows()` max limit of DataRows increased to 150,000
diff --git a/examples/basics/data_row_metadata.ipynb b/examples/basics/data_row_metadata.ipynb
@@ -39,7 +39,7 @@
       "source": [
         "# Data Row Metadata\n",
         "\n",
-        "Metadata is useful to be better understand data on the platform to help with labeling review, model diagnostics, and data selection. This **should not be confused with attachments**. Attachments provide additional context for labelers but is not searchable within Catalog."
+        "Metadata is useful to better understand data on the platform to help with labeling review, model diagnostics, and data selection. This **should not be confused with attachments**. Attachments provide additional context for labelers but is not searchable within Catalog."
       ]
     },
     {
@@ -261,21 +261,20 @@
       "source": [
         "# Construct a metadata field of string kind\n",
         "tag_metadata_field = DataRowMetadataField(\n",
-        "    schema_id=mdo.reserved_by_name[\"tag\"].uid,  # specify the schema id\n",
+        "    name=\"tag\",  # specify the schema name\n",
         "    value=\"tag_string\", # typed inputs\n",
         ")\n",
         "\n",
         "# Construct an metadata field of datetime kind\n",
         "capture_datetime_field = DataRowMetadataField(\n",
-        "    schema_id=mdo.reserved_by_name[\"captureDateTime\"].uid,  # specify the schema id\n",
+        "    name=\"captureDateTime\",  # specify the schema id\n",
         "    value=datetime.utcnow(), # typed inputs\n",
         ")\n",
         "\n",
         "# Construct a metadata field of Enums options\n",
-        "train_schema = mdo.reserved_by_name[\"split\"][\"train\"]\n",
         "split_metadta_field = DataRowMetadataField(\n",
-        "    schema_id=train_schema.parent,  # specify the schema id\n",
-        "    value=train_schema.uid, # typed inputs\n",
+        "    name=\"split\",  # specify the schema id\n",
+        "    value=\"train\", # typed inputs\n",
         ")"
       ]
     },
@@ -300,20 +299,20 @@
       "source": [
         "# Construct a dictionary of string metadata\n",
         "tag_metadata_field_dict = {\n",
-        "    \"schema_id\": mdo.reserved_by_name[\"tag\"].uid,\n",
+        "    \"name\": \"tag\",\n",
         "    \"value\": \"tag_string\",\n",
         "}\n",
         "\n",
         "# Construct a dictionary of datetime metadata\n",
         "capture_datetime_field_dict = {\n",
-        "    \"schema_id\": mdo.reserved_by_name[\"captureDateTime\"].uid,\n",
+        "    \"name\": \"captureDateTime\",\n",
         "    \"value\": datetime.utcnow(),\n",
         "}\n",
         "\n",
         "# Construct a dictionary of Enums options metadata\n",
         "split_metadta_field_dict = {\n",
-        "    \"schema_id\": mdo.reserved_by_name[\"split\"][\"train\"].parent,\n",
-        "    \"value\": mdo.reserved_by_name[\"split\"][\"train\"].uid,\n",
+        "    \"name\": \"split\",\n",
+        "    \"value\": \"train\",\n",
         "}"
       ]
     },
@@ -491,7 +490,7 @@
       "outputs": [],
       "source": [
         "# Select a dataset to use, or you can just use the 1-image dataset created above. \n",
-        "dataset_id = \"cl3ntfr7j7cmh07bmeqz3gfjt\"\n",
+        "dataset_id = dataset.uid\n",
         "dataset = client.get_dataset(dataset_id)"
       ]
     },
@@ -541,11 +540,11 @@
         "    # assign datarows a split\n",
         "    rnd = random.random()\n",
         "    if rnd < test:\n",
-        "        split = mdo.reserved_by_name[\"split\"][\"test\"]\n",
+        "        split = \"test\"\n",
         "    elif rnd < valid:\n",
-        "        split = mdo.reserved_by_name[\"split\"][\"valid\"]\n",
+        "        split = \"valid\"\n",
         "    else:\n",
-        "        split = mdo.reserved_by_name[\"split\"][\"train\"]\n",
+        "        split = \"train\"\n",
         "\n",
         "    embeddings.append(\n",
         "        list(model(processor(response.content), training=False)[0].numpy()))\n",
@@ -557,12 +556,11 @@
         "            data_row_id=datarow.uid,\n",
         "            fields=[\n",
         "                DataRowMetadataField(\n",
-        "                    schema_id=mdo.reserved_by_name[\"captureDateTime\"].uid,\n",
+        "                    name=\"captureDateTime\",\n",
         "                    value=dt,\n",
         "                ),\n",
-        "                DataRowMetadataField(schema_id=split.parent, value=split.uid),\n",
-        "                DataRowMetadataField(schema_id=mdo.reserved_by_name[\"tag\"].uid,\n",
-        "                                     value=message),\n",
+        "                DataRowMetadataField(name=\"split\", value=split),\n",
+        "                DataRowMetadataField(name=\"tag\", value=message),\n",
         "            ]))"
       ]
     },
@@ -620,7 +618,7 @@
         "for md, embd in zip(uploads, projected):\n",
         "    md.fields.append(\n",
         "        DataRowMetadataField(\n",
-        "            schema_id=mdo.reserved_by_name[\"embedding\"].uid,\n",
+        "            name=\"embedding\",\n",
         "            value=embd.tolist(),  # convert from numpy to list\n",
         "        ),)"
       ]
@@ -801,4 +799,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 5
-}
+}
diff --git a/labelbox/schema/data_row_metadata.py b/labelbox/schema/data_row_metadata.py
@@ -38,7 +38,10 @@ class DataRowMetadataSchema(BaseModel):
 
 # Metadata base class
 class DataRowMetadataField(_CamelCaseMixin):
-    schema_id: SchemaId
+    # One of `schema_id` or `name` must be provided. If `schema_id` is not provided, it is
+    # inferred from `name`
+    schema_id: Optional[SchemaId] = None
+    name: Optional[str] = None
     # value is of type `Any` so that we do not improperly coerce the value to the wrong tpye
     # Additional validation is performed before upload using the schema information
     value: Any
@@ -147,6 +150,16 @@ def _build_ontology(self):
             str, DataRowMetadataSchema] = self._make_normalized_name_index(
                 self.custom_fields)
 
+    @staticmethod
+    def _lookup_in_index_by_name(reserved_index, custom_index, name):
+        # search through reserved names first
+        if name in reserved_index:
+            return reserved_index[name]
+        elif name in custom_index:
+            return custom_index[name]
+        else:
+            raise KeyError(f"There is no metadata with name {name}")
+
     def get_by_name(
         self, name: str
     ) -> Union[DataRowMetadataSchema, Dict[str, DataRowMetadataSchema]]:
@@ -163,14 +176,17 @@ def get_by_name(
         Raises:
             KeyError: When provided name is not presented in neither reserved nor custom metadata list
         """
+        return self._lookup_in_index_by_name(self.reserved_by_name,
+                                             self.custom_by_name, name)
 
-        # search through reserved names first
-        if name in self.reserved_by_name:
-            return self.reserved_by_name[name]
-        elif name in self.custom_by_name:
-            return self.custom_by_name[name]
-        else:
-            raise KeyError(f"There is no metadata with name {name}")
+    def _get_by_name_normalized(self, name: str) -> DataRowMetadataSchema:
+        """ Get metadata by name. For options, it provides the option schema instead of list of 
+        options 
+        """
+        # using `normalized` indices to find options by name as well
+        return self._lookup_in_index_by_name(self.reserved_by_name_normalized,
+                                             self.custom_by_name_normalized,
+                                             name)
 
     @staticmethod
     def _make_name_index(
@@ -452,6 +468,8 @@ def parse_metadata_fields(
             else:
                 field = DataRowMetadataField(schema_id=schema.uid,
                                              value=f["value"])
+
+            field.name = schema.name
             parsed.append(field)
         return parsed
 
@@ -624,13 +642,17 @@ 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")):
+                if not "value" in metadata_field:
+                    raise ValueError(
+                        f"Custom metadata field '{metadata_field}' must have a 'value' key"
+                    )
+                if not "schema_id" in metadata_field and not "name" in metadata_field:
                     raise ValueError(
-                        f"Custom metadata field '{metadata_field}' must have 'schema_id' and 'value' keys"
+                        f"Custom metadata field '{metadata_field}' must have either 'schema_id' or 'name' key"
                     )
                 return DataRowMetadataField(
-                    schema_id=metadata_field["schema_id"],
+                    schema_id=metadata_field.get("schema_id"),
+                    name=metadata_field.get("name"),
                     value=metadata_field["value"])
             else:
                 raise ValueError(
@@ -664,11 +686,32 @@ def _upsert_schema(
         self.refresh_ontology()
         return _parse_metadata_schema(res)
 
+    def _load_option_by_name(self, metadatum: DataRowMetadataField):
+        is_value_a_valid_schema_id = metadatum.value in self.fields_by_id
+        if not is_value_a_valid_schema_id:
+            metadatum.value = self.get_by_name(
+                metadatum.name)[metadatum.value].uid
+
+    def _load_schema_id_by_name(self, metadatum: DataRowMetadataField):
+        """
+        Loads schema id by name for a metadata field including options schema id.
+        """
+        if metadatum.name is None:
+            return
+
+        if metadatum.schema_id is None:
+            schema = self._get_by_name_normalized(metadatum.name)
+            metadatum.schema_id = schema.uid
+            if schema.options:
+                self._load_option_by_name(metadatum)
+
     def _parse_upsert(
             self, metadatum: DataRowMetadataField
     ) -> List[_UpsertDataRowMetadataInput]:
         """Format for metadata upserts to GQL"""
 
+        self._load_schema_id_by_name(metadatum)
+
         if metadatum.schema_id not in self.fields_by_id:
             # Fetch latest metadata ontology if metadata can't be found
             self.refresh_ontology()
diff --git a/tests/integration/test_data_row_metadata.py b/tests/integration/test_data_row_metadata.py
@@ -18,6 +18,7 @@
 TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
 CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
 PRE_COMPUTED_EMBEDDINGS_ID = 'ckrzang79000008l6hb5s6za1'
+CUSTOM_TEXT_SCHEMA_NAME = 'custom_text'
 
 FAKE_NUMBER_FIELD = {
     "id": FAKE_SCHEMA_ID,
@@ -32,6 +33,7 @@ def mdo(client):
     mdo = client.get_data_row_metadata_ontology()
     for schema in mdo.custom_fields:
         mdo.delete_schema(schema.name)
+    mdo.create_schema(CUSTOM_TEXT_SCHEMA_NAME, DataRowMetadataKind.string)
     mdo._raw_ontology = mdo._get_ontology()
     mdo._raw_ontology.append(FAKE_NUMBER_FIELD)
     mdo._build_ontology()
@@ -69,6 +71,25 @@ def make_metadata(dr_id) -> DataRowMetadata:
     return metadata
 
 
+def make_named_metadata(dr_id) -> DataRowMetadata:
+    embeddings = [0.0] * 128
+    msg = "A message"
+    time = datetime.utcnow()
+
+    metadata = DataRowMetadata(data_row_id=dr_id,
+                               fields=[
+                                   DataRowMetadataField(name='split',
+                                                        value=TEST_SPLIT_ID),
+                                   DataRowMetadataField(name='captureDateTime',
+                                                        value=time),
+                                   DataRowMetadataField(
+                                       name=CUSTOM_TEXT_SCHEMA_NAME, value=msg),
+                                   DataRowMetadataField(name='embedding',
+                                                        value=embeddings),
+                               ])
+    return metadata
+
+
 def test_export_empty_metadata(configured_project_with_label):
     project, _, _, _ = configured_project_with_label
     # Wait for exporter to retrieve latest labels
@@ -81,7 +102,7 @@ def test_export_empty_metadata(configured_project_with_label):
 def test_get_datarow_metadata_ontology(mdo):
     assert len(mdo.fields)
     assert len(mdo.reserved_fields)
-    assert len(mdo.custom_fields) == 1
+    assert len(mdo.custom_fields) == 2
 
     split = mdo.reserved_by_name["split"]["train"]
 
@@ -129,6 +150,48 @@ def test_large_bulk_upsert_datarow_metadata(big_dataset, mdo):
         ]), metadata_lookup.get(data_row_id).fields
 
 
+def test_upsert_datarow_metadata_by_name(datarow, mdo):
+    metadata = [make_named_metadata(datarow.uid)]
+    errors = mdo.bulk_upsert(metadata)
+    assert len(errors) == 0
+
+    metadata_lookup = {
+        metadata.data_row_id: metadata
+        for metadata in mdo.bulk_export([datarow.uid])
+    }
+    assert len([
+        f for f in metadata_lookup.get(datarow.uid).fields
+        if f.schema_id != PRE_COMPUTED_EMBEDDINGS_ID
+    ]), metadata_lookup.get(datarow.uid).fields
+
+
+def test_upsert_datarow_metadata_option_by_name(datarow, mdo):
+    metadata = DataRowMetadata(data_row_id=datarow.uid,
+                               fields=[
+                                   DataRowMetadataField(name='split',
+                                                        value='test'),
+                               ])
+    errors = mdo.bulk_upsert([metadata])
+    assert len(errors) == 0
+
+    datarows = mdo.bulk_export([datarow.uid])
+    assert len(datarows[0].fields) == 1
+    metadata = datarows[0].fields[0]
+    assert metadata.schema_id == SPLIT_SCHEMA_ID
+    assert metadata.name == 'test'
+    assert metadata.value == TEST_SPLIT_ID
+
+
+def test_upsert_datarow_metadata_option_by_incorrect_name(datarow, mdo):
+    metadata = DataRowMetadata(data_row_id=datarow.uid,
+                               fields=[
+                                   DataRowMetadataField(name='split',
+                                                        value='test1'),
+                               ])
+    with pytest.raises(KeyError):
+        mdo.bulk_upsert([metadata])
+
+
 def test_bulk_delete_datarow_metadata(datarow, mdo):
     """test bulk deletes for all fields"""
     metadata = make_metadata(datarow.uid)
diff --git a/tests/integration/test_data_rows.py b/tests/integration/test_data_rows.py
