From bc9f852f4e95d8cfd8fe8b7b3ba4f8fbb43d6f91 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 23 Jun 2026 10:47:00 -0600 Subject: [PATCH 1/4] Cumulative-aggregate SLT: red baseline for parallel prefix scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW case to parallel_window.slt. ParallelWindow rule rejects this frame today (parallel_window.rs:146 rejects Rows; parallel_window.rs:158 rejects UNBOUNDED), so the plan collapses to Distribution::SinglePartition. Target plan introduces CarryExec — a pipeline-breaking N→N operator that buffers all input batches, derives each input partition's final cumulative value from the buffered batches (no separate state), and re-emits with partition i's rows offset by the prefix sum of prior finals. Mirrors the Ballista stage-shuffle model intra-node. EXPLAIN expected text is structural; statistics decoration needs an update-mode pass once CarryExec lands. --- .../test_files/parallel_window.slt | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/datafusion/sqllogictest/test_files/parallel_window.slt b/datafusion/sqllogictest/test_files/parallel_window.slt index c47050367007c..a46f0f7b1bb8c 100644 --- a/datafusion/sqllogictest/test_files/parallel_window.slt +++ b/datafusion/sqllogictest/test_files/parallel_window.slt @@ -136,6 +136,95 @@ 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 +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] +03)----CarryExec +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] +05)--------RangeRepartitionExec +06)----------SortExec: expr=[seq@0 ASC NULLS LAST], preserve_partitioning=[true] +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] + + +# 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 + # Reset session settings so this file doesn't leak config into the rest of the run. statement ok set datafusion.explain.show_statistics = false; From 7020f8e2e3afd2b3cf27c09503681de7951ccd29 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 23 Jun 2026 11:06:16 -0600 Subject: [PATCH 2/4] CarryExec passthrough stub + ParallelWindow cumulative branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit End-to-end plan-shape change visible in the SLT: cumulative ROWS UNBOUNDED PRECEDING / CURRENT ROW windows now plan as CarryExec(BWAG_parallel_aware(RangeRepartitionExec(SortExec))) instead of the BWAG-SinglePartition / SPM collapse. CarryExec is a passthrough — output equals input — so the EXPLAIN block, the LIMIT 10 result block (rows 0-9 all sit in partition 0 where carry-in is zero anyway), and the count sentinel all pass. Cross-partition-boundary result block goes RED as designed: at the boundaries (seq=24/49/74) each input partition's local cumulative sum restarts at zero — the prefix-sum offset that real CarryExec will apply is missing. Lands in the next commit. is_candidate_carry checks the frame shape via v.is_null() on the Preceding bound's ScalarValue rather than matching its concrete type, so we're robust to whatever datatype UNBOUNDED PRECEDING resolves to. --- .../physical-optimizer/src/parallel_window.rs | 172 +++++++++++++----- datafusion/physical-plan/src/carry.rs | 112 ++++++++++++ datafusion/physical-plan/src/lib.rs | 1 + .../test_files/parallel_window.slt | 40 +++- 4 files changed, 274 insertions(+), 51 deletions(-) create mode 100644 datafusion/physical-plan/src/carry.rs 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..6ff8f2e5eff66 --- /dev/null +++ b/datafusion/physical-plan/src/carry.rs @@ -0,0 +1,112 @@ +// 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: it buffers all batches per input +//! partition, derives each partition's final cumulative value from the +//! last row of the last batch (no separate state — the buffered batches +//! ARE the state), computes the prefix sum across partition finals, and +//! re-emits each partition's batches with that prefix added to the +//! aggregate column. Output partitioning matches input. +//! +//! Stub: currently a passthrough (no buffering, no offset). Exists to +//! anchor the plan shape so `ParallelWindow`'s cumulative branch has +//! somewhere to land. Real prefix-scan body is the next commit. + +use std::sync::Arc; + +use datafusion_common::Result; +use datafusion_execution::TaskContext; + +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; + +#[derive(Debug)] +pub struct CarryExec { + input: Arc, + /// Column index of the window aggregate output in the input schema. + /// Real implementation adds the prefix sum to this column; stub + /// ignores it. + agg_col: usize, + cache: Arc, +} + +impl CarryExec { + pub fn new(input: Arc, agg_col: usize) -> Self { + let cache = Arc::clone(input.properties()); + Self { + input, + agg_col, + cache, + } + } +} + +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> { + let input = children.swap_remove(0); + let cache = Arc::clone(input.properties()); + Ok(Arc::new(Self { + input, + agg_col: self.agg_col, + cache, + })) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + // Passthrough stub. Real implementation buffers and offsets. + self.input.execute(partition, context) + } +} 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 a46f0f7b1bb8c..b60f5534cf2e6 100644 --- a/datafusion/sqllogictest/test_files/parallel_window.slt +++ b/datafusion/sqllogictest/test_files/parallel_window.slt @@ -179,13 +179,13 @@ logical_plan 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 -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] -03)----CarryExec -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] -05)--------RangeRepartitionExec -06)----------SortExec: expr=[seq@0 ASC NULLS LAST], preserve_partitioning=[true] -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] +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 @@ -225,6 +225,32 @@ SELECT count(cumulative_sum) AS n FROM ( ---- 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; From 468f29334586abaa0235f242916a22c273570f4a Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 23 Jun 2026 11:33:42 -0600 Subject: [PATCH 3/4] CarryExec: real prefix-scan body via OnceCell-driven gather MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the passthrough stub with a poll-driven pipeline-break: the first output partition to poll triggers a single async gather over all input partitions; concurrent and subsequent polls await the same memoized result via tokio::sync::OnceCell. No spawned coordinator task, no per- partition oneshot channels, no State mutex — work happens on whichever executor task polls the stream. The gather drains every input partition into Vec, derives each partition's final cumulative value from the last row of its last batch (the buffered batches ARE the state), and computes the prefix sum across finals. Each output stream then re-emits its buffered batches with `add(agg_col, prefix)` applied. Error fan-out uses Arc rather than cloning DataFusionError (which doesn't implement Clone) — the message surfaces identically on every output partition's stream. Cross-partition-boundary SLT block flips from RED to GREEN; all four assertions in parallel_window.slt pass. --- datafusion/physical-plan/src/carry.rs | 164 ++++++++++++++++++++++---- 1 file changed, 141 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/carry.rs b/datafusion/physical-plan/src/carry.rs index 6ff8f2e5eff66..729361759c468 100644 --- a/datafusion/physical-plan/src/carry.rs +++ b/datafusion/physical-plan/src/carry.rs @@ -19,35 +19,55 @@ //! per-partition over `RangeRepartitionExec`-routed input. //! //! Each input partition produces a cumulative aggregate starting at zero. -//! `CarryExec` is pipeline-breaking: it buffers all batches per input -//! partition, derives each partition's final cumulative value from the -//! last row of the last batch (no separate state — the buffered batches -//! ARE the state), computes the prefix sum across partition finals, and -//! re-emits each partition's batches with that prefix added to the -//! aggregate column. Output partitioning matches input. +//! `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. //! -//! Stub: currently a passthrough (no buffering, no offset). Exists to -//! anchor the plan shape so `ParallelWindow`'s cumulative branch has -//! somewhere to land. Real prefix-scan body is the next commit. +//! Output partitioning equals input partitioning (N → N). use std::sync::Arc; -use datafusion_common::Result; +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, PlanProperties, + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, }; #[derive(Debug)] pub struct CarryExec { input: Arc, - /// Column index of the window aggregate output in the input schema. - /// Real implementation adds the prefix sum to this column; stub - /// ignores it. + /// 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 { @@ -57,6 +77,7 @@ impl CarryExec { input, agg_col, cache, + gathered: Arc::new(OnceCell::new()), } } } @@ -88,13 +109,7 @@ impl ExecutionPlan for CarryExec { self: Arc, mut children: Vec>, ) -> Result> { - let input = children.swap_remove(0); - let cache = Arc::clone(input.properties()); - Ok(Arc::new(Self { - input, - agg_col: self.agg_col, - cache, - })) + Ok(Arc::new(Self::new(children.swap_remove(0), self.agg_col))) } fn maintains_input_order(&self) -> Vec { @@ -106,7 +121,110 @@ impl ExecutionPlan for CarryExec { partition: usize, context: Arc, ) -> Result { - // Passthrough stub. Real implementation buffers and offsets. - self.input.execute(partition, context) + 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(); + let mut buffers: Vec> = Vec::with_capacity(n); + for k in 0..n { + let mut stream = input.execute(k, Arc::clone(&ctx))?; + let mut buf = Vec::new(); + while let Some(item) = stream.next().await { + buf.push(item?); + } + buffers.push(buf); + } + + // 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)?) } From b6dad41bbdcef6096717ff21051c61719e74afd1 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 30 Jun 2026 14:37:44 -0600 Subject: [PATCH 4/4] benchmark --- datafusion/core/Cargo.toml | 4 + .../parallel_window_cumulative_scaling.rs | 197 ++++++++++++++++++ ...plot_parallel_window_cumulative_scaling.py | 147 +++++++++++++ datafusion/physical-plan/src/carry.rs | 33 ++- 4 files changed, 374 insertions(+), 7 deletions(-) create mode 100644 datafusion/core/benches/parallel_window_cumulative_scaling.rs create mode 100644 datafusion/core/benches/scripts/plot_parallel_window_cumulative_scaling.py 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-plan/src/carry.rs b/datafusion/physical-plan/src/carry.rs index 729361759c468..99b42c7e3292c 100644 --- a/datafusion/physical-plan/src/carry.rs +++ b/datafusion/physical-plan/src/carry.rs @@ -166,14 +166,33 @@ async fn gather( 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 k in 0..n { - let mut stream = input.execute(k, Arc::clone(&ctx))?; - let mut buf = Vec::new(); - while let Some(item) = stream.next().await { - buf.push(item?); - } - buffers.push(buf); + 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