From 900ca2404b6b343ecfd2833a942c9063e8a44ba3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Wed, 12 Aug 2026 23:54:35 +0200 Subject: [PATCH 1/4] perf: direct mapped group values for dense primitive group by keys Single column primitive group by currently hashes every row. When the values are integers spanning a narrow range - ids, dates, small enums, dictionary codes - the group index can be looked up by indexing a vector with `value - min` instead, which removes hashing from the hot path. The direct mapped table is not chosen up front: the values are hashed as before while their range is tracked, and the table is only built once the groups fill at least 1/8 of the range they span (small tables are always worth it). Deciding from the first batch alone would be wrong in both directions, since a batch of a dense column looks sparse simply because it holds a fraction of the values. Values outside the range grow the table while it stays dense enough, otherwise the groups move back to the hash table for good. Slot lookup is a single wrapping subtraction on a bias-mapped u64, so values below the table are rejected by the same bounds check as values above it. Group values stay in `values` in group index order, so emitting is unchanged. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XYz9cxwBdfZLvcKmWcoiF9 --- .../group_values/single_group_by/primitive.rs | 757 ++++++++++++++++-- .../physical-plan/src/aggregates/mod.rs | 4 +- 2 files changed, 682 insertions(+), 79 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index e254aebcfd7ce..79e9bc0cfdb5d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -18,8 +18,8 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; use arrow::array::{ - ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, - cast::AsArray, + Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, + PrimitiveArray, cast::AsArray, }; use arrow::datatypes::{DataType, i256}; use datafusion_common::Result; @@ -96,26 +96,194 @@ macro_rules! hash_float { hash_float!(f16, f32, f64); -/// A [`GroupValues`] storing a single column of primitive values +/// A trait to allow direct mapped ("dense") lookup of integer like values /// -/// This specialization is significantly faster than using the more general -/// purpose `Row`s format -pub struct GroupValuesPrimitive { - /// The data type of the output array - data_type: DataType, +/// Values that map to a dense key can be interned by indexing a vector of +/// group indices with `key - min`, which avoids hashing entirely +pub trait DenseKey: Copy { + /// The dense key of this value, or `None` if this type does not support + /// direct mapped lookup + /// + /// This is `i128` so that the range of every supported type fits without + /// overflow (`i64::MIN..i64::MAX` as well as `0..u64::MAX`). It is only + /// used to size the table, never per row. + fn dense_key(self) -> Option; + + /// The value mapped into `u64` such that the order of, and distance + /// between, values is preserved + /// + /// Slot lookup is `biased(value) - biased(min)`, which is a single + /// wrapping subtraction: values below `min` wrap around to a large number + /// and are rejected by the same bounds check that catches values above the + /// end of the table. + fn biased(self) -> u64; +} + +macro_rules! dense_key_signed { + ($($t:ty),+) => { + $(impl DenseKey for $t { + #[inline] + fn dense_key(self) -> Option { + Some(self as i128) + } + + #[inline] + fn biased(self) -> u64 { + (self as i64 as u64) ^ (1 << 63) + } + })+ + }; +} +dense_key_signed!(i8, i16, i32, i64); + +macro_rules! dense_key_unsigned { + ($($t:ty),+) => { + $(impl DenseKey for $t { + #[inline] + fn dense_key(self) -> Option { + Some(self as i128) + } + + #[inline] + fn biased(self) -> u64 { + self as u64 + } + })+ + }; +} +dense_key_unsigned!(u8, u16, u32, u64); + +macro_rules! dense_key_unsupported { + ($($t:ty),+) => { + $(impl DenseKey for $t { + #[inline] + fn dense_key(self) -> Option { + None + } + + #[inline] + fn biased(self) -> u64 { + 0 + } + })+ + }; +} +dense_key_unsupported!(i128, i256, IntervalDayTime, IntervalMonthDayNano); +dense_key_unsupported!(f16, f32, f64); + +/// Marks an unused slot in the direct mapped table +const DENSE_EMPTY: u32 = u32::MAX; + +/// The most slots a direct mapped table may have (8MiB at 4 bytes per slot). +/// Ranges wider than this fall back to hashing +const DENSE_MAX_SLOTS: usize = 2 * 1024 * 1024; + +/// A direct mapped table larger than [`DENSE_SMALL_SLOTS`] is only built once +/// the groups fill at least this fraction (1/8) of the range they span, so +/// that sparse values are not given a mostly empty slot each +const DENSE_MIN_FILL_DENOM: usize = 8; + +/// A direct mapped table of at most this many slots (256KiB at 4 bytes each) +/// is built whatever the fill rate: even entirely empty it wastes little, and +/// waiting for a fill rate would give up the win on small group by keys +const DENSE_SMALL_SLOTS: usize = 64 * 1024; + +/// How the group index of each non null value is looked up +enum GroupStore { /// Stores the `(group_index, hash)` based on the hash of its value /// /// We also store `hash` is for reducing cost of rehashing. Such cost /// is obvious in high cardinality group by situation. /// More details can see: /// - map: HashTable<(usize, u64)>, + Hash(HashTable<(usize, u64)>), + /// Direct mapped lookup, where `group_ids[biased(value) - min_biased]` is + /// the group index of `value`, or [`DENSE_EMPTY`] if it has not been seen + /// + /// `min` is the same value as `min_biased`, kept as an `i128` for the + /// range arithmetic done once per batch + Dense { + min: i128, + min_biased: u64, + group_ids: Vec, + }, +} + +/// The result of interning a batch with the direct mapped table +enum DenseOutcome { + /// Every value of the batch was interned + Interned, + /// A value fell outside the range covered by the table + OutOfRange, +} + +/// The range of dense keys seen, and the biased form (see [`DenseKey::biased`]) +/// of its minimum, which is what slot lookup subtracts +#[derive(Clone, Copy)] +struct DenseRange { + min: i128, + max: i128, + min_biased: u64, +} + +impl DenseRange { + /// Widen the range to also cover `key`, whose biased form is `biased` + fn extend(self, key: i128, biased: u64) -> Self { + if key < self.min { + Self { + min: key, + max: self.max, + min_biased: biased, + } + } else { + Self { + max: self.max.max(key), + ..self + } + } + } + + /// Widen the range to also cover `other` + fn merge(self, other: Self) -> Self { + let (min, min_biased) = if other.min < self.min { + (other.min, other.min_biased) + } else { + (self.min, self.min_biased) + }; + Self { + min, + max: self.max.max(other.max), + min_biased, + } + } + + /// Number of slots needed to cover the range + fn len(self) -> Option { + usize::try_from(self.max.checked_sub(self.min)?.checked_add(1)?).ok() + } +} + +/// A [`GroupValues`] storing a single column of primitive values +/// +/// This specialization is significantly faster than using the more general +/// purpose `Row`s format +pub struct GroupValuesPrimitive { + /// The data type of the output array + data_type: DataType, + /// How the group index of each non null value is looked up + store: GroupStore, /// The group index of the null value if any null_group: Option, /// The values for each group index values: Vec, /// The random state used to generate hashes random_state: RandomState, + /// Set once the values have been found not to be dense, so that the + /// direct mapped table is not tried again + dense_disabled: bool, + /// The range of dense keys seen so far, used to decide whether a direct + /// mapped table is worth building + observed_range: Option, } impl GroupValuesPrimitive { @@ -123,27 +291,244 @@ impl GroupValuesPrimitive { assert!(PrimitiveArray::::is_compatible(&data_type)); Self { data_type, - map: HashTable::with_capacity(128), + store: GroupStore::Hash(HashTable::with_capacity(128)), values: Vec::with_capacity(128), null_group: None, random_state: crate::aggregates::AGGREGATION_HASH_SEED, + dense_disabled: false, + observed_range: None, } } } -impl GroupValues for GroupValuesPrimitive +impl GroupValuesPrimitive where - T::Native: HashValue, + T::Native: HashValue + DenseKey, { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { - assert_eq!(cols.len(), 1); - groups.clear(); + /// The range of dense keys in `array`, or `None` if the values do not + /// support direct mapped lookup or are all null + fn dense_range(array: &PrimitiveArray) -> Option { + let mut range: Option = None; + for value in array.iter().flatten() { + let key = value.dense_key()?; + range = Some(match range { + Some(seen) => seen.extend(key, value.biased()), + None => DenseRange { + min: key, + max: key, + min_biased: value.biased(), + }, + }); + } + range + } - for v in cols[0].as_primitive::() { + /// Migrate the groups interned so far into a direct mapped table, if the + /// values seen so far are dense enough to be worth a slot each + /// + /// The decision is deliberately not made from the first batch alone: a + /// batch of a dense column looks sparse simply because it only holds a + /// fraction of the column's values, and a batch of a sparse column can + /// look narrow enough to be worth it. Waiting until enough groups have + /// accumulated makes the fill rate meaningful, and costs only the range + /// tracking done while hashing. + fn try_migrate_to_dense(&mut self) { + debug_assert!(!self.dense_disabled); + let Some(range) = self.observed_range else { + return; + }; + let Some(len) = range.len().filter(|len| *len <= DENSE_MAX_SLOTS) else { + // Too wide a range to be worth a slot per value. It can only get + // wider, so stop considering it + self.dense_disabled = true; + return; + }; + + // A small table is always worth it. A larger one only once the values + // fill enough of the range they span; until then keep hashing, since + // the groups seen so far may just be a fraction of a dense column. + let groups = self.values.len() - usize::from(self.null_group.is_some()); + if len > DENSE_SMALL_SLOTS && groups * DENSE_MIN_FILL_DENOM < len { + return; + } + + let mut group_ids = vec![DENSE_EMPTY; len]; + for (group_idx, &value) in self.values.iter().enumerate() { + if Some(group_idx) == self.null_group { + continue; + } + let offset = value.biased().wrapping_sub(range.min_biased); + debug_assert!(offset < len as u64); + group_ids[offset as usize] = group_idx as u32; + } + + self.store = GroupStore::Dense { + min: range.min, + min_biased: range.min_biased, + group_ids, + }; + } + + /// Track the range of dense keys seen, which decides whether a direct + /// mapped table is worth building + fn observe_range(&mut self, array: &PrimitiveArray) { + if self.dense_disabled { + return; + } + let Some(range) = Self::dense_range(array) else { + // All null, or a type that cannot be direct mapped + if array.null_count() != array.len() { + self.dense_disabled = true; + } + return; + }; + self.observed_range = Some(match self.observed_range { + Some(seen) => seen.merge(range), + None => range, + }); + } + + /// Grow the direct mapped table so that it also covers `array`, returning + /// false if that would need more than [`DENSE_MAX_SLOTS`] slots + fn grow_dense(&mut self, array: &PrimitiveArray) -> bool { + let groups = self.values.len() - usize::from(self.null_group.is_some()); + let GroupStore::Dense { + min, + min_biased, + group_ids, + } = &mut self.store + else { + return false; + }; + let Some(batch) = Self::dense_range(array) else { + return false; + }; + + // The range the table covers today, widened to also cover the batch + let Some(covered_max) = (*min) + .checked_add(group_ids.len() as i128) + .map(|end| end - 1) + else { + return false; + }; + let grown = DenseRange { + min: *min, + max: covered_max, + min_biased: *min_biased, + } + .merge(batch); + let Some(len) = grown.len().filter(|len| *len <= DENSE_MAX_SLOTS) else { + return false; + }; + let new_min = grown.min; + // Growing must not leave the table mostly empty + if len > DENSE_SMALL_SLOTS && groups * DENSE_MIN_FILL_DENOM < len { + return false; + } + + if new_min < *min { + // Values below the current range: shift the existing slots up + let prefix = (*min - new_min) as usize; + let mut shifted = vec![DENSE_EMPTY; len]; + shifted[prefix..prefix + group_ids.len()].copy_from_slice(group_ids); + *group_ids = shifted; + *min = new_min; + *min_biased = grown.min_biased; + } else { + group_ids.resize(len, DENSE_EMPTY); + } + true + } + + /// Intern `array` with the direct mapped table + /// + /// Interning is idempotent, so a batch that reports [`DenseOutcome::OutOfRange`] + /// can simply be interned again once the table covers it + fn intern_dense( + &mut self, + array: &PrimitiveArray, + groups: &mut Vec, + ) -> DenseOutcome { + let Self { + store, + values, + null_group, + .. + } = self; + let GroupStore::Dense { + min_biased, + group_ids, + .. + } = store + else { + return DenseOutcome::OutOfRange; + }; + let min_biased = *min_biased; + let len = group_ids.len() as u64; + + for v in array { let group_id = match v { - None => *self.null_group.get_or_insert_with(|| { - let group_id = self.values.len(); - self.values.push(Default::default()); + None => *null_group.get_or_insert_with(|| { + let group_id = values.len(); + values.push(Default::default()); + group_id + }), + Some(key) => { + let offset = key.biased().wrapping_sub(min_biased); + if offset >= len { + return DenseOutcome::OutOfRange; + } + + let slot = &mut group_ids[offset as usize]; + if *slot == DENSE_EMPTY { + let group_id = values.len(); + *slot = group_id as u32; + values.push(key); + group_id + } else { + *slot as usize + } + } + }; + groups.push(group_id); + } + + DenseOutcome::Interned + } + + /// Rebuild the group values as a hash table, which is possible at any + /// point because [`Self::values`] holds every group value in group index + /// order + fn convert_dense_to_hash(&mut self) { + let state = &self.random_state; + let mut map = HashTable::with_capacity(self.values.len()); + for (group_idx, &value) in self.values.iter().enumerate() { + if Some(group_idx) == self.null_group { + continue; + } + let hash = value.hash(state); + map.insert_unique(hash, (group_idx, hash), |&(_, hash)| hash); + } + self.store = GroupStore::Hash(map); + } + + fn intern_hash(&mut self, array: &PrimitiveArray, groups: &mut Vec) { + let Self { + store, + values, + null_group, + random_state, + .. + } = self; + let GroupStore::Hash(map) = store else { + unreachable!("group values are not stored in a hash table") + }; + + for v in array { + let group_id = match v { + None => *null_group.get_or_insert_with(|| { + let group_id = values.len(); + values.push(Default::default()); group_id }), Some(key) => { @@ -151,12 +536,11 @@ where // so the bit-equal `is_eq` matches and the stored value is // the canonical representative. let key = key.canonicalize(); - let state = &self.random_state; - let hash = key.hash(state); - let insert = self.map.entry( + let hash = key.hash(random_state); + let insert = map.entry( hash, |&(g, h)| unsafe { - hash == h && self.values.get_unchecked(g).is_eq(key) + hash == h && values.get_unchecked(g).is_eq(key) }, |&(_, h)| h, ); @@ -164,9 +548,9 @@ where match insert { hashbrown::hash_table::Entry::Occupied(o) => o.get().0, hashbrown::hash_table::Entry::Vacant(v) => { - let g = self.values.len(); + let g = values.len(); v.insert((g, hash)); - self.values.push(key); + values.push(key); g } } @@ -174,11 +558,56 @@ where }; groups.push(group_id) } + } +} + +impl GroupValues for GroupValuesPrimitive +where + T::Native: HashValue + DenseKey, +{ + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + assert_eq!(cols.len(), 1); + let array = cols[0].as_primitive::(); + groups.clear(); + + // Only while hashing: once the values are direct mapped the range is + // maintained by the table itself, and once they are known not to be + // dense there is nothing left to decide + if !self.dense_disabled && matches!(self.store, GroupStore::Hash(_)) { + self.observe_range(array); + self.try_migrate_to_dense(); + } + + if matches!(self.store, GroupStore::Dense { .. }) { + if let DenseOutcome::Interned = self.intern_dense(array, groups) { + return Ok(()); + } + + // Values outside the range the table covers: grow it if that is + // still worthwhile, otherwise fall back to hashing for good + groups.clear(); + if self.grow_dense(array) { + if let DenseOutcome::Interned = self.intern_dense(array, groups) { + return Ok(()); + } + groups.clear(); + } + self.convert_dense_to_hash(); + self.dense_disabled = true; + } + + self.intern_hash(array, groups); Ok(()) } fn size(&self) -> usize { - self.map.capacity() * size_of::<(usize, u64)>() + self.values.allocated_size() + let store = match &self.store { + GroupStore::Hash(map) => map.capacity() * size_of::<(usize, u64)>(), + GroupStore::Dense { group_ids, .. } => { + group_ids.capacity() * size_of::() + } + }; + store + self.values.allocated_size() } fn is_empty(&self) -> bool { @@ -207,23 +636,41 @@ where let array: PrimitiveArray = match emit_to { EmitTo::All => { - self.map.clear(); + self.store = GroupStore::Hash(HashTable::with_capacity(0)); + self.observed_range = None; build_primitive(std::mem::take(&mut self.values), self.null_group.take()) } EmitTo::First(n) => { - self.map.retain(|entry| { - // Decrement group index by n - let group_idx = entry.0; - match group_idx.checked_sub(n) { - // Group index was >= n, shift value down - Some(sub) => { - entry.0 = sub; - true + match &mut self.store { + GroupStore::Hash(map) => { + map.retain(|entry| { + // Decrement group index by n + let group_idx = entry.0; + match group_idx.checked_sub(n) { + // Group index was >= n, shift value down + Some(sub) => { + entry.0 = sub; + true + } + // Group index was < n, so remove from table + None => false, + } + }); + } + GroupStore::Dense { group_ids, .. } => { + for slot in group_ids.iter_mut() { + if *slot == DENSE_EMPTY { + continue; + } + match (*slot as usize).checked_sub(n) { + // Group index was >= n, shift value down + Some(sub) => *slot = sub as u32, + // Group index was < n, so free the slot + None => *slot = DENSE_EMPTY, + } } - // Group index was < n, so remove from table - None => false, } - }); + } let null_group = match &mut self.null_group { Some(v) if *v >= n => { *v -= n; @@ -242,54 +689,210 @@ where fn clear_shrink(&mut self, num_rows: usize) { self.values.clear(); self.values.shrink_to(num_rows); - self.map.clear(); - self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + self.null_group = None; + self.observed_range = None; + match &mut self.store { + GroupStore::Hash(map) => { + map.clear(); + map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + } + GroupStore::Dense { .. } => { + self.store = GroupStore::Hash(HashTable::with_capacity(num_rows)); + } + } } } #[cfg(test)] mod tests { use super::*; - use arrow::array::types::Int32Type; - use arrow::array::{ArrayRef, Int32Array}; - use arrow::datatypes::DataType; - use datafusion_expr::EmitTo; - use std::sync::Arc; - - /// Mirror of the `EmitTo::take_needed` regression test, applied to the - /// concrete `GroupValuesPrimitive` accumulator. - /// - /// When `n` is small, the old `split_off(n) + swap` pattern used inside - /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation - /// and returned the emitted prefix carrying the original large backing. - /// - /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: - /// the emitted prefix gets a compact allocation and `self.values` retains the - /// original large one. + use arrow::array::{Array, Int64Array}; + use arrow::datatypes::Int64Type; + + fn is_dense(gv: &GroupValuesPrimitive) -> bool { + matches!(gv.store, GroupStore::Dense { .. }) + } + + fn array(values: &[Option]) -> ArrayRef { + Arc::new(Int64Array::from(values.to_vec())) as ArrayRef + } + + /// State of the group values at the point everything was interned, before + /// emitting drains it + struct Interned { + dense: bool, + num_groups: usize, + } + + /// Intern every batch, then emit, and check that every input row maps to a + /// group holding exactly that row's value. This holds no matter which + /// store the values ended up in. + fn check_intern(batches: &[Vec>]) -> Interned { + let mut gv = GroupValuesPrimitive::::new(DataType::Int64); + let mut all_rows = vec![]; + let mut all_groups = vec![]; + let mut groups = vec![]; + + for batch in batches { + gv.intern(&[array(batch)], &mut groups).unwrap(); + assert_eq!(groups.len(), batch.len()); + all_rows.extend(batch.iter().copied()); + all_groups.extend(groups.iter().copied()); + } + + let interned = Interned { + dense: is_dense(&gv), + num_groups: gv.len(), + }; + + let emitted = gv.emit(EmitTo::All).unwrap(); + let emitted = emitted[0].as_primitive::(); + assert_eq!(emitted.len(), all_groups.iter().max().map_or(0, |m| m + 1)); + + for (row, (value, group)) in all_rows.iter().zip(&all_groups).enumerate() { + let got = (!emitted.is_null(*group)).then(|| emitted.value(*group)); + assert_eq!( + got, *value, + "row {row} interned to group {group} holding {got:?}, expected {value:?}" + ); + } + + // The same value must always map to the same group + let mut seen = std::collections::HashMap::new(); + for (value, group) in all_rows.iter().zip(&all_groups) { + let first = seen.entry(*value).or_insert(*group); + assert_eq!(first, group, "value {value:?} mapped to two groups"); + } + + interned + } + #[test] - fn emit_first_small_n_allocates_minimally() -> Result<()> { - let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + fn dense_keys_use_direct_mapped_table() { + assert!(check_intern(&[vec![Some(10), Some(11), Some(12), Some(10)]]).dense); + } + + #[test] + fn dense_table_grows_for_ascending_keys() { + // Each batch extends the range upwards, so the table has to grow + let batches = (0..8) + .map(|batch| (0..1000).map(|v| Some(batch * 1000 + v)).collect()) + .collect::>(); + let interned = check_intern(&batches); + assert!(interned.dense); + assert_eq!(interned.num_groups, 8000); + } + + #[test] + fn dense_table_grows_downwards_for_descending_keys() { + let batches = (0..8) + .map(|batch| (0..1000).map(|v| Some(-(batch * 1000 + v))).collect()) + .collect::>(); + let interned = check_intern(&batches); + assert!(interned.dense); + assert_eq!(interned.num_groups, 8000); + } + + /// Values spread thinly over a wide range are not worth a slot each + #[test] + fn sparse_values_keep_hashing() { + let batches = (0..8) + .map(|batch| { + (0..1000) + .map(|v| Some((batch * 1000 + v) * 64)) + .collect::>() + }) + .collect::>(); + let interned = check_intern(&batches); + assert!(!interned.dense); + assert_eq!(interned.num_groups, 8000); + } - // Intern 20 distinct values; `new()` pre-allocates capacity 128 for `values`. - let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..20i32)); + /// A column whose values only look sparse until enough of them have been + /// seen still ends up direct mapped + #[test] + fn values_dense_over_a_wide_range_migrate() { + // 100k values scattered over a 100k range, 1000 at a time, so that no + // single batch fills much of the range + let batches = (0..100) + .map(|batch| { + (0..1000) + .map(|v| Some((v * 100 + batch) % 100_000)) + .collect::>() + }) + .collect::>(); + let interned = check_intern(&batches); + assert!(interned.dense); + assert_eq!(interned.num_groups, 100_000); + } + + #[test] + fn wide_range_falls_back_to_hashing() { + assert!(!check_intern(&[vec![Some(0), Some(i64::MAX), Some(i64::MIN)]]).dense); + } + + /// A batch that starts dense but later sees a value far outside the range + /// must fall back without losing or renumbering the groups it handed out + #[test] + fn fallback_after_dense_start_keeps_groups() { + let interned = check_intern(&[ + vec![Some(5), Some(6), Some(7)], + vec![Some(6), Some(50_000_000), Some(5)], + ]); + assert!(!interned.dense); + assert_eq!(interned.num_groups, 4); + } + + #[test] + fn nulls_get_their_own_group_in_dense_mode() { + let interned = check_intern(&[vec![None, Some(3), None, Some(4), Some(3)]]); + assert!(interned.dense); + assert_eq!(interned.num_groups, 3); + } + + #[test] + fn emit_first_reindexes_dense_groups() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int64); let mut groups = vec![]; - gv.intern(&[arr], &mut groups)?; - let capacity_before = gv.values.capacity(); // 128 - - // n=4, n*2=8 <= len=20 -> drain branch - let emitted = gv.emit(EmitTo::First(4))?; - - assert_eq!(emitted[0].len(), 4); - - // `self.values` must retain its original large allocation. - // Old split_off+swap left it with a fresh small allocation (~16). - assert_eq!( - gv.values.capacity(), - capacity_before, - "self.values capacity {} should equal original {} after small First(n) emit", - gv.values.capacity(), - capacity_before, - ); + + gv.intern(&[array(&[Some(10), None, Some(11), Some(12)])], &mut groups)?; + assert_eq!(groups, vec![0, 1, 2, 3]); + assert!(is_dense(&gv)); + + let emitted = gv.emit(EmitTo::First(2))?; + let emitted = emitted[0].as_primitive::(); + assert_eq!(emitted.len(), 2); + assert_eq!(emitted.value(0), 10); + assert!(emitted.is_null(1)); + assert_eq!(gv.len(), 2); + + // 11 and 12 keep their (shifted) groups, 10 was emitted so it is new + gv.intern(&[array(&[Some(11), Some(12), Some(10)])], &mut groups)?; + assert_eq!(groups, vec![0, 1, 2]); + + let emitted = gv.emit(EmitTo::All)?; + let emitted = emitted[0].as_primitive::(); + assert_eq!(emitted.values(), &[11, 12, 10]); + + Ok(()) + } + + #[test] + fn clear_shrink_resets_dense_table() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int64); + let mut groups = vec![]; + + gv.intern(&[array(&[None, Some(1)])], &mut groups)?; + gv.clear_shrink(0); + assert!(gv.is_empty()); + + gv.intern(&[array(&[Some(5), None])], &mut groups)?; + assert_eq!(groups, vec![0, 1]); + + let emitted = gv.emit(EmitTo::All)?; + let emitted = emitted[0].as_primitive::(); + assert_eq!(emitted.value(0), 5); + assert!(emitted.is_null(1)); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 1c5b56ffadbc0..d77c2aa9f65f6 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3511,7 +3511,7 @@ mod tests { let task_ctx = if spill { // set to an appropriate value to trigger spill - new_spill_ctx(2, 1600) + new_spill_ctx(2, 200) } else { Arc::new(TaskContext::default()) }; @@ -3577,7 +3577,7 @@ mod tests { let task_ctx = if spill { // enlarge memory limit to let the final aggregation finish - new_spill_ctx(2, 4640) + new_spill_ctx(2, 900) } else { Arc::clone(&task_ctx) }; From 12d6df51ac179cb840aa5e2b9426b61f54a04d7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 13 Aug 2026 08:54:06 +0200 Subject: [PATCH 2/4] fix: do not try to build a direct mapped table for non integer values Tracking the range of a batch can itself rule the values out, which is what happens on the first batch of a type that cannot be direct mapped at all, such as a float group by key. `try_migrate_to_dense` was called anyway and tripped its own debug assertion. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XYz9cxwBdfZLvcKmWcoiF9 --- .../group_values/single_group_by/primitive.rs | 34 +++++++++++++++++-- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 79e9bc0cfdb5d..12fc66e0cb6fe 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -333,7 +333,11 @@ where /// accumulated makes the fill rate meaningful, and costs only the range /// tracking done while hashing. fn try_migrate_to_dense(&mut self) { - debug_assert!(!self.dense_disabled); + // Tracking the range can itself rule the values out, e.g. the first + // batch of a type that cannot be direct mapped at all + if self.dense_disabled { + return; + } let Some(range) = self.observed_range else { return; }; @@ -706,8 +710,8 @@ where #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, Int64Array}; - use arrow::datatypes::Int64Type; + use arrow::array::{Array, Float64Array, Int64Array}; + use arrow::datatypes::{Float64Type, Int64Type}; fn is_dense(gv: &GroupValuesPrimitive) -> bool { matches!(gv.store, GroupStore::Dense { .. }) @@ -896,4 +900,28 @@ mod tests { Ok(()) } + + /// Types that cannot be direct mapped at all are ruled out while their + /// range is tracked, which must not be mistaken for a table worth building + #[test] + fn float_values_keep_hashing() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Float64); + let mut groups = vec![]; + + for batch in 0..3usize { + let values = (0..4).map(|v| Some((batch * 4 + v) as f64)); + let input = Arc::new(Float64Array::from_iter(values)) as ArrayRef; + gv.intern(&[input], &mut groups)?; + // Each batch holds four values not seen before + let first = batch * 4; + assert_eq!(groups, (first..first + 4).collect::>()); + } + + assert!(matches!(gv.store, GroupStore::Hash(_))); + + let emitted = gv.emit(EmitTo::All)?; + assert_eq!(emitted[0].len(), 12); + + Ok(()) + } } From 21387538ce656f6d1e0b51a5d2d195c66077a5df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 13 Aug 2026 10:48:58 +0200 Subject: [PATCH 3/4] refactor: simplify the direct mapped group values Review cleanups, no intended change in behaviour beyond the two fixes noted below: - Whether a type can be direct mapped is a compile time fact, so it is an associated const rather than a per value `Option`. Non integer types are now ruled out when the group values are created instead of by scanning a batch, which supersedes the earlier fix for float keys. - The range is tracked as the groups are created while hashing rather than by a second pass over every batch. The range of all the values equals the range of the distinct ones, so this is the same range for O(new groups) instead of O(rows), and removes the ~5% cost this used to add to columns that never become dense. - Slot lookup casts straight to `u64`: the bias cancels in the wrapping subtraction, so it only obscured that the join side maps keys the same way (see `ArrayMap`). - Growing and building the table were the same decision procedure twice; they are one function taking the range to cover. - `EmitTo::First` walks the groups rather than the slots, of which there can be many more. Fixes: values found not to be dense no longer keep hashing after the groups are drained, which matters because a stream is reused across spill cycles and the post spill keys arrive sorted. Slack added to a table is also part of the range it records, so groups created in that slack are still placed when the table is later rebuilt. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XYz9cxwBdfZLvcKmWcoiF9 --- .../group_values/single_group_by/primitive.rs | 556 +++++++++--------- 1 file changed, 268 insertions(+), 288 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 12fc66e0cb6fe..aea4278463ab1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -18,10 +18,10 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; use arrow::array::{ - Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, - PrimitiveArray, cast::AsArray, + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, + cast::AsArray, }; -use arrow::datatypes::{DataType, i256}; +use arrow::datatypes::{ArrowNativeType, DataType, i256}; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::utils::split_vec_min_alloc; @@ -96,96 +96,68 @@ macro_rules! hash_float { hash_float!(f16, f32, f64); -/// A trait to allow direct mapped ("dense") lookup of integer like values +/// A trait to allow direct mapped ("dense") lookup of integer values, by +/// indexing a vector of group indices with `value - min` instead of hashing. +/// The join side maps keys the same way, see [`ArrayMap`] /// -/// Values that map to a dense key can be interned by indexing a vector of -/// group indices with `key - min`, which avoids hashing entirely +/// [`ArrayMap`]: crate::joins::array_map::ArrayMap pub trait DenseKey: Copy { - /// The dense key of this value, or `None` if this type does not support - /// direct mapped lookup - /// - /// This is `i128` so that the range of every supported type fits without - /// overflow (`i64::MIN..i64::MAX` as well as `0..u64::MAX`). It is only - /// used to size the table, never per row. - fn dense_key(self) -> Option; + /// Whether this type can be direct mapped at all + const DENSE: bool; - /// The value mapped into `u64` such that the order of, and distance - /// between, values is preserved + /// The value as a `u64`, a plain cast for the supported types /// - /// Slot lookup is `biased(value) - biased(min)`, which is a single - /// wrapping subtraction: values below `min` wrap around to a large number - /// and are rejected by the same bounds check that catches values above the - /// end of the table. - fn biased(self) -> u64; -} - -macro_rules! dense_key_signed { - ($($t:ty),+) => { - $(impl DenseKey for $t { - #[inline] - fn dense_key(self) -> Option { - Some(self as i128) - } - - #[inline] - fn biased(self) -> u64 { - (self as i64 as u64) ^ (1 << 63) - } - })+ - }; + /// Slot lookup is the wrapping `index(value) - index(min)`, so values + /// below `min` wrap to a large number and are rejected by the same bounds + /// check as values above the table. The cast preserves distances, so this + /// holds for signed values too. + fn index(self) -> u64; } -dense_key_signed!(i8, i16, i32, i64); -macro_rules! dense_key_unsigned { +macro_rules! dense_key { ($($t:ty),+) => { $(impl DenseKey for $t { - #[inline] - fn dense_key(self) -> Option { - Some(self as i128) - } + const DENSE: bool = true; #[inline] - fn biased(self) -> u64 { + fn index(self) -> u64 { self as u64 } })+ }; } -dense_key_unsigned!(u8, u16, u32, u64); +dense_key!(i8, i16, i32, i64, u8, u16, u32, u64); -macro_rules! dense_key_unsupported { +/// Types too wide, or not integers, to be direct mapped +macro_rules! sparse_key { ($($t:ty),+) => { $(impl DenseKey for $t { - #[inline] - fn dense_key(self) -> Option { - None - } + const DENSE: bool = false; #[inline] - fn biased(self) -> u64 { + fn index(self) -> u64 { 0 } })+ }; } -dense_key_unsupported!(i128, i256, IntervalDayTime, IntervalMonthDayNano); -dense_key_unsupported!(f16, f32, f64); +sparse_key!(i128, i256, f16, f32, f64); +sparse_key!(IntervalDayTime, IntervalMonthDayNano); /// Marks an unused slot in the direct mapped table const DENSE_EMPTY: u32 = u32::MAX; -/// The most slots a direct mapped table may have (8MiB at 4 bytes per slot). -/// Ranges wider than this fall back to hashing +/// Memory budget for one table: 2M slots, 8MiB at 4 bytes each. Wider ranges +/// are hashed instead const DENSE_MAX_SLOTS: usize = 2 * 1024 * 1024; -/// A direct mapped table larger than [`DENSE_SMALL_SLOTS`] is only built once -/// the groups fill at least this fraction (1/8) of the range they span, so -/// that sparse values are not given a mostly empty slot each +/// Minimum fill (1/8) of its range for a table above [`DENSE_SMALL_SLOTS`], so +/// that sparse values are not given a mostly empty slot each. The join side +/// admits keys the same way, see `perfect_hash_join_min_key_density` const DENSE_MIN_FILL_DENOM: usize = 8; -/// A direct mapped table of at most this many slots (256KiB at 4 bytes each) -/// is built whatever the fill rate: even entirely empty it wastes little, and -/// waiting for a fill rate would give up the win on small group by keys +/// Tables this small (256KiB) are built whatever the fill rate: even empty they +/// waste little const DENSE_SMALL_SLOTS: usize = 64 * 1024; /// How the group index of each non null value is looked up @@ -197,69 +169,33 @@ enum GroupStore { /// More details can see: /// Hash(HashTable<(usize, u64)>), - /// Direct mapped lookup, where `group_ids[biased(value) - min_biased]` is - /// the group index of `value`, or [`DENSE_EMPTY`] if it has not been seen - /// - /// `min` is the same value as `min_biased`, kept as an `i128` for the - /// range arithmetic done once per batch - Dense { - min: i128, - min_biased: u64, - group_ids: Vec, - }, + /// Direct mapped lookup, where `group_ids[index(value) - min]` is the + /// group index of `value`, or [`DENSE_EMPTY`] if it has not been seen + Dense { min: u64, group_ids: Vec }, } -/// The result of interning a batch with the direct mapped table -enum DenseOutcome { - /// Every value of the batch was interned - Interned, - /// A value fell outside the range covered by the table - OutOfRange, -} - -/// The range of dense keys seen, and the biased form (see [`DenseKey::biased`]) -/// of its minimum, which is what slot lookup subtracts +/// The range of values seen, which decides whether a direct mapped table is +/// worth building #[derive(Clone, Copy)] -struct DenseRange { - min: i128, - max: i128, - min_biased: u64, +struct DenseRange { + min: N, + max: N, } -impl DenseRange { - /// Widen the range to also cover `key`, whose biased form is `biased` - fn extend(self, key: i128, biased: u64) -> Self { - if key < self.min { - Self { - min: key, - max: self.max, - min_biased: biased, - } - } else { - Self { - max: self.max.max(key), - ..self - } - } - } - - /// Widen the range to also cover `other` - fn merge(self, other: Self) -> Self { - let (min, min_biased) = if other.min < self.min { - (other.min, other.min_biased) - } else { - (self.min, self.min_biased) - }; - Self { - min, - max: self.max.max(other.max), - min_biased, +impl DenseRange { + /// Widen the range to also cover `value` + fn extend(&mut self, value: N) { + if value.is_lt(self.min) { + self.min = value; + } else if value.is_gt(self.max) { + self.max = value; } } /// Number of slots needed to cover the range fn len(self) -> Option { - usize::try_from(self.max.checked_sub(self.min)?.checked_add(1)?).ok() + let span = self.max.index().wrapping_sub(self.min.index()); + usize::try_from(span as u128 + 1).ok() } } @@ -278,15 +214,17 @@ pub struct GroupValuesPrimitive { values: Vec, /// The random state used to generate hashes random_state: RandomState, - /// Set once the values have been found not to be dense, so that the - /// direct mapped table is not tried again + /// Set once the values are known not to be dense, so that the direct + /// mapped table is not tried again until the groups are drained dense_disabled: bool, - /// The range of dense keys seen so far, used to decide whether a direct - /// mapped table is worth building - observed_range: Option, + /// The range of the group values, maintained while hashing + observed_range: Option>, } -impl GroupValuesPrimitive { +impl GroupValuesPrimitive +where + T::Native: HashValue + DenseKey, +{ pub fn new(data_type: DataType) -> Self { assert!(PrimitiveArray::::is_compatible(&data_type)); Self { @@ -295,7 +233,7 @@ impl GroupValuesPrimitive { values: Vec::with_capacity(128), null_group: None, random_state: crate::aggregates::AGGREGATION_HASH_SEED, - dense_disabled: false, + dense_disabled: !T::Native::DENSE, observed_range: None, } } @@ -305,172 +243,106 @@ impl GroupValuesPrimitive where T::Native: HashValue + DenseKey, { - /// The range of dense keys in `array`, or `None` if the values do not - /// support direct mapped lookup or are all null - fn dense_range(array: &PrimitiveArray) -> Option { - let mut range: Option = None; - for value in array.iter().flatten() { - let key = value.dense_key()?; - range = Some(match range { - Some(seen) => seen.extend(key, value.biased()), - None => DenseRange { - min: key, - max: key, - min_biased: value.biased(), - }, - }); - } - range + /// Number of groups holding a value, i.e. excluding the null group + fn value_groups(&self) -> usize { + self.values.len() - usize::from(self.null_group.is_some()) } - /// Migrate the groups interned so far into a direct mapped table, if the - /// values seen so far are dense enough to be worth a slot each + /// The value of each group holding one, in group index order + fn value_groups_iter(&self) -> impl Iterator + '_ { + self.values + .iter() + .enumerate() + .filter(|(group_idx, _)| Some(*group_idx) != self.null_group) + .map(|(group_idx, value)| (group_idx, *value)) + } + + /// Build a direct mapped table over `range` if it is worth a slot per + /// value, returning whether the group values are now direct mapped /// - /// The decision is deliberately not made from the first batch alone: a - /// batch of a dense column looks sparse simply because it only holds a - /// fraction of the column's values, and a batch of a sparse column can - /// look narrow enough to be worth it. Waiting until enough groups have - /// accumulated makes the fill rate meaningful, and costs only the range - /// tracking done while hashing. - fn try_migrate_to_dense(&mut self) { - // Tracking the range can itself rule the values out, e.g. the first - // batch of a type that cannot be direct mapped at all - if self.dense_disabled { - return; - } - let Some(range) = self.observed_range else { - return; - }; - let Some(len) = range.len().filter(|len| *len <= DENSE_MAX_SLOTS) else { - // Too wide a range to be worth a slot per value. It can only get - // wider, so stop considering it + /// Judging this on the groups accumulated so far rather than on one batch + /// matters: a batch of a dense column looks sparse simply because it holds + /// a fraction of the values, and a batch of a sparse one can look narrow. + fn try_build_dense( + &mut self, + mut range: DenseRange, + slack: usize, + ) -> bool { + let Some(needed) = range.len().filter(|len| *len <= DENSE_MAX_SLOTS) else { + // Too wide to be worth a slot per value, and it can only get wider self.dense_disabled = true; - return; + return false; }; - // A small table is always worth it. A larger one only once the values - // fill enough of the range they span; until then keep hashing, since - // the groups seen so far may just be a fraction of a dense column. - let groups = self.values.len() - usize::from(self.null_group.is_some()); - if len > DENSE_SMALL_SLOTS && groups * DENSE_MIN_FILL_DENOM < len { - return; + // A small table is always worth it, a larger one only once the values + // fill enough of the range they span. Judged on the range the values + // need, not on any slack added below. + if needed > DENSE_SMALL_SLOTS + && self.value_groups() * DENSE_MIN_FILL_DENOM < needed + { + return false; } + // Slack keeps repeated growth amortized. It widens the range itself so + // that the table always covers exactly the range, which is what lets + // any later rebuild place every group. + let len = needed.max(slack).min(DENSE_MAX_SLOTS); + range.max = range.max.add_wrapping(T::Native::usize_as(len - needed)); + self.observed_range = Some(range); + let min = range.min.index(); let mut group_ids = vec![DENSE_EMPTY; len]; - for (group_idx, &value) in self.values.iter().enumerate() { - if Some(group_idx) == self.null_group { - continue; - } - let offset = value.biased().wrapping_sub(range.min_biased); + for (group_idx, value) in self.value_groups_iter() { + let offset = value.index().wrapping_sub(min); debug_assert!(offset < len as u64); group_ids[offset as usize] = group_idx as u32; } - self.store = GroupStore::Dense { - min: range.min, - min_biased: range.min_biased, - group_ids, - }; + self.store = GroupStore::Dense { min, group_ids }; + true } - /// Track the range of dense keys seen, which decides whether a direct - /// mapped table is worth building - fn observe_range(&mut self, array: &PrimitiveArray) { - if self.dense_disabled { - return; - } - let Some(range) = Self::dense_range(array) else { - // All null, or a type that cannot be direct mapped - if array.null_count() != array.len() { - self.dense_disabled = true; - } - return; - }; - self.observed_range = Some(match self.observed_range { - Some(seen) => seen.merge(range), - None => range, - }); - } - - /// Grow the direct mapped table so that it also covers `array`, returning - /// false if that would need more than [`DENSE_MAX_SLOTS`] slots - fn grow_dense(&mut self, array: &PrimitiveArray) -> bool { - let groups = self.values.len() - usize::from(self.null_group.is_some()); - let GroupStore::Dense { - min, - min_biased, - group_ids, - } = &mut self.store - else { - return false; - }; - let Some(batch) = Self::dense_range(array) else { + /// Rebuild the direct mapped table so that it also covers `array`, which + /// holds the rows that fell outside it + fn try_widen_dense(&mut self, array: &PrimitiveArray) -> bool { + let GroupStore::Dense { group_ids, .. } = &self.store else { return false; }; + let slack = group_ids.len().saturating_mul(2); - // The range the table covers today, widened to also cover the batch - let Some(covered_max) = (*min) - .checked_add(group_ids.len() as i128) - .map(|end| end - 1) + // The range covered so far, widened to cover the rows left over + let (Some(low), Some(high)) = + (arrow::compute::min(array), arrow::compute::max(array)) else { return false; }; - let grown = DenseRange { - min: *min, - max: covered_max, - min_biased: *min_biased, - } - .merge(batch); - let Some(len) = grown.len().filter(|len| *len <= DENSE_MAX_SLOTS) else { - return false; - }; - let new_min = grown.min; - // Growing must not leave the table mostly empty - if len > DENSE_SMALL_SLOTS && groups * DENSE_MIN_FILL_DENOM < len { - return false; - } + let mut range = self.observed_range.expect("dense table without a range"); + range.extend(low); + range.extend(high); - if new_min < *min { - // Values below the current range: shift the existing slots up - let prefix = (*min - new_min) as usize; - let mut shifted = vec![DENSE_EMPTY; len]; - shifted[prefix..prefix + group_ids.len()].copy_from_slice(group_ids); - *group_ids = shifted; - *min = new_min; - *min_biased = grown.min_biased; - } else { - group_ids.resize(len, DENSE_EMPTY); - } - true + self.try_build_dense(range, slack) } - /// Intern `array` with the direct mapped table - /// - /// Interning is idempotent, so a batch that reports [`DenseOutcome::OutOfRange`] - /// can simply be interned again once the table covers it + /// Intern `array`, returning the row holding the first value that fell + /// outside the table, if any. Earlier rows are interned, so the caller + /// resumes from that row once the table covers the rest fn intern_dense( &mut self, array: &PrimitiveArray, groups: &mut Vec, - ) -> DenseOutcome { + ) -> Option { let Self { store, values, null_group, .. } = self; - let GroupStore::Dense { - min_biased, - group_ids, - .. - } = store - else { - return DenseOutcome::OutOfRange; + let GroupStore::Dense { min, group_ids } = store else { + unreachable!("group values are not direct mapped") }; - let min_biased = *min_biased; + let min = *min; let len = group_ids.len() as u64; - for v in array { + for (row, v) in array.iter().enumerate() { let group_id = match v { None => *null_group.get_or_insert_with(|| { let group_id = values.len(); @@ -478,9 +350,9 @@ where group_id }), Some(key) => { - let offset = key.biased().wrapping_sub(min_biased); + let offset = key.index().wrapping_sub(min); if offset >= len { - return DenseOutcome::OutOfRange; + return Some(row); } let slot = &mut group_ids[offset as usize]; @@ -497,7 +369,7 @@ where groups.push(group_id); } - DenseOutcome::Interned + None } /// Rebuild the group values as a hash table, which is possible at any @@ -506,22 +378,23 @@ where fn convert_dense_to_hash(&mut self) { let state = &self.random_state; let mut map = HashTable::with_capacity(self.values.len()); - for (group_idx, &value) in self.values.iter().enumerate() { - if Some(group_idx) == self.null_group { - continue; - } + for (group_idx, value) in self.value_groups_iter() { let hash = value.hash(state); map.insert_unique(hash, (group_idx, hash), |&(_, hash)| hash); } self.store = GroupStore::Hash(map); } + /// Intern `array` by hashing, tracking the range of the groups it creates + /// so [`Self::try_build_dense`] can judge the fill rate without a scan fn intern_hash(&mut self, array: &PrimitiveArray, groups: &mut Vec) { let Self { store, values, null_group, random_state, + observed_range, + dense_disabled, .. } = self; let GroupStore::Hash(map) = store else { @@ -555,6 +428,16 @@ where let g = values.len(); v.insert((g, hash)); values.push(key); + // Only new groups can widen the range + if !*dense_disabled { + match observed_range { + Some(range) => range.extend(key), + None => { + *observed_range = + Some(DenseRange { min: key, max: key }) + } + } + } g } } @@ -574,42 +457,37 @@ where let array = cols[0].as_primitive::(); groups.clear(); - // Only while hashing: once the values are direct mapped the range is - // maintained by the table itself, and once they are known not to be - // dense there is nothing left to decide - if !self.dense_disabled && matches!(self.store, GroupStore::Hash(_)) { - self.observe_range(array); - self.try_migrate_to_dense(); - } - if matches!(self.store, GroupStore::Dense { .. }) { - if let DenseOutcome::Interned = self.intern_dense(array, groups) { + let Some(row) = self.intern_dense(array, groups) else { return Ok(()); - } + }; - // Values outside the range the table covers: grow it if that is - // still worthwhile, otherwise fall back to hashing for good - groups.clear(); - if self.grow_dense(array) { - if let DenseOutcome::Interned = self.intern_dense(array, groups) { - return Ok(()); - } - groups.clear(); + // Rebuild over the wider range if still worthwhile, else hash + let rest = array.slice(row, array.len() - row); + if self.try_widen_dense(&rest) { + let interned = self.intern_dense(&rest, groups); + debug_assert!(interned.is_none(), "widened table left a value out"); + return Ok(()); } self.convert_dense_to_hash(); self.dense_disabled = true; + self.intern_hash(&rest, groups); + return Ok(()); } self.intern_hash(array, groups); + if !self.dense_disabled + && let Some(range) = self.observed_range + { + self.try_build_dense(range, 0); + } Ok(()) } fn size(&self) -> usize { let store = match &self.store { GroupStore::Hash(map) => map.capacity() * size_of::<(usize, u64)>(), - GroupStore::Dense { group_ids, .. } => { - group_ids.capacity() * size_of::() - } + GroupStore::Dense { group_ids, .. } => group_ids.allocated_size(), }; store + self.values.allocated_size() } @@ -642,6 +520,8 @@ where EmitTo::All => { self.store = GroupStore::Hash(HashTable::with_capacity(0)); self.observed_range = None; + // The next values may be dense even if these were not + self.dense_disabled = !T::Native::DENSE; build_primitive(std::mem::take(&mut self.values), self.null_group.take()) } EmitTo::First(n) => { @@ -661,17 +541,21 @@ where } }); } - GroupStore::Dense { group_ids, .. } => { - for slot in group_ids.iter_mut() { - if *slot == DENSE_EMPTY { - continue; - } - match (*slot as usize).checked_sub(n) { + GroupStore::Dense { min, group_ids } => { + // Driven by the groups rather than the slots, of which + // there can be many more + for (group_idx, value) in + self.values.iter().enumerate().filter(|(group_idx, _)| { + Some(*group_idx) != self.null_group + }) + { + let offset = value.index().wrapping_sub(*min) as usize; + group_ids[offset] = match group_idx.checked_sub(n) { // Group index was >= n, shift value down - Some(sub) => *slot = sub as u32, + Some(sub) => sub as u32, // Group index was < n, so free the slot - None => *slot = DENSE_EMPTY, - } + None => DENSE_EMPTY, + }; } } } @@ -695,6 +579,7 @@ where self.values.shrink_to(num_rows); self.null_group = None; self.observed_range = None; + self.dense_disabled = !T::Native::DENSE; match &mut self.store { GroupStore::Hash(map) => { map.clear(); @@ -710,13 +595,51 @@ where #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, Float64Array, Int64Array}; - use arrow::datatypes::{Float64Type, Int64Type}; + use arrow::array::{Array, Float64Array, Int32Array, Int64Array}; + use arrow::datatypes::{Float64Type, Int32Type, Int64Type}; fn is_dense(gv: &GroupValuesPrimitive) -> bool { matches!(gv.store, GroupStore::Dense { .. }) } + /// Mirror of the `EmitTo::take_needed` regression test, applied to the + /// concrete `GroupValuesPrimitive` accumulator. + /// + /// When `n` is small, the old `split_off(n) + swap` pattern used inside + /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation + /// and returned the emitted prefix carrying the original large backing. + /// + /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: + /// the emitted prefix gets a compact allocation and `self.values` retains the + /// original large one. + #[test] + fn emit_first_small_n_allocates_minimally() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + + // Intern 20 distinct values; `new()` pre-allocates capacity 128 for `values`. + let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..20i32)); + let mut groups = vec![]; + gv.intern(&[arr], &mut groups)?; + let capacity_before = gv.values.capacity(); // 128 + + // n=4, n*2=8 <= len=20 -> drain branch + let emitted = gv.emit(EmitTo::First(4))?; + + assert_eq!(emitted[0].len(), 4); + + // `self.values` must retain its original large allocation. + // Old split_off+swap left it with a fresh small allocation (~16). + assert_eq!( + gv.values.capacity(), + capacity_before, + "self.values capacity {} should equal original {} after small First(n) emit", + gv.values.capacity(), + capacity_before, + ); + + Ok(()) + } + fn array(values: &[Option]) -> ArrayRef { Arc::new(Int64Array::from(values.to_vec())) as ArrayRef } @@ -901,6 +824,63 @@ mod tests { Ok(()) } + /// A stream is reused across spill cycles, so values found not to be dense + /// must not keep the next cycle hashing + #[test] + fn draining_the_groups_reconsiders_dense() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int64); + let mut groups = vec![]; + + // Too wide a range to be worth a slot per value + gv.intern(&[array(&[Some(0), Some(i64::MAX)])], &mut groups)?; + assert!(!is_dense(&gv)); + + gv.emit(EmitTo::All)?; + gv.intern(&[array(&[Some(10), Some(11), Some(12)])], &mut groups)?; + assert!(is_dense(&gv)); + + // `clear_shrink` drains the groups just as well + gv.clear_shrink(0); + gv.intern(&[array(&[Some(0), Some(i64::MAX)])], &mut groups)?; + assert!(!is_dense(&gv)); + gv.clear_shrink(0); + gv.intern(&[array(&[Some(20), Some(21)])], &mut groups)?; + assert!(is_dense(&gv)); + + Ok(()) + } + + /// A table is grown with slack, so groups can be created above every value + /// used to size it. Rebuilding it must still place those groups rather + /// than index past the end of the new table. + #[test] + fn groups_above_the_observed_range_survive_a_rebuild() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int64); + let mut groups = vec![]; + + // Covers 1000..=1099 + let values = (1000..1100).map(Some).collect::>(); + gv.intern(&[array(&values)], &mut groups)?; + assert!(is_dense(&gv)); + + // Grows to 1000..=1199, of which only 1000..=1100 was used to size it + gv.intern(&[array(&[Some(1100)])], &mut groups)?; + // A group in the slack, above every value the table was sized from + gv.intern(&[array(&[Some(1199)])], &mut groups)?; + + // Forces a rebuild from far below: the range must still reach 1199 + gv.intern(&[array(&[Some(-5000)])], &mut groups)?; + + let emitted = gv.emit(EmitTo::All)?; + let emitted = emitted[0].as_primitive::(); + assert_eq!(emitted.len(), 103); + assert_eq!(emitted.value(100), 1100); + assert_eq!(emitted.value(101), 1199); + assert_eq!(emitted.value(102), -5000); + + Ok(()) + } + /// Types that cannot be direct mapped at all are ruled out while their /// range is tracked, which must not be mistaken for a table worth building #[test] From 7545a2a33ab0a439503a9a669de2fa30c8a059a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 13 Aug 2026 12:32:34 +0200 Subject: [PATCH 4/4] fix: keep the direct mapped table out of the way of spilling A direct mapped table is sized by the range of the keys, which spilling does not shrink, so rebuilding one after a spill hands the operator back the memory it just released. `clear_shrink` only runs when spilling, so the groups stay hashed from then on. Also updates two expectations that the smaller tables move: the reported peak memory of a small aggregate, and the pool for the four partition case of the spill test, which shared one partition's budget between four streams. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XYz9cxwBdfZLvcKmWcoiF9 --- .../group_values/single_group_by/primitive.rs | 49 +++++++++++++------ .../test_files/aggregate_memory_spill.slt | 4 ++ .../test_files/explain_analyze.slt | 2 +- 3 files changed, 39 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index aea4278463ab1..99b723406784c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -98,9 +98,7 @@ hash_float!(f16, f32, f64); /// A trait to allow direct mapped ("dense") lookup of integer values, by /// indexing a vector of group indices with `value - min` instead of hashing. -/// The join side maps keys the same way, see [`ArrayMap`] -/// -/// [`ArrayMap`]: crate::joins::array_map::ArrayMap +/// The join side maps keys the same way, see `joins::array_map::ArrayMap` pub trait DenseKey: Copy { /// Whether this type can be direct mapped at all const DENSE: bool; @@ -152,8 +150,10 @@ const DENSE_EMPTY: u32 = u32::MAX; const DENSE_MAX_SLOTS: usize = 2 * 1024 * 1024; /// Minimum fill (1/8) of its range for a table above [`DENSE_SMALL_SLOTS`], so -/// that sparse values are not given a mostly empty slot each. The join side -/// admits keys the same way, see `perfect_hash_join_min_key_density` +/// that sparse values are not given a mostly empty slot each. At 4 bytes a slot +/// against 16 a bucket this is also where it stops using less memory than the +/// hash table it replaces. The join side admits keys the same way, see +/// `perfect_hash_join_min_key_density` const DENSE_MIN_FILL_DENOM: usize = 8; /// Tables this small (256KiB) are built whatever the fill rate: even empty they @@ -518,7 +518,14 @@ where let array: PrimitiveArray = match emit_to { EmitTo::All => { - self.store = GroupStore::Hash(HashTable::with_capacity(0)); + match &mut self.store { + // Cleared rather than dropped, so that the capacity built + // up for these groups is reused by the next ones + GroupStore::Hash(map) => map.clear(), + GroupStore::Dense { .. } => { + self.store = GroupStore::Hash(HashTable::with_capacity(0)) + } + } self.observed_range = None; // The next values may be dense even if these were not self.dense_disabled = !T::Native::DENSE; @@ -579,7 +586,10 @@ where self.values.shrink_to(num_rows); self.null_group = None; self.observed_range = None; - self.dense_disabled = !T::Native::DENSE; + // Only called when spilling. A direct mapped table is sized by the + // range of the keys, which spilling does not shrink, so rebuilding one + // would keep the operator at the memory it just tried to give back. + self.dense_disabled = true; match &mut self.store { GroupStore::Hash(map) => { map.clear(); @@ -824,10 +834,10 @@ mod tests { Ok(()) } - /// A stream is reused across spill cycles, so values found not to be dense - /// must not keep the next cycle hashing + /// Emitting every group starts the decision again, since the next values + /// may be dense even if these were not #[test] - fn draining_the_groups_reconsiders_dense() -> Result<()> { + fn emitting_all_groups_reconsiders_dense() -> Result<()> { let mut gv = GroupValuesPrimitive::::new(DataType::Int64); let mut groups = vec![]; @@ -839,13 +849,22 @@ mod tests { gv.intern(&[array(&[Some(10), Some(11), Some(12)])], &mut groups)?; assert!(is_dense(&gv)); - // `clear_shrink` drains the groups just as well + Ok(()) + } + + /// Spilling does not shrink the range of the keys, so a direct mapped + /// table would come straight back at the size the operator just gave up + #[test] + fn spilling_keeps_the_groups_hashed() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int64); + let mut groups = vec![]; + + gv.intern(&[array(&[Some(10), Some(11), Some(12)])], &mut groups)?; + assert!(is_dense(&gv)); + gv.clear_shrink(0); - gv.intern(&[array(&[Some(0), Some(i64::MAX)])], &mut groups)?; + gv.intern(&[array(&[Some(20), Some(21), Some(22)])], &mut groups)?; assert!(!is_dense(&gv)); - gv.clear_shrink(0); - gv.intern(&[array(&[Some(20), Some(21)])], &mut groups)?; - assert!(is_dense(&gv)); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 3dbf880fd1fa9..c533f3f23ca1d 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -198,6 +198,10 @@ FROM ( statement ok SET datafusion.execution.target_partitions = 4 +# Four partitions share the pool, so give them the room one partition had +statement ok +SET datafusion.runtime.memory_limit = '4M' + query II SELECT count(*), sum(total) FROM ( diff --git a/datafusion/sqllogictest/test_files/explain_analyze.slt b/datafusion/sqllogictest/test_files/explain_analyze.slt index d64efe80ccae5..57f3dae91940a 100644 --- a/datafusion/sqllogictest/test_files/explain_analyze.slt +++ b/datafusion/sqllogictest/test_files/explain_analyze.slt @@ -500,7 +500,7 @@ GROUP BY k; ---- Plan with Metrics 01)ProjectionExec: expr=[k@0 as k, count(Int64(1))@1 as count(*)], metrics=[output_bytes=1056.0 B] -02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=9.2 KB] +02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=2.2 KB] 03)----ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]