Add support for labels on feature sets (feast-dev#707)

Joostrothweiler · web-flow · commit 6fd0874c609c · 2020-06-09T17:25:02.000+08:00
* Update Python SDK to support labels

* Format python code

* Fix equals comparison FeatureSet

* Fix bug labels returns presence and initialize labels by default as empty dict

* Propagate KeyErrors labels and fields and update e2e tests with labels in constructor

Co-authored-by: Joost Rothweiler &lt;=&gt;
diff --git a/sdk/python/feast/feature.py b/sdk/python/feast/feature.py
@@ -28,6 +28,7 @@ def to_proto(self) -> FeatureProto:
         return FeatureProto(
             name=self.name,
             value_type=value_type,
+            labels=self.labels,
             presence=self.presence,
             group_presence=self.group_presence,
             shape=self.shape,
@@ -57,7 +58,9 @@ def from_proto(cls, feature_proto: FeatureProto):
             Feature object
         """
         feature = cls(
-            name=feature_proto.name, dtype=ValueType(feature_proto.value_type),
+            name=feature_proto.name,
+            dtype=ValueType(feature_proto.value_type),
+            labels=feature_proto.labels,
         )
         feature.update_presence_constraints(feature_proto)
         feature.update_shape_type(feature_proto)
diff --git a/sdk/python/feast/feature_set.py b/sdk/python/feast/feature_set.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 import warnings
 from collections import OrderedDict
-from typing import Dict, List, Optional
+from typing import Dict, List, MutableMapping, Optional
 
 import pandas as pd
 from google.protobuf import json_format
@@ -56,6 +56,7 @@ def __init__(
         entities: List[Entity] = None,
         source: Source = None,
         max_age: Optional[Duration] = None,
+        labels: Optional[MutableMapping[str, str]] = None,
     ):
         self._name = name
         self._project = project
@@ -68,6 +69,10 @@ def __init__(
             self._source = None
         else:
             self._source = source
+        if labels is None:
+            self._labels = OrderedDict()
+        else:
+            self._labels = labels
         self._max_age = max_age
         self._status = None
         self._created_timestamp = None
@@ -84,7 +89,8 @@ def __eq__(self, other):
                 return False
 
         if (
-            self.name != other.name
+            self.labels != other.labels
+            or self.name != other.name
             or self.project != other.project
             or self.max_age != other.max_age
         ):
@@ -217,6 +223,21 @@ def max_age(self, max_age):
         """
         self._max_age = max_age
 
+    @property
+    def labels(self):
+        """
+        Returns the labels of this feature set. This is the user defined metadata
+        defined as a dictionary.
+        """
+        return self._labels
+
+    @labels.setter
+    def labels(self, labels: MutableMapping[str, str]):
+        """
+        Set the labels for this feature set
+        """
+        self._labels = labels
+
     @property
     def status(self):
         """
@@ -245,6 +266,18 @@ def created_timestamp(self, created_timestamp):
         """
         self._created_timestamp = created_timestamp
 
+    def set_label(self, key: str, value: str):
+        """
+        Sets the label value for a given key
+        """
+        self.labels[key] = value
+
+    def remove_label(self, key: str):
+        """
+        Removes a label based on key
+        """
+        del self.labels[key]
+
     def add(self, resource):
         """
         Adds a resource (Feature, Entity) to this Feature Set.
@@ -279,11 +312,7 @@ def drop(self, name: str):
         Args:
             name: Name of Feature or Entity to be removed
         """
-        if name not in self._fields:
-            raise ValueError("Could not find field " + name + ", no action taken")
-        if name in self._fields:
-            del self._fields[name]
-            return
+        del self._fields[name]
 
     def _add_fields(self, fields: List[Field]):
         """
@@ -796,6 +825,7 @@ def from_proto(cls, feature_set_proto: FeatureSetProto):
                 and feature_set_proto.spec.max_age.nanos == 0
                 else feature_set_proto.spec.max_age
             ),
+            labels=feature_set_proto.spec.labels,
             source=(
                 None
                 if feature_set_proto.spec.source.type == 0
@@ -825,6 +855,7 @@ def to_proto(self) -> FeatureSetProto:
             name=self.name,
             project=self.project,
             max_age=self.max_age,
+            labels=self.labels,
             source=self.source.to_proto() if self.source is not None else None,
             features=[
                 field.to_proto()
diff --git a/sdk/python/feast/field.py b/sdk/python/feast/field.py
@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Union
+from collections import OrderedDict
+from typing import MutableMapping, Optional, Union
 
 from feast.core.FeatureSet_pb2 import EntitySpec, FeatureSpec
 from feast.value_type import ValueType
@@ -24,11 +25,20 @@ class Field:
     features.
     """
 
-    def __init__(self, name: str, dtype: ValueType):
+    def __init__(
+        self,
+        name: str,
+        dtype: ValueType,
+        labels: Optional[MutableMapping[str, str]] = None,
+    ):
         self._name = name
         if not isinstance(dtype, ValueType):
             raise ValueError("dtype is not a valid ValueType")
         self._dtype = dtype
+        if labels is None:
+            self._labels = OrderedDict()
+        else:
+            self._labels = labels
         self._presence = None
         self._group_presence = None
         self._shape = None
@@ -47,7 +57,11 @@ def __init__(self, name: str, dtype: ValueType):
         self._time_of_day_domain = None
 
     def __eq__(self, other):
-        if self.name != other.name or self.dtype != other.dtype:
+        if (
+            self.name != other.name
+            or self.dtype != other.dtype
+            or self.labels != other.labels
+        ):
             return False
         return True
 
@@ -65,6 +79,13 @@ def dtype(self) -> ValueType:
         """
         return self._dtype
 
+    @property
+    def labels(self) -> MutableMapping[str, str]:
+        """
+        Getter for labels of this field
+        """
+        return self._labels
+
     @property
     def presence(self) -> schema_pb2.FeaturePresence:
         """
diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py
@@ -276,6 +276,7 @@ def test_get_feature_set(self, mocked_client, mocker):
                     spec=FeatureSetSpecProto(
                         name="my_feature_set",
                         max_age=Duration(seconds=3600),
+                        labels={"key1": "val1", "key2": "val2"},
                         features=[
                             FeatureSpecProto(
                                 name="my_feature_1",
@@ -308,6 +309,10 @@ def test_get_feature_set(self, mocked_client, mocker):
 
         assert (
             feature_set.name == "my_feature_set"
+            and "key1" in feature_set.labels
+            and feature_set.labels["key1"] == "val1"
+            and "key2" in feature_set.labels
+            and feature_set.labels["key2"] == "val2"
             and feature_set.fields["my_feature_1"].name == "my_feature_1"
             and feature_set.fields["my_feature_1"].dtype == ValueType.FLOAT
             and feature_set.fields["my_entity_1"].name == "my_entity_1"
diff --git a/sdk/python/tests/test_feature_set.py b/sdk/python/tests/test_feature_set.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pathlib
+from collections import OrderedDict
 from concurrent import futures
 from datetime import datetime
 
@@ -62,7 +63,7 @@ def test_add_remove_features_success(self):
         assert len(fs.features) == 1 and fs.features[0].name == "my-feature-2"
 
     def test_remove_feature_failure(self):
-        with pytest.raises(ValueError):
+        with pytest.raises(KeyError):
             fs = FeatureSet("my-feature-set")
             fs.drop(name="my-feature-1")
 
@@ -287,6 +288,98 @@ def make_tfx_schema_domain_info_inline(schema):
                 feature.int_domain.MergeFrom(domain_ref_to_int_domain[domain_ref])
 
 
+def test_feature_set_class_contains_labels():
+    fs = FeatureSet("my-feature-set", labels={"key1": "val1", "key2": "val2"})
+    assert "key1" in fs.labels.keys() and fs.labels["key1"] == "val1"
+    assert "key2" in fs.labels.keys() and fs.labels["key2"] == "val2"
+
+
+def test_feature_class_contains_labels():
+    fs = FeatureSet("my-feature-set", labels={"key1": "val1", "key2": "val2"})
+    fs.add(
+        Feature(
+            name="my-feature-1",
+            dtype=ValueType.INT64,
+            labels={"feature_key1": "feature_val1"},
+        )
+    )
+    assert "feature_key1" in fs.features[0].labels.keys()
+    assert fs.features[0].labels["feature_key1"] == "feature_val1"
+
+
+def test_feature_set_without_labels_empty_dict():
+    fs = FeatureSet("my-feature-set")
+    assert fs.labels == OrderedDict()
+    assert len(fs.labels) == 0
+
+
+def test_feature_without_labels_empty_dict():
+    f = Feature("my feature", dtype=ValueType.INT64)
+    assert f.labels == OrderedDict()
+    assert len(f.labels) == 0
+
+
+def test_set_label_feature_set():
+    fs = FeatureSet("my-feature-set")
+    fs.set_label("k1", "v1")
+    assert fs.labels["k1"] == "v1"
+
+
+def test_set_labels_overwrites_existing():
+    fs = FeatureSet("my-feature-set")
+    fs.set_label("k1", "v1")
+    fs.set_label("k1", "v2")
+    assert fs.labels["k1"] == "v2"
+
+
+def test_remove_labels_empty_failure():
+    fs = FeatureSet("my-feature-set")
+    with pytest.raises(KeyError):
+        fs.remove_label("key1")
+
+
+def test_remove_labels_invalid_key_failure():
+    fs = FeatureSet("my-feature-set")
+    fs.set_label("k1", "v1")
+    with pytest.raises(KeyError):
+        fs.remove_label("key1")
+
+
+def test_unequal_feature_based_on_labels():
+    f1 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "v1"})
+    f2 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "v1"})
+    assert f1 == f2
+    f3 = Feature(name="feature-1", dtype=ValueType.INT64)
+    assert f1 != f3
+    f4 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "notv1"})
+    assert f1 != f4
+
+
+def test_unequal_feature_set_based_on_labels():
+    fs1 = FeatureSet("my-feature-set")
+    fs2 = FeatureSet("my-feature-set")
+    assert fs1 == fs2
+    fs1.set_label("k1", "v1")
+    fs2.set_label("k1", "v1")
+    assert fs1 == fs2
+    fs2.set_label("k1", "unequal")
+    assert not fs1 == fs2
+
+
+def test_unequal_feature_set_other_has_no_labels():
+    fs1 = FeatureSet("my-feature-set")
+    fs2 = FeatureSet("my-feature-set")
+    assert fs1 == fs2
+    fs1.set_label("k1", "v1")
+    assert not fs1 == fs2
+
+
+def test_unequal_feature_other_has_no_labels():
+    f1 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "v1"})
+    f2 = Feature(name="feature-1", dtype=ValueType.INT64)
+    assert f1 != f2
+
+
 class TestFeatureSetRef:
     def test_from_feature_set(self):
         feature_set = FeatureSet("test", "test")
