feast-dev · piyush182004 · Aug 26, 2026
@@ -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,
     )

@@ -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"]