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/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 397b766f41697..2e5f92ed2dda6 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 @@ -18,13 +18,12 @@ use std::marker::PhantomData; use std::sync::Arc; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion_common::{Result, assert_eq_or_internal_err}; - use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{AggregateExec, evaluate_group_by}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 32eabfe8dc26c..19d3b94f8a72a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -31,12 +31,12 @@ mod single_group_by; use datafusion_physical_expr::binary_map::OutputType; use multi_group_by::{GroupValuesColumn, GroupValuesOrdered}; -pub(crate) use single_group_by::primitive::HashValue; +pub(crate) use single_group_by::primitive::{GroupValuesPrimitive, HashValue}; use crate::aggregates::{ group_values::single_group_by::{ boolean::GroupValuesBoolean, bytes::GroupValuesBytes, - bytes_view::GroupValuesBytesView, primitive::GroupValuesPrimitive, + bytes_view::GroupValuesBytesView, }, order::GroupOrdering, }; 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 da656e0a85c94..530e4eb407eef 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 3032536657803..7fb4f26c6a71c 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/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index b9547079aaf40..c58ce2a948e3e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -30,7 +30,7 @@ use datafusion_expr::GroupSelection; use hashbrown::{HashMap, hash_table::HashTable}; use std::marker::PhantomData; use std::mem::size_of; -use std::sync::Arc; +use std::sync::{Arc, Weak}; use crate::aggregates::AGGREGATION_HASH_SEED; @@ -63,10 +63,10 @@ pub struct DictionaryGroupValuesColumn val_to_inner: Vec, /// Hashes of the values array in `cached_values`, one per `val_idx`. val_hashes: Vec, - /// The values array that `val_hashes` and `val_to_inner` were built for. - /// When an incoming `dict.values()` is `Arc::ptr_eq` to this, both caches - /// are reused instead of being rebuilt. See [`Self::sync_value_cache`]. - cached_values: Option, + /// Weak reference to the values array that `val_hashes` and `val_to_inner` + /// were built for. A matching live array reuses both caches without + /// retaining an input dictionary after its batch is released. + cached_values: Option>, _phantom: PhantomData, } @@ -181,6 +181,12 @@ impl DictionaryGroupValuesColumn { .unwrap(); } + fn is_cached(&self, values: &ArrayRef) -> bool { + self.cached_values + .as_ref() + .is_some_and(|cached| std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values))) + } + /// Makes `val_hashes` and `val_to_inner` correspond to `dict_values`, /// reusing them when they were already built for this same values array. /// @@ -190,17 +196,13 @@ impl DictionaryGroupValuesColumn { /// rather than O(rows). `take_n` remaps `inner` and clears /// `cached_values`, forcing a rebuild. fn sync_value_cache(&mut self, dict_values: &ArrayRef) { - if self - .cached_values - .as_ref() - .is_some_and(|c| Arc::ptr_eq(c, dict_values)) - { + if self.is_cached(dict_values) { return; } self.hash_values(dict_values); self.val_to_inner.clear(); self.val_to_inner.resize(dict_values.len(), usize::MAX); - self.cached_values = Some(Arc::clone(dict_values)); + self.cached_values = Some(Arc::downgrade(dict_values)); } fn find_or_insert_value( @@ -360,10 +362,7 @@ impl GroupColumn let dict_keys = dict.keys(); let dict_values = dict.values(); - let cached = self - .cached_values - .as_ref() - .is_some_and(|c| Arc::ptr_eq(c, dict_values)); + let cached = self.is_cached(dict_values); if !cached && rhs_rows.len() < dict_values.len() { self.equal_to_per_row( lhs_rows, @@ -713,6 +712,29 @@ mod tests { (0..buf.len()).map(|i| buf.get_bit(i)).collect() } + #[test] + fn hash_cache_does_not_retain_dictionary_values() { + let mut column = utf8_col(); + let input = i32_dict(&[Some(0)], &[Some("retained")]); + let values = Arc::clone(input.as_dictionary::().values()); + let weak_values = Arc::downgrade(&values); + column.append_val(&input, 0).unwrap(); + assert!( + column + .cached_values + .as_ref() + .is_some_and(|cached| cached.ptr_eq(&weak_values)) + ); + let size_with_input = column.size(); + + // The cache is only an identity hint: it must not keep an input + // dictionary allocation alive after its batch is released. + drop(values); + drop(input); + assert!(weak_values.upgrade().is_none()); + assert_eq!(column.size(), size_with_input); + } + fn all_true(len: usize) -> BooleanBufferBuilder { let mut buf = BooleanBufferBuilder::new(len); buf.append_n(len, true); @@ -993,11 +1015,7 @@ mod tests { usize::MAX, "unreferenced value resolved" ); - assert!( - col.cached_values - .as_ref() - .is_some_and(|c| Arc::ptr_eq(c, &values)) - ); + assert!(col.is_cached(&values)); assert_eq!(col.inner.len(), 2); } @@ -1033,10 +1051,7 @@ mod tests { col.vectorized_append(&second, &[0]).unwrap(); assert!( - col.cached_values.as_ref().is_some_and(|c| Arc::ptr_eq( - c, - second.as_dictionary::().values() - )), + col.is_cached(second.as_dictionary::().values()), "cache should track the most recent values array" ); assert_eq!(col.inner.len(), 1, "'a' must dedup across both arrays"); @@ -1081,11 +1096,7 @@ mod tests { let mut col = utf8_col(); col.append_val(&batch, 0).unwrap(); - assert!( - col.cached_values - .as_ref() - .is_some_and(|c| Arc::ptr_eq(c, &values)) - ); + assert!(col.is_cached(&values)); col.vectorized_append(&batch, &[0, 1]).unwrap(); assert_eq!(col.inner.len(), 2, "'a' must not be duplicated"); @@ -1182,9 +1193,7 @@ mod tests { col.vectorized_equal_to(&[group_a], &only_a, &[0], &mut buf); assert_eq!(bool_vec(&buf), vec![true]); assert!( - col.cached_values - .as_ref() - .is_some_and(|c| Arc::ptr_eq(c, &old_values)), + col.is_cached(&old_values), "fallback must leave the cache on the previous values array" ); } 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 bd59e564c717d..521bc8e2f005a 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 b5dbdb38763f4..b61e1987d7ee2 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 9497d521e6321..68d56f00a9de5 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 @@ -103,7 +103,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 @@ -273,6 +273,23 @@ impl VectorizedOperationBuffers { self.equal_to_group_indices.clear(); self.remaining_row_indices.clear(); } + + fn clear_shrink(&mut self, num_rows: usize) { + self.clear(); + self.append_row_indices.shrink_to(num_rows); + self.equal_to_row_indices.shrink_to(num_rows); + self.equal_to_group_indices.shrink_to(num_rows); + self.equal_to_results = BooleanBufferBuilder::new(num_rows); + self.remaining_row_indices.shrink_to(num_rows); + } + + 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 { @@ -1067,8 +1084,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 { @@ -1217,8 +1251,10 @@ impl GroupValues for GroupValuesColumn { // Such structures are only used in `non-streaming` case if !STREAMING { self.group_index_lists.clear(); + self.group_index_lists.shrink_to(num_rows); self.emit_group_index_list_buffer.clear(); - self.vectorized_operation_buffers.clear(); + self.emit_group_index_list_buffer.shrink_to(num_rows); + self.vectorized_operation_buffers.clear_shrink(num_rows); } } } @@ -1258,7 +1294,7 @@ mod tests { compute::{concat_batches, take}, util::pretty::pretty_format_batches, }; - use datafusion_common::utils::proxy::HashTableAllocExt; + use datafusion_common::{Result, utils::proxy::HashTableAllocExt}; use datafusion_expr::{EmitTo, GroupSelection}; use crate::aggregates::group_values::{ @@ -1269,6 +1305,79 @@ mod tests { GroupIndexView, group_column_supported_type, make_group_column, supported_schema, }; + #[test] + fn collision_emit_and_vectorized_buffers_retain_capacity() -> 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)); + let baseline = group_values.size(); + + 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!(group_values.size() > baseline); + + 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!(group_values.size() > baseline); + + let input: ArrayRef = Arc::new(Int32Array::from_iter_values(0..32)); + group_values.intern(&[input], &mut vec![])?; + assert!(group_values.size() > baseline); + Ok(()) + } + + #[test] + fn clear_shrink_releases_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 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); + + assert!(group_values.size() > baseline); + group_values.clear_shrink(0); + + assert_eq!(group_values.size(), baseline); + 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 fe160b58dbe15..ebd1e071cf8d5 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 { 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 f488ef79dff06..c45942edac6ff 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 { diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 9604087c33334..8fb7652367ba8 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -180,13 +180,18 @@ impl GroupValues for GroupValuesRows { } fn size(&self) -> usize { - let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0); - // `HashTable::allocation_size` reports the complete hashbrown allocation, - // which includes the entry array, control bytes, and trailing group layout. - 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 inline row descriptors. + // `HashTable::allocation_size` includes entries, control bytes, and group layout. + size_of::() + self.row_converter.size() - size_of::() + group_values_size + self.map.allocation_size() + self.rows_buffer.size() + - size_of::() + self.hashes_buffer.allocated_size() } @@ -445,6 +450,65 @@ mod tests { use arrow::datatypes::{Field, Int32Type, Schema}; use std::mem::size_of; + 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.allocation_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_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); + let size_after_emit = group_values.size(); + assert_eq!( + group_values.rows_buffer.size() - size_of::(), + rows_buffer_size + ); + assert_eq!( + group_values.hashes_buffer.allocated_size(), + hashes_buffer_size + ); + + let input: ArrayRef = Arc::new(Int32Array::from_iter_values([256])); + group_values.intern(&[input], &mut vec![])?; + assert_eq!(group_values.len(), 256); + assert!(group_values.size() >= size_after_emit); + assert_eq!( + group_values.rows_buffer.size() - size_of::(), + rows_buffer_size + ); + assert_eq!( + group_values.hashes_buffer.allocated_size(), + hashes_buffer_size + ); + Ok(()) + } + #[test] fn preserving_nested_row_values() -> Result<()> { let field = Arc::new(Field::new_list_field(DataType::Int32, true)); @@ -495,12 +559,7 @@ mod tests { // 1. Initial state: map is unallocated assert_eq!(group_values.map.capacity(), 0); assert_eq!(group_values.map.allocation_size(), 0); - assert_eq!( - group_values.size(), - group_values.row_converter.size() - + group_values.rows_buffer.size() - + group_values.hashes_buffer.allocated_size() - ); + assert_eq!(group_values.size(), expected_size(&group_values)); // 2. Insert multi-column keys and observe table allocation growth let num_distinct = 500; @@ -526,16 +585,7 @@ mod tests { ); // GroupValuesRows::size() should accurately account for the map's allocation_size - let expected_size = group_values.row_converter.size() - + group_values - .group_values - .as_ref() - .map(|v| v.size()) - .unwrap_or(0) - + table_bytes - + group_values.rows_buffer.size() - + group_values.hashes_buffer.allocated_size(); - assert_eq!(group_values.size(), expected_size); + assert_eq!(group_values.size(), expected_size(&group_values)); // 3. Partial emit (EmitTo::First) - table retains entries and capacity let emit_count = 200; @@ -543,20 +593,9 @@ mod tests { assert_eq!(emitted[0].len(), emit_count); assert_eq!(group_values.len(), (num_distinct as usize) - emit_count); - // Retained capacity is still charged until the table releases it + // Retained capacity is still charged until the table releases it. assert_eq!(group_values.map.allocation_size(), table_bytes); - assert_eq!( - group_values.size(), - group_values.row_converter.size() - + group_values - .group_values - .as_ref() - .map(|v| v.size()) - .unwrap_or(0) - + table_bytes - + group_values.rows_buffer.size() - + group_values.hashes_buffer.allocated_size() - ); + assert_eq!(group_values.size(), expected_size(&group_values)); // 4. Emit all (EmitTo::All) - map is cleared but capacity remains allocated for reuse let remaining_count = group_values.len(); @@ -574,18 +613,7 @@ mod tests { shrunken_table_bytes < table_bytes, "clear_shrink should shrink table from {table_bytes} to {shrunken_table_bytes}" ); - assert_eq!( - group_values.size(), - group_values.row_converter.size() - + group_values - .group_values - .as_ref() - .map(|v| v.size()) - .unwrap_or(0) - + shrunken_table_bytes - + group_values.rows_buffer.size() - + group_values.hashes_buffer.allocated_size() - ); + assert_eq!(group_values.size(), expected_size(&group_values)); // 6. Reinsert new distinct values into the shrunken table let new_col_a = @@ -597,18 +625,7 @@ mod tests { group_values.intern(&[new_col_a, new_col_b], &mut groups)?; assert_eq!(groups.len(), num_distinct as usize); assert!(group_values.map.allocation_size() > shrunken_table_bytes); - assert_eq!( - group_values.size(), - group_values.row_converter.size() - + group_values - .group_values - .as_ref() - .map(|v| v.size()) - .unwrap_or(0) - + group_values.map.allocation_size() - + group_values.rows_buffer.size() - + group_values.hashes_buffer.allocated_size() - ); + assert_eq!(group_values.size(), expected_size(&group_values)); Ok(()) } @@ -635,18 +652,7 @@ mod tests { let table_bytes = group_values.map.allocation_size(); assert!(table_bytes > 0); - assert_eq!( - group_values.size(), - group_values.row_converter.size() - + group_values - .group_values - .as_ref() - .map(|v| v.size()) - .unwrap_or(0) - + table_bytes - + group_values.rows_buffer.size() - + group_values.hashes_buffer.allocated_size() - ); + assert_eq!(group_values.size(), expected_size(&group_values)); // Shrink to 0: table should completely deallocate group_values.clear_shrink(0); 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..65b2c71e9e3e7 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 { diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index b9a6d2833ab7b..a8c35ff774d6c 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -63,7 +63,9 @@ use datafusion_common::{ }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool, +}; use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; @@ -369,6 +371,9 @@ pub(crate) struct GroupedHashAggregateStream { // EXECUTION RESOURCES: // Fields related to managing execution resources and monitoring performance. // ======================================================================== + /// The pool shared by aggregate state and spill replay workspace. + merge_pool: Arc, + /// The memory reservation for this grouping reservation: MemoryReservation, @@ -551,12 +556,16 @@ impl GroupedHashAggregateStream { }; let group_values = new_group_values(group_schema, &group_ordering)?; - let reservation = MemoryConsumer::new(name) - // We interpret 'can spill' as 'can handle memory back pressure'. - // This value needs to be set to true for the default memory pool implementations - // to ensure fair application of back pressure amongst the memory consumers. - .with_can_spill(oom_mode != OutOfMemoryMode::ReportError) - .register(context.memory_pool()); + let merge_pool = Arc::new(MergeMemoryPool::new( + Arc::clone(context.memory_pool()), + MemoryConsumer::new(name) + // We interpret 'can spill' as 'can handle memory back pressure'. + // This value needs to be set to true for the default memory pool implementations + // to ensure fair application of back pressure amongst the memory consumers. + .with_can_spill(oom_mode != OutOfMemoryMode::ReportError), + )); + let reservation = MemoryConsumer::new("GroupedHashAggregateStream state") + .register(&(Arc::clone(&merge_pool) as Arc)); timer.done(); let exec_state = ExecutionState::ReadingInput; @@ -631,6 +640,7 @@ impl GroupedHashAggregateStream { aggregate_arguments, filter_expressions, group_by: agg_group_by, + merge_pool, reservation, oom_mode, group_values, @@ -1352,6 +1362,7 @@ impl GroupedHashAggregateStream { .with_metrics(self.baseline_metrics.clone()) .with_batch_size(self.batch_size) .with_reservation(self.reservation.new_empty()) + .with_merge_pool(Arc::clone(&self.merge_pool)) .with_replay_headroom() .build()?; self.input_done = false; diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2f1bb9b638a4c..12babfe59cf9f 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -428,73 +428,48 @@ impl PartialHashAggregateStream { } } - /// emit a materialized partial-state on memory pressure - /// batch in `batch_size`(from configuration) slices + /// Emit a materialized partial-state batch in `batch_size` slices. 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, 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 - } + Ok(()) => {} 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; - return Ok(()); } - Err(e) => return Err(e), + Err(error) => return Err(error), } 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)) .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(()) } @@ -889,7 +864,9 @@ mod tests { use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion_common::Result; use datafusion_execution::config::SessionConfig; - use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, + }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::{min_max::min_udaf, sum::sum_udaf}; @@ -1278,53 +1255,23 @@ mod tests { #[tokio::test] async fn test_partial_hash_stream_accounts_held_batch_on_memory_pressure_while_slicing() -> 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. - let batch_size = 1024; - // One row per group so the state batch is emitted in 4 slices let num_groups = 4 * batch_size; - - // 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) = 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 let first = tokio::time::timeout(Duration::from_secs(5), stream.next()) .await - .expect( - "did not get early emit due to OOM, this probably means that the \ - memory limit is too high to trigger the OOM", - ) + .expect("early emit should occur") .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" - ); + assert!(runtime.memory_pool.reserved() >= held_size); 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!( first .column(0) @@ -1338,65 +1285,34 @@ mod tests { .values() .inner() .data_ptr(), - "both batches should be slices of the same materialized state batch" + "batches should slice one materialized state batch" ); - // Let the input finish and drain the stream: no groups lost input.wait_finish().await; let mut total_rows = first.num_rows() + second.num_rows(); while let Some(batch) = stream.next().await { total_rows += batch?.num_rows(); } assert_eq!(total_rows, num_groups); - Ok(()) } #[tokio::test] async fn test_partial_hash_stream_emits_whole_batch_when_held_batch_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. - 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. let memory_limit = 32 * 1024; 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()) .await - .expect( - "did not get early emit due to OOM, this probably means that the \ - memory limit is too high to trigger the OOM", - ) + .expect("early emit should occur") .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!( - reserved < emitted_size, - "memory pool has {reserved} bytes reserved but the stream no longer \ - holds the emitted state batch of {emitted_size} bytes" - ); + assert!(first.get_array_memory_size() > memory_limit); + assert!(runtime.memory_pool.reserved() < first.get_array_memory_size()); input.wait_finish().await; let mut total_rows = first.num_rows(); @@ -1404,60 +1320,35 @@ mod tests { total_rows += batch?.num_rows(); } assert_eq!(total_rows, num_groups); - Ok(()) } #[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. - 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) = 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; for slice_idx in 0..num_slices { let slice = if slice_idx == 0 { tokio::time::timeout(Duration::from_secs(5), stream.next()) .await - .expect( - "did not get early emit due to OOM, this probably means that the \ - memory limit is too high to trigger the OOM", - ) + .expect("early emit should occur") .expect("stream ended early")? } else { stream.next().await.expect("stream ended early")? }; - 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" - ); + assert!(runtime.memory_pool.reserved() >= held_size); } 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!(runtime.memory_pool.reserved() < held_size); } } assert!(held_size > 0); @@ -1468,7 +1359,6 @@ mod tests { total_rows += batch?.num_rows(); } assert_eq!(total_rows, num_groups); - Ok(()) } @@ -1482,20 +1372,21 @@ mod tests { #[tokio::test] async fn final_hash_spill_replay_with_other_partitions_holding_state() -> Result<()> { for spills in [1, 2, 3] { - run_shared_pool_case(spills, 1024 * 1024, Finish::Collect).await?; + run_shared_pool_case(spills, Some(1024 * 1024), Finish::Collect).await?; } - // An unlimited pool produces the reference results without spilling. - run_shared_pool_case(0, 10 * 1024 * 1024, Finish::Collect).await + // Collision-mode scratch can exceed a fixed reference budget. Use an + // unbounded pool so this case always exercises output without spilling. + run_shared_pool_case(0, None, Finish::Collect).await } #[tokio::test] async fn final_hash_spill_replay_releases_memory_on_drop() -> Result<()> { - run_shared_pool_case(1, 1024 * 1024, Finish::DropDuringReplay).await + run_shared_pool_case(1, Some(1024 * 1024), Finish::DropDuringReplay).await } #[tokio::test] async fn final_hash_spill_releases_memory_on_input_error() -> Result<()> { - run_shared_pool_case(1, 1024 * 1024, Finish::InputError).await + run_shared_pool_case(1, Some(1024 * 1024), Finish::InputError).await } /// Partition 0 spills `spills` times and replays while partitions 1..3 @@ -1511,7 +1402,7 @@ mod tests { /// issue #25423. async fn run_shared_pool_case( spills: usize, - limit: usize, + limit: Option, finish: Finish, ) -> Result<()> { const PARTITIONS: usize = 4; @@ -1569,7 +1460,10 @@ mod tests { )?; assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); - let pool: Arc = Arc::new(GreedyMemoryPool::new(limit)); + let pool: Arc = match limit { + Some(limit) => Arc::new(GreedyMemoryPool::new(limit)), + None => Arc::new(UnboundedMemoryPool::default()), + }; let context = Arc::new( TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(128)) @@ -1579,9 +1473,7 @@ mod tests { .build_arc()?, ), ); - let mut streams = vec![]; - let mut senders = vec![]; - for partition in 0..PARTITIONS { + let make_stream = |partition| -> Result<_> { let (sender, receiver) = mpsc::unbounded(); let input = Box::pin(RecordBatchStreamAdapter::new( Arc::clone(&partial_schema), @@ -1591,6 +1483,15 @@ mod tests { &aggregate, &context, partition, input, )? .into_stream(); + Ok((sender, stream)) + }; + + // Construct the holders first. Their reservation is the baseline that + // must survive partition 0's spill/replay lifecycle. + let mut streams = vec![]; + let mut senders = vec![]; + for partition in 1..PARTITIONS { + let (sender, stream) = make_stream(partition)?; senders.push(sender); streams.push(stream); } @@ -1634,25 +1535,27 @@ mod tests { }; // Channel inputs return Pending after each supplied batch, so the // interleaving below does not depend on task scheduling. - let mut feed = |partition: usize, batch: i64| { - senders[partition] - .unbounded_send(Ok(make_batch(partition as i64, batch * 128))) - .unwrap(); - assert!(streams[partition].next().now_or_never().is_none()); - }; - // Keep the state of partitions 1..3 live while partition 0 spills and - // replays. let mut held_batches = 0; while pool.reserved() < HELD_BYTES { assert!(held_batches < MAX_BATCHES, "held state stays small"); - for partition in 1..PARTITIONS { - feed(partition, held_batches); + for (partition, (sender, stream)) in + senders.iter().zip(streams.iter_mut()).enumerate() + { + let partition = partition + 1; + sender + .unbounded_send(Ok(make_batch(partition as i64, held_batches * 128))) + .unwrap(); + assert!(stream.next().now_or_never().is_none()); } held_batches += 1; } let held = pool.reserved(); let spill_count = || aggregate.metrics().unwrap().spill_count().unwrap(); assert_eq!(spill_count(), 0); + + // Do not include partition 0's empty table in `held`: it is correctly + // released when that stream spills, errors, or is dropped. + let (partition_zero_sender, mut first) = make_stream(0)?; // Feed partition 0 until it has spilled `spills` times. Key (0, 0) // repeats in every batch, so replay must merge its sum across runs. let mut batches = 0; @@ -1666,24 +1569,29 @@ mod tests { break; } assert!(batches < MAX_BATCHES, "partition 0 did not spill"); - feed(0, batches); + partition_zero_sender + .unbounded_send(Ok(make_batch(0, batches * 128))) + .unwrap(); + assert!(first.next().now_or_never().is_none()); batches += 1; } // Add groups after the last spill so replay also merges the final // in-memory run. for _ in 0..8 { - feed(0, batches); + partition_zero_sender + .unbounded_send(Ok(make_batch(0, batches * 128))) + .unwrap(); + assert!(first.next().now_or_never().is_none()); batches += 1; } assert!(spill_count() >= spills); assert_eq!(spill_count() == 0, spills == 0); - let mut first = streams.remove(0); match finish { Finish::Collect => { - senders[0].close_channel(); + partition_zero_sender.close_channel(); let mut output = collect(first).await?; assert_eq!(pool.reserved(), held); - for sender in &senders[1..] { + for sender in &senders { sender.close_channel(); } for stream in streams.drain(..) { @@ -1726,14 +1634,14 @@ mod tests { assert_eq!(spill_count() == 0, spills == 0); } Finish::DropDuringReplay => { - senders[0].close_channel(); + partition_zero_sender.close_channel(); first.next().await.unwrap()?; assert!(pool.reserved() > held); drop(first); assert_eq!(pool.reserved(), held); } Finish::InputError => { - senders[0] + partition_zero_sender .unbounded_send(datafusion_common::exec_err!( "injected input failure" )) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6552bd95e06a5..7560032fd088a 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3643,6 +3643,7 @@ pub fn evaluate_group_by( #[cfg(test)] mod tests { use std::collections::HashSet; + use std::mem::size_of; use std::task::{Context, Poll}; use super::*; @@ -3924,7 +3925,8 @@ mod tests { const KEYS: usize = 64; const VALUES_PER_KEY: i64 = 64; - const MEMORY_LIMIT: usize = 8192; + // Leave only a small amount above 8 KiB for the empty aggregate state. + const MEMORY_LIMIT: usize = 8_384; let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Int64, false), Field::new("value", DataType::Int64, false), @@ -4328,8 +4330,8 @@ mod tests { )]; let task_ctx = if spill { - // adjust the max memory size to have the partial aggregate result for spill mode. - new_spill_ctx(4, 500) + // Includes the descriptor-aware empty grouping-set table state. + new_spill_ctx(4, 700) } else { Arc::new(TaskContext::default()) }; @@ -4347,6 +4349,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! { @@ -7751,9 +7761,11 @@ mod tests { )?); let batch_size = 2; - let memory_pool = Arc::new(FairSpillPool::new( - initial_reservation(&single_aggregate)? + 200, - )); + // Allow the already-retained initial table but no extra headroom. + // This fixture's groups can fit in that initial capacity, so successful + // execution need not spill. + let memory_pool = + Arc::new(FairSpillPool::new(initial_reservation(&single_aggregate)?)); let task_ctx = Arc::new( TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(batch_size)) @@ -7767,8 +7779,6 @@ mod tests { let result = collect(single_aggregate.execute(0, Arc::clone(&task_ctx))?).await; match result { Ok(result) => { - assert_spill_count_metric(true, single_aggregate); - allow_duplicates! { assert_snapshot!(batches_to_string(&result), @r" +---+---+--------+--------+ diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 199dceac8c1a9..249b20cf49f38 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -159,7 +159,7 @@ impl<'a> StreamingMergeBuilder<'a> { } /// Keep spill workspace until the final merge pass selects its buffer budget. - pub(super) fn with_merge_pool(mut self, pool: Arc) -> Self { + pub(crate) fn with_merge_pool(mut self, pool: Arc) -> Self { self.merge_pool = Some(pool); self } diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 48622127bce99..3c57b7fade5d3 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -195,6 +195,10 @@ FROM ( # --- Case G: partial/final aggregation under memory limit --- +# Four final aggregate partitions share the runtime pool; leave replay headroom. +statement ok +SET datafusion.runtime.memory_limit = '1M' + statement ok SET datafusion.execution.target_partitions = 4 diff --git a/datafusion/sqllogictest/test_files/explain_analyze.slt b/datafusion/sqllogictest/test_files/explain_analyze.slt index 511ba88ed3345..89567a920de68 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=[]