fix: split BoundedWindowAggExec output into batch_size batches

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<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
+        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<MetricsSet> {
@@ -1286,8 +1301,7 @@ impl Stream for BoundedWindowAggStream {
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Self::Item>> {
-        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<dyn ExecutionPlan>)?;
 
-        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

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: <slt:ignore>, mode=[Sorted], metrics=[output_rows=10, output_batches=6, batches_split=6]
+<slt:ignore>
+
+statement ok
+reset datafusion.explain.analyze_categories;
+
+statement ok
+DROP TABLE bwag_peers;
+
 statement ok
 reset datafusion.execution.batch_size;