avantgardnerio commented on code in PR #2169: URL: https://github.com/apache/datafusion-ballista/pull/2169#discussion_r3650181816
########## ballista/core/src/execution_plans/ordered_range_repartition.rs: ########## @@ -0,0 +1,931 @@ +// 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 N locally-sorted overlapping input partitions. +//! Redistributes them into K range-disjoint output partitions where each +//! output is fully sorted on the routing expression. Concatenating the K +//! outputs in order yields a globally-sorted stream. +//! +//! ```text +//! N sorted input outputs (K = 4, each sorted) +//! partitions +//! ─┐ ┌─▶ 0: sorted, key < c₁ +//! ─┼──▶ RuntimeStatsExec ──▶ SortExec/preserve ──▶ OrderedRRE ─┼─▶ 1: sorted, c₁ ≤ key < c₂ +//! ─┤ (T-Digest tap) (per-partition sort) (K = 4) ├─▶ 2: sorted, c₂ ≤ key < c₃ +//! ─┘ │ └─▶ 3: sorted, c₃ ≤ key +//! ▲ │ +//! └──────────── walker ◀─────────────────┘ +//! (descends past sort-preserving +//! nodes, matches on routing expr, +//! reads K−1 quantile cuts) +//! +//! Inside OrderedRRE: N × K channels + K k-way merges. +//! Each of N scatter tasks splits its sorted input by the K-1 cuts +//! (K sub-batches per batch) and forwards each sub-batch to its output's +//! sender. Each of K outputs runs a `StreamingMerge` across the N +//! sub-streams targeting it → the output stream is sorted. +//! ``` +//! +//! # Contrast with `UnorderedRangeRepartitionExec` +//! +//! The two operators share the same runtime discovery mechanism (walk child +//! subtree → find matching `RuntimeStatsExec` → quantile cuts from T-Digest, +//! see [`crate::execution_plans::range_repartition_common`]). They diverge +//! on execution model: +//! +//! - **Unordered**: N scatter tasks → K channels → K raw receivers. Rows +//! from different sources arrive interleaved, order is lost. +//! - **Ordered**: N scatter tasks → N × K channels → K [`StreamingMerge`] +//! streams. Each output p performs a heap-based k-way merge across N +//! sorted sub-streams (one per input source, pre-filtered to output p's +//! value range). Battle-tested merger via DataFusion's `SortPreservingMerge` +//! internals; no custom heap code here. +//! +//! # Discovery + fallback +//! +//! Identical to Unordered: walk `self.input`'s subtree through single-child +//! chains and the distribution-preserving whitelist; snapshot the matching +//! T-Digest on first batch; empty cut set = single-bucket fallback where all +//! rows land in output 0. +//! +//! # Ordering claim +//! +//! Constructor requires `input.output_ordering()` to lead with the routing +//! expression — otherwise the merger would produce garbled output. +//! [`PlanProperties::eq_properties`] declares each output partition sorted +//! on `order_by`, letting downstream operators (BWAG, HaloDrop) rely on the +//! claim without inserting a redundant `SortExec`. +//! +//! [`StreamingMerge`]: datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder + +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::runtime::SpawnedTask; +use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{ + Distribution, EquivalenceProperties, LexOrdering, OrderingRequirements, Partitioning, + PhysicalExpr, PhysicalSortExpr, +}; +use datafusion::physical_plan::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, +}; +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, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::stream::StreamExt; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; + +use crate::execution_plans::range_repartition_common::{ + discover_cuts, guarded_scatter, split_batch_by_range, +}; + +/// Per-output-partition channel capacity, per input source. Matches the +/// unordered variant's default; see the discussion there. Total buffered +/// batches at rest ≤ `N × K × CHANNEL_CAPACITY`. +const CHANNEL_CAPACITY: usize = 2; + +/// Value-range router that preserves the input's sort order. See the +/// module-level docs. +pub struct OrderedRangeRepartitionExec { + input: Arc<dyn ExecutionPlan>, + /// Lexicographic ORDER BY. `try_new` guarantees the first entry evaluates + /// to `Float64` and matches the input's declared output ordering. + order_by: Vec<PhysicalSortExpr>, + /// K — number of output partitions. + output_partitions: usize, + /// Lazy channel + merge-stream setup guarded by `Mutex`. First + /// `execute()` call spawns the N input-reader tasks, wires N × K + /// channels, and constructs K `StreamingMerge` streams — one per + /// output partition. Subsequent `execute(p)` calls take + /// `merged_streams[p]`. + state: Arc<Mutex<DispatchState>>, + properties: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +/// Per-exec-instance lazy state. +struct DispatchState { + /// K slots. Each holds `Some(stream)` after setup, `None` once its + /// output partition has been consumed. Each stream is a k-way merge + /// (via `StreamingMerge`) across N sorted sub-streams — one per input + /// source, filtered to the output's value range on the scatter side. + merged_streams: Vec<Option<SendableRecordBatchStream>>, + /// N scatter tasks (one per input partition). Each runs its scatter + /// body inside `guarded_scatter`, which converts panic/error/clean-end + /// into an explicit signal on every output channel. `SpawnedTask` + /// aborts its inner tokio task on drop, so holding these here ties the + /// background work's lifetime to this exec's — dropping the exec + /// cancels all scatter work. + _drop_helper: Vec<SpawnedTask<()>>, Review Comment: It's just to signify that no one reads this. It's just a struct member purely so it can get held until the struct is dropped. -- 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]
