complete streaming pipeline

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
feast-dev · feast-ci-bot · Oct 13, 2020 · Sep 16, 2020 · Oct 7, 2020 · Oct 7, 2020
commit 45b3600c2c1e56e2319c86abd50fc67aade68043
diff --git a/spark/ingestion/pom.xml b/spark/ingestion/pom.xml
@@ -36,7 +36,6 @@
         <spark.version>2.4.7</spark.version>
         <scala-maven-plugin.version>4.4.0</scala-maven-plugin.version>
         <maven-assembly-plugin.version>3.3.0</maven-assembly-plugin.version>
-        <project.version>0.7-SNAPSHOT</project.version>
     </properties>
 
     <dependencies>

diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala
@@ -17,7 +17,8 @@
 package feast.ingestion
 
 import org.apache.spark.SparkConf
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{Column, SparkSession}
+import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.streaming.StreamingQuery
 
 trait BasePipeline {
@@ -67,4 +68,27 @@ trait BasePipeline {
   }
 
   def createPipeline(sparkSession: SparkSession, config: IngestionJobConfig): Option[StreamingQuery]
+
+  /**
+    * Build column projection using custom mapping with fallback to feature|entity names.
+    */
+  def inputProjection(
+      source: Source,
+      features: Seq[Field],
+      entities: Seq[Field]
+  ): Array[Column] = {
+    val featureColumns = features
+      .filter(f => !source.mapping.contains(f.name))
+      .map(f => (f.name, f.name)) ++ source.mapping
+
+    val timestampColumn = Seq((source.timestampColumn, source.timestampColumn))
+    val entitiesColumns =
+      entities
+        .filter(e => !source.mapping.contains(e.name))
+        .map(e => (e.name, e.name))
+
+    (featureColumns ++ entitiesColumns ++ timestampColumn).map { case (alias, source) =>
+      col(source).alias(alias)
+    }.toArray
+  }
 }
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala
@@ -79,27 +79,4 @@ object BatchPipeline extends BasePipeline {
 
     None
   }
-
-  /**
-    * Build column projection using custom mapping with fallback to feature|entity names.
-    */
-  private def inputProjection(
-      source: Source,
-      features: Seq[Field],
-      entities: Seq[Field]
-  ): Array[Column] = {
-    val featureColumns = features
-      .filter(f => !source.mapping.contains(f.name))
-      .map(f => (f.name, f.name)) ++ source.mapping
-
-    val timestampColumn = Seq((source.timestampColumn, source.timestampColumn))
-    val entitiesColumns =
-      entities
-        .filter(e => !source.mapping.contains(e.name))
-        .map(e => (e.name, e.name))
-
-    (featureColumns ++ entitiesColumns ++ timestampColumn).map { case (alias, source) =>
-      col(source).alias(alias)
-    }.toArray
-  }
 }
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala
@@ -81,7 +81,7 @@ object IngestionJob {
             BatchPipeline.createPipeline(sparkSession, config)
           case Modes.Online =>
             val sparkSession = BatchPipeline.createSparkSession(config)
-            StreamingPipeline.createPipeline(sparkSession, config)
+            StreamingPipeline.createPipeline(sparkSession, config).get.awaitTermination
         }
       case None =>
         println("Parameters can't be parsed")

diff --git a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala
@@ -20,10 +20,13 @@ import java.sql
 
 import com.google.protobuf.Descriptors.{Descriptor, EnumValueDescriptor, FieldDescriptor}
 import com.google.protobuf.{AbstractMessage, ByteString, GeneratedMessageV3, Parser, Timestamp}
-import org.apache.spark.sql.{Row, SparkSession}
+import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}
 import org.apache.spark.sql.functions.udf
 import org.apache.spark.sql.types._
 import com.google.protobuf.Descriptors.FieldDescriptor.JavaType._
