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

Open
2010YOUY01 wants to merge 2 commits into
base: main
Choose a base branch
from
+140 −250
Open

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
2 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
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
100 changes: 100 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,106 @@ 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<F>(
&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: <https://github.com/apache/datafusion/issues/7065>
pub(super) fn next_output_batch_inner<F>(
&mut self,
mut materialize_accumulator_fn: F,
) -> Result<Option<RecordBatch>>
where
F: FnMut(&mut HashAggregateAccumulator, EmitTo, &mut Vec<ArrayRef>) -> 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)
Expand Down
98 changes: 13 additions & 85 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};

/// 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,24 @@ 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() {
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<RecordBatch> {
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<()> {
Expand Down
101 changes: 14 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, PartialReduceMarker};

/// Methods specific to the aggregate hash table used in the partial-reduce stage.
impl AggregateHashTable<PartialReduceMarker> {
Expand Down Expand Up @@ -55,91 +49,24 @@ 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(|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<()> {
Expand Down
Loading
Loading