Skip to content

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

Draft
avantgardnerio wants to merge 13 commits into
apache:mainfrom
avantgardnerio:brent/parallel-range-window
Draft

feat: parallel BWAG for the range-window shape (h2o Q8)#2223
avantgardnerio wants to merge 13 commits into
apache:mainfrom
avantgardnerio:brent/parallel-range-window

Conversation

@avantgardnerio

@avantgardnerio avantgardnerio commented Aug 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

Follow-ups (not this PR)

  • 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.
  • Range paritioned index files

avantgardnerio and others added 9 commits August 3, 2026 13:42
@avantgardnerio
add h2o benchmark suite (groupby / join / window)
cf32d26 
Ports the h2o `db-benchmark` runner from arrow-datafusion into
`benchmarks/src/bin/h2o.rs`, with `datafusion` and `ballista`
subcommands sharing the same query loader and table registration.

Suite is inferred from the queries file name:
- groupby.sql → registers `x`
- join.sql    → registers `x`, `small`, `medium`, `large`
- window.sql  → registers `large` (the 4th `--join-paths` entry)

Query files are copied verbatim from
`arrow-datafusion/benchmarks/queries/h2o/` so query indices align 1:1
between the two runners.

Data generation is unchanged from upstream — point `--path` /
`--join-paths` at h2o CSV or Parquet files produced by
`arrow-datafusion/benchmarks/bench.sh data h2o_small_window` (or the
equivalent for `h2o_small` / `h2o_small_join`).
@avantgardnerio @claude
feat(benchmarks): add --explain flag to the h2o binary
f072677 
Adds `--explain` to both `datafusion` and `ballista` subcommands: prints
the physical plan for each query in `--query`/`--queries-path` and exits
without executing. Convenient for iterating on planner rules against the
h2o window suite without spinning up an executor loop.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
refactor(core): replace PerPartitionFilterExec with RangeFilterExec
3a77f01 
PerPartitionFilterExec's only caller (BallistaAdapter above ShuffleReader
for hash-agg correctness) was already using it as a range-shaped filter.
Widening it into a general per-partition arbitrary-predicate op — with
halo-widening bolted on for the parallel-window rewrite — would leak a
range concept into an arbitrary-predicate contract.

