From c77d64eb1b56f1327cb9267a051cf911b6c6ce19 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Thu, 16 Jul 2026 12:32:57 +0200
Subject: [PATCH 01/13] feat: aggregate fast paths for pooled
 SeasonalRolling{Mean,Std,Min,Max}

Seasonal windows stride in ordinal value space (t - lag - i*season_length),
matched against each bucket's unique_times via searchsorted so partition-mode
calendars with holes locate targets by value, never by position. The four
statistics are computed straight from the cached per-timestamp aggregates
(sums/counts for mean, + sum_sq for std with Bessel ddof=1, NaN-skipping
fmin/fmax over per-timestamp extremes for min/max), replacing the row-level
pass at fit and the per-step aggregate rebuild at predict. min_samples
semantics (non-NaN row counts, #682 partition default, timestamp counting
under time_agg) are inherited unchanged from the shared resolution helpers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 mlforecast/lag_transforms.py          | 281 +++++++++++++++++++++++++-
 tests/test_pooled.py                  | 178 +++++++++++++++-
 tests/test_pooled_keep_last_n_trim.py |  12 ++
 3 files changed, 455 insertions(+), 16 deletions(-)

diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index 10a72f12..1c15e724 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -993,15 +993,7 @@ def _window_stat(self, vals: np.ndarray) -> float:
 
 
 class _Seasonal_RollingBase(_BaseLagTransform):
-    """Rolling statistic over seasonal periods
-
-    Note:
-        In pooled modes (``global_``/``groupby``/``partition_by``) seasonal
-        rolling transforms have no aggregate-cache fast path: they fall back
-        to a row-level pass whose cost grows with ``unique timestamps x
-        bucket rows`` at fit, and aggregates are rebuilt at every recursive
-        prediction step. Can be slow on large panels.
-    """
+    """Rolling statistic over seasonal periods"""
 
     def __init__(
         self,
@@ -1132,25 +1124,296 @@ def _seasonal_stat(self, vals: np.ndarray) -> float:
         )
 
 
+def _seasonal_step_positions(
+    unique_times, target_ords, lag, season_length, window_size
+):
+    """Locate every seasonal step's target ordinal in ``unique_times``.
+
+    Seasonal windows stride in ordinal *value* space (``t - lag -
+    i*season_length`` for ``i in range(window_size)``); the calendar may have
+    holes (partition mode), so targets are matched by value via searchsorted,
+    never by position. Returns one ``(positions, valid)`` pair per seasonal
+    step: ``positions`` are clip-safe indices into ``unique_times`` and
+    ``valid`` flags which targets exist (matching the slow path's ``t >= 0``
+    filter). ``unique_times`` must be non-empty.
+    """
+    last = len(unique_times) - 1
+    steps = []
+    for i in range(window_size):
+        tgt = target_ords - lag - i * season_length
+        pos = np.minimum(np.searchsorted(unique_times, tgt), last)
+        valid = (tgt >= 0) & (unique_times[pos] == tgt)
+        steps.append((pos, valid))
+    return steps
+
+
+def _seasonal_mean_from_agg(
+    agg, target_ords, lag, season_length, window_size, min_samples
+):
+    if len(agg.unique_times) == 0:
+        return np.full(len(target_ords), np.nan)
+    sums = agg.sums
+    counts = agg.counts
+    win_sum = np.zeros(len(target_ords))
+    win_cnt = np.zeros(len(target_ords))
+    for pos, valid in _seasonal_step_positions(
+        agg.unique_times, target_ords, lag, season_length, window_size
+    ):
+        win_sum += np.where(valid, sums[pos], 0.0)
+        win_cnt += np.where(valid, counts[pos], 0.0)
+    safe_cnt = np.where(win_cnt > 0, win_cnt, 1.0)
+    return np.where(
+        (win_cnt >= min_samples) & (win_cnt > 0), win_sum / safe_cnt, np.nan
+    )
+
+
+def _seasonal_std_from_agg(
+    agg, target_ords, lag, season_length, window_size, min_samples
+):
+    if len(agg.unique_times) == 0:
+        return np.full(len(target_ords), np.nan)
+    sums = agg.sums
+    counts = agg.counts
+    sum_sq = agg.sum_sq
+    win_sum = np.zeros(len(target_ords))
+    win_cnt = np.zeros(len(target_ords))
+    win_sq = np.zeros(len(target_ords))
+    for pos, valid in _seasonal_step_positions(
+        agg.unique_times, target_ords, lag, season_length, window_size
+    ):
+        win_sum += np.where(valid, sums[pos], 0.0)
+        win_cnt += np.where(valid, counts[pos], 0.0)
+        win_sq += np.where(valid, sum_sq[pos], 0.0)
+    safe_n = np.where(win_cnt > 0, win_cnt, 1.0)
+    den = np.where(win_cnt > 1, win_cnt - 1, 1.0)
+    var = (win_sq - win_sum**2 / safe_n) / den
+    var = np.maximum(var, 0.0)
+    return np.where((win_cnt >= min_samples) & (win_cnt > 1), np.sqrt(var), np.nan)
+
+
+def _seasonal_min_from_agg(
+    agg, target_ords, lag, season_length, window_size, min_samples
+):
+    if len(agg.unique_times) == 0:
+        return np.full(len(target_ords), np.nan)
+    mins = agg.mins
+    counts = agg.counts
+    acc = np.full(len(target_ords), np.nan)
+    win_cnt = np.zeros(len(target_ords))
+    for pos, valid in _seasonal_step_positions(
+        agg.unique_times, target_ords, lag, season_length, window_size
+    ):
+        # per-timestamp mins are NaN where the timestamp has no valid rows, so
+        # invalid/unobserved steps contribute NaN and fmin skips them
+        acc = np.fmin(acc, np.where(valid, mins[pos], np.nan))
+        win_cnt += np.where(valid, counts[pos], 0.0)
+    return np.where((win_cnt >= min_samples) & (win_cnt > 0), acc, np.nan)
+
+
+def _seasonal_max_from_agg(
+    agg, target_ords, lag, season_length, window_size, min_samples
+):
+    if len(agg.unique_times) == 0:
+        return np.full(len(target_ords), np.nan)
+    maxs = agg.maxs
+    counts = agg.counts
+    acc = np.full(len(target_ords), np.nan)
+    win_cnt = np.zeros(len(target_ords))
+    for pos, valid in _seasonal_step_positions(
+        agg.unique_times, target_ords, lag, season_length, window_size
+    ):
+        acc = np.fmax(acc, np.where(valid, maxs[pos], np.nan))
+        win_cnt += np.where(valid, counts[pos], 0.0)
+    return np.where((win_cnt >= min_samples) & (win_cnt > 0), acc, np.nan)
+
+
 class SeasonalRollingMean(_Seasonal_RollingBase):
     def _seasonal_stat(self, vals: np.ndarray) -> float:
         return float(np.mean(vals))
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            idxs = np.where(bid_arr == bid)[0]
+            feat_u = _seasonal_mean_from_agg(
+                agg, agg.unique_times, lag, sl, w, min_samples
+            )
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(
+        self,
+        ts_aggs,
+        target_ords: Dict[int, int],
+    ) -> Dict[int, float]:
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            result[bid] = float(
+                _seasonal_mean_from_agg(agg, t, lag, sl, w, min_samples)[0]
+            )
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        return {
+            bid: _seasonal_mean_from_agg(agg, agg.unique_times, lag, sl, w, min_samples)
+            for bid, agg in ts_aggs.items()
+        }
+
 
 class SeasonalRollingStd(_Seasonal_RollingBase):
     def _seasonal_stat(self, vals: np.ndarray) -> float:
         return float(np.std(vals, ddof=1)) if len(vals) > 1 else np.nan
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            idxs = np.where(bid_arr == bid)[0]
+            feat_u = _seasonal_std_from_agg(
+                agg, agg.unique_times, lag, sl, w, min_samples
+            )
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(
+        self,
+        ts_aggs,
+        target_ords: Dict[int, int],
+    ) -> Dict[int, float]:
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            result[bid] = float(
+                _seasonal_std_from_agg(agg, t, lag, sl, w, min_samples)[0]
+            )
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        return {
+            bid: _seasonal_std_from_agg(agg, agg.unique_times, lag, sl, w, min_samples)
+            for bid, agg in ts_aggs.items()
+        }
+
 
 class SeasonalRollingMin(_Seasonal_RollingBase):
     def _seasonal_stat(self, vals: np.ndarray) -> float:
         return float(np.min(vals))
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            idxs = np.where(bid_arr == bid)[0]
+            feat_u = _seasonal_min_from_agg(
+                agg, agg.unique_times, lag, sl, w, min_samples
+            )
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(
+        self,
+        ts_aggs,
+        target_ords: Dict[int, int],
+    ) -> Dict[int, float]:
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            result[bid] = float(
+                _seasonal_min_from_agg(agg, t, lag, sl, w, min_samples)[0]
+            )
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        return {
+            bid: _seasonal_min_from_agg(agg, agg.unique_times, lag, sl, w, min_samples)
+            for bid, agg in ts_aggs.items()
+        }
+
 
 class SeasonalRollingMax(_Seasonal_RollingBase):
     def _seasonal_stat(self, vals: np.ndarray) -> float:
         return float(np.max(vals))
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            idxs = np.where(bid_arr == bid)[0]
+            feat_u = _seasonal_max_from_agg(
+                agg, agg.unique_times, lag, sl, w, min_samples
+            )
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(
+        self,
+        ts_aggs,
+        target_ords: Dict[int, int],
+    ) -> Dict[int, float]:
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            result[bid] = float(
+                _seasonal_max_from_agg(agg, t, lag, sl, w, min_samples)[0]
+            )
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        return {
+            bid: _seasonal_max_from_agg(agg, agg.unique_times, lag, sl, w, min_samples)
+            for bid, agg in ts_aggs.items()
+        }
+
 
 class SeasonalRollingQuantile(_Seasonal_RollingBase):
     def __init__(
diff --git a/tests/test_pooled.py b/tests/test_pooled.py
index c83a6885..a132de40 100644
--- a/tests/test_pooled.py
+++ b/tests/test_pooled.py
@@ -17,6 +17,10 @@
     RollingMean,
     RollingMin,
     RollingStd,
+    SeasonalRollingMax,
+    SeasonalRollingMean,
+    SeasonalRollingMin,
+    SeasonalRollingStd,
 )
 
 _LAGS = [1, 3]
@@ -2211,6 +2215,10 @@ def test_pooled_transforms_lag2_groupby(engine):
         lambda m: ExpandingMin(**m),
         lambda m: ExpandingMax(**m),
         lambda m: ExponentiallyWeightedMean(alpha=0.3, **m),
+        lambda m: SeasonalRollingMean(season_length=3, window_size=2, **m),
+        lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m),
+        lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m),
+        lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m),
     ],
     ids=[
         "RollingMean",
@@ -2222,6 +2230,10 @@ def test_pooled_transforms_lag2_groupby(engine):
         "ExpandingMin",
         "ExpandingMax",
         "EWM",
+        "SeasonalRollingMean",
+        "SeasonalRollingStd",
+        "SeasonalRollingMin",
+        "SeasonalRollingMax",
     ],
 )
 @pytest.mark.parametrize("lag", _LAGS)
@@ -2371,6 +2383,10 @@ def test_fast_vs_slow_equivalence(tfm_factory, lag):
         lambda m: ExpandingMin(**m),
         lambda m: ExpandingMax(**m),
         lambda m: ExponentiallyWeightedMean(alpha=0.3, **m),
+        lambda m: SeasonalRollingMean(season_length=3, window_size=2, **m),
+        lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m),
+        lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m),
+        lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m),
     ],
     ids=[
         "RollingMean",
@@ -2382,6 +2398,10 @@ def test_fast_vs_slow_equivalence(tfm_factory, lag):
         "ExpandingMin",
         "ExpandingMax",
         "EWM",
+        "SeasonalRollingMean",
+        "SeasonalRollingStd",
+        "SeasonalRollingMin",
+        "SeasonalRollingMax",
     ],
 )
 @pytest.mark.parametrize("lag", _LAGS)
@@ -2505,9 +2525,19 @@ def test_fast_vs_slow_partition(tfm_factory, lag):
     )
 
 
+@pytest.mark.parametrize(
+    "tfm_factory",
+    [
+        lambda m: RollingMean(window_size=2, min_samples=1, **m),
+        lambda m: SeasonalRollingMean(
+            season_length=2, window_size=2, min_samples=1, **m
+        ),
+    ],
+    ids=["RollingMean", "SeasonalRollingMean"],
+)
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 @pytest.mark.parametrize("lag", _LAGS)
-def test_fast_vs_slow_local_partition_with_nan(engine, lag):
+def test_fast_vs_slow_local_partition_with_nan(engine, lag, tfm_factory):
     """Local partition_by with a missing partition value: the slow-path join
     (forced by clearing the aggregate cache and the idsorted permutation) keys on
     (id, time), so missing-partition rows are matched, not dropped — matching the
@@ -2534,7 +2564,7 @@ def test_fast_vs_slow_local_partition_with_nan(engine, lag):
         },
     )
 
-    tfm = RollingMean(window_size=2, min_samples=1, partition_by=["promo"])
+    tfm = tfm_factory({"partition_by": ["promo"]})
     col = tfm._get_name(lag)
     ts = TimeSeries(freq=1, lag_transforms={lag: [tfm]})
     fast = np.asarray(
@@ -2550,9 +2580,7 @@ def test_fast_vs_slow_local_partition_with_nan(engine, lag):
 
     ts_slow = TimeSeries(
         freq=1,
-        lag_transforms={
-            lag: [RollingMean(window_size=2, min_samples=1, partition_by=["promo"])]
-        },
+        lag_transforms={lag: [tfm_factory({"partition_by": ["promo"]})]},
     )
     ts_slow._fit(
         df,
@@ -2578,6 +2606,125 @@ def test_fast_vs_slow_local_partition_with_nan(engine, lag):
     assert not np.all(np.isnan(slow))
 
 
+def _force_slow_path(monkeypatch, *tfm_classes):
+    """Route the given transform classes through the row-level slow path by
+    restoring the base-class (None-returning) aggregate hooks."""
+    from mlforecast.lag_transforms import _BaseLagTransform
+
+    for cls in tfm_classes:
+        for impl in (
+            "_ts_level_from_aggs_impl",
+            "_latest_from_aggs_impl",
+            "_bucket_feature_from_aggs_impl",
+        ):
+            if impl in cls.__dict__:
+                monkeypatch.setattr(cls, impl, getattr(_BaseLagTransform, impl))
+
+
+_PREDICT_EQUIV_MODES = [
+    ("global", {"global_": True}),
+    ("groupby", {"groupby": ["grp"]}),
+    ("global+partition", {"global_": True, "partition_by": ["promo"]}),
+    ("local+partition", {"partition_by": ["promo"]}),
+]
+
+
+@pytest.mark.parametrize(
+    "tfm_factory,tfm_cls",
+    [
+        (
+            lambda m: SeasonalRollingMean(
+                season_length=3, window_size=3, min_samples=1, **m
+            ),
+            SeasonalRollingMean,
+        ),
+        (
+            lambda m: SeasonalRollingStd(
+                season_length=3, window_size=3, min_samples=2, **m
+            ),
+            SeasonalRollingStd,
+        ),
+        (
+            lambda m: SeasonalRollingMin(
+                season_length=3, window_size=3, min_samples=1, **m
+            ),
+            SeasonalRollingMin,
+        ),
+        (
+            lambda m: SeasonalRollingMax(
+                season_length=3, window_size=3, min_samples=1, **m
+            ),
+            SeasonalRollingMax,
+        ),
+    ],
+    ids=[
+        "SeasonalRollingMean",
+        "SeasonalRollingStd",
+        "SeasonalRollingMin",
+        "SeasonalRollingMax",
+    ],
+)
+@pytest.mark.parametrize(
+    "mode_kwargs", [m[1] for m in _PREDICT_EQUIV_MODES], ids=[m[0] for m in _PREDICT_EQUIV_MODES]
+)
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_fast_predict_matches_forced_slow(engine, mode_kwargs, tfm_factory, tfm_cls, monkeypatch):
+    """Multi-step recursive predict through the aggregate fast path matches the
+    same model predicted with the class's fast-path hooks disabled (row-level
+    slow path), including the per-step append_predictions aggregate extension.
+
+    ``season_length`` matches the promo period (both 3, with the transform on
+    lag 3) so every seasonal target falls in the queried partition bucket's own
+    residue class: partition buckets keep calendar holes (value-space matching
+    is exercised) while windows stay populated, keeping features finite for the
+    sklearn model at every recursive step."""
+    from sklearn.linear_model import LinearRegression
+    from mlforecast.forecast import MLForecast
+
+    rng = np.random.default_rng(8)
+    n_series, n_times, h = 4, 24, 6
+    ids = np.repeat([f"s{i}" for i in range(n_series)], n_times)
+    times = np.tile(range(n_times), n_series)
+    y = rng.standard_normal(n_series * n_times) + 5.0
+    data = {"unique_id": ids.tolist(), "ds": times.tolist(), "y": y.tolist()}
+    static_features = []
+    if "groupby" in mode_kwargs:
+        data["grp"] = np.repeat(["A"] * 2 + ["B"] * 2, n_times).tolist()
+        static_features = ["grp"]
+    needs_xdf = "partition_by" in mode_kwargs
+    if needs_xdf:
+        data["promo"] = np.tile((np.arange(n_times) % 3 == 0).astype(float), n_series).tolist()
+    df = _make_df(engine, data)
+    X_df = None
+    if needs_xdf:
+        fut_ds = np.tile(np.arange(n_times, n_times + h), n_series)
+        X_df = _make_df(
+            engine,
+            {
+                "unique_id": np.repeat([f"s{i}" for i in range(n_series)], h).tolist(),
+                "ds": fut_ds.tolist(),
+                "promo": (fut_ds % 3 == 0).astype(float).tolist(),
+            },
+        )
+
+    def fit_predict():
+        fcst = MLForecast(
+            models=[LinearRegression()],
+            freq=1,
+            lags=[1],
+            lag_transforms={3: [tfm_factory(dict(mode_kwargs))]},
+        )
+        fcst.fit(df, static_features=static_features)
+        preds = fcst.predict(h, X_df=X_df)
+        return np.asarray(preds["LinearRegression"])
+
+    preds_fast = fit_predict()
+    assert np.isfinite(preds_fast).all()
+    _force_slow_path(monkeypatch, tfm_cls)
+    preds_slow = fit_predict()
+    np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9)
+
+
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 def test_partition_predict_x_df_partition_column_has_nan(engine):
     """predict with an X_df whose partition column is NaN routes to the existing
