feat(core): scaffold PrefixMergeExec for cross-partition windowed-aggregate state merge

[avantgardnerio]

Summary

Introduces PrefixMergeExec — the downstream half of the AQE range-shuffle prefix-scan pipeline. Range-shuffle produces N ordered disjoint partitions; each executor task's local windowed aggregate is correct within its slice but not across slices. This operator corrects each partition by applying the scheduler-provided already-prefix-merged state from every earlier partition, aggregate-agnostic by construction.

Division of labor (fixed at design time):

• Scheduler (global step): collects each upstream task's finalized Accumulator::state() via task-status transport, computes the prefix-merge across all tasks. Requires cross-task visibility.
• Executor (this operator): receives the pre-merged state for its partition and applies it row-wise to correct the window-aggregate columns.

What's here

WindowApply enum describes how to correct each output column. Two shapes:

• Scalar { op, offset, output_column } — fast path via arrow kernels. No Accumulator reconstructed. Covers aggregates whose per-row output is a valid partial state:
  • ScalarOp::Add — SUM, COUNT, row_number (via arrow::compute::kernels::numeric::add)
  • ScalarOp::Min / Max — MIN / MAX (via cmp::lt_eq / gt_eq + zip::zip)
  • ScalarOp::Overwrite — first_value / last_value (constant-fill from the scheduler-picked global scalar)
• Aggregate { udf, args, output_column, window_expr_index } — seeded-Accumulator re-run. Constructs a fresh Accumulator via AggregateExprBuilder, seeds it via merge_batch with the pre-merged offset state, then replays each row through update_batch + evaluate to overwrite output_column. Covers aggregates whose per-row output isn't a valid partial state:
  • AVG (without decomposition — (sum, count) state can't be recovered from a mean scalar)
  • Sketch-backed windows: APPROX_DISTINCT, APPROX_QUANTILE, etc.
  • Statistical aggregates: STDDEV, VAR, correlation family
  • Concat aggregates: ARRAY_AGG, STRING_AGG

The operator picks the shape per-column from what the caller supplies — one uniform code path.

Explicitly out of scope (documented in the module header):

• lead / lag / nth_value — solved by halo rows in the shuffle layer, not per-row correction.
• rank / dense_rank / percent_rank / cume_dist / ntile — need a segment-tree-plus-broadcast design; separate infrastructure.

Assumes at most one PARTITION BY key per input partition (matches the AQE synthetic-PARTITION-BY pattern). Multi-key handling can grow when a workload needs it; today it errors explicitly rather than silently applying one key.

Relation to DataFusion

The upstream tasks' finalized state — which the scheduler prefix-merges before handing to this operator — is produced by BoundedWindowAggExec::finalized_partition_state, added in apache/datafusion#24007. A local FinalizedPartitionState type alias in this crate mirrors the type from that PR so this crate compiles against stable DataFusion 54 until the DF change lands.

Test plan

Six tests. Three exercise construction-time validation; three are end-to-end demos, one per aggregate class, showing the operator producing globally-correct output from mock BWAG output that is only locally correct.

• try_new_rejects_state_length_mismatch — per_partition_state.len() must match input partition count.
• try_new_rejects_scalar_offset_length_mismatch — Scalar variant's offset length must match input partition count.
• try_new_rejects_output_column_out_of_range — output_column past the schema errors at construction.
• sum_corrects_running_sum_across_partitions — SUM via Scalar::Add. Local [4, 9, 15] corrected to global [10, 15, 21].
• avg_corrects_running_mean_across_partitions — AVG via Aggregate (no decomposition). Local [40.0, 45.0, 50.0] corrected to global [25.0, 30.0, 35.0] after seeding with (sum=60, count=3).
• approx_distinct_corrects_running_distinct_across_partitions — APPROX_DISTINCT via Aggregate. Local [1, 2, 3] corrected to global [4, 5, 6] after seeding with the HLL state of {1, 2, 3}. This is the case that provably can't be handled by a per-row scalar ProjectionExec — recovering HLL registers from a running count is not tractable.

All test setup mocks the upstream BWAG output so nothing here depends on apache/datafusion#24007 landing first. Once that PR merges, the local FinalizedPartitionState alias can be swapped for the DataFusion type directly.

What's not yet here (follow-ups)

• Scheduler-side state collection (analog of collect_runtime_stats_reports) that gathers per-task finalized_partition_state and computes the prefix-merge.
• Downstream stage generation that injects PrefixMergeExec with per-task offsets baked in.
• Plan rule that recognizes AQE range-shuffled windowed aggregates and routes them through this operator (with canonicalization to BoundedWindowAggExec where the upstream is sorted).

[andygrove]

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]

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.