avantgardnerio commented on code in PR #2123: URL: https://github.com/apache/datafusion-ballista/pull/2123#discussion_r3629704398
########## 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: I would strongly prefer to upstream. Once this is all working and making a benchmark faster, I'd like to submit all of the upstreamable parts to DF for review. -- 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]
