diff --git a/sdk/python/feast/infra/offline_stores/dask.py b/sdk/python/feast/infra/offline_stores/dask.py index cd930f7ab5c..2e98a63d7f9 100644 --- a/sdk/python/feast/infra/offline_stores/dask.py +++ b/sdk/python/feast/infra/offline_stores/dask.py @@ -58,6 +58,13 @@ # See (https://github.com/dask/dask/issues/10881#issuecomment-1923327936) dask.config.set({"dataframe.convert-string": False}) +# Identifies each input entity_df row through the join/dedup pipeline below, so that +# rows sharing the same join key(s) and event timestamp - but differing in other +# entity_df columns - are not collapsed into one by _drop_duplicates. SQL-based +# offline stores (BigQuery, Snowflake, Redshift, Postgres, ...) avoid this the same +# way, via a per-row "entity_row_unique_id" carried through their generated queries. +_ENTITY_ROW_ID_COL = "__entity_row_unique_id__" + class DaskOfflineStoreConfig(FeastConfigBaseModel): """Offline store config for dask store""" @@ -214,6 +221,13 @@ def evaluate_historical_retrieval(): # Create a copy of entity_df to prevent modifying the original entity_df_with_features = entity_df.copy() + # Tag each input row with a unique id before it can be fanned out by the + # per-feature-view join below, so that distinct entity_df rows sharing a + # join key and event timestamp are never collapsed together by + # _drop_duplicates (see _ENTITY_ROW_ID_COL). + entity_df_with_features = entity_df_with_features.reset_index(drop=True) + entity_df_with_features[_ENTITY_ROW_ID_COL] = entity_df_with_features.index + entity_df_event_timestamp_col_type = entity_df_with_features.dtypes[ entity_df_event_timestamp_col ] @@ -338,6 +352,11 @@ def evaluate_historical_retrieval(): timestamp_field, created_timestamp_column, entity_df_event_timestamp_col, + # In non-entity mode there is one synthetic entity_df row shared + # by every real entity fanned out from the feature source, so + # dedup must stay keyed on (join key, timestamp) - not the + # single shared row id. + None if non_entity_mode else _ENTITY_ROW_ID_COL, ) entity_df_with_features = _drop_columns( @@ -347,6 +366,9 @@ def evaluate_historical_retrieval(): # Ensure that we delete dataframes to free up memory del df_to_join + entity_df_with_features = entity_df_with_features.drop( + columns=[_ENTITY_ROW_ID_COL] + ) return entity_df_with_features.persist() job = DaskRetrievalJob( @@ -1248,6 +1270,7 @@ def _drop_duplicates( timestamp_field: str, created_timestamp_column: str, entity_df_event_timestamp_col: str, + entity_row_id_col: Optional[str] = None, ) -> dd.DataFrame: column_order = df_to_join.columns @@ -1277,8 +1300,17 @@ def _drop_duplicates( ) df_to_join = df_to_join.persist() + # Deduplicate per original entity_df row (entity_row_id_col) rather than per + # (join key, event timestamp): two distinct input rows can legitimately share + # both, and must each keep their own matched feature values instead of being + # collapsed into one (see _ENTITY_ROW_ID_COL). + dedup_subset = ( + [entity_row_id_col] + if entity_row_id_col and entity_row_id_col in df_to_join.columns + else all_join_keys + [entity_df_event_timestamp_col] + ) df_to_join = df_to_join.drop_duplicates( - all_join_keys + [entity_df_event_timestamp_col], + dedup_subset, keep="last", ignore_index=True, ) diff --git a/sdk/python/tests/unit/infra/offline_stores/test_dask_entity_row_dedup.py b/sdk/python/tests/unit/infra/offline_stores/test_dask_entity_row_dedup.py new file mode 100644 index 00000000000..3e059eaadae --- /dev/null +++ b/sdk/python/tests/unit/infra/offline_stores/test_dask_entity_row_dedup.py @@ -0,0 +1,137 @@ +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock + +import dask.dataframe as dd +import pandas as pd + +from feast.entity import Entity +from feast.feature_view import FeatureView, Field +from feast.infra.offline_stores import dask as dask_mod +from feast.infra.offline_stores.dask import DaskOfflineStore, DaskOfflineStoreConfig +from feast.infra.offline_stores.file_source import FileSource +from feast.repo_config import RepoConfig +from feast.types import Float32, ValueType + + +def _mock_entity(): + return [ + Entity( + name="driver_id", + join_keys=["driver_id"], + value_type=ValueType.INT64, + ) + ] + + +def _mock_feature_view(): + return FeatureView( + name="driver_stats", + entities=_mock_entity(), + schema=[Field(name="conv_rate", dtype=Float32)], + source=FileSource(path="unused", timestamp_field="event_timestamp"), + ttl=timedelta(days=1), + ) + + +def _mock_repo_config(): + return RepoConfig( + project="proj", + registry="unused", + provider="local", + offline_store=DaskOfflineStoreConfig(type="dask"), + ) + + +class TestEntityRowDeduplication: + """ + get_historical_features must return exactly one output row per input + entity_df row. Two distinct entity_df rows sharing a join key and event + timestamp - e.g. two orders placed by the same customer in the same + logged second - are not duplicates of each other, even though the + feature-source join produces the same feature values for both. + + Regression test for a bug where such rows were collapsed by + _drop_duplicates, since its dedup key (join keys + event timestamp) did + not account for other, distinct entity_df columns. + """ + + def test_distinct_entity_rows_are_not_collapsed(self, monkeypatch): + ts = datetime(2024, 1, 1, tzinfo=timezone.utc) + src = pd.DataFrame( + { + "driver_id": [1], + "event_timestamp": [ts - timedelta(days=1)], + "conv_rate": [0.5], + } + ) + monkeypatch.setattr( + dask_mod, + "_read_datasource", + lambda ds, repo_path: dd.from_pandas(src, npartitions=1), + ) + + # Three distinct requests: two share (driver_id=1, ts), one is a + # different entity. Each carries a unique label that must survive. + entity_df = pd.DataFrame( + { + "driver_id": [1, 1, 2], + "event_timestamp": [ts, ts, ts], + "request_id": ["req-a", "req-b", "req-c"], + } + ) + + job = DaskOfflineStore.get_historical_features( + config=_mock_repo_config(), + feature_views=[_mock_feature_view()], + feature_refs=["driver_stats:conv_rate"], + entity_df=entity_df, + registry=MagicMock(), + project="proj", + full_feature_names=False, + ) + result = job.to_df() + + assert len(result) == len(entity_df) + assert sorted(result["request_id"]) == ["req-a", "req-b", "req-c"] + # driver_id=1's two distinct requests both get the feature value. + assert set(result.loc[result["driver_id"] == 1, "conv_rate"]) == {0.5} + + def test_row_count_matches_input_with_duplicated_join_key_and_timestamp( + self, monkeypatch + ): + """Same as above, phrased as a direct row-count invariant.""" + ts = datetime(2024, 1, 1, tzinfo=timezone.utc) + src = pd.DataFrame( + { + "driver_id": [1, 2], + "event_timestamp": [ts - timedelta(days=1)] * 2, + "conv_rate": [0.1, 0.2], + } + ) + monkeypatch.setattr( + dask_mod, + "_read_datasource", + lambda ds, repo_path: dd.from_pandas(src, npartitions=1), + ) + + entity_df = pd.DataFrame( + { + "driver_id": [1, 1, 1, 2], + "event_timestamp": [ts, ts, ts, ts], + "label": ["a", "b", "c", "d"], + } + ) + + job = DaskOfflineStore.get_historical_features( + config=_mock_repo_config(), + feature_views=[_mock_feature_view()], + feature_refs=["driver_stats:conv_rate"], + entity_df=entity_df, + registry=MagicMock(), + project="proj", + full_feature_names=False, + ) + result = job.to_df() + + assert len(result) == 4 + assert sorted(result["label"]) == ["a", "b", "c", "d"]