Bugfix: Re-applying of featuresets does not update label changes (#857)

* Fix applyFeatureSet to propagate label updates * Address PR comments Co-authored-by: Terence <terence.limxp@go-jek.com>
feast-dev · Jul 2, 2020 · 000982a · 000982a
1 parent 9a037dc
commit 000982a
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 4 deletions.
diff --git 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
@@ -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<String, FeatureSpec> 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
@@ -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<String, String> 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<String, String> updatedFsLabels = updatedFs.getSpec().getLabelsMap();
+
+    Map<String, String> updatedFeatureLabels = updatedFs.getSpec().getFeatures(0).getLabelsMap();
+    Map<String, String> emptyFeatureLabels = new HashMap<>();
+
+    assertEquals(featureSetLabels, updatedFsLabels);
+    assertEquals(emptyFeatureLabels, updatedFeatureLabels);
+  }
+
   @Test
   public void applyFeatureSetShouldAcceptFeatureSetLabels() throws InvalidProtocolBufferException {
     Map<String, String> featureSetLabels =