From 99716828e6beabd1c020be96c57e41faa4362a63 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 3 Aug 2026 15:45:03 -0600 Subject: [PATCH 01/11] refactor(core): replace PerPartitionFilterExec with RangeFilterExec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 #2204 direction) that a generic FilterExec can't take. Notable pieces: - `cuts: Arc>>>` + `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` 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) feat(scheduler): ParallelWindowRule — distributed range-shuffle for BWAG 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) 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) fix(scheduler): stop DE inserting Exchange between SPM and rule-planted 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) refactor(scheduler): size ParallelWindowRule's K from source partitions 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) feat(core, scheduler): PartitionedBoundedWindowAggExec — parallel BWAG 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) feat(core, scheduler): gate ParallelWindowRule behind ballista.planner.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`. notes feat(core, scheduler, executor): RangeShuffleReaderExec — ordered k-way 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) refactor(scheduler): audit-driven RangeShuffleReaderExec whitelist fills 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) docs(parallel-range-window): tick Ordered ShuffleReader, drop correctness-gap section The gap is closed — RangeShuffleReaderExec ships in d87321db + c8f3f20f, 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) perf(core): RangeFilterExec min/max fast paths + binary-search slice 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) fix(docs): rustdoc + prettier CI - rustdoc: three `[`resolve_cuts`]` references in `range_filter.rs` (module doc + two item docs) resolved to no target; qualify as `RangeFilterExec::resolve_cuts` / `Self::resolve_cuts` so cargo doc no longer errors on ballista-core. - prettier: `docs/developer/parallel-range-window.md` had two `*emphasis*` spans (`*shape*`, `*task-level*`) and a stray blank line — prettier wants `_emphasis_` + single blank. No content change. Co-Authored-By: Claude Opus 4.7 (1M context) feat(core): RangeFilterExec metrics — fast-path counters + baseline Was returning `None` from `metrics()`, so the operator was invisible in the scheduler's stage-metrics dump. Adds `ExecutionPlanMetricsSet` with `BaselineMetrics` (elapsed_compute, output_rows via record_poll) and five path counters: `fast_skip_batches`, `fast_pass_batches`, `fast_slice_batches`, `slow_batches`, `input_rows`. Timer scoped post-poll so upstream shuffle IO isn't billed to this op. Scout on h2o Q8 @ 1e8 confirms fast path is firing as intended (99%+ pass-through on the narrow filter, 85% skip on the wide one, zero slow-path fallbacks) — filter is not the perf bottleneck. Co-Authored-By: Claude Opus 4.7 (1M context) fix(docs): drop unresolved intra-doc link in parallel_window `resolve_range_filter_cuts` is private to the adapter module and not in scope from `parallel_window.rs`, so rustdoc rejects the intra-doc link under `-D warnings`. Keep it as plain inline code. Co-Authored-By: Claude Opus 4.7 (1M context) style(core): join split struct decl to satisfy rustfmt Co-Authored-By: Claude Opus 4.7 (1M context) refactor(core): slim RangeFilterExec — scheduler owns cuts/partition indices, RFE widens by halo Split the range-partitioning concerns out of RangeFilterExec so it looks like PerPartitionFilterExec's counterpart for sorted-key filtering: - RFE fields: input, routing_expr, halo_lo/hi (ScalarValue), raw_bounds (late-bound), sorted_on_key detection, metrics. - Gone: cuts, partition_indices, resolve_cuts, restrict_partitions, try_new_with_indices. Widening from cuts+halo to per-partition bounds moves scheduler-side (adapter builds raw_bounds from cuts; RFE widens by its own halos internally at resolve_bounds time). - task_builder RFE branch is now a plain "slice raw_bounds parallel to input restriction" — no partition_indices remap. - All APIs and proto fields are ScalarValue (arrow-primitive-generic); internal downcast to f64 with Err for non-Float64 until KLL widens. Halos are functional on RFE (widens raw→widened at resolve time), not write-only decoration. The scheduler-side cut_partitions also needs halo-widened overlap for correct file routing to RANGE-frame consumers; that's a separate cross-stage lookup left as a TODO here. Co-Authored-By: Claude Opus 4.7 (1M context) docs(core): fix RangeFilterExec intra-doc link at module scope `[\`Self::resolve_bounds\`]` on line 39 was in the module-level `//!` comment where `Self` is not defined. CI runs cargo doc with -D warnings so it fails; local runs pass silently. Use the fully-qualified path. Co-Authored-By: Claude Opus 4.7 (1M context) docs(scheduler): revert PPFE→RFE comment renames to reduce PR diff Three files carried only doc/comment renames from PerPartitionFilterExec to RangeFilterExec — no code changes. PPFE still exists in the tree, so the original phrasing remains accurate. Reverting shrinks the PR's review surface without touching semantics; a follow-up sweep can update these comments after PPFE is fully retired. - exchange.rs: 3 comment mentions of PPFE-as-cuts-consumer - test/coalesce_rule.rs: 1 test comment - test/range_repartition.rs: 2 test comments Co-Authored-By: Claude Opus 4.7 (1M context) fix(scheduler): plant RSE#1 below SortExec so cuts land before ORRE routes `ParallelWindowRule` used to leave `RSE#1` above `SortExec`, so the local sketch only started ingesting after Sort had fully materialized. ORRE consumed from RSE#1 as soon as Sort emitted, meaning the scheduler often handed ORRE a still-being-built sketch → approximate cuts → skewed shuffle files. Two changes to close this: 1. Move the rule to run *after* the DataFusion optimizer chain. At the old position the input was `BWAG → DataSource` (sources with `sort_order_for_reorder` satisfy BWAG's ordering natively, no Sort inserted yet); the SortExec placement we care about is only materialized once EnforceSorting / RepartitionFileScans have run. Running earlier also lets DF's later sort-pushdown move any Sort we plant down through the passthrough RSE#1, undoing the intended order. 2. Strip whatever DF planted for BWAG's SinglePartition + Sorted requirements (SPM and/or SortExec) and plant a fresh `SortExec → RSE#1 → source` chain below `ORRE`. The fresh Sort is the pipeline break: it consumes all input before emitting the first row, so RSE#1's sketch fully reports while Sort buffers. Q8 (h2o, SF=1e7, 8 vcores): - rule skipped (buggy pattern): 61s - RSE#1 above Sort (prior): 24s - RSE#1 below Sort (this): 17s ← ~1.4× speedup over prior Wide-RFE metrics on the new plan: input=18.84M / output=11.01M against a 10M-row dataset → 1.88× row-level read amplification, matching the theoretical (1 + halo/cut_width) ≈ 1.24× floor plus batch-granularity overhead from RangeShuffleReader. Co-Authored-By: Claude Opus 4.7 (1M context)  Conflicts:  ballista/core/proto/ballista.proto  ballista/core/src/execution_plans/mod.rs  ballista/core/src/execution_plans/range_filter.rs  ballista/core/src/serde/generated/ballista.rs --- ballista/core/src/config.rs | 18 + .../core/src/execution_plans/runtime_stats.rs | 12 +- ballista/executor/src/execution_engine.rs | 15 +- ballista/scheduler/src/cluster/mod.rs | 52 +- ballista/scheduler/src/planner.rs | 62 +- ballista/scheduler/src/state/aqe/adapter.rs | 274 +++++++-- .../optimizer_rule/distributed_exchange.rs | 71 ++- .../src/state/aqe/optimizer_rule/mod.rs | 2 + .../aqe/optimizer_rule/parallel_window.rs | 573 ++++++++++++++++++ ballista/scheduler/src/state/aqe/planner.rs | 12 +- ballista/scheduler/src/state/task_builder.rs | 136 +++-- docs/developer/parallel-range-window.md | 76 +++ docs/source/user-guide/configs.md | 1 + 13 files changed, 1177 insertions(+), 127 deletions(-) create mode 100644 ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs create mode 100644 docs/developer/parallel-range-window.md diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index b2d7ed6624..2995e94eab 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -136,6 +136,11 @@ pub const BALLISTA_COALESCE_ENABLED: &str = "ballista.planner.coalesce.enabled"; /// This could benefit the workload by injecting EmptyExec in the plan (i.e during joins) pub const BALLISTA_PROPAGATE_EMPTY_ENABLED: &str = "ballista.planner.propagate_empty.enabled"; +/// Configuration key to enable the AQE `ParallelWindowRule`, which rewrites +/// bounded-RANGE-frame windows into a distributed range-shuffle so BWAG's +/// single-partition constraint isn't a serial bottleneck. Opt-in. +pub const BALLISTA_PARALLEL_WINDOW_ENABLED: &str = + "ballista.planner.parallel_window.enabled"; /// Configuration key for the target post-coalesce partition byte size (bytes). /// Mirrors Spark's `spark.sql.adaptive.advisoryPartitionSizeInBytes`. pub const BALLISTA_COALESCE_TARGET_PARTITION_BYTES: &str = @@ -323,6 +328,14 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| of a join, allowing downstream work to be skipped.".to_string(), DataType::Boolean, Some(true.to_string())), + ConfigEntry::new(BALLISTA_PARALLEL_WINDOW_ENABLED.to_string(), + "Enables the AQE parallel-window rule (ParallelWindowRule), which \ + rewrites bounded-RANGE-frame windows into a distributed range-shuffle \ + so BoundedWindowAggExec's single-partition constraint is not a serial \ + bottleneck. Disabled by default — opt in when the workload contains \ + matching window shapes.".to_string(), + DataType::Boolean, + Some(false.to_string())), ConfigEntry::new( BALLISTA_COALESCE_TARGET_PARTITION_BYTES.to_string(), "Target post-coalesce partition size in bytes. Mirrors Spark's \ @@ -706,6 +719,11 @@ impl BallistaConfig { self.get_bool_setting(BALLISTA_PROPAGATE_EMPTY_ENABLED) } + /// Returns whether the AQE parallel-window rule is enabled. + pub fn parallel_window_enabled(&self) -> bool { + self.get_bool_setting(BALLISTA_PARALLEL_WINDOW_ENABLED) + } + /// Returns compression codec that will be used during write stage of shuffle pub fn shuffle_compression_codec( &self, diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index eff45c6f17..ea70731ef8 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -874,6 +874,16 @@ pub fn repartition_routing_expr( /// confirms the file is empty (`Some(0)`). If the file has rows or the /// row count is unknown (`None`), silently skipping would lose data — /// error out instead. +/// +/// TODO(halo-aware routing): when the downstream stage has a +/// `RangeFilterExec` with non-zero halo (bounded RANGE-frame windows), each +/// partition's *effective* read range is `[cuts[k-1] - halo_lo, cuts[k] + +/// halo_hi)`. This function currently uses the raw cut range, so files +/// straddling the halo boundary aren't routed to their halo-widened +/// consumer. Boundary rows near cuts can be missing from downstream +/// window sums — a correctness gap for RANGE frames that this refactor +/// does not resolve. Fix requires reaching across stages to read the +/// consumer RFE's halos and widening the overlap check here. pub fn cut_partitions( original_partitions: Vec>, reports: &[TaskRuntimeStats], @@ -1755,7 +1765,7 @@ mod overlap_remap_tests { /// A straddling sub-part — one whose sketched [min, max] spans the cut /// — appears in BOTH downstream partitions' lists. This is the case - /// PerPartitionFilterExec exists to clean up. + /// RangeFilterExec exists to clean up. #[test] fn overlap_remap_straddling_producer_appears_in_both_partitions() { // Producer 300 covers [5, 25) — straddles the cut at 15. diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index fcfd52514c..ea4dbffe17 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -23,7 +23,9 @@ use ballista_core::client_pool::BallistaClientPool; use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec; -use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec}; +use ballista_core::execution_plans::{ + RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec, +}; use ballista_core::serde::protobuf::ShuffleWritePartition; use ballista_core::serde::scheduler::PartitionStats; use ballista_core::{JobId, utils}; @@ -150,6 +152,17 @@ impl ExecutionEngine for DefaultExecutionEngine { reader.with_work_dir(work_dir.to_string()), ))), } + } else if let Some(reader) = p.downcast_ref::() { + match &self.client_pool { + Some(client_pool) => Ok(Transformed::yes(Arc::new( + reader + .with_work_dir(work_dir.to_string()) + .with_client_pool(client_pool.clone()), + ))), + None => Ok(Transformed::yes(Arc::new( + reader.with_work_dir(work_dir.to_string()), + ))), + } } else { // Scan restriction is scheduler-side (see // ballista/scheduler/src/state/task_builder.rs). The plan diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index 43de627535..46a3c4bd4f 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -24,7 +24,7 @@ use crate::state::execution_graph::{ use crate::state::task_manager::JobInfoCache; use ballista_core::config::BallistaConfig; use ballista_core::error::Result; -use ballista_core::execution_plans::ShuffleReaderExec; +use ballista_core::execution_plans::{RangeShuffleReaderExec, ShuffleReaderExec}; use ballista_core::serde::protobuf::{ AvailableVcores, ExecutorHeartbeat, JobStatus, job_status, }; @@ -369,15 +369,17 @@ pub trait JobState: Send + Sync { /// (e.g. `UnorderedRangeRepartitionExec`). Stops at leaves, multi-child /// operators (fan-in / joins), and stage boundaries. /// -/// The stage-boundary stop is currently a `ShuffleReaderExec` downcast — the -/// only kind of stage-boundary leaf that appears in a resolved stage plan. -/// The *general* rule is "stop at any stage boundary"; if new stage-boundary -/// operators appear, add them here (or, better, get `ExecutionPlan` upstream -/// to expose an `is_stage_boundary()` property so we don't keep -/// enumerating). +/// The stage-boundary stop enumerates the leaf readers that terminate a +/// resolved stage plan: `ShuffleReaderExec` (regular / broadcast / coalesced) +/// and `RangeShuffleReaderExec` (ordering-preserving). The *general* rule is +/// "stop at any stage boundary"; if new stage-boundary operators appear, add +/// them here (or, better, get `ExecutionPlan` upstream to expose an +/// `is_stage_boundary()` property so we don't keep enumerating). fn stage_has_input_collapse(plan_root: &Arc) -> bool { fn walk(node: &Arc) -> bool { - if node.downcast_ref::().is_some() { + if node.downcast_ref::().is_some() + || node.downcast_ref::().is_some() + { return false; } if node.properties().output_partitioning().partition_count() == 1 { @@ -665,6 +667,7 @@ mod test { ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, }; + use crate::cluster::stage_has_input_collapse; use crate::cluster::{BoundTask, bind_task_bias, bind_task_round_robin}; use crate::state::execution_graph::{ExecutionGraph, StaticExecutionGraph}; use crate::state::task_manager::JobInfoCache; @@ -891,4 +894,37 @@ mod test { }, ] } + + /// Both shuffle reader kinds are stage boundaries — walking through them + /// to detect an input collapse would mis-classify the *next* stage's leaf + /// as this stage's collapse. `stage_has_input_collapse` must return false + /// as soon as a reader is seen. Guard the range variant explicitly since + /// `UnknownPartitioning(1)` would otherwise trigger the single-partition + /// arm. + #[test] + fn stage_has_input_collapse_stops_at_range_reader() { + use ballista_core::execution_plans::RangeShuffleReaderExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("v", 0))); + let merge_ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + // Single output partition — the case where the `partition_count == 1` + // arm would fire without the reader guard. + let reader = Arc::new( + RangeShuffleReaderExec::try_new(1, vec![vec![]], schema, merge_ordering) + .unwrap(), + ) as Arc; + let root: Arc = Arc::new(CoalescePartitionsExec::new(reader)); + + assert!( + !stage_has_input_collapse(&root), + "a range-shuffle reader is a stage boundary, not an input collapse", + ); + } } diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 0b3af8fc00..765b181fee 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -27,8 +27,8 @@ use ballista_core::execution_plans::ShuffleWriter; use ballista_core::execution_plans::sort_shuffle::SortShuffleConfig; use ballista_core::{ execution_plans::{ - ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec, - UnresolvedShuffleExec, + RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec, + SortShuffleWriterExec, UnresolvedShuffleExec, }, serde::scheduler::PartitionLocation, }; @@ -785,6 +785,11 @@ pub fn remove_unresolved_shuffles( /// Rollback the ShuffleReaderExec to UnresolvedShuffleExec. /// Used when the input stages are finished but some partitions are missing due to executor lost. /// The entire stage need to be rolled back and rescheduled. +/// +/// `RangeShuffleReaderExec` rolls back to a plain `UnresolvedShuffleExec` — its +/// range-ness is a derived property of the child's declared ordering at plan +/// time, not intrinsic reader metadata. Re-planning walks the adapter, which +/// re-detects the ordering and plants a fresh `RangeShuffleReaderExec`. pub fn rollback_resolved_shuffles( stage: Arc, ) -> Result> { @@ -806,6 +811,14 @@ pub fn rollback_resolved_shuffles( )) }; new_children.push(unresolved); + } else if let Some(range_reader) = child.downcast_ref::() + { + let unresolved = Arc::new(UnresolvedShuffleExec::new( + range_reader.stage_id, + range_reader.schema(), + range_reader.properties().partitioning.clone(), + )); + new_children.push(unresolved); } else { new_children.push(rollback_resolved_shuffles(child.clone())?); } @@ -2041,6 +2054,51 @@ order by Ok(()) } + /// `RangeShuffleReaderExec` rolls back to a plain `UnresolvedShuffleExec` + /// (info-losing on the range-ness). Re-planning walks the adapter, which + /// re-detects the child's ordering and plants a fresh range reader. + #[tokio::test] + async fn rollback_resolved_shuffles_reduces_range_reader_to_plain_unresolved() + -> Result<(), BallistaError> { + use ballista_core::execution_plans::RangeShuffleReaderExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("v", 0))); + let merge_ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let reader = Arc::new( + RangeShuffleReaderExec::try_new( + 7, + vec![vec![]; 4], + schema.clone(), + merge_ordering, + ) + .unwrap(), + ) as Arc; + let parent: Arc = Arc::new( + datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec::new( + reader, + ), + ); + + let rolled_back = crate::planner::rollback_resolved_shuffles(parent)?; + let child = rolled_back.children()[0].clone(); + let unresolved = child + .downcast_ref::() + .expect("expected rolled-back UnresolvedShuffleExec"); + // The range-ness is derived at plan time; the rolled-back node carries + // no ordering, no broadcast, no coalesce. + assert!(!unresolved.broadcast); + assert!(unresolved.coalesce.is_none()); + assert_eq!(unresolved.stage_id, 7); + assert_eq!(unresolved.output_partition_count, 4); + + Ok(()) + } + #[tokio::test] async fn distributed_window_plan() -> Result<(), BallistaError> { let ctx = datafusion_test_context("testdata").await?; diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index b6ec54866b..4b8ab2deef 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -21,14 +21,16 @@ use crate::state::aqe::planner::AdaptiveStageInfo; use crate::state::execution_graph::StageOutput; use ballista_core::JobId; use ballista_core::execution_plans::{ - PerPartitionFilterExec, ShuffleReaderExec, range_partition_predicates, + RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, }; use datafusion::common::exec_err; use datafusion::config::ConfigOptions; use datafusion::error::DataFusionError; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning}; +use datafusion::scalar::ScalarValue; use datafusion::{ - common::tree_node::{Transformed, TreeNode}, + common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, physical_plan::ExecutionPlan, }; use log::debug; @@ -71,66 +73,81 @@ impl BallistaAdapter { self.inputs.insert(stage_id, stage_output); let partitioning = exchange.properties().partitioning.clone(); - let reader = match (exchange.coalesce(), exchange.broadcast) { - (Some(cp), false) => { - // Concatenate M-shape locations into K-shape per CoalescePlan.groups. - let k_shape: Vec> = cp - .groups - .iter() - .map(|pg| { - let mut concat = Vec::new(); - for &idx in &pg.upstream_indices { - if let Some(inner) = partitions.get(idx as usize) { - concat.extend_from_slice(inner); + let reader: Arc = + match (exchange.coalesce(), exchange.broadcast) { + (Some(cp), false) => { + // Concatenate M-shape locations into K-shape per CoalescePlan.groups. + let k_shape: Vec> = cp + .groups + .iter() + .map(|pg| { + let mut concat = Vec::new(); + for &idx in &pg.upstream_indices { + if let Some(inner) = partitions.get(idx as usize) { + concat.extend_from_slice(inner); + } } + concat + }) + .collect(); + let new_partitioning = match &partitioning { + Partitioning::Hash(keys, _m) => { + Partitioning::Hash(keys.clone(), cp.groups.len()) } - concat - }) - .collect(); - let new_partitioning = match &partitioning { - Partitioning::Hash(keys, _m) => { - Partitioning::Hash(keys.clone(), cp.groups.len()) + _ => Partitioning::UnknownPartitioning(cp.groups.len()), + }; + Arc::new(ShuffleReaderExec::try_new_coalesced( + stage_id, + k_shape, + (*cp).clone(), + schema, + new_partitioning, + )?) + } + (None, false) => { + // Ordered-writer path: when the child declared an output + // ordering, preserve it across the shuffle boundary with a + // k-way merge instead of the arrival-order concat that the + // regular reader does. Fixes the silent RANGE-frame + // corruption in `docs/developer/parallel-range-window.md`. + if let Some(ordering) = exchange.input().output_ordering() { + Arc::new(RangeShuffleReaderExec::try_new( + stage_id, + partitions, + schema, + ordering.clone(), + )?) + } else { + Arc::new(ShuffleReaderExec::try_new( + stage_id, + partitions, + schema, + partitioning, + )?) } - _ => Partitioning::UnknownPartitioning(cp.groups.len()), - }; - ShuffleReaderExec::try_new_coalesced( + } + (_, true) => Arc::new(ShuffleReaderExec::try_new_broadcast( stage_id, - k_shape, - (*cp).clone(), + exchange.shuffle_partitions_flattened(), schema, - new_partitioning, - )? - } - (None, false) => ShuffleReaderExec::try_new( - stage_id, - partitions, - schema, - partitioning, - )?, - (_, true) => ShuffleReaderExec::try_new_broadcast( - stage_id, - exchange.shuffle_partitions_flattened(), - schema, - exchange.input().output_partitioning().partition_count(), - )?, - }; - - let reader: Arc = Arc::new(reader); - // Without a per-partition filter, straddling sub-parts from a - // range-repartitioned upstream would feed multiple downstream - // partitions and `FinalPartitioned` would split their partial sums. - if let Some(routing) = exchange.range_repartition_routing() { - let predicates = - range_partition_predicates(routing.routing_expr, &routing.cuts); + exchange.input().output_partitioning().partition_count(), + )?), + }; + // The adapter no longer injects a `RangeFilterExec` above the + // reader. Rules that emit a range-repartition upstream (today: + // `ParallelWindowRule`) are also responsible for planting the + // read-side `RangeFilterExec`(s) at plan time with `cuts=None`; + // the scheduler-side `resolve_range_filter_cuts` walker fills + // in cuts once stage-0's sketches merge. If a range-routing + // ExchangeExec reaches this point without a rule-planted + // filter, `RangeFilterExec::execute()` fails loud rather than + // straddling sub-parts silently corrupting downstream partial + // aggregates. + if exchange.range_repartition_routing().is_some() { debug!( - "range-repartition: injecting PerPartitionFilterExec above \ - ShuffleReader for stage {} — {} predicates over {} cuts", - stage_id, - predicates.len(), - routing.cuts.len(), + "range-repartition: ExchangeExec has resolved routing for stage {stage_id}; \ + any rule-planted RangeFilterExec above should have been resolved by now" ); - let filtered = PerPartitionFilterExec::try_new(reader, predicates)?; - return Ok(Transformed::yes(Arc::new(filtered))); } Ok(Transformed::yes(reader)) } else { @@ -148,6 +165,7 @@ impl BallistaAdapter { ) -> datafusion::error::Result { if let Some(root) = plan.downcast_ref::() { let mut adapter = BallistaAdapter::default(); + resolve_range_filter_cuts(root.input())?; let plan = root .input() .clone() @@ -175,6 +193,7 @@ impl BallistaAdapter { }) } else if let Some(root) = plan.downcast_ref::() { let mut adapter = BallistaAdapter::default(); + resolve_range_filter_cuts(root.input())?; let plan = root .input() .clone() @@ -201,3 +220,148 @@ impl BallistaAdapter { } } } + +/// Walk `plan` and resolve every pending [`RangeFilterExec`]'s bounds from +/// its downstream [`ExchangeExec`]'s stored routing. Called at +/// `adapt_to_ballista` time, once stage-0's sketches have merged into +/// cuts and been parked on the boundary `ExchangeExec` via +/// `set_repartition_routing`. +/// +/// Cross-referencing is by `routing_expr` equality — the rule plants both +/// wide (below SPM) and narrow (above BWAG) filters with the same routing +/// expression, both pointing at the same shuffle boundary. +/// +/// RFE receives *unwidened* half-open ranges (`(cuts[k-1], cuts[k])` with ±∞ +/// sentinels at ends). RFE widens by its own halos internally at +/// `resolve_bounds` time — see the separation-of-concerns note on +/// [`RangeFilterExec`]. The scheduler stays halo-blind at this boundary. +/// +/// Errors if a `RangeFilterExec` is still pending after the walk: the +/// rule promised bounds and the scheduler didn't deliver, which is either +/// a routing_expr mismatch or a stage-progress bug. +fn resolve_range_filter_cuts( + plan: &Arc, +) -> Result<(), DataFusionError> { + let mut routings: Vec<(Arc, Vec)> = Vec::new(); + plan.apply(|node| { + if let Some(exchange) = node.downcast_ref::() + && let Some(routing) = exchange.range_repartition_routing() + { + routings.push((routing.routing_expr, routing.cuts)); + } + Ok(TreeNodeRecursion::Continue) + })?; + plan.apply(|node| { + let Some(rf) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + if rf.raw_bounds().is_some() { + return Ok(TreeNodeRecursion::Continue); + } + let rf_expr = rf.routing_expr(); + let cuts = routings + .iter() + .find(|(expr, _)| expr.eq(rf_expr)) + .map(|(_, cuts)| cuts.clone()) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "RangeFilterExec: no matching ExchangeExec routing for expr {rf_expr}" + )) + })?; + let raw_bounds = raw_bounds_from_cuts(&cuts); + rf.resolve_bounds(raw_bounds)?; + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) +} + +/// Project K-1 cuts to K half-open `(cuts[k-1], cuts[k])` ranges with `None` +/// sentinels at ±∞. This is the pure range-partitioning projection — no halo +/// arithmetic here (RFE widens internally at resolve time). +fn raw_bounds_from_cuts(cuts: &[f64]) -> Vec<(Option, Option)> { + let k = cuts.len() + 1; + (0..k) + .map(|i| { + let lo = i + .checked_sub(1) + .and_then(|j| cuts.get(j).copied()) + .map(|v| ScalarValue::Float64(Some(v))); + let hi = cuts.get(i).copied().map(|v| ScalarValue::Float64(Some(v))); + (lo, hi) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use ballista_core::execution_plans::RangeShuffleReaderExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::sorts::sort::SortExec; + + fn f64_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) + } + + fn asc(schema: &Arc, col: &str) -> PhysicalSortExpr { + let column = Column::new_with_schema(col, schema).unwrap(); + PhysicalSortExpr::new_default(Arc::new(column)) + } + + /// When the exchange's child declares an output ordering, the adapter + /// must plant `RangeShuffleReaderExec` so the shuffle boundary preserves + /// sortedness via k-way merge. + #[test] + fn plants_range_reader_when_child_declares_ordering() { + let schema = f64_schema(); + let empty: Vec> = vec![vec![]]; + let source = + MemorySourceConfig::try_new_exec(&empty, schema.clone(), None).unwrap(); + let sort_lex = LexOrdering::new(vec![asc(&schema, "v")]).unwrap(); + let sorted = + Arc::new(SortExec::new(sort_lex, source).with_preserve_partitioning(true)) + as Arc; + + let exchange = ExchangeExec::new(sorted, None, 0); + exchange.set_stage_id(1); + exchange.resolve_shuffle_partitions(vec![vec![]]); + let plan: Arc = Arc::new(exchange); + + let mut adapter = BallistaAdapter::default(); + let out = adapter.transform_children(plan).unwrap().data; + + assert!( + out.downcast_ref::().is_some(), + "expected RangeShuffleReaderExec but got {}", + out.name() + ); + } + + /// The unordered path must still plant the regular `ShuffleReaderExec`. + #[test] + fn plants_regular_reader_when_no_ordering() { + let schema = f64_schema(); + let empty: Vec> = vec![vec![]]; + let source = + MemorySourceConfig::try_new_exec(&empty, schema.clone(), None).unwrap(); + + let exchange = ExchangeExec::new(source, None, 0); + exchange.set_stage_id(1); + exchange.resolve_shuffle_partitions(vec![vec![]]); + let plan: Arc = Arc::new(exchange); + + let mut adapter = BallistaAdapter::default(); + let out = adapter.transform_children(plan).unwrap().data; + + assert!( + out.downcast_ref::().is_some(), + "expected ShuffleReaderExec but got {}", + out.name() + ); + assert!(out.downcast_ref::().is_none()); + } +} diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs index 80e325566f..c3a7b9962b 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs @@ -17,7 +17,8 @@ use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use ballista_core::execution_plans::{ - OrderedRangeRepartitionExec, UnorderedRangeRepartitionExec, preserves_partitioning, + OrderedRangeRepartitionExec, RangeFilterExec, UnorderedRangeRepartitionExec, + preserves_partitioning, }; use datafusion::common::plan_err; use datafusion::common::tree_node::{Transformed, TreeNode}; @@ -104,7 +105,7 @@ impl DistributedExchangeRule { execution_plan.downcast_ref::() { let input = sort_preserving_merge.input(); - if input.downcast_ref::().is_none() + if !is_stage_boundary(input) && !matches!(nearest_exchange_status(input), ExchangeStatus::Unresolved) { let exchange_exec = ExchangeExec::new( @@ -261,6 +262,22 @@ impl PhysicalOptimizerRule for DistributedExchangeRule { } } +/// True when `node` is (or transparently sits on) a stage boundary. +/// `RangeFilterExec` counts because we chose not to fold range-filtering +/// into `ShuffleReader`/`ExchangeExec` — the operator is part of the +/// boundary shape by design. +fn is_stage_boundary(node: &Arc) -> bool { + if node.is::() { + return true; + } + if node.is::() + && let [child] = node.children().as_slice() + { + return child.is::(); + } + false +} + /// Scans the subtree for the nearest `ExchangeExec` in each path and returns the /// aggregate status. Stops recursing at `ExchangeExec` boundaries so that only the /// shallowest exchange in each branch is considered. @@ -534,6 +551,52 @@ mod tests { ); } + #[test] + fn spm_skips_when_range_filter_covers_exchange() { + // ParallelWindowRule plants a RangeFilterExec directly on the + // resolved range-repartition ExchangeExec with SPM above. The + // filter must count as part of the boundary — otherwise DE + // inserts another ExchangeExec between SPM and the filter, + // collapsing K partitions into a single outer-stage task. + use ballista_core::execution_plans::RangeFilterExec; + use datafusion::scalar::ScalarValue; + + let rule = DistributedExchangeRule::default(); + let exchange = resolved_exchange(float_leaf_exec()); + let filter: Arc = Arc::new( + RangeFilterExec::try_new_pending( + exchange, + Arc::new(Column::new("v", 0)), + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + ) + .unwrap(), + ); + let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("v", 0))); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let spm: Arc = + Arc::new(SortPreservingMergeExec::new(ordering, filter)); + + let result = rule.optimize(spm, &config()).unwrap(); + + let adaptive = result.downcast_ref::().unwrap(); + let spm_out = adaptive + .input() + .downcast_ref::() + .unwrap(); + let below_spm = spm_out.children()[0]; + assert!( + below_spm.downcast_ref::().is_some(), + "SPM's direct child should remain RangeFilterExec, not a new ExchangeExec" + ); + assert!( + below_spm.children()[0] + .downcast_ref::() + .is_some(), + "resolved ExchangeExec should remain under the RangeFilterExec" + ); + } + // --- RepartitionExec --- #[test] @@ -830,7 +893,7 @@ mod tests { /// range-repartition-inserting rule emits, with nothing above it — /// must still get an `ExchangeExec` wrapped above it. Without it, /// `set_repartition_routing` has no parking slot for the recovered - /// cuts and downstream never gets a `PerPartitionFilterExec` to + /// cuts and downstream never gets a `RangeFilterExec` to /// trim straddler duplication. #[test] fn range_repartition_at_plan_root_gets_exchange_inserted() { @@ -920,7 +983,7 @@ mod tests { /// A `ProjectionExec` between (O/U)RRE and the boundary could /// reindex, drop, or shadow the routing expression's referenced - /// columns — the read-side `PerPartitionFilterExec` would evaluate + /// columns — the read-side `RangeFilterExec` would evaluate /// against the wrong column and silently misroute. DER rejects the /// shape at plan time; the fix will be revisited when arbitrary /// routing expressions replace the current single-column form. diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs index b4ae1fbc83..392c4bc373 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs @@ -19,9 +19,11 @@ pub mod chaos_exec; pub mod coalesce_partitions; pub mod distributed_exchange; pub mod join_selection; +pub mod parallel_window; pub mod propagate_empty; pub use coalesce_partitions::*; pub use distributed_exchange::*; pub use join_selection::*; +pub use parallel_window::*; pub use propagate_empty::*; diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs new file mode 100644 index 0000000000..282eeb8046 --- /dev/null +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs @@ -0,0 +1,573 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Rewrite bounded-RANGE-frame windows into a distributed range-shuffle so +//! `BoundedWindowAggExec`'s single-partition constraint isn't a serial +//! bottleneck. See [[parallel-range-window]] for the design. +//! +//! # Matched shape +//! +//! ```text +//! BoundedWindowAggExec [Sorted, RANGE frame, finite bounds] +//! SortPreservingMergeExec [ORDER BY] +//! SortExec (preserve_partitioning=true) [ORDER BY] +//! +//! ``` +//! +//! Restricted to: +//! - single window expression +//! - no PARTITION BY +//! - single-column ORDER BY on a physical `Column` (widening: multi-key, +//! computed exprs — separate rewrites) +//! - `RANGE` frame with finite `PRECEDING` / `FOLLOWING` / `CurrentRow` +//! bounds (UNBOUNDED goes down the prefix-scan path — see +//! [[project-prefix-scan-two-pass-rejected]]) +//! - ORDER BY column is `Float64` today (T-Digest restriction; lifts with +//! [[kll-sketch]]) +//! +//! # Rewrite +//! +//! ```text +//! RangeFilterExec (narrow, halo_lo=0, halo_hi=0, cuts=pending) +//! PartitionedBoundedWindowAggExec (wraps BWAG; declares UnspecifiedDistribution) +//! RangeFilterExec (wide, halo_lo, halo_hi, cuts=pending) +//! RuntimeStatsExec #2 (per-ORRE-output-partition sketch → scheduler) +//! OrderedRangeRepartitionExec (K outputs, walks child for RSE #1) +//! SortExec (planted here, preserve_partitioning=true) +//! RuntimeStatsExec #1 (local sketch — feeds ORRE's cut walker) +//! +//! ``` +//! +//! RSE#1 sits *below* SortExec so the local sketch ingests the whole +//! partition while Sort buffers, giving the scheduler full-fidelity cuts +//! to hand ORRE before it starts routing. RSE#1 above Sort would force +//! ORRE to route against a still-being-built sketch → skewed shuffle files. +//! +//! The rule runs *after* DF's optimizer chain so `EnforceSorting` / +//! `RepartitionFileScans` have already materialized the SortExec placement +//! we peel here. Running earlier hits two failure modes: (a) sources with +//! `sort_order_for_reorder` set have no SortExec yet at all, and (b) DF's +//! sort-pushdown later moves any Sort we plant down through the +//! passthrough RSE#1, undoing the intended order. +//! +//! Any SPM the DF planner inserted above BWAG for its `SinglePartition` +//! requirement is dropped: the wrapper flips that declaration to +//! `UnspecifiedDistribution`, and `EnforceDistribution` doesn't re-add one. +//! +//! Both `RangeFilterExec` operators are planted with `cuts=None`. The +//! scheduler-side `resolve_range_filter_cuts` walker fills them in once +//! stage-0's `RuntimeStatsExec` reports have been merged into cuts. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use ballista_core::execution_plans::{ + OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec, RangeFilterExec, + RuntimeStatsExec, +}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::logical_expr::{WindowFrameBound, WindowFrameUnits}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::windows::BoundedWindowAggExec; +use datafusion::scalar::ScalarValue; +use log::debug; + +/// Physical optimizer pass: match the parallel-window shape and rewrite each +/// hit to insert `RSE#1 → ORRE → RSE#2 → RangeFilterExec_wide` below the +/// existing SPM. +#[derive(Default, Debug)] +pub struct ParallelWindowRule; + +impl PhysicalOptimizerRule for ParallelWindowRule { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> datafusion::common::Result> { + let bc = config + .extensions + .get::() + .cloned() + .unwrap_or_default(); + if !bc.parallel_window_enabled() { + return Ok(plan); + } + // K = the number of range-disjoint output partitions the ORRE will + // produce. At rule-fire time DataFusion's initial physical plan is + // still "loose" — the DataSourceExec below BWAG has 1 file_group, + // not the eventual `target_partitions` split (RepartitionFileScans + // and friends run later in the AQE chain). So we can't ask the + // plan tree for the true source width; we use the config knob that + // those later rules also target. + let output_partitions = config.execution.target_partitions.max(2); + plan.transform_up(|node| { + let Some(candidate) = as_candidate(node.as_ref()) else { + return Ok(Transformed::no(node)); + }; + match rewrite_bwag(&node, &candidate, output_partitions) { + Ok(rewritten) => { + debug!( + "ParallelWindowRule: rewrote BWAG on `{}` (RANGE {} — {})", + candidate.order_key, + fmt_bound(&candidate.start_bound), + fmt_bound(&candidate.end_bound), + ); + Ok(Transformed::yes(rewritten)) + } + Err(e) => { + debug!( + "ParallelWindowRule: shape matched but rewrite skipped for `{}`: {e}", + candidate.order_key, + ); + Ok(Transformed::no(node)) + } + } + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "ParallelWindow" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Shape captured from a matching `BoundedWindowAggExec`. Everything the +/// rewrite needs to build the new subtree. +#[derive(Debug, Clone)] +struct WindowCandidate { + order_key: String, + sort_expr: PhysicalSortExpr, + start_bound: WindowFrameBound, + end_bound: WindowFrameBound, +} + +/// True if any descendant of `nodes` is an `OrderedRangeRepartitionExec` +/// or `RangeFilterExec`. Used as an idempotency guard: those ops are what +/// our own rewrite plants, so seeing them below a BWAG means we've already +/// rewritten this candidate on a previous optimizer pass. +fn subtree_contains_our_rewrite(nodes: &[&Arc]) -> bool { + for node in nodes { + if node.is::() || node.is::() { + return true; + } + if subtree_contains_our_rewrite(node.children().as_slice()) { + return true; + } + } + false +} + +fn as_candidate(node: &dyn ExecutionPlan) -> Option { + let window = node.downcast_ref::()?; + // Shape gates as slice patterns: 0 or 2+ elements simply don't match. + let [expr] = window.window_expr() else { + return None; + }; + let [] = expr.partition_by() else { + return None; + }; + let [order] = expr.order_by() else { + return None; + }; + let column = order.expr.downcast_ref::()?; + let frame = expr.get_window_frame(); + let WindowFrameUnits::Range = frame.units else { + return None; + }; + let (Some(start), Some(end)) = + (as_finite(&frame.start_bound), as_finite(&frame.end_bound)) + else { + return None; + }; + // Idempotency: if the BWAG's subtree already contains our own + // range-repartition machinery, we've already rewritten this window. + // Re-plans (AQE fires the optimizer chain again for stage N+1) would + // otherwise wrap another ORRE around the previous rewrite's + // RangeFilterExec+ShuffleReader — and that ORRE's child doesn't claim + // ordering, blowing up at execute-time. + if subtree_contains_our_rewrite(window.children().as_slice()) { + return None; + } + Some(WindowCandidate { + order_key: column.name().to_string(), + sort_expr: order.clone(), + start_bound: start.clone(), + end_bound: end.clone(), + }) +} + +/// Returns the bound unchanged when it's `CurrentRow` or a non-null scalar +/// offset. `UNBOUNDED PRECEDING/FOLLOWING` is represented as a typed-null +/// scalar and returns `None`. +fn as_finite(bound: &WindowFrameBound) -> Option<&WindowFrameBound> { + match bound { + WindowFrameBound::CurrentRow => Some(bound), + WindowFrameBound::Preceding(scalar) | WindowFrameBound::Following(scalar) + if !scalar.is_null() => + { + Some(bound) + } + _ => None, + } +} + +fn fmt_bound(bound: &WindowFrameBound) -> String { + match bound { + WindowFrameBound::CurrentRow => "CURRENT ROW".to_string(), + WindowFrameBound::Preceding(scalar) => format!("{scalar} PRECEDING"), + WindowFrameBound::Following(scalar) => format!("{scalar} FOLLOWING"), + } +} + +/// Extract the halo width in `f64` from a bound. `CurrentRow` → 0. +/// Errors on non-numeric scalar (e.g. Interval bounds — future work). +fn halo_from_bound(bound: &WindowFrameBound) -> datafusion::common::Result { + let scalar = match bound { + WindowFrameBound::CurrentRow => return Ok(0.0), + WindowFrameBound::Preceding(s) | WindowFrameBound::Following(s) => s, + }; + // Widen anything Int-ish or Float-ish to f64. Interval bounds (for + // time-typed ORDER BYs) are the widening TODO alongside KLL. + match scalar { + ScalarValue::Int8(Some(v)) => Ok(*v as f64), + ScalarValue::Int16(Some(v)) => Ok(*v as f64), + ScalarValue::Int32(Some(v)) => Ok(*v as f64), + ScalarValue::Int64(Some(v)) => Ok(*v as f64), + ScalarValue::UInt8(Some(v)) => Ok(*v as f64), + ScalarValue::UInt16(Some(v)) => Ok(*v as f64), + ScalarValue::UInt32(Some(v)) => Ok(*v as f64), + ScalarValue::UInt64(Some(v)) => Ok(*v as f64), + ScalarValue::Float32(Some(v)) => Ok(*v as f64), + ScalarValue::Float64(Some(v)) => Ok(*v), + other => datafusion::common::internal_err!( + "ParallelWindowRule: unsupported halo bound type {other:?}" + ), + } +} + +/// Wrap BWAG in a `PartitionedBoundedWindowAggExec` and splice +/// `SortExec → RSE#1 → source` below a fresh `ORRE → RSE#2 → RFE_wide` +/// chain. The rule runs after DF's optimizer chain, so BWAG's descendants +/// have the fully-materialized `SPM → SortExec → source` shape here — we +/// strip both since our rewrite overrides BWAG's distribution and takes +/// ownership of Sort placement (see module doc for why). +fn rewrite_bwag( + bwag: &Arc, + candidate: &WindowCandidate, + output_partitions: usize, +) -> datafusion::common::Result> { + let bwag_children = bwag.children(); + let [immediate] = bwag_children.as_slice() else { + return datafusion::common::internal_err!( + "ParallelWindowRule: BWAG must have exactly 1 child" + ); + }; + + // Strip whatever DF planted above the true source purely to satisfy + // BWAG's SinglePartition + Sorted requirements (SPM and SortExec) — + // our rewrite overrides both. Loop tolerates any order (SPM→Sort or + // Sort→SPM) or partial shapes (source that claims ordering natively + // via `sort_order_for_reorder` skips the Sort entirely). + let mut base_source: Arc = (*immediate).clone(); + while base_source.is::() || base_source.is::() { + let children = base_source.children(); + let [inner] = children.as_slice() else { + return datafusion::common::internal_err!( + "ParallelWindowRule: SPM/SortExec must have exactly 1 child" + ); + }; + base_source = (*inner).clone(); + } + let source_schema = base_source.schema(); + + // Route on the ORDER BY column. ORRE requires Float64 today. + let routing_type = candidate.sort_expr.expr.data_type(&source_schema)?; + if !matches!(routing_type, DataType::Float64) { + return datafusion::common::internal_err!( + "ParallelWindowRule: routing expression `{}` must be Float64, got {routing_type:?}", + candidate.sort_expr.expr + ); + } + + let sort_expr = normalize_sort_expr(&candidate.sort_expr); + let rse1: Arc = Arc::new(RuntimeStatsExec::try_new( + base_source, + Some(vec![sort_expr.clone()]), + )?); + // Plant a fresh SortExec above RSE#1 as the pipeline break: Sort + // consumes all input before emitting the first row, so RSE#1's sketch + // fully ingests and reports while Sort buffers — ORRE then routes + // against final cuts instead of approximate ones (which would produce + // skewed shuffle files). + let sort_lex = LexOrdering::new(vec![sort_expr.clone()]).ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "ParallelWindowRule: could not build LexOrdering from ORDER BY".into(), + ) + })?; + let sorted_over_rse1: Arc = + Arc::new(SortExec::new(sort_lex, rse1).with_preserve_partitioning(true)); + let orre: Arc = Arc::new(OrderedRangeRepartitionExec::try_new( + sorted_over_rse1, + vec![sort_expr.clone()], + output_partitions, + )?); + let rse2: Arc = Arc::new(RuntimeStatsExec::try_new( + orre, + Some(vec![sort_expr.clone()]), + )?); + let halo_lo = halo_from_bound(&candidate.start_bound)?; + let halo_hi = halo_from_bound(&candidate.end_bound)?; + let wide_filter: Arc = Arc::new(RangeFilterExec::try_new_pending( + rse2, + sort_expr.expr.clone(), + ScalarValue::Float64(Some(halo_lo)), + ScalarValue::Float64(Some(halo_hi)), + )?); + + // Wrap BWAG in PartitionedBoundedWindowAggExec instead of collapsing + // K→1 with SPM. The wrapper declares `UnspecifiedDistribution` so + // EnforceDistribution won't reinsert an SPM below, and BWAG's own + // per-partition execute() runs each of the K sub-ranges independently. + // See execution_plans::partitioned_bounded_window_agg for what makes + // this safe (range-repartition upstream + halo). + let bwag_ref = bwag.downcast_ref::().ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "ParallelWindowRule: rewrite_bwag caller passed non-BWAG".into(), + ) + })?; + let partitioned_bwag: Arc = + Arc::new(PartitionedBoundedWindowAggExec::try_new( + bwag_ref.window_expr().to_vec(), + wide_filter, + )?); + // Narrow filter above BWAG drops the halo rows the wide filter let in + // for BWAG's frame-context. `halo_lo == halo_hi == 0.0` collapses the + // predicate to `cuts[k-1] <= v < cuts[k]` — task k's own range. + let narrow_filter: Arc = + Arc::new(RangeFilterExec::try_new_pending( + partitioned_bwag, + sort_expr.expr.clone(), + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + )?); + Ok(narrow_filter) +} + +/// ORRE requires `nulls_first == false` today (T-Digest has no NULL slot). +/// The BWAG's `NULLS LAST` sort expressions arrive with `nulls_first: false` +/// already, but explicit sanitization keeps the invariant obvious to future +/// readers. +fn normalize_sort_expr(expr: &PhysicalSortExpr) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: expr.expr.clone(), + options: SortOptions { + descending: expr.options.descending, + nulls_first: false, + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::config::ExtensionOptions; + use datafusion::datasource::empty::EmptyTable; + use datafusion::physical_plan::displayable; + use datafusion::prelude::SessionContext; + + async fn plan(sql: &str) -> datafusion::common::Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id1", DataType::Int64, false), + Field::new("id2", DataType::Int64, false), + Field::new("id3", DataType::Int64, false), + Field::new("v2", DataType::Float64, false), + ])); + let ctx = SessionContext::new(); + ctx.register_table("large", Arc::new(EmptyTable::new(schema)))?; + ctx.sql(sql).await?.create_physical_plan().await + } + + fn optimize( + plan: Arc, + ) -> datafusion::common::Result> { + let mut config = ConfigOptions::default(); + config.execution.target_partitions = 8; + let mut bc = BallistaConfig::default(); + bc.set("planner.parallel_window.enabled", "true").unwrap(); + config.extensions.insert(bc); + ParallelWindowRule.optimize(plan, &config) + } + + #[tokio::test] + async fn disabled_by_default() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + // No BallistaConfig extension registered → default is `false`. + let mut config = ConfigOptions::default(); + config.execution.target_partitions = 8; + let out = ParallelWindowRule.optimize(plan.clone(), &config)?; + let rendered = format!("{}", displayable(out.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "flag off: rewrite must not fire:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn rewrites_q8_shape() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + // The rewrite must plant each of these ops. Cheap string contains — + // exhaustive plan-shape assertions in follow-up integration tests. + for expected in [ + "PartitionedBoundedWindowAggExec", + "BoundedWindowAggExec", + "RangeFilterExec", + "RuntimeStatsExec", + "OrderedRangeRepartitionExec", + "SortExec", + ] { + assert!( + rendered.contains(expected), + "expected `{expected}` in rewritten plan:\n{rendered}" + ); + } + // BWAG's SinglePartition collapse is what this whole rewrite + // avoids — any SPM in the output would defeat that. + assert!( + !rendered.contains("SortPreservingMergeExec"), + "SortPreservingMergeExec must NOT appear in the rewritten plan:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_rows_frame() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT avg(v2) OVER (ORDER BY id3 \ + ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "ROWS frames should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_partition_by() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (PARTITION BY id1 ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "PARTITION BY should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_unbounded_frame() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "UNBOUNDED PRECEDING should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[tokio::test] + async fn no_rewrite_on_non_float64_order_key() -> datafusion::common::Result<()> { + // id3 is Int64; ORRE requires Float64 today (T-Digest restriction). + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY id3 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("OrderedRangeRepartitionExec"), + "non-Float64 order key should not be rewritten:\n{rendered}" + ); + Ok(()) + } + + #[test] + fn halo_from_bound_reads_all_numeric_variants() { + assert_eq!(halo_from_bound(&WindowFrameBound::CurrentRow).unwrap(), 0.0); + assert_eq!( + halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Int64(Some(3)))) + .unwrap(), + 3.0 + ); + assert_eq!( + halo_from_bound(&WindowFrameBound::Following(ScalarValue::Float64(Some( + 2.5 + )))) + .unwrap(), + 2.5 + ); + assert!( + halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Utf8(Some( + "x".into() + )))) + .is_err() + ); + } +} diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 6dc08dab6a..99b44a0057 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -22,7 +22,7 @@ use crate::state::aqe::execution_plan::{ use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule; use crate::state::aqe::optimizer_rule::{ CoalescePartitionsRule, DelayJoinSelectionRule, DistributedExchangeRule, - PropagateEmptyExecRule, SelectJoinRule, + ParallelWindowRule, PropagateEmptyExecRule, SelectJoinRule, }; use crate::state::distributed_explain::handle_explain_plan; use crate::state::execution_stage::StageOutput; @@ -555,6 +555,16 @@ impl AdaptivePlanner { // physical_optimizers.extend(Self::datafusion_optimizers()); + // Rewrite bounded RANGE-frame windows into a range-shuffle so BWAG's + // single-partition constraint is not a serial bottleneck. Runs AFTER + // DataFusion's optimizer chain (EnforceSorting, RepartitionFileScans, + // …) so we see the fully-materialized SortExec placement — placement + // we peel and re-plant so RSE#1 sits *below* the pipeline-break Sort, + // letting the sketch fully report before ORRE routes. Must still run + // before DistributedExchangeRule — the rule emits an ORRE that DE + // picks up as the shuffle-boundary K-space source. + physical_optimizers.push(Arc::new(ParallelWindowRule)); + // `DistributedExchangeRule` should be the last plan mutator rule in the chain physical_optimizers .push(Arc::new(DistributedExchangeRule::new(plan_id_generator))); diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index 68811762e0..fc5367e579 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -36,7 +36,9 @@ //! flows from parent to descendants via function arguments, so sibling //! subtrees never share state and there's no traversal-order dependency. -use ballista_core::execution_plans::{PerPartitionFilterExec, ShuffleReaderExec}; +use ballista_core::execution_plans::{ + RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, +}; use datafusion::common::internal_err; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{ @@ -84,27 +86,40 @@ fn restrict( return Ok(rewritten); } - // PerPartitionFilterExec: its `predicates` vec is positionally aligned - // with the child's output partitions (predicates[k] filters - // input.execute(k)). When we restrict the child from K partitions to - // `partitions.len()`, the predicate vec must be sliced by the same - // indices in the same order - if !under_collect && let Some(ppf) = plan.downcast_ref::() { + // RangeFilterExec: raw_bounds is indexed by input partition; restriction + // slices bounds parallel to the input's partition subset. Halos + routing + // are carried over verbatim; RFE re-widens on the fresh operator. + if !under_collect && let Some(rf) = plan.downcast_ref::() { let children = plan.children(); let [child] = children.as_slice() else { return internal_err!( - "PerPartitionFilterExec must have exactly 1 child, got {}", + "RangeFilterExec must have exactly 1 child, got {}", children.len() ); }; let new_child = restrict((*child).clone(), partitions, false)?; - let new_predicates: Vec<_> = partitions + let raw_bounds = rf.raw_bounds().ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "RangeFilterExec: task-restriction before resolve_bounds()".into(), + ) + })?; + let sliced_bounds: Vec<_> = partitions .iter() - .map(|&part_idx| ppf.predicates()[part_idx].clone()) - .collect(); - return Ok(Arc::new(PerPartitionFilterExec::try_new( + .map(|&global| { + raw_bounds.get(global).cloned().ok_or_else(|| { + datafusion::common::DataFusionError::Internal(format!( + "RangeFilterExec: partition index {global} out of bounds ({} raw bounds)", + raw_bounds.len() + )) + }) + }) + .collect::>()?; + return Ok(Arc::new(RangeFilterExec::try_new_resolved( new_child, - new_predicates, + rf.routing_expr().clone(), + rf.halo_lo().clone(), + rf.halo_hi().clone(), + sliced_bounds, )?)); } @@ -288,6 +303,24 @@ fn select_output_partitions( return Some(Arc::new(restricted)); } + // RangeShuffleReaderExec: cross-stage inputs, ordering-preserving. Same + // partition-slice restriction as ShuffleReaderExec — carry the merge + // ordering unchanged. + if let Some(reader) = plan.downcast_ref::() { + let kept: Vec> = indices + .iter() + .filter_map(|&p| reader.partition.get(p).cloned()) + .collect(); + let restricted = RangeShuffleReaderExec::try_new( + reader.stage_id, + kept, + reader.schema(), + reader.merge_ordering().clone(), + ) + .ok()?; + return Some(Arc::new(restricted)); + } + // DataSourceExec: file-backed or in-memory scans. if let Some(exec) = plan.downcast_ref::() { let source: &dyn Any = exec.data_source().as_ref(); @@ -650,23 +683,19 @@ mod tests { ); } - /// A `PerPartitionFilterExec` restricted to a subset of partitions must - /// slice its `predicates` vector by the same indices, in the same order, - /// as its child. Otherwise the operator's construction invariant - /// (`predicates.len() == child.partition_count()`) breaks and - /// task-local partition `j` would filter through a global predicate - /// that no longer matches. + /// A `RangeFilterExec` restricted to a subset of partitions must slice + /// its `raw_bounds` by the same indices in the same order as its child. + /// Halos + routing_expr carry over unchanged. #[test] - fn per_partition_filter_predicates_are_sliced_with_partitions() { - use ballista_core::execution_plans::PerPartitionFilterExec; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::physical_expr::{Partitioning, PhysicalExpr}; - use datafusion::scalar::ScalarValue; - - // 4 upstream partitions, each with its own bespoke predicate so we - // can assert the slice ordering survives. - let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + fn range_filter_bounds_are_sliced_with_partitions() { + use ballista_core::execution_plans::RangeFilterExec; + use datafusion::physical_expr::Partitioning; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::Column; + + // 4 upstream partitions with a global 4-way range partition. + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); let partitions_locs: Vec> = (0..4).map(|i| vec![create_partition(i)]).collect(); let reader = ShuffleReaderExec::try_new( @@ -676,41 +705,38 @@ mod tests { Partitioning::UnknownPartitioning(4), ) .unwrap(); - let make_pred = |lo: i64| -> Arc { - Arc::new(BinaryExpr::new( - Arc::new(Column::new("v", 0)), - Operator::GtEq, - Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))), - )) - }; - let predicates: Vec> = - (0..4).map(|i| make_pred(i as i64 * 100)).collect(); + use datafusion::scalar::ScalarValue; + let routing_expr: Arc = Arc::new(Column::new("v", 0)); + // K=4 raw bounds derived from cuts [100, 200, 300]. + let sv = |v: f64| ScalarValue::Float64(Some(v)); + let raw_bounds: Vec<(Option, Option)> = vec![ + (None, Some(sv(100.0))), + (Some(sv(100.0)), Some(sv(200.0))), + (Some(sv(200.0)), Some(sv(300.0))), + (Some(sv(300.0)), None), + ]; let plan: Arc = Arc::new( - PerPartitionFilterExec::try_new( + RangeFilterExec::try_new_resolved( Arc::new(reader) as Arc, - predicates.clone(), + routing_expr, + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + raw_bounds.clone(), ) .unwrap(), ); let restricted = restrict_plan_to_partitions(plan, &[1, 3]).unwrap(); - let ppf = restricted - .downcast_ref::() - .expect("top must remain PerPartitionFilterExec"); - assert_eq!(ppf.predicates().len(), 2); - assert_eq!( - ppf.predicates()[0].to_string(), - predicates[1].to_string(), - "local partition 0 must carry the global-partition-1 predicate" - ); - assert_eq!( - ppf.predicates()[1].to_string(), - predicates[3].to_string(), - "local partition 1 must carry the global-partition-3 predicate" - ); + let rf = restricted + .downcast_ref::() + .expect("top must remain RangeFilterExec"); + let restricted_bounds = rf.raw_bounds().unwrap(); + assert_eq!(restricted_bounds.len(), 2); + assert_eq!(restricted_bounds[0], raw_bounds[1]); + assert_eq!(restricted_bounds[1], raw_bounds[3]); // Reader below must have been restricted in the same order. - let child = ppf.children()[0].clone(); + let child = rf.children()[0].clone(); let reader = child .downcast_ref::() .expect("child must be a ShuffleReaderExec"); diff --git a/docs/developer/parallel-range-window.md b/docs/developer/parallel-range-window.md new file mode 100644 index 0000000000..09914fc39a --- /dev/null +++ b/docs/developer/parallel-range-window.md @@ -0,0 +1,76 @@ + + +# Parallel bounded-RANGE-frame windows — design doc + +## The formula + +1. Envision the end state. +2. Figure out where you actually are. +3. Plot a course. +4. Create a series of steps. +5. Each step's vector must have a dot product > 0.9 with the target direction. +6. Plan the nearby step very well. +7. Plan the farthest step barely at all. +8. Interpolate planning along the way. +9. Revisit at each step. +10. **Rope-bridge principle.** Fire an arrow with a string; pull twine; pull rope; pull larger rope; pull a floor. The bridge fulfills "bridge" from day one. Ship the _shape_, then thicken. Correct-and-slow-and-limited is a viable arrow; the way to a full bridge is not to design the floor first. + +Corollary: a step is a candidate for skipping-and-back-filling if its output correctness holds without it. "Necessary for the final impl" ≠ "necessary for this step." Filtering files by ValueIndex range is important in the end state; it doesn't stop earlier steps from being correct without it. + +## The end state (step 1) + +On a cluster with E executors × V vcores each, scan the input across E×V partitions in parallel. Stage N collects per-partition stats via a runtime sketch and drives U/ORRE with `output_cnt == input_cnt`, keeping every vcore busy inside the task (local exchange is cheap). If the shuffle is ordered, write ValueIndex files so downstream doesn't have to over-sample or over-fetch. Cuts flow to the scheduler and become global. + +Stage N+1 uses multi-partition tasks: each task claims `vcores` input partitions plus a file-halo overlap with the neighbouring task. Overlap is a _task-level_ concern, not a partition-level concern — inside a task, adjacent partitions borrow context via local memory (free); across tasks, file-halo is fetched over shuffle via ValueIndex-based partial reads. The task k-way merges its inputs into `vcores` sorted DF partitions, using the global cuts to distribute evenly. It then filters with row-halo, runs `PartitionedBoundedWindowAggExec`, filters without row-halo, and writes `vcores` output files. + +Invariants: + +- **Width invariant.** Partition count stays at E×V across stages. No funnel except where explicitly planted (not usually needed, done in ballista client). +- **Ordered-shuffle propagation.** When the writer declared an output ordering, the reader preserves it via k-way merge. Every ordered-shuffle consumer (BWAG, SortMergeJoin build side, …) benefits from the same primitive. +- **Two-level halo.** Stage-level file-halo crosses task boundaries at shuffle cost (partial-file read); task-level row-halo crosses local-partition boundaries at memory cost. + +## The plan + +Ordered chronologically. Ticked items are landed (or mostly landed). Unticked items may be skipped and back-filled per the rope-bridge principle. + +- [x] **Multi-partition-task substrate.** #2038: `partition_slice` on task launch, K-drain ShuffleWriter, executor-side partition restriction. Every parallel operator downstream sits on this. +- [x] **T-Digest / KLL runtime stats.** #2180: sketch per partition, wire report to scheduler, merge into cuts. Foundation for any data-driven range op. +- [x] **URRE / ORRE.** #2169, #2196: N sorted overlapping → K sorted disjoint (ORRE via k-way merge internally) or unordered variant. +- [x] **RuntimeStatsExec, cut-discovery walker.** Same PR family. Late-binding cuts flow scheduler → downstream ops. +- [x] **RangeFilterExec + PartitionedBoundedWindowAggExec.** #2223: filter by resolved cuts + halo; BWAG wrapper that hides from tree walkers so `EnforceDistribution` doesn't collapse K→1. +- [x] **ParallelWindowRule.** #2223: match the `BWAG on Column Float64 ORDER BY, RANGE PRECEDING/FOLLOWING` shape, rewrite to insert `RuntimeStats → ORRE → RangeFilter(wide) → PartitionedBWAG → RangeFilter(narrow)`. +- [x] **Feature flag.** `ballista.planner.parallel_window.enabled=false` by default. Off is inert; on activates the rule. +- [x] **Ordered ShuffleReader.** New `RangeShuffleReaderExec` — keeps each upstream source as its own stream, feeds N into `StreamingMergeBuilder` on the child's declared ordering. Adapter picks it whenever `exchange.input().output_ordering().is_some()` (writer-driven gate). No permit governor and no per-source buffering; backpressure flows from the merge's demand through h2 / disk. Reusable — SortMergeJoin build side wants the same thing. Verified: h2o Q8 @ 1e7 SUM diff between `parallel_window.enabled=true` and `=false` agrees to 5e-14 relative (Float64 noise floor). If a wasted-merge cost surfaces in profiles later, tighten to demand-driven via a new consumer-side `requires_globally_sorted_input` bit. +- [x] **RangeFilter min/max fast-path + binary-search slice.** Post-ordered-ShuffleReader, batches are internally sorted. `min/max` fast paths (100% pass → Arc-clone; 0% overlap → skip) collapse the hot filter cost without touching correctness. Binary-search + `RecordBatch::slice` covers the mixed case with zero data copy. `sorted_on_key` derived at construction from `input.output_ordering()` — ascending on `routing_expr`, no config knob. Nullable routing columns fall back to `filter_record_batch` on a per-batch basis. Verified h2o Q8, 2 execs × 4 vcores, MPT=4: at 1e7 (2G cap) 7.6 s → 2.5 s = 3.0×; at 1e8 (4G cap) 143 s → 92 s = 1.55×. The 1e8 delta is smaller because the bottleneck shifts to shuffle/merge memory — see the next two items. +- [ ] **ValueIndex-based partial-file reads at shuffle-fetch time.** #2204 landed the write-side + reader primitives; consumer plumbing to translate value-range → byte-range at fetch has to hook in. Lets stage-N+1 halo reads pull only the halo slice from a neighbour file, not the whole file. +- [ ] **Per-task halo metadata on task-status.** New axis on task shipping: `own_files + halo_slice(file_ref, value_range)`. Composes with the ValueIndex plumbing — the scheduler emits `PartitionLocation`-with-value-range instead of `PartitionLocation`-whole-file. This is the rope. Enables inter-task halo without over-fetch. +- [ ] **Intra-task ORRE.** Once per-task input is one sorted merged stream (post-ordered-ShuffleReader), split it into `vcores` sub-partitions by task-local sub-cuts (derived from global cuts + vcore count). Task-level halo at the sub-partition boundaries is intra-task (local memory, free). Now cores stay busy inside every task without inter-task shuffle. +- [ ] **Two-level halo semantics in the rule.** RangeFilter at the stage boundary uses stage-level halo width (crosses tasks, shuffles); RangeFilter at the task-local boundary uses task-level halo width (crosses local partitions, memory). The rule plants both. +- [ ] **Symmetric halo (PRECEDING + FOLLOWING).** Generalize the halo direction on the shape. Plumbing extension; no new operators. +- [ ] **Range-partition invariant across stages.** E×V vcores → E×V in-flight partitions at every stage, funnel-free. Where an actual funnel is required, plant it explicitly. Everywhere else, keep width invariant. + +Out of scope for this end state: **unbounded PRECEDING** (running SUM). Halo degenerates to "all preceding tasks" → serialism. Sibling design (prefix scan). + +## Dot-product check on the near steps + +- **RangeFilter fast paths.** Perf. Correctness-preserving. Doesn't unlock or block anything else, but pays for itself immediately on the wide filter (which sees 100% pass every batch given ORRE's exact-routing). +- **ValueIndex plumbing.** Skippable at the cost of over-fetching halo files whole. Everything after it correctness-holds without it. +- **Per-task halo metadata.** Skippable at the cost of stage-level halo remaining a single-width parameter on the shape. Everything after correctness-holds without it. +- **Intra-task ORRE.** Skippable at the cost of not saturating vcores when a task's input is ≤ vcores partitions. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c0c721cf09..215fb6fb11 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -129,6 +129,7 @@ standard DataFusion settings. | ballista.planner.coalesce.merged_partition_factor | Float64 | 1.2 | Two adjacent partitions are merged when their combined size is below target_partition_bytes times this factor. Mirrors Spark's legacy coalesce semantics. | | ballista.planner.coalesce.small_partition_factor | Float64 | 0.2 | A coalesced partition smaller than target_partition_bytes times this factor counts as small and is merged into its neighbour. Mirrors Spark's legacy coalesce semantics. | | ballista.planner.coalesce.target_partition_bytes | UInt64 | 67108864 | Target post-coalesce partition size in bytes. Mirrors Spark's advisoryPartitionSizeInBytes. | +| ballista.planner.parallel_window.enabled | Boolean | false | Enables the AQE parallel-window rule (ParallelWindowRule), which rewrites bounded-RANGE-frame windows into a distributed range-shuffle so BoundedWindowAggExec's single-partition constraint is not a serial bottleneck. Disabled by default — opt in when the workload contains matching window shapes. | | ballista.planner.propagate_empty.enabled | Boolean | true | Enables the AQE propagate-empty-relation rule. Injects EmptyExec into the plan where an input is known to be empty, such as one side of a join, allowing downstream work to be skipped. | | ballista.scheduler.max_partitions_per_task | UInt64 | 1 | Upper bound on the number of input partitions packed into a single task's `partition_slice`. `1` (default) means one task per input partition. Raise to enable multi-partition tasks (fewer tasks, parallel-sort / parallel-join wins); `0` means unbounded — the scheduler fills each task up to the executor's free vcore count. Does not apply to collapse stages, which must pack their full pending queue into a single task for correctness. | | ballista.standalone.parallelism | UInt16 | number of available CPU cores | Number of concurrent tasks a standalone in-process executor will run. | From 90ae6db4b5aa363661deac2c1c293b2b3ab8bf6b Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 07:18:13 -0600 Subject: [PATCH 02/11] ci(h2o): enable parallel_window rule so Q8 exercises the parallel path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The h2o job was still running Q8 through the baseline serial pipeline — `ballista.planner.parallel_window.enabled` defaults to `false`, so the whole rule was inert here. Flipping it on lets `--verify` diff the parallel path against the local DataFusion oracle in CI, which is what we actually want before undrafting #2223. Other queries in the suite either don't match the rule's shape gates (no PARTITION BY + single-column Float64 ORDER BY + finite RANGE frame) and pass through untouched, or they do match and now get extra coverage. --- .github/workflows/h2o.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/h2o.yml b/.github/workflows/h2o.yml index 997935510e..4cc17f5506 100644 --- a/.github/workflows/h2o.yml +++ b/.github/workflows/h2o.yml @@ -171,6 +171,7 @@ jobs: --partitions 4 \ --verify \ -c ballista.planner.adaptive.enabled=true \ + -c ballista.planner.parallel_window.enabled=true \ -c ballista.scheduler.max_partitions_per_task=0 echo "::endgroup::" done From 8bf438ca230fc1b75d1475c18705c5676ef40646 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 09:34:36 -0600 Subject: [PATCH 03/11] WIP fix(scheduler): halo-aware range-repartition routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Missing rows in Q8 window sums — `cut_partitions` routed files by raw cut ranges, but the downstream `RangeFilterExec` widens each partition by `[halo_lo, halo_hi]`. A producer file whose sketch fell entirely on one side of a cut but within halo width of it was never delivered to the neighbouring partition, so RANGE-frame window sums lost those halo rows. - `cut_partitions` gains `(halo_lo, halo_hi)` f64 args; the b_lo/b_hi partition_points shift the sketch by the halos so bucket k's effective range becomes `[cuts[k-1] - halo_lo, cuts[k] + halo_hi)`. - `aqe/mod.rs::update_stage_progress` walks the full plan for the `RangeFilterExec` sitting directly on the boundary `ExchangeExec` (matched by `routing_expr` eq) and reads its halos. Absent RFE is an error — range-repartition writers produce straddler duplicates that only the reader-side filter can trim, so the shape is required. - Non-Float64 halo scalars also error (shape violation). Verified: h2o Q8 @ 1e7 with `parallel_window.enabled=true` now diffs to the DataFusion oracle at OK. WIP: reviewer TODOs left inline for follow-up: - multi-hit handling in `repartition_routing_expr` - test coverage for halo_hi and K=3 partitions - whether the walker should scope to the downstream stage - multi-legged plan behaviour - hard-error on routing_expr mismatch Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/runtime_stats.rs | 118 +++++++++++++----- ballista/scheduler/src/state/aqe/mod.rs | 112 +++++++++++++++-- .../src/state/aqe/test/range_repartition.rs | 3 +- 3 files changed, 196 insertions(+), 37 deletions(-) diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index ea70731ef8..4e72b2aa03 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -848,6 +848,7 @@ pub fn repartition_routing_expr( [] => internal_err!("OrderedRangeRepartitionExec has empty ORDER BY"), }; } + // TODO: are multiple chilren an error? for child in plan.children() { if let Some(expr) = repartition_routing_expr(child.as_ref())? { return Ok(Some(expr)); @@ -859,35 +860,46 @@ pub fn repartition_routing_expr( /// Rebuild a stage's `Vec>` under range-repartition /// overlap semantics: for each producer file in `original_partitions`, /// find its sketch (from `reports`), and route the file into every -/// downstream partition whose global cut range overlaps +/// downstream partition whose *halo-widened* range overlaps /// `[sketch.min(), sketch.max()]`. /// -/// Downstream partition ranges follow the half-open convention: -/// - `k = 0` → `(-∞, cuts[0])` -/// - `0 < k < K - 1` → `[cuts[k-1], cuts[k])` -/// - `k = K - 1` → `[cuts[K-2], +∞)` +/// Downstream partition ranges follow the half-open convention, widened +/// by the downstream `RangeFilterExec`'s halos on each side: +/// - `k = 0` → `(-∞, cuts[0] + halo_hi)` +/// - `0 < k < K - 1` → `[cuts[k-1] - halo_lo, cuts[k] + halo_hi)` +/// - `k = K - 1` → `[cuts[K-2] - halo_lo, +∞)` /// /// `[min, max]` overlaps `[lower, upper)` iff `max >= lower AND min < upper`. /// +/// `halo_lo`/`halo_hi` are `0.0` when the downstream stage has no halo +/// consumer (hash-agg, no-window range-repartition) — the check collapses +/// to raw cuts. When the downstream stage has a `RangeFilterExec` with +/// non-zero halo (bounded RANGE-frame windows), the caller passes the +/// widened halos so files straddling the halo band route to both sides. +/// Skipping this widening loses boundary rows from downstream window sums. +/// /// Files without a corresponding sketch (missing entirely, or present /// with `count == 0`) are safe to skip only when `partition_stats.num_rows` /// confirms the file is empty (`Some(0)`). If the file has rows or the /// row count is unknown (`None`), silently skipping would lose data — /// error out instead. /// -/// TODO(halo-aware routing): when the downstream stage has a -/// `RangeFilterExec` with non-zero halo (bounded RANGE-frame windows), each -/// partition's *effective* read range is `[cuts[k-1] - halo_lo, cuts[k] + -/// halo_hi)`. This function currently uses the raw cut range, so files -/// straddling the halo boundary aren't routed to their halo-widened -/// consumer. Boundary rows near cuts can be missing from downstream -/// window sums — a correctness gap for RANGE frames that this refactor -/// does not resolve. Fix requires reaching across stages to read the -/// consumer RFE's halos and widening the overlap check here. +/// # Arguments +/// +/// * `original_partitions` — passthrough shuffle output, `partitions[k]` +/// holds every file the writer produced for global partition `k`. +/// * `reports` — one per completed producer task; each carries the +/// per-sub-part sketches used for overlap lookup. +/// * `global_cuts` — K-1 monotone quantile cuts derived from merged +/// sketches; produce K downstream buckets. +/// * `halo_lo` / `halo_hi` — downstream `RangeFilterExec`'s halo widths +/// in the routing expression's units. `0.0` for non-halo consumers. pub fn cut_partitions( original_partitions: Vec>, reports: &[TaskRuntimeStats], global_cuts: &[f64], + halo_lo: f64, + halo_hi: f64, ) -> Result>> { use std::collections::HashMap; @@ -937,13 +949,15 @@ pub fn cut_partitions( } continue; }; - // Bucket i has (lower, upper) = (cuts[i-1], cuts[i]) with ±∞ at the - // ends, and matches iff `sketch_max >= lower && sketch_min < upper`. - // Monotone cuts → the set of matching buckets is a contiguous range - // [b_lo, b_hi], found by two partition_points over `global_cuts`. + // Bucket i has (lower, upper) = (cuts[i-1] - halo_lo, cuts[i] + + // halo_hi) with ±∞ at the ends, and matches iff `sketch_max + + // halo_lo >= cuts[i-1] && sketch_min - halo_hi < cuts[i]`. + // Monotone cuts → the set of matching buckets is a contiguous + // range [b_lo, b_hi], found by two partition_points over + // `global_cuts` with the sketch shifted by the halos. let (sketch_min, sketch_max) = (sketch.min(), sketch.max()); - let b_lo = global_cuts.partition_point(|&c| c <= sketch_min); - let b_hi = global_cuts.partition_point(|&c| c <= sketch_max); + let b_lo = global_cuts.partition_point(|&c| c <= sketch_min - halo_hi); + let b_hi = global_cuts.partition_point(|&c| c <= sketch_max + halo_lo); for bucket in &mut remapped[b_lo..=b_hi] { bucket.push(file.clone()); } @@ -1753,7 +1767,7 @@ mod overlap_remap_tests { // Passthrough map: both producers wrote to sub_part_id=0. let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2, "K = cuts.len() + 1"); // Partition 0: only producer 100. assert_eq!(remapped[0].len(), 1); @@ -1773,7 +1787,7 @@ mod overlap_remap_tests { let cuts = vec![15.0]; let original_partitions = vec![vec![location(0, 300)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert_eq!(remapped[0].len(), 1, "straddler in partition 0"); assert_eq!(remapped[0][0].file_id, Some(300)); @@ -1793,7 +1807,7 @@ mod overlap_remap_tests { bad.file_id = None; let original_partitions = vec![vec![bad]]; - let err = cut_partitions(original_partitions, &reports, &cuts) + let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) .expect_err("missing file_id must surface as an error"); assert!( err.to_string().contains("missing file_id"), @@ -1815,7 +1829,7 @@ mod overlap_remap_tests { let cuts = vec![10.0]; let original_partitions = vec![vec![location(0, 200)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1829,7 +1843,7 @@ mod overlap_remap_tests { let cuts = vec![10.0]; let original_partitions = vec![vec![location(0, 100)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1846,7 +1860,7 @@ mod overlap_remap_tests { orphan.partition_stats = PartitionStats::new(Some(5), None, None); let original_partitions = vec![vec![orphan]]; - let err = cut_partitions(original_partitions, &reports, &cuts) + let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) .expect_err("file with rows but no sketch must error"); let msg = err.to_string(); assert!( @@ -1886,7 +1900,7 @@ mod overlap_remap_tests { location(0, 6), ]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 4); let ids = |b: &[PartitionLocation]| { let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); @@ -1910,7 +1924,7 @@ mod overlap_remap_tests { orphan.partition_stats = PartitionStats::default(); // num_rows = None let original_partitions = vec![vec![orphan]]; - let err = cut_partitions(original_partitions, &reports, &cuts) + let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) .expect_err("file with unknown rows but no sketch must error"); let msg = err.to_string(); assert!( @@ -1918,4 +1932,52 @@ mod overlap_remap_tests { "unexpected error: {msg}" ); } + + /// A producer whose sketch lies entirely on one side of a cut but + /// within `halo_lo` of it must ALSO route to the partition on the + /// other side — the downstream `RangeFilterExec` widens each + /// partition's effective range by `[halo_lo, halo_hi]`, and any file + /// that could contain rows in the widened band must be visible to + /// that partition or the range-frame window sums it lose those rows. + /// + /// Matches h2o Q8's shape: `RANGE BETWEEN 3 PRECEDING AND CURRENT ROW` + /// → halo_lo = 3, halo_hi = 0. + // TODO: add a test that lovers 3 partitios including halo_hi + #[test] + fn overlap_remap_routes_halo_band_to_adjacent_partition() { + // Producer 100: sketch [13, 14] — below cut 15, but within + // halo_lo=3 of it. Must appear in BOTH partition 0 (its own bucket) + // AND partition 1 (via halo widening). + // Producer 200: sketch [16, 17] — above cut, no halo needed to + // route it to partition 1. + let reports = vec![ + sketch_report(100, vec![vec![13.0, 14.0]]), + sketch_report(200, vec![vec![16.0, 17.0]]), + ]; + let cuts = vec![15.0]; + let halo_lo = 3.0; + let halo_hi = 0.0; + let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; + + let remapped = + cut_partitions(original_partitions, &reports, &cuts, halo_lo, halo_hi) + .unwrap(); + let ids = |b: &[PartitionLocation]| { + let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); + v.sort(); + v + }; + assert_eq!(remapped.len(), 2); + assert_eq!( + ids(&remapped[0]), + vec![100u64], + "partition 0 sees only its own bucket (halo_hi=0)" + ); + assert_eq!( + ids(&remapped[1]), + vec![100u64, 200], + "partition 1's effective range is [15-3, +∞) = [12, +∞); \ + producer 100's [13,14] falls in the halo band and must route here", + ); + } } diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index dac8482214..5ed12113c4 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -18,7 +18,7 @@ use crate::display::print_stage_metrics; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::timestamp_millis; -use crate::state::aqe::execution_plan::RangeRepartitionRouting; +use crate::state::aqe::execution_plan::{ExchangeExec, RangeRepartitionRouting}; use crate::state::aqe::planner::{AdaptivePlanner, AdaptiveStageInfo}; use crate::state::execution_graph::{ ExecutionGraph, ExecutionGraphBox, ExecutionStage, ResolvedStage, RunningTaskInfo, @@ -29,7 +29,8 @@ use crate::state::task_manager::UpdatedStages; use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{ - cut_partitions, merge_runtime_stats_reports, repartition_routing_expr, + RangeFilterExec, cut_partitions, merge_runtime_stats_reports, + repartition_routing_expr, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; @@ -38,10 +39,13 @@ use ballista_core::serde::protobuf::{ job_status, task_status, }; use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; +use datafusion::scalar::ScalarValue; use log::{debug, error, info, warn}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -372,12 +376,28 @@ impl AdaptiveExecutionGraph { && let Some(routing) = Self::repartition_routing(stage, routing_expr)? { let reports = &stage.runtime_stats_reports; - let remapped = cut_partitions(partitions, reports, &routing.cuts) - .map_err(|err| { - BallistaError::General(format!( - "range-repartition stage {stage_id}: overlap remap failed: {err}" - )) - })?; + // Reader-side halos widen each partition's effective + // range to `[cuts[k-1] - halo_lo, cuts[k] + halo_hi)`; + // files straddling that band must route to both sides. + // Range-repartition boundaries always have a consuming + // RFE (writer produces straddler duplicates that only + // the reader-side filter can trim), so the walker errors + // if none is found. + // TODO: do we really mean to walk the whole plan? Or just stage N+1? + let (halo_lo, halo_hi) = + downstream_halos(&self.planner.plan, &routing.routing_expr) + .map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: halo lookup failed: {err}" + )) + })?; + let remapped = + cut_partitions(partitions, reports, &routing.cuts, halo_lo, halo_hi) + .map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: overlap remap failed: {err}" + )) + })?; // Save boundaries to ExchangeExec so they are there for resolve_stage_partitions self.planner.set_repartition_routing(stage_id, routing)?; remapped @@ -1429,6 +1449,82 @@ impl ExecutionGraph for AdaptiveExecutionGraph { } } +/// Find the halos of the downstream `RangeFilterExec` planted directly on +/// the boundary `ExchangeExec` for `routing_expr`. A range-repartition +/// boundary always has a consuming RFE — the writer produces straddler +/// duplicates that only the reader-side filter can trim — so absence is +/// a shape bug, not a runtime default. +/// +/// The rule may plant multiple RFEs (e.g. wide directly on the boundary +/// and narrow above the window operator). Only the one whose immediate +/// child is the boundary `ExchangeExec` describes the reader-visible +/// halo band and matters for straddler routing. +/// +/// # Arguments +/// +/// * `full_plan` — the AdaptivePlanner's current plan tree; the walker +/// descends the whole tree looking for the RFE on this boundary. +/// * `routing_expr` — the range-repartition op's routing expression; +/// matched by `PhysicalExpr::eq` against each candidate RFE's own. +fn downstream_halos( + full_plan: &Arc, + routing_expr: &Arc, +) -> datafusion::common::Result<(f64, f64)> { + let mut result: Option<(f64, f64)> = None; + let mut halo_err: Option = None; + full_plan.apply(|node| { + let Some(rf) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + let children = rf.children(); + let [child] = children.as_slice() else { + // TODO: is this correct for multi-legged plans? + // we probably need to start at the leaves and walk up? + return Ok(TreeNodeRecursion::Continue); + }; + if !child.is::() { + return Ok(TreeNodeRecursion::Continue); + } + if !rf.routing_expr().eq(routing_expr) { + // TODO: should hard error? + return Ok(TreeNodeRecursion::Continue); + } + match (scalar_to_f64(rf.halo_lo()), scalar_to_f64(rf.halo_hi())) { + (Ok(lo), Ok(hi)) => { + result = Some((lo, hi)); + Ok(TreeNodeRecursion::Stop) + } + (Err(e), _) | (_, Err(e)) => { + halo_err = Some(e); + Ok(TreeNodeRecursion::Stop) + } + } + })?; + if let Some(e) = halo_err { + return Err(e); + } + result.ok_or_else(|| { + datafusion::common::DataFusionError::Internal(format!( + "range-repartition boundary for expr `{routing_expr}` has no consuming \ + RangeFilterExec — straddler files would corrupt downstream partial \ + aggregates. Check the rule that planted this ORRE.", + )) + }) +} + +/// Halos travel through the RFE public API as `ScalarValue` for future +/// type widening (Interval, timestamps under KLL). Today the internal +/// routing math is `f64`; any other variant is a shape violation upstream +/// and we fail loud rather than silently zero-widen. +fn scalar_to_f64(sv: &ScalarValue) -> datafusion::common::Result { + match sv { + ScalarValue::Float64(Some(v)) => Ok(*v), + other => datafusion::common::internal_err!( + "only f64 halos are implemented, got: {other:?}" + ), + } +} + /// Checks is the plan same as expected string representation #[cfg(test)] #[macro_export] diff --git a/ballista/scheduler/src/state/aqe/test/range_repartition.rs b/ballista/scheduler/src/state/aqe/test/range_repartition.rs index 8d25a8b83e..129e90d2d3 100644 --- a/ballista/scheduler/src/state/aqe/test/range_repartition.rs +++ b/ballista/scheduler/src/state/aqe/test/range_repartition.rs @@ -149,7 +149,8 @@ async fn routing_parks_when_range_repartition_is_plan_root() }, }]; let cuts = vec![15.0]; - let remapped = cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts)?; + let remapped = + cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts, 0.0, 0.0)?; // `cut_partitions` must duplicate the straddler into both partitions — // the read-side filter is expected to trim on read. From 68f1c73b19369db1c79daa83cbdfb5e27671af9f Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 09:39:33 -0600 Subject: [PATCH 04/11] style: cargo fmt --- .../core/src/execution_plans/runtime_stats.rs | 15 +++++--- ballista/scheduler/src/state/aqe/mod.rs | 35 +++++++++++-------- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 4e72b2aa03..a3d1c86e5c 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -1767,7 +1767,8 @@ mod overlap_remap_tests { // Passthrough map: both producers wrote to sub_part_id=0. let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2, "K = cuts.len() + 1"); // Partition 0: only producer 100. assert_eq!(remapped[0].len(), 1); @@ -1787,7 +1788,8 @@ mod overlap_remap_tests { let cuts = vec![15.0]; let original_partitions = vec![vec![location(0, 300)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert_eq!(remapped[0].len(), 1, "straddler in partition 0"); assert_eq!(remapped[0][0].file_id, Some(300)); @@ -1829,7 +1831,8 @@ mod overlap_remap_tests { let cuts = vec![10.0]; let original_partitions = vec![vec![location(0, 200)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1843,7 +1846,8 @@ mod overlap_remap_tests { let cuts = vec![10.0]; let original_partitions = vec![vec![location(0, 100)]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1900,7 +1904,8 @@ mod overlap_remap_tests { location(0, 6), ]]; - let remapped = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = + cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); assert_eq!(remapped.len(), 4); let ids = |b: &[PartitionLocation]| { let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index 5ed12113c4..792e5d0683 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -384,20 +384,27 @@ impl AdaptiveExecutionGraph { // the reader-side filter can trim), so the walker errors // if none is found. // TODO: do we really mean to walk the whole plan? Or just stage N+1? - let (halo_lo, halo_hi) = - downstream_halos(&self.planner.plan, &routing.routing_expr) - .map_err(|err| { - BallistaError::General(format!( - "range-repartition stage {stage_id}: halo lookup failed: {err}" - )) - })?; - let remapped = - cut_partitions(partitions, reports, &routing.cuts, halo_lo, halo_hi) - .map_err(|err| { - BallistaError::General(format!( - "range-repartition stage {stage_id}: overlap remap failed: {err}" - )) - })?; + let (halo_lo, halo_hi) = downstream_halos( + &self.planner.plan, + &routing.routing_expr, + ) + .map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: halo lookup failed: {err}" + )) + })?; + let remapped = cut_partitions( + partitions, + reports, + &routing.cuts, + halo_lo, + halo_hi, + ) + .map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: overlap remap failed: {err}" + )) + })?; // Save boundaries to ExchangeExec so they are there for resolve_stage_partitions self.planner.set_repartition_routing(stage_id, routing)?; remapped From efaa6c8883ada0b9b54c3cf49be88fe29c51c51a Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 09:52:44 -0600 Subject: [PATCH 05/11] refactor(core): repartition_routing_expr walks the partition-preserving spine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous body descended into every child looking for the first (U/O)RRE, so a range-repartition op sitting below a join or union leg would be misattributed as driving this stage's output partitioning. In practice today only ParallelWindowRule plants RREs (single-legged shape), so the misattribution was inert — but the shape was ambiguous and needed to become explicit before further rules land. Descend only through nodes on the [`preserves_partitioning`] whitelist (Filter, Projection, Sort-with-preserve-partitioning, RSE, ShuffleWriter, Window, Buffer). Any partition-non-preserving barrier (join, union, hash-agg, unknown op) stops descent and returns `Ok(None)` — the RRE below such a barrier isn't visible in this stage's output partitioning anyway. Multi-child descent-through-node is now an internal_err since every whitelist entry is single-child by construction; the guard catches whitelist bugs early. No behaviour change on the current call sites; every existing `repartition_routing_expr_*` test still passes and h2o Q8 @ 1e7 still verifies OK against the DataFusion oracle. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/runtime_stats.rs | 40 ++++++++++++++----- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index a3d1c86e5c..31a6e7b2e1 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -827,12 +827,23 @@ pub struct TaskRuntimeStats { pub report: RuntimeStatsReport, } -/// Walk `plan` for the first `UnorderedRangeRepartitionExec` or -/// `OrderedRangeRepartitionExec` and return its routing expression -/// (`order_by[0].expr`). `Ok(None)` means no range-repartition operator -/// in the plan; `Err(_)` means one was found but its `order_by` was +/// Walk the partition-preserving spine of `plan` for the +/// `UnorderedRangeRepartitionExec` or `OrderedRangeRepartitionExec` that +/// drives this stage's output partitioning, and return its routing +/// expression (`order_by[0].expr`). +/// +/// The spine is the chain of partition-preserving ops (see +/// [`preserves_partitioning`]) between the stage root and the barrier +/// that sets the stage's output partitioning. Descent stops at any +/// non-preserving op (join, union, hash-agg, unknown node) — an RRE +/// below such a barrier drives a different logical partitioning that +/// this stage's output no longer carries. +/// +/// `Ok(None)` means no range-repartition op drives this stage's +/// partitioning; `Err(_)` means one was found but its `order_by` was /// empty (invariant break — a range repartition without a routing key -/// can't route anything). +/// can't route anything), or the spine hit a partition-preserving node +/// with more than one child (shape bug in the whitelist). pub fn repartition_routing_expr( plan: &dyn ExecutionPlan, ) -> Result>> { @@ -848,13 +859,20 @@ pub fn repartition_routing_expr( [] => internal_err!("OrderedRangeRepartitionExec has empty ORDER BY"), }; } - // TODO: are multiple chilren an error? - for child in plan.children() { - if let Some(expr) = repartition_routing_expr(child.as_ref())? { - return Ok(Some(expr)); - } + if !super::preserves_partitioning(plan) { + return Ok(None); + } + let children = plan.children(); + match children.as_slice() { + [] => Ok(None), + [child] => repartition_routing_expr(child.as_ref()), + _ => internal_err!( + "partition-preserving op `{}` has {} children — the whitelist \ + assumes single-child; expand the algorithm if this fires", + plan.name(), + children.len() + ), } - Ok(None) } /// Rebuild a stage's `Vec>` under range-repartition From 684da907a1ca4a17af4e249f872fd78efae85130 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 10:03:42 -0600 Subject: [PATCH 06/11] test(core): cover both halo sides + far-partition non-bleed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The K=2 halo test only proved `halo_lo` widens downward on a 2-partition layout — nothing verified `halo_hi` or that the halo band stays local to adjacent partitions. Replace with a K=5 layout using asymmetric halos (`halo_lo=1`, `halo_hi=2`) so a single fixture proves: - `halo_lo` widens each partition downward (200 → P2 via halo_lo). - `halo_hi` widens each partition upward (400 → P2 via halo_hi). - The middle partition sees siblings from BOTH halo bands simultaneously. - Halo does NOT bleed across two cut hops — P0 doesn't get 200's file, P4 doesn't get 400's file. Existing `disjoint` and `straddling` tests still cover the no-halo path (halo=0,0), so the raw-cut semantics stay separately verified. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/runtime_stats.rs | 82 +++++++++++++------ 1 file changed, 56 insertions(+), 26 deletions(-) diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 31a6e7b2e1..59ea829e68 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -1956,31 +1956,47 @@ mod overlap_remap_tests { ); } - /// A producer whose sketch lies entirely on one side of a cut but - /// within `halo_lo` of it must ALSO route to the partition on the - /// other side — the downstream `RangeFilterExec` widens each - /// partition's effective range by `[halo_lo, halo_hi]`, and any file - /// that could contain rows in the widened band must be visible to - /// that partition or the range-frame window sums it lose those rows. + /// Halo widening lets each partition see files that sit within + /// `[halo_lo, halo_hi]` of its raw cut range — the downstream + /// `RangeFilterExec`'s frame-context rows come from those files, and + /// missing any of them causes RANGE-frame window sums to drop rows + /// at boundaries. /// - /// Matches h2o Q8's shape: `RANGE BETWEEN 3 PRECEDING AND CURRENT ROW` - /// → halo_lo = 3, halo_hi = 0. - // TODO: add a test that lovers 3 partitios including halo_hi + /// The K=5 layout with `halo_lo != halo_hi` proves three things at + /// once: (a) `halo_lo` widens downward, (b) `halo_hi` widens upward, + /// (c) the halo band stays *local* — it does not bleed across two + /// cut hops to far-away partitions. #[test] - fn overlap_remap_routes_halo_band_to_adjacent_partition() { - // Producer 100: sketch [13, 14] — below cut 15, but within - // halo_lo=3 of it. Must appear in BOTH partition 0 (its own bucket) - // AND partition 1 (via halo widening). - // Producer 200: sketch [16, 17] — above cut, no halo needed to - // route it to partition 1. + fn overlap_remap_halo_band_widens_both_sides_without_bleeding_to_far_partitions() { + // K=5, asymmetric halos so we can tell halo_lo and halo_hi apart. + let cuts = vec![10.0, 20.0, 30.0, 40.0]; + let halo_lo = 1.0; + let halo_hi = 2.0; + // Effective partition ranges: + // P0: (-∞, 12) P1: [9, 22) P2: [19, 32) P3: [29, 42) P4: [39, +∞) let reports = vec![ - sketch_report(100, vec![vec![13.0, 14.0]]), - sketch_report(200, vec![vec![16.0, 17.0]]), + // 100 sits deep inside P0 — far from P1's halo, stays P0-only. + sketch_report(100, vec![vec![5.0, 6.0]]), + // 200 is entirely below cut 20 but within halo_lo=1 of it — + // routes to P1 (own bucket) AND P2 (halo band from below). + // Must NOT reach P0 (two cut hops away). + sketch_report(200, vec![vec![18.0, 19.0]]), + // 300 sits cleanly inside P2 — no halo participation. + sketch_report(300, vec![vec![25.0, 26.0]]), + // 400 is entirely above cut 30 but within halo_hi=2 of it — + // routes to P2 (halo band from above) AND P3 (own bucket). + // Must NOT reach P4 (two cut hops away). + sketch_report(400, vec![vec![31.0, 32.0]]), + // 500 sits deep inside P4 — far from P3's halo, stays P4-only. + sketch_report(500, vec![vec![45.0, 46.0]]), ]; - let cuts = vec![15.0]; - let halo_lo = 3.0; - let halo_hi = 0.0; - let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; + let original_partitions = vec![vec![ + location(0, 100), + location(0, 200), + location(0, 300), + location(0, 400), + location(0, 500), + ]]; let remapped = cut_partitions(original_partitions, &reports, &cuts, halo_lo, halo_hi) @@ -1990,17 +2006,31 @@ mod overlap_remap_tests { v.sort(); v }; - assert_eq!(remapped.len(), 2); + assert_eq!(remapped.len(), 5); assert_eq!( ids(&remapped[0]), vec![100u64], - "partition 0 sees only its own bucket (halo_hi=0)" + "P0 sees only its own bucket — 200's halo band belongs to P1/P2, not here", ); assert_eq!( ids(&remapped[1]), - vec![100u64, 200], - "partition 1's effective range is [15-3, +∞) = [12, +∞); \ - producer 100's [13,14] falls in the halo band and must route here", + vec![200u64], + "P1 sees its own straddler (200) below cut 20", + ); + assert_eq!( + ids(&remapped[2]), + vec![200u64, 300, 400], + "P2 (middle) sees siblings from BOTH halo bands — 200 via halo_lo, 400 via halo_hi — plus its own 300", + ); + assert_eq!( + ids(&remapped[3]), + vec![400u64], + "P3 sees its own straddler (400) above cut 30", + ); + assert_eq!( + ids(&remapped[4]), + vec![500u64], + "P4 sees only its own bucket — 400's halo band belongs to P2/P3, not here", ); } } From ac5fa0886f4c2f76027be5faa8f3155701c423c1 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 10:22:05 -0600 Subject: [PATCH 07/11] refactor(scheduler): halo walker matches by producer stage_id + hard-errors on shape violations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two concerns the previous walker punted on: - `PhysicalExpr::eq(routing_expr)` isn't unique across stages — two range-repartition boundaries could share the same `Column(name, idx)` shape after independent projections, and the top-down walk would return the first match rather than THIS boundary's RFE. - A `RangeFilterExec` with `children().len() != 1` is a shape violation, but the walker silently `Continue`d past it. Match by `ExchangeExec::stage_id() == producer_stage_id` — the caller already has the completing stage's id and stage_id is uniquely assigned per boundary, so no cross-stage collision. Any non-single- child RFE hard-errors via `internal_err!`. The manual `halo_err` sentinel is gone — `apply` propagates `Err` from the closure directly. Also fixes the `[`preserves_partitioning`]` intra-doc link that CI caught with `-D warnings` — the function lives in `super::` from the runtime_stats scope. Verified: h2o Q8 @ 1e7 with `parallel_window.enabled=true` still verifies OK against the DataFusion oracle. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/runtime_stats.rs | 2 +- ballista/scheduler/src/state/aqe/mod.rs | 69 +++++++++---------- 2 files changed, 32 insertions(+), 39 deletions(-) diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 59ea829e68..312041f7d5 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -833,7 +833,7 @@ pub struct TaskRuntimeStats { /// expression (`order_by[0].expr`). /// /// The spine is the chain of partition-preserving ops (see -/// [`preserves_partitioning`]) between the stage root and the barrier +/// [`super::preserves_partitioning`]) between the stage root and the barrier /// that sets the stage's output partitioning. Descent stops at any /// non-preserving op (join, union, hash-agg, unknown node) — an RRE /// below such a barrier drives a different logical partitioning that diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index 792e5d0683..0b70cf612a 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -42,7 +42,6 @@ use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; -use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion::scalar::ScalarValue; @@ -383,12 +382,8 @@ impl AdaptiveExecutionGraph { // RFE (writer produces straddler duplicates that only // the reader-side filter can trim), so the walker errors // if none is found. - // TODO: do we really mean to walk the whole plan? Or just stage N+1? - let (halo_lo, halo_hi) = downstream_halos( - &self.planner.plan, - &routing.routing_expr, - ) - .map_err(|err| { + let (halo_lo, halo_hi) = downstream_halos(&self.planner.plan, stage_id) + .map_err(|err| { BallistaError::General(format!( "range-repartition stage {stage_id}: halo lookup failed: {err}" )) @@ -1457,10 +1452,17 @@ impl ExecutionGraph for AdaptiveExecutionGraph { } /// Find the halos of the downstream `RangeFilterExec` planted directly on -/// the boundary `ExchangeExec` for `routing_expr`. A range-repartition -/// boundary always has a consuming RFE — the writer produces straddler -/// duplicates that only the reader-side filter can trim — so absence is -/// a shape bug, not a runtime default. +/// the boundary `ExchangeExec` produced by `producer_stage_id`. A +/// range-repartition boundary always has a consuming RFE — the writer +/// produces straddler duplicates that only the reader-side filter can +/// trim — so absence is a shape bug, not a runtime default. +/// +/// Boundaries are disambiguated by `ExchangeExec::stage_id()` (the id +/// of the stage that produces to the exchange), not by routing +/// expression: two range-repartition stages could share the same +/// `Column(name, index)` shape after independent projections, so +/// `PhysicalExpr::eq` isn't unique across stages. The producer stage id +/// is. /// /// The rule may plant multiple RFEs (e.g. wide directly on the boundary /// and narrow above the window operator). Only the one whose immediate @@ -1471,50 +1473,41 @@ impl ExecutionGraph for AdaptiveExecutionGraph { /// /// * `full_plan` — the AdaptivePlanner's current plan tree; the walker /// descends the whole tree looking for the RFE on this boundary. -/// * `routing_expr` — the range-repartition op's routing expression; -/// matched by `PhysicalExpr::eq` against each candidate RFE's own. +/// * `producer_stage_id` — the id of the completing upstream stage; +/// matches `ExchangeExec::stage_id()` on the boundary exchange. fn downstream_halos( full_plan: &Arc, - routing_expr: &Arc, + producer_stage_id: usize, ) -> datafusion::common::Result<(f64, f64)> { let mut result: Option<(f64, f64)> = None; - let mut halo_err: Option = None; full_plan.apply(|node| { let Some(rf) = node.downcast_ref::() else { return Ok(TreeNodeRecursion::Continue); }; let children = rf.children(); let [child] = children.as_slice() else { - // TODO: is this correct for multi-legged plans? - // we probably need to start at the leaves and walk up? - return Ok(TreeNodeRecursion::Continue); + return datafusion::common::internal_err!( + "RangeFilterExec must have exactly 1 child, got {}", + children.len() + ); }; - if !child.is::() { + // RFE above the window op (narrow, halo=[0,0]) has a + // `PartitionedBoundedWindowAggExec` child, not the boundary + // exchange — keep looking for the wide RFE below. + let Some(exchange) = child.downcast_ref::() else { return Ok(TreeNodeRecursion::Continue); - } - if !rf.routing_expr().eq(routing_expr) { - // TODO: should hard error? + }; + if exchange.stage_id() != Some(producer_stage_id) { return Ok(TreeNodeRecursion::Continue); } - match (scalar_to_f64(rf.halo_lo()), scalar_to_f64(rf.halo_hi())) { - (Ok(lo), Ok(hi)) => { - result = Some((lo, hi)); - Ok(TreeNodeRecursion::Stop) - } - (Err(e), _) | (_, Err(e)) => { - halo_err = Some(e); - Ok(TreeNodeRecursion::Stop) - } - } + result = Some((scalar_to_f64(rf.halo_lo())?, scalar_to_f64(rf.halo_hi())?)); + Ok(TreeNodeRecursion::Stop) })?; - if let Some(e) = halo_err { - return Err(e); - } result.ok_or_else(|| { datafusion::common::DataFusionError::Internal(format!( - "range-repartition boundary for expr `{routing_expr}` has no consuming \ - RangeFilterExec — straddler files would corrupt downstream partial \ - aggregates. Check the rule that planted this ORRE.", + "range-repartition boundary for producer stage {producer_stage_id} has \ + no consuming RangeFilterExec — straddler files would corrupt \ + downstream partial aggregates. Check the rule that planted this ORRE." )) }) } From fc8155565ef474c9c0dee18f61b4eb646bba01e6 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 10:43:55 -0600 Subject: [PATCH 08/11] less doc --- ballista/executor/src/execution_engine.rs | 4 ---- ballista/scheduler/src/state/aqe/adapter.rs | 19 +------------------ 2 files changed, 1 insertion(+), 22 deletions(-) diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index ea4dbffe17..e673a51c8e 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -137,10 +137,6 @@ impl ExecutionEngine for DefaultExecutionEngine { ) -> Result> { let plan = plan .transform(|p| { - // TODO: RangeShuffleReaderExec needs the same late-bind - // (with_work_dir + with_client_pool) once a planner rule - // plants it; without it, the first task carrying one will - // fail with "work dir should have been set by executor". if let Some(reader) = p.downcast_ref::() { match &self.client_pool { Some(client_pool) => Ok(Transformed::yes(Arc::new( diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 4b8ab2deef..338c3fb438 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -108,8 +108,7 @@ impl BallistaAdapter { // Ordered-writer path: when the child declared an output // ordering, preserve it across the shuffle boundary with a // k-way merge instead of the arrival-order concat that the - // regular reader does. Fixes the silent RANGE-frame - // corruption in `docs/developer/parallel-range-window.md`. + // regular reader does. if let Some(ordering) = exchange.input().output_ordering() { Arc::new(RangeShuffleReaderExec::try_new( stage_id, @@ -133,22 +132,6 @@ impl BallistaAdapter { exchange.input().output_partitioning().partition_count(), )?), }; - // The adapter no longer injects a `RangeFilterExec` above the - // reader. Rules that emit a range-repartition upstream (today: - // `ParallelWindowRule`) are also responsible for planting the - // read-side `RangeFilterExec`(s) at plan time with `cuts=None`; - // the scheduler-side `resolve_range_filter_cuts` walker fills - // in cuts once stage-0's sketches merge. If a range-routing - // ExchangeExec reaches this point without a rule-planted - // filter, `RangeFilterExec::execute()` fails loud rather than - // straddling sub-parts silently corrupting downstream partial - // aggregates. - if exchange.range_repartition_routing().is_some() { - debug!( - "range-repartition: ExchangeExec has resolved routing for stage {stage_id}; \ - any rule-planted RangeFilterExec above should have been resolved by now" - ); - } Ok(Transformed::yes(reader)) } else { Ok(Transformed::no(plan)) From 8dbd823028dc5b2a7fcbfcc6a29286837a768739 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 11:05:44 -0600 Subject: [PATCH 09/11] fix(scheduler): pair RFE with boundary by descent, not routing_expr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `resolve_range_filter_cuts` used `PhysicalExpr::eq` on routing_expr to pair each pending `RangeFilterExec` with its boundary `ExchangeExec`. That's ambiguous in multi-legged shapes (a future SMJ with range-repartition on both sides could have both boundaries route on `Column("v", 0)` after independent projections, silently wiring the wrong cuts to the wrong RFE). Same latent bug shape as the halo walker just fixed via `stage_id` disambiguation. Cure: descend the RFE's own single-child spine to the first descendant `ExchangeExec`. Descent is unique by construction — every RFE has exactly one child — so multi-legged shapes pair each RFE with its own leg's boundary. Then verify `rf.routing_expr().eq(&routing. routing_expr)` as a plant-time sanity check; disagreement means the rule wired the wrong RFE to the boundary. Shape violations (multi-child RFE, spine fork before hitting an exchange, unresolved descendant routing, routing_expr mismatch) all hard-error via `internal_err!` — plant-time or stage-progress bugs that should never be silently absorbed. No behaviour change on today's single-legged Q8 shape; h2o Q8 @ 1e7 with `parallel_window.enabled=true` still verifies OK against the DataFusion oracle. Co-Authored-By: Claude Opus 4.7 (1M context) --- ballista/scheduler/src/state/aqe/adapter.rs | 101 +++++++++++++------- 1 file changed, 69 insertions(+), 32 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 338c3fb438..29d137615d 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -16,7 +16,9 @@ // under the License. use crate::planner::create_shuffle_writer_with_config; -use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; +use crate::state::aqe::execution_plan::{ + AdaptiveDatafusionExec, ExchangeExec, RangeRepartitionRouting, +}; use crate::state::aqe::planner::AdaptiveStageInfo; use crate::state::execution_graph::StageOutput; use ballista_core::JobId; @@ -26,14 +28,12 @@ use ballista_core::execution_plans::{ use datafusion::common::exec_err; use datafusion::config::ConfigOptions; use datafusion::error::DataFusionError; -use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning}; use datafusion::scalar::ScalarValue; use datafusion::{ common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, physical_plan::ExecutionPlan, }; -use log::debug; use std::collections::HashMap; use std::sync::Arc; @@ -205,35 +205,34 @@ impl BallistaAdapter { } /// Walk `plan` and resolve every pending [`RangeFilterExec`]'s bounds from -/// its downstream [`ExchangeExec`]'s stored routing. Called at -/// `adapt_to_ballista` time, once stage-0's sketches have merged into -/// cuts and been parked on the boundary `ExchangeExec` via +/// its own descendant boundary `ExchangeExec`'s stored routing. Called at +/// `adapt_to_ballista` time, once the upstream stage's sketches have +/// merged into cuts and been parked on the boundary `ExchangeExec` via /// `set_repartition_routing`. /// -/// Cross-referencing is by `routing_expr` equality — the rule plants both -/// wide (below SPM) and narrow (above BWAG) filters with the same routing -/// expression, both pointing at the same shuffle boundary. +/// Pairing is by tree structure — each RFE descends its own subtree to +/// the first `ExchangeExec` and takes that exchange's routing. Using +/// `PhysicalExpr::eq` on routing exprs alone would collide in +/// multi-legged shapes (e.g. SMJ with range-repartition on both sides +/// sharing the same `Column(name, idx)` after independent projections); +/// descent is unique by construction because every RFE has exactly one +/// child. The RFE's own `routing_expr` is then cross-checked against +/// the descendant exchange's for a plant-time invariant assert — if they +/// disagree, the rule wired the wrong RFE to the boundary. /// /// RFE receives *unwidened* half-open ranges (`(cuts[k-1], cuts[k])` with ±∞ /// sentinels at ends). RFE widens by its own halos internally at /// `resolve_bounds` time — see the separation-of-concerns note on /// [`RangeFilterExec`]. The scheduler stays halo-blind at this boundary. /// -/// Errors if a `RangeFilterExec` is still pending after the walk: the -/// rule promised bounds and the scheduler didn't deliver, which is either -/// a routing_expr mismatch or a stage-progress bug. +/// Errors if descent hits a fork (multi-child op) before reaching an +/// `ExchangeExec`, if the RFE has anything other than 1 child, if the +/// descendant `ExchangeExec` has no resolved routing yet, or if its +/// routing_expr disagrees with the RFE's — all four are plant-time or +/// stage-progress bugs. fn resolve_range_filter_cuts( plan: &Arc, ) -> Result<(), DataFusionError> { - let mut routings: Vec<(Arc, Vec)> = Vec::new(); - plan.apply(|node| { - if let Some(exchange) = node.downcast_ref::() - && let Some(routing) = exchange.range_repartition_routing() - { - routings.push((routing.routing_expr, routing.cuts)); - } - Ok(TreeNodeRecursion::Continue) - })?; plan.apply(|node| { let Some(rf) = node.downcast_ref::() else { return Ok(TreeNodeRecursion::Continue); @@ -241,23 +240,61 @@ fn resolve_range_filter_cuts( if rf.raw_bounds().is_some() { return Ok(TreeNodeRecursion::Continue); } - let rf_expr = rf.routing_expr(); - let cuts = routings - .iter() - .find(|(expr, _)| expr.eq(rf_expr)) - .map(|(_, cuts)| cuts.clone()) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "RangeFilterExec: no matching ExchangeExec routing for expr {rf_expr}" - )) - })?; - let raw_bounds = raw_bounds_from_cuts(&cuts); + let children = rf.children(); + let [child] = children.as_slice() else { + return datafusion::common::internal_err!( + "RangeFilterExec must have exactly 1 child, got {}", + children.len() + ); + }; + let routing = descend_to_boundary_routing(child)?; + if !rf.routing_expr().eq(&routing.routing_expr) { + return datafusion::common::internal_err!( + "RangeFilterExec routing_expr `{}` disagrees with its descendant \ + boundary ExchangeExec's routing_expr `{}` — plant-time invariant \ + broken", + rf.routing_expr(), + routing.routing_expr + ); + } + let raw_bounds = raw_bounds_from_cuts(&routing.cuts); rf.resolve_bounds(raw_bounds)?; Ok(TreeNodeRecursion::Continue) })?; Ok(()) } +/// Descend the single-child spine below a `RangeFilterExec` until we +/// hit an `ExchangeExec`, and return its `range_repartition_routing`. +/// Forks in the spine are shape violations because a range-filter only +/// makes sense above a single-input boundary. +fn descend_to_boundary_routing( + start: &Arc, +) -> Result { + let mut node = Arc::clone(start); + loop { + if let Some(exchange) = node.downcast_ref::() { + return exchange.range_repartition_routing().ok_or_else(|| { + DataFusionError::Internal( + "RangeFilterExec's descendant ExchangeExec has no resolved \ + range-repartition routing yet — stage progress skipped a step" + .into(), + ) + }); + } + let children = node.children(); + let [child] = children.as_slice() else { + return datafusion::common::internal_err!( + "RangeFilterExec descent hit a fork at `{}` ({} children) — cannot \ + pair with a single boundary", + node.name(), + children.len() + ); + }; + node = Arc::clone(*child); + } +} + /// Project K-1 cuts to K half-open `(cuts[k-1], cuts[k])` ranges with `None` /// sentinels at ±∞. This is the pure range-partitioning projection — no halo /// arithmetic here (RFE widens internally at resolve time). From f6b6b2dd1353f72cb340f01b1afcf663662748b8 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 11:08:28 -0600 Subject: [PATCH 10/11] chore: move parallel-range-window design notes out of tree MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The doc was a running notebook between me and Claude — the "rope-bridge principle" formula, the end-state sketch, the ticked/unticked plan — not public developer reference material. Moving it to `.local/` (gitignored) keeps the notebook accessible locally without shipping it as `docs/developer/`. Also drops three memory-slug refs (`[[parallel-range-window]]`, `[[project-prefix-scan-two-pass-rejected]]`, `[[kll-sketch]]`) from `parallel_window.rs`'s module doc — those were private auto-memory links, they don't render as intra-doc references and readers can't follow them. Replaced with plain-English equivalents that stand on their own. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../aqe/optimizer_rule/parallel_window.rs | 9 +-- docs/developer/parallel-range-window.md | 76 ------------------- 2 files changed, 4 insertions(+), 81 deletions(-) delete mode 100644 docs/developer/parallel-range-window.md diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs index 282eeb8046..07379d7dc2 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs @@ -17,7 +17,7 @@ //! Rewrite bounded-RANGE-frame windows into a distributed range-shuffle so //! `BoundedWindowAggExec`'s single-partition constraint isn't a serial -//! bottleneck. See [[parallel-range-window]] for the design. +//! bottleneck. //! //! # Matched shape //! @@ -34,10 +34,9 @@ //! - single-column ORDER BY on a physical `Column` (widening: multi-key, //! computed exprs — separate rewrites) //! - `RANGE` frame with finite `PRECEDING` / `FOLLOWING` / `CurrentRow` -//! bounds (UNBOUNDED goes down the prefix-scan path — see -//! [[project-prefix-scan-two-pass-rejected]]) -//! - ORDER BY column is `Float64` today (T-Digest restriction; lifts with -//! [[kll-sketch]]) +//! bounds (UNBOUNDED frames go down a different path) +//! - ORDER BY column is `Float64` today (T-Digest restriction; lifts when +//! the sketch swaps to KLL) //! //! # Rewrite //! diff --git a/docs/developer/parallel-range-window.md b/docs/developer/parallel-range-window.md deleted file mode 100644 index 09914fc39a..0000000000 --- a/docs/developer/parallel-range-window.md +++ /dev/null @@ -1,76 +0,0 @@ - - -# Parallel bounded-RANGE-frame windows — design doc - -## The formula - -1. Envision the end state. -2. Figure out where you actually are. -3. Plot a course. -4. Create a series of steps. -5. Each step's vector must have a dot product > 0.9 with the target direction. -6. Plan the nearby step very well. -7. Plan the farthest step barely at all. -8. Interpolate planning along the way. -9. Revisit at each step. -10. **Rope-bridge principle.** Fire an arrow with a string; pull twine; pull rope; pull larger rope; pull a floor. The bridge fulfills "bridge" from day one. Ship the _shape_, then thicken. Correct-and-slow-and-limited is a viable arrow; the way to a full bridge is not to design the floor first. - -Corollary: a step is a candidate for skipping-and-back-filling if its output correctness holds without it. "Necessary for the final impl" ≠ "necessary for this step." Filtering files by ValueIndex range is important in the end state; it doesn't stop earlier steps from being correct without it. - -## The end state (step 1) - -On a cluster with E executors × V vcores each, scan the input across E×V partitions in parallel. Stage N collects per-partition stats via a runtime sketch and drives U/ORRE with `output_cnt == input_cnt`, keeping every vcore busy inside the task (local exchange is cheap). If the shuffle is ordered, write ValueIndex files so downstream doesn't have to over-sample or over-fetch. Cuts flow to the scheduler and become global. - -Stage N+1 uses multi-partition tasks: each task claims `vcores` input partitions plus a file-halo overlap with the neighbouring task. Overlap is a _task-level_ concern, not a partition-level concern — inside a task, adjacent partitions borrow context via local memory (free); across tasks, file-halo is fetched over shuffle via ValueIndex-based partial reads. The task k-way merges its inputs into `vcores` sorted DF partitions, using the global cuts to distribute evenly. It then filters with row-halo, runs `PartitionedBoundedWindowAggExec`, filters without row-halo, and writes `vcores` output files. - -Invariants: - -- **Width invariant.** Partition count stays at E×V across stages. No funnel except where explicitly planted (not usually needed, done in ballista client). -- **Ordered-shuffle propagation.** When the writer declared an output ordering, the reader preserves it via k-way merge. Every ordered-shuffle consumer (BWAG, SortMergeJoin build side, …) benefits from the same primitive. -- **Two-level halo.** Stage-level file-halo crosses task boundaries at shuffle cost (partial-file read); task-level row-halo crosses local-partition boundaries at memory cost. - -## The plan - -Ordered chronologically. Ticked items are landed (or mostly landed). Unticked items may be skipped and back-filled per the rope-bridge principle. - -- [x] **Multi-partition-task substrate.** #2038: `partition_slice` on task launch, K-drain ShuffleWriter, executor-side partition restriction. Every parallel operator downstream sits on this. -- [x] **T-Digest / KLL runtime stats.** #2180: sketch per partition, wire report to scheduler, merge into cuts. Foundation for any data-driven range op. -- [x] **URRE / ORRE.** #2169, #2196: N sorted overlapping → K sorted disjoint (ORRE via k-way merge internally) or unordered variant. -- [x] **RuntimeStatsExec, cut-discovery walker.** Same PR family. Late-binding cuts flow scheduler → downstream ops. -- [x] **RangeFilterExec + PartitionedBoundedWindowAggExec.** #2223: filter by resolved cuts + halo; BWAG wrapper that hides from tree walkers so `EnforceDistribution` doesn't collapse K→1. -- [x] **ParallelWindowRule.** #2223: match the `BWAG on Column Float64 ORDER BY, RANGE PRECEDING/FOLLOWING` shape, rewrite to insert `RuntimeStats → ORRE → RangeFilter(wide) → PartitionedBWAG → RangeFilter(narrow)`. -- [x] **Feature flag.** `ballista.planner.parallel_window.enabled=false` by default. Off is inert; on activates the rule. -- [x] **Ordered ShuffleReader.** New `RangeShuffleReaderExec` — keeps each upstream source as its own stream, feeds N into `StreamingMergeBuilder` on the child's declared ordering. Adapter picks it whenever `exchange.input().output_ordering().is_some()` (writer-driven gate). No permit governor and no per-source buffering; backpressure flows from the merge's demand through h2 / disk. Reusable — SortMergeJoin build side wants the same thing. Verified: h2o Q8 @ 1e7 SUM diff between `parallel_window.enabled=true` and `=false` agrees to 5e-14 relative (Float64 noise floor). If a wasted-merge cost surfaces in profiles later, tighten to demand-driven via a new consumer-side `requires_globally_sorted_input` bit. -- [x] **RangeFilter min/max fast-path + binary-search slice.** Post-ordered-ShuffleReader, batches are internally sorted. `min/max` fast paths (100% pass → Arc-clone; 0% overlap → skip) collapse the hot filter cost without touching correctness. Binary-search + `RecordBatch::slice` covers the mixed case with zero data copy. `sorted_on_key` derived at construction from `input.output_ordering()` — ascending on `routing_expr`, no config knob. Nullable routing columns fall back to `filter_record_batch` on a per-batch basis. Verified h2o Q8, 2 execs × 4 vcores, MPT=4: at 1e7 (2G cap) 7.6 s → 2.5 s = 3.0×; at 1e8 (4G cap) 143 s → 92 s = 1.55×. The 1e8 delta is smaller because the bottleneck shifts to shuffle/merge memory — see the next two items. -- [ ] **ValueIndex-based partial-file reads at shuffle-fetch time.** #2204 landed the write-side + reader primitives; consumer plumbing to translate value-range → byte-range at fetch has to hook in. Lets stage-N+1 halo reads pull only the halo slice from a neighbour file, not the whole file. -- [ ] **Per-task halo metadata on task-status.** New axis on task shipping: `own_files + halo_slice(file_ref, value_range)`. Composes with the ValueIndex plumbing — the scheduler emits `PartitionLocation`-with-value-range instead of `PartitionLocation`-whole-file. This is the rope. Enables inter-task halo without over-fetch. -- [ ] **Intra-task ORRE.** Once per-task input is one sorted merged stream (post-ordered-ShuffleReader), split it into `vcores` sub-partitions by task-local sub-cuts (derived from global cuts + vcore count). Task-level halo at the sub-partition boundaries is intra-task (local memory, free). Now cores stay busy inside every task without inter-task shuffle. -- [ ] **Two-level halo semantics in the rule.** RangeFilter at the stage boundary uses stage-level halo width (crosses tasks, shuffles); RangeFilter at the task-local boundary uses task-level halo width (crosses local partitions, memory). The rule plants both. -- [ ] **Symmetric halo (PRECEDING + FOLLOWING).** Generalize the halo direction on the shape. Plumbing extension; no new operators. -- [ ] **Range-partition invariant across stages.** E×V vcores → E×V in-flight partitions at every stage, funnel-free. Where an actual funnel is required, plant it explicitly. Everywhere else, keep width invariant. - -Out of scope for this end state: **unbounded PRECEDING** (running SUM). Halo degenerates to "all preceding tasks" → serialism. Sibling design (prefix scan). - -## Dot-product check on the near steps - -- **RangeFilter fast paths.** Perf. Correctness-preserving. Doesn't unlock or block anything else, but pays for itself immediately on the wide filter (which sees 100% pass every batch given ORRE's exact-routing). -- **ValueIndex plumbing.** Skippable at the cost of over-fetching halo files whole. Everything after it correctness-holds without it. -- **Per-task halo metadata.** Skippable at the cost of stage-level halo remaining a single-width parameter on the shape. Everything after correctness-holds without it. -- **Intra-task ORRE.** Skippable at the cost of not saturating vcores when a task's input is ≤ vcores partitions. From ea5108ae57deb7dc8bdfa5f5f6575854d6dc89f8 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 10 Aug 2026 11:32:05 -0600 Subject: [PATCH 11/11] refactor(scheduler): collapse ParallelWindow shape gate + rewrite into one fn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `as_candidate` and `rewrite_bwag` were split for no real reason: the former was a loose shape check that got thrown away and re-asserted by the latter (BWAG re-downcast with an "internal error: caller passed non-BWAG" path that could never fire). The split also blurred failure signaling — the non-Float64 order key and non-numeric halo bound checks came back as `Err("shape matched but skipped")` when they're really shape gates. Merged into `maybe_rewrite_bwag(node, K) -> Result>`: - `Ok(None)` = shape gate missed (not a BWAG, PARTITION BY, ROWS frame, UNBOUNDED bound, non-Float64 order key, non-numeric halo scalar, or already-rewritten subtree). Silent — no log spam on the hot path. - `Ok(Some(_))` = rewrite applied. - `Err(_)` = actual invariant violation (BWAG/SPM/SortExec with =/= 1 child, schema lookup on ORDER BY expr, constructor try_new). Drops the `WindowCandidate` shuttle struct and the dead BWAG re-downcast. `as_finite` -> `is_finite` (bool now that the passthrough-of-reference buys nothing); `halo_from_bound` -> `Option` matching the new "non-numeric = shape gate" model, with its test updated in kind. Verified: h2o Q8 @ 1e7 still produces the expected RFE_narrow -> PBWAG -> RFE_wide(halo=[3,0]) -> ORRE(->8 parts) plan and `--verify` passes vs the DataFusion oracle on 10M rows. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../partitioned_bounded_window_agg.rs | 10 +- .../aqe/optimizer_rule/parallel_window.rs | 252 ++++++++---------- 2 files changed, 113 insertions(+), 149 deletions(-) diff --git a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs index d8f044442e..2b01987740 100644 --- a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs +++ b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs @@ -74,11 +74,11 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; -// The rule's `as_candidate` gates guarantee no PARTITION BY + single Column -// ORDER BY over a sorted source, so `BWAG::try_new` is always invoked with -// `InputOrderMode::Sorted` and `can_repartition=false` (partition_keys() is -// empty either way when there's no PARTITION BY). Hardcode both to keep the -// wire and the type small. +// `maybe_rewrite_bwag`'s shape gates guarantee no PARTITION BY + single +// Column ORDER BY over a sorted source, so `BWAG::try_new` is always invoked +// with `InputOrderMode::Sorted` and `can_repartition=false` (partition_keys() +// is empty either way when there's no PARTITION BY). Hardcode both to keep +// the wire and the type small. const BWAG_INPUT_ORDER_MODE: InputOrderMode = InputOrderMode::Sorted; const BWAG_CAN_REPARTITION: bool = false; diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs index 07379d7dc2..1e8fe0cb6a 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs @@ -121,28 +121,9 @@ impl PhysicalOptimizerRule for ParallelWindowRule { // plan tree for the true source width; we use the config knob that // those later rules also target. let output_partitions = config.execution.target_partitions.max(2); - plan.transform_up(|node| { - let Some(candidate) = as_candidate(node.as_ref()) else { - return Ok(Transformed::no(node)); - }; - match rewrite_bwag(&node, &candidate, output_partitions) { - Ok(rewritten) => { - debug!( - "ParallelWindowRule: rewrote BWAG on `{}` (RANGE {} — {})", - candidate.order_key, - fmt_bound(&candidate.start_bound), - fmt_bound(&candidate.end_bound), - ); - Ok(Transformed::yes(rewritten)) - } - Err(e) => { - debug!( - "ParallelWindowRule: shape matched but rewrite skipped for `{}`: {e}", - candidate.order_key, - ); - Ok(Transformed::no(node)) - } - } + plan.transform_up(|node| match maybe_rewrite_bwag(&node, output_partitions)? { + Some(rewritten) => Ok(Transformed::yes(rewritten)), + None => Ok(Transformed::no(node)), }) .map(|t| t.data) } @@ -156,16 +137,6 @@ impl PhysicalOptimizerRule for ParallelWindowRule { } } -/// Shape captured from a matching `BoundedWindowAggExec`. Everything the -/// rewrite needs to build the new subtree. -#[derive(Debug, Clone)] -struct WindowCandidate { - order_key: String, - sort_expr: PhysicalSortExpr, - start_bound: WindowFrameBound, - end_bound: WindowFrameBound, -} - /// True if any descendant of `nodes` is an `OrderedRangeRepartitionExec` /// or `RangeFilterExec`. Used as an idempotency guard: those ops are what /// our own rewrite plants, so seeing them below a BWAG means we've already @@ -182,57 +153,14 @@ fn subtree_contains_our_rewrite(nodes: &[&Arc]) -> bool { false } -fn as_candidate(node: &dyn ExecutionPlan) -> Option { - let window = node.downcast_ref::()?; - // Shape gates as slice patterns: 0 or 2+ elements simply don't match. - let [expr] = window.window_expr() else { - return None; - }; - let [] = expr.partition_by() else { - return None; - }; - let [order] = expr.order_by() else { - return None; - }; - let column = order.expr.downcast_ref::()?; - let frame = expr.get_window_frame(); - let WindowFrameUnits::Range = frame.units else { - return None; - }; - let (Some(start), Some(end)) = - (as_finite(&frame.start_bound), as_finite(&frame.end_bound)) - else { - return None; - }; - // Idempotency: if the BWAG's subtree already contains our own - // range-repartition machinery, we've already rewritten this window. - // Re-plans (AQE fires the optimizer chain again for stage N+1) would - // otherwise wrap another ORRE around the previous rewrite's - // RangeFilterExec+ShuffleReader — and that ORRE's child doesn't claim - // ordering, blowing up at execute-time. - if subtree_contains_our_rewrite(window.children().as_slice()) { - return None; - } - Some(WindowCandidate { - order_key: column.name().to_string(), - sort_expr: order.clone(), - start_bound: start.clone(), - end_bound: end.clone(), - }) -} - -/// Returns the bound unchanged when it's `CurrentRow` or a non-null scalar -/// offset. `UNBOUNDED PRECEDING/FOLLOWING` is represented as a typed-null -/// scalar and returns `None`. -fn as_finite(bound: &WindowFrameBound) -> Option<&WindowFrameBound> { +/// True when the bound is `CurrentRow` or a non-null scalar offset. +/// `UNBOUNDED PRECEDING/FOLLOWING` is a typed-null scalar and returns `false`. +fn is_finite(bound: &WindowFrameBound) -> bool { match bound { - WindowFrameBound::CurrentRow => Some(bound), - WindowFrameBound::Preceding(scalar) | WindowFrameBound::Following(scalar) - if !scalar.is_null() => - { - Some(bound) + WindowFrameBound::CurrentRow => true, + WindowFrameBound::Preceding(scalar) | WindowFrameBound::Following(scalar) => { + !scalar.is_null() } - _ => None, } } @@ -244,55 +172,95 @@ fn fmt_bound(bound: &WindowFrameBound) -> String { } } -/// Extract the halo width in `f64` from a bound. `CurrentRow` → 0. -/// Errors on non-numeric scalar (e.g. Interval bounds — future work). -fn halo_from_bound(bound: &WindowFrameBound) -> datafusion::common::Result { +/// Extract the halo width in `f64` from a bound. `CurrentRow` → `Some(0.0)`. +/// Non-numeric scalars (e.g. Interval bounds) return `None` — a shape gate, +/// widened alongside KLL. +fn halo_from_bound(bound: &WindowFrameBound) -> Option { let scalar = match bound { - WindowFrameBound::CurrentRow => return Ok(0.0), + WindowFrameBound::CurrentRow => return Some(0.0), WindowFrameBound::Preceding(s) | WindowFrameBound::Following(s) => s, }; - // Widen anything Int-ish or Float-ish to f64. Interval bounds (for - // time-typed ORDER BYs) are the widening TODO alongside KLL. match scalar { - ScalarValue::Int8(Some(v)) => Ok(*v as f64), - ScalarValue::Int16(Some(v)) => Ok(*v as f64), - ScalarValue::Int32(Some(v)) => Ok(*v as f64), - ScalarValue::Int64(Some(v)) => Ok(*v as f64), - ScalarValue::UInt8(Some(v)) => Ok(*v as f64), - ScalarValue::UInt16(Some(v)) => Ok(*v as f64), - ScalarValue::UInt32(Some(v)) => Ok(*v as f64), - ScalarValue::UInt64(Some(v)) => Ok(*v as f64), - ScalarValue::Float32(Some(v)) => Ok(*v as f64), - ScalarValue::Float64(Some(v)) => Ok(*v), - other => datafusion::common::internal_err!( - "ParallelWindowRule: unsupported halo bound type {other:?}" - ), + ScalarValue::Int8(Some(v)) => Some(*v as f64), + ScalarValue::Int16(Some(v)) => Some(*v as f64), + ScalarValue::Int32(Some(v)) => Some(*v as f64), + ScalarValue::Int64(Some(v)) => Some(*v as f64), + ScalarValue::UInt8(Some(v)) => Some(*v as f64), + ScalarValue::UInt16(Some(v)) => Some(*v as f64), + ScalarValue::UInt32(Some(v)) => Some(*v as f64), + ScalarValue::UInt64(Some(v)) => Some(*v as f64), + ScalarValue::Float32(Some(v)) => Some(*v as f64), + ScalarValue::Float64(Some(v)) => Some(*v), + _ => None, } } -/// Wrap BWAG in a `PartitionedBoundedWindowAggExec` and splice -/// `SortExec → RSE#1 → source` below a fresh `ORRE → RSE#2 → RFE_wide` -/// chain. The rule runs after DF's optimizer chain, so BWAG's descendants -/// have the fully-materialized `SPM → SortExec → source` shape here — we -/// strip both since our rewrite overrides BWAG's distribution and takes -/// ownership of Sort placement (see module doc for why). -fn rewrite_bwag( - bwag: &Arc, - candidate: &WindowCandidate, +/// Match the parallel-window shape rooted at `node` and, if it fits, splice +/// `RFE_narrow → PBWAG(BWAG) → RFE_wide → RSE#2 → ORRE → SortExec → RSE#1 → +/// ` in place of the DF-planted `BWAG → SPM → SortExec → ` +/// subtree. +/// +/// - `Ok(None)`: shape gate missed (not a BWAG, PARTITION BY present, ROWS +/// frame, UNBOUNDED bound, non-Float64 ORDER BY, non-numeric bound scalar, +/// or subtree already rewritten). No log noise on the hot path. +/// - `Ok(Some(_))`: rewrite happened. +/// - `Err(_)`: an invariant the shape gates should have upheld didn't — BWAG +/// with ≠1 child, SPM/SortExec with ≠1 child, schema-lookup failure on the +/// ORDER BY expression, or a constructor `try_new` error. +/// +/// The rule runs after DF's optimizer chain, so BWAG's descendants have the +/// fully-materialized `SPM → SortExec → source` shape by the time we peel +/// here (see module doc for why the SPM and SortExec get stripped). +fn maybe_rewrite_bwag( + node: &Arc, output_partitions: usize, -) -> datafusion::common::Result> { - let bwag_children = bwag.children(); - let [immediate] = bwag_children.as_slice() else { +) -> datafusion::common::Result>> { + let Some(window) = node.downcast_ref::() else { + return Ok(None); + }; + // Shape gates as slice patterns: 0 or 2+ elements simply don't match. + let [expr] = window.window_expr() else { + return Ok(None); + }; + let [] = expr.partition_by() else { + return Ok(None); + }; + let [order] = expr.order_by() else { + return Ok(None); + }; + let Some(column) = order.expr.downcast_ref::() else { + return Ok(None); + }; + let frame = expr.get_window_frame(); + let WindowFrameUnits::Range = frame.units else { + return Ok(None); + }; + if !is_finite(&frame.start_bound) || !is_finite(&frame.end_bound) { + return Ok(None); + } + // Idempotency: re-plans (AQE fires the optimizer chain again for stage + // N+1) would otherwise wrap another ORRE around the previous rewrite's + // RangeFilterExec+ShuffleReader — and that ORRE's child doesn't claim + // ordering, blowing up at execute-time. + if subtree_contains_our_rewrite(window.children().as_slice()) { + return Ok(None); + } + let (Some(halo_lo), Some(halo_hi)) = ( + halo_from_bound(&frame.start_bound), + halo_from_bound(&frame.end_bound), + ) else { + return Ok(None); + }; + + let node_children = node.children(); + let [immediate] = node_children.as_slice() else { return datafusion::common::internal_err!( "ParallelWindowRule: BWAG must have exactly 1 child" ); }; - - // Strip whatever DF planted above the true source purely to satisfy - // BWAG's SinglePartition + Sorted requirements (SPM and SortExec) — - // our rewrite overrides both. Loop tolerates any order (SPM→Sort or - // Sort→SPM) or partial shapes (source that claims ordering natively - // via `sort_order_for_reorder` skips the Sort entirely). + // Loop tolerates any order (SPM→Sort or Sort→SPM) or partial shapes + // (source that claims ordering natively via `sort_order_for_reorder` + // skips the Sort entirely). let mut base_source: Arc = (*immediate).clone(); while base_source.is::() || base_source.is::() { let children = base_source.children(); @@ -305,16 +273,14 @@ fn rewrite_bwag( } let source_schema = base_source.schema(); - // Route on the ORDER BY column. ORRE requires Float64 today. - let routing_type = candidate.sort_expr.expr.data_type(&source_schema)?; + // Route on the ORDER BY column. ORRE requires Float64 today (T-Digest + // restriction; lifts when the sketch swaps to KLL). + let routing_type = order.expr.data_type(&source_schema)?; if !matches!(routing_type, DataType::Float64) { - return datafusion::common::internal_err!( - "ParallelWindowRule: routing expression `{}` must be Float64, got {routing_type:?}", - candidate.sort_expr.expr - ); + return Ok(None); } - let sort_expr = normalize_sort_expr(&candidate.sort_expr); + let sort_expr = normalize_sort_expr(order); let rse1: Arc = Arc::new(RuntimeStatsExec::try_new( base_source, Some(vec![sort_expr.clone()]), @@ -340,8 +306,6 @@ fn rewrite_bwag( orre, Some(vec![sort_expr.clone()]), )?); - let halo_lo = halo_from_bound(&candidate.start_bound)?; - let halo_hi = halo_from_bound(&candidate.end_bound)?; let wide_filter: Arc = Arc::new(RangeFilterExec::try_new_pending( rse2, sort_expr.expr.clone(), @@ -355,14 +319,9 @@ fn rewrite_bwag( // per-partition execute() runs each of the K sub-ranges independently. // See execution_plans::partitioned_bounded_window_agg for what makes // this safe (range-repartition upstream + halo). - let bwag_ref = bwag.downcast_ref::().ok_or_else(|| { - datafusion::common::DataFusionError::Internal( - "ParallelWindowRule: rewrite_bwag caller passed non-BWAG".into(), - ) - })?; let partitioned_bwag: Arc = Arc::new(PartitionedBoundedWindowAggExec::try_new( - bwag_ref.window_expr().to_vec(), + window.window_expr().to_vec(), wide_filter, )?); // Narrow filter above BWAG drops the halo rows the wide filter let in @@ -375,7 +334,14 @@ fn rewrite_bwag( ScalarValue::Float64(Some(0.0)), ScalarValue::Float64(Some(0.0)), )?); - Ok(narrow_filter) + + debug!( + "ParallelWindowRule: rewrote BWAG on `{}` (RANGE {} - {})", + column.name(), + fmt_bound(&frame.start_bound), + fmt_bound(&frame.end_bound), + ); + Ok(Some(narrow_filter)) } /// ORRE requires `nulls_first == false` today (T-Digest has no NULL slot). @@ -549,24 +515,22 @@ mod tests { #[test] fn halo_from_bound_reads_all_numeric_variants() { - assert_eq!(halo_from_bound(&WindowFrameBound::CurrentRow).unwrap(), 0.0); + assert_eq!(halo_from_bound(&WindowFrameBound::CurrentRow), Some(0.0)); assert_eq!( - halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Int64(Some(3)))) - .unwrap(), - 3.0 + halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Int64(Some(3)))), + Some(3.0) ); assert_eq!( halo_from_bound(&WindowFrameBound::Following(ScalarValue::Float64(Some( 2.5 - )))) - .unwrap(), - 2.5 + )))), + Some(2.5) ); - assert!( + assert_eq!( halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Utf8(Some( "x".into() - )))) - .is_err() + )))), + None ); } }