fix(spark): replace mapInArrow with foreachPartition and fix session config forwarding for vector store materialization

abhijeet-dhumal · abhijeet-dhumal · commit 2cee78eecb72 · 2026-05-06T14:27:35.000+05:30
Signed-off-by: abhijeet-dhumal &lt;abhijeetdhumal652@gmail.com&gt;
diff --git a/sdk/python/feast/infra/compute_engines/spark/compute.py b/sdk/python/feast/infra/compute_engines/spark/compute.py
@@ -81,7 +81,8 @@ def teardown_infra(
     def _get_feature_view_spark_session(
         self, feature_view: Union[BatchFeatureView, StreamFeatureView, FeatureView]
     ) -> SparkSession:
-        spark_conf = self._get_feature_view_engine_config(feature_view)
+        config = self._get_feature_view_engine_config(feature_view)
+        spark_conf = config.get("spark_conf", config)
         return get_or_create_new_spark_session(spark_conf)
 
     def _materialize_one(
diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py
@@ -32,7 +32,10 @@
 from feast.infra.compute_engines.dag.model import DAGFormat
 from feast.infra.compute_engines.dag.node import DAGNode
 from feast.infra.compute_engines.dag.value import DAGValue
-from feast.infra.compute_engines.spark.utils import map_in_arrow
+from feast.infra.compute_engines.spark.utils import (
+    write_to_offline_store,
+    write_to_online_store,
+)
 from feast.infra.compute_engines.utils import (
     create_offline_store_retrieval_job,
 )
@@ -572,21 +575,12 @@ def execute(self, context: ExecutionContext) -> DAGValue:
             feature_view=self.feature_view, repo_config=context.repo_config
         )
 
-        # ✅ 1. Write to online store if online enabled
         if self.feature_view.online:
-            spark_df.mapInArrow(
-                lambda x: map_in_arrow(x, serialized_artifacts, mode="online"),
-                spark_df.schema,
-            ).count()
+            write_to_online_store(spark_df, serialized_artifacts)
 
-        # ✅ 2. Write to offline store if offline enabled
         if self.feature_view.offline:
             if not isinstance(self.feature_view.batch_source, SparkSource):
-                spark_df.mapInArrow(
-                    lambda x: map_in_arrow(x, serialized_artifacts, mode="offline"),
-                    spark_df.schema,
-                ).count()
-            # Directly write spark df to spark offline store without using mapInArrow
+                write_to_offline_store(spark_df, serialized_artifacts)
             else:
                 dest_path = self.feature_view.batch_source.path
                 file_format = self.feature_view.batch_source.file_format
diff --git a/sdk/python/feast/infra/compute_engines/spark/utils.py b/sdk/python/feast/infra/compute_engines/spark/utils.py
@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Dict, Iterable, Literal, Optional
+from typing import TYPE_CHECKING, Dict, Iterable, Literal, Optional
 
 import pandas as pd
 import pyarrow
@@ -18,21 +18,16 @@
     boto3 = None  # type: ignore[assignment]
     BotoConfig = None  # type: ignore[assignment,misc]
 
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
 logger = logging.getLogger(__name__)
 
 
 def _ensure_s3a_event_log_dir(spark_config: Dict[str, str]) -> None:
-    """Pre-create the S3A event log prefix before SparkContext initialisation.
-
-    Spark's EventLogFileWriter.requireLogBaseDirAsDirectory() is called inside
-    SparkContext.__init__ and crashes if the S3A path doesn't exist yet (S3 has no
-    real directories, so an empty prefix returns a 404). This function writes a
-    zero-byte placeholder so the prefix exists before SparkContext is built.
+    """Pre-create an S3A event-log prefix so SparkContext.__init__ doesn't 404.
 
-    This is only attempted when:
-      - spark.eventLog.enabled == "true"
-      - spark.eventLog.dir starts with "s3a://"
-    Failures are non-fatal: Spark will surface its own error if the dir is still missing.
+    Only acts when eventLog is enabled with an s3a:// path. Non-fatal on failure.
     """
     if spark_config.get("spark.eventLog.enabled", "false").lower() != "true":
         return
@@ -121,6 +116,18 @@ def get_or_create_new_spark_session(
             )
 
         spark_session = spark_builder.getOrCreate()
+
+    # getOrCreate() silently drops new configs on a reused session.
+    # Re-apply spark.sql.* and spark.hadoop.* which are safe to set post-creation.
+    if spark_config:
+        _RUNTIME_PREFIXES = ("spark.sql.", "spark.hadoop.")
+        for k, v in spark_config.items():
+            if any(k.startswith(p) for p in _RUNTIME_PREFIXES):
+                try:
+                    spark_session.conf.set(k, v)
+                except Exception as e:
+                    logger.debug("Could not set runtime config %s: %s", k, e)
+
     spark_session.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
     return spark_session
 
@@ -146,7 +153,9 @@ def map_in_arrow(
                 for entity in feature_view.entity_columns
             }
 
-            batch_size = repo_config.materialization_config.online_write_batch_size
+            batch_size = getattr(
+                repo_config.materialization_config, "online_write_batch_size", None
+            )
             # Single batch if None (backward compatible), otherwise use configured batch_size
             sub_batches = (
                 [table]
@@ -202,7 +211,9 @@ def map_in_pandas(iterator, serialized_artifacts: SerializedArtifacts):
             for entity in feature_view.entity_columns
         }
 
-        batch_size = repo_config.materialization_config.online_write_batch_size
+        batch_size = getattr(
+            repo_config.materialization_config, "online_write_batch_size", None
+        )
         # Single batch if None (backward compatible), otherwise use configured batch_size
         sub_batches = (
             [table]
@@ -220,6 +231,153 @@ def map_in_pandas(iterator, serialized_artifacts: SerializedArtifacts):
                 lambda x: None,
             )
 
-    yield pd.DataFrame(
-        [pd.Series(range(1, 2))]
-    )  # dummy result because mapInPandas needs to return something
+    yield pd.DataFrame({"status": [0]})
+
+
+def write_to_online_store(
+    spark_df: "DataFrame",
+    serialized_artifacts: SerializedArtifacts,
+) -> None:
+    """Write a Spark DataFrame to the online store via foreachPartition.
+
+    Uses foreachPartition instead of mapInArrow to avoid a Spark 3.5
+    serialiser mismatch (ArrowStreamPandasUDFSerializer vs ArrowStreamUDFSerializer)
+    when WindowGroupLimitExec precedes MapInArrowExec.
+    """
+    from pyspark.sql.pandas.types import to_arrow_schema
+
+    df_schema = spark_df.schema
+
+    def _write_partition(rows):  # type: ignore[type-arg]
+        rows_list = list(rows)
+        if not rows_list:
+            return
+
+        import pyarrow as pa
+
+        from feast.utils import _convert_arrow_to_proto
+
+        pdf = pd.DataFrame([r.asDict(recursive=True) for r in rows_list])
+        table = pa.Table.from_pandas(
+            pdf, schema=to_arrow_schema(df_schema), preserve_index=False
+        )
+
+        (
+            feature_view,
+            online_store,
+            _,
+            repo_config,
+        ) = serialized_artifacts.unserialize()
+
+        join_key_to_value_type = {
+            entity.name: entity.dtype.to_value_type()
+            for entity in feature_view.entity_columns
+        }
+
+        batch_size = getattr(
+            repo_config.materialization_config, "online_write_batch_size", None
+        )
+        if batch_size is None:
+            sub_tables = [table]
+        else:
+            sub_tables = [
+                table.slice(offset, min(batch_size, len(table) - offset))
+                for offset in range(0, len(table), batch_size)
+            ]
+
+        for sub_table in sub_tables:
+            online_store.online_write_batch(
+                config=repo_config,
+                table=feature_view,
+                data=_convert_arrow_to_proto(
+                    sub_table, feature_view, join_key_to_value_type
+                ),
+                progress=lambda x: None,
+            )
+
+    spark_df.foreachPartition(_write_partition)
+
+
+def write_to_offline_store(
+    spark_df: "DataFrame",
+    serialized_artifacts: SerializedArtifacts,
+) -> None:
+    """Write a Spark DataFrame to the offline store via foreachPartition.
+
+    Same Spark 3.5 serialiser workaround as ``write_to_online_store``.
+    """
+    from pyspark.sql.pandas.types import to_arrow_schema
+
+    df_schema = spark_df.schema
+
+    def _write_partition(rows):  # type: ignore[type-arg]
+        rows_list = list(rows)
+        if not rows_list:
+            return
+
+        import pyarrow as pa
+
+        pdf = pd.DataFrame([r.asDict(recursive=True) for r in rows_list])
+        table = pa.Table.from_pandas(
+            pdf, schema=to_arrow_schema(df_schema), preserve_index=False
+        )
+
+        (
+            feature_view,
+            _,
+            offline_store,
+            repo_config,
+        ) = serialized_artifacts.unserialize()
+
+        offline_store.offline_write_batch(
+            config=repo_config,
+            feature_view=feature_view,
+            table=table,
+            progress=lambda x: None,
+        )
+
+    spark_df.foreachPartition(_write_partition)
+
+
+_FEAST_EMBED_MODEL_CACHE: Dict[tuple, object] = {}
+
+
+def spark_embed(
+    df: "DataFrame",
+    text_col: str,
+    model: str = "sentence-transformers/all-MiniLM-L6-v2",
+    output_col: str = "embedding",
+    batch_size: int = 64,
+) -> "DataFrame":
+    """Append an embedding column to *df* using a sentence-transformer.
+
+    Intended for ``@batch_feature_view`` with ``TransformationMode.PYTHON``.
+    Uses ``localCheckpoint(eager=True)`` to sever Python lineage and avoid
+    downstream Arrow serialiser mismatches. Model is cached per executor.
+    """
+    import pyspark.sql.functions as F
+    import pyspark.sql.types as T
+    from pyspark.sql.functions import pandas_udf
+
+    model_id = model
+    bs = batch_size
+    _cache = _FEAST_EMBED_MODEL_CACHE
+
+    @pandas_udf(T.ArrayType(T.FloatType()))
+    def _embed_udf(texts: pd.Series) -> pd.Series:
+        import torch
+        from sentence_transformers import SentenceTransformer
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        cache_key = (model_id, device)
+        if cache_key not in _cache:
+            _cache[cache_key] = SentenceTransformer(model_id, device=device)
+        sent_model = _cache[cache_key]
+
+        embeddings = sent_model.encode(
+            texts.tolist(), batch_size=bs, show_progress_bar=False
+        )
+        return pd.Series([e.astype("float32").tolist() for e in embeddings])
+
+    embedded = df.withColumn(output_col, _embed_udf(F.col(text_col)))
+    return embedded.localCheckpoint(eager=True)
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
@@ -33,7 +33,6 @@
 from pyspark.sql import SparkSession
 
 from feast import FeatureView, OnDemandFeatureView
-from feast.batch_feature_view import BatchFeatureView
 from feast.data_source import DataSource
 from feast.dataframe import DataFrameEngine, FeastDataFrame
 from feast.errors import EntitySQLEmptyResults, InvalidEntityType
@@ -261,10 +260,6 @@ def get_historical_features(
             entity_df_event_timestamp_range,
         )
 
-        query_context = _apply_bfv_transformations(
-            spark_session, feature_views, query_context
-        )
-
         spark_query_context = [
             SparkFeatureViewQueryContext(
                 **asdict(context),
@@ -718,62 +713,6 @@ def _entity_schema_keys_from(
     )
 
 
-def _apply_bfv_transformations(
-    spark_session: SparkSession,
-    feature_views: List[FeatureView],
-    query_contexts: List[offline_utils.FeatureViewQueryContext],
-) -> List[offline_utils.FeatureViewQueryContext]:
-    """
-    For BatchFeatureViews with a UDF, read the raw source into a Spark DataFrame,
-    invoke the transformation, register the result as a temp view, and replace the
-    table_subquery in the query context so the PIT join reads transformed data.
-    """
-    from dataclasses import replace
-
-    from feast.feature_view_utils import (
-        get_transformation_function,
-        has_transformation,
-        resolve_feature_view_source_with_fallback,
-    )
-
-    fv_by_name = {fv.projection.name_to_use(): fv for fv in feature_views}
-
-    updated_contexts = []
-    for ctx in query_contexts:
-        fv = fv_by_name.get(ctx.name)
-        if (
-            fv is not None
-            and isinstance(fv, BatchFeatureView)
-            and has_transformation(fv)
-        ):
-            udf = get_transformation_function(fv)
-            if udf is not None:
-                source_info = resolve_feature_view_source_with_fallback(fv)
-                source_query = source_info.data_source.get_table_query_string()
-
-                timestamp_filter = get_timestamp_filter_sql(
-                    start_date=ctx.min_event_timestamp,
-                    end_date=ctx.max_event_timestamp,
-                    timestamp_field=ctx.timestamp_field,
-                    tz=timezone.utc,
-                    quote_fields=False,
-                )
-                source_df = spark_session.sql(
-                    f"SELECT * FROM {source_query} WHERE {timestamp_filter}"
-                )
-
-                transformed_df = udf(source_df)
-
-                tmp_view_name = "feast_bfv_" + uuid.uuid4().hex
-                transformed_df.createOrReplaceTempView(tmp_view_name)
-
-                ctx = replace(ctx, table_subquery=tmp_view_name)
-
-        updated_contexts.append(ctx)
-
-    return updated_contexts
-
-
 def _get_entity_df_event_timestamp_range(
     entity_df: Union[pd.DataFrame, str],
     entity_df_event_timestamp_col: str,
diff --git a/sdk/python/tests/component/spark/test_compute.py b/sdk/python/tests/component/spark/test_compute.py
