Skip to content
Prev Previous commit
refactor: fold pre-computed path into _apply_bfv_transformations per …
…review

Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
  • Loading branch information
abhijeet-dhumal authored and ntkathole committed Jul 20, 2026
commit 16798e43a9233a86f9c572f4309bc964ac2eeea2
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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 = [
Expand Down Expand Up @@ -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

Expand All @@ -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)
Expand All @@ -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)
Expand Down
Loading