|
| 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