Skip to content

Commit d5a2a06

Browse files
author
zhilingc
committed
Merge remote-tracking branch 'origin/0.3-dev' into 0.3-dev
2 parents a5eb7c0 + 1085133 commit d5a2a06

25 files changed

Lines changed: 870 additions & 256 deletions

‎ingestion/pom.xml‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
<revision>0.3.0-SNAPSHOT</revision>
3030
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
3131

32-
<org.apache.beam.version>2.15.0</org.apache.beam.version>
32+
<org.apache.beam.version>2.16.0</org.apache.beam.version>
3333
<com.google.cloud.version>1.91.0</com.google.cloud.version>
3434
<grpcVersion>1.21.1</grpcVersion>
3535
<protocVersion>3.6.1</protocVersion>
@@ -93,7 +93,7 @@
9393
implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
9494
<transformer
9595
implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
96-
<mainClass>feast.ingestion.ImportJob</mainClass>
96+
<mainClass>feast.ingestion.ImportJobOld</mainClass>
9797
</transformer>
9898
<transformer
9999
implementation="org.apache.maven.plugins.shade.resource.AppendingTransformer">
@@ -162,9 +162,15 @@
162162
</dependency>
163163

164164
<dependency>
165-
<groupId>com.google.auto.service</groupId>
166-
<artifactId>auto-service</artifactId>
167-
<version>1.0-rc4</version>
165+
<groupId>com.google.auto.value</groupId>
166+
<artifactId>auto-value-annotations</artifactId>
167+
<version>1.6.6</version>
168+
</dependency>
169+
<dependency>
170+
<groupId>com.google.auto.value</groupId>
171+
<artifactId>auto-value</artifactId>
172+
<version>1.6.6</version>
173+
<scope>provided</scope>
168174
</dependency>
169175

170176
<!-- Provides FileSystemProvider for GCS. -->

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

Lines changed: 79 additions & 162 deletions
Original file line numberDiff line numberDiff line change
@@ -1,186 +1,103 @@
11
package feast.ingestion;
22

3-
import com.google.cloud.bigquery.BigQueryOptions;
4-
import com.google.common.collect.Sets;
5-
import com.google.protobuf.util.JsonFormat;
3+
import com.google.protobuf.InvalidProtocolBufferException;
64
import feast.core.FeatureSetProto.FeatureSetSpec;
7-
import feast.core.SourceProto.KafkaSourceConfig;
8-
import feast.core.SourceProto.Source;
9-
import feast.core.SourceProto.SourceType;
105
import feast.core.StoreProto.Store;
11-
import feast.ingestion.options.ImportJobPipelineOptions;
12-
import feast.ingestion.transform.FilterFeatureRow;
13-
import feast.ingestion.transform.ReadFeatureRow;
14-
import feast.ingestion.transform.ToFeatureRowExtended;
15-
import feast.ingestion.transform.WriteFeaturesTransform;
16-
import feast.ingestion.transform.metrics.WriteMetricsTransform;
17-
import feast.ingestion.util.StorageUtil;
18-
import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
19-
import java.io.IOException;
20-
import java.net.URISyntaxException;
21-
import java.util.HashMap;
6+
import feast.ingestion.options.ImportOptions;
7+
import feast.ingestion.transform.ReadFromSource;
8+
import feast.ingestion.transform.WriteFailedElementToBigQuery;
9+
import feast.ingestion.transform.WriteToStore;
10+
import feast.ingestion.utils.ResourceUtil;
11+
import feast.ingestion.utils.SpecUtil;
12+
import feast.ingestion.utils.StoreUtil;
13+
import feast.ingestion.values.FailedElement;
14+
import feast.types.FeatureRowProto.FeatureRow;
2215
import java.util.List;
23-
import java.util.Map;
24-
import java.util.Properties;
2516
import lombok.extern.slf4j.Slf4j;
2617
import org.apache.beam.sdk.Pipeline;
2718
import org.apache.beam.sdk.PipelineResult;
2819
import org.apache.beam.sdk.options.PipelineOptionsFactory;
2920
import org.apache.beam.sdk.options.PipelineOptionsValidator;
30-
import org.apache.beam.sdk.values.PCollection;
31-
import org.apache.kafka.clients.consumer.KafkaConsumer;
32-
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
33-
import org.apache.kafka.common.PartitionInfo;
34-
import org.apache.kafka.common.TopicPartition;
21+
import org.apache.beam.sdk.values.PCollectionTuple;
22+
import org.apache.beam.sdk.values.TupleTag;
3523

