Add entity join key and fix entity references (feast-dev#1429)

* Add entity join key Signed-off-by: Jacob Klegar <jacob@tecton.ai> * Fix test Signed-off-by: Jacob Klegar <jacob@tecton.ai> * Fix more tests Signed-off-by: Jacob Klegar <jacob@tecton.ai> * Simplify online retrieval code Signed-off-by: Willem Pienaar <git@willem.co> * Fix linting Signed-off-by: Willem Pienaar <git@willem.co> Co-authored-by: Willem Pienaar <git@willem.co>
woop · Apr 10, 2021 · ec9c4fd · ec9c4fd
1 parent f371c2d
commit ec9c4fd
Show file tree

Hide file tree

Showing 14 changed files with 326 additions and 256 deletions.
diff --git a/protos/feast/core/Entity.proto b/protos/feast/core/Entity.proto
@@ -44,6 +44,9 @@ message EntitySpecV2 {
     // Description of the entity.
     string description = 3;
 
+    // Join key for the entity (i.e. name of the column the entity maps to).
+    string join_key = 4;
+
     // User defined metadata
     map<string,string> labels = 8;
 }

diff --git a/protos/feast/types/EntityKey.proto b/protos/feast/types/EntityKey.proto
@@ -25,6 +25,6 @@ option java_outer_classname = "EntityKeyProto";
 option go_package = "github.com/feast-dev/feast/sdk/go/protos/feast/types";
 
 message EntityKey {
-    repeated string entity_names = 1;
+    repeated string join_keys = 1;
     repeated feast.types.Value entity_values = 2;
 }
diff --git a/sdk/python/feast/entity.py b/sdk/python/feast/entity.py
@@ -36,11 +36,17 @@ def __init__(
         name: str,
         value_type: ValueType,
         description: str = "",
+        join_key: Optional[str] = None,
         labels: Optional[MutableMapping[str, str]] = None,
     ):
         self._name = name
         self._description = description
         self._value_type = value_type
+        if join_key:
+            self._join_key = join_key
+        else:
+            self._join_key = name
+
         if labels is None:
             self._labels = dict()  # type: MutableMapping[str, str]
         else:
@@ -58,6 +64,7 @@ def __eq__(self, other):
             or self.name != other.name
             or self.description != other.description
             or self.value_type != other.value_type
+            or self.join_key != other.join_key
         ):
             return False
 
@@ -94,6 +101,20 @@ def description(self, description):
         """
         self._description = description
 
+    @property
+    def join_key(self):
+        """
+        Returns the join key of this entity
+        """
+        return self._join_key
+
+    @join_key.setter
+    def join_key(self, join_key):
+        """
+        Sets the join key of this entity
+        """
+        self._join_key = join_key
+
     @property
     def value_type(self) -> ValueType:
         """
@@ -197,6 +218,7 @@ def from_proto(cls, entity_proto: EntityV2Proto):
             description=entity_proto.spec.description,
             value_type=ValueType(entity_proto.spec.value_type),
             labels=entity_proto.spec.labels,
+            join_key=entity_proto.spec.join_key,
         )
 
         entity._created_timestamp = entity_proto.meta.created_timestamp
@@ -222,6 +244,7 @@ def to_proto(self) -> EntityV2Proto:
             description=self.description,
             value_type=self.value_type.value,
             labels=self.labels,
+            join_key=self.join_key,
         )
 
         return EntityV2Proto(spec=spec, meta=meta)
@@ -266,6 +289,7 @@ def to_spec_proto(self) -> EntitySpecProto:
             description=self.description,
             value_type=self.value_type.value,
             labels=self.labels,
+            join_key=self.join_key,
         )
 
         return spec
@@ -282,5 +306,6 @@ def _update_from_entity(self, entity):
         self.description = entity.description
         self.value_type = entity.value_type
         self.labels = entity.labels
+        self.join_key = entity.join_key
         self._created_timestamp = entity.created_timestamp
         self._last_updated_timestamp = entity.last_updated_timestamp