From e958da7b49431431c7333f2ba6f4617f06d9f94c Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 31 Oct 2026 10:16:47 -0700 Subject: [PATCH 1/2] fix: split `BoundedWindowAggExec` output into `batch_size` batches `BoundedWindowAggExec` emitted every row that became final in one batch, so frames that need later rows could exceed `batch_size`. Wrap its stream in `BatchSplitStream` and `ObservedStream`, as `UnnestExec` does, so the output is sliced into batches of at most `batch_size` rows and `output_batches` counts the slices. Closes #26959. --- .../src/windows/bounded_window_agg_exec.rs | 30 ++++++++++++++----- datafusion/sqllogictest/test_files/window.slt | 24 +++++++++++++++ 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 5c09fa2261d40..964739a353265 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -14,9 +14,11 @@ use std::sync::Arc; use std::task::{Context, Poll}; use super::utils::create_schema; -use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, SplitMetrics, +}; use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{BatchSplitStream, EmptyRecordBatchStream, ObservedStream}; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, window_equivalence_properties, @@ -506,18 +508,31 @@ impl ExecutionPlan for BoundedWindowAggExec { partition: usize, context: Arc, ) -> Result { + let batch_size = context.session_config().batch_size(); let input = self.input.execute(partition, context)?; let search_mode = self.get_search_algo()?; + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); let stream = Box::pin(BoundedWindowAggStream::new( Arc::clone(&self.schema), self.window_expr.clone(), input, - BaselineMetrics::new(&self.metrics, partition), + baseline_metrics.clone(), search_mode, partition, self.state_observer.clone(), )?); - Ok(stream) + // Rows that become final together are emitted as one batch, so split + // them to respect `batch_size`. + let stream = Box::pin(BatchSplitStream::new( + stream, + batch_size, + SplitMetrics::new(&self.metrics, partition), + )); + Ok(Box::pin(ObservedStream::new( + stream, + baseline_metrics, + None, + ))) } fn metrics(&self) -> Option { @@ -1278,8 +1293,7 @@ impl Stream for BoundedWindowAggStream { mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { - let poll = self.poll_next_inner(cx); - self.baseline_metrics.record_poll(poll) + self.poll_next_inner(cx) } } @@ -2313,7 +2327,9 @@ mod tests { ) .map(|e| Arc::new(e) as Arc)?; - let batches = collect(physical_plan.execute(0, task_context())?).await?; + // The default batch size keeps each emission in a single batch. + let task_ctx = Arc::new(TaskContext::default()); + let batches = collect(physical_plan.execute(0, task_ctx)?).await?; // The skip must not delay results that are ready to be finalized; // they stream out as soon as every window expression has produced diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index d81bdbddcf6c0..447dfb5e1ba53 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7801,6 +7801,30 @@ FROM range_nan_peers ORDER BY id; 4 5 3 5 1 3 +# BoundedWindowAggExec emits at most `batch_size` rows per batch, 2 here. Peer +# group k = 1 becomes final when the first k = 2 row arrives, and k = 2 at the +# end of the input. Each 5-row result is emitted as 2, 2 and 1 rows. +# https://github.com/apache/datafusion/issues/26959 +statement ok +CREATE TABLE bwag_peers(k INT, v INT) AS VALUES + (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (2, 6), (2, 7), (2, 8), (2, 9), (2, 10); + +statement ok +set datafusion.explain.analyze_categories = 'rows'; + +query TT +EXPLAIN ANALYZE SELECT k, v, sum(v) OVER (ORDER BY k) FROM bwag_peers; +---- +Plan with Metrics +01)BoundedWindowAggExec: , mode=[Sorted], metrics=[output_rows=10, output_batches=6, batches_split=6] + + +statement ok +reset datafusion.explain.analyze_categories; + +statement ok +DROP TABLE bwag_peers; + statement ok reset datafusion.execution.batch_size; From 1086fc376a522dd585ac9f01bc677af7a60323f7 Mon Sep 17 00:00:00 2001 From: Oleks V Date: Tue, 3 Nov 2026 08:05:14 -0800 Subject: [PATCH 2/2] Update datafusion/sqllogictest/test_files/window.slt Co-authored-by: Yongting You <2010youy01@gmail.com> --- datafusion/sqllogictest/test_files/window.slt | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 447dfb5e1ba53..bebf729212642 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7812,6 +7812,7 @@ CREATE TABLE bwag_peers(k INT, v INT) AS VALUES statement ok set datafusion.explain.analyze_categories = 'rows'; +# Metrics `output_rows` and `output_batches` satisfy the test goal above. query TT EXPLAIN ANALYZE SELECT k, v, sum(v) OVER (ORDER BY k) FROM bwag_peers; ----