This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23324-3089ace491701b1b1d0cb047df66d4821dfdf15e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 301e684b583fb2d3661f73c6489f81891f1150aa Author: Yongting You <[email protected]> AuthorDate: Thu Jul 9 03:08:53 2026 +0800 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 https://github.com/apache/datafusion/issues/22710 - An alternative for and closes https://github.com/apache/datafusion/pull/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 https://github.com/apache/datafusion/pull/23309#discussion_r3522932671 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. --> --- .../src/aggregates/aggregate_hash_table/common.rs | 121 +++++++++++++++++++++ .../aggregates/aggregate_hash_table/final_table.rs | 92 +--------------- .../aggregate_hash_table/partial_reduce_table.rs | 93 +--------------- .../aggregate_hash_table/partial_table.rs | 85 +-------------- 4 files changed, 139 insertions(+), 252 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 f29f3e7ff8..e6e690c4d1 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,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) 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 568b866b10..522cc9066b 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, 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<()> { 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 4d94c55943..d8e92c5928 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, 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<()> { 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 f11eef8c14..ffac42feaa 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,7 @@ impl AggregateHashTable<PartialMarker> { 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_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<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 + self.next_output_batch_inner(HashAggregateAccumulator::state) } pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool { @@ -161,31 +106,13 @@ impl AggregateHashTable<PartialMarker> { }) } + /// 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, HashAggregateAccumulator::update_batch) } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
