Skip to content

Commit 558e864

Browse files
YassinNouh21ntkathole
authored andcommitted
fix(postgres): Use end_date instead of start_date in synthetic entity_df for non-entity retrieval
The non-entity retrieval path created a synthetic entity_df using pd.date_range(start=start_date, ...)[:1], which placed start_date as the event_timestamp. Since PIT joins use MAX(entity_timestamp) as the upper bound for feature data filtering, using start_date made end_date unreachable — no features after start_date would be returned. Fix: use [end_date] directly, matching the ClickHouse implementation (PR feast-dev#6066) and the Dask offline store behavior. Signed-off-by: yassinnouh21 <yassinnouh21@gmail.com>
1 parent 6b31a43 commit 558e864

2 files changed

Lines changed: 61 additions & 7 deletions

File tree

  • sdk/python
    • feast/infra/offline_stores/contrib/postgres_offline_store
    • tests/unit/infra/offline_stores/contrib/postgres_offline_store

‎sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -159,13 +159,7 @@ def get_historical_features(
159159
else:
160160
start_date = make_tzaware(start_date)
161161

162-
entity_df = pd.DataFrame(
163-
{
164-
"event_timestamp": pd.date_range(
165-
start=start_date, end=end_date, freq="1s", tz=timezone.utc
166-
)[:1] # Just one row
167-
}
168-
)
162+
entity_df = pd.DataFrame({"event_timestamp": [end_date]})
169163

170164
entity_schema = _get_entity_schema(entity_df, config)
171165

‎sdk/python/tests/unit/infra/offline_stores/contrib/postgres_offline_store/test_postgres.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -615,6 +615,66 @@ def test_non_entity_mode_with_both_dates(self):
615615
assert "start_date" not in str(e)
616616
assert "end_date" not in str(e)
617617

618+
def test_non_entity_entity_df_uses_end_date(self):
619+
"""Test that the synthetic entity_df uses end_date, not start_date.
620+
621+
Regression test: the old code used pd.date_range(start=start_date, ...)[:1]
622+
which put start_date in the entity_df. Since PIT joins use
623+
MAX(entity_timestamp) as the upper bound, start_date made end_date
624+
unreachable. The fix uses [end_date] directly.
625+
"""
626+
test_repo_config = RepoConfig(
627+
project="test_project",
628+
registry="test_registry",
629+
provider="local",
630+
offline_store=_mock_offline_store_config(),
631+
)
632+
633+
feature_view = _mock_feature_view("test_fv", ttl=timedelta(days=1))
634+
start_date = datetime(2023, 1, 1, tzinfo=timezone.utc)
635+
end_date = datetime(2023, 1, 7, tzinfo=timezone.utc)
636+
637+
mock_get_entity_schema = MagicMock(
638+
return_value={"event_timestamp": "timestamp"}
639+
)
640+
641+
with patch.multiple(
642+
"feast.infra.offline_stores.contrib.postgres_offline_store.postgres",
643+
_get_conn=MagicMock(),
644+
_upload_entity_df=MagicMock(),
645+
_get_entity_schema=mock_get_entity_schema,
646+
_get_entity_df_event_timestamp_range=MagicMock(
647+
return_value=(start_date, end_date)
648+
),
649+
), patch(
650+
"feast.infra.offline_stores.contrib.postgres_offline_store.postgres.offline_utils.get_expected_join_keys",
651+
return_value=[],
652+
), patch(
653+
"feast.infra.offline_stores.contrib.postgres_offline_store.postgres.offline_utils.assert_expected_columns_in_entity_df",
654+
), patch(
655+
"feast.infra.offline_stores.contrib.postgres_offline_store.postgres.offline_utils.get_feature_view_query_context",
656+
return_value=[],
657+
):
658+
PostgreSQLOfflineStore.get_historical_features(
659+
config=test_repo_config,
660+
feature_views=[feature_view],
661+
feature_refs=["test_fv:feature1"],
662+
entity_df=None,
663+
registry=MagicMock(),
664+
project="test_project",
665+
start_date=start_date,
666+
end_date=end_date,
667+
)
668+
669+
# _get_entity_schema is called with the synthetic entity_df
670+
df = mock_get_entity_schema.call_args[0][0]
671+
assert len(df) == 1
672+
ts = df["event_timestamp"].iloc[0]
673+
# The entity_df must use end_date, not start_date
674+
assert ts == end_date, (
675+
f"entity_df timestamp should be end_date ({end_date}), got {ts}"
676+
)
677+
618678
def test_non_entity_mode_with_end_date_only(self):
619679
"""Test non-entity retrieval calculates start_date from TTL"""
620680
test_repo_config = RepoConfig(

0 commit comments

Comments
 (0)