phillipleblanc commented on code in PR #2123: URL: https://github.com/apache/datafusion-ballista/pull/2123#discussion_r3629913284
########## ballista/core/src/execution_plans/unordered_range_repartition.rs: ########## @@ -0,0 +1,859 @@ +// 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. + +//! Value-range router over unordered inputs. Reads `P` input partitions with +//! no sort assumption, evaluates the first ORDER BY expression per row, and +//! routes each row to one of `K` output partitions under the half-open +//! convention: partition `p` owns `[cut[p-1], cut[p])` with virtual `-∞`/`+∞` +//! sentinels on the ends. +//! +//! ```text +//! P input outputs (K = 4) +//! partitions +//! ─┐ ┌─▶ 0: key < c₁ +//! ─┼──▶ RuntimeStatsExec ──▶ BufferExec ──▶ UnorderedRRE ─┼─▶ 1: c₁ ≤ key < c₂ +//! ─┤ (T-Digest tap) (Dam, (K = 4) ├─▶ 2: c₂ ≤ key < c₃ +//! ─┘ whitelisted) │ └─▶ 3: c₃ ≤ key +//! ▲ │ +//! └────────── walker ◀────────────┘ +//! (descends past BufferExec, +//! matches on routing expr, +//! reads K−1 quantile cuts) +//! ``` +//! +//! # Dynamic discovery +//! +//! The point of this operator is that boundaries are discovered at +//! runtime, not baked in at plan time. On the first batch to arrive from +//! any input partition, the operator walks its own child subtree to find a +//! sibling [`RuntimeStatsExec`], snapshots its `merged_quantile_sketch()`, +//! and computes `K - 1` quantile cuts at `1/K, 2/K, ..., (K-1)/K`. All +//! batches (including the one that triggered the snapshot) then route +//! through those cuts. +//! +//! # Fallback +//! +//! If no sibling [`RuntimeStatsExec`] exists, or it's in row-count-only mode +//! (no sketch), or the sketch is empty, discovery returns an empty cut set +//! and every row lands in output partition 0 — the natural single-bucket +//! outcome of `boundaries.len() + 1 = 1`. Runtime routing must never crash; +//! degraded-but-alive beats the alternative, and downstream sees an empty +//! stream on the K-1 partitions that got no data. +//! +//! # Type generality +//! +//! The impl hardcodes Float64 downcast internally (that's what DataFusion's +//! T-Digest speaks). The public API and the sibling [`RuntimeStatsExec`] +//! stay type-agnostic; widening to other `Ord` `ScalarValue` types replaces +//! the downcast + boundary computation, no API break. +//! +//! Sibling `OrderedRangeRepartitionExec` (not yet built) handles the sorted +//! case (N sorted → M sorted range-disjoint via k-way merge). See +//! `docs/source/contributors-guide/parallel-window-kll-adaptive.md`. +//! +//! [`RuntimeStatsExec`]: crate::execution_plans::RuntimeStatsExec + +use std::fmt::{self, Debug, Formatter}; +use std::sync::{Arc, Mutex, OnceLock}; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{ + Distribution, EquivalenceProperties, OrderingRequirements, Partitioning, + PhysicalExpr, PhysicalSortExpr, +}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::stream::StreamExt; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; + +use crate::execution_plans::range_repartition_common::{ + broadcast_error, discover_cuts, split_batch_by_range, +}; + +/// Per-output-partition channel capacity. Small = tight backpressure; the +/// classic double-buffering shape (one batch in-flight while consumer works +/// on another). When a downstream drain lags, `send().await` suspends the +/// scatter task, which suspends its input read, which propagates upstream. +/// Bump if profiling shows scatter tasks blocked on `send` more than the +/// downstream is meaningfully doing work. +const CHANNEL_CAPACITY: usize = 2; + +/// Value-range router over unordered inputs. See the module-level docs. +pub struct UnorderedRangeRepartitionExec { + input: Arc<dyn ExecutionPlan>, + /// Lexicographic ORDER BY carried through from the wrapping window + /// operator. `try_new` guarantees at least one element; the first entry + /// (a `Float64` column, until we widen) drives routing. + order_by: Vec<PhysicalSortExpr>, + /// K — number of output partitions. K=1 collapses all P inputs to a + /// single bucket (the same shape discovery-failure fallback produces); + /// larger K spreads rows by value range. + output_partitions: usize, + /// Lazy channel setup guarded by `Mutex`. First `execute()` call spawns + /// the P input-reader tasks and creates the K channels; subsequent + /// `execute(p)` calls take `channels[p]`. + state: Arc<Mutex<DispatchState>>, + properties: Arc<PlanProperties>, +} + +/// Per-exec-instance lazy state. Owns the K receivers between setup and the +/// K per-partition takes. +struct DispatchState { + /// K slots. Each holds `Some(receiver)` after setup, `None` once its + /// output partition has been consumed. + receivers: Vec<Option<mpsc::Receiver<Result<RecordBatch>>>>, + initialized: bool, +} + +impl UnorderedRangeRepartitionExec { + /// Wrap `input`. `order_by` must be non-empty and its first entry must + /// evaluate to `Float64` against `input.schema()`. `output_partitions` + /// is K; any value works, K=1 gives a coalesce-shaped passthrough. + pub fn try_new( + input: Arc<dyn ExecutionPlan>, + order_by: Vec<PhysicalSortExpr>, + output_partitions: usize, + // TODO: support RANGE & ROW halos + ) -> Result<Self> { + let [routing, ..] = order_by.as_slice() else { + return internal_err!( + "UnorderedRangeRepartitionExec requires at least one ORDER BY expression" + ); + }; + let schema = input.schema(); + let routing_type = routing.expr.data_type(&schema)?; + if !matches!(routing_type, DataType::Float64) { + // TODO: support all continuous primitives + return internal_err!( + "UnorderedRangeRepartitionExec routing expression `{}` must be Float64, got {:?}", + routing.expr, + routing_type + ); + } + // TODO: fixed by KLL — a NULL-aware sketch lifts this restriction and + // lets `split_batch_by_range` honor SortOptions::nulls_first properly. + if routing.expr.nullable(&schema)? { + return internal_err!( + "UnorderedRangeRepartitionExec: routing expression `{}` must be non-nullable", + routing.expr + ); + } + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(output_partitions), + input.pipeline_behavior(), + input.boundedness(), + )); + let state = Arc::new(Mutex::new(DispatchState { + receivers: Vec::new(), + initialized: false, + })); + Ok(Self { + input, + order_by, + output_partitions, + state, + properties, + }) + } + + /// Full ORDER BY carried through from the wrapping window operator. + pub fn order_by(&self) -> &[PhysicalSortExpr] { + &self.order_by + } + + /// K — the fixed output partition count. + pub fn output_partitions(&self) -> usize { + self.output_partitions + } +} + +impl Debug for UnorderedRangeRepartitionExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("UnorderedRangeRepartitionExec") + .field("order_by", &self.order_by) + .field("output_partitions", &self.output_partitions) + .finish() + } +} + +impl DisplayAs for UnorderedRangeRepartitionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + let routing = &self.order_by[0]; + write!( + f, + "UnorderedRangeRepartitionExec: routing={} {} → {} partitions", + routing.expr, + if routing.options.descending { + "desc" + } else { + "asc" + }, + self.output_partitions, + ) + } +} + +impl ExecutionPlan for UnorderedRangeRepartitionExec { + fn name(&self) -> &str { + "UnorderedRangeRepartitionExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.properties + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + /// Input distribution is irrelevant — the operator re-routes every row + /// by value range regardless of how the child organized them. + fn required_input_distribution(&self) -> Vec<Distribution> { + vec![Distribution::UnspecifiedDistribution] + } + + /// "Unordered" in the name: no ordering requirement on the child. + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![None] + } + + /// Scattering by value range across K outputs definitively breaks the + /// input order — that's the whole point. + fn maintains_input_order(&self) -> Vec<bool> { + vec![false] + } + + /// Extra input repartitioning above us buys nothing — we already read + /// all P input partitions concurrently and scatter them to K outputs. + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + /// Row count is preserved but redistributed across K new output + /// partitions; per-partition breakdown depends on the runtime sketch, + /// which isn't known at plan time. + fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> { + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } + + /// Every input row is emitted exactly once. Overrides default `Unknown`. + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let [input] = children.as_slice() else { + return internal_err!( + "UnorderedRangeRepartitionExec expects exactly one child, got {}", + children.len() + ); + }; + Ok(Arc::new(UnorderedRangeRepartitionExec::try_new( + input.clone(), + self.order_by.clone(), + self.output_partitions, + )?)) + } + + fn execute( + &self, + partition: usize, + ctx: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let mut state = self + .state + .lock() + .map_err(|_| internal_datafusion_err!("dispatch state mutex poisoned"))?; + if !state.initialized { + let mut senders = Vec::with_capacity(self.output_partitions); + let mut receivers = Vec::with_capacity(self.output_partitions); + for _ in 0..self.output_partitions { + let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY); + senders.push(tx); + receivers.push(Some(rx)); + } + state.receivers = receivers; + state.initialized = true; + let senders: Arc<[mpsc::Sender<Result<RecordBatch>>]> = senders.into(); + let cuts_cell: Arc<OnceLock<Vec<f64>>> = Arc::new(OnceLock::new()); + let input_partitions = self.input.output_partitioning().partition_count(); + let routing_expr = self.order_by[0].expr.clone(); // TODO: KLL for multi-column? + for input_partition in 0..input_partitions { + let child = self.input.clone(); + let senders = senders.clone(); + let cuts_cell = cuts_cell.clone(); + let routing_expr = routing_expr.clone(); + let ctx = ctx.clone(); + let output_partitions = self.output_partitions; + tokio::spawn(async move { Review Comment: Makes sense to me. Agree with keeping things simple. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
