From 90d661749a31ce69ec8cd825dbefa9c9598006a1 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Sun, 5 Jul 2026 10:16:39 +0800 Subject: [PATCH 1/4] refactor(hash-aggr): Simplify aggregate hash table with tempated functions --- .../aggregates/aggregate_hash_table/common.rs | 96 +++++++++++++++++ .../aggregate_hash_table/final_table.rs | 98 +++-------------- .../partial_reduce_table.rs | 101 +++--------------- .../aggregate_hash_table/partial_table.rs | 91 +++------------- 4 files changed, 136 insertions(+), 250 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index a435360ca3364..39c3b7eb2a389 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -171,6 +171,102 @@ impl AggregateHashTable { }) } + /// Aggregates one input batch after selecting the mode-specific accumulator + /// operation. + pub(super) fn aggregate_batch_inner( + &mut self, + batch: &RecordBatch, + mut aggregate_fn: F, + ) -> Result<()> + where + F: FnMut( + &mut HashAggregateAccumulator, + &EvaluatedAccumulatorArgs, + &[usize], + usize, + ) -> 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 us [`GroupsAccumulator::state`]. + /// + /// This is a temporary solution until blocked state management is implemented: + /// Issue: + pub(super) fn next_output_batch_inner( + &mut self, + mut materialize_accumulator_fn: F, + ) -> Result> + where + F: FnMut(&mut HashAggregateAccumulator, EmitTo, &mut Vec) -> Result<()>, + { + 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() { + materialize_accumulator_fn(acc, emit_to, &mut columns)?; + } + 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) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 568b866b10517..63b773f452f3e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/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}; /// Implementation specific to final aggregation, where the table stores partial /// aggregate states and the input rows are also partial states. @@ -61,90 +55,24 @@ impl AggregateHashTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - 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 { - // 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() { + self.next_output_batch_inner(|acc, emit_to, output| { 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 { - let batch = output.next_batch(batch_size); - if output.is_exhausted() { - self.state = AggregateHashTableState::Done; - } else { - self.state = AggregateHashTableState::OutputtingMaterialized(output); - } - batch + Ok(()) + }) } + /// 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, + |acc, values, group_indices, total_num_groups| { + acc.merge_batch(values, group_indices, total_num_groups) + }, + ) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index 4d94c559436fb..b1dc213b3efd5 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/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, PartialReduceMarker}; /// Methods specific to the aggregate hash table used in the partial-reduce stage. impl AggregateHashTable { @@ -55,91 +49,24 @@ impl AggregateHashTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - 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 { - // `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 { - 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(|acc, emit_to, output| { + output.extend(acc.state(emit_to)?); + Ok(()) + }) } + /// 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, + |acc, values, group_indices, total_num_groups| { + acc.merge_batch(values, group_indices, total_num_groups) + }, + ) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index f11eef8c14277..32fad90938e49 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -22,8 +22,7 @@ use std::sync::Arc; 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 datafusion_common::{Result, assert_eq_or_internal_err}; use crate::aggregates::group_values::new_group_values; use crate::aggregates::order::GroupOrdering; @@ -31,8 +30,7 @@ use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, - EvaluatedAccumulatorArgs, HashAggregateAccumulator, MaterializedAggregateOutput, - PartialMarker, PartialSkipMarker, + EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, }; /// Implementation specific to partial aggregation, where the table stores @@ -67,60 +65,10 @@ impl AggregateHashTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - 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_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") - } - } - } - - fn materialize_partial_output( - &self, - mut state: AggregateHashTableBuffer, - output_schema: SchemaRef, - ) -> Result { - 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() { + self.next_output_batch_inner(|acc, emit_to, output| { 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 { - let batch = output.next_batch(batch_size); - if output.is_exhausted() { - self.state = AggregateHashTableState::Done; - } else { - self.state = AggregateHashTableState::OutputtingMaterialized(output); - } - batch + Ok(()) + }) } pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool { @@ -161,31 +109,18 @@ impl AggregateHashTable { }) } + /// Partial aggregation consumes raw input rows and updates 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.update_batch(values, group_indices, total_num_groups)?; - } - } - - Ok(()) + self.aggregate_batch_inner( + batch, + |acc, values, group_indices, total_num_groups| { + acc.update_batch(values, group_indices, total_num_groups) + }, + ) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { From a1c9b6a2d0bb5eaa1e529c81a21edc8646405039 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Sun, 5 Jul 2026 10:23:22 +0800 Subject: [PATCH 2/4] more comments --- .../src/aggregates/aggregate_hash_table/common.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 39c3b7eb2a389..71c297064d5d4 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -173,6 +173,10 @@ impl AggregateHashTable { /// 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, From 32691af0c8593694f2a051de3fc8aafed5bd9d90 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Wed, 8 Jul 2026 16:40:35 +0800 Subject: [PATCH 3/4] review: exract AggregateBatchFn to a separate type --- .../aggregates/aggregate_hash_table/common.rs | 30 ++++++++++++------- .../aggregate_hash_table/final_table.rs | 9 ++---- .../partial_reduce_table.rs | 9 ++---- .../aggregate_hash_table/partial_table.rs | 7 +---- 4 files changed, 24 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 71c297064d5d4..e9e0371859d4a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -177,19 +177,11 @@ impl AggregateHashTable { /// 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( + pub(super) fn aggregate_batch_inner( &mut self, batch: &RecordBatch, - mut aggregate_fn: F, - ) -> Result<()> - where - F: FnMut( - &mut HashAggregateAccumulator, - &EvaluatedAccumulatorArgs, - &[usize], - usize, - ) -> Result<()>, - { + aggregate_fn: AggregateBatchFn, + ) -> Result<()> { let evaluated_batch = self.evaluate_batch(batch)?; let state = self.state.building_mut(); @@ -341,6 +333,22 @@ 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<()>; + /// Evaluated aggregate arguments and filter for one input batch. /// /// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 63b773f452f3e..84c220410a970 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -21,7 +21,7 @@ use datafusion_common::Result; use crate::aggregates::AggregateExec; -use super::common::{AggregateHashTable, FinalMarker}; +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. @@ -67,12 +67,7 @@ impl AggregateHashTable { &mut self, batch: &RecordBatch, ) -> Result<()> { - self.aggregate_batch_inner( - batch, - |acc, values, group_indices, total_num_groups| { - acc.merge_batch(values, group_indices, total_num_groups) - }, - ) + self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index b1dc213b3efd5..2e119706f13c2 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs @@ -21,7 +21,7 @@ use datafusion_common::Result; use crate::aggregates::AggregateExec; -use super::common::{AggregateHashTable, PartialReduceMarker}; +use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker}; /// Methods specific to the aggregate hash table used in the partial-reduce stage. impl AggregateHashTable { @@ -61,12 +61,7 @@ impl AggregateHashTable { &mut self, batch: &RecordBatch, ) -> Result<()> { - self.aggregate_batch_inner( - batch, - |acc, values, group_indices, total_num_groups| { - acc.merge_batch(values, group_indices, total_num_groups) - }, - ) + self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 32fad90938e49..6007116f029a2 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -115,12 +115,7 @@ impl AggregateHashTable { &mut self, batch: &RecordBatch, ) -> Result<()> { - self.aggregate_batch_inner( - batch, - |acc, values, group_indices, total_num_groups| { - acc.update_batch(values, group_indices, total_num_groups) - }, - ) + self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { From 4445f34df9682319df90b0f43a058fffa210fbcc Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Wed, 8 Jul 2026 16:56:20 +0800 Subject: [PATCH 4/4] review: exract MaterializeAccumulatorFn to a separate type --- .../aggregates/aggregate_hash_table/common.rs | 29 ++++++++++++++----- .../aggregate_hash_table/final_table.rs | 5 +--- .../partial_reduce_table.rs | 5 +--- .../aggregate_hash_table/partial_table.rs | 5 +--- 4 files changed, 24 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index e9e0371859d4a..b6e8976e57506 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -210,17 +210,14 @@ impl AggregateHashTable { /// /// 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 us [`GroupsAccumulator::state`]. + /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`]. /// /// This is a temporary solution until blocked state management is implemented: /// Issue: - pub(super) fn next_output_batch_inner( + pub(super) fn next_output_batch_inner( &mut self, - mut materialize_accumulator_fn: F, - ) -> Result> - where - F: FnMut(&mut HashAggregateAccumulator, EmitTo, &mut Vec) -> Result<()>, - { + materialize_accumulator_fn: MaterializeAccumulatorFn, + ) -> Result> { let output_schema = Arc::clone(&self.output_schema); let batch_size = self.batch_size; @@ -237,7 +234,7 @@ impl AggregateHashTable { 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() { - materialize_accumulator_fn(acc, emit_to, &mut columns)?; + columns.extend(materialize_accumulator_fn(acc, emit_to)?); } drop(timer); @@ -349,6 +346,15 @@ pub(super) type AggregateBatchFn = fn( 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>; + /// Evaluated aggregate arguments and filter for one input batch. /// /// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` @@ -547,6 +553,13 @@ impl HashAggregateAccumulator { self.accumulator.evaluate(emit_to) } + pub(super) fn evaluate_to_columns( + &mut self, + emit_to: EmitTo, + ) -> Result> { + 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) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 84c220410a970..522cc9066b14b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -55,10 +55,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_inner(|acc, emit_to, output| { - output.push(acc.evaluate(emit_to)?); - Ok(()) - }) + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) } /// Final aggregation consumes partial aggregate states and merges them into diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index 2e119706f13c2..d8e92c5928b8a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs @@ -49,10 +49,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_inner(|acc, emit_to, output| { - output.extend(acc.state(emit_to)?); - Ok(()) - }) + self.next_output_batch_inner(HashAggregateAccumulator::state) } /// Partial-reduce aggregation consumes partial aggregate states and merges diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 6007116f029a2..ffac42feaa3b3 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -65,10 +65,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_inner(|acc, emit_to, output| { - output.extend(acc.state(emit_to)?); - Ok(()) - }) + self.next_output_batch_inner(HashAggregateAccumulator::state) } pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool {