avantgardnerio opened a new pull request, #2211: URL: https://github.com/apache/datafusion-ballista/pull/2211
## Summary Introduces `PrefixMergeExec` as the downstream half of the AQE range-shuffle prefix-scan pipeline. In a range-shuffled plan, each executor task runs a windowed aggregate over an ordered slice of the input; the per-row aggregate values are correct *within* the slice but not across slices (a running SUM in task `k` starts at zero, not at the sum of everything in partitions `[0..k)`). This operator closes that gap by merging the accumulated state from all earlier partitions into partition `k`'s output. **This commit is scaffolding only.** `execute` forwards batches unchanged; the merge itself lands in a follow-up. The type signature and constructor are committed now so surrounding plumbing (scheduler collection of per-task state via task-status transport, plan-rule placement, per-task state injection) can be built against a stable shape. ## What's here - `PrefixMergeExec` (new file `ballista/core/src/execution_plans/prefix_merge.rs`) — wraps a child input plan; carries one `FinalizedPartitionState` per input partition (the accumulated state from every earlier partition, ready to be merged into this partition's output). - Constructor validates that `per_partition_state.len()` matches the input's output partition count. - `execute(partition)` currently forwards `self.input.execute(partition)` verbatim; the eventual merge will fold `self.per_partition_state[partition]` into each row's window-aggregate columns via `Accumulator::merge_batch`-shaped composition — aggregate-agnostic by construction, working uniformly for scalar aggregates (SUM, COUNT, MIN, MAX) and sketch-backed aggregates (KLL, TDigest, HLL). - Local `FinalizedPartitionState` type alias mirrors the type of the same name added in [apache/datafusion#24007](https://github.com/apache/datafusion/pull/24007). Doc comment marks the correspondence explicitly so the swap to the DataFusion type is unambiguous once that PR lands. - Exported from `ballista_core::execution_plans` alongside `PerPartitionFilterExec` and other AQE-shuffle-adjacent operators. ## Relation to DataFusion The per-input-partition state this operator expects is produced by `BoundedWindowAggExec::finalized_partition_state`, added in [apache/datafusion#24007](https://github.com/apache/datafusion/pull/24007). Until that lands, callers can pass empty maps — the scaffold forwards batches regardless — to exercise the plumbing. ## Test plan - [x] `try_new_rejects_state_length_mismatch` — mismatched state length surfaces as an error, not a runtime panic. - [x] `scaffold_forwards_input_rows_unchanged` — end-to-end passthrough test on a two-partition memory source; each partition's rows arrive unchanged at the caller. - [x] `cargo clippy -p ballista-core --all-targets` — clean. - [x] `cargo test -p ballista-core --lib execution_plans::prefix_merge` — 2 passed. 🤖 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]
