Adjust tag propagation in EmbeddingOperator (#146)

This limits the propagation of tags from the input id column to the output embedding column.

Author: karlhigley

merlin/dataloader/ops/embeddings.py

@@ -98,10 +98,23 @@ def compute_output_schema(
             col_schemas.append(col_schema)
         id_schema = input_schema.column_schemas[self.lookup_key]
         embedding_dim = self.embeddings.shape[1]
+
+        new_tags = [Tags.EMBEDDING]
+        propagated_tags = [
+            Tags.LIST,
+            Tags.SEQUENCE,
+            Tags.USER,
+            Tags.ITEM,
+            Tags.CONTEXT,
+            Tags.SESSION,
+        ]
+        for tag in propagated_tags:
+            if tag in id_schema.tags:
+                new_tags.append(tag)
         col_schemas.append(
             ColumnSchema(
                 name=self.embedding_name,
-                tags=[Tags.EMBEDDING],
+                tags=new_tags,
                 dtype=self.embeddings.dtype,
                 dims=id_schema.shape.as_tuple + (embedding_dim,),
             )

tests/unit/dataloader/test_embeddings.py

@@ -25,6 +25,7 @@
 from merlin.dataloader.ops.padding import Padding
 from merlin.io import Dataset
 from merlin.schema import Tags
+from merlin.schema.tags import TagSet
 from merlin.table import TensorColumn, TensorTable
 
 
@@ -71,7 +72,7 @@ def test_embedding_np_mmap_dl_no_lookup(tmpdir, embedding_ids, np_embeddings_fro
     dataset = dataset.repartition(10)
     schema = dataset.schema
     for col_name in cat_names:
-        schema[col_name] = schema[col_name].with_tags([Tags.CATEGORICAL, Tags.EMBEDDING])
+        schema[col_name] = schema[col_name].with_tags([Tags.CATEGORICAL, Tags.ITEM, Tags.ID])
     dataset.schema = schema
     data_loader = Loader(
         dataset,
@@ -85,8 +86,10 @@ def test_embedding_np_mmap_dl_no_lookup(tmpdir, embedding_ids, np_embeddings_fro
     assert data_loader.output_schema.column_names == ["id", "embeddings"]
 
     embeddings_dim = 1024
-    embeddings_value_count = data_loader.output_schema["embeddings"].value_count
+    embedding_schema = data_loader.output_schema["embeddings"]
+    embeddings_value_count = embedding_schema.value_count
     assert embeddings_value_count.min == embeddings_value_count.max == embeddings_dim
+    assert embedding_schema.tags == TagSet([Tags.EMBEDDING, Tags.ITEM])
 
     full_len = 0
     for idx, batch in enumerate(data_loader):