phillipleblanc commented on code in PR #2123: URL: https://github.com/apache/datafusion-ballista/pull/2123#discussion_r3626759882
########## 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: This `JoinHandle` is being discarded, so if `scatter_input_partition` panics then the sender is silently dropped, which would result in a "successful" task that returned wrong results. I believe the DataFusion `RepartitionExec` has a good model we could follow: https://github.com/apache/datafusion/blob/45d943dfb8699dc9cb9ef2320e955b73e3e6c03b/datafusion/physical-plan/src/repartition/mod.rs#L528-L542 ########## 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( Review Comment: `PlanProperties::new` defaults to `EvaluationType::Lazy` and `SchedulingType::NonCooperative`. Since this spawns producers that eagerly drive each input partition, we should set `EvaluationType::Eager` and `SchedulingType::Cooperative` (similar to what DataFusion RepartitionExec does) ########## ballista/core/src/execution_plans/range_repartition_common.rs: ########## @@ -0,0 +1,375 @@ +// 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. + +//! Shared building blocks for the two range-repartition operators — +//! [`UnorderedRangeRepartitionExec`] and (soon) `OrderedRangeRepartitionExec`. +//! +//! The two operators disagree substantially on execution model — the +//! unordered variant is pure scatter, the ordered one is scatter + per-output +//! k-way merge — but they agree on: +//! +//! 1. **How to find the cut boundaries at runtime** — walk the child subtree +//! for a matching sibling [`RuntimeStatsExec`], snapshot its T-Digest, +//! compute quantile cuts. Only descend through whitelisted +//! distribution-preserving operators; refuse otherwise. +//! 2. **How to split one batch across K value ranges** — [`split_batch_by_range`]. +//! 3. **How to broadcast a terminal error to every output channel** — +//! [`broadcast_error`]. +//! +//! Everything in this module is `pub(super)` — visible to sibling +//! `execution_plans::*` modules that own the operators, invisible outside. +//! +//! [`UnorderedRangeRepartitionExec`]: super::UnorderedRangeRepartitionExec +//! [`RuntimeStatsExec`]: super::RuntimeStatsExec + +use std::sync::Arc; + +use datafusion::arrow::array::{Array, Float64Array, RecordBatch, UInt32Array}; +use datafusion::arrow::compute::take_arrays; +use datafusion::common::{Result, internal_datafusion_err}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; +use log::warn; +use tokio::sync::mpsc; + +use crate::execution_plans::{ + BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, +}; + +/// Walk `child`'s subtree for a [`RuntimeStatsExec`] that sketches on our +/// routing expression, snapshot its merged T-Digest, and compute `K - 1` +/// quantile cuts. Any failure to find a matching sketch returns an empty +/// `Vec` — the caller's `split_batch_by_range(&[])` produces a single +/// bucket and every row lands in output partition 0. Never crashes. +pub(super) fn discover_cuts( + child: &Arc<dyn ExecutionPlan>, + routing_expr: &dyn PhysicalExpr, + output_partitions: usize, +) -> Vec<f64> { + let Some(stats) = find_runtime_stats(child, routing_expr) else { + warn!( + "range-repartition: no matching RuntimeStatsExec found in child subtree — \ + single-bucket fallback" + ); + return Vec::new(); + }; + // Walker returned Some → stats.order_by()'s first entry matches our + // routing expression → RuntimeStatsExec's construction contract + // guarantees sketch is present. Belt-and-braces arms in case that + // invariant ever drifts, plus mutex-poisoning is theoretically possible. + let sketch = match stats.merged_quantile_sketch() { + Ok(Some(sketch)) => sketch, + Ok(None) => { + warn!( + "range-repartition: matching RuntimeStatsExec has no sketch \ + (RuntimeStatsExec contract broken?) — single-bucket fallback" + ); + return Vec::new(); + } + Err(e) => { + warn!( + "range-repartition: sketch snapshot failed ({e}) — single-bucket fallback" + ); + return Vec::new(); + } + }; + // `count()` is the sum of centroid weights — total observed row count + // that fed the digest. Zero means no samples arrived before the + // snapshot; degenerate cuts would follow. + if sketch.count() == 0.0 { + warn!( + "range-repartition: matching sketch has no samples yet — single-bucket fallback" + ); + return Vec::new(); + } + // K-1 cuts at 1/K, 2/K, ..., (K-1)/K. `estimate_quantile` is monotone by + // construction, so cuts are non-decreasing (ties possible on hot-value + // distributions — `split_batch_by_range` handles those correctly, it + // just skews the resulting distribution). + let k = output_partitions as f64; + (1..output_partitions) + .map(|i| sketch.estimate_quantile(i as f64 / k)) + .collect() +} + +/// Walks `plan`'s subtree through single-child chains only, returning the +/// first [`RuntimeStatsExec`] that sketches on `routing_expr`. +/// +/// Two invariants have to hold for a sketch to be trustworthy: +/// 1. **Expression match.** A sketch of column `foo` says nothing about +/// routing on column `bar`. A `RuntimeStatsExec` sketching on a different +/// expression is treated as a plain passthrough — the walker keeps +/// descending past it looking for a matching one deeper in the chain. +/// 2. **Distribution preservation.** Any operator between us and the stats +/// that drops rows (`FilterExec`, `LimitExec`), transforms the routing +/// value (`ProjectionExec` with a computed column), or duplicates rows +/// (`JoinExec`) makes the sketch stale — the count still holds but the +/// distribution has drifted. The walker consults [`preserves_distribution`] +/// and refuses to descend past anything it doesn't know is safe. +/// +/// Also stops at any branch (> 1 child) or leaf (0 children) — descending +/// into a join's sides would risk picking up a sketch of the wrong subtree. +pub(super) fn find_runtime_stats<'a>( + plan: &'a Arc<dyn ExecutionPlan>, + routing_expr: &dyn PhysicalExpr, +) -> Option<&'a RuntimeStatsExec> { + if let Some(stats) = plan.downcast_ref::<RuntimeStatsExec>() { + let matches = stats + .order_by() + .and_then(|order_by| order_by.first()) + .is_some_and(|first| first.expr.as_ref() == routing_expr); + if matches { + return Some(stats); + } + // Non-matching stats is still a passthrough for our purposes — fall + // through to the descent step. + } else if !preserves_distribution(plan.as_ref()) { + // Unrecognized node type — could change the row set or value + // distribution of the routing key. Refuse to descend. + return None; + } + let children = plan.children(); + let [only_child] = children.as_slice() else { + return None; + }; + find_runtime_stats(only_child, routing_expr) +} + +/// Whitelist of pass-through operator types the walker will descend through +/// on its way to a matching [`RuntimeStatsExec`]. Unlisted operators might +/// drop rows, duplicate rows, or transform the routing key's value — any of +/// which would make an upstream sketch stale by the time data reaches us. +/// +/// Being conservative is the safety net: unrecognized node → walker gives +/// up → single-bucket fallback. Extending this list requires positive +/// verification that the operator is a distribution-preserving passthrough +/// for the routing key. Absent an upstream `ExecutionPlan::affects_distribution()` +/// method (nice-to-have that hopefully lands one day), we maintain this by hand. +pub(super) fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool { Review Comment: We may want a way for library consumers to also register their own exec plans here that are passthrough nodes. But probably an upstream trait type is the correct approach. -- 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]
