Skip to content

Commit 630d9f8

Browse files
authored
feat(spark): Add compute-on-read support for BatchFeatureView in get_… (#6357)
* feat(spark): add compute-on-read support for BatchFeatureView in get_historical_features Signed-off-by: Siddhesh Khairnar <[email protected]> * Add None check for batch_source in _apply_bfv_transformations Signed-off-by: Siddhesh Khairnar <[email protected]> * refactor(spark): use feature_view_utils helpers for BFV transformation logic Signed-off-by: Siddhesh Khairnar <[email protected]> * fix(spark): apply time-range filter before BFV UDF and update temp view naming Signed-off-by: Siddhesh Khairnar <[email protected]> * refactor(spark): use resolve_feature_view_source_with_fallback for BFV source resolution Signed-off-by: Siddhesh Khairnar <[email protected]> --------- Signed-off-by: Siddhesh Khairnar <[email protected]>
1 parent 77ed779 commit 630d9f8

2 files changed

Lines changed: 283 additions & 0 deletions

File tree

‎sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
from pyspark.sql import SparkSession
3434

3535
from feast import FeatureView, OnDemandFeatureView
36+
from feast.batch_feature_view import BatchFeatureView
3637
from feast.data_source import DataSource
3738
from feast.dataframe import DataFrameEngine, FeastDataFrame
3839
from feast.errors import EntitySQLEmptyResults, InvalidEntityType
@@ -260,6 +261,10 @@ def get_historical_features(
260261
entity_df_event_timestamp_range,
261262
)
262263

264+
query_context = _apply_bfv_transformations(
265+
spark_session, feature_views, query_context
266+
)
267+
263268
spark_query_context = [
264269
SparkFeatureViewQueryContext(
265270
**asdict(context),
@@ -713,6 +718,62 @@ def _entity_schema_keys_from(
713718
)
714719

715720

721+
def _apply_bfv_transformations(
722+
spark_session: SparkSession,
723+
feature_views: List[FeatureView],
724+
query_contexts: List[offline_utils.FeatureViewQueryContext],
725+
) -> List[offline_utils.FeatureViewQueryContext]:
726+
"""
727+
For BatchFeatureViews with a UDF, read the raw source into a Spark DataFrame,
728+
invoke the transformation, register the result as a temp view, and replace the
729+
table_subquery in the query context so the PIT join reads transformed data.
730+
"""
731+
from dataclasses import replace
732+
733+
from feast.feature_view_utils import (
734+
get_transformation_function,
735+
has_transformation,
736+
resolve_feature_view_source_with_fallback,
737+
)
738+
739+
fv_by_name = {fv.projection.name_to_use(): fv for fv in feature_views}
740+
741+
updated_contexts = []
742+
for ctx in query_contexts:
743+
fv = fv_by_name.get(ctx.name)
744+
if (
745+
fv is not None
746+
and isinstance(fv, BatchFeatureView)
747+
and has_transformation(fv)
748+
):
749+
udf = get_transformation_function(fv)
750+
if udf is not None:
751+
source_info = resolve_feature_view_source_with_fallback(fv)
752+
source_query = source_info.data_source.get_table_query_string()
753+
754+
timestamp_filter = get_timestamp_filter_sql(
755+
start_date=ctx.min_event_timestamp,
756+
end_date=ctx.max_event_timestamp,
757+
timestamp_field=ctx.timestamp_field,
758+
tz=timezone.utc,
759+
quote_fields=False,
760+
)
761+
source_df = spark_session.sql(
762+
f"SELECT * FROM {source_query} WHERE {timestamp_filter}"
763+
)
764+
765+
transformed_df = udf(source_df)
766+
767+
tmp_view_name = "feast_bfv_" + uuid.uuid4().hex
768+
transformed_df.createOrReplaceTempView(tmp_view_name)
769+
770+
ctx = replace(ctx, table_subquery=tmp_view_name)
771+
772+
updated_contexts.append(ctx)
773+
774+
return updated_contexts
775+
776+
716777
def _get_entity_df_event_timestamp_range(
717778
entity_df: Union[pd.DataFrame, str],
718779
entity_df_event_timestamp_col: str,
Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
"""
2+
Unit tests for BFV compute-on-read in SparkOfflineStore.get_historical_features().
3+
4+
Verifies that BatchFeatureViews with a UDF have their transformation applied
5+
during get_historical_features(), with the transformed DataFrame registered as
6+
a temp view that replaces the raw table_subquery in the PIT join.
7+
"""
8+
9+
from dataclasses import replace
10+
from unittest.mock import MagicMock
11+
12+
import pytest
13+
14+
from feast.batch_feature_view import BatchFeatureView
15+
from feast.feature_view import FeatureView
16+
from feast.infra.offline_stores import offline_utils
17+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
18+
_apply_bfv_transformations,
19+
)
20+
from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import (
21+
SparkSource,
22+
)
23+
from feast.transformation.base import Transformation
24+
25+
26+
@pytest.fixture()
27+
def spark_session():
28+
mock = MagicMock()
29+
mock.sql.return_value = MagicMock(name="source_df")
30+
return mock
31+
32+
33+
@pytest.fixture()
34+
def spark_source():
35+
source = MagicMock(spec=SparkSource)
36+
source.get_table_query_string.return_value = "`raw_events`"
37+
return source
38+
39+
40+
@pytest.fixture()
41+
def base_query_context():
42+
return offline_utils.FeatureViewQueryContext(
43+
name="my_bfv",
44+
ttl=3600,
45+
entities=["user_id"],
46+
features=["avg_rating"],
47+
field_mapping={},
48+
timestamp_field="event_timestamp",
49+
created_timestamp_column=None,
50+
table_subquery="`raw_events`",
51+
entity_selections=["user_id AS user_id"],
52+
min_event_timestamp="2023-01-01T00:00:00",
53+
max_event_timestamp="2024-01-01T00:00:00",
54+
date_partition_column=None,
55+
)
56+
57+
58+
def _make_bfv(name: str, spark_source, has_udf: bool = True):
59+
"""Create a mock BatchFeatureView with optional UDF."""
60+
fv = MagicMock(spec=BatchFeatureView)
61+
fv.name = name
62+
fv.projection = MagicMock()
63+
fv.projection.name_to_use.return_value = name
64+
fv.batch_source = spark_source
65+
fv.source_views = []
66+
67+
if has_udf:
68+
transformation = MagicMock(spec=Transformation)
69+
transformed_df = MagicMock(name="transformed_df")
70+
transformation.udf = MagicMock(return_value=transformed_df)
71+
fv.feature_transformation = transformation
72+
fv.udf = transformation.udf
73+
else:
74+
fv.feature_transformation = None
75+
fv.udf = None
76+
77+
return fv
78+
79+
80+
def _make_plain_fv(name: str, spark_source):
81+
"""Create a mock plain FeatureView (not a BatchFeatureView)."""
82+
fv = MagicMock(spec=FeatureView)
83+
fv.name = name
84+
fv.projection = MagicMock()
85+
fv.projection.name_to_use.return_value = name
86+
fv.batch_source = spark_source
87+
fv.feature_transformation = None
88+
fv.udf = None
89+
return fv
90+
91+
92+
class TestApplyBfvTransformations:
93+
def test_bfv_with_udf_replaces_table_subquery(
94+
self, spark_session, spark_source, base_query_context
95+
):
96+
"""BFV with a UDF should have its table_subquery replaced with a temp view."""
97+
bfv = _make_bfv("my_bfv", spark_source)
98+
contexts = [base_query_context]
99+
100+
result = _apply_bfv_transformations(spark_session, [bfv], contexts)
101+
102+
assert len(result) == 1
103+
assert result[0].table_subquery != "`raw_events`"
104+
assert result[0].table_subquery.startswith("feast_bfv_")
105+
106+
def test_bfv_udf_is_invoked_with_source_df(
107+
self, spark_session, spark_source, base_query_context
108+
):
109+
"""The UDF should be called with the DataFrame read from the raw source."""
110+
bfv = _make_bfv("my_bfv", spark_source)
111+
contexts = [base_query_context]
112+
113+
_apply_bfv_transformations(spark_session, [bfv], contexts)
114+
115+
sql_arg = spark_session.sql.call_args[0][0]
116+
assert "SELECT * FROM `raw_events`" in sql_arg
117+
assert "WHERE" in sql_arg
118+
source_df = spark_session.sql.return_value
119+
bfv.feature_transformation.udf.assert_called_once_with(source_df)
120+
121+
def test_transformed_df_registered_as_temp_view(
122+
self, spark_session, spark_source, base_query_context
123+
):
124+
"""The transformed DataFrame should be registered as a temp view."""
125+
bfv = _make_bfv("my_bfv", spark_source)
126+
transformed_df = bfv.feature_transformation.udf.return_value
127+
contexts = [base_query_context]
128+
129+
result = _apply_bfv_transformations(spark_session, [bfv], contexts)
130+
131+
transformed_df.createOrReplaceTempView.assert_called_once()
132+
view_name = transformed_df.createOrReplaceTempView.call_args[0][0]
133+
assert view_name == result[0].table_subquery
134+
135+
def test_plain_feature_view_unchanged(
136+
self, spark_session, spark_source, base_query_context
137+
):
138+
"""Plain FeatureViews (not BFV) should pass through without modification."""
139+
fv = _make_plain_fv("my_bfv", spark_source)
140+
contexts = [base_query_context]
141+
142+
result = _apply_bfv_transformations(spark_session, [fv], contexts)
143+
144+
assert result[0].table_subquery == "`raw_events`"
145+
spark_session.sql.assert_not_called()
146+
147+
def test_bfv_without_udf_unchanged(
148+
self, spark_session, spark_source, base_query_context
149+
):
150+
"""BFV without a UDF should pass through without modification."""
151+
bfv = _make_bfv("my_bfv", spark_source, has_udf=False)
152+
contexts = [base_query_context]
153+
154+
result = _apply_bfv_transformations(spark_session, [bfv], contexts)
155+
156+
assert result[0].table_subquery == "`raw_events`"
157+
spark_session.sql.assert_not_called()
158+
159+
def test_mixed_views_only_transforms_bfvs(
160+
self, spark_session, spark_source, base_query_context
161+
):
162+
"""With mixed BFV and plain FVs, only BFVs with UDFs get transformed."""
163+
bfv = _make_bfv("my_bfv", spark_source)
164+
plain_fv = _make_plain_fv("plain_fv", spark_source)
165+
166+
ctx_bfv = base_query_context
167+
ctx_plain = replace(
168+
base_query_context,
169+
name="plain_fv",
170+
features=["some_feature"],
171+
)
172+
173+
result = _apply_bfv_transformations(
174+
spark_session, [bfv, plain_fv], [ctx_bfv, ctx_plain]
175+
)
176+
177+
assert result[0].table_subquery.startswith("feast_bfv_")
178+
assert result[1].table_subquery == "`raw_events`"
179+
180+
def test_time_range_filter_applied(
181+
self, spark_session, spark_source, base_query_context
182+
):
183+
"""Source query should include time bounds from the context."""
184+
bfv = _make_bfv("my_bfv", spark_source)
185+
contexts = [base_query_context]
186+
187+
_apply_bfv_transformations(spark_session, [bfv], contexts)
188+
189+
sql_arg = spark_session.sql.call_args[0][0]
190+
assert "2023-01-01" in sql_arg
191+
assert "2024-01-01" in sql_arg
192+
assert "event_timestamp" in sql_arg
193+
194+
def test_time_range_filter_with_none_min_timestamp(
195+
self, spark_session, spark_source, base_query_context
196+
):
197+
"""When min_event_timestamp is None (no TTL), query should still work."""
198+
bfv = _make_bfv("my_bfv", spark_source)
199+
ctx = replace(base_query_context, min_event_timestamp=None)
200+
201+
result = _apply_bfv_transformations(spark_session, [bfv], [ctx])
202+
203+
assert result[0].table_subquery.startswith("feast_bfv_")
204+
sql_arg = spark_session.sql.call_args[0][0]
205+
assert "2024-01-01" in sql_arg
206+
207+
def test_other_context_fields_preserved(
208+
self, spark_session, spark_source, base_query_context
209+
):
210+
"""All fields besides table_subquery should remain unchanged."""
211+
bfv = _make_bfv("my_bfv", spark_source)
212+
contexts = [base_query_context]
213+
214+
result = _apply_bfv_transformations(spark_session, [bfv], contexts)
215+
216+
assert result[0].name == base_query_context.name
217+
assert result[0].ttl == base_query_context.ttl
218+
assert result[0].entities == base_query_context.entities
219+
assert result[0].features == base_query_context.features
220+
assert result[0].timestamp_field == base_query_context.timestamp_field
221+
assert result[0].min_event_timestamp == base_query_context.min_event_timestamp
222+
assert result[0].max_event_timestamp == base_query_context.max_event_timestamp

0 commit comments

Comments
 (0)