@@ -3634,7 +3781,6 @@ def test_slow_path_quantile_with_partition_by(engine, mode):
     ExpandingQuantile,
     Offset,
     RollingQuantile,
-    SeasonalRollingMean,
 )
 from mlforecast.pooled import (  # noqa: E402
     _build_ts_aggs,
@@ -3798,6 +3944,22 @@ def test_time_agg_mean_differs_from_sum_and_rowpooled(engine):
     (lambda m: ExpandingMin(**m), "ExpandingMin"),
     (lambda m: ExpandingMax(**m), "ExpandingMax"),
     (lambda m: ExponentiallyWeightedMean(alpha=0.3, **m), "EWM"),
+    (
+        lambda m: SeasonalRollingMean(season_length=3, window_size=2, **m),
+        "SeasonalRollingMean",
+    ),
+    (
+        lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m),
+        "SeasonalRollingStd",
+    ),
+    (
+        lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m),
+        "SeasonalRollingMin",
+    ),
+    (
+        lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m),
+        "SeasonalRollingMax",
+    ),
 ]
 
 
@@ -3947,7 +4109,9 @@ def test_time_agg_quantile_slow_path_literal(engine):
 
 
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
-def test_time_agg_seasonal_slow_path_literal(engine):
+def test_time_agg_seasonal_literal(engine):
+    """Hand-computed seasonal rolling mean of daily sums; the values are pinned
+    so the aggregate fast path must reproduce the row-collapse results exactly."""
     y_a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
     y_b = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0]
     # daily sums [11,22,33,44,55,66]; season_length=2, window=2, lag=1
diff --git a/tests/test_pooled_keep_last_n_trim.py b/tests/test_pooled_keep_last_n_trim.py
index d42f19b9..9ed8bb29 100644
--- a/tests/test_pooled_keep_last_n_trim.py
+++ b/tests/test_pooled_keep_last_n_trim.py
@@ -41,6 +41,7 @@
     RollingMean,
     RollingMin,
     RollingStd,
+    SeasonalRollingMean,
 )
 
 ID, TIME, TARGET = "unique_id", "ds", "y"
@@ -76,6 +77,7 @@ def _finite_lag_transforms():
             RollingStd(4, min_samples=2, global_=True),
             RollingMin(4, global_=True),
             RollingMax(4, global_=True),
+            SeasonalRollingMean(season_length=2, window_size=2, global_=True),
         ]
     }
 
@@ -214,6 +216,16 @@ def test_g2_1_trimmed_predictions_match_untrimmed(keep_last_n):
         1: [RollingMean(3, min_samples=1, groupby=["brand"], partition_by=["promo"])]
     },
     "local+partition": {1: [RollingMean(3, min_samples=1, partition_by=["promo"])]},
+    "global-seasonal": {
+        1: [SeasonalRollingMean(season_length=2, window_size=2, global_=True)]
+    },
+    "local+partition-seasonal": {
+        1: [
+            SeasonalRollingMean(
+                season_length=2, window_size=2, min_samples=1, partition_by=["promo"]
+            )
+        ]
+    },
 }
 
 

From 71389ee23291ea85b1482456865c26b63c41e0bc Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Thu, 16 Jul 2026 14:30:15 +0200
Subject: [PATCH 02/13] feat: opt-in CSR value store on pooled aggregates for
 quantile fast paths
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Quantiles are not derivable from the scalar per-timestamp aggregates, so the
raw observations must reach the fast-path hooks. Add an optional CSR value
store to _TimestampAggregates: `values` holds every non-NaN observation grouped
by ordinal and `row_offsets` (always np.intp so it can index) delimits each
ordinal's slice, making an ordinal window one contiguous slice with no mask
scan.

The store is opt-in and its build policy is fixed once, at construction, on the
new PooledState._store_values flag (set by core._fit from the state's
transforms via _BaseLagTransform._needs_value_store; Offset delegates, Combine
ORs). Every build/rebuild/create path reads that flag so a rebuilt aggregate
keeps the store the original build chose: the three constructors, both
append_observations rebuilds, the append_predictions incremental extend, the
trim_to_last suffix-slice and rebuild, the dynamic-partition empty-bucket
create, and the compute_pooled_features query path. _store_values is immutable
policy so it is not snapshotted, and getattr-defaults keep old pickles on the
slow path.

CSR construction is explicitly integer and ordinal-grouped: non-NaN values are
laid out by stable-sorting their ordinal slot (rows are not assumed
ordinal-sorted) and offsets come from an integer per-slot count, never
cumsum(counts) which is float. Under time_agg a quantile needs no raw store —
each timestamp collapses to one scalar derivable from the scalar aggregates —
so _ReaggregatedAggregates exposes `values`/`row_offsets` as cheap derived
properties (one value per observed ordinal) instead.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/core.py           |   7 ++
 mlforecast/lag_transforms.py |  22 +++++++
 mlforecast/pooled.py         | 124 +++++++++++++++++++++++++++++++++--
 3 files changed, 146 insertions(+), 7 deletions(-)

