Skip to content

Commit

Permalink
Check the size of result against deduplicated request (#67)
Browse files Browse the repository at this point in the history
  • Loading branch information
pradithya authored and feast-ci-bot committed Jan 15, 2019
1 parent 3f97a8e commit c691191
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 11 deletions.
26 changes: 18 additions & 8 deletions core/src/main/java/feast/core/service/SpecService.java
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
package feast.core.service;

import com.google.common.base.Strings;
import com.google.common.collect.Sets;
import com.google.protobuf.util.JsonFormat;
import feast.core.dao.EntityInfoRepository;
import feast.core.dao.FeatureGroupInfoRepository;
Expand All @@ -38,6 +39,7 @@
import feast.specs.FeatureGroupSpecProto.FeatureGroupSpec;
import feast.specs.FeatureSpecProto.FeatureSpec;
import feast.specs.StorageSpecProto.StorageSpec;
import java.util.Set;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
Expand Down Expand Up @@ -83,8 +85,10 @@ public List<EntityInfo> getEntities(List<String> ids) {
if (ids.size() == 0) {
throw new IllegalArgumentException("ids cannot be empty");
}
List<EntityInfo> entityInfos = this.entityInfoRepository.findAllById(ids);
if (entityInfos.size() < ids.size()) {
Set<String> dedupIds = Sets.newHashSet(ids);

List<EntityInfo> entityInfos = this.entityInfoRepository.findAllById(dedupIds);
if (entityInfos.size() < dedupIds.size()) {
throw new RetrievalException(
"unable to retrieve all entities requested"); // TODO: check and return exactly which ones
}
Expand Down Expand Up @@ -113,8 +117,10 @@ public List<FeatureInfo> getFeatures(List<String> ids) {
if (ids.size() == 0) {
throw new IllegalArgumentException("ids cannot be empty");
}
List<FeatureInfo> featureInfos = this.featureInfoRepository.findAllById(ids);
if (featureInfos.size() < ids.size()) {
Set<String> dedupIds = Sets.newHashSet(ids);

List<FeatureInfo> featureInfos = this.featureInfoRepository.findAllById(dedupIds);
if (featureInfos.size() < dedupIds.size()) {
throw new RetrievalException(
"unable to retrieve all features requested"); // TODO: check and return exactly which ones
}
Expand Down Expand Up @@ -143,8 +149,10 @@ public List<FeatureGroupInfo> getFeatureGroups(List<String> ids) {
if (ids.size() == 0) {
throw new IllegalArgumentException("ids cannot be empty");
}
List<FeatureGroupInfo> featureGroupInfos = this.featureGroupInfoRepository.findAllById(ids);
if (featureGroupInfos.size() < ids.size()) {
Set<String> dedupIds = Sets.newHashSet(ids);

List<FeatureGroupInfo> featureGroupInfos = this.featureGroupInfoRepository.findAllById(dedupIds);
if (featureGroupInfos.size() < dedupIds.size()) {
throw new RetrievalException(
"unable to retrieve all feature groups requested"); // TODO: check and return exactly
// which ones
Expand Down Expand Up @@ -174,8 +182,10 @@ public List<StorageInfo> getStorage(List<String> ids) {
if (ids.size() == 0) {
throw new IllegalArgumentException("ids cannot be empty");
}
List<StorageInfo> storageInfos = this.storageInfoRepository.findAllById(ids);
if (storageInfos.size() < ids.size()) {
Set<String> dedupIds = Sets.newHashSet(ids);

List<StorageInfo> storageInfos = this.storageInfoRepository.findAllById(dedupIds);
if (storageInfos.size() < dedupIds.size()) {
throw new RetrievalException(
"unable to retrieve all storage requested"); // TODO: check and return exactly which ones
}
Expand Down
65 changes: 62 additions & 3 deletions core/src/test/java/feast/core/service/SpecServiceTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@

import static org.hamcrest.Matchers.equalTo;
import static org.junit.Assert.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.mockito.MockitoAnnotations.initMocks;
Expand Down Expand Up @@ -104,7 +105,26 @@ public void shouldGetEntitiesMatchingIds() {
EntityInfo entity2 = newTestEntityInfo("entity2");

ArrayList<String> ids = Lists.newArrayList("entity1", "entity2");
when(entityInfoRepository.findAllById(ids)).thenReturn(Lists.newArrayList(entity1, entity2));
when(entityInfoRepository.findAllById(any(Iterable.class))).thenReturn(Lists.newArrayList(entity1, entity2));
SpecService specService =
new SpecService(
entityInfoRepository,
featureInfoRepository,
storageInfoRepository,
featureGroupInfoRepository,
schemaManager);
List<EntityInfo> actual = specService.getEntities(ids);
List<EntityInfo> expected = Lists.newArrayList(entity1, entity2);
assertThat(actual, equalTo(expected));
}

@Test
public void shouldDeduplicateGetEntities() {
EntityInfo entity1 = newTestEntityInfo("entity1");
EntityInfo entity2 = newTestEntityInfo("entity2");

ArrayList<String> ids = Lists.newArrayList("entity1", "entity2", "entity2");
when(entityInfoRepository.findAllById(any(Iterable.class))).thenReturn(Lists.newArrayList(entity1, entity2));
SpecService specService =
new SpecService(
entityInfoRepository,
Expand Down Expand Up @@ -161,7 +181,26 @@ public void shouldGetFeaturesMatchingIds() {
FeatureInfo feature2 = newTestFeatureInfo("feature2");

ArrayList<String> ids = Lists.newArrayList("entity.none.feature1", "entity.none.feature2");
when(featureInfoRepository.findAllById(ids)).thenReturn(Lists.newArrayList(feature1, feature2));
when(featureInfoRepository.findAllById(any(Iterable.class))).thenReturn(Lists.newArrayList(feature1, feature2));
SpecService specService =
new SpecService(
entityInfoRepository,
featureInfoRepository,
storageInfoRepository,
featureGroupInfoRepository,
schemaManager);
List<FeatureInfo> actual = specService.getFeatures(ids);
List<FeatureInfo> expected = Lists.newArrayList(feature1, feature2);
assertThat(actual, equalTo(expected));
}

@Test
public void shouldDeduplicateGetFeature() {
FeatureInfo feature1 = newTestFeatureInfo("feature1");
FeatureInfo feature2 = newTestFeatureInfo("feature2");

ArrayList<String> ids = Lists.newArrayList("entity.none.feature1", "entity.none.feature2", "entity.none.feature2");
when(featureInfoRepository.findAllById(any(Iterable.class))).thenReturn(Lists.newArrayList(feature1, feature2));
SpecService specService =
new SpecService(
entityInfoRepository,
Expand Down Expand Up @@ -216,7 +255,27 @@ public void shouldGetStorageMatchingIds() {
StorageInfo bqStorage = newTestStorageInfo("BIGQUERY1", "BIGQUERY");

ArrayList<String> ids = Lists.newArrayList("REDIS1", "BIGQUERY1");
when(storageInfoRepository.findAllById(ids))
when(storageInfoRepository.findAllById(any(Iterable.class)))
.thenReturn(Lists.newArrayList(redisStorage, bqStorage));
SpecService specService =
new SpecService(
entityInfoRepository,
featureInfoRepository,
storageInfoRepository,
featureGroupInfoRepository,
schemaManager);
List<StorageInfo> actual = specService.getStorage(ids);
List<StorageInfo> expected = Lists.newArrayList(redisStorage, bqStorage);
assertThat(actual, equalTo(expected));
}

@Test
public void shouldDeduplicateGetStorage() {
StorageInfo redisStorage = newTestStorageInfo("REDIS1", "REDIS");
StorageInfo bqStorage = newTestStorageInfo("BIGQUERY1", "BIGQUERY");

ArrayList<String> ids = Lists.newArrayList("REDIS1", "BIGQUERY1", "BIGQUERY1");
when(storageInfoRepository.findAllById(any(Iterable.class)))
.thenReturn(Lists.newArrayList(redisStorage, bqStorage));
SpecService specService =
new SpecService(
Expand Down

0 comments on commit c691191

Please sign in to comment.