diff --git a/core/src/main/java/feast/core/model/Feature.java b/core/src/main/java/feast/core/model/Feature.java index 7229f13af87..b387f3403bb 100644 --- a/core/src/main/java/feast/core/model/Feature.java +++ b/core/src/main/java/feast/core/model/Feature.java @@ -235,7 +235,7 @@ public void archive() { } /** - * Update the feature object with a valid feature spec. Only schema changes are allowed. + * Update the feature object with a valid feature spec. * * @param featureSpec {@link FeatureSpec} containing schema changes. */ @@ -252,6 +252,7 @@ public void updateFromProto(FeatureSpec featureSpec) { "You are attempting to change the type of feature %s from %s to %s. This isn't allowed. Please create a new feature.", featureSpec.getName(), type, featureSpec.getValueType())); } + this.setLabels(TypeConversion.convertMapToJsonString(featureSpec.getLabelsMap())); updateSchema(featureSpec); } diff --git a/core/src/main/java/feast/core/model/FeatureSet.java b/core/src/main/java/feast/core/model/FeatureSet.java index 89d8c65440c..41008fa6625 100644 --- a/core/src/main/java/feast/core/model/FeatureSet.java +++ b/core/src/main/java/feast/core/model/FeatureSet.java @@ -314,9 +314,10 @@ public void updateFromProto(FeatureSetProto.FeatureSet featureSetProto) spec.getEntitiesList(), existingEntities)); } - // 4. Update max age and source. - maxAgeSeconds = spec.getMaxAge().getSeconds(); - source = Source.fromProto(spec.getSource()); + // 2. Update max age, source and labels. + this.maxAgeSeconds = spec.getMaxAge().getSeconds(); + this.source = Source.fromProto(spec.getSource()); + this.setLabels(TypeConversion.convertMapToJsonString(spec.getLabelsMap())); Map updatedFeatures = spec.getFeaturesList().stream().collect(Collectors.toMap(FeatureSpec::getName, fs -> fs)); diff --git a/core/src/test/java/feast/core/service/SpecServiceTest.java b/core/src/test/java/feast/core/service/SpecServiceTest.java index 3d0a3a75d5d..48572ddcecf 100644 --- a/core/src/test/java/feast/core/service/SpecServiceTest.java +++ b/core/src/test/java/feast/core/service/SpecServiceTest.java @@ -700,6 +700,42 @@ public void applyFeatureSetShouldAcceptFeatureLabels() throws InvalidProtocolBuf assertEquals(appliedFeatureSpecsLabels, featureLabels); } + @Test + public void applyFeatureSetShouldUpdateLabels() throws InvalidProtocolBufferException { + FeatureSpec updatedFeature = + FeatureSpec.newBuilder().setName("feature").setValueType(Enum.STRING).build(); + + FeatureSet featureSet = featureSets.get(0); + FeatureSetSpec featureSetSpec = featureSet.toProto().getSpec().toBuilder().build(); + Map featureSetLabels = + new HashMap<>() { + { + put("fsLabel1", "fsValue1"); + } + }; + + FeatureSetProto.FeatureSet incomingFeatureSet = + FeatureSetProto.FeatureSet.newBuilder() + .setSpec( + featureSetSpec + .toBuilder() + .setFeatures(0, updatedFeature) + .putAllLabels(featureSetLabels) + .build()) + .build(); + + ApplyFeatureSetResponse applyFeatureSetResponse = + specService.applyFeatureSet(incomingFeatureSet); + FeatureSetProto.FeatureSet updatedFs = applyFeatureSetResponse.getFeatureSet(); + Map updatedFsLabels = updatedFs.getSpec().getLabelsMap(); + + Map updatedFeatureLabels = updatedFs.getSpec().getFeatures(0).getLabelsMap(); + Map emptyFeatureLabels = new HashMap<>(); + + assertEquals(featureSetLabels, updatedFsLabels); + assertEquals(emptyFeatureLabels, updatedFeatureLabels); + } + @Test public void applyFeatureSetShouldAcceptFeatureSetLabels() throws InvalidProtocolBufferException { Map featureSetLabels =