diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 34a77310bac..13fd9f288bc 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -2341,7 +2341,7 @@ def _materialize_odfv( } for source_fv in source_fvs: - all_join_keys.update(source_fv.entities) + all_join_keys.update(source_fv.join_keys) if source_fv.batch_source: entity_timestamp_col_names.add(source_fv.batch_source.timestamp_field) @@ -2381,7 +2381,7 @@ def _materialize_odfv( job = provider.offline_store.pull_latest_from_table_or_query( config=self.config, data_source=source_fv.batch_source, - join_key_columns=source_fv.entities, + join_key_columns=source_fv.join_keys, feature_name_columns=[f.name for f in source_fv.features], timestamp_field=source_fv.batch_source.timestamp_field, created_timestamp_column=getattr( diff --git a/sdk/python/tests/unit/local_feast_tests/test_local_feature_store.py b/sdk/python/tests/unit/local_feast_tests/test_local_feature_store.py index 9b7660bf692..ec2513f0726 100644 --- a/sdk/python/tests/unit/local_feast_tests/test_local_feature_store.py +++ b/sdk/python/tests/unit/local_feast_tests/test_local_feature_store.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from tempfile import mkstemp from unittest.mock import AsyncMock, Mock, patch @@ -19,11 +19,13 @@ from feast.infra.offline_stores.file_source import FileSource from feast.infra.online_stores.dynamodb import DynamoDBOnlineStore from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig +from feast.on_demand_feature_view import OnDemandFeatureView from feast.permissions.action import AuthzedAction from feast.permissions.permission import Permission from feast.permissions.policy import RoleBasedPolicy from feast.repo_config import RegistryConfig, RepoConfig from feast.stream_feature_view import stream_feature_view +from feast.transformation.pandas_transformation import PandasTransformation from feast.types import Array, Bytes, Float32, Int64, String, ValueType, from_value_type from tests.universal.feature_repos.universal.feature_views import TAGS from tests.utils.cli_repo_creator import CliRunner, get_example_repo @@ -511,6 +513,69 @@ def test_reapply_feature_view(test_feature_store, dataframe_source): test_feature_store.teardown() +@pytest.mark.parametrize( + "test_feature_store", + [lazy_fixture("feature_store_with_local_registry")], +) +@pytest.mark.parametrize("dataframe_source", [lazy_fixture("simple_dataset_1")]) +def test_materialize_incremental_odfv_entity_name_differs_from_join_key( + test_feature_store, dataframe_source +): + """Regression test for https://github.com/feast-dev/feast/issues/5965. + + _materialize_odfv must resolve source feature views' join key *columns* + (FeatureView.join_keys), not their entity *names* (FeatureView.entities), + when building the entity_df and querying the offline store. Before the + fix, this failed whenever an Entity's name differed from its join_keys. + """ + with prep_file_source(df=dataframe_source, timestamp_field="ts_1") as file_source: + # Entity name ("id") intentionally differs from its join key + # ("id_join_key"), mirroring the exact repro from the issue. + e = Entity(name="id", join_keys=["id_join_key"]) + + source_fv = FeatureView( + name="my_feature_view_1", + schema=[Field(name="float_col", dtype=Float32)], + entities=[e], + source=file_source, + ttl=timedelta(days=3650), + ) + + def transform(features_df: pd.DataFrame) -> pd.DataFrame: + out = pd.DataFrame() + out["label"] = features_df["float_col"].apply( + lambda v: "high" if v >= 1 else "low" + ) + return out + + odfv = OnDemandFeatureView( + name="my_odfv", + entities=[e], + sources=[source_fv], + schema=[Field(name="label", dtype=String)], + feature_transformation=PandasTransformation( + udf=transform, udf_string="transform" + ), + write_to_online_store=True, + ) + + test_feature_store.apply([e, source_fv, odfv]) + + # Should not raise. Before the fix, this raised + # FeastJoinKeysDuringMaterialization because _materialize_odfv + # queried the offline store for a column named "id" (the entity + # name) instead of "id_join_key" (the actual join key column). + test_feature_store.materialize_incremental(end_date=datetime.now(timezone.utc)) + + response = test_feature_store.get_online_features( + features=["my_odfv:label"], + entity_rows=[{"id_join_key": 1}], + ).to_dict() + assert response["label"][0] is not None + + test_feature_store.teardown() + + def test_apply_conflicting_feature_view_names(feature_store_with_local_registry): """Test applying feature views with non-case-insensitively unique names""" driver = Entity(name="driver", join_keys=["driver_id"])