Allow `predict(ids=...)` with safe pooled/grouped lag-transform subsets

[simonez-tuidi]

Description

Problem

predict(ids=...) is currently rejected when nonlocal pooled lag transforms are present. This is correct for pure global transforms, because they pool across all series. However, it is overly restrictive for grouped pooled transforms.

For example, with RollingMean(..., groupby=["store_id"]), the transform only pools series within each store_id. Predicting a subset should be safe when the selected ids include every series in each selected store.

Desired Behavior

Allow subset prediction when the selected series preserve the pooled dependencies required by every lag transform.

The invariant should be:

A selected subset is valid only if every non-global pooled state can be evaluated without dropping part of any pooled group needed by the selected forecast population.

This applies across multiple grouped configurations simultaneously. If transforms use both groupby=["store_id"] and groupby=["region"], the selected subset must form complete groups for both groupings.

Partition By Compatibility

The feature/pooled_partition_by branch generalizes pooled state via _pooled_states keyed as (mode, group_cols, partition_cols). The validation should target that abstraction rather than special-casing only _pooled_groups.

Suggested rules:

Pooled state	Example key	Subset behavior
Pure global	("global", (), ())	Require all ids
Pure groupby	("groupby", ("brand",), ())	Allow complete selected groups
Local partition	("local", (), ("promo",))	Allow arbitrary ids
Global + partition	("nonlocal", (), ("promo",))	Require all ids
Groupby + partition	("nonlocal", ("brand",), ("promo",))	Allow complete selected groups

partition_by values do not need their own filter parameter. They are resolved from X_df or static features at prediction time. The subset safety condition is about preserving group completeness for the parent pooled population, not selecting partition values.

API Options

Option A: Infer From ids

Keep the current public API:

predict(..., ids=[...])

When ids is provided:

• reject if any pure global or global + partition pooled state is present;
• allow local partition states;
• for groupby and groupby + partition states, verify that selected ids form complete groups for every distinct group_cols configuration;
• raise a clear error if any grouping is partially selected.

This is the smallest API change and preserves current usage.

Option B: Add groupby_values

Add an optional group-based selector, for example:

predict(..., groupby_values={"store_id": ["A", "B"]})

This would select all ids belonging to the requested group values, then run the same pooled-state closure validation as Option A.

Important details:

• groupby_values should select ids only through groupby columns, not partition_by columns.
• It should still validate closure across all pooled states.
• If transforms use multiple groupby configurations, selecting values for one grouping may still be invalid if it partially selects another grouping.
• If both ids and groupby_values are supplied, either reject the combination or require their implied id sets to match exactly.

This gives users a clearer way to request complete groups without manually expanding group values into ids.

Recommended Direction

Implement Option A first, because it requires no new public API and fixes the current overly restrictive behavior.

Option B can be added as a convenience layer later. It should compile groupby_values into an id set and reuse the same pooled-state closure validator.

Implementation Notes

The existing _maybe_subset path already temporarily subsets ga, uids, static features, dates, target transforms, and lag transforms. It should also temporarily subset pooled state.

For grouped and grouped + partition states:

• use series_bucket_id to map ids to pooled buckets;
• check that selected ids include all series for every selected bucket;
• subset pooled arrays and metadata to selected buckets;
• compact bucket ids so temporary state remains internally consistent;
• restore the original pooled state after prediction.

For local partition_by, arbitrary id subsets are semantically safe, but the pooled state still needs temporary subsetting or remapping so state.series_bucket_id matches the narrowed self.uids.

Tests

Add pandas and polars coverage for:

• ids with pure global transforms still raises.
• ids with one groupby transform succeeds for complete groups.
• ids with one groupby transform raises for partial groups.
• multiple groupby configurations require closure across all groupings.
• local partition_by allows arbitrary id subsets.
• global + partition_by still requires all ids.
• groupby + partition_by allows complete selected groups.
• groupby_values selects complete groups if that API is added.
• groupby_values still fails when it breaks another groupby configuration.
• subset predictions match full predictions filtered to the same ids.
• prediction state is restored after successful or failed subset prediction.

Use case

I don't have a use case myself for now but it might be good to allow this as it is overly restrictive at this stage