feat: parallel BWAG for the range-window shape (h2o Q8)

[avantgardnerio]

Summary

Turns h2o's window-Q8 shape (sum(v2) OVER (ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW)) from a serial pipeline into a distributed range-shuffle.

The stack layers cleanly on top of the parallel-window primitives that have already landed (#2038, #2169, #2175, #2180, #2195, #2196):

• ParallelWindowRule matches BoundedWindowAggExec in the "no PARTITION BY + single Column ORDER BY on Float64 + finite RANGE frame" shape and rewrites as a completely parallel pipeline.
• is_stage_boundary in DistributedExchangeRule teaches the SPM branch that a RangeFilterExec sitting directly on a resolved ExchangeExec counts as part of the boundary: we chose not to fold range-filtering into ShuffleReader/ExchangeExec, so the filter is conceptually part of the boundary shape by design.
• PartitionedBoundedWindowAggExec is a Ballista-specific wrapper for DataFusion's BoundedWindowAggExec. It hides BWAG from tree walkers (children() returns only the wrapper's input) and declares Distribution::UnspecifiedDistribution, so EnforceDistribution doesn't insert an SPM(K→1) beneath. execute(i) delegates straight to BWAG::execute(i), which already processes each partition independently: DataFusion's BWAG algorithm has no cross-partition state. Safe because the rule's shape gates guarantee range-repartition upstream + halo covers frame boundaries.

Effectively, this is apache/datafusion#23026 (parallel-BWAG) implemented as a Ballista-side wrapper: one operator, no DF-internals fork.

Plan shape

Stage 1 (K MPT tasks in parallel):
  ProjectionExec
    RangeFilterExec (narrow, halo=[0,0])
      PartitionedBoundedWindowAggExec (wraps BWAG; declares UnspecifiedDistribution)
        RangeFilterExec (wide, halo_lo=frame_low, halo_hi=frame_high)
          ExchangeExec (range_repartition_cuts)              ← stage boundary