+import feast.ingestion.BatchPipeline.inputProjection
+import feast.ingestion.utils.ProtoReflection
+import feast.ingestion.validation.RowValidator
 import org.apache.spark.sql.streaming.StreamingQuery
 
 import scala.collection.convert.ImplicitConversions._
@@ -37,111 +40,82 @@ object StreamingPipeline extends BasePipeline with Serializable {
     import sparkSession.implicits._
 
     val featureTable = config.featureTable
+    val projection =
+      inputProjection(config.source, featureTable.features, featureTable.entities)
+    val validator = new RowValidator(featureTable)
+
+    val defaultInstance = defaultInstanceFromProtoClass(
+      config.source.asInstanceOf[StreamingSource].classpath
+    )
+    val protoParser = udf(
+      ProtoReflection.createMessageParser(defaultInstance),
+      inferSchemaFromProto(defaultInstance)
+    )
 
     val input = config.source match {
       case source: KafkaSource =>
         sparkSession.readStream
           .format("kafka")
           .option("kafka.bootstrap.servers", source.bootstrapServers)
           .option("subscribe", source.topic)
-          .option("startingOffsets", "earliest")
           .load()
     }
 
-    val klass = loadClass(config.source.asInstanceOf[StreamingSource].classpath)
-      .asInstanceOf[Class[GeneratedMessageV3]]
-
-    val defaultInstance =
-      klass.getMethod("getDefaultInstance").invoke(null).asInstanceOf[GeneratedMessageV3]
-
-    val schema = StructType(defaultInstance.getDescriptorForType.getFields.flatMap(structFieldFor))
-    print(schema)
-
-    val u = udf(createMessageParser(defaultInstance), schema)
-    val o = input.withColumn("content", u($"value")).select("content.*")
+    val projected = input
+      .withColumn("features", protoParser($"value"))
+      .select("features.*")
+      .select(projection: _*)
+
+    val query = projected.writeStream
+      .foreachBatch { (batchDF: DataFrame, batchID: Long) =>
+        batchDF.persist()
+
+        val validRows = batchDF
+          .filter(validator.checkAll)
+
+        validRows.write
+          .format("feast.ingestion.stores.redis")
+          .option("entity_columns", featureTable.entities.map(_.name).mkString(","))
+          .option("namespace", featureTable.name)
+          .option("project_name", featureTable.project)
+          .option("timestamp_column", config.source.timestampColumn)
+          .save()
+
+        config.deadLetterPath match {
+          case Some(path) =>
+            batchDF
+              .filter(!validator.checkAll)
+              .write
+              .format("parquet")
+              .mode(SaveMode.Append)
+              .save(path)
+          case _ =>
+            batchDF
+              .filter(!validator.checkAll)
+              .foreach(r => {
+                println(s"Row failed validation $r")
+              })
+        }
 
-    val query = o.writeStream
-      .format("feast.ingestion.stores.redis")
-      .option("entity_columns", featureTable.entities.map(_.name).mkString(","))
-      .option("namespace", featureTable.name)
-      .option("project_name", featureTable.project)
-      .option("timestamp_column", config.source.timestampColumn)
+        batchDF.unpersist()
+        () // return Unit to avoid compile error with overloaded foreachBatch
+      }
       .start()
 
     Some(query)
   }
 