diff --git a/tests/e2e/redis/basic-ingest-redis-serving.py b/tests/e2e/redis/basic-ingest-redis-serving.py
@@ -688,11 +688,16 @@ def get_feature_set(self, core_service_stub, name, project):
     @pytest.mark.run(order=51)
     def test_register_feature_set_with_labels(self, core_service_stub):
         feature_set_name = "test_feature_set_labels"
-        feature_set_proto = FeatureSet(feature_set_name, PROJECT_NAME).to_proto()
-        feature_set_proto.spec.labels[self.LABEL_KEY] = self.LABEL_VALUE
+        feature_set_proto = FeatureSet(
+            name=feature_set_name,
+            project=PROJECT_NAME,
+            labels={self.LABEL_KEY: self.LABEL_VALUE},
+        ).to_proto()
         self.apply_feature_set(core_service_stub, feature_set_proto)
 
-        retrieved_feature_set = self.get_feature_set(core_service_stub, feature_set_name, PROJECT_NAME)
+        retrieved_feature_set = self.get_feature_set(
+            core_service_stub, feature_set_name, PROJECT_NAME
+        )
 
         assert self.LABEL_KEY in retrieved_feature_set.spec.labels
         assert retrieved_feature_set.spec.labels[self.LABEL_KEY] == self.LABEL_VALUE
