allow spark expression in field mapping (#1122)

Signed-off-by: Oleksii Moskalenko <[email protected]>
feast-dev · Nov 2, 2020 · 8d1326f · 8d1326f
1 parent 4c1cc09
commit 8d1326f
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala
@@ -18,7 +18,7 @@ package feast.ingestion
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.{Column, SparkSession}
-import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.functions.expr
 import org.apache.spark.sql.streaming.StreamingQuery
 
 trait BasePipeline {
@@ -87,7 +87,7 @@ trait BasePipeline {
         .map(e => (e.name, e.name))
 
     (featureColumns ++ entitiesColumns ++ timestampColumn).map { case (alias, source) =>
-      col(source).alias(alias)
+      expr(source).alias(alias)
     }.toArray
   }
 }
diff --git a/spark/ingestion/src/test/scala/feast/ingestion/BatchPipelineIT.scala b/spark/ingestion/src/test/scala/feast/ingestion/BatchPipelineIT.scala
@@ -254,4 +254,38 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
       )
     })
   }
+
+  "Spark expressions" should "be allowed in mapping" in new Scope {
+    val gen  = rowGenerator(DateTime.parse("2020-08-01"), DateTime.parse("2020-09-01"))
+    val rows = generateDistinctRows(gen, 100, groupByEntity)
+
+    val tempPath = storeAsParquet(sparkSession, rows)
+
+    val configWithMapping = config.copy(
+      source = FileSource(
+        tempPath,
+        Map(
+          "feature1" -> "feature1 + 1",
+          "feature2" -> "feature1 + feature2 * 2"
+        ),
+        "eventTimestamp"
+      )
+    )
+
+    BatchPipeline.createPipeline(sparkSession, configWithMapping)
+
+    val featureKeyEncoder: String => String = encodeFeatureKey(config.featureTable)
+
+    rows.foreach(r => {
+      val storedValues =
+        jedis.hgetAll(encodeEntityKey(r, configWithMapping.featureTable)).asScala.toMap
+      storedValues should beStoredRow(
+        Map(
+          featureKeyEncoder("feature1") -> (r.feature1 + 1),
+          featureKeyEncoder("feature2") -> (r.feature1 + r.feature2 * 2),
+          "_ts:test-fs"                 -> r.eventTimestamp
+        )
+      )
+    })
+  }
 }