diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 99fceb91ea..85d62a9aa5 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -55,7 +55,9 @@ message BallistaPhysicalPlanNode { BufferExecNode buffer = 7; UnorderedRangeRepartitionExecNode unordered_range_repartition = 8; OrderedRangeRepartitionExecNode ordered_range_repartition = 9; - PerPartitionFilterExecNode per_partition_filter = 10; + RangeFilterExecNode range_filter = 10; + PartitionedBoundedWindowAggExecNode partitioned_bounded_window_agg = 11; + RangeShuffleReaderExecNode range_shuffle_reader = 12; } } @@ -125,12 +127,33 @@ message UnorderedRangeRepartitionExecNode { uint32 output_partitions = 2; } -// Filter with per-input-partition predicates. `predicates[k]` is the -// boolean expression applied to input partition `k`. Requires -// `predicates.len() == input_partition_count`. The child plan is -// plumbed by the framework as `inputs[0]` during decode. -message PerPartitionFilterExecNode { - repeated datafusion.PhysicalExprNode predicates = 1; +// Filter over ordered inputs with per-partition half-open range predicates +// derived from `cuts` + `halo_lo` / `halo_hi`. Zero halo recovers the exact +// range-repartition trim used above `ShuffleReaderExec`. Non-zero halo +// widens each partition's read range to include a boundary "context" band +// (bounded RANGE-frame windows). The child plan is plumbed by the framework +// as `inputs[0]` during decode. Serialization requires cuts to be resolved. +// +// `partition_indices` maps this operator's local partition index to the +// global partition index in the original K-shape defined by `cuts`. Under +// task-level restriction the input is sliced to a subset of the K global +// partitions; each entry stays < `cuts.len() + 1`. +message RangeFilterExecNode { + datafusion.PhysicalExprNode routing_expr = 1; + repeated datafusion_common.ScalarValue cuts = 2; + datafusion_common.ScalarValue halo_lo = 3; + datafusion_common.ScalarValue halo_hi = 4; + repeated uint32 partition_indices = 5; +} + +// Wrapper for `BoundedWindowAggExec` that overrides +// `required_input_distribution` to `Unspecified` — see the module doc on +// `execution_plans::partitioned_bounded_window_agg` for what makes that safe. +// The child plan is plumbed by the framework as `inputs[0]` during decode. +// `input_order_mode` and `can_repartition` are hardcoded on the decode side +// per the rule's shape gates; only `window_expr` needs to cross the wire. +message PartitionedBoundedWindowAggExecNode { + repeated datafusion.PhysicalWindowExprNode window_expr = 1; } message ChaosExecNode { @@ -190,6 +213,18 @@ message ShuffleReaderPartition { repeated PartitionLocation location = 1; } +// Ordering-preserving shuffle reader. Reuses `ShuffleReaderPartition` for the +// M-shape source layout. Partitioning is derived from `partition.len()` +// (always `UnknownPartitioning`, range-partitioned by `merge_ordering`). +message RangeShuffleReaderExecNode { + repeated ShuffleReaderPartition partition = 1; + datafusion_common.Schema schema = 2; + uint32 stage_id = 3; + // Sort key the reader's k-way merge preserves. Advertised on the reader's + // `PlanProperties.eq_properties` for downstream consumers. + repeated datafusion.PhysicalSortExprNode merge_ordering = 4; +} + // CoalescePartitionsRule output: groups upstream partitions into coalesced output partitions. // Empty when no coalesce is applied (the optional field on the parent message is absent). message CoalescePlan { diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index b3f41ee32b..8ad8b14945 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/mod.rs b/ballista/core/src/execution_plans/mod.rs index 5d9240f3a3..3df757e703 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -23,9 +23,11 @@ mod chaos_exec; mod distributed_explain_analyze; mod distributed_query; mod ordered_range_repartition; -mod per_partition_filter; +mod partitioned_bounded_window_agg; pub mod plan_algebra; +mod range_filter; mod range_repartition_common; +mod range_shuffle_reader; mod runtime_stats; mod shuffle_reader; mod shuffle_writer; @@ -42,8 +44,10 @@ use datafusion::common::exec_err; pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; pub use ordered_range_repartition::OrderedRangeRepartitionExec; -pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates}; +pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec; pub use plan_algebra::{preserves_distribution, preserves_partitioning}; +pub use range_filter::RangeFilterExec; +pub use range_shuffle_reader::RangeShuffleReaderExec; pub use runtime_stats::{ MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, collect_reports as collect_runtime_stats_reports, cut_partitions, diff --git a/ballista/core/src/execution_plans/ordered_range_repartition.rs b/ballista/core/src/execution_plans/ordered_range_repartition.rs index c040674da4..528b7df748 100644 --- a/ballista/core/src/execution_plans/ordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/ordered_range_repartition.rs @@ -176,25 +176,13 @@ impl OrderedRangeRepartitionExec { routing.expr ); } - // Input MUST claim to be sorted on our routing expression — otherwise - // the k-way merge produces garbled output. Sortedness of individual - // input partitions is enforced by the operator upstream (`SortExec` - // with `preserve_partitioning=true`); this check verifies the plan - // node declares that property. - let input_first_sort = input.output_ordering().map(|ordering| ordering.first()); - let Some(input_first) = input_first_sort else { - return internal_err!( - "OrderedRangeRepartitionExec requires sorted input — child plan claims no ordering" - ); - }; - if input_first.expr.as_ref() != routing.expr.as_ref() { - return internal_err!( - "OrderedRangeRepartitionExec: input's first sort key `{}` does not match \ - routing expression `{}`", - input_first.expr, - routing.expr - ); - } + // NB: input sortedness is NOT checked at try_new. The k-way merge in + // execute() needs each input partition sorted on the routing key, but + // that guarantee comes from `required_input_ordering()` below + + // `EnforceSorting`: DataFusion inserts a `SortExec` above the source + // when the declared requirement isn't satisfied. Checking at try_new + // races with rule-time construction (rules run before EnforceSorting), + // so the runtime check has moved to `execute()`. // Advertise each output partition as sorted on `order_by`. Downstream // operators (BWAG, HaloDrop) rely on this claim to skip redundant // Sort insertions. @@ -356,6 +344,30 @@ impl ExecutionPlan for OrderedRangeRepartitionExec { partition: usize, ctx: Arc, ) -> Result { + // Invariant: EnforceSorting must have satisfied our + // `required_input_ordering` — the k-way merge assumes each input + // partition is sorted on the routing key. Rules that emit ORRE run + // before EnforceSorting, so this check lives at execute-time rather + // than construction-time. + let input_first = self + .input + .output_ordering() + .map(|ordering| ordering.first()) + .ok_or_else(|| { + internal_datafusion_err!( + "OrderedRangeRepartitionExec: input claims no ordering at execute — \ + EnforceSorting should have planted a SortExec" + ) + })?; + let routing = &self.order_by[0]; + if input_first.expr.as_ref() != routing.expr.as_ref() { + return internal_err!( + "OrderedRangeRepartitionExec: input's first sort key `{}` does not \ + match routing expression `{}`", + input_first.expr, + routing.expr + ); + } let mut state = self .state .lock() @@ -686,17 +698,24 @@ mod tests { } #[test] - fn try_new_rejects_unsorted_input() { + fn execute_rejects_unsorted_input() { + // try_new no longer checks input ordering — EnforceSorting is + // trusted to plant a SortExec after rule-time construction. If + // the invariant is broken by the time we get to execute(), the + // runtime check fires. let schema = schema_v2_id(); - // MemorySourceConfig with no declared ordering — output_ordering() is None. - let err = OrderedRangeRepartitionExec::try_new( + let orre = OrderedRangeRepartitionExec::try_new( empty_input(&schema), vec![asc(&schema, "v2")], 4, ) - .expect_err("input without ordering claim must be rejected"); + .expect("construction succeeds; check moved to execute()"); + let ctx = datafusion::prelude::SessionContext::new().task_ctx(); + let Err(err) = orre.execute(0, ctx) else { + panic!("execute() should reject input without ordering claim"); + }; assert!( - err.to_string().contains("child plan claims no ordering"), + err.to_string().contains("input claims no ordering"), "got: {err}" ); } @@ -720,20 +739,24 @@ mod tests { } #[test] - fn try_new_rejects_mismatched_sort_key() { + fn execute_rejects_mismatched_sort_key() { let schema = schema_v2_id(); - // Sort input on `id` (Int64); DRR tries to route on `v2`. + // Sort input on `id` (Int64); ORRE tries to route on `v2`. let source = empty_input(&schema); let id_sort = LexOrdering::new(vec![asc(&schema, "id")]).unwrap(); let sorted_on_id = Arc::new(SortExec::new(id_sort, source).with_preserve_partitioning(true)) as Arc; - let err = OrderedRangeRepartitionExec::try_new( + let orre = OrderedRangeRepartitionExec::try_new( sorted_on_id, vec![asc(&schema, "v2")], 4, ) - .expect_err("mismatched sort key must be rejected"); + .expect("construction succeeds; check moved to execute()"); + let ctx = datafusion::prelude::SessionContext::new().task_ctx(); + let Err(err) = orre.execute(0, ctx) else { + panic!("execute() should reject mismatched sort key"); + }; assert!( err.to_string().contains("does not match routing"), "got: {err}" diff --git a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs new file mode 100644 index 0000000000..e686d5fad6 --- /dev/null +++ b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs @@ -0,0 +1,192 @@ +// 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. + +//! Wrap a `BoundedWindowAggExec` and override its +//! `required_input_distribution` to `Unspecified`. +//! +//! DataFusion's `BoundedWindowAggExec` declares +//! `Distribution::SinglePartition` when no PARTITION BY is present — a +//! correctness guard because a window frame's semantics span rows across +//! the whole input. With `ParallelWindowRule`'s range-repartition upstream, +//! each input partition IS globally range-disjoint (halo covers boundary +//! neighbours), so BWAG CAN run per-partition and produce K correct outputs. +//! +//! This wrapper flips only the distribution declaration. Everything else — +//! schema, ordering, per-partition `execute()` — delegates to a canonical +//! inner BWAG constructed at rule time. Because `children()` returns only +//! the wrapper's own input (not the inner BWAG), tree walkers like +//! `EnforceDistribution` never see BWAG and can't re-insert an +//! `SPM(K→1)` beneath it. +//! +//! # Constraints assumed by the caller +//! +//! The wrapper is safe iff the input is already range-repartitioned so that +//! each partition is a globally disjoint slice of the ORDER BY key + halo +//! for frame boundaries. Callers are responsible for this — the wrapper +//! itself doesn't (and can't) verify it. Wiring this wrapper below arbitrary +//! (non-range-partitioned) inputs will silently produce wrong window values. +//! +//! # Assumption on DataFusion internals +//! +//! `BoundedWindowAggExec::execute(i, ctx)` in DataFusion 54 processes +//! partition `i` of its input independently: no cross-partition state, no +//! spawned tasks touching sibling partitions. The wrapper depends on that +//! shape. If DataFusion ever changes BWAG to have cross-partition state, +//! the change would be visible in `BoundedWindowAggExec::execute`. + +use std::fmt; +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, Statistics, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{Distribution, OrderingRequirements}; +use datafusion::physical_plan::execution_plan::{CardinalityEffect, InputOrderMode}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::windows::BoundedWindowAggExec; +use datafusion::physical_plan::windows::WindowExpr; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + 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. +const BWAG_INPUT_ORDER_MODE: InputOrderMode = InputOrderMode::Sorted; +const BWAG_CAN_REPARTITION: bool = false; + +/// Wrap a `BoundedWindowAggExec` overriding its `required_input_distribution` +/// to `Unspecified`. See module docs for what makes this safe. +#[derive(Debug, Clone)] +pub struct PartitionedBoundedWindowAggExec { + /// Canonical inner BWAG built at construction. Not a plan-tree child — + /// `children()` returns only [`Self::input`], so tree walkers can't + /// reach it. + inner_bwag: Arc, + /// Multi-partition input; same `Arc` `inner_bwag.input()` holds. + input: Arc, +} + +impl PartitionedBoundedWindowAggExec { + /// Construct the wrapper. Builds a canonical inner + /// `BoundedWindowAggExec` from `window_expr` + `input` with the + /// hardcoded mode/repartition constants; failures propagate verbatim so + /// callers see the same error surface as constructing BWAG directly. + pub fn try_new( + window_expr: Vec>, + input: Arc, + ) -> Result { + let inner_bwag = Arc::new(BoundedWindowAggExec::try_new( + window_expr, + input.clone(), + BWAG_INPUT_ORDER_MODE, + BWAG_CAN_REPARTITION, + )?); + Ok(Self { inner_bwag, input }) + } + + /// The wrapped `BoundedWindowAggExec` — for accessors that don't exist + /// on `Self` and for wire-encoding. + pub fn inner_bwag(&self) -> &Arc { + &self.inner_bwag + } + + /// The window expressions carried by the wrapped BWAG. + pub fn window_expr(&self) -> &[Arc] { + self.inner_bwag.window_expr() + } + + /// Output schema (delegates to BWAG). + pub fn schema(&self) -> SchemaRef { + self.inner_bwag.schema() + } +} + +impl DisplayAs for PartitionedBoundedWindowAggExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PartitionedBoundedWindowAggExec: ")?; + self.inner_bwag.fmt_as(t, f) + } +} + +impl ExecutionPlan for PartitionedBoundedWindowAggExec { + fn name(&self) -> &'static str { + "PartitionedBoundedWindowAggExec" + } + + fn properties(&self) -> &Arc { + // BWAG's properties already advertise `input.output_partitioning()` + // as this operator's output partitioning — reuse them verbatim. + self.inner_bwag.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn required_input_distribution(&self) -> Vec { + // The whole point of this wrapper. + vec![Distribution::UnspecifiedDistribution] + } + + fn required_input_ordering(&self) -> Vec> { + self.inner_bwag.required_input_ordering() + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [new_input] = children.as_slice() else { + return internal_err!( + "PartitionedBoundedWindowAggExec expects exactly 1 child, got {}", + children.len() + ); + }; + Ok(Arc::new(Self::try_new( + self.window_expr().to_vec(), + new_input.clone(), + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner_bwag.execute(partition, context) + } + + fn metrics(&self) -> Option { + self.inner_bwag.metrics() + } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.inner_bwag.partition_statistics(partition) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + self.inner_bwag.cardinality_effect() + } +} diff --git a/ballista/core/src/execution_plans/per_partition_filter.rs b/ballista/core/src/execution_plans/per_partition_filter.rs deleted file mode 100644 index 9c93a7a6d1..0000000000 --- a/ballista/core/src/execution_plans/per_partition_filter.rs +++ /dev/null @@ -1,571 +0,0 @@ -// 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. - -//! Filter with a distinct predicate per input partition. -//! -//! `FilterExec` in DataFusion carries a single predicate applied to every -//! partition. That's wrong for the range-repartition-consuming reader in -//! the adaptive range shuffle: each downstream partition `k` needs a range -//! predicate `cuts[k-1] <= key < cuts[k]` unique to that partition, so -//! straddling sub-parts from the producer are trimmed to just partition -//! `k`'s slice. -//! -//! One-task-per-downstream-partition + plain `FilterExec` would work but -//! defeats vcore packing (`K` tasks instead of `K / vcores`). This operator -//! keeps packing: `predicates[k]` is applied to `input.execute(k)`, so a -//! single task consuming several partitions still gets each partition's -//! own predicate. -//! -//! Semantics per batch mirror `FilterExec`: evaluate the boolean expr -//! against the batch, then `filter_record_batch`. No projection, no -//! coalescing, no metrics — those can grow later if the wiring warrants. - -use std::fmt::{self, Debug, Formatter}; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::compute::filter_record_batch; -use datafusion::arrow::datatypes::{DataType, SchemaRef}; -use datafusion::common::cast::as_boolean_array; -use datafusion::common::{Result, Statistics, internal_err}; -use datafusion::execution::TaskContext; -use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; -use datafusion::physical_plan::execution_plan::CardinalityEffect; -use datafusion::physical_plan::stream::{ - EmptyRecordBatchStream, RecordBatchStreamAdapter, -}; -use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, -}; -use futures::{Stream, StreamExt, ready}; - -/// Filter with per-input-partition predicates. -/// -/// `predicates[k]` is applied to `input.execute(k)`. Requires -/// `predicates.len() == input.output_partitioning().partition_count()`. -pub struct PerPartitionFilterExec { - input: Arc, - predicates: Vec>, - properties: Arc, -} - -impl PerPartitionFilterExec { - /// Wrap `input` with a vector of predicates, one per input partition. - /// - /// Fails if the predicate count doesn't match the input partition count - /// or if any predicate does not evaluate to `Boolean` against the input - /// schema. - pub fn try_new( - input: Arc, - predicates: Vec>, - ) -> Result { - let partition_count = input.output_partitioning().partition_count(); - if predicates.len() != partition_count { - return internal_err!( - "PerPartitionFilterExec: predicate count {} does not match input partition count {}", - predicates.len(), - partition_count - ); - } - let schema = input.schema(); - for (k, predicate) in predicates.iter().enumerate() { - let dt = predicate.data_type(&schema)?; - if dt != DataType::Boolean { - return internal_err!( - "PerPartitionFilterExec: predicate[{k}] must evaluate to Boolean, got {dt}" - ); - } - } - let properties = Arc::new(PlanProperties::new( - input.equivalence_properties().clone(), - input.output_partitioning().clone(), - input.pipeline_behavior(), - input.boundedness(), - )); - Ok(Self { - input, - predicates, - properties, - }) - } - - /// The per-partition predicates. `predicates()[k]` corresponds to input partition `k`. - pub fn predicates(&self) -> &[Arc] { - &self.predicates - } -} - -impl Debug for PerPartitionFilterExec { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - f.debug_struct("PerPartitionFilterExec") - .field("num_predicates", &self.predicates.len()) - .finish() - } -} - -impl DisplayAs for PerPartitionFilterExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!( - f, - "PerPartitionFilterExec: predicates=[{}]", - self.predicates - .iter() - .map(|p| p.to_string()) - .collect::>() - .join(", ") - ) - } - DisplayFormatType::TreeRender => { - write!(f, "PerPartitionFilterExec") - } - } - } -} - -impl ExecutionPlan for PerPartitionFilterExec { - fn name(&self) -> &str { - "PerPartitionFilterExec" - } - - fn schema(&self) -> SchemaRef { - self.input.schema() - } - - fn properties(&self) -> &Arc { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - let [input] = children.as_slice() else { - return internal_err!( - "PerPartitionFilterExec expects exactly one child, got {}", - children.len() - ); - }; - Ok(Arc::new(PerPartitionFilterExec::try_new( - input.clone(), - self.predicates.clone(), - )?)) - } - - fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] - } - - fn required_input_ordering(&self) -> Vec> { - vec![None] - } - - fn maintains_input_order(&self) -> Vec { - vec![true] - } - - fn benefits_from_input_partitioning(&self) -> Vec { - vec![false] - } - - fn partition_statistics(&self, _partition: Option) -> Result> { - Ok(Arc::new(Statistics::new_unknown(&self.schema()))) - } - - fn cardinality_effect(&self) -> CardinalityEffect { - CardinalityEffect::LowerEqual - } - - fn execute( - &self, - partition: usize, - ctx: Arc, - ) -> Result { - let Some(predicate) = self.predicates.get(partition).cloned() else { - return internal_err!( - "PerPartitionFilterExec: partition {} out of bounds ({} predicates)", - partition, - self.predicates.len() - ); - }; - let schema = self.schema(); - let input = self.input.execute(partition, ctx)?; - let stream = PerPartitionFilterStream { - schema: schema.clone(), - predicate, - input, - }; - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) - } -} - -struct PerPartitionFilterStream { - schema: SchemaRef, - predicate: Arc, - input: SendableRecordBatchStream, -} - -impl Stream for PerPartitionFilterStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - loop { - match ready!(self.input.poll_next_unpin(cx)) { - Some(Ok(batch)) => { - let mask = self - .predicate - .evaluate(&batch) - .and_then(|v| v.into_array(batch.num_rows()))?; - let mask = as_boolean_array(&mask)?; - let filtered = filter_record_batch(&batch, mask)?; - if filtered.num_rows() == 0 { - // Nothing left after filtering; pull the next batch - // rather than emit an empty batch downstream. - continue; - } - return Poll::Ready(Some(Ok(filtered))); - } - Some(Err(e)) => return Poll::Ready(Some(Err(e))), - None => { - // Release the input pipeline's resources on EOS — - // mirrors DataFusion's FilterExec so the input's - // child chain doesn't linger on the heap until the - // outer stream is itself dropped. - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - return Poll::Ready(None); - } - } - } - } -} - -impl RecordBatchStream for PerPartitionFilterStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} - -/// Build the `K = cuts.len() + 1` half-open range predicates a -/// `PerPartitionFilterExec` needs to reproduce the range repartition's -/// write-side routing on the read side. -/// -/// Partition `i` receives the predicate -/// -/// ```text -/// i = 0 → routing_expr < cuts[0] -/// 0 < i < K-1 → cuts[i-1] <= routing_expr AND routing_expr < cuts[i] -/// i = K-1 → routing_expr >= cuts[K-2] -/// K = 1 → lit(true) // empty cuts, single-bucket range repartition -/// ``` -/// -/// Consistent with the private `range_repartition_common::split_batch_by_range` -/// helper, which uses the same half-open convention on the write side. -/// Callers pass the range repartition's routing expression verbatim -/// (`CAST(order_by[0] AS Float64)` today). -/// -/// Non-null routing expressions only. Both `UnorderedRangeRepartitionExec` -/// and `OrderedRangeRepartitionExec` refuse nullable routing exprs at -/// `try_new`, so any expression that reaches this helper via -/// `RangeRepartitionRouting` is guaranteed non-null — no `IS NULL` branch -/// needed. -pub fn range_partition_predicates( - routing_expr: Arc, - cuts: &[f64], -) -> Vec> { - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; - use datafusion::scalar::ScalarValue; - - let partition_count = cuts.len() + 1; - let lit = |v: f64| -> Arc { - Arc::new(Literal::new(ScalarValue::Float64(Some(v)))) - }; - let ge = |lo: f64| -> Arc { - Arc::new(BinaryExpr::new( - routing_expr.clone(), - Operator::GtEq, - lit(lo), - )) - }; - let lt = |hi: f64| -> Arc { - Arc::new(BinaryExpr::new(routing_expr.clone(), Operator::Lt, lit(hi))) - }; - (0..partition_count) - .map(|partition_idx| { - let lo = partition_idx - .checked_sub(1) - .and_then(|cut_idx| cuts.get(cut_idx).copied()); - let hi = cuts.get(partition_idx).copied(); - match (lo, hi) { - (None, None) => { - // K == 1: single bucket covers everything. - Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))) - as Arc - } - (None, Some(hi)) => lt(hi), - (Some(lo), None) => ge(lo), - (Some(lo), Some(hi)) => { - Arc::new(BinaryExpr::new(ge(lo), Operator::And, lt(hi))) - } - } - }) - .collect() -} - -#[cfg(test)] -mod tests { - use super::*; - use datafusion::arrow::array::Int64Array; - use datafusion::arrow::datatypes::{Field, Schema}; - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::physical_plan::repartition::RepartitionExec; - use datafusion::physical_plan::{ExecutionPlan, Partitioning}; - use datafusion::prelude::SessionContext; - use datafusion::scalar::ScalarValue; - use futures::TryStreamExt; - - fn one_col_schema() -> SchemaRef { - Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])) - } - - /// Memory source with `partitions` partitions, each carrying a single - /// batch of `[start .. start + rows_per)` where `start = k * rows_per`. - fn partitioned_source(partitions: usize, rows_per: usize) -> Arc { - let schema = one_col_schema(); - let mut per_partition: Vec> = Vec::with_capacity(partitions); - for k in 0..partitions { - let start = (k * rows_per) as i64; - let arr = Int64Array::from_iter_values(start..start + rows_per as i64); - let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(); - per_partition.push(vec![batch]); - } - let src = - MemorySourceConfig::try_new(&per_partition, schema, None).expect("mem src"); - Arc::new(DataSourceExec::new(Arc::new(src))) - } - - /// Predicate `v >= lo AND v < hi` against column `v`. - fn range_pred(lo: i64, hi: i64) -> Arc { - let col = Arc::new(Column::new("v", 0)); - let lo_lit = Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))); - let hi_lit = Arc::new(Literal::new(ScalarValue::Int64(Some(hi)))); - let ge: Arc = - Arc::new(BinaryExpr::new(col.clone(), Operator::GtEq, lo_lit)); - let lt: Arc = - Arc::new(BinaryExpr::new(col, Operator::Lt, hi_lit)); - Arc::new(BinaryExpr::new(ge, Operator::And, lt)) - } - - fn ctx() -> Arc { - SessionContext::new().task_ctx() - } - - async fn collect(plan: Arc, partition: usize) -> Result> { - let stream = plan.execute(partition, ctx())?; - let batches: Vec = stream.try_collect().await?; - let mut out = Vec::new(); - for b in batches { - let arr = b - .column(0) - .as_any() - .downcast_ref::() - .expect("Int64Array"); - out.extend(arr.iter().map(|v| v.unwrap())); - } - Ok(out) - } - - /// Each of three partitions carries `[k*100, k*100+100)`. With - /// per-partition predicates that each carve a five-row slice, every - /// partition emits its own five rows and nothing from another - /// partition leaks through. - #[tokio::test] - async fn per_partition_predicate_filters_only_that_partition() -> Result<()> { - let src = partitioned_source(3, 100); - let predicates = vec![ - range_pred(0, 5), // partition 0 → 0..5 - range_pred(105, 110), // partition 1 → 105..110 - range_pred(295, 300), // partition 2 → 295..300 - ]; - let ppf: Arc = - Arc::new(PerPartitionFilterExec::try_new(src, predicates)?); - assert_eq!(collect(ppf.clone(), 0).await?, (0..5).collect::>()); - assert_eq!( - collect(ppf.clone(), 1).await?, - (105..110).collect::>() - ); - assert_eq!(collect(ppf, 2).await?, (295..300).collect::>()); - Ok(()) - } - - /// A predicate that matches nothing yields an empty stream (no zero-row - /// batches surfaced to the caller). Regression pin — an earlier draft - /// forwarded empty batches, which some downstream operators dislike. - #[tokio::test] - async fn empty_predicate_yields_empty_stream() -> Result<()> { - let src = partitioned_source(1, 100); - let predicates = vec![range_pred(1_000_000, 2_000_000)]; - let ppf: Arc = - Arc::new(PerPartitionFilterExec::try_new(src, predicates)?); - assert_eq!(collect(ppf, 0).await?, Vec::::new()); - Ok(()) - } - - /// Predicate-count mismatch is rejected at construction time. - #[test] - fn rejects_predicate_count_mismatch() { - let src = partitioned_source(3, 10); - let err = - PerPartitionFilterExec::try_new(src, vec![range_pred(0, 5)]).unwrap_err(); - let msg = err.to_string(); - assert!( - msg.contains("predicate count 1") && msg.contains("input partition count 3"), - "unexpected error: {msg}" - ); - } - - /// Non-boolean predicate is rejected at construction time (the - /// expression must evaluate to `Boolean` against the input schema). - #[test] - fn rejects_non_boolean_predicate() { - let src = partitioned_source(1, 10); - // Just the column `v` — evaluates to Int64, not Boolean. - let bad: Arc = Arc::new(Column::new("v", 0)); - let err = PerPartitionFilterExec::try_new(src, vec![bad]).unwrap_err(); - assert!( - err.to_string().contains("Boolean"), - "unexpected error: {err}" - ); - } - - /// The K=4 range predicates cover every value under the half-open - /// convention, and each row lands in exactly one predicate. Random - /// probe values are routed through the predicates and expected to - /// match the same partition assignment as the range repartition's - /// write-side `split_batch_by_range` would produce. - #[test] - fn range_partition_predicates_partition_every_value_exactly_once() { - use datafusion::arrow::array::Float64Array; - use datafusion::arrow::datatypes::Field; - use datafusion::physical_expr::expressions::Column; - - let cuts = vec![10.0, 20.0, 30.0]; - let k = cuts.len() + 1; - let routing: Arc = Arc::new(Column::new("v", 0)); - let preds = range_partition_predicates(routing, &cuts); - assert_eq!(preds.len(), k); - - let schema = - Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); - let values: Vec = vec![ - -5.0, 0.0, 9.999, 10.0, 15.0, 19.999, 20.0, 25.0, 30.0, 100.0, - ]; - let arr = Float64Array::from_iter_values(values.iter().copied()); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(); - - // For each row, find the unique partition whose predicate accepts it. - for (row, &v) in values.iter().enumerate() { - let mut hits = 0; - for pred in &preds { - let mask = pred - .evaluate(&batch) - .and_then(|v| v.into_array(batch.num_rows())) - .unwrap(); - let mask = as_boolean_array(&mask).unwrap(); - if mask.value(row) { - hits += 1; - } - } - assert_eq!( - hits, 1, - "value {v} matched {hits} predicates, expected exactly 1" - ); - } - - // Expected assignment mirrors split_batch_by_range: `partition_point` - // returns the count of cuts `<= key`, which is the partition index - // under the half-open convention. - let expected: Vec = values - .iter() - .map(|v| cuts.partition_point(|&c| c <= *v)) - .collect(); - for (row, want) in expected.iter().enumerate() { - let mask = preds[*want] - .evaluate(&batch) - .and_then(|v| v.into_array(batch.num_rows())) - .unwrap(); - let mask = as_boolean_array(&mask).unwrap(); - assert!( - mask.value(row), - "value {} should have landed in partition {}", - values[row], - want - ); - } - } - - /// Degenerate K=1 (empty cuts) yields a single lit(true) predicate. - #[test] - fn range_partition_predicates_single_bucket_when_cuts_empty() { - use datafusion::physical_expr::expressions::Column; - - let routing: Arc = Arc::new(Column::new("v", 0)); - let preds = range_partition_predicates(routing, &[]); - assert_eq!(preds.len(), 1); - assert_eq!(preds[0].to_string(), "true"); - } - - /// `with_new_children` swaps the input while preserving the predicate - /// vector. Wrapping the original source in a `RepartitionExec` that - /// keeps the partition count (RoundRobin(3)) gives a valid child; the - /// filter still routes partition-`k` rows through `predicates[k]`. - #[tokio::test] - async fn with_new_children_preserves_predicates() -> Result<()> { - let src = partitioned_source(3, 100); - let predicates = - vec![range_pred(0, 3), range_pred(100, 103), range_pred(200, 203)]; - let ppf = Arc::new(PerPartitionFilterExec::try_new( - src.clone(), - predicates.clone(), - )?); - // Wrap the source in RoundRobin(3) — same partition count, different plan. - let repart: Arc = Arc::new(RepartitionExec::try_new( - src, - Partitioning::RoundRobinBatch(3), - )?); - let swapped: Arc = ppf.with_new_children(vec![repart])?; - // Just verify construction succeeded and the operator name survives. - assert_eq!(swapped.name(), "PerPartitionFilterExec"); - Ok(()) - } -} diff --git a/ballista/core/src/execution_plans/range_filter.rs b/ballista/core/src/execution_plans/range_filter.rs new file mode 100644 index 0000000000..8fb128d101 --- /dev/null +++ b/ballista/core/src/execution_plans/range_filter.rs @@ -0,0 +1,1151 @@ +// 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. + +//! Filter over ordered inputs with per-partition half-open range predicates. +//! +//! `execute(k)` applies the predicate +//! +//! ```text +//! cuts[k-1] - halo_lo <= routing_expr < cuts[k] + halo_hi +//! ``` +//! +//! to partition `k`, with virtual `-∞` / `+∞` sentinels on the ends. Zero +//! halo (`halo_lo == halo_hi == 0.0`) recovers the exact range-repartition +//! trim used above `ShuffleReaderExec` for hash-agg correctness; non-zero +//! halo widens each partition's read range to include a boundary "context" +//! band (`WindowFrame` PRECEDING/FOLLOWING for bounded RANGE frames — see +//! [[parallel-range-window]]). +//! +//! Ordering knowledge on the input opens the door to a future value-index +//! binary-search path (aligned with the [`ValueIndexReader`] direction from +//! PR #2204) that a generic `FilterExec` couldn't do — the predicate here is +//! monotone over `routing_expr` and the input is sorted on it, so the +//! partition's slice is a contiguous run in the input. +//! +//! # Late-binding cuts +//! +//! `cuts` is `Arc>>>` — the ParallelWindow rewrite rule +//! plants a `RangeFilterExec` at plan time, well before the runtime cuts are +//! known. The scheduler calls [`resolve_cuts`] after stage 0's +//! `RuntimeStatsExec` reports have been merged (mirrors +//! `ExchangeExec::resolve_range_repartition_routing`). `execute` refuses +//! while cuts are unresolved; serialization refuses too — over-the-wire +//! plans always ship with cuts bound. +//! +//! # Type generality +//! +//! Cuts and halo widths are `f64` today. This matches URRE/ORRE's Float64 +//! hardcode (T-Digest is Float64-only). Widening to other `Ord` +//! `ScalarValue` types is a KLL-migration follow-up ([[kll-sketch]]) — +//! see the type-generality note in [[parallel-range-window]]. +//! +//! [`ValueIndexReader`]: crate::execution_plans::ShuffleReaderExec + +use std::fmt::{self, Debug, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::{Array, RecordBatch}; +use datafusion::arrow::compute::filter_record_batch; +use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::cast::{as_boolean_array, as_float64_array}; +use datafusion::common::{Result, Statistics, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; +use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::stream::{ + EmptyRecordBatchStream, RecordBatchStreamAdapter, +}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, +}; +use datafusion::scalar::ScalarValue; +use futures::{Stream, StreamExt, ready}; +use parking_lot::Mutex; + +/// Filter over an ordered input, applying a per-partition half-open range +/// predicate derived from runtime-discovered cuts + optional halo widths. +/// +/// `partition_indices` maps this operator's local partition index (what +/// `execute(k)` receives) to the *global* partition index in the original +/// K-shape defined by `cuts`. Under [`Self::restrict_partitions`] the input +/// is sliced to a subset of the original K partitions; the RangeFilterExec +/// then reports the restricted count via `output_partitioning`, but still +/// applies the correct global predicate for each restricted slice. +/// +/// Invariant: `partition_indices.len() == input.output_partitioning().partition_count()`, +/// and when `cuts` is `Some`, every entry is `< cuts.len() + 1`. +pub struct RangeFilterExec { + input: Arc, + routing_expr: Arc, + /// Late-bound. `None` until the scheduler calls [`resolve_cuts`]; + /// `execute` and serialization both refuse while unresolved. + cuts: Arc>>>, + halo_lo: f64, + halo_hi: f64, + partition_indices: Vec, + /// True when `input.output_ordering()` leads with `routing_expr` in + /// ascending order. Enables the min/max fast paths + binary-search slice + /// in [`RangeFilterStream`] — with a sorted input, per-batch first/last + /// values bound the entire batch, so most batches never touch + /// `filter_record_batch`. + sorted_on_key: bool, + properties: Arc, +} + +impl RangeFilterExec { + /// Construct with cuts already resolved (adapter path — cuts arrive + /// from `ExchangeExec::range_repartition_routing()` at Stage-2 + /// planning time). + pub fn try_new_resolved( + input: Arc, + routing_expr: Arc, + cuts: Vec, + halo_lo: ScalarValue, + halo_hi: ScalarValue, + ) -> Result { + Self::try_new_inner(input, routing_expr, Some(cuts), halo_lo, halo_hi, None) + } + + /// Construct with an explicit `partition_indices` mapping. Wire-decoding + /// path — restrict_partitions is preferred at plan-rewrite time. + pub fn try_new_with_indices( + input: Arc, + routing_expr: Arc, + cuts: Vec, + halo_lo: ScalarValue, + halo_hi: ScalarValue, + partition_indices: Vec, + ) -> Result { + Self::try_new_inner( + input, + routing_expr, + Some(cuts), + halo_lo, + halo_hi, + Some(partition_indices), + ) + } + + /// Construct with cuts pending (rule path — the ParallelWindow rewrite + /// plants the operator at plan time; the scheduler resolves cuts after + /// stage 0's stats reports merge). + pub fn try_new_pending( + input: Arc, + routing_expr: Arc, + halo_lo: ScalarValue, + halo_hi: ScalarValue, + ) -> Result { + Self::try_new_inner(input, routing_expr, None, halo_lo, halo_hi, None) + } + + fn try_new_inner( + input: Arc, + routing_expr: Arc, + cuts: Option>, + halo_lo: ScalarValue, + halo_hi: ScalarValue, + partition_indices: Option>, + ) -> Result { + let schema = input.schema(); + let expr_type = routing_expr.data_type(&schema)?; + if !expr_type.is_numeric() { + return internal_err!( + "RangeFilterExec: routing_expr must be numeric, got {expr_type}" + ); + } + let partition_count = input.output_partitioning().partition_count(); + let partition_indices = + partition_indices.unwrap_or_else(|| (0..partition_count).collect()); + if partition_indices.len() != partition_count { + return internal_err!( + "RangeFilterExec: partition_indices.len() ({}) does not match input partition count ({})", + partition_indices.len(), + partition_count + ); + } + let cuts_f64 = cuts + .as_ref() + .map(|c| c.iter().map(as_f64).collect::>>()) + .transpose()?; + if let Some(cuts_f64) = &cuts_f64 { + let global_count = cuts_f64.len() + 1; + for &idx in &partition_indices { + if idx >= global_count { + return internal_err!( + "RangeFilterExec: partition_indices contains {idx} but only {global_count} global partitions exist" + ); + } + } + if !cuts_f64.windows(2).all(|w| w[0] <= w[1]) { + return internal_err!("RangeFilterExec: cuts must be monotone"); + } + } + let halo_lo_f64 = as_f64(&halo_lo)?; + let halo_hi_f64 = as_f64(&halo_hi)?; + if !halo_lo_f64.is_finite() || halo_lo_f64 < 0.0 { + return internal_err!( + "RangeFilterExec: halo_lo must be finite and non-negative, got {halo_lo_f64}" + ); + } + if !halo_hi_f64.is_finite() || halo_hi_f64 < 0.0 { + return internal_err!( + "RangeFilterExec: halo_hi must be finite and non-negative, got {halo_hi_f64}" + ); + } + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + let sorted_on_key = input + .output_ordering() + .map(|ord| { + let first = ord.first(); + first.expr.as_ref() == routing_expr.as_ref() && !first.options.descending + }) + .unwrap_or(false); + Ok(Self { + input, + routing_expr, + cuts: Arc::new(Mutex::new(cuts_f64)), + halo_lo: halo_lo_f64, + halo_hi: halo_hi_f64, + partition_indices, + sorted_on_key, + properties, + }) + } + + /// Build a restricted RangeFilterExec: the same routing_expr / cuts / + /// halo, but sliced to a subset of the input's partitions. + /// `restricted_input` must already have been restricted to the same + /// `task_partitions` by the caller; the RangeFilterExec's job here is + /// only to remap `partition_indices` so `execute(local_k)` still finds + /// the right global cut range. + pub fn restrict_partitions( + &self, + restricted_input: Arc, + task_partitions: &[usize], + ) -> Result { + let new_indices: Vec = task_partitions + .iter() + .map(|&local_j| self.partition_indices[local_j]) + .collect(); + let cuts_snapshot = self.cuts.lock().clone().map(|c| { + c.into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect() + }); + Self::try_new_inner( + restricted_input, + self.routing_expr.clone(), + cuts_snapshot, + ScalarValue::Float64(Some(self.halo_lo)), + ScalarValue::Float64(Some(self.halo_hi)), + Some(new_indices), + ) + } + + /// Idempotent overwrite. Matches `ExchangeExec::resolve_range_repartition_routing` + /// — called by the scheduler once stage-0 sketches are merged into cuts. + /// `cuts` describes the *global* K-shape, so `cuts.len() + 1` must be at + /// least `max(partition_indices) + 1`. Under restriction the operator's + /// local partition count can be smaller than K. + pub fn resolve_cuts(&self, cuts: Vec) -> Result<()> { + let cuts_f64: Vec = cuts.iter().map(as_f64).collect::>()?; + let global_count = cuts_f64.len() + 1; + if let Some(&max_idx) = self.partition_indices.iter().max() + && max_idx >= global_count + { + return internal_err!( + "RangeFilterExec::resolve_cuts: cuts describe {global_count} global partitions but partition_indices references {max_idx}" + ); + } + if !cuts_f64.windows(2).all(|w| w[0] <= w[1]) { + return internal_err!("RangeFilterExec::resolve_cuts: cuts must be monotone"); + } + self.cuts.lock().replace(cuts_f64); + Ok(()) + } + + /// Snapshot cuts. `None` before [`resolve_cuts`] fires; `Some` after. + pub fn cuts(&self) -> Option> { + self.cuts.lock().clone().map(|c| { + c.into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect() + }) + } + + /// The physical expression whose value each row is bucketed by. + pub fn routing_expr(&self) -> &Arc { + &self.routing_expr + } + + /// Halo-widening amount applied to each partition's lower bound. + pub fn halo_lo(&self) -> ScalarValue { + ScalarValue::Float64(Some(self.halo_lo)) + } + + /// Halo-widening amount applied to each partition's upper bound. + pub fn halo_hi(&self) -> ScalarValue { + ScalarValue::Float64(Some(self.halo_hi)) + } + + /// Task-local → global partition index mapping. See the struct-level + /// invariant note. + pub fn partition_indices(&self) -> &[usize] { + &self.partition_indices + } +} + +impl Debug for RangeFilterExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("RangeFilterExec") + .field("routing_expr", &self.routing_expr.to_string()) + .field("halo_lo", &self.halo_lo) + .field("halo_hi", &self.halo_hi) + .field("cuts", &self.cuts.lock()) + .finish() + } +} + +impl DisplayAs for RangeFilterExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let cuts = self.cuts.lock(); + let cut_display = match cuts.as_ref() { + Some(c) => format!("{:?}", c), + None => "pending".to_string(), + }; + write!( + f, + "RangeFilterExec: routing={}, halo=[{}, {}], cuts={}", + self.routing_expr, self.halo_lo, self.halo_hi, cut_display + ) + } + DisplayFormatType::TreeRender => write!(f, "RangeFilterExec"), + } + } +} + +impl ExecutionPlan for RangeFilterExec { + fn name(&self) -> &str { + "RangeFilterExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input] = children.as_slice() else { + return internal_err!( + "RangeFilterExec expects exactly one child, got {}", + children.len() + ); + }; + // Preserve the cuts slot across the rewrite so a pending + // RangeFilterExec that gets its child transformed doesn't lose + // the eventual scheduler resolution target. Preserve + // partition_indices too — with_new_children is a tree rewrite, + // not a partition restriction; if a rewriter needs to change the + // partition mapping it goes through restrict_partitions instead. + let cuts_snapshot = self.cuts.lock().clone().map(|c| { + c.into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect() + }); + Ok(Arc::new(Self::try_new_inner( + input.clone(), + self.routing_expr.clone(), + cuts_snapshot, + ScalarValue::Float64(Some(self.halo_lo)), + ScalarValue::Float64(Some(self.halo_hi)), + Some(self.partition_indices.clone()), + )?)) + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + fn required_input_ordering(&self) -> Vec> { + vec![None] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn partition_statistics(&self, _partition: Option) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + let cuts = self.cuts.lock().clone().ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "RangeFilterExec: execute() called before resolve_cuts()".into(), + ) + })?; + let Some(&global_partition) = self.partition_indices.get(partition) else { + return internal_err!( + "RangeFilterExec: partition {} out of bounds ({} local partitions)", + partition, + self.partition_indices.len() + ); + }; + let (lo, hi) = + partition_bounds(&cuts, global_partition, self.halo_lo, self.halo_hi); + let predicate = build_predicate_from_bounds(self.routing_expr.clone(), lo, hi); + let schema = self.schema(); + let input = self.input.execute(partition, ctx)?; + let fast_path = self.sorted_on_key.then(|| FastPathState { + routing_expr: self.routing_expr.clone(), + lo, + hi, + }); + let stream = RangeFilterStream { + schema: schema.clone(), + predicate, + input, + fast_path, + }; + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +/// Extract an `f64` from a `ScalarValue`. Restricted to `Float64` today +/// because the surrounding operators (URRE/ORRE, T-Digest) only understand +/// Float64. TODO widen with the KLL migration to accept any +/// `arrow::datatypes::ArrowPrimitiveType`. +fn as_f64(sv: &ScalarValue) -> Result { + match sv { + ScalarValue::Float64(Some(v)) => Ok(*v), + ScalarValue::Float64(None) => { + internal_err!("RangeFilterExec: null ScalarValue is not permitted") + } + other => internal_err!( + "RangeFilterExec: only Float64 ScalarValue supported today, got {other:?}" + ), + } +} + +/// Compute partition `k`'s half-open bounds `[lo, hi)`. `None` means unbounded +/// on that side (virtual ±∞). Partition 0 is `(-∞, cuts[0] + halo_hi)`, +/// partition K-1 is `[cuts[K-2] - halo_lo, +∞)`, K == 1 is `(-∞, +∞)`. +fn partition_bounds( + cuts: &[f64], + partition: usize, + halo_lo: f64, + halo_hi: f64, +) -> (Option, Option) { + let lo = partition + .checked_sub(1) + .and_then(|i| cuts.get(i).copied()) + .map(|c| c - halo_lo); + let hi = cuts.get(partition).copied().map(|c| c + halo_hi); + (lo, hi) +} + +/// Assemble the boolean `PhysicalExpr` predicate from `[lo, hi)`. Used both by +/// [`RangeFilterStream`]'s slow path and by the tests that inspect the +/// generated expression tree. +fn build_predicate_from_bounds( + routing_expr: Arc, + lo: Option, + hi: Option, +) -> Arc { + let lit = |v: f64| -> Arc { + Arc::new(Literal::new(ScalarValue::Float64(Some(v)))) + }; + let ge = |lo: f64| -> Arc { + Arc::new(BinaryExpr::new( + routing_expr.clone(), + Operator::GtEq, + lit(lo), + )) + }; + let lt = |hi: f64| -> Arc { + Arc::new(BinaryExpr::new(routing_expr.clone(), Operator::Lt, lit(hi))) + }; + match (lo, hi) { + (None, None) => Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))), + (None, Some(hi)) => lt(hi), + (Some(lo), None) => ge(lo), + (Some(lo), Some(hi)) => Arc::new(BinaryExpr::new(ge(lo), Operator::And, lt(hi))), + } +} + +/// Test-only shim preserving the pre-fast-path signature. Real code paths +/// call `partition_bounds` + `build_predicate_from_bounds` separately, since +/// the bounds are also fed to the fast path. +#[cfg(test)] +fn build_predicate( + routing_expr: Arc, + cuts: &[f64], + partition: usize, + halo_lo: f64, + halo_hi: f64, +) -> Arc { + let (lo, hi) = partition_bounds(cuts, partition, halo_lo, halo_hi); + build_predicate_from_bounds(routing_expr, lo, hi) +} + +/// Fast-path state. Present only when the input is sorted ascending on +/// `routing_expr` — then a batch's first and last routing values bound the +/// whole batch's value range, unlocking three shortcuts: +/// +/// - `last < lo` or `first >= hi` — batch is entirely outside the partition's +/// window. Drop it. +/// - `first >= lo && last < hi` — batch is entirely inside. Pass it through +/// unchanged (Arc-clone). +/// - Otherwise — the window covers a prefix, suffix, or interior slice. +/// Binary-search the routing column for the slice bounds and +/// `RecordBatch::slice` (zero-copy). +/// +/// Falls back to the general `filter_record_batch` path when the batch's +/// routing column contains nulls (Float64Array binary search would treat null +/// slots as garbage values). +struct FastPathState { + routing_expr: Arc, + lo: Option, + hi: Option, +} + +struct RangeFilterStream { + schema: SchemaRef, + predicate: Arc, + input: SendableRecordBatchStream, + fast_path: Option, +} + +impl RangeFilterStream { + /// Apply the general predicate to `batch` — used both by the non-sorted + /// fallback and by the sorted path when the routing column has nulls. + fn slow_filter(&self, batch: &RecordBatch) -> Result { + let mask = self + .predicate + .evaluate(batch) + .and_then(|v| v.into_array(batch.num_rows()))?; + let mask = as_boolean_array(&mask)?; + Ok(filter_record_batch(batch, mask)?) + } + + /// Try the sorted-input shortcuts. Returns `None` iff the batch is + /// entirely outside the window (drop). Returns `Some(batch)` with the + /// selected rows otherwise. + fn fast_filter( + &self, + state: &FastPathState, + batch: RecordBatch, + ) -> Result> { + let n = batch.num_rows(); + let arr = state + .routing_expr + .evaluate(&batch) + .and_then(|v| v.into_array(n))?; + let col = as_float64_array(&arr)?; + // Nulls in routing column: `values()` returns garbage for null slots, + // and NULL vs bound comparisons must be false. Slow path handles both. + if col.null_count() > 0 { + let filtered = self.slow_filter(&batch)?; + return Ok((filtered.num_rows() > 0).then_some(filtered)); + } + let first = col.value(0); + let last = col.value(n - 1); + // Skip: entire batch is outside the window. + if state.hi.is_some_and(|hi| first >= hi) || state.lo.is_some_and(|lo| last < lo) + { + return Ok(None); + } + // Pass-through: entire batch is inside the window. + let above_lo = state.lo.is_none_or(|lo| first >= lo); + let below_hi = state.hi.is_none_or(|hi| last < hi); + if above_lo && below_hi { + return Ok(Some(batch)); + } + // Mixed: partition the sorted column and slice. + let values = col.values(); + let start = state.lo.map_or(0, |lo| values.partition_point(|v| *v < lo)); + let end = state.hi.map_or(n, |hi| values.partition_point(|v| *v < hi)); + if start >= end { + return Ok(None); + } + Ok(Some(batch.slice(start, end - start))) + } +} + +impl Stream for RangeFilterStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() == 0 { + continue; + } + let filtered = match &self.fast_path { + Some(state) => match self.fast_filter(state, batch)? { + Some(b) => b, + None => continue, + }, + None => { + let out = self.slow_filter(&batch)?; + if out.num_rows() == 0 { + continue; + } + out + } + }; + return Poll::Ready(Some(Ok(filtered))); + } + Some(Err(e)) => return Poll::Ready(Some(Err(e))), + None => { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + return Poll::Ready(None); + } + } + } + } +} + +impl RecordBatchStream for RangeFilterStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +// Silence the unused import warning when the file is compiled without +// arrow — DataType is only exercised via `data_type(...)` return checks. +#[allow(dead_code)] +fn _touch_datatype(_: DataType) {} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::compute::SortOptions; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::LexOrdering; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + use datafusion::prelude::SessionContext; + + fn v_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) + } + + fn v_source(partitions: usize) -> Arc { + let schema = v_schema(); + let source: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(), + ))); + Arc::new( + RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(partitions)) + .unwrap(), + ) + } + + fn v_col() -> Arc { + Arc::new(Column::new_with_schema("v", v_schema().as_ref()).unwrap()) + } + + /// Single-partition memory source containing `batches`, declaring + /// `sort_information` on `v` — RangeFilterExec's `sorted_on_key` detection + /// picks this up and enables the fast path. + fn sorted_v_source( + batches: Vec, + options: SortOptions, + ) -> Arc { + let schema = v_schema(); + let sort_expr = PhysicalSortExpr::new(v_col(), options); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let cfg = MemorySourceConfig::try_new(&[batches], schema.clone(), None).unwrap(); + let cfg = cfg.try_with_sort_information(vec![ordering]).unwrap(); + Arc::new(DataSourceExec::new(Arc::new(cfg))) + } + + fn asc() -> SortOptions { + SortOptions { + descending: false, + nulls_first: false, + } + } + + fn batch(values: &[f64]) -> RecordBatch { + RecordBatch::try_new( + v_schema(), + vec![Arc::new(Float64Array::from(values.to_vec()))], + ) + .unwrap() + } + + fn sv(v: f64) -> ScalarValue { + ScalarValue::Float64(Some(v)) + } + + fn svs(vs: &[f64]) -> Vec { + vs.iter().map(|&v| sv(v)).collect() + } + + #[test] + fn default_partition_indices_must_fit_cut_count() { + // Default partition_indices = (0..input.partition_count()); with only 1 cut + // there are 2 global partitions, so index 2 is out of range. + let src = v_source(3); + let err = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[10.0]), + sv(0.0), + sv(0.0), + ) + .unwrap_err(); + assert!( + err.to_string().contains("only 2 global partitions"), + "got: {err}" + ); + } + + #[test] + fn cuts_must_be_monotone() { + let src = v_source(3); + let err = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[10.0, 5.0]), + sv(0.0), + sv(0.0), + ) + .unwrap_err(); + assert!(err.to_string().contains("monotone")); + } + + #[test] + fn halo_must_be_non_negative() { + let src = v_source(2); + let err = RangeFilterExec::try_new_resolved( + src.clone(), + v_col(), + svs(&[5.0]), + sv(-1.0), + sv(0.0), + ) + .unwrap_err(); + assert!(err.to_string().contains("halo_lo")); + let err = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[5.0]), + sv(0.0), + sv(f64::NAN), + ) + .unwrap_err(); + assert!(err.to_string().contains("halo_hi")); + } + + #[test] + fn non_float64_cuts_are_rejected() { + let src = v_source(2); + let cuts = vec![ScalarValue::Int64(Some(5))]; + let err = RangeFilterExec::try_new_resolved(src, v_col(), cuts, sv(0.0), sv(0.0)) + .unwrap_err(); + assert!(err.to_string().contains("only Float64"), "got: {err}"); + } + + #[test] + fn pending_construction_defers_check() { + let src = v_source(3); + // 5 cuts wouldn't fit 3 partitions if we required alignment eagerly. + let rf = + RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0)).unwrap(); + assert!(rf.cuts().is_none()); + } + + #[tokio::test] + async fn execute_before_resolve_errors() { + let src = v_source(2); + let rf = + RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0)).unwrap(); + let ctx = SessionContext::new().task_ctx(); + let Err(err) = rf.execute(0, ctx) else { + panic!("execute() should error before resolve_cuts") + }; + assert!(err.to_string().contains("before resolve_cuts")); + } + + #[test] + fn resolve_cuts_validates() { + let src = v_source(3); + let rf = + RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0)).unwrap(); + // 3 default partition_indices = [0,1,2]; 1 cut means 2 global partitions + // so index 2 is out of range. + let err = rf.resolve_cuts(svs(&[1.0])).unwrap_err(); + assert!( + err.to_string().contains("2 global partitions"), + "got: {err}" + ); + // Non-monotone. + let err = rf.resolve_cuts(svs(&[5.0, 1.0])).unwrap_err(); + assert!(err.to_string().contains("monotone")); + // Good. + rf.resolve_cuts(svs(&[1.0, 5.0])).unwrap(); + assert_eq!(rf.cuts().unwrap(), svs(&[1.0, 5.0])); + } + + #[test] + fn restrict_partitions_remaps_indices() { + // Original K=4 partitions with cuts [10, 20, 30]. Restrict to + // task-local [1, 3] — new operator has 2 local partitions that + // apply the predicates for global partitions 1 and 3. + let src = v_source(4); + let rf = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[10.0, 20.0, 30.0]), + sv(0.0), + sv(0.0), + ) + .unwrap(); + let restricted_input = v_source(2); + let restricted = rf.restrict_partitions(restricted_input, &[1, 3]).unwrap(); + assert_eq!(restricted.partition_indices(), &[1, 3]); + assert_eq!(restricted.cuts().unwrap(), svs(&[10.0, 20.0, 30.0])); + } + + #[test] + fn build_predicate_half_open_boundaries() { + // K = 3, cuts = [10, 20], no halo: + // part 0: v < 10 + // part 1: v >= 10 AND v < 20 + // part 2: v >= 20 + let cuts = vec![10.0, 20.0]; + let p0 = build_predicate(v_col(), &cuts, 0, 0.0, 0.0).to_string(); + let p1 = build_predicate(v_col(), &cuts, 1, 0.0, 0.0).to_string(); + let p2 = build_predicate(v_col(), &cuts, 2, 0.0, 0.0).to_string(); + assert!(p0.contains("<") && p0.contains("10")); + assert!(p1.contains("10") && p1.contains("20") && p1.contains("AND")); + assert!(p2.contains(">=") && p2.contains("20")); + } + + #[test] + fn build_predicate_halo_widens_boundaries() { + // K = 3, cuts = [10, 20], halo_lo=3, halo_hi=0: + // part 1: v >= (10 - 3) AND v < (20 - 0) => v >= 7 AND v < 20 + let cuts = vec![10.0, 20.0]; + let p1 = build_predicate(v_col(), &cuts, 1, 3.0, 0.0).to_string(); + assert!( + p1.contains("7") && p1.contains("20"), + "expected halo-widened lo bound, got: {p1}" + ); + } + + #[test] + fn build_predicate_k1_true_when_no_cuts() { + let expr = build_predicate(v_col(), &[], 0, 0.0, 0.0).to_string(); + assert!(expr.contains("true"), "expected lit(true), got: {expr}"); + } + + #[test] + fn sorted_on_key_detected_when_input_ascending_on_routing_expr() { + let src = sorted_v_source(vec![batch(&[1.0, 2.0, 3.0])], asc()); + let rf = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[2.0]), + sv(0.0), + sv(0.0), + ) + .unwrap(); + assert!(rf.sorted_on_key); + } + + #[test] + fn sorted_on_key_false_when_input_descending_on_routing_expr() { + let desc = SortOptions { + descending: true, + nulls_first: false, + }; + let src = sorted_v_source(vec![batch(&[3.0, 2.0, 1.0])], desc); + let rf = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[2.0]), + sv(0.0), + sv(0.0), + ) + .unwrap(); + // Fast path assumes ascending — reverse order would flip min/max. + assert!(!rf.sorted_on_key); + } + + #[test] + fn sorted_on_key_false_when_input_advertises_no_ordering() { + // RepartitionExec on a single partition drops ordering information. + let src = v_source(2); + let rf = RangeFilterExec::try_new_resolved( + src, + v_col(), + svs(&[2.0]), + sv(0.0), + sv(0.0), + ) + .unwrap(); + assert!(!rf.sorted_on_key); + } + + /// Helper to drain one partition of a resolved RangeFilterExec into a + /// concatenated `Vec` of the output routing values, so tests can + /// assert the exact rows that survived without needing to reason about + /// intermediate batch boundaries. + async fn drain(rf: Arc, partition: usize) -> Vec { + let ctx = SessionContext::new().task_ctx(); + let mut stream = rf.execute(partition, ctx).unwrap(); + let mut out = Vec::new(); + while let Some(res) = stream.next().await { + let b = res.unwrap(); + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..col.len() { + out.push(col.value(i)); + } + } + out + } + + /// Build a resolved 1-local-partition RangeFilterExec that applies the + /// predicate for global partition `global_k` of the K=cuts.len()+1 shape. + /// The input has 1 partition, so `execute(0)` runs the chosen global + /// predicate over the whole test batch. + fn one_partition_rf( + input: Arc, + cuts: &[f64], + global_k: usize, + ) -> Arc { + Arc::new( + RangeFilterExec::try_new_with_indices( + input, + v_col(), + svs(cuts), + sv(0.0), + sv(0.0), + vec![global_k], + ) + .unwrap(), + ) + } + + #[tokio::test] + async fn fast_path_skips_batch_entirely_below_partition_lo() { + // Global partition 1 window is [10, 20). Batch entirely below 10 — + // fast path detects `last < lo` and drops the whole batch. + let src = sorted_v_source(vec![batch(&[1.0, 2.0, 3.0])], asc()); + let rf = one_partition_rf(src, &[10.0, 20.0], 1); + assert!(rf.sorted_on_key); + let rows = drain(rf, 0).await; + assert!( + rows.is_empty(), + "expected zero surviving rows, got {rows:?}" + ); + } + + #[tokio::test] + async fn fast_path_skips_batch_entirely_above_partition_hi() { + // Global partition 0 window is (-inf, 10). Batch entirely at/above 10. + let src = sorted_v_source(vec![batch(&[10.0, 20.0, 30.0])], asc()); + let rf = one_partition_rf(src, &[10.0, 20.0], 0); + let rows = drain(rf, 0).await; + assert!( + rows.is_empty(), + "expected zero surviving rows, got {rows:?}" + ); + } + + #[tokio::test] + async fn fast_path_passes_batch_entirely_inside_window() { + // Global partition 1 window is [10, 20). Batch [11, 15, 19] is + // entirely inside. Fast path Arc-clones — verify identical rows out. + let input_rows = vec![11.0, 15.0, 19.0]; + let src = sorted_v_source(vec![batch(&input_rows)], asc()); + let rf = one_partition_rf(src, &[10.0, 20.0], 1); + let rows = drain(rf, 0).await; + assert_eq!(rows, input_rows); + } + + #[tokio::test] + async fn fast_path_slices_mixed_batch() { + // Global partition 1 window is [10, 20). Batch straddles both + // boundaries. Rows 10..20 should survive; 5, 25 dropped. + let src = + sorted_v_source(vec![batch(&[5.0, 10.0, 15.0, 19.0, 20.0, 25.0])], asc()); + let rf = one_partition_rf(src, &[10.0, 20.0], 1); + let rows = drain(rf, 0).await; + assert_eq!(rows, vec![10.0, 15.0, 19.0]); + } + + #[tokio::test] + async fn fast_path_slices_open_upper_bound() { + // Global partition K-1 has hi = None (open above). Verify the + // partition_point logic that treats missing `hi` as `n` still slices + // the low end. Partition 2 is [20, +inf): 25, 35 survive. + let src = sorted_v_source(vec![batch(&[5.0, 15.0, 25.0, 35.0])], asc()); + let rf = one_partition_rf(src, &[10.0, 20.0], 2); + let rows = drain(rf, 0).await; + assert_eq!(rows, vec![25.0, 35.0]); + } + + #[tokio::test] + async fn fast_path_falls_back_when_routing_column_has_nulls() { + // Nulls in the routing column — `Float64Array::values()` returns + // garbage for null slots, so partition_point would be meaningless. + // Verify the slow path handles it and still produces correct rows. + let nullable_schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, true)])); + let arr = Float64Array::from(vec![Some(5.0), None, Some(15.0), None, Some(25.0)]); + let b = + RecordBatch::try_new(nullable_schema.clone(), vec![Arc::new(arr)]).unwrap(); + let v_expr: Arc = + Arc::new(Column::new_with_schema("v", nullable_schema.as_ref()).unwrap()); + let sort_expr = PhysicalSortExpr::new(v_expr.clone(), asc()); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let cfg = MemorySourceConfig::try_new(&[vec![b]], nullable_schema.clone(), None) + .unwrap(); + let cfg = cfg.try_with_sort_information(vec![ordering]).unwrap(); + let src: Arc = Arc::new(DataSourceExec::new(Arc::new(cfg))); + let rf = Arc::new( + RangeFilterExec::try_new_with_indices( + src, + v_expr, + svs(&[10.0, 20.0]), + sv(0.0), + sv(0.0), + vec![1], + ) + .unwrap(), + ); + assert!(rf.sorted_on_key, "input still advertises ordering"); + // Global partition 1 = [10, 20). Only 15.0 qualifies; nulls compare false. + let rows = drain(rf, 0).await; + assert_eq!(rows, vec![15.0]); + } + + #[tokio::test] + async fn slow_path_matches_fast_path_when_input_unsorted() { + // Same batch, same predicate, but no advertised ordering: slow path + // via `filter_record_batch` must produce the same rows the fast path + // would. + let cfg = MemorySourceConfig::try_new( + &[vec![batch(&[5.0, 10.0, 15.0, 19.0, 20.0, 25.0])]], + v_schema(), + None, + ) + .unwrap(); + let mem_src: Arc = + Arc::new(DataSourceExec::new(Arc::new(cfg))); + // RepartitionExec to 1 partition drops the ordering claim without + // fanning the data out (RoundRobinBatch to 1 = identity). + let repart: Arc = Arc::new( + RepartitionExec::try_new(mem_src, Partitioning::RoundRobinBatch(1)).unwrap(), + ); + assert!(repart.output_ordering().is_none()); + let rf = one_partition_rf(repart, &[10.0, 20.0], 1); + assert!(!rf.sorted_on_key); + let rows = drain(rf, 0).await; + assert_eq!(rows, vec![10.0, 15.0, 19.0]); + } + + #[tokio::test] + async fn execute_filters_by_partition_range() { + // Build an input with 3 sorted partitions containing a batch each, + // apply resolved cuts [10, 20], verify each partition emits only its slice. + let schema = v_schema(); + let source: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new( + &[ + vec![batch(&[5.0, 15.0, 25.0])], + vec![batch(&[5.0, 15.0, 25.0])], + vec![batch(&[5.0, 15.0, 25.0])], + ], + schema.clone(), + None, + ) + .unwrap(), + ))); + let rf = Arc::new( + RangeFilterExec::try_new_resolved( + source, + v_col(), + svs(&[10.0, 20.0]), + sv(0.0), + sv(0.0), + ) + .unwrap(), + ); + let ctx = SessionContext::new().task_ctx(); + for (partition, expected) in [(0, 5.0), (1, 15.0), (2, 25.0)] { + let mut stream = rf.execute(partition, ctx.clone()).unwrap(); + let batch = stream.next().await.unwrap().unwrap(); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.len(), 1, "partition {partition}"); + assert_eq!(col.value(0), expected, "partition {partition}"); + } + } +} diff --git a/ballista/core/src/execution_plans/range_shuffle_reader.rs b/ballista/core/src/execution_plans/range_shuffle_reader.rs new file mode 100644 index 0000000000..96bfe1bdef --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle_reader.rs @@ -0,0 +1,558 @@ +// 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. + +//! Ordering-preserving shuffle reader. +//! +//! The regular [`ShuffleReaderExec`](super::ShuffleReaderExec) fans every +//! upstream source for an output partition into one mpsc channel, `try_flatten`s +//! them in arrival order, then coalesces batches. Each source file is +//! internally sorted, but the concatenation is not — the resulting stream +//! violates the monotonicity that RANGE-frame window operators (and +//! sort-merge-join build sides) require. +//! +//! `RangeShuffleReaderExec` keeps each upstream source alive as its own +//! stream and feeds all N into a `StreamingMerge` keyed on the child's +//! declared output ordering. Batches within an output partition are +//! globally sorted on the merge key. +//! +//! Trade-offs vs. the regular reader: +//! +//! - No permit-based fetch governor. Backpressure flows from the merge +//! consumer down through each per-source stream — the merge only polls +//! the source it needs next, so h2 (for remote) and disk (for local) +//! throttle naturally. +//! - No mid-body fetch retry. `fetch_partition_remote` streams directly +//! without buffering the whole source, so a transport error mid-body +//! fails the merged stream; the task-level retry re-executes. +//! - No coalesce or broadcast variants. Ordered outputs come from +//! `OrderedRangeRepartitionExec`, which is one-to-one and never fanned +//! into a broadcast. + +use crate::client_pool::BallistaClientPool; +use crate::execution_plans::shuffle_reader::{ + fetch_partition_local, fetch_partition_remote, local_remote_read_split, + stats_for_partition, +}; +use crate::extension::SessionConfigExt; +use crate::serde::scheduler::PartitionLocation; +use crate::utils::GrpcClientConfig; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, Statistics}; +use datafusion::error::DataFusionError; +use datafusion::execution::TaskContext; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, Partitioning}; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; +use futures::TryStreamExt; +use log::debug; +use std::sync::Arc; + +/// Ordering-preserving shuffle reader. See module docs. +#[derive(Debug)] +pub struct RangeShuffleReaderExec { + /// Upstream stage that produced these files. + pub stage_id: usize, + pub(crate) schema: SchemaRef, + /// M-shape: `partition[k]` lists the upstream partition locations that + /// output partition `k` needs to merge. + pub partition: Vec>, + /// Sort key the merge preserves. Advertised in `PlanProperties.eq_properties` + /// so downstream operators (BWAG, SMJ build side) see the output ordering. + merge_ordering: LexOrdering, + metrics: ExecutionPlanMetricsSet, + properties: Arc, + work_dir: Option, + client_pool: Option>, +} + +impl RangeShuffleReaderExec { + /// The output partition count is `partition.len()`. Partitioning is + /// range on the merge key — DF has no `Partitioning::Range`, so it's + /// reported as `UnknownPartitioning`. Downstream co-partitioning is + /// carried by the advertised ordering, not by partitioning kind. + pub fn try_new( + stage_id: usize, + partition: Vec>, + schema: SchemaRef, + merge_ordering: LexOrdering, + ) -> Result { + let output_partition_count = partition.len(); + let eq_properties = EquivalenceProperties::new_with_orderings( + schema.clone(), + vec![merge_ordering.clone()], + ); + let properties = Arc::new(PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(output_partition_count), + datafusion::physical_plan::execution_plan::EmissionType::Incremental, + datafusion::physical_plan::execution_plan::Boundedness::Bounded, + )); + Ok(Self { + stage_id, + schema, + partition, + merge_ordering, + metrics: ExecutionPlanMetricsSet::new(), + properties, + work_dir: None, + client_pool: None, + }) + } + + /// Late-bound by the executor. + pub fn with_work_dir(&self, work_dir: String) -> Self { + Self { + stage_id: self.stage_id, + schema: self.schema.clone(), + partition: self.partition.clone(), + merge_ordering: self.merge_ordering.clone(), + metrics: self.metrics.clone(), + properties: self.properties.clone(), + work_dir: Some(work_dir), + client_pool: self.client_pool.clone(), + } + } + + /// Late-bound by the executor. + pub fn with_client_pool(&self, client_pool: Arc) -> Self { + Self { + stage_id: self.stage_id, + schema: self.schema.clone(), + partition: self.partition.clone(), + merge_ordering: self.merge_ordering.clone(), + metrics: self.metrics.clone(), + properties: self.properties.clone(), + work_dir: self.work_dir.clone(), + client_pool: Some(client_pool), + } + } + + /// Sort key the merge preserves. + pub fn merge_ordering(&self) -> &LexOrdering { + &self.merge_ordering + } +} + +impl DisplayAs for RangeShuffleReaderExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "RangeShuffleReaderExec: upstream_stage: {}, partitions: {}, ordering: {}", + self.stage_id, + self.partition.len(), + self.merge_ordering, + ) + } + DisplayFormatType::TreeRender => { + writeln!(f, "upstream_stage={}", self.stage_id)?; + writeln!(f, "output_partitions={}", self.partition.len())?; + writeln!(f, "ordering={}", self.merge_ordering) + } + } + } +} + +impl ExecutionPlan for RangeShuffleReaderExec { + fn name(&self) -> &str { + "RangeShuffleReaderExec" + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return Err(DataFusionError::Plan( + "RangeShuffleReaderExec does not support children plans".to_owned(), + )); + } + Ok(Arc::new(Self { + stage_id: self.stage_id, + schema: self.schema.clone(), + partition: self.partition.clone(), + merge_ordering: self.merge_ordering.clone(), + metrics: ExecutionPlanMetricsSet::new(), + properties: self.properties.clone(), + work_dir: self.work_dir.clone(), + client_pool: self.client_pool.clone(), + })) + } + + fn execute( + &self, + output_partition: usize, + ctx: Arc, + ) -> Result { + let task_id = ctx + .task_id() + .unwrap_or_else(|| output_partition.to_string()); + debug!("RangeShuffleReaderExec::execute({task_id})"); + + let config = ctx.session_config(); + let work_dir = self.work_dir.as_ref().ok_or_else(|| { + DataFusionError::Configuration( + "RangeShuffleReaderExec work dir should have been set by executor" + .to_owned(), + ) + })?; + + let locations = self.partition[output_partition].clone(); + let (local_locations, remote_locations) = local_remote_read_split( + work_dir, + locations, + config.ballista_shuffle_reader_force_remote_read(), + ); + + let mut sub_streams: Vec = + Vec::with_capacity(local_locations.len() + remote_locations.len()); + + for loc in local_locations { + let stream = fetch_partition_local(work_dir, &loc) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + sub_streams.push(stream); + } + + if !remote_locations.is_empty() { + let grpc_config: Arc = + Arc::new((&config.ballista_config()).into()); + let customize_endpoint = + config.ballista_override_create_grpc_client_endpoint(); + let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight(); + let client_pool = self.client_pool.clone(); + + for loc in remote_locations { + let schema = self.schema.clone(); + let grpc_config = grpc_config.clone(); + let customize_endpoint = customize_endpoint.clone(); + let client_pool = client_pool.clone(); + + // Lazy connect: the first poll from the merge triggers the + // remote fetch; subsequent polls stream batches directly. No + // buffering, no retry — task-level retry covers transport + // failures. + let lazy = futures::stream::once(async move { + fetch_partition_remote( + &loc, + grpc_config, + prefer_flight, + customize_endpoint, + client_pool, + ) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) + }) + .try_flatten(); + + sub_streams.push(Box::pin(RecordBatchStreamAdapter::new(schema, lazy))); + } + } + + if sub_streams.is_empty() { + return Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + futures::stream::empty(), + ))); + } + + let baseline = BaselineMetrics::new(&self.metrics, output_partition); + let reservation = MemoryConsumer::new(format!( + "RangeShuffleReaderExec[stage={},out={}]", + self.stage_id, output_partition, + )) + .register(ctx.memory_pool()); + + let merged = StreamingMergeBuilder::new() + .with_streams(sub_streams) + .with_schema(self.schema.clone()) + .with_expressions(&self.merge_ordering) + .with_batch_size(config.batch_size()) + .with_metrics(baseline) + .with_reservation(reservation) + .build()?; + + Ok(merged) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn partition_statistics(&self, partition: Option) -> Result> { + if let Some(idx) = partition { + let partition_count = self.partition.len(); + if idx >= partition_count { + return datafusion::common::internal_err!( + "Invalid partition index: {}, the partition count is {}", + idx, + partition_count + ); + } + let stats = + stats_for_partition(idx, self.schema.fields().len(), &self.partition)?; + return Ok(Arc::new(stats)); + } + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::serde::scheduler::{ + ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, + PartitionId, PartitionStats, + }; + use datafusion::arrow::array::{ArrayRef, Float64Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::ipc::writer::StreamWriter; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_expr::expressions::Column; + use datafusion::prelude::SessionContext; + use std::fs::{File, create_dir_all}; + use tempfile::tempdir; + + /// Write an Arrow IPC stream containing `batches` at `path`. + fn write_ipc_stream(path: &std::path::Path, batches: &[RecordBatch]) { + create_dir_all(path.parent().unwrap()).unwrap(); + let file = File::create(path).unwrap(); + let mut writer = StreamWriter::try_new(file, &batches[0].schema()).unwrap(); + for b in batches { + writer.write(b).unwrap(); + } + writer.finish().unwrap(); + } + + fn sorted_batch(schema: SchemaRef, values: &[f64]) -> RecordBatch { + let col: ArrayRef = Arc::new(Float64Array::from(values.to_vec())); + RecordBatch::try_new(schema, vec![col]).unwrap() + } + + fn make_location(job: &str, stage: usize, partition: usize) -> PartitionLocation { + PartitionLocation { + map_partition_id: partition, + partition_id: PartitionId { + job_id: job.into(), + stage_id: stage, + partition_id: partition, + }, + executor_meta: ExecutorMetadata { + id: "test-executor".to_string(), + host: "127.0.0.1".to_string(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default(), + os_info: ExecutorOperatingSystemSpecification::default(), + }, + partition_stats: PartitionStats::default(), + file_id: None, + is_sort_shuffle: false, + } + } + + /// Two upstream sources, each internally sorted; the range reader must + /// interleave them into a single globally-sorted stream. This is the core + /// correctness property the regular reader violates. + #[tokio::test] + async fn merges_two_sorted_sources() { + let dir = tempdir().unwrap(); + let work_dir = dir.path(); + let job = "job-merges-two"; + let stage_id = 1usize; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + + // Source 0: [1.0, 3.0, 5.0] + // Source 1: [2.0, 4.0, 6.0] + // Merged: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + let src0 = sorted_batch(schema.clone(), &[1.0, 3.0, 5.0]); + let src1 = sorted_batch(schema.clone(), &[2.0, 4.0, 6.0]); + + let loc0 = make_location(job, stage_id, 0); + let loc1 = make_location(job, stage_id, 1); + let path0 = loc0.path(work_dir.to_str().unwrap()).expect("path0"); + let path1 = loc1.path(work_dir.to_str().unwrap()).expect("path1"); + write_ipc_stream(&path0, &[src0]); + write_ipc_stream(&path1, &[src1]); + + // Both sources land in output partition 0. + let partitions = vec![vec![loc0, loc1]]; + let merge_ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default( + Arc::new(Column::new("v", 0)), + )]) + .unwrap(); + + let reader = RangeShuffleReaderExec::try_new( + stage_id, + partitions, + schema.clone(), + merge_ordering, + ) + .unwrap() + .with_work_dir(work_dir.to_string_lossy().to_string()); + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let stream = reader.execute(0, task_ctx).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let all: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(all, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + } + + /// If the caller only has one source per output partition, the reader + /// must still work — StreamingMerge over 1 stream is a valid degenerate + /// case. + #[tokio::test] + async fn passes_through_single_source() { + let dir = tempdir().unwrap(); + let work_dir = dir.path(); + let job = "job-single-source"; + let stage_id = 2usize; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let src = sorted_batch(schema.clone(), &[10.0, 20.0, 30.0]); + + let loc = make_location(job, stage_id, 0); + let path = loc.path(work_dir.to_str().unwrap()).expect("path"); + write_ipc_stream(&path, &[src]); + + let partitions = vec![vec![loc]]; + let merge_ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default( + Arc::new(Column::new("v", 0)), + )]) + .unwrap(); + + let reader = RangeShuffleReaderExec::try_new( + stage_id, + partitions, + schema.clone(), + merge_ordering, + ) + .unwrap() + .with_work_dir(work_dir.to_string_lossy().to_string()); + + let ctx = SessionContext::new(); + let stream = reader.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let all: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(all, vec![10.0, 20.0, 30.0]); + } + + /// Empty partitions.len() > 0 with all-empty inner locations must return + /// an empty stream instead of tripping StreamingMerge on zero sub-streams. + #[tokio::test] + async fn handles_empty_partition() { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let merge_ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default( + Arc::new(Column::new("v", 0)), + )]) + .unwrap(); + + let reader = RangeShuffleReaderExec::try_new( + 7, + vec![vec![]], + schema.clone(), + merge_ordering, + ) + .unwrap() + .with_work_dir("/tmp".to_string()); + + let ctx = SessionContext::new(); + let stream = reader.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + assert!(batches.is_empty()); + } + + /// The reader must advertise its merge ordering so downstream operators + /// see the sortedness invariant (BWAG's RANGE-frame cursor, SMJ build side). + #[test] + fn advertises_merge_ordering() { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let merge_ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default( + Arc::new(Column::new("v", 0)), + )]) + .unwrap(); + + let reader = RangeShuffleReaderExec::try_new( + 3, + vec![vec![]; 4], + schema, + merge_ordering.clone(), + ) + .unwrap(); + + let out = reader + .properties() + .eq_properties + .oeq_class() + .iter() + .next() + .cloned(); + let expected = out.expect("expected an advertised ordering"); + assert_eq!(expected, merge_ordering); + + assert_eq!(reader.properties().partitioning.partition_count(), 4); + } +} diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 949241a163..172de2d30e 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -1722,7 +1722,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/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index d85f1978fa..6f4ec1a6a0 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/core/src/execution_plans/shuffle_reader.rs @@ -692,7 +692,7 @@ impl RecordBatchStream for GovernedStream { /// Local partitions are read directly from local Arrow IPC files, /// while remote partitions are fetched using the Arrow Flight client. /// If `force_remote_read` is true, all partitions are treated as remote. -fn local_remote_read_split( +pub(crate) fn local_remote_read_split( work_dir: &str, partition_locations: Vec, force_remote_read: bool, @@ -1037,7 +1037,7 @@ async fn new_ballista_client( .await } -async fn fetch_partition_remote( +pub(crate) async fn fetch_partition_remote( location: &PartitionLocation, config: Arc, prefer_flight: bool, @@ -1108,7 +1108,7 @@ async fn fetch_partition_remote( } } -fn fetch_partition_local( +pub(crate) fn fetch_partition_local( work_dir: &str, location: &PartitionLocation, ) -> result::Result { diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 5bf6964eb2..3e165d2bbe 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -60,7 +60,11 @@ pub mod ballista_physical_plan_node { #[prost(message, tag = "9")] OrderedRangeRepartition(super::OrderedRangeRepartitionExecNode), #[prost(message, tag = "10")] - PerPartitionFilter(super::PerPartitionFilterExecNode), + RangeFilter(super::RangeFilterExecNode), + #[prost(message, tag = "11")] + PartitionedBoundedWindowAgg(super::PartitionedBoundedWindowAggExecNode), + #[prost(message, tag = "12")] + RangeShuffleReader(super::RangeShuffleReaderExecNode), } } /// Value-range router over N locally-sorted overlapping input partitions. @@ -135,16 +139,44 @@ pub struct UnorderedRangeRepartitionExecNode { #[prost(uint32, tag = "2")] pub output_partitions: u32, } -/// Filter with per-input-partition predicates. `predicates\[k\]` is the -/// boolean expression applied to input partition `k`. Requires -/// `predicates.len() == input_partition_count`. The child plan is -/// plumbed by the framework as `inputs\[0\]` during decode. +/// Filter over ordered inputs with per-partition half-open range predicates +/// derived from `cuts` + `halo_lo` / `halo_hi`. Zero halo recovers the exact +/// range-repartition trim used above `ShuffleReaderExec`. Non-zero halo +/// widens each partition's read range to include a boundary "context" band +/// (bounded RANGE-frame windows). The child plan is plumbed by the framework +/// as `inputs\[0\]` during decode. Serialization requires cuts to be resolved. +/// +/// `partition_indices` maps this operator's local partition index to the +/// global partition index in the original K-shape defined by `cuts`. Under +/// task-level restriction the input is sliced to a subset of the K global +/// partitions; each entry stays \< `cuts.len() + 1`. #[derive(Clone, PartialEq, ::prost::Message)] -pub struct PerPartitionFilterExecNode { - #[prost(message, repeated, tag = "1")] - pub predicates: ::prost::alloc::vec::Vec< +pub struct RangeFilterExecNode { + #[prost(message, optional, tag = "1")] + pub routing_expr: ::core::option::Option< ::datafusion_proto::protobuf::PhysicalExprNode, >, + #[prost(message, repeated, tag = "2")] + pub cuts: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + #[prost(message, optional, tag = "3")] + pub halo_lo: ::core::option::Option<::datafusion_proto_common::ScalarValue>, + #[prost(message, optional, tag = "4")] + pub halo_hi: ::core::option::Option<::datafusion_proto_common::ScalarValue>, + #[prost(uint32, repeated, tag = "5")] + pub partition_indices: ::prost::alloc::vec::Vec, +} +/// Wrapper for `BoundedWindowAggExec` that overrides +/// `required_input_distribution` to `Unspecified` — see the module doc on +/// `execution_plans::partitioned_bounded_window_agg` for what makes that safe. +/// The child plan is plumbed by the framework as `inputs\[0\]` during decode. +/// `input_order_mode` and `can_repartition` are hardcoded on the decode side +/// per the rule's shape gates; only `window_expr` needs to cross the wire. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PartitionedBoundedWindowAggExecNode { + #[prost(message, repeated, tag = "1")] + pub window_expr: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalWindowExprNode, + >, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ChaosExecNode { @@ -234,6 +266,24 @@ pub struct ShuffleReaderPartition { #[prost(message, repeated, tag = "1")] pub location: ::prost::alloc::vec::Vec, } +/// Ordering-preserving shuffle reader. Reuses `ShuffleReaderPartition` for the +/// M-shape source layout. Partitioning is derived from `partition.len()` +/// (always `UnknownPartitioning`, range-partitioned by `merge_ordering`). +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeShuffleReaderExecNode { + #[prost(message, repeated, tag = "1")] + pub partition: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "2")] + pub schema: ::core::option::Option<::datafusion_proto_common::Schema>, + #[prost(uint32, tag = "3")] + pub stage_id: u32, + /// Sort key the reader's k-way merge preserves. Advertised on the reader's + /// `PlanProperties.eq_properties` for downstream consumers. + #[prost(message, repeated, tag = "4")] + pub merge_ordering: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalSortExprNode, + >, +} /// CoalescePartitionsRule output: groups upstream partitions into coalesced output partitions. /// Empty when no coalesce is applied (the optional field on the parent message is absent). #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 3ceaa2dfd6..00232a8f7c 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -57,9 +57,9 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ BufferExec, BufferMode, ChaosExec, CoalescePlan, OrderedRangeRepartitionExec, - PartitionGroup, PerPartitionFilterExec, RuntimeStatsExec, ShuffleReaderExec, - ShuffleWriterExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, - UnresolvedShuffleExec, + PartitionGroup, PartitionedBoundedWindowAggExec, RangeFilterExec, + RangeShuffleReaderExec, RuntimeStatsExec, ShuffleReaderExec, ShuffleWriterExec, + SortShuffleWriterExec, UnorderedRangeRepartitionExec, UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -507,6 +507,46 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { }; Ok(Arc::new(exec)) } + PhysicalPlanType::RangeShuffleReader(range_reader) => { + let stage_id = range_reader.stage_id as usize; + let schema: SchemaRef = Arc::new(convert_required!(range_reader.schema)?); + let partition_location: Vec> = range_reader + .partition + .iter() + .map(|p| { + p.location + .iter() + .map(|l| { + l.clone().try_into().map_err(|e| { + DataFusionError::Internal(format!( + "Fail to get partition location due to {e:?}" + )) + }) + }) + .collect::, _>>() + }) + .collect::, DataFusionError>>()?; + let merge_ordering_exprs = parse_physical_sort_exprs( + &range_reader.merge_ordering, + &decode_ctx, + schema.as_ref(), + &converter, + )?; + let merge_ordering = datafusion::physical_expr::LexOrdering::new( + merge_ordering_exprs, + ) + .ok_or_else(|| { + proto_error( + "RangeShuffleReaderExec: merge_ordering must be non-empty", + ) + })?; + Ok(Arc::new(RangeShuffleReaderExec::try_new( + stage_id, + partition_location, + schema, + merge_ordering, + )?)) + } PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => { let schema: SchemaRef = Arc::new(convert_required!(unresolved_shuffle.schema)?); @@ -637,29 +677,84 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { node.output_partitions as usize, )?)) } - PhysicalPlanType::PerPartitionFilter(node) => { + PhysicalPlanType::RangeFilter(node) => { let [input] = inputs else { return Err(DataFusionError::Internal(format!( - "PerPartitionFilterExec expects exactly 1 input, got {}", + "RangeFilterExec expects exactly 1 input, got {}", inputs.len() ))); }; let schema = input.schema(); - let predicates = node - .predicates + let routing_expr_proto = node.routing_expr.as_ref().ok_or_else(|| { + DataFusionError::Internal( + "RangeFilterExecNode missing routing_expr".into(), + ) + })?; + let routing_expr = + datafusion_proto::physical_plan::from_proto::parse_physical_expr( + routing_expr_proto, + ctx, + schema.as_ref(), + self, + )?; + let partition_indices: Vec = + node.partition_indices.iter().map(|&i| i as usize).collect(); + let sv_from_proto = |p: &datafusion_proto_common::ScalarValue| { + datafusion::scalar::ScalarValue::try_from(p).map_err(|e| { + DataFusionError::Internal(format!( + "RangeFilterExec: failed to decode ScalarValue: {e:?}" + )) + }) + }; + let cuts: Vec<_> = node.cuts.iter().map(sv_from_proto).collect::>( + )?; + let halo_lo_proto = node.halo_lo.as_ref().ok_or_else(|| { + DataFusionError::Internal( + "RangeFilterExecNode missing halo_lo".into(), + ) + })?; + let halo_lo = sv_from_proto(halo_lo_proto)?; + let halo_hi_proto = node.halo_hi.as_ref().ok_or_else(|| { + DataFusionError::Internal( + "RangeFilterExecNode missing halo_hi".into(), + ) + })?; + let halo_hi = sv_from_proto(halo_hi_proto)?; + Ok(Arc::new(RangeFilterExec::try_new_with_indices( + input.clone(), + routing_expr, + cuts, + halo_lo, + halo_hi, + partition_indices, + )?)) + } + PhysicalPlanType::PartitionedBoundedWindowAgg(node) => { + let [input] = inputs else { + return Err(DataFusionError::Internal(format!( + "PartitionedBoundedWindowAggExec expects exactly 1 input, got {}", + inputs.len() + ))); + }; + let input_schema = input.schema(); + let window_expr = node + .window_expr .iter() - .map(|p| { - datafusion_proto::physical_plan::from_proto::parse_physical_expr( - p, - ctx, - schema.as_ref(), - self, + .map(|we| { + datafusion_proto::physical_plan::from_proto::parse_physical_window_expr( + we, + &decode_ctx, + input_schema.as_ref(), + &converter, ) }) - .collect::, DataFusionError>>()?; - Ok(Arc::new(PerPartitionFilterExec::try_new( + .collect::, _>>()?; + Ok(Arc::new(PartitionedBoundedWindowAggExec::try_new( + window_expr, input.clone(), - predicates, )?)) } } @@ -779,6 +874,46 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; + Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + let stage_id = exec.stage_id as u32; + let mut partition = vec![]; + for location in &exec.partition { + partition.push(protobuf::ShuffleReaderPartition { + location: location + .iter() + .map(|l| { + l.clone().try_into().map_err(|e| { + DataFusionError::Internal(format!( + "Fail to get partition location due to {e:?}" + )) + }) + }) + .collect::, _>>()?, + }); + } + let converter = DefaultPhysicalProtoConverter {}; + let merge_ordering = serialize_physical_sort_exprs( + exec.merge_ordering().iter().cloned(), + self.default_codec.as_ref(), + &converter, + )?; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::RangeShuffleReader( + protobuf::RangeShuffleReaderExecNode { + stage_id, + partition, + schema: Some(exec.schema().as_ref().try_into()?), + merge_ordering, + }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode range shuffle reader execution plan: {e:?}" + )) + })?; + Ok(()) } else if let Some(exec) = node.downcast_ref::() { let converter = DefaultPhysicalProtoConverter {}; @@ -900,25 +1035,79 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; Ok(()) - } else if let Some(exec) = node.downcast_ref::() { - let predicates = exec - .predicates() + } else if let Some(exec) = node.downcast_ref::() { + let cuts = exec.cuts().ok_or_else(|| { + DataFusionError::Internal( + "RangeFilterExec: cannot serialize before resolve_cuts()".into(), + ) + })?; + let routing_expr = + datafusion_proto::physical_plan::to_proto::serialize_physical_expr( + exec.routing_expr(), + self.default_codec.as_ref(), + )?; + let partition_indices: Vec = + exec.partition_indices().iter().map(|&i| i as u32).collect(); + let cuts_proto = cuts .iter() - .map(|p| { - datafusion_proto::physical_plan::to_proto::serialize_physical_expr( - p, + .map(datafusion_proto_common::ScalarValue::try_from) + .collect::, _>>() + .map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode RangeFilterExec cuts: {e:?}" + )) + })?; + let halo_lo = datafusion_proto_common::ScalarValue::try_from(&exec.halo_lo()) + .map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode RangeFilterExec halo_lo: {e:?}" + )) + })?; + let halo_hi = datafusion_proto_common::ScalarValue::try_from(&exec.halo_hi()) + .map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode RangeFilterExec halo_hi: {e:?}" + )) + })?; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::RangeFilter( + protobuf::RangeFilterExecNode { + routing_expr: Some(routing_expr), + cuts: cuts_proto, + halo_lo: Some(halo_lo), + halo_hi: Some(halo_hi), + partition_indices, + }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode RangeFilterExec: {e:?}" + )) + })?; + Ok(()) + } else if let Some(exec) = node.downcast_ref::() + { + let converter = DefaultPhysicalProtoConverter {}; + let window_expr = exec + .window_expr() + .iter() + .map(|we| { + datafusion_proto::physical_plan::to_proto::serialize_physical_window_expr( + we, self.default_codec.as_ref(), + &converter, ) }) - .collect::, DataFusionError>>()?; + .collect::, _>>()?; let proto = protobuf::BallistaPhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::PerPartitionFilter( - protobuf::PerPartitionFilterExecNode { predicates }, + physical_plan_type: Some(PhysicalPlanType::PartitionedBoundedWindowAgg( + protobuf::PartitionedBoundedWindowAggExecNode { window_expr }, )), }; proto.encode(buf).map_err(|e| { DataFusionError::Internal(format!( - "failed to encode PerPartitionFilterExec: {e:?}" + "failed to encode PartitionedBoundedWindowAggExec: {e:?}" )) })?; Ok(()) @@ -1168,6 +1357,67 @@ mod test { ); } + /// `RangeShuffleReaderExec` carries its merge ordering across the wire — + /// the reader's k-way merge machinery is inert without it, and the + /// scheduler side plants the reader with the child's declared ordering. + #[tokio::test] + async fn test_range_shuffle_reader_exec_roundtrip() { + use datafusion::arrow::compute::SortOptions; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + + let schema = create_test_schema(); + let sort_expr = PhysicalSortExpr { + expr: col("id", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }; + let merge_ordering = LexOrdering::new(vec![sort_expr.clone()]).unwrap(); + + let original = RangeShuffleReaderExec::try_new( + 7, + vec![vec![]; 4], + schema.clone(), + merge_ordering.clone(), + ) + .unwrap(); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec.try_encode(Arc::new(original), &mut buf).unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); + let decoded = decoded_plan + .downcast_ref::() + .expect("Expected RangeShuffleReaderExec"); + + assert_eq!(decoded.stage_id, 7); + assert_eq!(decoded.schema().as_ref(), schema.as_ref()); + assert_eq!(decoded.partition.len(), 4, "partition shape must survive"); + assert_eq!( + decoded.merge_ordering().len(), + 1, + "merge_ordering must round-trip" + ); + assert_eq!( + decoded.merge_ordering().first().expr.to_string(), + sort_expr.expr.to_string(), + ); + // The ordering must land on `PlanProperties.eq_properties` — downstream + // consumers (BWAG, SMJ build side) read it there. + let advertised = decoded + .properties() + .eq_properties + .oeq_class() + .iter() + .next() + .cloned() + .expect("advertised ordering"); + assert_eq!(advertised, merge_ordering); + } + /// The sort shuffle writer's per-task memory budget must survive the /// protobuf round trip so a session override reaches the executor where the /// writer actually runs (issue #2089). A non-default value and the special @@ -1799,43 +2049,37 @@ mod test { ); } - /// `PerPartitionFilterExec` round-trips through the codec: three input - /// partitions with three distinct range predicates re-materialize - /// on the other side with the same predicate strings in the same - /// order. + /// `RangeFilterExec` round-trips through the codec: routing expr, cuts, + /// and halo widths reappear identical on the other side. #[tokio::test] - async fn test_per_partition_filter_exec_roundtrip() { - use crate::execution_plans::PerPartitionFilterExec; - use datafusion::logical_expr::Operator; + async fn test_range_filter_exec_roundtrip() { + use crate::execution_plans::RangeFilterExec; use datafusion::physical_expr::PhysicalExpr; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::repartition::RepartitionExec; - use datafusion::scalar::ScalarValue; - let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); let source: Arc = Arc::new( datafusion::physical_plan::empty::EmptyExec::new(schema.clone()), ); let input: Arc = Arc::new( RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(3)).unwrap(), ); - let range = |lo: i64, hi: i64| -> Arc { - let c: Arc = Arc::new(Column::new("v", 0)); - let ge: Arc = Arc::new(BinaryExpr::new( - c.clone(), - Operator::GtEq, - Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))), - )); - let lt: Arc = Arc::new(BinaryExpr::new( - c, - Operator::Lt, - Arc::new(Literal::new(ScalarValue::Int64(Some(hi)))), - )); - Arc::new(BinaryExpr::new(ge, Operator::And, lt)) - }; - let predicates = vec![range(0, 10), range(10, 20), range(20, 30)]; - let original = - PerPartitionFilterExec::try_new(input.clone(), predicates.clone()).unwrap(); + use datafusion::scalar::ScalarValue; + let routing_expr: Arc = Arc::new(Column::new("v", 0)); + let cuts: Vec = vec![ + ScalarValue::Float64(Some(10.0)), + ScalarValue::Float64(Some(20.0)), + ]; + let original = RangeFilterExec::try_new_resolved( + input.clone(), + routing_expr.clone(), + cuts.clone(), + ScalarValue::Float64(Some(3.0)), + ScalarValue::Float64(Some(0.0)), + ) + .unwrap(); let codec = BallistaPhysicalExtensionCodec::default(); let mut buf: Vec = vec![]; @@ -1844,16 +2088,45 @@ mod test { let ctx = SessionContext::new().task_ctx(); let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap(); let decoded = decoded_plan - .downcast_ref::() - .expect("Expected PerPartitionFilterExec"); - assert_eq!(decoded.predicates().len(), 3); - for (k, expected) in predicates.iter().enumerate() { - assert_eq!( - decoded.predicates()[k].to_string(), - expected.to_string(), - "predicate {k} mismatched after roundtrip", - ); - } + .downcast_ref::() + .expect("Expected RangeFilterExec"); + assert_eq!(decoded.cuts().unwrap(), cuts); + assert_eq!(decoded.halo_lo(), ScalarValue::Float64(Some(3.0))); + assert_eq!(decoded.halo_hi(), ScalarValue::Float64(Some(0.0))); + assert_eq!(decoded.routing_expr().to_string(), routing_expr.to_string()); + } + + /// A pending RangeFilterExec (unresolved cuts) refuses to serialize — + /// wire plans must always carry resolved cuts. + #[tokio::test] + async fn test_range_filter_exec_pending_refuses_serialization() { + use crate::execution_plans::RangeFilterExec; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::repartition::RepartitionExec; + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let source: Arc = Arc::new( + datafusion::physical_plan::empty::EmptyExec::new(schema.clone()), + ); + let input: Arc = Arc::new( + RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(2)).unwrap(), + ); + use datafusion::scalar::ScalarValue; + let routing_expr: Arc = Arc::new(Column::new("v", 0)); + let pending = RangeFilterExec::try_new_pending( + input, + routing_expr, + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + ) + .unwrap(); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + let err = codec.try_encode(Arc::new(pending), &mut buf).unwrap_err(); + assert!(err.to_string().contains("before resolve_cuts")); } /// `BufferExec` in `Dam` mode round-trips through the codec. diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 4f8e6663d8..e673a51c8e 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}; @@ -146,6 +148,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 539b45c6b8..da6fb413c2 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -20,14 +20,16 @@ use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use crate::state::aqe::planner::AdaptiveStageInfo; 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; @@ -64,66 +66,81 @@ impl BallistaAdapter { self.inputs.push(stage_id); 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 { @@ -141,6 +158,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() @@ -168,6 +186,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() @@ -194,3 +213,129 @@ impl BallistaAdapter { } } } + +/// Walk `plan` and resolve every pending [`RangeFilterExec`]'s cuts 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. +/// +/// Errors if a `RangeFilterExec` is still pending after the walk: the +/// rule promised cuts 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.cuts().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 cuts_sv: Vec = cuts + .into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect(); + rf.resolve_cuts(cuts_sv)?; + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) +} + +#[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/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index abb33a8414..48e74b4f4e 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -38,7 +38,7 @@ use std::sync::{Arc, atomic::AtomicI64}; /// of this exchange. Written after the range-repartition-producing stage /// completes and its runtime-stats sketches are merged; read at /// task-specialization time to build per-downstream-partition range filters -/// (see `PerPartitionFilterExec`). +/// (see `RangeFilterExec`). /// /// `cuts` are `K - 1` monotone `f64` boundaries expressed in the value space /// of `routing_expr`; downstream partition `k` owns `[cuts[k-1], cuts[k])` @@ -98,7 +98,7 @@ pub struct ExchangeExec { /// range-repartition op (URRE or ORRE). Stored when /// the range-repartition-producing stage completes and its per-sub-part /// quantile sketches have been merged. Read at task-specialization time - /// to build `PerPartitionFilterExec` predicates for downstream stage `N+1`. + /// to build `RangeFilterExec` predicates for downstream stage `N+1`. /// /// `None` on any exchange that isn't downstream of a range repartition range_repartition_routing: Arc>>, @@ -484,7 +484,7 @@ mod range_repartition_routing_tests { //! writes here at range-repartition-stage completion; task //! specialization reads it back at //! `BallistaAdapter::transform_children` time to wrap the - //! ShuffleReader in a `PerPartitionFilterExec`. Neither side is + //! ShuffleReader in a `RangeFilterExec`. Neither side is //! exercised end-to-end without the URRE-inserting rule (a follow-up //! PR), so tests here cover the slot itself: roundtrip through //! `resolve_range_repartition_routing` → `range_repartition_routing()`, 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..0920f29033 --- /dev/null +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs @@ -0,0 +1,555 @@ +// 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) +//! RuntimeStatsExec #1 (local sketch — feeds ORRE's cut walker) +//! SortExec (unchanged, preserve_partitioning=true) +//! +//! ``` +//! +//! Any SPM the DF planner had inserted above BWAG for its +//! `SinglePartition` requirement is dropped: the wrapper flips that +//! declaration to `UnspecifiedDistribution`, and `EnforceDistribution` +//! (running later in the AQE pass) 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::PhysicalSortExpr; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +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:?}" + ), + } +} + +/// Splice `RSE#1 → ORRE → RSE#2 → RangeFilterExec_wide` below BWAG's +/// sorted source and wrap the BWAG in a `PartitionedBoundedWindowAggExec`. +/// Accepts either shape: +/// +/// ```text +/// BWAG → SPM → (multi-partition input) +/// BWAG → (single-partition input) +/// ``` +/// +/// Any pre-existing SPM below the BWAG is dropped — the wrapper's +/// `UnspecifiedDistribution` declaration keeps `EnforceDistribution` from +/// putting one back. +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" + ); + }; + // Peel off any existing SPM to reveal the sorted source underneath; + // otherwise the immediate child *is* the sorted source. + let source: Arc = if immediate + .downcast_ref::() + .is_some() + { + let spm_children = immediate.children(); + let [inner] = spm_children.as_slice() else { + return datafusion::common::internal_err!( + "ParallelWindowRule: SPM must have exactly 1 child" + ); + }; + (*inner).clone() + } else { + (*immediate).clone() + }; + let source_schema = 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( + source, + Some(vec![sort_expr.clone()]), + )?); + let orre: Arc = Arc::new(OrderedRangeRepartitionExec::try_new( + 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 690a5f2bcd..2cb3b43b6a 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; @@ -533,6 +533,12 @@ impl AdaptivePlanner { // changing plan before other built in optimizers kick in physical_optimizers.push(Arc::new(PropagateEmptyExecRule::default())); + // Rewrite bounded RANGE-frame windows into a range-shuffle so BWAG's + // single-partition constraint is not a serial bottleneck. Must 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)); + // select actual join implementation based on current runtime information physical_optimizers .push(Arc::new(SelectJoinRule::new(plan_id_generator.clone()))); diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index c51cff0a1d..8ec34c60b7 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -234,7 +234,7 @@ async fn should_skip_coalesce_when_leaf_has_range_repartition_routing() // runnable cache — `finalise_stage_internal` will remove it. Mirrors // `SchedulerAqe`: `set_repartition_routing` runs before // `resolve_stage_partitions` at the boundary. 7 cuts → K=8, matching - // the hash-partitioning's M=8 so the adapter's `PerPartitionFilterExec` + // the hash-partitioning's M=8 so the adapter's `RangeFilterExec` // builds cleanly; the point of the test is the rule's bail decision, // not a realistic end-to-end range-repartition query. planner.set_repartition_routing( diff --git a/ballista/scheduler/src/state/aqe/test/range_repartition.rs b/ballista/scheduler/src/state/aqe/test/range_repartition.rs index 8d25a8b83e..cc15edfaee 100644 --- a/ballista/scheduler/src/state/aqe/test/range_repartition.rs +++ b/ballista/scheduler/src/state/aqe/test/range_repartition.rs @@ -19,7 +19,7 @@ //! through `AdaptivePlanner` — DER inserts the boundary //! `ExchangeExec`, `set_repartition_routing` parks the recovered //! cuts on it, and `cut_partitions` duplicates straddlers so -//! downstream can inject a `PerPartitionFilterExec` to trim them. +//! downstream can inject a `RangeFilterExec` to trim them. use crate::state::aqe::execution_plan::RangeRepartitionRouting; use crate::state::aqe::planner::AdaptivePlanner; @@ -169,7 +169,7 @@ async fn routing_parks_when_range_repartition_is_plan_root() assert!( plan_str.contains("range_repartition_cuts=1"), "cuts must be parked on the boundary ExchangeExec so downstream \ - gets a PerPartitionFilterExec — actual plan:\n{plan_str}" + gets a RangeFilterExec — actual plan:\n{plan_str}" ); Ok(()) diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index 68811762e0..ab597271d9 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,28 +86,21 @@ 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: cuts describe the K global partitions and are shared + // across all restricted views; only the operator's `partition_indices` + // mapping changes. We restrict the child first, then hand it + + // `partitions` (task-local → global mapping) to the operator, which + // re-derives its new mapping. + 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 - .iter() - .map(|&part_idx| ppf.predicates()[part_idx].clone()) - .collect(); - return Ok(Arc::new(PerPartitionFilterExec::try_new( - new_child, - new_predicates, - )?)); + return Ok(Arc::new(rf.restrict_partitions(new_child, partitions)?)); } // UnionExec: parent partition `p` maps to exactly one child's local @@ -288,6 +283,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 +663,20 @@ 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 remap + /// its `partition_indices` by the same indices in the same order as its + /// child. Cuts stay whole (they describe the K global partitions); + /// only the local→global mapping changes. #[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_partition_indices_are_remapped_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 +686,33 @@ 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)); + let cuts: Vec = vec![ + ScalarValue::Float64(Some(100.0)), + ScalarValue::Float64(Some(200.0)), + ScalarValue::Float64(Some(300.0)), + ]; let plan: Arc = Arc::new( - PerPartitionFilterExec::try_new( + RangeFilterExec::try_new_resolved( Arc::new(reader) as Arc, - predicates.clone(), + routing_expr, + cuts.clone(), + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), ) .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"); + assert_eq!(rf.partition_indices(), &[1, 3]); + assert_eq!(rf.cuts().unwrap(), cuts); // 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/benchmarks/queries/h2o/groupby.sql b/benchmarks/queries/h2o/groupby.sql new file mode 100644 index 0000000000..4fae7a1381 --- /dev/null +++ b/benchmarks/queries/h2o/groupby.sql @@ -0,0 +1,19 @@ +SELECT id1, SUM(v1) AS v1 FROM x GROUP BY id1; + +SELECT id1, id2, SUM(v1) AS v1 FROM x GROUP BY id1, id2; + +SELECT id3, SUM(v1) AS v1, AVG(v3) AS v3 FROM x GROUP BY id3; + +SELECT id4, AVG(v1) AS v1, AVG(v2) AS v2, AVG(v3) AS v3 FROM x GROUP BY id4; + +SELECT id6, SUM(v1) AS v1, SUM(v2) AS v2, SUM(v3) AS v3 FROM x GROUP BY id6; + +SELECT id4, id5, MEDIAN(v3) AS median_v3, STDDEV(v3) AS sd_v3 FROM x GROUP BY id4, id5; + +SELECT id3, MAX(v1) - MIN(v2) AS range_v1_v2 FROM x GROUP BY id3; + +SELECT id6, largest2_v3 FROM (SELECT id6, v3 AS largest2_v3, ROW_NUMBER() OVER (PARTITION BY id6 ORDER BY v3 DESC) AS order_v3 FROM x WHERE v3 IS NOT NULL) sub_query WHERE order_v3 <= 2; + +SELECT id2, id4, POWER(CORR(v1, v2), 2) AS r2 FROM x GROUP BY id2, id4; + +SELECT id1, id2, id3, id4, id5, id6, SUM(v3) AS v3, COUNT(*) AS count FROM x GROUP BY id1, id2, id3, id4, id5, id6; \ No newline at end of file diff --git a/benchmarks/queries/h2o/join.sql b/benchmarks/queries/h2o/join.sql new file mode 100644 index 0000000000..84cd661fdd --- /dev/null +++ b/benchmarks/queries/h2o/join.sql @@ -0,0 +1,9 @@ +SELECT x.id1, x.id2, x.id3, x.id4 as xid4, small.id4 as smallid4, x.id5, x.id6, x.v1, small.v2 FROM x INNER JOIN small ON x.id1 = small.id1; + +SELECT x.id1 as xid1, medium.id1 as mediumid1, x.id2, x.id3, x.id4 as xid4, medium.id4 as mediumid4, x.id5 as xid5, medium.id5 as mediumid5, x.id6, x.v1, medium.v2 FROM x INNER JOIN medium ON x.id2 = medium.id2; + +SELECT x.id1 as xid1, medium.id1 as mediumid1, x.id2, x.id3, x.id4 as xid4, medium.id4 as mediumid4, x.id5 as xid5, medium.id5 as mediumid5, x.id6, x.v1, medium.v2 FROM x LEFT JOIN medium ON x.id2 = medium.id2; + +SELECT x.id1 as xid1, medium.id1 as mediumid1, x.id2, x.id3, x.id4 as xid4, medium.id4 as mediumid4, x.id5 as xid5, medium.id5 as mediumid5, x.id6, x.v1, medium.v2 FROM x JOIN medium ON x.id5 = medium.id5; + +SELECT x.id1 as xid1, large.id1 as largeid1, x.id2 as xid2, large.id2 as largeid2, x.id3, x.id4 as xid4, large.id4 as largeid4, x.id5 as xid5, large.id5 as largeid5, x.id6 as xid6, large.id6 as largeid6, x.v1, large.v2 FROM x JOIN large ON x.id3 = large.id3; \ No newline at end of file diff --git a/benchmarks/queries/h2o/window.sql b/benchmarks/queries/h2o/window.sql new file mode 100644 index 0000000000..346a8e4713 --- /dev/null +++ b/benchmarks/queries/h2o/window.sql @@ -0,0 +1,150 @@ +-- Basic Window +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER () AS window_basic +FROM large; + +-- Sorted Window +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (ORDER BY id3) AS first_order_by, + row_number() OVER (ORDER BY id3) AS row_number_order_by +FROM large; + +-- PARTITION BY +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (PARTITION BY id1) AS sum_by_id1, + sum(v2) OVER (PARTITION BY id2) AS sum_by_id2, + sum(v2) OVER (PARTITION BY id3) AS sum_by_id3 +FROM large; + +-- PARTITION BY ORDER BY +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (PARTITION BY id2 ORDER BY id3) AS first_by_id2_ordered_by_id3 +FROM large; + +-- Lead and Lag +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (ORDER BY id3 ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING) AS my_lag, + first_value(v2) OVER (ORDER BY id3 ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) AS my_lead +FROM large; + +-- Moving Averages +SELECT + id1, + id2, + id3, + v2, + avg(v2) OVER (ORDER BY id3 ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) AS my_moving_average +FROM large; + +-- Rolling Sum +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS my_rolling_sum +FROM large; + +-- RANGE BETWEEN +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) AS my_range_between +FROM large; + +-- First PARTITION BY ROWS BETWEEN +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING) AS my_lag_by_id2, + first_value(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) AS my_lead_by_id2 +FROM large; + +-- Moving Averages PARTITION BY +SELECT + id1, + id2, + id3, + v2, + avg(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) AS my_moving_average_by_id2 +FROM large; + +-- Rolling Sum PARTITION BY +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS my_rolling_sum_by_id2 +FROM large; + +-- RANGE BETWEEN PARTITION BY +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (PARTITION BY id2 ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) AS my_range_between_by_id2 +FROM large; + +-- Window Top-N (ROW_NUMBER top-2 per partition) +SELECT id2, largest2_v2 FROM ( + SELECT id2, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N partition cardinality sweep (id3 % N gives N distinct partitions). +-- These exercise PartitionedTopKExec across cardinalities to validate it stays +-- competitive with the SortExec+Filter baseline as partition count grows. +-- Window Top-N: 100 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 100 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 100 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N: 1,000 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 1000 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 1000 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N: 10,000 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 10000 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 10000 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N: 100,000 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 100000 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 100000 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; diff --git a/benchmarks/src/bin/h2o.rs b/benchmarks/src/bin/h2o.rs new file mode 100644 index 0000000000..09498489fe --- /dev/null +++ b/benchmarks/src/bin/h2o.rs @@ -0,0 +1,451 @@ +// 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. + +//! H2O `db-benchmark` runner for Ballista. +//! +//! Ports `arrow-datafusion/benchmarks/src/h2o.rs` to Ballista. Supports both +//! single-process DataFusion and distributed Ballista execution over the h2o +//! groupby, join, and window query files. + +use ballista::extension::SessionConfigExt; +use ballista::prelude::SessionContextExt; +use ballista_core::object_store::{ + session_config_with_s3_support, session_state_with_s3_support, +}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::prelude::*; +use std::path::{Path, PathBuf}; +use std::time::Instant; +use structopt::StructOpt; + +#[cfg(feature = "mimalloc")] +#[global_allocator] +static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; + +#[derive(Debug, StructOpt, Clone)] +struct DataFusionBenchmarkOpt { + /// Query number (1..=max). If not specified, runs every query in the file. + #[structopt(short, long)] + query: Option, + + /// Path to the queries SQL file. Suite is inferred from the file name + /// (`groupby.sql`, `join.sql`, `window.sql`). + #[structopt( + parse(from_os_str), + short = "r", + long = "queries-path", + default_value = "benchmarks/queries/h2o/groupby.sql" + )] + queries_path: PathBuf, + + /// Path to the primary data file (used for `groupby.sql` / `window.sql`). + #[structopt( + parse(from_os_str), + short = "p", + long = "path", + default_value = "benchmarks/data/h2o/G1_1e7_1e7_100_0.csv" + )] + path: PathBuf, + + /// Comma-separated join data files, in order: x, small, medium, large. + /// Used for `join.sql` and `window.sql` (window uses the `large` file). + #[structopt( + short = "j", + long = "join-paths", + default_value = "benchmarks/data/h2o/J1_1e7_NA_0.csv,benchmarks/data/h2o/J1_1e7_1e1_0.csv,benchmarks/data/h2o/J1_1e7_1e4_0.csv,benchmarks/data/h2o/J1_1e7_1e7_NA.csv" + )] + join_paths: String, + + /// Iterations per query. + #[structopt(short = "i", long = "iterations", default_value = "1")] + iterations: usize, + + /// Number of partitions to process in parallel. + #[structopt(short = "n", long = "partitions", default_value = "2")] + partitions: usize, + + /// Batch size when reading CSV or Parquet files. + #[structopt(short = "s", long = "batch-size", default_value = "8192")] + batch_size: usize, + + /// Print plans and query text. + #[structopt(short, long)] + debug: bool, + + /// Print the physical plan and exit without running the query. + #[structopt(long)] + explain: bool, +} + +#[derive(Debug, StructOpt, Clone)] +struct BallistaBenchmarkOpt { + /// Query number (1..=max). If not specified, runs every query in the file. + #[structopt(short, long)] + query: Option, + + /// Path to the queries SQL file. Suite is inferred from the file name + /// (`groupby.sql`, `join.sql`, `window.sql`). + #[structopt( + parse(from_os_str), + short = "r", + long = "queries-path", + default_value = "benchmarks/queries/h2o/groupby.sql" + )] + queries_path: PathBuf, + + /// Path to the primary data file (used for `groupby.sql` / `window.sql`). + /// May be a local path or an object-store URL such as `s3://bucket/prefix`. + #[structopt(short = "p", long = "path")] + path: String, + + /// Comma-separated join data files, in order: x, small, medium, large. + /// Used for `join.sql` and `window.sql` (window uses the `large` file). + /// Paths may be local or object-store URLs. + #[structopt(short = "j", long = "join-paths", default_value = "")] + join_paths: String, + + /// Iterations per query. + #[structopt(short = "i", long = "iterations", default_value = "1")] + iterations: usize, + + /// Number of partitions to process in parallel. + #[structopt(short = "n", long = "partitions", default_value = "2")] + partitions: usize, + + /// Batch size when reading CSV or Parquet files. + #[structopt(short = "s", long = "batch-size", default_value = "8192")] + batch_size: usize, + + /// Ballista scheduler host. + #[structopt(long = "host")] + host: String, + + /// Ballista scheduler port. + #[structopt(long = "port")] + port: u16, + + /// Configuration overrides in `key=value` format. Repeatable. + #[structopt(short = "c", long = "config", number_of_values = 1)] + config_overrides: Vec, + + /// Print plans and query text. + #[structopt(short, long)] + debug: bool, + + /// Print the physical plan and exit without running the query. + #[structopt(long)] + explain: bool, +} + +#[derive(Debug, StructOpt)] +#[structopt(name = "h2o", about = "H2O db-benchmark for Ballista")] +enum H2oOpt { + #[structopt(name = "datafusion")] + DataFusion(DataFusionBenchmarkOpt), + #[structopt(name = "ballista")] + Ballista(BallistaBenchmarkOpt), +} + +#[tokio::main] +async fn main() -> Result<()> { + env_logger::init(); + match H2oOpt::from_args() { + H2oOpt::DataFusion(opt) => run_datafusion(opt).await, + H2oOpt::Ballista(opt) => run_ballista(opt).await, + } +} + +async fn run_datafusion(opt: DataFusionBenchmarkOpt) -> Result<()> { + println!("Running h2o (datafusion) with: {opt:?}"); + let queries = AllQueries::from_file(&opt.queries_path)?; + let query_range = query_range(&queries, opt.query); + + let config = SessionConfig::new() + .with_target_partitions(opt.partitions) + .with_batch_size(opt.batch_size); + let ctx = SessionContext::new_with_config(config); + + let suite = Suite::from_queries_path(&opt.queries_path)?; + let paths = SuitePaths { + primary: opt.path.to_string_lossy().into_owned(), + join_csv: opt.join_paths.clone(), + }; + register_h2o_tables(&ctx, suite, &paths).await?; + + if opt.explain { + return explain_queries(&ctx, &queries, query_range).await; + } + + run_queries( + &ctx, + &queries, + query_range, + opt.iterations, + opt.debug, + "datafusion", + ) + .await +} + +async fn run_ballista(opt: BallistaBenchmarkOpt) -> Result<()> { + println!("Running h2o (ballista) with: {opt:?}"); + let queries = AllQueries::from_file(&opt.queries_path)?; + let query_range = query_range(&queries, opt.query); + + let mut config = session_config_with_s3_support() + .with_target_partitions(opt.partitions) + .with_batch_size(opt.batch_size) + .with_ballista_job_name("h2o benchmark"); + + for kv in &opt.config_overrides { + match kv.split_once('=') { + Some((key, value)) => { + if let Err(err) = config.options_mut().set(key.trim(), value.trim()) { + println!("Warning: could not set config '{kv}': {err}"); + } + } + None => println!( + "Warning: ignoring invalid config override '{kv}'. Expected key=value" + ), + } + } + + let state = session_state_with_s3_support(config)?; + let address = format!("df://{}:{}", opt.host, opt.port); + let ctx = SessionContext::remote_with_state(&address, state).await?; + + let suite = Suite::from_queries_path(&opt.queries_path)?; + let paths = SuitePaths { + primary: opt.path.clone(), + join_csv: opt.join_paths.clone(), + }; + register_h2o_tables(&ctx, suite, &paths).await?; + + if opt.explain { + return explain_queries(&ctx, &queries, query_range).await; + } + + run_queries( + &ctx, + &queries, + query_range, + opt.iterations, + opt.debug, + "ballista", + ) + .await +} + +async fn explain_queries( + ctx: &SessionContext, + queries: &AllQueries, + query_range: std::ops::RangeInclusive, +) -> Result<()> { + use datafusion::physical_plan::displayable; + for query_id in query_range { + let sql = queries.get(query_id)?; + println!("Q{query_id}: {sql}"); + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + println!( + "=== Physical plan (Q{query_id}) ===\n{}", + displayable(plan.as_ref()).indent(true) + ); + } + Ok(()) +} + +async fn run_queries( + ctx: &SessionContext, + queries: &AllQueries, + query_range: std::ops::RangeInclusive, + iterations: usize, + debug: bool, + label: &str, +) -> Result<()> { + let mut total_secs = 0.0; + for query_id in query_range { + let sql = queries.get(query_id)?; + println!("Q{query_id}: {sql}"); + + let mut per_iter_secs = Vec::with_capacity(iterations); + let mut row_count = 0usize; + for iteration in 1..=iterations { + let start = Instant::now(); + let batches = ctx.sql(sql).await?.collect().await?; + let elapsed = start.elapsed().as_secs_f64(); + row_count = batches.iter().map(|batch| batch.num_rows()).sum(); + per_iter_secs.push(elapsed); + println!( + "Query {query_id} iteration {iteration} took {:.3} s and returned {row_count} rows ({label})", + elapsed + ); + } + let avg = per_iter_secs.iter().sum::() / per_iter_secs.len() as f64; + println!("Query {query_id} avg time: {avg:.3} s ({row_count} rows)"); + total_secs += avg; + + if debug { + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + println!("=== Optimized logical plan ===\n{plan:?}\n"); + } + } + println!("Total avg time across queries: {total_secs:.3} s"); + Ok(()) +} + +fn query_range( + queries: &AllQueries, + query: Option, +) -> std::ops::RangeInclusive { + match query { + Some(query_id) => query_id..=query_id, + None => queries.min_id()..=queries.max_id(), + } +} + +#[derive(Copy, Clone, Debug)] +enum Suite { + Groupby, + Join, + Window, +} + +impl Suite { + fn from_queries_path(path: &Path) -> Result { + let name = path + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or_default(); + match name { + "groupby.sql" => Ok(Suite::Groupby), + "join.sql" => Ok(Suite::Join), + "window.sql" => Ok(Suite::Window), + other => Err(DataFusionError::Plan(format!( + "unknown h2o suite {other:?} — expected groupby.sql, join.sql, or window.sql" + ))), + } + } +} + +struct SuitePaths { + primary: String, + join_csv: String, +} + +async fn register_h2o_tables( + ctx: &SessionContext, + suite: Suite, + paths: &SuitePaths, +) -> Result<()> { + match suite { + Suite::Groupby => register_table(ctx, "x", &paths.primary).await, + Suite::Join => { + let join_paths: Vec<&str> = paths.join_csv.split(',').collect(); + if join_paths.len() != 4 { + return Err(DataFusionError::Plan(format!( + "join suite needs 4 comma-separated paths, got {}", + join_paths.len() + ))); + } + for (table, path) in ["x", "small", "medium", "large"] + .iter() + .zip(join_paths.iter()) + { + register_table(ctx, table, path.trim()).await?; + } + Ok(()) + } + Suite::Window => { + // The window suite uses only the `large` table from the join dataset. + let large = paths.join_csv.split(',').nth(3).ok_or_else(|| { + DataFusionError::Plan( + "window suite: --join-paths must contain 4 comma-separated paths; \ + the fourth is registered as `large`" + .to_string(), + ) + })?; + register_table(ctx, "large", large.trim()).await + } + } +} + +async fn register_table(ctx: &SessionContext, table: &str, path: &str) -> Result<()> { + let extension = Path::new(path) + .extension() + .and_then(|s| s.to_str()) + .unwrap_or_default(); + match extension { + "csv" => ctx + .register_csv(table, path, CsvReadOptions::default()) + .await + .map_err(|err| { + DataFusionError::Context( + format!("registering table {table:?} from {path}"), + Box::new(err), + ) + }), + "parquet" => ctx + .register_parquet(table, path, ParquetReadOptions::default()) + .await + .map_err(|err| { + DataFusionError::Context( + format!("registering table {table:?} from {path}"), + Box::new(err), + ) + }), + other => Err(DataFusionError::Plan(format!( + "unsupported extension {other:?} for {path}" + ))), + } +} + +struct AllQueries { + queries: Vec, +} + +impl AllQueries { + fn from_file(path: &Path) -> Result { + let contents = std::fs::read_to_string(path).map_err(|err| { + DataFusionError::Execution(format!("reading {path:?}: {err}")) + })?; + // Queries are separated by blank lines — matches the datafusion h2o + // runner exactly so query indices line up 1:1 across both binaries. + let queries = contents.split("\n\n").map(str::to_owned).collect(); + Ok(Self { queries }) + } + + fn get(&self, query_id: usize) -> Result<&str> { + self.queries + .get(query_id - 1) + .map(String::as_str) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "invalid query id {query_id}. Must be between {} and {}", + self.min_id(), + self.max_id() + )) + }) + } + + fn min_id(&self) -> usize { + 1 + } + + fn max_id(&self) -> usize { + self.queries.len() + } +} diff --git a/docs/developer/parallel-range-window.md b/docs/developer/parallel-range-window.md new file mode 100644 index 0000000000..5f0e66e0eb --- /dev/null +++ b/docs/developer/parallel-range-window.md @@ -0,0 +1,77 @@ + + +# 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. |