avantgardnerio opened a new pull request, #24007: URL: https://github.com/apache/datafusion/pull/24007
## Summary Adds a getter on `BoundedWindowAggExec` that returns the finalized `Accumulator::state()` for each aggregate window expression at partition close, keyed by PARTITION BY tuple. This lets callers observe per-partition final aggregate state without emitting it as an output column on every row. The change is purely additive: - **New public type alias** `FinalizedPartitionState = HashMap<PartitionKey, Vec<Option<Vec<ScalarValue>>>>` — one `Accumulator::state()` per (PARTITION BY tuple, window expression) pair. Non-aggregate window functions (`row_number`, `rank`, `lead`/`lag`, ...) occupy `None` at their slot in the inner `Vec`. - **New field** `finalized_state: Arc<[Mutex<FinalizedPartitionState>]>` on `BoundedWindowAggExec` — one slot per output partition. - **New public method** `BoundedWindowAggExec::finalized_partition_state(partition)` — returns the current snapshot for a given output partition. - **Write site** at the top of `prune_state` in `BoundedWindowAggStream` — snapshots every entry whose `state.is_end` is set, immediately before either `prune_out_columns` or `prune_partition_batches` drops it. This covers both mid-stream partition close (as SQL PARTITION BY groups end) and end-of-stream with a single write path. No trait changes: the existing public `Accumulator::state()` method (already used by the group-by Partial/Final protocol) is the state source. Behavior without a reader is unchanged — the slots simply hold state that no one reads. ## Motivation Distributed executors that split a windowed aggregate across multiple tasks need each task's final per-partition state to compose results across tasks (for aggregates like `AVG` whose state is not recoverable from the output column, or for sketch-backed aggregates whose state is a merge-able structure). Exposing state at partition-close time keeps the observation off the output stream — no state emitted as row-level columns. ## Scope `BoundedWindowAggExec` only. `WindowAggExec` doesn't hold `Accumulator` boxes across its partition loop — its accumulators are constructed and discarded inside `WindowExpr::evaluate` — so the same shape there would need either a `WindowExpr` trait extension or a compute-side restructure. Deferring that until a caller actually needs it; distributed callers can canonicalize to `BoundedWindowAggExec` (with `InputOrderMode::Sorted`) when the input is already ordered, which is the shape-fitting choice regardless. ## Test plan - [x] New unit test `finalized_partition_state_captures_sum_per_partition_key`: runs cumulative `SUM(v) OVER (PARTITION BY pk)` on two partition keys, drains the stream, verifies each key's finalized state contains the expected sum. Also exercises the out-of-range partition error path. - [x] `cargo test -p datafusion-physical-plan --lib` — 1610 tests pass, no regressions. - [x] `cargo clippy -p datafusion-physical-plan --all-targets` — clean. - [x] `cargo check --workspace` — clean. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