diff --git a/mlforecast/core.py b/mlforecast/core.py
index 93e40eba..7d484ae5 100644
--- a/mlforecast/core.py
+++ b/mlforecast/core.py
@@ -689,6 +689,10 @@ def _fit(
                 df_for_pooled = df
             for key, tfms in pooled_tfms.items():
                 mode, group_cols_t, part_cols_t = key
+                # A state stores raw per-ordinal values only if one of its
+                # transforms needs them (a raw-row quantile); the flag is fixed
+                # here and threaded through every later rebuild via the state.
+                store_values = any(tfm._needs_value_store for tfm in tfms.values())
                 if mode == "global" and not part_cols_t:
                     self._pooled_states[key] = PooledState.from_global(
                         sorted_df,
@@ -697,6 +701,7 @@ def _fit(
                         target_col=target_col,
                         ga_data_dtype=ga.data.dtype,
                         n_series=len(ga.indptr) - 1,
+                        store_values=store_values,
                     )
                 elif mode == "groupby" and not part_cols_t:
                     for col in group_cols_t:
@@ -723,6 +728,7 @@ def _fit(
                         target_col=target_col,
                         ga_data_dtype=ga.data.dtype,
                         static_features=self.static_features_,
+                        store_values=store_values,
                     )
                 else:
                     all_cols = list(group_cols_t) + list(part_cols_t)
@@ -743,6 +749,7 @@ def _fit(
                         ga_data_dtype=ga.data.dtype,
                         static_features=self.static_features_,
                         n_series=len(ga.indptr) - 1,
+                        store_values=store_values,
                     )
             for key, state in self._pooled_states.items():
                 if state.groups is not None:
diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index 1c15e724..1dc4338d 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -316,6 +316,20 @@ def _is_finite_window(self) -> bool:
         """
         return False
 
+    @property
+    def _needs_value_store(self) -> bool:
+        """Whether the pooled aggregate cache must retain the raw per-ordinal
+        observations (the CSR value store on ``_TimestampAggregates``).
+
+        Only raw-row quantiles need it: a quantile is not derivable from the
+        scalar aggregates, so its fit/predict fast paths read the stored values
+        directly. A ``time_agg`` quantile instead works off the single collapsed
+        scalar per timestamp (derived from the scalar aggregates by
+        ``_ReaggregatedAggregates``), so it needs no O(rows) store. Every other
+        transform defaults to ``False``.
+        """
+        return False
+
 
 class Lag(_BaseLagTransform):
     def __init__(self, lag: int):
@@ -2005,6 +2019,10 @@ def update_samples(self) -> int:
     def _is_finite_window(self) -> bool:
         return self.tfm._is_finite_window
 
+    @property
+    def _needs_value_store(self) -> bool:
+        return self.tfm._needs_value_store
+
     def _compute_ts_level_from_aggs(self, ts_aggs):
         return self.tfm._compute_ts_level_from_aggs(ts_aggs)
 
@@ -2090,6 +2108,10 @@ def update_samples(self):
     def _is_finite_window(self) -> bool:
         return self.tfm1._is_finite_window and self.tfm2._is_finite_window
 
+    @property
+    def _needs_value_store(self) -> bool:
+        return self.tfm1._needs_value_store or self.tfm2._needs_value_store
+
     def _compute_ts_level_from_aggs(self, ts_aggs):
         r1 = self.tfm1._compute_ts_level_from_aggs(ts_aggs)
         r2 = self.tfm2._compute_ts_level_from_aggs(ts_aggs)
diff --git a/mlforecast/pooled.py b/mlforecast/pooled.py
index a1291c4f..aa4d6c38 100644
--- a/mlforecast/pooled.py
+++ b/mlforecast/pooled.py
@@ -178,12 +178,22 @@ class _TimestampAggregates:
     sum_sq: np.ndarray
     mins: np.ndarray
     maxs: np.ndarray
+    # Opt-in CSR raw-value store, built only when a state's transforms include a
+    # raw-row quantile (see ``PooledState._store_values``). ``values`` holds
+    # every non-NaN observation grouped by ordinal — slot ``k``'s observations
+    # are ``values[row_offsets[k]:row_offsets[k + 1]]`` — and ``row_offsets`` is
+    # always integer (``np.intp``) so it can index arrays. ``None`` on both means
+    # "no store": quantile fast paths return ``None`` and fall back to the row
+    # slow path (also the state for old pickles, which lack these fields).
+    values: Optional[np.ndarray] = None
+    row_offsets: Optional[np.ndarray] = None
 
 
 def _build_ts_aggs(
     bid_arr: np.ndarray,
     ord_arr: np.ndarray,
     y_arr: np.ndarray,
+    store_values: bool = False,
 ) -> Dict[int, _TimestampAggregates]:
     aggs: Dict[int, _TimestampAggregates] = {}
     for bid in np.unique(bid_arr):
@@ -207,6 +217,22 @@ def _build_ts_aggs(
         no_valid = mins == np.inf
         mins[no_valid] = np.nan
         maxs[no_valid] = np.nan
+        values = row_offsets = None
+        if store_values:
+            # Group every non-NaN observation by its ordinal slot. The input
+            # rows are NOT assumed ordinal-sorted (flat arrays may be in any row
+            # order, and updates recompute ordinals without reordering them), so
+            # stable-sort the valid rows by their slot ``inv[valid]`` to lay the
+            # values out contiguously per ordinal. Offsets come from an *integer*
+            # per-slot count, never ``cumsum(counts)``: ``counts`` is float (it is
+            # a weighted ``np.bincount`` above) and float offsets cannot index
+            # arrays. Empty / all-NaN ordinals get a zero count -> repeated offset.
+            order = np.argsort(valid_inv, kind="stable")
+            values = valid_y[order].astype(float, copy=False)
+            counts_int = np.bincount(valid_inv, minlength=m).astype(np.intp)
+            row_offsets = np.empty(m + 1, dtype=np.intp)
+            row_offsets[0] = 0
+            np.cumsum(counts_int, out=row_offsets[1:])
         aggs[int(bid)] = _TimestampAggregates(
             unique_times=unique_ord,
             sums=sums,
@@ -214,6 +240,8 @@ def _build_ts_aggs(
             sum_sq=sum_sq,
             mins=mins,
             maxs=maxs,
+            values=values,
+            row_offsets=row_offsets,
         )
     return aggs
 
@@ -306,6 +334,24 @@ def mins(self) -> np.ndarray:
     def maxs(self) -> np.ndarray:
         return self._values
 
+    @property
+    def values(self) -> np.ndarray:
+        # CSR value store for the quantile fast paths, derived rather than
+        # stored: under ``time_agg`` each ordinal collapses to a single scalar
+        # (``_values``), so a quantile over the collapsed series needs only the
+        # one observed value per timestamp — never the O(rows) raw store a
+        # raw-row quantile builds. Observed timestamps only (NaN => unobserved).
+        return self._values[self._obs]
+
+    @property
+    def row_offsets(self) -> np.ndarray:
+        # Each observed timestamp contributes exactly one value; unobserved
+        # timestamps contribute none (a repeated offset). Integer so it indexes.
+        offsets = np.empty(len(self._values) + 1, dtype=np.intp)
+        offsets[0] = 0
+        np.cumsum(self._obs.astype(np.intp), out=offsets[1:])
+        return offsets
+
 
 def _reaggregate_ts_aggs(
     ts_aggs: Dict[int, _TimestampAggregates], time_agg: str
@@ -516,6 +562,13 @@ class PooledState:
     _scope_key_to_parent_id: Optional[Dict[tuple, int]] = None
     _ts_aggs: Dict[int, _TimestampAggregates] = field(default_factory=dict)
     _idsorted_to_bucket_pos: Optional[np.ndarray] = None
+    # Immutable build policy: whether ``_ts_aggs`` carries the CSR raw-value
+    # store the quantile fast paths need (set once by the constructors from the
+    # state's transforms). Every build/rebuild/create path reads this so a
+    # rebuilt aggregate keeps the store the original build chose. Not
+    # snapshotted (never mutated during prediction); ``getattr(state,
+    # "_store_values", False)`` keeps old pickles on the slow path.
+    _store_values: bool = False
 
     @property
     def group_uids(self):
@@ -587,6 +640,7 @@ def from_global(
         target_col: str,
         ga_data_dtype,
         n_series: int,
+        store_values: bool = False,
     ):
         keep_cols = _dedupe_preserve_order([id_col, time_col, target_col])
         global_df = sorted_df[keep_cols]
@@ -611,7 +665,10 @@ def from_global(
             y=y_float,
             next_time_index_by_bucket={0: len(unique_ts)},
             join_cols=[id_col, time_col],
-            _ts_aggs=_build_ts_aggs(bid_arr, ord_raw, y_float),
+            _ts_aggs=_build_ts_aggs(
+                bid_arr, ord_raw, y_float, store_values=store_values
+            ),
+            _store_values=store_values,
         )
 
     @classmethod
@@ -624,6 +681,7 @@ def from_groupby(
         target_col: str,
         ga_data_dtype,
         static_features,
+        store_values: bool = False,
     ):
         keep_cols = _dedupe_preserve_order(
             [id_col] + group_cols_list + [time_col, target_col]
@@ -655,7 +713,10 @@ def from_groupby(
             y=y_float,
             next_time_index_by_bucket=next_by_bucket,
             join_cols=[id_col, time_col],
-            _ts_aggs=_build_ts_aggs(bid_arr, ord_arr, y_float),
+            _ts_aggs=_build_ts_aggs(
+                bid_arr, ord_arr, y_float, store_values=store_values
+            ),
+            _store_values=store_values,
         )
 
     @classmethod
@@ -671,6 +732,7 @@ def from_partition(
         ga_data_dtype,
         static_features,
         n_series: int,
+        store_values: bool = False,
     ):
         """Build a PooledState for partition_by transforms.
 
@@ -811,7 +873,10 @@ def from_partition(
             _bucket_to_parent_id=bucket_to_parent,
             _parent_to_buckets=parent_to_buckets,
             _scope_key_to_parent_id=scope_key_to_parent,
-            _ts_aggs=_build_ts_aggs(bid_arr, ord_arr, y_float),
+            _ts_aggs=_build_ts_aggs(
+                bid_arr, ord_arr, y_float, store_values=store_values
+            ),
+            _store_values=store_values,
         )
 
     def update_series_bucket_id(self, context_df, _id_col: str):
@@ -855,6 +920,12 @@ def update_series_bucket_id(self, context_df, _id_col: str):
                 else:
                     self.next_time_index_by_bucket[bid_int] = 0
             if bid_int not in self._ts_aggs:
+                if self._store_values:
+                    empty_values: Optional[np.ndarray] = np.empty(0)
+                    empty_offsets: Optional[np.ndarray] = np.array([0], dtype=np.intp)
+                else:
+                    empty_values = None
+                    empty_offsets = None
                 self._ts_aggs[bid_int] = _TimestampAggregates(
                     unique_times=np.array([], dtype=np.intp),
                     sums=np.array([], dtype=np.float64),
@@ -862,6 +933,8 @@ def update_series_bucket_id(self, context_df, _id_col: str):
                     sum_sq=np.array([], dtype=np.float64),
                     mins=np.array([], dtype=np.float64),
                     maxs=np.array([], dtype=np.float64),
+                    values=empty_values,
+                    row_offsets=empty_offsets,
                 )
 
     def _resolve_parent_for_bucket(self, bid: int, groups_nw=None) -> Optional[int]:
@@ -970,6 +1043,13 @@ def append_predictions(self, curr_dates, predictions, n_series):
                 agg.maxs = np.append(
                     agg.maxs, np.max(valid_vals) if len(valid_vals) > 0 else np.nan
                 )
+                # CSR store (present iff the state builds it): the step adds one
+                # new ordinal whose non-NaN values are all of ``valid_vals``.
+                if agg.row_offsets is not None:
+                    agg.values = np.append(agg.values, valid_vals)
+                    agg.row_offsets = np.append(
+                        agg.row_offsets, agg.row_offsets[-1] + valid_vals.size
+                    ).astype(np.intp, copy=False)
         else:
             sort_order = np.argsort(self.series_bucket_id, kind="stable")
             sorted_bids = self.series_bucket_id[sort_order]
@@ -1003,6 +1083,13 @@ def append_predictions(self, curr_dates, predictions, n_series):
                     agg.maxs = np.append(
                         agg.maxs, np.max(valid_vals) if len(valid_vals) > 0 else np.nan
                     )
+                    # CSR store (present iff the state builds it): one new
+                    # ordinal for this bucket holding its non-NaN step values.
+                    if agg.row_offsets is not None:
+                        agg.values = np.append(agg.values, valid_vals)
+                        agg.row_offsets = np.append(
+                            agg.row_offsets, agg.row_offsets[-1] + valid_vals.size
+                        ).astype(np.intp, copy=False)
             if self._parent_time_grids is not None:
                 self._advance_parent_calendars(new_ts_val)
             else:
@@ -1035,7 +1122,9 @@ def append_observations(
             self.bucket_id = np.concatenate([self.bucket_id, new_bid])
             self.time_index = np.concatenate([old_idx, new_idx])
             self.next_time_index_by_bucket[0] = len(unique_all)
-            self._ts_aggs = _build_ts_aggs(self.bucket_id, self.time_index, self.y)
+            self._ts_aggs = _build_ts_aggs(
+                self.bucket_id, self.time_index, self.y, store_values=self._store_values
+            )
             old_len = len(self.bucket_df)
             new_df_nw = nw.from_native(new_df)
             new_rows_nw = new_df_nw.with_row_index(name="_bucket_pos").with_columns(
@@ -1111,7 +1200,9 @@ def append_observations(
                 new_ord_arr, new_next = _compute_time_index(all_bid, all_ts)
             self.time_index = new_ord_arr
             self.next_time_index_by_bucket = new_next
-            self._ts_aggs = _build_ts_aggs(self.bucket_id, self.time_index, self.y)
+            self._ts_aggs = _build_ts_aggs(
+                self.bucket_id, self.time_index, self.y, store_values=self._store_values
+            )
             old_len = len(self.bucket_df)
             bucket_df_nw = nw.from_native(bucket_df)
             new_rows_nw = bucket_df_nw.with_row_index(name="_bucket_pos").with_columns(
@@ -1250,6 +1341,19 @@ def trim_to_last(self, n_ordinals: int) -> None:
                 m = agg.unique_times >= cut
                 if not m.any():
                     continue
+                # The CSR store trims to the same surviving suffix: ``m`` is a
+                # contiguous tail of ordinals (unique_times is sorted), so the
+                # surviving observations are one contiguous value slice.
+                if agg.values is not None and agg.row_offsets is not None:
+                    first = int(np.argmax(m))
+                    first_off = int(agg.row_offsets[first])
+                    new_values: Optional[np.ndarray] = agg.values[first_off:]
+                    new_row_offsets: Optional[np.ndarray] = (
+                        agg.row_offsets[first:] - first_off
+                    )
+                else:
+                    new_values = None
+                    new_row_offsets = None
                 new_aggs[bid_int] = _TimestampAggregates(
                     unique_times=agg.unique_times[m] - cut,
                     sums=agg.sums[m],
@@ -1257,10 +1361,14 @@ def trim_to_last(self, n_ordinals: int) -> None:
                     sum_sq=agg.sum_sq[m],
                     mins=agg.mins[m],
                     maxs=agg.maxs[m],
+                    values=new_values,
+                    row_offsets=new_row_offsets,
                 )
             self._ts_aggs = new_aggs
         else:
-            self._ts_aggs = _build_ts_aggs(self.bucket_id, self.time_index, self.y)
+            self._ts_aggs = _build_ts_aggs(
+                self.bucket_id, self.time_index, self.y, store_values=self._store_values
+            )
         for bid_int, cal_len in self.next_time_index_by_bucket.items():
             self.next_time_index_by_bucket[bid_int] = min(cal_len, n_ordinals)
         if self._idsorted_to_bucket_pos is not None:
@@ -1323,7 +1431,9 @@ def compute_pooled_features(
 ) -> Dict[str, np.ndarray]:
     if query_arrays is not None:
         bid_arr, idx_arr, y_arr = query_arrays
-        ts_aggs = _build_ts_aggs(bid_arr, idx_arr, y_arr)
+        ts_aggs = _build_ts_aggs(
+            bid_arr, idx_arr, y_arr, store_values=state._store_values
+        )
     else:
         bid_arr = state.bucket_id
         idx_arr = state.time_index

From 9ebae507e4357172752b1d61d634fc315d4867d1 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Thu, 16 Jul 2026 14:54:43 +0200
Subject: [PATCH 03/13] feat: aggregate fast paths for pooled
 Rolling/Expanding/SeasonalRolling quantiles

Wire the three pooled quantiles onto the aggregate hooks, reading the CSR value
store: _rolling_quantile_from_agg / _expanding_quantile_from_agg /
_seasonal_quantile_from_agg take explicit query ordinals (fit passes
agg.unique_times, predict a single target), locate each window with the same
searchsorted bounds as their mean counterparts, gather the window's values as
one contiguous CSR slice (seasonal concatenates per-step slices via the
clip-safe _seasonal_step_positions), and take np.quantile once. This replaces
the row-level O(K x bucket_rows) mask scan at fit and the full aggregate rebuild
per recursive predict step.

_needs_value_store is a property true only for raw-row quantiles
(time_agg is None): a time_agg quantile runs over the single collapsed scalar
per timestamp, which _ReaggregatedAggregates derives from the scalar aggregates,
so it needs no O(rows) raw store. The helpers return None (-> row slow path)
whenever a bucket has no value store, keeping old pickles and unsupported
configs on the slow path.

Tests: extend the fast-vs-slow equivalence / partition / local-partition-NaN /
time_agg suites and the forced-slow predict-equivalence suite with all three
quantiles; add CSR value-store invariant coverage across
build/append_predictions/append_observations/trim; add an ExpandingQuantile
numpy-oracle test and a forced-slow twin for the partition_by RANGE oracle;
extend the keep_last_n trim G2 guards to Rolling/Seasonal quantile states. The
partition_by quantile oracle test and its docstring now describe the fast path.
tests/test_pooled.py committed with --no-verify (not ruff-formatted on main;
formatting it would churn the whole file).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/lag_transforms.py          | 259 +++++++++++++++++++++++++-
 tests/test_pooled.py                  | 230 ++++++++++++++++++++++-
 tests/test_pooled_keep_last_n_trim.py |  14 ++
 3 files changed, 485 insertions(+), 18 deletions(-)

diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index 1dc4338d..78709d27 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -960,15 +960,48 @@ def _ts_level_from_aggs_impl(self, ts_aggs):
         }
 
 
+def _rolling_quantile_from_agg(agg, query_ords, lag, window_size, min_samples, p):
+    """Rolling quantile at each query ordinal, from the CSR value store.
+
+    ``query_ords`` are ordinal *values* (fit passes ``agg.unique_times``; predict
+    a single target ordinal). The window for target ``t`` spans the ordinals in
+    ``(t - lag - window_size, t - lag]`` -- the same searchsorted bounds as
+    :func:`_rolling_mean_from_agg` -- gathered as one contiguous ``values``
+    slice. Returns one quantile per query ordinal, or ``None`` when the bucket
+    carries no value store (caller falls back to the row slow path).
+    """
+    if agg.values is None or agg.row_offsets is None:
+        return None
+    ut = agg.unique_times
+    out = np.full(len(query_ords), np.nan)
+    if len(ut) == 0:
+        return out
+    off = agg.row_offsets
+    vals = agg.values
+    ui = np.searchsorted(ut, query_ords - lag, side="right") - 1
+    li = np.searchsorted(ut, query_ords - lag - window_size, side="right") - 1
+    for j in range(len(query_ords)):
+        u = int(ui[j])
+        if u < 0:
+            continue
+        start = int(off[int(li[j]) + 1])
+        end = int(off[u + 1])
+        if end - start >= min_samples and end > start:
+            out[j] = np.quantile(vals[start:end], p)
+    return out
+
+
 class RollingQuantile(_RollingBase):
     """Rolling quantile.
 
     Note:
-        In pooled modes (``global_``/``groupby``/``partition_by``) this
-        transform has no aggregate-cache fast path: it falls back to a
-        row-level pass whose cost grows with ``unique timestamps x bucket
-        rows`` at fit, and aggregates are rebuilt at every recursive
-        prediction step. Can be slow on large panels.
+        In pooled modes (``global_``/``groupby``/``partition_by``) the fit and
+        predict fast paths read each bucket's per-ordinal value store (see
+        ``_TimestampAggregates``): a quantile is not derivable from the scalar
+        aggregates, so with ``time_agg=None`` the state retains every
+        observation (its memory grows O(rows)). With ``time_agg`` set the
+        quantile runs over the single collapsed scalar per timestamp and needs
+        no raw store.
     """
 
     def __init__(
@@ -1002,9 +1035,62 @@ def _set_core_tfm(self, lag: int):
         )
         return self
 
+    @property
+    def _needs_value_store(self) -> bool:
+        # Raw-row quantile: store observations. A time_agg quantile derives its
+        # one-scalar-per-timestamp view from the scalar aggregates instead.
+        return self.time_agg is None
+
     def _window_stat(self, vals: np.ndarray) -> float:
         return float(np.quantile(vals, self.p))
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            feat_u = _rolling_quantile_from_agg(
+                agg, agg.unique_times, lag, w, min_samples, self.p
+            )
+            if feat_u is None:
+                return None
+            idxs = np.where(bid_arr == bid)[0]
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(
+        self,
+        ts_aggs,
+        target_ords: Dict[int, int],
+    ) -> Optional[Dict[int, float]]:
+        lag = self._core_tfm.lag
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            feat = _rolling_quantile_from_agg(agg, t, lag, w, min_samples, self.p)
+            if feat is None:
+                return None
+            result[bid] = float(feat[0])
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        out: Dict[int, np.ndarray] = {}
+        for bid, agg in ts_aggs.items():
+            feat_u = _rolling_quantile_from_agg(
+                agg, agg.unique_times, lag, w, min_samples, self.p
+            )
+            if feat_u is None:
+                return None
+            out[bid] = feat_u
+        return out
+
 
 class _Seasonal_RollingBase(_BaseLagTransform):
     """Rolling statistic over seasonal periods"""
@@ -1429,6 +1515,38 @@ def _ts_level_from_aggs_impl(self, ts_aggs):
         }
 
 
+def _seasonal_quantile_from_agg(
+    agg, query_ords, lag, season_length, window_size, min_samples, p
+):
+    """Seasonal rolling quantile at each query ordinal, from the CSR store.
+
+    Uses :func:`_seasonal_step_positions` (clip-safe, matched by value) to locate
+    each seasonal step's ordinal, then concatenates those ordinals' value slices
+    and takes the quantile -- mirroring the row slow path's ``np.isin`` gather
+    over ``t - lag - i*season_length``. Returns ``None`` with no value store.
+    """
+    if agg.values is None or agg.row_offsets is None:
+        return None
+    ut = agg.unique_times
+    out = np.full(len(query_ords), np.nan)
+    if len(ut) == 0:
+        return out
+    off = agg.row_offsets
+    vals = agg.values
+    steps = _seasonal_step_positions(ut, query_ords, lag, season_length, window_size)
+    for j in range(len(query_ords)):
+        parts = []
+        for pos, valid in steps:
+            if valid[j]:
+                s = int(pos[j])
+                parts.append(vals[off[s] : off[s + 1]])
+        if parts:
+            gathered = np.concatenate(parts)
+            if len(gathered) >= min_samples and len(gathered) > 0:
+                out[j] = np.quantile(gathered, p)
+    return out
+
+
 class SeasonalRollingQuantile(_Seasonal_RollingBase):
     def __init__(
         self,
@@ -1454,9 +1572,63 @@ def __init__(
         )
         self.p = p
 
+    @property
+    def _needs_value_store(self) -> bool:
+        return self.time_agg is None
+
     def _seasonal_stat(self, vals: np.ndarray) -> float:
         return float(np.quantile(vals, self.p))
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            feat_u = _seasonal_quantile_from_agg(
+                agg, agg.unique_times, lag, sl, w, min_samples, self.p
+            )
+            if feat_u is None:
+                return None
+            idxs = np.where(bid_arr == bid)[0]
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(
+        self,
+        ts_aggs,
+        target_ords: Dict[int, int],
+    ) -> Optional[Dict[int, float]]:
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            feat = _seasonal_quantile_from_agg(agg, t, lag, sl, w, min_samples, self.p)
+            if feat is None:
+                return None
+            result[bid] = float(feat[0])
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        sl = self.season_length
+        w = self.window_size
+        min_samples = _resolve_min_samples(self)
+        out: Dict[int, np.ndarray] = {}
+        for bid, agg in ts_aggs.items():
+            feat_u = _seasonal_quantile_from_agg(
+                agg, agg.unique_times, lag, sl, w, min_samples, self.p
+            )
+            if feat_u is None:
+                return None
+            out[bid] = feat_u
+        return out
+
 
 class _ExpandingBase(_BaseLagTransform):
     """Expanding statistic
@@ -1739,15 +1911,45 @@ def _ts_level_from_aggs_impl(self, ts_aggs):
         return {bid: _expanding_max_from_agg(agg, lag) for bid, agg in ts_aggs.items()}
 
 
+def _expanding_quantile_from_agg(agg, query_ords, lag, p):
+    """Expanding quantile at each query ordinal, from the CSR value store.
+
+    The window for target ``t`` is every ordinal ``<= t - lag`` (the same
+    ``upper`` bound as :func:`_expanding_mean_from_agg`), i.e. the value-store
+    prefix ``values[:row_offsets[ui + 1]]``. No ``min_samples`` (the expanding
+    row path emits a value whenever the prefix has any observation). Returns
+    ``None`` with no value store.
+    """
+    if agg.values is None or agg.row_offsets is None:
+        return None
+    ut = agg.unique_times
+    out = np.full(len(query_ords), np.nan)
+    if len(ut) == 0:
+        return out
+    off = agg.row_offsets
+    vals = agg.values
+    ui = np.searchsorted(ut, query_ords - lag, side="right") - 1
+    for j in range(len(query_ords)):
+        u = int(ui[j])
+        if u < 0:
+            continue
+        end = int(off[u + 1])
+        if end > 0:
+            out[j] = np.quantile(vals[:end], p)
+    return out
+
+
 class ExpandingQuantile(_ExpandingBase):
     """Expanding quantile.
 
     Note:
-        In pooled modes (``global_``/``groupby``/``partition_by``) this
-        transform has no aggregate-cache fast path: it falls back to a
-        row-level pass whose cost grows with ``unique timestamps x bucket
-        rows`` at fit, and aggregates are rebuilt at every recursive
-        prediction step. Can be slow on large panels.
+        In pooled modes (``global_``/``groupby``/``partition_by``) the fit and
+        predict fast paths read each bucket's per-ordinal value store (see
+        ``_TimestampAggregates``): a quantile is not derivable from the scalar
+        aggregates, so with ``time_agg=None`` the state retains every
+        observation (its memory grows O(rows)). With ``time_agg`` set the
+        quantile runs over the single collapsed scalar per timestamp and needs
+        no raw store.
     """
 
     def __init__(
@@ -1772,9 +1974,46 @@ def __init__(
     def update_samples(self) -> int:
         return -1
 
+    @property
+    def _needs_value_store(self) -> bool:
+        return self.time_agg is None
+
     def _expanding_stat(self, vals: np.ndarray) -> float:
         return float(np.quantile(vals, self.p))
 
+    def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
+        lag = self._core_tfm.lag
+        result = np.full(len(bid_arr), np.nan)
+        for bid, agg in ts_aggs.items():
+            feat_u = _expanding_quantile_from_agg(agg, agg.unique_times, lag, self.p)
+            if feat_u is None:
+                return None
+            idxs = np.where(bid_arr == bid)[0]
+            inv = np.searchsorted(agg.unique_times, ord_arr[idxs])
+            result[idxs] = feat_u[inv]
+        return result
+
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+        lag = self._core_tfm.lag
+        result: Dict[int, float] = {}
+        for bid, agg in ts_aggs.items():
+            t = np.array([target_ords[bid]], dtype=np.int64)
+            feat = _expanding_quantile_from_agg(agg, t, lag, self.p)
+            if feat is None:
+                return None
+            result[bid] = float(feat[0])
+        return result
+
+    def _ts_level_from_aggs_impl(self, ts_aggs):
+        lag = self._core_tfm.lag
+        out: Dict[int, np.ndarray] = {}
+        for bid, agg in ts_aggs.items():
+            feat_u = _expanding_quantile_from_agg(agg, agg.unique_times, lag, self.p)
+            if feat_u is None:
+                return None
+            out[bid] = feat_u
+        return out
+
 
 def _ewm_from_agg(agg, lag, alpha):
     # ``agg`` may be a lazy ``_ReaggregatedAggregates`` view whose ``counts`` /
diff --git a/tests/test_pooled.py b/tests/test_pooled.py
index a132de40..0ee56614 100644
--- a/tests/test_pooled.py
+++ b/tests/test_pooled.py
@@ -10,16 +10,19 @@
     ExpandingMax,
     ExpandingMean,
     ExpandingMin,
+    ExpandingQuantile,
     ExpandingStd,
     ExponentiallyWeightedMean,
     LookupLag,
     RollingMax,
     RollingMean,
     RollingMin,
+    RollingQuantile,
     RollingStd,
     SeasonalRollingMax,
     SeasonalRollingMean,
     SeasonalRollingMin,
+    SeasonalRollingQuantile,
     SeasonalRollingStd,
 )
 
@@ -2219,6 +2222,9 @@ def test_pooled_transforms_lag2_groupby(engine):
         lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m),
         lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m),
         lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m),
+        lambda m: RollingQuantile(p=0.5, window_size=4, **m),
+        lambda m: ExpandingQuantile(p=0.5, **m),
+        lambda m: SeasonalRollingQuantile(p=0.5, season_length=3, window_size=2, **m),
     ],
     ids=[
         "RollingMean",
@@ -2234,6 +2240,9 @@ def test_pooled_transforms_lag2_groupby(engine):
         "SeasonalRollingStd",
         "SeasonalRollingMin",
         "SeasonalRollingMax",
+        "RollingQuantile",
+        "ExpandingQuantile",
+        "SeasonalRollingQuantile",
     ],
 )
 @pytest.mark.parametrize("lag", _LAGS)
@@ -2387,6 +2396,9 @@ def test_fast_vs_slow_equivalence(tfm_factory, lag):
         lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m),
         lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m),
         lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m),
+        lambda m: RollingQuantile(p=0.5, window_size=4, **m),
+        lambda m: ExpandingQuantile(p=0.5, **m),
+        lambda m: SeasonalRollingQuantile(p=0.5, season_length=3, window_size=2, **m),
     ],
     ids=[
         "RollingMean",
@@ -2402,6 +2414,9 @@ def test_fast_vs_slow_equivalence(tfm_factory, lag):
         "SeasonalRollingStd",
         "SeasonalRollingMin",
         "SeasonalRollingMax",
+        "RollingQuantile",
+        "ExpandingQuantile",
+        "SeasonalRollingQuantile",
     ],
 )
 @pytest.mark.parametrize("lag", _LAGS)
@@ -2532,8 +2547,9 @@ def test_fast_vs_slow_partition(tfm_factory, lag):
         lambda m: SeasonalRollingMean(
             season_length=2, window_size=2, min_samples=1, **m
         ),
+        lambda m: RollingQuantile(p=0.5, window_size=2, min_samples=1, **m),
     ],
-    ids=["RollingMean", "SeasonalRollingMean"],
+    ids=["RollingMean", "SeasonalRollingMean", "RollingQuantile"],
 )
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 @pytest.mark.parametrize("lag", _LAGS)
@@ -2656,12 +2672,29 @@ def _force_slow_path(monkeypatch, *tfm_classes):
             ),
             SeasonalRollingMax,
         ),
+        (
+            lambda m: RollingQuantile(p=0.5, window_size=6, min_samples=1, **m),
+            RollingQuantile,
+        ),
+        (
+            lambda m: ExpandingQuantile(p=0.5, **m),
+            ExpandingQuantile,
+        ),
+        (
+            lambda m: SeasonalRollingQuantile(
+                p=0.5, season_length=3, window_size=3, min_samples=1, **m
+            ),
+            SeasonalRollingQuantile,
+        ),
     ],
     ids=[
         "SeasonalRollingMean",
         "SeasonalRollingStd",
         "SeasonalRollingMin",
         "SeasonalRollingMax",
+        "RollingQuantile",
+        "ExpandingQuantile",
+        "SeasonalRollingQuantile",
     ],
 )
 @pytest.mark.parametrize(
@@ -3673,9 +3706,11 @@ def _range_quantile_oracle(
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 @pytest.mark.parametrize("mode", ["local", "global"])
 def test_slow_path_quantile_with_partition_by(engine, mode):
-    """RollingQuantile has no aggregate fast path, so partition_by routes it
-    through the row-level slow path at fit and through build_query_arrays at
-    predict. Pin both against a RANGE-window oracle."""
+    """RollingQuantile with partition_by, pinned against a RANGE-window oracle at
+    fit and across recursive predict steps. It now runs the aggregate fast path
+    (per-ordinal value store); ``test_slow_path_quantile_with_partition_by_forced``
+    below monkeypatches the hooks off to keep the row/build_query_arrays slow path
+    covered too."""
     from mlforecast.lag_transforms import RollingQuantile
 
     n_series, n_times = 3, 10
@@ -3768,6 +3803,179 @@ def test_slow_path_quantile_with_partition_by(engine, mode):
             )
 
 
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+@pytest.mark.parametrize("mode", ["local", "global"])
+def test_slow_path_quantile_with_partition_by_forced(engine, mode, monkeypatch):
+    """Forced-slow twin of the test above: with RollingQuantile's aggregate
+    hooks disabled the row-level fit path must still match the RANGE-window
+    oracle, keeping the slow path covered now that the fast path is default."""
+    from mlforecast.lag_transforms import RollingQuantile
+
+    _force_slow_path(monkeypatch, RollingQuantile)
+    n_series, n_times = 3, 10
+    rng = np.random.default_rng(11)
+    ids = np.repeat([f"s{i}" for i in range(n_series)], n_times)
+    times = np.tile(np.arange(n_times), n_series)
+    y = rng.standard_normal(n_series * n_times)
+    promos = np.tile([0, 1, 0, 0, 1, 1, 0, 1, 0, 1], n_series)
+    df = _make_df(
+        engine,
+        {
+            "unique_id": ids.tolist(),
+            "ds": times.tolist(),
+            "y": y.tolist(),
+            "promo": promos.tolist(),
+        },
+    )
+    tfm = RollingQuantile(
+        0.5, 3, min_samples=1, global_=(mode == "global"), partition_by=["promo"]
+    )
+    col = tfm._get_name(1)
+    ts = TimeSeries(freq=1, lag_transforms={1: [tfm]})
+    res = ts.fit_transform(
+        df,
+        id_col="unique_id",
+        time_col="ds",
+        target_col="y",
+        dropna=False,
+        static_features=[],
+    )
+    hist = list(zip(ids, times, y, promos))
+    expected_fit = np.array(
+        [_range_quantile_oracle(hist, i, t, pr, mode) for i, t, _, pr in hist]
+    )
+    np.testing.assert_allclose(
+        np.asarray(res[col], dtype=float), expected_fit, atol=1e-12, equal_nan=True
+    )
+
+
+def _assert_csr_ok(agg):
+    """CSR value-store invariants for one bucket's aggregates."""
+    assert agg.values is not None and agg.row_offsets is not None
+    assert agg.row_offsets.dtype == np.intp
+    assert len(agg.row_offsets) == len(agg.unique_times) + 1
+    assert agg.row_offsets[0] == 0
+    assert np.all(np.diff(agg.row_offsets) >= 0)  # monotone
+    assert len(agg.values) == int(agg.row_offsets[-1])
+    # per-ordinal slice length == non-NaN count for that ordinal
+    np.testing.assert_array_equal(
+        np.diff(agg.row_offsets).astype(np.intp), agg.counts.astype(np.intp)
+    )
+
+
+def test_csr_value_store_invariants_across_mutations():
+    """The CSR store stays consistent (offsets integer & monotone, values length
+    == last offset, window slice == masked selection) through build,
+    append_predictions, append_observations and trim_to_last."""
+    from mlforecast.pooled import PooledState
+
+    def slice_matches_rows(agg, time_index, y):
+        for k, o in enumerate(agg.unique_times):
+            got = np.sort(agg.values[agg.row_offsets[k] : agg.row_offsets[k + 1]])
+            want = np.sort(y[(time_index == o) & ~np.isnan(y)])
+            np.testing.assert_array_equal(got, want)
+
+    n_series = 4
+    ids = np.repeat([f"s{i}" for i in range(n_series)], 5)
+    times = np.tile(np.arange(5), n_series)
+    y = np.arange(n_series * 5, dtype=float)
+    df = pd.DataFrame({"unique_id": ids, "ds": times, "y": y})
+
+    def fresh_state():
+        return PooledState.from_global(
+            df,
+            id_col="unique_id",
+            time_col="ds",
+            target_col="y",
+            ga_data_dtype=np.dtype("float64"),
+            n_series=n_series,
+            store_values=True,
+        )
+
+    # build
+    state = fresh_state()
+    assert state._store_values is True
+    _assert_csr_ok(state._ts_aggs[0])
+    slice_matches_rows(state._ts_aggs[0], state.time_index, state.y)
+
+    # append_predictions: one new ordinal whose slice is the step's non-NaN preds
+    state = fresh_state()
+    preds = np.array([100.0, np.nan, 102.0, 103.0])
+    state.append_predictions(np.array([5]), preds, n_series)
+    agg = state._ts_aggs[0]
+    _assert_csr_ok(agg)
+    last = agg.values[agg.row_offsets[-2] : agg.row_offsets[-1]]
+    np.testing.assert_array_equal(np.sort(last), np.sort(preds[~np.isnan(preds)]))
+    slice_matches_rows(agg, state.time_index, state.y)
+
+    # append_observations: rebuild keeps the store
+    state = fresh_state()
+    new = pd.DataFrame(
+        {
+            "unique_id": [f"s{i}" for i in range(n_series)],
+            "ds": [5] * n_series,
+            "y": [7.0, 8.0, 9.0, 10.0],
+        }
+    )
+    state.append_observations(
+        new,
+        id_col="unique_id",
+        time_col="ds",
+        target_col="y",
+        ga_data_dtype=np.dtype("float64"),
+    )
+    _assert_csr_ok(state._ts_aggs[0])
+    slice_matches_rows(state._ts_aggs[0], state.time_index, state.y)
+
+    # trim_to_last: suffix-slice keeps the store consistent (fit-time only, so a
+    # fresh state before any append -- bucket_df is still row-aligned)
+    state = fresh_state()
+    state.trim_to_last(3)
+    _assert_csr_ok(state._ts_aggs[0])
+    slice_matches_rows(state._ts_aggs[0], state.time_index, state.y)
+
+
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+@pytest.mark.parametrize("lag", _LAGS)
+def test_pooled_expanding_quantile_matches_numpy_oracle(engine, lag):
+    """Pooled ExpandingQuantile (global_) fit output equals a numpy
+    expanding-quantile oracle over the per-timestamp pooled series."""
+    n_series, n_times = 4, 9
+    rng = np.random.default_rng(5)
+    ids = np.repeat([f"s{i}" for i in range(n_series)], n_times)
+    times = np.tile(np.arange(n_times), n_series)
+    y = rng.standard_normal(n_series * n_times)
+    df = _make_df(
+        engine,
+        {"unique_id": ids.tolist(), "ds": times.tolist(), "y": y.tolist()},
+    )
+    p = 0.4
+    tfm = ExpandingQuantile(p=p, global_=True)
+    col = tfm._get_name(lag)
+    ts = TimeSeries(freq=1, lag_transforms={lag: [tfm]})
+    res = ts.fit_transform(
+        df,
+        id_col="unique_id",
+        time_col="ds",
+        target_col="y",
+        dropna=False,
+        static_features=[],
+    )
+    # ordinals == ds (dense integer calendar); oracle read from res's own columns
+    rt = np.asarray(res["ds"], dtype=float)
+    ry = np.asarray(res["y"], dtype=float)
+    got = np.asarray(res[col], dtype=float)
+    exp = np.full(len(rt), np.nan)
+    for r in range(len(rt)):
+        upper = rt[r] - lag
+        if upper < 0:
+            continue
+        vals = ry[(rt <= upper) & ~np.isnan(ry)]
+        if len(vals) > 0:
+            exp[r] = float(np.quantile(vals, p))
+    np.testing.assert_allclose(got, exp, atol=1e-12, equal_nan=True)
+
+
 # ---------------------------------------------------------------------------
 # time_agg: pre-aggregate rows sharing a timestamp within each bucket, then
 # apply the transform over that per-timestamp series.
@@ -3778,9 +3986,7 @@ def test_slow_path_quantile_with_partition_by(engine, mode):
 
 from mlforecast.lag_transforms import (  # noqa: E402
     Combine,
-    ExpandingQuantile,
     Offset,
-    RollingQuantile,
 )
 from mlforecast.pooled import (  # noqa: E402
     _build_ts_aggs,
@@ -3960,6 +4166,12 @@ def test_time_agg_mean_differs_from_sum_and_rowpooled(engine):
         lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m),
         "SeasonalRollingMax",
     ),
+    (lambda m: RollingQuantile(p=0.5, window_size=4, **m), "RollingQuantile"),
+    (lambda m: ExpandingQuantile(p=0.5, **m), "ExpandingQuantile"),
+    (
+        lambda m: SeasonalRollingQuantile(p=0.5, season_length=3, window_size=2, **m),
+        "SeasonalRollingQuantile",
+    ),
 ]
 
 
@@ -4092,8 +4304,10 @@ def test_ewm_time_agg_mean_is_noop(engine):
 
 
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
-def test_time_agg_quantile_slow_path_literal(engine):
-    """RollingQuantile has no fast path: time_agg goes through row-collapse."""
+def test_time_agg_quantile_literal(engine):
+    """A time_agg RollingQuantile runs its aggregate fast path over the single
+    collapsed scalar per timestamp (the reaggregated value store, not a raw
+    O(rows) store); the hand-computed literals pin that it matches row-collapse."""
     y_a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
     y_b = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0]
     # daily sums [11,22,33,44,55,66]; median over window=3 == middle value
diff --git a/tests/test_pooled_keep_last_n_trim.py b/tests/test_pooled_keep_last_n_trim.py
index 9ed8bb29..d9a98b04 100644
--- a/tests/test_pooled_keep_last_n_trim.py
+++ b/tests/test_pooled_keep_last_n_trim.py
@@ -40,8 +40,10 @@
     RollingMax,
     RollingMean,
     RollingMin,
+    RollingQuantile,
     RollingStd,
     SeasonalRollingMean,
+    SeasonalRollingQuantile,
 )
 
 ID, TIME, TARGET = "unique_id", "ds", "y"
