Visitar URL original
fix(spark): Use SELECT * when feature_name_columns is empty in pull_a… · feast-dev/feast@6c8b25c · GitHub
Skip to content

Commit 6c8b25c

Browse files
abhijeet-dhumalabhijeet-dhumal
committed
fix(spark): Use SELECT * when feature_name_columns is empty in pull_all_from_table_or_query
pull_all_from_table_or_query always builds an explicit SELECT projection from join_key_columns + feature_name_columns + timestamp_fields. When feature_name_columns=[] — the "read all source columns" signal used by FeatureBuilder.get_column_info for BatchFeatureView with TransformationMode.PYTHON, ray, and pandas — the generated SQL becomes: SELECT user_id, event_timestamp FROM source WHERE ... All raw feature columns (rating, text, helpful_vote, …) are silently dropped. The UDF receives a 2-column DataFrame and every aggregation returns null or fails. Fix: guard on feature_name_columns being non-empty before building the explicit projection; fall through to SELECT * when it is empty. Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
1 parent 835cda8 commit 6c8b25c

2 files changed

Lines changed: 200 additions & 5 deletions

File tree

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -387,12 +387,18 @@ def pull_all_from_table_or_query(
387387
timestamp_fields = [timestamp_field]
388388
if created_timestamp_column:
389389
timestamp_fields.append(created_timestamp_column)
390-
(fields_with_aliases, aliases) = _get_fields_with_aliases(
391-
fields=join_key_columns + feature_name_columns + timestamp_fields,
392-
field_mappings=data_source.field_mapping,
393-
)
394390

395-
fields_with_alias_string = ", ".join(fields_with_aliases)
391+
if feature_name_columns:
392+
(fields_with_aliases, _) = _get_fields_with_aliases(
393+
fields=join_key_columns + feature_name_columns + timestamp_fields,
394+
field_mappings=data_source.field_mapping,
395+
)
396+
fields_with_alias_string = ", ".join(fields_with_aliases)
397+
else:
398+
# Empty feature_name_columns signals "read all source columns".
399+
# Used by BatchFeatureView with TransformationMode.PYTHON/ray/pandas where
400+
# the UDF computes output features from raw input — don't project upfront.
401+
fields_with_alias_string = "*"
396402

397403
from_expression = data_source.get_table_query_string()
398404
timestamp_filter = get_timestamp_filter_sql(

‎sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_spark_offline_store_pull_all.py‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
"""
2+
Unit tests for SparkOfflineStore.pull_all_from_table_or_query SQL generation.
3+
4+
Covers the bug where feature_name_columns=[] (signalling "read all source
5+
columns" for BatchFeatureView UDF transformations) caused a bare
6+
SELECT user_id, event_timestamp FROM source
7+
instead of SELECT *, silently dropping all columns the UDF needs.
8+
"""
9+
10+
import sys
11+
from datetime import datetime, timezone
12+
from types import ModuleType
13+
from unittest.mock import MagicMock, patch
14+
15+
import pytest
16+
17+
18+
# ---------------------------------------------------------------------------
19+
# Stub pyspark so SparkOfflineStore can be imported without a Spark install.
20+
# Only the minimal surface needed for module-level imports in spark.py.
21+
# ---------------------------------------------------------------------------
22+
def _stub_pyspark():
23+
pyspark = ModuleType("pyspark")
24+
pyspark.SparkConf = MagicMock()
25+
sql = ModuleType("pyspark.sql")
26+
sql.SparkSession = MagicMock()
27+
sql.DataFrame = MagicMock()
28+
sql.functions = ModuleType("pyspark.sql.functions")
29+
sql.types = ModuleType("pyspark.sql.types")
30+
sql.types.StructType = MagicMock()
31+
pyspark.sql = sql
32+
for name, mod in {
33+
"pyspark": pyspark,
34+
"pyspark.sql": sql,
35+
"pyspark.sql.functions": sql.functions,
36+
"pyspark.sql.types": sql.types,
37+
}.items():
38+
sys.modules.setdefault(name, mod)
39+
40+
41+
_stub_pyspark()
42+
43+
# Feast imports must come after _stub_pyspark() so pyspark is already stubbed
44+
# in sys.modules before feast.infra.offline_stores.contrib.spark_offline_store
45+
# tries to import it at module level.
46+
from feast.infra.offline_stores.contrib.spark_offline_store.spark import ( # noqa: E402
47+
SparkOfflineStore,
48+
SparkOfflineStoreConfig,
49+
)
50+
from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import ( # noqa: E402
51+
SparkSource,
52+
)
53+
from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig # noqa: E402
54+
from feast.repo_config import RepoConfig # noqa: E402
55+
56+
# ---------------------------------------------------------------------------
57+
# Shared fixtures
58+
# ---------------------------------------------------------------------------
59+
60+
START = datetime(2023, 1, 1, tzinfo=timezone.utc)
61+
END = datetime(2024, 1, 1, tzinfo=timezone.utc)
62+
63+
64+
@pytest.fixture()
65+
def repo_config():
66+
return RepoConfig(
67+
registry="file:///tmp/registry.db",
68+
project="test",
69+
provider="local",
70+
online_store=SqliteOnlineStoreConfig(type="sqlite"),
71+
offline_store=SparkOfflineStoreConfig(type="spark"),
72+
)
73+
74+
75+
@pytest.fixture()
76+
def spark_source():
77+
src = SparkSource(
78+
name="raw_reviews",
79+
path="s3a://bucket/processed/reviews/",
80+
file_format="parquet",
81+
timestamp_field="event_timestamp",
82+
)
83+
return src
84+
85+
86+
def _run_pull_all(repo_config, spark_source, feature_name_columns):
87+
"""
88+
Call pull_all_from_table_or_query with a mocked SparkSession and return
89+
the SQL query string that would be issued.
90+
"""
91+
mock_spark = MagicMock()
92+
93+
with patch(
94+
"feast.infra.offline_stores.contrib.spark_offline_store.spark"
95+
".get_spark_session_or_start_new_with_repoconfig",
96+
return_value=mock_spark,
97+
):
98+
job = SparkOfflineStore.pull_all_from_table_or_query(
99+
config=repo_config,
100+
data_source=spark_source,
101+
join_key_columns=["user_id"],
102+
feature_name_columns=feature_name_columns,
103+
timestamp_field="event_timestamp",
104+
created_timestamp_column=None,
105+
start_date=START,
106+
end_date=END,
107+
)
108+
109+
return job.query.strip()
110+
111+
112+
def test_pull_all_with_empty_feature_cols_generates_select_star(
113+
repo_config, spark_source
114+
):
115+
"""
116+
feature_name_columns=[] must produce SELECT * so UDF-based
117+
BatchFeatureViews receive all raw source columns for aggregation.
118+
"""
119+
sql = _run_pull_all(repo_config, spark_source, feature_name_columns=[])
120+
121+
assert sql.startswith("SELECT *"), (
122+
"Expected 'SELECT *' when feature_name_columns=[], "
123+
f"got: {sql[:120]!r}\n\n"
124+
"BatchFeatureView UDFs need all raw source columns to compute "
125+
"aggregations — projecting only join key + timestamp silently "
126+
"drops rating, text, helpful_vote, etc."
127+
)
128+
assert "user_id" not in sql.split("FROM")[0], (
129+
"SELECT * must not also explicitly list join key columns"
130+
)
131+
132+
133+
def test_pull_all_with_feature_cols_generates_explicit_projection(
134+
repo_config, spark_source
135+
):
136+
"""
137+
When feature_name_columns is non-empty (normal FeatureView path),
138+
the query must project only the requested columns — not SELECT *.
139+
"""
140+
sql = _run_pull_all(
141+
repo_config,
142+
spark_source,
143+
feature_name_columns=["avg_rating", "review_count"],
144+
)
145+
146+
assert "SELECT *" not in sql, (
147+
"Non-empty feature_name_columns must produce explicit SELECT projection, not SELECT *"
148+
)
149+
assert "avg_rating" in sql
150+
assert "review_count" in sql
151+
assert "user_id" in sql
152+
assert "event_timestamp" in sql
153+
154+
155+
def test_pull_all_empty_feature_cols_upstream_regression(repo_config, spark_source):
156+
"""
157+
Regression guard: the upstream (unfixed) behaviour with feature_name_columns=[]
158+
produced a query that only selected join key + timestamp, dropping all columns
159+
the UDF needs. Verify the fixed code does NOT produce that broken query.
160+
161+
Broken upstream SQL looked like:
162+
SELECT user_id, event_timestamp FROM ... WHERE ...
163+
"""
164+
sql = _run_pull_all(repo_config, spark_source, feature_name_columns=[])
165+
166+
projection = sql.split("FROM")[0]
167+
assert "user_id" not in projection, (
168+
"Upstream bug: query projected only 'user_id, event_timestamp', "
169+
"silently dropping all columns needed by the BFV UDF. "
170+
"Fixed query should use SELECT *."
171+
)
172+
173+
174+
@pytest.mark.parametrize(
175+
"feature_cols,expect_star",
176+
[
177+
([], True),
178+
(["f1"], False),
179+
(["f1", "f2", "f3"], False),
180+
],
181+
)
182+
def test_pull_all_select_star_only_when_feature_cols_empty(
183+
repo_config, spark_source, feature_cols, expect_star
184+
):
185+
sql = _run_pull_all(repo_config, spark_source, feature_name_columns=feature_cols)
186+
has_star = sql.strip().upper().startswith("SELECT *")
187+
assert has_star == expect_star, (
188+
f"feature_cols={feature_cols!r}: expected SELECT *={expect_star}, got SQL: {sql[:100]!r}"
189+
)

0 commit comments

Comments
 (0)