Skip to content

Commit

Permalink
[FSTORE-1134] Enable updating embedding (#1496)
Browse files Browse the repository at this point in the history
(cherry picked from commit 2191bad)
  • Loading branch information
kennethmhc authored Feb 21, 2024
1 parent 52a5772 commit 8ca901d
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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 {
Expand Down

0 comments on commit 8ca901d

Please sign in to comment.