feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates by avantgardnerio · Pull Request #2211 · apache/datafusion-ballista

avantgardnerio · 2026-07-30T16:45:30Z

Summary

BoundedWindowAggExec declares SinglePartition when there is no PARTITION BY, because a window frame spans rows across the whole input. For an ever-expanding frame that means the entire window computation runs on one core, however wide the cluster.

This PR removes the collapse. Each task computes a partition-local running aggregate over a range-disjoint slice, and a downstream operator corrects it with the merged accumulator state of every prior partition.

Shape

PrefixMergeExec                        state applied row-wise
  ExchangeExec (partitioning: None)    boundary 2, planted by the rule
    PartitionedBoundedWindowAggExec    window runs per partition
      RangeFilterExec                  trim to this task's cut range
        ExchangeExec                   boundary 1, inserted by DistributedExchangeRule
          RuntimeStatsExec
            OrderedRangeRepartitionExec
              SortExec
                RuntimeStatsExec
                  <source>

Boundary 2 carries no repartition. It exists so the scheduler has a point where every upstream task has published its state.

State flow

DataFusion 55 exposes accumulator state through a callback (apache/datafusion#24035). A WindowStateCollector catches it, ShuffleWriterExec translates each capture's task-local partition index to a stage-global one, and it rides SuccessfulTask to the scheduler. The scheduler prefix-merges the reports into one carry-in per partition and binds them to the downstream PrefixMergeExec.

Merging goes through the aggregate's own merge_batch, so non-decomposable aggregates work: two approx_distinct HLL sketches combine correctly where two distinct counts could not.

One state's wire size is fixed by the sketch rather than by the rows it saw. An approx_distinct state encodes to 16393 bytes at one row and at 100k alike, because the HLL is a dense register array, so a task's payload is (partitions in its slice) x (aggregate window expressions) x 16 KiB with no term for input size. sketch_state_wire_size_does_not_grow_with_rows pins that.

Benchmarks

h2o window.sql Q7, the rolling sum:

SELECT id1, id2, id3, v2,
       sum(v2) OVER (ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING
                                           AND CURRENT ROW) AS my_rolling_sum
FROM large;

1e9 rows read from S3. Four executors at 48 vcores each, one per r6i.24xlarge,
pinned one-per-node by anti-affinity; the client runs on a fifth node so pulling
the result back does not contend with an executor. Partitions per task is set to
K/4 throughout, so each executor takes exactly one task and the shuffle reader's
merge width stays fixed at 4 while K varies. Every figure is a single run.

Speedup grows with width:

partitions	rule off	rule on	server-side	wall
16	290.2s	84.5s	3.43x	2.30x
32	286.1s	54.1s	5.29x	2.81x
64	305.1s	42.4s	7.19x	3.23x
128	318.5s	33.5s	9.50x	3.64x

Where the time goes, elapsed_compute per stage:

stage	K=16	K=32	K=64	K=128
0, sort plus ORRE	48.35s	31.11s	27.89s	24.43s
1, the window	22.51s	12.41s	9.34s	6.09s
2, `PrefixMergeExec` correction	13.62s	10.50s	5.19s	2.92s
the same window, rule off, 1 task	253.46s	269.29s	294.82s	313.57s

Taking stages 1 and 2 together against the serial window, the window work itself
goes 7.0x, 11.8x, 20.3x, then 34.8x.

Three things the tables show:

Both ends widen. The rewrite improves with width while the serial arm gets
monotonically worse (253s to 314s), because its SortPreservingMergeExec
collapses an ever-wider fan into one task. Adding cores actively hurts the
status quo, which is the shape of a serial bottleneck rather than a slow
operator.
The correction is cheaper than the window it corrects at every width, and
it kept scaling to the end, 13.62s down to 2.92s. The per-row evaluate()
replay is the cheapest of the three stages by K=128.
Stage 0 is the wall, not the correction. It flattens at roughly 24s and by
K=128 is 73% of the parallel plan, so the asymptote here is about 13x. Further
width buys progressively less until that read-sort-repartition preamble gets
cheaper.

What this measurement understates

SUM is the cheapest possible window aggregate. The design exists for aggregates whose per-row output is not a valid partial state: approx_distinct, quantile sketches, stddev. There each row costs a register scan or a quantile computation in evaluate(), so the window dominates and the sort and repartition tax rounds to noise. Q7 is the weakest showcase in the suite.
WindowApply::Scalar is implemented but not yet selected. For SUM, COUNT, MIN and MAX it replaces the correction's per-row accumulator replay with one arrow kernel per batch, which should remove most of stage 2 for exactly the query measured above.
The result set inflates both arms equally. Q7 emits one row per input row, so the client collects 1e9 rows in both configurations. That is a constant added to each side, which compresses the ratio: 9.50x server-side becomes 3.64x wall.

What flatters it

The serial baseline may be bandwidth-bound rather than CPU-bound. Its collapse stage pulls all 1e9 rows from S3 through a single task, so part of the 253-314s is fetch rather than window computation. The rewrite spreads that read across four executors, and some of the ratio is that rather than the window itself.

What works against it

The extra materialization makes the win contingent on shuffle throughput. Boundary 2 writes and reads the whole window output. Where that write path is slow relative to compute, the extra round trip can eat the gain.
The ORRE preamble is not free, and it is now the floor. Stage 0 pays for the sketch and the range repartition, and it stops improving with width: 48.35s at K=16 down to only 24.43s at K=128, by which point it is 73% of the total.
Three stages rather than two, so one more scheduler round trip and the state transport described above.
The correction is still O(rows) with an evaluate() per row. It measures well here and scales with cores, but it is linear work that the Scalar path avoids entirely where it applies.
Narrow trigger. See "Known limits".
Cost-model: for a user like Coralogix, using 16x the cores to get 4x the wallclock is a big win. For the Spark persona running a fixed-sized cluster of owned nodes and not latency sensitive, this rule might be a net loss - better for the job to take longer and free the cores for other jobs. It could solve OOMs for them though.

Correctness

ballista/client/tests/prefix_window.rs runs a distributed running sum against an independently computed oracle through the real scheduler, shuffle and executor, with max_partitions_per_task = 2 so tasks carry genuine multi-partition slices. Two routing keys are covered: a non-null Float64 one, and a nullable Int64 one where the NULL run has to reach the partition that holds it.

One fix outside the feature flag

Writing that nullable test surfaced a pre-existing bug in cut_partitions, which assigns each producer's shuffle files to consumers by value range. The reported extremes are value extremes, so a file holding both values and NULLs was delivered only to the consumers its values overlap, and its NULL run went wherever its values went. Every producer routes its own NULLs to its own last slot, so a stage with more than one producer task silently dropped all but one producer's share of the run: a short result with every surviving row individually correct.

Fixed by delivering a file carrying NULLs to the run's partition as well, skipped when its value range already reaches there. This is shared range-repartition machinery, reached by any ORRE/URRE stage rather than only by the rewrite added here, so it is the one change in this PR that is not behind the flag.

Safety

Gated behind ballista.planner.parallel_window.enabled, off by default. The rewrite fires only on: no PARTITION BY, a single ascending ORDER BY column whose type the sort-key sketch codec can encode, and an ever-expanding frame. Everything else falls through untouched.

Known limits

No PARTITION BY, single ORDER BY column, ASC only.
Fixed-width routing keys. The sketch codec packs a key into a u64, so variable-length types such as Utf8 fall through the gate. Nullable keys are supported.
WindowApply::Scalar is implemented but not yet selected; every apply currently routes through the accumulator path.

andygrove · 2026-08-08T15:01:47Z

CI is red on two jobs here, but only one of them is yours.

cargo doc is real:

error: public documentation for `FinalizedPartitionState` links to private item `self`
error: public documentation for `PrefixMergeExec` links to private item `self`

The [module-level docs][self] links resolve to mod prefix_merge, which is private, so rustdoc rejects them under -D warnings. Making it pub mod prefix_merge; in execution_plans/mod.rs is the smallest fix and matches what plan_algebra and sort_shuffle already do. Dropping the two intra doc links works too if you'd rather keep the module private.

test linux crates is not this PR. It's ballista-chaos::ha exhausted_retries_fail_the_job_and_leave_the_cluster_healthy::case_1_aqe_off failing at cluster startup with executor registration ConnectionRefused, so an infrastructure flake. Should clear on a rerun.

Design feedback coming in a separate comment.

andygrove · 2026-08-08T15:01:48Z

Design feedback, separate from the CI note above. The overall shape reads well to me, and the module header is genuinely good documentation. Splitting the global prefix merge onto the scheduler and leaving a row wise apply on the executor is the right decomposition, and the APPROX_DISTINCT case makes a convincing argument for why the Aggregate path has to exist. A few things I'd want settled before this grows more code on top of it.

The per row accumulator replay looks like a performance trap. AggregateApply::apply calls update_batch on a one row slice and then evaluate() once per row. For SUM that's just wasteful, but the motivating cases are sketches, and that's where it gets expensive. evaluate() on an HLL scans every register to produce a cardinality estimate, and on TDigest or KLL it runs a quantile computation. Doing that once per row turns a linear pass into something quite a bit worse, and it re derives work the upstream BWAG already did. The APPROX_DISTINCT test proves correctness on three rows, which is exactly the size that won't surface this. Could you run it over a realistic partition before we commit to the shape? If the numbers are bad there may be a middle path where the upstream emits partial state columns and the correction stays batch at a time.

Serde is deferred, and it's the hard part. #2255 landed its proto message and codec arm in the same PR, and I'd like this one to end up there too, since PrefixMergeExec can't reach an executor without them. No objection to a scaffold that defers it, I just want to flag that the remaining work isn't mechanical. Arc<AggregateUDF>, Vec<Arc<dyn PhysicalExpr>>, and Vec<ScalarValue> sketch state all have to cross the wire.

The related question is the transport the design picks. The description says upstream state reaches the scheduler over task status. That's a hot and frequent message, and HLL, KLL or TDigest state per task per window expression is not small. Since you describe the division of labor as fixed at design time, I'd rather pressure test that choice now than after the scheduler side is built on top of it.

Partition index coupling has no guard. Both per_partition_state[k] and Scalar.offset[k] are keyed by partition index, but the operator declares UnspecifiedDistribution, no required input ordering, and maintains_input_order: true, and with_new_children only re validates counts. So any rule that repartitions the input to the same partition count would silently attach each partition's offsets to the wrong rows. Wrong answers, no error. In practice the scheduler hands over a finished plan so it may never happen, but this is the "correct on one node, silently wrong once split across stages" shape that user-personas.md calls out for Persona 1, and I'd want at least a loud invariant comment on it.

Smaller things:

ScalarOp::Overwrite is documented as fitting first_value and last_value. first_value I follow. For the cumulative frame last_value is just the current row's value and needs no correction at all, so overwriting every row with a single scalar would be wrong. Which frame is that aimed at?
No metrics. PrefixMergeExec doesn't implement metrics() and ApplyStream has no BaselineMetrics. Given the first point above this is the operator you'd most want timings from, and Spark shaped users lean on per operator timings for skew debugging.
Type drift is only caught by accident. If numeric::add promotes, or evaluate() returns a different type than the column it replaces, RecordBatch::try_new fails with an opaque arrow error rather than something naming the offending applies[i].
ScalarOp and WindowApply are public enums that you say will grow. Marking both #[non_exhaustive] now costs nothing and saves a breaking change on the first new variant. Similarly, FinalizedPartitionState as a transparent pub type alias means swapping it for the real DataFusion type once feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec datafusion#24007 lands is a silent public API change. A newtype now would keep that swap internal.
Minor housekeeping, the PR description has the "Generated with Claude Code" footer, and CLAUDE.md in the repo asks us to keep that out of PRs.

avantgardnerio · 2026-08-14T17:37:14Z

+        for i in 0..num_rows {
+            let row_args: Vec<ArrayRef> =
+                arg_arrays.iter().map(|a| a.slice(i, 1)).collect();
+            self.accumulator.update_batch(&row_args)?;
+            new_values.push(self.accumulator.evaluate()?);
+        }


it re derives work the upstream BWAG already did

Yes, and that is exactly what the second pass is doing: the upstream window already advanced an accumulator through every row, and this walks it again from the carry-in.

The reason it is not avoidable today is where the state is exposed. apache/datafusion#24035 publishes accumulator state at PARTITION BY group close, one state per group, so there is no per row state to carry into a batch at a time correction.

there may be a middle path where the upstream emits partial state columns

I think that is worth pursuing, but it needs more than the API. DataFusion's dense HLL is 16 KiB per sketch (approx_distinct.rs#L257), so materializing per row state for a sketch is 16 KiB times row count as a column. A workable version probably needs a cheaper per row representation than raw accumulator state, which makes it a real piece of design rather than a switch to flip.

Worth being precise about the current state: this PR does not select the scalar path yet. WindowApply::Scalar exists on the operator and applies batch at a time via arrow kernels, but the rule currently builds an Aggregate apply for every window expression, including SUM. Selecting Scalar where it applies will land on this PR before merge, and for SUM, COUNT, MIN and MAX it removes this cost rather than halving it.

Could you run it over a realistic partition

Yes. A three row test does not surface any of this. The trade we are making is 2x local work divided across cores (theoretically 16x on 32, to be measured shortly). Benchmarks will land in this thread before anything else is built on the shape.

Benchmarks, as promised in this thread.

h2o Q7 at 1e9 rows, 2 executors, 16 partitions, release build, local NVMe:

stage rule off rule on

sort (plus ORRE when on) 40.82s 54.56s

window 122.71s, 1 task 12.56s

PrefixMergeExec correction n/a 8.25s

wall 179.5s 91.3s

The correction costs less than the window it corrects: 8.25s against 12.56s.

Both passes also scale with cores. Doubling from 8 to 16 partitions at fixed data: window 21.86s to 12.56s (1.74x), correction 14.84s to 8.25s (1.80x), while the serial window is flat at 118.53s and 122.71s because it is one task either way. End to end that moves 1.34x to 1.97x, so the advantage grows with cluster width rather than being a fixed constant.

One thing I said here needs correcting:

Selecting Scalar where it applies will land on this PR before merge

That no longer follows from the numbers. Scalar selection is worth doing and still removes the replay for SUM, COUNT, MIN and MAX, but it is an optimization on top of a win rather than the thing that makes this a win. I would rather land it separately than hold this PR for it.

Three caveats. These are single runs. SUM is the cheapest window aggregate there is, so this understates the sketch cases the design exists for, where each row costs a register scan in evaluate(). And both executors share one disk, which penalises the arm doing more shuffle IO, namely this one.

avantgardnerio · 2026-08-14T17:42:27Z

+  // TODO: watch the size of this. Task completion is a hot, frequent
+  // message, and sketch-backed aggregates make the payload unbounded in a
+  // way row counts and quantile sketches are not — an HLL or KLL state is
+  // kilobytes per window expression per partition, and a task covering a
+  // wide partition slice carries one of each. If it stops being small,
+  // write the state as a sidecar next to the shuffle files instead, the way
+  // sort-shuffle already writes `<data>.arrow.index` beside its data
+  // (`sort_shuffle::get_index_path`), and send only a reference here. That
+  // keeps the completion message fixed-size regardless of aggregate.


HLL, KLL or TDigest state per task per window expression is not small

Agreed, and it is worth sizing. DataFusion's dense HLL is 16 KiB per sketch (approx_distinct.rs#L257). A task carries one state per (partition in its slice, aggregate window expression), so max_partitions_per_task multiplies it: a 32 partition slice with a single sketch aggregate puts 512 KiB on that task's completion message.

Two things make me think the current shape is defensible as a starting point rather than a commitment. RuntimeStatsExec already ships quantile sketches over this exact path, so this adds a second payload of a class task status already carries rather than introducing one. And the escape hatch is cheap: the state can be written as a sidecar beside the shuffle files, the way sort shuffle already writes <data>.arrow.index next to its data, with only a reference on the message. That keeps completion fixed size regardless of aggregate, and it is a change to this one field rather than to the design around it.

What I have not done is measure it. The e2e only exercises a Float64 SUM, where the payload is a handful of bytes. If you would rather see a real number for approx_distinct before this merges, I can add a test that reports the encoded size and post it here.

I measured it and the above statement is correct.

avantgardnerio · 2026-08-14T18:05:17Z

Heads up that this is functional now rather than a scaffold. It runs end to end and ballista/client/tests/prefix_window.rs asserts a distributed running sum against the serial answer through a real cluster. I have replied to each of your points inline on the relevant lines so they can be resolved individually.

Leaving it in draft: benchmarks for the per-row replay are not in yet, and the scalar path selection lands before merge. Both are tracked in the threads.

…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

`approx_distinct`'s HLL is a dense register array, so one state encodes to 16393 bytes whatever the row count. Asserts the size is independent of cardinality, with a loose ceiling as a tripwire for a representation change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

avantgardnerio · 2026-08-20T20:21:04Z

@phillipleblanc FYI

`free_port` bound :0, read the port, then dropped the listener, so two consecutive calls could return the same number. `spawn_executor` asked twice and handed its child one port for both the flight and gRPC servers: the second bind failed AddrInUse, the executor exited during startup, and the harness timed out waiting for two registrations. Seen in CI as `port: 38571, grpc_port: 38571` on one executor's registration, with distinct ports on the other. Not reproducible locally (0 collisions in 1000 attempts at the old pattern), so `free_ports` fixes it by construction instead: hold every listener until all N ports have been read, which the OS cannot satisfy with a duplicate. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

milenkovicm · 2026-08-20T21:06:21Z

Thanks @avantgardnerio will have a look tomorrow and over the weekend

avantgardnerio requested a review from Dandandan July 30, 2026 18:05

avantgardnerio mentioned this pull request Aug 3, 2026

Expose accumulator state to allow prefix scanning apache/datafusion#24035

Merged

avantgardnerio force-pushed the brent/prefix-merge-scaffold branch 3 times, most recently from 8696f35 to f85b988 Compare August 14, 2026 16:56

avantgardnerio changed the title ~~feat(core): scaffold PrefixMergeExec for cross-partition windowed-aggregate state merge~~ feat(core): PrefixMergeExec for cross-partition windowed-aggregate state merge Aug 14, 2026