Skip to content

Commit 213a130

Browse files
MohammadHijjawi97ntkathole
MohammadHijjawi97
authored and
ntkathole
committed
fix: Compute std and var correctly for tiled aggregations
With enable_tiling=True, the Spark and Ray compute engines build time-windowed aggregations from cumulative tiles. For std/stddev and var/variance the tiles already carry sum, count and sum of squares, but both apply_sawtooth_window_tiling and convert_cumulative_to_windowed computed the final value as sum / count, so these features were materialized as the window mean instead of the standard deviation or variance. Compute the final value for holistic aggregations in one helper: the mean for avg/mean, and the sample variance (and its square root for std) from the sum of squares, matching the non-tiled pandas/Spark results. Windows with fewer than two values yield NaN, as they do without tiling. Signed-off-by: Mohammad Hijjawi <[email protected]>
1 parent d7cb4d0 commit 213a130

4 files changed

Lines changed: 118 additions & 15 deletions

File tree

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
from dataclasses import dataclass
99
from typing import Any, List, Optional, Tuple
1010

11+
import numpy as np
12+
1113
from feast.aggregation import Aggregation
1214

1315

@@ -95,3 +97,34 @@ def get_ir_metadata_for_aggregation(
9597
else:
9698
# Unknown aggregation: treat as algebraic
9799
return ([], IRMetadata(type="algebraic"))
100+
101+
102+
def compute_holistic_value_from_irs(
103+
agg: Aggregation, ir_sum: Any, ir_count: Any, ir_sum_sq: Any = None
104+
) -> Any:
105+
"""
106+
Compute the final value of a holistic aggregation from its IR components.
107+
108+
Works on scalars as well as array-likes (e.g. pandas Series).
109+
110+
- avg/mean: sum / count (0 when count is 0)
111+
- var/variance: sample variance (sum_sq - sum^2 / count) / (count - 1)
112+
- std/stddev: square root of the sample variance
113+
Sample variance and standard deviation are NaN when count < 2, matching
114+
the non-tiled aggregation.
115+
"""
116+
agg_type = agg.function.lower()
117+
count = np.asarray(ir_count, dtype="float64")
118+
total = np.asarray(ir_sum, dtype="float64")
119+
120+
with np.errstate(divide="ignore", invalid="ignore"):
121+
if agg_type in ["std", "stddev", "var", "variance"]:
122+
sum_sq = np.asarray(ir_sum_sq, dtype="float64")
123+
variance = (sum_sq - total**2 / count) / (count - 1)
124+
# Clamp tiny negative values caused by floating point cancellation
125+
variance = np.where(count > 1, np.maximum(variance, 0.0), np.nan)
126+
if agg_type in ["std", "stddev"]:
127+
return np.sqrt(variance)
128+
return variance
129+
130+
return np.where(count > 0, total / count, 0)

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

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,13 @@
88
from datetime import timedelta
99
from typing import Any, Callable, Dict, List, Tuple, Union
1010

11-
import numpy as np
1211
import pandas as pd
1312

1413
from feast.aggregation import Aggregation
15-
from feast.aggregation.tiling.base import get_ir_metadata_for_aggregation
14+
from feast.aggregation.tiling.base import (
15+
compute_holistic_value_from_irs,
16+
get_ir_metadata_for_aggregation,
17+
)
1618

1719

1820
def apply_sawtooth_window_tiling(
@@ -178,12 +180,16 @@ def apply_sawtooth_window_tiling(
178180
sum_col = ir_column_names[0]
179181
count_col = ir_column_names[1]
180182

183+
sum_sq_col = (
184+
ir_column_names[2] if len(ir_column_names) >= 3 else None
185+
)
186+
181187
if sum_col in result.columns and count_col in result.columns:
182-
# Compute avg = sum / count
183-
result[feature_name] = np.where(
184-
result[count_col] > 0,
185-
result[sum_col] / result[count_col],
186-
0,
188+
result[feature_name] = compute_holistic_value_from_irs(
189+
agg,
190+
result[sum_col],
191+
result[count_col],
192+
result[sum_sq_col] if sum_sq_col else None,
187193
)
188194

189195
return result

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

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@
1111
import pandas as pd
1212

1313
from feast.aggregation import Aggregation
14-
from feast.aggregation.tiling.base import get_ir_metadata_for_aggregation
14+
from feast.aggregation.tiling.base import (
15+
compute_holistic_value_from_irs,
16+
get_ir_metadata_for_aggregation,
17+
)
1518

1619

1720
def convert_cumulative_to_windowed(
@@ -124,15 +127,21 @@ def convert_cumulative_to_windowed(
124127
1
125128
] # e.g., "_tail_avg_amount_3600s_count"
126129

130+
sum_sq_col = (
131+
ir_column_names[2]
132+
if len(ir_column_names) >= 3
133+
else None
134+
)
135+
127136
if sum_col in windowed_row and count_col in windowed_row:
128-
count_val = windowed_row[count_col]
129-
if count_val <= 0:
130-
windowed_row[feature_name] = 0
131-
else:
132-
# avg = windowed_sum / windowed_count
133-
windowed_row[feature_name] = (
134-
windowed_row[sum_col] / count_val
137+
windowed_row[feature_name] = float(
138+
compute_holistic_value_from_irs(
139+
agg,
140+
windowed_row[sum_col],
141+
windowed_row[count_col],
142+
windowed_row.get(sum_sq_col),
135143
)
144+
)
136145

137146
# Set event_timestamp to the tile end time
138147
windowed_row[timestamp_col] = pd.to_datetime(tile_end, unit="ms")

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,3 +135,58 @@ def test_count_distinct_tiling_raises():
135135
agg = Aggregation(column="item_id", function="count_distinct")
136136
with pytest.raises(ValueError, match="count_distinct does not support tiling"):
137137
get_ir_metadata_for_aggregation(agg, "count_distinct_item_id")
138+
139+
140+
@pytest.mark.parametrize(
141+
"function, expected_fn",
142+
[
143+
("avg", lambda s: s.mean()),
144+
("std", lambda s: s.std()),
145+
("var", lambda s: s.var()),
146+
],
147+
)
148+
def test_tiled_holistic_aggregation_matches_pandas(function, expected_fn):
149+
"""Tiled avg/std/var over a sliding window match the pandas result."""
150+
from feast.aggregation.tiling.orchestrator import apply_sawtooth_window_tiling
151+
from feast.aggregation.tiling.tile_subtraction import (
152+
convert_cumulative_to_windowed,
153+
)
154+
155+
start = pd.Timestamp("2024-01-01")
156+
values = [1.0, 2.0, 3.0, 10.0, 4.0, 7.0, 5.0, 8.0]
157+
df = pd.DataFrame(
158+
{
159+
"driver_id": [1] * len(values),
160+
"event_timestamp": [
161+
start + timedelta(minutes=5 * i) for i in range(len(values))
162+
],
163+
"trips": values,
164+
}
165+
)
166+
window = timedelta(minutes=20)
167+
agg = Aggregation(column="trips", function=function, time_window=window)
168+
169+
tiles = apply_sawtooth_window_tiling(
170+
df=df.copy(),
171+
aggregations=[agg],
172+
group_by_keys=["driver_id"],
173+
timestamp_col="event_timestamp",
174+
window_size=window,
175+
hop_size=timedelta(minutes=5),
176+
)
177+
result = convert_cumulative_to_windowed(
178+
tiles_df=tiles,
179+
entity_keys=["driver_id"],
180+
timestamp_col="event_timestamp",
181+
window_size=window,
182+
aggregations=[agg],
183+
)
184+
185+
feature_name = agg.resolved_name(window)
186+
for _, row in result.iterrows():
187+
tile_end = row["event_timestamp"]
188+
in_window = df[
189+
(df["event_timestamp"] >= tile_end - window)
190+
& (df["event_timestamp"] < tile_end)
191+
]["trips"]
192+
assert row[feature_name] == pytest.approx(expected_fn(in_window), nan_ok=True)

0 commit comments

Comments
 (0)