diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index c682d7caf1b68..3b4c87e4f195c 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -255,6 +255,17 @@ pub trait GroupsAccumulator: Send + std::any::Any { not_impl_err!("Preserving grouped evaluation is not implemented") } + /// Number of groups per storage block, if the state is kept in blocks. + /// + /// When `Some(n)`, [`Self::evaluate`] and [`Self::state`] with + /// `EmitTo::First(n)` cost O(n) rather than O(remaining groups): the + /// first block is handed over whole and the rest keeps its layout. Hash + /// aggregation uses this to emit and free its output one block at a + /// time instead of materializing every group at once. + fn block_len(&self) -> Option { + None + } + /// Returns `true` if [`Self::evaluate_preserving`] is implemented. fn supports_evaluate_preserving(&self) -> bool { false diff --git a/datafusion/functions-aggregate-common/Cargo.toml b/datafusion/functions-aggregate-common/Cargo.toml index 0898e6b4c5522..86d0b641e54ee 100644 --- a/datafusion/functions-aggregate-common/Cargo.toml +++ b/datafusion/functions-aggregate-common/Cargo.toml @@ -57,3 +57,7 @@ rand = { workspace = true } [[bench]] harness = false name = "accumulate" + +[[bench]] +harness = false +name = "blocked_update" diff --git a/datafusion/functions-aggregate-common/benches/blocked_update.rs b/datafusion/functions-aggregate-common/benches/blocked_update.rs new file mode 100644 index 0000000000000..83be8a16fbc31 --- /dev/null +++ b/datafusion/functions-aggregate-common/benches/blocked_update.rs @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! How the per-group state is addressed in the accumulator update loop. +//! +//! `flat` is `values[g]` on one `Vec`. `naive_blocked` is +//! `blocks[g >> SHIFT][g & MASK]`, the dependent lookup that made earlier +//! blocked designs slower at high cardinality: the block pointer must load +//! before the state access can issue, so misses on the state serialize. +//! `gather_blocked` resolves a chunk of addresses in an independent pass +//! first, then updates through them. `accumulator` is the real +//! [`PrimitiveGroupsAccumulator`] (single block up to `1 << 20` groups, +//! address gather beyond). +//! +//! Group indices are uniformly random so nothing but the addressing differs +//! between the variants; the interesting sizes are the ones whose state no +//! longer fits in cache. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array}; +use arrow::datatypes::{DataType, Int64Type}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_expr_common::groups_accumulator::GroupsAccumulator; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::RESOLVE_CHUNK; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocks::BlockedVec; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; + +const BATCH_SIZE: usize = 8192; +const BATCHES: usize = 128; +/// Small blocks so every size is blocked; the block count does not matter. +const BENCH_SHIFT: u32 = 16; + +/// Group indices as `intern` hands them out: numbered in order of first +/// appearance, so every batch's new groups are dense and `NullState` takes +/// its fast path. `4 * num_groups` rows; the stream is replayed from an +/// empty accumulator each iteration. +fn dense_stream(num_groups: usize) -> (Vec>, Vec) { + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let mut ids = vec![usize::MAX; num_groups]; + let mut next = 0usize; + let mut batches = vec![]; + let mut totals = vec![]; + for _ in 0..(4 * num_groups).div_ceil(BATCH_SIZE) { + let batch: Vec = (0..BATCH_SIZE) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + let key = (state % num_groups as u64) as usize; + if ids[key] == usize::MAX { + ids[key] = next; + next += 1; + } + ids[key] + }) + .collect(); + batches.push(batch); + totals.push(next); + } + (batches, totals) +} + +fn group_batches(num_groups: usize) -> Vec> { + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % num_groups as u64) as usize + }; + (0..BATCHES) + .map(|_| (0..BATCH_SIZE).map(|_| next()).collect()) + .collect() +} + +fn flat(values: &mut [i64], groups: &[usize], input: &[i64]) { + for (&g, &v) in groups.iter().zip(input) { + values[g] += v; + } +} + +fn naive_blocked( + values: &mut BlockedVec, + groups: &[usize], + input: &[i64], +) { + for (&g, &v) in groups.iter().zip(input) { + *values.get_mut(g) += v; + } +} + +fn gather_blocked( + values: &mut BlockedVec, + groups: &[usize], + input: &[i64], +) { + let resolver = values.address_resolver(); + let mut addrs = [std::ptr::null_mut::(); RESOLVE_CHUNK]; + for (groups, input) in groups + .chunks(RESOLVE_CHUNK) + .zip(input.chunks(RESOLVE_CHUNK)) + { + let addrs = &mut addrs[..groups.len()]; + resolver.resolve(groups, addrs); + for (&addr, &v) in addrs.iter().zip(input) { + // SAFETY: resolved from live blocks the resolver borrows + unsafe { *addr += v }; + } + } +} + +fn criterion_benchmark(c: &mut Criterion) { + let input: Vec = (0..BATCH_SIZE as i64).collect(); + let input_array: ArrayRef = Arc::new(Int64Array::from(input.clone())); + let mut group = c.benchmark_group("blocked_update"); + group.sample_size(10); + group.throughput(Throughput::Elements((BATCHES * BATCH_SIZE) as u64)); + + for (label, num_groups) in [("100k", 100_000), ("1M", 1_000_000), ("10M", 10_000_000)] + { + let batches = group_batches(num_groups); + + let mut values = vec![0i64; num_groups]; + group.bench_with_input( + BenchmarkId::new("flat", label), + &batches, + |b, batches| { + b.iter(|| { + for groups in batches { + flat(&mut values, black_box(groups), black_box(&input)); + } + }) + }, + ); + + let mut values = BlockedVec::::new(); + values.resize(num_groups, 0); + group.bench_with_input( + BenchmarkId::new("naive_blocked", label), + &batches, + |b, batches| { + b.iter(|| { + for groups in batches { + naive_blocked(&mut values, black_box(groups), black_box(&input)); + } + }) + }, + ); + + let mut values = BlockedVec::::new(); + values.resize(num_groups, 0); + group.bench_with_input( + BenchmarkId::new("gather_blocked", label), + &batches, + |b, batches| { + b.iter(|| { + for groups in batches { + gather_blocked(&mut values, black_box(groups), black_box(&input)); + } + }) + }, + ); + + let mut accumulator = PrimitiveGroupsAccumulator::::new( + &DataType::Int64, + |current, value| *current += value, + ); + group.bench_with_input( + BenchmarkId::new("accumulator", label), + &batches, + |b, batches| { + b.iter(|| { + for groups in batches { + accumulator + .update_batch( + std::slice::from_ref(&input_array), + black_box(groups), + None, + num_groups, + ) + .unwrap(); + } + }) + }, + ); + } + group.finish(); + + // The partial stage's loop: dense new groups, no nulls, no filter. + let mut group = c.benchmark_group("blocked_update_dense"); + group.sample_size(10); + for (label, num_groups) in + [("1M", 1_000_000), ("10M", 10_000_000), ("20M", 20_000_000)] + { + let (batches, totals) = dense_stream(num_groups); + group.throughput(Throughput::Elements((batches.len() * BATCH_SIZE) as u64)); + group.bench_with_input( + BenchmarkId::new("accumulator", label), + &(batches, totals), + |b, (batches, totals)| { + b.iter(|| { + let mut accumulator = PrimitiveGroupsAccumulator::::new( + &DataType::Int64, + |current, value| *current += value, + ); + for (groups, &total) in batches.iter().zip(totals) { + accumulator + .update_batch( + std::slice::from_ref(&input_array), + black_box(groups), + None, + total, + ) + .unwrap(); + } + black_box(accumulator.size()) + }) + }, + ); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index ad01fbce5b506..87d0bb8458ce6 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -19,6 +19,7 @@ //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] pub mod accumulate; +pub mod blocks; pub mod bool_op; pub mod nulls; pub mod prim_op; diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs index 36ba0fba0bd34..76772f9d1f07e 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -212,6 +212,46 @@ impl NullState { }); } + /// [`Self::accumulate`] for state that is not addressed by group index + /// directly: `resolve` turns a chunk of group indices into the addresses + /// of their state and `value_fn(addr, group_index, value)` updates + /// through them. See [`accumulate_resolved`]. + pub fn accumulate_resolved( + &mut self, + group_indices: &[usize], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + resolve: R, + mut value_fn: F, + ) where + T: ArrowPrimitiveType + Send, + R: FnMut(&[usize], &mut [*mut S]), + F: FnMut(*mut S, usize, T::Native), + { + if opt_filter.is_none() + && values.null_count() == 0 + && let SeenValues::All { num_values } = &mut self.seen_values + && new_groups_are_dense(group_indices, *num_values, total_num_groups) + { + accumulate_resolved(group_indices, values, None, resolve, value_fn); + *num_values = total_num_groups; + return; + } + + let seen_values = self.seen_values.get_builder(total_num_groups); + accumulate_resolved( + group_indices, + values, + opt_filter, + resolve, + |addr, group_index, value| { + seen_values.set_bit(group_index, true); + value_fn(addr, group_index, value); + }, + ); + } + /// Invokes `value_fn(group_index, value)` for each non null, non /// filtered value in `values`, while tracking which groups have /// seen null inputs and which groups have seen any inputs, for @@ -580,6 +620,141 @@ pub fn accumulate_multiple( } } +/// Rows per address-resolution chunk of [`accumulate_resolved`]: small +/// enough that a chunk's group indices, input values and resolved addresses +/// all stay in L1 while the update loop runs over them, and that the cache +/// lines the resolve pass prefetches are still there when the loop reaches +/// them. +pub const RESOLVE_CHUNK: usize = 256; + +/// [`accumulate`] for state that is not addressed by group index directly. +/// +/// Rows are processed in chunks of [`RESOLVE_CHUNK`]. For each chunk, +/// `resolve(group_indices, addrs)` first fills `addrs[i]` with the address +/// of the state of `group_indices[i]`; the update loop then walks the +/// resolved addresses, calling `value_fn(addr, group_index, value)` for every +/// non null, non filtered row, exactly as [`accumulate`] calls its +/// `value_fn(group_index, value)`. +/// +/// The two passes keep the update loop's memory accesses independent of one +/// another. State kept in blocks (see `BlockedVec`) would otherwise be +/// reached through `blocks[g >> SHIFT][g & MASK]`, where the block pointer +/// has to load before the state access can be issued, and that dependency +/// serializes cache misses a flat `values[g]` loop overlaps; at high +/// cardinality the loop is bound by exactly that memory-level parallelism. +/// +/// `value_fn` receives raw pointers: `resolve` must return pointers that stay +/// valid for the whole call, and the caller dereferences them. +pub fn accumulate_resolved( + group_indices: &[usize], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + mut resolve: R, + mut value_fn: F, +) where + T: ArrowPrimitiveType + Send, + R: FnMut(&[usize], &mut [*mut S]), + F: FnMut(*mut S, usize, T::Native), +{ + let data: &[T::Native] = values.values(); + assert_eq!(data.len(), group_indices.len()); + if let Some(filter) = opt_filter { + assert_eq!(filter.len(), group_indices.len()); + } + let mut addrs = [std::ptr::null_mut::(); RESOLVE_CHUNK]; + + let chunks = group_indices + .chunks(RESOLVE_CHUNK) + .zip(data.chunks(RESOLVE_CHUNK)); + for (chunk, (groups, data)) in chunks.enumerate() { + let first_row = chunk * RESOLVE_CHUNK; + let addrs = &mut addrs[..groups.len()]; + resolve(groups, addrs); + let addrs = &*addrs; + let rows = addrs.iter().zip(groups).zip(data); + + match (values.null_count() > 0, opt_filter) { + // no nulls, no filter + (false, None) => { + rows.for_each(|((&addr, &group_index), &new_value)| { + value_fn(addr, group_index, new_value) + }); + } + // nulls, no filter: 64 rows per validity word, as in `accumulate` + (true, None) => { + let nulls = values + .nulls() + .unwrap() + .inner() + .slice(first_row, groups.len()); + let bit_chunks = nulls.bit_chunks(); + let addr_chunks = addrs.chunks_exact(64); + let group_chunks = groups.chunks_exact(64); + let data_chunks = data.chunks_exact(64); + let addr_remainder = addr_chunks.remainder(); + let group_remainder = group_chunks.remainder(); + let data_remainder = data_chunks.remainder(); + + addr_chunks + .zip(group_chunks) + .zip(data_chunks) + .zip(bit_chunks.iter()) + .for_each(|(((addrs, groups), data), mask)| { + let mut index_mask = 1; + addrs.iter().zip(groups).zip(data).for_each( + |((&addr, &group_index), &new_value)| { + if (mask & index_mask) != 0 { + value_fn(addr, group_index, new_value); + } + index_mask <<= 1; + }, + ) + }); + + let remainder_bits = bit_chunks.remainder_bits(); + addr_remainder + .iter() + .zip(group_remainder) + .zip(data_remainder) + .enumerate() + .for_each(|(i, ((&addr, &group_index), &new_value))| { + if remainder_bits & (1 << i) != 0 { + value_fn(addr, group_index, new_value); + } + }); + } + // no nulls, but a filter + (false, Some(filter)) => { + let filter = filter.slice(first_row, groups.len()); + rows.zip(filter.iter()).for_each( + |(((&addr, &group_index), &new_value), filter_value)| { + if filter_value == Some(true) { + value_fn(addr, group_index, new_value); + } + }, + ); + } + // both null values and filters + (true, Some(filter)) => { + let filter = filter.slice(first_row, groups.len()); + let values = values.slice(first_row, groups.len()); + addrs + .iter() + .zip(groups) + .zip(values.iter()) + .zip(filter.iter()) + .for_each(|(((&addr, &group_index), new_value), filter_value)| { + if filter_value == Some(true) + && let Some(new_value) = new_value + { + value_fn(addr, group_index, new_value) + } + }); + } + } + } +} + /// This function is called to update the accumulator state per row /// when the value is not needed (e.g. COUNT) /// @@ -763,6 +938,62 @@ mod test { .run() } + /// `accumulate_resolved` visits exactly the rows `accumulate` visits, in + /// the same order, for every null/filter combination and chunk boundary. + #[test] + fn accumulate_resolved_matches_accumulate() { + let len = 3 * RESOLVE_CHUNK + 77; + let group_indices: Vec = (0..len).map(|i| (i * 7919) % 1000).collect(); + let values_no_nulls = UInt32Array::from_iter_values((0..len).map(|i| i as u32)); + let values_with_nulls = UInt32Array::from_iter( + (0..len).map(|i| if i % 3 == 0 { None } else { Some(i as u32) }), + ); + let filter = BooleanArray::from_iter((0..len).map(|i| match i % 5 { + 0 => None, + 1 => Some(false), + _ => Some(true), + })); + let mut state = vec![0u64; 1000]; + + for values in [&values_no_nulls, &values_with_nulls] { + for opt_filter in [None, Some(&filter)] { + let mut expected = vec![]; + super::accumulate(&group_indices, values, opt_filter, |g, v| { + expected.push((g, v)) + }); + let mut actual = vec![]; + accumulate_resolved( + &group_indices, + values, + opt_filter, + |groups, addrs| { + for (addr, &g) in addrs.iter_mut().zip(groups) { + *addr = &raw mut state[g]; + } + }, + |addr, g, v| { + // SAFETY: `addr` points into `state`, which outlives the call + unsafe { *addr += 1 }; + actual.push((g, v)); + }, + ); + assert_eq!(actual, expected); + } + } + assert_eq!( + state.iter().sum::() as usize, + 4 * len - { + // rows skipped: nulls in two runs, filter in two runs, both in one + let nulls = values_with_nulls.null_count(); + let filtered = filter.iter().filter(|f| *f != Some(true)).count(); + let both = (0..len) + .filter(|&i| i % 3 == 0 || !filter.value(i) || filter.is_null(i)) + .count(); + nulls + filtered + both + } + ); + } + #[test] fn accumulate_fuzz() { let mut rng = rand::rng(); diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocks.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocks.rs new file mode 100644 index 0000000000000..ac5641c131186 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocks.rs @@ -0,0 +1,425 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`BlockedVec`]: per-group state stored in fixed-size blocks. + +use std::marker::PhantomData; +use std::mem::size_of; + +use datafusion_common::utils::split_vec_min_alloc; + +/// Number of groups per block, as a power of two: `1 << DEFAULT_BLOCK_SHIFT`. +/// +/// Every table with at most this many groups is a single block, which is +/// read and updated exactly like a flat `Vec`. Larger tables pay one extra +/// address-resolution pass per batch (see [`AddressResolver`]). +pub const DEFAULT_BLOCK_SHIFT: u32 = 20; + +/// Per-group values stored in blocks of `1 << SHIFT` elements: group `g` +/// lives in block `g >> SHIFT` at offset `g & ((1 << SHIFT) - 1)`. +/// +/// Blocked storage lets a hash aggregation grow, emit and free its state one +/// block at a time instead of as one contiguous allocation. The cost that +/// made earlier blocked designs slower is *how* the state is addressed in the +/// update loop, not the blocks themselves: `blocks[g >> SHIFT][g & MASK]` is +/// a dependent load, so the CPU cannot issue the next group's cache miss +/// until the block pointer has resolved, and the flat `Vec` wins on +/// memory-level parallelism alone. [`AddressResolver`] resolves the +/// addresses of a chunk of rows in a separate, independent pass first, after +/// which the update loop issues its misses exactly as the flat loop does. +#[derive(Debug)] +pub struct BlockedVec { + blocks: Vec>, + len: usize, +} + +impl Default for BlockedVec { + fn default() -> Self { + Self::new() + } +} + +impl BlockedVec { + /// Elements per block. + pub const BLOCK_LEN: usize = 1 << SHIFT; + const MASK: usize = Self::BLOCK_LEN - 1; + + pub const fn new() -> Self { + Self { + blocks: Vec::new(), + len: 0, + } + } + + /// Reserves room for `capacity` elements in the first block. + pub fn with_capacity(capacity: usize) -> Self { + Self { + blocks: vec![Vec::with_capacity(capacity.min(Self::BLOCK_LEN))], + len: 0, + } + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Number of blocks currently allocated. + pub fn num_blocks(&self) -> usize { + self.blocks.len() + } + + /// Bytes allocated by the blocks, not including `self`. + pub fn allocated_size(&self) -> usize { + self.blocks + .iter() + .map(|b| b.capacity() * size_of::()) + .sum() + } + + #[inline] + fn block_of(index: usize) -> (usize, usize) { + (index >> SHIFT, index & Self::MASK) + } + + /// The whole storage as one slice when it fits in a single block, so the + /// common case is read and updated exactly like a flat `Vec`. + #[inline] + pub fn as_single_block_mut(&mut self) -> Option<&mut [T]> { + match self.blocks.as_mut_slice() { + [] => Some(&mut []), + [block] => Some(block.as_mut_slice()), + _ => None, + } + } + + /// Reads the element at `index`. + #[inline] + pub fn get(&self, index: usize) -> &T { + let (block, offset) = Self::block_of(index); + &self.blocks[block][offset] + } + + /// Mutable access through the dependent `block -> offset` lookup. For + /// batch updates prefer [`Self::address_resolver`]. + #[inline] + pub fn get_mut(&mut self, index: usize) -> &mut T { + let (block, offset) = Self::block_of(index); + &mut self.blocks[block][offset] + } + + pub fn push(&mut self, value: T) { + if self.len & Self::MASK == 0 && self.len >> SHIFT == self.blocks.len() { + self.blocks.push(Vec::new()); + } + self.blocks[self.len >> SHIFT].push(value); + self.len += 1; + } + + /// Removes every element, keeping the first block's allocation. + pub fn clear(&mut self) { + self.blocks.truncate(1); + if let Some(block) = self.blocks.first_mut() { + block.clear(); + } + self.len = 0; + } + + /// Shrinks the retained allocation to `capacity` elements at most. + pub fn shrink_to(&mut self, capacity: usize) { + if let Some(block) = self.blocks.first_mut() { + block.shrink_to(capacity.min(Self::BLOCK_LEN)); + } + } + + /// Takes the elements out in order. A single block is moved without a + /// copy; several blocks are concatenated. + pub fn take_all(&mut self) -> Vec { + self.len = 0; + match self.blocks.len() { + 0 => Vec::new(), + 1 => std::mem::take(&mut self.blocks[0]), + _ => { + let mut all = + Vec::with_capacity(self.blocks.iter().map(|b| b.len()).sum()); + for block in self.blocks.drain(..) { + all.extend(block); + } + all + } + } + } + + /// Takes the first `n` elements, shifting the rest down so the element + /// at `n` becomes element `0`. + pub fn take_first(&mut self, n: usize) -> Vec { + if n == self.len { + return self.take_all(); + } + if self.blocks.len() > 1 && n == Self::BLOCK_LEN { + // A whole first block is moved out; the rest is already packed. + self.len -= n; + return self.blocks.remove(0); + } + if self.blocks.len() <= 1 { + let taken = match self.blocks.first_mut() { + Some(block) => split_vec_min_alloc(block, n), + None => Vec::new(), + }; + self.len -= taken.len(); + return taken; + } + // Whole blocks come first, so they are moved rather than copied. + let whole = n >> SHIFT; + let mut taken: Vec = Vec::with_capacity(n); + for block in self.blocks.drain(..whole) { + taken.extend(block); + } + let rest = n & Self::MASK; + if rest > 0 { + taken.extend(self.blocks[0].drain(..rest)); + } + self.len -= n; + // Blocks after the first are now offset by `rest`; re-pack. + if rest > 0 && self.blocks.len() > 1 { + let mut remaining = std::mem::take(&mut self.blocks); + let mut carry: Vec = remaining.drain(1..).flatten().collect(); + let first = &mut remaining[0]; + let fill = Self::BLOCK_LEN - first.len(); + first.extend(carry.drain(..fill.min(carry.len()))); + self.blocks = remaining; + let mut carry = carry.into_iter(); + loop { + let block: Vec = carry.by_ref().take(Self::BLOCK_LEN).collect(); + if block.is_empty() { + break; + } + self.blocks.push(block); + } + } + taken + } + + /// Resolves group indices to element addresses for a batch update; see + /// [`AddressResolver`]. + pub fn address_resolver(&mut self) -> AddressResolver<'_, T, SHIFT> { + AddressResolver { + bases: self.blocks.iter_mut().map(|b| b.as_mut_ptr()).collect(), + len: self.len, + _lifetime: PhantomData, + } + } +} + +impl BlockedVec { + /// Grows to `new_len` elements, filling with `value`; blocks are added + /// as needed and the storage never shrinks. + /// + /// Every block grows by doubling like a `Vec` until it is full, so only + /// the last block ever carries growth slack and the allocated size stays + /// proportional to the number of groups. + pub fn resize(&mut self, new_len: usize, value: T) { + while self.len < new_len { + let block = self.len >> SHIFT; + if block == self.blocks.len() { + self.blocks.push(Vec::new()); + } + let block = &mut self.blocks[block]; + let target = (new_len - (self.len - block.len())).min(Self::BLOCK_LEN); + block.resize(target, value); + self.len = (self.len & !Self::MASK) + block.len(); + } + } +} + +/// Turns group indices into the addresses of their elements, for +/// `accumulate_resolved` in the `accumulate` module. +/// +/// The block base table is a few pointers and stays in L1, so resolving a +/// chunk of rows is one short independent pass; the update loop that follows +/// walks the resolved addresses, and its cache misses are as independent of +/// one another as a flat `values[g]` loop's are. That is what the dependent +/// `blocks[g >> SHIFT][g & MASK]` lookup of [`BlockedVec::get_mut`] loses. +pub struct AddressResolver<'a, T, const SHIFT: u32> { + bases: Vec<*mut T>, + len: usize, + _lifetime: PhantomData<&'a mut T>, +} + +impl AddressResolver<'_, T, SHIFT> { + const MASK: usize = (1 << SHIFT) - 1; + + /// Fills `addrs[i]` with the address of element `group_indices[i]` and + /// starts fetching each element's cache line, so that by the time the + /// update loop reaches a row its state is on the way in. + /// + /// Every group index must be below the vector's length, as for any + /// group index passed to an accumulator. The pointers stay valid for the + /// resolver's lifetime, during which it borrows the blocks mutably. + #[inline] + pub fn resolve(&self, group_indices: &[usize], addrs: &mut [*mut T]) { + assert!(addrs.len() >= group_indices.len()); + for (addr, &group) in addrs.iter_mut().zip(group_indices) { + debug_assert!(group < self.len, "group index {group} out of bounds"); + // SAFETY: `group < len`, so its block exists and, as only the + // last block can be partially filled, `group & MASK` is within it. + let element = unsafe { + self.bases + .get_unchecked(group >> SHIFT) + .add(group & Self::MASK) + }; + prefetch_for_write(element); + *addr = element; + } + } +} + +/// Hints the CPU to fetch the cache line holding `ptr` for a coming write. +/// +/// Issued from the resolve pass, where the address is known without any +/// dependent load, this overlaps the state's memory latency with the rest of +/// the pass; it is a hint only and never faults. A no-op on targets without +/// a stable prefetch instruction. +#[inline(always)] +fn prefetch_for_write(ptr: *mut T) { + #[cfg(target_arch = "x86_64")] + // SAFETY: `_mm_prefetch` is a hint and does not access memory. + unsafe { + std::arch::x86_64::_mm_prefetch::<{ std::arch::x86_64::_MM_HINT_T0 }>( + ptr as *const i8, + ) + } + #[cfg(target_arch = "aarch64")] + // SAFETY: `prfm` is a hint that reads no memory and has no side effects + // on registers or flags. + unsafe { + std::arch::asm!( + "prfm pstl1keep, [{0}]", + in(reg) ptr, + options(nostack, preserves_flags, readonly) + ); + } + #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] + let _ = ptr; +} + +#[cfg(test)] +mod tests { + use super::*; + + type Small = BlockedVec; // four elements per block + + fn filled(n: usize) -> Small { + let mut v = Small::new(); + for i in 0..n as u32 { + v.push(i); + } + v + } + + fn contents(v: &Small) -> Vec { + (0..v.len()).map(|i| *v.get(i)).collect() + } + + #[test] + fn push_resize_and_index_across_blocks() { + let mut v = filled(6); + assert_eq!(v.num_blocks(), 2); + assert_eq!(contents(&v), (0..6).collect::>()); + v.resize(11, 9); + assert_eq!(v.len(), 11); + assert_eq!(v.num_blocks(), 3); + assert_eq!(contents(&v), vec![0, 1, 2, 3, 4, 5, 9, 9, 9, 9, 9]); + *v.get_mut(10) = 7; + assert_eq!(*v.get(10), 7); + assert!(v.as_single_block_mut().is_none()); + assert!(filled(4).as_single_block_mut().is_some()); + } + + #[test] + fn blocks_grow_by_doubling_not_by_block_len() { + let mut v = BlockedVec::::new(); + v.resize(1000, 0); + assert_eq!(v.num_blocks(), 1); + assert!(v.allocated_size() < 2 * 1000 * size_of::()); + v.resize(BlockedVec::::BLOCK_LEN + 1, 0); + assert_eq!(v.num_blocks(), 2); + assert!( + v.allocated_size() < (BlockedVec::::BLOCK_LEN + 8) * size_of::() + ); + } + + #[test] + fn take_first_block_moves_it_without_copying() { + let mut v = Small::new(); + for i in 0..10u32 { + v.push(i); + } + let before = v.allocated_size(); + let first = v.take_first(Small::BLOCK_LEN); + assert_eq!(first, vec![0, 1, 2, 3]); + assert_eq!(v.len(), 6); + assert_eq!(contents(&v), vec![4, 5, 6, 7, 8, 9]); + assert_eq!( + v.allocated_size(), + before - first.capacity() * size_of::() + ); + // The last, partial block is handed over whole as well. + assert_eq!(v.take_first(4), vec![4, 5, 6, 7]); + assert_eq!(v.take_first(2), vec![8, 9]); + assert!(v.is_empty()); + } + + #[test] + fn take_all_and_take_first_keep_order() { + assert_eq!(filled(10).take_all(), (0..10).collect::>()); + for n in 0..=10 { + let mut v = filled(10); + assert_eq!(v.take_first(n), (0..n as u32).collect::>()); + assert_eq!(v.len(), 10 - n); + assert_eq!(contents(&v), (n as u32..10).collect::>()); + // The layout is packed again: pushes land after the last element. + v.push(99); + assert_eq!(*v.get(10 - n), 99); + } + } + + #[test] + fn resolver_addresses_every_block() { + let groups: Vec = (0..5000).map(|i| (i * 7) % 11).collect(); + let mut v = Small::new(); + v.resize(11, 0); + let resolver = v.address_resolver(); + let mut addrs = vec![std::ptr::null_mut(); 64]; + for groups in groups.chunks(64) { + resolver.resolve(groups, &mut addrs); + for &addr in &addrs[..groups.len()] { + // SAFETY: resolved from live blocks the resolver borrows + unsafe { *addr += 1 }; + } + } + drop(resolver); + let mut expected = vec![0u32; 11]; + for &g in &groups { + expected[g] += 1; + } + assert_eq!(contents(&v), expected); + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index 0453124528a0f..e62af48e1a4ee 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::mem::size_of; use std::sync::Arc; use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray}; @@ -29,6 +28,7 @@ use datafusion_expr_common::groups_accumulator::{ }; use super::accumulate::NullState; +use super::blocks::BlockedVec; /// An accumulator that implements a single operation over /// [`ArrowPrimitiveType`] where the accumulated state is the same as @@ -46,7 +46,7 @@ where F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, { /// values per group, stored as the native type - values: Vec, + values: BlockedVec, /// The output type (needed for Decimal precision and scale) data_type: DataType, @@ -68,7 +68,7 @@ where { pub fn new(data_type: &DataType, prim_fn: F) -> Self { Self { - values: vec![], + values: BlockedVec::new(), data_type: data_type.clone(), null_state: NullState::new(), starting_value: T::default_value(), @@ -102,23 +102,45 @@ where self.values.resize(total_num_groups, self.starting_value); // NullState dispatches / handles tracking nulls and groups that saw no values - self.null_state.accumulate( - group_indices, - values, - opt_filter, - total_num_groups, - |group_index, new_value| { - // SAFETY: group_index is guaranteed to be in bounds - let value = unsafe { self.values.get_unchecked_mut(group_index) }; - (self.prim_fn)(value, new_value); - }, - ); + if let Some(values_block) = self.values.as_single_block_mut() { + self.null_state.accumulate( + group_indices, + values, + opt_filter, + total_num_groups, + |group_index, new_value| { + // SAFETY: group_index is guaranteed to be in bounds + let value = unsafe { values_block.get_unchecked_mut(group_index) }; + (self.prim_fn)(value, new_value); + }, + ); + } else { + // Several blocks: resolve each chunk's addresses ahead of its + // update loop so the loop's cache misses stay independent. + let resolver = self.values.address_resolver(); + self.null_state.accumulate_resolved( + group_indices, + values, + opt_filter, + total_num_groups, + |groups, addrs| resolver.resolve(groups, addrs), + |value, _, new_value| { + // SAFETY: `resolve` returned the address of a live + // element, and the resolver borrows the blocks for the + // whole call. + (self.prim_fn)(unsafe { &mut *value }, new_value); + }, + ); + } Ok(()) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - let values = emit_to.take_needed(&mut self.values); + let values = match emit_to { + EmitTo::All => self.values.take_all(), + EmitTo::First(n) => self.values.take_first(n), + }; let nulls = self.null_state.build(emit_to); let values = PrimitiveArray::::new(values.into(), nulls) // no copy .with_data_type(self.data_type.clone()); @@ -127,8 +149,10 @@ where fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { selection.validate_num_groups(self.values.len())?; - let values: Vec = - selection.iter().map(|index| self.values[index]).collect(); + let values: Vec = selection + .iter() + .map(|index| *self.values.get(index)) + .collect(); let nulls = self.null_state.build_preserving(selection)?; let values = PrimitiveArray::::new(values.into(), nulls) .with_data_type(self.data_type.clone()); @@ -139,6 +163,10 @@ where true } + fn block_len(&self) -> Option { + Some(BlockedVec::::BLOCK_LEN) + } + fn state(&mut self, emit_to: EmitTo) -> Result> { self.evaluate(emit_to).map(|arr| vec![arr]) } @@ -217,16 +245,64 @@ where Ok(vec![Arc::new(state_values)]) } fn size(&self) -> usize { - self.values.capacity() * size_of::() + self.null_state.size() + self.values.allocated_size() + self.null_state.size() } } #[cfg(test)] mod tests { use super::*; - use arrow::array::Int64Array; + use arrow::array::{Array, Int64Array}; use arrow::datatypes::Int64Type; + /// Groups that span several blocks are updated through resolved + /// addresses and emitted in group order, with and without nulls and + /// filters, and for both emit modes. + #[test] + fn multi_block_groups_update_and_emit_in_order() -> Result<()> { + const BLOCK: usize = BlockedVec::::BLOCK_LEN; + let total = BLOCK + 5; + let mut accumulator = PrimitiveGroupsAccumulator::::new( + &DataType::Int64, + |current, value| *current += value, + ); + // Rows alternate between the two blocks; row 3 is null and row 5 is + // filtered out, so groups 1 and 3 never see a value. + let groups = vec![0, BLOCK, BLOCK + 4, 1, 0, 3, BLOCK + 4]; + let values: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + Some(2), + Some(3), + None, + Some(4), + Some(5), + Some(6), + ])); + let filter = BooleanArray::from(vec![true, true, true, true, true, false, true]); + accumulator.update_batch(&[values], &groups, Some(&filter), total)?; + + let all = accumulator.evaluate_preserving(GroupSelection::all(total))?; + let all = all.as_primitive::(); + assert_eq!(all.len(), total); + assert_eq!(all.value(0), 5); + assert!(all.is_null(1)); + assert!(all.is_null(3)); + assert_eq!(all.value(BLOCK), 2); + assert_eq!(all.value(BLOCK + 4), 9); + assert_eq!(all.null_count(), total - 3); + + // First(n) across the block boundary keeps the rest addressable. + let first = accumulator.evaluate(EmitTo::First(BLOCK + 1))?; + let first = first.as_primitive::(); + assert_eq!(first.len(), BLOCK + 1); + assert_eq!(first.value(BLOCK), 2); + let rest = accumulator.evaluate(EmitTo::All)?; + let rest = rest.as_primitive::(); + assert_eq!(rest.len(), 4); + assert_eq!(rest.value(3), 9); + Ok(()) + } + #[test] fn preserving_reads_do_not_change_accumulator_state() -> Result<()> { let mut accumulator = PrimitiveGroupsAccumulator::::new( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 5f101950870d2..949db3dfba518 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -265,33 +265,54 @@ impl AggregateHashTable { let batch_size = self.batch_size; let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - let mut output = + let group_by_metrics = self.group_by_metrics.clone(); + + // `state` is `Some` while groups are handed over one block at a time. + let (mut output, state) = match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { AggregateHashTableState::Outputting(mut state) => { if state.group_values.is_empty() { return Ok(None); } - - // Accumulator output consumes internal state. Materialize all - // groups once, then slice the materialized batch on later polls. - let emit_to = EmitTo::All; - let columns = self.group_by_metrics.time_emitting(|| { - let mut columns = state.group_values.emit(emit_to)?; - for (idx, acc) in state.accumulators.iter_mut().enumerate() { - columns.extend(accumulator_metrics.time( - idx, - accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), - )?); - } - Ok::<_, datafusion_common::DataFusionError>(columns) - })?; - - let batch = RecordBatch::try_new(output_schema, columns)?; - debug_assert!(batch.num_rows() > 0); - MaterializedAggregateOutput::new(batch) + if let Some(block_len) = emit_block_len(&state) { + // Every column keeps its state in blocks of the same + // length: emit one block, so its state is freed while + // the rest of the table is still waiting to be drained. + let block = emit_block( + &mut state, + block_len, + materialize_accumulator_fn, + accumulator_phase, + &output_schema, + &accumulator_metrics, + &group_by_metrics, + )?; + (MaterializedAggregateOutput::new(block), Some(state)) + } else { + // Accumulator output consumes internal state. Materialize all + // groups once, then slice the materialized batch on later polls. + let emit_to = EmitTo::All; + let columns = group_by_metrics.time_emitting(|| { + let mut columns = state.group_values.emit(emit_to)?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + columns.extend(accumulator_metrics.time( + idx, + accumulator_phase, + || materialize_accumulator_fn(acc, emit_to), + )?); + } + Ok::<_, datafusion_common::DataFusionError>(columns) + })?; + + let batch = RecordBatch::try_new(output_schema, columns)?; + debug_assert!(batch.num_rows() > 0); + (MaterializedAggregateOutput::new(batch), None) + } + } + AggregateHashTableState::OutputtingMaterialized(output) => (output, None), + AggregateHashTableState::OutputtingBlock { state, block } => { + (block, Some(state)) } - AggregateHashTableState::OutputtingMaterialized(output) => output, AggregateHashTableState::Done => return Ok(None), AggregateHashTableState::Building(_) => { return internal_err!( @@ -301,11 +322,18 @@ impl AggregateHashTable { }; let batch = output.next_batch(batch_size); - if output.is_exhausted() { - self.state = AggregateHashTableState::Done; - } else { - self.state = AggregateHashTableState::OutputtingMaterialized(output); - } + self.state = match (state, output.is_exhausted()) { + // The next poll emits the next block. + (Some(state), true) if !state.group_values.is_empty() => { + AggregateHashTableState::Outputting(state) + } + (Some(state), false) => AggregateHashTableState::OutputtingBlock { + state, + block: output, + }, + (None, false) => AggregateHashTableState::OutputtingMaterialized(output), + _ => AggregateHashTableState::Done, + }; Ok(batch) } @@ -325,6 +353,15 @@ impl AggregateHashTable { AggregateHashTableState::OutputtingMaterialized(output) => { output.memory_size() } + AggregateHashTableState::OutputtingBlock { state, block } => { + state + .accumulators + .iter() + .map(|acc| acc.accumulator.size()) + .sum::() + + state.group_values.size() + + block.memory_size() + } AggregateHashTableState::Done => 0, } } @@ -602,9 +639,54 @@ pub(super) enum AggregateHashTableState { /// Note this is a temporary solution until the `GroupValues` issue is solved: /// Issue: OutputtingMaterialized(MaterializedAggregateOutput), + /// One block of groups has been emitted from `state` and is being sliced + /// out of `block`; the groups still in `state` follow, a block at a time. + /// Only used when every column keeps its state in blocks of one length + /// (see [`GroupsAccumulator::block_len`]). + OutputtingBlock { + state: AggregateHashTableBuffer, + block: MaterializedAggregateOutput, + }, Done, } +/// The block length shared by the group values and every accumulator, if +/// they all keep their state in blocks of the same length. +fn emit_block_len(state: &AggregateHashTableBuffer) -> Option { + let block_len = state.group_values.block_len()?; + state + .accumulators + .iter() + .all(|acc| acc.accumulator.block_len() == Some(block_len)) + .then_some(block_len) +} + +/// Emits the first block of groups (all remaining ones if fewer than +/// `block_len`) as one batch, freeing their state. +fn emit_block( + state: &mut AggregateHashTableBuffer, + block_len: usize, + materialize_accumulator_fn: MaterializeAccumulatorFn, + accumulator_phase: AccumulatorPhase, + output_schema: &SchemaRef, + accumulator_metrics: &AggregateAccumulatorMetrics, + group_by_metrics: &GroupByMetrics, +) -> Result { + let emit_to = EmitTo::First(block_len.min(state.group_values.len())); + let columns = group_by_metrics.time_emitting(|| { + let mut columns = state.group_values.emit_block()?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + columns.extend(accumulator_metrics.time(idx, accumulator_phase, || { + materialize_accumulator_fn(acc, emit_to) + })?); + } + Ok::<_, datafusion_common::DataFusionError>(columns) + })?; + let batch = RecordBatch::try_new(Arc::clone(output_schema), columns)?; + debug_assert!(batch.num_rows() > 0); + Ok(batch) +} + /// Fully evaluated aggregate output and the next row offset to emit. /// /// Final aggregate evaluation consumes accumulator state, and partial terminal diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index b764fd7792b39..9b7673c8cac41 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -116,6 +116,21 @@ pub trait GroupValues: Send { /// Emits the group values fn emit(&mut self, emit_to: EmitTo) -> Result>; + /// Number of groups per storage block, if the values are kept in blocks; + /// see [`Self::emit_block`]. + fn block_len(&self) -> Option { + None + } + + /// Emits the first [`Self::block_len`] groups (all remaining ones if + /// fewer), renumbering the rest down as `EmitTo::First` does, and drops + /// the lookup index: no [`Self::intern`] may follow. Unlike + /// `emit(EmitTo::First(n))`, which keeps the index usable by renumbering + /// every entry, this costs O(block). + fn emit_block(&mut self) -> Result> { + not_impl_err!("Block emission is not implemented") + } + /// Materializes selected group values without changing the stored values or /// their group indices. /// diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 21b62457e3831..6a6a8a4cfd6ad 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 @@ -24,9 +24,8 @@ use arrow::array::{ use arrow::datatypes::{DataType, i256}; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; -use datafusion_common::utils::split_vec_min_alloc; -use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupSelection}; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocks::BlockedVec; use half::f16; use hashbrown::hash_table::HashTable; #[cfg(not(feature = "force_hash_collisions"))] @@ -103,17 +102,18 @@ hash_float!(f16, f32, f64); pub struct GroupValuesPrimitive { /// The data type of the output array data_type: DataType, - /// Stores the `(group_index, hash)` based on the hash of its value + /// Stores `(group_index, value)` under the hash of the 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)>, + /// The value is kept in the entry so a probe compares against the + /// bucket it already loaded and never reads `values`, which is one + /// cache miss per row less than looking the group's value up by index + /// (and what lets `values` live in blocks at no cost to the probe). + /// Rehashing recomputes hashes from the stored values. + map: HashTable<(usize, T::Native)>, /// The group index of the null value if any null_group: Option, /// The values for each group index - values: Vec, + values: BlockedVec, /// The random state used to generate hashes random_state: RandomState, } @@ -124,13 +124,28 @@ impl GroupValuesPrimitive { Self { data_type, map: HashTable::with_capacity(128), - values: Vec::with_capacity(128), + values: BlockedVec::with_capacity(128), null_group: None, random_state: crate::aggregates::AGGREGATION_HASH_SEED, } } } +fn build_primitive( + values: Vec, + null_idx: Option, +) -> PrimitiveArray { + let nulls = null_idx.map(|null_idx| { + let mut buffer = NullBufferBuilder::new(values.len()); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(values.len() - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }); + PrimitiveArray::::new(values.into(), nulls) +} + impl GroupValues for GroupValuesPrimitive where T::Native: HashValue, @@ -155,17 +170,15 @@ where let hash = key.hash(state); let insert = self.map.entry( hash, - |&(g, h)| unsafe { - hash == h && self.values.get_unchecked(g).is_eq(key) - }, - |&(_, h)| h, + |&(_, k)| k.is_eq(key), + |&(_, k)| k.hash(state), ); match insert { hashbrown::hash_table::Entry::Occupied(o) => o.get().0, hashbrown::hash_table::Entry::Vacant(v) => { let g = self.values.len(); - v.insert((g, hash)); + v.insert((g, key)); self.values.push(key); g } @@ -178,7 +191,8 @@ where } fn size(&self) -> usize { - self.map.capacity() * size_of::<(usize, u64)>() + self.values.allocated_size() + self.map.capacity() * size_of::<(usize, T::Native)>() + + self.values.allocated_size() } fn is_empty(&self) -> bool { @@ -190,25 +204,10 @@ where } fn emit(&mut self, emit_to: EmitTo) -> Result> { - fn build_primitive( - values: Vec, - null_idx: Option, - ) -> PrimitiveArray { - let nulls = null_idx.map(|null_idx| { - let mut buffer = NullBufferBuilder::new(values.len()); - buffer.append_n_non_nulls(null_idx); - buffer.append_null(); - buffer.append_n_non_nulls(values.len() - null_idx - 1); - // NOTE: The inner builder must be constructed as there is at least one null - buffer.finish().unwrap() - }); - PrimitiveArray::::new(values.into(), nulls) - } - let array: PrimitiveArray = match emit_to { EmitTo::All => { self.map.clear(); - build_primitive(std::mem::take(&mut self.values), self.null_group.take()) + build_primitive(self.values.take_all(), self.null_group.take()) } EmitTo::First(n) => { self.map.retain(|entry| { @@ -232,20 +231,43 @@ where Some(_) => self.null_group.take(), None => None, }; - build_primitive(split_vec_min_alloc(&mut self.values, n), null_group) + build_primitive(self.values.take_first(n), null_group) } }; Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) } + fn block_len(&self) -> Option { + Some(BlockedVec::::BLOCK_LEN) + } + + fn emit_block(&mut self) -> Result> { + // Output only from here on: the index is dropped, not renumbered. + self.map.clear(); + self.map.shrink_to(0, |_| 0); + let n = self.values.len().min(BlockedVec::::BLOCK_LEN); + let null_group = match &mut self.null_group { + Some(v) if *v >= n => { + *v -= n; + None + } + Some(_) => self.null_group.take(), + None => None, + }; + let array = build_primitive::(self.values.take_first(n), null_group); + Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) + } + fn values_preserving( &mut self, selection: GroupSelection<'_>, ) -> Result> { selection.validate_num_groups(self.values.len())?; - let values: Vec = - selection.iter().map(|index| self.values[index]).collect(); + let values: Vec = selection + .iter() + .map(|index| *self.values.get(index)) + .collect(); let nulls = if let Some(null_group) = self.null_group { let mut nulls = NullBufferBuilder::new(values.len()); for index in selection.iter() { @@ -280,7 +302,7 @@ where mod tests { use super::*; use arrow::array::types::Int32Type; - use arrow::array::{ArrayRef, Int32Array}; + use arrow::array::{Array, ArrayRef, Int32Array}; use arrow::datatypes::DataType; use datafusion_expr::EmitTo; use std::sync::Arc; @@ -303,7 +325,7 @@ mod tests { 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 + let capacity_before = gv.values.allocated_size(); // 128 values // n=4, n*2=8 <= len=20 -> drain branch let emitted = gv.emit(EmitTo::First(4))?; @@ -313,13 +335,92 @@ mod tests { // `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(), + gv.values.allocated_size(), capacity_before, "self.values capacity {} should equal original {} after small First(n) emit", - gv.values.capacity(), + gv.values.allocated_size(), capacity_before, ); Ok(()) } + + /// `emit_block` hands over one block at a time in group order, keeps the + /// null group's position right across blocks, and frees the index. + #[test] + fn emit_block_walks_the_blocks_in_order() -> Result<()> { + const BLOCK: usize = BlockedVec::::BLOCK_LEN; + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + let n = BLOCK + 3; + let mut groups = vec![]; + let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..n as i32)); + gv.intern(&[arr], &mut groups)?; + gv.intern( + &[Arc::new(Int32Array::from(vec![None::]))], + &mut groups, + )?; + assert_eq!(groups, vec![n]); + assert_eq!(gv.block_len(), Some(BLOCK)); + let size_before = gv.size(); + + let first = gv.emit_block()?; + let first = first[0].as_primitive::(); + assert_eq!(first.len(), BLOCK); + assert_eq!(first.null_count(), 0); + assert_eq!(first.value(BLOCK - 1), BLOCK as i32 - 1); + assert_eq!(gv.len(), 4); + assert!(gv.size() < size_before / 2, "index and block freed"); + + let rest = gv.emit_block()?; + let rest = rest[0].as_primitive::(); + assert_eq!(rest.len(), 4); + assert_eq!(rest.value(0), BLOCK as i32); + assert!(rest.is_null(3)); + assert!(gv.is_empty()); + Ok(()) + } + + /// Values interned past one block are found again, emitted in group + /// order, and `First(n)` across the block boundary renumbers the rest. + #[test] + fn intern_and_emit_across_blocks() -> Result<()> { + const BLOCK: usize = BlockedVec::::BLOCK_LEN; + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + let n = BLOCK + 3; + let mut groups = vec![]; + let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..n as i32)); + gv.intern(&[arr], &mut groups)?; + assert_eq!(groups, (0..n).collect::>()); + assert_eq!(gv.values.num_blocks(), 2); + + // Re-interning finds every existing group without adding any. + let again: ArrayRef = + Arc::new(Int32Array::from(vec![Some(BLOCK as i32), None, Some(0)])); + gv.intern(&[again], &mut groups)?; + assert_eq!(groups, vec![BLOCK, n, 0]); + assert_eq!(gv.len(), n + 1); + + let first = gv.emit(EmitTo::First(BLOCK + 1))?; + assert_eq!(first[0].len(), BLOCK + 1); + assert_eq!( + first[0].as_primitive::().value(BLOCK), + BLOCK as i32 + ); + + // The null group and the value `BLOCK + 2` were renumbered down. + gv.intern( + &[Arc::new(Int32Array::from(vec![ + None, + Some(BLOCK as i32 + 2), + ]))], + &mut groups, + )?; + assert_eq!(groups, vec![2, 1]); + let rest = gv.emit(EmitTo::All)?; + let rest = rest[0].as_primitive::(); + assert_eq!(rest.len(), 3); + assert_eq!(rest.value(1), BLOCK as i32 + 2); + assert!(rest.is_null(2)); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 152082d88dd90..5afd7e3d07826 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -1543,6 +1543,106 @@ mod tests { Ok(()) } + /// When the group values and every accumulator keep their state in + /// blocks, the output is handed over one block at a time: a block's + /// state is released as soon as it is emitted instead of the whole + /// table staying reserved until the last slice goes out. + #[tokio::test] + async fn test_single_hash_stream_releases_state_block_by_block() -> Result<()> { + use crate::aggregates::single_stream::SingleHashAggregateStream; + use arrow::datatypes::Int64Type; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocks::BlockedVec; + + const BLOCK: usize = BlockedVec::::BLOCK_LEN; + let batch_size = 8192; + // Two blocks, the second one partial; one row per group. + let num_groups = BLOCK + 4 * batch_size; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int64, false), + Field::new("value_col", DataType::Int64, false), + ])); + let batches = (0..num_groups) + .step_by(batch_size) + .map(|start| { + let end = (start + batch_size).min(num_groups); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from_iter_values(start as i64..end as i64)), + Arc::new(Int64Array::from(vec![1i64; end - start])), + ], + ) + }) + .collect::, _>>()?; + + let runtime = RuntimeEnvBuilder::default().build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(Arc::clone(&runtime)); + let session_config = task_ctx.session_config().clone().set( + "datafusion.execution.batch_size", + &datafusion_common::ScalarValue::UInt64(Some(batch_size as u64)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum_value") + .build()?, + )]; + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + input, + Arc::clone(&schema), + )?; + let mut stream = SingleHashAggregateStream::new(&aggregate_exec, &task_ctx, 0)?; + + let mut rows = 0; + let mut reserved_during_first_block = None; + let mut reserved_during_second_block = None; + while let Some(batch) = stream.next().await { + let batch = batch?; + let keys = batch.column(0).as_primitive::(); + let sums = batch.column(1).as_primitive::(); + // Group order is kept across the block boundary. + assert_eq!(keys.value(0), rows as i64); + assert!(sums.iter().all(|sum| sum == Some(1))); + rows += batch.num_rows(); + + let reserved = runtime.memory_pool.reserved(); + if reserved_during_first_block.is_none() { + reserved_during_first_block = Some(reserved); + } else if rows > BLOCK && reserved_during_second_block.is_none() { + reserved_during_second_block = Some(reserved); + } + } + assert_eq!(rows, num_groups); + + // While the first block is sliced, its batch and the second block's + // state are reserved; once the second block is emitted, only its + // (much smaller) batch is. + let first = reserved_during_first_block.unwrap(); + let second = reserved_during_second_block.unwrap(); + assert!( + first >= BLOCK * 2 * size_of::(), + "the first block's batch ({first} bytes) is not reserved" + ); + assert!( + second * 4 <= first, + "the first block's state was not released: {second} bytes still reserved \ + after it was drained, {first} during" + ); + assert_eq!(runtime.memory_pool.reserved(), 0); + + Ok(()) + } + #[tokio::test] async fn test_partial_hash_stream_releases_held_batch_after_last_slice() -> Result<()> {