Skip to content

Commit 029edb4

Browse files
timsfeast-ci-bot
tims
authored and
feast-ci-bot
committed
Only lookup storage specs that we actually need (#52)
* Only lookup storage specs that we actually need * Fix error handling and do not use empty string for tuple tags * ensure that rows that have no write transform still get returned in the split main output, so they can be written downstream (eg to the warehouse) * Fix failing unit test
1 parent 5ac4b9c commit 029edb4

21 files changed

Lines changed: 597 additions & 318 deletions

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

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

2020
import com.google.api.services.bigquery.model.TableRow;
21+
import com.google.common.collect.Lists;
2122
import com.google.inject.Guice;
2223
import com.google.inject.Inject;
2324
import com.google.inject.Injector;
@@ -42,6 +43,7 @@
4243
import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
4344
import feast.types.FeatureRowProto.FeatureRow;
4445
import java.util.Arrays;
46+
import java.util.List;
4547
import java.util.Random;
4648
import lombok.extern.slf4j.Slf4j;
4749
import org.apache.beam.runners.dataflow.DataflowPipelineJob;
@@ -53,13 +55,15 @@
5355
import org.apache.beam.sdk.extensions.protobuf.ProtoCoder;
5456
import org.apache.beam.sdk.io.gcp.bigquery.TableRowJsonCoder;
5557
import org.apache.beam.sdk.options.PipelineOptionsFactory;
58+
import org.apache.beam.sdk.transforms.Flatten;
5659
import org.apache.beam.sdk.transforms.ParDo;
5760
import org.apache.beam.sdk.transforms.Sample;
5861
import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
5962
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
6063
import org.apache.beam.sdk.transforms.windowing.Window;
6164
import org.apache.beam.sdk.values.PCollection;
6265
import org.apache.beam.sdk.values.PCollection.IsBounded;
66+
import org.apache.beam.sdk.values.PCollectionList;
6367
import org.apache.beam.sdk.values.TypeDescriptor;
6468
import org.apache.commons.codec.digest.DigestUtils;
6569
import org.joda.time.DateTime;
@@ -68,6 +72,7 @@
6872

6973
@Slf4j
7074
public class ImportJob {
75+
7176
private static Random random = new Random(System.currentTimeMillis());
7277

7378
private final Pipeline pipeline;
@@ -164,17 +169,23 @@ public void expand() {
164169
"Round event timestamps to granularity",
165170
ParDo.of(new RoundEventTimestampsDoFn())),
166171
pFeatureRows.getErrors());
172+
167173
if (!dryRun) {
174+
List<PCollection<FeatureRowExtended>> errors = Lists.newArrayList();
168175
pFeatureRows = pFeatureRows.apply("Write to Serving Stores", servingStoreTransform);
169-
pFeatureRows.getErrors().apply("Write serving errors", errorsStoreTransform);
176+
errors.add(pFeatureRows.getErrors());
177+
pFeatureRows = PFeatureRows.of(pFeatureRows.getMain());
170178

171179
log.info(
172180
"A sample of any 2 rows from each of MAIN, RETRIES and ERRORS will logged for convenience");
173181
logNRows(pFeatureRows, "Output sample", 2);
174182

175183
PFeatureRows.of(pFeatureRows.getMain())
176184
.apply("Write to Warehouse Stores", warehouseStoreTransform);
177-
pFeatureRows.getErrors().apply("Write warehouse errors", errorsStoreTransform);
185+
errors.add(pFeatureRows.getErrors());
186+
187+
PCollectionList.of(errors).apply("flatten errors", Flatten.pCollections())
188+
.apply("Write serving errors", errorsStoreTransform);
178189
}
179190
}
180191

‎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/main/java/feast/ingestion/transform/ErrorsStoreTransform.java‎

Lines changed: 26 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -17,72 +17,66 @@
1717

1818
package feast.ingestion.transform;
1919

20+
import static com.google.common.base.Preconditions.checkArgument;
2021
import static feast.ingestion.util.JsonUtil.convertJsonStringToMap;
2122

