Skip to content

feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec - #24007

Draft
avantgardnerio wants to merge 2 commits into
apache:mainfrom
avantgardnerio:brent/bwag-finalized-state
Draft

feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec#24007
avantgardnerio wants to merge 2 commits into
apache:mainfrom
avantgardnerio:brent/bwag-finalized-state

Conversation

@avantgardnerio

@avantgardnerio avantgardnerio commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

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.

🤖 Generated with Claude Code

…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>
@codecov-commenter

codecov-commenter commented Jul 30, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 80.74074% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.85%. Comparing base (62cfc0c) to head (61c2aa8).
⚠️ Report is 6 commits behind head on main.

Files with missing lines Patch % Lines
...ysical-plan/src/windows/bounded_window_agg_exec.rs 80.74% 11 Missing and 15 partials ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

… 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-expr Changes to the physical-expr crates physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants