feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec - #24007
Draft
avantgardnerio wants to merge 2 commits into
Draft
feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec#24007avantgardnerio wants to merge 2 commits into
avantgardnerio wants to merge 2 commits into
Conversation
…dowAggExec Add per-output-partition slots that hold each aggregate window function's final Accumulator::state() at partition-close, plus a public getter that distributed executors can call once a task has drained to ship state via a side-channel for cross-partition prefix scans without a two-pass halo-row scheme. The stream snapshots state at the top of prune_state — the last moment `state.is_end` entries are live before either prune_out_columns or prune_partition_batches drops them — so both mid-stream partition close (as SQL PARTITION BY groups end) and EOS are covered by one write site. Non-aggregate window functions (row_number, rank, lead/lag, ...) occupy `None` at their slot in the per-partition-key Vec. No trait changes: this uses the existing Accumulator::state() surface that the group-by Partial/Final protocol already relies on. Behavior without a reader is unchanged — the slots simply hold state that no one reads. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
6 tasks
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24007 +/- ##
========================================
Coverage 80.85% 80.85%
========================================
Files 1096 1099 +3
Lines 373911 374473 +562
Branches 373911 374473 +562
========================================
+ Hits 302329 302798 +469
- Misses 53547 53611 +64
- Partials 18035 18064 +29 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
… group The caller for this getter — cross-DF-partition prefix scan — only applies when at most one PARTITION BY group is active per DF output partition (either no PARTITION BY, or a synthetic single-key one). Queries with a real multi-group PARTITION BY are handled by DataFusion's normal key-partitioned distribution and don't need this API. The previous shape leaked BWAG's internal PARTITION BY tuple keying into the public type, forcing every caller to project it away. The getter now returns Vec<Option<Vec<ScalarValue>>> directly — outer Vec indexed by window expression, inner Vec the aggregate's state — and errors if BWAG ever observed more than one group. Internal storage mirrors that: a 3-state FinalStateSlot enum (Empty / Single / Multi) rather than a HashMap that grows with N groups. Split the existing test into single-group-success and multi-group-error to cover both paths under the new shape. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
avantgardnerio
added a commit
to avantgardnerio/arrow-ballista
that referenced
this pull request
Jul 30, 2026
…plification Follows apache/datafusion#24007's shape simplification: the DF getter now returns Vec<Option<Vec<ScalarValue>>> directly rather than HashMap<PartitionKey, ...>, since the cross-DF-partition prefix-scan use case is by construction scoped to at most one PARTITION BY group per DF partition. The DF side handles the scoping (Empty/Single/Multi slot internally, error on multi). Ballista's local FinalizedPartitionState alias matches — dropping the outer HashMap layer means execute() no longer has to project it away or error on multi-key (both handled upstream), and the demo tests construct state directly as Vec<Option<Vec<ScalarValue>>> rather than via a HashMap keyed by an empty PartitionKey. No semantic change; the aggregate re-run tests still cover APPROX_DISTINCT and AVG, the scalar test still covers SUM, all 6 pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a getter on
BoundedWindowAggExecthat returns the finalizedAccumulator::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:
FinalizedPartitionState = HashMap<PartitionKey, Vec<Option<Vec<ScalarValue>>>>— oneAccumulator::state()per (PARTITION BY tuple, window expression) pair. Non-aggregate window functions (row_number,rank,lead/lag, ...) occupyNoneat their slot in the innerVec.finalized_state: Arc<[Mutex<FinalizedPartitionState>]>onBoundedWindowAggExec— one slot per output partition.BoundedWindowAggExec::finalized_partition_state(partition)— returns the current snapshot for a given output partition.prune_stateinBoundedWindowAggStream— snapshots every entry whosestate.is_endis set, immediately before eitherprune_out_columnsorprune_partition_batchesdrops 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.🤖 Generated with Claude Code