feat: support staging for spark materialization ([#5671](#5671))

jfw-ppi · jfw-ppi · commit b3ba61166c94 · 2025-12-31T22:02:21.000+01:00
Signed-off-by: Jacob Weinhold &lt;29459386+jfw-ppi@users.noreply.github.com&gt;
diff --git a/docs/reference/offline-stores/spark.md b/docs/reference/offline-stores/spark.md
@@ -34,7 +34,6 @@ offline_store:
         spark.sql.execution.arrow.pyspark.enabled: "true"
     # Optional: spill large materializations to the staging location instead of collecting in the driver
     staging_location: "s3://my-bucket/tmp/feast"
-    staging_allow_materialize: true
 online_store:
     path: data/online_store.db
 ```
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py
@@ -28,7 +28,7 @@
 import pyarrow.dataset as ds
 import pyarrow.parquet as pq
 import pyspark
-from pydantic import StrictBool, StrictStr
+from pydantic import StrictStr
 from pyspark import SparkConf
 from pyspark.sql import SparkSession
 
@@ -68,9 +68,6 @@ class SparkOfflineStoreConfig(FeastConfigBaseModel):
     staging_location: Optional[StrictStr] = None
     """ Remote path for batch materialization jobs"""
 
-    staging_allow_materialize: StrictBool = False
-    """ Enable use of staging_location during materialization to avoid driver OOM """
-
     region: Optional[StrictStr] = None
     """ AWS Region if applicable for s3-based staging locations"""
 
@@ -459,7 +456,6 @@ def _should_use_staging_for_arrow(self) -> bool:
         offline_store = getattr(self._config, "offline_store", None)
         return bool(
             isinstance(offline_store, SparkOfflineStoreConfig)
-            and getattr(offline_store, "staging_allow_materialize", False)
             and getattr(offline_store, "staging_location", None)
         )
 
diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_spark.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_spark.py
@@ -518,7 +518,6 @@ def test_to_arrow_uses_staging_when_enabled(monkeypatch, tmp_path):
         offline_store=SparkOfflineStoreConfig(
             type="spark",
             staging_location=str(tmp_path),
-            staging_allow_materialize=True,
         ),
     )
 
@@ -555,7 +554,6 @@ def test_to_arrow_normalizes_local_staging_paths(monkeypatch, tmp_path):
         offline_store=SparkOfflineStoreConfig(
             type="spark",
             staging_location=str(tmp_path / "local"),
-            staging_allow_materialize=True,
         ),
     )
 
@@ -590,8 +588,7 @@ def test_to_arrow_falls_back_to_pandas_when_staging_disabled(monkeypatch):
         provider="local",
         offline_store=SparkOfflineStoreConfig(
             type="spark",
-            staging_location="/tmp",
-            staging_allow_materialize=False,
+            staging_location=None,
         ),
     )
 
