add checkpoint in join loop

pyalex · pyalex · commit d591297f5464 · 2021-08-18T14:53:37.000+08:00
Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
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
 
 

Original file line number	Diff line number	Diff line change
`@@ -539,6 +539,9 @@ def join_entity_to_feature_tables(`
`539`	`539`	`joined_df = as_of_join(`
`540`	`540`	`joined_df, entity_event_timestamp_column, feature_table_df, feature_table,`
`541`	`541`	`)`
	`542`	`+ if SparkContext._active_spark_context._jsc.sc().getCheckpointDir().nonEmpty():`
	`543`	`+ joined_df = joined_df.checkpoint()`
	`544`	`+`
`542`	`545`	`return joined_df`
`543`	`546`
`544`	`547`