-  private def loadClass(className: String) =
-    Class.forName(className, true, getClass.getClassLoader)
-
-  def structFieldFor(fd: FieldDescriptor): Option[StructField] = {
-    val dataType = fd.getJavaType match {
-      case INT         => Some(IntegerType)
-      case LONG        => Some(LongType)
-      case FLOAT       => Some(FloatType)
-      case DOUBLE      => Some(DoubleType)
-      case BOOLEAN     => Some(BooleanType)
-      case STRING      => Some(StringType)
-      case BYTE_STRING => Some(BinaryType)
-      case ENUM        => Some(StringType)
-      case MESSAGE =>
-        fd.getMessageType.getFullName match {
-          case "google.protobuf.Timestamp" => Some(TimestampType)
-          case _ =>
-            Option(fd.getMessageType.getFields.flatMap(structFieldFor))
-              .filter(_.nonEmpty)
-              .map(StructType.apply)
-        }
-    }
+  private def defaultInstanceFromProtoClass(className: String): GeneratedMessageV3 =
+    Class
+      .forName(className, true, getClass.getClassLoader)
+      .asInstanceOf[Class[GeneratedMessageV3]]
+      .getMethod("getDefaultInstance")
+      .invoke(null)
+      .asInstanceOf[GeneratedMessageV3]
 
-    dataType.map(dt =>
-      StructField(
-        fd.getName,
-        if (fd.isRepeated) ArrayType(dt, containsNull = false) else dt,
-        nullable = !fd.isRequired && !fd.isRepeated
-      )
+  private def inferSchemaFromProto(defaultInstance: GeneratedMessageV3) =
+    StructType(
+      defaultInstance.getDescriptorForType.getFields.flatMap(ProtoReflection.structFieldFor)
     )
-  }
-
-  /**
-    * @param defaultInstance except protobuf message instance because it's serializable
-    *                        and both parser & fields descriptor can be extracted from it
-    */
-  def createMessageParser(defaultInstance: GeneratedMessageV3): Array[Byte] => Row = {
-    def toRowData(fd: FieldDescriptor, obj: AnyRef): AnyRef = {
-      fd.getJavaType match {
-        case BYTE_STRING => obj.asInstanceOf[ByteString].toByteArray
-        case ENUM        => obj.asInstanceOf[EnumValueDescriptor].getName
-        case MESSAGE =>
-          fd.getMessageType.getFullName match {
-            case "google.protobuf.Timestamp" =>
-              new sql.Timestamp(obj.asInstanceOf[Timestamp].getSeconds * 1000)
-            case _ => messageToRow(obj.asInstanceOf[AbstractMessage])
-          }
-
-        case _ => obj
-      }
-    }
-
-    def messageToRow(message: AbstractMessage) = {
-      val fields = message.getAllFields
-
-      Row(defaultInstance.getDescriptorForType.getFields.map { fd =>
-        if (fields.containsKey(fd)) {
-          val obj = fields.get(fd)
-          if (fd.isRepeated) {
-            obj.asInstanceOf[java.util.List[Object]].map(toRowData(fd, _))
-          } else {
-            toRowData(fd, obj)
-          }
-        } else if (fd.isRepeated) {
-          Seq()
-        } else null
-      }: _*)
-    }
-
-    bytes: Array[Byte] =>
-      messageToRow(defaultInstance.getParserForType.parseFrom(bytes).asInstanceOf[AbstractMessage])
-  }
 
 }
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/DefaultSource.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/DefaultSource.scala
@@ -16,16 +16,14 @@
  */
 package feast.ingestion.stores.redis
 
-import org.apache.spark.sql.execution.streaming.Sink
 import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
-import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider, StreamSinkProvider}
-import org.apache.spark.sql.streaming.OutputMode
+import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider}
 
 /**
   * Entrypoint to Redis Storage. Implements only `CreatableRelationProvider` since it's only possible write to Redis.
   * Here we parse configuration from spark parameters & provide SparkRedisConfig to `RedisSinkRelation`
   */
-class RedisRelationProvider extends CreatableRelationProvider with StreamSinkProvider {
+class RedisRelationProvider extends CreatableRelationProvider {
   override def createRelation(
       sqlContext: SQLContext,
       mode: SaveMode,
@@ -39,13 +37,6 @@ class RedisRelationProvider extends CreatableRelationProvider with StreamSinkPro
 
     relation
   }
-
-  override def createSink(
-      sqlContext: SQLContext,
-      parameters: Map[String, String],
-      partitionColumns: Seq[String],
-      outputMode: OutputMode
-  ): Sink = {}
 }
 
 class DefaultSource extends RedisRelationProvider