refactor(hash-aggr): Simplify aggregate hash table with tempated functions (#23324)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Part of #22710

- An alternative for and closes
#23309


## Rationale for this change

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
See #23309 and
#23309 (comment)
for background.

I prefer this approach, but I want to point out the tradeoff for this
PR's approach: the shared utility includes a complex lambda function
argument, this makes them harder to extend. But I think it's okay since
most functionality has been implemented for the refactor, and there are
not likely to have new functional requirements.

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

## Are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

Author: 2010YOUY01

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

@@ -171,6 +171,95 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         })
     }
 
+    /// Aggregates one input batch after selecting the mode-specific accumulator
+    /// operation.
+    ///
+    /// Each aggregation mode chooses a different `aggregate_fn` according to its
+    /// semantics. For example, partial aggregation takes raw inputs, and update them
+    /// into stored partial states, so [`GroupsAccumulator::update_batch`] is used.
+    pub(super) fn aggregate_batch_inner(
+        &mut self,
+        batch: &RecordBatch,
+        aggregate_fn: AggregateBatchFn,
+    ) -> Result<()> {
+        let evaluated_batch = self.evaluate_batch(batch)?;
+        let state = self.state.building_mut();
+
+        let _timer = self.group_by_metrics.aggregation_time.timer();
+        for group_values in &evaluated_batch.grouping_set_args {
+            state
+                .group_values
+                .intern(group_values, &mut state.batch_group_indices)?;
+            let group_indices = &state.batch_group_indices;
+            let total_num_groups = state.group_values.len();
+
+            for (acc, values) in state
+                .accumulators
+                .iter_mut()
+                .zip(evaluated_batch.accumulator_args.iter())
+            {
+                aggregate_fn(acc, values, group_indices, total_num_groups)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Materializes the full output once, then returns it downstream incrementally
+    /// by slicing it into `batch_size` chunks.
+    ///
+    /// Each aggregation mode chooses a different `materialize_accumulator_fn`
+    /// according to its semantics. For example, partial aggregation emits
+    /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`].
+    ///
+    /// This is a temporary solution until blocked state management is implemented:
+    /// Issue: <https://github.com/apache/datafusion/issues/7065>
+    pub(super) fn next_output_batch_inner(
+        &mut self,
+        materialize_accumulator_fn: MaterializeAccumulatorFn,
+    ) -> Result<Option<RecordBatch>> {
+        let output_schema = Arc::clone(&self.output_schema);
+        let batch_size = self.batch_size;
+
+        let mut output =
+            match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
+                AggregateHashTableState::Outputting(mut state) => {
+                    if state.group_values.is_empty() {
+                        return Ok(None);
+                    }
+
+                    // Accumulator output consumes internal state. Materialize all
+                    // groups once, then slice the materialized batch on later polls.
+                    let emit_to = EmitTo::All;
+                    let timer = self.group_by_metrics.emitting_time.timer();
+                    let mut columns = state.group_values.emit(emit_to)?;
+                    for acc in state.accumulators.iter_mut() {
+                        columns.extend(materialize_accumulator_fn(acc, emit_to)?);
+                    }
+                    drop(timer);
+
+                    let batch = RecordBatch::try_new(output_schema, columns)?;
+                    debug_assert!(batch.num_rows() > 0);
+                    MaterializedAggregateOutput::new(batch)
+                }
+                AggregateHashTableState::OutputtingMaterialized(output) => output,
+                AggregateHashTableState::Done => return Ok(None),
+                AggregateHashTableState::Building(_) => {
+                    return internal_err!(
+                        "next_output_batch must be called in the outputting state"
+                    );
+                }
+            };
+
+        let batch = output.next_batch(batch_size);
+        if output.is_exhausted() {
+            self.state = AggregateHashTableState::Done;
+        } else {
+            self.state = AggregateHashTableState::OutputtingMaterialized(output);
+        }
+        Ok(batch)
+    }
+
     pub(in crate::aggregates) fn memory_size(&self) -> usize {
         match &self.state {
             AggregateHashTableState::Building(state)
@@ -241,6 +330,31 @@ pub(super) struct HashAggregateAccumulator {
 
 pub(super) type AggregateAccumulator = HashAggregateAccumulator;
 
+/// Function used by [`AggregateHashTable::aggregate_batch_inner`] to update one
+/// accumulator with one evaluated input batch.
+///
+/// Arguments:
+/// * accumulator to update.
+/// * accumulator's evaluated arguments and optional filter.
+/// * one group index per input row, mapping each row to its interned group.
+/// * total number of groups currently interned in that buffer, including newly
+///   interned groups.
+pub(super) type AggregateBatchFn = fn(
+    &mut AggregateAccumulator,
+    &EvaluatedAccumulatorArgs,
+    &[usize],
+    usize,
+) -> Result<()>;
+
+/// Function used by [`AggregateHashTable::next_output_batch_inner`] to
+/// materialize one accumulator's output columns.
+///
+/// Arguments:
+/// * accumulator to materialize.
+/// * group range to emit from the accumulator.
+pub(super) type MaterializeAccumulatorFn =
+    fn(&mut AggregateAccumulator, EmitTo) -> Result<Vec<ArrayRef>>;
+
 /// Evaluated aggregate arguments and filter for one input batch.
 ///
 /// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1`
@@ -439,6 +553,13 @@ impl HashAggregateAccumulator {
         self.accumulator.evaluate(emit_to)
     }
 
+    pub(super) fn evaluate_to_columns(
+        &mut self,
+        emit_to: EmitTo,
+    ) -> Result<Vec<ArrayRef>> {
+        Ok(vec![self.evaluate(emit_to)?])
+    }
+
     /// Evaluating partial aggregate results according to `EmitTo`, and reset inner
     /// states. (e.g. after `state(EmitTo::All)`, it returns all accumulated groups
     /// , and clear the inner buffers)

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

@@ -15,19 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
-
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
-use datafusion_common::{Result, internal_err};
-use datafusion_expr::EmitTo;
+use datafusion_common::Result;
 
 use crate::aggregates::AggregateExec;
 
-use super::common::{
-    AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker,
-    MaterializedAggregateOutput,
-};
+use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator};
 
 /// Implementation specific to final aggregation, where the table stores partial
 /// aggregate states and the input rows are also partial states.
@@ -61,90 +55,16 @@ impl AggregateHashTable<FinalMarker> {
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
-        let output_schema = Arc::clone(&self.output_schema);
-        let batch_size = self.batch_size;
-        // 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() {
-                    return Ok(None);
-                }
-
-                let output = self.materialize_final_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")
-            }
-        }
-    }
-
-    fn materialize_final_output(
-        &self,
-        mut state: AggregateHashTableBuffer,
-        output_schema: SchemaRef,
-    ) -> 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;
-        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(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
+        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
     }
 
