implement data validation scala to avoid python dependencies (#148)

keshav2211 · web-flow · commit 311db4669fc6 · 2022-06-15T11:10:58.000+08:00
* implement data validation scala to avoid python dependencies

Signed-off-by: KeshavSharma &lt;keshav.sharma@gojek.com&gt;

* fix formatting

Signed-off-by: KeshavSharma &lt;keshav.sharma@gojek.com&gt;

* reuse spark session from existing class

Signed-off-by: KeshavSharma &lt;keshav.sharma@gojek.com&gt;

* dont parse expectation config in rowvalidator

Signed-off-by: KeshavSharma &lt;keshav.sharma@gojek.com&gt;

Co-authored-by: KeshavSharma &lt;keshav.sharma@gojek.com&gt;
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala
@@ -40,8 +40,9 @@ object BatchPipeline extends BasePipeline {
       config: IngestionJobConfig
   ): Option[StreamingQuery] = {
     val featureTable = config.featureTable
-    val rowValidator = new RowValidator(featureTable, config.source.eventTimestampColumn)
-    val metrics      = new IngestionPipelineMetrics
+    val rowValidator =
+      new RowValidator(featureTable, config.source.eventTimestampColumn, config.expectationSpec)
+    val metrics = new IngestionPipelineMetrics
 
     val input = config.source match {
       case source: BQSource =>
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala
@@ -67,7 +67,9 @@ object IngestionJob {
           featureTable = ft,
           streamingTriggeringSecs = ft.labels.getOrElse("_streaming_trigger_secs", "0").toInt,
           validationConfig =
-            ft.labels.get("_validation").map(parseJSON(_).camelizeKeys.extract[ValidationConfig])
+            ft.labels.get("_validation").map(parseJSON(_).camelizeKeys.extract[ValidationConfig]),
+          expectationSpec =
+            ft.labels.get("_expectations").map(parseJSON(_).camelizeKeys.extract[ExpectationSpec])
         )
       })
       .required()
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala
@@ -120,6 +120,15 @@ case class ValidationConfig(
     includeArchivePath: String
 )
 
+case class ExpectationSpec(
+    expectations: List[Expectation]
+)
+
+case class Expectation(
+    expectationType: String,
+    kwargs: Map[String, String]
+)
+
 case class IngestionJobConfig(
     mode: Modes = Modes.Offline,
     featureTable: FeatureTable = null,
@@ -132,6 +141,7 @@ case class IngestionJobConfig(
     stencilURL: Option[String] = None,
     streamingTriggeringSecs: Int = 0,
     validationConfig: Option[ValidationConfig] = None,
+    expectationSpec: Option[ExpectationSpec] = None,
     doNotIngestInvalidRows: Boolean = false,
     checkpointPath: Option[String] = None
 )
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala
@@ -54,8 +54,9 @@ object StreamingPipeline extends BasePipeline with Serializable {
   ): Option[StreamingQuery] = {
     import sparkSession.implicits._
 
-    val featureTable     = config.featureTable
-    val rowValidator     = new RowValidator(featureTable, config.source.eventTimestampColumn)
+    val featureTable = config.featureTable
+    val rowValidator =
+      new RowValidator(featureTable, config.source.eventTimestampColumn, config.expectationSpec)
     val metrics          = new IngestionPipelineMetrics
     val streamingMetrics = new StreamingMetrics
 
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/validation/RowValidator.scala b/spark/ingestion/src/main/scala/feast/ingestion/validation/RowValidator.scala
@@ -16,11 +16,16 @@
  */
 package feast.ingestion.validation
 
-import feast.ingestion.FeatureTable
+import feast.ingestion.{FeatureTable, ExpectationSpec, Expectation}
 import org.apache.spark.sql.Column
-import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.functions.{col, lit}
+
+class RowValidator(
+    featureTable: FeatureTable,
+    timestampColumn: String,
+    expectationSpec: Option[ExpectationSpec]
+) extends Serializable {
 
-class RowValidator(featureTable: FeatureTable, timestampColumn: String) extends Serializable {
   def allEntitiesPresent: Column =
     featureTable.entities.map(e => col(e.name).isNotNull).reduce(_.&&(_))
 
@@ -30,6 +35,36 @@ class RowValidator(featureTable: FeatureTable, timestampColumn: String) extends
   def timestampPresent: Column =
     col(timestampColumn).isNotNull
 
+  def expectColumnValuesToBeBetween(expectation: Expectation): Column = {
+    val minValue: Option[String] = expectation.kwargs.get("minValue")
+    val maxValue: Option[String] = expectation.kwargs.get("maxValue")
+
+    (minValue, maxValue) match {
+      case (Some(min), Some(max)) => col(expectation.kwargs("column")).between(min, max)
+      case (Some(min), None)      => col(expectation.kwargs("column")).>=(min)
+      case (None, Some(max))      => col(expectation.kwargs("column")).<=(max)
+      case _                      => lit(true)
+    }
+  }
+
+  def validate(expectation: Expectation): Column = {
+    expectation.expectationType match {
+      case "expect_column_values_to_not_be_null" => col(expectation.kwargs("column")).isNotNull
+      case "expect_column_values_to_be_between"  => expectColumnValuesToBeBetween(expectation)
+      case _                                     => lit(true)
+    }
+  }
+
+  def validationChecks: Column = {
+
+    expectationSpec match {
+      case Some(value) if value.expectations.isEmpty => lit(true)
+      case Some(value) =>
+        value.expectations.map(expectation => validate(expectation)).reduce(_.&&(_))
+      case None => lit(true)
+    }
+  }
+
   def allChecks: Column =
-    allEntitiesPresent && timestampPresent
+    allEntitiesPresent && timestampPresent && validationChecks
 }
diff --git a/spark/ingestion/src/test/scala/feast/ingestion/RowValidatorTest.scala b/spark/ingestion/src/test/scala/feast/ingestion/RowValidatorTest.scala
