From 7f68c238d8de5a3a8eb461504cf25c8b636c2940 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Fri, 3 Jul 2026 21:12:07 +0800 Subject: [PATCH 1/3] refactor(hash-aggr): consolidate aggregate hash table --- .../aggregates/aggregate_hash_table/common.rs | 392 ++++++++++++++++-- .../aggregate_hash_table/final_table.rs | 154 ------- .../aggregates/aggregate_hash_table/mod.rs | 7 +- .../partial_reduce_table.rs | 149 ------- .../aggregate_hash_table/partial_table.rs | 298 ------------- .../src/aggregates/hash_aggregate.rs | 68 ++- .../src/aggregates/partial_reduce_stream.rs | 35 +- 7 files changed, 411 insertions(+), 692 deletions(-) delete mode 100644 datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs delete mode 100644 datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs delete mode 100644 datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs 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..a90a67f6fd358 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -15,13 +15,13 @@ // specific language governing permissions and limitations // under the License. -use std::marker::PhantomData; +use std::collections::HashMap; use std::sync::Arc; -use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::array::{ArrayRef, AsArray, BooleanArray, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{Result, internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; @@ -31,19 +31,31 @@ use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_val use crate::aggregates::order::GroupOrdering; use crate::aggregates::row_hash::create_group_accumulator; use crate::aggregates::{ - AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, group_id_array, max_duplicate_ordinal, }; -/// Marker for raw rows -> partial state aggregation. +/// Semantic mode for the aggregate hash table. +/// +/// See [`AggregateHashTable`] comment's 'Mode' section for details. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub(in crate::aggregates) enum AggregateTableMode { + /// Raw rows -> partial aggregate state rows. + Partial, + /// Partial aggregate state rows -> final aggregate value rows. + Final, + /// Partial aggregate state rows -> merged partial aggregate state rows. + PartialReduce, + /// Raw rows -> final aggregate value rows. + Single, +} + +/// Marker for ordered raw rows -> partial state aggregation. pub(in crate::aggregates) struct PartialMarker; -/// Marker for partial state -> partial state aggregation. -pub(in crate::aggregates) struct PartialReduceMarker; -/// Marker for raw rows -> partial state conversion without aggregation. -pub(in crate::aggregates) struct PartialSkipMarker; -/// Marker for partial state -> final value aggregation. +/// Marker for ordered partial state -> final value aggregation. pub(in crate::aggregates) struct FinalMarker; -/// Grouped hash table shared by the partial and final paths. +/// Grouped hash table shared by different hash aggregation modes. /// /// While building, it consumes input batches and updates group / accumulator /// state. While outputting, it incrementally drains that state into output @@ -60,19 +72,36 @@ pub(in crate::aggregates) struct FinalMarker; /// [`GroupsAccumulator`]. Both use columnar storage so aggregation can stay /// vectorized. /// -/// # Marker Type -/// `AggrMode` selects the aggregate semantics. +/// # Mode +/// +/// [`AggregateTableMode`] controls how input batches update accumulator state +/// and how output batches are materialized. /// -/// e.g. `AggregateHashTable::::new(...)` creates an aggregate hash table -/// for the partial hash aggregate stage, the input schema is raw rows and output -/// schema is intermediate states. +/// ```text +/// Example: `AVG(x) GROUP BY k`. /// -/// It is a zero-sized compile-time marker, so each stage keeps its update logic -/// in a separate impl block, to make the behavior difference explicit. -pub(in crate::aggregates) struct AggregateHashTable { +/// In `Partial` mode, the table stores partial state: +/// k, sum(x), count(x) +/// The input batch contains raw values: +/// k, x +/// The output batch also contains partial state: +/// k, sum(x), count(x) +/// ``` +/// +/// So input uses [`GroupsAccumulator::update_batch`], and output uses +/// [`GroupsAccumulator::state`]. +/// +/// Other modes use different input/output combinations: +/// - `Final`: merge_batch + evaluate +/// - `PartialReduce`: merge_batch + state +/// - `Single`: update_batch + evaluate +pub(in crate::aggregates) struct AggregateHashTable { /// Grouping and accumulator-specific timing metrics. pub(super) group_by_metrics: GroupByMetrics, + /// Semantic behavior for this table. + pub(super) mode: AggregateTableMode, + /// Raw input schema, used to evaluate expressions and synthesize empty /// grouping-set rows. pub(super) input_schema: SchemaRef, @@ -85,29 +114,63 @@ pub(in crate::aggregates) struct AggregateHashTable { /// Lifecycle-specific state: building stage / outputting stage. pub(super) state: AggregateHashTableState, - - pub(super) _mode: PhantomData, } /// Methods shared by all aggregate hash table modes. -impl AggregateHashTable { - pub(super) fn new_with_filters( +impl AggregateHashTable { + pub(in crate::aggregates) fn new( agg: &AggregateExec, partition: usize, output_schema: SchemaRef, batch_size: usize, - filters: Vec>>, ) -> Result { if batch_size == 0 { return internal_err!("AggregateHashTable requires config batch_size >= 1"); } + // Infer the internal `AggregateTableMode` based on `AggregateExec`'s mode + // + // TODO(simplification): `AggregateMode` seems bloated for aggregate hash + // table semantics. Consider remove `AggregateMode` and only use `AggregateTableMode` + // after the refactor has finished. + // + // Issue: + let mode = match agg.mode { + AggregateMode::Partial => AggregateTableMode::Partial, + AggregateMode::Final | AggregateMode::FinalPartitioned => { + AggregateTableMode::Final + } + AggregateMode::PartialReduce => AggregateTableMode::PartialReduce, + AggregateMode::Single | AggregateMode::SinglePartitioned => { + AggregateTableMode::Single + } + }; + let input_schema = agg.input().schema(); + let aggregate_mode = match mode { + AggregateTableMode::Partial => AggregateMode::Partial, + AggregateTableMode::Final => AggregateMode::Final, + AggregateTableMode::PartialReduce => AggregateMode::PartialReduce, + AggregateTableMode::Single => AggregateMode::Single, + }; let aggregate_arguments = aggregate_expressions( &agg.aggr_expr, - &agg.mode, + &aggregate_mode, agg.group_by.num_group_exprs(), )?; + + // Filters apply only when the table consumes raw input rows. Final and + // partial-reduce modes consume partial states, so their filters are not + // applicable. + let filters = match mode { + AggregateTableMode::Partial | AggregateTableMode::Single => { + agg.filter_expr.iter().cloned().collect() + } + AggregateTableMode::Final | AggregateTableMode::PartialReduce => { + vec![None; agg.aggr_expr.len()] + } + }; + let accumulators: Vec<_> = agg .aggr_expr .iter() @@ -129,6 +192,7 @@ impl AggregateHashTable { Ok(Self { group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + mode, input_schema, output_schema, batch_size, @@ -138,10 +202,13 @@ impl AggregateHashTable { batch_group_indices: Default::default(), accumulators, }), - _mode: PhantomData, }) } + pub(in crate::aggregates) fn mode(&self) -> AggregateTableMode { + self.mode + } + /// See comments in [`EvaluatedAggregateBatch`] pub(super) fn evaluate_batch( &self, @@ -196,6 +263,46 @@ impl AggregateHashTable { self.state.building().group_values.len() } + pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool { + self.state + .building() + .accumulators + .iter() + .all(|acc| acc.supports_convert_to_state()) + } + + /// In skip-partial-aggregation optimization, when a decision has been made + /// to skip partial stage, build a table only for converting raw input rows + /// into partial aggregate state rows. + pub(in crate::aggregates) fn partial_skip_table( + &self, + ) -> Result { + let state = self.state.building(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let accumulators = state + .accumulators + .iter() + .map(HashAggregateAccumulator::empty_like) + .collect::>>()?; + + Ok(PartialSkipHashTable { + table: AggregateHashTable { + group_by_metrics: self.group_by_metrics.clone(), + mode: AggregateTableMode::Partial, + input_schema: Arc::clone(&self.input_schema), + output_schema: Arc::clone(&self.output_schema), + batch_size: self.batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&state.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + }, + }) + } + pub(in crate::aggregates) fn is_building(&self) -> bool { matches!(self.state, AggregateHashTableState::Building(_)) } @@ -214,6 +321,239 @@ impl AggregateHashTable { state.batch_group_indices = Vec::new(); self.state = AggregateHashTableState::Outputting(state); } + + /// Aggregates one input batch according to this table's input semantics. + /// + /// `Partial` and `Single` update accumulator state from raw input rows. + /// `Final` and `PartialReduce` merge accumulator state emitted by an + /// earlier aggregate stage. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + let mode = self.mode; + 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()) + { + match mode { + AggregateTableMode::Partial | AggregateTableMode::Single => { + acc.update_batch(values, group_indices, total_num_groups)? + } + AggregateTableMode::Final | AggregateTableMode::PartialReduce => { + acc.merge_batch(values, group_indices, total_num_groups)? + } + } + } + } + drop(timer); + + Ok(()) + } + + /// Emits the next output batch according to this table's output semantics. + /// + /// `Partial` and `PartialReduce` emit accumulator states for downstream + /// aggregate stages. `Final` and `Single` evaluate final aggregate values + /// for the query result. + 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_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_output( + &self, + mut state: AggregateHashTableBuffer, + output_schema: SchemaRef, + ) -> Result { + // Final evaluation and partial state emission both consume accumulator + // state. Materialize all groups once, then slice 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() { + match self.mode { + AggregateTableMode::Partial | AggregateTableMode::PartialReduce => { + output.extend(acc.state(emit_to)?) + } + AggregateTableMode::Final | AggregateTableMode::Single => { + 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 + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + if matches!( + self.mode, + AggregateTableMode::Partial | AggregateTableMode::Single + ) { + self.init_empty_grouping_sets()?; + } + self.start_outputting(); + Ok(()) + } + + /// Creates the required empty grouping-set rows when the input is empty. + /// + /// For example, this query must still produce one grand-total group even if + /// `t` has no rows: + /// + /// ```sql + /// SELECT COUNT(v) + /// FROM t + /// GROUP BY GROUPING SETS (()); + /// ``` + /// + /// The synthetic row is filtered out before accumulator update so aggregates + /// see the same state they would see for an empty input, rather than a real + /// null-valued row. + pub(super) fn init_empty_grouping_sets(&mut self) -> Result<()> { + let state = self.state.building_mut(); + if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + return Ok(()); + } + + let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let n_expr = state.group_by.expr().len(); + let mut any_interned = false; + + for group in state.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let ordinal = *entry; + *entry += 1; + ordinal + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + let mut cols: Vec = group_schema + .fields() + .iter() + .take(n_expr) + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + state + .group_values + .intern(&cols, &mut state.batch_group_indices)?; + any_interned = true; + } + + if any_interned { + let total_groups = state.group_values.len(); + let false_filter = BooleanArray::from(vec![false]); + for acc in state.accumulators.iter_mut() { + let null_args = acc.null_arguments(&self.input_schema)?; + let values = EvaluatedAccumulatorArgs { + arguments: null_args, + filter: Some(Arc::new(false_filter.clone())), + }; + acc.update_batch(&values, &[0], total_groups)?; + } + } + + Ok(()) + } +} + +/// Hash table used only for converting raw input rows directly into partial +/// aggregate state rows after partial aggregation has been skipped. +pub(in crate::aggregates) struct PartialSkipHashTable { + pub(super) table: AggregateHashTable, +} + +impl PartialSkipHashTable { + pub(in crate::aggregates) fn convert_batch_to_state( + &mut self, + batch: &RecordBatch, + ) -> Result { + let evaluated_batch = self.table.evaluate_batch(batch)?; + + assert_eq_or_internal_err!( + evaluated_batch.grouping_set_args.len(), + 1, + "group_values expected to have single element" + ); + let mut output = evaluated_batch + .grouping_set_args + .into_iter() + .next() + .unwrap_or_default(); + + let state = self.table.state.building_mut(); + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + output.extend(acc.convert_to_state(values)?); + } + + Ok(RecordBatch::try_new( + Arc::clone(&self.table.output_schema), + output, + )?) + } } /// State and argument information for a single Aggregate 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 deleted file mode 100644 index 568b866b10517..0000000000000 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ /dev/null @@ -1,154 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// 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 crate::aggregates::AggregateExec; - -use super::common::{ - AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker, - MaterializedAggregateOutput, -}; - -/// Implementation specific to final aggregation, where the table stores partial -/// aggregate states and the input rows are also partial states. -/// -/// Example: `AVG(x) GROUP BY k` -/// -/// - Aggregate table stores: `k, sum(x), count(x)` -/// - Input rows: `k, sum(x), count(x)` -impl AggregateHashTable { - pub(in crate::aggregates) fn new( - agg: &AggregateExec, - partition: usize, - output_schema: SchemaRef, - batch_size: usize, - ) -> Result { - Self::new_with_filters( - agg, - partition, - output_schema, - batch_size, - vec![None; agg.aggr_expr.len()], - ) - } - - /// Emits the next batch of aggregated group keys and final aggregate values. - /// - /// The output batch size is determined by `self.batch_size`. - /// - /// Returns `Some(batch)` for each emitted batch, `None` when output is - /// exhausted, and an internal error if polled in the `Building` state. - 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() { - 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 - } - - 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(()) - } - - pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { - self.start_outputting(); - Ok(()) - } -} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs index 0d2495a1b556c..ed52881e60fe2 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -17,14 +17,11 @@ mod common; mod common_ordered; -mod final_table; mod ordered_final_table; mod ordered_partial_table; -mod partial_reduce_table; -mod partial_table; pub(super) use common::{ - AggregateHashTable, FinalMarker, PartialMarker, PartialReduceMarker, - PartialSkipMarker, + AggregateHashTable, AggregateTableMode, FinalMarker, PartialMarker, + PartialSkipHashTable, }; pub(super) use common_ordered::OrderedAggregateTable; 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 deleted file mode 100644 index 4d94c559436fb..0000000000000 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ /dev/null @@ -1,149 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// 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 crate::aggregates::AggregateExec; - -use super::common::{ - AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, - MaterializedAggregateOutput, PartialReduceMarker, -}; - -/// Methods specific to the aggregate hash table used in the partial-reduce stage. -impl AggregateHashTable { - pub(in crate::aggregates) fn new( - agg: &AggregateExec, - partition: usize, - output_schema: SchemaRef, - batch_size: usize, - ) -> Result { - Self::new_with_filters( - agg, - partition, - output_schema, - batch_size, - vec![None; agg.aggr_expr.len()], - ) - } - - /// Emits the next batch of aggregated group keys and aggregate states. - /// - /// The output batch size is determined by `self.batch_size`. - /// - /// Returns `Some(batch)` for each emitted batch, `None` when output is - /// exhausted, and an internal error if polled in the `Building` state. - 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 - } - - 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(()) - } - - pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { - self.start_outputting(); - Ok(()) - } -} 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 deleted file mode 100644 index f11eef8c14277..0000000000000 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ /dev/null @@ -1,298 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::marker::PhantomData; -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 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, MaterializedAggregateOutput, - PartialMarker, PartialSkipMarker, -}; - -/// Implementation specific to partial aggregation, where the table stores -/// partial aggregate states and the input rows are raw rows. -/// -/// Example: `AVG(x) GROUP BY k` -/// -/// - Aggregate table stores: `k, sum(x), count(x)` -/// - Input rows: `k, x` -impl AggregateHashTable { - pub(in crate::aggregates) fn new( - agg: &AggregateExec, - partition: usize, - output_schema: SchemaRef, - batch_size: usize, - ) -> Result { - Self::new_with_filters( - agg, - partition, - output_schema, - batch_size, - agg.filter_expr.iter().cloned().collect(), - ) - } - - /// Emits the next batch of aggregated group keys and aggregate states. - /// - /// The output batch size is determined by `self.batch_size`. - /// - /// Returns `Some(batch)` for each emitted batch, `None` when output is - /// exhausted, and an internal error if polled in the `Building` state. - 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() { - 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 - } - - pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool { - self.state - .building() - .accumulators - .iter() - .all(|acc| acc.supports_convert_to_state()) - } - - /// In skip-partial-aggregation optimization, when a decision has been made to skip - /// partial stage, build a typed hash table only for aggregation state conversion - /// row-by-row. - pub(in crate::aggregates) fn partial_skip_table( - &self, - ) -> Result> { - let state = self.state.building(); - let group_schema = state.group_by.group_schema(&self.input_schema)?; - let group_values = new_group_values(group_schema, &GroupOrdering::None)?; - let accumulators = state - .accumulators - .iter() - .map(HashAggregateAccumulator::empty_like) - .collect::>>()?; - - Ok(AggregateHashTable { - group_by_metrics: self.group_by_metrics.clone(), - input_schema: Arc::clone(&self.input_schema), - output_schema: Arc::clone(&self.output_schema), - batch_size: self.batch_size, - state: AggregateHashTableState::Building(AggregateHashTableBuffer { - group_by: Arc::clone(&state.group_by), - group_values, - batch_group_indices: Default::default(), - accumulators, - }), - _mode: PhantomData, - }) - } - - 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(()) - } - - pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { - self.init_empty_grouping_sets()?; - self.start_outputting(); - Ok(()) - } - - /// Creates the required empty grouping-set rows when the input is empty. - /// - /// For example, this query must still produce one grand-total group even if - /// `t` has no rows: - /// - /// ```sql - /// SELECT COUNT(v) - /// FROM t - /// GROUP BY GROUPING SETS (()); - /// ``` - /// - /// The synthetic row is filtered out before accumulator update so aggregates - /// see the same state they would see for an empty input, rather than a real - /// null-valued row. - fn init_empty_grouping_sets(&mut self) -> Result<()> { - let state = self.state.building_mut(); - if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { - return Ok(()); - } - - let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); - let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); - let group_schema = state.group_by.group_schema(&self.input_schema)?; - let n_expr = state.group_by.expr().len(); - let mut any_interned = false; - - for group in state.group_by.groups() { - let ordinal = { - let entry = ordinals.entry(group.as_slice()).or_insert(0); - let ordinal = *entry; - *entry += 1; - ordinal - }; - - if !group.iter().all(|&is_null| is_null) { - continue; - } - - let mut cols: Vec = group_schema - .fields() - .iter() - .take(n_expr) - .map(|field| new_null_array(field.data_type(), 1)) - .collect(); - cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); - - state - .group_values - .intern(&cols, &mut state.batch_group_indices)?; - any_interned = true; - } - - if any_interned { - let total_groups = state.group_values.len(); - let false_filter = BooleanArray::from(vec![false]); - for acc in state.accumulators.iter_mut() { - let null_args = acc.null_arguments(&self.input_schema)?; - let values = EvaluatedAccumulatorArgs { - arguments: null_args, - filter: Some(Arc::new(false_filter.clone())), - }; - acc.update_batch(&values, &[0], total_groups)?; - } - } - - Ok(()) - } -} - -impl AggregateHashTable { - pub(in crate::aggregates) fn convert_batch_to_state( - &mut self, - batch: &RecordBatch, - ) -> Result { - let evaluated_batch = self.evaluate_batch(batch)?; - - assert_eq_or_internal_err!( - evaluated_batch.grouping_set_args.len(), - 1, - "group_values expected to have single element" - ); - let mut output = evaluated_batch - .grouping_set_args - .into_iter() - .next() - .unwrap_or_default(); - - let state = self.state.building_mut(); - for (acc, values) in state - .accumulators - .iter_mut() - .zip(evaluated_batch.accumulator_args.iter()) - { - output.extend(acc.convert_to_state(values)?); - } - - Ok(RecordBatch::try_new( - Arc::clone(&self.output_schema), - output, - )?) - } -} diff --git a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs index 4c8756c0e865c..99c10c1d727e1 100644 --- a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs +++ b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs @@ -31,14 +31,14 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{Result, assert_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{ - AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, + AggregateHashTable, AggregateTableMode, PartialSkipHashTable, }; use super::skip_partial::SkipAggregationProbe; use crate::metrics::{ @@ -140,18 +140,18 @@ pub(crate) struct PartialHashAggregateStream { /// States for partial hash aggregation processing. enum PartialHashAggregateState { ReadingInput { - hash_table: AggregateHashTable, + hash_table: AggregateHashTable, }, ProducingOutput { - hash_table: AggregateHashTable, + hash_table: AggregateHashTable, /// If `None`, partial skip was never triggered and this state will /// finish in `Done`. If `Some`, partial skip has triggered and the /// stream will move to `SkippingAggregation` after these accumulated /// groups are emitted. - skip_hash_table: Option>, + skip_hash_table: Option, }, SkippingAggregation { - hash_table: AggregateHashTable, + hash_table: PartialSkipHashTable, }, Done, } @@ -163,7 +163,7 @@ type PartialHashAggregateStateTransition = ControlFlow< >; impl PartialHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { + fn hash_table(&self) -> &AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table, .. } => hash_table, @@ -173,7 +173,7 @@ impl PartialHashAggregateState { } } - fn hash_table_mut(&mut self) -> &mut AggregateHashTable { + fn hash_table_mut(&mut self) -> &mut AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table, .. } => hash_table, @@ -213,12 +213,8 @@ pub(crate) struct FinalHashAggregateStream { // The typestate pattern is used in case the inner logic becomes more complex in // the future. enum FinalHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - }, - ProducingOutput { - hash_table: AggregateHashTable, - }, + ReadingInput { hash_table: AggregateHashTable }, + ProducingOutput { hash_table: AggregateHashTable }, Done, } @@ -229,7 +225,7 @@ type FinalHashAggregateStateTransition = ControlFlow< >; impl FinalHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { + fn hash_table(&self) -> &AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table @@ -238,7 +234,7 @@ impl FinalHashAggregateState { } } - fn hash_table_mut(&mut self) -> &mut AggregateHashTable { + fn hash_table_mut(&mut self) -> &mut AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table @@ -247,7 +243,7 @@ impl FinalHashAggregateState { } } - fn into_hash_table(self) -> AggregateHashTable { + fn into_hash_table(self) -> AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table @@ -287,12 +283,12 @@ impl PartialHashAggregateStream { .with_type(metrics::MetricType::Summary) .ratio_metrics("reduction_factor", partition); - let hash_table = AggregateHashTable::::new( - agg, - partition, - Arc::clone(&schema), - batch_size, - )?; + let hash_table = + AggregateHashTable::new(agg, partition, Arc::clone(&schema), batch_size)?; + assert_or_internal_err!( + hash_table.mode() == AggregateTableMode::Partial, + "PartialHashAggregateStream expected partial aggregate hash table" + ); let can_skip_aggregation = agg.group_by.is_single() && hash_table.can_skip_aggregation(); let skip_aggregation_probe = if can_skip_aggregation { @@ -334,10 +330,7 @@ impl PartialHashAggregateStream { } /// See comments in [`Self::group_values_soft_limit`] for details. - fn hit_soft_group_limit( - &self, - hash_table: &AggregateHashTable, - ) -> bool { + fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { self.group_values_soft_limit .is_some_and(|limit| limit <= hash_table.building_group_count()) } @@ -359,7 +352,7 @@ impl PartialHashAggregateStream { fn start_output( &mut self, - hash_table: &mut AggregateHashTable, + hash_table: &mut AggregateHashTable, close_input: bool, ) -> Result<()> { if close_input { @@ -762,12 +755,12 @@ impl FinalHashAggregateStream { // Preserve the existing aggregate metric surface for this plan node. let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); - let hash_table = AggregateHashTable::::new( - agg, - partition, - Arc::clone(&schema), - batch_size, - )?; + let hash_table = + AggregateHashTable::new(agg, partition, Arc::clone(&schema), batch_size)?; + assert_or_internal_err!( + hash_table.mode() == AggregateTableMode::Final, + "FinalHashAggregateStream expected final aggregate hash table" + ); let reservation = MemoryConsumer::new(format!("FinalHashAggregateStream[{partition}]")) @@ -784,15 +777,12 @@ impl FinalHashAggregateStream { } /// See comments in [`Self::group_values_soft_limit`] for details. - fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { + fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { self.group_values_soft_limit .is_some_and(|limit| limit <= hash_table.building_group_count()) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { + fn start_output(&mut self, hash_table: &mut AggregateHashTable) -> Result<()> { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); hash_table.start_output() diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 1a4980c89851a..14c4496dbc8dd 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -28,13 +28,13 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{Result, assert_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; -use super::aggregate_hash_table::{AggregateHashTable, PartialReduceMarker}; +use super::aggregate_hash_table::{AggregateHashTable, AggregateTableMode}; use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -90,12 +90,8 @@ pub(crate) struct PartialReduceHashAggregateStream { // The typestate pattern mirrors the final stream and keeps the input/output // semantics explicit for this mode. enum PartialReduceHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - }, - ProducingOutput { - hash_table: AggregateHashTable, - }, + ReadingInput { hash_table: AggregateHashTable }, + ProducingOutput { hash_table: AggregateHashTable }, Done, } @@ -109,7 +105,7 @@ type PartialReduceHashAggregateStateTransition = ControlFlow< >; impl PartialReduceHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { + fn hash_table(&self) -> &AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table @@ -118,7 +114,7 @@ impl PartialReduceHashAggregateState { } } - fn hash_table_mut(&mut self) -> &mut AggregateHashTable { + fn hash_table_mut(&mut self) -> &mut AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table @@ -127,7 +123,7 @@ impl PartialReduceHashAggregateState { } } - fn into_hash_table(self) -> AggregateHashTable { + fn into_hash_table(self) -> AggregateHashTable { match self { Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { hash_table @@ -164,12 +160,12 @@ impl PartialReduceHashAggregateStream { // Preserve the existing aggregate metric surface for this plan node. let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); - let hash_table = AggregateHashTable::::new( - agg, - partition, - Arc::clone(&schema), - batch_size, - )?; + let hash_table = + AggregateHashTable::new(agg, partition, Arc::clone(&schema), batch_size)?; + assert_or_internal_err!( + hash_table.mode() == AggregateTableMode::PartialReduce, + "PartialReduceHashAggregateStream expected partial-reduce aggregate hash table" + ); let reservation = MemoryConsumer::new(format!("PartialReduceHashAggregateStream[{partition}]")) @@ -184,10 +180,7 @@ impl PartialReduceHashAggregateStream { }) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { + fn start_output(&mut self, hash_table: &mut AggregateHashTable) -> Result<()> { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); hash_table.start_output() From 77532c62a4a377489ec04aad93fef5f21e6f00ac Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Sat, 4 Jul 2026 13:18:02 +0800 Subject: [PATCH 2/3] fix wrong link --- .../physical-plan/src/aggregates/aggregate_hash_table/common.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 a90a67f6fd358..45fddc6928556 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -134,7 +134,7 @@ impl AggregateHashTable { // table semantics. Consider remove `AggregateMode` and only use `AggregateTableMode` // after the refactor has finished. // - // Issue: + // Issue: let mode = match agg.mode { AggregateMode::Partial => AggregateTableMode::Partial, AggregateMode::Final | AggregateMode::FinalPartitioned => { From 6871dc74b7b0314d7b8efc9b55582ca597212284 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Sat, 4 Jul 2026 13:24:21 +0800 Subject: [PATCH 3/3] review: swap inner/outer loop for mode --- .../aggregates/aggregate_hash_table/common.rs | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 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 45fddc6928556..0d2616548ef4b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -99,7 +99,7 @@ pub(in crate::aggregates) struct AggregateHashTable { /// Grouping and accumulator-specific timing metrics. pub(super) group_by_metrics: GroupByMetrics, - /// Semantic behavior for this table. + /// Semantic behavior for this table. See struct comments for details. pub(super) mode: AggregateTableMode, /// Raw input schema, used to evaluate expressions and synthesize empty @@ -322,11 +322,9 @@ impl AggregateHashTable { self.state = AggregateHashTableState::Outputting(state); } - /// Aggregates one input batch according to this table's input semantics. + /// Aggregates one input batch. The behavior depends on the aggregate mode. /// - /// `Partial` and `Single` update accumulator state from raw input rows. - /// `Final` and `PartialReduce` merge accumulator state emitted by an - /// earlier aggregate stage. + /// See comments at [`AggregateHashTable`] for details. pub(in crate::aggregates) fn aggregate_batch( &mut self, batch: &RecordBatch, @@ -343,16 +341,22 @@ impl AggregateHashTable { 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()) - { - match mode { - AggregateTableMode::Partial | AggregateTableMode::Single => { + match mode { + AggregateTableMode::Partial | AggregateTableMode::Single => { + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { acc.update_batch(values, group_indices, total_num_groups)? } - AggregateTableMode::Final | AggregateTableMode::PartialReduce => { + } + AggregateTableMode::Final | AggregateTableMode::PartialReduce => { + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { acc.merge_batch(values, group_indices, total_num_groups)? } } @@ -393,7 +397,9 @@ impl AggregateHashTable { } } } - + /// Materialize the aggregated output. The behavior depends on the aggregate mode. + /// + /// See comments at [`AggregateHashTable`] for details. fn materialize_output( &self, mut state: AggregateHashTableBuffer, @@ -405,12 +411,14 @@ impl AggregateHashTable { 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() { - match self.mode { - AggregateTableMode::Partial | AggregateTableMode::PartialReduce => { + match self.mode { + AggregateTableMode::Partial | AggregateTableMode::PartialReduce => { + for acc in state.accumulators.iter_mut() { output.extend(acc.state(emit_to)?) } - AggregateTableMode::Final | AggregateTableMode::Single => { + } + AggregateTableMode::Final | AggregateTableMode::Single => { + for acc in state.accumulators.iter_mut() { output.push(acc.evaluate(emit_to)?) } }