From d591297f54645ccd16eca252ebb9d2cbc0cb4305 Mon Sep 17 00:00:00 2001 From: Oleksii Moskalenko Date: Wed, 18 Aug 2021 14:53:37 +0800 Subject: [PATCH] add checkpoint in join loop Signed-off-by: Oleksii Moskalenko --- python/feast_spark/pyspark/historical_feature_retrieval_job.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/feast_spark/pyspark/historical_feature_retrieval_job.py b/python/feast_spark/pyspark/historical_feature_retrieval_job.py index 094b94ac..3dc642a5 100644 --- 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