Avoid repeated EmitTo::First in partial hash aggregate output (apache#23250)

## Which issue does this PR close?

- Closes apache#23249.

## Rationale for this change

The migrated partial hash aggregate output path still used
`EmitTo::First(batch_size)` when draining grouped aggregate state in
batches.

For terminal output this is unnecessary and can be expensive:
`EmitTo::First` is not just slicing the first N rows, it also shifts
remaining group indexes and maintains `GroupValues` lookup state. For
high-cardinality partial aggregate output, this can cause repeated work
during output draining.

The final hash aggregate path already avoids this by materializing
output once with `EmitTo::All` and then slicing the resulting
`RecordBatch`. This PR applies the same approach to partial hash
aggregate output.

## What changes are included in this PR?

- Remove the helper that selected `EmitTo::First(batch_size)` for hash
aggregate terminal output.
- Change migrated partial hash aggregate output to:
  - materialize grouped keys and aggregate state once with `EmitTo::All`
- slice the materialized `RecordBatch` into `batch_size` chunks across
output polls
- Rename the shared materialized-output state/type to mode-neutral names
because it is now used by both final and partial output paths.
- Add a regression test with a custom `GroupsAccumulator` that fails if
partial terminal output calls `EmitTo::First(_)`.
- Strengthen the regression test to verify both batch slicing and
emitted key/state values.

## Are these changes tested?

Yes.

Local targeted tests:
```bash
cargo test -p datafusion-physical-plan partial_grouped_aggregate_materializes_before_slicing -- --nocapture
cargo test -p datafusion-physical-plan materialized_aggregate_output_slices_batches_until_exhausted -- --nocapture
git diff --check
```
Additional local verification run during development:

```bash
cargo test -p datafusion-physical-plan materialized_final_output_slices_batches_until_exhausted -- --nocapture
cargo test -p datafusion-physical-plan partial_grouped_aggregate_uses_raw_partial_stream -- --nocapture
```
The new regression test was also applied to the pre-fix baseline and
failed with the expected internal error when the partial output path
used `EmitTo::First`.
Local benchmark evidence was collected against the implementation commit
before the final test/naming polish commit.
ClickBench full 43-query run, 5 iterations, 24 cores skip partial
aggregation probe ratio `0.8`:

mode | total warm time | geomean warm time
-- | -- | --
baseline migrated aggregate | 128509.47 ms | 352.79 ms
patched migrated aggregate | 19652.37 ms | 180.65 ms
baseline old aggregate path | 19774.70 ms | 181.25 ms


Largest patched/current wins included:
q33: 32961.02ms -> 1642.08ms
q34: 32739.34ms -> 1635.07ms
q18: 25673.25ms -> 1767.25ms
q16: 5949.82ms -> 810.17ms
q17: 5906.51ms -> 807.10ms
TPC-DS SF10 full99, 10 rounds:
Failures: 0
Aggregate geomean current/main: 0.982817
Aggregate current speedup: 1.748%

## Are there any user-facing changes?
No. This is an internal physical execution change for hash aggregate
output draining. There are no public API or documented behavior changes.

---------

Co-authored-by: Qiwei Huang <qiwei.huang@jsessh.com>
Co-authored-by: kamille <kamille@apache.org>

Author: 3 people

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

@@ -182,7 +182,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                 acc + state.group_values.size()
                     + state.batch_group_indices.allocated_size()
             }
-            AggregateHashTableState::OutputtingMaterializedFinal(output) => {
+            AggregateHashTableState::OutputtingMaterialized(output) => {
                 output.memory_size()
             }
             AggregateHashTableState::Done => 0,
@@ -214,15 +214,6 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
     }
 }
 
-pub(super) fn emit_to_for_batch_size(batch_size: usize, group_count: usize) -> EmitTo {
-    debug_assert!(batch_size > 0);
-    if group_count <= batch_size {
-        EmitTo::All
-    } else {
-        EmitTo::First(batch_size)
-    }
-}
-
 /// State and argument information for a single Aggregate
 ///
 /// For example, for `SELECT COUNT(x), SUM(y WHERE z > 10) ...`  there would be two