2223
import com.google.inject.Inject;
2324
import feast.ingestion.model.Specs;
2425
import feast.ingestion.options.ImportJobOptions;
2526
import feast.ingestion.transform.FeatureIO.Write;
26-
import feast.ingestion.transform.fn.LoggerDoFn;
2727
import feast.specs.StorageSpecProto.StorageSpec;
2828
import feast.storage.ErrorsStore;
2929
import feast.storage.noop.NoOpIO;
3030
import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
3131
import java.util.List;
3232
import lombok.extern.slf4j.Slf4j;
33-
import org.apache.beam.sdk.transforms.ParDo;
3433
import org.apache.beam.sdk.values.PCollection;
3534
import org.apache.beam.sdk.values.PDone;
36-
import org.slf4j.event.Level;
35+
import org.apache.hadoop.hbase.util.Strings;
3736

3837
@Slf4j
3938
public class ErrorsStoreTransform extends FeatureIO.Write {
4039

41-
public static final String ERRORS_STORE_STDERR = "stderr";
42-
public static final String ERRORS_STORE_STDOUT = "stdout";
43-
public static final String ERRORS_STORE_JSON = "file.json";
44-
4540
private String errorsStoreType;
4641
private StorageSpec errorsStoreSpec;
47-
private ErrorsStore errorsStore;
4842
private Specs specs;
43+
private List<ErrorsStore> errorsStores;
4944

5045
@Inject
5146
public ErrorsStoreTransform(
5247
ImportJobOptions options, Specs specs, List<ErrorsStore> errorsStores) {
5348
this.specs = specs;
49+
this.errorsStores = errorsStores;
5450
this.errorsStoreType = options.getErrorsStoreType();
5551

56-
for (ErrorsStore errorsStore : errorsStores) {
57-
if (errorsStore.getType().equals(errorsStoreType)) {
58-
this.errorsStore = errorsStore;
59-
}
52+
if (!Strings.isEmpty(errorsStoreType)) {
53+
this.errorsStoreSpec =
54+
StorageSpec.newBuilder()
55+
.setType(errorsStoreType)
56+
.putAllOptions(convertJsonStringToMap(options.getErrorsStoreOptions()))
57+
.build();
6058
}
61-
62-
this.errorsStoreSpec =
63-
StorageSpec.newBuilder()
64-
.setType(errorsStoreType)
65-
.putAllOptions(convertJsonStringToMap(options.getErrorsStoreOptions()))
66-
.build();
6759
}
6860

6961
@Override
7062
public PDone expand(PCollection<FeatureRowExtended> input) {
71-
switch (errorsStoreType) {
72-
case ERRORS_STORE_STDOUT:
73-
input.apply("Log errors to STDOUT", ParDo.of(new LoggerDoFn(Level.INFO)));
74-
break;
75-
case ERRORS_STORE_STDERR:
76-
input.apply("Log errors to STDERR", ParDo.of(new LoggerDoFn(Level.ERROR)));
77-
break;
78-
default:
79-
if (errorsStore == null) {
80-
log.warn("No valid errors store specified, errors will be discarded");
81-
return input.apply(new NoOpIO.Write());
82-
}
83-
Write write = errorsStore.create(this.errorsStoreSpec, specs);
84-
return input.apply(write);
63+
Write write;
64+
if (Strings.isEmpty(errorsStoreType)) {
65+
write = new NoOpIO.Write();
66+
} else {
67+
write = getErrorStore().create(this.errorsStoreSpec, specs);
8568
}
69+
input.apply("errors to " + String.valueOf(errorsStoreType), write);
8670
return PDone.in(input.getPipeline());
8771
}
72+
73+
ErrorsStore getErrorStore() {
74+
checkArgument(!errorsStoreType.isEmpty(), "Errors store type not provided");
75+
for (ErrorsStore errorsStore : errorsStores) {
76+
if (errorsStore.getType().equals(errorsStoreType)) {
77+
return errorsStore;
78+
}
79+
}
80+
throw new IllegalArgumentException("Errors store type not found");
81+
}
8882
}

0 commit comments

Comments
 (0)