Expose accumulator state to allow prefix scanning - #24035
Conversation
ae37c0b to
65014d5
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24035 +/- ##
==========================================
+ Coverage 80.99% 81.01% +0.02%
==========================================
Files 1106 1106
Lines 383352 384686 +1334
Branches 383352 384686 +1334
==========================================
+ Hits 310488 311665 +1177
- Misses 54544 54620 +76
- Partials 18320 18401 +81 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Thanks for this contribution @avantgardnerio ! I'd find it helpful if you'd elaborate a little bit more about the motivation for this change in the PR description. For example, some intended use-cases, what kind of performance improvement this unlocks, etc. |
|
@neilconway I'm trying to speed up window functions using parallel prefix scans. I am presently incubating this in Ballista, and this is the minimum API exposure that I need to do it for non-decomposable operations like The jury is still out about re-partition cost vs performance benefit, but the signs are hopeful:
And at least from a big-O time perspective (table 1) it should be optimal for some queries ( |
0767f73 to
13023b2
Compare
|
@gene-bordegaray and @JSOD11 you guys might be interested as well. |
alamb
left a comment
There was a problem hiding this comment.
Thanks @avantgardnerio and @neilconway -- I left some comments
alamb
left a comment
There was a problem hiding this comment.
I may not fully understand prefix scanning, but it seems to me like this API will only give you access to the window state for the single last row in each partition.
Don't you potentially need access to the window state for the last N rows in a partition (e.g the HALO rows) 🤔
|
FWIW claude claims this doesn't get run with windows like UNBOUNDED PRECEDING → CURRENT ROW |
|
I think it would also be super useful to add some sort of example / test that shows how you intend to use this API (for exmple some simple example for computing a window function in parallel or something 🤔 that way we could see the API in action |
Thanks @alamb ! That was a critical bug that would have defeated the whole point. It is now fixed and asserted in |
Yes, this is exactly what is required.
No, not for prefix scanning. (answer below)
The HPC "halo" term is a good fit for bounded preceding/following (surrounding cells, in 1D) but doesn't extend cleanly to "last row of every other partition." Regardless of the name, this PR doesn't take that approach - because although it works for SUM, and decomposes for AVG (sum+count), it fails by the time you get to arbitrary accumulators like approx_distinct.
Which is exactly where (the newly added) Edit: added the qualifier (DF) partition to distinguish between the ambiguously named (SQL window) partition. |
|
Pseudo code, stripped directly from the |
Match the field's type so with_new_children collapses to a single chained call and the setter can also clear a previously-installed observer. Addresses apache#24035 review comment 3738803158. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The method now always mutates when called and takes the observer as a required argument; the "is observer installed?" check moves to the caller in `compute_aggregates`. Removes the "&mut self that only mutates when observer is set" shape. Addresses apache#24035 review comment 3738822701. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Rename the trait method to `finalize_window_aggregate` and split its signature so the callback fires once per aggregate window expression per closing PARTITION BY group, receiving that expression's Arc and its own `Accumulator::state` directly. Non-aggregate window functions no longer fire the callback at all. Removes the per-partition-key `Vec<Option<Vec<ScalarValue>>>` wrapper allocation, and gives the observer the window-expression context needed to disambiguate calls when the exec carries multiple window expressions. Leaves room to add a peer `finalize_window_function` later for built-in (non-aggregate) window functions. Addresses apache#24035 review comments 3738790091 and 3738816488. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…close tests The two `test_finalized_state_observer_*` tests were structurally identical apart from the window frame. Fold their common setup and assertions into a single async helper that takes the frame, so each test body is now just the frame construction + a one-line comment explaining which causality regime it exercises. Addresses apache#24035 review comment 3738845164. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Thank you @avantgardnerio and @timsaucer |
|
run benchmark bounded_window |
|
run benchmark window_query_sql |
|
run benchmark h2o_small_window |
|
(just running some benchmarks to make sure this doesn't change performance noticably) |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark window_query_sqlResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowCPU Details (lscpu)Details
Resource Usagebounded_window — base (merge-base)
bounded_window — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowCPU Details (lscpu)Details
Resource Usageh2o_small_window — base (merge-base)
h2o_small_window — branch
File an issue against this benchmark runner |
|
run benchmark h2o_small_window |
|
run benchmark bounded_window |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowCPU Details (lscpu)Details
Resource Usagebounded_window — base (merge-base)
bounded_window — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowCPU Details (lscpu)Details
Resource Usageh2o_small_window — base (merge-base)
h2o_small_window — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark window_query_sqlCPU Details (lscpu)Details
Resource Usagewindow_query_sql — base (merge-base)
window_query_sql — branch
File an issue against this benchmark runner |
## Summary
Expose state of aggregate streams within BWAG so downstream prefix
scanning can take place.
## API
```rust
// physical-plan/src/windows/bounded_window_agg_exec.rs
pub type FinalizedWindowStateObserver = Arc<
dyn Fn(usize, &PartitionKey, &[Option<Vec<ScalarValue>>]) -> Result<()>
+ Send + Sync,
>;
impl BoundedWindowAggExec {
pub fn with_finalized_state_observer(mut self, obs: FinalizedWindowStateObserver) -> Self { … }
}
// physical-expr/src/window/window_expr.rs
impl WindowState {
/// `Accumulator::state()` if this is an aggregate window function, `None` otherwise.
pub fn aggregate_state(&mut self) -> Result<Option<Vec<ScalarValue>>> { … }
}
```
|
BTW I tried to reproduce the slowdown that seems tohappen above cargo bench --profile=profiling --bench window_query_sql -- "window empty over, aggregate functions"However, I didn't see any evidence of this new code being involved |
|
@alamb oh, thank you! I wondered what you were doing. I saw you running a bunch of these, and as of the last ones I looked at I saw +/-2%. Do you need me to profile or something? |
nope! I am just trying to make sure this didn't add any measurable overhead (I don't think it did) |
## Which issue does this PR close? - follow on to apache#24035 ## Rationale for this change This is a minor thing I found while doing some performance profiling for apache#24035 ## What changes are included in this PR? Avoid a few clone calls ## Are these changes tested? By CI ## Are there any user-facing changes? realistically nothing that someone will measure
Vec<Option<Vec<ScalarValue>>> appeared in this operator's public signatures, where it is neither readable nor searchable, and it spelled "no state for this window expression" two ways: a missing index, and a None at a present index. Every caller handled both. slot(window_expr_index) collapses them into one answer, and a later change to the representation now stays internal. Also removes two comments claiming DataFusion guarantees at most one PARTITION BY group per partition. It does not — apache/datafusion#24035 shipped a callback keyed by group, so that invariant is ours, and the scheduler enforces it by rejecting any report carrying a key. A window that does have a PARTITION BY needs nothing from this operator anyway: BoundedWindowAggExec asks for KeyPartitioned input, so each partition's window is already independent. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…e state merge Introduces PrefixMergeExec as the downstream half of the AQE range-shuffle prefix-scan pipeline: it takes per-input-partition window-aggregate state that the scheduler has already prefix-merged and applies it row-wise to the current partition's output, so cross-partition running aggregates come out correct. Both apply paths are implemented: - WindowApply::Aggregate builds a fresh Accumulator per partition, seeds it via merge_batch from the offset state, and replays each row through update_batch + evaluate to overwrite the output column. - WindowApply::Scalar applies the ScalarOp batch-at-a-time via arrow kernels: numeric::add for Add, cmp::lt_eq/gt_eq + zip for Min/Max, and a constant fill for Overwrite. Purely additive: nothing in-tree constructs a PrefixMergeExec. The remaining work is the state source — collecting each upstream task's finalized accumulator state out of BoundedWindowAggExec and transporting it to the scheduler — which lands separately. FinalizedPartitionState is defined locally, indexed by window-expression position, as the shape this operator consumes. wip(core,scheduler): prefix-window rewrite plants the shape, collector captures state Follows the data flow end to end for the AQE prefix-scan pipeline. Stages 0 and 1 run on a real cluster; stage 2 is blocked on PrefixMergeExec serde. PrefixWindowRule: sibling of ParallelWindowRule for UNBOUNDED PRECEDING frames, gating on start_bound.is_unbounded() where that rule gates on is_finite() — complementary, so no plan matches both. Plants the ORRE preamble, a zero-halo RangeFilterExec trim, PBWAG, a passthrough ExchangeExec for the state round trip, and PrefixMergeExec. Accepts ROWS as well as RANGE units, which for an unbounded start differ only in tie handling. Module docs record the rule's actual captured input rather than an assumed one, including that AQE re-plans and calls optimize three times — hence the idempotency guard. WindowStateCollector: implements DataFusion 55's WindowStateObserver and retains each finalized accumulator state. Retention rather than polling because Accumulator::state is a destructive read fired at most once per group. PBWAG installs one exactly when every frame is ever-expanding, the same condition with_state_observer enforces, so the halo shape is untouched and the wire format needs no new field. PartitionSliceable: operators carrying data indexed by global input partition now implement their own slicing next to the fields being sliced, replacing two bespoke arms in the scheduler's task builder. RangeFilterExec's bounds and PrefixMergeExec's state/offsets both slice when a task is restricted to a partition subset — without which PrefixMergeExec attaches each partition's offsets to the wrong rows. Tests: a client-side e2e asserting the running sum against a computed oracle through the real distributed path. Red until the transport lands, since PrefixMergeExec is a passthrough with no state. The rule's unit tests pin why h2o Q7 does not rewrite today — it orders by an Int64 column and ORRE routes on a Float64-only T-Digest until KLL. Known gaps: no serde for PrefixMergeExec; no transport from collector to scheduler; observed partition_idx is task-local and needs pairing with the task's global partition ids. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> wip(core,executor): ShuffleWriter translates window state to global partition ids Continues following the data: BWAG accumulator -> collector -> PBWAG getter -> ShuffleWriter -> executor. Logged at task completion rather than transported, so the path is exercised end to end before anything is built on it. The local-to-global translation lives on the writer, not on the operator that captured the state. A task's plan is restricted to a partition slice, so an operator mid-plan only ever sees local indices; the writer is the node the scheduler hands global_output_partition_ids to. Reassembling downstream instead would have the scheduler re-derive a mapping it already computed, and a prefix scan fed a permuted order is wrong with nothing to show for it. Verified on the client e2e: two tasks each covering two partitions previously both reported local 0 and 1; they now report globals 0/1 and 2/3, with states 10/26/42/58 over input 1..16. collect_window_state joins collect_plan_metrics and collect_runtime_stats_reports as a task-completion peer on QueryStageExecutor, wired into both the pull (execution_loop) and push (executor_server) task paths. PBWAG grows observed_window_state() and keeps the TODO that the install site moves when the wrapper collapses. Neither the collector nor the writer-side walk depends on the wrapper beyond one downcast. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> wip(core,scheduler,executor): transport window state to the scheduler Adds WindowStateReport to SuccessfulTask, completing the path from a task's BWAG accumulator to the scheduler: collector -> PBWAG -> ShuffleWriter -> wire -> RunningStage::window_state_reports. Verified on the client e2e, where four range-disjoint partitions over input 1..16 arrive as globals 0..3 with states 10/26/42/58. State and partition key cross as datafusion_common.ScalarValue rather than a numeric field, so sketch-backed aggregates (approx_distinct's HLL blob) work unchanged. The proto carries a TODO on payload size: if it stops being small, write the state as a sidecar beside the shuffle files the way sort-shuffle already writes <data>.arrow.index, and send only a reference. Failures fail the task rather than dropping a report. Unlike runtime stats, which are an optimization input, this state is load-bearing: the downstream prefix merge is arithmetically wrong without every partition's contribution, and wrong in a way nothing later detects. Collection is skipped entirely when execution already failed. Reports are tagged with their producer task and purged on reset, in both reset_task_info and reset_tasks. A retried task re-runs its slice and reports the same global partitions again; without the purge the stage would hold two states for one partition and the prefix merge would double-count them. The file-addressing reason RuntimeStats needs its tag does not transfer — the writer already stamped stage-global ids — but the purge reason does. Both scheduler task-status paths (classic execution_graph and AQE) and both executor task paths (pull execution_loop and push executor_server) are wired; each pair are peer implementations that need every completion hook. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> wip(core): prefix-scan the accumulated window state on the scheduler prefix_merge_window_state turns per-partition finalized states into one carry-in per partition: out[0] empty, out[k] the merge of every partition before k. That is what a downstream PrefixMergeExec adds to each partition's local running aggregate to make it global. Merging goes through the aggregate's own Accumulator::merge_batch rather than arithmetic here, which is what lets non-decomposable aggregates work — two approx_distinct HLL sketches combine correctly where two distinct counts could not. The accumulator comes from PlainAggregateWindowExpr, the type an ever-expanding frame always produces; a sliding expression reaching this is an error rather than a silently wrong answer. Built incrementally, out[k] = merge(out[k-1], state[k-1]), so two merges per partition rather than merging every prior from scratch. A fresh accumulator per partition is still required because Accumulator::state is a destructive read and must not be called twice; seeding it from the previous carry-in is the same round trip two-phase aggregation makes. Enforces here what stopped being DataFusion's guarantee when the API turned out to be push-shaped: a report carrying a PARTITION BY key is rejected, because FinalizedPartitionState has no key dimension and a second group in one partition would have nowhere to go. Tests cover the carry-in arithmetic, independence from report arrival order (reports arrive in close order and tasks complete in any order), rejection of a duplicate partition state (only reachable if the producer-task purge failed, and would double-count), and carrying across an empty partition that closes no group and so publishes nothing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> wip(core,scheduler): bake the prefix-merged state into PrefixMergeExec Closes the loop from reports to operator, all in scheduler memory. When a stage completes, its accumulated window-state reports are prefix-merged and bound to the PrefixMergeExec waiting on it downstream. Verified on the client e2e, where four partitions reporting 10/26/42/58 resolve to carry-ins [None, 10, 36, 78]. State is late-bound, mirroring RangeFilterExec's cuts: try_new_pending for the rule's plant-time path, try_new_resolved for wire decode and task restriction, resolve_state as the setter. Both execute() and slice_to_partitions refuse while unresolved rather than treating an absent carry-in as zero — that would emit partition-local aggregates, which look plausible and are wrong. The scheduler hooks update_stage_progress on completion: walk the plan for the PrefixMergeExec whose state-sync boundary carries this stage id, recover the window expressions from the operator below that boundary (the exchange retains its input subtree even once resolved), prefix-merge, resolve. A stage that reported state with no consumer to bind it to is an error rather than a skip; the state exists because something downstream cannot be correct without it. Still passthrough: `applies` is empty, so the operator carries the state without applying it. The descriptors that turn state into corrected columns, and the serde that lets the operator reach an executor, are next. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> wip(core): serde for PrefixMergeExec Stage 2 now reaches an executor and runs. The e2e's remaining failure is the expected one: every partition's running sum is off by exactly its carry-in (5+10=15, 9+36=45, 13+78=91), so the state that crossed the wire is provably correct and nothing is applying it yet. Both WindowApply shapes cross. The aggregate arm carries its UDAF by name, resolved from the executor's function registry on decode, with args as PhysicalExprNodes. State crosses as ScalarValue so sketch-backed aggregates work unchanged, and an absent slot stays distinct from a present-but-empty one — a non-aggregate window function publishes no state, which is not the same as publishing nothing. Encoding refuses while state is unresolved, matching RangeFilterExec's refusal on unresolved bounds. An executor has no way to obtain prefix state, so a plan reaching the wire without it could only produce partition-local aggregates. Round-trip test covers both arms, the UDAF-by-name resolution, and the None-vs-empty distinction in the state slots. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> feat(scheduler,core): apply descriptors — parallel prefix scan is correct The client e2e passes: a running sum computed across four range-disjoint partitions matches the serial answer through the real distributed engine, scheduler and shuffle and executor included. One WindowApply per aggregate window expression. BWAG appends its window columns after the input's, so expression i lands at input_field_count + i and the input columns keep their indices — which is why the aggregate's own argument expressions carry over unchanged despite being resolved against the input schema. Non-aggregate window functions get no apply; they publish no state to merge. This is the last of the stubs each earlier step stood on: with `applies` empty the operator was a passthrough, which is why the state was provably correct and the output was still partition-local. SUM goes through the Aggregate path even though the cheaper Scalar path covers it. Seeding an accumulator and replaying rows is the shape non-decomposable aggregates need, and exercising it where the answer is independently checkable beats the arrow-kernel shortcut. Choosing Scalar where it applies is a later optimization, worth measuring. Also flattens the codec arms added in the previous commit. The prefix-state encoder was three nested maps around a transpose; it is now six named helpers built from plain loops, taking the decode arm from ~85 lines to 13 and the encode arm from ~95 to 24. Option handling is an explicit match rather than `.map(..).transpose()?`, which puts the absent-versus-empty distinction where a reader can see it, and each helper names what it was converting so failures read as "failed to encode prefix state" rather than an anonymous try_from. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> docs: correct status claims that went stale when the pipeline went green Five places still described a half-built system. Each was accurate when written and became a lie at a different commit: - the client e2e's "Fails today" doc, now stating the quiet failure it guards against rather than predicting one - the rule's "Status: shape only", which claimed the rewrite corrected nothing; it now records what it is correct for, and that h2o Q7 is blocked on the sketch's Float64 restriction rather than on this rule - prefix_merge's "Nothing in-tree collects it yet", which now points at the collector that does - the collector's and the stage's scaffolding-log comments, both explaining themselves as stand-ins for consumers that now exist Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> fix(core): review follow-ups — Overwrite doc, non_exhaustive, named type drift Three small ones from review. ScalarOp::Overwrite claimed to fit last_value. Over an ever-expanding frame last_value is the current row's own value and needs no correction, so overwriting every row with one scalar would be wrong. The doc now says first_value only, and says why last_value is excluded. ScalarOp and WindowApply are #[non_exhaustive]. Both are expected to grow — the ranking family needs a segment-tree broadcast shape — and each addition would otherwise be a breaking change for anyone matching on them. Construction is unaffected, so the rule still builds an Aggregate apply. Type drift now names the apply responsible. Arrow reports a mismatch at a column index and nothing about which correction produced it, which is the wrong half when several applies rewrite one batch. rebuild_batch reports "applies[3] produced Int64 for column 1, which the schema declares as Float64", using the apply_index AggregateApply already carried for exactly this and did not use. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> feat(core): metrics on PrefixMergeExec, split by apply path BaselineMetrics for elapsed_compute and output_rows, per partition because execute(partition) builds a stream each. Plus a rows_corrected counter and a separate timer per apply path. The split is the point. WindowApply::Scalar is an arrow kernel over a whole batch; WindowApply::Aggregate seeds an accumulator and replays every row through it. A sketch-heavy query pays the second and a SUM-heavy one need not, which a single total would hide. On the client e2e the operator now reports aggregate_apply_time=234.34us against elapsed_compute=238.60us, so the replay is 98% of its time on a trivial SUM — a shape rather than a magnitude at 16 rows, but it makes the replay cost something the operator reports in production rather than something only a benchmark can see. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> refactor(core): FinalizedPartitionState as a newtype Vec<Option<Vec<ScalarValue>>> appeared in this operator's public signatures, where it is neither readable nor searchable, and it spelled "no state for this window expression" two ways: a missing index, and a None at a present index. Every caller handled both. slot(window_expr_index) collapses them into one answer, and a later change to the representation now stays internal. Also removes two comments claiming DataFusion guarantees at most one PARTITION BY group per partition. It does not — apache/datafusion#24035 shipped a callback keyed by group, so that invariant is ours, and the scheduler enforces it by rejecting any report carrying a key. A window that does have a PARTITION BY needs nothing from this operator anyway: BoundedWindowAggExec asks for KeyPartitioned input, so each partition's window is already independent. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> � Conflicts: � ballista/core/src/execution_plans/mod.rs � ballista/core/src/serde/mod.rs � ballista/scheduler/src/state/aqe/mod.rs � ballista/scheduler/src/state/aqe/planner.rs � ballista/scheduler/src/state/task_builder.rs

Summary
Expose state of aggregate streams within BWAG so downstream prefix scanning can take place.
API