@@ -304,24 +295,25 @@ pub(super) enum AggregateHashTableState {
     Building(AggregateHashTableBuffer),
     /// Emitting results directly from group keys and aggregate state.
     Outputting(AggregateHashTableBuffer),
-    /// Materialize all the output results, and then incrementally output in the `OutputtingMaterializedFinal` state.
+    /// Materialize all the output results, and then incrementally output in the `OutputtingMaterialized` state.
     ///
     /// Note this is a temporary solution until the `GroupValues` issue is solved:
     /// Issue: <https://github.com/apache/datafusion/issues/23178>
-    OutputtingMaterializedFinal(MaterializedFinalOutput),
+    OutputtingMaterialized(MaterializedAggregateOutput),
     Done,
 }
 
-/// Fully evaluated final aggregate output and the next row offset to emit.
+/// Fully evaluated aggregate output and the next row offset to emit.
 ///
-/// Final aggregate evaluation consumes accumulator state, so final output is
-/// materialized once and then sliced to honor `batch_size` across output polls.
-pub(super) struct MaterializedFinalOutput {
+/// Final aggregate evaluation consumes accumulator state, and partial terminal
+/// output should not repeatedly renumber group values with `EmitTo::First`.
+/// Materialize once and then slice to honor `batch_size` across output polls.
+pub(super) struct MaterializedAggregateOutput {
     batch: RecordBatch,
     offset: usize,
 }
 
-impl MaterializedFinalOutput {
+impl MaterializedAggregateOutput {
     pub(super) fn new(batch: RecordBatch) -> Self {
         Self { batch, offset: 0 }
     }
@@ -496,7 +488,7 @@ mod tests {
     use super::*;
 
     #[test]
-    fn materialized_final_output_slices_batches_until_exhausted() -> Result<()> {
+    fn materialized_aggregate_output_slices_batches_until_exhausted() -> Result<()> {
         let schema = Arc::new(Schema::new(vec![Field::new(
             "group_col",
             DataType::Int32,
@@ -506,7 +498,7 @@ mod tests {
             schema,
             vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
         )?;
-        let mut output = MaterializedFinalOutput::new(batch);
+        let mut output = MaterializedAggregateOutput::new(batch);
 
         assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]);
         assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]);

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

@@ -26,7 +26,7 @@ use crate::aggregates::AggregateExec;
 
 use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker,
-    MaterializedFinalOutput,
+    MaterializedAggregateOutput,
 };
 
 /// Methods specific to the aggregate hash table used in the final aggregation stage.
@@ -57,8 +57,8 @@ impl AggregateHashTable<FinalMarker> {
     ) -> Result<Option<RecordBatch>> {
         let output_schema = Arc::clone(&self.output_schema);
         let batch_size = self.batch_size;
-        // Take ownership of the output state. Note `emit_next_materialized_batch`
-        // updates state after it emits a materialized slice.
+        // Take ownership of the output state. `emit_next_materialized_batch`
+        // restores `self.state` to `OutputtingMaterialized` or `Done`.
         match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
             AggregateHashTableState::Outputting(state) => {
                 if state.group_values.is_empty() {
@@ -68,7 +68,7 @@ impl AggregateHashTable<FinalMarker> {
                 let output = self.materialize_final_output(state, output_schema)?;
                 Ok(self.emit_next_materialized_batch(output, batch_size))
             }
-            AggregateHashTableState::OutputtingMaterializedFinal(output) => {
+            AggregateHashTableState::OutputtingMaterialized(output) => {
                 Ok(self.emit_next_materialized_batch(output, batch_size))
             }
             AggregateHashTableState::Done => Ok(None),
@@ -82,7 +82,7 @@ impl AggregateHashTable<FinalMarker> {
         &self,
         mut state: AggregateHashTableBuffer,
         output_schema: SchemaRef,
-    ) -> Result<MaterializedFinalOutput> {
+    ) -> Result<MaterializedAggregateOutput> {
         // Final aggregate evaluation consumes accumulator state. Evaluate all
         // groups once, then slice the materialized batch on subsequent polls.
         let emit_to = EmitTo::All;
@@ -96,19 +96,19 @@ impl AggregateHashTable<FinalMarker> {
 
         let batch = RecordBatch::try_new(output_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
-        Ok(MaterializedFinalOutput::new(batch))
+        Ok(MaterializedAggregateOutput::new(batch))
     }
 
     fn emit_next_materialized_batch(
         &mut self,
-        mut output: MaterializedFinalOutput,
+        mut output: MaterializedAggregateOutput,
         batch_size: usize,
     ) -> Option<RecordBatch> {
         let batch = output.next_batch(batch_size);
         if output.is_exhausted() {
             self.state = AggregateHashTableState::Done;
         } else {
-            self.state = AggregateHashTableState::OutputtingMaterializedFinal(output);
+            self.state = AggregateHashTableState::OutputtingMaterialized(output);
         }
         batch
     }

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

@@ -23,15 +23,16 @@ use arrow::array::{ArrayRef, BooleanArray, new_null_array};
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, assert_eq_or_internal_err, internal_err};
+use datafusion_expr::EmitTo;
 
 use crate::aggregates::group_values::new_group_values;
 use crate::aggregates::order::GroupOrdering;
 use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal};
 
 use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState,
-    EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker,
-    emit_to_for_batch_size,
+    EvaluatedAccumulatorArgs, HashAggregateAccumulator, MaterializedAggregateOutput,
+    PartialMarker, PartialSkipMarker,
 };
 
 /// Methods specific to the aggregate hash table used in the partial aggregation stage.
@@ -62,43 +63,60 @@ impl AggregateHashTable<PartialMarker> {
     ) -> Result<Option<RecordBatch>> {
         let output_schema = Arc::clone(&self.output_schema);
         let batch_size = self.batch_size;
-        match &mut self.state {
+        // Take ownership of the output state. `emit_next_materialized_batch`
+        // restores `self.state` to `OutputtingMaterialized` or `Done`.
+        match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
             AggregateHashTableState::Outputting(state) => {
                 if state.group_values.is_empty() {
-                    self.state = AggregateHashTableState::Done;
                     return Ok(None);
                 }
 
-                let emit_to =
-                    emit_to_for_batch_size(batch_size, state.group_values.len());
-                let timer = self.group_by_metrics.emitting_time.timer();
-                let mut output = state.group_values.emit(emit_to)?;
-
-                for acc in state.accumulators.iter_mut() {
-                    output.extend(acc.state(emit_to)?);
-                }
-                let done = state.group_values.is_empty();
-                drop(timer);
-
-                let batch = RecordBatch::try_new(output_schema, output)?;
-                debug_assert!(batch.num_rows() > 0);
-                if done {
-                    self.state = AggregateHashTableState::Done;
-                }
-                Ok(Some(batch))
+                let output = self.materialize_partial_output(state, output_schema)?;
+                Ok(self.emit_next_materialized_batch(output, batch_size))
+            }
+            AggregateHashTableState::OutputtingMaterialized(output) => {
+                Ok(self.emit_next_materialized_batch(output, batch_size))
             }
             AggregateHashTableState::Done => Ok(None),
             AggregateHashTableState::Building(_) => {
                 internal_err!("next_output_batch must be called in the outputting state")
             }
-            AggregateHashTableState::OutputtingMaterializedFinal(_) => {
-                internal_err!(
-                    "partial aggregate output should not materialize final output"
-                )
-            }
         }
     }
 
+    fn materialize_partial_output(
+        &self,
+        mut state: AggregateHashTableBuffer,
+        output_schema: SchemaRef,
+    ) -> Result<MaterializedAggregateOutput> {
+        let emit_to = EmitTo::All;
+        let timer = self.group_by_metrics.emitting_time.timer();
+        let mut output = state.group_values.emit(emit_to)?;
+
+        for acc in state.accumulators.iter_mut() {
+            output.extend(acc.state(emit_to)?);
+        }
+        drop(timer);
+
+        let batch = RecordBatch::try_new(output_schema, output)?;
+        debug_assert!(batch.num_rows() > 0);
+        Ok(MaterializedAggregateOutput::new(batch))
+    }
+
+    fn emit_next_materialized_batch(
+        &mut self,
+        mut output: MaterializedAggregateOutput,
+        batch_size: usize,
+    ) -> Option<RecordBatch> {
+        let batch = output.next_batch(batch_size);
+        if output.is_exhausted() {
+            self.state = AggregateHashTableState::Done;
+        } else {
+            self.state = AggregateHashTableState::OutputtingMaterialized(output);
+        }
+        batch
+    }
+
     pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool {
         self.state
             .building()