@@ -701,12 +706,22 @@ def test_register_feature_set_with_labels(self, core_service_stub):
     @pytest.mark.run(order=52)
     def test_register_feature_with_labels(self, core_service_stub):
         feature_set_name = "test_feature_labels"
-        feature_set_proto = FeatureSet(feature_set_name, PROJECT_NAME, features=[Feature("rating", ValueType.INT64)]) \
-            .to_proto()
-        feature_set_proto.spec.features[0].labels[self.LABEL_KEY] = self.LABEL_VALUE
+        feature_set_proto = FeatureSet(
+            name=feature_set_name,
+            project=PROJECT_NAME,
+            features=[
+                Feature(
+                    name="rating",
+                    dtype=ValueType.INT64,
+                    labels={self.LABEL_KEY: self.LABEL_VALUE},
+                )
+            ],
+        ).to_proto()
         self.apply_feature_set(core_service_stub, feature_set_proto)
 
-        retrieved_feature_set = self.get_feature_set(core_service_stub, feature_set_name, PROJECT_NAME)
+        retrieved_feature_set = self.get_feature_set(
+            core_service_stub, feature_set_name, PROJECT_NAME
+        )
         retrieved_feature = retrieved_feature_set.spec.features[0]
 
         assert self.LABEL_KEY in retrieved_feature.labels
