diff --git a/datafusion/expr-common/src/blocked_groups_accumulator.rs b/datafusion/expr-common/src/blocked_groups_accumulator.rs new file mode 100644 index 0000000000000..6c58a609c295c --- /dev/null +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -0,0 +1,326 @@ +// 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. + +//! Vectorized [`BlockedGroupsAccumulator`]: [`GroupsAccumulator`] with state +//! stored in blocks. +//! +//! Storing per-group state in fixed-size blocks lets hash aggregation emit and +//! free its state one block at a time instead of materializing every group +//! into one large array. See . +//! +//! [`GroupsAccumulator`]: crate::groups_accumulator::GroupsAccumulator + +use std::any::Any; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray}; +use datafusion_common::Result; + +use crate::accumulator::AggregateMetrics; +use crate::groups_accumulator::EmitTo; + +/// Identifies one group in blocked storage: the block that holds the group +/// and the group's position inside that block. +/// +/// This is the only place that knows how a group index is laid out; storage +/// and accumulators only use the accessors below. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Default)] +pub struct BlocksIndex(u64); + +impl BlocksIndex { + /// Creates the index of the `index_in_block`-th group of block + /// `block_index`. + #[inline] + pub const fn new(block_index: usize, index_in_block: usize) -> Self { + debug_assert!(block_index <= u32::MAX as usize); + debug_assert!(index_in_block <= u32::MAX as usize); + Self(((block_index as u64) << 32) | index_in_block as u64) + } + + /// Block that holds this group. + #[inline] + pub const fn block_index(self) -> usize { + (self.0 >> 32) as usize + } + + /// Position of this group inside its block. + #[inline] + pub const fn index_in_block(self) -> usize { + self.0 as u32 as usize + } + + /// Creates the index of the group at flat index `flat`, counting from the + /// first group, when every block holds `block_size` groups. + #[inline] + pub const fn from_flat(flat: usize, block_size: usize) -> Self { + Self::new(flat / block_size, flat % block_size) + } + + /// Flat index of this group, counting from the first group, when every + /// block holds `block_size` groups. + #[inline] + pub const fn flat(self, block_size: usize) -> usize { + self.block_index() * block_size + self.index_in_block() + } + + /// Returns `true` if every index in `indices` points to an existing group + /// in storage that holds `len` groups in blocks of `block_size`, where + /// every block is full except maybe the last one. + /// + /// For example, with `block_size = 4` and `len = 10` the blocks hold 4, 4 + /// and 2 groups, so `(2, 1)` is valid, but `(2, 2)`, `(3, 0)` and `(0, 4)` + /// are not. + pub fn all_in_bounds(indices: &[Self], block_size: usize, len: usize) -> bool { + let Some(last_group) = len.checked_sub(1) else { + return indices.is_empty(); + }; + let last = Self::from_flat(last_group, block_size); + + // A single pass without a branch per index, so it is vectorized + let (max, or) = indices + .iter() + .fold((0, 0), |(max, or), index| (max.max(index.0), or | index.0)); + + // Indices order by block, then position, so an index is valid exactly + // when it is not after the last group and its position is below + // `block_size` + let positions_in_block = if block_size.is_power_of_two() { + (or as u32 as usize) < block_size + } else { + indices + .iter() + .all(|index| index.index_in_block() < block_size) + }; + max <= last.0 && positions_in_block + } + + /// Writes [`Self::from_flat`] of every index in `flat` to `out`, + /// replacing its contents. + pub fn from_flat_slice(flat: &[usize], block_size: usize, out: &mut Vec) { + out.clear(); + if block_size.is_power_of_two() { + let shift = block_size.trailing_zeros(); + let mask = block_size - 1; + out.extend(flat.iter().map(|&f| Self::new(f >> shift, f & mask))); + } else { + out.extend(flat.iter().map(|&f| Self::from_flat(f, block_size))); + } + } + + /// Writes [`Self::flat`] of every index in `indices` to `out`, replacing + /// its contents. + pub fn to_flat_slice(indices: &[Self], block_size: usize, out: &mut Vec) { + out.clear(); + if block_size.is_power_of_two() { + let shift = block_size.trailing_zeros(); + out.extend( + indices + .iter() + .map(|i| (i.block_index() << shift) | i.index_in_block()), + ); + } else { + out.extend(indices.iter().map(|i| i.flat(block_size))); + } + } +} + +/// Which groups to emit from a [`BlockedGroupsAccumulator`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BlockedEmitTo { + /// Emit every group, as one array per block, and reset the state. + All, + /// Emit the first block (or nothing if there are no groups). + /// + /// The remaining groups move down by one block, like [`EmitTo::First`] + /// with the first block's length. To emit every group, use [`Self::All`]. + NextBlock, + /// Emit the first `n` groups, where `0 < n < block_size`, and shift the + /// remaining groups down by `n`, like [`EmitTo::First`]. + First(usize), +} + +impl BlockedEmitTo { + /// Splits a flat [`EmitTo`] into block-sized emits. + /// + /// `EmitTo::First(n)` becomes `n / block_size` [`Self::NextBlock`]s + /// followed by one [`Self::First`] with the remainder, if any. Each emit + /// shifts the remaining groups, so they must all be applied, in order, + /// before any new group is added. + pub fn from_emit_to(emit_to: EmitTo, block_size: usize) -> Vec { + match emit_to { + EmitTo::All => vec![Self::All], + EmitTo::First(n) => { + let mut emits = vec![Self::NextBlock; n / block_size]; + if !n.is_multiple_of(block_size) { + emits.push(Self::First(n % block_size)); + } + emits + } + } + } +} + +/// Like [`GroupsAccumulator`], but the per-group state is stored in blocks of +/// [`Self::block_size`] groups, and emitting returns one array per block +/// instead of one large array. +/// +/// Hash aggregation uses a `BlockedGroupsAccumulator` for every aggregate +/// that has one, and a [`GroupsAccumulator`] for the others, in the same +/// aggregation. +/// +/// [`GroupsAccumulator`]: crate::groups_accumulator::GroupsAccumulator +pub trait BlockedGroupsAccumulator: Send + Any { + /// Number of groups per block. Fixed for the accumulator's lifetime. + fn block_size(&self) -> usize; + + /// See `GroupsAccumulator::set_metrics`. + fn set_metrics(&mut self, _metrics: Arc) {} + + /// Updates the state of each group with the values of its rows. + /// + /// Same contract as `GroupsAccumulator::update_batch`, with the group of + /// each row given as a [`BlocksIndex`]. + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()>; + + /// Merges intermediate state produced by [`Self::state`]. + /// + /// Same contract as `GroupsAccumulator::merge_batch`. + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()>; + + /// Returns the final value of the groups selected by `emit_to`, and + /// removes their state. + /// + /// * [`BlockedEmitTo::All`]: one array per block + /// * [`BlockedEmitTo::NextBlock`]: zero or one arrays + /// * [`BlockedEmitTo::First`]: exactly one array + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result>; + + /// Returns the intermediate state of the groups selected by `emit_to`, + /// indexed `[block][state column]`, and removes their state. + /// + /// Blocks follow the same rules as [`Self::evaluate`]. + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>>; + + /// Converts an input batch directly to intermediate state. + /// + /// Same contract as `GroupsAccumulator::convert_to_state`. + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result>; + + /// Amount of memory used to store the state, in bytes. + fn size(&self) -> usize; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn all_in_bounds() { + let in_bounds = |indices: &[(usize, usize)], block_size, len| { + let indices: Vec<_> = indices + .iter() + .map(|&(b, i)| BlocksIndex::new(b, i)) + .collect(); + BlocksIndex::all_in_bounds(&indices, block_size, len) + }; + + // Blocks of 4, 4 and 2 groups + assert!(in_bounds(&[(0, 0), (1, 3), (2, 1)], 4, 10)); + // After the last group of the last block + assert!(!in_bounds(&[(0, 0), (2, 2)], 4, 10)); + // After the last block + assert!(!in_bounds(&[(3, 0)], 4, 10)); + // Position past the end of a full block + assert!(!in_bounds(&[(0, 4)], 4, 10)); + + // Block size that is not a power of two: blocks of 3, 3 and 1 groups + assert!(in_bounds(&[(1, 2), (2, 0)], 3, 7)); + assert!(!in_bounds(&[(0, 3)], 3, 7)); + + // No groups + assert!(in_bounds(&[], 4, 0)); + assert!(!in_bounds(&[(0, 0)], 4, 0)); + } + + #[test] + fn from_emit_to_splits_into_blocks() { + assert_eq!( + BlockedEmitTo::from_emit_to(EmitTo::All, 4), + vec![BlockedEmitTo::All] + ); + assert_eq!( + BlockedEmitTo::from_emit_to(EmitTo::First(3), 4), + vec![BlockedEmitTo::First(3)] + ); + assert_eq!( + BlockedEmitTo::from_emit_to(EmitTo::First(4), 4), + vec![BlockedEmitTo::NextBlock] + ); + assert_eq!( + BlockedEmitTo::from_emit_to(EmitTo::First(9), 4), + vec![ + BlockedEmitTo::NextBlock, + BlockedEmitTo::NextBlock, + BlockedEmitTo::First(1) + ] + ); + assert!(BlockedEmitTo::from_emit_to(EmitTo::First(0), 4).is_empty()); + } + + #[test] + fn blocks_index_slice_conversions() { + for block_size in [4, 6] { + let flat = vec![0, 3, 4, 13, 7]; + let mut blocked = vec![BlocksIndex::new(9, 9)]; + BlocksIndex::from_flat_slice(&flat, block_size, &mut blocked); + let expected: Vec<_> = flat + .iter() + .map(|&f| BlocksIndex::from_flat(f, block_size)) + .collect(); + assert_eq!(blocked, expected); + let mut back = vec![42]; + BlocksIndex::to_flat_slice(&blocked, block_size, &mut back); + assert_eq!(back, flat); + } + } + + #[test] + fn blocks_index_round_trips() { + let index = BlocksIndex::new(3, 7); + assert_eq!(index.block_index(), 3); + assert_eq!(index.index_in_block(), 7); + assert_eq!(index.flat(16), 55); + assert_eq!(BlocksIndex::from_flat(55, 16), index); + assert!(BlocksIndex::new(0, 9) < BlocksIndex::new(1, 0)); + } +} diff --git a/datafusion/expr-common/src/lib.rs b/datafusion/expr-common/src/lib.rs index c9a95fd294503..86a2c0f61e715 100644 --- a/datafusion/expr-common/src/lib.rs +++ b/datafusion/expr-common/src/lib.rs @@ -34,6 +34,7 @@ #![cfg_attr(test, allow(clippy::needless_pass_by_value))] pub mod accumulator; +pub mod blocked_groups_accumulator; pub mod casts; pub mod columnar_value; pub mod dyn_eq; diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 8a21dc669346f..95ddb43e6607d 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -110,6 +110,9 @@ pub use datafusion_doc::{ pub use datafusion_expr_common::accumulator::{ Accumulator, AggregateMetric, AggregateMetricRecorder, AggregateMetrics, }; +pub use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, +}; pub use datafusion_expr_common::columnar_value::ColumnarValue; pub use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 7358568f0afe1..38ecf32fea47a 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -44,6 +44,7 @@ use crate::udf_eq::UdfEq; use crate::utils::{AggregateOrderSensitivity, format_state_name, ordering_state_fields}; use crate::{Accumulator, Expr, expr_vec_fmt}; use crate::{Documentation, Signature}; +use datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator; /// Logical representation of a user-defined [aggregate function] (UDAF). /// @@ -270,6 +271,21 @@ impl AggregateUDF { self.inner.create_groups_accumulator(args) } + /// See [`AggregateUDFImpl::blocked_groups_accumulator_supported`] for more details. + pub fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + self.inner.blocked_groups_accumulator_supported(args) + } + + /// See [`AggregateUDFImpl::create_blocked_groups_accumulator`] for more details. + pub fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + self.inner + .create_blocked_groups_accumulator(args, block_size) + } + pub fn create_sliding_accumulator( &self, args: AccumulatorArgs, @@ -669,6 +685,24 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { not_impl_err!("GroupsAccumulator hasn't been implemented for {self:?} yet") } + /// If this aggregate has a [`BlockedGroupsAccumulator`]. When `true`, + /// hash aggregation stores this aggregate's group state in blocks and + /// releases it block by block. Defaults to `false`. + fn blocked_groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + false + } + + /// Creates a [`BlockedGroupsAccumulator`] storing `block_size` groups per + /// block. Only called when [`Self::blocked_groups_accumulator_supported`] + /// returned `true`. + fn create_blocked_groups_accumulator( + &self, + _args: AccumulatorArgs, + _block_size: usize, + ) -> Result> { + not_impl_err!("BlockedGroupsAccumulator hasn't been implemented for {self:?} yet") + } + /// Sliding accumulator is an alternative accumulator that can be used for /// window functions. It has retract method to revert the previous update. /// @@ -1622,6 +1656,19 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl { self.inner.create_groups_accumulator(args) } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + self.inner.blocked_groups_accumulator_supported(args) + } + + fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + self.inner + .create_blocked_groups_accumulator(args, block_size) + } + fn create_sliding_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index dbf2af512b9f7..1d8cb5997a142 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 blocked_vec; 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 bf5f37475f30e..1b3d5b8b13b89 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -25,6 +25,7 @@ use arrow::datatypes::ArrowPrimitiveType; use crate::aggregate::groups_accumulator::nulls::filter_to_validity; use datafusion_common::Result; +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; use datafusion_expr_common::groups_accumulator::{EmitTo, GroupSelection}; /// If the input has nulls, then the accumulator must potentially @@ -592,9 +593,39 @@ pub fn accumulate_indices( group_indices: &[usize], nulls: Option<&NullBuffer>, opt_filter: Option<&BooleanArray>, - mut index_fn: F, + index_fn: F, ) where F: FnMut(usize) + Send, +{ + accumulate_group_indices(group_indices, nulls, opt_filter, index_fn) +} + +/// [`accumulate_indices`] for blocked group indices. +/// +/// `F`: Invoked like `value_fn(group_index)` for all non null values +/// passing the filter. +pub fn accumulate_blocked_indices( + group_indices: &[BlocksIndex], + nulls: Option<&NullBuffer>, + opt_filter: Option<&BooleanArray>, + index_fn: F, +) where + F: FnMut(BlocksIndex), +{ + accumulate_group_indices(group_indices, nulls, opt_filter, index_fn) +} + +/// Shared implementation of [`accumulate_indices`] and +/// [`accumulate_blocked_indices`]. +#[inline(always)] +fn accumulate_group_indices( + group_indices: &[I], + nulls: Option<&NullBuffer>, + opt_filter: Option<&BooleanArray>, + mut index_fn: F, +) where + I: Copy, + F: FnMut(I), { match (nulls, opt_filter) { (None, None) => { diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs new file mode 100644 index 0000000000000..305c3d50610e5 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs @@ -0,0 +1,735 @@ +// 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`]: group state stored in fixed size blocks + +use std::mem::size_of; + +use arrow::array::BooleanArray; +use arrow::buffer::NullBuffer; + +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; + +use super::accumulate::accumulate_blocked_indices; + +/// Rows per chunk whose element addresses are resolved before any of them is +/// updated, see [`BlockedVec::update`]. +const RESOLVE_CHUNK: usize = 256; + +/// A growable vector of per group state stored as a list of blocks, used by +/// [`BlockedGroupsAccumulator`] implementations. +/// +/// It stays a single flat `Vec` (growing by doubling) until it reaches +/// `block_size` elements. That `Vec` then becomes block 0 without any copy. +/// Every following block also grows by doubling, up to `block_size`, so a +/// block that only holds a few groups does not allocate a whole block. +/// Accumulators update their state with [`Self::update`] and +/// [`Self::update_with`], which pick the fastest loop for the current layout +/// once per batch. +/// +/// Emitting the first block ([`Self::take_next_block`]) or all blocks +/// ([`Self::take_all`]) moves the block `Vec`s out without copying. +/// +/// `FIXED_BLOCK_SIZE = false` is reserved for the values of nested children +/// (e.g. list elements), where block boundaries are driven by the parent +/// through a future `start_new_block()`. It is not implemented yet, but the +/// layout (one `Vec` per block, the block length being the `Vec` length) +/// already supports it. +/// +/// # Implementation Notes +/// +/// ## Why `T: Copy`? +/// +/// 1. So the [`BlockedVec::allocated_size`] will be accurate since the size of T is known, and not include heap allocations (like `String` or `Vec`) that are not part of the allocated size of the `BlockedVec` +/// 2. So we can provide mutable access to the items (e.g. `IndexMut`) since if `T` is not `Copy` (like when `T` is a `Vec`) we could change the size of it without the [`BlockedVec::allocated_size`] changing, which would be confusing and lead to bugs +/// +/// [`BlockedGroupsAccumulator`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator +#[derive(Debug)] +pub struct BlockedVec { + /// Every block except the last one holds exactly `block_size` elements + blocks: Vec>, + /// Total number of elements over all blocks + len: usize, + block_size: usize, + /// Running total of the bytes allocated by the blocks' buffers + allocated: usize, +} + +impl BlockedVec { + /// Creates an empty vector. + /// + /// # Panics + /// If `block_size` is 0. + pub fn new(block_size: usize) -> Self { + if !FIXED_BLOCK_SIZE { + // Reserved for nested child values, whose block sizes are driven + // by the parent via a future `start_new_block()`. + unimplemented!("BlockedVec with dynamic block sizes"); + } + assert!(block_size > 0, "block_size must be positive"); + Self { + blocks: Vec::new(), + len: 0, + block_size, + allocated: 0, + } + } + + /// Maximum number of elements in a block. + pub fn block_size(&self) -> usize { + self.block_size + } + + /// Total number of elements. + #[inline] + pub fn len(&self) -> usize { + self.len + } + + /// Returns `true` if there are no elements. + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Number of blocks, including a partially filled last block. + #[inline] + pub fn num_blocks(&self) -> usize { + self.blocks.len() + } + + /// Grows the vector to `new_len` elements, filling with `value`. + /// Never shrinks: does nothing if `new_len <= self.len()`. + pub fn grow_to(&mut self, new_len: usize, value: T) { + while self.len < new_len { + let last_len = self.last_block_for_append(); + let additional = (new_len - self.len).min(self.block_size - last_len); + let last = self.reserve_last(last_len + additional); + last.resize(last_len + additional, value); + self.len += additional; + } + } + + /// Appends `value` and returns its index. + #[inline] + pub fn push(&mut self, value: T) -> BlocksIndex { + self.len += 1; + let num_blocks = self.blocks.len(); + match self.blocks.last_mut() { + // Block capacities never exceed `block_size`, so this never + // reallocates + Some(last) if last.len() < last.capacity() => { + last.push(value); + BlocksIndex::new(num_blocks - 1, last.len() - 1) + } + _ => self.push_slow(value), + } + } + + #[cold] + fn push_slow(&mut self, value: T) -> BlocksIndex { + let last_len = self.last_block_for_append(); + self.reserve_last(last_len + 1).push(value); + BlocksIndex::new(self.blocks.len() - 1, last_len) + } + + /// Makes sure the last block has room for at least one more element, + /// adding a new block if needed, and returns its length. + fn last_block_for_append(&mut self) -> usize { + match self.blocks.last() { + Some(last) if last.len() < self.block_size => last.len(), + // Every block, including block 0 (the flat phase), grows by + // doubling in `reserve_last` + _ => { + self.blocks.push(Vec::new()); + 0 + } + } + } + + /// Makes sure the last block can hold `required <= block_size` elements, + /// growing it by doubling (capped at `block_size`), and returns it. + fn reserve_last(&mut self, required: usize) -> &mut Vec { + let block_size = self.block_size; + let last = self.blocks.last_mut().expect("at least one block"); + let old_capacity = last.capacity(); + if old_capacity < required { + let new_capacity = (old_capacity * 2).max(required).min(block_size); + last.reserve_exact(new_capacity - last.len()); + self.allocated += (last.capacity() - old_capacity) * size_of::(); + } + last + } + + /// Returns the element at `index`. + /// + /// # Panics + /// If `index` is out of bounds. + #[inline] + pub fn get(&self, index: BlocksIndex) -> T { + self.blocks[index.block_index()][index.index_in_block()] + } + + /// Returns a mutable reference to the element at `index`. + /// + /// # Safety + /// `index` must point to an element of this vector. + #[inline] + pub unsafe fn get_unchecked_mut(&mut self, index: BlocksIndex) -> &mut T { + // SAFETY: guaranteed by the caller + unsafe { + self.blocks + .get_unchecked_mut(index.block_index()) + .get_unchecked_mut(index.index_in_block()) + } + } + + /// Returns all elements as one slice if there is at most one block, so + /// hot loops can run on a flat slice exactly like today's accumulators. + #[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, + } + } + + /// Returns the address of every block, for resolving many indices + /// without going through the outer `Vec` each time. + /// + /// The returned [`BlockPtrs`] is invalidated by any later use of `self`. + fn block_ptrs_mut(&mut self) -> BlockPtrs { + BlockPtrs { + ptrs: self.blocks.iter_mut().map(|b| b.as_mut_ptr()).collect(), + } + } + + /// Grows the vector to `total_num_groups` elements (new ones are + /// `starting_value`), then calls `update_fn` on the element of every + /// row's group, skipping rows that are null in `nulls` or not selected by + /// `opt_filter`. + /// + /// Chooses the loop once per call: with a single block it is a plain flat + /// loop; with several blocks and no nulls or filter, the addresses of each + /// chunk of rows are resolved before any of them is updated, so the cache + /// misses of the updates don't wait on each other. + /// + /// `group_indices` must only point to the first `total_num_groups` + /// groups, which is the contract of + /// [`BlockedGroupsAccumulator::update_batch`] and + /// [`BlockedGroupsAccumulator::merge_batch`]. + /// + /// [`BlockedGroupsAccumulator::update_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::update_batch + /// [`BlockedGroupsAccumulator::merge_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::merge_batch + /// + /// # Panics + /// If an index in `group_indices` is not below `total_num_groups`. + /// if there is an index in block that is larger than block size + pub fn update( + &mut self, + total_num_groups: usize, + starting_value: T, + group_indices: &[BlocksIndex], + nulls: Option<&NullBuffer>, + opt_filter: Option<&BooleanArray>, + update_fn: F, + ) where + F: FnMut(&mut T), + { + self.grow_to(total_num_groups, starting_value); + self.assert_in_bounds(group_indices); + + // SAFETY: grown to `total_num_groups` and checked above + unsafe { + self.update_unchecked(group_indices, nulls, opt_filter, update_fn); + } + } + + /// Same as [`Self::update`], but does not grow the vector and the caller guarantees that indices are in bounds. + /// + /// # Safety + /// The behavior is undefined if any of the following are true: + /// 1. `groups_indices` contain block index outside the number of blocks (e.g. you did not call `grow_to` first to the expected number of values) + /// 2. `groups_indices` index in block is larger than block size + /// 3. `groups_indices` index in block for last blocks is larger than the number of elements in the last block + pub unsafe fn update_unchecked( + &mut self, + group_indices: &[BlocksIndex], + nulls: Option<&NullBuffer>, + opt_filter: Option<&BooleanArray>, + mut update_fn: F, + ) where + F: FnMut(&mut T), + { + if let Some(block) = self.as_single_block_mut() { + accumulate_blocked_indices(group_indices, nulls, opt_filter, |index| { + // SAFETY: indices are in bounds, guaranteed by the caller + update_fn(unsafe { block.get_unchecked_mut(index.index_in_block()) }) + }); + return; + } + + let ptrs = self.block_ptrs_mut(); + if nulls.is_none() && opt_filter.is_none() { + let mut addrs = [std::ptr::null_mut::(); RESOLVE_CHUNK]; + for chunk in group_indices.chunks(RESOLVE_CHUNK) { + for (addr, &index) in addrs.iter_mut().zip(chunk) { + // SAFETY: indices are in bounds, guaranteed by the caller + *addr = unsafe { ptrs.ptr(index) }; + } + for &addr in &addrs[..chunk.len()] { + // SAFETY: resolved above from a block that is not + // touched while `ptrs` lives + update_fn(unsafe { &mut *addr }); + } + } + } else { + accumulate_blocked_indices(group_indices, nulls, opt_filter, |index| { + // SAFETY: indices are in bounds, guaranteed by the caller + update_fn(unsafe { &mut *ptrs.ptr(index) }) + }); + } + } + + /// Grows the vector to `total_num_groups` elements (new ones are + /// `starting_value`), then calls `update_fn` with the element of every + /// row's group and the row's value in `values`. + /// + /// Chooses the loop once per call, like [`Self::update`]. + /// + /// `group_indices` must only point to the first `total_num_groups` + /// groups, which is the contract of + /// [`BlockedGroupsAccumulator::update_batch`] and + /// [`BlockedGroupsAccumulator::merge_batch`]. + /// + /// [`BlockedGroupsAccumulator::update_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::update_batch + /// [`BlockedGroupsAccumulator::merge_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::merge_batch + /// + /// # Panics + /// If `values` and `group_indices` have different lengths, or an index in + /// `group_indices` is not below `total_num_groups`. + pub fn update_with( + &mut self, + total_num_groups: usize, + starting_value: T, + group_indices: &[BlocksIndex], + values: &[V], + update_fn: F, + ) where + V: Copy, + F: FnMut(&mut T, V), + { + self.grow_to(total_num_groups, starting_value); + self.assert_in_bounds(group_indices); + + // SAFETY: grown to `total_num_groups` and checked above + unsafe { self.update_with_unchecked(group_indices, values, update_fn) } + } + + /// Same as [`Self::update_with`], but does not grow the vector and the caller guarantees that indices are in bounds. + /// + /// # Safety + /// The behavior is undefined if any of the following are true: + /// 1. `groups_indices` contain block index outside the number of blocks (e.g. you did not call `grow_to` first to the expected number of values) + /// 2. `groups_indices` index in block is larger than block size + /// 3. `groups_indices` index in block for last block is larger than the number of elements in the last block (e.g. you did not call `grow_to` first to the expected number of values) + pub unsafe fn update_with_unchecked( + &mut self, + group_indices: &[BlocksIndex], + values: &[V], + mut update_fn: F, + ) where + V: Copy, + F: FnMut(&mut T, V), + { + assert_eq!(group_indices.len(), values.len()); + if let Some(block) = self.as_single_block_mut() { + for (&index, &value) in group_indices.iter().zip(values) { + // SAFETY: indices are in bounds, guaranteed by the caller + update_fn( + unsafe { block.get_unchecked_mut(index.index_in_block()) }, + value, + ); + } + return; + } + + let ptrs = self.block_ptrs_mut(); + let mut addrs = [std::ptr::null_mut::(); RESOLVE_CHUNK]; + for (indices, values) in group_indices + .chunks(RESOLVE_CHUNK) + .zip(values.chunks(RESOLVE_CHUNK)) + { + for (addr, &index) in addrs.iter_mut().zip(indices) { + // SAFETY: indices are in bounds, guaranteed by the caller + *addr = unsafe { ptrs.ptr(index) }; + } + for (&addr, &value) in addrs.iter().zip(values) { + // SAFETY: resolved above from a block that is not touched + // while `ptrs` lives + update_fn(unsafe { &mut *addr }, value); + } + } + } + + /// Checks that every index points to an element, so the update loops + /// can index without bounds checks. Every block but the last holds + /// `block_size` elements, see [`BlocksIndex::all_in_bounds`]. + #[inline] + fn assert_in_bounds(&self, group_indices: &[BlocksIndex]) { + assert!( + BlocksIndex::all_in_bounds(group_indices, self.block_size, self.len), + "group index out of bounds" + ); + } + + /// Removes and returns the first block, or `None` if empty. Elements + /// after it move down by the removed block's length. + pub fn take_next_block(&mut self) -> Option> { + if self.blocks.is_empty() { + return None; + } + // There are few blocks, so shifting the outer `Vec` is cheap + let block = self.blocks.remove(0); + self.len -= block.len(); + self.allocated -= block.capacity() * size_of::(); + Some(block) + } + + /// Removes and returns all blocks, leaving the vector empty. + pub fn take_all(&mut self) -> Vec> { + self.len = 0; + self.allocated = 0; + std::mem::take(&mut self.blocks) + } + + /// Removes and returns the first `n` elements; the remaining elements + /// move down by `n`, like `EmitTo::First` on a flat `Vec`. + /// + /// # Panics + /// If `n > self.len()`. + pub fn take_first(&mut self, n: usize) -> Vec { + assert!(n <= self.len, "take_first({n}) with len {}", self.len); + if n == 0 { + return Vec::new(); + } + + // Move out whole leading blocks + let mut taken = Vec::new(); + while let Some(first) = self.blocks.first() + && n - taken.len() >= first.len() + { + let block = self.blocks.remove(0); + if taken.is_empty() { + taken = block; + taken.reserve_exact(n - taken.len()); + } else { + taken.extend(block); + } + } + + // Take the rest from the (now) first block, and shift every following + // block down by `rest` so all blocks but the last stay full + let rest = n - taken.len(); + if rest > 0 { + taken.extend(self.blocks[0].drain(..rest)); + for i in 1..self.blocks.len() { + let (previous, current) = self.blocks.split_at_mut(i); + let current = &mut current[0]; + let moved = rest.min(current.len()); + previous[i - 1].extend(current.drain(..moved)); + } + if self.blocks.last().is_some_and(|b| b.is_empty()) { + self.blocks.pop(); + } + } + + self.len -= n; + if self.blocks.is_empty() { + self.blocks = Vec::new(); + } + // Rare path: recompute instead of tracking each block + self.allocated = self + .blocks + .iter() + .map(|b| b.capacity() * size_of::()) + .sum(); + taken + } + + /// Bytes allocated by this vector, O(1). + #[inline] + pub fn allocated_size(&self) -> usize { + self.allocated + self.blocks.capacity() * size_of::>() + } +} + +/// Raw addresses of the blocks of a [`BlockedVec`], see +/// [`BlockedVec::block_ptrs_mut`]. +/// +/// The pointers are invalidated by any mutation (or other use) of the +/// [`BlockedVec`] they were created from, so only keep a `BlockPtrs` for the +/// duration of one hot loop. +#[derive(Debug)] +struct BlockPtrs { + ptrs: Vec<*mut T>, +} + +impl BlockPtrs { + /// Returns the address of the element at `index`. + /// + /// # Safety + /// `index` must point to an element of the [`BlockedVec`] when this + /// `BlockPtrs` was created, and that [`BlockedVec`] must not have been + /// used since. + #[inline] + unsafe fn ptr(&self, index: BlocksIndex) -> *mut T { + // SAFETY: guaranteed by the caller + unsafe { + self.ptrs + .get_unchecked(index.block_index()) + .add(index.index_in_block()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stays_flat_until_first_block_fills() { + let mut v = BlockedVec::::new(4); + v.grow_to(3, 0); + assert_eq!(v.num_blocks(), 1); + assert!(v.as_single_block_mut().is_some()); + v.grow_to(4, 0); + assert_eq!(v.num_blocks(), 1); + v.grow_to(5, 0); + assert_eq!(v.num_blocks(), 2); + assert!(v.as_single_block_mut().is_none()); + } + + #[test] + fn indexing_across_blocks() { + let mut v = BlockedVec::::new(4); + for i in 0..10 { + assert_eq!(v.push(i as i64), BlocksIndex::from_flat(i, 4)); + } + for i in 0..10 { + assert_eq!(v.get(BlocksIndex::from_flat(i, 4)), i as i64); + } + let ptrs = v.block_ptrs_mut(); + unsafe { *ptrs.ptr(BlocksIndex::from_flat(9, 4)) += 100 }; + assert_eq!(v.get(BlocksIndex::from_flat(9, 4)), 109); + } + + #[test] + fn take_next_block_and_all() { + let mut v = BlockedVec::::new(4); + (0..10).for_each(|i| { + v.push(i); + }); + assert_eq!(v.take_next_block(), Some(vec![0, 1, 2, 3])); + assert_eq!(v.len(), 6); + assert_eq!(v.take_all(), vec![vec![4, 5, 6, 7], vec![8, 9]]); + assert!(v.is_empty()); + assert_eq!(v.take_next_block(), None); + assert_eq!(v.allocated_size(), 0); + } + + #[test] + fn take_first_shifts_across_blocks() { + let mut v = BlockedVec::::new(4); + (0..10).for_each(|i| { + v.push(i); + }); + assert_eq!(v.take_first(3), vec![0, 1, 2]); + assert_eq!(v.len(), 7); + let rest: Vec = (0..7) + .map(|i| v.get(BlocksIndex::from_flat(i, 4))) + .collect(); + assert_eq!(rest, vec![3, 4, 5, 6, 7, 8, 9]); + assert_eq!(v.num_blocks(), 2); + } + + #[test] + fn allocated_size_tracks_blocks() { + let mut v = BlockedVec::::new(4); + v.grow_to(9, 0); + assert!(v.allocated_size() >= 9 * 8); + let before = v.allocated_size(); + v.take_next_block(); + assert!(v.allocated_size() < before); + } + + #[test] + fn grow_to_spills_into_several_blocks_at_once() { + let mut v = BlockedVec::::new(4); + v.grow_to(2, 7); + assert_eq!(v.num_blocks(), 1); + v.grow_to(11, 1); + assert_eq!(v.len(), 11); + assert_eq!(v.num_blocks(), 3); + let all: Vec = (0..11) + .map(|i| v.get(BlocksIndex::from_flat(i, 4))) + .collect(); + assert_eq!(all, vec![7, 7, 1, 1, 1, 1, 1, 1, 1, 1, 1]); + // Shrinking is a no-op + v.grow_to(3, 0); + assert_eq!(v.len(), 11); + assert!(v.allocated_size() >= 3 * 4 * size_of::()); + let blocks = v.take_all(); + assert_eq!(blocks, vec![vec![7, 7, 1, 1], vec![1; 4], vec![1; 3]]); + // Blocks grow by doubling up to the block size, so the last one does + // not allocate a whole block + assert!(blocks[..2].iter().all(|b| b.capacity() == 4)); + assert_eq!(blocks[2].capacity(), 3); + } + + #[test] + fn take_first_whole_first_block() { + let mut v = BlockedVec::::new(4); + (0..10).for_each(|i| { + v.push(i); + }); + assert_eq!(v.take_first(4), vec![0, 1, 2, 3]); + assert_eq!(v.len(), 6); + assert_eq!(v.num_blocks(), 2); + let rest: Vec = (0..6) + .map(|i| v.get(BlocksIndex::from_flat(i, 4))) + .collect(); + assert_eq!(rest, vec![4, 5, 6, 7, 8, 9]); + // All remaining elements, emptying the vec + assert_eq!(v.take_first(6), vec![4, 5, 6, 7, 8, 9]); + assert!(v.is_empty()); + assert_eq!(v.num_blocks(), 0); + assert_eq!(v.allocated_size(), 0); + // Still usable afterwards + assert_eq!(v.push(42), BlocksIndex::new(0, 0)); + assert_eq!(v.get(BlocksIndex::from_flat(0, 4)), 42); + } + + #[test] + fn take_first_within_single_block() { + let mut v = BlockedVec::::new(8); + (0..5).for_each(|i| { + v.push(i); + }); + assert_eq!(v.take_first(2), vec![0, 1]); + assert_eq!(v.as_single_block_mut().unwrap(), &mut [2, 3, 4]); + } + + #[test] + fn allocated_size_zero_after_take_all() { + let mut v = BlockedVec::::new(4); + v.grow_to(3, 0); + assert!(v.allocated_size() >= 3 * 8); + let _ = v.take_all(); + assert_eq!(v.allocated_size(), 0); + v.grow_to(9, 0); + assert_eq!(v.num_blocks(), 3); + let _ = v.take_all(); + assert_eq!(v.allocated_size(), 0); + assert_eq!(v.num_blocks(), 0); + } + + fn indices(flats: &[usize]) -> Vec { + flats + .iter() + .map(|&f| BlocksIndex::from_flat(f, 4)) + .collect() + } + + fn values(v: &BlockedVec) -> Vec { + (0..v.len()) + .map(|f| v.get(BlocksIndex::from_flat(f, 4))) + .collect() + } + + #[test] + fn update_single_and_multi_block() { + let mut v = BlockedVec::::new(4); + v.update(3, 0, &indices(&[0, 2, 2]), None, None, |x| *x += 1); + assert_eq!(values(&v), vec![1, 0, 2]); + + // several blocks, more rows than one resolve chunk + let rows: Vec = (0..RESOLVE_CHUNK * 2 + 3).map(|i| i % 10).collect(); + v.update(10, 0, &indices(&rows), None, None, |x| *x += 1); + let expected: Vec = (0..10) + .map(|g| rows.iter().filter(|&&r| r == g).count() as i64) + .zip([1, 0, 2, 0, 0, 0, 0, 0, 0, 0]) + .map(|(a, b)| a + b) + .collect(); + assert_eq!(values(&v), expected); + } + + #[test] + fn update_skips_nulls_and_filtered_rows() { + let mut v = BlockedVec::::new(4); + let nulls = NullBuffer::from(vec![true, false, true, true]); + let filter = BooleanArray::from(vec![true, true, false, true]); + v.update( + 9, + 0, + &indices(&[0, 5, 8, 8]), + Some(&nulls), + Some(&filter), + |x| *x += 1, + ); + assert_eq!(values(&v), vec![1, 0, 0, 0, 0, 0, 0, 0, 1]); + } + + #[test] + fn update_with_values_across_blocks() { + let mut v = BlockedVec::::new(4); + let rows: Vec = (0..RESOLVE_CHUNK + 5).map(|i| i % 10).collect(); + let vals: Vec = (0..rows.len() as i64).collect(); + v.update_with(10, 0, &indices(&rows), &vals, |x, val| *x += val); + let expected: Vec = (0..10) + .map(|g| { + rows.iter() + .zip(&vals) + .filter(|(r, _)| **r == g) + .map(|(_, val)| val) + .sum() + }) + .collect(); + assert_eq!(values(&v), expected); + } + + #[test] + #[should_panic(expected = "block_size must be positive")] + fn zero_block_size_panics() { + let _ = BlockedVec::::new(0); + } + + #[test] + #[should_panic(expected = "group index out of bounds")] + fn update_out_of_bounds_panics() { + let mut v = BlockedVec::::new(4); + v.update(1, 0, &[BlocksIndex::new(0, 1)], None, None, |v| *v += 1); + } + + #[test] + #[should_panic(expected = "group index out of bounds")] + fn update_with_out_of_bounds_panics() { + let mut v = BlockedVec::::new(4); + v.update_with(6, 0, &[BlocksIndex::new(0, 4)], &[1], |v, x| *v += x); + } +} diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 7c3bfdceefdf7..bd58f55f38ea2 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -35,9 +35,10 @@ use datafusion_common::{ stats::Precision, utils::expr::COUNT_STAR_EXPANSION, }; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupSelection, - GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, - TypeSignature, Volatility, WindowFunctionDefinition, + Accumulator, AggregateUDFImpl, BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, + Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, ReversedUDAF, + SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility, + WindowFunctionDefinition, expr::WindowFunction, function::{AccumulatorArgs, StateFieldsArgs}, utils::{AggregateOrderSensitivity, format_state_name}, @@ -54,6 +55,7 @@ use datafusion_functions_aggregate_common::aggregate::{ count_distinct::FloatDistinctCountAccumulator, count_distinct::PrimitiveDistinctCountAccumulator, groups_accumulator::accumulate::accumulate_indices, + groups_accumulator::blocked_vec::BlockedVec, }; use datafusion_macros::user_doc; use datafusion_physical_expr::expressions; @@ -392,6 +394,18 @@ impl AggregateUDFImpl for Count { create_distinct_count_groups_accumulator(&args) } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + !args.is_distinct && self.groups_accumulator_supported(args) + } + + fn create_blocked_groups_accumulator( + &self, + _args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + Ok(Box::new(BlockedCountGroupsAccumulator::new(block_size))) + } + fn reverse_expr(&self) -> ReversedUDAF { ReversedUDAF::Identical } @@ -828,6 +842,125 @@ impl GroupsAccumulator for CountGroupsAccumulator { } } +/// [`CountGroupsAccumulator`] with the counts stored in blocks of +/// `block_size` groups, so growing past a block never reallocates and +/// copies the existing counts. +#[derive(Debug)] +struct BlockedCountGroupsAccumulator { + /// Count per group, see [`CountGroupsAccumulator::counts`]. + counts: BlockedVec, +} + +impl BlockedCountGroupsAccumulator { + fn new(block_size: usize) -> Self { + Self { + counts: BlockedVec::new(block_size), + } + } + + fn counts_to_array(counts: Vec) -> ArrayRef { + // zero copy, count is never null + Arc::new(Int64Array::new(counts.into(), None)) + } + + fn blocks_to_arrays(blocks: impl IntoIterator>) -> Vec { + blocks.into_iter().map(Self::counts_to_array).collect() + } + + fn take(&mut self, emit_to: BlockedEmitTo) -> Vec { + match emit_to { + BlockedEmitTo::All => Self::blocks_to_arrays(self.counts.take_all()), + BlockedEmitTo::NextBlock => { + Self::blocks_to_arrays(self.counts.take_next_block()) + } + BlockedEmitTo::First(n) => { + Self::blocks_to_arrays([self.counts.take_first(n)]) + } + } + } +} + +impl BlockedGroupsAccumulator for BlockedCountGroupsAccumulator { + fn block_size(&self) -> usize { + self.counts.block_size() + } + + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + assert_eq!(values.len(), 1, "single argument to update_batch"); + let values = &values[0]; + let nulls = values.logical_nulls().filter(|n| n.null_count() > 0); + + self.counts.grow_to(total_num_groups, 0); + + // Add one to each group's counter for each non null, non filtered value + // SAFETY: group_index is guaranteed to be in bounds and less than total_num_groups + unsafe { + self.counts.update_unchecked( + group_indices, + nulls.as_ref(), + opt_filter, + |count| *count += 1, + ); + } + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()> { + assert_eq!(values.len(), 1, "one argument to merge_batch"); + let partial_counts = values[0].as_primitive::(); + + // intermediate counts are always created as non null + assert_eq!(partial_counts.null_count(), 0); + let partial_counts = partial_counts.values(); + + // Adds the counts with the partial counts + self.counts.update_with( + total_num_groups, + 0, + group_indices, + partial_counts, + |count, partial_count| *count += partial_count, + ); + + Ok(()) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + Ok(self.take(emit_to)) + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + Ok(self + .take(emit_to) + .into_iter() + .map(|counts| vec![counts]) + .collect()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + CountGroupsAccumulator::new().convert_to_state(values, opt_filter) + } + + fn size(&self) -> usize { + self.counts.allocated_size() + } +} + /// count null values for multiple columns /// for each row if one column value is null, then null_count + 1 fn null_count_for_multiple_cols(values: &[ArrayRef]) -> usize { @@ -999,6 +1132,181 @@ mod tests { Ok(()) } + #[test] + fn blocked_count_update_merge_emit_across_blocks() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(4); + let values: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + Some(4), + Some(5), + Some(6), + ])); + let groups = blocks_index(&[0, 0, 5, 5, 1, 9], 4); + acc.update_batch(&[values], &groups, None, 10)?; + + let partial: ArrayRef = Arc::new(Int64Array::from(vec![10, 20])); + acc.merge_batch(&[partial], &blocks_index(&[1, 9], 4), 10)?; + + let first = acc.evaluate(BlockedEmitTo::NextBlock)?; + assert_eq!(first.len(), 1); + assert_eq!( + first[0].as_primitive::().values(), + &[1, 11, 0, 0] + ); + let rest = acc.evaluate(BlockedEmitTo::All)?; + assert_eq!(rest.len(), 2); + assert_eq!(rest[0].as_primitive::().values(), &[0, 2, 0, 0]); + assert_eq!(rest[1].as_primitive::().values(), &[0, 21]); + assert!(acc.evaluate(BlockedEmitTo::NextBlock)?.is_empty()); + Ok(()) + } + + fn blocks_index(groups: &[usize], block_size: usize) -> Vec { + groups + .iter() + .map(|&flat| BlocksIndex::from_flat(flat, block_size)) + .collect() + } + + fn blocked_values(arrays: &[ArrayRef]) -> Vec> { + arrays + .iter() + .map(|a| a.as_primitive::().values().to_vec()) + .collect() + } + + #[test] + fn blocked_count_nulls_and_filter_across_blocks() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(4); + let groups = blocks_index(&[0, 5, 9, 9, 5, 0, 6], 4); + let values: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + ])); + + // nulls only + acc.update_batch(&[Arc::clone(&values)], &groups, None, 10)?; + // filter (with a null filter entry) and nulls + let filter = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(false), + None, + Some(true), + Some(true), + Some(false), + ]); + acc.update_batch(&[Arc::clone(&values)], &groups, Some(&filter), 10)?; + // filter only + let no_nulls: ArrayRef = Arc::new(Int32Array::from(vec![0; 7])); + acc.update_batch(&[no_nulls], &groups, Some(&filter), 10)?; + + assert_eq!( + blocked_values(&acc.evaluate(BlockedEmitTo::All)?), + vec![vec![6, 0, 0, 0], vec![0, 4, 1, 0], vec![0, 2]] + ); + Ok(()) + } + + #[test] + fn blocked_count_multi_block_chunks() -> Result<()> { + // more rows than one address-resolve chunk, spread over 3 blocks + let mut acc = BlockedCountGroupsAccumulator::new(4); + let rows = 256 * 2 + 7; + let groups: Vec = (0..rows).map(|i| (i * 7) % 11).collect(); + let values: ArrayRef = Arc::new(Int32Array::from(vec![1; rows])); + acc.update_batch(&[values], &blocks_index(&groups, 4), None, 11)?; + + let mut expected = vec![0i64; 11]; + for &g in &groups { + expected[g] += 1; + } + let actual: Vec = blocked_values(&acc.evaluate(BlockedEmitTo::All)?) + .into_iter() + .flatten() + .collect(); + assert_eq!(actual, expected); + Ok(()) + } + + #[test] + fn blocked_count_single_block() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(8); + let values: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4)])); + acc.update_batch(&[values], &blocks_index(&[2, 2, 0, 2], 8), None, 3)?; + assert!(acc.counts.as_single_block_mut().is_some()); + let partial: ArrayRef = Arc::new(Int64Array::from(vec![5])); + acc.merge_batch(&[partial], &blocks_index(&[1], 8), 3)?; + assert_eq!( + blocked_values(&acc.evaluate(BlockedEmitTo::All)?), + vec![vec![1, 5, 2]] + ); + Ok(()) + } + + #[test] + fn blocked_count_emit_first() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(4); + let groups: Vec = (0..10).collect(); + let values: ArrayRef = Arc::new(Int32Array::from(vec![1; 10])); + acc.update_batch(&[values], &blocks_index(&groups, 4), None, 10)?; + let partial: ArrayRef = Arc::new(Int64Array::from((0..10).collect::>())); + acc.merge_batch(&[partial], &blocks_index(&groups, 4), 10)?; + + assert_eq!( + blocked_values(&acc.evaluate(BlockedEmitTo::First(3))?), + vec![vec![1, 2, 3]] + ); + // remaining groups shift down by 3 + let values: ArrayRef = Arc::new(Int32Array::from(vec![1])); + acc.update_batch(&[values], &blocks_index(&[0], 4), None, 7)?; + assert_eq!( + blocked_values(&acc.evaluate(BlockedEmitTo::All)?), + vec![vec![5, 5, 6, 7], vec![8, 9, 10]] + ); + Ok(()) + } + + #[test] + fn blocked_count_state_is_one_column_per_block() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(4); + let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + acc.update_batch(&[values], &blocks_index(&[0, 4, 8], 4), None, 9)?; + + let state = acc.state(BlockedEmitTo::All)?; + assert_eq!(state.len(), 3); + assert!(state.iter().all(|columns| columns.len() == 1)); + let columns: Vec = state.into_iter().flatten().collect(); + assert_eq!( + blocked_values(&columns), + vec![vec![1, 0, 0, 0], vec![1, 0, 0, 0], vec![1]] + ); + assert!(acc.state(BlockedEmitTo::NextBlock)?.is_empty()); + Ok(()) + } + + #[test] + fn blocked_count_size_tracks_blocks() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(4); + assert_eq!(acc.size(), 0); + let values: ArrayRef = Arc::new(Int32Array::from(vec![1; 9])); + let groups: Vec = (0..9).collect(); + acc.update_batch(&[values], &blocks_index(&groups, 4), None, 9)?; + assert_eq!(acc.size(), acc.counts.allocated_size()); + assert!(acc.size() >= 9 * size_of::()); + acc.evaluate(BlockedEmitTo::All)?; + assert_eq!(acc.size(), 0); + Ok(()) + } + #[test] fn count_accumulator_nulls() -> Result<()> { let mut accumulator = CountAccumulator::new(); diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index 6d95d8ea12bd8..7bd3eeee2a263 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -54,6 +54,7 @@ use datafusion_expr::expr::{ use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{AggregateUDF, Expr, ReversedUDAF, SetMonotonicity}; use datafusion_expr_common::accumulator::{Accumulator, AggregateMetrics}; +use datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator; use datafusion_expr_common::groups_accumulator::GroupsAccumulator; use datafusion_expr_common::type_coercion::aggregates::check_arg_count; use datafusion_functions_aggregate_common::accumulator::{ @@ -953,6 +954,44 @@ impl AggregateFunctionExpr { Ok(accumulator) } + /// If the aggregate expression has a [`BlockedGroupsAccumulator`] + /// implementation. If this returns true, + /// [`Self::create_blocked_groups_accumulator`] may be called. + pub fn blocked_groups_accumulator_supported(&self) -> bool { + let args = AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + expr_fields: &self.arg_fields, + ignore_nulls: self.ignore_nulls, + order_bys: self.order_bys.as_ref(), + is_distinct: self.is_distinct, + name: &self.name, + is_reversed: self.is_reversed, + exprs: &self.args, + }; + self.fun.blocked_groups_accumulator_supported(args) + } + + /// Return a [`BlockedGroupsAccumulator`] that manages state for all + /// groups in blocks of `block_size` groups. + pub fn create_blocked_groups_accumulator( + &self, + block_size: usize, + ) -> Result> { + let args = AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + expr_fields: &self.arg_fields, + ignore_nulls: self.ignore_nulls, + order_bys: self.order_bys.as_ref(), + is_distinct: self.is_distinct, + name: &self.name, + is_reversed: self.is_reversed, + exprs: &self.args, + }; + self.fun.create_blocked_groups_accumulator(args, block_size) + } + /// Construct an expression that calculates the aggregate in reverse. /// Typically the "reverse" expression is itself (e.g. SUM, COUNT). /// For aggregates that do not support calculation in reverse, 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 b25caf815eed1..3480f96c129e6 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -16,7 +16,7 @@ // under the License. use std::borrow::Cow; -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::marker::PhantomData; use std::sync::Arc; @@ -27,7 +27,6 @@ use arrow::compute::{filter_record_batch, prep_null_mask_filter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; -use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{AggregateMetrics, EmitTo, GroupsAccumulator}; use datafusion_physical_expr::GroupsAccumulatorAdapter; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; @@ -36,7 +35,7 @@ use log::debug; use crate::PhysicalExpr; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, GroupValues, new_group_values, + GroupByMetrics, }; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -45,6 +44,10 @@ use crate::aggregates::{ }; use super::AggregateTableMetrics; +use super::storage::{ + AccumulatorStorage, Emit, GroupIndices, GroupKeys, MaterializedBatch, block_size, + materialize_batches, use_blocked_keys, +}; /// Marker for raw rows -> partial state aggregation. pub(in crate::aggregates) struct PartialMarker; @@ -80,6 +83,27 @@ pub(in crate::aggregates) fn create_group_accumulator( } } +/// Create the state storage for `agg_expr`: a [`BlockedGroupsAccumulator`] +/// with `block_size` groups per block if the aggregate has one, otherwise the +/// flat accumulator from [`create_group_accumulator`]. +/// +/// [`BlockedGroupsAccumulator`]: datafusion_expr::BlockedGroupsAccumulator +pub(super) fn create_accumulator_storage( + agg_expr: &Arc, + metrics: Arc, + block_size: usize, +) -> Result { + if agg_expr.blocked_groups_accumulator_supported() { + let mut accumulator = agg_expr.create_blocked_groups_accumulator(block_size)?; + accumulator.set_metrics(metrics); + Ok(AccumulatorStorage::Blocked(accumulator)) + } else { + Ok(AccumulatorStorage::Flat(create_group_accumulator( + agg_expr, metrics, + )?)) + } +} + /// Grouped hash table shared by the partial and final paths. /// /// While building, it consumes input batches and updates group / accumulator @@ -97,6 +121,8 @@ pub(in crate::aggregates) fn create_group_accumulator( /// [`GroupsAccumulator`]. Both use columnar storage so aggregation can stay /// vectorized. /// +/// [`GroupValues`]: crate::aggregates::group_values::GroupValues +/// /// # Marker Type /// `AggrMode` selects the aggregate semantics. /// @@ -160,15 +186,21 @@ impl AggregateHashTable { &agg.mode, agg.group_by().num_group_exprs(), )?; - let accumulators: Vec<_> = agg + let group_schema = agg.group_by().group_schema(&input_schema)?; + let block_size = block_size(batch_size); + + let accumulators: Vec = agg .aggr_expr() .iter() .zip(aggregate_arguments) .zip(filters) .zip(metrics.submetrics.iter()) .map(|(((agg_expr, arguments), filter), submetrics)| { - let accumulator = - create_group_accumulator(agg_expr, Arc::clone(submetrics))?; + let accumulator = create_accumulator_storage( + agg_expr, + Arc::clone(submetrics), + block_size, + )?; Ok(HashAggregateAccumulator::new( Arc::clone(agg_expr), arguments, @@ -178,9 +210,16 @@ impl AggregateHashTable { )) }) .collect::>()?; - - let group_schema = agg.group_by().group_schema(&input_schema)?; - let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let keys = GroupKeys::try_new( + Arc::clone(&group_schema), + &GroupOrdering::None, + use_blocked_keys(agg, &group_schema, block_size), + block_size, + &accumulators + .iter() + .map(|acc| acc.storage()) + .collect::>(), + )?; Ok(Self { group_by_metrics: metrics.group_by, @@ -193,8 +232,7 @@ impl AggregateHashTable { batch_size, state: AggregateHashTableState::Building(AggregateHashTableBuffer { group_by: Arc::clone(agg.group_by()), - group_values, - batch_group_indices: Default::default(), + keys, accumulators, }), _mode: PhantomData, @@ -249,16 +287,13 @@ impl AggregateHashTable { let state = self.state.building_mut(); for group_values in &evaluated_batch.grouping_set_args { - group_by_metrics.time_group_key_preparation(|| { - state - .group_values - .intern(group_values, &mut state.batch_group_indices) - })?; + group_by_metrics + .time_group_key_preparation(|| state.keys.intern(group_values))?; // Register groups from the full input. Each filtered aggregate compacts // this row-aligned vector independently immediately before its update. - let group_indices = &state.batch_group_indices; - let total_num_groups = state.group_values.len(); + let keys = &state.keys; + let total_num_groups = keys.len(); group_by_metrics.time_aggregation(|| { for (idx, (acc, values)) in state .accumulators @@ -266,6 +301,7 @@ impl AggregateHashTable { .zip(evaluated_batch.accumulator_args.iter()) .enumerate() { + let group_indices = keys.indices_for(acc.storage()); accumulator_metrics.time(idx, accumulator_phase, || { aggregate_fn(acc, values, group_indices, total_num_groups) })?; @@ -277,15 +313,16 @@ impl AggregateHashTable { Ok(()) } - /// Materializes the full output once, then returns it downstream incrementally - /// by slicing it into `batch_size` chunks. + /// Materializes the full output once, then returns it downstream + /// incrementally by slicing it into `batch_size` chunks. + /// + /// Flat storage materializes one batch. Blocked storage materializes one + /// batch per block without copying, so each block's memory is released as + /// soon as its last slice is handed out. /// /// Each aggregation mode chooses a different `materialize_accumulator_fn` /// according to its semantics. For example, partial aggregation emits /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`]. - /// - /// This is a temporary solution until blocked state management is implemented: - /// Issue: pub(super) fn next_output_batch_inner( &mut self, materialize_accumulator_fn: MaterializeAccumulatorFn, @@ -298,28 +335,25 @@ impl AggregateHashTable { let mut output = match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { AggregateHashTableState::Outputting(mut state) => { - if state.group_values.is_empty() { + if state.keys.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) + // groups once, then slice the materialized batches on later polls. + let batches = self.group_by_metrics.time_emitting(|| { + materialize_batches( + &mut state.keys, + &mut state.accumulators, + EmitTo::All, + materialize_accumulator_fn, + accumulator_phase, + &accumulator_metrics, + &output_schema, + ) })?; - - let batch = RecordBatch::try_new(output_schema, columns)?; - debug_assert!(batch.num_rows() > 0); - MaterializedAggregateOutput::new(batch) + debug_assert!(batches.iter().all(|b| b.batch.num_rows() > 0)); + MaterializedAggregateOutput::new(batches) } AggregateHashTableState::OutputtingMaterialized(output) => output, AggregateHashTableState::Done => return Ok(None), @@ -342,16 +376,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn memory_size(&self) -> usize { match &self.state { AggregateHashTableState::Building(state) - | AggregateHashTableState::Outputting(state) => { - let acc = state - .accumulators - .iter() - .map(|acc| acc.accumulator.size()) - .sum::(); - - acc + state.group_values.size() - + state.batch_group_indices.allocated_size() - } + | AggregateHashTableState::Outputting(state) => state.size(), AggregateHashTableState::OutputtingMaterialized(output) => { output.memory_size() } @@ -361,7 +386,7 @@ impl AggregateHashTable { /// Returns the number of distinct groups accumulated so far. pub(in crate::aggregates) fn building_group_count(&self) -> usize { - self.state.building().group_values.len() + self.state.building().keys.len() } /// Takes every intermediate aggregate state and resets the table so it can @@ -370,40 +395,39 @@ impl AggregateHashTable { /// Unlike normal single aggregation output, this materializes intermediate /// states rather than final values. The states can therefore be merged after /// spilling without finalizing the same group more than once. - pub(in crate::aggregates) fn take_state_batch( + /// + /// Returns no batch if there are no groups, one batch for flat storage, + /// and one batch per block for blocked storage. + pub(in crate::aggregates) fn take_state_batches( &mut self, - ) -> Result> { + ) -> Result> { let state_schema = Arc::clone(&self.state_schema); let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); let group_by_metrics = self.group_by_metrics.clone(); let state = self.state.building_mut(); - if state.group_values.is_empty() { - return Ok(None); + if state.keys.is_empty() { + return Ok(vec![]); } - let output = group_by_metrics.time_emitting(|| { - let mut output = state.group_values.emit(EmitTo::All)?; - for (idx, acc) in state.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( - idx, - AccumulatorPhase::State, - || acc.state(EmitTo::All), - )?); - } - Ok::<_, datafusion_common::DataFusionError>(output) + let batches = group_by_metrics.time_emitting(|| { + materialize_batches( + &mut state.keys, + &mut state.accumulators, + EmitTo::All, + HashAggregateAccumulator::state, + AccumulatorPhase::State, + &accumulator_metrics, + &state_schema, + ) })?; + debug_assert!(batches.iter().all(|b| b.batch.num_rows() > 0)); - let batch = RecordBatch::try_new(state_schema, output)?; - debug_assert!(batch.num_rows() > 0); - - // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // Emitting all groups resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released // before the batch is sorted for spilling. - state.group_values.clear_shrink(0); - state.batch_group_indices.clear(); - state.batch_group_indices.shrink_to_fit(); + state.keys.clear_shrink(0); - Ok(Some(batch)) + Ok(batches) } pub(in crate::aggregates) fn is_building(&self) -> bool { @@ -421,7 +445,7 @@ impl AggregateHashTable { unreachable!("hash aggregate table is not building") }; - state.batch_group_indices = Vec::new(); + state.keys.free_indices(); self.state = AggregateHashTableState::Outputting(state); } @@ -447,7 +471,7 @@ impl AggregateHashTable { pub(super) fn init_empty_grouping_sets(&mut self) -> Result<()> { let group_by_metrics = self.group_by_metrics.clone(); let state = self.state.building_mut(); - if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + if !state.group_by.has_grouping_set() || !state.keys.is_empty() { return Ok(()); } @@ -479,16 +503,14 @@ impl AggregateHashTable { .collect(); cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); - state - .group_values - .intern(&cols, &mut state.batch_group_indices)?; + state.keys.intern(&cols)?; any_interned = true; } Ok::<_, datafusion_common::DataFusionError>(any_interned) })?; if any_interned { - let total_groups = state.group_values.len(); + let total_groups = state.keys.len(); let values = state .accumulators .iter() @@ -503,8 +525,9 @@ impl AggregateHashTable { for (idx, (acc, values)) in state.accumulators.iter_mut().zip(values.iter()).enumerate() { + let no_groups = acc.storage().no_groups(); accumulator_metrics.time(idx, AccumulatorPhase::Update, || { - acc.update_batch(values, &[], total_groups) + acc.update_batch(values, no_groups, total_groups) })?; } Ok::<(), datafusion_common::DataFusionError>(()) @@ -535,7 +558,7 @@ pub(super) struct HashAggregateAccumulator { filter: Option>, /// Accumulator state for all groups for one aggregate expression. - accumulator: Box, + accumulator: AccumulatorStorage, /// Optional internal metrics owned by this aggregate expression. submetrics: Arc, @@ -555,7 +578,7 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator; pub(super) type AggregateBatchFn = fn( &mut AggregateAccumulator, &CompactedAccumulatorArgs, - &[usize], + GroupIndices<'_>, usize, ) -> Result<()>; @@ -565,8 +588,11 @@ pub(super) type AggregateBatchFn = fn( /// Arguments: /// * accumulator to materialize. /// * group range to emit from the accumulator. +/// +/// Returns the output columns indexed `[block][column]`; flat storage always +/// returns one block. pub(super) type MaterializeAccumulatorFn = - fn(&mut AggregateAccumulator, EmitTo) -> Result>; + fn(&mut AggregateAccumulator, Emit) -> Result>>; /// Aggregate arguments compacted according to one aggregate's `FILTER`. pub(super) struct CompactedAccumulatorArgs { @@ -605,18 +631,17 @@ pub(super) struct EvaluatedAggregateBatch { /// /// [`GroupValues`] stores the physical group-key layout, while /// [`GroupsAccumulator`] stores per-group aggregate state. +/// +/// [`GroupValues`]: crate::aggregates::group_values::GroupValues pub(super) struct AggregateHashTableBuffer { /// GROUP BY expressions evaluated for each input batch. pub(super) group_by: Arc, - /// Interned group keys. Accumulator state is stored separately by group index. - pub(super) group_values: Box, - - /// Group index for each row in the current input batch. - /// - /// Each value indexes into `group_values`, and the same index is used by every - /// accumulator to update that group's aggregate state. - pub(super) batch_group_indices: Vec, + /// Interned group keys, and the group index of each row in the current + /// input batch. Accumulator state is stored separately by group index, and + /// the same index is used by every accumulator to update that group's + /// aggregate state. + pub(super) keys: GroupKeys, /// One item per aggregate expression. /// @@ -625,6 +650,18 @@ pub(super) struct AggregateHashTableBuffer { pub(super) accumulators: Vec, } +impl AggregateHashTableBuffer { + /// Memory used by the group keys and every accumulator's state. + pub(super) fn size(&self) -> usize { + self.keys.size() + + self + .accumulators + .iter() + .map(|acc| acc.accumulator.size()) + .sum::() + } +} + pub(super) enum AggregateHashTableState { /// Accumulating input rows into group keys and aggregate state. Building(AggregateHashTableBuffer), @@ -643,34 +680,47 @@ pub(super) enum AggregateHashTableState { /// Final aggregate evaluation consumes accumulator state, and partial terminal /// output should not repeatedly renumber group values with `EmitTo::First`. /// Materialize once and then slice to honor `batch_size` across output polls. +/// +/// Holds one batch for flat storage and one batch per block for blocked +/// storage. A batch is dropped once all of it was handed out, so the memory of +/// each block is released block by block. pub(super) struct MaterializedAggregateOutput { - batch: RecordBatch, + batches: VecDeque, + /// Rows of the front batch already handed out. offset: usize, } impl MaterializedAggregateOutput { - pub(super) fn new(batch: RecordBatch) -> Self { - Self { batch, offset: 0 } + pub(super) fn new(batches: Vec) -> Self { + Self { + batches: batches.into(), + offset: 0, + } } + /// Returns the next slice of at most `batch_size` rows. Never spans two + /// batches, so the last slice of a block may be shorter. pub(super) fn next_batch(&mut self, batch_size: usize) -> Option { debug_assert!(batch_size > 0); - if self.is_exhausted() { - return None; - } - - let length = batch_size.min(self.batch.num_rows() - self.offset); - let batch = self.batch.slice(self.offset, length); + let front = &self.batches.front()?.batch; + let length = batch_size.min(front.num_rows() - self.offset); + let batch = front.slice(self.offset, length); self.offset += length; + if self.offset >= front.num_rows() { + self.batches.pop_front(); + self.offset = 0; + } Some(batch) } pub(super) fn is_exhausted(&self) -> bool { - self.offset >= self.batch.num_rows() + self.batches.is_empty() } + /// Memory of the batches not fully handed out yet, see + /// [`MaterializedBatch::memory_size`]. pub(super) fn memory_size(&self) -> usize { - self.batch.get_array_memory_size() + self.batches.iter().map(|b| b.memory_size).sum() } } @@ -679,10 +729,10 @@ impl MaterializedAggregateOutput { /// Returns `None` when every row is selected so callers can reuse the input /// slice without allocating. At high selectivity, copying contiguous selected /// ranges avoids branching once per row. -fn compact_group_indices( - group_indices: &[usize], +fn compact_group_indices( + group_indices: &[T], filter: &BooleanArray, -) -> Option> { +) -> Option> { debug_assert_eq!(group_indices.len(), filter.len()); let filter = match filter.null_count() { @@ -715,7 +765,7 @@ impl HashAggregateAccumulator { aggregate_expr: Arc, arguments: Vec>, filter: Option>, - accumulator: Box, + accumulator: AccumulatorStorage, submetrics: Arc, ) -> Self { Self { @@ -727,11 +777,18 @@ impl HashAggregateAccumulator { } } + /// Storage of this accumulator's state. + pub(super) fn storage(&self) -> &AccumulatorStorage { + &self.accumulator + } + /// Construct a new accumulator with the same definition, but with empty internal - /// state buffers (empty [`GroupsAccumulator`]). + /// state buffers (empty flat [`GroupsAccumulator`]). pub(super) fn empty_like(&self) -> Result { - let accumulator = - create_group_accumulator(&self.aggregate_expr, Arc::clone(&self.submetrics))?; + let accumulator = AccumulatorStorage::Flat(create_group_accumulator( + &self.aggregate_expr, + Arc::clone(&self.submetrics), + )?); Ok(Self::new( Arc::clone(&self.aggregate_expr), self.arguments.clone(), @@ -835,52 +892,87 @@ impl HashAggregateAccumulator { pub(super) fn update_batch( &mut self, values: &CompactedAccumulatorArgs, - group_indices: &[usize], + group_indices: GroupIndices<'_>, total_num_groups: usize, ) -> Result<()> { - let filtered_group_indices = values - .selection - .as_ref() - .and_then(|selection| compact_group_indices(group_indices, selection)); - let group_indices = filtered_group_indices.as_deref().unwrap_or(group_indices); - self.accumulator.update_batch( - &values.arguments, - group_indices, - None, - total_num_groups, - ) + match (&mut self.accumulator, group_indices) { + (AccumulatorStorage::Flat(acc), GroupIndices::Flat(group_indices)) => { + let filtered_group_indices = + values.selection.as_ref().and_then(|selection| { + compact_group_indices(group_indices, selection) + }); + let group_indices = + filtered_group_indices.as_deref().unwrap_or(group_indices); + acc.update_batch(&values.arguments, group_indices, None, total_num_groups) + } + (AccumulatorStorage::Blocked(acc), GroupIndices::Blocked(group_indices)) => { + let filtered_group_indices = + values.selection.as_ref().and_then(|selection| { + compact_group_indices(group_indices, selection) + }); + let group_indices = + filtered_group_indices.as_deref().unwrap_or(group_indices); + acc.update_batch(&values.arguments, group_indices, None, total_num_groups) + } + _ => mismatched_storage(), + } } pub(super) fn merge_batch( &mut self, values: &CompactedAccumulatorArgs, - group_indices: &[usize], + group_indices: GroupIndices<'_>, total_num_groups: usize, ) -> Result<()> { debug_assert!(values.selection.is_none()); - self.accumulator - .merge_batch(&values.arguments, group_indices, total_num_groups) + match (&mut self.accumulator, group_indices) { + (AccumulatorStorage::Flat(acc), GroupIndices::Flat(group_indices)) => { + acc.merge_batch(&values.arguments, group_indices, total_num_groups) + } + (AccumulatorStorage::Blocked(acc), GroupIndices::Blocked(group_indices)) => { + acc.merge_batch(&values.arguments, group_indices, total_num_groups) + } + _ => mismatched_storage(), + } } - /// Evaluating final aggregate results according to `EmitTo`, and reset inner + /// Evaluating final aggregate results according to `emit`, and reset inner /// states. (e.g. after `evaluate(EmitTo::All)`, it returns all accumulated groups /// , and clear the inner buffers) - pub(super) fn evaluate(&mut self, emit_to: EmitTo) -> Result { - self.accumulator.evaluate(emit_to) - } - + /// + /// Returns one final-value column per block. pub(super) fn evaluate_to_columns( &mut self, - emit_to: EmitTo, - ) -> Result> { - Ok(vec![self.evaluate(emit_to)?]) + emit: Emit, + ) -> Result>> { + match (&mut self.accumulator, emit) { + (AccumulatorStorage::Flat(acc), Emit::Flat(emit_to)) => { + Ok(vec![vec![acc.evaluate(emit_to)?]]) + } + (AccumulatorStorage::Blocked(acc), Emit::Blocked(emit_to)) => Ok(acc + .evaluate(emit_to)? + .into_iter() + .map(|array| vec![array]) + .collect()), + _ => mismatched_storage(), + } } - /// Evaluating partial aggregate results according to `EmitTo`, and reset inner + /// Evaluating partial aggregate results according to `emit`, and reset inner /// states. (e.g. after `state(EmitTo::All)`, it returns all accumulated groups /// , and clear the inner buffers) - pub(super) fn state(&mut self, emit_to: EmitTo) -> Result> { - self.accumulator.state(emit_to) + /// + /// Returns the state columns indexed `[block][column]`. + pub(super) fn state(&mut self, emit: Emit) -> Result>> { + match (&mut self.accumulator, emit) { + (AccumulatorStorage::Flat(acc), Emit::Flat(emit_to)) => { + Ok(vec![acc.state(emit_to)?]) + } + (AccumulatorStorage::Blocked(acc), Emit::Blocked(emit_to)) => { + acc.state(emit_to) + } + _ => mismatched_storage(), + } } /// Converts evaluated row-aligned arguments directly to partial state. @@ -888,8 +980,14 @@ impl HashAggregateAccumulator { &self, values: &RowAlignedAccumulatorArgs, ) -> Result> { - self.accumulator - .convert_to_state(&values.arguments, values.filter.as_ref()) + match &self.accumulator { + AccumulatorStorage::Flat(acc) => { + acc.convert_to_state(&values.arguments, values.filter.as_ref()) + } + AccumulatorStorage::Blocked(acc) => { + acc.convert_to_state(&values.arguments, values.filter.as_ref()) + } + } } pub(super) fn null_arguments( @@ -907,6 +1005,12 @@ impl HashAggregateAccumulator { } } +/// Group keys and accumulators are always created with the same storage, so a +/// mismatch is a bug. +fn mismatched_storage() -> Result { + internal_err!("accumulator storage does not match the group key storage") +} + impl AggregateHashTableState { pub(super) fn building(&self) -> &AggregateHashTableBuffer { let Self::Building(state) = self else { @@ -969,7 +1073,11 @@ mod tests { schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], )?; - let mut output = MaterializedAggregateOutput::new(batch); + let memory_size = batch.get_array_memory_size(); + let mut output = MaterializedAggregateOutput::new(vec![MaterializedBatch { + batch, + memory_size, + }]); assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]); assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]); @@ -1048,8 +1156,10 @@ mod tests { .alias("SUM(value)") .build()?, ); - let accumulator = - create_group_accumulator(&aggregate_expr, Arc::clone(&submetrics))?; + let accumulator = AccumulatorStorage::Flat(create_group_accumulator( + &aggregate_expr, + Arc::clone(&submetrics), + )?); Ok(HashAggregateAccumulator::new( aggregate_expr, vec![argument], diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 86c052009faea..ba32393a1bc16 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -44,6 +44,7 @@ use super::common::{ AggregateAccumulator, AggregateBatchFn, AggregateHashTable, EvaluatedAggregateBatch, MaterializeAccumulatorFn, create_group_accumulator, }; +use super::storage::{AccumulatorStorage, Emit, GroupIndices}; #[derive(Clone)] pub(in crate::aggregates) struct OrderedAggregateTableMetrics { @@ -210,7 +211,7 @@ impl OrderedAggregateTable { Arc::clone(agg_expr), arguments, filter, - accumulator, + AccumulatorStorage::Flat(accumulator), Arc::clone(submetrics), )) }) @@ -328,11 +329,14 @@ impl OrderedAggregateTable { let output = self.group_by_metrics.time_emitting(|| { let mut output = self.buffer.group_values.emit(EmitTo::All)?; for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( - idx, - AccumulatorPhase::State, - || acc.state(EmitTo::All), - )?); + output.extend( + accumulator_metrics + .time(idx, AccumulatorPhase::State, || { + acc.state(Emit::Flat(EmitTo::All)) + })? + .into_iter() + .flatten(), + ); } Ok::<_, datafusion_common::DataFusionError>(output) })?; @@ -395,7 +399,7 @@ impl OrderedAggregateTable { aggregate_fn( acc, values, - &self.buffer.group_indices, + GroupIndices::Flat(&self.buffer.group_indices), total_num_groups, ) })?; @@ -425,11 +429,14 @@ impl OrderedAggregateTable { } for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( - idx, - accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), - )?); + output.extend( + accumulator_metrics + .time(idx, accumulator_phase, || { + materialize_accumulator_fn(acc, Emit::Flat(emit_to)) + })? + .into_iter() + .flatten(), + ); } Ok::<_, datafusion_common::DataFusionError>(output) })?; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs index 960b7498f1eb2..104902d6d913d 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -24,6 +24,7 @@ mod ordered_single_table; mod partial_reduce_table; mod partial_table; mod single_table; +mod storage; use std::sync::Arc; @@ -102,6 +103,7 @@ pub(super) use common::{ PartialSkipMarker, SingleMarker, create_group_accumulator, }; pub(super) use common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; +pub(super) use storage::MaterializedBatch; #[cfg(test)] mod tests { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 397b766f41697..abe09ebec7734 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -22,7 +22,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; -use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; +use crate::aggregates::group_values::AccumulatorPhase; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{AggregateExec, evaluate_group_by}; @@ -30,6 +30,7 @@ use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, }; +use super::storage::{GroupKeys, block_size}; /// Implementation specific to partial aggregation, where the table stores /// partial aggregate states and the input rows are raw rows. @@ -78,12 +79,23 @@ impl AggregateHashTable { ) -> Result> { let state = self.state.building(); let group_schema = state.group_by.group_schema(&self.input_schema)?; - let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + // Skipping aggregation only converts input rows to state, which needs + // no per-group storage, so the skip table always uses flat storage. let accumulators = state .accumulators .iter() .map(HashAggregateAccumulator::empty_like) .collect::>>()?; + let keys = GroupKeys::try_new( + group_schema, + &GroupOrdering::None, + false, + block_size(self.batch_size), + &accumulators + .iter() + .map(|acc| acc.storage()) + .collect::>(), + )?; Ok(AggregateHashTable { group_by_metrics: self.group_by_metrics.clone(), @@ -98,8 +110,7 @@ impl AggregateHashTable { batch_size: self.batch_size, state: AggregateHashTableState::Building(AggregateHashTableBuffer { group_by: Arc::clone(&state.group_by), - group_values, - batch_group_indices: Default::default(), + keys, accumulators, }), _mode: PhantomData, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs new file mode 100644 index 0000000000000..676318a758eb3 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs @@ -0,0 +1,573 @@ +// 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. + +//! Flat or blocked storage for the group keys and accumulator states of an +//! aggregate table. +//! +//! Every part of a table chooses its storage on its own, once, when the table +//! is created: the group keys use [`BlockedGroupValues`] when their type +//! supports it (see [`use_blocked_keys`]), and each aggregate uses a +//! [`BlockedGroupsAccumulator`] when it has one. Everything else keeps today's +//! flat [`GroupValues`] and [`GroupsAccumulator`], unchanged, so any mix of +//! flat and blocked parts works. +//! +//! The keys choose independently of the aggregates, so blocked keys can be +//! mixed with only flat aggregates (e.g. `SELECT k, sum(v) GROUP BY k` with a +//! primitive `k`). Only a table whose keys and aggregates are all flat runs +//! exactly the flat code. +//! +//! Mixing needs two things: +//! * Group indices in both layouts: the keys produce their own layout, and +//! [`GroupKeys`] converts them once per batch for the accumulators that use +//! the other layout (see [`GroupKeys::indices_for`]). +//! * Aligned output: blocked parts emit one array per block, and flat parts +//! emit one array that is sliced, without copying, at the same block +//! boundaries (see [`materialize_batches`]). + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, EmitTo, GroupsAccumulator, +}; + +use crate::aggregates::AggregateExec; +use crate::aggregates::group_values::blocked::{ + BlockedGroupValues, new_blocked_group_values, +}; +use crate::aggregates::group_values::{ + AccumulatorPhase, AggregateAccumulatorMetrics, GroupValues, new_group_values, +}; +use crate::aggregates::order::GroupOrdering; + +use super::common::{AggregateAccumulator, MaterializeAccumulatorFn}; + +/// Number of groups per block for every blocked part of a table. +/// +/// `batch_size`, so every block is emitted as one output batch. +pub(in crate::aggregates) fn block_size(batch_size: usize) -> usize { + batch_size +} + +/// Returns `true` if the group keys of `agg` can be stored in blocks: a single +/// grouping set whose key has a [`BlockedGroupValues`] implementation. +pub(in crate::aggregates) fn use_blocked_keys( + agg: &AggregateExec, + group_schema: &SchemaRef, + block_size: usize, +) -> bool { + agg.group_by().is_single() + && new_blocked_group_values(group_schema, block_size).is_some() +} + +/// Group indices of one input batch, in the layout of one accumulator. +#[derive(Clone, Copy)] +pub(in crate::aggregates) enum GroupIndices<'a> { + Flat(&'a [usize]), + Blocked(&'a [BlocksIndex]), +} + +/// Which groups to materialize, in the layout of one part of the table. +#[derive(Clone, Copy, Debug)] +pub(in crate::aggregates) enum Emit { + Flat(EmitTo), + Blocked(BlockedEmitTo), +} + +enum KeyStorage { + Flat(Box), + Blocked(Box), +} + +/// Group keys, and the group index of every row of the current input batch +/// in each layout the accumulators need. +pub(in crate::aggregates) struct GroupKeys { + storage: KeyStorage, + block_size: usize, + /// Flat indices: written by flat keys, or converted for flat accumulators + flat_indices: Vec, + /// Blocked indices: written by blocked keys, or converted for blocked + /// accumulators + blocked_indices: Vec, + /// Whether some accumulator uses the layout the keys do not produce + convert: bool, +} + +impl GroupKeys { + /// Creates blocked keys when `blocked` is set, flat keys otherwise. + /// `accumulators` are the storages of the table's accumulators. + pub(super) fn try_new( + group_schema: SchemaRef, + group_ordering: &GroupOrdering, + blocked: bool, + block_size: usize, + accumulators: &[&AccumulatorStorage], + ) -> Result { + let storage = if blocked { + match new_blocked_group_values(&group_schema, block_size) { + Some(group_values) => KeyStorage::Blocked(group_values), + None => { + return internal_err!( + "blocked group values are not supported for {group_schema:?}" + ); + } + } + } else { + KeyStorage::Flat(new_group_values(group_schema, group_ordering)?) + }; + let convert = accumulators + .iter() + .any(|acc| matches!(acc, AccumulatorStorage::Blocked(_)) != blocked); + Ok(Self { + storage, + block_size, + flat_indices: Vec::new(), + blocked_indices: Vec::new(), + convert, + }) + } + + /// Assigns a group to each row of `cols`, see [`GroupValues::intern`], + /// then converts the indices once for the accumulators using the other + /// layout. + pub(super) fn intern(&mut self, cols: &[ArrayRef]) -> Result<()> { + match &mut self.storage { + KeyStorage::Flat(group_values) => { + group_values.intern(cols, &mut self.flat_indices)?; + if self.convert { + BlocksIndex::from_flat_slice( + &self.flat_indices, + self.block_size, + &mut self.blocked_indices, + ); + } + } + KeyStorage::Blocked(group_values) => { + group_values.intern(cols, &mut self.blocked_indices)?; + if self.convert { + BlocksIndex::to_flat_slice( + &self.blocked_indices, + self.block_size, + &mut self.flat_indices, + ); + } + } + } + Ok(()) + } + + /// Group of each row of the last interned batch, in the layout of + /// `accumulator`. + pub(super) fn indices_for( + &self, + accumulator: &AccumulatorStorage, + ) -> GroupIndices<'_> { + match accumulator { + AccumulatorStorage::Flat(_) => GroupIndices::Flat(&self.flat_indices), + AccumulatorStorage::Blocked(_) => { + GroupIndices::Blocked(&self.blocked_indices) + } + } + } + + pub(super) fn len(&self) -> usize { + match &self.storage { + KeyStorage::Flat(group_values) => group_values.len(), + KeyStorage::Blocked(group_values) => group_values.len(), + } + } + + pub(super) fn is_empty(&self) -> bool { + match &self.storage { + KeyStorage::Flat(group_values) => group_values.is_empty(), + KeyStorage::Blocked(group_values) => group_values.is_empty(), + } + } + + /// Memory used by the keys and the scratch indices. + pub(super) fn size(&self) -> usize { + let group_values = match &self.storage { + KeyStorage::Flat(group_values) => group_values.size(), + KeyStorage::Blocked(group_values) => group_values.size(), + }; + group_values + + self.flat_indices.allocated_size() + + self.blocked_indices.allocated_size() + } + + /// Clears the keys and releases their memory, see + /// [`GroupValues::clear_shrink`], and frees the scratch indices. + pub(super) fn clear_shrink(&mut self, num_rows: usize) { + match &mut self.storage { + KeyStorage::Flat(group_values) => group_values.clear_shrink(num_rows), + KeyStorage::Blocked(group_values) => group_values.clear_shrink(num_rows), + } + self.free_indices(); + } + + /// Frees the scratch indices, which are not needed once input is done. + pub(super) fn free_indices(&mut self) { + self.flat_indices = Vec::new(); + self.blocked_indices = Vec::new(); + } +} + +/// Per-group state of one aggregate expression. +pub(in crate::aggregates) enum AccumulatorStorage { + Flat(Box), + Blocked(Box), +} + +impl AccumulatorStorage { + pub(super) fn size(&self) -> usize { + match self { + Self::Flat(acc) => acc.size(), + Self::Blocked(acc) => acc.size(), + } + } + + /// No group indices, in this accumulator's layout. + pub(super) fn no_groups(&self) -> GroupIndices<'static> { + match self { + Self::Flat(_) => GroupIndices::Flat(&[]), + Self::Blocked(_) => GroupIndices::Blocked(&[]), + } + } + + fn is_blocked(&self) -> bool { + matches!(self, Self::Blocked(_)) + } +} + +/// One materialized output batch, and the memory released once all of it has +/// been handed out. +pub(in crate::aggregates) struct MaterializedBatch { + pub(in crate::aggregates) batch: RecordBatch, + /// Buffers only this batch holds, plus, for the last batch, the flat + /// arrays all batches slice (they are released only with the last slice). + /// Summed over all batches, this is the memory of the output, each buffer + /// counted once. + pub(in crate::aggregates) memory_size: usize, +} + +/// Emits of a part of the table, as columns indexed `[block][column]`. +fn emit_part( + blocked: bool, + emit_to: EmitTo, + block_size: usize, + mut emit: impl FnMut(Emit) -> Result>>, +) -> Result>> { + if !blocked { + return emit(Emit::Flat(emit_to)); + } + let mut blocks = Vec::new(); + for blocked_emit in BlockedEmitTo::from_emit_to(emit_to, block_size) { + blocks.extend(emit(Emit::Blocked(blocked_emit))?); + } + Ok(blocks) +} + +/// Removes the groups selected by `emit_to` from the keys and every +/// accumulator, and returns them as one [`MaterializedBatch`] per block. +/// +/// Blocked parts emit one array per block. Flat parts emit a single array, +/// sliced without copying at the same block boundaries. With no blocked part +/// this returns one batch, exactly like the flat code. Blocks are never +/// concatenated. +pub(super) fn materialize_batches( + keys: &mut GroupKeys, + accumulators: &mut [AggregateAccumulator], + emit_to: EmitTo, + materialize_accumulator_fn: MaterializeAccumulatorFn, + accumulator_phase: AccumulatorPhase, + accumulator_metrics: &AggregateAccumulatorMetrics, + schema: &SchemaRef, +) -> Result> { + let block_size = keys.block_size; + // One entry per part (keys first): whether it is blocked, and its columns + // indexed `[block][column]` (a single "block" for flat parts) + let mut parts = Vec::with_capacity(accumulators.len() + 1); + let keys_blocked = matches!(keys.storage, KeyStorage::Blocked(_)); + parts.push(( + keys_blocked, + emit_part(keys_blocked, emit_to, block_size, |emit| { + match (&mut keys.storage, emit) { + (KeyStorage::Flat(group_values), Emit::Flat(emit_to)) => { + Ok(vec![group_values.emit(emit_to)?]) + } + (KeyStorage::Blocked(group_values), Emit::Blocked(emit_to)) => { + group_values.emit(emit_to) + } + (_, emit) => { + internal_err!("{emit:?} does not match the group key storage") + } + } + })?, + )); + for (idx, acc) in accumulators.iter_mut().enumerate() { + let blocked = acc.storage().is_blocked(); + let columns = accumulator_metrics.time(idx, accumulator_phase, || { + emit_part(blocked, emit_to, block_size, |emit| { + materialize_accumulator_fn(acc, emit) + }) + })?; + parts.push((blocked, columns)); + } + + // Row count of every output batch: the blocks of any blocked part (they + // all have the same boundaries), or one batch if every part is flat + let batch_lengths: Vec = match parts.iter().find(|(blocked, _)| *blocked) { + Some((_, blocks)) => blocks.iter().map(|columns| columns[0].len()).collect(), + None => vec![parts[0].1[0][0].len()], + }; + + let mut batches: Vec<(Vec, usize)> = batch_lengths + .iter() + .map(|_| (Vec::with_capacity(schema.fields().len()), 0)) + .collect(); + let mut shared_memory = 0; + for (blocked, blocks) in parts { + if blocked { + if blocks.len() != batch_lengths.len() { + return internal_err!( + "blocked parts emitted {} and {} blocks", + blocks.len(), + batch_lengths.len() + ); + } + for ((columns, memory), block) in batches.iter_mut().zip(blocks) { + *memory += block + .iter() + .map(|a| a.get_array_memory_size()) + .sum::(); + columns.extend(block); + } + } else { + let [flat] = <[_; 1]>::try_from(blocks).map_err(|blocks| { + datafusion_common::internal_datafusion_err!( + "flat part emitted {} blocks", + blocks.len() + ) + })?; + shared_memory += flat + .iter() + .map(|a| a.get_array_memory_size()) + .sum::(); + let mut offset = 0; + for ((columns, _), &len) in batches.iter_mut().zip(&batch_lengths) { + columns.extend(flat.iter().map(|array| array.slice(offset, len))); + offset += len; + } + } + } + if let Some((_, memory)) = batches.last_mut() { + *memory += shared_memory; + } + + batches + .into_iter() + .map(|(columns, memory_size)| { + Ok(MaterializedBatch { + batch: RecordBatch::try_new(Arc::clone(schema), columns)?, + memory_size, + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{AsArray, Int64Array}; + use arrow::datatypes::{DataType, Field, Int64Type, Schema}; + use datafusion_expr::AggregateUDF; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::min_max::min_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + + use super::*; + use crate::aggregates::aggregate_hash_table::{AggregateHashTable, SingleMarker}; + use crate::aggregates::{AggregateMode, PhysicalGroupBy, create_schema}; + use crate::test::TestMemoryExec; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])) + } + + /// `SELECT (k) FROM t GROUP BY k` in single mode. + fn aggregate_exec( + input: Vec, + aggregates: &[Arc], + ) -> Result { + let schema = schema(); + let aggr_expr = aggregates + .iter() + .map(|udaf| { + AggregateExprBuilder::new(Arc::clone(udaf), vec![col("k", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(udaf.name()) + .build() + .map(Arc::new) + }) + .collect::>>()?; + let exec = TestMemoryExec::try_new(&[input], Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".to_string())]), + aggr_expr.clone(), + vec![None; aggr_expr.len()], + exec, + schema, + ) + } + + fn single_table( + agg: &AggregateExec, + batch_size: usize, + ) -> Result> { + let state_schema = Arc::new(create_schema( + &schema(), + agg.group_by(), + agg.aggr_expr(), + AggregateMode::Partial, + )?); + AggregateHashTable::::new( + agg, + 0, + Arc::clone(&agg.schema), + state_schema, + batch_size, + ) + } + + #[test] + fn each_part_chooses_its_own_storage() -> Result<()> { + let agg = aggregate_exec(vec![], &[count_udaf(), min_udaf()])?; + let table = single_table(&agg, 8192)?; + let state = table.state.building(); + assert!(matches!(state.keys.storage, KeyStorage::Blocked(_))); + // count has a blocked accumulator, min does not (yet) + assert!(state.accumulators[0].storage().is_blocked()); + assert!(!state.accumulators[1].storage().is_blocked()); + // min needs flat indices while the keys produce blocked ones + assert!(state.keys.convert); + assert_eq!(block_size(8192), 8192); + Ok(()) + } + + /// Blocked `count` and flat `min` over blocked keys: flat columns are + /// sliced at the block boundaries, and their memory is counted once. + #[test] + fn mixed_storage_output_is_aligned_and_counted_once() -> Result<()> { + let batch_size = 100; + let num_groups = batch_size + 10; + let input = RecordBatch::try_new( + schema(), + vec![Arc::new(Int64Array::from_iter_values(0..num_groups as i64))], + )?; + let agg = aggregate_exec(vec![], &[count_udaf(), min_udaf()])?; + let mut table = single_table(&agg, batch_size)?; + table.aggregate_batch(&input)?; + + let batches = table.take_state_batches()?; + assert_eq!( + batches + .iter() + .map(|b| b.batch.num_rows()) + .collect::>(), + vec![batch_size, 10] + ); + for b in &batches { + let keys = b.batch.column(0).as_primitive::().values(); + let counts = b.batch.column(1).as_primitive::().values(); + let mins = b.batch.column(2).as_primitive::().values(); + assert!(counts.iter().all(|&c| c == 1)); + assert_eq!(keys, mins); + } + + // Keys and counts are owned by each block, the flat mins are shared + // by both batches and counted once, on the last batch + let owned = |b: &MaterializedBatch| { + b.batch.column(0).get_array_memory_size() + + b.batch.column(1).get_array_memory_size() + }; + let shared = batches[0].batch.column(2).get_array_memory_size(); + assert_eq!(batches[0].memory_size, owned(&batches[0])); + assert_eq!(batches[1].memory_size, owned(&batches[1]) + shared); + Ok(()) + } + + /// Output is materialized one block at a time, and each block's memory is + /// released once all its slices are handed out. + #[test] + fn output_releases_state_block_by_block() -> Result<()> { + let batch_size = 100; + let num_groups = 4 * batch_size + 10; + let input = RecordBatch::try_new( + schema(), + vec![Arc::new(Int64Array::from_iter_values(0..num_groups as i64))], + )?; + let agg = aggregate_exec(vec![], &[count_udaf()])?; + let mut table = single_table(&agg, batch_size)?; + table.aggregate_batch(&input)?; + table.start_output()?; + let memory_while_building = table.memory_size(); + + let mut keys = vec![]; + let mut memory_per_batch = vec![]; + while let Some(batch) = table.next_output_batch()? { + assert!(batch.num_rows() <= batch_size); + keys.extend( + batch + .column(0) + .as_primitive::() + .values() + .to_vec(), + ); + assert!( + batch + .column(1) + .as_primitive::() + .values() + .iter() + .all(|c| *c == 1) + ); + memory_per_batch.push(table.memory_size()); + } + + keys.sort_unstable(); + assert_eq!(keys, (0..num_groups as i64).collect::>()); + // Each output batch is one block, released once handed out, so only + // the last, small block (10 groups) remains after the full blocks. + assert_eq!(memory_per_batch.len(), 5); + assert!(memory_per_batch.is_sorted_by(|a, b| a > b)); + assert!( + memory_per_batch[3] < memory_while_building / 4, + "memory after the full blocks: {}, while building: {memory_while_building}", + memory_per_batch[3] + ); + assert_eq!(*memory_per_batch.last().unwrap(), 0); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs new file mode 100644 index 0000000000000..75ad490d418ab --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs @@ -0,0 +1,155 @@ +// 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. + +//! [`BlockedGroupValues`]: group keys stored in fixed size blocks + +use arrow::array::ArrayRef; +use arrow::array::downcast_primitive; +use arrow::array::types::{ + Date32Type, Date64Type, Decimal128Type, Time32MillisecondType, Time32SecondType, + Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, +}; +use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion_common::Result; +use datafusion_expr::{BlockedEmitTo, BlocksIndex}; + +mod primitive; +use primitive::BlockedGroupValuesPrimitive; + +/// Like [`GroupValues`], but the group keys are stored in blocks of +/// `block_size` groups, matching the blocks of the +/// [`BlockedGroupsAccumulator`]s of the same hash table, so each block of +/// groups can be emitted (and its memory freed) on its own. +/// +/// [`GroupValues`]: super::GroupValues +/// [`BlockedGroupsAccumulator`]: datafusion_expr::BlockedGroupsAccumulator +pub(crate) trait BlockedGroupValues: Send { + /// Calculates the group index of each row of `cols`, adding new groups + /// as needed. Same contract as [`GroupValues::intern`]. + /// + /// [`GroupValues::intern`]: super::GroupValues::intern + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()>; + + /// Returns the number of bytes of memory used. + fn size(&self) -> usize; + + /// Returns true if there are no groups. + fn is_empty(&self) -> bool; + + /// The number of groups. + fn len(&self) -> usize; + + /// Emits the group values, indexed `[block][group column]`; same block + /// rules as `BlockedGroupsAccumulator::evaluate`. + fn emit(&mut self, emit_to: BlockedEmitTo) -> Result>>; + + /// Clear the contents and shrink the capacity to the size of the batch + /// (free up memory usage). + fn clear_shrink(&mut self, num_rows: usize); +} + +/// Returns a [`BlockedGroupValues`] for `schema`, or `None` if blocked group +/// values are not supported for it. +/// +/// Only a single group column of a type that `GroupValuesPrimitive` handles +/// in [`new_group_values`] is supported. +/// +/// [`new_group_values`]: super::new_group_values +pub(crate) fn new_blocked_group_values( + schema: &SchemaRef, + block_size: usize, +) -> Option> { + if schema.fields.len() != 1 { + return None; + } + let d = schema.fields[0].data_type(); + + macro_rules! downcast_helper { + ($t:ty, $d:ident) => { + return Some(Box::new(BlockedGroupValuesPrimitive::<$t>::new( + $d.clone(), + block_size, + ))) + }; + } + + downcast_primitive! { + d => (downcast_helper, d), + _ => {} + } + + match d { + DataType::Date32 => downcast_helper!(Date32Type, d), + DataType::Date64 => downcast_helper!(Date64Type, d), + DataType::Time32(t) => match t { + TimeUnit::Second => downcast_helper!(Time32SecondType, d), + TimeUnit::Millisecond => downcast_helper!(Time32MillisecondType, d), + _ => None, + }, + DataType::Time64(t) => match t { + TimeUnit::Microsecond => downcast_helper!(Time64MicrosecondType, d), + TimeUnit::Nanosecond => downcast_helper!(Time64NanosecondType, d), + _ => None, + }, + DataType::Timestamp(t, _tz) => match t { + TimeUnit::Second => downcast_helper!(TimestampSecondType, d), + TimeUnit::Millisecond => downcast_helper!(TimestampMillisecondType, d), + TimeUnit::Microsecond => downcast_helper!(TimestampMicrosecondType, d), + TimeUnit::Nanosecond => downcast_helper!(TimestampNanosecondType, d), + }, + DataType::Decimal128(_, _) => downcast_helper!(Decimal128Type, d), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::datatypes::{Field, Schema}; + + use super::*; + + fn supported(data_type: DataType) -> bool { + let schema = Arc::new(Schema::new(vec![Field::new("a", data_type, true)])); + new_blocked_group_values(&schema, 4).is_some() + } + + #[test] + fn supports_same_types_as_group_values_primitive() { + assert!(supported(DataType::Int32)); + assert!(supported(DataType::UInt8)); + assert!(supported(DataType::Float64)); + assert!(supported(DataType::Date32)); + assert!(supported(DataType::Time32(TimeUnit::Second))); + assert!(supported(DataType::Time64(TimeUnit::Nanosecond))); + assert!(supported(DataType::Timestamp( + TimeUnit::Microsecond, + Some("UTC".into()) + ))); + assert!(supported(DataType::Decimal128(10, 2))); + + assert!(!supported(DataType::Utf8)); + assert!(!supported(DataType::Boolean)); + let two_columns = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + assert!(new_blocked_group_values(&two_columns, 4).is_none()); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked/primitive.rs new file mode 100644 index 0000000000000..de13ff4b67770 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked/primitive.rs @@ -0,0 +1,412 @@ +// 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. + +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, + cast::AsArray, +}; +use arrow::datatypes::DataType; +use datafusion_common::Result; +use datafusion_common::hash_utils::RandomState; +use datafusion_expr::{BlockedEmitTo, BlocksIndex}; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::BlockedVec; +use hashbrown::hash_table::HashTable; + +use super::BlockedGroupValues; +use crate::aggregates::group_values::HashValue; + +/// A [`BlockedGroupValues`] storing a single column of primitive values. +/// +/// Like `GroupValuesPrimitive`, but the values are stored in a [`BlockedVec`] +/// and the hash table entries hold `(group_index, value)` instead of +/// `(group_index, hash)`, so probing never reads the (blocked) values. +pub(crate) struct BlockedGroupValuesPrimitive { + /// The data type of the output array + data_type: DataType, + /// Stores `(group_index, value)` based on the hash of the value + /// + /// Storing the value instead of its hash means comparing keys never goes + /// through `values`, and rehashing is cheap for primitive values. + map: HashTable<(BlocksIndex, T::Native)>, + /// The group index of the null value if any + null_group: Option, + /// The values for each group index + values: BlockedVec, + /// The random state used to generate hashes + random_state: RandomState, +} + +impl BlockedGroupValuesPrimitive { + pub(crate) fn new(data_type: DataType, block_size: usize) -> Self { + assert!(PrimitiveArray::::is_compatible(&data_type)); + Self { + data_type, + map: HashTable::with_capacity(128), + values: BlockedVec::new(block_size), + null_group: None, + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + } + } + + /// Builds the output array of one block whose first group is `start` + fn build_block(&self, values: Vec, start: usize) -> ArrayRef { + let len = values.len(); + let block_size = self.values.block_size(); + let nulls = self + .null_group + .map(|null_group| null_group.flat(block_size)) + .filter(|null_group| (start..start + len).contains(null_group)) + .map(|null_group| { + let null_idx = null_group - start; + let mut buffer = NullBufferBuilder::new(len); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(len - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }); + Arc::new( + PrimitiveArray::::new(values.into(), nulls) + .with_data_type(self.data_type.clone()), + ) + } + + /// Updates the null group after the first `n` groups were emitted + fn shift_null_group(&mut self, n: usize) { + let block_size = self.values.block_size(); + self.null_group = self + .null_group + .and_then(|null_group| null_group.flat(block_size).checked_sub(n)) + .map(|flat| BlocksIndex::from_flat(flat, block_size)); + } +} + +impl BlockedGroupValues for BlockedGroupValuesPrimitive +where + T::Native: HashValue, +{ + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + assert_eq!(cols.len(), 1); + groups.clear(); + + for v in cols[0].as_primitive::() { + let group_id = match v { + None => *self + .null_group + .get_or_insert_with(|| self.values.push(Default::default())), + Some(key) => { + // Fold equivalence-class duplicates (e.g. `-0.0` → `+0.0`) + // 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( + hash, + |&(_, 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.push(key); + v.insert((g, key)); + g + } + } + } + }; + groups.push(group_id) + } + Ok(()) + } + + fn size(&self) -> usize { + self.map.capacity() * size_of::<(BlocksIndex, T::Native)>() + + self.values.allocated_size() + } + + fn is_empty(&self) -> bool { + self.values.is_empty() + } + + fn len(&self) -> usize { + self.values.len() + } + + fn emit(&mut self, emit_to: BlockedEmitTo) -> Result>> { + let blocks = match emit_to { + BlockedEmitTo::All => { + self.map.clear(); + let mut start = 0; + let blocks = self + .values + .take_all() + .into_iter() + .map(|block| { + let len = block.len(); + let array = self.build_block(block, start); + start += len; + vec![array] + }) + .collect(); + self.null_group = None; + blocks + } + BlockedEmitTo::NextBlock => { + self.map.retain(|entry| match entry.0.block_index() { + // Group was in the emitted block, so remove from table + 0 => false, + // Group moves down by one block + block_index => { + entry.0 = + BlocksIndex::new(block_index - 1, entry.0.index_in_block()); + true + } + }); + match self.values.take_next_block() { + Some(block) => { + let len = block.len(); + let array = self.build_block(block, 0); + self.shift_null_group(len); + vec![vec![array]] + } + None => vec![], + } + } + BlockedEmitTo::First(n) => { + let block_size = self.values.block_size(); + self.map.retain(|entry| { + // Decrement group index by n + match entry.0.flat(block_size).checked_sub(n) { + // Group index was >= n, shift value down + Some(sub) => { + entry.0 = BlocksIndex::from_flat(sub, block_size); + true + } + // Group index was < n, so remove from table + None => false, + } + }); + let block = self.values.take_first(n); + let array = self.build_block(block, 0); + self.shift_null_group(n); + vec![vec![array]] + } + }; + Ok(blocks) + } + + fn clear_shrink(&mut self, num_rows: usize) { + self.values = BlockedVec::new(self.values.block_size()); + self.null_group = None; + self.map.clear(); + self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Float64Array, Int32Array}; + use arrow::datatypes::{Float64Type, Int32Type}; + + fn intern( + gv: &mut BlockedGroupValuesPrimitive, + v: Vec>, + ) -> Vec { + let mut groups = vec![]; + gv.intern(&[Arc::new(Int32Array::from(v)) as ArrayRef], &mut groups) + .unwrap(); + groups.into_iter().map(|g| g.flat(4)).collect() + } + + #[test] + fn intern_and_emit_blocks_in_order() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + assert_eq!( + intern(&mut gv, (0..6).map(Some).collect()), + vec![0, 1, 2, 3, 4, 5] + ); + assert_eq!(intern(&mut gv, vec![Some(5), Some(0)]), vec![5, 0]); + let first = gv.emit(BlockedEmitTo::NextBlock)?; + assert_eq!( + first[0][0].as_primitive::().values(), + &[0, 1, 2, 3] + ); + let rest = gv.emit(BlockedEmitTo::All)?; + assert_eq!(rest.len(), 1); + assert_eq!(rest[0][0].as_primitive::().values(), &[4, 5]); + assert!(gv.is_empty()); + Ok(()) + } + + #[test] + fn null_group_in_second_block() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + let mut v: Vec> = (0..5).map(Some).collect(); + v.push(None); + assert_eq!(intern(&mut gv, v), vec![0, 1, 2, 3, 4, 5]); + let first = gv.emit(BlockedEmitTo::NextBlock)?; + assert_eq!(first[0][0].null_count(), 0); + let second = gv.emit(BlockedEmitTo::NextBlock)?; + let arr = second[0][0].as_primitive::(); + assert_eq!(arr.len(), 2); + assert!(arr.is_null(1)); + Ok(()) + } + + #[test] + fn first_n_renumbers_remaining_groups() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + intern(&mut gv, (0..6).map(Some).collect()); + let out = gv.emit(BlockedEmitTo::First(2))?; + assert_eq!(out[0][0].as_primitive::().values(), &[0, 1]); + assert_eq!( + intern(&mut gv, vec![Some(5), Some(2), Some(9)]), + vec![3, 0, 4] + ); + Ok(()) + } + + #[test] + fn first_n_shifts_null_group() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + intern(&mut gv, vec![Some(0), None, Some(2), Some(3), Some(4)]); + // Null group emitted + let out = gv.emit(BlockedEmitTo::First(2))?; + assert!(out[0][0].is_null(1)); + // A new null group is created after the emit + assert_eq!(intern(&mut gv, vec![None, Some(4)]), vec![3, 2]); + // Null group kept and shifted down + let out = gv.emit(BlockedEmitTo::First(1))?; + assert_eq!(out[0][0].null_count(), 0); + assert_eq!(intern(&mut gv, vec![None]), vec![2]); + Ok(()) + } + + #[test] + fn all_with_null_group_in_last_block() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + let mut v: Vec> = (0..9).map(Some).collect(); + v.push(None); + intern(&mut gv, v); + let blocks = gv.emit(BlockedEmitTo::All)?; + assert_eq!(blocks.len(), 3); + assert_eq!( + blocks.iter().map(|b| b[0].len()).collect::>(), + vec![4, 4, 2] + ); + assert_eq!(blocks[0][0].null_count(), 0); + assert_eq!(blocks[1][0].null_count(), 0); + let last = blocks[2][0].as_primitive::(); + assert_eq!(last.value(0), 8); + assert!(last.is_null(1)); + assert!(gv.is_empty()); + // State fully reset: a new null group starts again at 0 + assert_eq!(intern(&mut gv, vec![None, Some(8)]), vec![0, 1]); + Ok(()) + } + + #[test] + fn float_negative_and_positive_zero_are_one_group() -> Result<()> { + let mut gv = + BlockedGroupValuesPrimitive::::new(DataType::Float64, 4); + let mut groups = vec![]; + let input = Arc::new(Float64Array::from(vec![-0.0, 0.0, 1.0, -0.0])); + gv.intern(&[input as ArrayRef], &mut groups)?; + let groups: Vec = groups.into_iter().map(|g| g.flat(4)).collect(); + assert_eq!(groups, vec![0, 0, 1, 0]); + let out = gv.emit(BlockedEmitTo::All)?; + let values = out[0][0].as_primitive::(); + // Stored value is the canonical +0.0 + assert_eq!(values.value(0).to_bits(), 0.0f64.to_bits()); + assert_eq!(values.value(1), 1.0); + Ok(()) + } + + #[test] + fn keeps_output_data_type() -> Result<()> { + let data_type = DataType::Timestamp( + arrow::datatypes::TimeUnit::Nanosecond, + Some("UTC".into()), + ); + let mut gv = BlockedGroupValuesPrimitive::< + arrow::datatypes::TimestampNanosecondType, + >::new(data_type.clone(), 4); + let input = Arc::new( + arrow::array::TimestampNanosecondArray::from(vec![1, 2, 3, 4, 5]) + .with_data_type(data_type.clone()), + ); + gv.intern(&[input as ArrayRef], &mut vec![])?; + let blocks = gv.emit(BlockedEmitTo::All)?; + assert_eq!(blocks.len(), 2); + assert!(blocks.iter().all(|b| b[0].data_type() == &data_type)); + Ok(()) + } + + #[test] + fn clear_shrink_resets_state() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + intern(&mut gv, vec![Some(1), None, Some(2), Some(3), Some(4)]); + gv.clear_shrink(2); + assert!(gv.is_empty()); + assert_eq!(gv.len(), 0); + assert_eq!(intern(&mut gv, vec![Some(3), None, Some(3)]), vec![0, 1, 0]); + Ok(()) + } + + #[test] + fn next_block_keeps_interning_working() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + intern(&mut gv, (0..6).map(Some).collect()); + let first = gv.emit(BlockedEmitTo::NextBlock)?; + assert_eq!( + first[0][0].as_primitive::().values(), + &[0, 1, 2, 3] + ); + // 4 and 5 moved down one block; 0 was emitted so it is a new group + assert_eq!( + intern(&mut gv, vec![Some(5), Some(4), Some(0)]), + vec![1, 0, 2] + ); + let rest = gv.emit(BlockedEmitTo::All)?; + assert_eq!(rest[0][0].as_primitive::().values(), &[4, 5, 0]); + Ok(()) + } + + #[test] + fn size_decreases_while_draining() -> Result<()> { + let mut gv = BlockedGroupValuesPrimitive::::new(DataType::Int32, 4); + intern(&mut gv, (0..10).map(Some).collect()); + let mut previous = gv.size(); + assert!(previous >= 10 * size_of::()); + while !gv.is_empty() { + previous = gv.size(); + assert_eq!(gv.emit(BlockedEmitTo::NextBlock)?.len(), 1); + assert!(gv.size() < previous); + } + assert!(gv.emit(BlockedEmitTo::NextBlock)?.is_empty()); + // Only the (empty) outer `Vec` of blocks of `BlockedVec` is left + assert!(gv.values.allocated_size() <= 4 * size_of::>()); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 32eabfe8dc26c..4c8f6d3fc3539 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -23,6 +23,7 @@ use datafusion_common::{Result, not_impl_err}; use datafusion_expr::{EmitTo, GroupSelection}; +pub(crate) mod blocked; pub mod multi_group_by; mod row; diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2f1bb9b638a4c..0db42be8a99ef 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -34,8 +34,8 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{ - AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker, - PartialSkipMarker, + AggregateHashTable, FinalMarker, MaterializedBatch, OrderedAggregateTableMetrics, + PartialMarker, PartialSkipMarker, }; use super::skip_partial::SkipAggregationProbe; use super::spill::AggregateSpill; @@ -316,20 +316,44 @@ impl PartialHashAggregateStream { break; } HandleInputResult::OOM => { - let materialized_group_states = hash_table.take_state_batch()?.ok_or_else(|| { - internal_datafusion_err!( + let materialized_group_states = + hash_table.take_state_batches()?; + if materialized_group_states.is_empty() { + return Err(internal_datafusion_err!( "Partial hash aggregate ran out of memory with no aggregated groups" - ) - })?; + )); + } self.early_emit_count.add(1); timer.done(); - self.emit_on_memory_pressure( - materialized_group_states, - &mut emitter, - hash_table.memory_size(), - ) - .await?; + + // Blocked storage returns one batch per block, moved out of + // the table without copying. Emit them in turn, keeping the + // batches not emitted yet in the reservation. + let pending_memory: usize = materialized_group_states + .iter() + // Don't include the first batch since we will emit it right away and release its memory + // if it needs slicing then a we will try to hold on that reservation while slicing + .skip(1) + .map(|b| b.memory_size) + .sum(); + + // Make sure we can hold on the hash tables and all the batches that need to be emitted (except the first one) + // if we can't hold it than we can't do anything about it. + self.reservation.try_resize( + hash_table.memory_size() + // The batches that need to be held while emitting each one + + pending_memory, + )?; + + for (i, batch) in + materialized_group_states.into_iter().enumerate() + { + if i != 0 { + self.reservation.try_shrink(batch.memory_size)?; + } + self.emit_on_memory_pressure(batch, &mut emitter).await?; + } } } } @@ -434,63 +458,62 @@ impl PartialHashAggregateStream { &mut self, // After each incremental emitting step, the `remaining_groups` will be updated // with batch slicing. - mut remaining_groups: RecordBatch, + remaining_groups: MaterializedBatch, emitter: &mut TryEmitter, - hash_table_mem_size: usize, ) -> Result<()> { - let remaining_groups_memory = remaining_groups.get_array_memory_size(); - - // Emitting clears the aggregate table and releases its - // accumulated memory. Update the reservation accordingly. - // We account here for the remaining groups memory to see if we can return batch size states - // if there is not enough memory, fallback to emit large batch - match self - .reservation - .try_resize(hash_table_mem_size + remaining_groups_memory) - { - Ok(_) => { - // Continue with slicing - } - Err(DataFusionError::ResourcesExhausted(_)) => { - // Fail to reserve memory for the hash table + state batch while slicing so emit a huge batch - - // Try resize without holding the state batch, if it fails there is nothing we can do - self.reservation.try_resize(hash_table_mem_size)?; + let MaterializedBatch { + batch: mut remaining_groups, + memory_size: remaining_groups_memory, + } = remaining_groups; + + // If we should slice the batch + if remaining_groups.num_rows() > self.batch_size { + // Emitting clears the aggregate table and releases its + // accumulated memory. Update the reservation accordingly. + // We account here for the remaining groups memory to see if we can return batch size states + // if there is not enough memory, fallback to emit large batch + match self.reservation.try_grow(remaining_groups_memory) { + Ok(_) => { + // Continue with slicing + } + Err(DataFusionError::ResourcesExhausted(_)) => { + // Fail to reserve memory for the hash table + state batch while slicing so emit a huge batch - self.reduction_factor.add_part(remaining_groups.num_rows()); - emitter - .emit(remaining_groups.record_output(&self.baseline_metrics)) - .await; + self.reduction_factor.add_part(remaining_groups.num_rows()); + emitter + .emit(remaining_groups.record_output(&self.baseline_metrics)) + .await; - return Ok(()); + return Ok(()); + } + Err(e) => return Err(e), } - Err(e) => return Err(e), - } - while remaining_groups.num_rows() > self.batch_size { - // More batch to output, continue in the current state. - let output = remaining_groups.slice(0, self.batch_size); + while remaining_groups.num_rows() > self.batch_size { + // More batch to output, continue in the current state. + let output = remaining_groups.slice(0, self.batch_size); - remaining_groups = remaining_groups.slice( - self.batch_size, - remaining_groups.num_rows() - self.batch_size, - ); + remaining_groups = remaining_groups.slice( + self.batch_size, + remaining_groups.num_rows() - self.batch_size, + ); - self.reduction_factor.add_part(output.num_rows()); - debug_assert!(output.num_rows() > 0); + self.reduction_factor.add_part(output.num_rows()); + debug_assert!(output.num_rows() > 0); - emitter - .emit(output.record_output(&self.baseline_metrics)) - .await; + emitter + .emit(output.record_output(&self.baseline_metrics)) + .await; + } + + // We are no longer holding on the batch while slicing, so release the memory. + // The memory will now equal to the hash table size + self.reservation.try_shrink(remaining_groups_memory)?; } self.reduction_factor.add_part(remaining_groups.num_rows()); debug_assert!(remaining_groups.num_rows() > 0); - // We are no longer holding on the batch while slicing, so release the memory. - // The memory will now equal to the hash table size - self.reservation.try_shrink(remaining_groups_memory)?; - emitter .emit(remaining_groups.record_output(&self.baseline_metrics)) .await; @@ -757,8 +780,8 @@ impl FinalHashAggregateStream { // Go to the next state to perform spilling the aggregated // groups so far. let result = hash_table - .take_state_batch() - .and_then(|batch| spill_context.sort_and_spill(batch)); + .take_state_batches() + .and_then(|batch| spill_context.sort_and_spill_batches(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -796,8 +819,8 @@ impl FinalHashAggregateStream { // Input was exhausted after spilling. Spill the last in-memory run hash_table - .take_state_batch() - .and_then(|batch| spill_context.sort_and_spill(batch))?; + .take_state_batches() + .and_then(|batch| spill_context.sort_and_spill_batches(batch))?; // Construct the ordered input used to merge all spill files. let mut output_stream = diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d5057d53cf19d..eebc9e12cb8e6 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -17,6 +17,7 @@ //! Partial-reduce hash aggregation stream implementation. +use std::collections::VecDeque; use std::ops::ControlFlow; use std::sync::Arc; use std::task::{Context, Poll}; @@ -29,7 +30,9 @@ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; -use super::aggregate_hash_table::{AggregateHashTable, PartialReduceMarker}; +use super::aggregate_hash_table::{ + AggregateHashTable, MaterializedBatch, PartialReduceMarker, +}; use crate::metrics::{BaselineMetrics, Count, MetricBuilder, RecordOutput, SpillMetrics}; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -103,13 +106,14 @@ enum PartialReduceHashAggregateState { ReadingInput { hash_table: AggregateHashTable, }, - /// A fully materialized partial-state batch being emitted incrementally - /// because the table ran out of memory while reading input. + /// Fully materialized partial-state batches (one per block for blocked + /// storage) being emitted incrementally because the table ran out of + /// memory while reading input. EmittingOnMemoryPressure { hash_table: AggregateHashTable, - // After each incremental emitting step, `remaining_groups` is updated - // with batch slicing. - remaining_groups: RecordBatch, + // After each incremental emitting step, the front of `remaining_groups` + // is updated with batch slicing, and removed once fully emitted. + remaining_groups: VecDeque, }, ProducingOutput { hash_table: AggregateHashTable, @@ -296,8 +300,8 @@ impl PartialReduceHashAggregateStream { /// /// The reservation is left at its pre-emission size while the states are being /// emitted, because the cleared states are still held in memory as - /// `remaining_groups`. The reservation will be reset after exiting the - /// `EmittingOnMemoryPressure` state. + /// `remaining_groups`. It shrinks as each batch is fully emitted, and is + /// reset after exiting the `EmittingOnMemoryPressure` state. /// /// # Implementation Note /// All accumulated states are materialized at once, and then sliced into @@ -321,21 +325,23 @@ impl PartialReduceHashAggregateStream { Err(e) => return Self::break_with_err(e), }; - let state_batch_result = original_state.hash_table_mut().take_state_batch(); + let state_batch_result = original_state.hash_table_mut().take_state_batches(); match state_batch_result { - Ok(Some(remaining_groups)) => { + // No accumulated group to emit, so early emission cannot release any + // memory: report the original error. + Ok(remaining_groups) if remaining_groups.is_empty() => { + Self::break_with_err(oom) + } + Ok(remaining_groups) => { self.early_emit_count.add(1); ControlFlow::Continue( PartialReduceHashAggregateState::EmittingOnMemoryPressure { hash_table: original_state.into_hash_table(), - remaining_groups, + remaining_groups: remaining_groups.into(), }, ) } - // No accumulated group to emit, so early emission cannot release any - // memory: report the original error. - Ok(None) => Self::break_with_err(oom), Err(e) => Self::break_with_err(e), } } @@ -353,30 +359,45 @@ impl PartialReduceHashAggregateStream { ) -> PartialReduceHashAggregateStateTransition { let PartialReduceHashAggregateState::EmittingOnMemoryPressure { hash_table, - remaining_groups: batch, + mut remaining_groups, } = original_state else { unreachable!("expected the EmittingOnMemoryPressure state") }; + let Some(MaterializedBatch { batch, memory_size }) = remaining_groups.pop_front() + else { + unreachable!("EmittingOnMemoryPressure always holds a batch") + }; - let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + let output_batch = if batch.num_rows() <= self.batch_size { + // Last slice of this batch, release its memory. Never grow here: + // the pre-emission reservation may be smaller than what is held. + let held = hash_table.memory_size() + + remaining_groups + .iter() + .map(|b| b.memory_size) + .sum::(); + if held < self.reservation.size() { + self.reservation.shrink(self.reservation.size() - held); + } + batch + } else { + // More rows in this batch, keep the rest at the front. + remaining_groups.push_front(MaterializedBatch { + batch: batch.slice(self.batch_size, batch.num_rows() - self.batch_size), + memory_size, + }); + batch.slice(0, self.batch_size) + }; + let next_state = if remaining_groups.is_empty() { // Go back to `ReadingInput` - ( - batch, - PartialReduceHashAggregateState::ReadingInput { hash_table }, - ) + PartialReduceHashAggregateState::ReadingInput { hash_table } } else { // More batches to output, continue in the current state. - let remaining = - batch.slice(self.batch_size, batch.num_rows() - self.batch_size); - let output = batch.slice(0, self.batch_size); - ( - output, - PartialReduceHashAggregateState::EmittingOnMemoryPressure { - hash_table, - remaining_groups: remaining, - }, - ) + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups, + } }; debug_assert!(output_batch.num_rows() > 0); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index f27017784bfd3..9723f71ff33cc 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -364,7 +364,8 @@ impl Aggregating { 0, "Single hash aggregate ran out of memory with no aggregated groups" ); - spill_context.sort_and_spill(self.hash_table.take_state_batch()?)?; + spill_context + .sort_and_spill_batches(self.hash_table.take_state_batches()?)?; reservation .try_resize(self.hash_table.memory_size()) .map_err(|e| { @@ -393,8 +394,8 @@ impl Aggregating { // Input was exhausted after spilling. Spill the last in-memory run. let mut hash_table = self.hash_table; hash_table - .take_state_batch() - .and_then(|batch| spill_context.sort_and_spill(batch))?; + .take_state_batches() + .and_then(|batch| spill_context.sort_and_spill_batches(batch))?; // Construct the replay stream: an ordered final aggregate stream // over the sort-preserving merge of all spill runs. diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 887be41855996..e06ce1024036d 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -19,6 +19,8 @@ use std::sync::Arc; +use arrow::array::Array; +use arrow::compute::{SortColumn, concat, interleave_record_batch, lexsort_to_indices}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; @@ -28,7 +30,7 @@ use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; -use super::aggregate_hash_table::OrderedAggregateTableMetrics; +use super::aggregate_hash_table::{MaterializedBatch, OrderedAggregateTableMetrics}; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, AggregateMode}; use crate::metrics::{BaselineMetrics, SpillMetrics}; @@ -228,6 +230,73 @@ impl AggregateSpill { self.spill_expr.clone(), self.batch_size, ); + self.spill_sorted(sorted_iter, max_batch_rows) + } + + /// [`Self::sort_and_spill`] for the state batches of a hash aggregate + /// table: one batch for flat storage, one per block for blocked storage. + /// All batches are written as one sorted spill file. + pub(super) fn sort_and_spill_batches( + &mut self, + state_batches: Vec, + ) -> Result<()> { + let mut batches: Vec = + state_batches.into_iter().map(|b| b.batch).collect(); + if batches.len() <= 1 { + return self.sort_and_spill(batches.pop()); + } + + // Sort the key columns of all blocks together, then gather each + // output batch from the blocks, so all blocks form one sorted run + // while only the sort keys are concatenated + // + // TODO: the concatenated sort keys and the sort indices are temporary + // memory that is not reserved, allocated while we are already out of + // memory for another batch. With many group columns (e.g. 130 columns + // and 10M groups) this is large, and concatenating byte arrays can + // overflow their i32 offsets. `sort_and_spill` on main has the same + // problem for a single batch. Reserve memory for this sort, sort across + // multiple arrays without concat, or consider using `ExternalSorter`, + // which already handles sorting in memory and spilling. + let sort_columns = self + .spill_expr + .iter() + .map(|expr| { + let arrays = batches + .iter() + .map(|b| expr.expr.evaluate(b)?.into_array(b.num_rows())) + .collect::>>()?; + let arrays: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect(); + Ok(SortColumn { + values: concat(&arrays)?, + options: Some(expr.options), + }) + }) + .collect::>>()?; + let sorted = lexsort_to_indices(&sort_columns, None)?; + drop(sort_columns); + + // Every block but the last holds the same number of rows + let block_rows = batches[0].num_rows(); + let batch_refs: Vec<&RecordBatch> = batches.iter().collect(); + let max_batch_rows = sorted.len().min(self.batch_size); + let sorted_iter = sorted.values().chunks(self.batch_size).map(|chunk| { + let indices: Vec<(usize, usize)> = chunk + .iter() + .map(|&i| (i as usize / block_rows, i as usize % block_rows)) + .collect(); + Ok(interleave_record_batch(&batch_refs, &indices)?) + }); + self.spill_sorted(sorted_iter, max_batch_rows) + } + + /// Writes the sorted batches of `sorted_iter` as one spill file, whose + /// largest batch has `max_batch_rows` rows. + fn spill_sorted( + &mut self, + sorted_iter: impl Iterator>, + max_batch_rows: usize, + ) -> Result<()> { let spill_file = self .spill_manager .spill_record_batch_iter_and_return_max_batch_memory(