fix: respect batch size in ordered aggregates

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -378,6 +378,9 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str
         .await
         .unwrap();
     assert!(collected_running.len() > 2);
+    // Running should produce more chunk than the usual AggregateExec.
+    // Otherwise it means that we cannot generate result in running mode.
+    assert!(collected_running.len() > collected_usual.len());
     // compare
     let usual_formatted = pretty_format_batches(&collected_usual).unwrap().to_string();
     let running_formatted = pretty_format_batches(&collected_running)

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs

@@ -23,7 +23,7 @@ use std::sync::Arc;
 
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
-use datafusion_common::Result;
+use datafusion_common::{Result, internal_err};
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use datafusion_expr::EmitTo;
 
@@ -64,6 +64,9 @@ pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> {
     /// Output schema: group columns followed by aggregate state or final values.
     pub(super) output_schema: SchemaRef,
 
+    /// Maximum rows per emitted output batch, from config `batch_size`.
+    pub(super) batch_size: usize,
+
     /// Grouping and accumulator-specific timing metrics.
     pub(super) group_by_metrics: GroupByMetrics,
 
@@ -103,15 +106,26 @@ pub(super) struct OrderedAggregateTableBuffer {
 
 /// Methods shared by all aggregate modes
 impl<AggrMode> OrderedAggregateTable<AggrMode> {
+    #[expect(
+        clippy::too_many_arguments,
+        reason = "keeps ordered partial and final table construction explicit"
+    )]
     pub(super) fn new_for_mode(
         agg: &AggregateExec,
         partition: usize,
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
+        batch_size: usize,
         input_order_mode: &InputOrderMode,
         aggregate_mode: &AggregateMode,
         filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
     ) -> Result<Self> {
+        if batch_size == 0 {
+            return internal_err!(
+                "OrderedAggregateTable requires config batch_size >= 1"
+            );
+        }
+
         let group_ordering = GroupOrdering::try_new(input_order_mode)?;
         let group_schema = agg.group_by.group_schema(input_schema)?;
         let group_values = new_group_values(group_schema, &group_ordering)?;
@@ -138,6 +152,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
 
         Ok(Self {
             output_schema,
+            batch_size,
             group_by_metrics: GroupByMetrics::new(&agg.metrics, partition),
             buffer: OrderedAggregateTableBuffer {
                 group_by: Arc::clone(&agg.group_by),
@@ -202,6 +217,24 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
             + self.buffer.group_ordering.size()
             + self.buffer.group_indices.allocated_size()
     }
+
+    /// Returns the [`EmitTo`], clamped to the specified batch size
+    ///
+    /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number
+    /// of groups to emit from `GroupValues` / accumulators, and
+    /// `should_remove_groups` indicates whether `GroupOrdering` must also shift
+    /// its tracked indexes.
+    pub(super) fn clamp_emit_to(
+        &self,
+        group_count: usize,
+        emit_to: EmitTo,
+    ) -> (EmitTo, bool) {
+        match emit_to {
+            EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true),
+            EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false),
+            EmitTo::All => (EmitTo::First(self.batch_size), false),
+        }
+    }
 }
 
 pub(super) fn remove_emitted_groups(group_ordering: &mut GroupOrdering, emit_to: EmitTo) {

datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs

@@ -45,13 +45,15 @@ impl OrderedAggregateTable<OrderedFinalMarker> {
         partition: usize,
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
+        batch_size: usize,
         input_order_mode: &InputOrderMode,
     ) -> Result<Self> {
         Self::new_for_mode(
             agg,
             partition,
             input_schema,
             output_schema,
+            batch_size,
             input_order_mode,
             &AggregateMode::Final,
             vec![None; agg.aggr_expr.len()],
@@ -109,10 +111,14 @@ impl OrderedAggregateTable<OrderedFinalMarker> {
         let Some(emit_to) = self.buffer.group_ordering.emit_to() else {
             return Ok(None);
         };
+        let (emit_to, should_remove_groups) =
+            self.clamp_emit_to(self.buffer.group_values.len(), emit_to);
 
         let timer = self.group_by_metrics.emitting_time.timer();
         let mut output = self.buffer.group_values.emit(emit_to)?;
-        remove_emitted_groups(&mut self.buffer.group_ordering, emit_to);
+        if should_remove_groups {
+            remove_emitted_groups(&mut self.buffer.group_ordering, emit_to);
+        }
 
         for acc in &mut self.buffer.accumulators {
             output.push(acc.evaluate(emit_to)?);

datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs

@@ -53,13 +53,15 @@ impl OrderedAggregateTable<OrderedPartialMarker> {
         agg: &AggregateExec,
         partition: usize,
         output_schema: SchemaRef,
+        batch_size: usize,
     ) -> Result<Self> {
         let input_schema = agg.input().schema();
         Self::new_for_mode(
             agg,
             partition,
             &input_schema,
             output_schema,
+            batch_size,
             &agg.input_order_mode,
             &AggregateMode::Partial,
             agg.filter_expr.iter().cloned().collect(),
@@ -127,10 +129,14 @@ impl OrderedAggregateTable<OrderedPartialMarker> {
         let Some(emit_to) = self.buffer.group_ordering.emit_to() else {
             return Ok(None);
         };
+        let (emit_to, should_remove_groups) =
+            self.clamp_emit_to(self.buffer.group_values.len(), emit_to);
 
         let timer = self.group_by_metrics.emitting_time.timer();
         let mut output = self.buffer.group_values.emit(emit_to)?;
-        remove_emitted_groups(&mut self.buffer.group_ordering, emit_to);
+        if should_remove_groups {
+            remove_emitted_groups(&mut self.buffer.group_ordering, emit_to);
+        }
 
         for acc in &mut self.buffer.accumulators {
             output.extend(acc.state(emit_to)?);

datafusion/physical-plan/src/aggregates/ordered_final_stream.rs

@@ -95,6 +95,7 @@ impl OrderedFinalAggregateStream {
 
         let schema = Arc::clone(&agg.schema);
         let input_schema = input.schema();
+        let batch_size = context.session_config().batch_size();
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
 
         // Preserve the existing aggregate metric surface for this plan node.
@@ -105,6 +106,7 @@ impl OrderedFinalAggregateStream {
             partition,
             &input_schema,
             Arc::clone(&schema),
+            batch_size,
             input_order_mode,
         )?;
         let reservation =

datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs

@@ -114,6 +114,7 @@ impl OrderedPartialAggregateStream {
 
         let schema = Arc::clone(&agg.schema);
         let input = agg.input.execute(partition, Arc::clone(context))?;
+        let batch_size = context.session_config().batch_size();
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
 
         // Preserve the existing aggregate metric surface for this plan node.
@@ -126,6 +127,7 @@ impl OrderedPartialAggregateStream {
             agg,
             partition,
             Arc::clone(&schema),
+            batch_size,
         )?;
         let reservation =
             MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]"))