From 5b8d63e22f25934e78a994435ef1961f1b4f2428 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 20:00:08 +0800 Subject: [PATCH 01/15] fix GroupValues*::size() - GroupValues::size() charges owner + retained buffers. - Column: added group-index, emit, vectorized, and vec backing. - Added capacity/reuse tests. - Updated tight spill test pools. --- .../group_values/multi_group_by/mod.rs | 108 +++++++++++++++++- .../src/aggregates/group_values/row.rs | 34 +++++- .../group_values/single_group_by/primitive.rs | 21 +++- .../physical-plan/src/aggregates/mod.rs | 4 +- 4 files changed, 159 insertions(+), 8 deletions(-) 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..f3df6a5242fb5 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 @@ -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,80 @@ 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_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(()) + } + /// 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/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 01e9f3eaa71ee..222bba74495be 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -186,7 +186,8 @@ 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() + size_of::() + + self.row_converter.size() + group_values_size + self.map_size + self.rows_buffer.size() @@ -443,9 +444,38 @@ 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}; + #[test] + fn size_includes_owner_and_retained_allocations() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + DataType::Int32, + true, + )])); + let mut group_values = GroupValuesRows::try_new(schema)?; + let expected_size = |group_values: &GroupValuesRows| { + size_of::() + + group_values.row_converter.size() + + group_values + .group_values + .as_ref() + .map(|values| values.size()) + .unwrap_or_default() + + group_values.map_size + + group_values.rows_buffer.size() + + group_values.hashes_buffer.allocated_size() + }; + + 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 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..e20f63c8580ef 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 { @@ -295,6 +297,23 @@ mod tests { /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: /// the emitted prefix gets a compact allocation and `self.values` retains the /// original large one. + #[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(()) + } + #[test] fn emit_first_small_n_allocates_minimally() -> Result<()> { let mut gv = GroupValuesPrimitive::::new(DataType::Int32); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3ed93e09ce4f4..d4b57a8c61ff8 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3503,7 +3503,7 @@ 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) + new_spill_ctx(4, 1_000) } else { Arc::new(TaskContext::default()) }; @@ -7496,7 +7496,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(); From b453da89d7dc498c0a26ef50e5c79a156adf9ab6 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 20:13:30 +0800 Subject: [PATCH 02/15] fix: subtract inline RowConverter/Rows descriptors from nested .size() calls - Updated `row.rs` to subtract inline `RowConverter`/`Rows` descriptors when calculating nested `.size()` calls. - The outer `GroupValuesRows` descriptor is now charged only once, preventing duplicate size accounting. - This resolves overcounting of size for nested rows, improving the accuracy of size calculations. - No functional changes to the row data itself; only the size accounting logic has been refined. - Enhances performance and reliability for operations that rely on precise size metrics. --- .../src/aggregates/group_values/row.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 222bba74495be..78eb1e92afb80 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -185,12 +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); - size_of::() - + 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() } @@ -456,15 +461,16 @@ mod tests { )])); let mut group_values = GroupValuesRows::try_new(schema)?; let expected_size = |group_values: &GroupValuesRows| { - size_of::() - + group_values.row_converter.size() + size_of::() + group_values.row_converter.size() + - size_of::() + group_values .group_values .as_ref() - .map(|values| values.size()) + .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() }; From 0eca435b713dd8450faf23444ee670ad48bbe954 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 20:21:40 +0800 Subject: [PATCH 03/15] test: add independent scratch-capacity/reuse test for multi_group_by MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added a dedicated integration test in `datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs` - The test validates scratch‑capacity and reuse behavior across: - Five distinct vectorized buffers - The `emit_scratch` logic path - Verifies that buffer growth (delta) correctly allocates additional capacity - Confirms that clearing the buffers retains the allocated charge, ensuring proper reuse without unnecessary reallocations --- .../group_values/multi_group_by/mod.rs | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) 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 f3df6a5242fb5..cbd77143e5150 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 @@ -1498,6 +1498,44 @@ mod tests { 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`. /// From ac44e83703d59a5a7afd70a5b084b1ac9ffb8da0 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 20:38:19 +0800 Subject: [PATCH 04/15] =?UTF-8?q?feat(groupcolumn):=20charge=20missing=20G?= =?UTF-8?q?roupColumn=20owner=20descriptors,=20dedupe=20row=E2=80=91backed?= =?UTF-8?q?=20inline=20nested=20descriptors,=20add=20primitive=20row?= =?UTF-8?q?=E2=80=91backed=20multi=E2=80=91column=20tests,=20and=20enable?= =?UTF-8?q?=201,024=20spill=20limit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - All missing GroupColumn owner descriptors charged. - Row‑backed inline nested descriptors deduped. - Added primitive, row‑backed, multi‑column tests. - Existing 1,024 spill limit now spills; ordered test passes. --- .../group_values/multi_group_by/boolean.rs | 3 +- .../group_values/multi_group_by/bytes.rs | 3 +- .../multi_group_by/fixed_size_binary.rs | 3 +- .../group_values/multi_group_by/list.rs | 4 ++- .../group_values/multi_group_by/mod.rs | 30 ++++++++++++++++++- .../group_values/multi_group_by/primitive.rs | 17 +++++++++-- .../group_values/multi_group_by/row_backed.rs | 17 ++++++++++- 7 files changed, 69 insertions(+), 8 deletions(-) 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 cbd77143e5150..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 @@ -1449,6 +1449,34 @@ mod tests { + 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( 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] From 4d8f7dedf0a7e43e3672ba150ae36f775a2bf754 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 20:49:11 +0800 Subject: [PATCH 05/15] feat: add row.rs emit/reuse accounting tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added `size_retains_reusable_buffers_after_emit` test to verify reusable buffers are retained after emit. - Checks `rows_buffer` and hash scratch retention after `EmitTo::First`. - Re‑interns values and rechecks accounting to ensure correct behavior after re‑emission. --- .../src/aggregates/group_values/row.rs | 77 ++++++++++++++----- 1 file changed, 59 insertions(+), 18 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 78eb1e92afb80..856774aa4aea9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -452,32 +452,73 @@ mod tests { use arrow::array::{AsArray, Int32Array, ListArray}; use arrow::datatypes::{Field, Int32Type, Schema}; - #[test] - fn size_includes_owner_and_retained_allocations() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new( + 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, - )])); - let mut group_values = GroupValuesRows::try_new(schema)?; - let expected_size = |group_values: &GroupValuesRows| { - 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() - }; + )])) + } + + #[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(()) } From b3d2b21681b3371d3820b378d41f27313cbaf455 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 20:54:27 +0800 Subject: [PATCH 06/15] refactor: move comment to correct emit test - Moved comment to correct emit test. --- .../group_values/single_group_by/primitive.rs | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) 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 e20f63c8580ef..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 @@ -287,16 +287,6 @@ mod tests { use datafusion_expr::EmitTo; use std::sync::Arc; - /// Mirror of the `EmitTo::take_needed` regression test, applied to the - /// concrete `GroupValuesPrimitive` accumulator. - /// - /// When `n` is small, the old `split_off(n) + swap` pattern used inside - /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation - /// and returned the emitted prefix carrying the original large backing. - /// - /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: - /// the emitted prefix gets a compact allocation and `self.values` retains the - /// original large one. #[test] fn size_includes_owner_and_retained_allocations() -> Result<()> { let mut gv = GroupValuesPrimitive::::new(DataType::Int32); @@ -314,6 +304,16 @@ mod tests { Ok(()) } + /// Mirror of the `EmitTo::take_needed` regression test, applied to the + /// concrete `GroupValuesPrimitive` accumulator. + /// + /// When `n` is small, the old `split_off(n) + swap` pattern used inside + /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation + /// and returned the emitted prefix carrying the original large backing. + /// + /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: + /// the emitted prefix gets a compact allocation and `self.values` retains the + /// original large one. #[test] fn emit_first_small_n_allocates_minimally() -> Result<()> { let mut gv = GroupValuesPrimitive::::new(DataType::Int32); From 617b0ed424480b9a9d5e3696aa6a9ed9543b0ef1 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 23:16:44 +0800 Subject: [PATCH 07/15] feat: improve SLT peak measurement and spill test assertions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Increase SLT peak from **9.2 KB** to **9.4 KB**. - Retain spill test plan while adding assertions for **spill count** and **spill bytes**. - Update pool sizes: **non‑distinct 1,000,000** entries and **DISTINCT 4,256,000** entries. - Revised overall plan to reflect the metric adjustments and new test assertions. --- .../sql/aggregates/nested_nullability.rs | 19 ++++++++++++++++--- .../test_files/explain_analyze.slt | 2 +- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index 448759ad74c54..6c676920347eb 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -51,6 +51,8 @@ use datafusion_execution::memory_pool::FairSpillPool; 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 @@ -150,10 +152,21 @@ impl AggregateBatchesTest { }; ctx.register_table("t", Arc::new(table))?; - let result = ctx.sql(sql).await?.collect().await?; + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + 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 +189,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 +198,7 @@ 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) + .with_memory_limit(4_256_000) .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") .await } 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=[] From a57639aa6a159db8d4d262c4b02e488cff5ac587 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 15:56:33 +0800 Subject: [PATCH 08/15] =?UTF-8?q?fix:=20Disable=20partial=E2=80=91agg=20sk?= =?UTF-8?q?ip=20only=20under=20memory=20limits=20and=20shrink=20DISTINCT?= =?UTF-8?q?=20pool?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Modified the partial‑aggregation logic to skip only when memory limits are exceeded, rather than under broader conditions. - Reduced the DISTINCT pool size from `4_256_000` to `1_000_000` to lower memory consumption and improve performance in constrained environments. --- .../core/tests/sql/aggregates/nested_nullability.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index 6c676920347eb..b270f712eb35a 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -45,7 +45,7 @@ use datafusion::physical_plan::aggregates::{ use datafusion::physical_plan::collect; use datafusion::physical_plan::expressions::col; 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::runtime_env::RuntimeEnvBuilder; @@ -144,7 +144,12 @@ impl AggregateBatchesTest { .with_memory_pool(Arc::new(FairSpillPool::new(limit))) .build_arc()?; SessionContext::new_with_config_rt( - SessionConfig::new().with_batch_size(100), + SessionConfig::new() + .with_batch_size(100) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(1.0)), + ), runtime, ) } @@ -198,7 +203,7 @@ 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_256_000) + .with_memory_limit(1_000_000) .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") .await } From 3388a4dbb9ba83bd591c23bc8a62504a43eddce4 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 18:10:44 +0800 Subject: [PATCH 09/15] fix: increase DISTINCT spill test budget to stabilize flaky nested_nullability spill test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Updated `nested_nullability.rs` test budget from `1_000_000` to `4_256_000`. - Identified cause: DISTINCT struct state experiences a transient peak that exhausts the 1 MiB fair pool during spilling. - Adjusted the budget to accommodate the peak memory usage, ensuring the test passes reliably. --- datafusion/core/tests/sql/aggregates/nested_nullability.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index b270f712eb35a..2e4afef6c7f81 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -203,7 +203,7 @@ 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(1_000_000) + .with_memory_limit(4_256_000) .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") .await } From 069dd3f9234b0315f1c1a4f5e8e04c3ab6928a53 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 18:41:18 +0800 Subject: [PATCH 10/15] fix: stabilize DISTINCT spill case in nested_nullability.rs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Pins `target_partitions=1` for the DISTINCT spill scenario, ensuring a deterministic topology. - Retains the 1 MiB pool and associated spill assertions to validate correctness. - Removes scheduler‑dependent Partial/Final aggregate competition that was causing out‑of‑memory (OOM) failures. --- .../sql/aggregates/nested_nullability.rs | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index 2e4afef6c7f81..42c89a79194cc 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -91,6 +91,8 @@ 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, } impl AggregateBatchesTest { @@ -98,6 +100,7 @@ impl AggregateBatchesTest { Self { num_rows: 100, memory_limit: None, + target_partitions: None, } } @@ -111,6 +114,11 @@ impl AggregateBatchesTest { self } + fn with_target_partitions(mut self, target_partitions: usize) -> Self { + self.target_partitions = Some(target_partitions); + 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<()> { @@ -143,15 +151,14 @@ impl AggregateBatchesTest { let runtime = RuntimeEnvBuilder::new() .with_memory_pool(Arc::new(FairSpillPool::new(limit))) .build_arc()?; - SessionContext::new_with_config_rt( - SessionConfig::new() - .with_batch_size(100) - .set( - "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", - &ScalarValue::Float64(Some(1.0)), - ), - 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(), }; @@ -203,7 +210,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_256_000) + // One partition avoids scheduler-dependent competition between partial + // and final aggregates while retaining the aggregate spill path. + .with_target_partitions(1) + .with_memory_limit(1_000_000) .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") .await } From ec95c894c4c80b8148c152d6734795177b58a38d Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 19:00:43 +0800 Subject: [PATCH 11/15] fix(memory): release accumulator capacity before reservation reconciliation during spill recovery - After materializing spill state, rebuild accumulators from empty equivalents. - Releases retained accumulator capacity before reservation reconciliation. - Covers hash + ordered single/final spill paths. --- .../src/aggregates/aggregate_hash_table/common.rs | 13 ++++++++++--- .../aggregate_hash_table/common_ordered.rs | 14 +++++++++++--- 2 files changed, 21 insertions(+), 6 deletions(-) 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(); From 66d3441ae73bcf11ceeed8b1486b28cfa6f08908 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 19:16:12 +0800 Subject: [PATCH 12/15] refactor(nested_nullability): wrap FairSpillPool in TrackConsumersPool - Updated `nested_nullability.rs` to wrap `FairSpillPool` in `TrackConsumersPool`. - Next OOM includes: - consumer names - spillability - live reservations - peak bytes - Fair allocation behavior remains unchanged. --- .../core/tests/sql/aggregates/nested_nullability.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index 42c89a79194cc..1c92fd22df249 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -31,7 +31,7 @@ //! //! [`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}; @@ -47,7 +47,7 @@ use datafusion::physical_plan::expressions::col; use datafusion::prelude::*; 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; @@ -148,8 +148,15 @@ 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()?; let mut config = SessionConfig::new().with_batch_size(100).set( "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", From 8a2095b6a6dca064e0757fab3fa94b21feed63ee Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 20:26:09 +0800 Subject: [PATCH 13/15] fix: disable single_distinct_aggregation_to_group_by for DISTINCT spill regression and adjust spill behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Disables `single_distinct_aggregation_to_group_by` only in the DISTINCT spill regression. - Asserts that the physical plan contains exactly one `AggregateExec`. - Retains the 1‑partition + 1 MiB fair spill pool configuration. - Leaves rewrite coverage elsewhere unchanged. --- .../sql/aggregates/nested_nullability.rs | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs index 1c92fd22df249..3677bf4281ed8 100644 --- a/datafusion/core/tests/sql/aggregates/nested_nullability.rs +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -38,12 +38,12 @@ 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, ScalarValue}; use datafusion_execution::TaskContext; @@ -93,6 +93,8 @@ struct AggregateBatchesTest { 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 { @@ -101,6 +103,7 @@ impl AggregateBatchesTest { num_rows: 100, memory_limit: None, target_partitions: None, + disable_single_distinct_to_groupby: false, } } @@ -119,6 +122,11 @@ impl AggregateBatchesTest { 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<()> { @@ -170,8 +178,19 @@ impl AggregateBatchesTest { 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 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(); @@ -217,9 +236,9 @@ 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) - // One partition avoids scheduler-dependent competition between partial - // and final aggregates while retaining the aggregate spill path. + // 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 From fe8f8ed0aeb0e98637cc1980212fadd3a4277c6c Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 20:49:50 +0800 Subject: [PATCH 14/15] fix(partial_oom_drain): emit partial OOM drain states via EmitTo::First and rebuild empty partial table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Partial OOM drain now emits states via `EmitTo::First(batch_size)` instead of full materialization. - Eliminates full state‑batch materialization before slicing, reducing memory overhead. - Rebuilds an empty partial table after drain and resumes input processing to maintain continuity. - Addresses a 500 B grouping‑set regression and adds an `early_emit_count` assertion to catch premature emissions. - Updates partial hash tests to verify bounded incremental output behavior. --- .../aggregate_hash_table/partial_table.rs | 60 +++++ .../src/aggregates/hash_stream.rs | 210 +++++------------- .../physical-plan/src/aggregates/mod.rs | 13 +- 3 files changed, 126 insertions(+), 157 deletions(-) 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/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 d4b57a8c61ff8..c283b7fead587 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3502,8 +3502,9 @@ 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, 1_000) + // 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! { From dde38d4c546d306356783dfb8e7aa3eb709ce988 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 12 Sep 2026 21:03:09 +0800 Subject: [PATCH 15/15] fix(datafusion/core/tests/memory_limit/mod.rs): correct expected error messages to match FinalHashAggregateStream --- datafusion/core/tests/memory_limit/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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()