Skip to content

Commit eb93b08

Browse files
LuisFigueroaGshuchu
LuisFigueroaG
and
shuchu
authored
fix: Break event timestamp ties by created timestamp in the Dask offline store (#6952)
The Dask offline store sorted by created_timestamp_column and then by timestamp_field in a second pass, relying on that second sort to keep the created timestamp order for rows sharing an event timestamp. sort_values is not stable, so with enough rows the older row could end up last and win the keep="last" deduplication, both in historical retrieval and in pull_latest_from_table_or_query. Sort once by [timestamp_field, created_timestamp_column] instead. Signed-off-by: LuisFigueroaG <[email protected]> Co-authored-by: Shuchu Han <[email protected]>
1 parent d520584 commit eb93b08

2 files changed

Lines changed: 131 additions & 26 deletions

File tree

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

Lines changed: 16 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -476,23 +476,19 @@ def evaluate_offline_job(
476476

477477
# try-catch block is added to deal with this issue https://github.com/dask/dask/issues/8939.
478478
# TODO(kevjumba): remove try catch when fix is merged upstream in Dask.
479-
try:
480-
if created_timestamp_column:
481-
source_df = source_df.sort_values(
482-
by=created_timestamp_column,
483-
)
479+
# Sort by both columns at once: sort_values is not stable, so a separate
480+
# sort by created_timestamp_column would not reliably break ties.
481+
sort_columns = [timestamp_field]
482+
if created_timestamp_column:
483+
sort_columns.append(created_timestamp_column)
484484

485-
source_df = source_df.sort_values(by=timestamp_field)
485+
try:
486+
source_df = source_df.sort_values(by=sort_columns)
486487

487488
except ZeroDivisionError:
488489
# Use 1 partition to get around case where everything in timestamp column is the same so the partition algorithm doesn't
489490
# try to divide by zero.
490-
if created_timestamp_column:
491-
source_df = source_df.sort_values(
492-
by=created_timestamp_column, npartitions=1
493-
)
494-
495-
source_df = source_df.sort_values(by=timestamp_field, npartitions=1)
491+
source_df = source_df.sort_values(by=sort_columns, npartitions=1)
496492

497493
# TODO: The old implementation is inclusive of start_date and exclusive of end_date.
498494
# Which is inconsistent with other offline stores.
@@ -1281,27 +1277,21 @@ def _drop_duplicates(
12811277

12821278
# try-catch block is added to deal with this issue https://github.com/dask/dask/issues/8939.
12831279
# TODO(kevjumba): remove try catch when fix is merged upstream in Dask.
1284-
try:
1285-
if created_timestamp_column:
1286-
df_to_join = df_to_join.sort_values(
1287-
by=created_timestamp_column, na_position="first"
1288-
)
1289-
df_to_join = df_to_join.persist()
1280+
# Sort by both columns at once: sort_values is not stable, so a separate
1281+
# sort by created_timestamp_column would not reliably break ties.
1282+
sort_columns = [timestamp_field]
1283+
if created_timestamp_column:
1284+
sort_columns.append(created_timestamp_column)
12901285

1291-
df_to_join = df_to_join.sort_values(by=timestamp_field, na_position="first")
1286+
try:
1287+
df_to_join = df_to_join.sort_values(by=sort_columns, na_position="first")
12921288
df_to_join = df_to_join.persist()
12931289

12941290
except ZeroDivisionError:
12951291
# Use 1 partition to get around case where everything in timestamp column is the same so the partition algorithm doesn't
12961292
# try to divide by zero.
1297-
if created_timestamp_column:
1298-
df_to_join = df_to_join[column_order].sort_values(
1299-
by=created_timestamp_column, na_position="first", npartitions=1
1300-
)
1301-
df_to_join = df_to_join.persist()
1302-
13031293
df_to_join = df_to_join[column_order].sort_values(
1304-
by=timestamp_field, na_position="first", npartitions=1
1294+
by=sort_columns, na_position="first", npartitions=1
13051295
)
13061296
df_to_join = df_to_join.persist()
13071297

‎sdk/python/tests/unit/infra/offline_stores/test_dask_created_timestamp_tiebreak.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
from datetime import timedelta
2+
from unittest.mock import MagicMock
3+
4+
import dask.dataframe as dd
5+
import pandas as pd
6+
7+
from feast.entity import Entity
8+
from feast.feature_view import FeatureView, Field
9+
from feast.infra.offline_stores import dask as dask_mod
10+
from feast.infra.offline_stores.dask import DaskOfflineStore, DaskOfflineStoreConfig
11+
from feast.infra.offline_stores.file_source import FileSource
12+
from feast.repo_config import RepoConfig
13+
from feast.types import Float32, Int64, ValueType
14+
15+
# Enough rows spread over several partitions that an unstable sort reorders
16+
# rows sharing an event timestamp; small inputs tend to keep their order.
17+
NUM_DRIVERS = 100
18+
NUM_PARTITIONS = 4
19+
20+
EVENT_TS = pd.Timestamp("2025-01-01T10:00:00Z")
21+
ENTITY_TS = pd.Timestamp("2025-01-02T00:00:00Z")
22+
23+
24+
def _source_df() -> pd.DataFrame:
25+
# Each driver has two rows with the same event timestamp. The row created
26+
# later carries the correction and must win.
27+
driver_ids = list(range(NUM_DRIVERS))
28+
event_ts = [EVENT_TS + timedelta(minutes=i % 10) for i in driver_ids]
29+
return pd.DataFrame(
30+
{
31+
"driver_id": driver_ids + driver_ids,
32+
"event_timestamp": event_ts + event_ts,
33+
"created_ts": event_ts + [ts + timedelta(hours=1) for ts in event_ts],
34+
"conv_rate": [0.0] * NUM_DRIVERS + [1.0] * NUM_DRIVERS,
35+
}
36+
)
37+
38+
39+
def _use_source_df(monkeypatch):
40+
ddf = dd.from_pandas(_source_df(), npartitions=NUM_PARTITIONS)
41+
monkeypatch.setattr(dask_mod, "_read_datasource", lambda ds, repo_path: ddf)
42+
43+
44+
def _config() -> RepoConfig:
45+
return RepoConfig(
46+
project="test_project",
47+
registry="test_registry",
48+
provider="local",
49+
offline_store=DaskOfflineStoreConfig(type="dask"),
50+
)
51+
52+
53+
def _file_source() -> FileSource:
54+
return FileSource(
55+
path="dummy.parquet", # not read in this test
56+
timestamp_field="event_timestamp",
57+
created_timestamp_column="created_ts",
58+
)
59+
60+
61+
def test_historical_retrieval_breaks_event_timestamp_ties_by_created_timestamp(
62+
monkeypatch,
63+
):
64+
_use_source_df(monkeypatch)
65+
fv = FeatureView(
66+
name="driver_stats",
67+
entities=[
68+
Entity(
69+
name="driver_id",
70+
join_keys=["driver_id"],
71+
value_type=ValueType.INT64,
72+
)
73+
],
74+
schema=[Field(name="conv_rate", dtype=Float32)],
75+
source=_file_source(),
76+
ttl=timedelta(days=7),
77+
)
78+
fv.entity_columns = [Field(name="driver_id", dtype=Int64)]
79+
registry = MagicMock()
80+
registry.list_on_demand_feature_views.return_value = []
81+
entity_df = pd.DataFrame(
82+
{
83+
"driver_id": list(range(NUM_DRIVERS)),
84+
"event_timestamp": [ENTITY_TS] * NUM_DRIVERS,
85+
}
86+
)
87+
88+
df = DaskOfflineStore.get_historical_features(
89+
config=_config(),
90+
feature_views=[fv],
91+
feature_refs=["driver_stats:conv_rate"],
92+
entity_df=entity_df,
93+
registry=registry,
94+
project="test_project",
95+
full_feature_names=False,
96+
).to_df()
97+
98+
assert len(df) == NUM_DRIVERS
99+
assert df["conv_rate"].tolist() == [1.0] * NUM_DRIVERS
100+
101+
102+
def test_pull_latest_breaks_event_timestamp_ties_by_created_timestamp(monkeypatch):
103+
_use_source_df(monkeypatch)
104+
105+
df = DaskOfflineStore.pull_latest_from_table_or_query(
106+
config=_config(),
107+
data_source=_file_source(),
108+
join_key_columns=["driver_id"],
109+
feature_name_columns=["conv_rate"],
110+
timestamp_field="event_timestamp",
111+
created_timestamp_column="created_ts",
112+
).to_df()
113+
114+
assert len(df) == NUM_DRIVERS
115+
assert df["conv_rate"].tolist() == [1.0] * NUM_DRIVERS

0 commit comments

Comments
 (0)