Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
2010YOUY01 merged 4 commits into from
Jul 9, 2026
+139 −252
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
90d6617
refactor(hash-aggr): Simplify aggregate hash table with tempated func…
2010YOUY01 Jul 5, 2026
a1c9b6a
more comments
2010YOUY01 Jul 5, 2026
32691af
review: exract AggregateBatchFn to a separate type
2010YOUY01 Jul 8, 2026
4445f34
review: exract MaterializeAccumulatorFn to a separate type
2010YOUY01 Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 121 additions & 0 deletions datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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`
Expand Down Expand Up @@ -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)
Expand Down
92 changes: 6 additions & 86 deletions datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<()> {
Expand Down
93 changes: 6 additions & 87 deletions datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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<()> {
Expand Down
Loading
Loading