Skip to content

Commit a87c070

Browse files
MohammadHijjawi97ntkathole
MohammadHijjawi97
authored and
ntkathole
committed
fix: Keep entity rows apart in the ibis point-in-time join
The ibis point-in-time join (used by the DuckDB, MSSQL and Oracle offline stores) identifies each entity row by concatenating the string values of its join keys and event timestamp. Different rows could produce the same id, e.g. driver_id="1", customer_id="11" and driver_id="11", customer_id="1", so one row got the other's feature values. A null key also made the whole id null, so the row lost the features of every other feature view, not just the one keyed on the null column. Length-prefix each value and give nulls their own marker so the id is only shared by rows with equal values. Signed-off-by: Mohammad Hijjawi <[email protected]>
1 parent ce37029 commit a87c070

2 files changed

Lines changed: 84 additions & 13 deletions

File tree

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

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import uuid
44
from datetime import datetime, timedelta, timezone
55
from pathlib import Path
6-
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
6+
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
77

88
import ibis
99
import numpy as np
@@ -120,6 +120,27 @@ def _to_utc(entity_df: pd.DataFrame, event_timestamp_col):
120120
return entity_df
121121

122122

123+
def _entity_row_id(entity_table: Table, columns: Iterable[str]):
124+
"""Build an id that is the same only for rows with equal values in columns.
125+
126+
Each value is length-prefixed so that different rows cannot concatenate to
127+
the same id (e.g. "1" + "11" and "11" + "1"), and a null value gets its own
128+
marker instead of turning the whole id into null.
129+
"""
130+
r = ibis.literal("_")
131+
132+
for e in columns:
133+
value = entity_table[e].cast("string")
134+
r = r.concat(
135+
value.isnull().ifelse(
136+
"N",
137+
value.length().cast("string").concat(":", value),
138+
)
139+
)
140+
141+
return r
142+
143+
123144
def _generate_row_id(
124145
entity_table: Table, feature_views: List[FeatureView], event_timestamp_col
125146
) -> Table:
@@ -130,12 +151,9 @@ def _generate_row_id(
130151
else:
131152
all_entities.extend([e.name for e in fv.entity_columns])
132153

133-
r = ibis.literal("_")
134-
135-
for e in set(all_entities):
136-
r = r.concat(entity_table[e].cast("string")) # type: ignore
137-
138-
entity_table = entity_table.mutate(entity_row_id=r)
154+
entity_table = entity_table.mutate(
155+
entity_row_id=_entity_row_id(entity_table, set(all_entities))
156+
)
139157

140158
return entity_table
141159

@@ -400,12 +418,9 @@ def point_in_time_join(
400418
) in feature_tables:
401419
all_entities.extend(join_key_map.values())
402420

403-
r = ibis.literal("_")
404-
405-
for e in set(all_entities):
406-
r = r.concat(entity_table[e].cast("string")) # type: ignore
407-
408-
entity_table = entity_table.mutate(entity_row_id=r)
421+
entity_table = entity_table.mutate(
422+
entity_row_id=_entity_row_id(entity_table, set(all_entities))
423+
)
409424

410425
acc_table = entity_table
411426

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,3 +170,59 @@ def test_point_in_time_join():
170170
).to_pyarrow()
171171

172172
assert tables_equal_ignore_order(actual, expected)
173+
174+
175+
def test_point_in_time_join_keeps_entity_rows_apart():
176+
"""Entity rows whose key values concatenate to the same string, or that have
177+
a null key for another feature view, still get their own feature values."""
178+
entity_table = pa.table(
179+
{
180+
"driver_id": ["1", "11", "7"],
181+
"customer_id": ["11", "1", None],
182+
"event_timestamp": [pa_datetime(2024, 1, 2)] * 3,
183+
}
184+
)
185+
drivers = pa.table(
186+
{
187+
"driver_id": ["1", "11", "7"],
188+
"event_timestamp": [pa_datetime(2024, 1, 1)] * 3,
189+
"created": [pa_datetime(2024, 1, 1)] * 3,
190+
"driver_rating": [1.0, 11.0, 7.0],
191+
}
192+
)
193+
customers = pa.table(
194+
{
195+
"customer_id": ["11", "1"],
196+
"event_timestamp": [pa_datetime(2024, 1, 1)] * 2,
197+
"created": [pa_datetime(2024, 1, 1)] * 2,
198+
"customer_score": [110.0, 10.0],
199+
}
200+
)
201+
202+
actual = point_in_time_join(
203+
ibis.memtable(entity_table),
204+
feature_tables=[
205+
(
206+
ibis.memtable(drivers),
207+
"event_timestamp",
208+
"created",
209+
{"driver_id": "driver_id"},
210+
["driver_rating"],
211+
timedelta(days=10),
212+
),
213+
(
214+
ibis.memtable(customers),
215+
"event_timestamp",
216+
"created",
217+
{"customer_id": "customer_id"},
218+
["customer_score"],
219+
timedelta(days=10),
220+
),
221+
],
222+
).to_pyarrow()
223+
224+
rows = {
225+
row["driver_id"]: (row["driver_rating"], row["customer_score"])
226+
for row in actual.to_pylist()
227+
}
228+
assert rows == {"1": (1.0, 110.0), "11": (11.0, 10.0), "7": (7.0, None)}

0 commit comments

Comments
 (0)