[PERF] Amortized predict-time compute for pooled lag transforms

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
@@ -689,6 +690,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 +702,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 +729,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 +750,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:
@@ -1340,10 +1348,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:
@@ -1590,13 +1605,20 @@ 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:
             self.ga = ga
             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