Skip to content

Commit 39d76a8

Browse files
timstims
committed
Only lookup storage specs that we actually need
1 parent 5ac4b9c commit 39d76a8

13 files changed

Lines changed: 360 additions & 196 deletions

File tree

‎ingestion/src/main/java/feast/ingestion/ImportJob.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@
6868

6969
@Slf4j
7070
public class ImportJob {
71+
7172
private static Random random = new Random(System.currentTimeMillis());
7273

7374
private final Pipeline pipeline;

‎ingestion/src/main/java/feast/ingestion/model/Specs.java‎

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
package feast.ingestion.model;
1919

2020
import com.google.common.base.Preconditions;
21+
import com.google.common.collect.Lists;
2122
import feast.ingestion.service.SpecService;
2223
import feast.specs.EntitySpecProto.EntitySpec;
2324
import feast.specs.FeatureSpecProto.FeatureSpec;
@@ -31,10 +32,15 @@
3132
import java.util.Map.Entry;
3233
import lombok.Builder;
3334
import lombok.Getter;
35+
import lombok.ToString;
36+
import lombok.extern.slf4j.Slf4j;
3437

3538
@Builder
3639
@Getter
40+
@Slf4j
41+
@ToString
3742
public class Specs implements Serializable {
43+
3844
private String jobName;
3945
private ImportSpec importSpec;
4046
private Map<String, EntitySpec> entitySpecs;
@@ -57,14 +63,23 @@ public static Specs of(String jobName, ImportSpec importSpec, SpecService specSe
5763
specsBuilder.featureSpecs(specService.getFeatureSpecs(featureIds));
5864

5965
List<String> entityNames = importSpec.getEntitiesList();
66+
List<String> storageIds = Lists.newArrayList();
6067
for (FeatureSpec featureSpec : specsBuilder.featureSpecs.values()) {
6168
Preconditions.checkArgument(
6269
entityNames.contains(featureSpec.getEntity()),
6370
"Feature has entity not listed in import spec featureSpec=" + featureSpec.toString());
71+
String servingId = featureSpec.getDataStores().getServing().getId();
72+
if (!servingId.isEmpty()) {
73+
storageIds.add(servingId);
74+
}
75+
String warehouseId = featureSpec.getDataStores().getWarehouse().getId();
76+
if (!warehouseId.isEmpty()) {
77+
storageIds.add(warehouseId);
78+
}
6479
}
6580
specsBuilder.entitySpecs(specService.getEntitySpecs(entityNames));
6681

67-
specsBuilder.storageSpecs(specService.getAllStorageSpecs());
82+
specsBuilder.storageSpecs(specService.getStorageSpecs(storageIds));
6883

6984
return specsBuilder.build();
7085
} catch (RuntimeException e) {
@@ -79,13 +94,19 @@ public void validate() {
7994

8095
// Sanity checks that our maps are built correctly
8196
for (Entry<String, FeatureSpec> entry : featureSpecs.entrySet()) {
82-
Preconditions.checkArgument(entry.getKey().equals(entry.getValue().getId()));
97+
Preconditions.checkArgument(entry.getKey().equals(entry.getValue().getId()),
98+
String.format("Feature id does not match spec %s!=%s", entry.getKey(),
99+
entry.getValue().getId()));
83100
}
84101
for (Entry<String, EntitySpec> entry : entitySpecs.entrySet()) {
85-
Preconditions.checkArgument(entry.getKey().equals(entry.getValue().getName()));
102+
Preconditions.checkArgument(entry.getKey().equals(entry.getValue().getName()),
103+
String.format("Entity name does not match spec %s!=%s", entry.getKey(),
104+
entry.getValue().getName()));
86105
}
87106
for (Entry<String, StorageSpec> entry : storageSpecs.entrySet()) {
88-
Preconditions.checkArgument(entry.getKey().equals(entry.getValue().getId()));
107+
Preconditions.checkArgument(entry.getKey().equals(entry.getValue().getId()),
108+
String.format("Storage id does not match spec %s!=%s", entry.getKey(),
109+
entry.getValue().getId()));
89110
}
90111

91112
for (FeatureSpec featureSpec : featureSpecs.values()) {
@@ -96,17 +117,21 @@ public void validate() {
96117
"Feature %s references unknown entity %s",
97118
featureSpec.getId(), featureSpec.getEntity()));
98119
// Check that feature has a matching serving store
99-
Preconditions.checkArgument(
100-
storageSpecs.containsKey(featureSpec.getDataStores().getServing().getId()),
101-
String.format(
102-
"Feature %s references unknown serving store %s",
103-
featureSpec.getId(), featureSpec.getDataStores().getServing().getId()));
120+
if (!featureSpec.getDataStores().getServing().getId().isEmpty()) {
121+
Preconditions.checkArgument(
122+
storageSpecs.containsKey(featureSpec.getDataStores().getServing().getId()),
123+
String.format(
124+
"Feature %s references unknown serving store %s",
125+
featureSpec.getId(), featureSpec.getDataStores().getServing().getId()));
126+
}
104127
// Check that feature has a matching warehouse store
105-
Preconditions.checkArgument(
106-
storageSpecs.containsKey(featureSpec.getDataStores().getWarehouse().getId()),
107-
String.format(
108-
"Feature %s references unknown warehouse store %s",
109-
featureSpec.getId(), featureSpec.getDataStores().getWarehouse().getId()));
128+
if (!featureSpec.getDataStores().getWarehouse().getId().isEmpty()) {
129+
Preconditions.checkArgument(
130+
storageSpecs.containsKey(featureSpec.getDataStores().getWarehouse().getId()),
131+
String.format(
132+
"Feature %s references unknown warehouse store %s",
133+
featureSpec.getId(), featureSpec.getDataStores().getWarehouse().getId()));
134+
}
110135
}
111136
}
112137

‎ingestion/src/main/java/feast/ingestion/service/FileSpecService.java‎

Lines changed: 6 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@
2121
import com.google.protobuf.InvalidProtocolBufferException;
2222
import com.google.protobuf.Message;
2323
import com.google.protobuf.util.JsonFormat;
24+
import feast.ingestion.exceptions.SpecNotFound;
25+
import feast.ingestion.util.PathUtil;
26+
import feast.specs.EntitySpecProto.EntitySpec;
27+
import feast.specs.FeatureSpecProto.FeatureSpec;
28+
import feast.specs.StorageSpecProto.StorageSpec;
2429
import java.io.FileNotFoundException;
2530
import java.io.IOException;
2631
import java.io.InputStream;
@@ -35,11 +40,6 @@
3540
import java.util.function.Function;
3641
import java.util.stream.Collectors;
3742
import lombok.AllArgsConstructor;
38-
import feast.ingestion.exceptions.SpecNotFound;
39-
import feast.ingestion.util.PathUtil;
40-
import feast.specs.EntitySpecProto.EntitySpec;
41-
import feast.specs.FeatureSpecProto.FeatureSpec;
42-
import feast.specs.StorageSpecProto.StorageSpec;
4343

4444
@AllArgsConstructor
4545
public class FileSpecService implements SpecService {
@@ -112,31 +112,16 @@ public Map<String, EntitySpec> getEntitySpecs(Iterable<String> entityIds) {
112112
return getSpecs(EntitySpec.getDefaultInstance(), ENTITY_SPEC, entityIds);
113113
}
114114

115-
@Override
116-
public Map<String, EntitySpec> getAllEntitySpecs() {
117-
return getAllSpecs(EntitySpec.getDefaultInstance(), ENTITY_SPEC, EntitySpec::getName);
118-
}
119-
120115
@Override
121116
public Map<String, FeatureSpec> getFeatureSpecs(Iterable<String> featureIds) {
122117
return getSpecs(FeatureSpec.getDefaultInstance(), FEATURE_SPEC, featureIds);
123118
}
124119

125-
@Override
126-
public Map<String, FeatureSpec> getAllFeatureSpecs() {
127-
return getAllSpecs(FeatureSpec.getDefaultInstance(), FEATURE_SPEC, FeatureSpec::getId);
128-
}
129-
130120
@Override
131121
public Map<String, StorageSpec> getStorageSpecs(Iterable<String> storageIds) {
132122
return getSpecs(StorageSpec.getDefaultInstance(), STORAGE_SPEC, storageIds);
133123
}
134124

135-
@Override
136-
public Map<String, StorageSpec> getAllStorageSpecs() {
137-
return getAllSpecs(StorageSpec.getDefaultInstance(), STORAGE_SPEC, StorageSpec::getId);
138-
}
139-
140125
private <T extends Message> void putSpecs(
141126
String type, Function<T, String> keyFunc, Iterable<T> specs) {
142127
for (T spec : specs) {
@@ -169,6 +154,7 @@ public void putStorageSpecs(Iterable<StorageSpec> storageSpecs) {
169154

170155
@AllArgsConstructor
171156
public static class Builder implements SpecService.Builder {
157+
172158
private String basePath;
173159

174160
@Override

‎ingestion/src/main/java/feast/ingestion/service/SpecService.java‎

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -38,13 +38,6 @@ public interface SpecService {
3838
*/
3939
Map<String, EntitySpec> getEntitySpecs(Iterable<String> entityIds);
4040

41-
/**
42-
* Get all {@link EntitySpec} from Core API.
43-
*
44-
* @return map of {@link EntitySpec}, where the key is the entity name.
45-
*/
46-
Map<String, EntitySpec> getAllEntitySpecs();
47-
4841
/**
4942
* Get a map of {@link FeatureSpec} from Core API, given a collection of featureId.
5043
*
@@ -54,13 +47,6 @@ public interface SpecService {
5447
*/
5548
Map<String, FeatureSpec> getFeatureSpecs(Iterable<String> featureIds);
5649

57-
/**
58-
* Get all {@link FeatureSpec} available in Core API.
59-
*
60-
* @return map of {@link FeatureSpec}, where the key is feature id.
61-
*/
62-
Map<String, FeatureSpec> getAllFeatureSpecs();
63-
6450
/**
6551
* Get map of {@link StorageSpec} from Core API, given a collection of storageId.
6652
*
@@ -70,13 +56,6 @@ public interface SpecService {
7056
*/
7157
Map<String, StorageSpec> getStorageSpecs(Iterable<String> storageIds);
7258

73-
/**
74-
* Get all {@link StorageSpec} from Core API.
75-
*
76-
* @return map of {@link StorageSpec}, where the key is storage id.
77-
*/
78-
Map<String, StorageSpec> getAllStorageSpecs();
79-
8059
interface Builder extends Serializable {
8160
SpecService build();
8261
}

‎ingestion/src/test/java/feast/ingestion/ImportJobCSVTest.java‎

Lines changed: 19 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,11 @@
1717

1818
package feast.ingestion;
1919

20+
import static feast.FeastMatchers.hasCount;
21+
import static feast.ToOrderedFeatureRows.orderedFeatureRow;
22+
import static feast.storage.MockErrorsStore.MOCK_ERRORS_STORE_TYPE;
23+
import static org.junit.Assert.assertEquals;
24+
2025
import com.google.common.base.Charsets;
2126
import com.google.common.collect.Lists;
2227
import com.google.common.io.Files;
@@ -30,6 +35,7 @@
3035
import feast.ingestion.model.Features;
3136
import feast.ingestion.model.Values;
3237
import feast.ingestion.options.ImportJobOptions;
38+
import feast.ingestion.service.SpecRetrievalException;
3339
import feast.ingestion.util.ProtoUtil;
3440
import feast.specs.ImportSpecProto.ImportSpec;
3541
import feast.storage.MockErrorsStore;
@@ -41,6 +47,11 @@
4147
import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
4248
import feast.types.FeatureRowProto.FeatureRow;
4349
import feast.types.GranularityProto.Granularity;
50+
import java.io.File;
51+
import java.io.IOException;
52+
import java.nio.file.Path;
53+
import java.nio.file.Paths;
54+
import java.util.List;
4455
import lombok.extern.slf4j.Slf4j;
4556
import org.apache.beam.sdk.options.PipelineOptionsFactory;
4657
import org.apache.beam.sdk.testing.PAssert;
@@ -52,24 +63,14 @@
5263
import org.junit.Test;
5364
import org.junit.rules.TemporaryFolder;
5465

55-
import java.io.File;
56-
import java.io.IOException;
57-
import java.nio.file.Path;
58-
import java.nio.file.Paths;
59-
import java.util.List;
60-
61-
import static feast.FeastMatchers.hasCount;
62-
import static feast.ToOrderedFeatureRows.orderedFeatureRow;
63-
import static feast.storage.MockErrorsStore.MOCK_ERRORS_STORE_TYPE;
64-
import static org.junit.Assert.assertEquals;
65-
import static org.junit.Assert.fail;
66-
6766
@Slf4j
6867
public class ImportJobCSVTest {
6968

70-
@Rule public TemporaryFolder folder = new TemporaryFolder();
69+
@Rule
70+
public TemporaryFolder folder = new TemporaryFolder();
7171

72-
@Rule public TestPipeline testPipeline = TestPipeline.create();
72+
@Rule
73+
public TestPipeline testPipeline = TestPipeline.create();
7374

7475
public ImportSpec initImportSpec(ImportSpec importSpec, String dataFile) {
7576
return importSpec.toBuilder().putOptions("path", dataFile).build();
@@ -125,7 +126,7 @@ public void testImportCSV() throws IOException {
125126

126127
PCollection<FeatureRowExtended> writtenToWarehouse =
127128
PCollectionList.of(
128-
WarehouseStoreService.get(MockWarehouseStore.class).getWrite().getInputs())
129+
WarehouseStoreService.get(MockWarehouseStore.class).getWrite().getInputs())
129130
.apply("flatten warehouse input", Flatten.pCollections());
130131

131132
PCollection<FeatureRowExtended> writtenToErrors =
@@ -173,7 +174,7 @@ public void testImportCSV() throws IOException {
173174
testPipeline.run();
174175
}
175176

176-
@Test
177+
@Test(expected = SpecRetrievalException.class)
177178
public void testImportCSVUnknownServingStoreError() throws IOException {
178179
ImportSpec importSpec =
179180
ProtoUtil.decodeProtoYaml(
@@ -189,8 +190,7 @@ public void testImportCSVUnknownServingStoreError() throws IOException {
189190
+ " timestampValue: 2018-09-25T00:00:00.000Z\n"
190191
+ " fields:\n"
191192
+ " - name: id\n"
192-
+ " - featureId: testEntity.none.redisInt32\n" // Redis is not available by
193-
// default from the json specs
193+
+ " - featureId: testEntity.none.unknownInt32\n" // Unknown store is not available
194194
+ " - featureId: testEntity.none.testString\n"
195195
+ "\n",
196196
ImportSpec.getDefaultInstance());
@@ -209,14 +209,7 @@ public void testImportCSVUnknownServingStoreError() throws IOException {
209209
injector.getInstance(ImportJob.class);
210210

211211
// Job should fail during expand(), so we don't even need to start the pipeline.
212-
try {
213-
job.expand();
214-
fail("Should not reach here, we should have thrown an exception");
215-
} catch (IllegalArgumentException e) {
216-
assertEquals(
217-
"Feature testEntity.none.redisInt32 references unknown serving store REDIS1",
218-
e.getMessage());
219-
}
212+
job.expand();
220213
}
221214

222215
@Test

0 commit comments

Comments
 (0)