diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 2cda54d678250..f6e2e576a54c5 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -123,7 +123,7 @@ async fn group_by_hash() { .with_query("select count(*) from t GROUP BY service, host, pod, container") .with_expected_errors(vec![ "Resources exhausted: Additional allocation failed", - "for PartialHashAggregateStream[0]", + "for FinalHashAggregateStream[0]", ]) .with_memory_limit(1_000) .run() @@ -747,7 +747,7 @@ async fn oom_grouped_hash_aggregate() { .with_query("SELECT COUNT(*), SUM(request_bytes) FROM t GROUP BY host") .with_expected_errors(vec![ "Failed to allocate additional", - "for PartialHashAggregateStream[0]", + "for FinalHashAggregateStream[0]", ]) .with_memory_limit(1_000) .run() diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index 448759ad74c54..3677bf4281ed8 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -31,26 +31,28 @@ //! //! [`Schema::contains`]: arrow::datatypes::Schema::contains -use std::sync::Arc; +use std::{num::NonZeroUsize, sync::Arc}; use arrow::array::{BooleanArray, RecordBatch, StructArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; use datafusion::datasource::MemTable; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::physical_expr::aggregate::AggregateExprBuilder; -use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; use datafusion::physical_plan::collect; use datafusion::physical_plan::expressions::col; +use datafusion::physical_plan::{ExecutionPlan, displayable}; use datafusion::prelude::*; -use datafusion_common::Result; +use datafusion_common::{Result, ScalarValue}; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::array_agg::array_agg_udaf; +use crate::helper::plan_metrics::{plan_spill_count, plan_spilled_bytes}; + /// Returns the fields of the struct column `b`: a single `colA Boolean`. /// /// `col_a_nullable` controls whether `colA` is declared nullable — the only @@ -89,6 +91,10 @@ struct AggregateBatchesTest { /// If set, the context uses a [`FairSpillPool`] of this size (and a small /// batch size) so the aggregation is forced to spill. memory_limit: Option, + /// If set, fixes aggregate parallelism for deterministic memory pressure. + target_partitions: Option, + /// If set, test native DISTINCT aggregation rather than its group-by rewrite. + disable_single_distinct_to_groupby: bool, } impl AggregateBatchesTest { @@ -96,6 +102,8 @@ impl AggregateBatchesTest { Self { num_rows: 100, memory_limit: None, + target_partitions: None, + disable_single_distinct_to_groupby: false, } } @@ -109,6 +117,16 @@ impl AggregateBatchesTest { self } + fn with_target_partitions(mut self, target_partitions: usize) -> Self { + self.target_partitions = Some(target_partitions); + self + } + + fn without_single_distinct_to_groupby(mut self) -> Self { + self.disable_single_distinct_to_groupby = true; + self + } + /// Runs `sql` against the table described above and asserts the result /// has one output row per group (i.e. [`Self::num_rows`] rows in total). async fn run(self, sql: &str) -> Result<()> { @@ -138,22 +156,55 @@ impl AggregateBatchesTest { let ctx = match self.memory_limit { Some(limit) => { + // Include live consumers and peaks in any memory-pool failure. + // The FairSpillPool limit alone does not identify which concurrent + // spillable reservations divided its per-consumer allocation. + let memory_pool = TrackConsumersPool::new( + FairSpillPool::new(limit), + NonZeroUsize::new(10).unwrap(), + ); let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(FairSpillPool::new(limit))) + .with_memory_pool(Arc::new(memory_pool)) .build_arc()?; - SessionContext::new_with_config_rt( - SessionConfig::new().with_batch_size(100), - runtime, - ) + let mut config = SessionConfig::new().with_batch_size(100).set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(1.0)), + ); + if let Some(target_partitions) = self.target_partitions { + config = config.with_target_partitions(target_partitions); + } + SessionContext::new_with_config_rt(config, runtime) } None => SessionContext::new(), }; ctx.register_table("t", Arc::new(table))?; + if self.disable_single_distinct_to_groupby { + assert!(ctx.remove_optimizer_rule("single_distinct_aggregation_to_group_by")); + } - let result = ctx.sql(sql).await?.collect().await?; + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + if self.disable_single_distinct_to_groupby { + let plan = displayable(plan.as_ref()).indent(true).to_string(); + assert_eq!( + plan.matches("AggregateExec").count(), + 1, + "expected native DISTINCT aggregation:\n{plan}" + ); + } + let result = collect(Arc::clone(&plan), ctx.task_ctx()).await?; let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); assert_eq!(total_rows, self.num_rows as usize); + if self.memory_limit.is_some() { + assert!( + plan_spill_count(plan.as_ref()) > 0, + "expected aggregation to spill" + ); + assert!( + plan_spilled_bytes(plan.as_ref()) > 0, + "expected aggregation to spill bytes" + ); + } Ok(()) } } @@ -176,7 +227,7 @@ async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> { async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { AggregateBatchesTest::new() .with_num_rows(10_000) - .with_memory_limit(4_000_000) + .with_memory_limit(1_000_000) .run("SELECT a, array_agg(b) FROM t GROUP BY a") .await } @@ -185,7 +236,10 @@ async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> { AggregateBatchesTest::new() .with_num_rows(10_000) - .with_memory_limit(4_000_000) + // One partition keeps the native aggregate's memory pressure deterministic. + .with_target_partitions(1) + .without_single_distinct_to_groupby() + .with_memory_limit(1_000_000) .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") .await } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 5f101950870d2..854bbdafe3a04 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -366,9 +366,16 @@ impl AggregateHashTable { let batch = RecordBatch::try_new(state_schema, output)?; debug_assert!(batch.num_rows() > 0); - // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the - // key/index buffers too so the memory reservation can be released - // before the batch is sorted for spilling. + // State emission should reset accumulators, but spill recovery must + // release every emitted allocation even for an accumulator that retains + // capacity. Rebuild the accumulator set before returning the state batch. + state.accumulators = state + .accumulators + .iter() + .map(HashAggregateAccumulator::empty_like) + .collect::>()?; + // Explicitly shrink key/index buffers too so the memory reservation can + // be released before the batch is sorted for spilling. state.group_values.clear_shrink(0); state.batch_group_indices.clear(); state.batch_group_indices.shrink_to_fit(); 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..15d8ef69bf843 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 @@ -341,9 +341,17 @@ impl OrderedAggregateTable { let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; debug_assert!(batch.num_rows() > 0); - // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the - // key/index buffers too so the memory reservation can be released - // before the batch is passed downstream or sorted for spilling. + // State emission should reset accumulators, but spill recovery must + // release every emitted allocation even for an accumulator that retains + // capacity. Rebuild the accumulator set before returning the state batch. + self.buffer.accumulators = self + .buffer + .accumulators + .iter() + .map(AggregateAccumulator::empty_like) + .collect::>()?; + // Explicitly shrink key/index buffers too so the memory reservation can + // be released before the batch is passed downstream or sorted for spilling. self.buffer.group_values.clear_shrink(0); self.buffer.group_indices.clear(); self.buffer.group_indices.shrink_to_fit(); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 377abdae71f4c..fd2beb4318ece 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_expr::EmitTo; use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; use crate::aggregates::order::GroupOrdering; @@ -105,6 +106,65 @@ impl AggregateHashTable { }) } + /// Starts a bounded-memory drain of partial aggregate states. + pub(in crate::aggregates) fn start_early_emit(&mut self) { + self.start_outputting(); + } + + /// Emits at most one output batch while releasing its groups from the table. + /// + /// Unlike terminal output, this must not materialize all states: early + /// emission can be triggered precisely because the complete state does not + /// fit in the memory pool. Once drained, rebuild an empty table so raw input + /// aggregation can resume. + pub(in crate::aggregates) fn next_early_emit_batch( + &mut self, + ) -> Result> { + let state_schema = Arc::clone(&self.state_schema); + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let group_by_metrics = self.group_by_metrics.clone(); + let AggregateHashTableState::Outputting(mut state) = + std::mem::replace(&mut self.state, AggregateHashTableState::Done) + else { + return Ok(None); + }; + + let emit_to = EmitTo::First(self.batch_size.min(state.group_values.len())); + let columns = group_by_metrics.time_emitting(|| { + let mut columns = state.group_values.emit(emit_to)?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + columns.extend(accumulator_metrics.time( + idx, + AccumulatorPhase::State, + || acc.state(emit_to), + )?); + } + Ok::<_, datafusion_common::DataFusionError>(columns) + })?; + let batch = RecordBatch::try_new(state_schema, columns)?; + debug_assert!(batch.num_rows() > 0); + + if state.group_values.is_empty() { + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let accumulators = state + .accumulators + .iter() + .map(HashAggregateAccumulator::empty_like) + .collect::>>()?; + self.state = AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: state.group_by, + group_values, + batch_group_indices: Vec::new(), + accumulators, + }); + } else { + self.state = AggregateHashTableState::Outputting(state); + } + + Ok(Some(batch)) + } + /// Partial aggregation consumes raw input rows and updates the table's /// partial-state accumulators. pub(in crate::aggregates) fn aggregate_batch( diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs index a90c4379c1a0c..e9dbfe4d9b8d9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::mem::size_of; use std::sync::Arc; use crate::aggregates::group_values::multi_group_by::Nulls; @@ -164,7 +165,7 @@ impl GroupColumn for BooleanGroupValueBuilder { } fn size(&self) -> usize { - self.buffer.capacity() / 8 + self.nulls.allocated_size() + size_of::() + self.buffer.capacity() / 8 + self.nulls.allocated_size() } fn build(self: Box) -> ArrayRef { diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index 94d2e48ad34c4..619a3ac10f9a7 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -343,7 +343,8 @@ where } fn size(&self) -> usize { - self.buffer.capacity() * size_of::() + size_of::() + + self.buffer.capacity() * size_of::() + self.offsets.allocated_size() + self.nulls.allocated_size() } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs index 26dd7f8e05c78..f86e94160c1f1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs @@ -28,6 +28,7 @@ use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, exec_datafusion_err}; use datafusion_expr::GroupSelection; +use std::mem::size_of; use std::sync::Arc; /// An implementation of [`GroupColumn`] for `FixedSizeBinary` values @@ -212,7 +213,7 @@ impl GroupColumn for FixedSizeBinaryGroupValueBuilder { } fn size(&self) -> usize { - self.buffer.allocated_size() + self.nulls.allocated_size() + size_of::() + self.buffer.allocated_size() + self.nulls.allocated_size() } fn build(self: Box) -> ArrayRef { diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs index 8c681698be283..cfa6fffa0414d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs @@ -40,6 +40,7 @@ use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::GroupSelection; +use std::mem::size_of; use std::sync::Arc; /// A [`GroupColumn`] for `List` (`O = i32`) and `LargeList` (`O = i64`). @@ -184,7 +185,8 @@ impl GroupColumn for ListGroupValueBuilder { } fn size(&self) -> usize { - self.offsets.allocated_size() + size_of::() + + self.offsets.allocated_size() + self.outer_nulls.allocated_size() + self.child.size() } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 6c1926b402cee..d46d44b6711f9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -105,7 +105,7 @@ pub trait GroupColumn: Send + Sync { self.len() == 0 } - /// Returns the number of bytes used by this [`GroupColumn`] + /// Returns this column's concrete owner descriptor and retained allocations. fn size(&self) -> usize; /// Builds a new array from all of the stored rows @@ -275,6 +275,14 @@ impl VectorizedOperationBuffers { self.equal_to_group_indices.clear(); self.remaining_row_indices.clear(); } + + fn size(&self) -> usize { + self.append_row_indices.allocated_size() + + self.equal_to_row_indices.allocated_size() + + self.equal_to_group_indices.allocated_size() + + self.equal_to_results.capacity() / 8 + + self.remaining_row_indices.allocated_size() + } } impl GroupValuesColumn { @@ -1196,8 +1204,25 @@ impl GroupValues for GroupValuesColumn { } fn size(&self) -> usize { - let group_values_size: usize = self.group_values.iter().map(|v| v.size()).sum(); - group_values_size + self.map_size + self.hashes_buffer.allocated_size() + let group_values_size = self.group_values.allocated_size() + + self + .group_values + .iter() + .map(|value| value.size()) + .sum::(); + let group_index_lists_size = self.group_index_lists.allocated_size() + + self + .group_index_lists + .iter() + .map(VecAllocExt::allocated_size) + .sum::(); + size_of::() + + group_values_size + + self.map_size + + group_index_lists_size + + self.emit_group_index_list_buffer.allocated_size() + + self.vectorized_operation_buffers.size() + + self.hashes_buffer.allocated_size() } fn is_empty(&self) -> bool { @@ -1385,7 +1410,10 @@ mod tests { compute::{concat_batches, take}, util::pretty::pretty_format_batches, }; - use datafusion_common::utils::proxy::HashTableAllocExt; + use datafusion_common::{ + Result, + utils::proxy::{HashTableAllocExt, VecAllocExt}, + }; use datafusion_expr::{EmitTo, GroupSelection}; use crate::aggregates::group_values::{ @@ -1396,6 +1424,146 @@ mod tests { GroupIndexView, group_column_supported_type, make_group_column, supported_schema, }; + fn expected_size(group_values: &GroupValuesColumn) -> usize { + let buffers = &group_values.vectorized_operation_buffers; + size_of::>() + + group_values.group_values.allocated_size() + + group_values + .group_values + .iter() + .map(|value| value.size()) + .sum::() + + group_values.map_size + + group_values.hashes_buffer.allocated_size() + + group_values.group_index_lists.allocated_size() + + group_values + .group_index_lists + .iter() + .map(VecAllocExt::allocated_size) + .sum::() + + group_values.emit_group_index_list_buffer.allocated_size() + + buffers.append_row_indices.allocated_size() + + buffers.equal_to_row_indices.allocated_size() + + buffers.equal_to_group_indices.allocated_size() + + buffers.equal_to_results.capacity() / 8 + + buffers.remaining_row_indices.allocated_size() + } + + #[test] + fn size_includes_boxed_primitive_and_row_backed_owners() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("primitive", DataType::Int32, false), + Field::new( + "nested", + DataType::Struct( + vec![Arc::new(Field::new("child", DataType::Int32, false))].into(), + ), + false, + ), + ])); + let group_values = GroupValuesColumn::::try_new(schema)?; + + // This schema builds a primitive column and a `RowsGroupColumn`. + assert_eq!( + group_values.size(), + size_of::>() + + group_values.group_values.allocated_size() + + group_values + .group_values + .iter() + .map(|value| value.size()) + .sum::() + ); + Ok(()) + } + + #[test] + fn size_includes_collision_emit_and_vectorized_buffers() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + DataType::Int32, + false, + )])); + let mut group_values = GroupValuesColumn::::try_new(schema).unwrap(); + let input: ArrayRef = Arc::new(Int32Array::from_iter_values(0..32)); + + assert_eq!(group_values.size(), expected_size(&group_values)); + group_values.intern(&[Arc::clone(&input)], &mut vec![])?; + group_values.intern(&[input], &mut vec![])?; + assert!( + group_values + .vectorized_operation_buffers + .append_row_indices + .capacity() + > 0 + ); + assert!( + group_values + .vectorized_operation_buffers + .equal_to_row_indices + .capacity() + > 0 + ); + assert!( + group_values + .vectorized_operation_buffers + .equal_to_results + .capacity() + / 8 + > 0 + ); + assert_eq!(group_values.size(), expected_size(&group_values)); + + insert_non_inline_group_index_view(&mut group_values, u64::MAX, vec![1, 2]); + group_values.emit(EmitTo::First(1))?; + assert!(!group_values.group_index_lists.is_empty()); + assert!(group_values.emit_group_index_list_buffer.capacity() > 0); + assert_eq!(group_values.size(), expected_size(&group_values)); + + let input: ArrayRef = Arc::new(Int32Array::from_iter_values(0..32)); + group_values.intern(&[input], &mut vec![])?; + assert_eq!(group_values.size(), expected_size(&group_values)); + Ok(()) + } + + #[test] + fn size_retains_vectorized_and_emit_scratch_capacity() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + DataType::Int32, + false, + )])); + let mut group_values = GroupValuesColumn::::try_new(schema)?; + let baseline = group_values.size(); + + let scratch_size = { + let buffers = &mut group_values.vectorized_operation_buffers; + buffers.append_row_indices.push(0); + buffers.equal_to_row_indices.push(0); + buffers.equal_to_group_indices.push(0); + buffers.equal_to_results.append(true); + buffers.remaining_row_indices.push(0); + group_values.emit_group_index_list_buffer.push(0); + + buffers.append_row_indices.allocated_size() + + buffers.equal_to_row_indices.allocated_size() + + buffers.equal_to_group_indices.allocated_size() + + buffers.equal_to_results.capacity() / 8 + + buffers.remaining_row_indices.allocated_size() + + group_values.emit_group_index_list_buffer.allocated_size() + }; + assert_eq!(group_values.size(), baseline + scratch_size); + + group_values.vectorized_operation_buffers.clear(); + group_values + .vectorized_operation_buffers + .equal_to_results + .truncate(0); + group_values.emit_group_index_list_buffer.clear(); + assert_eq!(group_values.size(), baseline + scratch_size); + Ok(()) + } + /// A mixed group-by key of several native columns plus one nested column /// that has no type-specialized `GroupColumn`. /// diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index 441d4b489dc3b..621ded222db3b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -33,6 +33,7 @@ use datafusion_common::utils::split_vec_min_alloc; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::GroupSelection; use std::iter; +use std::mem::size_of; use std::sync::Arc; /// An implementation of [`GroupColumn`] for primitive values @@ -259,7 +260,9 @@ where } fn size(&self) -> usize { - self.group_values.allocated_size() + self.nulls.allocated_size() + size_of::() + + self.group_values.allocated_size() + + self.nulls.allocated_size() } fn build(self: Box) -> ArrayRef { @@ -309,7 +312,7 @@ where #[cfg(test)] mod tests { - use std::sync::Arc; + use std::{mem::size_of, sync::Arc}; use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; use arrow::array::{ @@ -330,6 +333,16 @@ mod tests { (0..buf.len()).map(|i| buf.get_bit(i)).collect() } + #[test] + fn size_includes_boxed_owner_descriptor() { + let builder = + PrimitiveGroupValueBuilder::::new(DataType::Int32); + assert_eq!( + builder.size(), + size_of::>() + ); + } + #[test] fn test_nullable_primitive_equal_to() { let append = |builder: &mut PrimitiveGroupValueBuilder, diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index 2fa01e1ae24cd..9f8f397110f73 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -54,6 +54,7 @@ use arrow::datatypes::DataType; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::{DataFusionError, Result}; use datafusion_expr::GroupSelection; +use std::mem::size_of; /// A [`GroupColumn`] that stores group values for a single column in the arrow /// [row format], backed by a single-field [`RowConverter`]. @@ -288,7 +289,9 @@ impl GroupColumn for RowsGroupColumn { } fn size(&self) -> usize { - self.row_converter.size() + self.group_values.size() + size_of::() + self.row_converter.size() - size_of::() + + self.group_values.size() + - size_of::() } fn build(self: Box) -> ArrayRef { @@ -365,6 +368,18 @@ mod tests { )) } + #[test] + fn size_includes_boxed_owner_descriptor() { + let column = RowsGroupColumn::try_new(DataType::Int32).unwrap(); + assert_eq!( + column.size(), + size_of::() + column.row_converter.size() + - size_of::() + + column.group_values.size() + - size_of::() + ); + } + /// The generic column must agree with a per-row reference for equality, /// including inner-null and outer-null rows, on a `FixedSizeList`. #[test] diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 01e9f3eaa71ee..856774aa4aea9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -185,11 +185,17 @@ impl GroupValues for GroupValuesRows { } fn size(&self) -> usize { - let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0); - self.row_converter.size() + let group_values_size = self + .group_values + .as_ref() + .map(|values| values.size() - size_of::()) + .unwrap_or_default(); + // `size_of::()` already accounts for these inline descriptors. + size_of::() + self.row_converter.size() - size_of::() + group_values_size + self.map_size + self.rows_buffer.size() + - size_of::() + self.hashes_buffer.allocated_size() } @@ -443,9 +449,80 @@ pub(crate) fn encode_array_if_necessary( #[cfg(test)] mod tests { use super::*; - use arrow::array::{AsArray, ListArray}; + use arrow::array::{AsArray, Int32Array, ListArray}; use arrow::datatypes::{Field, Int32Type, Schema}; + fn expected_size(group_values: &GroupValuesRows) -> usize { + size_of::() + group_values.row_converter.size() + - size_of::() + + group_values + .group_values + .as_ref() + .map(|values| values.size() - size_of::()) + .unwrap_or_default() + + group_values.map_size + + group_values.rows_buffer.size() + - size_of::() + + group_values.hashes_buffer.allocated_size() + } + + fn int32_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new( + "group", + DataType::Int32, + true, + )])) + } + + #[test] + fn size_includes_owner_and_retained_allocations() -> Result<()> { + let mut group_values = GroupValuesRows::try_new(int32_schema())?; + + assert_eq!(group_values.size(), expected_size(&group_values)); + + let input: ArrayRef = Arc::new(Int32Array::from_iter_values(0..256)); + group_values.intern(&[input], &mut vec![])?; + assert_eq!(group_values.size(), expected_size(&group_values)); + Ok(()) + } + + #[test] + fn size_retains_reusable_buffers_after_emit() -> Result<()> { + let mut group_values = GroupValuesRows::try_new(int32_schema())?; + let input: ArrayRef = Arc::new(Int32Array::from_iter_values(0..256)); + group_values.intern(&[input], &mut vec![])?; + + let rows_buffer_size = group_values.rows_buffer.size() - size_of::(); + let hashes_buffer_size = group_values.hashes_buffer.allocated_size(); + + let output = group_values.emit(EmitTo::First(1))?; + assert_eq!(output[0].len(), 1); + assert_eq!(group_values.len(), 255); + assert_eq!( + group_values.rows_buffer.size() - size_of::(), + rows_buffer_size + ); + assert_eq!( + group_values.hashes_buffer.allocated_size(), + hashes_buffer_size + ); + assert_eq!(group_values.size(), expected_size(&group_values)); + + let input: ArrayRef = Arc::new(Int32Array::from_iter_values([256])); + group_values.intern(&[input], &mut vec![])?; + assert_eq!(group_values.len(), 256); + assert_eq!( + group_values.rows_buffer.size() - size_of::(), + rows_buffer_size + ); + assert_eq!( + group_values.hashes_buffer.allocated_size(), + hashes_buffer_size + ); + assert_eq!(group_values.size(), expected_size(&group_values)); + Ok(()) + } + #[test] fn preserving_nested_row_values() -> Result<()> { let field = Arc::new(Field::new_list_field(DataType::Int32, true)); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 21b62457e3831..f7b71452be38c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -178,7 +178,9 @@ where } fn size(&self) -> usize { - self.map.capacity() * size_of::<(usize, u64)>() + self.values.allocated_size() + size_of::() + + self.map.capacity() * size_of::<(usize, u64)>() + + self.values.allocated_size() } fn is_empty(&self) -> bool { @@ -285,6 +287,23 @@ mod tests { use datafusion_expr::EmitTo; use std::sync::Arc; + #[test] + fn size_includes_owner_and_retained_allocations() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + let expected_size = |gv: &GroupValuesPrimitive| { + size_of::>() + + gv.map.capacity() * size_of::<(usize, u64)>() + + gv.values.allocated_size() + }; + + assert_eq!(gv.size(), expected_size(&gv)); + + let input: ArrayRef = Arc::new(Int32Array::from_iter_values(0..256)); + gv.intern(&[input], &mut vec![])?; + assert_eq!(gv.size(), expected_size(&gv)); + Ok(()) + } + /// Mirror of the `EmitTo::take_needed` regression test, applied to the /// concrete `GroupValuesPrimitive` accumulator. /// diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 152082d88dd90..15d60db2f24f4 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -151,9 +151,6 @@ pub(crate) struct PartialHashAggregateStream { /// Input batches containing raw rows, not partial aggregate state. input: SendableRecordBatchStream, - /// Target output batch size from configuration. - batch_size: usize, - /// Memory reservation for group keys and accumulators. reservation: MemoryReservation, @@ -434,7 +431,6 @@ impl PartialHashAggregateStream { Ok(Self { schema, input, - batch_size, baseline_metrics, reservation, reduction_factor, @@ -476,20 +472,16 @@ impl PartialHashAggregateStream { break; } HandleInputResult::OOM => { - let materialized_group_states = hash_table.take_state_batch()?.ok_or_else(|| { - internal_datafusion_err!( + if hash_table.building_group_count() == 0 { + return Err(internal_datafusion_err!( "Partial hash aggregate ran out of memory with no aggregated groups" - ) - })?; + )); + } self.early_emit_count.add(1); timer.done(); - self.emit_on_memory_pressure( - materialized_group_states, - &mut emitter, - hash_table.memory_size(), - ) - .await?; + self.emit_on_memory_pressure(&mut hash_table, &mut emitter) + .await?; } } } @@ -588,74 +580,41 @@ impl PartialHashAggregateStream { } } - /// emit a materialized partial-state on memory pressure - /// batch in `batch_size`(from configuration) slices + /// Drain partial aggregate states in bounded output batches after memory + /// pressure. Each batch is removed from the table before the next one is + /// materialized, so this path never requires the complete state to fit. async fn emit_on_memory_pressure( &mut self, - // After each incremental emitting step, the `remaining_groups` will be updated - // with batch slicing. - mut remaining_groups: RecordBatch, + hash_table: &mut AggregateHashTable, emitter: &mut TryEmitter, - hash_table_mem_size: usize, ) -> Result<()> { - let remaining_groups_memory = remaining_groups.get_array_memory_size(); - - // Emitting clears the aggregate table and releases its - // accumulated memory. Update the reservation accordingly. - // We account here for the remaining groups memory to see if we can return batch size states - // if there is not enough memory, fallback to emit large batch - match self - .reservation - .try_resize(hash_table_mem_size + remaining_groups_memory) - { - Ok(_) => { - // Continue with slicing - } - Err(DataFusionError::ResourcesExhausted(_)) => { - // Fail to reserve memory for the hash table + state batch while slicing so emit a huge batch - - // Try resize without holding the state batch, if it fails there is nothing we can do - self.reservation.try_resize(hash_table_mem_size)?; - - self.reduction_factor.add_part(remaining_groups.num_rows()); - emitter - .emit(remaining_groups.record_output(&self.baseline_metrics)) - .await; + hash_table.start_early_emit(); + loop { + let batch = hash_table.next_early_emit_batch()?.ok_or_else(|| { + internal_datafusion_err!( + "Partial hash aggregate exhausted early-emission state unexpectedly" + ) + })?; - return Ok(()); + self.reduction_factor.add_part(batch.num_rows()); + // The reservation may already be above the pool limit that caused + // early emission. As with terminal output, make progress by + // releasing table state rather than requiring this output batch to + // fit alongside all remaining groups. A failed resize is expected + // until enough groups have been released; no materialized output + // batch is retained across the next iteration. + match self.reservation.try_resize(hash_table.memory_size()) { + Ok(()) | Err(DataFusionError::ResourcesExhausted(_)) => {} + Err(error) => return Err(error), } - Err(e) => return Err(e), - } - - while remaining_groups.num_rows() > self.batch_size { - // More batch to output, continue in the current state. - let output = remaining_groups.slice(0, self.batch_size); - - remaining_groups = remaining_groups.slice( - self.batch_size, - remaining_groups.num_rows() - self.batch_size, - ); - - self.reduction_factor.add_part(output.num_rows()); - debug_assert!(output.num_rows() > 0); - emitter - .emit(output.record_output(&self.baseline_metrics)) + .emit(batch.record_output(&self.baseline_metrics)) .await; - } - self.reduction_factor.add_part(remaining_groups.num_rows()); - debug_assert!(remaining_groups.num_rows() > 0); - - // We are no longer holding on the batch while slicing, so release the memory. - // The memory will now equal to the hash table size - self.reservation.try_shrink(remaining_groups_memory)?; - - emitter - .emit(remaining_groups.record_output(&self.baseline_metrics)) - .await; - - Ok(()) + if hash_table.is_building() { + return Ok(()); + } + } } /// emit partial aggregate state batches. @@ -1411,14 +1370,10 @@ mod tests { } #[tokio::test] - async fn test_partial_hash_stream_accounts_held_batch_on_memory_pressure_while_slicing() + async fn test_partial_hash_stream_incrementally_emits_on_memory_pressure() -> Result<()> { - // When memory pressure triggers early emission, the materialized state - // batch is held while it is sliced into `batch_size` outputs. The - // stream must keep that held batch accounted for in its memory - // reservation until the last slice is emitted; before the fix the - // reservation was resized down to just the (emptied) hash table size, - // leaving the held batch unaccounted. + // Early emission materializes each output batch directly. It must not + // retain one complete state batch while slicing it. let batch_size = 1024; // One row per group so the state batch is emitted in 4 slices @@ -1427,11 +1382,10 @@ mod tests { // Smaller than the building hash table (so pressure triggers) but large // enough to hold the materialized state batch (so slicing can proceed) let memory_limit = 100 * 1024; - let (mut stream, input, runtime) = + let (mut stream, input, _runtime) = partial_stream_under_memory_limit(memory_limit, batch_size, num_groups)?; - // The first output batch must be a pressure-emitted slice, with the rest - // of the materialized state batch still held by the stream + // The first output batch must be pressure-emitted. let first = tokio::time::timeout(Duration::from_secs(5), stream.next()) .await .expect( @@ -1441,26 +1395,12 @@ mod tests { .expect("stream ended early")?; assert_eq!(first.num_rows(), batch_size); - // The emitted slice shares buffers with the held state batch, so its - // array memory size reflects the full held allocation - let held_size = first.get_array_memory_size(); - let reserved = runtime.memory_pool.reserved(); - assert!( - reserved >= held_size, - "memory pool has {reserved} bytes reserved but the stream is \ - holding a materialized state batch of {held_size} bytes" - ); - let second = stream.next().await.expect("stream ended early")?; assert_eq!(second.num_rows(), batch_size); - // Make sure the state batch is really being sliced (and not emitted whole by the fallback path): - // the second output must share the same underlying buffer as the first - // - // If you changed the code and this fail because - // - you now deep copy `batch_size` from the full state batch, please update this assertion to something else - // - you only take batch size from the hash table, you can remove the test - assert_eq!( + // Each output is independently materialized from the table, rather + // than a slice of one retained complete state batch. + assert_ne!( first .column(0) .as_primitive::() @@ -1488,22 +1428,17 @@ mod tests { } #[tokio::test] - async fn test_partial_hash_stream_emits_whole_batch_when_held_batch_does_not_fit() + async fn test_partial_hash_stream_emits_bounded_batch_when_full_state_does_not_fit() -> Result<()> { - // When memory pressure triggers early emission but the materialized - // state batch itself does not fit in the reservation, the stream must - // not fail with a resources exhausted error. Instead it gives up on - // slicing and emits the whole state batch at once. + // When the complete state does not fit, early emission must still + // produce a bounded batch rather than failing or emitting all states. let batch_size = 1024; let num_groups = 4 * batch_size; - // Smaller than the materialized state batch (4096 rows of Int32 group - // keys plus Int64 counts is at least 48 KiB), so the reservation for - // hash table held batch fails. The emptied hash table itself is tiny - // and still fits. + // Smaller than the full 4096-row state batch. let memory_limit = 32 * 1024; - let (mut stream, input, runtime) = + let (mut stream, input, _runtime) = partial_stream_under_memory_limit(memory_limit, batch_size, num_groups)?; let first = tokio::time::timeout(Duration::from_secs(5), stream.next()) @@ -1514,23 +1449,10 @@ mod tests { ) .expect("stream ended early")?; - // The whole state batch is emitted at once instead of `batch_size` slices - assert_eq!(first.num_rows(), num_groups); - assert!( - first.get_array_memory_size() > memory_limit, - "test setup is wrong: the state batch fits within the memory limit, \ - so the slicing path would have been taken" - ); - - // Unlike the slicing path, the stream does not hold on to the emitted - // batch, so it must not be accounted for in the reservation. Only the - // (emptied) hash table remains reserved - let emitted_size = first.get_array_memory_size(); - let reserved = runtime.memory_pool.reserved(); + assert_eq!(first.num_rows(), batch_size); assert!( - reserved < emitted_size, - "memory pool has {reserved} bytes reserved but the stream no longer \ - holds the emitted state batch of {emitted_size} bytes" + first.get_array_memory_size() < memory_limit, + "test setup is wrong: one output batch does not fit within the memory limit" ); input.wait_finish().await; @@ -1544,23 +1466,20 @@ mod tests { } #[tokio::test] - async fn test_partial_hash_stream_releases_held_batch_after_last_slice() -> Result<()> - { - // While the pressure-emitted state batch is sliced, the stream holds - // the remaining groups and keeps them reserved. Once the last slice is - // handed out nothing is held anymore, so the reservation must drop - // back to just the (emptied) hash table before the input is resumed. + async fn test_partial_hash_stream_releases_groups_after_incremental_emit() + -> Result<()> { + // Draining all early-emitted batches must release groups and resume + // input without retaining a complete materialized state batch. let batch_size = 1024; let num_slices = 4; let num_groups = num_slices * batch_size; let memory_limit = 100 * 1024; - let (mut stream, input, runtime) = + let (mut stream, input, _runtime) = partial_stream_under_memory_limit(memory_limit, batch_size, num_groups)?; - // The input has not finished, so all of these are pressure-emitted slices - let mut held_size = 0; + // The input has not finished, so all are pressure-emitted batches. for slice_idx in 0..num_slices { let slice = if slice_idx == 0 { tokio::time::timeout(Duration::from_secs(5), stream.next()) @@ -1576,27 +1495,8 @@ mod tests { assert_eq!(slice.num_rows(), batch_size); - // Every slice shares buffers with the held state batch, so this is - // the size of the full held allocation - held_size = slice.get_array_memory_size(); - let reserved = runtime.memory_pool.reserved(); - - if slice_idx + 1 < num_slices { - assert!( - reserved >= held_size, - "after slice {slice_idx} the stream still holds {held_size} \ - bytes but only {reserved} bytes are reserved" - ); - } else { - assert!( - reserved < held_size, - "after the last slice nothing is held anymore but {reserved} \ - bytes are still reserved (held batch was {held_size} bytes)" - ); - } + assert!(slice.get_array_memory_size() < memory_limit); } - assert!(held_size > 0); - input.wait_finish().await; let mut total_rows = num_groups; while let Some(batch) = stream.next().await { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3ed93e09ce4f4..c283b7fead587 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3502,7 +3502,8 @@ mod tests { )]; let task_ctx = if spill { - // adjust the max memory size to have the partial aggregate result for spill mode. + // Smaller than the complete grouping-set state. Partial early + // emission must materialize and release state incrementally. new_spill_ctx(4, 500) } else { Arc::new(TaskContext::default()) @@ -3521,6 +3522,14 @@ mod tests { collect(partial_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; if spill { + let early_emit_count = partial_aggregate + .metrics() + .unwrap() + .sum_by_name("early_emit_count") + .unwrap() + .as_usize(); + assert!(early_emit_count > 0); + // In spill mode, we test with the limited memory, if the mem usage exceeds, // we trigger the early emit rule, which turns out the partial aggregate result. allow_duplicates! { @@ -7496,7 +7505,7 @@ mod tests { Arc::clone(&schema), )?); - let task_ctx = new_migrated_spill_ctx(1, 600); + let task_ctx = new_migrated_spill_ctx(1, 1_024); let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await?; assert_spill_count_metric(true, Arc::clone(&aggr)); let metrics = aggr.metrics().unwrap(); diff --git a/datafusion/sqllogictest/test_files/explain_analyze.slt b/datafusion/sqllogictest/test_files/explain_analyze.slt index 81b274d6c1377..2d04ac20f66da 100644 --- a/datafusion/sqllogictest/test_files/explain_analyze.slt +++ b/datafusion/sqllogictest/test_files/explain_analyze.slt @@ -500,7 +500,7 @@ GROUP BY k; ---- Plan with Metrics 01)ProjectionExec: expr=[k@0 as k, count(Int64(1))@1 as count(*)], metrics=[output_bytes=1056.0 B] -02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=9.2 KB] +02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=9.4 KB] 03)----ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]