+    /// Final aggregation consumes partial aggregate states and merges them into
+    /// the table's partial-state accumulators.
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,
     ) -> Result<()> {
-        let evaluated_batch = self.evaluate_batch(batch)?;
-        let state = self.state.building_mut();
-
-        let timer = self.group_by_metrics.aggregation_time.timer();
-        for group_values in &evaluated_batch.grouping_set_args {
-            state
-                .group_values
-                .intern(group_values, &mut state.batch_group_indices)?;
-            let group_indices = &state.batch_group_indices;
-            let total_num_groups = state.group_values.len();
-
-            for (acc, values) in state
-                .accumulators
-                .iter_mut()
-                .zip(evaluated_batch.accumulator_args.iter())
-            {
-                acc.merge_batch(values, group_indices, total_num_groups)?;
-            }
-        }
-        drop(timer);
-
-        Ok(())
+        self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch)
     }
 
     pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {

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

@@ -15,19 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
-
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
-use datafusion_common::{Result, internal_err};
-use datafusion_expr::EmitTo;
+use datafusion_common::Result;
 
 use crate::aggregates::AggregateExec;
 
-use super::common::{
-    AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState,
-    MaterializedAggregateOutput, PartialReduceMarker,
-};
+use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker};
 
 /// Methods specific to the aggregate hash table used in the partial-reduce stage.
 impl AggregateHashTable<PartialReduceMarker> {
@@ -55,91 +49,16 @@ impl AggregateHashTable<PartialReduceMarker> {
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
     ) -> 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.
-        match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
-            AggregateHashTableState::Outputting(state) => {
-                if state.group_values.is_empty() {
-                    return Ok(None);
-                }
-
-                let output =
-                    self.materialize_partial_reduce_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")
-            }
-        }
-    }
-
-    fn materialize_partial_reduce_output(
-        &self,
-        mut state: AggregateHashTableBuffer,
-        output_schema: SchemaRef,
-    ) -> Result<MaterializedAggregateOutput> {
-        // `state(EmitTo::All)` consumes accumulator state. Emit all groups once,
-        // then slice the materialized batch on subsequent polls.
-        let emit_to_all = EmitTo::All;
-        let timer = self.group_by_metrics.emitting_time.timer();
-        let mut output = state.group_values.emit(emit_to_all)?;
-
-        for acc in state.accumulators.iter_mut() {
-            output.extend(acc.state(emit_to_all)?);
-        }
-        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
+        self.next_output_batch_inner(HashAggregateAccumulator::state)
     }
 
+    /// Partial-reduce aggregation consumes partial aggregate states and merges
+    /// them into the table's partial-state accumulators.
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,
     ) -> Result<()> {
-        let evaluated_batch = self.evaluate_batch(batch)?;
-        let state = self.state.building_mut();
-
-        let timer = self.group_by_metrics.aggregation_time.timer();
-        for group_values in &evaluated_batch.grouping_set_args {
-            state
-                .group_values
-                .intern(group_values, &mut state.batch_group_indices)?;
-            let group_indices = &state.batch_group_indices;
-            let total_num_groups = state.group_values.len();
-
-            for (acc, values) in state
-                .accumulators
-                .iter_mut()
-                .zip(evaluated_batch.accumulator_args.iter())
-            {
-                acc.merge_batch(values, group_indices, total_num_groups)?;
-            }
-        }
-        drop(timer);
-
-        Ok(())
+        self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch)
     }
 
     pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {