Skip to content

Commit aaa39aa

Browse files
pyalexpyalex
committed
verify column types match feature types
Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent 6b59ec0 commit aaa39aa

4 files changed

Lines changed: 112 additions & 3 deletions

File tree

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

Lines changed: 7 additions & 1 deletion
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

@@ -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 Left(error) =>
61+
throw new RuntimeException(s"Dataframes columns don't match expected feature types: $error")
62+
case _ => ()
63+
}
64+
5965
val validRows = projected
6066
.filter(validator.checkAll)
6167

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,18 @@ import feast.ingestion.registry.proto.ProtoRegistryFactory
2020
import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}
2121
import org.apache.spark.sql.functions.udf
2222
import feast.ingestion.utils.ProtoReflection
23-
import feast.ingestion.validation.RowValidator
23+
import feast.ingestion.validation.{RowValidator, TypeCheck}
2424
import org.apache.spark.sql.streaming.StreamingQuery
2525

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+
*/
2635
object StreamingPipeline extends BasePipeline with Serializable {
2736
override def createPipeline(
2837
sparkSession: SparkSession,
@@ -52,6 +61,12 @@ object StreamingPipeline extends BasePipeline with Serializable {
5261
.select("features.*")
5362
.select(projection: _*)
5463

64+
TypeCheck.allTypesMatch(projected.schema, featureTable) match {
65+
case Left(error) =>
66+
throw new RuntimeException(s"Dataframes columns don't match expected feature types: $error")
67+
case _ => ()
68+
}
69+
5570
val query = projected.writeStream
5671
.foreachBatch { (batchDF: DataFrame, batchID: Long) =>
5772
batchDF.persist()

‎spark/ingestion/src/main/scala/feast/ingestion/validation/TypeCheck.scala‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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.validation
18+
19+
import feast.ingestion.FeatureTable
20+
import feast.proto.types.ValueProto.ValueType
21+
import org.apache.spark.sql.types._
22+
23+
object TypeCheck {
24+
def typesMatch(defType: ValueType.Enum, columnType: DataType): Boolean =
25+
(defType, columnType) match {
26+
case (ValueType.Enum.BOOL, BooleanType) => true
27+
case (ValueType.Enum.INT32, IntegerType) => true
28+
case (ValueType.Enum.INT64, LongType) => true
29+
case (ValueType.Enum.FLOAT, FloatType) => true
30+
case (ValueType.Enum.DOUBLE, DoubleType) => true
31+
case (ValueType.Enum.STRING, StringType) => true
32+
case (ValueType.Enum.BYTES, BinaryType) => true
33+
case (ValueType.Enum.BOOL_LIST, ArrayType(_: BooleanType, _)) => true
34+
case (ValueType.Enum.INT32_LIST, ArrayType(_: IntegerType, _)) => true
35+
case (ValueType.Enum.INT64_LIST, ArrayType(_: LongType, _)) => true
36+
case (ValueType.Enum.FLOAT_LIST, ArrayType(_: FloatType, _)) => true
37+
case (ValueType.Enum.DOUBLE_LIST, ArrayType(_: DoubleType, _)) => true
38+
case (ValueType.Enum.STRING_LIST, ArrayType(_: StringType, _)) => true
39+
case (ValueType.Enum.BYTES_LIST, ArrayType(_: BinaryType, _)) => true
40+
case _ => false
41+
}
42+
43+
/**
44+
* Verify whether types declared in FeatureTable match correspondent columns
45+
* @param schema Spark's dataframe columns
46+
* @param featureTable definition of expected schema
47+
* @return true if all match
48+
*/
49+
def allTypesMatch(schema: StructType, featureTable: FeatureTable): Either[String, Boolean] = {
50+
val typeByField =
51+
(featureTable.entities ++ featureTable.features).map(f => f.name -> f.`type`).toMap
52+
53+
schema.fields
54+
.map(f =>
55+
if (typeByField.contains(f.name) && !typesMatch(typeByField(f.name), f.dataType)) {
56+
Left(
57+
s"Feature ${f.name} has different type ${f.dataType} instead of expected ${typeByField(f.name)}"
58+
)
59+
} else Right(true)
60+
)
61+
.foldLeft[Either[String, Boolean]](Right(true)) {
62+
case (Left(error), _) => Left(error)
63+
case (Right(_), Left(error)) => Left(error)
64+
case _ => Right(true)
65+
}
66+
}
67+
}

‎spark/ingestion/src/test/scala/feast/ingestion/StreamingPipelineIT.scala‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ class StreamingPipelineIT extends SparkSpec with ForAllTestContainer {
6969
name = "driver-fs",
7070
project = "default",
7171
entities = Seq(
72-
Field("s2_id", ValueType.Enum.INT32),
72+
Field("s2_id", ValueType.Enum.INT64),
7373
Field("vehicle_type", ValueType.Enum.STRING)
7474
),
7575
features = Seq(
@@ -256,4 +256,25 @@ class StreamingPipelineIT extends SparkSpec with ForAllTestContainer {
256256
)
257257
)
258258
}
259+
260+
"Expected feature types" should "match source types" in new Scope {
261+
val configWithKafka = config.copy(
262+
source = kafkaSource,
263+
featureTable = FeatureTable(
264+
name = "driver-fs",
265+
project = "default",
266+
entities = Seq(
267+
Field("s2_id", ValueType.Enum.STRING),
268+
Field("vehicle_type", ValueType.Enum.INT32)
269+
),
270+
features = Seq(
271+
Field("unique_drivers", ValueType.Enum.FLOAT)
272+
)
273+
)
274+
)
275+
276+
assertThrows[RuntimeException] {
277+
StreamingPipeline.createPipeline(sparkSession, configWithKafka).get
278+
}
279+
}
259280
}

0 commit comments

Comments
 (0)