diff --git a/core/src/main/java/feast/core/model/FeatureSet.java b/core/src/main/java/feast/core/model/FeatureSet.java index e1b50098b90..67a21e0b250 100644 --- a/core/src/main/java/feast/core/model/FeatureSet.java +++ b/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 labelsFilter) { + Map 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 getLabelsMap() { + return TypeConversion.convertJsonStringToMap(this.getLabels()); + } + @Override public int hashCode() { HashCodeBuilder hcb = new HashCodeBuilder(); diff --git a/core/src/main/java/feast/core/service/SpecService.java b/core/src/main/java/feast/core/service/SpecService.java index 9000ee983ee..a4f1202858d 100644 --- a/core/src/main/java/feast/core/service/SpecService.java +++ b/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. * *

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,6 +137,9 @@ 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 * + *

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 */ @@ -142,6 +147,7 @@ public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest.Filter fil throws InvalidProtocolBufferException { String name = filter.getFeatureSetName(); String project = filter.getProject(); + Map 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()); } diff --git a/core/src/test/java/feast/core/service/SpecServiceTest.java b/core/src/test/java/feast/core/service/SpecServiceTest.java index 9ec219b6c97..20a84edbdaa 100644 --- a/core/src/test/java/feast/core/service/SpecServiceTest.java +++ b/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 singleFeatureSetLabels = + new HashMap<>() { + { + put("fsLabel1", "fsValue1"); + } + }; + Map 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 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))); diff --git a/protos/feast/core/CoreService.proto b/protos/feast/core/CoreService.proto index 05755b701d8..a6ac313b2a5 100644 --- a/protos/feast/core/CoreService.proto +++ b/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 labels = 4; } } diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py index 4db41f92f18..34a3f600e15 100644 --- a/sdk/python/feast/cli.py +++ b/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 diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py index 316573ce93d..2ec90e434db 100644 --- a/sdk/python/feast/client.py +++ b/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( diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py index ee89b83737a..85fea31dbb1 100644 --- a/sdk/python/tests/test_client.py +++ b/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")], diff --git a/tests/e2e/redis/basic-ingest-redis-serving.py b/tests/e2e/redis/basic-ingest-redis-serving.py index 63a3c794a5a..1424936f0c8 100644 --- a/tests/e2e/redis/basic-ingest-redis-serving.py +++ b/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()