fix: Normalize Dask timestamps to UTC without a row-wise apply - #6818
Open
Daksha1611 wants to merge 1 commit into
Open
fix: Normalize Dask timestamps to UTC without a row-wise apply#6818Daksha1611 wants to merge 1 commit into
Daksha1611 wants to merge 1 commit into
Conversation
_normalize_timestamp made timestamp columns tz-aware with a row-wise apply
declaring meta="datetime64[ns, UTC]". meta only declares the dtype; the
partition gets whatever the lambda returns. With zero rows the lambda never
runs, so the computed column stayed tz-naive while meta claimed UTC, and the
tz-aware comparison in _filter_ttl then raised:
TypeError: Invalid comparison between dtype=datetime64[ns] and DatetimeArray
A zero-row entity_df is a normal case in batch scoring, when the upstream
query matched nothing for that run, so get_historical_features crashed rather
than returning an empty result.
A column carrying a non-UTC timezone diverged from the declared meta the same
way, because the lambda returned any value that already had tzinfo untouched.
Replace both applies with dd.to_datetime(..., utc=True), which localizes
tz-naive values, converts tz-aware ones, and yields the correct dtype for an
empty frame. It is also vectorized, so it avoids a Python-level call per row
on non-empty frames.
Signed-off-by: Daksha1611 <mehtadaksha1611@gmail.com>
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.
What this PR does / why we need it:
_normalize_timestampmade timestamp columns tz-aware with a row-wiseapplythatdeclared
meta=(timestamp_field, "datetime64[ns, UTC]").metaonly declares thedtype — the partition gets whatever the lambda returns. With zero rows the lambda never
runs, so the computed column stayed tz-naive while
metaclaimed UTC, and the tz-awarecomparison in
_filter_ttlthen raised:A zero-row
entity_dfis a normal case in batch scoring — the upstream query matchednothing for that run — so
get_historical_featurescrashed instead of returning anempty result.
While fixing this I found the same mismatch on a second path: a column carrying a
non-UTC timezone also diverged from the declared
meta, because the lambda returned anyvalue that already had
tzinfountouched. That is covered by a test here too.Both applies are replaced with
dd.to_datetime(..., utc=True), which localizes tz-naivevalues, converts tz-aware ones, and yields the correct dtype for an empty frame. It is
also vectorized, so it drops a Python-level call per row on non-empty frames.
Which issue(s) this PR fixes:
Fixes #6817
Checks
git commit -s)Testing Strategy
New
sdk/python/tests/unit/infra/offline_stores/test_dask_empty_entity_df.pycoversnormalization and TTL filtering at 0 and 1 rows, plus the non-UTC case. Three of its
five assertions fail on
masterand pass here; the two 1-row cases pass both ways,pinning the non-empty behaviour as unchanged.
Verified end to end as well:
get_historical_features(...).to_df()against afileoffline store with an otherwise identical entity frame raised
TypeErrorat 0 rowsbefore this change, and now returns an empty frame with the expected columns while the
1-row case is unchanged.
sdk/python/tests/unit/infra/offline_storesplustest_unit_feature_store.py: 235passed. The 20 errors in that run are MongoDB testcontainers failing to start locally
and are present on
mastertoo.ruff check,ruff format --checkandmypyare clean on both changed files.Misc
Point-in-time correctness is unaffected — I re-ran the
filter_by_created_timestampscenario (a row backfilled after the entity timestamp) and it still excludes the late
row with the flag and includes it without.