[PoC] Parallel prefix scan for cumulative-aggregate window functions

[avantgardnerio]

Status: PoC, not for merge. Stacked on top of apache#23026's branch. Open as a draft because this is just kicking the tires on how distribution could work in Ballista.

Edit: I think this is actually upstreamable in either parallel or streaming forms.

Based on: https://www.cs.cmu.edu/~guyb/papers/Ble93.pdf

What shape this parallelizes

Cumulative aggregates: ORDER BY x ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW. No PARTITION BY. Canonical log-analytics shape — "running sum over time."

SELECT seq, SUM(amount) OVER (
    ORDER BY seq
    ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
) AS cumulative_sum
FROM events

Today this hits the same Distribution::SinglePartition collapse as the bounded-RANGE case the parent PR addresses. But HaloDropExec's tactic doesn't apply: halo means "give me N units of context past my bucket boundary," and cumulative means "give me every prior row," which is exactly the same as serializing.

What the PoC does instead: parallel prefix scan

Two-stage. Stage 1 is the existing RangeRepartitionExec (sans halo) + per-partition BoundedWindowAggExec with parallel_aware = true — each partition computes its own local cumulative sum starting at zero, in parallel. Stage 2 is a new CarryExec that pipeline-breaks: it gathers all partitions' batches, derives each partition's final cumulative value from the last row of its last batch, computes the prefix sum across finals, and re-emits each partition's batches with that prefix added to the agg column.

Conceptually this mirrors the Ballista inter-stage shuffle model, just intra-node: stage 1's outputs are fully materialized; stage 2 reads stage 1's totals as a small "broadcast" and re-emits with offsets applied.

Plan transformation:

# Today
ProjectionExec
  GlobalLimitExec
    BoundedWindowAggExec(SinglePartition)
      SortPreservingMergeExec [seq]
        SortExec preserve_partitioning=true
          DataSourceExec(N partitions)

# With this PR
SortPreservingMergeExec [seq]
  CarryExec                              # <- new; gathers + offsets
    BoundedWindowAggExec(parallel_aware) # <- runs per partition
      RangeRepartitionExec               # <- no halo
        SortExec preserve_partitioning=true
          DataSourceExec(N partitions)

CarryExec mechanics

"First poll wins" via tokio::sync::OnceCell — no spawned coordinator task, no per-partition oneshot channels. Whichever output partition polls first runs the gather; concurrent and subsequent polls await the same memoized result. Since Carry is pipeline-breaking, downstream can't emit before gather completes anyway, so we lose no parallelism by gathering inline.

The buffered batches are the finals state — partitions[i].last().last_row()[agg_col] is partition i's final. No separate Vec<ScalarValue>. Prefix is just a running scalar accumulated via ScalarValue::add.

Empty partitions contribute the additive identity (zero in the agg type); their prefix equals the running total at that point. "No data anywhere" passes through with a null prefix (unchanged).

TDD via SLT

parallel_window.slt gains four assertions for the cumulative case:

1. EXPLAIN — locks the target plan shape.
2. LIMIT 10 result — rows 0..9 sit entirely in partition 0 where carry-in is zero anyway; correct on a passthrough Carry too, so this catches plan-shape regressions, not carry correctness.
3. count(cumulative_sum) sentinel — proves Carry neither drops nor duplicates rows.
4. Cross-partition-boundary result — queries seq in {24, 25, 26, 49, 50, 51, 74, 75, 76} which straddle the range-partition cutoffs (~25, 50, 75). This is the real correctness gate: a passthrough Carry returns local-cumsum-restarting-at-0 at each boundary; the real Carry's offsets line up with the serial baseline.

The middle commit (7020f8e) wires the rule branch with a passthrough Carry stub and shows assertion 4 going red; the final commit (468f293) implements the prefix-scan body and flips it green. Useful as a self-contained "first see the wrong thing, then see the right thing" history.

Same as parent: don't merge, don't chase CI green (this branch inherits the parent's failing checks). Discussion only.

[avantgardnerio]

Just add the "prefix" from prior partitions to each row.

[avantgardnerio]

Calculate the prefix to add

[avantgardnerio]

@2010YOUY01 and @milenkovicm - I can't seem to add you as reviewers, since this is a stacked PR on a fork, but perhaps you will see my mention.

[avantgardnerio]

Bench (16 cores / 32 threads):