refactor: fold pre-computed path into _apply_bfv_transformations per review

abhijeet-dhumal · ntkathole · commit 16798e43a923 · 2026-07-20T14:15:49.000+05:30
Signed-off-by: abhijeet-dhumal &lt;abhijeetdhumal652@gmail.com&gt;
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
@@ -92,87 +92,6 @@ class SparkFeatureViewQueryContext(offline_utils.FeatureViewQueryContext):
     max_date_partition: Optional[str]
 
 
-def _apply_bfv_transformations_for_historical(
-    spark_session: SparkSession,
-    feature_views: List[FeatureView],
-    query_context: List[offline_utils.FeatureViewQueryContext],
-) -> List[offline_utils.FeatureViewQueryContext]:
-    """
-    For BatchFeatureViews, redirect get_historical_features to read from the
-    pre-materialized offline store (batch_source.path) when available, avoiding
-    expensive UDF re-execution on raw data.
-
-    Precedence:
-      1. offline=True + batch_source.path set -> read pre-computed parquet
-      2. Python/pandas UDF present -> execute UDF on raw source (fallback)
-      3. Otherwise -> pass through unchanged
-    """
-    from dataclasses import replace
-
-    fv_by_name = {fv.projection.name_to_use(): fv for fv in feature_views}
-    new_contexts = []
-
-    for ctx in query_context:
-        fv = fv_by_name.get(ctx.name)
-        if fv is None or not isinstance(fv, BatchFeatureView):
-            new_contexts.append(ctx)
-            continue
-
-        if (
-            getattr(fv, "offline", False)
-            and isinstance(fv.batch_source, SparkSource)
-            and fv.batch_source.path
-        ):
-            tmp_view = f"__feast_offline_{ctx.name}_{uuid.uuid4().hex[:8]}"
-            file_format = fv.batch_source.file_format or "parquet"
-            try:
-                df = spark_session.read.format(file_format).load(fv.batch_source.path)
-                df.createOrReplaceTempView(tmp_view)
-                ctx = replace(ctx, table_subquery=tmp_view)
-                new_contexts.append(ctx)
-                continue
-            except (FileNotFoundError, PermissionError) as e:
-                warnings.warn(
-                    f"Offline path '{fv.batch_source.path}' not accessible for "
-                    f"'{ctx.name}': {e}; falling back to source query.",
-                    RuntimeWarning,
-                    stacklevel=2,
-                )
-            except Exception as e:
-                warnings.warn(
-                    f"Unexpected error loading offline path '{fv.batch_source.path}' "
-                    f"for '{ctx.name}': {e}; falling back to source query.",
-                    RuntimeWarning,
-                    stacklevel=2,
-                )
-
-        if (
-            hasattr(fv, "feature_transformation")
-            and fv.feature_transformation is not None
-            and (
-                getattr(fv.feature_transformation, "mode", None) in ("python", "pandas")
-                or getattr(
-                    getattr(fv.feature_transformation, "mode", None), "value", None
-                )
-                in ("python", "pandas")
-            )
-        ):
-            udf = getattr(fv.feature_transformation, "udf", None) or getattr(
-                fv, "udf", None
-            )
-            if udf is not None:
-                temp_view_name = f"__feast_bfv_{ctx.name}_{uuid.uuid4().hex[:8]}"
-                spark_session.conf.set("spark.sql.runSQLOnFiles", "true")
-                raw_df = spark_session.sql(f"SELECT * FROM {ctx.table_subquery}")
-                transformed_df = udf(raw_df)
-                transformed_df.createOrReplaceTempView(temp_view_name)
-                ctx = replace(ctx, table_subquery=temp_view_name)
-
-        new_contexts.append(ctx)
-
-    return new_contexts
-
-
 class SparkOfflineStore(OfflineStore):
     @staticmethod
     def pull_latest_from_table_or_query(
@@ -379,10 +298,10 @@ def get_historical_features(
             entity_df_event_timestamp_range,
         )
 
-        query_context = _apply_bfv_transformations_for_historical(
+        query_context = _apply_bfv_transformations(
             spark_session=spark_session,
             feature_views=feature_views,
-            query_context=query_context,
+            query_contexts=query_context,
         )
 
         spark_query_context = [
@@ -1489,9 +1408,16 @@ def _apply_bfv_transformations(
     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.
+    For BatchFeatureViews, update each query context in one of two ways:
+
+    1. Pre-computed path shortcut: if ``offline=True`` and
+       ``batch_source.path`` is set, read the pre-materialized parquet
+       directly — avoids re-running the UDF on every training call.
+    2. UDF execution: if the BFV has a transformation, run it against
+       the raw source and register the result as a temp view.
+
+    Plain FeatureViews and BFVs with neither a path nor a UDF pass
+    through unchanged.
     """
     from dataclasses import replace
 
@@ -1506,11 +1432,41 @@ def _apply_bfv_transformations(
     updated_contexts = []
     for ctx in query_contexts:
         fv = fv_by_name.get(ctx.name)
+        if fv is None or not isinstance(fv, BatchFeatureView):
+            updated_contexts.append(ctx)
+            continue
+
+        # 1. Pre-computed path shortcut
         if (
-            fv is not None
-            and isinstance(fv, BatchFeatureView)
-            and has_transformation(fv)
+            getattr(fv, "offline", False)
+            and isinstance(fv.batch_source, SparkSource)
+            and fv.batch_source.path
         ):
+            tmp_view = f"__feast_offline_{ctx.name}_{uuid.uuid4().hex[:8]}"
+            file_format = fv.batch_source.file_format or "parquet"
+            try:
+                df = spark_session.read.format(file_format).load(fv.batch_source.path)
+                df.createOrReplaceTempView(tmp_view)
+                updated_contexts.append(replace(ctx, table_subquery=tmp_view))
+                continue
+            except (FileNotFoundError, PermissionError) as e:
+                warnings.warn(
+                    f"Offline path '{fv.batch_source.path}' not accessible for "
+                    f"'{ctx.name}': {e}; falling back to source query.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+            except Exception as e:
+                warnings.warn(
+                    f"Unexpected error loading offline path "
+                    f"'{fv.batch_source.path}' for '{ctx.name}': {e}; "
+                    f"falling back to source query.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+
+        # 2. UDF execution fallback
+        if has_transformation(fv):
             udf = get_transformation_function(fv)
             if udf is not None:
                 source_info = resolve_feature_view_source_with_fallback(fv)
@@ -1526,12 +1482,9 @@ def _apply_bfv_transformations(
                 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)
