diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs index f597c9c59a099..f9e7f2e10e789 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs @@ -378,6 +378,9 @@ async fn run_aggregate_test(input1: Vec, group_by_columns: Vec<&str .await .unwrap(); assert!(collected_running.len() > 2); + // Running should produce more chunk than the usual AggregateExec. + // Otherwise it means that we cannot generate result in running mode. + assert!(collected_running.len() > collected_usual.len()); // compare let usual_formatted = pretty_format_batches(&collected_usual).unwrap().to_string(); let running_formatted = pretty_format_batches(&collected_running) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 9f8e116ed6b42..c5ec1fc4e96c4 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -23,7 +23,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{Result, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::EmitTo; @@ -64,6 +64,9 @@ pub(in crate::aggregates) struct OrderedAggregateTable { /// Output schema: group columns followed by aggregate state or final values. pub(super) output_schema: SchemaRef, + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + /// Grouping and accumulator-specific timing metrics. pub(super) group_by_metrics: GroupByMetrics, @@ -103,15 +106,26 @@ pub(super) struct OrderedAggregateTableBuffer { /// Methods shared by all aggregate modes impl OrderedAggregateTable { + #[expect( + clippy::too_many_arguments, + reason = "keeps ordered partial and final table construction explicit" + )] pub(super) fn new_for_mode( agg: &AggregateExec, partition: usize, input_schema: &SchemaRef, output_schema: SchemaRef, + batch_size: usize, input_order_mode: &InputOrderMode, aggregate_mode: &AggregateMode, filters: Vec>>, ) -> Result { + if batch_size == 0 { + return internal_err!( + "OrderedAggregateTable requires config batch_size >= 1" + ); + } + let group_ordering = GroupOrdering::try_new(input_order_mode)?; let group_schema = agg.group_by.group_schema(input_schema)?; let group_values = new_group_values(group_schema, &group_ordering)?; @@ -138,6 +152,7 @@ impl OrderedAggregateTable { Ok(Self { output_schema, + batch_size, group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), buffer: OrderedAggregateTableBuffer { group_by: Arc::clone(&agg.group_by), @@ -202,6 +217,24 @@ impl OrderedAggregateTable { + self.buffer.group_ordering.size() + self.buffer.group_indices.allocated_size() } + + /// Returns the [`EmitTo`], clamped to the specified batch size + /// + /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number + /// of groups to emit from `GroupValues` / accumulators, and + /// `should_remove_groups` indicates whether `GroupOrdering` must also shift + /// its tracked indexes. + pub(super) fn clamp_emit_to( + &self, + group_count: usize, + emit_to: EmitTo, + ) -> (EmitTo, bool) { + match emit_to { + EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), + EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), + EmitTo::All => (EmitTo::First(self.batch_size), false), + } + } } pub(super) fn remove_emitted_groups(group_ordering: &mut GroupOrdering, emit_to: EmitTo) { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index 91c54888c44f7..03aa812607b88 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -45,6 +45,7 @@ impl OrderedAggregateTable { partition: usize, input_schema: &SchemaRef, output_schema: SchemaRef, + batch_size: usize, input_order_mode: &InputOrderMode, ) -> Result { Self::new_for_mode( @@ -52,6 +53,7 @@ impl OrderedAggregateTable { partition, input_schema, output_schema, + batch_size, input_order_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], @@ -109,10 +111,14 @@ impl OrderedAggregateTable { let Some(emit_to) = self.buffer.group_ordering.emit_to() else { return Ok(None); }; + let (emit_to, should_remove_groups) = + self.clamp_emit_to(self.buffer.group_values.len(), emit_to); let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.buffer.group_values.emit(emit_to)?; - remove_emitted_groups(&mut self.buffer.group_ordering, emit_to); + if should_remove_groups { + remove_emitted_groups(&mut self.buffer.group_ordering, emit_to); + } for acc in &mut self.buffer.accumulators { output.push(acc.evaluate(emit_to)?); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index c4fa380853be7..5ee6be4745a51 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -53,6 +53,7 @@ impl OrderedAggregateTable { agg: &AggregateExec, partition: usize, output_schema: SchemaRef, + batch_size: usize, ) -> Result { let input_schema = agg.input().schema(); Self::new_for_mode( @@ -60,6 +61,7 @@ impl OrderedAggregateTable { partition, &input_schema, output_schema, + batch_size, &agg.input_order_mode, &AggregateMode::Partial, agg.filter_expr.iter().cloned().collect(), @@ -127,10 +129,14 @@ impl OrderedAggregateTable { let Some(emit_to) = self.buffer.group_ordering.emit_to() else { return Ok(None); }; + let (emit_to, should_remove_groups) = + self.clamp_emit_to(self.buffer.group_values.len(), emit_to); let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.buffer.group_values.emit(emit_to)?; - remove_emitted_groups(&mut self.buffer.group_ordering, emit_to); + if should_remove_groups { + remove_emitted_groups(&mut self.buffer.group_ordering, emit_to); + } for acc in &mut self.buffer.accumulators { output.extend(acc.state(emit_to)?); diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index d76b020a8e2ac..ecc6f6afd6259 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -95,6 +95,7 @@ impl OrderedFinalAggregateStream { let schema = Arc::clone(&agg.schema); let input_schema = input.schema(); + 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. @@ -105,6 +106,7 @@ impl OrderedFinalAggregateStream { partition, &input_schema, Arc::clone(&schema), + batch_size, input_order_mode, )?; let reservation = diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index dc8b2a3668661..f68814ceb6fd8 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -114,6 +114,7 @@ impl OrderedPartialAggregateStream { 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. @@ -126,6 +127,7 @@ impl OrderedPartialAggregateStream { agg, partition, Arc::clone(&schema), + batch_size, )?; let reservation = MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]"))