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).

🤖 Generated with Claude Code