From adedbb4a8ca5132241246fbfd566520ba63f21bf Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Tue, 15 Sep 2026 11:36:09 +0800 Subject: [PATCH 1/3] perf: Avoid copying when materializing output in OrderedPartialAggregateStream --- .../aggregate_hash_table/common_ordered.rs | 54 +- .../ordered_partial_table.rs | 31 +- .../src/aggregates/ordered_partial_stream.rs | 601 +++++++++++++----- 3 files changed, 484 insertions(+), 202 deletions(-) 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 97ef898b51f4d..a5a81bc511e24 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 @@ -352,24 +352,6 @@ impl OrderedAggregateTable { Ok(Some(batch)) } - /// 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), - } - } - /// Aggregates one evaluated input batch after selecting the mode-specific /// accumulator operation. /// @@ -445,20 +427,34 @@ 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 emit_to = match emit_to { + EmitTo::First(n) => EmitTo::First(n.min(self.batch_size)), + EmitTo::All if self.num_groups() > self.batch_size => { + EmitTo::First(self.batch_size) + } + EmitTo::All => EmitTo::All, + }; + self.materialize_groups(emit_to, materialize_accumulator_fn, accumulator_phase) + .map(Some) + } + /// Removes the selected groups once and materializes their output columns. + /// The caller chooses the completed prefix and any output-size limit. + pub(super) fn materialize_groups( + &mut self, + emit_to: EmitTo, + materialize_accumulator_fn: MaterializeAccumulatorFn, + accumulator_phase: AccumulatorPhase, + ) -> Result { let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); let output = self.group_by_metrics.time_emitting(|| { let mut output = self.buffer.group_values.emit(emit_to)?; - if should_remove_groups { - match emit_to { - EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), - // `EmitTo::All` is only used after `input_done`, when all - // buffered groups are known complete and the ordering state is - // no longer needed. - EmitTo::All => {} - } + // EOF can also emit a prefix when a caller limits its batch size, + // but the completed ordering state no longer tracks group indexes. + if let EmitTo::First(n) = emit_to + && matches!(self.buffer.group_ordering.emit_to(), Some(EmitTo::First(_))) + { + self.buffer.group_ordering.remove_groups(n); } for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { @@ -474,6 +470,6 @@ impl OrderedAggregateTable { let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; debug_assert!(batch.num_rows() > 0); - Ok(Some(batch)) + Ok(batch) } } 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 6ed93e59f3296..a7b79c0508021 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 @@ -89,26 +89,27 @@ impl OrderedAggregateTable { ) } - /// Emits the next batch of partial state rows for groups proven complete by - /// the input ordering. + /// Materializes all groups proven complete by the input ordering, leaving + /// the active ordered-key range in the table. /// - /// For example, when the query is `GROUP BY a` and the input is ordered by - /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` - /// are complete and safe to emit. - /// - /// Key steps: - /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly. - /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and - /// all `GroupsAccumulator`s. - /// - /// This may output small batches. Avoiding tiny batches is left to future - /// ordered-aggregation optimizations. - pub(in crate::aggregates) fn next_output_batch( + /// For `GROUP BY a, b` ordered by `a`, seeing a new `a` completes every group + /// with the previous `a`. Remove that entire prefix once: removing only + /// `batch_size` groups at a time repeatedly shifts the remaining hash table + /// and accumulator indexes. The stream slices the materialized batch instead. + pub(in crate::aggregates) fn take_completed_state_batch( &mut self, ) -> Result> { - self.next_output_batch_inner( + if self.is_empty() { + return Ok(None); + } + let Some(emit_to) = self.group_ordering().emit_to() else { + return Ok(None); + }; + self.materialize_groups( + emit_to, HashAggregateAccumulator::state, AccumulatorPhase::State, ) + .map(Some) } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 8c2315588ad7a..5b3e4b65b5a6b 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -24,14 +24,14 @@ use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; -use futures::stream::{Stream, StreamExt}; +use futures::stream::StreamExt; use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; use crate::aggregates::order::GroupOrdering; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; -use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Partial aggregate stream for `InputOrderMode::Sorted` and @@ -66,7 +66,9 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// After each input batch, check whether any groups can be emitted eagerly to /// improve memory efficiency. For example, if the last group key seen is /// `k = 100`, it is safe to emit all groups with keys less than 100 because the -/// input is ordered. +/// input is ordered. Materialize that entire completed prefix once, then emit +/// slices of it before reading more input. This avoids repeatedly removing small +/// batches of groups and shifting the remaining hash table and accumulator state. /// /// # Memory Pressure and Spilling /// @@ -103,12 +105,36 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// remaining states, performs a sort-preserving merge of all runs, and feeds the /// merged input into a fully ordered final aggregate stream. pub(crate) struct OrderedPartialAggregateStream { - schema: SchemaRef, - input: SendableRecordBatchStream, reservation: MemoryReservation, + context: OrderedPartialAggregateContext, + stage: ExecutionStage, +} + +/// Execution stages described in [`OrderedPartialAggregateStream::into_stream`]. +enum ExecutionStage { + Aggregating(Aggregating), + Outputting(Outputting), +} + +struct Aggregating { + input: SendableRecordBatchStream, + table: OrderedAggregateTable, +} + +struct Outputting { + /// Materialized aggregate states. Each iteration emits the first `batch_size` + /// rows and replaces this batch with the remaining slice. + batch: RecordBatch, + /// Aggregation stage to resume after output; `None` after EOF. + resume: Option, +} + +/// Immutable execution context shared by aggregation and output emission. +struct OrderedPartialAggregateContext { + schema: SchemaRef, + batch_size: usize, baseline_metrics: BaselineMetrics, reduction_factor: metrics::RatioMetrics, - table: Option>, } impl OrderedPartialAggregateStream { @@ -146,199 +172,458 @@ impl OrderedPartialAggregateStream { .register(context.memory_pool()); Ok(Self { - schema, - input, reservation, - baseline_metrics, - reduction_factor, - table: Some(table), + context: OrderedPartialAggregateContext { + schema, + batch_size, + baseline_metrics, + reduction_factor, + }, + stage: ExecutionStage::Aggregating(Aggregating { input, table }), }) } - pub(crate) fn into_stream(self) -> SendableRecordBatchStream { - let schema_clone = Arc::clone(&self.schema); - - let cloned_metrics = self.baseline_metrics.clone(); - let stream = Box::pin(RecordBatchStreamAdapter::new( - schema_clone, - self.create_stream(), - )); - - Box::pin(ObservedStream::new(stream, cloned_metrics, None)) - } - - /// Entry point for the ordered partial aggregate state machine. - /// - /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas. + /// Entry point for the ordered partial aggregate execution stages. /// - /// State transitions are implemented using the generator pattern; see the comments in [`async_try_stream`]. + /// See [`OrderedPartialAggregateStream`] for high-level ideas. /// - /// Conceptual state-transition graph: + /// # Stage transition graph: /// /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling ordered input and aggregating batches - /// into the ordered partial aggregate table. + /// +----[2]----+ +----[5]----+ + /// | | | | + /// v | v | + /// +-------------------+ +-------------------+ + /// | | | | + /// (start)-[1]->| Aggregating |-----[3]---->| Outputting | + /// | |<----[6]-----| | + /// +-------------------+ +-------------------+ + /// | [4] | [7] + /// | | + /// +----------------+----------------+ + /// | + /// v + /// +---------+ + /// | Done |--[8]--> (end) + /// +---------+ + /// ``` + /// + /// ## Stages + /// + /// - [`Aggregating`]: Aggregates raw input and materializes one batch of partial + /// states. + /// - [`Outputting`]: Emits slices of one materialized batch. If the materialized + /// buffers cannot be reserved while slicing, hand off the whole batch, then + /// resume aggregation or finish as described below. /// - /// ReadingInput - /// -> ReadingInput - /// Aggregate one input batch. If the ordering proves some groups are - /// complete, yield one partial-state batch immediately, then continue - /// reading input. Otherwise continue directly with the next input batch. - /// -> DrainingFinal - /// Input was exhausted. Mark the table input as done so every remaining - /// group is safe to emit. + /// ### Incremental output /// - /// DrainingFinal - /// -> DrainingFinal - /// One remaining partial-state batch was yielded; repeat to continue - /// draining the table. - /// -> Done - /// All remaining groups were emitted. + /// Consider this query with input ordered only by `k1`: /// - /// Done - /// -> (end) + /// ```sql + /// SELECT k1, k2, AVG(v) + /// FROM table_with_order_k1 + /// GROUP BY k1, k2 /// ``` - fn create_stream(mut self) -> impl Stream> { - async_try_stream(|mut emitter| async move { - let mut table = self - .table - .take() - .expect("OrderedPartialAggregateStream state should not be None"); - - self.handle_reading_input(&mut table, &mut emitter).await?; - - // Input has exhausted, move to the final draining stage. - self.close_input(); - table.input_done(); - - self.handle_draining_final(table, &mut emitter).await?; - + /// + /// Suppose one `k1` value spans 1M rows with distinct, unordered `k2` + /// values. Ordering only proves these 1M `(k1, k2)` groups complete when + /// `k1` changes, so a single early emission can produce far more than + /// `batch_size` rows. + /// + /// Emitting those groups in small batches through [EmitTo::First] would + /// repeatedly remove a prefix from [`GroupValues`]. Because the group + /// values are stored contiguously, each removal copies the remaining values + /// and updates their group indexes. + /// + /// To avoid repeating that work, this stream: + /// + /// 1. Materializes all completed groups into one large batch. + /// 2. Emits `batch_size` slices that share the batch's buffers. + /// + /// Blocked aggregate state management may simplify this approach: + /// + /// + /// [`GroupValues`]: crate::aggregates::group_values::GroupValues + /// [EmitTo::First]: datafusion_expr::EmitTo::First + /// + /// + /// ## Transition Edges + /// + /// 1. Start. + /// 2. Aggregate one input batch. If memory fits and no groups are complete, + /// continue reading input. + /// 3. Prepare output: + /// - Ordering proves a prefix complete: materialize the entire prefix once, + /// retaining the input and active groups to resume aggregation. + /// - On memory pressure with partial ordering, materialize all current + /// states instead, including incomplete groups, and reset the table. + /// - At EOF, materialize all remaining states and prepare to output. + /// 4. Input was exhausted with no remaining groups, directly end. + /// 5. Yield one slice without materializing the table again. Keep the shared + /// buffers reserved until handing off the last slice. + /// 6. The batch was fully emitted and retained aggregation can resume. + /// 7. The output batch was fully emitted. + /// 8. End. + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let Self { + reservation, + context, + stage, + } = self; + let schema = Arc::clone(&context.schema); + let metrics = context.baseline_metrics.clone(); + let stream = async_try_stream(|mut emitter| async move { + let mut stage = Some(stage); + while let Some(current_stage) = stage { + stage = match current_stage { + ExecutionStage::Aggregating(aggregating) => { + aggregating.handle_stage(&context, &reservation).await? + } + ExecutionStage::Outputting(outputting) => { + outputting + .handle_stage(&context, &reservation, &mut emitter) + .await? + } + }; + } Ok(()) - }) - } - - fn close_input(&mut self) { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + }); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + Box::pin(ObservedStream::new(stream, metrics, None)) } +} - /// Consumes one ordered input batch, then immediately emits completed groups - /// if the ordering proves any group is ready. +impl Aggregating { + /// Aggregates raw input and materializes one batch of partial states. /// - /// See comments at [`Self::create_stream`] for details. - async fn handle_reading_input( - &mut self, - table: &mut OrderedAggregateTable, - emitter: &mut TryEmitter, - ) -> Result<()> { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + /// See [`OrderedPartialAggregateStream::into_stream`] for stage transitions. + async fn handle_stage( + mut self, + context: &OrderedPartialAggregateContext, + reservation: &MemoryReservation, + ) -> Result> { + let elapsed_compute = context.baseline_metrics.elapsed_compute(); while let Some(batch) = self.input.next().await.transpose()? { - let input_rows = batch.num_rows(); - self.reduction_factor.add_total(input_rows); - + context.reduction_factor.add_total(batch.num_rows()); let timer = elapsed_compute.timer(); - - table.aggregate_batch(&batch)?; - - // Check memory reservation. See function comments for details. - if let Some(batch) = self.resize_or_take_state_batch(table)? { - self.reduction_factor.add_part(batch.num_rows()); - drop(timer); - emitter.emit(batch).await; - continue; - } - - let Some(batch) = table.next_output_batch()? else { - // Can't do early emit, continue aggregating. + self.table.aggregate_batch(&batch)?; + + let output = match reservation.try_resize(self.table.memory_size()) { + Ok(()) => self.table.take_completed_state_batch()?, + Err(oom @ DataFusionError::ResourcesExhausted(_)) => { + // Partial ordering may have an unbounded active key range. + // The final stage can merge incomplete states emitted here. + if matches!(self.table.group_ordering(), GroupOrdering::Full(_)) { + return Err(oom); + } + let Some(batch) = self.table.take_state_batch()? else { + return Err(oom); + }; + Some(batch) + } + Err(e) => return Err(e), + }; + let Some(batch) = output else { continue; }; - self.reduction_factor.add_part(batch.num_rows()); - self.reservation.try_resize(table.memory_size())?; + timer.done(); - drop(timer); - emitter.emit(batch).await; + // OOM, do early emit next, and go back to the current state to continue + // aggregating + return Ok(Some(ExecutionStage::Outputting(Outputting { + batch, + resume: Some(self), + }))); } - Ok(()) + // Release upstream resources before draining the remaining states. + drop(self.input); + self.table.input_done(); + let timer = elapsed_compute.timer(); + let output = self.table.take_completed_state_batch()?; + drop(self.table); + timer.done(); + + let Some(batch) = output else { + reservation.try_resize(0)?; + return Ok(None); + }; + Ok(Some(ExecutionStage::Outputting(Outputting { + batch, + resume: None, + }))) } +} - /// Update the memory reservation, and: - /// - If memory reservation succeed, returns `Ok(None)` - /// - If memory reservation failed, - /// - If input is partially ordered, materialize all the output, and - /// directly send them to the final aggregation stage. - /// Returns `Ok(Some(batch))` - /// - If input is fully ordered, directly return error. It's not - /// expected to use more than constant memory. - /// Returns `Err(..)` - /// - /// # Implementation Note - /// Incrementally output it after the blocked state management is ready, keep - /// it simple for now. +impl Outputting { + /// Emits slices of one materialized batch without touching the hash table. /// - /// Issue: - fn resize_or_take_state_batch( - &mut self, - table: &mut OrderedAggregateTable, - ) -> Result> { - let oom = match self.reservation.try_resize(table.memory_size()) { - Ok(()) => return Ok(None), - Err(e @ DataFusionError::ResourcesExhausted(_)) => e, - Err(e) => return Err(e), + /// See [`OrderedPartialAggregateStream::into_stream`] for stage transitions + /// and output memory accounting. + async fn handle_stage( + self, + context: &OrderedPartialAggregateContext, + reservation: &MemoryReservation, + emitter: &mut TryEmitter, + ) -> Result> { + let Self { mut batch, resume } = self; + let elapsed_compute = context.baseline_metrics.elapsed_compute(); + let mut timer = elapsed_compute.timer(); + let (table_memory, next_stage) = match resume { + Some(aggregating) => ( + aggregating.table.memory_size(), + Some(ExecutionStage::Aggregating(aggregating)), + ), + None => (0, None), }; + let batch_memory = batch.get_array_memory_size(); + match reservation.try_resize(table_memory + batch_memory) { + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(_)) => { + // If we cannot hold the batch while slicing, hand it off whole. + // Only the retained table needs to remain reserved. + reservation.try_resize(table_memory)?; + context.reduction_factor.add_part(batch.num_rows()); + timer.done(); + emitter.emit(batch).await; + return Ok(next_stage); + } + Err(e) => return Err(e), + } - if matches!(table.group_ordering(), GroupOrdering::Full(_)) { - return Err(oom); + while batch.num_rows() > context.batch_size { + // 1. Emit first `batch_size` rows from `batch` + // 2. Update `batch`` with the remaining tail + let output = batch.slice(0, context.batch_size); + batch = + batch.slice(context.batch_size, batch.num_rows() - context.batch_size); + context.reduction_factor.add_part(output.num_rows()); + timer.done(); + emitter.emit(output).await; + timer = elapsed_compute.timer(); } - let Some(batch) = table.take_state_batch()? else { - return Err(oom); - }; - self.reservation.try_resize(table.memory_size())?; - Ok(Some(batch)) + // The final slice transfers ownership of the buffers to the consumer. + reservation.try_shrink(batch_memory)?; + context.reduction_factor.add_part(batch.num_rows()); + timer.done(); + emitter.emit(batch).await; + Ok(next_stage) } +} - /// Emits one batch after input is exhausted. - /// - /// `table.input_done()` has already made every remaining group safe to emit, - /// so this state keeps draining until the table is empty. - /// - /// See comments at [`Self::create_stream`] for details. - /// - async fn handle_draining_final( - &mut self, - mut table: OrderedAggregateTable, - emitter: &mut TryEmitter, - ) -> Result<()> { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let mut timer = elapsed_compute.timer(); - - while let Some(batch) = table.next_output_batch()? { - self.reduction_factor.add_part(batch.num_rows()); +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Int32Type, Int64Type, Schema}; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::LexOrdering; + + use crate::aggregates::PhysicalGroupBy; + use crate::test::TestMemoryExec; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("sort_key", DataType::Int32, false), + Field::new("group_key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])) + } - if table.is_empty() { - // Clear memory before emitting last batch so we don't have to wait for next poll to clear - drop(table); - let _ = self.reservation.try_resize(0); - drop(timer); + fn batch(sort_keys: Vec, group_keys: Vec) -> Result { + let values = vec![1; sort_keys.len()]; + Ok(RecordBatch::try_new( + schema(), + vec![ + Arc::new(Int32Array::from(sort_keys)), + Arc::new(Int32Array::from(group_keys)), + Arc::new(Int64Array::from(values)), + ], + )?) + } - emitter.emit(batch).await; + fn aggregate(batches: Vec, fully_sorted: bool) -> Result { + let schema = schema(); + let mut sort_exprs = + vec![PhysicalSortExpr::new_default(col("sort_key", &schema)?)]; + if fully_sorted { + sort_exprs.push(PhysicalSortExpr::new_default(col("group_key", &schema)?)); + } + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![LexOrdering::new(sort_exprs).unwrap()])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![ + (col("sort_key", &schema)?, "sort_key".to_string()), + (col("group_key", &schema)?, "group_key".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )], + vec![None], + input, + schema, + ) + } - return Ok(()); + #[tokio::test] + async fn materializes_completed_prefix_and_eof_once() -> Result<()> { + // The first input completes eight groups and leaves eight active groups. + // The second input updates only those active groups before EOF. + let aggregate = aggregate( + vec![ + batch( + [vec![1; 8], vec![2; 8]].concat(), + (0..8).cycle().take(16).collect(), + )?, + batch(vec![2; 8], (0..8).collect())?, + ], + false, + )?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)), + ); + let partial = OrderedPartialAggregateStream::new(&aggregate, &context, 0)?; + let reduction = partial.context.reduction_factor.clone(); + let mut stream = partial.into_stream(); + + for range in 1..=2 { + let mut first_buffer = None; + let mut first_reservation = 0; + for slice in 0..4 { + let output = stream.next().await.unwrap()?; + assert_eq!(output.num_rows(), 2); + assert_eq!( + output + .column(0) + .as_primitive::() + .values() + .as_ref(), + &[range, range] + ); + assert_eq!( + output + .column(1) + .as_primitive::() + .values() + .as_ref(), + &[slice * 2, slice * 2 + 1] + ); + assert_eq!( + output + .column(2) + .as_primitive::() + .values() + .as_ref(), + &[i64::from(range); 2] + ); + + // Every slice of this completed range must share the same + // allocation, rather than re-materializing the hash table. + let buffer = output + .column(1) + .as_primitive::() + .values() + .inner() + .data_ptr(); + assert_eq!(*first_buffer.get_or_insert(buffer), buffer); + let reserved = context.memory_pool().reserved(); + let held = output.get_array_memory_size(); + if slice == 0 { + first_reservation = reserved; + } + if slice < 3 { + assert_eq!(reserved, first_reservation); + assert!(reserved >= held); + } else { + assert_eq!(reserved, first_reservation - held); + if range == 2 { + assert_eq!( + reserved, 0, + "release all memory before the last EOF slice" + ); + } + } } + } + assert!(stream.next().await.is_none()); + assert_eq!(reduction.part(), 16); + assert_eq!(reduction.total(), 24); + Ok(()) + } - self.reservation.try_resize(table.memory_size())?; - - timer.done(); - emitter.emit(batch).await; - timer = elapsed_compute.timer(); + #[tokio::test] + async fn memory_pressure_slices_or_hands_off_whole_batch() -> Result<()> { + let num_groups = 4096; + let input = batch(vec![1; num_groups], (0..num_groups as i32).collect())?; + let aggregate = aggregate(vec![input.clone(), input], false)?; + + for (memory_limit, output_size) in [(128 * 1024, 1024), (32 * 1024, num_groups)] { + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(memory_limit, 1.0) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(SessionConfig::new().with_batch_size(1024)), + ); + let mut stream = OrderedPartialAggregateStream::new(&aggregate, &context, 0)? + .into_stream(); + let mut rows = 0; + while let Some(output) = stream.next().await.transpose()? { + assert_eq!(output.num_rows(), output_size); + // Each input must emit before the next input is aggregated, + // otherwise the counts would be 2 rather than 1. + assert!( + output + .column(2) + .as_primitive::() + .values() + .iter() + .all(|&count| count == 1) + ); + if output_size == num_groups { + assert!( + context.memory_pool().reserved() < output.get_array_memory_size() + ); + } + rows += output.num_rows(); + } + assert_eq!(rows, 2 * num_groups); + assert_eq!(context.memory_pool().reserved(), 0); } + Ok(()) + } - // was empty + #[tokio::test] + async fn fully_ordered_memory_pressure_still_errors() -> Result<()> { + let aggregate = aggregate(vec![batch(vec![1; 8], (0..8).collect())?], true)?; + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1, 1.0) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + let mut stream = + OrderedPartialAggregateStream::new(&aggregate, &context, 0)?.into_stream(); + assert!(matches!( + stream.next().await.unwrap(), + Err(DataFusionError::ResourcesExhausted(_)) + )); + drop(stream); + assert_eq!(context.memory_pool().reserved(), 0); Ok(()) } } From 0f165d0129820a3bd66641609d9efb954c0cc018 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Tue, 15 Sep 2026 11:37:51 +0800 Subject: [PATCH 2/3] remove unnecessary UTs by AI --- .../src/aggregates/ordered_partial_stream.rs | 213 ------------------ 1 file changed, 213 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 5b3e4b65b5a6b..dd845353c77eb 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -414,216 +414,3 @@ impl Outputting { Ok(next_stage) } } - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{AsArray, Int32Array, Int64Array}; - use arrow::datatypes::{DataType, Field, Int32Type, Int64Type, Schema}; - use datafusion_execution::config::SessionConfig; - use datafusion_execution::runtime_env::RuntimeEnvBuilder; - use datafusion_functions_aggregate::count::count_udaf; - use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::aggregate::AggregateExprBuilder; - use datafusion_physical_expr::expressions::col; - use datafusion_physical_expr_common::sort_expr::LexOrdering; - - use crate::aggregates::PhysicalGroupBy; - use crate::test::TestMemoryExec; - - fn schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("sort_key", DataType::Int32, false), - Field::new("group_key", DataType::Int32, false), - Field::new("value", DataType::Int64, false), - ])) - } - - fn batch(sort_keys: Vec, group_keys: Vec) -> Result { - let values = vec![1; sort_keys.len()]; - Ok(RecordBatch::try_new( - schema(), - vec![ - Arc::new(Int32Array::from(sort_keys)), - Arc::new(Int32Array::from(group_keys)), - Arc::new(Int64Array::from(values)), - ], - )?) - } - - fn aggregate(batches: Vec, fully_sorted: bool) -> Result { - let schema = schema(); - let mut sort_exprs = - vec![PhysicalSortExpr::new_default(col("sort_key", &schema)?)]; - if fully_sorted { - sort_exprs.push(PhysicalSortExpr::new_default(col("group_key", &schema)?)); - } - let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? - .try_with_sort_information(vec![LexOrdering::new(sort_exprs).unwrap()])?; - let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); - AggregateExec::try_new( - AggregateMode::Partial, - PhysicalGroupBy::new_single(vec![ - (col("sort_key", &schema)?, "sort_key".to_string()), - (col("group_key", &schema)?, "group_key".to_string()), - ]), - vec![Arc::new( - AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("count_value") - .build()?, - )], - vec![None], - input, - schema, - ) - } - - #[tokio::test] - async fn materializes_completed_prefix_and_eof_once() -> Result<()> { - // The first input completes eight groups and leaves eight active groups. - // The second input updates only those active groups before EOF. - let aggregate = aggregate( - vec![ - batch( - [vec![1; 8], vec![2; 8]].concat(), - (0..8).cycle().take(16).collect(), - )?, - batch(vec![2; 8], (0..8).collect())?, - ], - false, - )?; - let context = Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().with_batch_size(2)), - ); - let partial = OrderedPartialAggregateStream::new(&aggregate, &context, 0)?; - let reduction = partial.context.reduction_factor.clone(); - let mut stream = partial.into_stream(); - - for range in 1..=2 { - let mut first_buffer = None; - let mut first_reservation = 0; - for slice in 0..4 { - let output = stream.next().await.unwrap()?; - assert_eq!(output.num_rows(), 2); - assert_eq!( - output - .column(0) - .as_primitive::() - .values() - .as_ref(), - &[range, range] - ); - assert_eq!( - output - .column(1) - .as_primitive::() - .values() - .as_ref(), - &[slice * 2, slice * 2 + 1] - ); - assert_eq!( - output - .column(2) - .as_primitive::() - .values() - .as_ref(), - &[i64::from(range); 2] - ); - - // Every slice of this completed range must share the same - // allocation, rather than re-materializing the hash table. - let buffer = output - .column(1) - .as_primitive::() - .values() - .inner() - .data_ptr(); - assert_eq!(*first_buffer.get_or_insert(buffer), buffer); - let reserved = context.memory_pool().reserved(); - let held = output.get_array_memory_size(); - if slice == 0 { - first_reservation = reserved; - } - if slice < 3 { - assert_eq!(reserved, first_reservation); - assert!(reserved >= held); - } else { - assert_eq!(reserved, first_reservation - held); - if range == 2 { - assert_eq!( - reserved, 0, - "release all memory before the last EOF slice" - ); - } - } - } - } - assert!(stream.next().await.is_none()); - assert_eq!(reduction.part(), 16); - assert_eq!(reduction.total(), 24); - Ok(()) - } - - #[tokio::test] - async fn memory_pressure_slices_or_hands_off_whole_batch() -> Result<()> { - let num_groups = 4096; - let input = batch(vec![1; num_groups], (0..num_groups as i32).collect())?; - let aggregate = aggregate(vec![input.clone(), input], false)?; - - for (memory_limit, output_size) in [(128 * 1024, 1024), (32 * 1024, num_groups)] { - let runtime = RuntimeEnvBuilder::default() - .with_memory_limit(memory_limit, 1.0) - .build_arc()?; - let context = Arc::new( - TaskContext::default() - .with_runtime(runtime) - .with_session_config(SessionConfig::new().with_batch_size(1024)), - ); - let mut stream = OrderedPartialAggregateStream::new(&aggregate, &context, 0)? - .into_stream(); - let mut rows = 0; - while let Some(output) = stream.next().await.transpose()? { - assert_eq!(output.num_rows(), output_size); - // Each input must emit before the next input is aggregated, - // otherwise the counts would be 2 rather than 1. - assert!( - output - .column(2) - .as_primitive::() - .values() - .iter() - .all(|&count| count == 1) - ); - if output_size == num_groups { - assert!( - context.memory_pool().reserved() < output.get_array_memory_size() - ); - } - rows += output.num_rows(); - } - assert_eq!(rows, 2 * num_groups); - assert_eq!(context.memory_pool().reserved(), 0); - } - Ok(()) - } - - #[tokio::test] - async fn fully_ordered_memory_pressure_still_errors() -> Result<()> { - let aggregate = aggregate(vec![batch(vec![1; 8], (0..8).collect())?], true)?; - let runtime = RuntimeEnvBuilder::default() - .with_memory_limit(1, 1.0) - .build_arc()?; - let context = Arc::new(TaskContext::default().with_runtime(runtime)); - let mut stream = - OrderedPartialAggregateStream::new(&aggregate, &context, 0)?.into_stream(); - assert!(matches!( - stream.next().await.unwrap(), - Err(DataFusionError::ResourcesExhausted(_)) - )); - drop(stream); - assert_eq!(context.memory_pool().reserved(), 0); - Ok(()) - } -} From 1950fb2b0bd0271200097bc8f9332cad013465a2 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Tue, 15 Sep 2026 11:59:41 +0800 Subject: [PATCH 3/3] remove unnecessary comments --- .../aggregates/aggregate_hash_table/ordered_partial_table.rs | 5 ----- 1 file changed, 5 deletions(-) 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 a7b79c0508021..b53012d8bd1e4 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 @@ -91,11 +91,6 @@ impl OrderedAggregateTable { /// Materializes all groups proven complete by the input ordering, leaving /// the active ordered-key range in the table. - /// - /// For `GROUP BY a, b` ordered by `a`, seeing a new `a` completes every group - /// with the previous `a`. Remove that entire prefix once: removing only - /// `batch_size` groups at a time repeatedly shifts the remaining hash table - /// and accumulator indexes. The stream slices the materialized batch instead. pub(in crate::aggregates) fn take_completed_state_batch( &mut self, ) -> Result> {