fix: Preserve distinct entity_df rows sharing a join key and timestamp in Dask offline store - #6786
Open
piyush182004 wants to merge 1 commit into
Open
Conversation
…p in Dask offline store get_historical_features() on the Dask/File offline store returned fewer rows than were passed in whenever the entity_df contained two or more rows that shared a join key and event timestamp but differed in other columns - for example two orders placed by the same customer in the same logged second, each with its own order id and label. _drop_duplicates() deduplicated the joined result on (join keys, event timestamp) only. Since that pair does not uniquely identify an *input* row, distinct entity_df rows sharing it were collapsed into one by drop_duplicates(keep="last"), silently discarding every other column of the dropped rows - including labels a training pipeline would join on. Reproduced with a 4-row entity_df (two rows sharing driver_id and timestamp): before this change, only 2 of 4 rows survived. The DuckDB offline store, and every SQL-based store (BigQuery, Snowflake, Redshift, Postgres, Spark, ClickHouse, Trino, Couchbase, Athena), already avoid this by carrying a per-row "entity_row_unique_id" through their generated queries; Dask/File had no equivalent. The fix tags each input row with a synthetic unique id before the per-feature-view join, dedupes on that id instead of (join keys, timestamp), and drops the id column before returning. The non-entity mode (entity_df=None, used to snapshot all entities as of a timestamp) is unaffected: it still dedupes on (join keys, timestamp) since there is a single synthetic entity_df row shared by every real entity in that mode. Related but distinct from feast-dev#3360 (closed wontfix), which is about tie-breaking between 3+ created_timestamp versions on the *feature source* side; this is about the *entity_df* (request) side never being fanned out correctly in the first place. Signed-off-by: Piyush Mondal <piyushmondal182004@gmail.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Preserves distinct entity rows during Dask historical feature retrieval.
Changes:
- Adds synthetic row IDs for deduplication.
- Adds regression tests for duplicate join keys and timestamps.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
sdk/python/feast/infra/offline_stores/dask.py |
Implements row-ID-based deduplication. |
sdk/python/tests/unit/infra/offline_stores/test_dask_entity_row_dedup.py |
Tests row preservation and counts. |
Suppressed comments (1)
sdk/python/feast/infra/offline_stores/dask.py:371
- This unconditionally deletes a caller column named
__entity_row_unique_id__: line 229 first overwrites that column, and this cleanup then removes it from the result. Since entity dataframe payload columns are expected to survive retrieval, choose a collision-free temporary name (or preserve and restore the original column) and propagate that name through deduplication.
entity_df_with_features = entity_df_with_features.drop(
columns=[_ENTITY_ROW_ID_COL]
)
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+228
to
+229
| 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 |
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #6786 +/- ##
=======================================
Coverage 47.08% 47.09%
=======================================
Files 419 419
Lines 51878 51883 +5
Branches 7525 7525
=======================================
+ Hits 24429 24434 +5
Misses 25700 25700
Partials 1749 1749
Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
ISSUE:- #6787
What this PR does / why we need it:
get_historical_features()on the Dask/File offline store returned fewer rows than were passed in whenever theentity_dfcontained two or more rows that shared a join key and event timestamp but differed in other columns — for example, two orders placed by the same customer in the same logged second, each with its own order id and label._drop_duplicates()deduplicated the joined result on(join keys, event timestamp)only. Since that pair does not uniquely identify an input row, distinctentity_dfrows sharing it were collapsed into one bydrop_duplicates(keep="last"), silently discarding every other column of the dropped rows — including labels a training pipeline would join on.Reproduced with a 4-row
entity_df(two rows sharing a join key and timestamp): before this change, only 2 of 4 rows survived. The DuckDB offline store, and every SQL-based store (BigQuery, Snowflake, Redshift, Postgres, Spark, ClickHouse, Trino, Couchbase, Athena), already avoid this by carrying a per-rowentity_row_unique_idthrough their generated queries; Dask/File had no equivalent.The fix tags each input row with a synthetic unique id before the per-feature-view join, dedupes on that id instead of
(join keys, timestamp), and drops the id column before returning. Non-entity mode (entity_df=None, used to snapshot all entities as of a timestamp) is unaffected and covered by an existing test.Verification
Full offline-store + local-feast-tests suite: 260 passed, 0 failed. Added 2 regression tests covering the collapse scenario and the exact row-count invariant.
Which issue(s) this PR fixes:
Checks
git commit -s)Testing Strategy