Stage 0 (K MPT tasks in parallel):
          RuntimeStatsExec (post-ORRE per-partition sketch → scheduler)
            OrderedRangeRepartitionExec (K outputs, sorted, range-disjoint)
              SortExec (preserve_partitioning=true)
                RuntimeStatsExec (local sketch → feeds ORRE's cut walker)
                  <source>

**note** indexed reading in the picture above is in [proposal stage](https://github.com//pull/2204).

MPTs come from the existing ballista.scheduler.max_partitions_per_task knob; on a 2×4-vcore cluster with max=4 we get 2 tasks per stage, one per exec.

Results

Measured on EKS (2 pods, each on its own r6i.24xlarge node, 8 vCPU × 64 GiB per pod, gp3 shuffle work-dir, h2o parquet on S3 in eu-west-1). K=8, MPT=4 → 2 tasks/pod/stage. Sweep runs the full 17-query h2o window.sql suite at 1e8 rows (100M), once with parallel_window.enabled=false, once with =true.

Q8 headline (the shape this PR targets)

SELECT sum(v2) OVER (ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) FROM large

Config	Wall time
Baseline (parallel_window.enabled=false) : SPM collapse → 1-partition BWAG	288.4 s
Parallel (parallel_window.enabled=true) : K=8, 4 partitions/pod	60.4 s
Speedup	4.78×

Per-stage breakdown of the parallel run (from scheduler /api/job/{id}/stages):

Stage	Op tree	Elapsed compute	Task shape
0	ShuffleWriter → RuntimeStats → ORRE → RuntimeStats → SortExec → DataSource(s3)	6.9 s aggregate	2 tasks
1	ShuffleWriter → Projection → RangeFilter(narrow) → PartitionedBWAG → RangeFilter(wide) → RangeShuffleReader	46.7 s aggregate	2 tasks

Stage 0 is dominated by parquet scan + sort; ORRE scatter is cheap (~6 ns/row across 100M rows). Stage 1 is BWAG-dominated as expected; each pod runs 4 output partitions concurrently over its 8 physical vCPUs.

Full 17-query null-test sweep

The rule is deliberately narrow (Q8's exact shape). The other 16 queries are the null-test: same on/off wall time confirms no regression on shapes the rule doesn't target.

Q	prior (s)	post (s)	ratio	rule status
1	25.6	24.5	1.05×	PREFIX - sum() OVER (), could range partition / prefix scan
2	45.5	45.8	0.99×	PREFIX - first_value/row_number ORDER BY id3 (unbounded frame)
3	64.0	64.9	0.99×	HASH - PARTITION BY id1/id2/id3 already hash-parallel
4	20.7	19.7	1.05×	HASH - PARTITION BY id2
5	26.3	26.2	1.01×	ROWS - ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING (index-based)
6	43.4	43.9	0.99×	ROWS - ROWS BETWEEN 100 PRECEDING AND CURRENT ROW
7	38.5	38.2	1.01×	ROWS - ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
8	288.4	60.4	4.78×	← target shape
9	18.6	18.4	1.01×	HASH - PARTITION BY id2 ORDER BY id3 ROWS ...
10	19.2	19.6	0.98×	HASH
11	19.3	20.3	0.95×	HASH
12	29.7	30.3	0.98×	HASH - exact Q8 RANGE shape but with PARTITION BY id2
13	8.7	8.6	1.00×	HASH - ROW_NUMBER PARTITION BY id2 top-N
14	10.9	13.6	0.80×	HASH - top-N (variance dominates at ~10 s)
15	10.6	9.9	1.07×	HASH
16	13.5	9.5	1.42×	HASH - variance
17	9.0	13.0	0.69×	HASH - variance

Why Q8 is the only mover: it's the only query in the suite where the ordered window can't be parallelized by any existing mechanism.

• HASH (11 queries): PARTITION BY gives a natural K-space; EnforceDistribution → Hash(key) → per-partition BWAG already runs in parallel across K executors. Nothing for this rule to add.
• ROWS (3 queries): frame semantics are index-based (row counts), not value-based. Needs a different distributed algorithm (routing by row-index halo, not by value halo).
• PREFIX (2 queries): unbounded frame (row_number, first_value default). Needs a prefix-scan design (cross-partition running state).
• Q8 has neither PARTITION BY (so no natural K-space) nor a cheap-per-row window op (RANGE frame does value-band search per row, ~2.9 μs/row vs. ~380 ns/row for Q7's ROWS UNBOUNDED cumulative sum). That intersection - "must be single-partition AND single-partition is genuinely expensive" - is the niche this rule fills.

Workload-shape caveat: the 4.78× win assumes BWAG's per-row cost is comparable to what's around it. Heavier per-row work (a cryptographic hash, an expensive UDF, a regexp_replace chain) would scale even better, closer to the 8× vcore count. Lighter than a running sum is hard to construct, so 4.78× is roughly the floor for the shape this PR handles.

Correctness: .github/workflows/h2o.yml runs each h2o query against a single-process DataFusion oracle under --verify and diffs the row sets. Q8's parallel-path aggregate sum matches baseline to within Float64 associativity noise (relative delta ~5×10⁻¹⁴, under the √N × ε floor for 100M-row parallel summation).

Follow-ups (not this PR)

• KLL migration lifts the Float64/non-nullable restriction on the routing expression.
• Cross-stage cut coordination for SMJ / Union legs with range-repartition on each side.
• Range paritioned index files.

[avantgardnerio]

@andygrove @phillipleblanc we now have our first real (EKS) benchmark results for halo-based parallel windows. h2o is kind of a worst case, because as the data keeps getting bigger, the value range stays the same, so the cost of duplicated halos starts to dominate:

  ┌─────┬─────────────┬──────────────────────────┬─────────────────────────────────────┐
  │  K  │ Theoretical │ With #2204 range-pruning │           Without (today)           │
  ├─────┼─────────────┼──────────────────────────┼─────────────────────────────────────┤
  │ 8   │ 6.5×        │ 6.5×                     │ ~4.7× observed                      │
  ├─────┼─────────────┼──────────────────────────┼─────────────────────────────────────┤
  │ 16  │ 10.8×       │ ~10×                     │ probably ~6× (I/O amp bites harder) │
  ├─────┼─────────────┼──────────────────────────┼─────────────────────────────────────┤
  │ 32  │ 16.3×       │ ~14×                     │ probably ~7× (I/O amp dominates)    │
  ├─────┼─────────────┼──────────────────────────┼─────────────────────────────────────┤
  │ 100 │ ~25×        │ ~20×                     │ collapses (300% halo)               │
  └─────┴─────────────┴──────────────────────────┴─────────────────────────────────────┘

but even if we cap it at 8 or 32 cores, 5-15x wins seem worthwhile? (not to mention OOM reduction). I promised "this time" I was just going to deliver a PR with e2e abilities, but TBH it's grown too big again. Given that it now proves it works, you guys cool with me splitting off lots of little supporting PRs from this?

[avantgardnerio]

@phillipleblanc reviews are always appreciated if you have time. I finally got the diff down on this one to a reviewable size by merging all its dependencies. I think prefix scanning or extending this to cover ROWS BETWEEN might fix our next bottleneck.

[milenkovicm]

thanks @avantgardnerio , will need some time to review this

[phillipleblanc]

This makes sense as an improvement - I have one question about a potential wrong-answer issue:

• For ascending order, 3 PRECEDING means smaller numeric values, so widening the lower side of a range is correct.
• For descending order, 3 PRECEDING means larger numeric values. The PR still widens the lower side, so rows near a boundary may not receive all the values required for their window.

I believe this means the query can succeed while returning incorrect results. We could initially restrict this optimization to ascending order, which will prevent the issue, but the correct fix would be to reverse the halo directions for descending order.

[avantgardnerio]

thanks @phillipleblanc ! That was indeed a bug. I've added the ASC restriction for now, a test, and a TODO. I think it would be better to tackle it in a follow up PR to keep this one from growing any bigger.

[andygrove]

This is epic! Thanks @avantgardnerio and thanks @phillipleblanc for the review