Repository: feast-dev/feast

Issue:
Update Python SDK to support labels
This issue follows #463 which has a partial implementation in #536. 

One additional PR is needed to get #436 into a consistent and complete state for end users. Currently #536 makes changes to Feast Core by adding support for labels to feature sets and features, but it doesn't add support for labels to feature sets in the Python SDK.

The task here is to add this support and update any end-to-end tests that test both the Python SDK and Feast Core w.r.t label metadata.


Discussion:
Unless there is any objection then I think we can push this out to 0.6. We'd be introducing new API methods and not taking anything away? I don't see a compelling reason to block 0.5 over this.
Sounds reasonable, I also think this could be a great community contribution if there's someone wishing to use this functionality from the Python SDK who can prioritize it before the maintainers are able to.
I would be happy to pick this up somewhere next week. From #463:
> with the Python SDK having a set_label(key, value) method and a remove_label(key) method on the Feature class. list_feature_sets() would print out this information as well, but filtering will be left for a future release.

The idea is still to implement the 2 suggested functions and update `list_feature_sets()` right?
Hi @Joostrothweiler, that would be great!

Correct, it's about adding those to the `feature set` as opposed to the `feature`. Methods like `list_feature_sets()` should also be tested to see that they expose this information.

Resolving patch:
diff --git a/sdk/python/feast/feature.py b/sdk/python/feast/feature.py
--- 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
--- 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
--- 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:
         """


Test patch:
diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py
--- 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
--- 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
--- 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