[FSTORE-1134] Enable updating embedding (#1495)

(cherry picked from commit 2191bad)
logicalclocks · Feb 21, 2024 · e78f833 · e78f833
1 parent 71d009a
commit e78f833
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 8 deletions.
diff --git a/...pi/src/main/java/io/hops/hopsworks/api/featurestore/featuregroup/FeaturegroupService.java b/...pi/src/main/java/io/hops/hopsworks/api/featurestore/featuregroup/FeaturegroupService.java
@@ -332,6 +332,7 @@ public Response getFeatureGroupForOnlinefs(
     verifyIdProvided(featuregroupId);
     Featuregroup featuregroup = featuregroupController.getFeaturegroupById(featurestore, featuregroupId);
     FeaturegroupDTO featuregroupDTO = new FeaturegroupDTO(featuregroup);
+    featuregroupDTO.setFeatures(featuregroupController.getPrimaryKey(featuregroup));
     GenericEntity<FeaturegroupDTO> featuregroupGeneric =
         new GenericEntity<FeaturegroupDTO>(featuregroupDTO) {};
     return noCacheResponse.getNoCacheResponseBuilder(Response.Status.OK).entity(featuregroupGeneric).build();

diff --git a/.../main/java/io/hops/hopsworks/common/featurestore/featuregroup/FeaturegroupController.java b/.../main/java/io/hops/hopsworks/common/featurestore/featuregroup/FeaturegroupController.java
@@ -381,6 +381,35 @@ public FeaturegroupDTO convertFeaturegrouptoDTO(Featuregroup featuregroup, Proje
     }
   }
 
+  public List<FeatureGroupFeatureDTO> getPrimaryKey(Featuregroup featuregroup) throws FeaturestoreException {
+    switch (featuregroup.getFeaturegroupType()) {
+      case CACHED_FEATURE_GROUP:
+        return cachedFeaturegroupController.getPrimaryKeys(featuregroup)
+            .stream().map(pk -> {
+              FeatureGroupFeatureDTO featureDTO = new FeatureGroupFeatureDTO();
+              featureDTO.setName(pk);
+              featureDTO.setPrimary(true);
+              return featureDTO;
+            }).collect(Collectors.toList());
+      case STREAM_FEATURE_GROUP:
+        return streamFeatureGroupController.getPrimaryKeys(featuregroup)
+            .stream().map(pk -> {
+              FeatureGroupFeatureDTO featureDTO = new FeatureGroupFeatureDTO();
+              featureDTO.setName(pk);
+              featureDTO.setPrimary(true);
+              return featureDTO;
+            }).collect(Collectors.toList());
+      case ON_DEMAND_FEATURE_GROUP:
+        return onDemandFeaturegroupController.getFeaturesDTO(featuregroup)
+            .stream().filter(FeatureGroupFeatureDTO::getPrimary).collect(Collectors.toList());
+      default:
+        throw new IllegalArgumentException(RESTCodes.FeaturestoreErrorCode.ILLEGAL_FEATUREGROUP_TYPE.getMessage()
+            + ", Recognized Feature group types are: " + FeaturegroupType.ON_DEMAND_FEATURE_GROUP + "," +
+            FeaturegroupType.STREAM_FEATURE_GROUP + ",  and: " + FeaturegroupType.CACHED_FEATURE_GROUP +
+            ". The provided feature group type was not recognized: " + featuregroup.getFeaturegroupType());
+    }
+  }
+
   /**
    * Retrieves a list of feature groups with a specific name from a specific feature store
    *

diff --git a/.../hops/hopsworks/common/featurestore/featuregroup/cached/CachedFeaturegroupController.java b/.../hops/hopsworks/common/featurestore/featuregroup/cached/CachedFeaturegroupController.java
@@ -210,10 +210,7 @@ public CachedFeaturegroupDTO convertCachedFeaturegroupToDTO(Featuregroup feature
 
   public List<FeatureGroupFeatureDTO> getFeaturesDTO(Featuregroup featuregroup, Project project, Users user)
           throws FeaturestoreException {
-    Set<String> primaryKeys = featuregroup.getCachedFeaturegroup().getFeaturesExtraConstraints().stream()
-        .filter(CachedFeatureExtraConstraints::getPrimary)
-        .map(CachedFeatureExtraConstraints::getName)
-        .collect(Collectors.toSet());
+    Set<String> primaryKeys = getPrimaryKeys(featuregroup);
 
     Set<String> precombineKeys = featuregroup.getCachedFeaturegroup().getFeaturesExtraConstraints().stream()
         .filter(CachedFeatureExtraConstraints::getHudiPrecombineKey)
@@ -242,6 +239,13 @@ public List<FeatureGroupFeatureDTO> getFeaturesDTO(Featuregroup featuregroup, Pr
     return featureGroupFeatures;
   }
 
+  public Set<String> getPrimaryKeys(Featuregroup featuregroup) {
+    return featuregroup.getCachedFeaturegroup().getFeaturesExtraConstraints().stream()
+        .filter(CachedFeatureExtraConstraints::getPrimary)
+        .map(CachedFeatureExtraConstraints::getName)
+        .collect(Collectors.toSet());
+  }
+
   public List<FeatureGroupFeatureDTO> getFeaturesDTOOnlineChecked(Featuregroup featuregroup,
                                                                   Project project, Users user)
           throws FeaturestoreException {

diff --git a/.../hops/hopsworks/common/featurestore/featuregroup/stream/StreamFeatureGroupController.java b/.../hops/hopsworks/common/featurestore/featuregroup/stream/StreamFeatureGroupController.java
@@ -101,10 +101,7 @@ public StreamFeatureGroupDTO convertStreamFeatureGroupToDTO(Featuregroup feature
 
   public List<FeatureGroupFeatureDTO> getFeaturesDTO(Featuregroup featuregroup,
                                                      Project project, Users user) throws FeaturestoreException {
-    Set<String> primaryKeys = featuregroup.getStreamFeatureGroup().getFeaturesExtraConstraints().stream()
-        .filter(CachedFeatureExtraConstraints::getPrimary)
-        .map(CachedFeatureExtraConstraints::getName)
-        .collect(Collectors.toSet());
+    Set<String> primaryKeys = getPrimaryKeys(featuregroup);
 
     Set<String> precombineKeys = featuregroup.getStreamFeatureGroup().getFeaturesExtraConstraints().stream()
         .filter(CachedFeatureExtraConstraints::getHudiPrecombineKey)
@@ -131,6 +128,13 @@ public List<FeatureGroupFeatureDTO> getFeaturesDTO(Featuregroup featuregroup,
     return featureGroupFeatures;
   }
 
+  public Set<String> getPrimaryKeys(Featuregroup featuregroup) {
+    return featuregroup.getStreamFeatureGroup().getFeaturesExtraConstraints().stream()
+        .filter(CachedFeatureExtraConstraints::getPrimary)
+        .map(CachedFeatureExtraConstraints::getName)
+        .collect(Collectors.toSet());
+  }
+
   public List<FeatureGroupFeatureDTO> getFeaturesDTOOnlineChecked(Featuregroup featuregroup,
                                                                   Project project, Users user)
       throws FeaturestoreException {