From ac5661b394b1b1e32fbc100e2dc6e9ad0094c46f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:14:49 -0500 Subject: [PATCH 1/7] fix: report the real hashbrown allocation from ArrowBytesMap and ArrowBytesViewMap Both maps tracked their hash table footprint in a `map_size` field that was only ever incremented by `HashTableAllocExt::insert_accounted`, which charges `capacity * size_of::()` on growth and nothing else. That undercounts in two ways. `ArrowBytesViewMap::new` seeded `map_size` with `capacity() * size_of::>()`, which ignores the control bytes and the trailing group that hashbrown allocates alongside the entry array, so the reported size was roughly half the real allocation. `ArrowBytesMap::new` seeded `map_size` with 0 despite pre-allocating a table for 128 entries. Since `insert_accounted` only charges when the table grows, any map holding fewer entries than the pre-allocated capacity reported its hash table as free forever. Drop the field and ask hashbrown for the exact figure with `HashTable::allocation_size`, which covers entries, control bytes and the trailing group. It is a constant time layout calculation, so `size()` stays cheap, and it cannot drift out of sync with the table the way an incrementally maintained counter can. --- .../physical-expr-common/src/binary_map.rs | 26 +++++-------- .../src/binary_view_map.rs | 38 +++++++++++-------- 2 files changed, 31 insertions(+), 33 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 7543e6b29732..ceae72d4ccaf 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -28,7 +28,7 @@ use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::any::type_name; use std::fmt::Debug; @@ -217,8 +217,6 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, - /// Total size of the map in bytes - map_size: usize, /// In progress buffer containing all values buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used @@ -248,7 +246,6 @@ where Self { output_type, map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - map_size: 0, buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), @@ -415,11 +412,7 @@ where offset_or_inline: inline, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + self.map.insert_unique(hash, new_header, |header| header.hash); payload } } @@ -457,11 +450,7 @@ where offset_or_inline: offset, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + self.map.insert_unique(hash, new_header, |header| header.hash); payload } }; @@ -486,7 +475,6 @@ where let Self { output_type, map: _, - map_size: _, offsets, buffer, random_state: _, @@ -591,7 +579,11 @@ where /// Return the total size, in bytes, of memory used to store the data in /// this set, not including `self` pub fn size(&self) -> usize { - self.map_size + // `HashTable::allocation_size` reports the whole hashbrown allocation, + // which is the entry array plus the control bytes plus the trailing + // group, so it is larger than `capacity() * size_of::>()`. + // It is a constant time layout calculation, not a walk of the table. + self.map.allocation_size() + self.buffer.capacity() * size_of::() + self.offsets.allocated_size() + self.hashes_buffer.allocated_size() @@ -615,7 +607,7 @@ where fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ArrowBytesMap") .field("map", &"") - .field("map_size", &self.map_size) + .field("map_allocation_size", &self.map.allocation_size()) .field("buffer", &self.buffer) .field("random_state", &self.random_state) .field("hashes_buffer", &self.hashes_buffer) diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 0457825decb9..304b32187046 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -28,10 +28,9 @@ use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType}; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::fmt::Debug; -use std::mem::size_of; use std::sync::Arc; /// HashSet optimized for storing string or binary values that can produce that @@ -129,8 +128,6 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, - /// Total size of the map in bytes - map_size: usize, /// Views for all stored values (in insertion order) views: Vec, @@ -159,13 +156,9 @@ where V: Debug + PartialEq + Eq + Clone + Copy + Default, { pub fn new(output_type: OutputType) -> Self { - let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY); - let map_size = map.capacity() * size_of::>(); - Self { output_type, - map, - map_size, + map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -374,8 +367,7 @@ where payload, }; - self.map - .insert_accounted(new_header, |h| h.hash, &mut self.map_size); + self.map.insert_unique(hash, new_header, |h| h.hash); payload }; observe_payload_fn(payload); @@ -540,13 +532,18 @@ where pub fn size(&self) -> usize { // All fields below own their allocations. Count retained capacity rather // than used length because this value drives memory accounting. + // + // `HashTable::allocation_size` reports the whole hashbrown allocation, + // which is the entry array plus the control bytes plus the trailing + // group, so it is larger than `capacity() * size_of::>()`. It + // is a constant time layout calculation, not a walk of the table. let views_size = self.views.allocated_size(); let in_progress_size = self.in_progress.allocated_size(); let completed_size = self.completed.allocated_size() + self.completed.iter().map(Buffer::capacity).sum::(); let nulls_size = self.nulls.allocated_size(); - self.map_size + self.map.allocation_size() + views_size + in_progress_size + completed_size @@ -562,7 +559,7 @@ where fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ArrowBytesMap") .field("map", &"") - .field("map_size", &self.map_size) + .field("map_allocation_size", &self.map.allocation_size()) .field("views_len", &self.views.len()) .field("completed_buffers", &self.completed.len()) .field("random_state", &self.random_state) @@ -597,6 +594,7 @@ where mod tests { use arrow::array::{GenericByteViewArray, StringViewArray}; use datafusion_common::HashMap; + use std::mem::size_of; use super::*; @@ -789,7 +787,15 @@ mod tests { fn test_size_counts_initial_hash_table_capacity() { let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); - assert_eq!(map.size(), map.map.capacity() * size_of::>()); + assert_eq!(map.size(), map.map.allocation_size()); + // The reported size covers the control bytes as well as the entries, so + // it is strictly larger than the entry array on its own. + assert!( + map.size() > map.map.capacity() * size_of::>(), + "expected {} to exceed {}", + map.size(), + map.map.capacity() * size_of::>() + ); } #[test] @@ -822,7 +828,7 @@ mod tests { .any(|buffer| buffer.capacity() > buffer.len()) ); - let expected_size = map.map_size + let expected_size = map.map.allocation_size() + map.views.allocated_size() + map.in_progress.allocated_size() + map.completed.allocated_size() @@ -832,7 +838,7 @@ mod tests { assert_eq!(map.size(), expected_size); // Verify the retained-capacity delta independently from the production formula. - let legacy_size = map.map_size + let legacy_size = map.map.allocation_size() + map.views.len() * size_of::() + map.in_progress.capacity() + map.completed.iter().map(Buffer::len).sum::() From 7e0081d0cf42890633c8677398afa87fd8da9448 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:17:54 -0500 Subject: [PATCH 2/7] perf: stop pre-allocating a hash table per COUNT(DISTINCT) group `ArrowBytesMap` and `ArrowBytesViewMap` always pre-allocated their hash table, and `ArrowBytesMap` also pre-allocated an 8 KiB value buffer. That is the right trade for the single map that backs a `GROUP BY` on one string column, which goes on to hold every group value in the query. It is the wrong trade for `BytesDistinctCountAccumulator` and `BytesViewDistinctCountAccumulator`, because `GroupsAccumulatorAdapter` creates one accumulator per group: a grouped `COUNT(DISTINCT)` over a high cardinality key holds hundreds of thousands of them at once, and most see only a handful of values, so the pre-allocation dwarfs the data. Split the constructors. `new` no longer allocates anything, and `with_capacity` keeps the previous behavior for the callers that want it. The capacity is stored so `take` re-creates the map the way it was built. The `GroupValuesBytes` and `GroupValuesBytesView` call sites move to `with_capacity`; the two distinct-count accumulators stay on `new`. The `arrow_bytes_map` benchmark also moves to `with_capacity`: its `long_low_cardinality` case is defined by the distinct values fitting inside the pre-allocated buffer. --- .../src/aggregate/count_distinct/bytes.rs | 6 + .../benches/arrow_bytes_map.rs | 12 +- .../physical-expr-common/src/binary_map.rs | 134 ++++++++++++++++-- .../src/binary_view_map.rs | 103 ++++++++++++-- .../group_values/single_group_by/bytes.rs | 8 +- .../single_group_by/bytes_view.rs | 8 +- 6 files changed, 247 insertions(+), 24 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs index f6df4182a879..3aa60f6f3b6a 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs @@ -39,6 +39,10 @@ use std::mem::size_of_val; pub struct BytesDistinctCountAccumulator(ArrowBytesSet); impl BytesDistinctCountAccumulator { + /// The set deliberately does not pre-allocate. `GroupsAccumulatorAdapter` + /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a + /// high cardinality key holds hundreds of thousands of these at once and + /// most of them see only a handful of values. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesSet::new(output_type)) } @@ -100,6 +104,8 @@ impl Accumulator for BytesDistinctCountAccumulator { pub struct BytesViewDistinctCountAccumulator(ArrowBytesViewSet); impl BytesViewDistinctCountAccumulator { + /// See [`BytesDistinctCountAccumulator::new`] for why the set does not + /// pre-allocate. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesViewSet::new(output_type)) } diff --git a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs index 7c8cdc3b4c50..68351a839554 100644 --- a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs +++ b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs @@ -17,7 +17,9 @@ use arrow::array::{ArrayRef, StringArray}; use criterion::{Criterion, Throughput, criterion_group, criterion_main}; -use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_map::{ + ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType, +}; use std::hint::black_box; use std::sync::Arc; @@ -57,7 +59,13 @@ fn bench_arrow_bytes_map(c: &mut Criterion) { for (name, values) in cases { group.bench_function(name, |b| { b.iter(|| { - let mut map = ArrowBytesMap::::new(OutputType::Utf8); + // The `long_low_cardinality` case is defined by the distinct + // values fitting in the pre-allocated buffer, so this benchmark + // measures the pre-allocating constructor. + let mut map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); let mut next_payload = 0; map.insert_if_new( &values, diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ceae72d4ccaf..61b1e6e008c7 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -55,10 +55,21 @@ pub enum OutputType { pub struct ArrowBytesSet(ArrowBytesMap); impl ArrowBytesSet { + /// Creates a set that does not pre-allocate its hash table or value buffer. + /// + /// See [`ArrowBytesMap::new`] for when to prefer this over + /// [`Self::with_capacity`]. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesMap::new(output_type)) } + /// Creates a set with room for `map_capacity` entries. + /// + /// See [`ArrowBytesMap::with_capacity`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self(ArrowBytesMap::with_capacity(output_type, map_capacity)) + } + /// Return the contents of this set and replace it with a new empty /// set with the same output type pub fn take(&mut self) -> Self { @@ -217,6 +228,13 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, + /// Hash table capacity to re-create the map with in [`Self::take`], so a + /// map built with [`Self::with_capacity`] keeps its pre-allocation when it + /// is emptied and reused + initial_map_capacity: usize, + /// Value buffer capacity to re-create the buffer with in [`Self::take`], + /// for the same reason as `initial_map_capacity` + initial_buffer_capacity: usize, /// In progress buffer containing all values buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used @@ -234,19 +252,49 @@ where null: Option<(V, usize)>, } -/// The size, in number of entries, of the initial hash table -const INITIAL_MAP_CAPACITY: usize = 128; -/// The initial size, in bytes, of the string data +/// The size, in number of entries, of the hash table pre-allocated by +/// [`ArrowBytesMap::with_capacity`]. It is a warm up size for maps that go on +/// to hold many values, not a bound on what the map can hold. +pub const INITIAL_MAP_CAPACITY: usize = 128; +/// The size, in bytes, of the string data buffer pre-allocated by +/// [`ArrowBytesMap::with_capacity`] pub const INITIAL_BUFFER_CAPACITY: usize = 8 * 1024; impl ArrowBytesMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, { + /// Creates a map that does not pre-allocate its hash table or value buffer. + /// + /// Use this when maps are created in large numbers and most of them stay + /// small, such as the per group `COUNT(DISTINCT)` accumulators that + /// `GroupsAccumulatorAdapter` creates one of per group. There the + /// pre-allocation dwarfs the values the map actually holds. pub fn new(output_type: OutputType) -> Self { + Self::new_inner(output_type, 0, 0) + } + + /// Creates a map whose hash table is pre-allocated for `map_capacity` + /// entries and whose value buffer is pre-allocated with + /// [`INITIAL_BUFFER_CAPACITY`] bytes. + /// + /// Use this for the few long lived maps that are each expected to hold many + /// values, such as the single map backing a `GROUP BY` on one string + /// column. The capacities are preserved across [`Self::take`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self::new_inner(output_type, map_capacity, INITIAL_BUFFER_CAPACITY) + } + + fn new_inner( + output_type: OutputType, + map_capacity: usize, + buffer_capacity: usize, + ) -> Self { Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, + initial_buffer_capacity: buffer_capacity, + buffer: Vec::with_capacity(buffer_capacity), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), hashes_buffer: vec![], @@ -257,7 +305,11 @@ where /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.output_type); + let mut new_self = Self::new_inner( + self.output_type, + self.initial_map_capacity, + self.initial_buffer_capacity, + ); swap(self, &mut new_self); new_self } @@ -412,7 +464,8 @@ where offset_or_inline: inline, payload, }; - self.map.insert_unique(hash, new_header, |header| header.hash); + self.map + .insert_unique(hash, new_header, |header| header.hash); payload } } @@ -450,7 +503,8 @@ where offset_or_inline: offset, payload, }; - self.map.insert_unique(hash, new_header, |header| header.hash); + self.map + .insert_unique(hash, new_header, |header| header.hash); payload } }; @@ -475,6 +529,8 @@ where let Self { output_type, map: _, + initial_map_capacity: _, + initial_buffer_capacity: _, offsets, buffer, random_state: _, @@ -655,6 +711,68 @@ mod tests { use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; use std::collections::HashMap; + /// The bytes a hashbrown table of `buckets` buckets must allocate for + /// entries of type `T`, ignoring the alignment padding and the trailing + /// group. Derived independently of the production accounting so it can + /// bracket it. + fn min_table_bytes(buckets: usize) -> usize { + // One entry slot plus one control byte per bucket. + buckets * (size_of::() + 1) + } + + #[test] + fn map_new_does_not_allocate() { + let map = ArrowBytesMap::::new(OutputType::Utf8); + + assert_eq!(map.map.capacity(), 0); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.buffer.capacity(), 0); + // Only the single leading zero offset is allocated. + assert!(map.size() < 128, "expected {} to be tiny", map.size()); + } + + #[test] + fn map_with_capacity_reports_the_real_hash_table_allocation() { + let map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + + // Before this accounting was corrected the map reported its hash table + // as costing zero bytes until the table grew past its pre-allocation. + let table_bytes = map.map.allocation_size(); + let lower_bound = min_table_bytes::>(map.map.capacity()); + assert!( + table_bytes >= lower_bound, + "expected {table_bytes} to be at least {lower_bound}" + ); + assert!( + table_bytes <= 2 * lower_bound + 64, + "expected {table_bytes} to be within a small factor of {lower_bound}" + ); + assert!(map.size() >= table_bytes + INITIAL_BUFFER_CAPACITY); + } + + #[test] + fn take_preserves_the_capacity_the_map_was_built_with() { + let mut preallocated = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + let capacity = preallocated.map.capacity(); + preallocated.take(); + assert_eq!(preallocated.map.capacity(), capacity); + assert_eq!(preallocated.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + lazy.take(); + assert_eq!(lazy.map.capacity(), 0); + assert_eq!(lazy.buffer.capacity(), 0); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 304b32187046..9351d9209298 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -39,10 +39,21 @@ use std::sync::Arc; pub struct ArrowBytesViewSet(ArrowBytesViewMap<()>); impl ArrowBytesViewSet { + /// Creates a set that does not pre-allocate its hash table. + /// + /// See [`ArrowBytesViewMap::new`] for when to prefer this over + /// [`Self::with_capacity`]. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesViewMap::new(output_type)) } + /// Creates a set with room for `map_capacity` entries. + /// + /// See [`ArrowBytesViewMap::with_capacity`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self(ArrowBytesViewMap::with_capacity(output_type, map_capacity)) + } + /// Inserts each value from `values` into the set pub fn insert(&mut self, values: &ArrayRef) { fn make_payload_fn(_value: Option<&[u8]>) {} @@ -54,9 +65,7 @@ impl ArrowBytesViewSet { /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.0.output_type); - std::mem::swap(self, &mut new_self); - new_self + Self(self.0.take()) } /// Converts this set into a `StringViewArray` or `BinaryViewArray` @@ -128,6 +137,10 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, + /// Hash table capacity to re-create the map with in [`Self::take`], so a + /// map built with [`Self::with_capacity`] keeps its pre-allocation when it + /// is emptied and reused + initial_map_capacity: usize, /// Views for all stored values (in insertion order) views: Vec, @@ -148,17 +161,36 @@ where null: Option<(V, usize)>, } -/// The size, in number of entries, of the initial hash table -const INITIAL_MAP_CAPACITY: usize = 512; +/// The size, in number of entries, of the hash table pre-allocated by +/// [`ArrowBytesViewMap::with_capacity`]. It is a warm up size for maps that go +/// on to hold many values, not a bound on what the map can hold. +pub const INITIAL_MAP_CAPACITY: usize = 512; impl ArrowBytesViewMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, { + /// Creates a map that does not pre-allocate its hash table. + /// + /// Use this when maps are created in large numbers and most of them stay + /// small, such as the per group `COUNT(DISTINCT)` accumulators that + /// `GroupsAccumulatorAdapter` creates one of per group. There the + /// pre-allocation dwarfs the values the map actually holds. pub fn new(output_type: OutputType) -> Self { + Self::with_capacity(output_type, 0) + } + + /// Creates a map whose hash table is pre-allocated for `map_capacity` + /// entries. + /// + /// Use this for the few long lived maps that are each expected to hold many + /// values, such as the single map backing a `GROUP BY` on one string + /// column. The capacity is preserved across [`Self::take`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -172,7 +204,8 @@ where /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.output_type); + let mut new_self = + Self::with_capacity(self.output_type, self.initial_map_capacity); std::mem::swap(self, &mut new_self); new_self } @@ -783,13 +816,48 @@ mod tests { assert_eq!(set.len(), 10); } + /// The bytes a hashbrown table of `buckets` buckets must allocate for + /// entries of type `T`, ignoring the alignment padding and the trailing + /// group. Derived independently of the production accounting so it can + /// bracket it. + fn min_table_bytes(buckets: usize) -> usize { + // One entry slot plus one control byte per bucket. + buckets * (size_of::() + 1) + } + #[test] - fn test_size_counts_initial_hash_table_capacity() { + fn map_new_does_not_allocate() { let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); + assert_eq!(map.map.capacity(), 0); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.size(), 0); + } + + #[test] + fn test_size_counts_initial_hash_table_capacity() { + let map = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); assert_eq!(map.size(), map.map.allocation_size()); - // The reported size covers the control bytes as well as the entries, so - // it is strictly larger than the entry array on its own. + + // Before this accounting was corrected the map reported exactly + // `capacity() * size_of::>()`, which leaves out the control + // bytes and undercounts the real allocation by roughly half. + let lower_bound = min_table_bytes::>(map.map.capacity()); + assert!( + map.size() >= lower_bound, + "expected {} to be at least {lower_bound}", + map.size() + ); + assert!( + map.size() <= 2 * lower_bound + 64, + "expected {} to be within a small factor of {lower_bound}", + map.size() + ); assert!( map.size() > map.map.capacity() * size_of::>(), "expected {} to exceed {}", @@ -798,6 +866,21 @@ mod tests { ); } + #[test] + fn take_preserves_the_capacity_the_map_was_built_with() { + let mut preallocated = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + let capacity = preallocated.map.capacity(); + preallocated.take(); + assert_eq!(preallocated.map.capacity(), capacity); + + let mut lazy = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); + lazy.take(); + assert_eq!(lazy.map.capacity(), 0); + } + #[test] fn test_size_counts_retained_buffer_capacities() { let first = "a".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 34ec36be31d2..87eeba489ecc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -22,7 +22,9 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use datafusion_common::Result; use datafusion_expr::{EmitTo, GroupSelection}; -use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_map::{ + ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType, +}; /// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values /// @@ -38,7 +40,9 @@ pub struct GroupValuesBytes { impl GroupValuesBytes { pub fn new(output_type: OutputType) -> Self { Self { - map: ArrowBytesMap::new(output_type), + // One map holds every group value for the whole query, so it is + // worth pre-allocating the hash table and the value buffer. + map: ArrowBytesMap::with_capacity(output_type, INITIAL_MAP_CAPACITY), num_groups: 0, } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 997a7ce166a7..f5183ba329b3 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -19,7 +19,9 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef}; use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; -use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; +use datafusion_physical_expr_common::binary_view_map::{ + ArrowBytesViewMap, INITIAL_MAP_CAPACITY, +}; use std::mem::size_of; /// A [`GroupValues`] storing single column of Utf8View/BinaryView values @@ -36,7 +38,9 @@ pub struct GroupValuesBytesView { impl GroupValuesBytesView { pub fn new(output_type: OutputType) -> Self { Self { - map: ArrowBytesViewMap::new(output_type), + // One map holds every group value for the whole query, so it is + // worth pre-allocating the hash table. + map: ArrowBytesViewMap::with_capacity(output_type, INITIAL_MAP_CAPACITY), num_groups: 0, } } From eb8ad022bde712760c647d67fcd8810b8632f138 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:46:45 -0500 Subject: [PATCH 3/7] refactor: name the hash table term in ArrowBytesViewMap::size Keep the comment about what `HashTable::allocation_size` covers next to the value it describes, and say what the test helper's lower bound is derived from. --- datafusion/physical-expr-common/src/binary_map.rs | 12 +++++------- .../physical-expr-common/src/binary_view_map.rs | 15 +++++++-------- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 61b1e6e008c7..0b59bcf5937b 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -711,13 +711,11 @@ mod tests { use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; use std::collections::HashMap; - /// The bytes a hashbrown table of `buckets` buckets must allocate for - /// entries of type `T`, ignoring the alignment padding and the trailing - /// group. Derived independently of the production accounting so it can - /// bracket it. - fn min_table_bytes(buckets: usize) -> usize { - // One entry slot plus one control byte per bucket. - buckets * (size_of::() + 1) + /// A lower bound on the bytes a hashbrown table holding `entries` entries + /// of type `T` must allocate: one entry slot and one control byte each. + /// Derived independently of the production accounting so it can bracket it. + fn min_table_bytes(entries: usize) -> usize { + entries * (size_of::() + 1) } #[test] diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9351d9209298..35f45e8e9a12 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -570,13 +570,14 @@ where // which is the entry array plus the control bytes plus the trailing // group, so it is larger than `capacity() * size_of::>()`. It // is a constant time layout calculation, not a walk of the table. + let map_size = self.map.allocation_size(); let views_size = self.views.allocated_size(); let in_progress_size = self.in_progress.allocated_size(); let completed_size = self.completed.allocated_size() + self.completed.iter().map(Buffer::capacity).sum::(); let nulls_size = self.nulls.allocated_size(); - self.map.allocation_size() + map_size + views_size + in_progress_size + completed_size @@ -816,13 +817,11 @@ mod tests { assert_eq!(set.len(), 10); } - /// The bytes a hashbrown table of `buckets` buckets must allocate for - /// entries of type `T`, ignoring the alignment padding and the trailing - /// group. Derived independently of the production accounting so it can - /// bracket it. - fn min_table_bytes(buckets: usize) -> usize { - // One entry slot plus one control byte per bucket. - buckets * (size_of::() + 1) + /// A lower bound on the bytes a hashbrown table holding `entries` entries + /// of type `T` must allocate: one entry slot and one control byte each. + /// Derived independently of the production accounting so it can bracket it. + fn min_table_bytes(entries: usize) -> usize { + entries * (size_of::() + 1) } #[test] From 84f07dadabe68e25900b59909c9c0c28e99c7c56 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:06:54 -0500 Subject: [PATCH 4/7] fix: release the group values map allocations in `clear_shrink` `GroupValuesBytes::clear_shrink` and `GroupValuesBytesView::clear_shrink` reset their map with `take()`, which restores the capacity the map was configured with so the emptied map stays warm. That is what the emit path wants, but `clear_shrink` exists to hand memory back before spilling and before the spilled batch is sorted, so it left roughly 16 KiB (string and binary) and 34 KiB (view) reserved instead of releasing it. Add `clear_and_release` to `ArrowBytesMap` and `ArrowBytesViewMap`, which empties the map and drops its allocations while remembering the configured capacities so a later `take()` still warms the map up, and call it from the two `clear_shrink` implementations. The pre-allocation stays at construction, where the hot single column string `GROUP BY` path earns it. --- .../physical-expr-common/src/binary_map.rs | 55 ++++++++++++++++ .../src/binary_view_map.rs | 52 +++++++++++++++ .../group_values/single_group_by/bytes.rs | 65 ++++++++++++++++++- .../single_group_by/bytes_view.rs | 64 +++++++++++++++++- 4 files changed, 230 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 0b59bcf5937b..bc7d98717377 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -314,6 +314,22 @@ where new_self } + /// Empties this map and releases every allocation it holds, so + /// [`Self::size`] drops to approximately zero. + /// + /// This is the difference from [`Self::take`]: `take` restores the + /// capacities the map was configured with so the emptied map stays warm for + /// continued use, which is what emitting wants, whereas here the point is + /// to hand the memory back, as before spilling or before a downstream sort. + /// The configured capacities are remembered, so the next [`Self::take`] + /// warms the map up again. + pub fn clear_and_release(&mut self) { + let mut released = Self::new_inner(self.output_type, 0, 0); + released.initial_map_capacity = self.initial_map_capacity; + released.initial_buffer_capacity = self.initial_buffer_capacity; + *self = released; + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -771,6 +787,45 @@ mod tests { assert_eq!(lazy.buffer.capacity(), 0); } + #[test] + fn clear_and_release_frees_the_preallocation_that_take_keeps() { + let mut map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1_000).map(|i| format!("distinct value number {i}")), + )); + map.insert_if_new(&values, |_| (), |_| ()); + + let populated_size = map.size(); + assert!(populated_size > INITIAL_BUFFER_CAPACITY); + + // `take` deliberately keeps the map warm, so it does not release the + // configured capacities. + map.take(); + let taken_size = map.size(); + assert!( + taken_size > INITIAL_BUFFER_CAPACITY, + "expected take to retain the warm up allocations, got {taken_size}" + ); + + map.clear_and_release(); + let released_size = map.size(); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.buffer.capacity(), 0); + assert!( + released_size < 128, + "expected the released map to report approximately zero bytes, got {released_size}" + ); + + // The configured capacities survive, so the map warms back up when it + // is emitted from again. + map.take(); + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 35f45e8e9a12..29f4014c5f9a 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -210,6 +210,21 @@ where new_self } + /// Empties this map and releases every allocation it holds, so + /// [`Self::size`] drops to approximately zero. + /// + /// This is the difference from [`Self::take`]: `take` restores the capacity + /// the map was configured with so the emptied map stays warm for continued + /// use, which is what emitting wants, whereas here the point is to hand the + /// memory back, as before spilling or before a downstream sort. The + /// configured capacity is remembered, so the next [`Self::take`] warms the + /// map up again. + pub fn clear_and_release(&mut self) { + let mut released = Self::with_capacity(self.output_type, 0); + released.initial_map_capacity = self.initial_map_capacity; + *self = released; + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -880,6 +895,43 @@ mod tests { assert_eq!(lazy.map.capacity(), 0); } + #[test] + fn clear_and_release_frees_the_preallocation_that_take_keeps() { + let mut map = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1_000).map(|i| format!("distinct value number {i}")), + )); + map.insert_if_new(&values, |_| (), |_| ()); + + let warm_size = map.map.allocation_size(); + assert!(warm_size > 0); + + // `take` deliberately keeps the map warm, so it does not release the + // configured capacity. + map.take(); + let taken_size = map.size(); + assert!( + taken_size >= min_table_bytes::>(INITIAL_MAP_CAPACITY), + "expected take to retain the warm up allocation, got {taken_size}" + ); + + map.clear_and_release(); + assert_eq!(map.map.allocation_size(), 0); + let released_size = map.size(); + assert!( + released_size < 128, + "expected the released map to report approximately zero bytes, got {released_size}" + ); + + // The configured capacity survives, so the map warms back up when it is + // emitted from again. + map.take(); + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + } + #[test] fn test_size_counts_retained_buffer_capacities() { let first = "a".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 87eeba489ecc..9f7b4b4e91cb 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -137,8 +137,67 @@ impl GroupValues for GroupValuesBytes { } fn clear_shrink(&mut self, _num_rows: usize) { - // in theory we could potentially avoid this reallocation and clear the - // contents of the maps, but for now we just reset the map from the beginning - self.map.take(); + // Callers use this to hand memory back before spilling or sorting, so + // release the map's allocations rather than restoring the warm up + // capacities that `take` keeps for the emit path. + self.map.clear_and_release(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::StringArray; + use datafusion_physical_expr_common::binary_map::INITIAL_BUFFER_CAPACITY; + + /// `clear_shrink` is how the aggregate stream hands memory back before it + /// spills and before the spilled batch is sorted, so the memory it releases + /// has to actually show up in the size it reports afterwards. + #[test] + fn clear_shrink_releases_the_reported_memory() { + let mut group_values = GroupValuesBytes::::new(OutputType::Utf8); + let empty = size_of::>(); + + // The map is pre-allocated at construction, so it is already well above + // its own struct size before a single row is interned. + let warm_size = group_values.size(); + assert!( + warm_size > empty + INITIAL_BUFFER_CAPACITY, + "expected the pre-allocated map to report more than {} bytes, got {warm_size}", + empty + INITIAL_BUFFER_CAPACITY + ); + + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1_000).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + let populated_size = group_values.size(); + assert!(populated_size > warm_size); + + group_values.clear_shrink(0); + + // Everything the map held is gone: what remains is the struct itself + // plus the single leading zero offset. + let released_size = group_values.size(); + assert!( + released_size < empty + 128, + "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}" + ); + assert!( + released_size * 10 < populated_size, + "expected {released_size} to be far below {populated_size}" + ); + + // The map still works, and warms back up on the next emit. + group_values.intern(&[values], &mut groups).unwrap(); + assert!(group_values.size() > released_size); + group_values.emit(EmitTo::All).unwrap(); + assert!(group_values.size() > empty + INITIAL_BUFFER_CAPACITY); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index f5183ba329b3..23ea4e7ed3f8 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -139,8 +139,66 @@ impl GroupValues for GroupValuesBytesView { } fn clear_shrink(&mut self, _num_rows: usize) { - // in theory we could potentially avoid this reallocation and clear the - // contents of the maps, but for now we just reset the map from the beginning - self.map.take(); + // Callers use this to hand memory back before spilling or sorting, so + // release the map's allocations rather than restoring the warm up + // capacity that `take` keeps for the emit path. + self.map.clear_and_release(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::StringViewArray; + + /// `clear_shrink` is how the aggregate stream hands memory back before it + /// spills and before the spilled batch is sorted, so the memory it releases + /// has to actually show up in the size it reports afterwards. + #[test] + fn clear_shrink_releases_the_reported_memory() { + let mut group_values = GroupValuesBytesView::new(OutputType::Utf8View); + let empty = size_of::(); + + // The hash table is pre-allocated at construction, so the map is + // already well above its own struct size before a single row is + // interned. + let warm_size = group_values.size(); + assert!( + warm_size > empty + INITIAL_MAP_CAPACITY, + "expected the pre-allocated map to report more than {} bytes, got {warm_size}", + empty + INITIAL_MAP_CAPACITY + ); + + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1_000).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + let populated_size = group_values.size(); + assert!(populated_size > warm_size); + + group_values.clear_shrink(0); + + // Everything the map held is gone: what remains is the struct itself. + let released_size = group_values.size(); + assert!( + released_size < empty + 128, + "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}" + ); + assert!( + released_size * 10 < populated_size, + "expected {released_size} to be far below {populated_size}" + ); + + // The map still works, and warms back up on the next emit. + group_values.intern(&[values], &mut groups).unwrap(); + assert!(group_values.size() > released_size); + group_values.emit(EmitTo::All).unwrap(); + assert!(group_values.size() > empty + INITIAL_MAP_CAPACITY); } } From 54a0229ae8df904da340102b8aed5c4453d9c415 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:09:57 -0500 Subject: [PATCH 5/7] test: memory limit test for grouped `COUNT(DISTINCT )` A grouped `COUNT(DISTINCT )` gets one accumulator per group, and each of those owns a hash set of the distinct values it has seen. Those sets were created pre-allocated, so the query's memory use tracked the number of groups rather than the amount of data. Add two `memory_limit` tests that turn that into a binary observable, one for `Utf8` and one for `Utf8View`, over a new scenario of 4,000 groups holding 2 distinct values each. Measured against this branch's base commit with spilling disabled and `target_partitions` pinned to 1: | value column | budget needed before | budget needed after | | ------------ | -------------------- | ------------------- | | `Utf8` | ~35.5 MB | ~1.9 MB | | `Utf8View` | ~123 MB | ~2.7 MB | The tests run at 8 MB and 16 MB respectively, so each sits at least 4x above what the branch needs and at least 4x below what the base needs. Both fail on the base commit with `Resources exhausted` and pass here. --- datafusion/core/tests/memory_limit/mod.rs | 125 +++++++++++++++++++++- 1 file changed, 124 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 15b224d200bf..1527eb39c3a7 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,7 +26,9 @@ mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; -use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray}; +use arrow::array::{ + ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringArray, StringViewArray, +}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; @@ -125,6 +127,60 @@ async fn group_by_hash() { .await } +/// A grouped `COUNT(DISTINCT )` gets one accumulator per group, and +/// each of those owns a hash set of the distinct values it has seen. Those +/// sets used to be created pre-allocated, which costs far more than the +/// handful of values a group typically holds, so the query's memory use +/// tracked the number of groups rather than the amount of data. +/// +/// With 4,000 groups holding 2 distinct values each, this query needed about +/// 35.5 MB of budget before the per group pre-allocation was removed and +/// about 1.9 MB after, so an 8 MB limit is a failure before the change and a +/// success after it. Spilling is disabled, so completing means the query +/// genuinely fit in the budget. +/// +/// The `count(*)` is load bearing: without it +/// `single_distinct_aggregation_to_group_by` rewrites the distinct aggregate +/// into a plain two stage `GROUP BY`, which does not use these accumulators +/// at all. +#[tokio::test] +async fn group_by_count_distinct_utf8() { + TestCase::new() + .with_query( + "select group_key, count(distinct value), count(*) from t group by group_key", + ) + .with_scenario(Scenario::GroupedDistinctStrings { + groups: 4_000, + string_view: false, + }) + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_memory_limit(8_000_000) + .with_expected_success() + .run() + .await +} + +/// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering +/// the separate view flavoured hash set. The same query over a `Utf8View` +/// column needed about 123 MB of budget before the change and about 2.7 MB +/// after, so 16 MB separates the two. +#[tokio::test] +async fn group_by_count_distinct_utf8_view() { + TestCase::new() + .with_query( + "select group_key, count(distinct value), count(*) from t group by group_key", + ) + .with_scenario(Scenario::GroupedDistinctStrings { + groups: 4_000, + string_view: true, + }) + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_memory_limit(16_000_000) + .with_expected_success() + .run() + .await +} + #[tokio::test] async fn join_by_key_multiple_partitions() { let config = SessionConfig::new().with_target_partitions(2); @@ -982,6 +1038,14 @@ enum Scenario { /// If true, splits all input batches into 1 row each single_row_batches: bool, }, + + /// `groups` distinct integer keys paired with a short string value, for + /// grouped aggregates that build one accumulator per group. + GroupedDistinctStrings { + groups: usize, + /// If true, the value column is `Utf8View` rather than `Utf8` + string_view: bool, + }, } impl Scenario { @@ -1056,6 +1120,15 @@ impl Scenario { let table = SortedTableProvider::new(batches, sort_information); Arc::new(table) } + Self::GroupedDistinctStrings { + groups, + string_view, + } => { + let batches = grouped_distinct_string_batches(*groups, *string_view); + let table = + MemTable::try_new(batches[0].schema(), vec![batches]).unwrap(); + Arc::new(table) + } } } @@ -1079,10 +1152,60 @@ impl Scenario { // Use default rules None } + Self::GroupedDistinctStrings { .. } => { + // Disable the rules that would add a repartition, so the test + // measures the aggregate's budget rather than a repartition's + Some(vec![Arc::new(JoinSelection::new())]) + } } } } +/// Number of distinct string values held by every group produced by +/// [`grouped_distinct_string_batches`] +const DISTINCT_VALUES_PER_GROUP: usize = 2; + +/// Returns batches of 1024 rows with `groups` distinct keys in `group_key`, +/// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings +/// in `value`. The values are `Utf8View` if `string_view` is set, `Utf8` +/// otherwise. +fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec { + let value_type = if string_view { + DataType::Utf8View + } else { + DataType::Utf8 + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_key", DataType::Int32, false), + Field::new("value", value_type, false), + ])); + + const ROWS_PER_BATCH: usize = 1024; + + let rows = groups * DISTINCT_VALUES_PER_GROUP; + let mut keys = Vec::with_capacity(rows); + let mut values = Vec::with_capacity(rows); + for value in 0..DISTINCT_VALUES_PER_GROUP { + for group in 0..groups { + keys.push(group as i32); + values.push(format!("value-{value}")); + } + } + + keys.chunks(ROWS_PER_BATCH) + .zip(values.chunks(ROWS_PER_BATCH)) + .map(|(keys, values)| { + let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec())); + let values: ArrayRef = if string_view { + Arc::new(StringViewArray::from_iter_values(values)) + } else { + Arc::new(StringArray::from_iter_values(values)) + }; + RecordBatch::try_new(Arc::clone(&schema), vec![keys, values]).unwrap() + }) + .collect() +} + fn access_log_batches() -> Vec { AccessLogGenerator::new() .with_row_limit(1000) From 37d3ac299cefd6fb923718b6e629e0f5e9d44633 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:40:23 -0500 Subject: [PATCH 6/7] test: keep the memory limit tests off the single-distinct rewrite The two grouped `COUNT(DISTINCT )` memory limit tests only reach the per group accumulators while `single_distinct_aggregation_to_group_by` declines to rewrite the query. They leant on `count(*)` for that, which the rule rejects only because `count` is missing from the `sum`/`min`/`max` allow list. apache/datafusion#24859 proposes adding `count` to that list, which would rewrite the query, remove the accumulators, and leave both tests passing at any memory limit while still looking like they test something. Aggregate `avg(payload)` over a new `Int64` column instead. `avg` cannot be added to that list: the rule re-aggregates its own partial results over the deduplicated inner group by, and averaging per group averages of different sizes gives the wrong answer. That is why ClickBench Q9 keeps its distinct aggregate under #24859. Verified from the physical plan with #24859 cherry-picked on top of this branch: the `avg` query still plans as `aggr=[count(DISTINCT t.value), avg(t.payload)]`, while the `count(*)` query becomes `aggr=[count(alias1), sum(alias2)]` over an inner `GROUP BY group_key, value`, and drops from needing ~1.9 MB to ~0.9 MB. Re-swept both thresholds against the base commit. `Utf8` needs ~35.5 MB before and ~1.9 MB after; `Utf8View` needs ~123 MB before and ~2.5 MB after, so the 8 MB and 16 MB limits keep at least 4x margin on each side and are unchanged. --- datafusion/core/tests/memory_limit/mod.rs | 35 ++++++++++++++++------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 1527eb39c3a7..bc0bc827358a 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -27,7 +27,8 @@ mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; use arrow::array::{ - ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringArray, StringViewArray, + ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringArray, + StringViewArray, }; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; @@ -139,15 +140,23 @@ async fn group_by_hash() { /// success after it. Spilling is disabled, so completing means the query /// genuinely fit in the budget. /// -/// The `count(*)` is load bearing: without it -/// `single_distinct_aggregation_to_group_by` rewrites the distinct aggregate -/// into a plain two stage `GROUP BY`, which does not use these accumulators -/// at all. +/// The `avg(payload)` is load bearing, and `avg` specifically. Without a +/// second aggregate, `single_distinct_aggregation_to_group_by` rewrites the +/// distinct aggregate into a plain two stage `GROUP BY` that does not use +/// these accumulators at all. That rule tolerates a non-distinct `sum`, `min` +/// or `max` beside the distinct aggregate, because it re-aggregates its own +/// partial results over the deduplicated inner group by, and those three +/// compose with themselves. `avg` does not: averaging per group averages of +/// different sizes gives the wrong answer, so the rule can never accept it. +/// That is why ClickBench Q9 keeps its distinct aggregate. Do not replace +/// this with `count(*)`: `count` is only incidentally rejected today, and +/// proposes accepting it, +/// which would rewrite the query and leave this test passing by construction. #[tokio::test] async fn group_by_count_distinct_utf8() { TestCase::new() .with_query( - "select group_key, count(distinct value), count(*) from t group by group_key", + "select group_key, count(distinct value), avg(payload) from t group by group_key", ) .with_scenario(Scenario::GroupedDistinctStrings { groups: 4_000, @@ -162,13 +171,13 @@ async fn group_by_count_distinct_utf8() { /// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering /// the separate view flavoured hash set. The same query over a `Utf8View` -/// column needed about 123 MB of budget before the change and about 2.7 MB +/// column needed about 123 MB of budget before the change and about 2.5 MB /// after, so 16 MB separates the two. #[tokio::test] async fn group_by_count_distinct_utf8_view() { TestCase::new() .with_query( - "select group_key, count(distinct value), count(*) from t group by group_key", + "select group_key, count(distinct value), avg(payload) from t group by group_key", ) .with_scenario(Scenario::GroupedDistinctStrings { groups: 4_000, @@ -1167,8 +1176,8 @@ const DISTINCT_VALUES_PER_GROUP: usize = 2; /// Returns batches of 1024 rows with `groups` distinct keys in `group_key`, /// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings -/// in `value`. The values are `Utf8View` if `string_view` is set, `Utf8` -/// otherwise. +/// in `value` and an `Int64` `payload` to aggregate over. The values are +/// `Utf8View` if `string_view` is set, `Utf8` otherwise. fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec { let value_type = if string_view { DataType::Utf8View @@ -1178,6 +1187,7 @@ fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec Vec Date: Wed, 2 Sep 2026 17:54:06 -0500 Subject: [PATCH 7/7] fix: grow the value buffer on a power of two ladder `ArrowBytesMap::new` starts its value buffer empty and `ArrowBytesMap::with_capacity` starts it at `INITIAL_BUFFER_CAPACITY`. `Vec` then doubles from wherever its first allocation landed, so the two sit on different ladders and can hold the same values at capacities differing by up to 2x, in either direction depending on the value lengths. Measured on 500,000 distinct 28 byte values, the lazily grown map reported 52,494,344 bytes against 45,154,312 for a pre-allocated one, 16% more for identical contents. That matters because the ungrouped `COUNT(DISTINCT )` accumulator is the caller that had a use for the warm up: it builds one map and grows it to hold every distinct value in the input. Rounding every buffer growth up to a power of two puts both constructors on one ladder, so a lazily allocated map is never larger than a pre-allocated one holding the same values. Growth stays geometric, so appending is still amortized constant time. `ArrowBytesViewMap` has no such buffer and is unaffected. Two new tests cover the ungrouped path, which had none: `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set` and its `Utf8View` counterpart drive an accumulator to 0 through 500,000 distinct values and assert it is strictly cheaper than a pre-allocated set at per group cardinalities and exactly equal at ungrouped ones. The `Utf8` one fails without this change, at 1,000 distinct values, with the lazy set reporting 110,408 bytes against 96,072. Two map level tests pin the ladder itself. --- .../src/aggregate/count_distinct/bytes.rs | 149 ++++++++++++++++++ .../physical-expr-common/src/binary_map.rs | 76 ++++++++- 2 files changed, 223 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs index 3aa60f6f3b6a..d955d343ad62 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs @@ -43,6 +43,14 @@ impl BytesDistinctCountAccumulator { /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a /// high cardinality key holds hundreds of thousands of these at once and /// most of them see only a handful of values. + /// + /// The ungrouped path builds one of these and grows it to hold every + /// distinct value in the input, so it is the caller that had a use for the + /// warm up. It loses nothing here: the set grows into exactly the + /// capacities a pre-allocated one reaches, which + /// `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set` + /// pins. That is why this constructor needs no signal distinguishing the + /// two callers. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesSet::new(output_type)) } @@ -157,3 +165,144 @@ impl Accumulator for BytesViewDistinctCountAccumulator { size_of_val(self) + self.0.size() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{StringArray, StringViewArray}; + use datafusion_physical_expr_common::binary_map::INITIAL_MAP_CAPACITY; + use datafusion_physical_expr_common::binary_view_map::INITIAL_MAP_CAPACITY as INITIAL_VIEW_MAP_CAPACITY; + use std::sync::Arc; + + /// The batch size the aggregate stream drives an ungrouped accumulator with. + const BATCH_SIZE: usize = 8192; + + /// Distinct value counts spanning both sides of the warm up capacity, up to + /// ones where the two constructors have converged. + const CARDINALITIES: [usize; 7] = [0, 1, 100, 1_000, 10_000, 100_000, 500_000]; + + /// Cardinalities small enough that the warm up dominates what the set + /// holds. This is the per group population, where `GroupsAccumulatorAdapter` + /// holds one accumulator per group and most see a handful of values. + const PER_GROUP_SCALE: usize = 100; + + /// Cardinalities at which a lazily built set has grown into exactly the + /// capacities a pre-allocated one reaches. This is the ungrouped + /// population, one set holding every distinct value in the input. + const UNGROUPED_SCALE: usize = 10_000; + + /// Longer than the map's inline value length, so the value lands in the + /// value buffer rather than inside the hash table entry. + fn distinct_value(i: usize) -> String { + format!("distinct value number {i}") + } + + fn batches(distinct_values: usize, view: bool) -> Vec { + (0..distinct_values) + .step_by(BATCH_SIZE) + .map(|start| { + let values = (start..(start + BATCH_SIZE).min(distinct_values)) + .map(distinct_value); + if view { + Arc::new(StringViewArray::from_iter_values(values)) as ArrayRef + } else { + Arc::new(StringArray::from_iter_values(values)) as ArrayRef + } + }) + .collect() + } + + /// The property that decides whether the ungrouped path can afford to share + /// the lazy constructor with the per group path. + fn assert_lazy_is_not_worse( + distinct_values: usize, + lazy_size: usize, + pre_allocated_size: usize, + ) { + // The guarantee that lets both paths share one lazy constructor. + assert!( + lazy_size <= pre_allocated_size, + "at {distinct_values} distinct values the lazy set reported \ + {lazy_size} bytes against the {pre_allocated_size} bytes the \ + pre-allocated one reported" + ); + + if distinct_values <= PER_GROUP_SCALE { + // The warm up is pure overhead here, and removing it is the whole + // point of the change. + assert!( + lazy_size < pre_allocated_size, + "at {distinct_values} distinct values the lazy set should be \ + strictly cheaper, but reported {lazy_size} bytes against \ + {pre_allocated_size}" + ); + } + + if distinct_values >= UNGROUPED_SCALE { + // The hash table's bucket count is a power of two fixed by the + // number of entries, and the value buffer grows on a power of two + // ladder, so a set that starts empty lands on exactly the + // capacities a pre-allocated one reaches. The ungrouped path gives + // up nothing by starting empty, which is why these accumulators + // need no signal telling them apart from the per group ones. + assert_eq!( + lazy_size, pre_allocated_size, + "at {distinct_values} distinct values the lazy and pre-allocated \ + sets should have converged" + ); + } + } + + /// An ungrouped `COUNT(DISTINCT )` builds a single accumulator that + /// grows to hold every distinct value in its input, which is the population + /// the warm up existed for. It must be no more expensive without one. + #[test] + fn ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set() { + for distinct_values in CARDINALITIES { + let mut accumulator = + BytesDistinctCountAccumulator::::new(OutputType::Utf8); + let mut pre_allocated = ArrowBytesSet::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + for batch in batches(distinct_values, false) { + accumulator.update_batch(&[Arc::clone(&batch)]).unwrap(); + pre_allocated.insert(&batch); + } + + assert_eq!(accumulator.0.non_null_len(), distinct_values); + assert_lazy_is_not_worse( + distinct_values, + accumulator.0.size(), + pre_allocated.size(), + ); + } + } + + /// The `Utf8View` counterpart of + /// [`ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`]. + #[test] + fn ungrouped_utf8_view_accumulator_is_never_worse_than_a_pre_allocated_set() { + for distinct_values in CARDINALITIES { + let mut accumulator = + BytesViewDistinctCountAccumulator::new(OutputType::Utf8View); + let mut pre_allocated = ArrowBytesViewSet::with_capacity( + OutputType::Utf8View, + INITIAL_VIEW_MAP_CAPACITY, + ); + + for batch in batches(distinct_values, true) { + accumulator.update_batch(&[Arc::clone(&batch)]).unwrap(); + pre_allocated.insert(&batch); + } + + assert_eq!(accumulator.0.non_null_len(), distinct_values); + assert_lazy_is_not_worse( + distinct_values, + accumulator.0.size(), + pre_allocated.size(), + ); + } + } +} diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index bc7d98717377..4028520c776d 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -259,6 +259,28 @@ pub const INITIAL_MAP_CAPACITY: usize = 128; /// The size, in bytes, of the string data buffer pre-allocated by /// [`ArrowBytesMap::with_capacity`] pub const INITIAL_BUFFER_CAPACITY: usize = 8 * 1024; + +/// Appends `value` to a map's value buffer, growing the buffer on a power of +/// two ladder. +/// +/// `Vec` doubles from wherever its first allocation landed, so a buffer started +/// empty by [`ArrowBytesMap::new`] and one started at +/// [`INITIAL_BUFFER_CAPACITY`] by [`ArrowBytesMap::with_capacity`] sit on +/// different ladders, and can hold the same values at capacities differing by +/// up to 2x in either direction depending on the value lengths. Rounding every +/// growth up to a power of two puts both on one ladder, which is what makes a +/// lazily allocated map never larger than a pre-allocated one holding the same +/// values. Growth stays geometric, so appending is still amortized constant +/// time. +fn push_value_bytes(buffer: &mut Vec, value: &[u8]) { + let required = buffer.len() + value.len(); + if required > buffer.capacity() { + let target = required.checked_next_power_of_two().unwrap_or(required); + buffer.reserve_exact(target - buffer.len()); + } + buffer.extend_from_slice(value); +} + impl ArrowBytesMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, @@ -471,7 +493,7 @@ where // Put the small values into buffer and offsets so it appears // the output array, but store the actual bytes inline for // comparison - self.buffer.extend_from_slice(value); + push_value_bytes(&mut self.buffer, value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); let new_header = Entry { @@ -509,7 +531,7 @@ where // appears the output array, and store that offset // so the bytes can be compared if needed let offset = self.buffer.len(); // offset of start for data - self.buffer.extend_from_slice(value); + push_value_bytes(&mut self.buffer, value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); @@ -826,6 +848,56 @@ mod tests { assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); } + #[test] + fn lazy_and_pre_allocated_buffers_grow_on_the_same_ladder() { + // Value lengths chosen so the buffer requirement lands between powers + // of two, which is where the two ladders used to diverge. + for value_len in [9usize, 13, 24, 37] { + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + let mut pre_allocated = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + for batch in 0..8 { + let values: ArrayRef = + Arc::new(StringArray::from_iter_values((0..1_000).map(|i| { + let value = format!("{}:{i}", batch * 1_000 + i); + format!("{value:value_len$}") + }))); + lazy.insert_if_new(&values, |_| (), |_| ()); + pre_allocated.insert_if_new(&values, |_| (), |_| ()); + + assert_eq!(lazy.buffer.len(), pre_allocated.buffer.len()); + assert_eq!( + lazy.buffer.capacity(), + pre_allocated.buffer.capacity(), + "value length {value_len}, batch {batch}: a buffer that \ + started empty reached {} bytes of capacity against {} for \ + one that started at INITIAL_BUFFER_CAPACITY", + lazy.buffer.capacity(), + pre_allocated.buffer.capacity(), + ); + } + } + } + + #[test] + fn a_lazy_buffer_stays_below_the_pre_allocated_floor_while_it_is_small() { + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..10).map(|i| format!("distinct value number {i}")), + )); + lazy.insert_if_new(&values, |_| (), |_| ()); + + assert!( + lazy.buffer.capacity() < INITIAL_BUFFER_CAPACITY, + "expected a small lazy buffer to stay under the {INITIAL_BUFFER_CAPACITY} \ + byte pre-allocation, got {}", + lazy.buffer.capacity() + ); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8);