You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/08 01:32:26 UTC
(datafusion) branch main updated: refactor: Split hash aggregation logic into separated streams (#22729)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6fdef65172 refactor: Split hash aggregation logic into separated streams (#22729)
6fdef65172 is described below
commit 6fdef651723817738c4b8cc9611dd7279008d168
Author: Yongting You <20...@gmail.com>
AuthorDate: Mon Jun 8 09:18:37 2026 +0800
refactor: Split hash aggregation logic into separated streams (#22729)
## 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
- 1/N of https://github.com/apache/datafusion/pull/22712
## 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 issues.
This PR split out partial and final aggregate strem from
`GroupsHashAggregateStream`
To fully migrate hash aggregation, we have to
- Port this optimization back
https://github.com/apache/datafusion/pull/11627
- Support spilling
I think they should be leave to follow up PRs
Todo in this PR:
- [x] Add a temporary configuration `enable_migration_aggregate` to turn
off this path
Since it should be a regression if the above features are not added, it
also helps if to prevent potential regressions from the migration of
other aggregate streams.
## 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.
-->
Split out the streams from `GroupsHashAggregateStream`
1. Partial stage of hash aggregation
2. Final stage of hash aggregation
## 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
3. 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.
-->
---
datafusion/common/src/config.rs | 11 +
.../aggregation_fuzzer/context_generator.rs | 14 +
.../src/aggregates/group_values/metrics.rs | 9 +-
.../physical-plan/src/aggregates/hash_aggregate.rs | 345 ++++++++++++
.../physical-plan/src/aggregates/hash_table.rs | 619 +++++++++++++++++++++
datafusion/physical-plan/src/aggregates/mod.rs | 213 ++++++-
.../sqllogictest/test_files/information_schema.slt | 2 +
docs/source/user-guide/configs.md | 1 +
8 files changed, 1208 insertions(+), 6 deletions(-)
diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs
index 4025157cef..b10761a5fe 100644
--- a/datafusion/common/src/config.rs
+++ b/datafusion/common/src/config.rs
@@ -596,6 +596,17 @@ config_namespace! {
/// the new schema verification step.
pub skip_physical_aggregate_schema_check: bool, default = false
+ /// Temporary switch for aggregate stream implementations that are being
+ /// migrated from `GroupedHashAggregateStream`.
+ ///
+ /// When set to true, DataFusion tries the migrated implementations when
+ /// their preconditions are satisfied. When set to false, grouped
+ /// aggregation falls back to `GroupedHashAggregateStream`. This option
+ /// will be removed after the migration is finished.
+ ///
+ /// See <https://github.com/apache/datafusion/issues/22710> for details.
+ pub enable_migration_aggregate: bool, default = false
+
/// Sets the compression codec used when spilling data to disk.
///
/// Since datafusion writes spill files using the Arrow IPC Stream format,
diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs
index fe31098622..3579c6af84 100644
--- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs
+++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs
@@ -41,6 +41,7 @@ use crate::fuzz_cases::aggregation_fuzzer::data_generator::Dataset;
/// - `batch_size`
/// - `target_partitions`
/// - `skip_partial parameters`
+/// - `enable_migration_aggregate`
/// - hint `sorted` or not
/// - `spilling` or not (TODO, I think a special `MemoryPool` may be needed
/// to support this)
@@ -96,11 +97,13 @@ impl SessionContextGenerator {
let batch_size = self.max_batch_size;
let target_partitions = 1;
let skip_partial_params = SkipPartialParams::ensure_not_trigger();
+ let enable_migration_aggregate = false;
let builder = GeneratedSessionContextBuilder {
batch_size,
target_partitions,
skip_partial_params,
+ enable_migration_aggregate,
sort_hint: false,
table_name: self.table_name.clone(),
table_provider: Arc::new(provider),
@@ -120,6 +123,7 @@ impl SessionContextGenerator {
// - `batch_size`, from range: [1, `total_rows_num`]
// - `target_partitions`, from range: [1, cpu_num]
// - `skip_partial`, trigger or not trigger currently for simplicity
+ // - `enable_migration_aggregate`, true or false
// - `sorted`, if found a sorted dataset, will or will not push down this information
// - `spilling`(TODO)
let batch_size = rng.random_range(1..=self.max_batch_size);
@@ -131,6 +135,8 @@ impl SessionContextGenerator {
let skip_partial_params =
self.candidate_skip_partial_params[skip_partial_params_idx];
+ let enable_migration_aggregate = rng.random_bool(0.5);
+
let (provider, sort_hint) =
if rng.random_bool(0.5) && !self.dataset.sort_keys.is_empty() {
// Sort keys exist and random to push down
@@ -150,6 +156,7 @@ impl SessionContextGenerator {
target_partitions,
sort_hint,
skip_partial_params,
+ enable_migration_aggregate,
table_name: self.table_name.clone(),
table_provider: Arc::new(provider),
};
@@ -173,6 +180,7 @@ struct GeneratedSessionContextBuilder {
target_partitions: usize,
sort_hint: bool,
skip_partial_params: SkipPartialParams,
+ enable_migration_aggregate: bool,
table_name: String,
table_provider: Arc<dyn TableProvider>,
}
@@ -197,6 +205,10 @@ impl GeneratedSessionContextBuilder {
"datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
&ScalarValue::Float64(Some(self.skip_partial_params.ratio_threshold)),
);
+ session_config = session_config.set_bool(
+ "datafusion.execution.enable_migration_aggregate",
+ self.enable_migration_aggregate,
+ );
let ctx = SessionContext::new_with_config(session_config);
ctx.register_table(self.table_name, self.table_provider)?;
@@ -206,6 +218,7 @@ impl GeneratedSessionContextBuilder {
target_partitions: self.target_partitions,
sort_hint: self.sort_hint,
skip_partial_params: self.skip_partial_params,
+ enable_migration_aggregate: self.enable_migration_aggregate,
};
Ok(SessionContextWithParams { ctx, params })
@@ -220,6 +233,7 @@ pub struct SessionContextParams {
target_partitions: usize,
sort_hint: bool,
skip_partial_params: SkipPartialParams,
+ enable_migration_aggregate: bool,
}
/// Partial skipping parameters
diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
index b6c32204e8..a0934b976e 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
@@ -59,6 +59,7 @@ mod tests {
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use datafusion_execution::TaskContext;
+ use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_functions_aggregate::count::count_udaf;
use datafusion_functions_aggregate::sum::sum_udaf;
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -135,7 +136,13 @@ mod tests {
schema,
)?);
- let task_ctx = Arc::new(TaskContext::default());
+ // This test is for `GroupByMetrics`, which are maintained by
+ // `GroupedHashAggregateStream`. Use a finite memory pool so the partial
+ // aggregate does not take the initial-partial stream path.
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(10 * 1024 * 1024, 1.0)
+ .build_arc()?;
+ let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime));
let _result =
collect(Arc::clone(&aggregate_exec) as _, Arc::clone(&task_ctx)).await?;
diff --git a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs
new file mode 100644
index 0000000000..f25299631a
--- /dev/null
+++ b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs
@@ -0,0 +1,345 @@
+// 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.
+
+//! 2-stage hash aggregation stream implementation.
+//!
+//! See comments in [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`]
+//! for details.
+//!
+//! Note these streams are an incremental migration of the existing
+//! [`crate::aggregates::row_hash::GroupedHashAggregateStream`].
+//!
+//! See issue for details: <https://github.com/apache/datafusion/issues/22710>
+
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use datafusion_common::Result;
+use datafusion_execution::TaskContext;
+use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use futures::stream::{Stream, StreamExt};
+
+use super::AggregateExec;
+use super::hash_table::{AggregateHashTable, Final, Partial};
+use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput, SpillMetrics};
+use crate::stream::EmptyRecordBatchStream;
+use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics};
+
+/// Hash aggregation uses a 2-stage (partial and final) hash aggregation, this stream
+/// is for the partial stage.
+///
+/// # Example
+///
+/// select k, avg(v) from t group by k;
+///
+/// ## Plan
+/// AggregateExec(stage=final)
+/// -- RepartitionExec(hash(k))
+/// ---- AggregateExec(stage=partial)
+///
+/// ## Partial Stage Behavior
+/// Input: raw rows
+/// Output: partial states for all groups (e.g. for avg(x), it's sum(x), count(x))
+///
+/// ## Final Stage Behavior
+/// Input: partial states
+/// Output: results for all groups (e.g. for avg(x), it's avg(x) calculated from the state)
+pub(crate) struct PartialHashAggregateStream {
+ /// Output schema: group columns followed by partial aggregate state columns.
+ schema: SchemaRef,
+
+ /// Input batches containing raw rows, not partial aggregate state.
+ input: SendableRecordBatchStream,
+
+ /// Hash table state for this aggregate stream.
+ hash_table: AggregateHashTable<Partial>,
+
+ /// Memory reservation for group keys and accumulators.
+ reservation: MemoryReservation,
+
+ /// Execution metrics shared with the aggregate plan node.
+ baseline_metrics: BaselineMetrics,
+
+ /// Tracks partial aggregation row reduction, matching `GroupedHashAggregateStream`.
+ reduction_factor: metrics::RatioMetrics,
+}
+
+/// Hash aggregation uses a 2-stage (partial and final) hash aggregation, this stream
+/// is for the final stage.
+///
+/// See [`PartialHashAggregateStream`] for details.
+pub(crate) struct FinalHashAggregateStream {
+ /// Output schema: group columns followed by final aggregate value columns.
+ schema: SchemaRef,
+
+ /// Input batches containing partial aggregate state rows.
+ input: SendableRecordBatchStream,
+
+ /// Hash table state for this aggregate stream.
+ hash_table: AggregateHashTable<Final>,
+
+ /// Execution metrics shared with the aggregate plan node.
+ baseline_metrics: BaselineMetrics,
+
+ /// Memory reservation for group keys and accumulators.
+ reservation: MemoryReservation,
+}
+
+impl PartialHashAggregateStream {
+ pub fn new(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ ) -> Result<Self> {
+ debug_assert_eq!(agg.mode, super::AggregateMode::Partial);
+ debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear);
+
+ let schema = Arc::clone(&agg.schema);
+ let input = agg.input.execute(partition, Arc::clone(context))?;
+ let batch_size = context.session_config().batch_size();
+ let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
+
+ // Preserve the existing aggregate metric surface for this plan node.
+ let _spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+ let reduction_factor = MetricBuilder::new(&agg.metrics)
+ .with_type(metrics::MetricType::Summary)
+ .ratio_metrics("reduction_factor", partition);
+
+ let hash_table = AggregateHashTable::<Partial>::new(
+ agg,
+ partition,
+ Arc::clone(&schema),
+ batch_size,
+ )?;
+
+ let reservation =
+ MemoryConsumer::new(format!("PartialHashAggregateStream[{partition}]"))
+ .register(context.memory_pool());
+
+ Ok(Self {
+ schema,
+ input,
+ hash_table,
+ baseline_metrics,
+ reservation,
+ reduction_factor,
+ })
+ }
+}
+
+impl Stream for PartialHashAggregateStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+
+ loop {
+ if self.hash_table.is_done() {
+ let _ = self.reservation.try_resize(0);
+ return Poll::Ready(None);
+ } else if self.hash_table.is_building() {
+ match self.input.poll_next_unpin(cx) {
+ Poll::Pending => return Poll::Pending,
+ Poll::Ready(Some(Ok(batch))) => {
+ let timer = elapsed_compute.timer();
+ self.reduction_factor.add_total(batch.num_rows());
+ let result = self.hash_table.aggregate_batch(&batch);
+ timer.done();
+
+ if let Err(e) = result {
+ return Poll::Ready(Some(Err(e)));
+ }
+
+ // TODO: impl memory-limited aggr, when OOM directly send
+ // partial state to final aggregate stage
+ if let Err(e) =
+ self.reservation.try_resize(self.hash_table.memory_size())
+ {
+ return Poll::Ready(Some(Err(e)));
+ }
+ }
+ Poll::Ready(Some(Err(e))) => {
+ return Poll::Ready(Some(Err(e)));
+ }
+ Poll::Ready(None) => {
+ let input_schema = self.input.schema();
+ self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
+
+ let timer = elapsed_compute.timer();
+ let result = self.hash_table.start_output();
+ timer.done();
+
+ if let Err(e) = result {
+ return Poll::Ready(Some(Err(e)));
+ }
+ }
+ }
+ } else {
+ let timer = elapsed_compute.timer();
+ let result = self.hash_table.next_output_batch();
+ timer.done();
+
+ match result {
+ Ok(Some(batch)) => {
+ let _ =
+ self.reservation.try_resize(self.hash_table.memory_size());
+ self.reduction_factor.add_part(batch.num_rows());
+ debug_assert!(batch.num_rows() > 0);
+ return Poll::Ready(Some(Ok(
+ batch.record_output(&self.baseline_metrics)
+ )));
+ }
+ Ok(None) => {
+ let _ = self.reservation.try_resize(0);
+ return Poll::Ready(None);
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
+ }
+ }
+ }
+}
+
+impl RecordBatchStream for PartialHashAggregateStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl FinalHashAggregateStream {
+ pub fn new(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ ) -> Result<Self> {
+ debug_assert!(matches!(
+ agg.mode,
+ super::AggregateMode::Final | super::AggregateMode::FinalPartitioned
+ ));
+ debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear);
+
+ let schema = Arc::clone(&agg.schema);
+ let input = agg.input.execute(partition, Arc::clone(context))?;
+ let batch_size = context.session_config().batch_size();
+ let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
+
+ // Preserve the existing aggregate metric surface for this plan node.
+ let _spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+
+ let hash_table = AggregateHashTable::<Final>::new(
+ agg,
+ partition,
+ Arc::clone(&schema),
+ batch_size,
+ )?;
+
+ let reservation =
+ MemoryConsumer::new(format!("FinalHashAggregateStream[{partition}]"))
+ .register(context.memory_pool());
+
+ Ok(Self {
+ schema,
+ input,
+ hash_table,
+ baseline_metrics,
+ reservation,
+ })
+ }
+}
+
+impl Stream for FinalHashAggregateStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+
+ loop {
+ if self.hash_table.is_done() {
+ let _ = self.reservation.try_resize(0);
+ return Poll::Ready(None);
+ } else if self.hash_table.is_building() {
+ match self.input.poll_next_unpin(cx) {
+ Poll::Pending => return Poll::Pending,
+ Poll::Ready(Some(Ok(batch))) => {
+ let timer = elapsed_compute.timer();
+ let result = self.hash_table.aggregate_batch(&batch);
+ timer.done();
+
+ if let Err(e) = result {
+ return Poll::Ready(Some(Err(e)));
+ }
+
+ if let Err(e) =
+ self.reservation.try_resize(self.hash_table.memory_size())
+ {
+ return Poll::Ready(Some(Err(e)));
+ }
+ }
+ Poll::Ready(Some(Err(e))) => {
+ return Poll::Ready(Some(Err(e)));
+ }
+ Poll::Ready(None) => {
+ let input_schema = self.input.schema();
+ self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
+
+ let timer = elapsed_compute.timer();
+ let result = self.hash_table.start_output();
+ timer.done();
+
+ if let Err(e) = result {
+ return Poll::Ready(Some(Err(e)));
+ }
+ }
+ }
+ } else {
+ let timer = elapsed_compute.timer();
+ let result = self.hash_table.next_output_batch();
+ timer.done();
+
+ match result {
+ Ok(Some(batch)) => {
+ let _ =
+ self.reservation.try_resize(self.hash_table.memory_size());
+ debug_assert!(batch.num_rows() > 0);
+ return Poll::Ready(Some(Ok(
+ batch.record_output(&self.baseline_metrics)
+ )));
+ }
+ Ok(None) => {
+ let _ = self.reservation.try_resize(0);
+ return Poll::Ready(None);
+ }
+ Err(e) => return Poll::Ready(Some(Err(e))),
+ }
+ }
+ }
+ }
+}
+
+impl RecordBatchStream for FinalHashAggregateStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
diff --git a/datafusion/physical-plan/src/aggregates/hash_table.rs b/datafusion/physical-plan/src/aggregates/hash_table.rs
new file mode 100644
index 0000000000..278689d23f
--- /dev/null
+++ b/datafusion/physical-plan/src/aggregates/hash_table.rs
@@ -0,0 +1,619 @@
+// 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, AsArray, BooleanArray, new_null_array};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{Result, internal_err};
+use datafusion_execution::memory_pool::proxy::VecAllocExt;
+use datafusion_expr::{EmitTo, GroupsAccumulator};
+
+use super::group_values::{GroupByMetrics, GroupValues, new_group_values};
+use super::order::GroupOrdering;
+use super::row_hash::create_group_accumulator;
+use super::{
+ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
+ group_id_array, max_duplicate_ordinal,
+};
+use crate::PhysicalExpr;
+use crate::metrics::{MetricBuilder, MetricCategory};
+
+/// Marker for raw rows -> partial state aggregation.
+pub(super) struct Partial;
+/// Marker for partial state -> final value aggregation.
+pub(super) struct Final;
+
+/// Grouped hash table shared by the partial and final paths.
+///
+/// While building, it consumes input batches and updates group / accumulator
+/// state. While outputting, it incrementally output the materialized batches.
+///
+/// # Marker Type
+/// `AggrMode` selects the aggregate semantics.
+///
+/// e.g. `AggregateHashTable::<Partial>::new(...)` creates an aggregate hash table
+/// for the partial hash aggregate stage, the input schema is raw rows and output
+/// schema is intermediate states.
+///
+/// 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(super) struct AggregateHashTable<AggrMode> {
+ /// Grouping and accumulator-specific timing metrics.
+ group_by_metrics: GroupByMetrics,
+
+ /// Raw input schema, used to evaluate expressions and synthesize empty
+ /// grouping-set rows.
+ input_schema: SchemaRef,
+
+ /// Output schema: group columns followed by aggregate state or final values.
+ output_schema: SchemaRef,
+
+ /// Maximum rows per emitted output batch.
+ batch_size: usize,
+
+ /// Lifecycle-specific state: building stage / outputting stage
+ state: AggregateHashTableState,
+
+ _mode: PhantomData<AggrMode>,
+}
+
+struct HashAggregateAccumulator {
+ /// Arguments to pass to this accumulator.
+ ///
+ /// Example: `CORR(x, y)` stores two expressions here, while `SUM(x)` stores one.
+ arguments: Vec<Arc<dyn PhysicalExpr>>,
+
+ /// Optional `FILTER` expression for this accumulator.
+ ///
+ /// Example: `SUM(x) FILTER (WHERE x > 10)` stores the `x > 10` predicate.
+ filter: Option<Arc<dyn PhysicalExpr>>,
+
+ /// Accumulator state for all groups for one aggregate expression.
+ accumulator: Box<dyn GroupsAccumulator>,
+}
+
+struct EvaluatedHashAggregateAccumulator {
+ arguments: Vec<ArrayRef>,
+ filter: Option<ArrayRef>,
+}
+
+/// Evaluated all group by keys and accumulator args.
+///
+/// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function evaluates
+/// `k+1`, `v*v`
+struct EvaluatedAggregateBatch {
+ /// One entry per grouping set; each entry contains all evaluated group key
+ /// arrays for the current input batch.
+ grouping_set_args: Vec<Vec<ArrayRef>>,
+
+ /// Evaluated arguments and filters, one entry per aggregate expression.
+ accumulator_args: Vec<EvaluatedHashAggregateAccumulator>,
+}
+
+/// Hash table state while grouped aggregation is consuming input.
+///
+/// This owns the coupled state for:
+/// - evaluating group keys,
+/// - interning each distinct group,
+/// - mapping each input row to its group index,
+/// - evaluating aggregate inputs,
+/// - updating per-group accumulator state.
+struct BuildingHashTableState {
+ /// GROUP BY expressions evaluated for each input batch.
+ group_by: Arc<PhysicalGroupBy>,
+
+ /// Interned group keys. Accumulator state is stored separately by group index.
+ group_values: Box<dyn GroupValues>,
+
+ /// Group index for each row in the current input batch.
+ ///
+ /// Each value indexes into `group_values`, and the same index is used by every
+ /// accumulator to update that group's aggregate state.
+ batch_group_indices: Vec<usize>,
+
+ /// One item per aggregate expression.
+ ///
+ /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input
+ /// expressions, optional filter, and accumulator state for all groups.
+ accumulators: Vec<HashAggregateAccumulator>,
+}
+
+enum AggregateHashTableState {
+ Building(BuildingHashTableState),
+ Outputting {
+ output_batch: Option<RecordBatch>,
+ output_batch_offset: usize,
+ },
+ Done,
+}
+
+impl HashAggregateAccumulator {
+ fn new(
+ arguments: Vec<Arc<dyn PhysicalExpr>>,
+ filter: Option<Arc<dyn PhysicalExpr>>,
+ accumulator: Box<dyn GroupsAccumulator>,
+ ) -> Self {
+ Self {
+ arguments,
+ filter,
+ accumulator,
+ }
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<EvaluatedHashAggregateAccumulator> {
+ let arguments = self
+ .arguments
+ .iter()
+ .map(|expr| {
+ expr.evaluate(batch)
+ .and_then(|value| value.into_array(batch.num_rows()))
+ })
+ .collect::<Result<_>>()?;
+
+ let filter = self
+ .filter
+ .as_ref()
+ .map(|filter| {
+ filter
+ .evaluate(batch)
+ .and_then(|value| value.into_array(batch.num_rows()))
+ })
+ .transpose()?;
+
+ Ok(EvaluatedHashAggregateAccumulator { arguments, filter })
+ }
+
+ fn update_batch(
+ &mut self,
+ values: &EvaluatedHashAggregateAccumulator,
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> Result<()> {
+ let filter = values.filter.as_ref().map(|filter| filter.as_boolean());
+ self.accumulator.update_batch(
+ &values.arguments,
+ group_indices,
+ filter,
+ total_num_groups,
+ )
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &EvaluatedHashAggregateAccumulator,
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> Result<()> {
+ debug_assert!(values.filter.is_none());
+ self.accumulator.merge_batch(
+ &values.arguments,
+ group_indices,
+ None,
+ total_num_groups,
+ )
+ }
+
+ fn evaluate_final(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ self.accumulator.evaluate(emit_to)
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ self.accumulator.state(emit_to)
+ }
+
+ fn supports_convert_to_state(&self) -> bool {
+ self.accumulator.supports_convert_to_state()
+ }
+
+ fn null_arguments(&self, input_schema: &SchemaRef) -> Result<Vec<ArrayRef>> {
+ self.arguments
+ .iter()
+ .map(|expr| {
+ let data_type = expr.data_type(input_schema)?;
+ Ok(new_null_array(&data_type, 1))
+ })
+ .collect()
+ }
+}
+
+impl AggregateHashTableState {
+ fn building(&self) -> &BuildingHashTableState {
+ let Self::Building(state) = self else {
+ unreachable!("hash aggregate table is not building")
+ };
+ state
+ }
+
+ fn building_mut(&mut self) -> &mut BuildingHashTableState {
+ let Self::Building(state) = self else {
+ unreachable!("hash aggregate table is not building")
+ };
+ state
+ }
+}
+
+impl<Mode> AggregateHashTable<Mode> {
+ fn new_with_filters(
+ agg: &AggregateExec,
+ partition: usize,
+ output_schema: SchemaRef,
+ batch_size: usize,
+ filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
+ ) -> Result<Self> {
+ let input_schema = agg.input().schema();
+ let aggregate_arguments = aggregate_expressions(
+ &agg.aggr_expr,
+ &agg.mode,
+ agg.group_by.num_group_exprs(),
+ )?;
+ let accumulators: Vec<_> = agg
+ .aggr_expr
+ .iter()
+ .zip(aggregate_arguments)
+ .zip(filters)
+ .map(|((agg_expr, arguments), filter)| {
+ let accumulator = create_group_accumulator(agg_expr)?;
+ Ok(HashAggregateAccumulator::new(
+ arguments,
+ filter,
+ accumulator,
+ ))
+ })
+ .collect::<Result<_>>()?;
+
+ let group_schema = agg.group_by.group_schema(&input_schema)?;
+ let group_values = new_group_values(group_schema, &GroupOrdering::None)?;
+
+ Ok(Self {
+ group_by_metrics: GroupByMetrics::new(&agg.metrics, partition),
+ input_schema,
+ output_schema,
+ batch_size,
+ state: AggregateHashTableState::Building(BuildingHashTableState {
+ group_by: Arc::clone(&agg.group_by),
+ group_values,
+ batch_group_indices: Default::default(),
+ accumulators,
+ }),
+ _mode: PhantomData,
+ })
+ }
+
+ /// See comments in [`EvaluatedAggregateBatch`]
+ fn evaluate_batch(&self, batch: &RecordBatch) -> Result<EvaluatedAggregateBatch> {
+ let state = self.state.building();
+ let timer = self.group_by_metrics.time_calculating_group_ids.timer();
+ // outer vec: one per each grouping set
+ // inner vec: all group by exprs for the current grouping set
+ let grouping_set_args = evaluate_group_by(&state.group_by, batch)?;
+ drop(timer);
+
+ let timer = self.group_by_metrics.aggregate_arguments_time.timer();
+ // The evaluated args for each accumulator
+ let accumulator_args = self
+ .state
+ .building()
+ .accumulators
+ .iter()
+ .map(|acc| acc.evaluate(batch))
+ .collect::<Result<Vec<_>>>()?;
+ drop(timer);
+
+ Ok(EvaluatedAggregateBatch {
+ grouping_set_args,
+ accumulator_args,
+ })
+ }
+
+ pub(super) fn memory_size(&self) -> usize {
+ match &self.state {
+ AggregateHashTableState::Building(state) => {
+ let acc = state
+ .accumulators
+ .iter()
+ .map(|acc| acc.accumulator.size())
+ .sum::<usize>();
+
+ acc + state.group_values.size()
+ + state.batch_group_indices.allocated_size()
+ }
+ AggregateHashTableState::Outputting { output_batch, .. } => {
+ output_batch_memory_size(output_batch)
+ }
+ AggregateHashTableState::Done => 0,
+ }
+ }
+
+ pub(super) fn is_building(&self) -> bool {
+ matches!(self.state, AggregateHashTableState::Building(_))
+ }
+
+ pub(super) fn is_done(&self) -> bool {
+ matches!(self.state, AggregateHashTableState::Done)
+ }
+
+ fn set_output_batch(&mut self, output_batch: Option<RecordBatch>) {
+ self.state = AggregateHashTableState::Outputting {
+ output_batch,
+ output_batch_offset: 0,
+ };
+ }
+
+ pub(super) fn next_output_batch(&mut self) -> Result<Option<RecordBatch>> {
+ match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
+ AggregateHashTableState::Outputting {
+ output_batch,
+ mut output_batch_offset,
+ } => {
+ let Some(batch) = output_batch.as_ref() else {
+ return Ok(None);
+ };
+
+ let num_rows = batch.num_rows();
+ if output_batch_offset >= num_rows {
+ return Ok(None);
+ }
+
+ debug_assert!(self.batch_size > 0);
+ let output_len =
+ self.batch_size.max(1).min(num_rows - output_batch_offset);
+ let output = batch.slice(output_batch_offset, output_len);
+ output_batch_offset += output_len;
+
+ if output_batch_offset == num_rows {
+ self.state = AggregateHashTableState::Done;
+ } else {
+ self.state = AggregateHashTableState::Outputting {
+ output_batch,
+ output_batch_offset,
+ };
+ }
+
+ debug_assert!(output.num_rows() > 0);
+ debug_assert!(output.num_rows() <= self.batch_size.max(1));
+ Ok(Some(output))
+ }
+ _ => {
+ self.state = AggregateHashTableState::Done;
+ internal_err!("next_output_batch must be called in the outputting state")
+ }
+ }
+ }
+}
+
+impl AggregateHashTable<Partial> {
+ pub(super) fn new(
+ agg: &AggregateExec,
+ partition: usize,
+ output_schema: SchemaRef,
+ batch_size: usize,
+ ) -> Result<Self> {
+ let table = Self::new_with_filters(
+ agg,
+ partition,
+ output_schema,
+ batch_size,
+ agg.filter_expr.iter().cloned().collect(),
+ )?;
+
+ if table
+ .state
+ .building()
+ .accumulators
+ .iter()
+ .all(|acc| acc.supports_convert_to_state())
+ {
+ let _skipped_aggregation_rows = MetricBuilder::new(&agg.metrics)
+ .with_category(MetricCategory::Rows)
+ .counter("skipped_aggregation_rows", partition);
+ }
+
+ Ok(table)
+ }
+
+ pub(super) 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)?;
+ }
+ }
+ drop(timer);
+
+ Ok(())
+ }
+
+ pub(super) fn start_output(&mut self) -> Result<()> {
+ self.init_empty_grouping_sets()?;
+ let state = self.state.building_mut();
+
+ let output_batch = if state.group_values.is_empty() {
+ None
+ } else {
+ let timer = self.group_by_metrics.emitting_time.timer();
+ let mut output = state.group_values.emit(EmitTo::All)?;
+
+ for acc in state.accumulators.iter_mut() {
+ output.extend(acc.state(EmitTo::All)?);
+ }
+
+ let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?;
+ debug_assert!(batch.num_rows() > 0);
+ drop(timer);
+ Some(batch)
+ };
+
+ self.set_output_batch(output_batch);
+ 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<ArrayRef> = 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 = EvaluatedHashAggregateAccumulator {
+ arguments: null_args,
+ filter: Some(Arc::new(false_filter.clone())),
+ };
+ acc.update_batch(&values, &[0], total_groups)?;
+ }
+ }
+
+ Ok(())
+ }
+}
+
+impl AggregateHashTable<Final> {
+ pub(super) fn new(
+ agg: &AggregateExec,
+ partition: usize,
+ output_schema: SchemaRef,
+ batch_size: usize,
+ ) -> Result<Self> {
+ Self::new_with_filters(
+ agg,
+ partition,
+ output_schema,
+ batch_size,
+ vec![None; agg.aggr_expr.len()],
+ )
+ }
+
+ pub(super) 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(super) fn start_output(&mut self) -> Result<()> {
+ let state = self.state.building_mut();
+ let output_batch = if state.group_values.is_empty() {
+ None
+ } else {
+ let timer = self.group_by_metrics.emitting_time.timer();
+ let mut output = state.group_values.emit(EmitTo::All)?;
+
+ for acc in state.accumulators.iter_mut() {
+ output.push(acc.evaluate_final(EmitTo::All)?);
+ }
+
+ let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?;
+ debug_assert!(batch.num_rows() > 0);
+ drop(timer);
+ Some(batch)
+ };
+
+ self.set_output_batch(output_batch);
+ Ok(())
+ }
+}
+
+fn output_batch_memory_size(output_batch: &Option<RecordBatch>) -> usize {
+ output_batch
+ .as_ref()
+ .map(RecordBatch::get_array_memory_size)
+ .unwrap_or_default()
+}
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 5a2080990e..67327abea3 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -22,7 +22,9 @@ use std::sync::Arc;
use super::{DisplayAs, ExecutionPlanProperties, PlanProperties};
use crate::aggregates::{
- no_grouping::AggregateStream, row_hash::GroupedHashAggregateStream,
+ hash_aggregate::{FinalHashAggregateStream, PartialHashAggregateStream},
+ no_grouping::AggregateStream,
+ row_hash::GroupedHashAggregateStream,
topk_stream::GroupedTopKAggregateStream,
};
use crate::execution_plan::{CardinalityEffect, EmissionType};
@@ -50,6 +52,7 @@ use datafusion_common::{
internal_err, not_impl_err,
};
use datafusion_execution::TaskContext;
+use datafusion_execution::memory_pool::MemoryLimit;
use datafusion_expr::{Accumulator, Aggregate};
use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
use datafusion_physical_expr::equivalence::ProjectionMapping;
@@ -69,6 +72,8 @@ use topk::hash_table::is_supported_hash_key_type;
use topk::heap::is_supported_heap_type;
pub mod group_values;
+mod hash_aggregate;
+mod hash_table;
mod no_grouping;
pub mod order;
mod row_hash;
@@ -496,10 +501,39 @@ impl PartialEq for PhysicalGroupBy {
}
}
+/// Streams used by [`AggregateExec`].
+///
+/// # Stream Variant Schema Notation
+/// For example, `SELECT g, AVG(x) FROM t GROUP BY g` uses these schemas:
+///
+/// ```text
+/// initial input: [g, x]
+/// partial state: [g, AVG(x) state columns, e.g. sum/count]
+/// final result: [g, AVG(x)]
+/// ```
#[expect(clippy::large_enum_variant)]
enum StreamType {
+ /// Single group (no group by) aggregate stream.
+ /// Input output scheme: initial input -> final result
AggregateStream(AggregateStream),
+ /// Partial stage of the hash aggregation
+ /// Input output scheme: initial input -> partial state
+ PartialHash(PartialHashAggregateStream),
+ /// Final stage of the hash aggregation
+ /// Input output scheme: partial state -> final result
+ FinalHash(FinalHashAggregateStream),
+ /// Hash aggregation reused for multiple stages
+ ///
+ /// Note this is being incrementally migrated to dedicated streams like
+ /// [`StreamType::PartialHash`] and [`StreamType::FinalHash`]
+ ///
+ /// See issue for details: <https://github.com/apache/datafusion/issues/22710>
GroupedHash(GroupedHashAggregateStream),
+ /// Grouped TopK aggregate stream.
+ /// Input output scheme: initial input -> final result
+ ///
+ /// Used for grouped aggregation with LIMIT / ordering, where the stream keeps
+ /// only the top groups required by the query.
GroupedPriorityQueue(GroupedTopKAggregateStream),
}
@@ -507,6 +541,8 @@ impl From<StreamType> for SendableRecordBatchStream {
fn from(stream: StreamType) -> Self {
match stream {
StreamType::AggregateStream(stream) => Box::pin(stream),
+ StreamType::PartialHash(stream) => Box::pin(stream),
+ StreamType::FinalHash(stream) => Box::pin(stream),
StreamType::GroupedHash(stream) => Box::pin(stream),
StreamType::GroupedPriorityQueue(stream) => Box::pin(stream),
}
@@ -964,12 +1000,59 @@ impl AggregateExec {
));
}
+ if context
+ .session_config()
+ .options()
+ .execution
+ .enable_migration_aggregate
+ {
+ if self.should_use_partial_hash_stream(context) {
+ return Ok(StreamType::PartialHash(PartialHashAggregateStream::new(
+ self, context, partition,
+ )?));
+ }
+
+ if self.should_use_final_hash_stream(context) {
+ return Ok(StreamType::FinalHash(FinalHashAggregateStream::new(
+ self, context, partition,
+ )?));
+ }
+ }
+
// grouping by something else and we need to just materialize all results
Ok(StreamType::GroupedHash(GroupedHashAggregateStream::new(
self, context, partition,
)?))
}
+ fn should_use_partial_hash_stream(&self, context: &TaskContext) -> bool {
+ // TODO: implement memory-limited path and remove this limitation
+ if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
+ return false;
+ }
+
+ self.mode == AggregateMode::Partial
+ && self.limit_options.is_none()
+ && self.input_order_mode == InputOrderMode::Linear
+ && !self.group_by.is_true_no_grouping()
+ && self.group_by.is_single()
+ }
+
+ fn should_use_final_hash_stream(&self, context: &TaskContext) -> bool {
+ // TODO: implement memory-limited path and remove this limitation
+ if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
+ return false;
+ }
+
+ matches!(
+ self.mode,
+ AggregateMode::Final | AggregateMode::FinalPartitioned
+ ) && self.limit_options.is_none()
+ && self.input_order_mode == InputOrderMode::Linear
+ && !self.group_by.is_true_no_grouping()
+ && self.group_by.is_single()
+ }
+
/// Finds the DataType and SortDirection for this Aggregate, if there is one
pub fn get_minmax_desc(&self) -> Option<(FieldRef, bool)> {
let agg_expr = self.aggr_expr.iter().exactly_one().ok()?;
@@ -2180,6 +2263,32 @@ pub(crate) fn max_duplicate_ordinal(groups: &[Vec<bool>]) -> usize {
/// The outer Vec appears to be for grouping sets
/// The inner Vec contains the results per expression
/// The inner-inner Array contains the results per row
+///
+/// For example, for `GROUP BY GROUPING SETS ((a, b), (a))` with input:
+///
+/// ```text
+/// a b
+/// 1 1
+/// 1 2
+/// 2 1
+/// ```
+///
+/// The output is:
+///
+/// ```text
+/// [
+/// [
+/// a: [1, 1, 2]
+/// b: [1, 2, 1]
+/// grouping_id: [0, 0, 0]
+/// ],
+/// [
+/// a: [1, 1, 2]
+/// b: [NULL, NULL, NULL]
+/// grouping_id: [1, 1, 1]
+/// ]
+/// ]
+/// ```
pub fn evaluate_group_by(
group_by: &PhysicalGroupBy,
batch: &RecordBatch,
@@ -2954,6 +3063,94 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn partial_grouped_aggregate_uses_raw_partial_stream() -> Result<()> {
+ let (schema, batches) = some_data();
+ let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?;
+ let group_by =
+ PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]);
+ let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new(
+ vec![DataType::Float64],
+ vec![DataType::Int32],
+ DataType::Int64,
+ )));
+ let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new(
+ AggregateExprBuilder::new(udaf, vec![col("b", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("input_type_asserting(b)")
+ .build()?,
+ )];
+
+ let partial_aggregate = Arc::new(AggregateExec::try_new(
+ AggregateMode::Partial,
+ group_by.clone(),
+ aggregates.clone(),
+ vec![None],
+ input,
+ Arc::clone(&schema),
+ )?);
+ let task_ctx = Arc::new(
+ TaskContext::default().with_session_config(
+ SessionConfig::new()
+ .with_batch_size(2)
+ .set_bool("datafusion.execution.enable_migration_aggregate", true),
+ ),
+ );
+
+ let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?;
+ assert!(matches!(partial_stream, StreamType::PartialHash(_)));
+
+ let fallback_task_ctx = Arc::new(
+ TaskContext::default().with_session_config(
+ SessionConfig::new()
+ .with_batch_size(2)
+ .set_bool("datafusion.execution.enable_migration_aggregate", false),
+ ),
+ );
+ let stream = partial_aggregate.execute_typed(0, &fallback_task_ctx)?;
+ assert!(matches!(stream, StreamType::GroupedHash(_)));
+
+ let stream: SendableRecordBatchStream = partial_stream.into();
+ let batches = collect(stream).await?;
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2, 1]
+ );
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
+
+ let merge = Arc::new(CoalescePartitionsExec::new(partial_aggregate));
+ let final_aggregate = AggregateExec::try_new(
+ AggregateMode::Final,
+ group_by.as_final(),
+ aggregates,
+ vec![None],
+ merge,
+ Arc::clone(&schema),
+ )?;
+
+ let final_stream = final_aggregate.execute_typed(0, &task_ctx)?;
+ assert!(matches!(final_stream, StreamType::FinalHash(_)));
+
+ let stream = final_aggregate.execute_typed(0, &fallback_task_ctx)?;
+ assert!(matches!(stream, StreamType::GroupedHash(_)));
+
+ let stream: SendableRecordBatchStream = final_stream.into();
+ let batches = collect(stream).await?;
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2, 1]
+ );
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_drop_cancel_without_groups() -> Result<()> {
let task_ctx = Arc::new(TaskContext::default());
@@ -3682,8 +3879,11 @@ mod tests {
&ScalarValue::Float64(Some(0.1)),
);
- let ctx = TaskContext::default().with_session_config(session_config);
- let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?;
+ let ctx = Arc::new(TaskContext::default().with_session_config(session_config));
+ let stream: SendableRecordBatchStream = Box::pin(
+ GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?,
+ );
+ let output = collect(stream).await?;
allow_duplicates! {
assert_snapshot!(batches_to_string(&output), @r"
@@ -3769,8 +3969,11 @@ mod tests {
&ScalarValue::Float64(Some(0.1)),
);
- let ctx = TaskContext::default().with_session_config(session_config);
- let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?;
+ let ctx = Arc::new(TaskContext::default().with_session_config(session_config));
+ let stream: SendableRecordBatchStream = Box::pin(
+ GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?,
+ );
+ let output = collect(stream).await?;
allow_duplicates! {
assert_snapshot!(batches_to_string(&output), @r"
diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt
index 840bff6ea6..8d334d8433 100644
--- a/datafusion/sqllogictest/test_files/information_schema.slt
+++ b/datafusion/sqllogictest/test_files/information_schema.slt
@@ -218,6 +218,7 @@ datafusion.execution.batch_size 8192
datafusion.execution.coalesce_batches true
datafusion.execution.collect_statistics true
datafusion.execution.enable_ansi_mode false
+datafusion.execution.enable_migration_aggregate false
datafusion.execution.enable_recursive_ctes true
datafusion.execution.enforce_batch_size_in_joins false
datafusion.execution.hash_join_buffering_capacity 0
@@ -374,6 +375,7 @@ datafusion.execution.batch_size 8192 Default batch size while creating new batch
datafusion.execution.coalesce_batches true When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting
datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Applies to the default `ListingTableProvider` in DataFusion. Defaults to true.
datafusion.execution.enable_ansi_mode false Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, [...]
+datafusion.execution.enable_migration_aggregate false Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, DataFusion tries the migrated implementations when their preconditions are satisfied. When set to false, grouped aggregation falls back to `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See <https://github.com/apache/datafusion/issues/22710> for details.
datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs
datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower.
datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe sid [...]
diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md
index cc679549de..7c6756a096 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -120,6 +120,7 @@ The following configuration settings are available:
| datafusion.execution.parquet.content_defined_chunking.norm_level | 0 | Normalization level. Increasing this improves deduplication ratio but increases fragmentation. Recommended range is [-3, 3], default is 0. [...]
| datafusion.execution.planning_concurrency | 0 | Fan-out during initial physical planning. This is mostly use to plan `UNION` children in parallel. Defaults to the number of CPU cores on the system [...]
| datafusion.execution.skip_physical_aggregate_schema_check | false | When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. [...]
+| datafusion.execution.enable_migration_aggregate | false | Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, DataFusion tries the migrated implementations when their preconditions are satisfied. When set to false, grouped aggregation falls back to `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See <https://github.com/apa [...]
| datafusion.execution.spill_compression | uncompressed | Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression rati [...]
| datafusion.execution.sort_spill_reservation_bytes | 10485760 | Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's [...]
| datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. [...]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org