apache
diff --git a/‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs‎
Lines changed: 25 additions & 29 deletions b/‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs‎
Lines changed: 25 additions & 29 deletions
diff --git a/‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs‎
Lines changed: 12 additions & 16 deletions b/‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs‎
Lines changed: 12 additions & 16 deletions
@@ -362,24 +362,6 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         Ok(Some(batch))
     }
 
-    /// 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),
-        }
-    }
-
     /// Aggregates one evaluated input batch after selecting the mode-specific
     /// accumulator operation.
     ///
@@ -455,20 +437,34 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         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 emit_to = match emit_to {
+            EmitTo::First(n) => EmitTo::First(n.min(self.batch_size)),
+            EmitTo::All if self.num_groups() > self.batch_size => {
+                EmitTo::First(self.batch_size)
+            }
+            EmitTo::All => EmitTo::All,
+        };
+        self.materialize_groups(emit_to, materialize_accumulator_fn, accumulator_phase)
+            .map(Some)
+    }
 
+    /// Removes the selected groups once and materializes their output columns.
+    /// The caller chooses the completed prefix and any output-size limit.
+    pub(super) fn materialize_groups(
+        &mut self,
+        emit_to: EmitTo,
+        materialize_accumulator_fn: MaterializeAccumulatorFn,
+        accumulator_phase: AccumulatorPhase,
+    ) -> Result<RecordBatch> {
         let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics);
         let output = self.group_by_metrics.time_emitting(|| {
             let mut output = self.buffer.group_values.emit(emit_to)?;
-            if should_remove_groups {
-                match emit_to {
-                    EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n),
-                    // `EmitTo::All` is only used after `input_done`, when all
-                    // buffered groups are known complete and the ordering state is
-                    // no longer needed.
-                    EmitTo::All => {}
-                }
+            // EOF can also emit a prefix when a caller limits its batch size,
+            // but the completed ordering state no longer tracks group indexes.
+            if let EmitTo::First(n) = emit_to
+                && matches!(self.buffer.group_ordering.emit_to(), Some(EmitTo::First(_)))
+            {
+                self.buffer.group_ordering.remove_groups(n);
             }
 
             for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() {
@@ -484,6 +480,6 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?;
         debug_assert!(batch.num_rows() > 0);
 
-        Ok(Some(batch))
+        Ok(batch)
     }
 }
@@ -89,26 +89,22 @@ impl OrderedAggregateTable<PartialMarker> {
         )
     }
 
-    /// Emits the next batch of partial state rows for groups proven complete by
-    /// the input ordering.
-    ///
-    /// For example, when the query is `GROUP BY a` and the input is ordered by
-    /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3`
-    /// are complete and safe to emit.
-    ///
-    /// Key steps:
-    /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly.
-    /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and
-    ///    all `GroupsAccumulator`s.
-    ///
-    /// This may output small batches. Avoiding tiny batches is left to future
-    /// ordered-aggregation optimizations.
-    pub(in crate::aggregates) fn next_output_batch(
+    /// Materializes all groups proven complete by the input ordering, leaving
+    /// the active ordered-key range in the table.
+    pub(in crate::aggregates) fn take_completed_state_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_inner(
+        if self.is_empty() {
+            return Ok(None);
+        }
+        let Some(emit_to) = self.group_ordering().emit_to() else {
+            return Ok(None);
+        };
+        self.materialize_groups(
+            emit_to,
             HashAggregateAccumulator::state,
             AccumulatorPhase::State,
         )
+        .map(Some)
     }
 }