Fix final hash aggregate output regression by materializing once (apache#23182)

## Which issue does this PR close?
 
- part of apache#22710
- short-term solution for apache#23178.
- closes apache#23178

## Rationale for this change

PR apache#23055 changed final hash aggregate output to emit groups
incrementally with
`EmitTo::First(batch_size)`. For terminal final aggregate output, this
can cause
the group value state to be repeatedly compacted while output batches
are being
produced. On TPC-DS q23 this showed up as a significant regression.

This PR implements the short-term approach discussed in apache#23178:
materialize the
final aggregate output once, then return slices of that materialized
`RecordBatch` according to `batch_size`.

This avoids changing the `GroupValues` API while preserving bounded
downstream
batch sizes.

## What changes are included in this PR?

- Adds an `OutputtingMaterialized` hash aggregate state.
- Adds `MaterializedOutput`, a small wrapper around a `RecordBatch` plus
output
  offset.
- Changes final hash aggregate output to:
  - emit all final groups once,
  - evaluate all final aggregate values once,
  - slice the materialized batch for subsequent output polling.
- Leaves partial aggregate output behavior unchanged.
- Adds focused tests for materialized output slicing and final hash
aggregate
  output state transitions.

## Performance
TPC-DS SF10 full 99 queries, 10 rounds:

- Total runtime ratio: `0.857051`
- Geomean ratio: `0.976652` (~2.4% faster)
- q23 ratio: `0.313770` (~218.7% faster), faster in `10/10` rounds

Regressions over 5% were observed in 10 queries. Most have small
absolute
deltas, but the largest slowdowns were:

- q67: `1.055907`, +170.996 ms
- q39: `1.060436`, +98.544 ms
- q9: `1.050135`, +37.858 ms
- q70: `1.061124`, +11.848 ms
- q35: `1.052392`, +9.386 ms
- q33: `1.063655`, +6.995 ms
- q98: `1.071688`, +6.515 ms
- q91: `1.109819`, +5.362 ms
- q15: `1.058356`, +5.072 ms
- q27: `1.057686`, +0.815 ms

Overall, this recovers the q23 regression strongly and improves
full-query
geomean, but q39 and q67 are worth calling out as residual per-query
slowdowns.

## Testing

- `cargo fmt --all -- --check`
- `cargo test -p datafusion-physical-plan materializ`
- `cargo test -p datafusion-physical-plan aggregates::`
- TPC-DS SF10 q23, 3 rounds
- TPC-DS SF10 full 99 queries, 10 rounds

---------

Co-authored-by: Qiwei Huang <qiwei.huang@jsessh.com>
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: 3 people

datafusion/physical-plan/Cargo.toml

@@ -123,6 +123,7 @@ required-features = ["test_utils"]
 [[bench]]
 harness = false
 name = "aggregate_vectorized"
+required-features = ["test_utils"]
 
 [[bench]]
 harness = false

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

@@ -182,6 +182,9 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                 acc + state.group_values.size()
                     + state.batch_group_indices.allocated_size()
             }
+            AggregateHashTableState::OutputtingMaterializedFinal(output) => {
+                output.memory_size()
+            }
             AggregateHashTableState::Done => 0,
         }
     }
@@ -297,11 +300,53 @@ pub(super) struct AggregateHashTableBuffer {
 }
 
 pub(super) enum AggregateHashTableState {
+    /// Accumulating input rows into group keys and aggregate state.
     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.
+    ///
+    /// Note this is a temporary solution until the `GroupValues` issue is solved:
+    /// Issue: <https://github.com/apache/datafusion/issues/23178>
+    OutputtingMaterializedFinal(MaterializedFinalOutput),
     Done,
 }
 
+/// Fully evaluated final 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 {
+    batch: RecordBatch,
+    offset: usize,
+}
+
+impl MaterializedFinalOutput {
+    pub(super) fn new(batch: RecordBatch) -> Self {
+        Self { batch, offset: 0 }
+    }
+
+    pub(super) fn next_batch(&mut self, batch_size: usize) -> Option<RecordBatch> {
+        debug_assert!(batch_size > 0);
+        if self.is_exhausted() {
+            return None;
+        }
+
+        let length = batch_size.min(self.batch.num_rows() - self.offset);
+        let batch = self.batch.slice(self.offset, length);
+        self.offset += length;
+        Some(batch)
+    }
+
+    pub(super) fn is_exhausted(&self) -> bool {
+        self.offset >= self.batch.num_rows()
+    }
+
+    pub(super) fn memory_size(&self) -> usize {
+        self.batch.get_array_memory_size()
+    }
+}
+
 impl HashAggregateAccumulator {
     fn new(
         aggregate_expr: Arc<AggregateFunctionExpr>,
@@ -440,3 +485,44 @@ impl AggregateHashTableState {
         state
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::array::{Array, Int32Array};
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use super::*;
+
+    #[test]
+    fn materialized_final_output_slices_batches_until_exhausted() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "group_col",
+            DataType::Int32,
+            false,
+        )]));
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
+        )?;
+        let mut output = MaterializedFinalOutput::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]);
+        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]);
+        assert!(output.next_batch(2).is_none());
+        assert!(output.is_exhausted());
+
+        Ok(())
+    }
+
+    fn int32_values(batch: &RecordBatch, column: usize) -> Vec<i32> {
+        let array = batch
+            .column(column)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        (0..array.len()).map(|idx| array.value(idx)).collect()
+    }
+}

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

@@ -20,11 +20,13 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, internal_err};
+use datafusion_expr::EmitTo;
 
 use crate::aggregates::AggregateExec;
 
 use super::common::{
-    AggregateHashTable, AggregateHashTableState, FinalMarker, emit_to_for_batch_size,
+    AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker,
+    MaterializedFinalOutput,
 };
 
 /// Methods specific to the aggregate hash table used in the final aggregation stage.
@@ -55,30 +57,19 @@ impl AggregateHashTable<FinalMarker> {
     ) -> 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. Note `emit_next_materialized_batch`
+        // updates state after it emits a materialized slice.
+        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.push(acc.evaluate(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_final_output(state, output_schema)?;
+                Ok(self.emit_next_materialized_batch(output, batch_size))
+            }
+            AggregateHashTableState::OutputtingMaterializedFinal(output) => {
+                Ok(self.emit_next_materialized_batch(output, batch_size))
             }
             AggregateHashTableState::Done => Ok(None),
             AggregateHashTableState::Building(_) => {
@@ -87,6 +78,41 @@ impl AggregateHashTable<FinalMarker> {
         }
     }
 
+    fn materialize_final_output(
+        &self,
+        mut state: AggregateHashTableBuffer,
+        output_schema: SchemaRef,
+    ) -> Result<MaterializedFinalOutput> {
+        // Final aggregate evaluation consumes accumulator state. Evaluate all
+        // groups once, then slice the materialized batch on subsequent polls.
+        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.push(acc.evaluate(emit_to)?);
+        }
+        drop(timer);
+
+        let batch = RecordBatch::try_new(output_schema, output)?;
+        debug_assert!(batch.num_rows() > 0);
+        Ok(MaterializedFinalOutput::new(batch))
+    }
+
+    fn emit_next_materialized_batch(
+        &mut self,
+        mut output: MaterializedFinalOutput,
+        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);
+        }
+        batch
+    }
+
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,

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

@@ -91,6 +91,11 @@ impl AggregateHashTable<PartialMarker> {
             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"
+                )
+            }
         }
     }