Add persistent checkpoint in historical retrieval (#91)

Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>

python/feast_spark/pyspark/abc.py

@@ -151,6 +151,7 @@ def __init__(
         entity_source: Dict,
         destination: Dict,
         extra_packages: Optional[List[str]] = None,
+        checkpoint_path: Optional[str] = None,
     ):
         """
         Args:
@@ -265,6 +266,7 @@ def __init__(
         self._entity_source = entity_source
         self._destination = destination
         self._extra_packages = extra_packages if extra_packages else []
+        self._checkpoint_path = checkpoint_path
 
     def get_name(self) -> str:
         all_feature_tables_names = [ft["name"] for ft in self._feature_tables]
@@ -285,7 +287,7 @@ def get_arguments(self) -> List[str]:
         def json_b64_encode(obj) -> str:
             return b64encode(json.dumps(obj).encode("utf8")).decode("ascii")
 
-        return [
+        args = [
             "--feature-tables",
             json_b64_encode(self._feature_tables),
             "--feature-tables-sources",
@@ -295,6 +297,9 @@ def json_b64_encode(obj) -> str:
             "--destination",
             json_b64_encode(self._destination),
         ]
+        if self._checkpoint_path:
+            args.extend(["--checkpoint", self._checkpoint_path])
+        return args
 
     def get_destination_path(self) -> str:
         return self._destination["path"]

python/feast_spark/pyspark/historical_feature_retrieval_job.py

@@ -8,6 +8,7 @@
 from logging.config import dictConfig
 from typing import Any, Dict, List, NamedTuple, Optional
 
+from pyspark import SparkContext
 from pyspark.sql import DataFrame, SparkSession, Window
 from pyspark.sql import functions as func
 from pyspark.sql.functions import (
@@ -602,8 +603,9 @@ def filter_feature_table_by_time_range(
         )
         .where(col("distance") == col("min_distance"))
         .select(time_range_filtered_df.columns + [ENTITY_EVENT_TIMESTAMP_ALIAS])
-        .localCheckpoint()
     )
+    if SparkContext._active_spark_context._jsc.sc().getCheckpointDir().nonEmpty():
+        time_range_filtered_df = time_range_filtered_df.checkpoint()
 
     return time_range_filtered_df
 
@@ -848,6 +850,7 @@ def _get_args():
     parser.add_argument(
         "--destination", type=str, help="Retrieval result destination in json string"
     )
+    parser.add_argument("--checkpoint", type=str, help="Spark Checkpoint location")
     return parser.parse_args()
 
 
@@ -876,6 +879,9 @@ def json_b64_decode(s: str) -> Any:
     feature_tables_sources_conf = json_b64_decode(args.feature_tables_sources)
     entity_source_conf = json_b64_decode(args.entity_source)
     destination_conf = json_b64_decode(args.destination)
+    if args.checkpoint:
+        spark.sparkContext.setCheckpointDir(args.checkpoint)
+
     try:
         start_job(
             spark,

python/feast_spark/pyspark/launcher.py

@@ -233,6 +233,7 @@ def start_historical_feature_retrieval_job(
             ],
             destination={"format": output_format, "path": output_path},
             extra_packages=extra_packages,
+            checkpoint_path=client.config.get(opt.CHECKPOINT_PATH),
         )
     )