From ca7a86a42bfe3f9922a591694ae5736146ecc9ab 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/8] 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 7543e6b297329..ceae72d4ccafe 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 0457825decb96..304b321870467 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 8b2f71217fea8fc3b3425f93a9b70524f1e1bcae 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/8] 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 f6df4182a879b..3aa60f6f3b6ad 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 7c8cdc3b4c50e..68351a839554a 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 ceae72d4ccafe..61b1e6e008c7c 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 304b321870467..9351d92092984 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 34ec36be31d2e..87eeba489ecc9 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 997a7ce166a71..f5183ba329b30 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 7cccdd244f2e3e62b34e93a1e4699b2970284b26 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/8] 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 61b1e6e008c7c..0b59bcf5937bc 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 9351d92092984..35f45e8e9a12e 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 a4f43dc83369f238730ee0202c54fcc821313d4f 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/8] 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 0b59bcf5937bc..bc7d987173770 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 35f45e8e9a12e..29f4014c5f9a4 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 87eeba489ecc9..9f7b4b4e91cba 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 f5183ba329b30..23ea4e7ed3f88 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 b0cd6fba5927059a58cb0609c26d3479c2eae3d6 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/8] 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 9172e34a87c79..274ace182bf13 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}; @@ -220,6 +222,60 @@ mod count_distinct_spill { } } +/// 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); @@ -1077,6 +1133,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 { @@ -1151,6 +1215,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) + } } } @@ -1174,10 +1247,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 e55a3aa15343bd4083bea7e5c73e86286222a5aa 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/8] 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 274ace182bf13..f873f956010e1 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}; @@ -234,15 +235,23 @@ mod count_distinct_spill { /// 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, @@ -257,13 +266,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, @@ -1262,8 +1271,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 @@ -1273,6 +1282,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/8] 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 3aa60f6f3b6ad..d955d343ad629 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 bc7d987173770..4028520c776d0 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); From 7a1a4df754b078abc0371905ee7099725f2b47c9 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 6 Sep 2026 14:53:48 -0700 Subject: [PATCH 8/8] test: boundary tests for `push_value_bytes` Pin the buffer growth helper's behavior directly, rather than only through the higher level convergence test that compares capacities after 1000 value batches: an empty append, filling exactly to a power of two capacity, appending one byte past it, a single append larger than the 8192 byte pre-allocation, and a buffer that already has room not being reallocated. Co-Authored-By: Claude Opus 5 --- .../physical-expr-common/src/binary_map.rs | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 4028520c776d0..16e00b6549fb4 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -848,6 +848,91 @@ mod tests { assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); } + #[test] + fn push_value_bytes_ignores_an_empty_value() { + let mut buffer = Vec::new(); + push_value_bytes(&mut buffer, b""); + assert_eq!(buffer.capacity(), 0); + assert!(buffer.is_empty()); + + push_value_bytes(&mut buffer, b"abcd"); + let capacity = buffer.capacity(); + push_value_bytes(&mut buffer, b""); + assert_eq!(buffer, b"abcd"); + assert_eq!(buffer.capacity(), capacity); + } + + #[test] + fn push_value_bytes_grows_only_past_a_power_of_two() { + let mut buffer = Vec::new(); + + // A first append that is not itself a power of two still lands on the + // ladder: a plain `Vec` would have stopped at a capacity of 9. + push_value_bytes(&mut buffer, &[1u8; 9]); + assert_eq!(buffer.len(), 9); + assert_eq!(buffer.capacity(), 16); + + // Filling exactly to the capacity does not grow it. + push_value_bytes(&mut buffer, &[2u8; 7]); + assert_eq!(buffer.len(), 16); + assert_eq!(buffer.capacity(), 16); + + // The next byte moves up one rung. + push_value_bytes(&mut buffer, &[3u8; 1]); + assert_eq!(buffer.len(), 17); + assert_eq!(buffer.capacity(), 32); + + // So does an append that overshoots several rungs at once. + push_value_bytes(&mut buffer, &[4u8; 100]); + assert_eq!(buffer.len(), 117); + assert_eq!(buffer.capacity(), 128); + + let mut expected = vec![1u8; 9]; + expected.extend_from_slice(&[2u8; 7]); + expected.extend_from_slice(&[3u8; 1]); + expected.extend_from_slice(&[4u8; 100]); + assert_eq!(buffer, expected); + } + + #[test] + fn push_value_bytes_does_not_grow_a_buffer_with_room() { + let mut buffer = Vec::with_capacity(INITIAL_BUFFER_CAPACITY); + let start = buffer.as_ptr(); + + let value = vec![7u8; 512]; + for _ in 0..16 { + push_value_bytes(&mut buffer, &value); + } + + assert_eq!(buffer.len(), INITIAL_BUFFER_CAPACITY); + assert_eq!(buffer.capacity(), INITIAL_BUFFER_CAPACITY); + assert!( + std::ptr::eq(buffer.as_ptr(), start), + "buffer was reallocated" + ); + assert!(buffer.iter().all(|&byte| byte == 7)); + } + + #[test] + fn push_value_bytes_rounds_a_single_oversized_append_up_to_a_power_of_two() { + let value = vec![9u8; INITIAL_BUFFER_CAPACITY + 1]; + let rounded = (INITIAL_BUFFER_CAPACITY + 1).next_power_of_two(); + + // A value larger than the pre-allocation lands on the same rung + // whether the buffer started empty or pre-allocated, which is what + // keeps a lazily allocated map from outgrowing a pre-allocated one. + let mut lazy = Vec::new(); + push_value_bytes(&mut lazy, &value); + assert_eq!(lazy.len(), value.len()); + assert_eq!(lazy.capacity(), rounded); + assert_eq!(lazy, value); + + let mut pre_allocated = Vec::with_capacity(INITIAL_BUFFER_CAPACITY); + push_value_bytes(&mut pre_allocated, &value); + assert_eq!(pre_allocated.capacity(), rounded); + assert_eq!(pre_allocated, value); + } + #[test] fn lazy_and_pre_allocated_buffers_grow_on_the_same_ladder() { // Value lengths chosen so the buffer requirement lands between powers