Skip to content

Commit 2cd019c

Browse files
pyalexpyalex
authored
Streaming Ingestion Pipeline with Spark (#1027)
* simple streaming pipeline Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * complete streaming pipeline Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * clean up Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * skip surefire Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * all types test Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * verify column types match feature types Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * spelling Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * stencil test Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * refactor type checker Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent 42bf26a commit 2cd019c

23 files changed

Lines changed: 1226 additions & 152 deletions

‎spark/ingestion/pom.xml‎

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,9 @@
3636
<spark.version>2.4.7</spark.version>
3737
<scala-maven-plugin.version>4.4.0</scala-maven-plugin.version>
3838
<maven-assembly-plugin.version>3.3.0</maven-assembly-plugin.version>
39-
<project.version>0.7-SNAPSHOT</project.version>
39+
<protobuf.version>3.12.2</protobuf.version>
4040
</properties>
4141

42-
4342
<dependencies>
4443
<dependency>
4544
<groupId>dev.feast</groupId>
@@ -56,7 +55,13 @@
5655
<dependency>
5756
<groupId>com.google.protobuf</groupId>
5857
<artifactId>protobuf-java</artifactId>
59-
<version>3.12.2</version>
58+
<version>${protobuf.version}</version>
59+
</dependency>
60+
61+
<dependency>
62+
<groupId>com.gojek</groupId>
63+
<artifactId>stencil</artifactId>
64+
<version>4.2.0</version>
6065
</dependency>
6166

6267
<dependency>
@@ -102,7 +107,12 @@
102107
<groupId>org.apache.spark</groupId>
103108
<artifactId>spark-sql-kafka-0-10_${scala.version}</artifactId>
104109
<version>${spark.version}</version>
105-
<scope>provided</scope>
110+
</dependency>
111+
112+
<dependency>
113+
<groupId>org.apache.kafka</groupId>
114+
<artifactId>kafka-clients</artifactId>
115+
<version>2.5.1</version>
106116
</dependency>
107117

108118
<dependency>
@@ -169,12 +179,30 @@
169179
<scope>test</scope>
170180
</dependency>
171181

182+
<dependency>
183+
<groupId>com.dimafeng</groupId>
184+
<artifactId>testcontainers-scala-kafka_${scala.version}</artifactId>
185+
<version>0.38.3</version>
186+
<scope>test</scope>
187+
</dependency>
188+
189+
<dependency>
190+
<groupId>com.github.tomakehurst</groupId>
191+
<artifactId>wiremock-jre8</artifactId>
192+
<version>2.27.2</version>
193+
<scope>test</scope>
194+
</dependency>
172195
</dependencies>
173196

174197

175198
<build>
176199
<sourceDirectory>src/main/scala</sourceDirectory>
177200
<testSourceDirectory>src/test/scala</testSourceDirectory>
201+
<testResources>
202+
<testResource>
203+
<directory>src/test/resources</directory>
204+
</testResource>
205+
</testResources>
178206
<plugins>
179207
<plugin>
180208
<groupId>net.alchim31.maven</groupId>
@@ -280,6 +308,46 @@
280308
</execution>
281309
</executions>
282310
</plugin>
311+
<plugin>
312+
<groupId>com.github.os72</groupId>
313+
<artifactId>protoc-jar-maven-plugin</artifactId>
314+
<version>3.11.4</version>
315+
<executions>
316+
<execution>
317+
<phase>generate-test-sources</phase>
318+
<goals>
319+
<goal>run</goal>
320+
</goals>
321+
<configuration>
322+
<protocArtifact>com.google.protobuf:protoc:${protobuf.version}</protocArtifact>
323+
<protocVersion>${protobuf.version}</protocVersion>
324+
<inputDirectories>
325+
<include>src/test/proto</include>
326+
</inputDirectories>
327+
<includeStdTypes>true</includeStdTypes>
328+
<outputTargets>
329+
<outputTarget>
330+
<type>java</type>
331+
<addSources>none</addSources>
332+
<outputDirectory>src/test/scala</outputDirectory>
333+
</outputTarget>
334+
<outputTarget>
335+
<type>descriptor</type>
336+
<addSources>none</addSources>
337+
<outputDirectory>src/test/resources/stencil/__files</outputDirectory>
338+
</outputTarget>
339+
</outputTargets>
340+
</configuration>
341+
</execution>
342+
</executions>
343+
</plugin>
344+
<plugin>
345+
<groupId>org.apache.maven.plugins</groupId>
346+
<artifactId>maven-surefire-plugin</artifactId>
347+
<configuration>
348+
<skipTests>true</skipTests>
349+
</configuration>
350+
</plugin>
283351
</plugins>
284352
<pluginManagement>
285353
<plugins>

‎spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@
1717
package feast.ingestion
1818

1919
import org.apache.spark.SparkConf
20-
import org.apache.spark.sql.SparkSession
20+
import org.apache.spark.sql.{Column, SparkSession}
21+
import org.apache.spark.sql.functions.col
22+
import org.apache.spark.sql.streaming.StreamingQuery
2123

2224
trait BasePipeline {
2325
def createSparkSession(jobConfig: IngestionJobConfig): SparkSession = {
@@ -66,5 +68,28 @@ trait BasePipeline {
6668
.getOrCreate()
6769
}
6870

69-
def createPipeline(sparkSession: SparkSession, config: IngestionJobConfig): Unit
71+
def createPipeline(sparkSession: SparkSession, config: IngestionJobConfig): Option[StreamingQuery]
72+
73+
/**
74+
* Build column projection using custom mapping with fallback to feature|entity names.
75+
*/
76+
def inputProjection(
77+
source: Source,
78+
features: Seq[Field],
79+
entities: Seq[Field]
80+
): Array[Column] = {
81+
val featureColumns = features
82+
.filter(f => !source.mapping.contains(f.name))
83+
.map(f => (f.name, f.name)) ++ source.mapping
84+
85+
val timestampColumn = Seq((source.timestampColumn, source.timestampColumn))
86+
val entitiesColumns =
87+
entities
88+
.filter(e => !source.mapping.contains(e.name))
89+
.map(e => (e.name, e.name))
90+
91+
(featureColumns ++ entitiesColumns ++ timestampColumn).map { case (alias, source) =>
92+
col(source).alias(alias)
93+
}.toArray
94+
}
7095
}

‎spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala‎

Lines changed: 10 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ package feast.ingestion
1818

1919
import feast.ingestion.sources.bq.BigQueryReader
2020
import feast.ingestion.sources.file.FileReader
21-
import feast.ingestion.validation.RowValidator
21+
import feast.ingestion.validation.{RowValidator, TypeCheck}
2222
import org.apache.spark.sql.{Column, SparkSession}
2323
import org.apache.spark.sql.functions.col
2424

@@ -31,11 +31,11 @@ import org.apache.spark.sql.functions.col
3131
* 5. Store invalid rows in parquet format at `deadletter` destination
3232
*/
3333
object BatchPipeline extends BasePipeline {
34-
override def createPipeline(sparkSession: SparkSession, config: IngestionJobConfig): Unit = {
34+
override def createPipeline(sparkSession: SparkSession, config: IngestionJobConfig) = {
3535
val featureTable = config.featureTable
3636
val projection =
3737
inputProjection(config.source, featureTable.features, featureTable.entities)
38-
val validator = new RowValidator(featureTable)
38+
val validator = new RowValidator(featureTable, config.source.timestampColumn)
3939

4040
val input = config.source match {
4141
case source: BQSource =>
@@ -56,6 +56,12 @@ object BatchPipeline extends BasePipeline {
5656

5757
val projected = input.select(projection: _*).cache()
5858

59+
TypeCheck.allTypesMatch(projected.schema, featureTable) match {
60+
case Some(error) =>
61+
throw new RuntimeException(s"Dataframe columns don't match expected feature types: $error")
62+
case _ => ()
63+
}
64+
5965
val validRows = projected
6066
.filter(validator.checkAll)
6167

@@ -77,28 +83,6 @@ object BatchPipeline extends BasePipeline {
7783
case _ => None
7884
}
7985

80-
}
81-
82-
/**
83-
* Build column projection using custom mapping with fallback to feature|entity names.
84-
*/
85-
private def inputProjection(
86-
source: Source,
87-
features: Seq[Field],
88-
entities: Seq[Field]
89-
): Array[Column] = {
90-
val featureColumns = features
91-
.filter(f => !source.mapping.contains(f.name))
92-
.map(f => (f.name, f.name)) ++ source.mapping
93-
94-
val timestampColumn = Seq((source.timestampColumn, source.timestampColumn))
95-
val entitiesColumns =
96-
entities
97-
.filter(e => !source.mapping.contains(e.name))
98-
.map(e => (e.name, e.name))
99-
100-
(featureColumns ++ entitiesColumns ++ timestampColumn).map { case (alias, source) =>
101-
col(source).alias(alias)
102-
}.toArray
86+
None
10387
}
10488
}

‎spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,9 @@ object IngestionJob {
7979
case Modes.Offline =>
8080
val sparkSession = BatchPipeline.createSparkSession(config)
8181
BatchPipeline.createPipeline(sparkSession, config)
82+
case Modes.Online =>
83+
val sparkSession = BatchPipeline.createSparkSession(config)
84+
StreamingPipeline.createPipeline(sparkSession, config).get.awaitTermination
8285
}
8386
case None =>
8487
println("Parameters can't be parsed")

‎spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ abstract class Source {
4040

4141
abstract class BatchSource extends Source
4242

43-
abstract class StreamingSource extends Source
43+
abstract class StreamingSource extends Source {
44+
def classpath: String
45+
}
4446

4547
case class FileSource(
4648
path: String,
@@ -59,6 +61,7 @@ case class BQSource(
5961
case class KafkaSource(
6062
bootstrapServers: String,
6163
topic: String,
64+
override val classpath: String,
6265
override val mapping: Map[String, String],
6366
override val timestampColumn: String
6467
) extends StreamingSource

‎spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.ingestion
18+
19+
import feast.ingestion.registry.proto.ProtoRegistryFactory
20+
import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}
21+
import org.apache.spark.sql.functions.udf
22+
import feast.ingestion.utils.ProtoReflection
23+
import feast.ingestion.validation.{RowValidator, TypeCheck}
24+
import org.apache.spark.sql.streaming.StreamingQuery
25+
26+
/**
27+
* Streaming pipeline (currently in micro-batches mode only, since we need to have multiple sinks: redis & deadletters).
28+
* Flow:
29+
* 1. Read from streaming source (currently only Kafka)
30+
* 2. Parse bytes from streaming source into Row with schema inferenced from provided class (Protobuf)
31+
* 3. Map columns according to provided mapping rules
32+
* 4. Validate
33+
* 5. (In batches) store to redis valid rows / write to deadletter (parquet) invalid
34+
*/
35+
object StreamingPipeline extends BasePipeline with Serializable {
36+
override def createPipeline(
37+
sparkSession: SparkSession,
38+
config: IngestionJobConfig
39+
): Option[StreamingQuery] = {
40+
import sparkSession.implicits._
41+
42+
val featureTable = config.featureTable
43+
val projection =
44+
inputProjection(config.source, featureTable.features, featureTable.entities)
45+
val validator = new RowValidator(featureTable, config.source.timestampColumn)
46+
47+
val messageParser =
48+
protoParser(sparkSession, config.source.asInstanceOf[StreamingSource].classpath)
49+
50+
val input = config.source match {
51+
case source: KafkaSource =>
52+
sparkSession.readStream
53+
.format("kafka")
54+
.option("kafka.bootstrap.servers", source.bootstrapServers)
55+
.option("subscribe", source.topic)
56+
.load()
57+
}
58+
59+
val projected = input
60+
.withColumn("features", messageParser($"value"))
61+
.select("features.*")
62+
.select(projection: _*)
63+
64+
TypeCheck.allTypesMatch(projected.schema, featureTable) match {
65+
case Some(error) =>
66+
throw new RuntimeException(s"Dataframe columns don't match expected feature types: $error")
67+
case _ => ()
68+
}
69+
70+
val query = projected.writeStream
71+
.foreachBatch { (batchDF: DataFrame, batchID: Long) =>
72+
batchDF.persist()
73+
74+
val validRows = batchDF
75+
.filter(validator.checkAll)
76+
77+
validRows.write
78+
.format("feast.ingestion.stores.redis")
79+
.option("entity_columns", featureTable.entities.map(_.name).mkString(","))
80+
.option("namespace", featureTable.name)
81+
.option("project_name", featureTable.project)
82+
.option("timestamp_column", config.source.timestampColumn)
83+
.save()
84+
85+
config.deadLetterPath match {
86+
case Some(path) =>
87+
batchDF
88+
.filter(!validator.checkAll)
89+
.write
90+
.format("parquet")
91+
.mode(SaveMode.Append)
92+
.save(path)
93+
case _ =>
94+
batchDF
95+
.filter(!validator.checkAll)
96+
.foreach(r => {
97+
println(s"Row failed validation $r")
98+
})
99+
}
100+
101+
batchDF.unpersist()
102+
() // return Unit to avoid compile error with overloaded foreachBatch
103+
}
104+
.start()
105+
106+
Some(query)
107+
}
108+
109+
private def protoParser(sparkSession: SparkSession, className: String) = {
110+
val protoRegistry = ProtoRegistryFactory.resolveProtoRegistry(sparkSession)
111+
112+
val parser: Array[Byte] => Row = ProtoReflection.createMessageParser(protoRegistry, className)
113+
114+
udf(parser, ProtoReflection.inferSchema(protoRegistry.getProtoDescriptor(className)))
115+
}
116+
}

0 commit comments

Comments
 (0)