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

(cherry picked from commit 2191bad)
logicalclocks · Feb 21, 2024 · 8ca901d · 8ca901d
1 parent 52a5772
commit 8ca901d
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 {