Skip to content

Commit b8989cb

Browse files
MohammadHijjawi97ntkathole
authored andcommitted
fix: Handle non-nanosecond timestamps in tiled aggregations
apply_sawtooth_window_tiling turned the timestamp column into milliseconds by reading its raw int64 values and dividing by 10**6, which is only right for datetime64[ns]. With a datetime64[us] column (what pyarrow gives pandas 2 for timestamp[us] data) every event fell into one hop and the tile timestamps landed in January 1970. Measure the time since the epoch instead, which works for any unit and for timezone-aware columns. Signed-off-by: Mohammad Hijjawi <mohammad.hijjawi1997@gmail.com>
1 parent 023614b commit b8989cb

2 files changed

Lines changed: 55 additions & 5 deletions

File tree

‎sdk/python/feast/aggregation/tiling/orchestrator.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,12 @@ def apply_sawtooth_window_tiling(
4949
# Step 1: Add hop interval column
5050
hop_size_ms = int(hop_size.total_seconds() * 1000)
5151

52-
# Convert timestamp to milliseconds
53-
if pd.api.types.is_datetime64_any_dtype(df[timestamp_col]):
54-
timestamp_ms = df[timestamp_col].astype("int64") // 10**6
55-
else:
56-
timestamp_ms = pd.to_datetime(df[timestamp_col]).astype("int64") // 10**6
52+
# Convert timestamp to milliseconds. Measure from the epoch instead of
53+
# reading the raw integers, which are only nanoseconds for datetime64[ns]
54+
# (Arrow-backed frames are often datetime64[us]).
55+
timestamps = pd.to_datetime(df[timestamp_col])
56+
epoch = pd.Timestamp(0, tz=timestamps.dt.tz)
57+
timestamp_ms = (timestamps - epoch) // pd.Timedelta(milliseconds=1)
5758

5859
# Compute hop interval (inclusive lower boundaries)
5960
df["_hop_interval"] = (timestamp_ms // hop_size_ms) * hop_size_ms

‎sdk/python/tests/unit/test_aggregation_ops.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,3 +239,52 @@ def test_tiled_max_min_aggregation_matches_pandas(function):
239239
& (df["event_timestamp"] < tile_end)
240240
]["trips"]
241241
assert row[feature_name] == getattr(in_window, function)()
242+
243+
244+
@pytest.mark.parametrize("dtype", ["datetime64[us]", "datetime64[ms, UTC]"])
245+
def test_tiled_aggregation_does_not_depend_on_timestamp_resolution(dtype):
246+
"""Timestamps stored in a non-nanosecond unit give the same tiles."""
247+
from feast.aggregation.tiling.orchestrator import apply_sawtooth_window_tiling
248+
from feast.aggregation.tiling.tile_subtraction import (
249+
convert_cumulative_to_windowed,
250+
)
251+
252+
start = pd.Timestamp("2024-01-01")
253+
minutes = [0, 5, 10, 30, 35]
254+
df = pd.DataFrame(
255+
{
256+
"driver_id": [1] * len(minutes),
257+
"event_timestamp": [start + timedelta(minutes=m) for m in minutes],
258+
"trips": [1.0, 2.0, 3.0, 4.0, 5.0],
259+
}
260+
)
261+
window = timedelta(minutes=20)
262+
agg = Aggregation(column="trips", function="sum", time_window=window)
263+
264+
def windowed(frame):
265+
tiles = apply_sawtooth_window_tiling(
266+
df=frame,
267+
aggregations=[agg],
268+
group_by_keys=["driver_id"],
269+
timestamp_col="event_timestamp",
270+
window_size=window,
271+
hop_size=timedelta(minutes=5),
272+
)
273+
return convert_cumulative_to_windowed(
274+
tiles_df=tiles,
275+
entity_keys=["driver_id"],
276+
timestamp_col="event_timestamp",
277+
window_size=window,
278+
aggregations=[agg],
279+
)
280+
281+
expected = windowed(df.copy())
282+
ts = df["event_timestamp"]
283+
if "UTC" in dtype:
284+
ts = ts.dt.tz_localize("UTC")
285+
result = windowed(df.assign(event_timestamp=ts.astype(dtype)))
286+
287+
feature_name = agg.resolved_name(window)
288+
assert len(expected) == 8
289+
assert result["event_timestamp"].tolist() == expected["event_timestamp"].tolist()
290+
assert result[feature_name].tolist() == expected[feature_name].tolist()

0 commit comments

Comments
 (0)