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 b2372cc182922..fa8568ea17c93 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 { @@ -1286,8 +1301,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) } } @@ -2321,7 +2335,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 221439a4cb3d8..8508ca6ff66bb 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -7801,6 +7801,31 @@ 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'; + +# 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; +---- +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;