RangeFilterExec is the honest shape: routing_expr + cuts + halo_lo /
halo_hi. Per-partition semantics fall out of the local partition index,
not from a Vec of independent predicates. Ordering knowledge on the
input opens the door to a future ValueIndexReader-driven binary-search
path (PR apache#2204 direction) that a generic FilterExec can't take.

Notable pieces:

- `cuts: Arc<Mutex<Option<Vec<f64>>>>` + `resolve_cuts` API mirror
  `ExchangeExec::range_repartition_routing()` — the ParallelWindow rule
  plants a pending RangeFilterExec at plan time; the scheduler resolves
  cuts after stage 0's RuntimeStatsExec reports merge. `execute` and
  serialization both refuse while cuts are unresolved.
- `partition_indices: Vec<usize>` maps local → global partition index.
  Restrict slices this mapping without touching cuts (cuts stay whole;
  they describe the K global partitions). Replaces PPFE's per-partition
  predicate-vec slicing in task_builder's restrict path.
- Public API + proto speak `ScalarValue` (not `f64`) per the type-
  generality rule for the range-repartition family: the outer contract
  is type-agnostic so KLL can widen internal storage later without an
  API break. Internal downcast to `f64` today; non-Float64 inputs error
  with a clear message.
- Adapter builds RangeFilterExec with `halo=0` for the existing hash-agg
  case; the parallel-window rule will build it with non-zero halo.

Migration: delete `PerPartitionFilterExec`, migrate all callers, rename
proto `PerPartitionFilterExecNode` → `RangeFilterExecNode`, update doc
comments. Full test suite (597 tests) passes.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
feat(scheduler): ParallelWindowRule — distributed range-shuffle for BWAG
fcb3152 
Adds `ParallelWindowRule` to the AQE default_optimizers chain (position 2,
before `SelectJoinRule` and DF's own optimizers, before `DistributedExchangeRule`).
The rule matches bounded RANGE-frame windows with no PARTITION BY and a
single-column Float64 ORDER BY, and rewrites them into a range-shuffle
so BoundedWindowAggExec's SinglePartition requirement isn't a serial
bottleneck. Shape:

  RangeFilterExec (narrow, halo=0, cuts=pending)
    BoundedWindowAggExec
      SortPreservingMergeExec
        RangeFilterExec (wide, halo=frame bounds, cuts=pending)
          RuntimeStatsExec (post-ORRE per-partition sketch → scheduler)
            OrderedRangeRepartitionExec (K sorted disjoint outputs)
              RuntimeStatsExec (local sketch; feeds ORRE's cut walker)
                SortExec (preserve_partitioning=true)
                  <source>

Both RangeFilterExecs are planted with cuts=None. After stage 0's tasks
complete and their RSE reports are merged into K-1 quantile cuts, the
scheduler's `resolve_range_filter_cuts` walker (in `adapt_to_ballista`)
finds every pending RangeFilterExec in the downstream stage's plan and
resolves it against the matching ExchangeExec's routing_expr. Adapter no
longer injects RangeFilterExec — the rule is the sole planter, single
source of truth. Idempotency guard on the rule bails when the BWAG's
subtree already contains our own ORRE/RangeFilter (AQE re-plans fire the
chain again on the already-rewritten plan).

Also relaxes `ORRE::try_new` — the child-claims-sortedness check moved
from construction to `execute()`. Rule-time construction races with
`EnforceSorting` (which planted a SortExec on ORRE's declared
`required_input_ordering` *after* the rule ran), so refusing at try_new
was too strict. The runtime check at execute() still catches invariant
breaks; two tests moved from `try_new_rejects_*` to `execute_rejects_*`.

Verified on h2o Q8 at 1e7 under a 2G/exec cgroup cap: stage 0 (8 tasks)
and stage 1 (8 tasks) both parallelize across both executors, no OOM.
Stage 2 still collapses to a single task doing SPM+BWAG+narrow — DE
inserts a shuffle boundary below the SPM, putting BWAG in the final
stage. That collapse is the next follow-up; the machinery for the
range-shuffle itself is in.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
fix(scheduler): stop DE inserting Exchange between SPM and rule-plant…
a522a73 
…ed RangeFilterExec

`ParallelWindowRule` plants a `RangeFilterExec` directly on the resolved
range-repartition `ExchangeExec`, with `SortPreservingMergeExec` above.
`DistributedExchangeRule`'s SPM branch was checking whether SPM's
immediate child was an `ExchangeExec` — seeing the `RangeFilterExec`, it
injected another `ExchangeExec`, cutting the plan into an extra collapse
stage.

Introduce `is_stage_boundary` and treat a `RangeFilterExec` sitting
directly on an `ExchangeExec` as part of the boundary. That is a
conscious design shape — we chose not to fold range-filtering into
`ShuffleReader`/`ExchangeExec`, so the filter is part of the boundary
by construction. This matches the pre-`fcb31520` behaviour where the
adapter injected the filter after DE had already run.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
refactor(scheduler): size ParallelWindowRule's K from source partitions
9a317cc 
K was `config.execution.target_partitions.max(2)` — a placeholder
chosen while writing the rule. The natural sizing is
`source.output_partitioning().partition_count()`: ORRE re-slices each
input partition into a range-disjoint output partition, so K = input
partitions is the 1:1 rearrangement.

No behaviour change on h2o Q8 (`target_partitions` and source
partitions both settle at 8), but the rule no longer depends on the
config knob or its `.max(2)` fallback.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
feat(core, scheduler): PartitionedBoundedWindowAggExec — parallel BWA…
b9db3e8 
…G for the range-window shape

DataFusion's `BoundedWindowAggExec` declares `SinglePartition` when no
PARTITION BY is present, forcing `EnforceDistribution` to collapse K→1
via `SortPreservingMergeExec`. With `ParallelWindowRule`'s range-
repartition upstream, each ORRE output partition is a globally
range-disjoint slice + halo — BWAG can safely run per-partition on those
K slices and produce K correct outputs. `PartitionedBoundedWindowAggExec`
wraps BWAG, exposes only the input as its plan-tree child (BWAG itself
is hidden from tree walkers), and overrides `required_input_distribution`
to `UnspecifiedDistribution`. `execute(i)` delegates to the wrapped BWAG,
which already processes each partition independently.

- `ballista_core::execution_plans::partitioned_bounded_window_agg`: the
  new operator. `InputOrderMode` and `can_repartition` are hardcoded
  (`Sorted` / `false`) per the rule's shape gates.
- `BallistaPhysicalPlanNode::PartitionedBoundedWindowAgg`: proto message
  carrying only `window_expr` — the rest is implicit from the rule's
  invariants. Round-trip goes through DF's
  `serialize_physical_window_expr` / `parse_physical_window_expr`.
- `ParallelWindowRule::rewrite_bwag`: drops the SPM the previous
  rewrite planted between BWAG and the wide `RangeFilterExec`, and
  swaps BWAG for `PartitionedBoundedWindowAggExec`. K is sourced from
  `config.execution.target_partitions.max(2)` — at rule-fire time
  `DataSourceExec` still has 1 file_group (splits happen later in the
  AQE chain), so the plan tree can't yet tell us the true source
  width. Reverts the "size K from source" refactor.
- `rewrites_q8_shape` test now asserts `PartitionedBoundedWindowAggExec`
  and NO `SortPreservingMergeExec` in the output.

On h2o Q8 @ 1e7 under a 2G/exec cgroup cap, 2 execs × 4 vcores,
`ballista.scheduler.max_partitions_per_task=4`: 41 s (down from 155 s)
and returns the full 10M rows (previous runs returned only 1.55M —
the K→1 collapse dropped ~87% of the output because the narrow
`RangeFilterExec` above the collapsed BWAG kept only partition-0's
range). Both stages run 2 MPT tasks (one per exec).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio
feat(core, scheduler): gate ParallelWindowRule behind ballista.planne…
94cd20c 
…r.parallel_window.enabled

Adds an opt-in config flag so users of AQE don't inherit the range-window
rewrite by default. Matches the shape of `ballista.planner.coalesce.enabled`:
new AQE rule → new opt-in flag. Default `false`.

- `BALLISTA_PARALLEL_WINDOW_ENABLED` + registry entry + getter on
  `BallistaConfig`.
- Guard clause at the top of `ParallelWindowRule::optimize` returns the
  plan untouched when the flag is off.
- Existing shape tests keep the rule enabled through the local `optimize`
  helper; a new `disabled_by_default` test asserts the rewrite is inert
  without the extension registered.
- Regenerated `docs/source/user-guide/configs.md`.
@avantgardnerio
notes
5b6f7a6
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Aug 4, 2026
@avantgardnerio
avantgardnerio force-pushed the brent/parallel-range-window branch from 2d9e92a to 5b6f7a6 Compare August 4, 2026 18:43
avantgardnerio and others added 4 commits August 4, 2026 13:54
@avantgardnerio @claude
feat(core, scheduler, executor): RangeShuffleReaderExec — ordered k-w…
d87321d 
…ay merge at stage boundary [WIP]

Closes the RANGE-frame correctness gap: the regular ShuffleReaderExec
concatenates upstream sources in arrival order, breaking the monotonicity
BWAG's Range-frame cursor assumes. RangeShuffleReaderExec keeps each source
alive as its own stream and feeds them all into StreamingMerge on the
child's declared ordering.

- new RangeShuffleReaderExec (fetch reuses shuffle_reader helpers, backpressure
  via merge demand; no permit governor, no per-source buffering)
- adapter plants it whenever exchange.input().output_ordering().is_some()
- proto + codec round-trip; executor work_dir/client_pool late binding;
  task_builder partition-slice restriction

h2o Q8 @ 1e7 SUM diff: parallel_window=true/false now agree to 5e-14 relative
(FP noise floor). Previously diverged at run boundaries.

Follow-ups (see next-session TODO): planner.rs::rollback_resolved_shuffles
falls through, cluster/mod.rs::stage_has_input_collapse falls through.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
refactor(scheduler): audit-driven RangeShuffleReaderExec whitelist fills
c8f3f20 
Fills two production downcast sites that fell through the RangeShuffleReaderExec
shape, plus fmt fallout from the initial slice.

- planner::rollback_resolved_shuffles: rolls range readers back to plain
  UnresolvedShuffleExec. Range-ness is derived at plan time from the child's
  ordering, so a re-plan's adapter walk re-plants a fresh range reader — no
  proto extension needed.
- cluster::stage_has_input_collapse: range reader is a stage boundary; the
  walker must stop there, else a single-output-partition range reader
  spuriously trips the `partition_count == 1` collapse arm.

Tests:
- rollback_resolved_shuffles_reduces_range_reader_to_plain_unresolved
- stage_has_input_collapse_stops_at_range_reader

Fmt: adapter's #[cfg(test)] mod moved below resolve_range_filter_cuts to
satisfy clippy::items_after_test_module.

Follow-up still open: execution_graph_dot.rs graphviz — will render generic
node label for the range reader. Diagnostic only, safe to punt.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
docs(parallel-range-window): tick Ordered ShuffleReader, drop correct…
1dcc7b4 
…ness-gap section

The gap is closed — RangeShuffleReaderExec ships in d87321d + c8f3f20, and
the h2o Q8 SUM diff between parallel_window=true/false lands at 5e-14 relative
(Float64 noise floor). Rewrite the ticked line to describe the landed shape
and note the writer-vs-demand-driven follow-up. Drop the "Dot-product check"
bullet for the reader (now landed) and the whole correctness-gap section.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio @claude
perf(core): RangeFilterExec min/max fast paths + binary-search slice
260678b 
When `input.output_ordering()` leads with `routing_expr` ascending, take one
of three shortcuts on each batch before `filter_record_batch`:

  - `last < lo` or `first >= hi` → drop the whole batch (skip).
  - `first >= lo && last < hi`   → pass the batch through unchanged (Arc-clone).
  - mixed                         → `partition_point` on the Float64Array values
                                    for lo/hi indices + `RecordBatch::slice`
                                    (zero-copy view).

Nullable routing columns fall back to `filter_record_batch` on a per-batch
basis (Float64Array::values() returns garbage for null slots, breaking
partition_point). `sorted_on_key` is derived at construction — no config knob.

h2o Q8 with 2 execs × 4 vcores × MPT=4:

  scale  cap  parallel_window=false  parallel_window=true  speedup
  1e7    2G   7.6 s                  2.5 s                 3.0×
  1e8    4G   143 s                  92 s                  1.55×

The 1e8 delta is smaller because the bottleneck shifts to shuffle IO /
whole-file merge memory — the ValueIndex + per-task halo work next.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

documentation Improvements or additions to documentation

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@avantgardnerio