Optimize memory footprint for Spark Ingestion Job (#1265)

pyalex · web-flow · commit eaaf23372db7 · 2021-01-14T16:04:25.000+08:00
Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala
@@ -69,7 +69,7 @@ object BatchPipeline extends BasePipeline {
     }
 
     val validRows = projected
-      .mapPartitions(metrics.incrementRead)
+      .map(metrics.incrementRead)
       .filter(rowValidator.allChecks)
 
     validRows.write
@@ -85,7 +85,7 @@ object BatchPipeline extends BasePipeline {
       case Some(path) =>
         projected
           .filter(!rowValidator.allChecks)
-          .mapPartitions(metrics.incrementDeadLetters)
+          .map(metrics.incrementDeadLetters)
           .write
           .format("parquet")
           .mode(SaveMode.Append)
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala
@@ -107,7 +107,7 @@ object StreamingPipeline extends BasePipeline with Serializable {
         implicit def rowEncoder: Encoder[Row] = RowEncoder(rowsAfterValidation.schema)
 
         rowsAfterValidation
-          .mapPartitions(metrics.incrementRead)
+          .map(metrics.incrementRead)
           .filter(if (config.doNotIngestInvalidRows) expr("_isValid") else rowValidator.allChecks)
           .write
           .format("feast.ingestion.stores.redis")
@@ -122,7 +122,7 @@ object StreamingPipeline extends BasePipeline with Serializable {
           case Some(path) =>
             rowsAfterValidation
               .filter("!_isValid")
-              .mapPartitions(metrics.incrementDeadLetters)
+              .map(metrics.incrementDeadLetters)
               .write
               .format("parquet")
               .mode(SaveMode.Append)
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/metrics/IngestionPipelineMetrics.scala b/spark/ingestion/src/main/scala/feast/ingestion/metrics/IngestionPipelineMetrics.scala
@@ -22,20 +22,28 @@ import org.apache.spark.sql.Row
 
 class IngestionPipelineMetrics extends Serializable {
 
-  def incrementDeadLetters(rowIterator: Iterator[Row]): Iterator[Row] = {
-    val materialized = rowIterator.toArray
+  def incrementDeadLetters(row: Row): Row = {
     if (metricSource.nonEmpty)
-      metricSource.get.METRIC_DEADLETTER_ROWS_INSERTED.inc(materialized.length)
+      metricSource.get.METRIC_DEADLETTER_ROWS_INSERTED.inc()
 
-    materialized.toIterator
+    row
   }
 
-  def incrementRead(rowIterator: Iterator[Row]): Iterator[Row] = {
-    val materialized = rowIterator.toArray
+  def incrementRead(row: Row): Row = {
     if (metricSource.nonEmpty)
-      metricSource.get.METRIC_ROWS_READ_FROM_SOURCE.inc(materialized.length)
+      metricSource.get.METRIC_ROWS_READ_FROM_SOURCE.inc()
 
-    materialized.toIterator
+    row
+  }
+
+  def incrementRead(inc: Long): Unit = {
+    if (metricSource.nonEmpty)
+      metricSource.get.METRIC_ROWS_READ_FROM_SOURCE.inc(inc)
+  }
+
+  def incrementDeadLetters(inc: Long): Unit = {
+    if (metricSource.nonEmpty)
+      metricSource.get.METRIC_DEADLETTER_ROWS_INSERTED.inc(inc)
   }
 
   private lazy val metricSource: Option[IngestionPipelineMetricSource] = {
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/RedisSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/RedisSinkRelation.scala
@@ -71,7 +71,7 @@ class RedisSinkRelation(override val sqlContext: SQLContext, config: SparkRedisC
     // repartition for deduplication
     val dataToStore =
       if (config.repartitionByEntity)
-        data.repartition(config.entityColumns.map(col): _*)
+        data.repartition(config.entityColumns.map(col): _*).localCheckpoint()
       else data
 
     dataToStore.foreachPartition { partition: Iterator[Row] =>
