diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 60cff658a6a97..105ee555ec341 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -294,6 +294,10 @@ harness = false name = "preserve_file_partitioning" required-features = ["parquet"] +[[bench]] +harness = false +name = "parallel_window_cumulative_scaling" + [[bench]] harness = false name = "reset_plan_states" diff --git a/datafusion/core/benches/parallel_window_cumulative_scaling.rs b/datafusion/core/benches/parallel_window_cumulative_scaling.rs new file mode 100644 index 0000000000000..b8af88f59a474 --- /dev/null +++ b/datafusion/core/benches/parallel_window_cumulative_scaling.rs @@ -0,0 +1,197 @@ +// 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. + +//! Weak-scaling sweep for the cumulative-aggregate branch of the +//! `ParallelWindow` optimizer rule (prefix-scan via `CarryExec`). +//! +//! For each "cores" setting `N`, builds a fresh table with `N` +//! partitions of `ROWS_PER_CORE` rows each, sets `target_partitions = N`, +//! and runs a cumulative window aggregate +//! (`ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`, no `PARTITION BY`) +//! twice: once with `ParallelWindow` filtered out of the physical +//! optimizer chain (single-partition baseline) and once with the rule +//! enabled (parallel prefix-scan via `CarryExec`). Emits one CSV row +//! per iteration to stdout. +//! +//! Under linear scaling the PoC's wall-clock stays roughly constant +//! across the sweep while the baseline grows linearly with cores — +//! same shape as the bounded-RANGE bench, validated against the +//! prefix-scan path. CarryExec is pipeline-breaking, so its sequential +//! gather + offset cost shows up as the slope on the PoC line at high +//! core counts. +//! +//! `CarryExec` offsets a single aggregate column (the last column BWAG +//! appends to the input schema), so the SQL uses one `SUM(v) OVER ...`; +//! multiple cumulative aggregates would silently produce wrong values +//! for every aggregate but the last. +//! +//! Run: +//! cargo bench --bench parallel_window_cumulative_scaling \ +//! > cumulative.csv + +use arrow::array::{Float64Array, Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_optimizer::optimizer::PhysicalOptimizer; +use datafusion::prelude::{SessionConfig, SessionContext}; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand_distr::{Distribution, Uniform}; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Instant; +use tokio::runtime::Runtime; + +/// Weak-scaling design: rows scale linearly with cores so total work +/// grows proportionally to the parallelism budget. The PoC line stays +/// flat under linear scaling; the baseline grows linearly with cores +/// because the cumulative aggregate serializes through one partition. +/// Single `SUM` per row is cheap, so the per-core row count is larger +/// than the bounded-RANGE bench's to keep the baseline measurable at 1 +/// core. +const ROWS_PER_CORE: usize = 2_500_000; +const BATCH_SIZE: usize = 8 * 1024; +const ITERATIONS: usize = 3; +const CORE_SETTINGS: &[usize] = &[1, 2, 4, 8, 16, 32]; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("v", DataType::Float64, false), + ])) +} + +/// Build `num_partitions` partitions of `(ts, v)` rows, with `ts` +/// monotonically increasing within each partition AND between +/// partitions. Same shape as the bounded-RANGE bench's fixture — +/// keeps SortExec cheap so the bench measures BWAG + repartition + +/// CarryExec cost. +fn make_partitions(num_partitions: usize) -> Vec> { + let mut rng = StdRng::seed_from_u64(0xC0FFEE_C0FFEE); + let v_dist = Uniform::new(0.0f64, 1.0).unwrap(); + let schema = schema(); + (0..num_partitions) + .map(|part| { + let mut batches = Vec::new(); + let part_start = (part * ROWS_PER_CORE) as i64; + let mut next_ts = part_start; + let mut remaining = ROWS_PER_CORE; + while remaining > 0 { + let len = remaining.min(BATCH_SIZE); + let ts: Vec = (0..len as i64).map(|i| next_ts + i).collect(); + next_ts += len as i64; + let v: Vec = (0..len).map(|_| v_dist.sample(&mut rng)).collect(); + batches.push( + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(ts)), + Arc::new(Float64Array::from(v)), + ], + ) + .unwrap(), + ); + remaining -= len; + } + batches + }) + .collect() +} + +fn make_ctx( + data: &[Vec], + target_partitions: usize, + with_parallel_window: bool, +) -> SessionContext { + let table = MemTable::try_new(schema(), data.to_vec()).unwrap(); + let config = SessionConfig::new() + .with_target_partitions(target_partitions) + .with_batch_size(BATCH_SIZE); + + let mut builder = SessionStateBuilder::new() + .with_default_features() + .with_config(config); + if !with_parallel_window { + let rules: Vec<_> = PhysicalOptimizer::new() + .rules + .into_iter() + .filter(|r| r.name() != "ParallelWindow") + .collect(); + builder = builder.with_physical_optimizer_rules(rules); + } + let state = builder.build(); + let ctx = SessionContext::new_with_state(state); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx +} + +fn run_once(ctx: &SessionContext, rt: &Runtime, sql: &str) -> usize { + let df = rt.block_on(ctx.sql(sql)).unwrap(); + let batches = rt.block_on(df.collect()).unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + black_box(batches); + rows +} + +fn main() { + let rt = Runtime::new().unwrap(); + let sql = "SELECT SUM(v) OVER \ + (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) \ + FROM t"; + + if std::env::var("EXPLAIN_PLAN").is_ok() { + let cores: usize = std::env::var("EXPLAIN_CORES") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(32); + let data = make_partitions(cores); + let ctx = make_ctx(&data, cores, true); + let df = rt.block_on(ctx.sql(&format!("EXPLAIN {sql}"))).unwrap(); + let batches = rt.block_on(df.collect()).unwrap(); + for b in batches { + eprintln!( + "{}", + arrow::util::pretty::pretty_format_batches(&[b]).unwrap() + ); + } + return; + } + + println!("cores,with_poc,iter,seconds,rows"); + for &cores in CORE_SETTINGS { + let data = make_partitions(cores); + for &with_poc in &[false, true] { + let ctx = make_ctx(&data, cores, with_poc); + let warmup_rows = run_once(&ctx, &rt, sql); + for iter in 0..ITERATIONS { + let t = Instant::now(); + let rows = run_once(&ctx, &rt, sql); + let secs = t.elapsed().as_secs_f64(); + assert_eq!( + rows, warmup_rows, + "row count drifted: warmup={warmup_rows} run={rows}" + ); + println!("{cores},{with_poc},{iter},{secs:.6},{rows}"); + eprintln!( + "cores={cores:>2} poc={with_poc:<5} iter={iter} \ + secs={secs:>6.3} rows={rows}" + ); + } + } + } +} diff --git a/datafusion/core/benches/scripts/plot_parallel_window_cumulative_scaling.py b/datafusion/core/benches/scripts/plot_parallel_window_cumulative_scaling.py new file mode 100644 index 0000000000000..e013be6638d81 --- /dev/null +++ b/datafusion/core/benches/scripts/plot_parallel_window_cumulative_scaling.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +# 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. + +"""Plot the weak-scaling sweep emitted by +`parallel_window_cumulative_scaling.rs`. + +The bench emits one CSV row per iteration with header +`cores,with_poc,iter,seconds,rows`. Rows scale linearly with cores, +so under linear scaling the PoC's wall-clock stays constant and its +throughput grows linearly with cores; the baseline's wall-clock grows +linearly with cores (no parallelism ⇒ pure serial work) and its +throughput stays flat at single-core capacity. + +Usage: + cargo bench --bench parallel_window_cumulative_scaling > cumulative.csv + python3 plot_parallel_window_cumulative_scaling.py cumulative.csv \\ + [--out file.png] +""" + +import argparse +import csv +import statistics +import sys +from collections import defaultdict + +import matplotlib + +matplotlib.use("Agg") # headless PNG only +import matplotlib.pyplot as plt + + +def read_rows(source): + reader = csv.DictReader(source) + seconds = defaultdict(list) + rows_for = {} + for row in reader: + cores = int(row["cores"]) + with_poc = row["with_poc"].lower() == "true" + seconds[(cores, with_poc)].append(float(row["seconds"])) + rows_for[cores] = int(row["rows"]) + return seconds, rows_for + + +def medians(seconds): + out = defaultdict(dict) + for (cores, with_poc), samples in seconds.items(): + out[with_poc][cores] = statistics.median(samples) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("csv", help="CSV path; use '-' for stdin") + ap.add_argument( + "--out", + default="parallel_window_cumulative_scaling.png", + help="output PNG path", + ) + args = ap.parse_args() + + source = sys.stdin if args.csv == "-" else open(args.csv) + seconds, rows_for = read_rows(source) + if args.csv != "-": + source.close() + + by_poc = medians(seconds) + baseline = by_poc[False] + parallel = by_poc[True] + cores = sorted(set(baseline) | set(parallel)) + + fig, (ax_time, ax_throughput) = plt.subplots( + 1, 2, figsize=(13, 5), constrained_layout=True + ) + + bx = sorted(baseline) + by = [baseline[c] for c in bx] + px = sorted(parallel) + py = [parallel[c] for c in px] + ax_time.plot( + bx, + [by[0] * (c / bx[0]) for c in bx], + linestyle="--", + color="grey", + label="ideal serial (y = x · t₁)", + ) + ax_time.plot(bx, by, marker="o", color="C0", label="ParallelWindow off (baseline)") + ax_time.plot(px, py, marker="s", color="C1", label="ParallelWindow on (this PR)") + ax_time.set_xscale("log", base=2) + ax_time.set_xticks(cores) + ax_time.set_xticklabels([str(c) for c in cores]) + ax_time.set_xlabel("cores (= target_partitions = input partitions)") + ax_time.set_ylabel("wall-clock (seconds)") + ax_time.set_title("Wall-clock vs cores (weak scaling)") + ax_time.grid(True, which="both", alpha=0.3) + ax_time.legend(loc="upper left") + + bt = [rows_for[c] / baseline[c] / 1e6 for c in bx] + pt = [rows_for[c] / parallel[c] / 1e6 for c in px] + ax_throughput.plot( + px, + [pt[0] * (c / px[0]) for c in px], + linestyle="--", + color="grey", + label="ideal linear scaling (y = x · t₁)", + ) + ax_throughput.plot( + bx, bt, marker="o", color="C0", label="ParallelWindow off (baseline)" + ) + ax_throughput.plot( + px, pt, marker="s", color="C1", label="ParallelWindow on (this PR)" + ) + ax_throughput.set_xscale("log", base=2) + ax_throughput.set_yscale("log", base=2) + ax_throughput.set_xticks(cores) + ax_throughput.set_xticklabels([str(c) for c in cores]) + ax_throughput.set_xlabel("cores (= target_partitions = input partitions)") + ax_throughput.set_ylabel("throughput (million rows / second, log₂)") + ax_throughput.set_title("Throughput vs cores") + ax_throughput.grid(True, which="both", alpha=0.3) + ax_throughput.legend(loc="upper left") + + rows_per_core = next(iter(rows_for.values())) // cores[0] if cores else 0 + fig.suptitle( + f"Prefix-scan cumulative window — {rows_per_core:,} rows per core, " + f"SUM(v) OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)" + ) + fig.savefig(args.out, dpi=120) + print(f"wrote {args.out}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/datafusion/physical-optimizer/src/parallel_window.rs b/datafusion/physical-optimizer/src/parallel_window.rs index 4216a63987db1..1d7dae5df8c85 100644 --- a/datafusion/physical-optimizer/src/parallel_window.rs +++ b/datafusion/physical-optimizer/src/parallel_window.rs @@ -44,6 +44,7 @@ use datafusion_expr::{WindowFrameBound, WindowFrameUnits}; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::expressions::Column; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::carry::CarryExec; use datafusion_physical_plan::halo_drop::HaloDropExec; use datafusion_physical_plan::range_repartition::RangeRepartitionExec; use datafusion_physical_plan::windows::BoundedWindowAggExec; @@ -69,51 +70,36 @@ impl PhysicalOptimizerRule for ParallelWindow { let Some(window) = node.downcast_ref::() else { return Ok(Transformed::no(node)); }; - let Some((halo_preceding, halo_following)) = candidate_halo(window) - else { - return Ok(Transformed::no(node)); - }; - info!( - "ParallelWindow: candidate BoundedWindowAggExec (RANGE frame, no PARTITION BY); \ - halo: {halo_preceding} preceding, {halo_following} following" - ); - // `candidate_halo` already verified order_by.len()==1. - let sort_key = window.window_expr()[0].order_by()[0].clone(); - let lex = LexOrdering::new(vec![sort_key]) - .expect("candidate_halo guarantees one sort key"); - let original_input = Arc::clone(&node.children()[0]); - // Don't pre-insert a SortExec; RangeRepartitionExec now declares - // its required input ordering, so EnsureRequirements will plant - // the pipeline-breaking sort beneath us. Doing both would just - // produce a redundant SortExec that the optimizer collapses. - let range = Arc::new(RangeRepartitionExec::new( - original_input, - lex.clone(), - halo_preceding, - halo_following, - )); - // `parallel_aware = true` flips BWAG's required_input_distribution - // to UnspecifiedDistribution, so EnsureRequirements won't wrap - // us in an SPM. `can_repartition` is vacuous because - // candidate_halo already required partition_keys empty. - let new_window: Arc = Arc::new( - BoundedWindowAggExec::try_new( - window.window_expr().to_vec(), - range, - window.input_order_mode.clone(), + if let Some((halo_preceding, halo_following)) = candidate_halo(window) { + info!( + "ParallelWindow: candidate BoundedWindowAggExec (RANGE frame, no PARTITION BY); \ + halo: {halo_preceding} preceding, {halo_following} following" + ); + let drop_halo = build_halo_plan( + window, + &node, + halo_preceding, + halo_following, + )?; + // Jump past the result's children: the BWAG we just emitted is + // still a candidate by shape (RANGE frame, no PARTITION BY) and + // `transform_down` would otherwise re-wrap it forever. + return Ok(Transformed::new( + drop_halo, true, - )? - .with_parallel_aware(true), - ); - // Drop halo rows above the per-partition window. HaloDropExec - // reads its primary range from `input.runtime_partition_extremes`, - // which BWAG passes through and RangeRepartitionExec populates. - let drop_halo: Arc = - Arc::new(HaloDropExec::try_new(new_window, &lex)?); - // Jump past the result's children: the BWAG we just emitted is - // still a candidate by shape (RANGE frame, no PARTITION BY) and - // `transform_down` would otherwise re-wrap it forever. - Ok(Transformed::new(drop_halo, true, TreeNodeRecursion::Jump)) + TreeNodeRecursion::Jump, + )); + } + if is_candidate_carry(window) { + info!( + "ParallelWindow: candidate BoundedWindowAggExec \ + (cumulative ROWS UNBOUNDED PRECEDING, no PARTITION BY)" + ); + let carry = build_carry_plan(window, &node)?; + // Same jump-recursion concern as the halo branch. + return Ok(Transformed::new(carry, true, TreeNodeRecursion::Jump)); + } + Ok(Transformed::no(node)) })?; Ok(out.data) } @@ -127,6 +113,78 @@ impl PhysicalOptimizerRule for ParallelWindow { } } +/// Build the parallel plan for a bounded-RANGE-frame candidate: +/// `HaloDropExec(BWAG_parallel_aware(RangeRepartitionExec(input)))`. +fn build_halo_plan( + window: &BoundedWindowAggExec, + node: &Arc, + halo_preceding: i64, + halo_following: i64, +) -> datafusion_common::Result> { + // `candidate_halo` already verified order_by.len()==1. + let sort_key = window.window_expr()[0].order_by()[0].clone(); + let lex = + LexOrdering::new(vec![sort_key]).expect("candidate_halo guarantees one sort key"); + let original_input = Arc::clone(&node.children()[0]); + // Don't pre-insert a SortExec; RangeRepartitionExec declares its + // required input ordering, so EnsureRequirements plants the + // pipeline-breaking sort beneath us. Doing both would just produce + // a redundant SortExec that the optimizer collapses. + let range = Arc::new(RangeRepartitionExec::new( + original_input, + lex.clone(), + halo_preceding, + halo_following, + )); + // `parallel_aware = true` flips BWAG's required_input_distribution to + // UnspecifiedDistribution, so EnsureRequirements won't wrap us in an + // SPM. `can_repartition` is vacuous because candidate_halo already + // required partition_keys empty. + let new_window: Arc = Arc::new( + BoundedWindowAggExec::try_new( + window.window_expr().to_vec(), + range, + window.input_order_mode.clone(), + true, + )? + .with_parallel_aware(true), + ); + // Drop halo rows above the per-partition window. HaloDropExec reads + // its primary range from `input.runtime_partition_extremes`, which + // BWAG passes through and RangeRepartitionExec populates. + Ok(Arc::new(HaloDropExec::try_new(new_window, &lex)?)) +} + +/// Build the parallel plan for a cumulative ROWS-UNBOUNDED-PRECEDING +/// candidate: `CarryExec(BWAG_parallel_aware(RangeRepartitionExec(input)))` +/// with no halo distances. +fn build_carry_plan( + window: &BoundedWindowAggExec, + node: &Arc, +) -> datafusion_common::Result> { + let sort_key = window.window_expr()[0].order_by()[0].clone(); + let lex = LexOrdering::new(vec![sort_key]) + .expect("is_candidate_carry guarantees one sort key"); + let original_input = Arc::clone(&node.children()[0]); + // No halo: cumulative frames need no boundary context — partition i's + // local cumsum is complete on its own; CarryExec stitches the global + // offset across partitions. + let range = Arc::new(RangeRepartitionExec::new(original_input, lex, 0, 0)); + let new_window: Arc = Arc::new( + BoundedWindowAggExec::try_new( + window.window_expr().to_vec(), + range, + window.input_order_mode.clone(), + true, + )? + .with_parallel_aware(true), + ); + // BWAG appends the window aggregate column at the end of its input + // schema; that's the column CarryExec offsets by the prefix sum. + let agg_col = new_window.schema().fields().len() - 1; + Ok(Arc::new(CarryExec::new(new_window, agg_col))) +} + /// Returns `(halo_preceding, halo_following)` if the window matches the /// v1 shape we know how to parallelize: no PARTITION BY, a single /// `Column` sort key, RANGE frame, finite Int64 bounds (or CurrentRow). @@ -166,3 +224,29 @@ fn i64_halo(start: &WindowFrameBound, end: &WindowFrameBound) -> Option<(i64, i6 }; Some((preceding, following)) } + +/// Matches the cumulative-aggregate shape we parallelize via prefix scan: +/// no PARTITION BY, single `Column` ORDER BY, ROWS frame with +/// `UNBOUNDED PRECEDING` start and `CURRENT ROW` end. UNBOUNDED PRECEDING +/// is `Preceding()` regardless of the scalar's data type +/// (UInt64 in default cases; we don't depend on the type). +fn is_candidate_carry(window: &BoundedWindowAggExec) -> bool { + if !window.partition_keys().is_empty() { + return false; + } + let order_by = window.window_expr()[0].order_by(); + if order_by.len() != 1 { + return false; + } + if order_by[0].expr.downcast_ref::().is_none() { + return false; + } + let frame = window.window_expr()[0].get_window_frame(); + if frame.units != WindowFrameUnits::Rows { + return false; + } + let unbounded_start = + matches!(&frame.start_bound, WindowFrameBound::Preceding(v) if v.is_null()); + let current_end = matches!(&frame.end_bound, WindowFrameBound::CurrentRow); + unbounded_start && current_end +} diff --git a/datafusion/physical-plan/src/carry.rs b/datafusion/physical-plan/src/carry.rs new file mode 100644 index 0000000000000..99b42c7e3292c --- /dev/null +++ b/datafusion/physical-plan/src/carry.rs @@ -0,0 +1,249 @@ +// 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. + +//! Parallel prefix-scan carry above a `BoundedWindowAggExec` running +//! per-partition over `RangeRepartitionExec`-routed input. +//! +//! Each input partition produces a cumulative aggregate starting at zero. +//! `CarryExec` is pipeline-breaking: the first output partition to poll +//! drains every input partition fully into per-partition `Vec`, +//! derives each partition's final cumulative value from the last row of +//! the last batch (no separate finals state — the buffered batches ARE +//! the state), and computes the prefix sum across partition finals. +//! Concurrent and subsequent output-partition polls await the same +//! memoized result via `OnceCell`. Each output stream re-emits its +//! buffered batches with `prefix` added to the aggregate column. +//! +//! Output partitioning equals input partitioning (N → N). + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch}; +use arrow::compute::kernels::numeric::add; +use datafusion_common::{Result, ScalarValue, internal_datafusion_err}; +use datafusion_execution::TaskContext; +use futures::StreamExt; +use tokio::sync::OnceCell; + +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; + +#[derive(Debug)] +pub struct CarryExec { + input: Arc, + /// Column index in the input schema whose values are offset by the + /// prefix sum of prior partitions' finals. + agg_col: usize, + cache: Arc, + /// First output-partition poll runs the gather; concurrent polls await + /// its completion; later polls read the cached result. The error path + /// stores a stringified message because `DataFusionError` isn't + /// `Clone` and the same error must surface on every output partition. + gathered: Arc>, +} + +type GatherResult = std::result::Result>, Arc>; + +#[derive(Debug)] +struct PartitionPayload { + batches: Vec, + /// Already-prefix-summed offset to add to the agg column on every row + /// of every batch. `prefix[0]` is the additive identity for the agg + /// column's data type. + prefix: ScalarValue, +} + +impl CarryExec { + pub fn new(input: Arc, agg_col: usize) -> Self { + let cache = Arc::clone(input.properties()); + Self { + input, + agg_col, + cache, + gathered: Arc::new(OnceCell::new()), + } + } +} + +impl DisplayAs for CarryExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "CarryExec") + } +} + +impl ExecutionPlan for CarryExec { + fn name(&self) -> &'static str { + "CarryExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new(children.swap_remove(0), self.agg_col))) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let gathered = Arc::clone(&self.gathered); + let input = Arc::clone(&self.input); + let agg_col = self.agg_col; + let schema = self.schema(); + + let body = async move { + let payloads = gathered + .get_or_init(|| async { + gather(input, context, agg_col) + .await + .map(Arc::new) + .map_err(|e| Arc::new(e.to_string())) + }) + .await; + let payloads = match payloads { + Ok(p) => p, + Err(msg) => return Err(internal_datafusion_err!("{}", msg)), + }; + let payload = &payloads[partition]; + let prefix = payload.prefix.clone(); + // RecordBatch::clone is cheap (Arc'd columns). + let batches: Vec = payload.batches.clone(); + Ok(futures::stream::iter( + batches + .into_iter() + .map(move |batch| offset_batch(&batch, agg_col, &prefix)), + )) + }; + + use futures::stream::{TryStreamExt, once}; + let stream = once(body).try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +/// Drain every input partition fully, derive each partition's final from +/// the last row of its last batch, and compute the prefix sum across +/// finals. Empty partitions contribute the additive identity; their +/// prefix equals the running total at that point. +async fn gather( + input: Arc, + ctx: Arc, + agg_col: usize, +) -> Result> { + let n = input.output_partitioning().partition_count(); + + // Drain every output partition on its own tokio task. `try_join_all` + // would also be concurrent but only on the calling task — every poll + // would happen on a single worker thread, capping the consumer side + // at one CPU. Spawning puts each drain on the multi-threaded runtime + // and lets the upstream routers actually fan out. Sequential drain + // (the original) deadlocks because the upstream coordinator blocks + // sending to any unread bucket and stalls routing to every bucket. + let handles: Vec<_> = (0..n) + .map(|k| { + let mut stream = input.execute(k, Arc::clone(&ctx))?; + Ok::<_, datafusion_common::DataFusionError>(tokio::spawn(async move { + let mut buf: Vec = Vec::new(); + while let Some(item) = stream.next().await { + buf.push(item?); + } + Ok::<_, datafusion_common::DataFusionError>(buf) + })) + }) + .collect::>>()?; + let mut buffers: Vec> = Vec::with_capacity(n); + for handle in handles { + buffers.push( + handle + .await + .map_err(|e| internal_datafusion_err!("drain task panicked: {e}"))??, + ); + } + + // Derive each partition's final + cumulative prefix. `running` starts + // as the zero scalar in the agg column's data type, taken from the + // first non-empty batch we find. + let agg_type = buffers + .iter() + .flat_map(|b| b.iter()) + .next() + .map(|b| b.column(agg_col).data_type().clone()); + let Some(agg_type) = agg_type else { + // No data anywhere — every partition gets an empty payload with a + // null prefix (offset_batch passes through unchanged on null). + return Ok((0..n) + .map(|_| PartitionPayload { + batches: Vec::new(), + prefix: ScalarValue::Null, + }) + .collect()); + }; + let mut running = ScalarValue::new_zero(&agg_type)?; + let mut payloads = Vec::with_capacity(n); + for batches in buffers { + let prefix = running.clone(); + if let Some(last) = batches.last() { + let final_i = + ScalarValue::try_from_array(last.column(agg_col), last.num_rows() - 1)?; + running = running.add(&final_i)?; + } + payloads.push(PartitionPayload { batches, prefix }); + } + Ok(payloads) +} + +/// Replace the agg column with `agg + prefix` (broadcast scalar add). +fn offset_batch( + batch: &RecordBatch, + agg_col: usize, + prefix: &ScalarValue, +) -> Result { + if prefix.is_null() { + // Only happens when there's no data anywhere; pass through. + return Ok(batch.clone()); + } + let agg = batch.column(agg_col); + // Replicate the prefix to match batch length. `arrow::array::Scalar` + // would let us broadcast a single-element array as a Datum, but the + // replicate cost is negligible (one scalar per batch). + let prefix_array = prefix.to_array_of_size(batch.num_rows())?; + let new_agg: ArrayRef = add(&agg.as_ref(), &prefix_array.as_ref())?; + let mut cols = batch.columns().to_vec(); + cols[agg_col] = new_agg; + Ok(RecordBatch::try_new(batch.schema(), cols)?) +} diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 6f0aeda26cda7..e155b7819f2bc 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -64,6 +64,7 @@ pub mod aggregates; pub mod analyze; pub mod async_func; pub mod buffer; +pub mod carry; pub mod coalesce; pub mod coalesce_batches; pub mod coalesce_partitions; diff --git a/datafusion/sqllogictest/test_files/parallel_window.slt b/datafusion/sqllogictest/test_files/parallel_window.slt index c47050367007c..b60f5534cf2e6 100644 --- a/datafusion/sqllogictest/test_files/parallel_window.slt +++ b/datafusion/sqllogictest/test_files/parallel_window.slt @@ -136,6 +136,121 @@ SELECT count(rolling_sum) AS n FROM ( ---- 100 +# --------------------------------------------------------------------------- +# Cumulative aggregate: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW. +# Equivalent to DQL `window orderby $m.timestamp rows running() calc sum(...)`. +# Canonical log-analytics shape (cumulative metric over time, no PARTITION BY). +# +# Today this collapses to Distribution::SinglePartition because: +# 1. ParallelWindow only matches RANGE frames with finite preceding/following +# bounds (parallel_window.rs:146 rejects Rows; parallel_window.rs:158 +# rejects UNBOUNDED). +# 2. BoundedWindowAggExec then falls back to SinglePartition +# (bounded_window_agg_exec.rs:352). +# +# Goal: parallel prefix scan. Sketch of target plan: +# SortPreservingMergeExec [seq] +# CarryExec -- broadcasts each partition's final +# cumulative value as carry-in to +# the next; pipelined, no overlap. +# BoundedWindowAggExec(parallel_aware) -- per-partition cumulative from 0 +# RangeRepartitionExec -- range-partition by seq, NO halo +# SortExec [seq] +# DataSourceExec +# --------------------------------------------------------------------------- + +# EXPLAIN — RED today (single-partition fallback). Target plan is the parallel +# prefix-scan shape above. Expected text is structural; statistics decoration +# will need an update-mode pass once CarryExec prints itself. +query TT +EXPLAIN SELECT + seq, + SUM(amount) OVER ( + ORDER BY seq + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS cumulative_sum +FROM events +ORDER BY seq +LIMIT 10; +---- +logical_plan +01)Sort: events.seq ASC NULLS LAST, fetch=10 +02)--Projection: events.seq, sum(events.amount) ORDER BY [events.seq ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS cumulative_sum +03)----WindowAggr: windowExpr=[[sum(events.amount) ORDER BY [events.seq ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: events projection=[seq, amount] +physical_plan +01)SortPreservingMergeExec: [seq@0 ASC NULLS LAST], fetch=10, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:)]] +02)--ProjectionExec: expr=[seq@0 as seq, sum(events.amount) ORDER BY [events.seq ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as cumulative_sum], statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:)]] +03)----CarryExec, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] +04)------BoundedWindowAggExec: wdw=[sum(events.amount) ORDER BY [events.seq ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(events.amount) ORDER BY [events.seq ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted], statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] +05)--------RangeRepartitionExec, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:)]] +06)----------SortExec: expr=[seq@0 ASC NULLS LAST], preserve_partitioning=[true], statistics=[Rows=Exact(100), Bytes=Exact(1600), [(Col[0]: Min=Exact(Int64(0)) Max=Exact(Int64(99)) Null=Exact(0) ScanBytes=Exact(800)),(Col[1]: Min=Exact(Int64(0)) Max=Exact(Int64(6)) Null=Exact(0) ScanBytes=Exact(800))]] +07)------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parallel_window/events/0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parallel_window/events/1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parallel_window/events/2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parallel_window/events/3.parquet]]}, projection=[seq, amount], file_type=parquet, sort_order_for_reorder=[seq@0 ASC NULLS LAST], statistics=[Rows=Exact(100), Bytes=Exact(1600), [(Col[0]: Min=Exact(Int64(0)) Max=Exact(Int64(99)) Null=Exact(0) ScanBytes=Exact(800)),(Col[1]: Min=Exact(Int64(0)) Max=Exact(Int64(6)) Null=Exact(0) ScanBytes=Exact(800))]] + + +# Result — GREEN today (BoundedWindowAggExec is correct, just running on one +# core). After implementation this must remain row-for-row identical. +query II +SELECT + seq, + SUM(amount) OVER ( + ORDER BY seq + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS cumulative_sum +FROM events +ORDER BY seq +LIMIT 10; +---- +0 0 +1 1 +2 3 +3 6 +4 10 +5 15 +6 21 +7 21 +8 22 +9 24 + +# Sentinel: count(cumulative_sum) forces the window to materialize and proves +# CarryExec neither drops nor duplicates rows. +query I +SELECT count(cumulative_sum) AS n FROM ( + SELECT seq, SUM(amount) OVER ( + ORDER BY seq + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS cumulative_sum + FROM events +) t; +---- +100 + +# Cross-partition-boundary check — RED on the passthrough stub because each +# input partition's local cumulative sum starts at 0; the prefix-sum offset +# isn't applied. Once the real CarryExec body lands, these values match the +# serial reference. Manual reference values: amount(seq) = seq % 7, so +# cumulative_sum(N) = Σ_{i=0..N} (i % 7). +query II +SELECT seq, cumulative_sum FROM ( + SELECT seq, SUM(amount) OVER ( + ORDER BY seq + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS cumulative_sum + FROM events +) t +WHERE seq IN (24, 25, 26, 49, 50, 51, 74, 75, 76) +ORDER BY seq; +---- +24 69 +25 73 +26 78 +49 147 +50 148 +51 150 +74 220 +75 225 +76 231 + # Reset session settings so this file doesn't leak config into the rest of the run. statement ok set datafusion.explain.show_statistics = false;