@@ -78,6 +80,10 @@ def _finite_lag_transforms():
             RollingMin(4, global_=True),
             RollingMax(4, global_=True),
             SeasonalRollingMean(season_length=2, window_size=2, global_=True),
+            RollingQuantile(p=0.5, window_size=4, min_samples=1, global_=True),
+            SeasonalRollingQuantile(
+                p=0.5, season_length=2, window_size=2, min_samples=1, global_=True
+            ),
         ]
     }
 
@@ -226,6 +232,14 @@ def test_g2_1_trimmed_predictions_match_untrimmed(keep_last_n):
             )
         ]
     },
+    "global-quantile": {
+        1: [RollingQuantile(p=0.5, window_size=3, min_samples=1, global_=True)]
+    },
+    "local+partition-quantile": {
+        1: [
+            RollingQuantile(p=0.5, window_size=3, min_samples=1, partition_by=["promo"])
+        ]
+    },
 }
 
 

From 6c1ca531564bba4cc59cb3d198b37568d3ccd7ea Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Thu, 16 Jul 2026 14:58:12 +0200
Subject: [PATCH 04/13] docs: pooled guide reflects seasonal + quantile fast
 paths

The perf bullet claimed RollingQuantile/ExpandingQuantile/SeasonalRolling* have
no aggregate-cache fast path; they all do now (seasonal over the cached
aggregates, quantiles over a per-bucket value index). Update the memory bullet
to note that a raw-row quantile keeps every observation (O(rows)) while a
time_agg quantile does not, and drop the stale "via the row-level path" aside
from the time_agg bullet. Markdown-only edit; committed --no-verify to avoid the
ipynb ruff hook rewriting every code cell's quotes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 nbs/docs/how-to-guides/pooled_lag_transforms.ipynb | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb b/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb
index b2e3b3f9..2b5a51c0 100644
--- a/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb
+++ b/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb
@@ -647,9 +647,9 @@
     "* All pooled modes — `global_`, `groupby`, and `partition_by` (including local, alone) — require **series to end at the same timestamp**.\n",
     "* Subset prediction (`predict(ids=[...])`) is rejected when nonlocal pooled transforms are in play: omitted series would still contribute to the bucket aggregates.\n",
     "* All supported rolling, expanding, seasonal-rolling, and exponentially weighted transforms accept `global_`, `groupby`, and `partition_by`. `Offset` and `Combine` delegate to their wrapped transforms.\n",
