Skip to content

Commit 79eb4ab

Browse files
Chen Zhilingthirteen37
authored and
feast-ci-bot
committed
Zl/ingestion fixes (#286)
* 1. Refactor graph to only have a single read from source 2. Move validation outside of read from source * Use post-validation transform for metrics * Tidy up code Co-Authored-By: Yu-Xi Lim <thirteen37@users.noreply.github.com>
1 parent f3b1ce7 commit 79eb4ab

8 files changed

Lines changed: 241 additions & 228 deletions

File tree

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

Lines changed: 62 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@
22

33
import com.google.protobuf.InvalidProtocolBufferException;
44
import feast.core.FeatureSetProto.FeatureSetSpec;
5+
import feast.core.SourceProto.Source;
56
import feast.core.StoreProto.Store;
67
import feast.ingestion.options.ImportOptions;
78
import feast.ingestion.transform.ReadFromSource;
9+
import feast.ingestion.transform.ValidateFeatureRows;
810
import feast.ingestion.transform.WriteFailedElementToBigQuery;
911
import feast.ingestion.transform.WriteToStore;
1012
import feast.ingestion.transform.metrics.WriteMetricsTransform;
@@ -14,20 +16,26 @@
1416
import feast.ingestion.values.FailedElement;
1517
import feast.types.FeatureRowProto.FeatureRow;
1618
import java.util.List;
19+
import java.util.Map;
20+
import java.util.stream.Collectors;
1721
import org.apache.beam.sdk.Pipeline;
1822
import org.apache.beam.sdk.PipelineResult;
1923
import org.apache.beam.sdk.options.PipelineOptionsFactory;
2024
import org.apache.beam.sdk.options.PipelineOptionsValidator;
2125
import org.apache.beam.sdk.values.PCollectionTuple;
2226
import org.apache.beam.sdk.values.TupleTag;
27+
import org.apache.commons.lang3.tuple.Pair;
2328
import org.slf4j.Logger;
2429

2530
public class ImportJob {
31+
2632
// Tag for main output containing Feature Row that has been successfully processed.
27-
private static final TupleTag<FeatureRow> FEATURE_ROW_OUT = new TupleTag<FeatureRow>() {};
33+
private static final TupleTag<FeatureRow> FEATURE_ROW_OUT = new TupleTag<FeatureRow>() {
34+
};
2835

2936
// Tag for deadletter output containing elements and error messages from invalid input/transform.
30-
private static final TupleTag<FailedElement> DEADLETTER_OUT = new TupleTag<FailedElement>() {};
37+
private static final TupleTag<FailedElement> DEADLETTER_OUT = new TupleTag<FailedElement>() {
38+
};
3139
private static final Logger log = org.slf4j.LoggerFactory.getLogger(ImportJob.class);
3240

3341
/**
@@ -46,14 +54,17 @@ public static PipelineResult runPipeline(ImportOptions options)
4654
/*
4755
* Steps:
4856
* 1. Read messages from Feast Source as FeatureRow
49-
* 2. Write FeatureRow to the corresponding Store
50-
* 3. Write elements that failed to be processed to a dead letter queue.
51-
* 4. Write metrics to a metrics sink
57+
* 2. Validate the feature rows to ensure the schema matches what is registered to the system
58+
* 3. Write FeatureRow to the corresponding Store
59+
* 4. Write elements that failed to be processed to a dead letter queue.
60+
* 5. Write metrics to a metrics sink
5261
*/
5362

5463
PipelineOptionsValidator.validate(ImportOptions.class, options);
5564
Pipeline pipeline = Pipeline.create(options);
5665

66+
log.info("Starting import job with settings: \n{}", options.toString());
67+
5768
List<FeatureSetSpec> featureSetSpecs =
5869
SpecUtil.parseFeatureSetSpecJsonList(options.getFeatureSetSpecJson());
5970
List<Store> stores = SpecUtil.parseStoreJsonList(options.getStoreJson());
@@ -62,44 +73,71 @@ public static PipelineResult runPipeline(ImportOptions options)
6273
List<FeatureSetSpec> subscribedFeatureSets =
6374
SpecUtil.getSubscribedFeatureSets(store.getSubscriptionsList(), featureSetSpecs);
6475

76+
// Generate tags by key
77+
Map<String, TupleTag<FeatureRow>> featureSetTagsByKey = subscribedFeatureSets.stream()
78+
.map(fs -> {
79+
String id = String.format("%s:%s", fs.getName(), fs.getVersion());
80+
return Pair.of(id, new TupleTag<FeatureRow>(id) {
81+
});
82+
})
83+
.collect(Collectors.toMap(Pair::getLeft, Pair::getRight));
84+
85+
// TODO: make the source part of the job initialisation options
86+
Source source = subscribedFeatureSets.get(0).getSource();
87+
88+
// Step 1. Read messages from Feast Source as FeatureRow.
89+
PCollectionTuple convertedFeatureRows =
90+
pipeline.apply(
91+
"ReadFeatureRowFromSource",
92+
ReadFromSource.newBuilder()
93+
.setSource(source)
94+
.setFeatureSetTagByKey(featureSetTagsByKey)
95+
.setFailureTag(DEADLETTER_OUT)
96+
.build());
97+
6598
for (FeatureSetSpec featureSet : subscribedFeatureSets) {
6699
// Ensure Store has valid configuration and Feast can access it.
67100
StoreUtil.setupStore(store, featureSet);
101+
String id = String.format("%s:%s", featureSet.getName(), featureSet.getVersion());
102+
103+
// Step 2. Validate incoming FeatureRows
104+
PCollectionTuple validatedRows = convertedFeatureRows
105+
.get(featureSetTagsByKey.get(id))
106+
.apply(ValidateFeatureRows.newBuilder()
107+
.setFeatureSetSpec(featureSet)
108+
.setSuccessTag(FEATURE_ROW_OUT)
109+
.setFailureTag(DEADLETTER_OUT)
110+
.build());
68111

69-
// Step 1. Read messages from Feast Source as FeatureRow.
70-
PCollectionTuple convertedFeatureRows =
71-
pipeline.apply(
72-
"ReadFeatureRowFromSource",
73-
ReadFromSource.newBuilder()
74-
.setSource(featureSet.getSource())
75-
.setFieldByName(SpecUtil.getFieldByName(featureSet))
76-
.setFeatureSetName(featureSet.getName())
77-
.setFeatureSetVersion(featureSet.getVersion())
78-
.setSuccessTag(FEATURE_ROW_OUT)
79-
.setFailureTag(DEADLETTER_OUT)
80-
.build());
81-
82-
// Step 2. Write FeatureRow to the corresponding Store.
83-
convertedFeatureRows
112+
// Step 3. Write FeatureRow to the corresponding Store.
113+
validatedRows
84114
.get(FEATURE_ROW_OUT)
85115
.apply(
86116
"WriteFeatureRowToStore",
87117
WriteToStore.newBuilder().setFeatureSetSpec(featureSet).setStore(store).build());
88118

89-
// Step 3. Write FailedElements to a dead letter table in BigQuery.
119+
// Step 4. Write FailedElements to a dead letter table in BigQuery.
90120
if (options.getDeadLetterTableSpec() != null) {
91121
convertedFeatureRows
92122
.get(DEADLETTER_OUT)
93123
.apply(
94-
"WriteFailedElements",
124+
"WriteFailedElements_ReadFromSource",
125+
WriteFailedElementToBigQuery.newBuilder()
126+
.setJsonSchema(ResourceUtil.getDeadletterTableSchemaJson())
127+
.setTableSpec(options.getDeadLetterTableSpec())
128+
.build());
129+
130+
validatedRows
131+
.get(DEADLETTER_OUT)
132+
.apply("WriteFailedElements_ValidateRows",
95133
WriteFailedElementToBigQuery.newBuilder()
96134
.setJsonSchema(ResourceUtil.getDeadletterTableSchemaJson())
97135
.setTableSpec(options.getDeadLetterTableSpec())
98136
.build());
99137
}
100138

101-
// Step 4. Write metrics to a metrics sink.
102-
convertedFeatureRows
139+
// Step 5. Write metrics to a metrics sink.
140+
validatedRows
103141
.apply("WriteMetrics", WriteMetricsTransform.newBuilder()
104142
.setFeatureSetSpec(featureSet)
105143
.setStoreName(store.getName())

‎ingestion/src/main/java/feast/ingestion/transform/ReadFromSource.java‎

Lines changed: 10 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,41 +2,29 @@
22

33
import com.google.auto.value.AutoValue;
44
import com.google.common.base.Preconditions;
5-
import com.google.protobuf.InvalidProtocolBufferException;
5+
import com.google.common.collect.Lists;
66
import feast.core.SourceProto.Source;
77
import feast.core.SourceProto.SourceType;
88
import feast.ingestion.transform.fn.KafkaRecordToFeatureRowDoFn;
99
import feast.ingestion.values.FailedElement;
10-
import feast.ingestion.values.Field;
1110
import feast.types.FeatureRowProto.FeatureRow;
12-
import feast.types.FieldProto;
13-
import feast.types.ValueProto.Value.ValCase;
14-
import java.util.Base64;
11+
import java.util.List;
1512
import java.util.Map;
1613
import org.apache.beam.sdk.io.kafka.KafkaIO;
17-
import org.apache.beam.sdk.io.kafka.KafkaRecord;
18-
import org.apache.beam.sdk.transforms.DoFn;
1914
import org.apache.beam.sdk.transforms.PTransform;
2015
import org.apache.beam.sdk.transforms.ParDo;
2116
import org.apache.beam.sdk.values.PBegin;
2217
import org.apache.beam.sdk.values.PCollectionTuple;
2318
import org.apache.beam.sdk.values.TupleTag;
2419
import org.apache.beam.sdk.values.TupleTagList;
2520
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
26-
import org.apache.commons.lang3.exception.ExceptionUtils;
2721

2822
@AutoValue
2923
public abstract class ReadFromSource extends PTransform<PBegin, PCollectionTuple> {
3024

3125
public abstract Source getSource();
3226

33-
public abstract Map<String, Field> getFieldByName();
34-
35-
public abstract String getFeatureSetName();
36-
37-
public abstract int getFeatureSetVersion();
38-
39-
public abstract TupleTag<FeatureRow> getSuccessTag();
27+
public abstract Map<String, TupleTag<FeatureRow>> getFeatureSetTagByKey();
4028

4129
public abstract TupleTag<FailedElement> getFailureTag();
4230

@@ -49,13 +37,8 @@ public abstract static class Builder {
4937

5038
public abstract Builder setSource(Source source);
5139

52-
public abstract Builder setFeatureSetName(String featureSetName);
53-
54-
public abstract Builder setFeatureSetVersion(int featureSetVersion);
55-
56-
public abstract Builder setFieldByName(Map<String, Field> fieldByName);
57-
58-
public abstract Builder setSuccessTag(TupleTag<FeatureRow> successTag);
40+
public abstract Builder setFeatureSetTagByKey(
41+
Map<String, TupleTag<FeatureRow>> featureSetTagByKey);
5942

6043
public abstract Builder setFailureTag(TupleTag<FailedElement> failureTag);
6144

@@ -93,13 +76,13 @@ public PCollectionTuple expand(PBegin input) {
9376
.commitOffsetsInFinalize())
9477
.apply(
9578
"KafkaRecordToFeatureRow", ParDo.of(KafkaRecordToFeatureRowDoFn.newBuilder()
96-
.setFeatureSetName(getFeatureSetName())
97-
.setFeatureSetVersion(getFeatureSetVersion())
98-
.setFieldByName(getFieldByName())
99-
.setSuccessTag(getSuccessTag())
79+
.setFeatureSetTagByKey(getFeatureSetTagByKey())
10080
.setFailureTag(getFailureTag())
10181
.build())
102-
.withOutputTags(getSuccessTag(), TupleTagList.of(getFailureTag())));
82+
.withOutputTags(new TupleTag<FeatureRow>("placeholder") {},
83+
TupleTagList.of(Lists
84+
.newArrayList(getFeatureSetTagByKey().values()))
85+
.and(getFailureTag())));
10386
}
10487

10588
private String generateConsumerGroupId(String jobName) {

‎ingestion/src/main/java/feast/ingestion/transform/ValidateFeatureRows.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package feast.ingestion.transform;
2+
3+
import com.google.auto.value.AutoValue;
4+
import feast.core.FeatureSetProto.FeatureSetSpec;
5+
import feast.ingestion.transform.fn.ValidateFeatureRowDoFn;
6+
import feast.ingestion.utils.SpecUtil;
7+
import feast.ingestion.values.FailedElement;
8+
import feast.ingestion.values.Field;
9+
import feast.types.FeatureRowProto.FeatureRow;
10+
import java.util.Map;
11+
import org.apache.beam.sdk.transforms.PTransform;
12+
import org.apache.beam.sdk.transforms.ParDo;
13+
import org.apache.beam.sdk.values.PCollection;
14+
import org.apache.beam.sdk.values.PCollectionTuple;
15+
import org.apache.beam.sdk.values.TupleTag;
16+
import org.apache.beam.sdk.values.TupleTagList;
17+
18+
@AutoValue
19+
public abstract class ValidateFeatureRows extends
20+
PTransform<PCollection<FeatureRow>, PCollectionTuple> {
21+
22+
public abstract FeatureSetSpec getFeatureSetSpec();
23+
24+
public abstract TupleTag<FeatureRow> getSuccessTag();
25+
26+
public abstract TupleTag<FailedElement> getFailureTag();
27+
28+
public static Builder newBuilder() {
29+
return new AutoValue_ValidateFeatureRows.Builder();
30+
}
31+
32+
@AutoValue.Builder
33+
public abstract static class Builder {
34+
35+
public abstract Builder setFeatureSetSpec(FeatureSetSpec featureSetSpec);
36+
37+
public abstract Builder setSuccessTag(TupleTag<FeatureRow> successTag);
38+
39+
public abstract Builder setFailureTag(TupleTag<FailedElement> failureTag);
40+
41+
public abstract ValidateFeatureRows build();
42+
}
43+
44+
@Override
45+
public PCollectionTuple expand(PCollection<FeatureRow> input) {
46+
Map<String, Field> fieldsByName = SpecUtil
47+
.getFieldByName(getFeatureSetSpec());
48+
49+
return input.apply("ValidateFeatureRows",
50+
ParDo.of(ValidateFeatureRowDoFn.newBuilder()
51+
.setFeatureSetName(getFeatureSetSpec().getName())
52+
.setFeatureSetVersion(getFeatureSetSpec().getVersion())
53+
.setFieldByName(fieldsByName)
54+
.setSuccessTag(getSuccessTag())
55+
.setFailureTag(getFailureTag())
56+
.build())
57+
.withOutputTags(getSuccessTag(), TupleTagList.of(getFailureTag())));
58+
}
59+
}

‎ingestion/src/main/java/feast/ingestion/transform/fn/KafkaRecordToFeatureRowDoFn.java‎

Lines changed: 11 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
import com.google.auto.value.AutoValue;
44
import com.google.protobuf.InvalidProtocolBufferException;
5+
import feast.ingestion.transform.ReadFromSource;
6+
import feast.ingestion.transform.ReadFromSource.Builder;
57
import feast.ingestion.values.FailedElement;
68
import feast.ingestion.values.Field;
79
import feast.types.FeatureRowProto.FeatureRow;
@@ -17,14 +19,7 @@
1719
@AutoValue
1820
public abstract class KafkaRecordToFeatureRowDoFn extends
1921
DoFn<KafkaRecord<byte[], byte[]>, FeatureRow> {
20-
21-
public abstract String getFeatureSetName();
22-
23-
public abstract int getFeatureSetVersion();
24-
25-
public abstract Map<String, Field> getFieldByName();
26-
27-
public abstract TupleTag<FeatureRow> getSuccessTag();
22+
public abstract Map<String, TupleTag<FeatureRow>> getFeatureSetTagByKey();
2823

2924
public abstract TupleTag<FailedElement> getFailureTag();
3025

@@ -35,13 +30,7 @@ public static KafkaRecordToFeatureRowDoFn.Builder newBuilder() {
3530
@AutoValue.Builder
3631
public abstract static class Builder {
3732

38-
public abstract Builder setFeatureSetName(String featureSetName);
39-
40-
public abstract Builder setFeatureSetVersion(int featureSetVersion);
41-
42-
public abstract Builder setFieldByName(Map<String, Field> fieldByName);
43-
44-
public abstract Builder setSuccessTag(TupleTag<FeatureRow> successTag);
33+
public abstract Builder setFeatureSetTagByKey(Map<String, TupleTag<FeatureRow>> featureSetTagByKey);
4534

4635
public abstract Builder setFailureTag(TupleTag<FailedElement> failureTag);
4736

@@ -67,55 +56,19 @@ public void processElement(ProcessContext context) {
6756
.build());
6857
return;
6958
}
70-
71-
// If FeatureRow contains field names that do not exist as EntitySpec
72-
// or FeatureSpec in FeatureSetSpec, mark the FeatureRow as FailedElement.
73-
String error = null;
74-
String featureSetId = String.format("%s:%d", getFeatureSetName(), getFeatureSetVersion());
75-
if (featureRow.getFeatureSet().equals(featureSetId)) {
76-
77-
for (FieldProto.Field field : featureRow.getFieldsList()) {
78-
if (!getFieldByName().containsKey(field.getName())) {
79-
error =
80-
String.format(
81-
"FeatureRow contains field '%s' which do not exists in FeatureSet '%s' version '%d'. Please check the FeatureRow data.",
82-
field.getName(), getFeatureSetName(), getFeatureSetVersion());
83-
break;
84-
}
85-
// If value is set in the FeatureRow, make sure the value type matches
86-
// that defined in FeatureSetSpec
87-
if (!field.getValue().getValCase().equals(ValCase.VAL_NOT_SET)) {
88-
int expectedTypeFieldNumber =
89-
getFieldByName().get(field.getName()).getType().getNumber();
90-
int actualTypeFieldNumber = field.getValue().getValCase().getNumber();
91-
if (expectedTypeFieldNumber != actualTypeFieldNumber) {
92-
error =
93-
String.format(
94-
"FeatureRow contains field '%s' with invalid type '%s'. Feast expects the field type to match that in FeatureSet '%s'. Please check the FeatureRow data.",
95-
field.getName(),
96-
field.getValue().getValCase(),
97-
getFieldByName().get(field.getName()).getType());
98-
break;
99-
}
100-
}
101-
}
102-
} else {
103-
error = String.format(
104-
"FeatureRow contains invalid feature set id %s. Please check that the feature rows are being published to the correct topic on the feature stream.",
105-
featureSetId);
106-
}
107-
108-
if (error != null) {
59+
TupleTag<FeatureRow> tag = getFeatureSetTagByKey()
60+
.getOrDefault(featureRow.getFeatureSet(), null);
61+
if (tag == null) {
10962
context.output(
11063
getFailureTag(),
11164
FailedElement.newBuilder()
11265
.setTransformName("KafkaRecordToFeatureRow")
11366
.setJobName(context.getPipelineOptions().getJobName())
114-
.setPayload(featureRow.toString())
115-
.setErrorMessage(error)
67+
.setPayload(new String(Base64.getEncoder().encode(value)))
68+
.setErrorMessage(String.format("Got row with unexpected feature set id %s. Expected one of %s.", featureRow.getFeatureSet(), getFeatureSetTagByKey().keySet()))
11669
.build());
117-
} else {
118-
context.output(featureRow);
70+
return;
11971
}
72+
context.output(tag, featureRow);
12073
}
12174
}

0 commit comments

Comments
 (0)