diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 15b224d200bf4..bc0bc827358ab 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,7 +26,10 @@ 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, Int64Array, RecordBatch, StringArray, + StringViewArray, +}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; @@ -125,6 +128,68 @@ 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 `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), avg(payload) 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.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), avg(payload) 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 +1047,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 +1129,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,8 +1161,62 @@ 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` 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 + } else { + DataType::Utf8 + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_key", DataType::Int32, false), + Field::new("value", value_type, false), + Field::new("payload", DataType::Int64, 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 payloads: ArrayRef = + Arc::new(Int64Array::from_iter_values(keys.iter().map(|k| *k as i64))); + 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, payloads]) + .unwrap() + }) + .collect() } fn access_log_batches() -> Vec { 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..d955d343ad629 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,18 @@ 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. + /// + /// 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)) } @@ -100,6 +112,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)) } @@ -151,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/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 7543e6b297329..4028520c776d0 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; @@ -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,8 +228,13 @@ 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, + /// 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 @@ -236,20 +252,71 @@ 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; + +/// 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, { + /// 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), - map_size: 0, - 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![], @@ -260,11 +327,31 @@ 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 } + /// 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. @@ -406,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 { @@ -415,11 +502,8 @@ 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 } } @@ -447,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)); @@ -457,11 +541,8 @@ 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 +567,8 @@ where let Self { output_type, map: _, - map_size: _, + initial_map_capacity: _, + initial_buffer_capacity: _, offsets, buffer, random_state: _, @@ -591,7 +673,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 +701,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) @@ -663,6 +749,155 @@ mod tests { use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; use std::collections::HashMap; + /// 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] + 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 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 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); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 0457825decb96..29f4014c5f9a4 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 @@ -40,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]>) {} @@ -55,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` @@ -129,8 +137,10 @@ 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, + /// 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, @@ -151,21 +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 { - let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY); - let map_size = map.capacity() * size_of::>(); + 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, - map_size, + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -179,11 +204,27 @@ 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 } + /// 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. @@ -374,8 +415,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 +580,19 @@ 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 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_size + map_size + views_size + in_progress_size + completed_size @@ -562,7 +608,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 +643,7 @@ where mod tests { use arrow::array::{GenericByteViewArray, StringViewArray}; use datafusion_common::HashMap; + use std::mem::size_of; use super::*; @@ -785,11 +832,104 @@ mod tests { assert_eq!(set.len(), 10); } + /// 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] - fn test_size_counts_initial_hash_table_capacity() { + fn map_new_does_not_allocate() { let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); - assert_eq!(map.size(), map.map.capacity() * size_of::>()); + 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()); + + // 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 {}", + map.size(), + map.map.capacity() * size_of::>() + ); + } + + #[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 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] @@ -822,7 +962,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 +972,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::() 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..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 @@ -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, } } @@ -133,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 997a7ce166a71..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 @@ -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, } } @@ -135,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); } }