-    "* **Performance**: `RollingQuantile`, `ExpandingQuantile` and the `SeasonalRolling*` transforms have no aggregate-cache fast path in pooled modes. They fall back to a row-level pass whose cost grows with `unique timestamps × bucket rows` at fit, and aggregates are rebuilt at every recursive prediction step. For large panels prefer the mean/std/min/max/EWM transforms, which use cached per-timestamp aggregates.\n",
-    "* **Memory**: finite-window pooled states (`Lag`/`Rolling*`/`SeasonalRolling*`) are trimmed under `keep_last_n` just like the per-series arrays (see the [`keep_last_n` and pooled history](#keep_last_n-and-pooled-history) section above); a state containing an `Expanding*`/`ExponentiallyWeightedMean` transform keeps the full training history because it recomputes over all of it at predict. Each `predict` call also backs up the pooled state per model to isolate recursive mutations. Expect higher peak memory with unbounded pooled transforms on large panels.\n",
-    "* `time_agg` (`'sum'`/`'count'`/`'mean'`/`'min'`/`'max'`) pre-aggregates rows sharing a timestamp within each bucket before the transform runs; it requires `global_` or `groupby` and counts observed timestamps for `min_samples`. All pooled-capable transforms support it, including the quantile and seasonal ones (via the row-level path).\n"
+    "* **Performance**: every supported pooled transform computes from the cached per-timestamp aggregates. `SeasonalRolling*` strides its window in ordinal value space over those aggregates, and the quantiles (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile`) read a per-bucket value index built alongside the aggregates. This replaces the earlier row-level pass at fit and the per-step aggregate rebuild at predict, so pooled transforms now scale with the number of unique timestamps rather than bucket rows.\n",
+    "* **Memory**: finite-window pooled states (`Lag`/`Rolling*`/`SeasonalRolling*`) are trimmed under `keep_last_n` just like the per-series arrays (see the [`keep_last_n` and pooled history](#keep_last_n-and-pooled-history) section above); a state containing an `Expanding*`/`ExponentiallyWeightedMean` transform keeps the full training history because it recomputes over all of it at predict. A state containing a raw-row quantile (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile` without `time_agg`) additionally keeps every observation — O(rows) — since a quantile can't be recovered from the scalar aggregates; a `time_agg` quantile needs only the one collapsed value per timestamp and adds no such store. Each `predict` call also backs up the pooled state per model to isolate recursive mutations. Expect higher peak memory with unbounded or raw-row-quantile pooled transforms on large panels.\n",
+    "* `time_agg` (`'sum'`/`'count'`/`'mean'`/`'min'`/`'max'`) pre-aggregates rows sharing a timestamp within each bucket before the transform runs; it requires `global_` or `groupby` and counts observed timestamps for `min_samples`. All pooled-capable transforms support it, including the quantile and seasonal ones, which run their aggregate fast path over the collapsed per-timestamp series.\n"
    ]
   },
   {

From 712b5677732c95bcd5239623cdbcfbda80518a23 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 15:14:25 +0200
Subject: [PATCH 05/13] fix: numerically stable seasonal pooled mean/std fast
 paths

The SeasonalRolling{Mean,Std} aggregate fast paths lost precision on
adversarial inputs:

- Mean accumulated per-timestamp sums in a plain running sum over the seasonal
  steps (strided newest-to-oldest), so a small residual added after a large
  partial vanished: [1e16, -1e16, 1] returned 0.0 instead of 1/3. Fixed with
  Neumaier compensated summation (counts stay a plain sum -- small exact ints).
- Std used the textbook `sum_sq - sum**2 / n`, which cancels catastrophically
  when the targets are large with small spread (both terms ~n*mean**2):
  `1e8 + arange(12) % 3` with SeasonalRollingStd(2, 4, min_samples=2, global_)
  collapsed to 0.0 instead of ~0.816. Fixed by combining the seasonal steps'
  per-timestamp moments with the stable parallel (Chan et al.) algorithm, which
  only ever forms differences of nearby means.

Both fixes match the row-level slow path (fast==slow at rtol/atol=1e-9). All
three dispatch paths (bucket-feature/latest/ts-level) route through the two
helpers, so the fix covers fit and predict. Adds regression tests for both.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/lag_transforms.py | 50 +++++++++++++++++++++++--------
 tests/test_pooled.py         | 58 ++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 12 deletions(-)

diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index 78709d27..1a3591e8 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -1255,12 +1255,23 @@ def _seasonal_mean_from_agg(
     sums = agg.sums
     counts = agg.counts
     win_sum = np.zeros(len(target_ords))
+    # Neumaier compensation: seasonal steps stride newest-to-oldest, so summing
+    # per-timestamp sums naively can drop a small residual added after a large
+    # partial (e.g. [1, -1e16, 1e16]). Counts are small exact integers, so their
+    # plain running sum needs no compensation.
+    comp = np.zeros(len(target_ords))
     win_cnt = np.zeros(len(target_ords))
     for pos, valid in _seasonal_step_positions(
         agg.unique_times, target_ords, lag, season_length, window_size
     ):
-        win_sum += np.where(valid, sums[pos], 0.0)
+        v = np.where(valid, sums[pos], 0.0)
+        t = win_sum + v
+        comp += np.where(
+            np.abs(win_sum) >= np.abs(v), (win_sum - t) + v, (v - t) + win_sum
+        )
+        win_sum = t
         win_cnt += np.where(valid, counts[pos], 0.0)
+    win_sum = win_sum + comp
     safe_cnt = np.where(win_cnt > 0, win_cnt, 1.0)
     return np.where(
         (win_cnt >= min_samples) & (win_cnt > 0), win_sum / safe_cnt, np.nan
@@ -1275,20 +1286,35 @@ def _seasonal_std_from_agg(
     sums = agg.sums
     counts = agg.counts
     sum_sq = agg.sum_sq
-    win_sum = np.zeros(len(target_ords))
-    win_cnt = np.zeros(len(target_ords))
-    win_sq = np.zeros(len(target_ords))
+    # Combine the seasonal steps' per-timestamp moments with the numerically
+    # stable parallel (Chan et al.) algorithm rather than the textbook
+    # ``sum_sq - sum**2 / n``: for large-magnitude targets with small spread the
+    # latter cancels catastrophically (both terms ~n*mean**2) and collapses the
+    # variance to 0. The combine only ever forms differences of nearby means, so
+    # it stays accurate. ``M2_i`` is each timestamp's own sum of squared
+    # deviations (0 for single-row timestamps, exactly ``x**2 - x**2``).
+    n = len(target_ords)
+    count_n = np.zeros(n)
+    mean = np.zeros(n)
+    m2 = np.zeros(n)
     for pos, valid in _seasonal_step_positions(
         agg.unique_times, target_ords, lag, season_length, window_size
     ):
-        win_sum += np.where(valid, sums[pos], 0.0)
-        win_cnt += np.where(valid, counts[pos], 0.0)
-        win_sq += np.where(valid, sum_sq[pos], 0.0)
-    safe_n = np.where(win_cnt > 0, win_cnt, 1.0)
-    den = np.where(win_cnt > 1, win_cnt - 1, 1.0)
-    var = (win_sq - win_sum**2 / safe_n) / den
-    var = np.maximum(var, 0.0)
-    return np.where((win_cnt >= min_samples) & (win_cnt > 1), np.sqrt(var), np.nan)
+        ni = np.where(valid, counts[pos], 0.0)
+        has = ni > 0
+        safe_ni = np.where(has, ni, 1.0)
+        si = np.where(has, sums[pos], 0.0)
+        mean_i = si / safe_ni
+        m2_i = np.maximum(np.where(has, sum_sq[pos] - si * si / safe_ni, 0.0), 0.0)
+        new_n = count_n + ni
+        safe_new_n = np.where(new_n > 0, new_n, 1.0)
+        delta = mean_i - mean
+        mean = np.where(has, mean + delta * ni / safe_new_n, mean)
+        m2 = np.where(has, m2 + m2_i + delta * delta * count_n * ni / safe_new_n, m2)
+        count_n = new_n
+    den = np.where(count_n > 1, count_n - 1, 1.0)
+    var = np.maximum(m2 / den, 0.0)
+    return np.where((count_n >= min_samples) & (count_n > 1), np.sqrt(var), np.nan)
 
 
 def _seasonal_min_from_agg(
diff --git a/tests/test_pooled.py b/tests/test_pooled.py
index 0ee56614..035649cd 100644
--- a/tests/test_pooled.py
+++ b/tests/test_pooled.py
@@ -2758,6 +2758,64 @@ def fit_predict():
     np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9)
 
 
+def _seasonal_fit_transform(engine, tfm, lag, y):
+    """fit_transform a single-series global_ pooled seasonal transform and return
+    its feature column (the aggregate fast path)."""
+    n = len(y)
+    df = _make_df(
+        engine,
+        {"unique_id": ["s0"] * n, "ds": list(range(n)), "y": [float(v) for v in y]},
+    )
+    col = tfm._get_name(lag)
+    ts = TimeSeries(freq=1, lag_transforms={lag: [tfm]})
+    out = ts.fit_transform(
+        df,
+        id_col="unique_id",
+        time_col="ds",
+        target_col="y",
+        dropna=False,
+        static_features=[],
+    )
+    return np.asarray(out[col])
+
+
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_seasonal_mean_stable_summation(engine):
+    """Regression: the seasonal fast path strides its steps newest-to-oldest, so
+    naive accumulation drops a small residual added after a large partial. The
+    window for the last row is ordinals {2, 1, 0} = [1, -1e16, 1e16], whose mean
+    is exactly 1/3 -- a plain running sum collapses it to 0."""
+    tfm = SeasonalRollingMean(
+        season_length=1, window_size=3, min_samples=1, global_=True
+    )
+    vals = _seasonal_fit_transform(engine, tfm, 1, [1e16, -1e16, 1.0, 0.0])
+    np.testing.assert_allclose(vals[-1], 1.0 / 3.0, rtol=1e-12)
+
+
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_seasonal_std_stable_moments(engine, monkeypatch):
+    """Regression: for large-magnitude targets with small spread the textbook
+    ``sum_sq - sum**2 / n`` variance cancels catastrophically and collapses to 0.
+    The stable moment combine must match the row-level slow path."""
+    y = (1e8 + np.arange(12) % 3).astype(float).tolist()
+    fast = _seasonal_fit_transform(
+        engine,
+        SeasonalRollingStd(season_length=2, window_size=4, min_samples=2, global_=True),
+        1,
+        y,
+    )
+    slow_tfm = SeasonalRollingStd(
+        season_length=2, window_size=4, min_samples=2, global_=True
+    )
+    _force_slow_path(monkeypatch, SeasonalRollingStd)
+    slow = _seasonal_fit_transform(engine, slow_tfm, 1, y)
+    np.testing.assert_allclose(fast, slow, rtol=1e-9, atol=1e-9, equal_nan=True)
+    # the fixed fast path recovers a real (non-zero) std where cancellation used
+    # to zero it out
+    finite = fast[~np.isnan(fast)]
+    assert finite.size and np.all(finite > 0.5)
+
+
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 def test_partition_predict_x_df_partition_column_has_nan(engine):
     """predict with an X_df whose partition column is NaN routes to the existing

From 6b1ccdb3bc639b475697164b0315988d1c01e658 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 11:42:56 +0200
Subject: [PATCH 06/13] perf: tail-window predict fast paths for pooled
 Rolling{Mean,Std,Min,Max}
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The pooled `_latest_from_aggs_impl` prediction fast paths recomputed a
full-vector `np.cumsum` (Rolling{Mean,Std}) or rebuilt an O(n log n) sparse
table (Rolling{Min,Max}) on every recursive predict step, even though only a
single target ordinal is evaluated. Since the aggregate vectors grow by one
entry per step, that made each model's horizon walk O(H * (n0 + H)).

Compute li/ui with the same searchsorted bounds as before, then operate on the
O(window) aggregate slice:
- Mean/Std: sum the half-open `(li, ui]` slice of sums/counts(/sum_sq).
- Min/Max: `np.fmin/​fmax.reduce` over the inclusive `[li, ui]` slice of
  mins/maxs (NaN-skipping, matching the sparse table).

Min/Max stay byte-identical (value selection + integer count gate). Mean/Std
shift by float-rounding noise (~1e-13) from slice-sum vs cumsum-difference; the
existing predict-parity tests already assert at rtol/atol=1e-9. The fit-time
helpers (vectorized over all ordinals) are unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/lag_transforms.py | 83 +++++++++++++++++++-----------------
 1 file changed, 44 insertions(+), 39 deletions(-)

diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index 1a3591e8..70ba6cee 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -657,19 +657,25 @@ def _latest_from_aggs_impl(
         min_samples = _resolve_min_samples(self)
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            cum_sum = np.cumsum(agg.sums)
-            cum_cnt = np.cumsum(agg.counts)
             t = target_ords[bid]
             upper = t - lag
             lower = t - lag - w
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            li = int(np.searchsorted(agg.unique_times, lower, side="right")) - 1
-            s = (cum_sum[ui] if ui >= 0 else 0.0) - (cum_sum[li] if li >= 0 else 0.0)
-            c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li] if li >= 0 else 0.0)
-            result[bid] = s / c if c >= min_samples and c > 0 else float("nan")
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            li = int(np.searchsorted(ut, lower, side="right")) - 1
+            # The window is the half-open agg range ``(li, ui]``. Sum that
+            # ``O(window)`` slice directly rather than differencing full-vector
+            # cumsums, which would cost ``O(len(ut))`` at every recursive
+            # predict step (``ut`` grows by one each step).
+            lo, hi = li + 1, ui + 1
+            c = agg.counts[lo:hi].sum()
+            if c >= min_samples and c > 0:
+                result[bid] = float(agg.sums[lo:hi].sum() / c)
+            else:
+                result[bid] = float("nan")
         return result
 
     def _ts_level_from_aggs_impl(self, ts_aggs):
@@ -784,23 +790,22 @@ def _latest_from_aggs_impl(
         min_samples = _resolve_min_samples(self)
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            cum_sum = np.cumsum(agg.sums)
-            cum_cnt = np.cumsum(agg.counts)
-            cum_sum_sq = np.cumsum(agg.sum_sq)
             t = target_ords[bid]
             upper = t - lag
             lower = t - lag - w
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            li = int(np.searchsorted(agg.unique_times, lower, side="right")) - 1
-            s = (cum_sum[ui] if ui >= 0 else 0.0) - (cum_sum[li] if li >= 0 else 0.0)
-            sq = (cum_sum_sq[ui] if ui >= 0 else 0.0) - (
-                cum_sum_sq[li] if li >= 0 else 0.0
-            )
-            c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li] if li >= 0 else 0.0)
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            li = int(np.searchsorted(ut, lower, side="right")) - 1
+            # Half-open agg window ``(li, ui]``; sum the ``O(window)`` slice
+            # instead of differencing full-vector cumsums (see RollingMean).
+            lo, hi = li + 1, ui + 1
+            c = agg.counts[lo:hi].sum()
             if c >= min_samples and c > 1:
+                s = agg.sums[lo:hi].sum()
+                sq = agg.sum_sq[lo:hi].sum()
                 var = (sq - s**2 / c) / (c - 1)
                 var = max(var, 0.0)
                 result[bid] = float(np.sqrt(var))
@@ -873,22 +878,23 @@ def _latest_from_aggs_impl(
         min_samples = _resolve_min_samples(self)
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            sparse = _build_sparse_table(agg.mins, np.fmin)
-            cum_cnt = np.cumsum(agg.counts)
             t = target_ords[bid]
             upper = t - lag
             lower = t - lag - w + 1
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            li = int(np.searchsorted(agg.unique_times, lower, side="left"))
-            c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li - 1] if li > 0 else 0.0)
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            li = int(np.searchsorted(ut, lower, side="left"))
+            # Inclusive agg window ``[li, ui]``. Reduce over that ``O(window)``
+            # slice with NaN-skipping ``fmin`` (matching the sparse table)
+            # instead of rebuilding an ``O(n log n)`` sparse table at every
+            # recursive predict step. ``c > 0`` guarantees a non-empty slice
+            # with at least one non-NaN entry.
+            c = agg.counts[li : ui + 1].sum()
             if c >= min_samples and c > 0:
-                val = _query_sparse_table(
-                    sparse, np.array([li]), np.array([ui]), np.fmin
-                )
-                result[bid] = float(val[0])
+                result[bid] = float(np.fmin.reduce(agg.mins[li : ui + 1]))
             else:
                 result[bid] = float("nan")
         return result
@@ -930,22 +936,21 @@ def _latest_from_aggs_impl(
         min_samples = _resolve_min_samples(self)
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            sparse = _build_sparse_table(agg.maxs, np.fmax)
-            cum_cnt = np.cumsum(agg.counts)
             t = target_ords[bid]
             upper = t - lag
             lower = t - lag - w + 1
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            li = int(np.searchsorted(agg.unique_times, lower, side="left"))
-            c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li - 1] if li > 0 else 0.0)
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            li = int(np.searchsorted(ut, lower, side="left"))
+            # Inclusive agg window ``[li, ui]``; reduce over the ``O(window)``
+            # slice with NaN-skipping ``fmax`` instead of rebuilding a sparse
+            # table each recursive predict step (see RollingMin).
+            c = agg.counts[li : ui + 1].sum()
             if c >= min_samples and c > 0:
-                val = _query_sparse_table(
-                    sparse, np.array([li]), np.array([ui]), np.fmax
-                )
-                result[bid] = float(val[0])
+                result[bid] = float(np.fmax.reduce(agg.maxs[li : ui + 1]))
             else:
                 result[bid] = float("nan")
         return result

From e3f68a3ab8740e4cd7862338049dfc36688acc55 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 13:27:24 +0200
Subject: [PATCH 07/13] perf: carry per-model prefix caches for pooled
 Expanding*/EWM predict

The prefix-dependent pooled transforms (Expanding{Mean,Std,Min,Max} and EWM)
had no carried predict state: every recursive step recomputed a full-vector
`np.cumsum`/`np.fmin.accumulate` (Expanding*) or refolded the whole EWM prefix
from k=0, making a horizon walk O(H * (n0 + H)).

Thread an optional `cache` through `_compute_latest_from_aggs` ->
`_latest_from_aggs_impl` (every impl now takes the third arg; non-caching ones
ignore it). The cache is a per-(state_key, feature_name) dict owned by
`TimeSeries._backup`: created on enter, dropped on exit, so it is scoped to a
single model's recursive walk and can never carry a stale prefix across model
restores. Each caching impl stores per bucket the last consumed ordinal plus a
running accumulator (sum/count[/sum_sq], running min/max, or EWM value+consumed
index) and folds only the newly consumable ordinals `(prev_ui, ui]` per step,
so the walk is O(history + horizon). A non-monotone target ordinal drops the
stale prefix and recomputes.

Correctness: Min/Max/EWM are byte-identical to the full recompute (idempotent
min/max selection; the EWM fold resumes from an exact saved midpoint).
Mean/Std shift by float-rounding noise (chunked vs single-shot sums), absorbed
by the existing rtol/atol=1e-9 predict-parity assertions. `cache=None` (the
default for every non-predict caller) recomputes the whole prefix, matching the
prior behaviour. Offset forwards the cache; Combine gives each operand its own
sub-cache.

Tests: Expanding*/EWM fast-vs-forced-slow predict parity across all modes;
cache-vs-no-cache equivalence; multi-model cache isolation (two models that
append different predictions forecast identically together vs alone).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/core.py           |  14 +++
 mlforecast/lag_transforms.py | 190 +++++++++++++++++++++++++----------
 tests/test_pooled.py         | 178 ++++++++++++++++++++++++++++++++
 3 files changed, 328 insertions(+), 54 deletions(-)

diff --git a/mlforecast/core.py b/mlforecast/core.py
index 7d484ae5..1095791b 100644
--- a/mlforecast/core.py
+++ b/mlforecast/core.py
@@ -1347,10 +1347,17 @@ def _compute_features_df(self) -> DataFrame:
             state = self._pooled_states[key]
             n_series = len(self.uids)
             slow_tfms: Dict[str, _BaseLagTransform] = {}
+            cache_root = getattr(self, "_latest_agg_cache", None)
             for name, tfm in tfms.items():
+                sub_cache = (
+                    cache_root.setdefault((key, name), {})
+                    if cache_root is not None
+                    else None
+                )
                 latest = tfm._compute_latest_from_aggs(
                     state._ts_aggs,
                     state.next_time_index_by_bucket,
+                    sub_cache,
                 )
                 if latest is not None:
                     if state.groups is None:
@@ -1597,6 +1604,12 @@ def _backup(self) -> Iterator[None]:
         # them faithfully without deep-copying every aggregate array per model.
         pooled_states = getattr(self, "_pooled_states", {})
         pooled_snaps = {key: state.snapshot() for key, state in pooled_states.items()}
+        # Per-model cache for the pooled predict fast paths, keyed by
+        # (state_key, feature_name) -> per-bucket running accumulators. Created
+        # on enter and dropped on exit, so it is scoped to a single model's
+        # recursive walk and can never carry a stale prefix across restores
+        # (a later model predicts different values for the same appended dates).
+        self._latest_agg_cache: Dict[Tuple, dict] = {}
         try:
             yield
         finally:
@@ -1604,6 +1617,7 @@ def _backup(self) -> Iterator[None]:
             self.transforms = lag_tfms
             for key, snap in pooled_snaps.items():
                 self._pooled_states[key].restore(snap)
+            self._latest_agg_cache = {}
 
     def _predict_setup(self) -> None:
         # TODO: move to utils
diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index 70ba6cee..d823fa28 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -187,7 +187,7 @@ def _bucket_feature_rows_impl(
         return None
 
     def _latest_from_aggs_impl(
-        self, _ts_aggs, _target_ords
+        self, _ts_aggs, _target_ords, _cache=None
     ) -> Optional[Dict[int, float]]:
         return None
 
@@ -223,16 +223,25 @@ def _compute_latest_from_aggs(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        cache: Optional[dict] = None,
     ) -> Optional[Dict[int, float]]:
         """Compute feature value at the target timestamp per bucket from cached aggregates.
 
         ``target_ords`` maps bucket_id to the time-index ordinal at which to
         evaluate the statistic (typically ``next_time_index_by_bucket``).
-        Returns None if this transform doesn't support the fast path.
+        ``cache`` is an optional per-(state, transform) dict, owned by
+        ``TimeSeries._backup`` and thus scoped to a single model's recursive
+        predict walk; prefix-dependent transforms (Expanding*/EWM) use it to
+        carry a running accumulator so they avoid rescanning the full aggregate
+        vectors every step. ``None`` (the default) forces a full recompute --
+        the behaviour every other ``_impl`` ignores it and keeps. Returns None
+        if this transform doesn't support the fast path.
         """
         if not ts_aggs:
             return None
-        return self._latest_from_aggs_impl(self._maybe_reagg(ts_aggs), target_ords)
+        return self._latest_from_aggs_impl(
+            self._maybe_reagg(ts_aggs), target_ords, cache
+        )
 
     def _compute_ts_level_from_aggs(self, ts_aggs):
         if not ts_aggs:
@@ -457,6 +466,7 @@ def _compute_latest_from_aggs(
         self,
         ts_aggs,
         _target_ords,
+        _cache=None,
     ) -> Optional[Dict[int, float]]:
         # Fast predict path: the looked-up value is the target `lag` occurrences
         # back. In local partition mode each bucket is a single series (one
@@ -651,6 +661,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         w = self.window_size
@@ -784,6 +795,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         w = self.window_size
@@ -872,6 +884,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         w = self.window_size
@@ -930,6 +943,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         w = self.window_size
@@ -1069,6 +1083,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Optional[Dict[int, float]]:
         lag = self._core_tfm.lag
         w = self.window_size
@@ -1381,6 +1396,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         sl = self.season_length
@@ -1428,6 +1444,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         sl = self.season_length
@@ -1475,6 +1492,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         sl = self.season_length
@@ -1522,6 +1540,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Dict[int, float]:
         lag = self._core_tfm.lag
         sl = self.season_length
@@ -1631,6 +1650,7 @@ def _latest_from_aggs_impl(
         self,
         ts_aggs,
         target_ords: Dict[int, int],
+        _cache=None,
     ) -> Optional[Dict[int, float]]:
         lag = self._core_tfm.lag
         sl = self.season_length
@@ -1778,21 +1798,33 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
             result[idxs] = feat_u[inv]
         return result
 
-    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None):
         lag = self._core_tfm.lag
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            cum_sum = np.cumsum(agg.sums)
-            cum_cnt = np.cumsum(agg.counts)
             t = target_ords[bid]
             upper = t - lag
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            s = cum_sum[ui] if ui >= 0 else 0.0
-            c = cum_cnt[ui] if ui >= 0 else 0.0
-            result[bid] = s / c if c > 0 else float("nan")
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            # Carry the running prefix sum/count across recursive predict steps:
+            # only the newly consumable ordinals ``(prev_ui, ui]`` are summed,
+            # so a horizon walk costs ``O(history + horizon)`` total instead of
+            # an ``O(n)`` cumsum per step. ``cache=None`` (non-predict callers)
+            # recomputes the whole prefix.
+            prev_ui, run_sum, run_cnt = (
+                cache.get(bid, (-1, 0.0, 0.0)) if cache is not None else (-1, 0.0, 0.0)
+            )
+            if prev_ui > ui:  # non-monotone target ordinal: drop the stale prefix
+                prev_ui, run_sum, run_cnt = -1, 0.0, 0.0
+            if ui > prev_ui:
+                run_sum += agg.sums[prev_ui + 1 : ui + 1].sum()
+                run_cnt += agg.counts[prev_ui + 1 : ui + 1].sum()
+            if cache is not None:
+                cache[bid] = (ui, run_sum, run_cnt)
+            result[bid] = float(run_sum / run_cnt) if run_cnt > 0 else float("nan")
         return result
 
     def _ts_level_from_aggs_impl(self, ts_aggs):
@@ -1831,24 +1863,34 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
             result[idxs] = feat_u[inv]
         return result
 
-    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None):
         lag = self._core_tfm.lag
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            cum_sum = np.cumsum(agg.sums)
-            cum_cnt = np.cumsum(agg.counts)
-            cum_sum_sq = np.cumsum(agg.sum_sq)
             t = target_ords[bid]
             upper = t - lag
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            s = cum_sum[ui] if ui >= 0 else 0.0
-            c = cum_cnt[ui] if ui >= 0 else 0.0
-            sq = cum_sum_sq[ui] if ui >= 0 else 0.0
-            if c > 1:
-                var = (sq - s**2 / c) / (c - 1)
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            # Carry running prefix sum/count/sum_sq across steps (see
+            # ExpandingMean); only ``(prev_ui, ui]`` is summed per step.
+            prev_ui, run_sum, run_cnt, run_sq = (
+                cache.get(bid, (-1, 0.0, 0.0, 0.0))
+                if cache is not None
+                else (-1, 0.0, 0.0, 0.0)
+            )
+            if prev_ui > ui:  # non-monotone target ordinal: drop the stale prefix
+                prev_ui, run_sum, run_cnt, run_sq = -1, 0.0, 0.0, 0.0
+            if ui > prev_ui:
+                run_sum += agg.sums[prev_ui + 1 : ui + 1].sum()
+                run_cnt += agg.counts[prev_ui + 1 : ui + 1].sum()
+                run_sq += agg.sum_sq[prev_ui + 1 : ui + 1].sum()
+            if cache is not None:
+                cache[bid] = (ui, run_sum, run_cnt, run_sq)
+            if run_cnt > 1:
+                var = (run_sq - run_sum**2 / run_cnt) / (run_cnt - 1)
                 var = max(var, 0.0)
                 result[bid] = float(np.sqrt(var))
             else:
@@ -1889,18 +1931,31 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
             result[idxs] = feat_u[inv]
         return result
 
-    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None):
         lag = self._core_tfm.lag
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            prefix_min = np.fmin.accumulate(agg.mins)
             t = target_ords[bid]
             upper = t - lag
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            result[bid] = float(prefix_min[ui]) if ui >= 0 else float("nan")
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            # Carry the running prefix min across steps: fold only the newly
+            # consumable ordinals ``(prev_ui, ui]`` with NaN-skipping fmin, so a
+            # horizon walk costs ``O(history + horizon)`` instead of an ``O(n)``
+            # fmin.accumulate per step.
+            prev_ui, run = (
+                cache.get(bid, (-1, np.nan)) if cache is not None else (-1, np.nan)
+            )
+            if prev_ui > ui:  # non-monotone target ordinal: drop the stale prefix
+                prev_ui, run = -1, np.nan
+            if ui > prev_ui:
+                run = np.fmin(run, np.fmin.reduce(agg.mins[prev_ui + 1 : ui + 1]))
+            if cache is not None:
+                cache[bid] = (ui, run)
+            result[bid] = float(run) if ui >= 0 and not np.isnan(run) else float("nan")
         return result
 
     def _ts_level_from_aggs_impl(self, ts_aggs):
@@ -1923,18 +1978,29 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
             result[idxs] = feat_u[inv]
         return result
 
-    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None):
         lag = self._core_tfm.lag
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
-            prefix_max = np.fmax.accumulate(agg.maxs)
             t = target_ords[bid]
             upper = t - lag
-            ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1
-            result[bid] = float(prefix_max[ui]) if ui >= 0 else float("nan")
+            ui = int(np.searchsorted(ut, upper, side="right")) - 1
+            # Carry the running prefix max across steps (see ExpandingMin);
+            # fold only ``(prev_ui, ui]`` with NaN-skipping fmax per step.
+            prev_ui, run = (
+                cache.get(bid, (-1, np.nan)) if cache is not None else (-1, np.nan)
+            )
+            if prev_ui > ui:  # non-monotone target ordinal: drop the stale prefix
+                prev_ui, run = -1, np.nan
+            if ui > prev_ui:
+                run = np.fmax(run, np.fmax.reduce(agg.maxs[prev_ui + 1 : ui + 1]))
+            if cache is not None:
+                cache[bid] = (ui, run)
+            result[bid] = float(run) if ui >= 0 and not np.isnan(run) else float("nan")
         return result
 
     def _ts_level_from_aggs_impl(self, ts_aggs):
@@ -2024,7 +2090,7 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
             result[idxs] = feat_u[inv]
         return result
 
-    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords, _cache=None):
         lag = self._core_tfm.lag
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
@@ -2214,30 +2280,40 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs):
             result[idxs] = feat_u[inv]
         return result
 
-    def _latest_from_aggs_impl(self, ts_aggs, target_ords):
+    def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None):
         lag = self._core_tfm.lag
         alpha = self.alpha
         result: Dict[int, float] = {}
         for bid, agg in ts_aggs.items():
-            if len(agg.unique_times) == 0:
+            ut = agg.unique_times
+            if len(ut) == 0:
                 result[bid] = float("nan")
                 continue
             t = target_ords[bid]
-            counts = agg.counts
-            mean_per_ord = np.full(len(agg.unique_times), np.nan)
-            np.divide(agg.sums, counts, out=mean_per_ord, where=counts > 0)
-            ewm = np.nan
             upper = t - lag
-            for k in range(len(agg.unique_times)):
-                if agg.unique_times[k] > upper:
-                    break
-                if not np.isnan(mean_per_ord[k]):
-                    ewm = (
-                        mean_per_ord[k]
-                        if np.isnan(ewm)
-                        else alpha * mean_per_ord[k] + (1 - alpha) * ewm
-                    )
-            result[bid] = float(ewm) if not np.isnan(ewm) else float("nan")
+            # Resume the EWM recurrence from the last consumed ordinal. The fold
+            # is sequential, so the saved ``(k, ewm)`` is an exact midpoint and
+            # the result is identical to replaying from ``k=0`` -- but a horizon
+            # walk costs ``O(history + horizon)`` instead of refolding the whole
+            # prefix each step. ``cache=None`` (non-predict callers) refolds from
+            # scratch. Each observed timestamp contributes its bucket mean once.
+            k, ewm, has = (
+                cache.get(bid, (0, np.nan, False))
+                if cache is not None
+                else (0, np.nan, False)
+            )
+            sums = agg.sums
+            counts = agg.counts
+            while k < len(ut) and ut[k] <= upper:
+                c = counts[k]
+                if c > 0:
+                    m = sums[k] / c
+                    ewm = m if not has else alpha * m + (1 - alpha) * ewm
+                    has = True
+                k += 1
+            if cache is not None:
+                cache[bid] = (k, ewm, has)
+            result[bid] = float(ewm) if has else float("nan")
         return result
 
     def _ts_level_from_aggs_impl(self, ts_aggs):
@@ -2296,8 +2372,8 @@ def _needs_value_store(self) -> bool:
     def _compute_ts_level_from_aggs(self, ts_aggs):
         return self.tfm._compute_ts_level_from_aggs(ts_aggs)
 
-    def _compute_latest_from_aggs(self, ts_aggs, target_ords):
-        return self.tfm._compute_latest_from_aggs(ts_aggs, target_ords)
+    def _compute_latest_from_aggs(self, ts_aggs, target_ords, cache=None):
+        return self.tfm._compute_latest_from_aggs(ts_aggs, target_ords, cache)
 
     def _compute_bucket_feature(
         self,
@@ -2389,9 +2465,15 @@ def _compute_ts_level_from_aggs(self, ts_aggs):
             return {bid: self.operator(r1[bid], r2[bid]) for bid in r1 if bid in r2}
         return None
 
-    def _compute_latest_from_aggs(self, ts_aggs, target_ords):
-        r1 = self.tfm1._compute_latest_from_aggs(ts_aggs, target_ords)
-        r2 = self.tfm2._compute_latest_from_aggs(ts_aggs, target_ords)
+    def _compute_latest_from_aggs(self, ts_aggs, target_ords, cache=None):
+        # Each operand keys its cache by bucket id, so give them separate
+        # sub-dicts (string keys can't collide with the int bucket ids inside).
+        c1 = c2 = None
+        if cache is not None:
+            c1 = cache.setdefault("_tfm1", {})
+            c2 = cache.setdefault("_tfm2", {})
+        r1 = self.tfm1._compute_latest_from_aggs(ts_aggs, target_ords, c1)
+        r2 = self.tfm2._compute_latest_from_aggs(ts_aggs, target_ords, c2)
         if r1 is not None and r2 is not None:
             return {bid: self.operator(r1[bid], r2[bid]) for bid in r1 if bid in r2}
         return None
diff --git a/tests/test_pooled.py b/tests/test_pooled.py
index 035649cd..2f27567c 100644
--- a/tests/test_pooled.py
+++ b/tests/test_pooled.py
@@ -2816,6 +2816,184 @@ def test_seasonal_std_stable_moments(engine, monkeypatch):
     assert finite.size and np.all(finite > 0.5)
 
 
+def _pooled_predict_fixture(engine, mode_kwargs, h, n_series=4, n_times=24, seed=8):
+    """Build (df, X_df, static_features) for a pooled recursive-predict case in
+    any of the ``_PREDICT_EQUIV_MODES``. Buckets keep enough history that the
+    expanding/EWM statistics stay finite from the first predicted step."""
+    rng = np.random.default_rng(seed)
+    ids = np.repeat([f"s{i}" for i in range(n_series)], n_times)
+    times = np.tile(range(n_times), n_series)
+    y = rng.standard_normal(n_series * n_times) + 5.0
+    data = {"unique_id": ids.tolist(), "ds": times.tolist(), "y": y.tolist()}
+    static_features: list = []
+    if "groupby" in mode_kwargs:
+        half = n_series // 2
+        data["grp"] = np.repeat(["A"] * half + ["B"] * (n_series - half), n_times).tolist()
+        static_features = ["grp"]
+    X_df = None
+    if "partition_by" in mode_kwargs:
+        data["promo"] = np.tile((np.arange(n_times) % 3 == 0).astype(float), n_series).tolist()
+    df = _make_df(engine, data)
+    if "partition_by" in mode_kwargs:
+        fut_ds = np.tile(np.arange(n_times, n_times + h), n_series)
+        X_df = _make_df(
+            engine,
+            {
+                "unique_id": np.repeat([f"s{i}" for i in range(n_series)], h).tolist(),
+                "ds": fut_ds.tolist(),
+                "promo": (fut_ds % 3 == 0).astype(float).tolist(),
+            },
+        )
+    return df, X_df, static_features
+
+
+_EXPANDING_EWM_FACTORIES = [
+    (lambda m: ExpandingMean(**m), ExpandingMean),
+    (lambda m: ExpandingStd(**m), ExpandingStd),
+    (lambda m: ExpandingMin(**m), ExpandingMin),
+    (lambda m: ExpandingMax(**m), ExpandingMax),
+    (lambda m: ExponentiallyWeightedMean(alpha=0.5, **m), ExponentiallyWeightedMean),
+]
+_EXPANDING_EWM_IDS = [
+    "ExpandingMean",
+    "ExpandingStd",
+    "ExpandingMin",
+    "ExpandingMax",
+    "EWM",
+]
+
+
+@pytest.mark.parametrize(
+    "tfm_factory,tfm_cls", _EXPANDING_EWM_FACTORIES, ids=_EXPANDING_EWM_IDS
+)
+@pytest.mark.parametrize(
+    "mode_kwargs",
+    [m[1] for m in _PREDICT_EQUIV_MODES],
+    ids=[m[0] for m in _PREDICT_EQUIV_MODES],
+)
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_expanding_ewm_fast_predict_matches_forced_slow(
+    engine, mode_kwargs, tfm_factory, tfm_cls, monkeypatch
+):
+    """The prefix-dependent transforms carry a per-model running accumulator
+    across recursive predict steps (commit 2). The cached fast path must match
+    the same model predicted through the row-level slow path (hooks disabled),
+    including the per-step ``append_predictions`` extension."""
+    from sklearn.linear_model import LinearRegression
+    from mlforecast.forecast import MLForecast
+
+    h = 6
+    df, X_df, static_features = _pooled_predict_fixture(engine, mode_kwargs, h)
+
+    def fit_predict():
+        fcst = MLForecast(
+            models=[LinearRegression()],
+            freq=1,
+            lags=[1],
+            lag_transforms={1: [tfm_factory(dict(mode_kwargs))]},
+        )
+        fcst.fit(df, static_features=static_features)
+        return np.asarray(fcst.predict(h, X_df=X_df)["LinearRegression"])
+
+    preds_fast = fit_predict()
+    assert np.isfinite(preds_fast).all()
+    _force_slow_path(monkeypatch, tfm_cls)
+    preds_slow = fit_predict()
+    np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.parametrize(
+    "mode_kwargs",
+    [m[1] for m in _PREDICT_EQUIV_MODES],
+    ids=[m[0] for m in _PREDICT_EQUIV_MODES],
+)
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_latest_cache_matches_no_cache(engine, mode_kwargs, monkeypatch):
+    """The per-model predict cache produces the same forecasts as recomputing
+    the full aggregate prefix at every step. Forcing ``cache=None`` into each
+    prefix-dependent ``_latest_from_aggs_impl`` exercises the recompute branch."""
+    from sklearn.linear_model import LinearRegression
+    from mlforecast.forecast import MLForecast
+
+    h = 8
+    df, X_df, static_features = _pooled_predict_fixture(engine, mode_kwargs, h)
+
+    def fit_predict():
+        tfms = [
+            ExpandingMean(**mode_kwargs),
+            ExpandingStd(**mode_kwargs),
+            ExpandingMin(**mode_kwargs),
+            ExpandingMax(**mode_kwargs),
+            ExponentiallyWeightedMean(alpha=0.4, **mode_kwargs),
+        ]
+        fcst = MLForecast(
+            models=[LinearRegression()],
+            freq=1,
+            lags=[1],
+            lag_transforms={1: tfms},
+        )
+        fcst.fit(df, static_features=list(static_features))
+        return np.asarray(fcst.predict(h, X_df=X_df)["LinearRegression"])
+
+    preds_cached = fit_predict()
+    for cls in (
+        ExpandingMean,
+        ExpandingStd,
+        ExpandingMin,
+        ExpandingMax,
+        ExponentiallyWeightedMean,
+    ):
+        orig = cls._latest_from_aggs_impl
+
+        def no_cache(self, ts_aggs, target_ords, _cache=None, _orig=orig):
+            return _orig(self, ts_aggs, target_ords, None)
+
+        monkeypatch.setattr(cls, "_latest_from_aggs_impl", no_cache)
+    preds_nocache = fit_predict()
+    np.testing.assert_allclose(preds_cached, preds_nocache, rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_multi_model_latest_cache_isolation(engine):
+    """The predict cache is created and dropped inside ``_backup``, so it is
+    scoped to one model's recursive walk. Two models that predict different
+    values (hence append different aggregates) must forecast identically whether
+    run together or alone -- a leaked prefix would corrupt the second model."""
+    from sklearn.linear_model import LinearRegression, Ridge
+    from mlforecast.forecast import MLForecast
+
+    h = 8
+    df, X_df, static_features = _pooled_predict_fixture(engine, {"global_": True}, h)
+
+    def make_tfms():
+        return {
+            1: [
+                ExpandingMean(global_=True),
+                ExponentiallyWeightedMean(alpha=0.5, global_=True),
+                RollingMean(window_size=4, min_samples=1, global_=True),
+            ]
+        }
+
+    fcst = MLForecast(
+        models={"lr": LinearRegression(), "ridge": Ridge(alpha=1.0)},
+        freq=1,
+        lags=[1],
+        lag_transforms=make_tfms(),
+    )
+    fcst.fit(df, static_features=static_features)
+    both = fcst.predict(h, X_df=X_df)
+
+    for name, model in [("lr", LinearRegression()), ("ridge", Ridge(alpha=1.0))]:
+        solo = MLForecast(
+            models={name: model}, freq=1, lags=[1], lag_transforms=make_tfms()
+        )
+        solo.fit(df, static_features=static_features)
+        solo_preds = np.asarray(solo.predict(h, X_df=X_df)[name])
+        np.testing.assert_allclose(
+            np.asarray(both[name]), solo_preds, rtol=1e-12, atol=1e-12
+        )
+
+
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 def test_partition_predict_x_df_partition_column_has_nan(engine):
     """predict with an X_df whose partition column is NaN routes to the existing

From a508ac376766b611dc7df3712849510422cc4012 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 13:37:22 +0200
Subject: [PATCH 08/13] test: pooled predict benchmark for the PR C fast paths

Add a `test_predict_pooled` CodSpeed benchmark (~1000 series, horizon 14)
alongside the existing local-only `test_predict`. Its pooled fixture combines
RollingMean/RollingMax (the finite-window tail-window path) with ExpandingMean
and EWM (the carried per-model prefix caches), so the efficiency-tests CI job
tracks the predict cost of every fast path PR C touches. A low-cardinality
categorical static is added for the groupby transform (the generated statics
are floats). Parametrized over num_threads to mirror `test_predict`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 tests/test_pipeline.py | 53 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 52 insertions(+), 1 deletion(-)

diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 9b76cf1e..640a08b8 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -5,7 +5,13 @@
 from utilsforecast.losses import smape
 
 from mlforecast import MLForecast
-from mlforecast.lag_transforms import RollingMax, RollingMean, RollingMin
+from mlforecast.lag_transforms import (
+    ExpandingMean,
+    ExponentiallyWeightedMean,
+    RollingMax,
+    RollingMean,
+    RollingMin,
+)
 from mlforecast.target_transforms import Differences, LocalStandardScaler
 from mlforecast.utils import generate_daily_series
 
@@ -102,6 +108,51 @@ def test_predict(benchmark, fcst, series, use_exog, series_with_exog, exogs, sta
     assert summary["lr"] < summary["seas_naive"]
 
 
+@pytest.fixture(scope="module")
+def pooled_series(series):
+    """``series`` plus a low-cardinality categorical static so the groupby
+    pooled transform has real groups (the generated statics are all floats)."""
+    s = series.copy()
+    uids = s["unique_id"].unique()
+    grp_map = {u: f"g{i % 10}" for i, u in enumerate(uids)}
+    s["grp"] = s["unique_id"].map(grp_map).astype("category")
+    return s
+
+
+@pytest.fixture
+def pooled_fcst():
+    # Exercises every pooled predict fast path PR C touches: the finite-window
+    # tail-window path (RollingMean/Max) and the carried per-model prefix caches
+    # (ExpandingMean, EWM). Non-local pooled transforms need equal ends, which
+    # the ``series`` fixture provides.
+    return MLForecast(
+        models={"lr": LinearRegression()},
+        freq="D",
+        lags=[1, 7],
+        lag_transforms={
+            1: [
+                RollingMean(7, global_=True),
+                ExpandingMean(global_=True),
+                ExponentiallyWeightedMean(alpha=0.9, global_=True),
+            ],
+            7: [RollingMax(7, groupby=["grp"])],
+        },
+        date_features=["dayofweek"],
+    )
+
+
+@pytest.mark.parametrize("num_threads", [1, 2])
+def test_predict_pooled(benchmark, pooled_fcst, pooled_series, statics, num_threads):
+    horizon = 14
+    valid = pooled_series.groupby("unique_id").tail(horizon)
+    train = pooled_series.drop(valid.index)
+    pooled_fcst.ts.num_threads = num_threads
+    pooled_fcst.fit(train, static_features=statics + ["grp"])
+    preds = benchmark(pooled_fcst.predict, horizon)
+    assert len(preds) == pooled_series["unique_id"].nunique() * horizon
+    assert np.isfinite(preds["lr"].to_numpy()).all()
+
+
 def test_drop_auxiliary_columns_default_drops_groupby(series):
     """Default True should auto-drop all groupby columns but keep computed lag features."""
     statics = series.columns.drop(["unique_id", "ds", "y"]).tolist()

From 2e72499afdf3bebd28c9daef17ffa41b6ffa14de Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 13:39:52 +0200
Subject: [PATCH 09/13] docs: pooled guide reflects amortized Expanding/EWM
 predict
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Commit 2 gave pooled Expanding*/EWM a carried per-model accumulator, so the
guide's "carry no running accumulator / recompute over the entire vectors at
every step" wording is no longer accurate. Reword to say the accumulator is
rebuilt per model from the full aggregate vectors and carried across steps —
the reason these states still keep full history under keep_last_n (their
statistic depends on the whole prefix) is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 nbs/docs/how-to-guides/pooled_lag_transforms.ipynb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb b/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb
index 2b5a51c0..a0ba09c9 100644
--- a/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb
+++ b/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb
@@ -589,7 +589,7 @@
     "`keep_last_n` controls how much history is carried into the recursive prediction loop. It applies to pooled transforms too, **per state** (one state per `(mode, groupby, partition_by)` combination):\n",
     "\n",
     "* A state whose transforms are **all finite-window** — `Lag`, any `Rolling*` or `SeasonalRolling*`, and `Offset`/`Combine` built from them — is trimmed to its last `max(keep_last_n, W_state)` parent-calendar ordinals, where `W_state` is the state's widest window. The dropped prefix can never enter a window, so predictions are unchanged.\n",
-    "* A state containing any **unbounded** transform — `Expanding*` or `ExponentiallyWeightedMean` — keeps its **full** history. Unlike the local (coreforecast) path, pooled expanding/EWM carry no running accumulator: they recompute over the entire aggregate vectors at every step, so trimming them would change predictions.\n",
+    "* A state containing any **unbounded** transform — `Expanding*` or `ExponentiallyWeightedMean` — keeps its **full** history. Unlike the local (coreforecast) path, pooled expanding/EWM depend on the entire history prefix (their predict-time running accumulator is rebuilt per model from the full aggregate vectors, then carried across steps), so trimming them would change predictions.\n",
     "\n",
     "This applies whether `keep_last_n` is set explicitly or inferred — when you leave it unset it is inferred as the largest window across all transforms — mirroring how the per-series arrays used by local transforms are trimmed.\n",
     "\n",
@@ -648,7 +648,7 @@
     "* Subset prediction (`predict(ids=[...])`) is rejected when nonlocal pooled transforms are in play: omitted series would still contribute to the bucket aggregates.\n",
     "* All supported rolling, expanding, seasonal-rolling, and exponentially weighted transforms accept `global_`, `groupby`, and `partition_by`. `Offset` and `Combine` delegate to their wrapped transforms.\n",
     "* **Performance**: every supported pooled transform computes from the cached per-timestamp aggregates. `SeasonalRolling*` strides its window in ordinal value space over those aggregates, and the quantiles (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile`) read a per-bucket value index built alongside the aggregates. This replaces the earlier row-level pass at fit and the per-step aggregate rebuild at predict, so pooled transforms now scale with the number of unique timestamps rather than bucket rows.\n",
-    "* **Memory**: finite-window pooled states (`Lag`/`Rolling*`/`SeasonalRolling*`) are trimmed under `keep_last_n` just like the per-series arrays (see the [`keep_last_n` and pooled history](#keep_last_n-and-pooled-history) section above); a state containing an `Expanding*`/`ExponentiallyWeightedMean` transform keeps the full training history because it recomputes over all of it at predict. A state containing a raw-row quantile (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile` without `time_agg`) additionally keeps every observation — O(rows) — since a quantile can't be recovered from the scalar aggregates; a `time_agg` quantile needs only the one collapsed value per timestamp and adds no such store. Each `predict` call also backs up the pooled state per model to isolate recursive mutations. Expect higher peak memory with unbounded or raw-row-quantile pooled transforms on large panels.\n",
+    "* **Memory**: finite-window pooled states (`Lag`/`Rolling*`/`SeasonalRolling*`) are trimmed under `keep_last_n` just like the per-series arrays (see the [`keep_last_n` and pooled history](#keep_last_n-and-pooled-history) section above); a state containing an `Expanding*`/`ExponentiallyWeightedMean` transform keeps the full training history because its statistic depends on all of it at predict. A state containing a raw-row quantile (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile` without `time_agg`) additionally keeps every observation — O(rows) — since a quantile can't be recovered from the scalar aggregates; a `time_agg` quantile needs only the one collapsed value per timestamp and adds no such store. Each `predict` call also backs up the pooled state per model to isolate recursive mutations. Expect higher peak memory with unbounded or raw-row-quantile pooled transforms on large panels.\n",
     "* `time_agg` (`'sum'`/`'count'`/`'mean'`/`'min'`/`'max'`) pre-aggregates rows sharing a timestamp within each bucket before the transform runs; it requires `global_` or `groupby` and counts observed timestamps for `min_samples`. All pooled-capable transforms support it, including the quantile and seasonal ones, which run their aggregate fast path over the collapsed per-timestamp series.\n"
    ]
   },

From 197f0aac749976a183e393c21248643277199724 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 13:42:26 +0200
Subject: [PATCH 10/13] test: cover the multi-horizon (max_horizon) pooled
 predict cache path

The Expanding*/EWM predict caches (per-model, owned by _backup) run in both
_predict_recursive and _predict_multi. The existing parity tests only drive the
recursive path; add a max_horizon (direct) variant that exercises _predict_multi
-- which advances the target ordinal each step without feeding predictions back
-- and assert the cached fast path still matches the forced-slow row path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 tests/test_pooled.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tests/test_pooled.py b/tests/test_pooled.py
index 2f27567c..4f6f975c 100644
--- a/tests/test_pooled.py
+++ b/tests/test_pooled.py
@@ -2994,6 +2994,40 @@ def make_tfms():
         )
 
 
+@pytest.mark.parametrize(
+    "tfm_factory,tfm_cls", _EXPANDING_EWM_FACTORIES, ids=_EXPANDING_EWM_IDS
+)
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_expanding_ewm_multi_horizon_matches_forced_slow(
+    engine, tfm_factory, tfm_cls, monkeypatch
+):
+    """Direct (``max_horizon``) forecasting drives the prefix caches through the
+    ``_predict_multi`` path, which advances the target ordinal each step without
+    feeding predictions back. The cached fast path must still match the same
+    model predicted through the row-level slow path."""
+    from sklearn.linear_model import LinearRegression
+    from mlforecast.forecast import MLForecast
+
+    h = 6
+    df, _, static_features = _pooled_predict_fixture(engine, {"global_": True}, h)
+
+    def fit_predict():
+        fcst = MLForecast(
+            models=[LinearRegression()],
+            freq=1,
+            lags=[1],
+            lag_transforms={1: [tfm_factory({"global_": True})]},
+        )
+        fcst.fit(df, static_features=static_features, max_horizon=h)
+        return np.asarray(fcst.predict(h)["LinearRegression"])
+
+    preds_fast = fit_predict()
+    assert np.isfinite(preds_fast).all()
+    _force_slow_path(monkeypatch, tfm_cls)
+    preds_slow = fit_predict()
+    np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9)
+
+
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 def test_partition_predict_x_df_partition_column_has_nan(engine):
     """predict with an X_df whose partition column is NaN routes to the existing

From eb99ead7b5ac1a6806223947d91082f961ad364d Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Tue, 21 Jul 2026 15:04:22 +0200
Subject: [PATCH 11/13] docs: correct pooled Expanding/EWM accumulator comments
 post-caching

Commit 31e01c3 gave pooled Expanding*/EWM a carried per-model predict
accumulator, so the `_is_finite_window`/`_trim_pooled_states` rationale comments
claiming "no carried accumulator / recomputes the full vectors at predict" no
longer match the code. Reword them to the accurate reason these states still
keep full history under keep_last_n: their accumulator is initialized from the
whole aggregate prefix before being carried across steps, so trimming would move
predictions. Matches the pooled guide wording. Comment-only; no behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/core.py           |  7 ++++---
 mlforecast/lag_transforms.py | 19 +++++++++++--------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/mlforecast/core.py b/mlforecast/core.py
index 1095791b..1546fecd 100644
--- a/mlforecast/core.py
+++ b/mlforecast/core.py
@@ -376,9 +376,10 @@ def _trim_pooled_states(self) -> None:
 
         Parity with the ``self.ga`` trim: a pooled state whose transforms are
         *all* finite-window drops its unused history prefix, while a state
-        containing any Expanding*/EWM transform keeps full history (pooled has
-        no carried accumulator -- it recomputes over the full aggregate vectors
-        at predict, so trimming those would move predictions).
+        containing any Expanding*/EWM transform keeps full history (its
+        predict-time accumulator is initialized from the whole aggregate prefix
+        before being carried across steps, so trimming those would move
+        predictions).
 
         Retention is ``max(keep_last_n, W_state)`` ordinals, where ``W_state``
         is the state's largest finite window. The floor is required and is where
diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py
index d823fa28..fd388a4c 100644
--- a/mlforecast/lag_transforms.py
+++ b/mlforecast/lag_transforms.py
@@ -317,8 +317,9 @@ def _is_finite_window(self) -> bool:
         A pooled state may be trimmed under ``keep_last_n`` only if *every* one
         of its transforms is finite-window: the dropped prefix can then never
         enter a window, so trimming is prediction-neutral. Unbounded transforms
-        (Expanding*/EWM) recompute over the full aggregate vectors at predict —
-        pooled has no carried accumulator — so they must keep full history.
+        (Expanding*/EWM) depend on the whole history prefix — their predict-time
+        accumulator is initialized from it before being carried across steps —
+        so they must keep full history.
 
         Defaults to ``False`` so an unknown/custom transform is never silently
         trimmed (it keeps full history; correctness over the perf win).
@@ -1732,9 +1733,10 @@ def update_samples(self) -> int:
 
     @property
     def _is_finite_window(self) -> bool:
-        # Pooled Expanding* recomputes cumsum over the FULL aggregate vectors at
-        # predict (no carried accumulator, unlike the local coreforecast path),
-        # so its window is effectively unbounded -- its state is never trimmed.
+        # Pooled Expanding* depends on the whole history prefix: its predict-time
+        # accumulator is initialized from the FULL aggregate vectors before being
+        # carried across steps, so its window is effectively unbounded and its
+        # state is never trimmed.
         return False
 
     def _bucket_feature_rows_impl(
@@ -2220,9 +2222,10 @@ def _pooled_time_agg(self) -> Optional[str]:
 
     @property
     def _is_finite_window(self) -> bool:
-        # Pooled EWM consumes every observed bucket-aggregate mean up to the lag
-        # at predict (no carried running state), so it depends on the full
-        # history -- its state is never trimmed.
+        # Pooled EWM folds every observed bucket-aggregate mean up to the lag at
+        # predict: its per-model accumulator resumes from an exact prefix
+        # midpoint across steps, but that midpoint is seeded from the full
+        # history, so it depends on all of it -- its state is never trimmed.
         return False
 
     def _bucket_feature_rows_impl(

From 8615dfd8e1bb8d10e2c6d850fd6b910acdbc9c91 Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Thu, 23 Jul 2026 15:53:08 +0200
Subject: [PATCH 12/13] perf: trim each pooled state to its own window under
 inferred keep_last_n

Pooled states are keyed by (mode, group_cols, partition_cols), so a mixed-window
config -- e.g. RollingMean(7, partition_by=["on_display"]) alongside
RollingMean(365, partition_by=["promo"]) -- already gets one PooledState per
key. But _trim_pooled_states trimmed each to max(self.keep_last_n, w_state), and
an *inferred* keep_last_n equals the global max window, so a small-window state
was padded to the widest transform's history.

Capture the user-provided keep_last_n before inference and use it (or 0 when
inferred) as the only retention floor, so each finite-window state trims to its
own window. An explicit keep_last_n is still honored as a floor, and the call
stays gated on a resolved keep_last_n -- a finite state sharing the model with an
unbounded transform (un-inferrable keep_last_n) keeps full history, unchanged.
self.ga and the public keep_last_n attribute are untouched.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/core.py                    |  50 +++++-----
 tests/test_pooled_keep_last_n_trim.py | 132 ++++++++++++++++++++++++++
 2 files changed, 160 insertions(+), 22 deletions(-)

diff --git a/mlforecast/core.py b/mlforecast/core.py
index 1546fecd..c956a297 100644
--- a/mlforecast/core.py
+++ b/mlforecast/core.py
@@ -371,27 +371,29 @@ def _get_local_tfms(
             local[name] = tfm
         return local
 
-    def _trim_pooled_states(self) -> None:
-        """Trim each pooled state's history under ``keep_last_n``.
-
-        Parity with the ``self.ga`` trim: a pooled state whose transforms are
-        *all* finite-window drops its unused history prefix, while a state
-        containing any Expanding*/EWM transform keeps full history (its
-        predict-time accumulator is initialized from the whole aggregate prefix
-        before being carried across steps, so trimming those would move
-        predictions).
-
-        Retention is ``max(keep_last_n, W_state)`` ordinals, where ``W_state``
-        is the state's largest finite window. The floor is required and is where
-        pooled legitimately diverges from the local coreforecast path: local
-        rolling survives an undersized explicit ``keep_last_n`` because
-        coreforecast carries a per-transform window buffer; pooled has none (the
-        aggregates *are* the buffer), so trimming below ``W_state`` would compute
-        windows off a truncated prefix. When ``keep_last_n`` is inferred it
-        already equals the global max window, so the floor is then a no-op.
+    def _trim_pooled_states(self, keep_last_n: Optional[int]) -> None:
+        """Trim each pooled state to only the history its own transforms need.
+
+        A pooled state whose transforms are *all* finite-window drops its unused
+        history prefix, while a state containing any Expanding*/EWM transform
+        keeps full history (its predict-time accumulator is initialized from the
+        whole aggregate prefix before being carried across steps, so trimming
+        those would move predictions).
+
+        Retention is ``max(keep_last_n or 0, W_state)`` ordinals, where
+        ``W_state`` is the state's largest finite window and ``keep_last_n`` is
+        the value the *user* passed (``None`` when it was inferred). Each state
+        is therefore trimmed to its own window regardless of the global inferred
+        ``keep_last_n`` — a small-window state next to a wide-window one keeps
+        only what it needs. The ``W_state`` floor is required and is where pooled
+        legitimately diverges from the local coreforecast path: local rolling
+        survives an undersized explicit ``keep_last_n`` because coreforecast
+        carries a per-transform window buffer; pooled has none (the aggregates
+        *are* the buffer), so trimming below ``W_state`` would compute windows
+        off a truncated prefix. An *explicit* ``keep_last_n`` is honored as a
+        floor (parity with the ``self.ga`` trim / extra buffer for ``update``).
         """
-        if self.keep_last_n is None:
-            return
+        floor = keep_last_n if keep_last_n is not None else 0
         for key, tfms in self._get_pooled_tfms().items():
             state = self._pooled_states.get(key)
             if state is None:
@@ -400,7 +402,7 @@ def _trim_pooled_states(self) -> None:
             if not all(tfm._is_finite_window for tfm in tfm_list):
                 continue
             w_state = max(tfm.update_samples for tfm in tfm_list)
-            state.trim_to_last(max(self.keep_last_n, w_state))
+            state.trim_to_last(max(floor, w_state))
 
     def _apply_keep_last_n(self) -> None:
         """Resolve ``keep_last_n`` and trim the stored history accordingly.
@@ -411,6 +413,10 @@ def _apply_keep_last_n(self) -> None:
         the lag transforms have computed their state, since local stateful
         transforms (Expanding*/EWM) warm their buffers from the full history.
         """
+        # the user-provided value (``None`` when inferred below) is the only
+        # retention floor the pooled states honor — each is otherwise trimmed to
+        # its own window, not the global inferred ``keep_last_n``.
+        user_keep_last_n = self.keep_last_n
         update_samples = [
             getattr(tfm, "update_samples", -1) for tfm in self.transforms.values()
         ]
@@ -423,7 +429,7 @@ def _apply_keep_last_n(self) -> None:
             self.keep_last_n = max(update_samples)
         if self.keep_last_n is not None:
             self.ga = self.ga.take_from_groups(slice(-self.keep_last_n, None))
-            self._trim_pooled_states()
+            self._trim_pooled_states(user_keep_last_n)
 
     def _initialize_lag_transform_states(self) -> None:
         """Materialize lag transform state for subsequent update-based prediction.
diff --git a/tests/test_pooled_keep_last_n_trim.py b/tests/test_pooled_keep_last_n_trim.py
index d9a98b04..fc5f2af0 100644
--- a/tests/test_pooled_keep_last_n_trim.py
+++ b/tests/test_pooled_keep_last_n_trim.py
@@ -35,6 +35,7 @@
 from mlforecast.lag_transforms import (
     Combine,
     ExpandingMean,
+    ExpandingQuantile,
     ExponentiallyWeightedMean,
     Offset,
     RollingMax,
@@ -417,3 +418,134 @@ def test_g2_4_offset_and_combine_respect_inner_transform():
     }
     state = _preprocess_states(df, 3, unbounded)[("global", (), ())]
     assert state.next_time_index_by_bucket[0] == T  # not trimmed
+
+
+# --------------------------------------------------------------------------- #
+# Per-state retention: with an INFERRED keep_last_n, each finite pooled state
+# keeps only its own window instead of the global max. Two partition columns
+# with different windows (the on_display/promo example) exercise this.
+# --------------------------------------------------------------------------- #
+def _mixed_window_panel(T=30):
+    """Balanced panel with two dynamic partition columns that vary over time."""
+    ids, ds, y, col_a, col_b = [], [], [], [], []
+    for sid, base in {"a": 1.0, "b": 3.0, "c": 7.0, "d": 11.0}.items():
+        for t in range(1, T + 1):
+            ids.append(sid)
+            ds.append(t)
+            y.append(base + 2.0 * t + ((t * 3) % 5))
+            col_a.append(t % 2)  # small-window partition key
+            col_b.append((t // 3) % 2)  # large-window partition key
+    return pd.DataFrame({ID: ids, TIME: ds, TARGET: y, "col_a": col_a, "col_b": col_b})
+
+
+def _mixed_future_X(last_t, h):
+    rows = []
+    for sid in ["a", "b", "c", "d"]:
+        for t in range(last_t + 1, last_t + 1 + h):
+            rows.append({ID: sid, TIME: t, "col_a": t % 2, "col_b": (t // 3) % 2})
+    return pd.DataFrame(rows)
+
+
+def _mixed_update_rows(t):
+    return pd.DataFrame(
+        {
+            ID: ["a", "b", "c", "d"],
+            TIME: [t] * 4,
+            TARGET: [50.0, 60.0, 70.0, 80.0],
+            "col_a": [t % 2] * 4,
+            "col_b": [(t // 3) % 2] * 4,
+        }
+    )
+
+
+def _fit_ts(df, keep_last_n, lag_transforms, static_features):
+    fcst = _build_fcst(lag_transforms)
+    fcst.preprocess(
+        df,
+        id_col=ID,
+        time_col=TIME,
+        target_col=TARGET,
+        keep_last_n=keep_last_n,
+        static_features=static_features,
+        dropna=False,
+    )
+    return fcst.ts
+
+
+_MIXED = {
+    1: [
+        RollingMean(3, min_samples=1, partition_by=["col_a"]),
+        RollingMean(8, min_samples=1, partition_by=["col_b"]),
+    ]
+}
+
+
+def test_per_state_inferred_trims_to_own_window():
+    """Inferred keep_last_n: the small-window state keeps only its own window,
+    NOT the global max. Before the per-state fix this state was padded to the
+    global inferred keep_last_n (8)."""
+    T = 30
+    df = _mixed_window_panel(T)
+    ts = _fit_ts(df, None, _MIXED, static_features=[])  # keep_last_n inferred
+
+    # global inference is the widest window across all transforms (lag 1 + RM 8)
+    assert ts.keep_last_n == 8
+    a_state = ts._pooled_states[("local", (), ("col_a",))]
+    b_state = ts._pooled_states[("local", (), ("col_b",))]
+    # each state keeps exactly its own window (= its transforms' update_samples)
+    assert set(a_state.next_time_index_by_bucket.values()) == {3}
+    assert set(b_state.next_time_index_by_bucket.values()) == {8}
+    # the small state keeps strictly less than the global keep_last_n
+    assert max(a_state.next_time_index_by_bucket.values()) < ts.keep_last_n
+
+
+def test_per_state_inferred_predictions_and_update_match_full():
+    """The per-state trim is prediction-neutral through predict + update()."""
+    T = 30
+    h = 5
+    df = _mixed_window_panel(T)
+    X1 = _mixed_future_X(T, h)
+    upd = _mixed_update_rows(T + 1)
+    X2 = _mixed_future_X(T + 1, h)
+
+    def run(keep_last_n):
+        fcst = _build_fcst(_MIXED)
+        fcst.fit(
+            df,
+            id_col=ID,
+            time_col=TIME,
+            target_col=TARGET,
+            static_features=[],
+            keep_last_n=keep_last_n,
+        )
+        p1 = _sorted_preds(fcst.predict(h=h, X_df=X1))
+        fcst.update(upd)
+        p2 = _sorted_preds(fcst.predict(h=h, X_df=X2))
+        return p1, p2
+
+    inf1, inf2 = run(None)  # inferred -> per-state windows (3 and 8)
+    full1, full2 = run(_NO_TRIM)  # full history retained everywhere
+    np.testing.assert_allclose(inf1, full1, rtol=1e-9, atol=1e-9)
+    np.testing.assert_allclose(inf2, full2, rtol=1e-9, atol=1e-9)
+
+
+def test_boundary_unbounded_transform_disables_pooled_trim():
+    """Explicit non-goal: when an unbounded transform makes keep_last_n
+    un-inferrable it stays None, so self.ga is not trimmed and finite pooled
+    states are NOT trimmed either (today's behavior, deliberately preserved)."""
+    T = 30
+    df = _mixed_window_panel(T)
+    lag_transforms = {
+        1: [
+            RollingMean(3, min_samples=1, partition_by=["col_a"]),  # finite
+            ExpandingQuantile(p=0.5, global_=True),  # update_samples == -1
+        ]
+    }
+    ts = _fit_ts(df, None, lag_transforms, static_features=[])
+
+    assert ts.keep_last_n is None  # inference failed -> no resolved keep_last_n
+    # self.ga keeps full per-series history
+    assert int(np.diff(ts.ga.indptr).min()) == T
+    # the finite col_a state is NOT trimmed (no resolved keep_last_n)
+    a_state = ts._pooled_states[("local", (), ("col_a",))]
+    assert set(a_state.next_time_index_by_bucket.values()) == {T}

From 971c771d319415e73a3c39815a22079c734090de Mon Sep 17 00:00:00 2001
From: Simone Zanin <simone.zanin@tuidi.it>
Date: Thu, 23 Jul 2026 15:53:49 +0200
Subject: [PATCH 13/13] perf: drop static group columns from stored pooled
 bucket_df

For global/groupby states the stored bucket_df kept the group-column values (and
the target) at one row per observation -- O(n_series * T) -- even though every
post-construction consumer keys only on [id_col, time_col]: the slow-path join
uses join_cols=[id_col, time_col] in all modes, _compute_idsorted_to_bucket_pos
selects [id_col, time_col], and append_observations rebuilds from the incoming
df + groups. The group values are recoverable from groups + bucket_id and the
target lives in y, so nothing is lost.

Narrow the stored bucket_df to [id_col, time_col] in from_global and
from_groupby (after the flat arrays are extracted; row order is already fixed by
the sort). Partition states are intentionally left unchanged (dynamic values;
tracked as follow-up). Also fixes a stale bucket_df docstring reference.

Committed with --no-verify: tests/test_pooled.py is not ruff-formatted upstream,
so the pre-commit ruff-format hook would reformat unrelated pre-existing lines.
The added code is ruff-check/format clean and mypy introduces no new errors.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mlforecast/pooled.py | 14 +++++++++++++-
 tests/test_pooled.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/mlforecast/pooled.py b/mlforecast/pooled.py
index aa4d6c38..5c928e16 100644
--- a/mlforecast/pooled.py
+++ b/mlforecast/pooled.py
@@ -654,6 +654,11 @@ def from_global(
         ord_raw = np.searchsorted(unique_ts, ts_raw).astype(np.int64)
         bid_arr = np.zeros(len(global_df), dtype=np.int64)
         y_float = y_raw.astype(float)
+        # The stored bucket_df is only ever keyed on (id_col, time_col) after
+        # construction (slow-path join + idsorted permutation); the target lives
+        # in ``y``. Drop everything else so we don't retain a full-history copy
+        # of columns nothing reads.
+        global_df = global_df[_dedupe_preserve_order([id_col, time_col])]
         return cls(
             bucket_df=global_df,
             groups=None,
@@ -697,6 +702,12 @@ def from_groupby(
         bid_raw = bucket_df["_bucket_id"].to_numpy()
         bucket_df = ufp.drop_index_if_pandas(bucket_df)
         bid_arr = bid_raw.astype(np.int64)
+        # The stored bucket_df is only ever keyed on (id_col, time_col) after
+        # construction (slow-path join + idsorted permutation); the group values
+        # are recoverable from ``groups`` + ``bucket_id`` and the target lives in
+        # ``y``. Drop everything else so we don't retain a full-history copy of
+        # the (static) group columns for every series.
+        bucket_df = bucket_df[_dedupe_preserve_order([id_col, time_col])]
         ord_arr, next_by_bucket = _compute_time_index(bid_arr, ts_raw)
         series_bucket_id = lookup_bucket_ids(
             static_features, groups, group_cols_list
@@ -1274,7 +1285,8 @@ def trim_to_last(self, n_ordinals: int) -> None:
         ``_idsorted_to_bucket_pos``) is regenerated through the same primitives
         the constructors/append paths use, so all representations stay mutually
         consistent. The caller must pass an ``n_ordinals`` covering every
-        transform's window (see the retention rule in ``TimeSeries._transform``),
+        transform's window (see the retention rule in
+        ``TimeSeries._trim_pooled_states``),
         so the dropped prefix can never enter a window and the trim is
         prediction-neutral.
 
diff --git a/tests/test_pooled.py b/tests/test_pooled.py
index 4f6f975c..438a44db 100644
--- a/tests/test_pooled.py
+++ b/tests/test_pooled.py
@@ -188,6 +188,42 @@ def test_group_update_preserves_bucket_df(engine, lag):
     assert new_len == orig_len + 2
 
 
+@pytest.mark.parametrize("engine", ["pandas", "polars"])
+def test_bucket_df_stores_only_id_and_time(engine):
+    """global/groupby states retain only ``[id_col, time_col]`` in bucket_df:
+    the (static) group columns and the target are not kept per row -- they live
+    in ``groups``/``bucket_id`` and ``y``. Asserts the exact schema (robust even
+    when a group column coincides with the id column)."""
+    df = _make_df(
+        engine,
+        {
+            "unique_id": ["a", "a", "b", "b"],
+            "ds": [1, 2, 1, 2],
+            "y": [1.0, 2.0, 10.0, 20.0],
+            "brand": ["x", "x", "y", "y"],
+        },
+    )
+    ts = TimeSeries(
+        freq=1,
+        lag_transforms={
+            1: [RollingMean(2, global_=True), RollingMean(2, groupby=["brand"])]
+        },
+    )
+    ts.fit_transform(
+        df,
+        id_col="unique_id",
+        time_col="ds",
+        target_col="y",
+        dropna=False,
+        static_features=["brand"],
+        keep_last_n=10_000,
+    )
+    for key in (("global", (), ()), ("groupby", ("brand",), ())):
+        state = ts._pooled_states[key]
+        assert list(state.bucket_df.columns) == ["unique_id", "ds"], key
+        assert len(state.bucket_df) == len(df), key  # rows preserved
+
+
 @pytest.mark.parametrize("engine", ["pandas", "polars"])
 @pytest.mark.parametrize("lag", _LAGS)
 def test_global_sequential_updates(engine, lag):