3624
@Slf4j
3725
public class ImportJob {
26+
// Tag for main output containing Feature Row that has been successfully processed.
27+
private static final TupleTag<FeatureRow> FEATURE_ROW_OUT = new TupleTag<FeatureRow>() {};
3828

39-
/**
40-
* Create and run a Beam pipeline with PipelineOptions passed as a list of string arguments.
41-
*
42-
* <p>The arguments will be passed to Beam {@code PipelineOptionsFactory} to create {@code
43-
* ImportJobPipelineOptions}.
44-
*
45-
* <p>The returned PipelineResult object can be used to check the state of the pipeline e.g. if
46-
* it is running, done or cancelled.
47-
*
48-
* @param args command line arguments, typically come from the main() method
49-
* @return PipelineResult
50-
* @throws IOException if importJobSpecsUri specified in args is not accessible
51-
*/
52-
public static PipelineResult runPipeline(String[] args) throws IOException, URISyntaxException {
53-
ImportJobPipelineOptions pipelineOptions =
54-
PipelineOptionsFactory.fromArgs(args).withValidation().as(ImportJobPipelineOptions.class);
55-
return runPipeline(pipelineOptions);
56-
}
57-
58-
/**
59-
* Create and run a Beam pipeline from {@code ImportJobPipelineOptions}.
60-
*
61-
* <p>The returned PipelineResult object can be used to check the state of the pipeline e.g. if
62-
* it is running, done or cancelled.
63-
*
64-
* @param pipelineOptions configuration for the pipeline
65-
* @return PipelineResult
66-
* @throws IOException if importJobSpecsUri is not accessible
67-
*/
68-
public static PipelineResult runPipeline(ImportJobPipelineOptions pipelineOptions)
69-
throws IOException {
70-
pipelineOptions =
71-
PipelineOptionsValidator.validate(ImportJobPipelineOptions.class, pipelineOptions);
72-
Pipeline pipeline = Pipeline.create(pipelineOptions);
73-
74-
for (String storeJson : pipelineOptions.getStoreJson()) {
75-
Store.Builder storeBuilder = Store.newBuilder();
76-
JsonFormat.parser().merge(storeJson, storeBuilder);
77-
Store store = storeBuilder.build();
78-
79-
for (String featureSetSpecJson : pipelineOptions.getFeatureSetSpecJson()) {
80-
FeatureSetSpec.Builder featureSetSpecBuilder = FeatureSetSpec.newBuilder();
81-
JsonFormat.parser().merge(featureSetSpecJson, featureSetSpecBuilder);
82-
FeatureSetSpec featureSetSpec = featureSetSpecBuilder.build();
83-
84-
setupSource(pipelineOptions.getJobName(), featureSetSpec.getSource());
85-
setupStore(store, featureSetSpec);
86-
87-
PCollection<FeatureRowExtended> featureRows = pipeline
88-
.apply("Read FeatureRow", new ReadFeatureRow(featureSetSpec))
89-
.apply("Filter FeatureRow", new FilterFeatureRow(featureSetSpec))
90-
.apply("Create FeatureRowExtended from FeatureRow", new ToFeatureRowExtended());
91-
92-
featureRows
93-
.apply("Write metrics", new WriteMetricsTransform(store.getName(), featureSetSpec));
94-
featureRows
95-
.apply("Write FeatureRowExtended", new WriteFeaturesTransform(store, featureSetSpec));
96-
}
97-
}
98-
99-
return pipeline.run();
100-
}
29+
// Tag for deadletter output containing elements and error messages from invalid input/transform.
30+
private static final TupleTag<FailedElement> DEADLETTER_OUT = new TupleTag<FailedElement>() {};
10131

10232
/**
103-
* Configures the storage for Feast.
104-
*
105-
* <p>This method ensures that the storage backend is running and accessible by the import job,
106-
* and it also ensures that the storage backend has the necessary schema and configuration so the
107-
* import job can start writing Feature Row.
108-
*
109-
* <p>For example, when using BigQuery as the storage backend, this method ensures that, given a
110-
* list of features, the corresponding BigQuery dataset and table are created.
111-
*
112-
* @param store Store specification, refer to {@code feast.core.Store.proto}
33+
* @param args arguments to be passed to Beam pipeline
34+
* @throws InvalidProtocolBufferException if options passed to the pipeline are invalid
11335
*/
114-
private static void setupStore(Store store, FeatureSetSpec featureSetSpec) {
115-
switch (store.getType()) {
116-
case REDIS:
117-
StorageUtil.checkRedisConnection(store.getRedisConfig());
118-
break;
119-
case BIGQUERY:
120-
StorageUtil.setupBigQuery(
121-
featureSetSpec,
122-
store.getBigqueryConfig().getProjectId(),
123-
store.getBigqueryConfig().getDatasetId(),
124-
BigQueryOptions.getDefaultInstance().getService());
125-
break;
126-
default:
127-
throw new UnsupportedOperationException(
128-
String.format("Store type: %s not implemented yet", store.getType()));
129-
}
36+
public static void main(String[] args) throws InvalidProtocolBufferException {
37+
ImportOptions options =
38+
PipelineOptionsFactory.fromArgs(args).withValidation().create().as(ImportOptions.class);
39+
runPipeline(options);
13040
}
13141

132-
/**
133-
* TODO: Update documentation
134-
*
135-
* <p>Manually sets the consumer group offset for this job's consumer group to the offset at the
136-
* time at which we provision the ingestion job.
137-
*
138-
* <p>This is necessary because the setup time for certain runners (e.g. Dataflow) might cause
139-
* the worker to miss the messages that were emitted into the stream prior to the workers being
140-
* ready.
141-
*/
142-
private static void setupSource(String jobName, Source source) {
143-
if (!source.getType().equals(SourceType.KAFKA)) {
144-
throw new UnsupportedOperationException(
145-
String.format("Source type: %s not implemented yet", source.getType()));
146-
}
147-
148-
KafkaSourceConfig kafkaSourceConfig = source.getKafkaSourceConfig();
149-
150-
Properties consumerProperties = new Properties();
151-
consumerProperties.setProperty("group.id", jobName);
152-
consumerProperties.setProperty("bootstrap.servers", kafkaSourceConfig.getBootstrapServers());
153-
consumerProperties.setProperty(
154-
"key.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");
155-
consumerProperties.setProperty(
156-
"value.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");
157-
KafkaConsumer kafkaConsumer = new KafkaConsumer(consumerProperties);
158-
159-
String[] topics = {kafkaSourceConfig.getTopic()};
160-
long timestamp = System.currentTimeMillis();
161-
Map<TopicPartition, Long> timestampsToSearch = new HashMap<>();
162-
for (String topic : topics) {
163-
List<PartitionInfo> partitionInfos = kafkaConsumer.partitionsFor(topic);
164-
for (PartitionInfo partitionInfo : partitionInfos) {
165-
TopicPartition topicPartition = new TopicPartition(topic, partitionInfo.partition());
166-
timestampsToSearch.put(topicPartition, timestamp);
42+
@SuppressWarnings("UnusedReturnValue")
43+
public static PipelineResult runPipeline(ImportOptions options)
44+
throws InvalidProtocolBufferException {
45+
/*
46+
* Steps:
47+
* 1. Read messages from Feast Source as FeatureRow
48+
* 2. Write FeatureRow to the corresponding Store
49+
* 3. Write elements that failed to be processed to a dead letter queue.
50+
*/
51+
52+
PipelineOptionsValidator.validate(ImportOptions.class, options);
53+
Pipeline pipeline = Pipeline.create(options);
54+
55+
List<FeatureSetSpec> featureSetSpecs =
56+
SpecUtil.parseFeatureSetSpecJsonList(options.getFeatureSetSpecJson());
57+
List<Store> stores = SpecUtil.parseStoreJsonList(options.getStoreJson());
58+
59+
for (Store store : stores) {
60+
List<FeatureSetSpec> subscribedFeatureSets =
61+
SpecUtil.getSubscribedFeatureSets(store.getSubscriptionsList(), featureSetSpecs);
62+
63+
for (FeatureSetSpec featureSet : subscribedFeatureSets) {
64+
// Ensure Store has valid configuration and Feast can access it.
65+
StoreUtil.setupStore(store, featureSet);
66+
67+
// Step 1. Read messages from Feast Source as FeatureRow.
68+
PCollectionTuple convertedFeatureRows =
69+
pipeline.apply(
70+
"ReadFeatureRowFromSource",
71+
ReadFromSource.newBuilder()
72+
.setSource(featureSet.getSource())
73+
.setFieldByName(SpecUtil.getFieldByName(featureSet))
74+
.setFeatureSetName(featureSet.getName())
75+
.setFeatureSetVersion(featureSet.getVersion())
76+
.setSuccessTag(FEATURE_ROW_OUT)
77+
.setFailureTag(DEADLETTER_OUT)
78+
.build());
79+
80+
// Step 2. Write FeatureRow to the corresponding Store.
81+
convertedFeatureRows
82+
.get(FEATURE_ROW_OUT)
83+
.apply(
84+
"WriteFeatureRowToStore",
85+
WriteToStore.newBuilder().setFeatureSetSpec(featureSet).setStore(store).build());
86+
87+
// Step 3. Write FailedElements to a dead letter table in BigQuery.
88+
if (options.getDeadLetterTableSpec() != null) {
89+
convertedFeatureRows
90+
.get(DEADLETTER_OUT)
91+
.apply(
92+
"WriteFailedElements",
93+
WriteFailedElementToBigQuery.newBuilder()
94+
.setJsonSchema(ResourceUtil.getDeadletterTableSchemaJson())
95+
.setTableSpec(options.getDeadLetterTableSpec())
96+
.build());
97+
}
16798
}
16899
}
169-
Map<TopicPartition, OffsetAndTimestamp> offsets =
170-
kafkaConsumer.offsetsForTimes(timestampsToSearch);
171100

172-
kafkaConsumer.assign(offsets.keySet());
173-
kafkaConsumer.poll(1000);
174-
kafkaConsumer.commitSync();
175-
176-
offsets.forEach(
177-
(topicPartition, offset) -> {
178-
if (offset != null) {
179-
kafkaConsumer.seek(topicPartition, offset.offset());
180-
} else {
181-
kafkaConsumer.seekToBeginning(Sets.newHashSet(topicPartition));
182-
}
183-
});
101+
return pipeline.run();
184102
}
185-
186103
}

ingestion/src/main/java/feast/ingestion/options/ImportJobPipelineOptions.java renamed to ingestion/src/main/java/feast/ingestion/options/ImportOptions.java

Lines changed: 42 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,24 +17,17 @@
1717

1818
package feast.ingestion.options;
1919

20-
import com.google.auto.service.AutoService;
21-
import java.util.Collections;
2220
import java.util.List;
2321
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
2422
import org.apache.beam.runners.direct.DirectOptions;
2523
import org.apache.beam.sdk.options.Default;
2624
import org.apache.beam.sdk.options.Default.Boolean;
2725
import org.apache.beam.sdk.options.Description;
2826
import org.apache.beam.sdk.options.PipelineOptions;
29-
import org.apache.beam.sdk.options.PipelineOptionsRegistrar;
3027
import org.apache.beam.sdk.options.Validation.Required;
3128

32-
/**
33-
* Options passed to Beam to influence the job's execution environment
34-
*/
35-
public interface ImportJobPipelineOptions
36-
extends PipelineOptions, DataflowPipelineOptions, DirectOptions {
37-
29+
/** Options passed to Beam to influence the job's execution environment */
30+
public interface ImportOptions extends PipelineOptions, DataflowPipelineOptions, DirectOptions {
3831
@Required
3932
@Description(
4033
"JSON string representation of the FeatureSetSpec that the import job will process."
@@ -62,7 +55,21 @@ public interface ImportJobPipelineOptions
6255
void setStoreJson(List<String> storeJson);
6356

6457
@Description(
65-
"MetricsAccumulator exporter type to instantiate." //TODO: expound
58+
"(Optional) Deadletter elements will be written to this BigQuery table."
59+
+ "Table spec must follow this format PROJECT_ID:DATASET_ID.PROJECT_ID"
60+
+ "The table will be created if not exists.")
61+
String getDeadLetterTableSpec();
62+
63+
/**
64+
* @param deadLetterTableSpec (Optional) BigQuery table for storing elements that failed to be
65+
* processed. Table spec must follow this format
66+
* PROJECT_ID:DATASET_ID.PROJECT_ID
67+
*/
68+
void setDeadLetterTableSpec(String deadLetterTableSpec);
69+
70+
// TODO: expound
71+
@Description(
72+
"MetricsAccumulator exporter type to instantiate."
6673
)
6774
@Default.String("none")
6875
String getMetricsExporterType();
@@ -77,18 +84,34 @@ public interface ImportJobPipelineOptions
7784

7885
void setPrometheusExporterAddress(String prometheusExporterAddress);
7986

87+
@Description("Limit of rows to sample and output for debugging")
88+
@Default.Integer(0)
89+
int getSampleLimit();
90+
91+
void setSampleLimit(int value);
92+
93+
@Description(
94+
"Enable coalesce rows, merges feature rows within a time window to output only the latest value")
95+
@Default.Boolean(false)
96+
boolean isCoalesceRowsEnabled();
97+
98+
void setCoalesceRowsEnabled(boolean value);
99+
100+
@Description("Delay in seconds to wait for newer values to coalesce on key before emitting")
101+
@Default.Integer(10)
102+
int getCoalesceRowsDelaySeconds();
103+
104+
void setCoalesceRowsDelaySeconds(int value);
105+
106+
@Description("Time in seconds to retain feature rows to merge with newer records")
107+
@Default.Integer(30)
108+
int getCoalesceRowsTimeoutSeconds();
109+
110+
void setCoalesceRowsTimeoutSeconds(int value);
111+
80112
@Description("If dry run is set, execute up to feature row validation")
81113
@Default.Boolean(false)
82114
Boolean isDryRun();
83115

84116
void setDryRun(Boolean value);
85-
86-
@AutoService(PipelineOptionsRegistrar.class)
87-
class ImportJobPipelineOptionsRegistrar implements PipelineOptionsRegistrar {
88-
89-
@Override
90-
public Iterable<Class<? extends PipelineOptions>> getPipelineOptions() {
91-
return Collections.singleton(ImportJobPipelineOptions.class);
92-
}
93-
}
94117
}

0 commit comments

Comments
 (0)