Skip to content

Commit b1f2168

Browse files
sohailshkshuchu
andauthored
fix: Use date-only bounds for partition filters built from string timestamps (#7116)
get_timestamp_filter_sql formats datetime bounds as dates for the date_partition_column predicate, but returned string bounds unchanged. Athena's pull_all_from_table_or_query passes strings such as '2026-02-01 01:00:00.000', so the partition predicate became dt >= '2026-02-01 01:00:00.000'. Against a 'YYYY-MM-DD' partition value that compares false, which drops the whole first day. Truncate string bounds to their date part. Signed-off-by: sohailshk <[email protected]> Co-authored-by: Shuchu Han <[email protected]>
1 parent 8fdcee7 commit b1f2168

2 files changed

Lines changed: 56 additions & 1 deletion

File tree

‎sdk/python/feast/infra/offline_stores/offline_utils.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -412,7 +412,8 @@ def format_date(val: Union[str, datetime]) -> str:
412412
if tz:
413413
val = val.astimezone(tz)
414414
return val.strftime("%Y-%m-%d")
415-
return val
415+
# Partition columns hold dates; drop any time-of-day from ISO strings.
416+
return val[:10]
416417

417418
ts_field = quote_column_if_needed(timestamp_field)
418419
dp_field = quote_column_if_needed(date_partition_column)
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
from datetime import datetime, timezone
2+
from unittest.mock import MagicMock, patch
3+
4+
from feast.infra.offline_stores.contrib.athena_offline_store.athena import (
5+
AthenaOfflineStore,
6+
AthenaOfflineStoreConfig,
7+
)
8+
from feast.infra.offline_stores.contrib.athena_offline_store.athena_source import (
9+
AthenaSource,
10+
)
11+
from feast.infra.utils import aws_utils
12+
from feast.repo_config import RepoConfig
13+
14+
15+
@patch.object(aws_utils, "get_s3_resource", MagicMock())
16+
@patch.object(aws_utils, "get_athena_data_client", MagicMock())
17+
def test_pull_all_partition_filter_uses_dates_only():
18+
config = RepoConfig(
19+
registry="registry",
20+
project="project",
21+
provider="local",
22+
offline_store=AthenaOfflineStoreConfig(
23+
type="athena",
24+
data_source="AwsDataCatalog",
25+
region="us-west-2",
26+
database="db",
27+
workgroup="wg",
28+
s3_staging_location="s3://bucket/path",
29+
),
30+
entity_key_serialization_version=3,
31+
)
32+
source = AthenaSource(
33+
name="source",
34+
table="table",
35+
timestamp_field="ts",
36+
date_partition_column="dt",
37+
)
38+
39+
job = AthenaOfflineStore.pull_all_from_table_or_query(
40+
config=config,
41+
data_source=source,
42+
join_key_columns=["id"],
43+
feature_name_columns=["f"],
44+
timestamp_field="ts",
45+
start_date=datetime(2026, 2, 1, 1, tzinfo=timezone.utc),
46+
end_date=datetime(2026, 2, 2, 2, tzinfo=timezone.utc),
47+
)
48+
49+
with job._query_generator() as query:
50+
pass
51+
assert "ts BETWEEN TIMESTAMP '2026-02-01 01:00:00.000'" in query
52+
# The partition column holds dates, so the bounds must not carry a time of
53+
# day: '2026-02-01' < '2026-02-01 01:00:00.000' would drop the first day.
54+
assert "dt >= '2026-02-01' AND dt <= '2026-02-02'" in query

0 commit comments

Comments
 (0)