Add checkpoint in each iteration of join loop (historical retrieval j…

…ob) to truncate logical plan (#92) Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
feast-dev · Aug 18, 2021 · 4b553d5 · 4b553d5
1 parent 7c6878a
commit 4b553d5
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/feast_spark/pyspark/historical_feature_retrieval_job.py b/python/feast_spark/pyspark/historical_feature_retrieval_job.py
@@ -539,6 +539,9 @@ def join_entity_to_feature_tables(
         joined_df = as_of_join(
             joined_df, entity_event_timestamp_column, feature_table_df, feature_table,
         )
+        if SparkContext._active_spark_context._jsc.sc().getCheckpointDir().nonEmpty():
+            joined_df = joined_df.checkpoint()
+
     return joined_df