Add API for listing feature sets using labels (#785)

* Update CoreService proto

* Update SpecService and add unit test

* Update Python sdk and cli

* Fix python lint error

* Add more tests

* Black formatting changes

* Revert "Black formatting changes"

This reverts commit d82200f.

* Refactor FeatureSet labels filter logic

* Set default labels to empty dict

* Refactor feature_set_list cli

* Refactor brittle test

* Implement helper function and not convert to proto

* Update cli

* Add docstrings and improve variable naming

* Use getter method instead

* Fix java lint

* Add description for labels filter proto

Co-authored-by: Terence <terence.limxp@go-jek.com>

Author: terryyylim

core/src/main/java/feast/core/model/FeatureSet.java

@@ -128,6 +128,23 @@ private String getProjectName() {
     }
   }
 
+  /**
+   * Return a boolean to facilitate streaming elements on the basis of given predicate.
+   *
+   * @param labelsFilter labels contain key-value mapping for labels attached to the FeatureSet
+   * @return boolean True if FeatureSet contains all labels in the labelsFilter
+   */
+  public boolean hasAllLabels(Map<String, String> labelsFilter) {
+    Map<String, String> featureSetLabelsMap = this.getLabelsMap();
+    for (String key : labelsFilter.keySet()) {
+      if (!featureSetLabelsMap.containsKey(key)
+          || !featureSetLabelsMap.get(key).equals(labelsFilter.get(key))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   public void setProject(Project project) {
     this.project = project;
   }
@@ -293,6 +310,10 @@ public FeatureSetProto.FeatureSet toProto() throws InvalidProtocolBufferExceptio
     return FeatureSetProto.FeatureSet.newBuilder().setMeta(meta).setSpec(spec).build();
   }
 
+  public Map<String, String> getLabelsMap() {
+    return TypeConversion.convertJsonStringToMap(this.getLabels());
+  }
+
   @Override
   public int hashCode() {
     HashCodeBuilder hcb = new HashCodeBuilder();

core/src/main/java/feast/core/service/SpecService.java

@@ -46,7 +46,9 @@
 import feast.proto.core.StoreProto.Store.Subscription;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -123,9 +125,9 @@ public GetFeatureSetResponse getFeatureSet(GetFeatureSetRequest request)
   }
 
   /**
-   * Return a list of feature sets matching the feature set name and project provided in the filter.
-   * All fields are requried. Use '*' for all arguments in order to return all feature sets in all
-   * projects.
+   * Return a list of feature sets matching the feature set name, project and labels provided in the
+   * filter. All fields are required. Use '*' in feature set name and project, and empty map in
+   * labels in order to return all feature sets in all projects.
    *
    * <p>Project name can be explicitly provided, or an asterisk can be provided to match all
    * projects. It is not possible to provide a combination of asterisks/wildcards and text. If the
@@ -135,13 +137,17 @@ public GetFeatureSetResponse getFeatureSet(GetFeatureSetRequest request)
    * sets will be returned. Regex is not supported. Explicitly defining a feature set name is not
    * possible if a project name is not set explicitly
    *
+   * <p>The labels in the filter accepts a map. All feature sets which contain every provided label
+   * will be returned.
+   *
    * @param filter filter containing the desired featureSet name
    * @return ListFeatureSetsResponse with list of featureSets found matching the filter
    */
   public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest.Filter filter)
       throws InvalidProtocolBufferException {
     String name = filter.getFeatureSetName();
     String project = filter.getProject();
+    Map<String, String> labelsFilter = filter.getLabelsMap();
 
     if (name.isEmpty()) {
       throw new IllegalArgumentException(
@@ -197,6 +203,10 @@ public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest.Filter fil
 
     ListFeatureSetsResponse.Builder response = ListFeatureSetsResponse.newBuilder();
     if (featureSets.size() > 0) {
+      featureSets =
+          featureSets.stream()
+              .filter(featureSet -> featureSet.hasAllLabels(labelsFilter))
+              .collect(Collectors.toList());
       for (FeatureSet featureSet : featureSets) {
         response.addFeatureSets(featureSet.toProto());
       }

core/src/test/java/feast/core/service/SpecServiceTest.java

@@ -106,7 +106,7 @@ public class SpecServiceTest {
   // TODO: Updates update features in place, so if tests follow the wrong order they might break.
   // Refactor this maybe?
   @Before
-  public void setUp() {
+  public void setUp() throws InvalidProtocolBufferException {
     initMocks(this);
     defaultSource = TestObjectFactory.defaultSource;
 
@@ -121,7 +121,50 @@ public void setUp() {
             "f3", "project1", Arrays.asList(f3e1), Arrays.asList(f3f2, f3f1));
 
     FeatureSet featureSet4 = newDummyFeatureSet("f4", Project.DEFAULT_NAME);
-    featureSets = Arrays.asList(featureSet1, featureSet2, featureSet3, featureSet4);
+    Map<String, String> singleFeatureSetLabels =
+        new HashMap<>() {
+          {
+            put("fsLabel1", "fsValue1");
+          }
+        };
+    Map<String, String> duoFeatureSetLabels =
+        new HashMap<>() {
+          {
+            put("fsLabel1", "fsValue1");
+            put("fsLabel2", "fsValue2");
+          }
+        };
+    FeatureSet featureSet5 = newDummyFeatureSet("f5", Project.DEFAULT_NAME);
+    FeatureSet featureSet6 = newDummyFeatureSet("f6", Project.DEFAULT_NAME);
+    FeatureSetSpec featureSetSpec5 = featureSet5.toProto().getSpec().toBuilder().build();
+    FeatureSetSpec featureSetSpec6 = featureSet6.toProto().getSpec().toBuilder().build();
+    FeatureSetProto.FeatureSet fs5 =
+        FeatureSetProto.FeatureSet.newBuilder()
+            .setSpec(
+                featureSetSpec5
+                    .toBuilder()
+                    .setSource(defaultSource.toProto())
+                    .putAllLabels(singleFeatureSetLabels)
+                    .build())
+            .build();
+    FeatureSetProto.FeatureSet fs6 =
+        FeatureSetProto.FeatureSet.newBuilder()
+            .setSpec(
+                featureSetSpec6
+                    .toBuilder()
+                    .setSource(defaultSource.toProto())
+                    .putAllLabels(duoFeatureSetLabels)
+                    .build())
+            .build();
+
+    featureSets =
+        Arrays.asList(
+            featureSet1,
+            featureSet2,
+            featureSet3,
+            featureSet4,
+            FeatureSet.fromProto(fs5),
+            FeatureSet.fromProto(fs6));
 
     when(featureSetRepository.findAll()).thenReturn(featureSets);
     when(featureSetRepository.findAllByOrderByNameAsc()).thenReturn(featureSets);
@@ -713,6 +756,35 @@ public void applyFeatureSetShouldAcceptFeatureSetLabels() throws InvalidProtocol
     assertEquals(featureSetLabels, appliedLabels);
   }
 
+  @Test
+  public void shouldFilterByFeatureSetLabels() throws InvalidProtocolBufferException {
+    List<FeatureSetProto.FeatureSet> list = new ArrayList<>();
+    ListFeatureSetsResponse actual1 =
+        specService.listFeatureSets(
+            Filter.newBuilder()
+                .setFeatureSetName("*")
+                .setProject("*")
+                .putLabels("fsLabel2", "fsValue2")
+                .build());
+    list.add(featureSets.get(5).toProto());
+    ListFeatureSetsResponse expected1 =
+        ListFeatureSetsResponse.newBuilder().addAllFeatureSets(list).build();
+
+    ListFeatureSetsResponse actual2 =
+        specService.listFeatureSets(
+            Filter.newBuilder()
+                .setFeatureSetName("*")
+                .setProject("*")
+                .putLabels("fsLabel1", "fsValue1")
+                .build());
+    list.add(0, featureSets.get(4).toProto());
+    ListFeatureSetsResponse expected2 =
+        ListFeatureSetsResponse.newBuilder().addAllFeatureSets(list).build();
+
+    assertThat(actual1, equalTo(expected1));
+    assertThat(actual2, equalTo(expected2));
+  }
+
   @Test
   public void shouldUpdateStoreIfConfigChanges() throws InvalidProtocolBufferException {
     when(storeRepository.findById("SERVING")).thenReturn(Optional.of(stores.get(0)));

protos/feast/core/CoreService.proto

@@ -140,6 +140,10 @@ message ListFeatureSetsRequest {
         // - my-feature-set* can be used to match all features prefixed by "my-feature-set"
         // - my-feature-set-6 can be used to select a single feature set
         string feature_set_name = 1;
+
+        // User defined metadata for feature set.
+        // Feature sets with all matching labels will be returned.
+        map<string,string> labels = 4;
     }
 }
 

sdk/python/feast/cli.py

@@ -120,15 +120,55 @@ def feature_set():
     pass
 
 
+def _get_labels_dict(label_str: str):
+    """
+    Converts CLI input labels string to dictionary format if provided string is valid.
+    """
+    labels_dict = {}
+    labels_kv = label_str.split(",")
+    if label_str == "":
+        return labels_dict
+    if len(labels_kv) % 2 == 1:
+        raise ValueError("Uneven key-value label pairs were entered")
+    for k, v in zip(labels_kv[0::2], labels_kv[1::2]):
+        labels_dict[k] = v
+    return labels_dict
+
+
 @feature_set.command(name="list")
-def feature_set_list():
+@click.option(
+    "--project",
+    "-p",
+    help="Project that feature set belongs to",
+    type=click.STRING,
+    default="*",
+)
+@click.option(
+    "--name",
+    "-n",
+    help="Filters feature sets by name. Wildcards (*) may be included to match multiple feature sets",
+    type=click.STRING,
+    default="*",
+)
+@click.option(
+    "--labels",
+    "-l",
+    help="Labels to filter for feature sets",
+    type=click.STRING,
+    default="",
+)
+def feature_set_list(project: str, name: str, labels: str):
     """
     List all feature sets
     """
     feast_client = Client()  # type: Client
 
+    labels_dict = _get_labels_dict(labels)
+
     table = []
-    for fs in feast_client.list_feature_sets(project="*", name="*"):
+    for fs in feast_client.list_feature_sets(
+        project=project, name=name, labels=labels_dict
+    ):
         table.append([fs.name, repr(fs)])
 
     from tabulate import tabulate

sdk/python/feast/client.py

@@ -425,7 +425,7 @@ def _apply_feature_set(self, feature_set: FeatureSet):
         feature_set._update_from_feature_set(applied_fs)
 
     def list_feature_sets(
-        self, project: str = None, name: str = None,
+        self, project: str = None, name: str = None, labels: Dict[str, str] = dict()
     ) -> List[FeatureSet]:
         """
         Retrieve a list of feature sets from Feast Core
@@ -448,7 +448,9 @@ def list_feature_sets(
         if name is None:
             name = "*"
 
-        filter = ListFeatureSetsRequest.Filter(project=project, feature_set_name=name)
+        filter = ListFeatureSetsRequest.Filter(
+            project=project, feature_set_name=name, labels=labels
+        )
 
         # Get latest feature sets from Feast Core
         feature_set_protos = self._core_service_stub.ListFeatureSets(

sdk/python/tests/test_client.py

@@ -34,6 +34,7 @@
 from feast.core.CoreService_pb2 import (
     GetFeastCoreVersionResponse,
     GetFeatureSetResponse,
+    ListFeatureSetsResponse,
     ListIngestionJobsResponse,
 )
 from feast.core.FeatureSet_pb2 import EntitySpec as EntitySpecProto
@@ -321,6 +322,67 @@ def test_get_feature_set(self, mocked_client, mocker):
             and len(feature_set.entities) == 1
         )
 
+    @pytest.mark.parametrize(
+        "mocked_client",
+        [pytest.lazy_fixture("mock_client"), pytest.lazy_fixture("secure_mock_client")],
+    )
+    def test_list_feature_sets(self, mocked_client, mocker):
+        mocker.patch.object(
+            mocked_client,
+            "_core_service_stub",
+            return_value=Core.CoreServiceStub(grpc.insecure_channel("")),
+        )
+
+        feature_set_1_proto = FeatureSetProto(
+            spec=FeatureSetSpecProto(
+                project="test",
+                name="driver_car",
+                max_age=Duration(seconds=3600),
+                labels={"key1": "val1", "key2": "val2"},
+                features=[
+                    FeatureSpecProto(
+                        name="feature_1", value_type=ValueProto.ValueType.FLOAT
+                    )
+                ],
+            )
+        )
+        feature_set_2_proto = FeatureSetProto(
+            spec=FeatureSetSpecProto(
+                project="test",
+                name="driver_ride",
+                max_age=Duration(seconds=3600),
+                labels={"key1": "val1"},
+                features=[
+                    FeatureSpecProto(
+                        name="feature_1", value_type=ValueProto.ValueType.FLOAT
+                    )
+                ],
+            )
+        )
+
+        mocker.patch.object(
+            mocked_client._core_service_stub,
+            "ListFeatureSets",
+            return_value=ListFeatureSetsResponse(
+                feature_sets=[feature_set_1_proto, feature_set_2_proto]
+            ),
+        )
+
+        feature_sets = mocked_client.list_feature_sets(labels={"key1": "val1"})
+        assert len(feature_sets) == 2
+
+        feature_set = feature_sets[0]
+        assert (
+            feature_set.name == "driver_car"
+            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["feature_1"].name == "feature_1"
+            and feature_set.fields["feature_1"].dtype == ValueType.FLOAT
+            and len(feature_set.features) == 1
+        )
+
     @pytest.mark.parametrize(
         "mocked_client",
         [pytest.lazy_fixture("mock_client"), pytest.lazy_fixture("secure_mock_client")],

tests/e2e/redis/basic-ingest-redis-serving.py

@@ -126,6 +126,34 @@ def test_basic_register_feature_set_success(client):
                                                   project=PROJECT_NAME)
     assert cust_trans_fs_actual == cust_trans_fs_expected
 
+    # Register feature set with labels
+    driver_unlabelled_fs = FeatureSet(
+        "driver_unlabelled",
+        features=[
+            Feature("rating", ValueType.FLOAT),
+            Feature("cost", ValueType.FLOAT)
+        ],
+        entities=[Entity("entity_id", ValueType.INT64)],
+        max_age=Duration(seconds=100)
+    )
+    driver_labeled_fs_expected = FeatureSet(
+        "driver_labeled",
+        features=[
+            Feature("rating", ValueType.FLOAT),
+            Feature("cost", ValueType.FLOAT)
+        ],
+        entities=[Entity("entity_id", ValueType.INT64)],
+        max_age=Duration(seconds=100),
+        labels={"key1":"val1"}
+    )
+    client.set_project(PROJECT_NAME)
+    client.apply(driver_unlabelled_fs)
+    client.apply(driver_labeled_fs_expected)
+    driver_fs_actual = client.list_feature_sets(
+        project=PROJECT_NAME, labels={"key1": "val1"}
+    )[0]
+    assert driver_fs_actual == driver_labeled_fs_expected
+
     # reset client's project for other tests
     client.set_project()