avantgardnerio opened a new pull request, #2223:
URL: https://github.com/apache/datafusion-ballista/pull/2223
## 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. On h2o Q8 @ 1e7 under a 2 GB/exec cgroup cap, 2 execs × 4
vcores, this stack takes **41 s and returns the full 10M rows** — vs. 155 s /
1.55M rows for the SPM-collapsed baseline (the previous shape wasn't just slow,
it was silently dropping ~87% of the output because the K→1 SPM+BWAG collapse
only kept partition 0's range).
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 to `RangeFilterExec(narrow) → PartitionedBoundedWindowAggExec →
RangeFilterExec(wide) → RuntimeStatsExec → OrderedRangeRepartitionExec →
RuntimeStatsExec → SortExec → source`. Position 2 in the AQE chain, before
`DistributedExchangeRule`.
- **`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, no protocol
changes.
## 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)
RuntimeStatsExec (local sketch — feeds ORRE's cut walker)
SortExec (preserve_partitioning=true)
<source>
```
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.
## Test plan
- [ ] `cargo test -p ballista-core -p ballista-scheduler` — all pass locally
(315 scheduler + 235 core).
- [ ] `rewrites_q8_shape` unit test asserts
`PartitionedBoundedWindowAggExec` and no `SortPreservingMergeExec` in the
rewritten plan.
- [ ] `spm_skips_when_range_filter_covers_exchange` unit test asserts DE
doesn't inject an Exchange between SPM and a rule-planted `RangeFilterExec`.
- [ ] End-to-end h2o Q8 @ 1e7 completes in ~41 s under a 2 GB/exec cgroup
cap, returns 10M rows, both stages run 2 MPT tasks in parallel.
## Follow-ups (not this PR)
- Data-distribution skew on h2o Q8: task 0 finishes ~11 s before task 1
(16.7 s vs 27.7 s). Cuts equalize row count per partition; per-row compute in
the RANGE frame scales with local density. Options later: cut on estimated
*work*, higher K, or dynamic re-balance.
- KLL migration ([[kll-sketch]]) lifts the Float64/non-nullable restriction
on the routing expression.
- Cross-stage cut coordination for SMJ / Union legs with range-repartition
on each side.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
