From 1a742259f131521e34ca30a77eee756021b8bc18 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:53:33 +0300 Subject: [PATCH 01/10] feat(agg): blocked hash aggregation PoC Blocked storage for group keys and accumulators (count, count distinct, sum, avg, min/max bytes), mixed with flat storage per part. --- .../src/blocked_groups_accumulator.rs | 392 +++++++++ datafusion/expr-common/src/lib.rs | 1 + datafusion/expr/src/lib.rs | 3 + datafusion/expr/src/udaf.rs | 48 ++ .../src/aggregate/count_distinct.rs | 2 + .../count_distinct/blocked_groups.rs | 347 ++++++++ .../src/aggregate/count_distinct/groups.rs | 45 +- .../src/aggregate/groups_accumulator.rs | 3 + .../groups_accumulator/accumulate.rs | 63 +- .../groups_accumulator/blocked_null_state.rs | 385 +++++++++ .../groups_accumulator/blocked_prim_op.rs | 245 ++++++ .../groups_accumulator/blocked_vec.rs | 801 ++++++++++++++++++ .../aggregate/groups_accumulator/prim_op.rs | 104 ++- datafusion/functions-aggregate/src/average.rs | 487 +++++++++-- datafusion/functions-aggregate/src/count.rs | 402 ++++++++- datafusion/functions-aggregate/src/min_max.rs | 44 +- .../src/min_max/min_max_bytes.rs | 753 +++++++++++----- datafusion/functions-aggregate/src/sum.rs | 25 +- datafusion/physical-expr/src/aggregate.rs | 39 + .../aggregates/aggregate_hash_table/common.rs | 386 ++++++--- .../aggregate_hash_table/common_ordered.rs | 157 ++-- .../aggregates/aggregate_hash_table/mod.rs | 2 + .../ordered_final_table.rs | 13 +- .../ordered_partial_table.rs | 10 +- .../aggregate_hash_table/partial_table.rs | 19 +- .../aggregate_hash_table/storage.rs | 595 +++++++++++++ .../aggregates/group_values/blocked/mod.rs | 155 ++++ .../group_values/blocked/primitive.rs | 412 +++++++++ .../src/aggregates/group_values/mod.rs | 1 + .../src/aggregates/hash_stream.rs | 47 +- .../physical-plan/src/aggregates/order/mod.rs | 28 +- .../src/aggregates/order/partial.rs | 32 +- .../src/aggregates/ordered_final_stream.rs | 68 +- .../src/aggregates/ordered_partial_stream.rs | 87 +- .../src/aggregates/ordered_single_stream.rs | 4 +- .../src/aggregates/partial_reduce_stream.rs | 64 +- .../src/aggregates/single_stream.rs | 4 +- .../physical-plan/src/aggregates/spill.rs | 61 +- .../test_files/aggregate_blocked.slt | 160 ++++ 39 files changed, 5773 insertions(+), 721 deletions(-) create mode 100644 datafusion/expr-common/src/blocked_groups_accumulator.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs create mode 100644 datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs create mode 100644 datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates/group_values/blocked/primitive.rs create mode 100644 datafusion/sqllogictest/test_files/aggregate_blocked.slt 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..fa07ad494de07 --- /dev/null +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -0,0 +1,392 @@ +// 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, exec_err, not_impl_err}; + +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() + } + + /// 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 + } + } + } +} + +/// Selects groups for a non-destructive read of blocked state, like +/// [`GroupSelection`] for flat state. +/// +/// Selections created by [`Self::try_from_indices`] preserve the requested +/// order and support duplicate indices. +/// +/// [`GroupSelection`]: crate::groups_accumulator::GroupSelection +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlockedGroupSelection<'a> { + total_num_groups: usize, + block_size: usize, + indices: Option<&'a [BlocksIndex]>, +} + +impl<'a> BlockedGroupSelection<'a> { + /// Selects all `total_num_groups` groups, stored in blocks of + /// `block_size` groups, in group order. + pub fn all(total_num_groups: usize, block_size: usize) -> Self { + Self { + total_num_groups, + block_size, + indices: None, + } + } + + /// Selects groups in the order specified by `indices`. + /// + /// Returns an error if an index is not one of the first + /// `total_num_groups` groups. Empty selections are valid, and duplicate + /// indices are preserved. + pub fn try_from_indices( + indices: &'a [BlocksIndex], + total_num_groups: usize, + block_size: usize, + ) -> Result { + if let Some(index) = indices + .iter() + .find(|index| index.flat(block_size) >= total_num_groups) + { + return exec_err!( + "Group index {index:?} is out of bounds for {total_num_groups} groups" + ); + } + Ok(Self { + total_num_groups, + block_size, + indices: Some(indices), + }) + } + + /// Returns the group count against which this selection was constructed. + pub fn total_num_groups(self) -> usize { + self.total_num_groups + } + + /// Ensures this selection is being applied to the same number of groups + /// against which it was constructed, see + /// `GroupSelection::validate_num_groups`. + pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> { + if actual_num_groups != self.total_num_groups { + return exec_err!( + "Group selection was constructed for {} groups but applied to {actual_num_groups} groups", + self.total_num_groups + ); + } + Ok(()) + } + + /// Returns the number of selected groups. + pub fn len(self) -> usize { + self.indices + .map_or(self.total_num_groups, |indices| indices.len()) + } + + /// Returns `true` if no groups are selected. + pub fn is_empty(self) -> bool { + self.len() == 0 + } + + /// Returns the selected groups in output order. + pub fn iter(self) -> impl Iterator + 'a { + let block_size = self.block_size; + let (all, indices): (_, &'a [BlocksIndex]) = match self.indices { + None => (0..self.total_num_groups, &[]), + Some(indices) => (0..0, indices), + }; + all.map(move |flat| BlocksIndex::from_flat(flat, block_size)) + .chain(indices.iter().copied()) + } +} + +/// 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` only when every +/// aggregate in the query (and the group keys) support blocked storage; +/// otherwise the whole aggregation uses [`GroupsAccumulator`]. +/// +/// [`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 final values without changing the state or the group indices, + /// see `GroupsAccumulator::evaluate_preserving`. + /// + /// Rows are returned in the order of `selection`, in arrays of at most + /// [`Self::block_size`] rows. An empty selection returns no arrays. + fn evaluate_preserving( + &mut self, + _selection: BlockedGroupSelection<'_>, + ) -> Result> { + not_impl_err!("Preserving grouped evaluation is not implemented") + } + + /// Returns `true` if [`Self::evaluate_preserving`] is implemented. + fn supports_evaluate_preserving(&self) -> bool { + false + } + + /// 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>>; + + /// Returns intermediate state without changing the state or the group + /// indices, see `GroupsAccumulator::state_preserving`. + /// + /// Indexed `[chunk][state column]`, where chunks follow the same rules as + /// [`Self::evaluate_preserving`]. + fn state_preserving( + &mut self, + _selection: BlockedGroupSelection<'_>, + ) -> Result>> { + not_impl_err!("Preserving grouped state is not implemented") + } + + /// Returns `true` if [`Self::state_preserving`] is implemented. + fn supports_state_preserving(&self) -> bool { + false + } + + /// 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 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 ed9b4df4ce2bd..9c3cf0947a1c0 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -104,6 +104,9 @@ pub use datafusion_doc::{ pub use datafusion_expr_common::accumulator::{ Accumulator, AggregateMetric, AggregateMetricRecorder, AggregateMetrics, }; +pub use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, 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..d53ea58d02552 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,25 @@ 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 every + /// aggregate in a hash aggregation returns `true` (and the group key is + /// supported), group state is stored in blocks and released 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 +1657,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/count_distinct.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs index bb706aa614dbc..bdf4fc0eebd42 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs @@ -15,11 +15,13 @@ // specific language governing permissions and limitations // under the License. +mod blocked_groups; mod bytes; mod dict; mod groups; mod native; +pub use blocked_groups::BlockedPrimitiveDistinctCountGroupsAccumulator; pub use bytes::BytesDistinctCountAccumulator; pub use bytes::BytesViewDistinctCountAccumulator; pub use dict::DictionaryCountAccumulator; diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs new file mode 100644 index 0000000000000..04b124eb1979d --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs @@ -0,0 +1,347 @@ +// 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. + +//! [`BlockedPrimitiveDistinctCountGroupsAccumulator`]: distinct count with +//! state stored in blocks. + +use std::hash::Hash; +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, PrimitiveArray, +}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{ArrowPrimitiveType, Field}; +use datafusion_common::HashSet; +use datafusion_common::hash_utils::RandomState; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, +}; + +use super::groups::distinct_convert_to_state; +use crate::aggregate::groups_accumulator::accumulate::accumulate_blocked; +use crate::aggregate::groups_accumulator::blocked_vec::BlockedVec; + +/// Distinct values seen by the groups of one block, keyed by the group's +/// position in the block. +type BlockSeen = HashSet<(u32, V), RandomState>; + +/// [`PrimitiveDistinctCountGroupsAccumulator`] with state stored in blocks. +/// +/// The distinct values are kept in one set per block (instead of one set for +/// all groups), so a block is emitted by draining only its own set, and its +/// memory is released with it. +/// +/// [`PrimitiveDistinctCountGroupsAccumulator`]: super::PrimitiveDistinctCountGroupsAccumulator +pub struct BlockedPrimitiveDistinctCountGroupsAccumulator +where + T::Native: Eq + Hash, +{ + /// One set per block of `counts` + seen: Vec>, + counts: BlockedVec, +} + +impl BlockedPrimitiveDistinctCountGroupsAccumulator +where + T::Native: Eq + Hash, +{ + pub fn new(block_size: usize) -> Self { + Self { + seen: vec![], + counts: BlockedVec::new(block_size), + } + } + + /// Grows `counts` to `total_num_groups` and adds a set for every new block. + fn grow_to(&mut self, total_num_groups: usize) { + self.counts.grow_to(total_num_groups, 0); + self.seen + .resize_with(self.counts.num_blocks(), BlockSeen::default); + } + + /// Adds `value` to the group at `index`, counting it if it is new. + #[inline] + fn insert( + seen: &mut [BlockSeen], + counts: &mut BlockedVec, + index: BlocksIndex, + value: T::Native, + ) { + if seen[index.block_index()].insert((index.index_in_block() as u32, value)) { + // SAFETY: `index` was registered by `grow_to` + unsafe { *counts.get_unchecked_mut(index) += 1 }; + } + } + + /// Removes the counts and distinct values of the groups selected by + /// `emit_to`, one entry per emitted block. + fn take(&mut self, emit_to: BlockedEmitTo) -> Vec<(Vec, BlockSeen)> { + match emit_to { + BlockedEmitTo::All => { + let counts = self.counts.take_all(); + let seen = std::mem::take(&mut self.seen); + counts.into_iter().zip(seen).collect() + } + BlockedEmitTo::NextBlock => match self.counts.take_next_block() { + Some(counts) => vec![(counts, self.seen.remove(0))], + None => vec![], + }, + BlockedEmitTo::First(n) => { + // Rare path (ordered input): renumber every remaining group + let block_size = self.counts.block_size(); + let counts = self.counts.take_first(n); + let mut emitted = BlockSeen::default(); + let mut remaining: Vec> = + (0..self.counts.num_blocks()) + .map(|_| BlockSeen::default()) + .collect(); + for (block_index, block) in + std::mem::take(&mut self.seen).into_iter().enumerate() + { + for (index_in_block, value) in block { + let flat = BlocksIndex::new(block_index, index_in_block as usize) + .flat(block_size); + match flat.checked_sub(n) { + None => { + emitted.insert((index_in_block, value)); + } + Some(flat) => { + let index = BlocksIndex::from_flat(flat, block_size); + remaining[index.block_index()] + .insert((index.index_in_block() as u32, value)); + } + } + } + } + self.seen = remaining; + vec![(counts, emitted)] + } + } + } +} + +impl BlockedGroupsAccumulator + for BlockedPrimitiveDistinctCountGroupsAccumulator +where + T::Native: Eq + Hash, +{ + 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, + ) -> datafusion_common::Result<()> { + debug_assert_eq!(values.len(), 1); + self.grow_to(total_num_groups); + let arr = values[0].as_primitive::(); + let (seen, counts) = (&mut self.seen, &mut self.counts); + accumulate_blocked(group_indices, arr, opt_filter, |index, value| { + Self::insert(seen, counts, index, value) + }); + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> datafusion_common::Result<()> { + debug_assert_eq!(values.len(), 1); + self.grow_to(total_num_groups); + let list_array = values[0].as_list::(); + let inner_values = list_array.values().as_primitive::().values(); + let offsets = list_array.offsets(); + + for (row_idx, &index) in group_indices.iter().enumerate() { + let start = offsets[row_idx] as usize; + let end = offsets[row_idx + 1] as usize; + for &value in &inner_values[start..end] { + Self::insert(&mut self.seen, &mut self.counts, index, value); + } + } + Ok(()) + } + + fn evaluate( + &mut self, + emit_to: BlockedEmitTo, + ) -> datafusion_common::Result> { + Ok(self + .take(emit_to) + .into_iter() + // The emitted sets are dropped here, releasing their memory + .map(|(counts, _)| Arc::new(Int64Array::from(counts)) as ArrayRef) + .collect()) + } + + fn state( + &mut self, + emit_to: BlockedEmitTo, + ) -> datafusion_common::Result>> { + Ok(self + .take(emit_to) + .into_iter() + .map(|(counts, seen)| { + // Prefix-sum the counts of the block into list offsets, then + // move every distinct value to its group's list + let mut offsets = Vec::with_capacity(counts.len() + 1); + offsets.push(0i32); + let mut total = 0i32; + for &count in &counts { + total += count as i32; + offsets.push(total); + } + let mut values = vec![T::Native::default(); total as usize]; + let mut cursors: Vec = offsets[..counts.len()].to_vec(); + for (index_in_block, value) in seen { + let cursor = &mut cursors[index_in_block as usize]; + values[*cursor as usize] = value; + *cursor += 1; + } + let values = + Arc::new(PrimitiveArray::::new(ScalarBuffer::from(values), None)); + vec![Arc::new(ListArray::new( + Arc::new(Field::new_list_field(T::DATA_TYPE, true)), + OffsetBuffer::new(offsets.into()), + values, + None, + )) as ArrayRef] + }) + .collect()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> datafusion_common::Result> { + distinct_convert_to_state::(values, opt_filter) + } + + fn size(&self) -> usize { + let entry = size_of::<(u32, T::Native)>() + size_of::(); + size_of::() + + self + .seen + .iter() + .map(|s| s.capacity() * entry) + .sum::() + + self.counts.allocated_size() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Int32Array}; + use arrow::datatypes::Int32Type; + use datafusion_common::Result; + + fn indices(flats: &[usize]) -> Vec { + flats + .iter() + .map(|&f| BlocksIndex::from_flat(f, 4)) + .collect() + } + + fn counts(arrays: &[ArrayRef]) -> Vec> { + arrays + .iter() + .map(|a| { + a.as_primitive::() + .values() + .to_vec() + }) + .collect() + } + + fn accumulator() -> Result> + { + let mut acc = BlockedPrimitiveDistinctCountGroupsAccumulator::::new(4); + let values: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), + Some(1), + Some(2), + None, + Some(7), + Some(7), + Some(8), + ])); + acc.update_batch(&[values], &indices(&[0, 0, 0, 1, 5, 5, 5]), None, 6)?; + Ok(acc) + } + + #[test] + fn counts_distinct_values_per_block() -> Result<()> { + let mut acc = accumulator()?; + assert_eq!( + counts(&acc.evaluate(BlockedEmitTo::NextBlock)?), + vec![vec![2, 0, 0, 0]] + ); + assert_eq!(acc.seen.len(), 1, "the emitted block's set is released"); + assert_eq!(counts(&acc.evaluate(BlockedEmitTo::All)?), vec![vec![0, 2]]); + Ok(()) + } + + #[test] + fn state_round_trips_through_merge() -> Result<()> { + let mut acc = accumulator()?; + let state = acc.state(BlockedEmitTo::All)?; + assert_eq!(state.len(), 2); + let mut merged = + BlockedPrimitiveDistinctCountGroupsAccumulator::::new(4); + for (block_index, columns) in state.iter().enumerate() { + let rows = columns[0].len(); + let groups: Vec = (0..rows) + .map(|i| BlocksIndex::new(block_index, i)) + .collect(); + merged.merge_batch(columns, &groups, 6)?; + } + // merging the same values again does not change the distinct counts + merged.merge_batch(&state[1], &indices(&[4, 5]), 6)?; + assert_eq!( + counts(&merged.evaluate(BlockedEmitTo::All)?), + vec![vec![2, 0, 0, 0], vec![0, 2]] + ); + Ok(()) + } + + #[test] + fn first_n_renumbers_across_blocks() -> Result<()> { + let mut acc = accumulator()?; + assert_eq!( + counts(&acc.evaluate(BlockedEmitTo::First(1))?), + vec![vec![2]] + ); + // group 5 moved to position 4 (block 1, index 0) + let values: ArrayRef = Arc::new(Int32Array::from(vec![7, 9])); + acc.update_batch(&[values], &indices(&[4, 4]), None, 5)?; + assert_eq!( + counts(&acc.evaluate(BlockedEmitTo::All)?), + vec![vec![0, 0, 0, 0], vec![3]] + ); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 38f6d26200ec9..9b39ebc0fd103 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -211,24 +211,7 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> datafusion_common::Result> { - debug_assert_eq!(values.len(), 1); - let arr = values[0].as_primitive::(); - - let values_builder = PrimitiveBuilder::::with_capacity(arr.len()); - let mut builder = ListBuilder::new(values_builder) - .with_field(Arc::new(Field::new_list_field(T::DATA_TYPE, true))); - - for row in 0..arr.len() { - let included = arr.is_valid(row) - && opt_filter - .is_none_or(|filter| filter.is_valid(row) && filter.value(row)); - if included { - builder.values().append_value(arr.value(row)); - } - builder.append(true); - } - - Ok(vec![Arc::new(builder.finish())]) + distinct_convert_to_state::(values, opt_filter) } fn size(&self) -> usize { size_of::() @@ -237,6 +220,32 @@ where } } +/// Converts an input batch directly to state: a list per row with the row's +/// value, or an empty list for null or filtered rows. Shared by the flat and +/// blocked distinct count accumulators. +pub(super) fn distinct_convert_to_state( + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, +) -> datafusion_common::Result> { + debug_assert_eq!(values.len(), 1); + let arr = values[0].as_primitive::(); + + let values_builder = PrimitiveBuilder::::with_capacity(arr.len()); + let mut builder = ListBuilder::new(values_builder) + .with_field(Arc::new(Field::new_list_field(T::DATA_TYPE, true))); + + for row in 0..arr.len() { + let included = arr.is_valid(row) + && opt_filter.is_none_or(|filter| filter.is_valid(row) && filter.value(row)); + if included { + builder.values().append_value(arr.value(row)); + } + builder.append(true); + } + + Ok(vec![Arc::new(builder.finish())]) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index dbf2af512b9f7..bf0f24f122ef5 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -19,6 +19,9 @@ //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] pub mod accumulate; +pub mod blocked_null_state; +pub mod blocked_prim_op; +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..c912a12c5f3da 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 @@ -422,10 +423,38 @@ pub fn accumulate( group_indices: &[usize], values: &PrimitiveArray, opt_filter: Option<&BooleanArray>, - mut value_fn: F, + value_fn: F, ) where T: ArrowPrimitiveType + Send, F: FnMut(usize, T::Native) + Send, +{ + accumulate_values(group_indices, values, opt_filter, value_fn) +} + +/// [`accumulate`] for blocked group indices. +pub fn accumulate_blocked( + group_indices: &[BlocksIndex], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + value_fn: F, +) where + T: ArrowPrimitiveType, + F: FnMut(BlocksIndex, T::Native), +{ + accumulate_values(group_indices, values, opt_filter, value_fn) +} + +/// Shared implementation of [`accumulate`] and [`accumulate_blocked`]. +#[inline(always)] +fn accumulate_values( + group_indices: &[I], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + mut value_fn: F, +) where + I: Copy, + T: ArrowPrimitiveType, + F: FnMut(I, T::Native), { let data: &[T::Native] = values.values(); assert_eq!(data.len(), group_indices.len()); @@ -592,9 +621,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_null_state.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs new file mode 100644 index 0000000000000..04d358c3d0f70 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs @@ -0,0 +1,385 @@ +// 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. + +//! [`BlockedNullState`]: [`NullState`] for blocked group indices. +//! +//! [`NullState`]: super::accumulate::NullState + +use arrow::array::{ + Array, ArrowPrimitiveType, BooleanArray, BooleanBufferBuilder, PrimitiveArray, +}; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; + +use super::blocked_vec::{BlockedVec, SeenGroups}; + +/// Tracks which groups have seen at least one non null, non filtered value, +/// like [`NullState`], for state stored in a [`BlockedVec`]. +/// +/// The seen bits are stored in one bitmap per block, with the same block +/// boundaries as the [`BlockedVec`], so each emitted block gets its own +/// [`NullBuffer`]. +/// +/// [`NullState`]: super::accumulate::NullState +#[derive(Debug)] +pub struct BlockedNullState { + block_size: usize, + seen: Seen, +} + +#[derive(Debug)] +enum Seen { + /// Every group seen so far has seen a value, so no bitmap is needed + All { num_groups: usize }, + /// Every bitmap but the last holds exactly `block_size` bits + Some { blocks: Vec }, +} + +impl BlockedNullState { + pub fn new(block_size: usize) -> Self { + Self { + block_size, + seen: Seen::All { num_groups: 0 }, + } + } + + /// Size of all buffers allocated by this null state, not including self. + pub fn size(&self) -> usize { + match &self.seen { + Seen::All { .. } => 0, + Seen::Some { blocks } => blocks.iter().map(|b| b.capacity() / 8).sum(), + } + } + + /// Updates `state` with every non null, non filtered value of `values` + /// through [`BlockedVec::update_values`], and marks the groups of those + /// values as seen. + /// + /// Like [`NullState::accumulate`], no bits are tracked while every value + /// is valid and every new group is present in `group_indices`. + /// + /// [`NullState::accumulate`]: super::accumulate::NullState::accumulate + #[expect(clippy::too_many_arguments, reason = "mirrors `NullState::accumulate`")] + pub fn accumulate( + &mut self, + state: &mut BlockedVec, + total_num_groups: usize, + starting_value: T, + group_indices: &[BlocksIndex], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + update_fn: F, + ) where + T: Copy + Default, + V: ArrowPrimitiveType, + F: FnMut(&mut T, V::Native), + { + if opt_filter.is_none() + && values.null_count() == 0 + && let Seen::All { num_groups } = &mut self.seen + && new_groups_are_dense( + group_indices, + *num_groups, + total_num_groups, + self.block_size, + ) + { + state.update_values( + total_num_groups, + starting_value, + group_indices, + values, + None, + &mut (), + update_fn, + ); + *num_groups = total_num_groups; + return; + } + + let mut seen = SeenBits { + blocks: self.bitmaps(total_num_groups), + }; + state.update_values( + total_num_groups, + starting_value, + group_indices, + values, + opt_filter, + &mut seen, + update_fn, + ); + } + + /// Returns the bitmaps, grown to `total_num_groups` bits (new groups have + /// not seen a value), switching from [`Seen::All`] if needed. + fn bitmaps(&mut self, total_num_groups: usize) -> &mut Vec { + if let Seen::All { num_groups } = self.seen { + let mut blocks = vec![]; + append_bits(&mut blocks, self.block_size, num_groups, true); + self.seen = Seen::Some { blocks }; + } + let Seen::Some { blocks } = &mut self.seen else { + unreachable!("switched to `Seen::Some` above") + }; + let len = bits_len(blocks); + if total_num_groups > len { + append_bits(blocks, self.block_size, total_num_groups - len, false); + } + blocks + } + + /// Builds the null buffer of every block emitted by `emit_to`, `None` + /// when a block has no nulls, and removes them from the state. + /// + /// Returns one entry per emitted block, following the same rules as + /// `BlockedGroupsAccumulator::evaluate`. + pub fn build(&mut self, emit_to: BlockedEmitTo) -> Vec> { + let block_size = self.block_size; + match &mut self.seen { + Seen::All { num_groups } => { + let emitted_blocks = match emit_to { + BlockedEmitTo::All => { + let blocks = num_groups.div_ceil(block_size); + *num_groups = 0; + blocks + } + BlockedEmitTo::NextBlock => { + let blocks = usize::from(*num_groups > 0); + *num_groups -= (*num_groups).min(block_size); + blocks + } + BlockedEmitTo::First(n) => { + *num_groups -= n; + 1 + } + }; + vec![None; emitted_blocks] + } + Seen::Some { blocks } => match emit_to { + BlockedEmitTo::All => { + let nulls = blocks.iter_mut().map(|b| Some(finish(b))).collect(); + self.seen = Seen::All { num_groups: 0 }; + nulls + } + BlockedEmitTo::NextBlock => { + if blocks.is_empty() { + return vec![]; + } + vec![Some(finish(&mut blocks.remove(0)))] + } + BlockedEmitTo::First(n) => { + // Rare path (ordered input): re-chunk the remaining bits + let bits: Vec = + blocks.iter_mut().map(|b| b.finish()).collect(); + let first = bits[0].slice(0, n); + let mut remaining = vec![]; + for (i, block) in bits.iter().enumerate() { + let skip = if i == 0 { n } else { 0 }; + let block = block.slice(skip, block.len() - skip); + append_buffer(&mut remaining, block_size, &block); + } + *blocks = remaining; + vec![Some(NullBuffer::new(first))] + } + }, + } + } +} + +/// Marks groups as seen in the per-block bitmaps. +struct SeenBits<'a> { + blocks: &'a mut Vec, +} + +impl SeenGroups for SeenBits<'_> { + #[inline] + fn mark(&mut self, index: BlocksIndex) { + self.blocks[index.block_index()].set_bit(index.index_in_block(), true); + } +} + +/// Returns true when all newly registered groups are present in +/// `group_indices`, see the flat `new_groups_are_dense`. +fn new_groups_are_dense( + group_indices: &[BlocksIndex], + first_new_group: usize, + total_num_groups: usize, + block_size: usize, +) -> bool { + if first_new_group == total_num_groups { + return true; + } + + let mut next_new_group = first_new_group; + for &group_index in group_indices { + let group_index = group_index.flat(block_size); + if group_index == next_new_group { + next_new_group += 1; + } else if group_index > next_new_group { + return false; + } + } + next_new_group == total_num_groups +} + +fn bits_len(blocks: &[BooleanBufferBuilder]) -> usize { + blocks.iter().map(|b| b.len()).sum() +} + +/// Appends `n` bits of `value`, filling the last bitmap up to `block_size` +/// bits before starting a new one. +fn append_bits( + blocks: &mut Vec, + block_size: usize, + mut n: usize, + value: bool, +) { + while n > 0 { + let last = match blocks.last_mut() { + Some(last) if last.len() < block_size => last, + _ => { + blocks.push(BooleanBufferBuilder::new(n.min(block_size))); + blocks.last_mut().expect("just pushed") + } + }; + let count = n.min(block_size - last.len()); + last.append_n(count, value); + n -= count; + } +} + +/// Appends the bits of `buffer`, like [`append_bits`]. +fn append_buffer( + blocks: &mut Vec, + block_size: usize, + buffer: &BooleanBuffer, +) { + let mut offset = 0; + while offset < buffer.len() { + let last = match blocks.last_mut() { + Some(last) if last.len() < block_size => last, + _ => { + blocks.push(BooleanBufferBuilder::new( + (buffer.len() - offset).min(block_size), + )); + blocks.last_mut().expect("just pushed") + } + }; + let count = (buffer.len() - offset).min(block_size - last.len()); + last.append_buffer(&buffer.slice(offset, count)); + offset += count; + } +} + +fn finish(builder: &mut BooleanBufferBuilder) -> NullBuffer { + NullBuffer::new(builder.finish()) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::Int64Type; + + fn indices(flats: &[usize]) -> Vec { + flats + .iter() + .map(|&f| BlocksIndex::from_flat(f, 4)) + .collect() + } + + fn nulls_of(nulls: &[Option]) -> Vec>> { + nulls + .iter() + .map(|n| n.as_ref().map(|n| n.iter().collect())) + .collect() + } + + #[test] + fn all_seen_needs_no_bitmap() { + let mut state = BlockedVec::::new(4); + let mut nulls = BlockedNullState::new(4); + let values = Int64Array::from(vec![1, 2, 3, 4, 5, 6]); + nulls.accumulate( + &mut state, + 6, + 0, + &indices(&[0, 1, 2, 3, 4, 5]), + &values, + None, + |s, v| *s += v, + ); + assert_eq!(nulls.size(), 0); + assert_eq!(nulls.build(BlockedEmitTo::All), vec![None, None]); + } + + #[test] + fn tracks_unseen_groups_across_blocks() { + let mut state = BlockedVec::::new(4); + let mut nulls = BlockedNullState::new(4); + // group 1 only gets a null, group 5 is filtered out + let values = + Int64Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]); + let filter = BooleanArray::from(vec![true, true, true, true, false, true]); + nulls.accumulate( + &mut state, + 7, + 0, + &indices(&[0, 1, 2, 3, 5, 6]), + &values, + Some(&filter), + |s, v| *s += v, + ); + let expected = vec![ + Some(vec![true, false, true, true]), + Some(vec![false, false, true]), + ]; + assert_eq!(nulls_of(&nulls.build(BlockedEmitTo::All)), expected); + } + + #[test] + fn next_block_and_first_n_shift_bits() { + let mut state = BlockedVec::::new(4); + let mut nulls = BlockedNullState::new(4); + let values = + Int64Array::from(vec![Some(1), None, Some(3), Some(4), None, Some(6)]); + nulls.accumulate( + &mut state, + 6, + 0, + &indices(&[0, 1, 2, 3, 4, 5]), + &values, + None, + |s, v| *s += v, + ); + assert_eq!( + nulls_of(&nulls.build(BlockedEmitTo::First(1))), + vec![Some(vec![true])] + ); + // remaining: [f, t, t, f | t] + assert_eq!( + nulls_of(&nulls.build(BlockedEmitTo::NextBlock)), + vec![Some(vec![false, true, true, false])] + ); + assert_eq!( + nulls_of(&nulls.build(BlockedEmitTo::All)), + vec![Some(vec![true])] + ); + let _ = Int64Type::DATA_TYPE; + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs new file mode 100644 index 0000000000000..6b52acc9fb9e2 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs @@ -0,0 +1,245 @@ +// 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. + +//! [`BlockedPrimitiveGroupsAccumulator`]: [`PrimitiveGroupsAccumulator`] with +//! state stored in blocks. +//! +//! [`PrimitiveGroupsAccumulator`]: super::prim_op::PrimitiveGroupsAccumulator + +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray}; +use arrow::datatypes::{ArrowPrimitiveType, DataType}; +use datafusion_common::Result; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, +}; + +use super::blocked_null_state::BlockedNullState; +use super::blocked_vec::BlockedVec; +use super::prim_op::convert_to_state; + +/// [`PrimitiveGroupsAccumulator`] with per-group values stored in a +/// [`BlockedVec`], for aggregates whose state is the same as the input type +/// (such as `SUM`). +/// +/// F: The function to apply to two elements. The first argument is the +/// existing value and should be updated with the second value (e.g. +/// [`BitAndAssign`] style). +/// +/// [`PrimitiveGroupsAccumulator`]: super::prim_op::PrimitiveGroupsAccumulator +/// [`BitAndAssign`]: std::ops::BitAndAssign +#[derive(Debug)] +pub struct BlockedPrimitiveGroupsAccumulator +where + T: ArrowPrimitiveType + Send, + F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, +{ + /// values per group, stored as the native type + values: BlockedVec, + + /// The output type (needed for Decimal precision and scale) + data_type: DataType, + + /// The starting value for new groups + starting_value: T::Native, + + /// Track nulls in the input / filters + null_state: BlockedNullState, + + /// Function that computes the primitive result + prim_fn: F, +} + +impl BlockedPrimitiveGroupsAccumulator +where + T: ArrowPrimitiveType + Send, + F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, +{ + pub fn new(data_type: &DataType, block_size: usize, prim_fn: F) -> Self { + Self { + values: BlockedVec::new(block_size), + data_type: data_type.clone(), + null_state: BlockedNullState::new(block_size), + starting_value: T::default_value(), + prim_fn, + } + } + + /// Set the starting values for new groups + pub fn with_starting_value(mut self, starting_value: T::Native) -> Self { + self.starting_value = starting_value; + self + } + + fn take(&mut self, emit_to: BlockedEmitTo) -> Vec> { + match emit_to { + BlockedEmitTo::All => self.values.take_all(), + BlockedEmitTo::NextBlock => { + self.values.take_next_block().into_iter().collect() + } + BlockedEmitTo::First(n) => vec![self.values.take_first(n)], + } + } +} + +impl BlockedGroupsAccumulator for BlockedPrimitiveGroupsAccumulator +where + T: ArrowPrimitiveType + Send, + F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, +{ + fn block_size(&self) -> usize { + self.values.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].as_primitive::(); + let prim_fn = &self.prim_fn; + self.null_state.accumulate( + &mut self.values, + total_num_groups, + self.starting_value, + group_indices, + values, + opt_filter, + |value, new_value| prim_fn(value, new_value), + ); + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()> { + // update / merge are the same + self.update_batch(values, group_indices, None, total_num_groups) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + let blocks = self.take(emit_to); + let nulls = self.null_state.build(emit_to); + debug_assert_eq!(blocks.len(), nulls.len()); + Ok(blocks + .into_iter() + .zip(nulls) + .map(|(values, nulls)| { + // no copy + Arc::new( + PrimitiveArray::::new(values.into(), nulls) + .with_data_type(self.data_type.clone()), + ) as ArrayRef + }) + .collect()) + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + Ok(self + .evaluate(emit_to)? + .into_iter() + .map(|values| vec![values]) + .collect()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + convert_to_state::( + values, + opt_filter, + self.starting_value, + &self.prim_fn, + &self.data_type, + ) + } + + fn size(&self) -> usize { + self.values.allocated_size() + self.null_state.size() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Int64Array}; + use arrow::datatypes::Int64Type; + + fn sum() -> BlockedPrimitiveGroupsAccumulator { + BlockedPrimitiveGroupsAccumulator::new(&DataType::Int64, 4, |x, y| *x += y) + } + + fn indices(flats: &[usize]) -> Vec { + flats + .iter() + .map(|&f| BlocksIndex::from_flat(f, 4)) + .collect() + } + + fn values(arrays: &[ArrayRef]) -> Vec>> { + arrays + .iter() + .map(|a| a.as_primitive::().iter().collect()) + .collect() + } + + #[test] + fn sums_across_blocks_with_nulls_for_unseen_groups() -> Result<()> { + let mut acc = sum(); + let input: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + Some(2), + None, + Some(4), + Some(5), + ])); + acc.update_batch(&[input], &indices(&[0, 5, 1, 5, 0]), None, 6)?; + let partial: ArrayRef = Arc::new(Int64Array::from(vec![10, 20])); + acc.merge_batch(&[partial], &indices(&[2, 5]), 6)?; + + assert_eq!( + values(&acc.evaluate(BlockedEmitTo::All)?), + vec![vec![Some(6), None, Some(10), None], vec![None, Some(26)],] + ); + assert_eq!(acc.size(), 0); + Ok(()) + } + + #[test] + fn all_valid_input_has_no_nulls() -> Result<()> { + let mut acc = sum(); + let input: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])); + acc.update_batch(&[input], &indices(&[0, 1, 2, 3, 4]), None, 5)?; + let out = acc.evaluate(BlockedEmitTo::NextBlock)?; + assert_eq!(out.len(), 1); + assert_eq!(out[0].null_count(), 0); + assert_eq!(values(&out), vec![vec![Some(1), Some(2), Some(3), Some(4)]]); + let state = acc.state(BlockedEmitTo::All)?; + assert_eq!(state.len(), 1); + assert_eq!(values(&state[0]), vec![vec![Some(5)]]); + Ok(()) + } +} 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..341fbe744e6c3 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs @@ -0,0 +1,801 @@ +// 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::{Array, ArrowPrimitiveType, BooleanArray, PrimitiveArray}; +use arrow::buffer::NullBuffer; + +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; + +use super::accumulate::{accumulate_blocked, accumulate_blocked_indices}; + +/// Told about the group of every row whose value is used by +/// [`BlockedVec::update_values`], e.g. to track which groups saw a value. +pub trait SeenGroups { + fn mark(&mut self, index: BlocksIndex); +} + +/// Tracks nothing, so marking compiles away. +impl SeenGroups for () { + #[inline(always)] + fn mark(&mut self, _index: BlocksIndex) {} +} + +/// 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 +/// +/// ## Items that own heap memory (`T` is not `Copy`) +/// +/// `T` only needs to be `Clone` (for [`Self::grow_to`]), so it can own heap +/// memory, like the `Option>` of `min` / `max` on bytes. Then: +/// +/// 1. [`BlockedVec::allocated_size`] only counts the slots (`size_of::()` +/// each), not the heap allocations the items own, which the caller must +/// track itself. +/// 2. The mutable access given by [`Self::update`] and friends can change the +/// heap memory of an item without [`BlockedVec::allocated_size`] changing, +/// which is another reason the caller tracks that memory. +/// 3. Reading an item's heap memory is one more dependent load after the +/// block and the slot, which the address-resolve loops cannot hide. +/// Accumulators should read the stored items as rarely as possible, e.g. +/// once per group per batch instead of once per row. +/// +/// [`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 { + /// 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()] + } +} + +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.clone()); + 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 a reference to the element at `index`. + /// + /// # Panics + /// If `index` is out of bounds. + #[inline] + pub fn get_ref(&self, index: BlocksIndex) -> &T { + &self.blocks[index.block_index()][index.index_in_block()] + } + + /// Returns a mutable reference to the element at `index`. + /// + /// # Panics + /// If `index` is out of bounds. + #[inline] + pub fn get_mut(&mut self, index: BlocksIndex) -> &mut T { + &mut 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`]. It is only checked in debug + /// builds, like the flat accumulators' unchecked indexing. + /// + /// [`BlockedGroupsAccumulator::update_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::update_batch + /// [`BlockedGroupsAccumulator::merge_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::merge_batch + pub fn update( + &mut self, + total_num_groups: usize, + starting_value: T, + group_indices: &[BlocksIndex], + nulls: Option<&NullBuffer>, + opt_filter: Option<&BooleanArray>, + mut update_fn: F, + ) where + F: FnMut(&mut T), + { + self.grow_to(total_num_groups, starting_value); + self.debug_assert_in_bounds(group_indices); + if let Some(block) = self.as_single_block_mut() { + accumulate_blocked_indices(group_indices, nulls, opt_filter, |index| { + // SAFETY: indices are in bounds, see the method docs + 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, see the method docs + *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, see the method docs + 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`]. It is only checked in debug + /// builds, like the flat accumulators' unchecked indexing. + /// + /// [`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. + pub fn update_with( + &mut self, + total_num_groups: usize, + starting_value: T, + group_indices: &[BlocksIndex], + values: &[V], + mut update_fn: F, + ) where + V: Copy, + F: FnMut(&mut T, V), + { + assert_eq!(group_indices.len(), values.len()); + self.grow_to(total_num_groups, starting_value); + self.debug_assert_in_bounds(group_indices); + if let Some(block) = self.as_single_block_mut() { + for (&index, &value) in group_indices.iter().zip(values) { + // SAFETY: indices are in bounds, see the method docs + 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, see the method docs + *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); + } + } + } + + /// Grows the vector to `total_num_groups` elements (new ones are + /// `starting_value`), then, for every row of `values` that is not null + /// and is selected by `opt_filter`, calls `seen.mark` with the row's group + /// and `update_fn` with the group's element and the row's value. + /// + /// Chooses the loop once per call, like [`Self::update`]. + /// + /// `group_indices` must only point to the first `total_num_groups` + /// groups, like for [`Self::update`]. + /// + /// # Panics + /// If `values` and `group_indices` have different lengths. + #[expect(clippy::too_many_arguments, reason = "mirrors `NullState::accumulate`")] + pub fn update_values( + &mut self, + total_num_groups: usize, + starting_value: T, + group_indices: &[BlocksIndex], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + seen: &mut S, + mut update_fn: F, + ) where + V: ArrowPrimitiveType, + S: SeenGroups, + F: FnMut(&mut T, V::Native), + { + assert_eq!(group_indices.len(), values.len()); + self.grow_to(total_num_groups, starting_value); + self.debug_assert_in_bounds(group_indices); + let all_valid = values.null_count() == 0 && opt_filter.is_none(); + + if let Some(block) = self.as_single_block_mut() { + if all_valid { + for (&index, &value) in group_indices.iter().zip(values.values()) { + seen.mark(index); + // SAFETY: indices are in bounds, see the method docs + update_fn( + unsafe { block.get_unchecked_mut(index.index_in_block()) }, + value, + ); + } + } else { + accumulate_blocked(group_indices, values, opt_filter, |index, value| { + seen.mark(index); + // SAFETY: indices are in bounds, see the method docs + update_fn( + unsafe { block.get_unchecked_mut(index.index_in_block()) }, + value, + ) + }); + } + return; + } + + let ptrs = self.block_ptrs_mut(); + if all_valid { + let mut addrs = [std::ptr::null_mut::(); RESOLVE_CHUNK]; + for (indices, values) in group_indices + .chunks(RESOLVE_CHUNK) + .zip(values.values().chunks(RESOLVE_CHUNK)) + { + for (addr, &index) in addrs.iter_mut().zip(indices) { + // SAFETY: indices are in bounds, see the method docs + *addr = unsafe { ptrs.ptr(index) }; + } + for ((&addr, &index), &value) in addrs.iter().zip(indices).zip(values) { + seen.mark(index); + // SAFETY: resolved above from a block that is not touched + // while `ptrs` lives + update_fn(unsafe { &mut *addr }, value); + } + } + } else { + accumulate_blocked(group_indices, values, opt_filter, |index, value| { + seen.mark(index); + // SAFETY: indices are in bounds, see the method docs + update_fn(unsafe { &mut *ptrs.ptr(index) }, value) + }); + } + } + + /// Checks, in debug builds only, that every index points to an element. + #[inline] + fn debug_assert_in_bounds(&self, group_indices: &[BlocksIndex]) { + debug_assert!( + group_indices.iter().all(|index| { + self.blocks + .get(index.block_index()) + .is_some_and(|block| index.index_in_block() < block.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); + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index 0453124528a0f..e0422c1646a2b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -174,53 +174,75 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> Result> { - let values = values[0].as_primitive::().clone(); - - // Initializing state with starting values - let initial_state = - PrimitiveArray::::from_value(self.starting_value, values.len()); - - // Recalculating values in case there is filter - let values = match opt_filter { - None => values, - Some(filter) => { - let (filter_values, filter_nulls) = filter.clone().into_parts(); - // Calculating filter mask as a result of bitand of filter, and converting it to null buffer - let filter_bool = match filter_nulls { - Some(filter_nulls) => filter_nulls.inner() & &filter_values, - None => filter_values, - }; - let filter_nulls = NullBuffer::from(filter_bool); - - // Rebuilding input values with a new nulls mask, which is equal to - // the union of original nulls and filter mask - let (dt, values_buf, original_nulls) = values.into_parts(); - let nulls_buf = - NullBuffer::union(original_nulls.as_ref(), Some(&filter_nulls)); - PrimitiveArray::::new(values_buf, nulls_buf).with_data_type(dt) - } - }; - - let state_values = compute::binary_mut(initial_state, &values, |mut x, y| { - (self.prim_fn)(&mut x, y); - x - }); - let state_values = state_values - .map_err(|_| { - internal_datafusion_err!( - "initial_values underlying buffer must not be shared" - ) - })? - .map_err(DataFusionError::from)? - .with_data_type(self.data_type.clone()); - - Ok(vec![Arc::new(state_values)]) + convert_to_state::( + values, + opt_filter, + self.starting_value, + &self.prim_fn, + &self.data_type, + ) } fn size(&self) -> usize { self.values.capacity() * size_of::() + self.null_state.size() } } +/// Converts an input batch directly to the state of a primitive accumulator +/// (flat or blocked): `prim_fn` applied to `starting_value` and every non +/// null, non filtered value, and null otherwise. +pub(crate) fn convert_to_state( + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + starting_value: T::Native, + prim_fn: &F, + data_type: &DataType, +) -> Result> +where + T: ArrowPrimitiveType, + F: Fn(&mut T::Native, T::Native), +{ + let values = values[0].as_primitive::().clone(); + + // Initializing state with starting values + let initial_state = PrimitiveArray::::from_value(starting_value, values.len()); + + // Recalculating values in case there is filter + let values = match opt_filter { + None => values, + Some(filter) => { + let (filter_values, filter_nulls) = filter.clone().into_parts(); + // Calculating filter mask as a result of bitand of filter, and converting it to null buffer + let filter_bool = match filter_nulls { + Some(filter_nulls) => filter_nulls.inner() & &filter_values, + None => filter_values, + }; + let filter_nulls = NullBuffer::from(filter_bool); + + // Rebuilding input values with a new nulls mask, which is equal to + // the union of original nulls and filter mask + let (dt, values_buf, original_nulls) = values.into_parts(); + let nulls_buf = + NullBuffer::union(original_nulls.as_ref(), Some(&filter_nulls)); + PrimitiveArray::::new(values_buf, nulls_buf).with_data_type(dt) + } + }; + + let state_values = compute::binary_mut(initial_state, &values, |mut x, y| { + prim_fn(&mut x, y); + x + }); + let state_values = state_values + .map_err(|_| { + internal_datafusion_err!( + "initial_values underlying buffer must not be shared" + ) + })? + .map_err(DataFusionError::from)? + .with_data_type(data_type.clone()); + + Ok(vec![Arc::new(state_values)]) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index 85b63a9a17970..6137010e2b386 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -38,14 +38,16 @@ use datafusion_common::{ use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Coercion, Documentation, EmitTo, Expr, GroupSelection, - GroupsAccumulator, ReversedUDAF, Signature, TypeSignature, TypeSignatureClass, - Volatility, + Accumulator, AggregateUDFImpl, BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, + Coercion, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, + ReversedUDAF, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_functions_aggregate_common::aggregate::avg_distinct::{ DecimalDistinctAvgAccumulator, Float64DistinctAvgAccumulator, }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::NullState; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_null_state::BlockedNullState; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::BlockedVec; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{ filtered_null_mask, set_nulls, }; @@ -484,6 +486,63 @@ impl AggregateUDFImpl for Avg { } } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + matches!( + ( + args.expr_fields[0].data_type(), + args.return_field.data_type() + ), + (DataType::Float64, DataType::Float64) + | (DataType::Duration(_), DataType::Duration(_)) + ) && !args.is_distinct + } + + fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + use DataType::*; + + let data_type = args.expr_fields[0].data_type(); + let return_type = args.return_field.data_type(); + let duration_avg_fn = |sum: i64, count: u64| Ok(sum / count as i64); + macro_rules! duration { + ($t:ty) => { + Ok(Box::new(BlockedAvgGroupsAccumulator::<$t, _>::new( + data_type, + return_type, + block_size, + duration_avg_fn, + ))) + }; + } + match (data_type, return_type) { + (Float64, Float64) => Ok(Box::new(BlockedAvgGroupsAccumulator::< + Float64Type, + _, + >::new( + data_type, + return_type, + block_size, + |sum: f64, count: u64| Ok(sum / count as f64), + ))), + (Duration(TimeUnit::Second), Duration(_)) => duration!(DurationSecondType), + (Duration(TimeUnit::Millisecond), Duration(_)) => { + duration!(DurationMillisecondType) + } + (Duration(TimeUnit::Microsecond), Duration(_)) => { + duration!(DurationMicrosecondType) + } + (Duration(TimeUnit::Nanosecond), Duration(_)) => { + duration!(DurationNanosecondType) + } + _ => not_impl_err!( + "BlockedAvgGroupsAccumulator for ({data_type} --> {return_type})" + ), + } + } + fn aliases(&self) -> &[String] { &self.aliases } @@ -617,6 +676,113 @@ where ))) } +/// Computes the averages of groups, `null` for groups whose `nulls` bit is +/// unset. Shared by [`AvgGroupsAccumulator`] and +/// [`BlockedAvgGroupsAccumulator`]. +fn evaluate_avg_values( + counts: Vec, + sums: Vec, + nulls: Option, + avg_fn: &F, + return_data_type: &DataType, +) -> Result +where + O: ArrowNumericType, + S: ArrowNumericType, + F: Fn(S::Native, u64) -> Result, +{ + if let Some(nulls) = &nulls { + assert_eq!(nulls.len(), sums.len()); + } + assert_eq!(counts.len(), sums.len()); + + // Don't evaluate averages with null inputs to avoid errors on null values. + let array: PrimitiveArray = if let Some(nulls) = &nulls + && nulls.null_count() > 0 + { + let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) + .with_data_type(return_data_type.clone()); + let iter = sums.into_iter().zip(counts).zip(nulls.iter()); + + for ((sum, count), is_valid) in iter { + if is_valid { + builder.append_value(avg_fn(sum, count)?) + } else { + builder.append_null(); + } + } + builder.finish() + } else { + let averages: Vec = sums + .into_iter() + .zip(counts) + .map(|(sum, count)| avg_fn(sum, count)) + .collect::>>()?; + PrimitiveArray::new(averages.into(), nulls) + .with_data_type(return_data_type.clone()) + }; + + Ok(Arc::new(array)) +} + +/// The `[counts, sums]` state of groups. Shared by [`AvgGroupsAccumulator`] +/// and [`BlockedAvgGroupsAccumulator`]. +fn avg_state_values( + counts: Vec, + sums: Vec, + nulls: Option, + sum_data_type: &DataType, +) -> Vec { + let counts = UInt64Array::new(counts.into(), nulls.clone()); + let sums = PrimitiveArray::::new(sums.into(), nulls) + .with_data_type(sum_data_type.clone()); + vec![Arc::new(counts), Arc::new(sums)] +} + +/// Converts an input batch directly to `[counts, sums]` state. Shared by +/// [`AvgGroupsAccumulator`] and [`BlockedAvgGroupsAccumulator`]. +fn avg_convert_to_state( + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + sum_data_type: &DataType, +) -> Result> +where + I: ArrowNumericType, + S: ArrowNumericType + Send, + I::Native: Into, +{ + // When the sum type equals the input type (`I == S`: `Float64`, + // `Duration`, `Decimal256`, and any decimal whose precision already + // leaves [`avg_sum_data_type`] enough headroom) the input is already a + // valid sum array and is reused as is; the downcast is by Rust type, so + // it succeeds even when precision differs. Otherwise every value is + // widened. + let sums = match values[0].as_any().downcast_ref::>() { + Some(sums) => sums.clone().with_data_type(sum_data_type.clone()), + None => { + let values = values[0].as_primitive::(); + // Values under null slots are widened too rather than branching per + // element; `set_nulls` below masks them out again. + let sums: Vec = values + .values() + .iter() + .map(|value| (*value).into()) + .collect(); + PrimitiveArray::::new(sums.into(), values.nulls().cloned()) + .with_data_type(sum_data_type.clone()) + } + }; + let counts = UInt64Array::from_value(1, sums.len()); + + let nulls = filtered_null_mask(opt_filter, &sums); + + // set nulls on the arrays + let counts = set_nulls(counts, nulls.clone()); + let sums = set_nulls(sums, nulls); + + Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)]) +} + /// An accumulator to compute the average #[derive(Debug, Default)] pub struct AvgAccumulator { @@ -1040,38 +1206,13 @@ where sums: Vec, nulls: Option, ) -> Result { - if let Some(nulls) = &nulls { - assert_eq!(nulls.len(), sums.len()); - } - assert_eq!(counts.len(), sums.len()); - - // Don't evaluate averages with null inputs to avoid errors on null values. - let array: PrimitiveArray = if let Some(nulls) = &nulls - && nulls.null_count() > 0 - { - let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) - .with_data_type(self.return_data_type.clone()); - let iter = sums.into_iter().zip(counts).zip(nulls.iter()); - - for ((sum, count), is_valid) in iter { - if is_valid { - builder.append_value((self.avg_fn)(sum, count)?) - } else { - builder.append_null(); - } - } - builder.finish() - } else { - let averages: Vec = sums - .into_iter() - .zip(counts) - .map(|(sum, count)| (self.avg_fn)(sum, count)) - .collect::>>()?; - PrimitiveArray::new(averages.into(), nulls) - .with_data_type(self.return_data_type.clone()) - }; - - Ok(Arc::new(array)) + evaluate_avg_values::( + counts, + sums, + nulls, + &self.avg_fn, + &self.return_data_type, + ) } fn state_values( @@ -1080,10 +1221,7 @@ where sums: Vec, nulls: Option, ) -> Vec { - let counts = UInt64Array::new(counts.into(), nulls.clone()); - let sums = PrimitiveArray::::new(sums.into(), nulls) - .with_data_type(self.sum_data_type.clone()); - vec![Arc::new(counts), Arc::new(sums)] + avg_state_values::(counts, sums, nulls, &self.sum_data_type) } } @@ -1216,36 +1354,7 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> Result> { - // When the sum type equals the input type (`I == S`: `Float64`, - // `Duration`, `Decimal256`, and any decimal whose precision already - // leaves [`avg_sum_data_type`] enough headroom) the input is already a - // valid sum array and is reused as is; the downcast is by Rust type, so - // it succeeds even when precision differs. Otherwise every value is - // widened. - let sums = match values[0].as_any().downcast_ref::>() { - Some(sums) => sums.clone().with_data_type(self.sum_data_type.clone()), - None => { - let values = values[0].as_primitive::(); - // Values under null slots are widened too rather than branching per - // element; `set_nulls` below masks them out again. - let sums: Vec = values - .values() - .iter() - .map(|value| (*value).into()) - .collect(); - PrimitiveArray::::new(sums.into(), values.nulls().cloned()) - .with_data_type(self.sum_data_type.clone()) - } - }; - let counts = UInt64Array::from_value(1, sums.len()); - - let nulls = filtered_null_mask(opt_filter, &sums); - - // set nulls on the arrays - let counts = set_nulls(counts, nulls.clone()); - let sums = set_nulls(sums, nulls); - - Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)]) + avg_convert_to_state::(values, opt_filter, &self.sum_data_type) } fn size(&self) -> usize { // Heap buffers @@ -1256,9 +1365,249 @@ where } } +/// [`AvgGroupsAccumulator`] with sums and counts stored in blocks. +/// +/// Sums and counts are kept in two [`BlockedVec`]s and updated in two passes, +/// so each pass keeps the fast loops of [`BlockedVec`]. +#[derive(Debug)] +struct BlockedAvgGroupsAccumulator +where + I: ArrowNumericType + Send, + O: ArrowNumericType + Send, + S: ArrowNumericType + Send, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + 'static, +{ + /// The type of the internal sum + sum_data_type: DataType, + + /// The type of the returned average + return_data_type: DataType, + + /// Count per group (use u64 to make UInt64Array) + counts: BlockedVec, + + /// Sums per group, stored as the native type + sums: BlockedVec, + + /// Track nulls in the input / filters + null_state: BlockedNullState, + + /// Function that computes the final average (value / count) + avg_fn: F, + + _phantom: PhantomData<(I, O)>, +} + +impl BlockedAvgGroupsAccumulator +where + I: ArrowNumericType + Send, + O: ArrowNumericType + Send, + S: ArrowNumericType + Send, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + 'static, +{ + fn new( + sum_data_type: &DataType, + return_data_type: &DataType, + block_size: usize, + avg_fn: F, + ) -> Self { + Self { + return_data_type: return_data_type.clone(), + sum_data_type: sum_data_type.clone(), + counts: BlockedVec::new(block_size), + sums: BlockedVec::new(block_size), + null_state: BlockedNullState::new(block_size), + avg_fn, + _phantom: PhantomData, + } + } + + /// Removes the groups selected by `emit_to`, as `(counts, sums, nulls)` + /// of every emitted block. + fn take( + &mut self, + emit_to: BlockedEmitTo, + ) -> impl Iterator, Vec, Option)> + use + { + fn take( + v: &mut BlockedVec, + emit_to: BlockedEmitTo, + ) -> Vec> { + match emit_to { + BlockedEmitTo::All => v.take_all(), + BlockedEmitTo::NextBlock => v.take_next_block().into_iter().collect(), + BlockedEmitTo::First(n) => vec![v.take_first(n)], + } + } + let counts = take(&mut self.counts, emit_to); + let sums = take(&mut self.sums, emit_to); + let nulls = self.null_state.build(emit_to); + debug_assert_eq!(counts.len(), sums.len()); + debug_assert_eq!(counts.len(), nulls.len()); + counts + .into_iter() + .zip(sums) + .zip(nulls) + .map(|((counts, sums), nulls)| (counts, sums, nulls)) + } +} + +impl BlockedGroupsAccumulator for BlockedAvgGroupsAccumulator +where + I: ArrowNumericType + Send, + O: ArrowNumericType + Send, + S: ArrowNumericType + Send, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + 'static, +{ + fn block_size(&self) -> usize { + self.sums.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].as_primitive::(); + + let nulls = values.logical_nulls().filter(|n| n.null_count() > 0); + self.counts.update( + total_num_groups, + 0, + group_indices, + nulls.as_ref(), + opt_filter, + |count| *count += 1, + ); + self.null_state.accumulate( + &mut self.sums, + total_num_groups, + S::default_value(), + group_indices, + values, + opt_filter, + |sum, new_value| *sum = add_avg_sum::(*sum, new_value), + ); + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()> { + assert_eq!(values.len(), 2, "two arguments to merge_batch"); + // first batch is counts, second is partial sums + let partial_counts = values[0].as_primitive::(); + let partial_sums = values[1].as_primitive::(); + self.null_state.accumulate( + &mut self.counts, + total_num_groups, + 0, + group_indices, + partial_counts, + None, + |count, partial_count| *count += partial_count, + ); + self.null_state.accumulate( + &mut self.sums, + total_num_groups, + S::default_value(), + group_indices, + partial_sums, + None, + |sum, partial_sum| *sum = add_avg_sum::(*sum, partial_sum), + ); + Ok(()) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + let blocks: Vec<_> = self.take(emit_to).collect(); + blocks + .into_iter() + .map(|(counts, sums, nulls)| { + evaluate_avg_values::( + counts, + sums, + nulls, + &self.avg_fn, + &self.return_data_type, + ) + }) + .collect() + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + let blocks: Vec<_> = self.take(emit_to).collect(); + Ok(blocks + .into_iter() + .map(|(counts, sums, nulls)| { + avg_state_values::(counts, sums, nulls, &self.sum_data_type) + }) + .collect()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + avg_convert_to_state::(values, opt_filter, &self.sum_data_type) + } + + fn size(&self) -> usize { + self.counts.allocated_size() + self.sums.allocated_size() + self.null_state.size() + } +} + #[cfg(test)] mod tests { use super::*; + + #[test] + fn blocked_avg_across_blocks_with_nulls_and_merge() -> Result<()> { + let mut acc = BlockedAvgGroupsAccumulator::::new( + &DataType::Float64, + &DataType::Float64, + 4, + |sum: f64, count: u64| Ok(sum / count as f64), + ); + let groups = |flats: &[usize]| -> Vec { + flats + .iter() + .map(|&f| BlocksIndex::from_flat(f, 4)) + .collect() + }; + let values: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(3.0), + None, + Some(10.0), + ])); + acc.update_batch(&[values], &groups(&[0, 0, 1, 5]), None, 6)?; + // partial state for group 5: count 2, sum 20 + let counts: ArrayRef = Arc::new(UInt64Array::from(vec![2])); + let sums: ArrayRef = Arc::new(Float64Array::from(vec![20.0])); + acc.merge_batch(&[counts, sums], &groups(&[5]), 6)?; + + let out = acc.evaluate(BlockedEmitTo::All)?; + let out: Vec>> = out + .iter() + .map(|a| a.as_primitive::().iter().collect()) + .collect(); + assert_eq!( + out, + vec![vec![Some(2.0), None, None, None], vec![None, Some(10.0)]] + ); + Ok(()) + } use arrow::array::{ Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, DurationSecondArray, Float64Array, diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index ed2abe88644ae..8e69bfa22f294 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -35,14 +35,18 @@ use datafusion_common::{ stats::Precision, utils::expr::COUNT_STAR_EXPANSION, }; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupSelection, + Accumulator, AggregateUDFImpl, BlockedEmitTo, BlockedGroupSelection, + BlockedGroupsAccumulator, BlocksIndex, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility, WindowFunctionDefinition, expr::WindowFunction, function::{AccumulatorArgs, StateFieldsArgs}, utils::{AggregateOrderSensitivity, format_state_name}, }; -use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator; +use datafusion_functions_aggregate_common::aggregate::count_distinct::{ + BlockedPrimitiveDistinctCountGroupsAccumulator, + PrimitiveDistinctCountGroupsAccumulator, +}; use datafusion_functions_aggregate_common::aggregate::{ count_distinct::Bitmap65536DistinctCountAccumulator, count_distinct::Bitmap65536DistinctCountAccumulatorI16, @@ -54,6 +58,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 +397,40 @@ impl AggregateUDFImpl for Count { create_distinct_count_groups_accumulator(&args) } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + self.groups_accumulator_supported(args) + } + + fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + if !args.is_distinct { + return Ok(Box::new(BlockedCountGroupsAccumulator::new(block_size))); + } + macro_rules! distinct { + ($t:ty) => { + Ok(Box::new( + BlockedPrimitiveDistinctCountGroupsAccumulator::<$t>::new(block_size), + )) + }; + } + match args.expr_fields[0].data_type() { + DataType::Int8 => distinct!(Int8Type), + DataType::Int16 => distinct!(Int16Type), + DataType::Int32 => distinct!(Int32Type), + DataType::Int64 => distinct!(Int64Type), + DataType::UInt8 => distinct!(UInt8Type), + DataType::UInt16 => distinct!(UInt16Type), + DataType::UInt32 => distinct!(UInt32Type), + DataType::UInt64 => distinct!(UInt64Type), + data_type => not_impl_err!( + "BlockedGroupsAccumulator not supported for COUNT(DISTINCT) with {data_type}" + ), + } + } + fn reverse_expr(&self) -> ReversedUDAF { ReversedUDAF::Identical } @@ -824,6 +863,167 @@ 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); + + // Add one to each group's counter for each non null, non + // filtered value + self.counts.update( + total_num_groups, + 0, + 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 evaluate_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result> { + selection.validate_num_groups(self.counts.len())?; + // Build each output chunk in its own `Vec`, which becomes the array's + // buffer without copying + let block_size = self.counts.block_size(); + let mut remaining = selection.len(); + let mut arrays = Vec::with_capacity(remaining.div_ceil(block_size)); + let mut chunk = Vec::with_capacity(remaining.min(block_size)); + for index in selection.iter() { + chunk.push(self.counts.get(index)); + if chunk.len() == block_size { + remaining -= block_size; + let next = Vec::with_capacity(remaining.min(block_size)); + arrays.push(Self::counts_to_array(std::mem::replace(&mut chunk, next))); + } + } + if !chunk.is_empty() { + arrays.push(Self::counts_to_array(chunk)); + } + Ok(arrays) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + Ok(self + .take(emit_to) + .into_iter() + .map(|counts| vec![counts]) + .collect()) + } + + fn state_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result>> { + Ok(self + .evaluate_preserving(selection)? + .into_iter() + .map(|counts| vec![counts]) + .collect()) + } + + fn supports_state_preserving(&self) -> bool { + true + } + + 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 { @@ -995,6 +1195,204 @@ 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_preserving_reads_keep_state() -> Result<()> { + let mut acc = BlockedCountGroupsAccumulator::new(4); + let values: ArrayRef = Arc::new(Int32Array::from(vec![1; 7])); + acc.update_batch(&[values], &blocks_index(&[0, 1, 1, 5, 5, 5, 6], 4), None, 7)?; + + let all = acc.evaluate_preserving(BlockedGroupSelection::all(7, 4))?; + assert_eq!(blocked_values(&all), vec![vec![1, 2, 0, 0], vec![0, 3, 1]]); + + let indices = blocks_index(&[5, 0, 5], 4); + let selected = BlockedGroupSelection::try_from_indices(&indices, 7, 4)?; + let state = acc.state_preserving(selected)?; + assert_eq!(state.len(), 1); + assert_eq!(blocked_values(&state[0]), vec![vec![3, 1, 3]]); + + // state is unchanged + assert_eq!( + blocked_values(&acc.evaluate(BlockedEmitTo::All)?), + vec![vec![1, 2, 0, 0], vec![0, 3, 1]] + ); + 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/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs index 94814a802d971..50e5f26611c2a 100644 --- a/datafusion/functions-aggregate/src/min_max.rs +++ b/datafusion/functions-aggregate/src/min_max.rs @@ -42,7 +42,9 @@ use arrow::datatypes::{ TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; -use crate::min_max::min_max_bytes::MinMaxBytesAccumulator; +use crate::min_max::min_max_bytes::{ + BlockedMinMaxBytesAccumulator, MinMaxBytesAccumulator, +}; use crate::min_max::min_max_struct::MinMaxStructAccumulator; use datafusion_common::ScalarValue; use datafusion_expr::DistinctHandling; @@ -50,7 +52,7 @@ use datafusion_expr::{ Accumulator, AggregateUDFImpl, Documentation, SetMonotonicity, Signature, Volatility, function::AccumulatorArgs, }; -use datafusion_expr::{GroupsAccumulator, StatisticsArgs}; +use datafusion_expr::{BlockedGroupsAccumulator, GroupsAccumulator, StatisticsArgs}; use datafusion_macros::user_doc; use half::f16; use std::collections::VecDeque; @@ -278,6 +280,25 @@ impl AggregateUDFImpl for Max { ) } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + use DataType::*; + matches!( + args.return_field.data_type(), + Utf8 | LargeUtf8 | Utf8View | Binary | LargeBinary | BinaryView + ) + } + + fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + Ok(Box::new(BlockedMinMaxBytesAccumulator::new_max( + args.return_field.data_type().clone(), + block_size, + ))) + } + fn create_groups_accumulator( &self, args: AccumulatorArgs, @@ -577,6 +598,25 @@ impl AggregateUDFImpl for Min { ) } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + use DataType::*; + matches!( + args.return_field.data_type(), + Utf8 | LargeUtf8 | Utf8View | Binary | LargeBinary | BinaryView + ) + } + + fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + Ok(Box::new(BlockedMinMaxBytesAccumulator::new_min( + args.return_field.data_type().clone(), + block_size, + ))) + } + fn create_groups_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index 7d26143eb5790..29e1d4ec478b7 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -22,7 +22,11 @@ use arrow::array::{ use arrow::datatypes::DataType; use datafusion_common::hash_map::Entry; use datafusion_common::{HashMap, Result, internal_err}; -use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; +use datafusion_expr::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, EmitTo, GroupSelection, + GroupsAccumulator, +}; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::BlockedVec; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; use std::mem::size_of; use std::sync::Arc; @@ -69,93 +73,101 @@ impl MinMaxBytesAccumulator { num_values: usize, data_capacity: usize, ) -> Result { - let result: ArrayRef = match self.inner.data_type { - DataType::Utf8 => { - let mut builder = StringBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - // SAFETY: update_batch only accepts the configured input type. - Some(value) => builder.append_value(unsafe { - std::str::from_utf8_unchecked(value) - }), - } + build_bytes_array(&self.inner.data_type, min_maxes, num_values, data_capacity) + } +} + +/// Builds an array of `data_type` from the min/max values of groups. Shared +/// by the flat and blocked accumulators. +fn build_bytes_array<'a>( + data_type: &DataType, + min_maxes: impl Iterator>, + num_values: usize, + data_capacity: usize, +) -> Result { + let result: ArrayRef = match *data_type { + DataType::Utf8 => { + let mut builder = StringBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder + .append_value(unsafe { std::str::from_utf8_unchecked(value) }), } - Arc::new(builder.finish()) } - DataType::LargeUtf8 => { - let mut builder = - LargeStringBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - // SAFETY: update_batch only accepts the configured input type. - Some(value) => builder.append_value(unsafe { - std::str::from_utf8_unchecked(value) - }), - } + Arc::new(builder.finish()) + } + DataType::LargeUtf8 => { + let mut builder = + LargeStringBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder + .append_value(unsafe { std::str::from_utf8_unchecked(value) }), } - Arc::new(builder.finish()) } - DataType::Utf8View => { - let block_size = capacity_to_view_block_size(data_capacity); - let mut builder = StringViewBuilder::with_capacity(num_values) - .with_fixed_block_size(block_size); - for value in min_maxes { - match value { - None => builder.append_null(), - // SAFETY: update_batch only accepts the configured input type. - Some(value) => builder.append_value(unsafe { - std::str::from_utf8_unchecked(value) - }), - } + Arc::new(builder.finish()) + } + DataType::Utf8View => { + let block_size = capacity_to_view_block_size(data_capacity); + let mut builder = StringViewBuilder::with_capacity(num_values) + .with_fixed_block_size(block_size); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder + .append_value(unsafe { std::str::from_utf8_unchecked(value) }), } - Arc::new(builder.finish()) } - DataType::Binary => { - let mut builder = BinaryBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - Some(value) => builder.append_value(value), - } + Arc::new(builder.finish()) + } + DataType::Binary => { + let mut builder = BinaryBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), } - Arc::new(builder.finish()) } - DataType::LargeBinary => { - let mut builder = - LargeBinaryBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - Some(value) => builder.append_value(value), - } + Arc::new(builder.finish()) + } + DataType::LargeBinary => { + let mut builder = + LargeBinaryBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), } - Arc::new(builder.finish()) } - DataType::BinaryView => { - let block_size = capacity_to_view_block_size(data_capacity); - let mut builder = BinaryViewBuilder::with_capacity(num_values) - .with_fixed_block_size(block_size); - for value in min_maxes { - match value { - None => builder.append_null(), - Some(value) => builder.append_value(value), - } + Arc::new(builder.finish()) + } + DataType::BinaryView => { + let block_size = capacity_to_view_block_size(data_capacity); + let mut builder = BinaryViewBuilder::with_capacity(num_values) + .with_fixed_block_size(block_size); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), } - Arc::new(builder.finish()) } - _ => { - return internal_err!( - "Unexpected data type for MinMaxBytesAccumulator: {:?}", - self.inner.data_type - ); - } - }; + Arc::new(builder.finish()) + } + _ => { + return internal_err!( + "Unexpected data type for MinMaxBytesAccumulator: {:?}", + *data_type + ); + } + }; - assert_eq!(&self.inner.data_type, result.data_type()); - Ok(result) - } + assert_eq!(data_type, result.data_type()); + Ok(result) } impl GroupsAccumulator for MinMaxBytesAccumulator { @@ -173,126 +185,13 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { // apply filter if needed let array = apply_filter_as_nulls(array, opt_filter)?; - // dispatch to appropriate kernel / specialized implementation - fn string_min(a: &[u8], b: &[u8]) -> bool { - // safety: only called from this function, which ensures a and b come - // from an array with valid utf8 data - unsafe { - let a = std::str::from_utf8_unchecked(a); - let b = std::str::from_utf8_unchecked(b); - a < b - } - } - fn string_max(a: &[u8], b: &[u8]) -> bool { - // safety: only called from this function, which ensures a and b come - // from an array with valid utf8 data - unsafe { - let a = std::str::from_utf8_unchecked(a); - let b = std::str::from_utf8_unchecked(b); - a > b - } - } - fn binary_min(a: &[u8], b: &[u8]) -> bool { - a < b - } - - fn binary_max(a: &[u8], b: &[u8]) -> bool { - a > b - } - - fn str_to_bytes<'a>( - it: impl Iterator>, - ) -> impl Iterator> { - it.map(|s| s.map(|s| s.as_bytes())) - } - - match (self.is_min, &self.inner.data_type) { - // Utf8/LargeUtf8/Utf8View Min - (true, &DataType::Utf8) => self.inner.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_min, - ), - (true, &DataType::LargeUtf8) => self.inner.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_min, - ), - (true, &DataType::Utf8View) => self.inner.update_batch( - str_to_bytes(array.as_string_view().iter()), - group_indices, - total_num_groups, - string_min, - ), - - // Utf8/LargeUtf8/Utf8View Max - (false, &DataType::Utf8) => self.inner.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_max, - ), - (false, &DataType::LargeUtf8) => self.inner.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_max, - ), - (false, &DataType::Utf8View) => self.inner.update_batch( - str_to_bytes(array.as_string_view().iter()), - group_indices, - total_num_groups, - string_max, - ), - - // Binary/LargeBinary/BinaryView Min - (true, &DataType::Binary) => self.inner.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_min, - ), - (true, &DataType::LargeBinary) => self.inner.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_min, - ), - (true, &DataType::BinaryView) => self.inner.update_batch( - array.as_binary_view().iter(), - group_indices, - total_num_groups, - binary_min, - ), - - // Binary/LargeBinary/BinaryView Max - (false, &DataType::Binary) => self.inner.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_max, - ), - (false, &DataType::LargeBinary) => self.inner.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_max, - ), - (false, &DataType::BinaryView) => self.inner.update_batch( - array.as_binary_view().iter(), - group_indices, - total_num_groups, - binary_max, - ), - - _ => internal_err!( - "Unexpected combination for MinMaxBytesAccumulator: ({:?}, {:?})", - self.is_min, - self.inner.data_type - ), - } + dispatch_update( + &mut self.inner, + self.is_min, + &array, + group_indices, + total_num_groups, + ) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { @@ -363,6 +262,159 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { } } +/// Updates the min/max values of a flat or blocked state, see +/// [`MinMaxBytesUpdate`]. +fn dispatch_update( + state: &mut S, + is_min: bool, + array: &ArrayRef, + group_indices: &[Idx], + total_num_groups: usize, +) -> Result<()> +where + S: MinMaxBytesUpdate, +{ + let data_type = state.data_type().clone(); + // dispatch to appropriate kernel / specialized implementation + fn string_min(a: &[u8], b: &[u8]) -> bool { + // safety: only called from this function, which ensures a and b come + // from an array with valid utf8 data + unsafe { + let a = std::str::from_utf8_unchecked(a); + let b = std::str::from_utf8_unchecked(b); + a < b + } + } + fn string_max(a: &[u8], b: &[u8]) -> bool { + // safety: only called from this function, which ensures a and b come + // from an array with valid utf8 data + unsafe { + let a = std::str::from_utf8_unchecked(a); + let b = std::str::from_utf8_unchecked(b); + a > b + } + } + fn binary_min(a: &[u8], b: &[u8]) -> bool { + a < b + } + + fn binary_max(a: &[u8], b: &[u8]) -> bool { + a > b + } + + fn str_to_bytes<'a>( + it: impl Iterator>, + ) -> impl Iterator> { + it.map(|s| s.map(|s| s.as_bytes())) + } + + match (is_min, &data_type) { + // Utf8/LargeUtf8/Utf8View Min + (true, &DataType::Utf8) => state.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_min, + ), + (true, &DataType::LargeUtf8) => state.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_min, + ), + (true, &DataType::Utf8View) => state.update_batch( + str_to_bytes(array.as_string_view().iter()), + group_indices, + total_num_groups, + string_min, + ), + + // Utf8/LargeUtf8/Utf8View Max + (false, &DataType::Utf8) => state.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_max, + ), + (false, &DataType::LargeUtf8) => state.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_max, + ), + (false, &DataType::Utf8View) => state.update_batch( + str_to_bytes(array.as_string_view().iter()), + group_indices, + total_num_groups, + string_max, + ), + + // Binary/LargeBinary/BinaryView Min + (true, &DataType::Binary) => state.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_min, + ), + (true, &DataType::LargeBinary) => state.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_min, + ), + (true, &DataType::BinaryView) => state.update_batch( + array.as_binary_view().iter(), + group_indices, + total_num_groups, + binary_min, + ), + + // Binary/LargeBinary/BinaryView Max + (false, &DataType::Binary) => state.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_max, + ), + (false, &DataType::LargeBinary) => state.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_max, + ), + (false, &DataType::BinaryView) => state.update_batch( + array.as_binary_view().iter(), + group_indices, + total_num_groups, + binary_max, + ), + + _ => internal_err!( + "Unexpected combination for MinMaxBytesAccumulator: ({:?}, {:?})", + is_min, + data_type + ), + } +} + +/// Min/max state updated by [`dispatch_update`]. +trait MinMaxBytesUpdate { + fn data_type(&self) -> &DataType; + + /// `cmp` is called like `cmp(new_val, existing_val)` and returns true if + /// `new_val` should replace `existing_val` + fn update_batch<'a, F, I>( + &mut self, + iter: I, + group_indices: &[Idx], + total_num_groups: usize, + cmp: F, + ) -> Result<()> + where + F: FnMut(&[u8], &[u8]) -> bool + Send + Sync, + I: IntoIterator>; +} + /// Returns the block size in (contiguous buffer size) to use /// for a given data capacity (total string length) /// @@ -453,6 +505,42 @@ impl MinMaxBytesState { } } + /// Emits the specified min_max values + /// + /// Returns (data_capacity, min_maxes), updating the current value of total_data_bytes + /// + /// - `data_capacity`: the total length of all strings and their contents, + /// - `min_maxes`: the actual min/max values for each group + fn emit_to(&mut self, emit_to: EmitTo) -> (usize, Vec>>) { + match emit_to { + EmitTo::All => { + ( + std::mem::take(&mut self.total_data_bytes), // reset total bytes and min_max + std::mem::take(&mut self.min_max), + ) + } + EmitTo::First(n) => { + let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); + let first_data_capacity: usize = first_min_maxes + .iter() + .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) + .sum(); + self.total_data_bytes -= first_data_capacity; + (first_data_capacity, first_min_maxes) + } + } + } + + fn size(&self) -> usize { + self.total_data_bytes + self.min_max.len() * size_of::>>() + } +} + +impl MinMaxBytesUpdate for MinMaxBytesState { + fn data_type(&self) -> &DataType { + &self.data_type + } + /// Updates the min/max values for the given string values /// /// `cmp` is the comparison function to use, called like `cmp(new_val, existing_val)` @@ -506,41 +594,272 @@ impl MinMaxBytesState { Ok(()) } +} - /// Emits the specified min_max values - /// - /// Returns (data_capacity, min_maxes), updating the current value of total_data_bytes - /// - /// - `data_capacity`: the total length of all strings and their contents, - /// - `min_maxes`: the actual min/max values for each group - fn emit_to(&mut self, emit_to: EmitTo) -> (usize, Vec>>) { - match emit_to { - EmitTo::All => { - ( - std::mem::take(&mut self.total_data_bytes), // reset total bytes and min_max - std::mem::take(&mut self.min_max), - ) +/// [`MinMaxBytesState`] with the value of each group stored in a +/// [`BlockedVec`]. +/// +/// Reading a stored value is a load from the heap after the block and the +/// slot, which the address-resolve loops of [`BlockedVec`] cannot hide. +/// [`Self::update_batch`] follows the same algorithm as +/// [`MinMaxBytesState::update_batch`]. +#[derive(Debug)] +struct BlockedMinMaxBytesState { + /// The minimum/maximum value for each group + min_max: BlockedVec>>, + data_type: DataType, + /// The total bytes of the string data (for pre-allocating the final array, + /// and tracking memory usage) + total_data_bytes: usize, +} + +impl BlockedMinMaxBytesState { + fn new(data_type: DataType, block_size: usize) -> Self { + Self { + min_max: BlockedVec::new(block_size), + data_type, + total_data_bytes: 0, + } + } + + /// Removes the groups selected by `emit_to`, one `Vec` per block. + fn emit_to(&mut self, emit_to: BlockedEmitTo) -> Vec>>> { + let blocks = match emit_to { + BlockedEmitTo::All => self.min_max.take_all(), + BlockedEmitTo::NextBlock => { + self.min_max.take_next_block().into_iter().collect() } - EmitTo::First(n) => { - let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); - let first_data_capacity: usize = first_min_maxes - .iter() - .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) - .sum(); - self.total_data_bytes -= first_data_capacity; - (first_data_capacity, first_min_maxes) + BlockedEmitTo::First(n) => vec![self.min_max.take_first(n)], + }; + for value in blocks.iter().flatten() { + self.total_data_bytes -= value.as_ref().map_or(0, Vec::len); + } + blocks + } + + fn size(&self) -> usize { + self.min_max.allocated_size() + self.total_data_bytes + } +} + +impl MinMaxBytesUpdate for BlockedMinMaxBytesState { + fn data_type(&self) -> &DataType { + &self.data_type + } + + fn update_batch<'a, F, I>( + &mut self, + iter: I, + group_indices: &[BlocksIndex], + total_num_groups: usize, + mut cmp: F, + ) -> Result<()> + where + F: FnMut(&[u8], &[u8]) -> bool + Send + Sync, + I: IntoIterator>, + { + self.min_max.grow_to(total_num_groups, None); + // Minimize value copies by calculating the new min/maxes for each group + // in this batch (either the existing min/max or the new input value) + // and updating the owned values in `self.min_max` at most once + let mut locations = + HashMap::::with_capacity(group_indices.len()); + + // Figure out the new min value for each group + for (new_val, &group_index) in iter.into_iter().zip(group_indices) { + let Some(new_val) = new_val else { + continue; // skip nulls + }; + + match locations.entry(group_index) { + Entry::Occupied(mut occupied_entry) => { + if cmp(new_val, occupied_entry.get()) { + occupied_entry.insert(new_val); + } + } + Entry::Vacant(vacant_entry) => { + if let Some(old_val) = self.min_max.get_ref(group_index) { + if cmp(new_val, old_val) { + vacant_entry.insert(new_val); + } + } else { + vacant_entry.insert(new_val); + } + } + } + } + + // Update `self.min_max` with any new min/max values found in the input + for (group_index, new_val) in locations { + match self.min_max.get_mut(group_index) { + existing @ None => { + *existing = Some(new_val.to_vec()); + self.total_data_bytes += new_val.len(); + } + Some(existing_val) => { + // Copy data over to avoid re-allocating + self.total_data_bytes -= existing_val.len(); + self.total_data_bytes += new_val.len(); + existing_val.clear(); + existing_val.extend_from_slice(new_val); + } } } + Ok(()) + } +} + +/// [`MinMaxBytesAccumulator`] with state stored in blocks. +#[derive(Debug)] +pub(crate) struct BlockedMinMaxBytesAccumulator { + inner: BlockedMinMaxBytesState, + /// if true, is `MIN` otherwise is `MAX` + is_min: bool, +} + +impl BlockedMinMaxBytesAccumulator { + /// Create a new accumulator for computing `min(val)` + pub fn new_min(data_type: DataType, block_size: usize) -> Self { + Self { + inner: BlockedMinMaxBytesState::new(data_type, block_size), + is_min: true, + } + } + + /// Create a new accumulator for computing `max(val)` + pub fn new_max(data_type: DataType, block_size: usize) -> Self { + Self { + inner: BlockedMinMaxBytesState::new(data_type, block_size), + is_min: false, + } + } +} + +impl BlockedGroupsAccumulator for BlockedMinMaxBytesAccumulator { + fn block_size(&self) -> usize { + self.inner.min_max.block_size() + } + + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + let array = &values[0]; + assert_eq!(array.len(), group_indices.len()); + assert_eq!(array.data_type(), &self.inner.data_type); + + // apply filter if needed + let array = apply_filter_as_nulls(array, opt_filter)?; + + dispatch_update( + &mut self.inner, + self.is_min, + &array, + group_indices, + total_num_groups, + ) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()> { + // min/max are their own states (no transition needed) + self.update_batch(values, group_indices, None, total_num_groups) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + let data_type = self.inner.data_type.clone(); + self.inner + .emit_to(emit_to) + .into_iter() + .map(|block| { + let data_capacity = block + .iter() + .map(|value| value.as_ref().map_or(0, Vec::len)) + .sum(); + build_bytes_array( + &data_type, + block.iter().map(|value| value.as_deref()), + block.len(), + data_capacity, + ) + }) + .collect() + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + // min/max are their own states (no transition needed) + Ok(self + .evaluate(emit_to)? + .into_iter() + .map(|array| vec![array]) + .collect()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + // Min/max do not change the values as they are their own states + // apply the filter by combining with the null mask, if any + let output = apply_filter_as_nulls(&values[0], opt_filter)?; + Ok(vec![output]) } fn size(&self) -> usize { - self.total_data_bytes + self.min_max.len() * size_of::>>() + self.inner.size() } } #[cfg(test)] mod tests { use super::*; + + #[test] + fn blocked_min_max_bytes_across_blocks() -> Result<()> { + let long = "z".repeat(30); + let indices = |flats: &[usize]| -> Vec { + flats + .iter() + .map(|&f| BlocksIndex::from_flat(f, 4)) + .collect() + }; + let mut acc = BlockedMinMaxBytesAccumulator::new_max(DataType::Utf8, 4); + let values: ArrayRef = Arc::new(StringArray::from(vec![ + Some("a"), + Some("c"), + None, + Some(long.as_str()), + Some("b"), + ])); + acc.update_batch(&[values], &indices(&[0, 0, 1, 5, 5]), None, 6)?; + assert!(acc.size() >= long.len(), "values are accounted"); + // second batch: a smaller value must not replace, a larger one must + let values: ArrayRef = Arc::new(StringArray::from(vec!["b", "d"])); + acc.update_batch(&[values], &indices(&[0, 2]), None, 6)?; + + let out = acc.evaluate(BlockedEmitTo::All)?; + let out: Vec>> = out + .iter() + .map(|a| a.as_string::().iter().collect()) + .collect(); + assert_eq!( + out, + vec![ + vec![Some("c"), None, Some("d"), None], + vec![None, Some(long.as_str())], + ] + ); + assert_eq!(acc.inner.total_data_bytes, 0); + Ok(()) + } use arrow::array::StringArray; #[test] diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs index 1228b458a82b9..9e7a67a9ce1d3 100644 --- a/datafusion/functions-aggregate/src/sum.rs +++ b/datafusion/functions-aggregate/src/sum.rs @@ -40,10 +40,12 @@ use datafusion_expr::expr_fn::cast; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, GroupsAccumulator, + Accumulator, AggregateUDFImpl, BlockedGroupsAccumulator, Coercion, Documentation, + Expr, GroupsAccumulator, Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility, }; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_prim_op::BlockedPrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator; use datafusion_macros::user_doc; @@ -317,6 +319,27 @@ impl AggregateUDFImpl for Sum { downcast_sum!(args, helper) } + fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + self.groups_accumulator_supported(args) + } + + fn create_blocked_groups_accumulator( + &self, + args: AccumulatorArgs, + block_size: usize, + ) -> Result> { + macro_rules! helper { + ($t:ty, $dt:expr) => { + Ok(Box::new(BlockedPrimitiveGroupsAccumulator::<$t, _>::new( + &$dt, + block_size, + |x, y| *x = x.add_wrapping(y), + ))) + }; + } + downcast_sum!(args, helper) + } + fn create_sliding_accumulator( &self, args: AccumulatorArgs, 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..3fad1c3bdf5f9 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 @@ -160,15 +184,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 +208,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 +230,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 +285,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 +299,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 +311,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 +333,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 +374,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 +384,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 +393,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 +443,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 +469,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 +501,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 +523,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 +556,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 +576,7 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator; pub(super) type AggregateBatchFn = fn( &mut AggregateAccumulator, &CompactedAccumulatorArgs, - &[usize], + GroupIndices<'_>, usize, ) -> Result<()>; @@ -565,8 +586,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 { @@ -609,14 +633,11 @@ 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 +646,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 +676,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 +725,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 +761,7 @@ impl HashAggregateAccumulator { aggregate_expr: Arc, arguments: Vec>, filter: Option>, - accumulator: Box, + accumulator: AccumulatorStorage, submetrics: Arc, ) -> Self { Self { @@ -727,11 +773,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 +888,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 +976,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 +1001,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 +1069,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 +1152,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 fd9e92635fc65..d95c6405f8fea 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 @@ -23,16 +23,15 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; use datafusion_common::assert_or_internal_err; -use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_common::{Result, internal_err}; use datafusion_expr::{AggregateMetrics, EmitTo}; use crate::InputOrderMode; 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::{ @@ -43,7 +42,10 @@ use crate::aggregates::{ use super::AggregateTableMetrics; use super::common::{ AggregateAccumulator, AggregateBatchFn, AggregateHashTable, EvaluatedAggregateBatch, - MaterializeAccumulatorFn, create_group_accumulator, + HashAggregateAccumulator, MaterializeAccumulatorFn, create_accumulator_storage, +}; +use super::storage::{ + GroupKeys, MaterializedBatch, block_size, materialize_batches, use_blocked_keys, }; #[derive(Clone)] @@ -164,11 +166,9 @@ pub(super) struct OrderedAggregateTableBuffer { /// Tracks how far ordered input allows this table to drain safely. pub(super) group_ordering: GroupOrdering, - /// Interned group keys, in the same group-id order used by accumulators. - pub(super) group_values: Box, - - /// Scratch group id vector for the current input batch. - pub(super) group_indices: Vec, + /// Interned group keys, in the same group-id order used by accumulators, + /// and the scratch group id vector for the current input batch. + pub(super) keys: GroupKeys, /// One item per aggregate expression. /// @@ -201,21 +201,24 @@ impl OrderedAggregateTable { let group_ordering = GroupOrdering::try_new(input_order_mode)?; let group_schema = agg.group_by().group_schema(input_schema)?; - let group_values = new_group_values(group_schema, &group_ordering)?; + let block_size = block_size(batch_size); let aggregate_arguments = aggregate_expressions( agg.aggr_expr(), aggregate_mode, agg.group_by().num_group_exprs(), )?; - let accumulators = agg + 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(AggregateAccumulator::new( Arc::clone(agg_expr), arguments, @@ -225,6 +228,16 @@ impl OrderedAggregateTable { )) }) .collect::>()?; + let keys = GroupKeys::try_new( + Arc::clone(&group_schema), + &group_ordering, + use_blocked_keys(agg, &group_schema, block_size), + block_size, + &accumulators + .iter() + .map(|acc| acc.storage()) + .collect::>(), + )?; Ok(Self { output_schema, @@ -237,8 +250,7 @@ impl OrderedAggregateTable { buffer: OrderedAggregateTableBuffer { group_by: Arc::clone(agg.group_by()), group_ordering, - group_values, - group_indices: vec![], + keys, accumulators, }, _mode: PhantomData, @@ -292,7 +304,7 @@ impl OrderedAggregateTable { /// Number of groups currently buffered. pub(in crate::aggregates) fn num_groups(&self) -> usize { - self.buffer.group_values.len() + self.buffer.keys.len() } /// Check if there is zero groups accumulated so far. @@ -307,9 +319,8 @@ impl OrderedAggregateTable { .iter() .map(|acc| acc.size()) .sum::() - + self.buffer.group_values.size() + + self.buffer.keys.size() + self.buffer.group_ordering.size() - + self.buffer.group_indices.allocated_size() } pub(in crate::aggregates) fn metrics(&self) -> OrderedAggregateTableMetrics { @@ -328,38 +339,37 @@ impl OrderedAggregateTable { /// active (incomplete) groups. Partial aggregation can pass those states to /// its final stage, while single and final aggregation sort and spill them /// before replay. - 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> { - if self.buffer.group_values.is_empty() { - return Ok(None); + ) -> Result> { + if self.buffer.keys.is_empty() { + return Ok(vec![]); } let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - 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), - )?); - } - Ok::<_, datafusion_common::DataFusionError>(output) + let batches = self.group_by_metrics.time_emitting(|| { + materialize_batches( + &mut self.buffer.keys, + &mut self.buffer.accumulators, + EmitTo::All, + HashAggregateAccumulator::state, + AccumulatorPhase::State, + &accumulator_metrics, + &self.state_schema, + ) })?; + debug_assert!(batches.iter().all(|b| b.batch.num_rows() > 0)); - let batch = RecordBatch::try_new(Arc::clone(&self.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 passed downstream or sorted for spilling. - self.buffer.group_values.clear_shrink(0); - self.buffer.group_indices.clear(); - self.buffer.group_indices.shrink_to_fit(); + self.buffer.keys.clear_shrink(0); self.buffer.group_ordering.reset(); - Ok(Some(batch)) + Ok(batches) } /// Aggregates one evaluated input batch after selecting the mode-specific @@ -379,16 +389,13 @@ impl OrderedAggregateTable { let group_by_metrics = self.group_by_metrics.clone(); for group_values in &evaluated_batch.grouping_set_args { let total_num_groups = group_by_metrics.time_group_key_preparation(|| { - let starting_num_groups = self.buffer.group_values.len(); - self.buffer - .group_values - .intern(group_values, &mut self.buffer.group_indices)?; - let total_num_groups = self.buffer.group_values.len(); + let starting_num_groups = self.buffer.keys.len(); + self.buffer.keys.intern(group_values)?; + let total_num_groups = self.buffer.keys.len(); if total_num_groups > starting_num_groups { - self.buffer.group_ordering.new_groups( + self.buffer.keys.record_new_groups( + &mut self.buffer.group_ordering, group_values, - &self.buffer.group_indices, - total_num_groups, )?; } Ok::<_, datafusion_common::DataFusionError>(total_num_groups) @@ -402,13 +409,9 @@ impl OrderedAggregateTable { .zip(evaluated_batch.accumulator_args.iter()) .enumerate() { + let group_indices = self.buffer.keys.indices_for(acc.storage()); accumulator_metrics.time(idx, accumulator_phase, || { - aggregate_fn( - acc, - values, - &self.buffer.group_indices, - total_num_groups, - ) + aggregate_fn(acc, values, group_indices, total_num_groups) })?; } Ok::<(), datafusion_common::DataFusionError>(()) @@ -430,7 +433,7 @@ impl OrderedAggregateTable { materialize_accumulator_fn: MaterializeAccumulatorFn, accumulator_phase: AccumulatorPhase, ) -> Result> { - if self.buffer.group_values.is_empty() { + if self.buffer.keys.is_empty() { return Ok(None); } @@ -444,21 +447,32 @@ impl OrderedAggregateTable { } EmitTo::All => EmitTo::All, }; - self.materialize_groups(emit_to, materialize_accumulator_fn, accumulator_phase) - .map(Some) + // `emit_to` is capped at `batch_size`, at most one block, so this is + // one batch. + let mut batches = self.materialize_groups( + emit_to, + materialize_accumulator_fn, + accumulator_phase, + )?; + let Some(batch) = batches.pop().filter(|_| batches.is_empty()) else { + return internal_err!("{emit_to:?} must materialize exactly one batch"); + }; + Ok(Some(batch.batch)) } /// Removes the selected groups once and materializes their output columns. /// The caller chooses the completed prefix and any output-size limit. + /// + /// Returns one batch if every part is flat, otherwise one batch per + /// emitted block, see [`materialize_batches`]. pub(super) fn materialize_groups( &mut self, emit_to: EmitTo, materialize_accumulator_fn: MaterializeAccumulatorFn, accumulator_phase: AccumulatorPhase, - ) -> Result { + ) -> Result> { let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - let output = self.group_by_metrics.time_emitting(|| { - let mut output = self.buffer.group_values.emit(emit_to)?; + let batches = self.group_by_metrics.time_emitting(|| { // EOF can also emit a prefix when a caller limits its batch size, // but the completed ordering state no longer tracks group indexes. if let EmitTo::First(n) = emit_to @@ -467,19 +481,18 @@ impl OrderedAggregateTable { self.buffer.group_ordering.remove_groups(n); } - for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( - idx, - accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), - )?); - } - Ok::<_, datafusion_common::DataFusionError>(output) + materialize_batches( + &mut self.buffer.keys, + &mut self.buffer.accumulators, + emit_to, + materialize_accumulator_fn, + accumulator_phase, + &accumulator_metrics, + &self.output_schema, + ) })?; + debug_assert!(batches.iter().all(|b| b.batch.num_rows() > 0)); - let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; - debug_assert!(batch.num_rows() > 0); - - Ok(batch) + Ok(batches) } } 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/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index f70fc0bf0d961..621319750164c 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -31,6 +31,7 @@ use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorP use super::common::HashAggregateAccumulator; use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; +use super::storage::MaterializedBatch; /// Implementation specific to final aggregation, where the table stores partial /// aggregate states and the input rows are also partial states. @@ -83,21 +84,21 @@ impl OrderedAggregateTable { /// Materializes final results for all groups proven complete by the input /// ordering, leaving the active ordered-key range in the table. /// - /// Returns None if there are no completed groups. - pub(in crate::aggregates) fn take_completed_result_batch( + /// Returns no batch if there are no completed groups, and one batch per + /// emitted chunk for blocked storage. + pub(in crate::aggregates) fn take_completed_result_batches( &mut self, - ) -> Result> { + ) -> Result> { if self.is_empty() { - return Ok(None); + return Ok(vec![]); } let Some(emit_to) = self.group_ordering().emit_to() else { - return Ok(None); + return Ok(vec![]); }; self.materialize_groups( emit_to, HashAggregateAccumulator::evaluate_to_columns, AccumulatorPhase::Evaluate, ) - .map(Some) } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index 104855c5f0d8f..011ddf518691e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -42,6 +42,7 @@ use crate::aggregates::{ use super::common::HashAggregateAccumulator; use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; +use super::storage::MaterializedBatch; /// Implementation specific to partial aggregation, where the table stores /// partial aggregate states and the input rows are raw rows. @@ -91,20 +92,19 @@ impl OrderedAggregateTable { /// Materializes all groups proven complete by the input ordering, leaving /// the active ordered-key range in the table. - pub(in crate::aggregates) fn take_completed_state_batch( + pub(in crate::aggregates) fn take_completed_state_batches( &mut self, - ) -> Result> { + ) -> Result> { if self.is_empty() { - return Ok(None); + return Ok(vec![]); } let Some(emit_to) = self.group_ordering().emit_to() else { - return Ok(None); + return Ok(vec![]); }; self.materialize_groups( emit_to, HashAggregateAccumulator::state, AccumulatorPhase::State, ) - .map(Some) } } 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..c6a58c8b57a95 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs @@ -0,0 +1,595 @@ +// 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, and a table with no blocked part 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}; + +/// Smallest block, in groups. Tables with fewer groups stay in a single, +/// flat block and run exactly the flat hot loops. +const MIN_BLOCK_SIZE: usize = 1 << 18; + +/// Number of groups per block for every blocked part of a table. +/// +/// At least `batch_size`, so an output batch never spans two blocks. +pub(in crate::aggregates) fn block_size(batch_size: usize) -> usize { + batch_size.max(MIN_BLOCK_SIZE).next_power_of_two() +} + +/// 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) + } + } + } + + /// Tells `group_ordering` about the groups added by the last interned + /// batch `cols`, see [`GroupOrdering::new_groups`]. + pub(super) fn record_new_groups( + &self, + group_ordering: &mut GroupOrdering, + cols: &[ArrayRef], + ) -> Result<()> { + match &self.storage { + KeyStorage::Flat(_) => { + group_ordering.new_groups(cols, &self.flat_indices, self.len()) + } + KeyStorage::Blocked(_) => group_ordering.new_blocked_groups( + cols, + &self.blocked_indices, + self.block_size, + self.len(), + ), + } + } + + 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), MIN_BLOCK_SIZE); + // The block always holds at least one output batch. + assert_eq!(block_size(MIN_BLOCK_SIZE + 1), MIN_BLOCK_SIZE * 2); + 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 num_groups = MIN_BLOCK_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, 8192)?; + table.aggregate_batch(&input)?; + + let batches = table.take_state_batches()?; + assert_eq!( + batches + .iter() + .map(|b| b.batch.num_rows()) + .collect::>(), + vec![MIN_BLOCK_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 num_groups = MIN_BLOCK_SIZE + 10; + let batch_size = 8192; + 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::>()); + // While the first block is sliced, the table still holds the second + // block; once the first block is handed out, only the small second + // block (10 groups) remains. + let first_block_batches = MIN_BLOCK_SIZE / batch_size; + assert!(memory_per_batch[0] <= memory_while_building); + assert!( + memory_per_batch[first_block_batches] < memory_while_building / 4, + "memory after the first block: {}, while building: {memory_while_building}", + memory_per_batch[first_block_batches] + ); + 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..2371b0d78d7ab --- /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 +/// [`Self::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 80e6b335e3188..a9808507d5e74 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -28,6 +28,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 2ef7a6463f933..fa820725bd2de 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; @@ -313,20 +313,32 @@ 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 mut pending_memory: usize = materialized_group_states + .iter() + .map(|b| b.memory_size) + .sum(); + for batch in materialized_group_states { + pending_memory -= batch.memory_size; + self.emit_on_memory_pressure( + batch, + &mut emitter, + hash_table.memory_size() + pending_memory, + ) + .await?; + } } } } @@ -431,11 +443,14 @@ 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(); + let MaterializedBatch { + batch: mut remaining_groups, + memory_size: remaining_groups_memory, + } = remaining_groups; // Emitting clears the aggregate table and releases its // accumulated memory. Update the reservation accordingly. @@ -751,7 +766,7 @@ impl FinalHashAggregateStream { // Go to the next state to perform spilling the aggregated // groups so far. let result = hash_table - .take_state_batch() + .take_state_batches() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -790,7 +805,7 @@ impl FinalHashAggregateStream { // Input was exhausted after spilling. Spill the last in-memory run hash_table - .take_state_batch() + .take_state_batches() .and_then(|batch| spill_context.sort_and_spill(batch))?; // Construct the ordered input used to merge all spill files. diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 147e4e0d6f18d..8968f29d6fd5f 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -19,7 +19,7 @@ use std::mem::size_of; use arrow::array::ArrayRef; use datafusion_common::Result; -use datafusion_expr::EmitTo; +use datafusion_expr::{BlocksIndex, EmitTo}; mod full; mod partial; @@ -147,6 +147,32 @@ impl GroupOrdering { Ok(()) } + /// [`Self::new_groups`] for blocked group indices, with `block_size` + /// groups per block. + pub(crate) fn new_blocked_groups( + &mut self, + batch_group_values: &[ArrayRef], + group_indices: &[BlocksIndex], + block_size: usize, + total_num_groups: usize, + ) -> Result<()> { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => { + partial.new_blocked_groups( + batch_group_values, + group_indices, + block_size, + total_num_groups, + )?; + } + GroupOrdering::Full(full) => { + full.new_groups(total_num_groups); + } + } + Ok(()) + } + /// Returns the size of memory used by the ordering state, in bytes. pub fn size(&self) -> usize { size_of::() diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 1603bb6d079be..0038501aa5479 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -25,7 +25,7 @@ use arrow_ord::partition::partition; use datafusion_common::utils::{compare_rows, get_row_at_idx}; use datafusion_common::{Result, ScalarValue}; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::EmitTo; +use datafusion_expr::{BlocksIndex, EmitTo}; /// Tracks grouping state when the data is ordered by some subset of /// the group keys. @@ -216,6 +216,34 @@ impl GroupOrderingPartial { batch_group_values: &[ArrayRef], group_indices: &[usize], total_num_groups: usize, + ) -> Result<()> { + self.new_groups_with(batch_group_values, total_num_groups, |row| { + group_indices[row] + }) + } + + /// [`Self::new_groups`] for blocked group indices, with `block_size` + /// groups per block. + pub(crate) fn new_blocked_groups( + &mut self, + batch_group_values: &[ArrayRef], + group_indices: &[BlocksIndex], + block_size: usize, + total_num_groups: usize, + ) -> Result<()> { + self.new_groups_with(batch_group_values, total_num_groups, |row| { + group_indices[row].flat(block_size) + }) + } + + /// Shared implementation of [`Self::new_groups`] and + /// [`Self::new_blocked_groups`]. `group_position(row)` is the position, + /// counting from the first group, of the group of `row`. + fn new_groups_with( + &mut self, + batch_group_values: &[ArrayRef], + total_num_groups: usize, + group_position: impl Fn(usize) -> usize, ) -> Result<()> { assert!(total_num_groups > 0); assert!(!batch_group_values.is_empty()); @@ -242,7 +270,7 @@ impl GroupOrderingPartial { let ranges = partition(&sort_keys)?.ranges(); let last_range = ranges.last().unwrap(); - let range_current_sort = group_indices[last_range.start]; + let range_current_sort = group_position(last_range.start); let range_sort_key = get_row_at_idx(&sort_keys, last_range.start)?; let (current_sort, sort_key) = if last_range.start == 0 { diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 4a70ef2df6909..1fcafc9102a27 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -28,7 +28,7 @@ use futures::stream::StreamExt; use super::AggregateExec; use super::aggregate_hash_table::{ - FinalMarker, OrderedAggregateTable, OrderedAggregateTableMetrics, + FinalMarker, MaterializedBatch, OrderedAggregateTable, OrderedAggregateTableMetrics, }; use super::spill::AggregateSpill; use crate::aggregates::AggregateMode; @@ -76,8 +76,9 @@ struct Aggregating { } struct Outputting { - /// Materialized final results, emitted in slices of `batch_size` rows. - batch: RecordBatch, + /// Materialized final results: one batch for flat storage, one per chunk + /// for blocked storage. Each batch is emitted in slices of `batch_size` rows. + batches: Vec, /// Aggregation stage to resume after output; `None` after EOF. resume: Option, } @@ -338,7 +339,7 @@ impl Aggregating { if self.table.is_empty() { return Err(oom); } - spill_context.sort_and_spill(self.table.take_state_batch()?)?; + spill_context.sort_and_spill(self.table.take_state_batches()?)?; reservation .try_resize(self.table.memory_size()) .map_err(|e| { @@ -356,12 +357,13 @@ impl Aggregating { // through replay before any more final results can be emitted. continue; } - let Some(batch) = self.table.take_completed_result_batch()? else { + let batches = self.table.take_completed_result_batches()?; + if batches.is_empty() { continue; - }; + } timer.done(); return Ok(Some(ExecutionStage::Outputting(Outputting { - batch, + batches, resume: Some(self), }))); } @@ -370,7 +372,7 @@ impl Aggregating { drop(self.input); let timer = elapsed_compute.timer(); if let Some(mut spill_context) = self.spill_context.filter(|s| s.has_spills()) { - spill_context.sort_and_spill(self.table.take_state_batch()?)?; + spill_context.sort_and_spill(self.table.take_state_batches()?)?; let metrics = self.table.metrics(); drop(self.table); reservation.try_resize(0)?; @@ -385,29 +387,30 @@ impl Aggregating { } self.table.input_done(); - let output = self.table.take_completed_result_batch()?; + let batches = self.table.take_completed_result_batches()?; drop(self.table); timer.done(); - let Some(batch) = output else { + if batches.is_empty() { reservation.try_resize(0)?; return Ok(None); - }; + } Ok(Some(ExecutionStage::Outputting(Outputting { - batch, + batches, resume: None, }))) } } impl Outputting { - /// Emits slices of one materialized batch without touching the aggregate table. + /// Emits slices of the materialized batches without touching the aggregate + /// table. async fn handle_stage( self, context: &OrderedFinalAggregateContext, reservation: &MemoryReservation, emitter: &mut TryEmitter, ) -> Result> { - let Self { mut batch, resume } = self; + let Self { batches, resume } = self; let elapsed_compute = context.baseline_metrics.elapsed_compute(); let mut timer = elapsed_compute.timer(); let (table_memory, next_stage) = match resume { @@ -417,32 +420,43 @@ impl Outputting { ), None => (0, None), }; - let batch_memory = batch.get_array_memory_size(); - match reservation.try_resize(table_memory + batch_memory) { + let batches_memory: usize = batches.iter().map(|b| b.memory_size).sum(); + match reservation.try_resize(table_memory + batches_memory) { Ok(()) => {} Err(DataFusionError::ResourcesExhausted(_)) => { - // If we cannot hold the batch while slicing, hand it off whole. + // If we cannot hold the batches while slicing, hand them off whole. reservation.try_resize(table_memory)?; timer.done(); - emitter.emit(batch).await; + for MaterializedBatch { batch, .. } in batches { + emitter.emit(batch).await; + } return Ok(next_stage); } Err(e) => return Err(e), } - while batch.num_rows() > context.batch_size { - let output = batch.slice(0, context.batch_size); - batch = - batch.slice(context.batch_size, batch.num_rows() - context.batch_size); + for MaterializedBatch { + mut batch, + memory_size: batch_memory, + } in batches + { + while batch.num_rows() > context.batch_size { + let output = batch.slice(0, context.batch_size); + batch = batch + .slice(context.batch_size, batch.num_rows() - context.batch_size); + timer.done(); + emitter.emit(output).await; + timer = elapsed_compute.timer(); + } + + // The final slice transfers ownership of the buffers to the + // consumer, so this batch's memory is released batch by batch. + reservation.try_shrink(batch_memory)?; timer.done(); - emitter.emit(output).await; + emitter.emit(batch).await; timer = elapsed_compute.timer(); } - - // The final slice transfers ownership of the buffers to the consumer. - reservation.try_shrink(batch_memory)?; timer.done(); - emitter.emit(batch).await; Ok(next_stage) } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index dd845353c77eb..3f33364805d5a 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -27,7 +27,9 @@ use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::StreamExt; use super::AggregateExec; -use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; +use super::aggregate_hash_table::{ + MaterializedBatch, OrderedAggregateTable, PartialMarker, +}; use crate::aggregates::AggregateMode; use crate::aggregates::order::GroupOrdering; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; @@ -122,9 +124,10 @@ struct Aggregating { } struct Outputting { - /// Materialized aggregate states. Each iteration emits the first `batch_size` - /// rows and replaces this batch with the remaining slice. - batch: RecordBatch, + /// Materialized aggregate states: one batch for flat storage, one per + /// block or chunk for blocked storage. Each batch is emitted in slices of + /// `batch_size` rows. + batches: Vec, /// Aggregation stage to resume after output; `None` after EOF. resume: Option, } @@ -310,31 +313,32 @@ impl Aggregating { let timer = elapsed_compute.timer(); self.table.aggregate_batch(&batch)?; - let output = match reservation.try_resize(self.table.memory_size()) { - Ok(()) => self.table.take_completed_state_batch()?, + let batches = match reservation.try_resize(self.table.memory_size()) { + Ok(()) => self.table.take_completed_state_batches()?, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { // Partial ordering may have an unbounded active key range. // The final stage can merge incomplete states emitted here. if matches!(self.table.group_ordering(), GroupOrdering::Full(_)) { return Err(oom); } - let Some(batch) = self.table.take_state_batch()? else { + let batches = self.table.take_state_batches()?; + if batches.is_empty() { return Err(oom); - }; - Some(batch) + } + batches } Err(e) => return Err(e), }; - let Some(batch) = output else { + if batches.is_empty() { continue; - }; + } timer.done(); // OOM, do early emit next, and go back to the current state to continue // aggregating return Ok(Some(ExecutionStage::Outputting(Outputting { - batch, + batches, resume: Some(self), }))); } @@ -343,23 +347,23 @@ impl Aggregating { drop(self.input); self.table.input_done(); let timer = elapsed_compute.timer(); - let output = self.table.take_completed_state_batch()?; + let batches = self.table.take_completed_state_batches()?; drop(self.table); timer.done(); - let Some(batch) = output else { + if batches.is_empty() { reservation.try_resize(0)?; return Ok(None); - }; + } Ok(Some(ExecutionStage::Outputting(Outputting { - batch, + batches, resume: None, }))) } } impl Outputting { - /// Emits slices of one materialized batch without touching the hash table. + /// Emits slices of the materialized batches without touching the hash table. /// /// See [`OrderedPartialAggregateStream::into_stream`] for stage transitions /// and output memory accounting. @@ -369,7 +373,7 @@ impl Outputting { reservation: &MemoryReservation, emitter: &mut TryEmitter, ) -> Result> { - let Self { mut batch, resume } = self; + let Self { batches, resume } = self; let elapsed_compute = context.baseline_metrics.elapsed_compute(); let mut timer = elapsed_compute.timer(); let (table_memory, next_stage) = match resume { @@ -379,38 +383,49 @@ impl Outputting { ), None => (0, None), }; - let batch_memory = batch.get_array_memory_size(); - match reservation.try_resize(table_memory + batch_memory) { + let batches_memory: usize = batches.iter().map(|b| b.memory_size).sum(); + match reservation.try_resize(table_memory + batches_memory) { Ok(()) => {} Err(DataFusionError::ResourcesExhausted(_)) => { - // If we cannot hold the batch while slicing, hand it off whole. + // If we cannot hold the batches while slicing, hand them off whole. // Only the retained table needs to remain reserved. reservation.try_resize(table_memory)?; - context.reduction_factor.add_part(batch.num_rows()); timer.done(); - emitter.emit(batch).await; + for MaterializedBatch { batch, .. } in batches { + context.reduction_factor.add_part(batch.num_rows()); + emitter.emit(batch).await; + } return Ok(next_stage); } Err(e) => return Err(e), } - while batch.num_rows() > context.batch_size { - // 1. Emit first `batch_size` rows from `batch` - // 2. Update `batch`` with the remaining tail - let output = batch.slice(0, context.batch_size); - batch = - batch.slice(context.batch_size, batch.num_rows() - context.batch_size); - context.reduction_factor.add_part(output.num_rows()); + for MaterializedBatch { + mut batch, + memory_size: batch_memory, + } in batches + { + while batch.num_rows() > context.batch_size { + // 1. Emit first `batch_size` rows from `batch` + // 2. Update `batch`` with the remaining tail + let output = batch.slice(0, context.batch_size); + batch = batch + .slice(context.batch_size, batch.num_rows() - context.batch_size); + context.reduction_factor.add_part(output.num_rows()); + timer.done(); + emitter.emit(output).await; + timer = elapsed_compute.timer(); + } + + // The final slice transfers ownership of the buffers to the + // consumer, so this batch's memory is released batch by batch. + reservation.try_shrink(batch_memory)?; + context.reduction_factor.add_part(batch.num_rows()); timer.done(); - emitter.emit(output).await; + emitter.emit(batch).await; timer = elapsed_compute.timer(); } - - // The final slice transfers ownership of the buffers to the consumer. - reservation.try_shrink(batch_memory)?; - context.reduction_factor.add_part(batch.num_rows()); timer.done(); - emitter.emit(batch).await; Ok(next_stage) } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 97afab093a872..2c253e10f416e 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -409,7 +409,7 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let mut result = table - .take_state_batch() + .take_state_batches() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -456,7 +456,7 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let replay = match table - .take_state_batch() + .take_state_batches() .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d8ea2f711eb6c..d9001f4054751 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}; @@ -103,13 +104,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, @@ -318,21 +320,26 @@ 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_iter() + .map(|b| b.batch) + .collect(), }, ) } - // 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), } } @@ -350,30 +357,33 @@ impl PartialReduceHashAggregateStream { ) -> PartialReduceHashAggregateStateTransition { let PartialReduceHashAggregateState::EmittingOnMemoryPressure { hash_table, - remaining_groups: batch, + mut remaining_groups, } = original_state else { unreachable!("expected the EmittingOnMemoryPressure state") }; + let Some(batch) = 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 { + batch + } else { + // More rows in this batch, keep the rest at the front. + remaining_groups.push_front( + batch.slice(self.batch_size, batch.num_rows() - self.batch_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 288bf025e4e39..fee7e349dedf5 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -440,7 +440,7 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let mut result = hash_table - .take_state_batch() + .take_state_batches() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -488,7 +488,7 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let replay = match hash_table - .take_state_batch() + .take_state_batches() .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index a1d52382dead1..7ad58acbf55ce 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -20,7 +20,6 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryReservation; @@ -28,7 +27,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}; @@ -41,7 +40,7 @@ use crate::{InputOrderMode, SendableRecordBatchStream}; /// /// Every aggregation stream that spills does so the same way. Each spill event /// drains all currently buffered groups as intermediate state (see -/// `take_state_batch` on the aggregate tables), sorts them by the full group +/// `take_state_batches` on the aggregate tables), sorts them by the full group /// key, and writes them to one spill file. After the original input ends, all /// files are merged and replayed through an [`OrderedFinalAggregateStream`], /// which merges the states and evaluates the final aggregate values. @@ -208,37 +207,41 @@ impl AggregateSpill { !self.spills.is_empty() } - /// Sorts `state_batch`, the intermediate state of all currently buffered - /// groups (`None` if there are no groups), and writes it as one spill file. - /// Memory reservation should be updated by the caller. + /// Sorts `state_batches`, the intermediate state of all currently buffered + /// groups (empty if there are no groups), and writes each batch as one + /// spill file. Memory reservation should be updated by the caller. + /// + /// Flat aggregate storage produces one batch; blocked storage produces one + /// batch per block. pub(super) fn sort_and_spill( &mut self, - state_batch: Option, + state_batches: Vec, ) -> Result<()> { - let Some(state_batch) = state_batch else { - return Ok(()); - }; + // TODO: sort across multiple arrays without concat (see #24928 + // sort.rs/rank.rs) so all blocks are written as one sorted run instead + // of one run per block. + for state_batch in state_batches { + let sorted_iter = IncrementalSortIterator::new( + state_batch.batch, + self.spill_expr.clone(), + self.batch_size, + ); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + self.label, + )?; - let sorted_iter = IncrementalSortIterator::new( - state_batch, - self.spill_expr.clone(), - self.batch_size, - ); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - self.label, - )?; + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("{}: produced an empty spill", self.label); + }; - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("{}: produced an empty spill", self.label); - }; - - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + } Ok(()) } diff --git a/datafusion/sqllogictest/test_files/aggregate_blocked.slt b/datafusion/sqllogictest/test_files/aggregate_blocked.slt new file mode 100644 index 0000000000000..3657eb042fbdb --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_blocked.slt @@ -0,0 +1,160 @@ +# 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. + +# Blocked hash aggregation: `GROUP BY ` where every aggregate is +# `count` keeps its state in blocks of 2^18 groups and emits block by block. +# +# 600000 groups span 3 blocks. Keys are scrambled with `(v * 7) % 600000` +# (gcd(7, 600000) = 1, so still 600000 distinct keys) because generate_series +# is sorted, which would take the ordered path instead of the hash path. +# +# Each query summarizes the grouped result into one row. + +statement ok +SET datafusion.execution.target_partitions = 1 + +statement ok +CREATE TABLE t AS +SELECT (v * 7) % 600000 AS k, CASE WHEN v % 5 = 0 THEN NULL ELSE v END AS v, + CAST(v AS DOUBLE) / 4 AS f +FROM generate_series(0, 1799999) AS s(v) + +# count per group across several blocks +query IIII +SELECT count(*), sum(c), min(c), max(c) +FROM (SELECT k, count(v) AS c FROM t GROUP BY k) +---- +600000 1440000 0 3 + +# count(*) per group +query II +SELECT count(*), sum(c) +FROM (SELECT k, count(*) AS c FROM t GROUP BY k) +---- +600000 1800000 + +# NULL group key, interned after the first block is full +query III +SELECT count(*), sum(c), sum(CASE WHEN k2 IS NULL THEN c END) +FROM ( + SELECT CASE WHEN k >= 599990 THEN NULL ELSE k END AS k2, count(v) AS c + FROM t GROUP BY k2 +) +---- +599991 1440000 24 + +# FILTER on count across blocks +query II +SELECT count(*), sum(c) +FROM (SELECT k, count(v) FILTER (WHERE v % 2 = 0) AS c FROM t GROUP BY k) +---- +600000 720000 + +# count together with a non-blocked aggregate takes today's path +query III +SELECT count(*), sum(c), sum(s) +FROM (SELECT k, count(v) AS c, sum(v) AS s FROM t GROUP BY k) +---- +600000 1440000 1296000000000 + +# blocked sum +query IIII +SELECT count(*), sum(s), min(s), max(s) +FROM (SELECT k, sum(v) AS s FROM t GROUP BY k) +---- +600000 1296000000000 1800003 3599997 + +# blocked avg +query IRRR +SELECT count(*), sum(a), min(a), max(a) +FROM (SELECT k, avg(f) AS a FROM t GROUP BY k) +---- +600000 134999925000 150000 299999.75 + +# count, sum and avg, all blocked +query IIIR +SELECT count(*), sum(c), sum(s), sum(a) +FROM (SELECT k, count(v) AS c, sum(v) AS s, avg(f) AS a FROM t GROUP BY k) +---- +600000 1440000 1296000000000 134999925000 + +# blocked count(DISTINCT): one set of distinct values per block +query IIII +SELECT count(*), sum(d), min(d), max(d) +FROM (SELECT k, count(DISTINCT v % 2) AS d FROM t GROUP BY k) +---- +600000 480000 0 1 + +# blocked min / max on strings, both inline (short) and heap (long) values +query ITTI +SELECT count(*), min(mn), max(mx), sum(length(mx)) +FROM ( + SELECT k, min(CAST(v AS VARCHAR)) AS mn, max(repeat(CAST(v AS VARCHAR), 5)) AS mx + FROM t GROUP BY k +) +---- +600000 1 999999999999999999999999999999 14229605 + +# flat (string) keys with a blocked accumulator +query II +SELECT count(*), sum(s) +FROM (SELECT CAST(k AS VARCHAR) AS ks, sum(v) AS s FROM t GROUP BY ks) +---- +600000 1296000000000 + +# multiple partitions: partial + final aggregation both blocked +statement ok +SET datafusion.execution.target_partitions = 4 + +query II +SELECT count(*), sum(c) +FROM (SELECT k, count(v) AS c FROM t GROUP BY k) +---- +600000 1440000 + +statement ok +SET datafusion.execution.target_partitions = 1 + +# ordered input (generate_series is sorted on v) with more groups than a block +query II +SELECT count(*), sum(c) +FROM (SELECT v, count(*) AS c FROM generate_series(0, 599999) AS s(v) GROUP BY v) +---- +600000 600000 + +# spill with more groups than a block +statement ok +SET datafusion.runtime.memory_limit = '8M' + +query II +SELECT count(*), sum(c) +FROM (SELECT k, count(v) AS c FROM t GROUP BY k) +---- +600000 1440000 + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +DROP TABLE t + +# Restore settings to slt runner defaults +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema From bb9a3549326bbdc1c325a12d8a1bdf948b509a8e Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:04:49 +0300 Subject: [PATCH 02/10] refactor(agg): keep only blocked count and its building blocks --- .../src/aggregate/count_distinct.rs | 2 - .../count_distinct/blocked_groups.rs | 347 -------- .../src/aggregate/count_distinct/groups.rs | 45 +- .../src/aggregate/groups_accumulator.rs | 2 - .../groups_accumulator/accumulate.rs | 30 +- .../groups_accumulator/blocked_null_state.rs | 385 --------- .../groups_accumulator/blocked_prim_op.rs | 245 ------ .../groups_accumulator/blocked_vec.rs | 143 +--- .../aggregate/groups_accumulator/prim_op.rs | 104 +-- datafusion/functions-aggregate/src/average.rs | 487 ++--------- datafusion/functions-aggregate/src/count.rs | 33 +- datafusion/functions-aggregate/src/min_max.rs | 44 +- .../src/min_max/min_max_bytes.rs | 753 +++++------------- datafusion/functions-aggregate/src/sum.rs | 25 +- .../test_files/aggregate_blocked.slt | 54 +- 15 files changed, 370 insertions(+), 2329 deletions(-) delete mode 100644 datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs delete mode 100644 datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs delete mode 100644 datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs index bdf4fc0eebd42..bb706aa614dbc 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs @@ -15,13 +15,11 @@ // specific language governing permissions and limitations // under the License. -mod blocked_groups; mod bytes; mod dict; mod groups; mod native; -pub use blocked_groups::BlockedPrimitiveDistinctCountGroupsAccumulator; pub use bytes::BytesDistinctCountAccumulator; pub use bytes::BytesViewDistinctCountAccumulator; pub use dict::DictionaryCountAccumulator; diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs deleted file mode 100644 index 04b124eb1979d..0000000000000 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/blocked_groups.rs +++ /dev/null @@ -1,347 +0,0 @@ -// 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. - -//! [`BlockedPrimitiveDistinctCountGroupsAccumulator`]: distinct count with -//! state stored in blocks. - -use std::hash::Hash; -use std::mem::size_of; -use std::sync::Arc; - -use arrow::array::{ - ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, PrimitiveArray, -}; -use arrow::buffer::{OffsetBuffer, ScalarBuffer}; -use arrow::datatypes::{ArrowPrimitiveType, Field}; -use datafusion_common::HashSet; -use datafusion_common::hash_utils::RandomState; -use datafusion_expr_common::blocked_groups_accumulator::{ - BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, -}; - -use super::groups::distinct_convert_to_state; -use crate::aggregate::groups_accumulator::accumulate::accumulate_blocked; -use crate::aggregate::groups_accumulator::blocked_vec::BlockedVec; - -/// Distinct values seen by the groups of one block, keyed by the group's -/// position in the block. -type BlockSeen = HashSet<(u32, V), RandomState>; - -/// [`PrimitiveDistinctCountGroupsAccumulator`] with state stored in blocks. -/// -/// The distinct values are kept in one set per block (instead of one set for -/// all groups), so a block is emitted by draining only its own set, and its -/// memory is released with it. -/// -/// [`PrimitiveDistinctCountGroupsAccumulator`]: super::PrimitiveDistinctCountGroupsAccumulator -pub struct BlockedPrimitiveDistinctCountGroupsAccumulator -where - T::Native: Eq + Hash, -{ - /// One set per block of `counts` - seen: Vec>, - counts: BlockedVec, -} - -impl BlockedPrimitiveDistinctCountGroupsAccumulator -where - T::Native: Eq + Hash, -{ - pub fn new(block_size: usize) -> Self { - Self { - seen: vec![], - counts: BlockedVec::new(block_size), - } - } - - /// Grows `counts` to `total_num_groups` and adds a set for every new block. - fn grow_to(&mut self, total_num_groups: usize) { - self.counts.grow_to(total_num_groups, 0); - self.seen - .resize_with(self.counts.num_blocks(), BlockSeen::default); - } - - /// Adds `value` to the group at `index`, counting it if it is new. - #[inline] - fn insert( - seen: &mut [BlockSeen], - counts: &mut BlockedVec, - index: BlocksIndex, - value: T::Native, - ) { - if seen[index.block_index()].insert((index.index_in_block() as u32, value)) { - // SAFETY: `index` was registered by `grow_to` - unsafe { *counts.get_unchecked_mut(index) += 1 }; - } - } - - /// Removes the counts and distinct values of the groups selected by - /// `emit_to`, one entry per emitted block. - fn take(&mut self, emit_to: BlockedEmitTo) -> Vec<(Vec, BlockSeen)> { - match emit_to { - BlockedEmitTo::All => { - let counts = self.counts.take_all(); - let seen = std::mem::take(&mut self.seen); - counts.into_iter().zip(seen).collect() - } - BlockedEmitTo::NextBlock => match self.counts.take_next_block() { - Some(counts) => vec![(counts, self.seen.remove(0))], - None => vec![], - }, - BlockedEmitTo::First(n) => { - // Rare path (ordered input): renumber every remaining group - let block_size = self.counts.block_size(); - let counts = self.counts.take_first(n); - let mut emitted = BlockSeen::default(); - let mut remaining: Vec> = - (0..self.counts.num_blocks()) - .map(|_| BlockSeen::default()) - .collect(); - for (block_index, block) in - std::mem::take(&mut self.seen).into_iter().enumerate() - { - for (index_in_block, value) in block { - let flat = BlocksIndex::new(block_index, index_in_block as usize) - .flat(block_size); - match flat.checked_sub(n) { - None => { - emitted.insert((index_in_block, value)); - } - Some(flat) => { - let index = BlocksIndex::from_flat(flat, block_size); - remaining[index.block_index()] - .insert((index.index_in_block() as u32, value)); - } - } - } - } - self.seen = remaining; - vec![(counts, emitted)] - } - } - } -} - -impl BlockedGroupsAccumulator - for BlockedPrimitiveDistinctCountGroupsAccumulator -where - T::Native: Eq + Hash, -{ - 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, - ) -> datafusion_common::Result<()> { - debug_assert_eq!(values.len(), 1); - self.grow_to(total_num_groups); - let arr = values[0].as_primitive::(); - let (seen, counts) = (&mut self.seen, &mut self.counts); - accumulate_blocked(group_indices, arr, opt_filter, |index, value| { - Self::insert(seen, counts, index, value) - }); - Ok(()) - } - - fn merge_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[BlocksIndex], - total_num_groups: usize, - ) -> datafusion_common::Result<()> { - debug_assert_eq!(values.len(), 1); - self.grow_to(total_num_groups); - let list_array = values[0].as_list::(); - let inner_values = list_array.values().as_primitive::().values(); - let offsets = list_array.offsets(); - - for (row_idx, &index) in group_indices.iter().enumerate() { - let start = offsets[row_idx] as usize; - let end = offsets[row_idx + 1] as usize; - for &value in &inner_values[start..end] { - Self::insert(&mut self.seen, &mut self.counts, index, value); - } - } - Ok(()) - } - - fn evaluate( - &mut self, - emit_to: BlockedEmitTo, - ) -> datafusion_common::Result> { - Ok(self - .take(emit_to) - .into_iter() - // The emitted sets are dropped here, releasing their memory - .map(|(counts, _)| Arc::new(Int64Array::from(counts)) as ArrayRef) - .collect()) - } - - fn state( - &mut self, - emit_to: BlockedEmitTo, - ) -> datafusion_common::Result>> { - Ok(self - .take(emit_to) - .into_iter() - .map(|(counts, seen)| { - // Prefix-sum the counts of the block into list offsets, then - // move every distinct value to its group's list - let mut offsets = Vec::with_capacity(counts.len() + 1); - offsets.push(0i32); - let mut total = 0i32; - for &count in &counts { - total += count as i32; - offsets.push(total); - } - let mut values = vec![T::Native::default(); total as usize]; - let mut cursors: Vec = offsets[..counts.len()].to_vec(); - for (index_in_block, value) in seen { - let cursor = &mut cursors[index_in_block as usize]; - values[*cursor as usize] = value; - *cursor += 1; - } - let values = - Arc::new(PrimitiveArray::::new(ScalarBuffer::from(values), None)); - vec![Arc::new(ListArray::new( - Arc::new(Field::new_list_field(T::DATA_TYPE, true)), - OffsetBuffer::new(offsets.into()), - values, - None, - )) as ArrayRef] - }) - .collect()) - } - - fn convert_to_state( - &self, - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - ) -> datafusion_common::Result> { - distinct_convert_to_state::(values, opt_filter) - } - - fn size(&self) -> usize { - let entry = size_of::<(u32, T::Native)>() + size_of::(); - size_of::() - + self - .seen - .iter() - .map(|s| s.capacity() * entry) - .sum::() - + self.counts.allocated_size() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{Array, Int32Array}; - use arrow::datatypes::Int32Type; - use datafusion_common::Result; - - fn indices(flats: &[usize]) -> Vec { - flats - .iter() - .map(|&f| BlocksIndex::from_flat(f, 4)) - .collect() - } - - fn counts(arrays: &[ArrayRef]) -> Vec> { - arrays - .iter() - .map(|a| { - a.as_primitive::() - .values() - .to_vec() - }) - .collect() - } - - fn accumulator() -> Result> - { - let mut acc = BlockedPrimitiveDistinctCountGroupsAccumulator::::new(4); - let values: ArrayRef = Arc::new(Int32Array::from(vec![ - Some(1), - Some(1), - Some(2), - None, - Some(7), - Some(7), - Some(8), - ])); - acc.update_batch(&[values], &indices(&[0, 0, 0, 1, 5, 5, 5]), None, 6)?; - Ok(acc) - } - - #[test] - fn counts_distinct_values_per_block() -> Result<()> { - let mut acc = accumulator()?; - assert_eq!( - counts(&acc.evaluate(BlockedEmitTo::NextBlock)?), - vec![vec![2, 0, 0, 0]] - ); - assert_eq!(acc.seen.len(), 1, "the emitted block's set is released"); - assert_eq!(counts(&acc.evaluate(BlockedEmitTo::All)?), vec![vec![0, 2]]); - Ok(()) - } - - #[test] - fn state_round_trips_through_merge() -> Result<()> { - let mut acc = accumulator()?; - let state = acc.state(BlockedEmitTo::All)?; - assert_eq!(state.len(), 2); - let mut merged = - BlockedPrimitiveDistinctCountGroupsAccumulator::::new(4); - for (block_index, columns) in state.iter().enumerate() { - let rows = columns[0].len(); - let groups: Vec = (0..rows) - .map(|i| BlocksIndex::new(block_index, i)) - .collect(); - merged.merge_batch(columns, &groups, 6)?; - } - // merging the same values again does not change the distinct counts - merged.merge_batch(&state[1], &indices(&[4, 5]), 6)?; - assert_eq!( - counts(&merged.evaluate(BlockedEmitTo::All)?), - vec![vec![2, 0, 0, 0], vec![0, 2]] - ); - Ok(()) - } - - #[test] - fn first_n_renumbers_across_blocks() -> Result<()> { - let mut acc = accumulator()?; - assert_eq!( - counts(&acc.evaluate(BlockedEmitTo::First(1))?), - vec![vec![2]] - ); - // group 5 moved to position 4 (block 1, index 0) - let values: ArrayRef = Arc::new(Int32Array::from(vec![7, 9])); - acc.update_batch(&[values], &indices(&[4, 4]), None, 5)?; - assert_eq!( - counts(&acc.evaluate(BlockedEmitTo::All)?), - vec![vec![0, 0, 0, 0], vec![3]] - ); - Ok(()) - } -} diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 9b39ebc0fd103..38f6d26200ec9 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -211,7 +211,24 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> datafusion_common::Result> { - distinct_convert_to_state::(values, opt_filter) + debug_assert_eq!(values.len(), 1); + let arr = values[0].as_primitive::(); + + let values_builder = PrimitiveBuilder::::with_capacity(arr.len()); + let mut builder = ListBuilder::new(values_builder) + .with_field(Arc::new(Field::new_list_field(T::DATA_TYPE, true))); + + for row in 0..arr.len() { + let included = arr.is_valid(row) + && opt_filter + .is_none_or(|filter| filter.is_valid(row) && filter.value(row)); + if included { + builder.values().append_value(arr.value(row)); + } + builder.append(true); + } + + Ok(vec![Arc::new(builder.finish())]) } fn size(&self) -> usize { size_of::() @@ -220,32 +237,6 @@ where } } -/// Converts an input batch directly to state: a list per row with the row's -/// value, or an empty list for null or filtered rows. Shared by the flat and -/// blocked distinct count accumulators. -pub(super) fn distinct_convert_to_state( - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, -) -> datafusion_common::Result> { - debug_assert_eq!(values.len(), 1); - let arr = values[0].as_primitive::(); - - let values_builder = PrimitiveBuilder::::with_capacity(arr.len()); - let mut builder = ListBuilder::new(values_builder) - .with_field(Arc::new(Field::new_list_field(T::DATA_TYPE, true))); - - for row in 0..arr.len() { - let included = arr.is_valid(row) - && opt_filter.is_none_or(|filter| filter.is_valid(row) && filter.value(row)); - if included { - builder.values().append_value(arr.value(row)); - } - builder.append(true); - } - - Ok(vec![Arc::new(builder.finish())]) -} - #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index bf0f24f122ef5..1d8cb5997a142 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -19,8 +19,6 @@ //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] pub mod accumulate; -pub mod blocked_null_state; -pub mod blocked_prim_op; pub mod blocked_vec; pub mod bool_op; pub mod nulls; 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 c912a12c5f3da..1b3d5b8b13b89 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -423,38 +423,10 @@ pub fn accumulate( group_indices: &[usize], values: &PrimitiveArray, opt_filter: Option<&BooleanArray>, - value_fn: F, + mut value_fn: F, ) where T: ArrowPrimitiveType + Send, F: FnMut(usize, T::Native) + Send, -{ - accumulate_values(group_indices, values, opt_filter, value_fn) -} - -/// [`accumulate`] for blocked group indices. -pub fn accumulate_blocked( - group_indices: &[BlocksIndex], - values: &PrimitiveArray, - opt_filter: Option<&BooleanArray>, - value_fn: F, -) where - T: ArrowPrimitiveType, - F: FnMut(BlocksIndex, T::Native), -{ - accumulate_values(group_indices, values, opt_filter, value_fn) -} - -/// Shared implementation of [`accumulate`] and [`accumulate_blocked`]. -#[inline(always)] -fn accumulate_values( - group_indices: &[I], - values: &PrimitiveArray, - opt_filter: Option<&BooleanArray>, - mut value_fn: F, -) where - I: Copy, - T: ArrowPrimitiveType, - F: FnMut(I, T::Native), { let data: &[T::Native] = values.values(); assert_eq!(data.len(), group_indices.len()); diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs deleted file mode 100644 index 04d358c3d0f70..0000000000000 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_null_state.rs +++ /dev/null @@ -1,385 +0,0 @@ -// 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. - -//! [`BlockedNullState`]: [`NullState`] for blocked group indices. -//! -//! [`NullState`]: super::accumulate::NullState - -use arrow::array::{ - Array, ArrowPrimitiveType, BooleanArray, BooleanBufferBuilder, PrimitiveArray, -}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; - -use super::blocked_vec::{BlockedVec, SeenGroups}; - -/// Tracks which groups have seen at least one non null, non filtered value, -/// like [`NullState`], for state stored in a [`BlockedVec`]. -/// -/// The seen bits are stored in one bitmap per block, with the same block -/// boundaries as the [`BlockedVec`], so each emitted block gets its own -/// [`NullBuffer`]. -/// -/// [`NullState`]: super::accumulate::NullState -#[derive(Debug)] -pub struct BlockedNullState { - block_size: usize, - seen: Seen, -} - -#[derive(Debug)] -enum Seen { - /// Every group seen so far has seen a value, so no bitmap is needed - All { num_groups: usize }, - /// Every bitmap but the last holds exactly `block_size` bits - Some { blocks: Vec }, -} - -impl BlockedNullState { - pub fn new(block_size: usize) -> Self { - Self { - block_size, - seen: Seen::All { num_groups: 0 }, - } - } - - /// Size of all buffers allocated by this null state, not including self. - pub fn size(&self) -> usize { - match &self.seen { - Seen::All { .. } => 0, - Seen::Some { blocks } => blocks.iter().map(|b| b.capacity() / 8).sum(), - } - } - - /// Updates `state` with every non null, non filtered value of `values` - /// through [`BlockedVec::update_values`], and marks the groups of those - /// values as seen. - /// - /// Like [`NullState::accumulate`], no bits are tracked while every value - /// is valid and every new group is present in `group_indices`. - /// - /// [`NullState::accumulate`]: super::accumulate::NullState::accumulate - #[expect(clippy::too_many_arguments, reason = "mirrors `NullState::accumulate`")] - pub fn accumulate( - &mut self, - state: &mut BlockedVec, - total_num_groups: usize, - starting_value: T, - group_indices: &[BlocksIndex], - values: &PrimitiveArray, - opt_filter: Option<&BooleanArray>, - update_fn: F, - ) where - T: Copy + Default, - V: ArrowPrimitiveType, - F: FnMut(&mut T, V::Native), - { - if opt_filter.is_none() - && values.null_count() == 0 - && let Seen::All { num_groups } = &mut self.seen - && new_groups_are_dense( - group_indices, - *num_groups, - total_num_groups, - self.block_size, - ) - { - state.update_values( - total_num_groups, - starting_value, - group_indices, - values, - None, - &mut (), - update_fn, - ); - *num_groups = total_num_groups; - return; - } - - let mut seen = SeenBits { - blocks: self.bitmaps(total_num_groups), - }; - state.update_values( - total_num_groups, - starting_value, - group_indices, - values, - opt_filter, - &mut seen, - update_fn, - ); - } - - /// Returns the bitmaps, grown to `total_num_groups` bits (new groups have - /// not seen a value), switching from [`Seen::All`] if needed. - fn bitmaps(&mut self, total_num_groups: usize) -> &mut Vec { - if let Seen::All { num_groups } = self.seen { - let mut blocks = vec![]; - append_bits(&mut blocks, self.block_size, num_groups, true); - self.seen = Seen::Some { blocks }; - } - let Seen::Some { blocks } = &mut self.seen else { - unreachable!("switched to `Seen::Some` above") - }; - let len = bits_len(blocks); - if total_num_groups > len { - append_bits(blocks, self.block_size, total_num_groups - len, false); - } - blocks - } - - /// Builds the null buffer of every block emitted by `emit_to`, `None` - /// when a block has no nulls, and removes them from the state. - /// - /// Returns one entry per emitted block, following the same rules as - /// `BlockedGroupsAccumulator::evaluate`. - pub fn build(&mut self, emit_to: BlockedEmitTo) -> Vec> { - let block_size = self.block_size; - match &mut self.seen { - Seen::All { num_groups } => { - let emitted_blocks = match emit_to { - BlockedEmitTo::All => { - let blocks = num_groups.div_ceil(block_size); - *num_groups = 0; - blocks - } - BlockedEmitTo::NextBlock => { - let blocks = usize::from(*num_groups > 0); - *num_groups -= (*num_groups).min(block_size); - blocks - } - BlockedEmitTo::First(n) => { - *num_groups -= n; - 1 - } - }; - vec![None; emitted_blocks] - } - Seen::Some { blocks } => match emit_to { - BlockedEmitTo::All => { - let nulls = blocks.iter_mut().map(|b| Some(finish(b))).collect(); - self.seen = Seen::All { num_groups: 0 }; - nulls - } - BlockedEmitTo::NextBlock => { - if blocks.is_empty() { - return vec![]; - } - vec![Some(finish(&mut blocks.remove(0)))] - } - BlockedEmitTo::First(n) => { - // Rare path (ordered input): re-chunk the remaining bits - let bits: Vec = - blocks.iter_mut().map(|b| b.finish()).collect(); - let first = bits[0].slice(0, n); - let mut remaining = vec![]; - for (i, block) in bits.iter().enumerate() { - let skip = if i == 0 { n } else { 0 }; - let block = block.slice(skip, block.len() - skip); - append_buffer(&mut remaining, block_size, &block); - } - *blocks = remaining; - vec![Some(NullBuffer::new(first))] - } - }, - } - } -} - -/// Marks groups as seen in the per-block bitmaps. -struct SeenBits<'a> { - blocks: &'a mut Vec, -} - -impl SeenGroups for SeenBits<'_> { - #[inline] - fn mark(&mut self, index: BlocksIndex) { - self.blocks[index.block_index()].set_bit(index.index_in_block(), true); - } -} - -/// Returns true when all newly registered groups are present in -/// `group_indices`, see the flat `new_groups_are_dense`. -fn new_groups_are_dense( - group_indices: &[BlocksIndex], - first_new_group: usize, - total_num_groups: usize, - block_size: usize, -) -> bool { - if first_new_group == total_num_groups { - return true; - } - - let mut next_new_group = first_new_group; - for &group_index in group_indices { - let group_index = group_index.flat(block_size); - if group_index == next_new_group { - next_new_group += 1; - } else if group_index > next_new_group { - return false; - } - } - next_new_group == total_num_groups -} - -fn bits_len(blocks: &[BooleanBufferBuilder]) -> usize { - blocks.iter().map(|b| b.len()).sum() -} - -/// Appends `n` bits of `value`, filling the last bitmap up to `block_size` -/// bits before starting a new one. -fn append_bits( - blocks: &mut Vec, - block_size: usize, - mut n: usize, - value: bool, -) { - while n > 0 { - let last = match blocks.last_mut() { - Some(last) if last.len() < block_size => last, - _ => { - blocks.push(BooleanBufferBuilder::new(n.min(block_size))); - blocks.last_mut().expect("just pushed") - } - }; - let count = n.min(block_size - last.len()); - last.append_n(count, value); - n -= count; - } -} - -/// Appends the bits of `buffer`, like [`append_bits`]. -fn append_buffer( - blocks: &mut Vec, - block_size: usize, - buffer: &BooleanBuffer, -) { - let mut offset = 0; - while offset < buffer.len() { - let last = match blocks.last_mut() { - Some(last) if last.len() < block_size => last, - _ => { - blocks.push(BooleanBufferBuilder::new( - (buffer.len() - offset).min(block_size), - )); - blocks.last_mut().expect("just pushed") - } - }; - let count = (buffer.len() - offset).min(block_size - last.len()); - last.append_buffer(&buffer.slice(offset, count)); - offset += count; - } -} - -fn finish(builder: &mut BooleanBufferBuilder) -> NullBuffer { - NullBuffer::new(builder.finish()) -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::Int64Array; - use arrow::datatypes::Int64Type; - - fn indices(flats: &[usize]) -> Vec { - flats - .iter() - .map(|&f| BlocksIndex::from_flat(f, 4)) - .collect() - } - - fn nulls_of(nulls: &[Option]) -> Vec>> { - nulls - .iter() - .map(|n| n.as_ref().map(|n| n.iter().collect())) - .collect() - } - - #[test] - fn all_seen_needs_no_bitmap() { - let mut state = BlockedVec::::new(4); - let mut nulls = BlockedNullState::new(4); - let values = Int64Array::from(vec![1, 2, 3, 4, 5, 6]); - nulls.accumulate( - &mut state, - 6, - 0, - &indices(&[0, 1, 2, 3, 4, 5]), - &values, - None, - |s, v| *s += v, - ); - assert_eq!(nulls.size(), 0); - assert_eq!(nulls.build(BlockedEmitTo::All), vec![None, None]); - } - - #[test] - fn tracks_unseen_groups_across_blocks() { - let mut state = BlockedVec::::new(4); - let mut nulls = BlockedNullState::new(4); - // group 1 only gets a null, group 5 is filtered out - let values = - Int64Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]); - let filter = BooleanArray::from(vec![true, true, true, true, false, true]); - nulls.accumulate( - &mut state, - 7, - 0, - &indices(&[0, 1, 2, 3, 5, 6]), - &values, - Some(&filter), - |s, v| *s += v, - ); - let expected = vec![ - Some(vec![true, false, true, true]), - Some(vec![false, false, true]), - ]; - assert_eq!(nulls_of(&nulls.build(BlockedEmitTo::All)), expected); - } - - #[test] - fn next_block_and_first_n_shift_bits() { - let mut state = BlockedVec::::new(4); - let mut nulls = BlockedNullState::new(4); - let values = - Int64Array::from(vec![Some(1), None, Some(3), Some(4), None, Some(6)]); - nulls.accumulate( - &mut state, - 6, - 0, - &indices(&[0, 1, 2, 3, 4, 5]), - &values, - None, - |s, v| *s += v, - ); - assert_eq!( - nulls_of(&nulls.build(BlockedEmitTo::First(1))), - vec![Some(vec![true])] - ); - // remaining: [f, t, t, f | t] - assert_eq!( - nulls_of(&nulls.build(BlockedEmitTo::NextBlock)), - vec![Some(vec![false, true, true, false])] - ); - assert_eq!( - nulls_of(&nulls.build(BlockedEmitTo::All)), - vec![Some(vec![true])] - ); - let _ = Int64Type::DATA_TYPE; - } -} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs deleted file mode 100644 index 6b52acc9fb9e2..0000000000000 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_prim_op.rs +++ /dev/null @@ -1,245 +0,0 @@ -// 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. - -//! [`BlockedPrimitiveGroupsAccumulator`]: [`PrimitiveGroupsAccumulator`] with -//! state stored in blocks. -//! -//! [`PrimitiveGroupsAccumulator`]: super::prim_op::PrimitiveGroupsAccumulator - -use std::sync::Arc; - -use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray}; -use arrow::datatypes::{ArrowPrimitiveType, DataType}; -use datafusion_common::Result; -use datafusion_expr_common::blocked_groups_accumulator::{ - BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, -}; - -use super::blocked_null_state::BlockedNullState; -use super::blocked_vec::BlockedVec; -use super::prim_op::convert_to_state; - -/// [`PrimitiveGroupsAccumulator`] with per-group values stored in a -/// [`BlockedVec`], for aggregates whose state is the same as the input type -/// (such as `SUM`). -/// -/// F: The function to apply to two elements. The first argument is the -/// existing value and should be updated with the second value (e.g. -/// [`BitAndAssign`] style). -/// -/// [`PrimitiveGroupsAccumulator`]: super::prim_op::PrimitiveGroupsAccumulator -/// [`BitAndAssign`]: std::ops::BitAndAssign -#[derive(Debug)] -pub struct BlockedPrimitiveGroupsAccumulator -where - T: ArrowPrimitiveType + Send, - F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, -{ - /// values per group, stored as the native type - values: BlockedVec, - - /// The output type (needed for Decimal precision and scale) - data_type: DataType, - - /// The starting value for new groups - starting_value: T::Native, - - /// Track nulls in the input / filters - null_state: BlockedNullState, - - /// Function that computes the primitive result - prim_fn: F, -} - -impl BlockedPrimitiveGroupsAccumulator -where - T: ArrowPrimitiveType + Send, - F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, -{ - pub fn new(data_type: &DataType, block_size: usize, prim_fn: F) -> Self { - Self { - values: BlockedVec::new(block_size), - data_type: data_type.clone(), - null_state: BlockedNullState::new(block_size), - starting_value: T::default_value(), - prim_fn, - } - } - - /// Set the starting values for new groups - pub fn with_starting_value(mut self, starting_value: T::Native) -> Self { - self.starting_value = starting_value; - self - } - - fn take(&mut self, emit_to: BlockedEmitTo) -> Vec> { - match emit_to { - BlockedEmitTo::All => self.values.take_all(), - BlockedEmitTo::NextBlock => { - self.values.take_next_block().into_iter().collect() - } - BlockedEmitTo::First(n) => vec![self.values.take_first(n)], - } - } -} - -impl BlockedGroupsAccumulator for BlockedPrimitiveGroupsAccumulator -where - T: ArrowPrimitiveType + Send, - F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, -{ - fn block_size(&self) -> usize { - self.values.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].as_primitive::(); - let prim_fn = &self.prim_fn; - self.null_state.accumulate( - &mut self.values, - total_num_groups, - self.starting_value, - group_indices, - values, - opt_filter, - |value, new_value| prim_fn(value, new_value), - ); - Ok(()) - } - - fn merge_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[BlocksIndex], - total_num_groups: usize, - ) -> Result<()> { - // update / merge are the same - self.update_batch(values, group_indices, None, total_num_groups) - } - - fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { - let blocks = self.take(emit_to); - let nulls = self.null_state.build(emit_to); - debug_assert_eq!(blocks.len(), nulls.len()); - Ok(blocks - .into_iter() - .zip(nulls) - .map(|(values, nulls)| { - // no copy - Arc::new( - PrimitiveArray::::new(values.into(), nulls) - .with_data_type(self.data_type.clone()), - ) as ArrayRef - }) - .collect()) - } - - fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { - Ok(self - .evaluate(emit_to)? - .into_iter() - .map(|values| vec![values]) - .collect()) - } - - fn convert_to_state( - &self, - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - ) -> Result> { - convert_to_state::( - values, - opt_filter, - self.starting_value, - &self.prim_fn, - &self.data_type, - ) - } - - fn size(&self) -> usize { - self.values.allocated_size() + self.null_state.size() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{Array, Int64Array}; - use arrow::datatypes::Int64Type; - - fn sum() -> BlockedPrimitiveGroupsAccumulator { - BlockedPrimitiveGroupsAccumulator::new(&DataType::Int64, 4, |x, y| *x += y) - } - - fn indices(flats: &[usize]) -> Vec { - flats - .iter() - .map(|&f| BlocksIndex::from_flat(f, 4)) - .collect() - } - - fn values(arrays: &[ArrayRef]) -> Vec>> { - arrays - .iter() - .map(|a| a.as_primitive::().iter().collect()) - .collect() - } - - #[test] - fn sums_across_blocks_with_nulls_for_unseen_groups() -> Result<()> { - let mut acc = sum(); - let input: ArrayRef = Arc::new(Int64Array::from(vec![ - Some(1), - Some(2), - None, - Some(4), - Some(5), - ])); - acc.update_batch(&[input], &indices(&[0, 5, 1, 5, 0]), None, 6)?; - let partial: ArrayRef = Arc::new(Int64Array::from(vec![10, 20])); - acc.merge_batch(&[partial], &indices(&[2, 5]), 6)?; - - assert_eq!( - values(&acc.evaluate(BlockedEmitTo::All)?), - vec![vec![Some(6), None, Some(10), None], vec![None, Some(26)],] - ); - assert_eq!(acc.size(), 0); - Ok(()) - } - - #[test] - fn all_valid_input_has_no_nulls() -> Result<()> { - let mut acc = sum(); - let input: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])); - acc.update_batch(&[input], &indices(&[0, 1, 2, 3, 4]), None, 5)?; - let out = acc.evaluate(BlockedEmitTo::NextBlock)?; - assert_eq!(out.len(), 1); - assert_eq!(out[0].null_count(), 0); - assert_eq!(values(&out), vec![vec![Some(1), Some(2), Some(3), Some(4)]]); - let state = acc.state(BlockedEmitTo::All)?; - assert_eq!(state.len(), 1); - assert_eq!(values(&state[0]), vec![vec![Some(5)]]); - Ok(()) - } -} 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 index 341fbe744e6c3..a982d29f841ce 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs @@ -19,24 +19,12 @@ use std::mem::size_of; -use arrow::array::{Array, ArrowPrimitiveType, BooleanArray, PrimitiveArray}; +use arrow::array::BooleanArray; use arrow::buffer::NullBuffer; use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; -use super::accumulate::{accumulate_blocked, accumulate_blocked_indices}; - -/// Told about the group of every row whose value is used by -/// [`BlockedVec::update_values`], e.g. to track which groups saw a value. -pub trait SeenGroups { - fn mark(&mut self, index: BlocksIndex); -} - -/// Tracks nothing, so marking compiles away. -impl SeenGroups for () { - #[inline(always)] - fn mark(&mut self, _index: BlocksIndex) {} -} +use super::accumulate::accumulate_blocked_indices; /// Rows per chunk whose element addresses are resolved before any of them is /// updated, see [`BlockedVec::update`]. @@ -64,21 +52,10 @@ const RESOLVE_CHUNK: usize = 256; /// /// # Implementation Notes /// -/// ## Items that own heap memory (`T` is not `Copy`) -/// -/// `T` only needs to be `Clone` (for [`Self::grow_to`]), so it can own heap -/// memory, like the `Option>` of `min` / `max` on bytes. Then: +/// ## Why `T: Copy`? /// -/// 1. [`BlockedVec::allocated_size`] only counts the slots (`size_of::()` -/// each), not the heap allocations the items own, which the caller must -/// track itself. -/// 2. The mutable access given by [`Self::update`] and friends can change the -/// heap memory of an item without [`BlockedVec::allocated_size`] changing, -/// which is another reason the caller tracks that memory. -/// 3. Reading an item's heap memory is one more dependent load after the -/// block and the slot, which the address-resolve loops cannot hide. -/// Accumulators should read the stored items as rarely as possible, e.g. -/// once per group per batch instead of once per row. +/// 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)] @@ -93,17 +70,6 @@ pub struct BlockedVec { } impl BlockedVec { - /// 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()] - } -} - -impl BlockedVec { /// Creates an empty vector. /// /// # Panics @@ -153,7 +119,7 @@ impl BlockedVec { 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.clone()); + last.resize(last_len + additional, value); self.len += additional; } } @@ -209,22 +175,13 @@ impl BlockedVec { last } - /// Returns a reference to the element at `index`. - /// - /// # Panics - /// If `index` is out of bounds. - #[inline] - pub fn get_ref(&self, index: BlocksIndex) -> &T { - &self.blocks[index.block_index()][index.index_in_block()] - } - - /// Returns a mutable reference to the element at `index`. + /// Returns the element at `index`. /// /// # Panics /// If `index` is out of bounds. #[inline] - pub fn get_mut(&mut self, index: BlocksIndex) -> &mut T { - &mut self.blocks[index.block_index()][index.index_in_block()] + 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`. @@ -383,88 +340,6 @@ impl BlockedVec { } } - /// Grows the vector to `total_num_groups` elements (new ones are - /// `starting_value`), then, for every row of `values` that is not null - /// and is selected by `opt_filter`, calls `seen.mark` with the row's group - /// and `update_fn` with the group's element and the row's value. - /// - /// Chooses the loop once per call, like [`Self::update`]. - /// - /// `group_indices` must only point to the first `total_num_groups` - /// groups, like for [`Self::update`]. - /// - /// # Panics - /// If `values` and `group_indices` have different lengths. - #[expect(clippy::too_many_arguments, reason = "mirrors `NullState::accumulate`")] - pub fn update_values( - &mut self, - total_num_groups: usize, - starting_value: T, - group_indices: &[BlocksIndex], - values: &PrimitiveArray, - opt_filter: Option<&BooleanArray>, - seen: &mut S, - mut update_fn: F, - ) where - V: ArrowPrimitiveType, - S: SeenGroups, - F: FnMut(&mut T, V::Native), - { - assert_eq!(group_indices.len(), values.len()); - self.grow_to(total_num_groups, starting_value); - self.debug_assert_in_bounds(group_indices); - let all_valid = values.null_count() == 0 && opt_filter.is_none(); - - if let Some(block) = self.as_single_block_mut() { - if all_valid { - for (&index, &value) in group_indices.iter().zip(values.values()) { - seen.mark(index); - // SAFETY: indices are in bounds, see the method docs - update_fn( - unsafe { block.get_unchecked_mut(index.index_in_block()) }, - value, - ); - } - } else { - accumulate_blocked(group_indices, values, opt_filter, |index, value| { - seen.mark(index); - // SAFETY: indices are in bounds, see the method docs - update_fn( - unsafe { block.get_unchecked_mut(index.index_in_block()) }, - value, - ) - }); - } - return; - } - - let ptrs = self.block_ptrs_mut(); - if all_valid { - let mut addrs = [std::ptr::null_mut::(); RESOLVE_CHUNK]; - for (indices, values) in group_indices - .chunks(RESOLVE_CHUNK) - .zip(values.values().chunks(RESOLVE_CHUNK)) - { - for (addr, &index) in addrs.iter_mut().zip(indices) { - // SAFETY: indices are in bounds, see the method docs - *addr = unsafe { ptrs.ptr(index) }; - } - for ((&addr, &index), &value) in addrs.iter().zip(indices).zip(values) { - seen.mark(index); - // SAFETY: resolved above from a block that is not touched - // while `ptrs` lives - update_fn(unsafe { &mut *addr }, value); - } - } - } else { - accumulate_blocked(group_indices, values, opt_filter, |index, value| { - seen.mark(index); - // SAFETY: indices are in bounds, see the method docs - update_fn(unsafe { &mut *ptrs.ptr(index) }, value) - }); - } - } - /// Checks, in debug builds only, that every index points to an element. #[inline] fn debug_assert_in_bounds(&self, group_indices: &[BlocksIndex]) { diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index e0422c1646a2b..0453124528a0f 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -174,75 +174,53 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> Result> { - convert_to_state::( - values, - opt_filter, - self.starting_value, - &self.prim_fn, - &self.data_type, - ) + let values = values[0].as_primitive::().clone(); + + // Initializing state with starting values + let initial_state = + PrimitiveArray::::from_value(self.starting_value, values.len()); + + // Recalculating values in case there is filter + let values = match opt_filter { + None => values, + Some(filter) => { + let (filter_values, filter_nulls) = filter.clone().into_parts(); + // Calculating filter mask as a result of bitand of filter, and converting it to null buffer + let filter_bool = match filter_nulls { + Some(filter_nulls) => filter_nulls.inner() & &filter_values, + None => filter_values, + }; + let filter_nulls = NullBuffer::from(filter_bool); + + // Rebuilding input values with a new nulls mask, which is equal to + // the union of original nulls and filter mask + let (dt, values_buf, original_nulls) = values.into_parts(); + let nulls_buf = + NullBuffer::union(original_nulls.as_ref(), Some(&filter_nulls)); + PrimitiveArray::::new(values_buf, nulls_buf).with_data_type(dt) + } + }; + + let state_values = compute::binary_mut(initial_state, &values, |mut x, y| { + (self.prim_fn)(&mut x, y); + x + }); + let state_values = state_values + .map_err(|_| { + internal_datafusion_err!( + "initial_values underlying buffer must not be shared" + ) + })? + .map_err(DataFusionError::from)? + .with_data_type(self.data_type.clone()); + + Ok(vec![Arc::new(state_values)]) } fn size(&self) -> usize { self.values.capacity() * size_of::() + self.null_state.size() } } -/// Converts an input batch directly to the state of a primitive accumulator -/// (flat or blocked): `prim_fn` applied to `starting_value` and every non -/// null, non filtered value, and null otherwise. -pub(crate) fn convert_to_state( - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - starting_value: T::Native, - prim_fn: &F, - data_type: &DataType, -) -> Result> -where - T: ArrowPrimitiveType, - F: Fn(&mut T::Native, T::Native), -{ - let values = values[0].as_primitive::().clone(); - - // Initializing state with starting values - let initial_state = PrimitiveArray::::from_value(starting_value, values.len()); - - // Recalculating values in case there is filter - let values = match opt_filter { - None => values, - Some(filter) => { - let (filter_values, filter_nulls) = filter.clone().into_parts(); - // Calculating filter mask as a result of bitand of filter, and converting it to null buffer - let filter_bool = match filter_nulls { - Some(filter_nulls) => filter_nulls.inner() & &filter_values, - None => filter_values, - }; - let filter_nulls = NullBuffer::from(filter_bool); - - // Rebuilding input values with a new nulls mask, which is equal to - // the union of original nulls and filter mask - let (dt, values_buf, original_nulls) = values.into_parts(); - let nulls_buf = - NullBuffer::union(original_nulls.as_ref(), Some(&filter_nulls)); - PrimitiveArray::::new(values_buf, nulls_buf).with_data_type(dt) - } - }; - - let state_values = compute::binary_mut(initial_state, &values, |mut x, y| { - prim_fn(&mut x, y); - x - }); - let state_values = state_values - .map_err(|_| { - internal_datafusion_err!( - "initial_values underlying buffer must not be shared" - ) - })? - .map_err(DataFusionError::from)? - .with_data_type(data_type.clone()); - - Ok(vec![Arc::new(state_values)]) -} - #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index 6137010e2b386..85b63a9a17970 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -38,16 +38,14 @@ use datafusion_common::{ use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, - Coercion, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, - ReversedUDAF, Signature, TypeSignature, TypeSignatureClass, Volatility, + Accumulator, AggregateUDFImpl, Coercion, Documentation, EmitTo, Expr, GroupSelection, + GroupsAccumulator, ReversedUDAF, Signature, TypeSignature, TypeSignatureClass, + Volatility, }; use datafusion_functions_aggregate_common::aggregate::avg_distinct::{ DecimalDistinctAvgAccumulator, Float64DistinctAvgAccumulator, }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::NullState; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_null_state::BlockedNullState; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::BlockedVec; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{ filtered_null_mask, set_nulls, }; @@ -486,63 +484,6 @@ impl AggregateUDFImpl for Avg { } } - fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - matches!( - ( - args.expr_fields[0].data_type(), - args.return_field.data_type() - ), - (DataType::Float64, DataType::Float64) - | (DataType::Duration(_), DataType::Duration(_)) - ) && !args.is_distinct - } - - fn create_blocked_groups_accumulator( - &self, - args: AccumulatorArgs, - block_size: usize, - ) -> Result> { - use DataType::*; - - let data_type = args.expr_fields[0].data_type(); - let return_type = args.return_field.data_type(); - let duration_avg_fn = |sum: i64, count: u64| Ok(sum / count as i64); - macro_rules! duration { - ($t:ty) => { - Ok(Box::new(BlockedAvgGroupsAccumulator::<$t, _>::new( - data_type, - return_type, - block_size, - duration_avg_fn, - ))) - }; - } - match (data_type, return_type) { - (Float64, Float64) => Ok(Box::new(BlockedAvgGroupsAccumulator::< - Float64Type, - _, - >::new( - data_type, - return_type, - block_size, - |sum: f64, count: u64| Ok(sum / count as f64), - ))), - (Duration(TimeUnit::Second), Duration(_)) => duration!(DurationSecondType), - (Duration(TimeUnit::Millisecond), Duration(_)) => { - duration!(DurationMillisecondType) - } - (Duration(TimeUnit::Microsecond), Duration(_)) => { - duration!(DurationMicrosecondType) - } - (Duration(TimeUnit::Nanosecond), Duration(_)) => { - duration!(DurationNanosecondType) - } - _ => not_impl_err!( - "BlockedAvgGroupsAccumulator for ({data_type} --> {return_type})" - ), - } - } - fn aliases(&self) -> &[String] { &self.aliases } @@ -676,113 +617,6 @@ where ))) } -/// Computes the averages of groups, `null` for groups whose `nulls` bit is -/// unset. Shared by [`AvgGroupsAccumulator`] and -/// [`BlockedAvgGroupsAccumulator`]. -fn evaluate_avg_values( - counts: Vec, - sums: Vec, - nulls: Option, - avg_fn: &F, - return_data_type: &DataType, -) -> Result -where - O: ArrowNumericType, - S: ArrowNumericType, - F: Fn(S::Native, u64) -> Result, -{ - if let Some(nulls) = &nulls { - assert_eq!(nulls.len(), sums.len()); - } - assert_eq!(counts.len(), sums.len()); - - // Don't evaluate averages with null inputs to avoid errors on null values. - let array: PrimitiveArray = if let Some(nulls) = &nulls - && nulls.null_count() > 0 - { - let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) - .with_data_type(return_data_type.clone()); - let iter = sums.into_iter().zip(counts).zip(nulls.iter()); - - for ((sum, count), is_valid) in iter { - if is_valid { - builder.append_value(avg_fn(sum, count)?) - } else { - builder.append_null(); - } - } - builder.finish() - } else { - let averages: Vec = sums - .into_iter() - .zip(counts) - .map(|(sum, count)| avg_fn(sum, count)) - .collect::>>()?; - PrimitiveArray::new(averages.into(), nulls) - .with_data_type(return_data_type.clone()) - }; - - Ok(Arc::new(array)) -} - -/// The `[counts, sums]` state of groups. Shared by [`AvgGroupsAccumulator`] -/// and [`BlockedAvgGroupsAccumulator`]. -fn avg_state_values( - counts: Vec, - sums: Vec, - nulls: Option, - sum_data_type: &DataType, -) -> Vec { - let counts = UInt64Array::new(counts.into(), nulls.clone()); - let sums = PrimitiveArray::::new(sums.into(), nulls) - .with_data_type(sum_data_type.clone()); - vec![Arc::new(counts), Arc::new(sums)] -} - -/// Converts an input batch directly to `[counts, sums]` state. Shared by -/// [`AvgGroupsAccumulator`] and [`BlockedAvgGroupsAccumulator`]. -fn avg_convert_to_state( - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - sum_data_type: &DataType, -) -> Result> -where - I: ArrowNumericType, - S: ArrowNumericType + Send, - I::Native: Into, -{ - // When the sum type equals the input type (`I == S`: `Float64`, - // `Duration`, `Decimal256`, and any decimal whose precision already - // leaves [`avg_sum_data_type`] enough headroom) the input is already a - // valid sum array and is reused as is; the downcast is by Rust type, so - // it succeeds even when precision differs. Otherwise every value is - // widened. - let sums = match values[0].as_any().downcast_ref::>() { - Some(sums) => sums.clone().with_data_type(sum_data_type.clone()), - None => { - let values = values[0].as_primitive::(); - // Values under null slots are widened too rather than branching per - // element; `set_nulls` below masks them out again. - let sums: Vec = values - .values() - .iter() - .map(|value| (*value).into()) - .collect(); - PrimitiveArray::::new(sums.into(), values.nulls().cloned()) - .with_data_type(sum_data_type.clone()) - } - }; - let counts = UInt64Array::from_value(1, sums.len()); - - let nulls = filtered_null_mask(opt_filter, &sums); - - // set nulls on the arrays - let counts = set_nulls(counts, nulls.clone()); - let sums = set_nulls(sums, nulls); - - Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)]) -} - /// An accumulator to compute the average #[derive(Debug, Default)] pub struct AvgAccumulator { @@ -1206,13 +1040,38 @@ where sums: Vec, nulls: Option, ) -> Result { - evaluate_avg_values::( - counts, - sums, - nulls, - &self.avg_fn, - &self.return_data_type, - ) + if let Some(nulls) = &nulls { + assert_eq!(nulls.len(), sums.len()); + } + assert_eq!(counts.len(), sums.len()); + + // Don't evaluate averages with null inputs to avoid errors on null values. + let array: PrimitiveArray = if let Some(nulls) = &nulls + && nulls.null_count() > 0 + { + let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) + .with_data_type(self.return_data_type.clone()); + let iter = sums.into_iter().zip(counts).zip(nulls.iter()); + + for ((sum, count), is_valid) in iter { + if is_valid { + builder.append_value((self.avg_fn)(sum, count)?) + } else { + builder.append_null(); + } + } + builder.finish() + } else { + let averages: Vec = sums + .into_iter() + .zip(counts) + .map(|(sum, count)| (self.avg_fn)(sum, count)) + .collect::>>()?; + PrimitiveArray::new(averages.into(), nulls) + .with_data_type(self.return_data_type.clone()) + }; + + Ok(Arc::new(array)) } fn state_values( @@ -1221,7 +1080,10 @@ where sums: Vec, nulls: Option, ) -> Vec { - avg_state_values::(counts, sums, nulls, &self.sum_data_type) + let counts = UInt64Array::new(counts.into(), nulls.clone()); + let sums = PrimitiveArray::::new(sums.into(), nulls) + .with_data_type(self.sum_data_type.clone()); + vec![Arc::new(counts), Arc::new(sums)] } } @@ -1354,7 +1216,36 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> Result> { - avg_convert_to_state::(values, opt_filter, &self.sum_data_type) + // When the sum type equals the input type (`I == S`: `Float64`, + // `Duration`, `Decimal256`, and any decimal whose precision already + // leaves [`avg_sum_data_type`] enough headroom) the input is already a + // valid sum array and is reused as is; the downcast is by Rust type, so + // it succeeds even when precision differs. Otherwise every value is + // widened. + let sums = match values[0].as_any().downcast_ref::>() { + Some(sums) => sums.clone().with_data_type(self.sum_data_type.clone()), + None => { + let values = values[0].as_primitive::(); + // Values under null slots are widened too rather than branching per + // element; `set_nulls` below masks them out again. + let sums: Vec = values + .values() + .iter() + .map(|value| (*value).into()) + .collect(); + PrimitiveArray::::new(sums.into(), values.nulls().cloned()) + .with_data_type(self.sum_data_type.clone()) + } + }; + let counts = UInt64Array::from_value(1, sums.len()); + + let nulls = filtered_null_mask(opt_filter, &sums); + + // set nulls on the arrays + let counts = set_nulls(counts, nulls.clone()); + let sums = set_nulls(sums, nulls); + + Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)]) } fn size(&self) -> usize { // Heap buffers @@ -1365,249 +1256,9 @@ where } } -/// [`AvgGroupsAccumulator`] with sums and counts stored in blocks. -/// -/// Sums and counts are kept in two [`BlockedVec`]s and updated in two passes, -/// so each pass keeps the fast loops of [`BlockedVec`]. -#[derive(Debug)] -struct BlockedAvgGroupsAccumulator -where - I: ArrowNumericType + Send, - O: ArrowNumericType + Send, - S: ArrowNumericType + Send, - I::Native: Into, - F: Fn(S::Native, u64) -> Result + Send + 'static, -{ - /// The type of the internal sum - sum_data_type: DataType, - - /// The type of the returned average - return_data_type: DataType, - - /// Count per group (use u64 to make UInt64Array) - counts: BlockedVec, - - /// Sums per group, stored as the native type - sums: BlockedVec, - - /// Track nulls in the input / filters - null_state: BlockedNullState, - - /// Function that computes the final average (value / count) - avg_fn: F, - - _phantom: PhantomData<(I, O)>, -} - -impl BlockedAvgGroupsAccumulator -where - I: ArrowNumericType + Send, - O: ArrowNumericType + Send, - S: ArrowNumericType + Send, - I::Native: Into, - F: Fn(S::Native, u64) -> Result + Send + 'static, -{ - fn new( - sum_data_type: &DataType, - return_data_type: &DataType, - block_size: usize, - avg_fn: F, - ) -> Self { - Self { - return_data_type: return_data_type.clone(), - sum_data_type: sum_data_type.clone(), - counts: BlockedVec::new(block_size), - sums: BlockedVec::new(block_size), - null_state: BlockedNullState::new(block_size), - avg_fn, - _phantom: PhantomData, - } - } - - /// Removes the groups selected by `emit_to`, as `(counts, sums, nulls)` - /// of every emitted block. - fn take( - &mut self, - emit_to: BlockedEmitTo, - ) -> impl Iterator, Vec, Option)> + use - { - fn take( - v: &mut BlockedVec, - emit_to: BlockedEmitTo, - ) -> Vec> { - match emit_to { - BlockedEmitTo::All => v.take_all(), - BlockedEmitTo::NextBlock => v.take_next_block().into_iter().collect(), - BlockedEmitTo::First(n) => vec![v.take_first(n)], - } - } - let counts = take(&mut self.counts, emit_to); - let sums = take(&mut self.sums, emit_to); - let nulls = self.null_state.build(emit_to); - debug_assert_eq!(counts.len(), sums.len()); - debug_assert_eq!(counts.len(), nulls.len()); - counts - .into_iter() - .zip(sums) - .zip(nulls) - .map(|((counts, sums), nulls)| (counts, sums, nulls)) - } -} - -impl BlockedGroupsAccumulator for BlockedAvgGroupsAccumulator -where - I: ArrowNumericType + Send, - O: ArrowNumericType + Send, - S: ArrowNumericType + Send, - I::Native: Into, - F: Fn(S::Native, u64) -> Result + Send + 'static, -{ - fn block_size(&self) -> usize { - self.sums.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].as_primitive::(); - - let nulls = values.logical_nulls().filter(|n| n.null_count() > 0); - self.counts.update( - total_num_groups, - 0, - group_indices, - nulls.as_ref(), - opt_filter, - |count| *count += 1, - ); - self.null_state.accumulate( - &mut self.sums, - total_num_groups, - S::default_value(), - group_indices, - values, - opt_filter, - |sum, new_value| *sum = add_avg_sum::(*sum, new_value), - ); - Ok(()) - } - - fn merge_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[BlocksIndex], - total_num_groups: usize, - ) -> Result<()> { - assert_eq!(values.len(), 2, "two arguments to merge_batch"); - // first batch is counts, second is partial sums - let partial_counts = values[0].as_primitive::(); - let partial_sums = values[1].as_primitive::(); - self.null_state.accumulate( - &mut self.counts, - total_num_groups, - 0, - group_indices, - partial_counts, - None, - |count, partial_count| *count += partial_count, - ); - self.null_state.accumulate( - &mut self.sums, - total_num_groups, - S::default_value(), - group_indices, - partial_sums, - None, - |sum, partial_sum| *sum = add_avg_sum::(*sum, partial_sum), - ); - Ok(()) - } - - fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { - let blocks: Vec<_> = self.take(emit_to).collect(); - blocks - .into_iter() - .map(|(counts, sums, nulls)| { - evaluate_avg_values::( - counts, - sums, - nulls, - &self.avg_fn, - &self.return_data_type, - ) - }) - .collect() - } - - fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { - let blocks: Vec<_> = self.take(emit_to).collect(); - Ok(blocks - .into_iter() - .map(|(counts, sums, nulls)| { - avg_state_values::(counts, sums, nulls, &self.sum_data_type) - }) - .collect()) - } - - fn convert_to_state( - &self, - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - ) -> Result> { - avg_convert_to_state::(values, opt_filter, &self.sum_data_type) - } - - fn size(&self) -> usize { - self.counts.allocated_size() + self.sums.allocated_size() + self.null_state.size() - } -} - #[cfg(test)] mod tests { use super::*; - - #[test] - fn blocked_avg_across_blocks_with_nulls_and_merge() -> Result<()> { - let mut acc = BlockedAvgGroupsAccumulator::::new( - &DataType::Float64, - &DataType::Float64, - 4, - |sum: f64, count: u64| Ok(sum / count as f64), - ); - let groups = |flats: &[usize]| -> Vec { - flats - .iter() - .map(|&f| BlocksIndex::from_flat(f, 4)) - .collect() - }; - let values: ArrayRef = Arc::new(Float64Array::from(vec![ - Some(1.0), - Some(3.0), - None, - Some(10.0), - ])); - acc.update_batch(&[values], &groups(&[0, 0, 1, 5]), None, 6)?; - // partial state for group 5: count 2, sum 20 - let counts: ArrayRef = Arc::new(UInt64Array::from(vec![2])); - let sums: ArrayRef = Arc::new(Float64Array::from(vec![20.0])); - acc.merge_batch(&[counts, sums], &groups(&[5]), 6)?; - - let out = acc.evaluate(BlockedEmitTo::All)?; - let out: Vec>> = out - .iter() - .map(|a| a.as_primitive::().iter().collect()) - .collect(); - assert_eq!( - out, - vec![vec![Some(2.0), None, None, None], vec![None, Some(10.0)]] - ); - Ok(()) - } use arrow::array::{ Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, DurationSecondArray, Float64Array, diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 8e69bfa22f294..b8a7ce85466c2 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -43,10 +43,7 @@ use datafusion_expr::{ function::{AccumulatorArgs, StateFieldsArgs}, utils::{AggregateOrderSensitivity, format_state_name}, }; -use datafusion_functions_aggregate_common::aggregate::count_distinct::{ - BlockedPrimitiveDistinctCountGroupsAccumulator, - PrimitiveDistinctCountGroupsAccumulator, -}; +use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::{ count_distinct::Bitmap65536DistinctCountAccumulator, count_distinct::Bitmap65536DistinctCountAccumulatorI16, @@ -398,37 +395,15 @@ impl AggregateUDFImpl for Count { } fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - self.groups_accumulator_supported(args) + !args.is_distinct && self.groups_accumulator_supported(args) } fn create_blocked_groups_accumulator( &self, - args: AccumulatorArgs, + _args: AccumulatorArgs, block_size: usize, ) -> Result> { - if !args.is_distinct { - return Ok(Box::new(BlockedCountGroupsAccumulator::new(block_size))); - } - macro_rules! distinct { - ($t:ty) => { - Ok(Box::new( - BlockedPrimitiveDistinctCountGroupsAccumulator::<$t>::new(block_size), - )) - }; - } - match args.expr_fields[0].data_type() { - DataType::Int8 => distinct!(Int8Type), - DataType::Int16 => distinct!(Int16Type), - DataType::Int32 => distinct!(Int32Type), - DataType::Int64 => distinct!(Int64Type), - DataType::UInt8 => distinct!(UInt8Type), - DataType::UInt16 => distinct!(UInt16Type), - DataType::UInt32 => distinct!(UInt32Type), - DataType::UInt64 => distinct!(UInt64Type), - data_type => not_impl_err!( - "BlockedGroupsAccumulator not supported for COUNT(DISTINCT) with {data_type}" - ), - } + Ok(Box::new(BlockedCountGroupsAccumulator::new(block_size))) } fn reverse_expr(&self) -> ReversedUDAF { diff --git a/datafusion/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs index 50e5f26611c2a..94814a802d971 100644 --- a/datafusion/functions-aggregate/src/min_max.rs +++ b/datafusion/functions-aggregate/src/min_max.rs @@ -42,9 +42,7 @@ use arrow::datatypes::{ TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; -use crate::min_max::min_max_bytes::{ - BlockedMinMaxBytesAccumulator, MinMaxBytesAccumulator, -}; +use crate::min_max::min_max_bytes::MinMaxBytesAccumulator; use crate::min_max::min_max_struct::MinMaxStructAccumulator; use datafusion_common::ScalarValue; use datafusion_expr::DistinctHandling; @@ -52,7 +50,7 @@ use datafusion_expr::{ Accumulator, AggregateUDFImpl, Documentation, SetMonotonicity, Signature, Volatility, function::AccumulatorArgs, }; -use datafusion_expr::{BlockedGroupsAccumulator, GroupsAccumulator, StatisticsArgs}; +use datafusion_expr::{GroupsAccumulator, StatisticsArgs}; use datafusion_macros::user_doc; use half::f16; use std::collections::VecDeque; @@ -280,25 +278,6 @@ impl AggregateUDFImpl for Max { ) } - fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - use DataType::*; - matches!( - args.return_field.data_type(), - Utf8 | LargeUtf8 | Utf8View | Binary | LargeBinary | BinaryView - ) - } - - fn create_blocked_groups_accumulator( - &self, - args: AccumulatorArgs, - block_size: usize, - ) -> Result> { - Ok(Box::new(BlockedMinMaxBytesAccumulator::new_max( - args.return_field.data_type().clone(), - block_size, - ))) - } - fn create_groups_accumulator( &self, args: AccumulatorArgs, @@ -598,25 +577,6 @@ impl AggregateUDFImpl for Min { ) } - fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - use DataType::*; - matches!( - args.return_field.data_type(), - Utf8 | LargeUtf8 | Utf8View | Binary | LargeBinary | BinaryView - ) - } - - fn create_blocked_groups_accumulator( - &self, - args: AccumulatorArgs, - block_size: usize, - ) -> Result> { - Ok(Box::new(BlockedMinMaxBytesAccumulator::new_min( - args.return_field.data_type().clone(), - block_size, - ))) - } - fn create_groups_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index 29e1d4ec478b7..7d26143eb5790 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -22,11 +22,7 @@ use arrow::array::{ use arrow::datatypes::DataType; use datafusion_common::hash_map::Entry; use datafusion_common::{HashMap, Result, internal_err}; -use datafusion_expr::{ - BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, EmitTo, GroupSelection, - GroupsAccumulator, -}; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::BlockedVec; +use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; use std::mem::size_of; use std::sync::Arc; @@ -73,101 +69,93 @@ impl MinMaxBytesAccumulator { num_values: usize, data_capacity: usize, ) -> Result { - build_bytes_array(&self.inner.data_type, min_maxes, num_values, data_capacity) - } -} - -/// Builds an array of `data_type` from the min/max values of groups. Shared -/// by the flat and blocked accumulators. -fn build_bytes_array<'a>( - data_type: &DataType, - min_maxes: impl Iterator>, - num_values: usize, - data_capacity: usize, -) -> Result { - let result: ArrayRef = match *data_type { - DataType::Utf8 => { - let mut builder = StringBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - // SAFETY: update_batch only accepts the configured input type. - Some(value) => builder - .append_value(unsafe { std::str::from_utf8_unchecked(value) }), + let result: ArrayRef = match self.inner.data_type { + DataType::Utf8 => { + let mut builder = StringBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder.append_value(unsafe { + std::str::from_utf8_unchecked(value) + }), + } } + Arc::new(builder.finish()) } - Arc::new(builder.finish()) - } - DataType::LargeUtf8 => { - let mut builder = - LargeStringBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - // SAFETY: update_batch only accepts the configured input type. - Some(value) => builder - .append_value(unsafe { std::str::from_utf8_unchecked(value) }), + DataType::LargeUtf8 => { + let mut builder = + LargeStringBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder.append_value(unsafe { + std::str::from_utf8_unchecked(value) + }), + } } + Arc::new(builder.finish()) } - Arc::new(builder.finish()) - } - DataType::Utf8View => { - let block_size = capacity_to_view_block_size(data_capacity); - let mut builder = StringViewBuilder::with_capacity(num_values) - .with_fixed_block_size(block_size); - for value in min_maxes { - match value { - None => builder.append_null(), - // SAFETY: update_batch only accepts the configured input type. - Some(value) => builder - .append_value(unsafe { std::str::from_utf8_unchecked(value) }), + DataType::Utf8View => { + let block_size = capacity_to_view_block_size(data_capacity); + let mut builder = StringViewBuilder::with_capacity(num_values) + .with_fixed_block_size(block_size); + for value in min_maxes { + match value { + None => builder.append_null(), + // SAFETY: update_batch only accepts the configured input type. + Some(value) => builder.append_value(unsafe { + std::str::from_utf8_unchecked(value) + }), + } } + Arc::new(builder.finish()) } - Arc::new(builder.finish()) - } - DataType::Binary => { - let mut builder = BinaryBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - Some(value) => builder.append_value(value), + DataType::Binary => { + let mut builder = BinaryBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), + } } + Arc::new(builder.finish()) } - Arc::new(builder.finish()) - } - DataType::LargeBinary => { - let mut builder = - LargeBinaryBuilder::with_capacity(num_values, data_capacity); - for value in min_maxes { - match value { - None => builder.append_null(), - Some(value) => builder.append_value(value), + DataType::LargeBinary => { + let mut builder = + LargeBinaryBuilder::with_capacity(num_values, data_capacity); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), + } } + Arc::new(builder.finish()) } - Arc::new(builder.finish()) - } - DataType::BinaryView => { - let block_size = capacity_to_view_block_size(data_capacity); - let mut builder = BinaryViewBuilder::with_capacity(num_values) - .with_fixed_block_size(block_size); - for value in min_maxes { - match value { - None => builder.append_null(), - Some(value) => builder.append_value(value), + DataType::BinaryView => { + let block_size = capacity_to_view_block_size(data_capacity); + let mut builder = BinaryViewBuilder::with_capacity(num_values) + .with_fixed_block_size(block_size); + for value in min_maxes { + match value { + None => builder.append_null(), + Some(value) => builder.append_value(value), + } } + Arc::new(builder.finish()) } - Arc::new(builder.finish()) - } - _ => { - return internal_err!( - "Unexpected data type for MinMaxBytesAccumulator: {:?}", - *data_type - ); - } - }; + _ => { + return internal_err!( + "Unexpected data type for MinMaxBytesAccumulator: {:?}", + self.inner.data_type + ); + } + }; - assert_eq!(data_type, result.data_type()); - Ok(result) + assert_eq!(&self.inner.data_type, result.data_type()); + Ok(result) + } } impl GroupsAccumulator for MinMaxBytesAccumulator { @@ -185,13 +173,126 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { // apply filter if needed let array = apply_filter_as_nulls(array, opt_filter)?; - dispatch_update( - &mut self.inner, - self.is_min, - &array, - group_indices, - total_num_groups, - ) + // dispatch to appropriate kernel / specialized implementation + fn string_min(a: &[u8], b: &[u8]) -> bool { + // safety: only called from this function, which ensures a and b come + // from an array with valid utf8 data + unsafe { + let a = std::str::from_utf8_unchecked(a); + let b = std::str::from_utf8_unchecked(b); + a < b + } + } + fn string_max(a: &[u8], b: &[u8]) -> bool { + // safety: only called from this function, which ensures a and b come + // from an array with valid utf8 data + unsafe { + let a = std::str::from_utf8_unchecked(a); + let b = std::str::from_utf8_unchecked(b); + a > b + } + } + fn binary_min(a: &[u8], b: &[u8]) -> bool { + a < b + } + + fn binary_max(a: &[u8], b: &[u8]) -> bool { + a > b + } + + fn str_to_bytes<'a>( + it: impl Iterator>, + ) -> impl Iterator> { + it.map(|s| s.map(|s| s.as_bytes())) + } + + match (self.is_min, &self.inner.data_type) { + // Utf8/LargeUtf8/Utf8View Min + (true, &DataType::Utf8) => self.inner.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_min, + ), + (true, &DataType::LargeUtf8) => self.inner.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_min, + ), + (true, &DataType::Utf8View) => self.inner.update_batch( + str_to_bytes(array.as_string_view().iter()), + group_indices, + total_num_groups, + string_min, + ), + + // Utf8/LargeUtf8/Utf8View Max + (false, &DataType::Utf8) => self.inner.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_max, + ), + (false, &DataType::LargeUtf8) => self.inner.update_batch( + str_to_bytes(array.as_string::().iter()), + group_indices, + total_num_groups, + string_max, + ), + (false, &DataType::Utf8View) => self.inner.update_batch( + str_to_bytes(array.as_string_view().iter()), + group_indices, + total_num_groups, + string_max, + ), + + // Binary/LargeBinary/BinaryView Min + (true, &DataType::Binary) => self.inner.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_min, + ), + (true, &DataType::LargeBinary) => self.inner.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_min, + ), + (true, &DataType::BinaryView) => self.inner.update_batch( + array.as_binary_view().iter(), + group_indices, + total_num_groups, + binary_min, + ), + + // Binary/LargeBinary/BinaryView Max + (false, &DataType::Binary) => self.inner.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_max, + ), + (false, &DataType::LargeBinary) => self.inner.update_batch( + array.as_binary::().iter(), + group_indices, + total_num_groups, + binary_max, + ), + (false, &DataType::BinaryView) => self.inner.update_batch( + array.as_binary_view().iter(), + group_indices, + total_num_groups, + binary_max, + ), + + _ => internal_err!( + "Unexpected combination for MinMaxBytesAccumulator: ({:?}, {:?})", + self.is_min, + self.inner.data_type + ), + } } fn evaluate(&mut self, emit_to: EmitTo) -> Result { @@ -262,159 +363,6 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { } } -/// Updates the min/max values of a flat or blocked state, see -/// [`MinMaxBytesUpdate`]. -fn dispatch_update( - state: &mut S, - is_min: bool, - array: &ArrayRef, - group_indices: &[Idx], - total_num_groups: usize, -) -> Result<()> -where - S: MinMaxBytesUpdate, -{ - let data_type = state.data_type().clone(); - // dispatch to appropriate kernel / specialized implementation - fn string_min(a: &[u8], b: &[u8]) -> bool { - // safety: only called from this function, which ensures a and b come - // from an array with valid utf8 data - unsafe { - let a = std::str::from_utf8_unchecked(a); - let b = std::str::from_utf8_unchecked(b); - a < b - } - } - fn string_max(a: &[u8], b: &[u8]) -> bool { - // safety: only called from this function, which ensures a and b come - // from an array with valid utf8 data - unsafe { - let a = std::str::from_utf8_unchecked(a); - let b = std::str::from_utf8_unchecked(b); - a > b - } - } - fn binary_min(a: &[u8], b: &[u8]) -> bool { - a < b - } - - fn binary_max(a: &[u8], b: &[u8]) -> bool { - a > b - } - - fn str_to_bytes<'a>( - it: impl Iterator>, - ) -> impl Iterator> { - it.map(|s| s.map(|s| s.as_bytes())) - } - - match (is_min, &data_type) { - // Utf8/LargeUtf8/Utf8View Min - (true, &DataType::Utf8) => state.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_min, - ), - (true, &DataType::LargeUtf8) => state.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_min, - ), - (true, &DataType::Utf8View) => state.update_batch( - str_to_bytes(array.as_string_view().iter()), - group_indices, - total_num_groups, - string_min, - ), - - // Utf8/LargeUtf8/Utf8View Max - (false, &DataType::Utf8) => state.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_max, - ), - (false, &DataType::LargeUtf8) => state.update_batch( - str_to_bytes(array.as_string::().iter()), - group_indices, - total_num_groups, - string_max, - ), - (false, &DataType::Utf8View) => state.update_batch( - str_to_bytes(array.as_string_view().iter()), - group_indices, - total_num_groups, - string_max, - ), - - // Binary/LargeBinary/BinaryView Min - (true, &DataType::Binary) => state.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_min, - ), - (true, &DataType::LargeBinary) => state.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_min, - ), - (true, &DataType::BinaryView) => state.update_batch( - array.as_binary_view().iter(), - group_indices, - total_num_groups, - binary_min, - ), - - // Binary/LargeBinary/BinaryView Max - (false, &DataType::Binary) => state.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_max, - ), - (false, &DataType::LargeBinary) => state.update_batch( - array.as_binary::().iter(), - group_indices, - total_num_groups, - binary_max, - ), - (false, &DataType::BinaryView) => state.update_batch( - array.as_binary_view().iter(), - group_indices, - total_num_groups, - binary_max, - ), - - _ => internal_err!( - "Unexpected combination for MinMaxBytesAccumulator: ({:?}, {:?})", - is_min, - data_type - ), - } -} - -/// Min/max state updated by [`dispatch_update`]. -trait MinMaxBytesUpdate { - fn data_type(&self) -> &DataType; - - /// `cmp` is called like `cmp(new_val, existing_val)` and returns true if - /// `new_val` should replace `existing_val` - fn update_batch<'a, F, I>( - &mut self, - iter: I, - group_indices: &[Idx], - total_num_groups: usize, - cmp: F, - ) -> Result<()> - where - F: FnMut(&[u8], &[u8]) -> bool + Send + Sync, - I: IntoIterator>; -} - /// Returns the block size in (contiguous buffer size) to use /// for a given data capacity (total string length) /// @@ -505,42 +453,6 @@ impl MinMaxBytesState { } } - /// Emits the specified min_max values - /// - /// Returns (data_capacity, min_maxes), updating the current value of total_data_bytes - /// - /// - `data_capacity`: the total length of all strings and their contents, - /// - `min_maxes`: the actual min/max values for each group - fn emit_to(&mut self, emit_to: EmitTo) -> (usize, Vec>>) { - match emit_to { - EmitTo::All => { - ( - std::mem::take(&mut self.total_data_bytes), // reset total bytes and min_max - std::mem::take(&mut self.min_max), - ) - } - EmitTo::First(n) => { - let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); - let first_data_capacity: usize = first_min_maxes - .iter() - .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) - .sum(); - self.total_data_bytes -= first_data_capacity; - (first_data_capacity, first_min_maxes) - } - } - } - - fn size(&self) -> usize { - self.total_data_bytes + self.min_max.len() * size_of::>>() - } -} - -impl MinMaxBytesUpdate for MinMaxBytesState { - fn data_type(&self) -> &DataType { - &self.data_type - } - /// Updates the min/max values for the given string values /// /// `cmp` is the comparison function to use, called like `cmp(new_val, existing_val)` @@ -594,272 +506,41 @@ impl MinMaxBytesUpdate for MinMaxBytesState { Ok(()) } -} - -/// [`MinMaxBytesState`] with the value of each group stored in a -/// [`BlockedVec`]. -/// -/// Reading a stored value is a load from the heap after the block and the -/// slot, which the address-resolve loops of [`BlockedVec`] cannot hide. -/// [`Self::update_batch`] follows the same algorithm as -/// [`MinMaxBytesState::update_batch`]. -#[derive(Debug)] -struct BlockedMinMaxBytesState { - /// The minimum/maximum value for each group - min_max: BlockedVec>>, - data_type: DataType, - /// The total bytes of the string data (for pre-allocating the final array, - /// and tracking memory usage) - total_data_bytes: usize, -} - -impl BlockedMinMaxBytesState { - fn new(data_type: DataType, block_size: usize) -> Self { - Self { - min_max: BlockedVec::new(block_size), - data_type, - total_data_bytes: 0, - } - } - - /// Removes the groups selected by `emit_to`, one `Vec` per block. - fn emit_to(&mut self, emit_to: BlockedEmitTo) -> Vec>>> { - let blocks = match emit_to { - BlockedEmitTo::All => self.min_max.take_all(), - BlockedEmitTo::NextBlock => { - self.min_max.take_next_block().into_iter().collect() - } - BlockedEmitTo::First(n) => vec![self.min_max.take_first(n)], - }; - for value in blocks.iter().flatten() { - self.total_data_bytes -= value.as_ref().map_or(0, Vec::len); - } - blocks - } - - fn size(&self) -> usize { - self.min_max.allocated_size() + self.total_data_bytes - } -} - -impl MinMaxBytesUpdate for BlockedMinMaxBytesState { - fn data_type(&self) -> &DataType { - &self.data_type - } - - fn update_batch<'a, F, I>( - &mut self, - iter: I, - group_indices: &[BlocksIndex], - total_num_groups: usize, - mut cmp: F, - ) -> Result<()> - where - F: FnMut(&[u8], &[u8]) -> bool + Send + Sync, - I: IntoIterator>, - { - self.min_max.grow_to(total_num_groups, None); - // Minimize value copies by calculating the new min/maxes for each group - // in this batch (either the existing min/max or the new input value) - // and updating the owned values in `self.min_max` at most once - let mut locations = - HashMap::::with_capacity(group_indices.len()); - // Figure out the new min value for each group - for (new_val, &group_index) in iter.into_iter().zip(group_indices) { - let Some(new_val) = new_val else { - continue; // skip nulls - }; - - match locations.entry(group_index) { - Entry::Occupied(mut occupied_entry) => { - if cmp(new_val, occupied_entry.get()) { - occupied_entry.insert(new_val); - } - } - Entry::Vacant(vacant_entry) => { - if let Some(old_val) = self.min_max.get_ref(group_index) { - if cmp(new_val, old_val) { - vacant_entry.insert(new_val); - } - } else { - vacant_entry.insert(new_val); - } - } - } - } - - // Update `self.min_max` with any new min/max values found in the input - for (group_index, new_val) in locations { - match self.min_max.get_mut(group_index) { - existing @ None => { - *existing = Some(new_val.to_vec()); - self.total_data_bytes += new_val.len(); - } - Some(existing_val) => { - // Copy data over to avoid re-allocating - self.total_data_bytes -= existing_val.len(); - self.total_data_bytes += new_val.len(); - existing_val.clear(); - existing_val.extend_from_slice(new_val); - } + /// Emits the specified min_max values + /// + /// Returns (data_capacity, min_maxes), updating the current value of total_data_bytes + /// + /// - `data_capacity`: the total length of all strings and their contents, + /// - `min_maxes`: the actual min/max values for each group + fn emit_to(&mut self, emit_to: EmitTo) -> (usize, Vec>>) { + match emit_to { + EmitTo::All => { + ( + std::mem::take(&mut self.total_data_bytes), // reset total bytes and min_max + std::mem::take(&mut self.min_max), + ) } - } - Ok(()) - } -} - -/// [`MinMaxBytesAccumulator`] with state stored in blocks. -#[derive(Debug)] -pub(crate) struct BlockedMinMaxBytesAccumulator { - inner: BlockedMinMaxBytesState, - /// if true, is `MIN` otherwise is `MAX` - is_min: bool, -} - -impl BlockedMinMaxBytesAccumulator { - /// Create a new accumulator for computing `min(val)` - pub fn new_min(data_type: DataType, block_size: usize) -> Self { - Self { - inner: BlockedMinMaxBytesState::new(data_type, block_size), - is_min: true, - } - } - - /// Create a new accumulator for computing `max(val)` - pub fn new_max(data_type: DataType, block_size: usize) -> Self { - Self { - inner: BlockedMinMaxBytesState::new(data_type, block_size), - is_min: false, - } - } -} - -impl BlockedGroupsAccumulator for BlockedMinMaxBytesAccumulator { - fn block_size(&self) -> usize { - self.inner.min_max.block_size() - } - - fn update_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[BlocksIndex], - opt_filter: Option<&BooleanArray>, - total_num_groups: usize, - ) -> Result<()> { - let array = &values[0]; - assert_eq!(array.len(), group_indices.len()); - assert_eq!(array.data_type(), &self.inner.data_type); - - // apply filter if needed - let array = apply_filter_as_nulls(array, opt_filter)?; - - dispatch_update( - &mut self.inner, - self.is_min, - &array, - group_indices, - total_num_groups, - ) - } - - fn merge_batch( - &mut self, - values: &[ArrayRef], - group_indices: &[BlocksIndex], - total_num_groups: usize, - ) -> Result<()> { - // min/max are their own states (no transition needed) - self.update_batch(values, group_indices, None, total_num_groups) - } - - fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { - let data_type = self.inner.data_type.clone(); - self.inner - .emit_to(emit_to) - .into_iter() - .map(|block| { - let data_capacity = block + EmitTo::First(n) => { + let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); + let first_data_capacity: usize = first_min_maxes .iter() - .map(|value| value.as_ref().map_or(0, Vec::len)) + .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) .sum(); - build_bytes_array( - &data_type, - block.iter().map(|value| value.as_deref()), - block.len(), - data_capacity, - ) - }) - .collect() - } - - fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { - // min/max are their own states (no transition needed) - Ok(self - .evaluate(emit_to)? - .into_iter() - .map(|array| vec![array]) - .collect()) - } - - fn convert_to_state( - &self, - values: &[ArrayRef], - opt_filter: Option<&BooleanArray>, - ) -> Result> { - // Min/max do not change the values as they are their own states - // apply the filter by combining with the null mask, if any - let output = apply_filter_as_nulls(&values[0], opt_filter)?; - Ok(vec![output]) + self.total_data_bytes -= first_data_capacity; + (first_data_capacity, first_min_maxes) + } + } } fn size(&self) -> usize { - self.inner.size() + self.total_data_bytes + self.min_max.len() * size_of::>>() } } #[cfg(test)] mod tests { use super::*; - - #[test] - fn blocked_min_max_bytes_across_blocks() -> Result<()> { - let long = "z".repeat(30); - let indices = |flats: &[usize]| -> Vec { - flats - .iter() - .map(|&f| BlocksIndex::from_flat(f, 4)) - .collect() - }; - let mut acc = BlockedMinMaxBytesAccumulator::new_max(DataType::Utf8, 4); - let values: ArrayRef = Arc::new(StringArray::from(vec![ - Some("a"), - Some("c"), - None, - Some(long.as_str()), - Some("b"), - ])); - acc.update_batch(&[values], &indices(&[0, 0, 1, 5, 5]), None, 6)?; - assert!(acc.size() >= long.len(), "values are accounted"); - // second batch: a smaller value must not replace, a larger one must - let values: ArrayRef = Arc::new(StringArray::from(vec!["b", "d"])); - acc.update_batch(&[values], &indices(&[0, 2]), None, 6)?; - - let out = acc.evaluate(BlockedEmitTo::All)?; - let out: Vec>> = out - .iter() - .map(|a| a.as_string::().iter().collect()) - .collect(); - assert_eq!( - out, - vec![ - vec![Some("c"), None, Some("d"), None], - vec![None, Some(long.as_str())], - ] - ); - assert_eq!(acc.inner.total_data_bytes, 0); - Ok(()) - } use arrow::array::StringArray; #[test] diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs index 9e7a67a9ce1d3..1228b458a82b9 100644 --- a/datafusion/functions-aggregate/src/sum.rs +++ b/datafusion/functions-aggregate/src/sum.rs @@ -40,12 +40,10 @@ use datafusion_expr::expr_fn::cast; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, BlockedGroupsAccumulator, Coercion, Documentation, - Expr, GroupsAccumulator, + Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, GroupsAccumulator, Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility, }; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_prim_op::BlockedPrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator; use datafusion_macros::user_doc; @@ -319,27 +317,6 @@ impl AggregateUDFImpl for Sum { downcast_sum!(args, helper) } - fn blocked_groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - self.groups_accumulator_supported(args) - } - - fn create_blocked_groups_accumulator( - &self, - args: AccumulatorArgs, - block_size: usize, - ) -> Result> { - macro_rules! helper { - ($t:ty, $dt:expr) => { - Ok(Box::new(BlockedPrimitiveGroupsAccumulator::<$t, _>::new( - &$dt, - block_size, - |x, y| *x = x.add_wrapping(y), - ))) - }; - } - downcast_sum!(args, helper) - } - fn create_sliding_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/sqllogictest/test_files/aggregate_blocked.slt b/datafusion/sqllogictest/test_files/aggregate_blocked.slt index 3657eb042fbdb..86777be1c305a 100644 --- a/datafusion/sqllogictest/test_files/aggregate_blocked.slt +++ b/datafusion/sqllogictest/test_files/aggregate_blocked.slt @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. -# Blocked hash aggregation: `GROUP BY ` where every aggregate is -# `count` keeps its state in blocks of 2^18 groups and emits block by block. +# Blocked hash aggregation: `count` and `GROUP BY ` keys keep their +# state in blocks of 2^18 groups and emit block by block; other aggregates and +# key types keep flat storage in the same table. # # 600000 groups span 3 blocks. Keys are scrambled with `(v * 7) % 600000` # (gcd(7, 600000) = 1, so still 600000 distinct keys) because generate_series @@ -29,8 +30,7 @@ SET datafusion.execution.target_partitions = 1 statement ok CREATE TABLE t AS -SELECT (v * 7) % 600000 AS k, CASE WHEN v % 5 = 0 THEN NULL ELSE v END AS v, - CAST(v AS DOUBLE) / 4 AS f +SELECT (v * 7) % 600000 AS k, CASE WHEN v % 5 = 0 THEN NULL ELSE v END AS v FROM generate_series(0, 1799999) AS s(v) # count per group across several blocks @@ -64,57 +64,19 @@ FROM (SELECT k, count(v) FILTER (WHERE v % 2 = 0) AS c FROM t GROUP BY k) ---- 600000 720000 -# count together with a non-blocked aggregate takes today's path +# blocked count together with a flat aggregate (sum) query III SELECT count(*), sum(c), sum(s) FROM (SELECT k, count(v) AS c, sum(v) AS s FROM t GROUP BY k) ---- 600000 1440000 1296000000000 -# blocked sum -query IIII -SELECT count(*), sum(s), min(s), max(s) -FROM (SELECT k, sum(v) AS s FROM t GROUP BY k) ----- -600000 1296000000000 1800003 3599997 - -# blocked avg -query IRRR -SELECT count(*), sum(a), min(a), max(a) -FROM (SELECT k, avg(f) AS a FROM t GROUP BY k) ----- -600000 134999925000 150000 299999.75 - -# count, sum and avg, all blocked -query IIIR -SELECT count(*), sum(c), sum(s), sum(a) -FROM (SELECT k, count(v) AS c, sum(v) AS s, avg(f) AS a FROM t GROUP BY k) ----- -600000 1440000 1296000000000 134999925000 - -# blocked count(DISTINCT): one set of distinct values per block -query IIII -SELECT count(*), sum(d), min(d), max(d) -FROM (SELECT k, count(DISTINCT v % 2) AS d FROM t GROUP BY k) ----- -600000 480000 0 1 - -# blocked min / max on strings, both inline (short) and heap (long) values -query ITTI -SELECT count(*), min(mn), max(mx), sum(length(mx)) -FROM ( - SELECT k, min(CAST(v AS VARCHAR)) AS mn, max(repeat(CAST(v AS VARCHAR), 5)) AS mx - FROM t GROUP BY k -) ----- -600000 1 999999999999999999999999999999 14229605 - # flat (string) keys with a blocked accumulator query II -SELECT count(*), sum(s) -FROM (SELECT CAST(k AS VARCHAR) AS ks, sum(v) AS s FROM t GROUP BY ks) +SELECT count(*), sum(c) +FROM (SELECT CAST(k AS VARCHAR) AS ks, count(v) AS c FROM t GROUP BY ks) ---- -600000 1296000000000 +600000 1440000 # multiple partitions: partial + final aggregation both blocked statement ok From ee198dec114845d0eee7ec741a773d59c903048a Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:16:20 +0300 Subject: [PATCH 03/10] refactor(agg): drop blocked ordered-stream support and preserving reads --- .../src/blocked_groups_accumulator.rs | 127 +------------- datafusion/expr/src/lib.rs | 2 +- datafusion/functions-aggregate/src/count.rs | 75 +------- .../aggregate_hash_table/common_ordered.rs | 166 +++++++++--------- .../ordered_final_table.rs | 13 +- .../ordered_partial_table.rs | 10 +- .../aggregate_hash_table/storage.rs | 20 --- .../src/aggregates/hash_stream.rs | 4 +- .../physical-plan/src/aggregates/order/mod.rs | 28 +-- .../src/aggregates/order/partial.rs | 32 +--- .../src/aggregates/ordered_final_stream.rs | 68 +++---- .../src/aggregates/ordered_partial_stream.rs | 87 ++++----- .../src/aggregates/ordered_single_stream.rs | 4 +- .../src/aggregates/single_stream.rs | 4 +- .../physical-plan/src/aggregates/spill.rs | 69 +++++--- .../test_files/aggregate_blocked.slt | 122 ------------- 16 files changed, 210 insertions(+), 621 deletions(-) delete mode 100644 datafusion/sqllogictest/test_files/aggregate_blocked.slt diff --git a/datafusion/expr-common/src/blocked_groups_accumulator.rs b/datafusion/expr-common/src/blocked_groups_accumulator.rs index fa07ad494de07..624ba73ca8bbb 100644 --- a/datafusion/expr-common/src/blocked_groups_accumulator.rs +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -28,7 +28,7 @@ use std::any::Any; use std::sync::Arc; use arrow::array::{ArrayRef, BooleanArray}; -use datafusion_common::{Result, exec_err, not_impl_err}; +use datafusion_common::Result; use crate::accumulator::AggregateMetrics; use crate::groups_accumulator::EmitTo; @@ -144,97 +144,6 @@ impl BlockedEmitTo { } } -/// Selects groups for a non-destructive read of blocked state, like -/// [`GroupSelection`] for flat state. -/// -/// Selections created by [`Self::try_from_indices`] preserve the requested -/// order and support duplicate indices. -/// -/// [`GroupSelection`]: crate::groups_accumulator::GroupSelection -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct BlockedGroupSelection<'a> { - total_num_groups: usize, - block_size: usize, - indices: Option<&'a [BlocksIndex]>, -} - -impl<'a> BlockedGroupSelection<'a> { - /// Selects all `total_num_groups` groups, stored in blocks of - /// `block_size` groups, in group order. - pub fn all(total_num_groups: usize, block_size: usize) -> Self { - Self { - total_num_groups, - block_size, - indices: None, - } - } - - /// Selects groups in the order specified by `indices`. - /// - /// Returns an error if an index is not one of the first - /// `total_num_groups` groups. Empty selections are valid, and duplicate - /// indices are preserved. - pub fn try_from_indices( - indices: &'a [BlocksIndex], - total_num_groups: usize, - block_size: usize, - ) -> Result { - if let Some(index) = indices - .iter() - .find(|index| index.flat(block_size) >= total_num_groups) - { - return exec_err!( - "Group index {index:?} is out of bounds for {total_num_groups} groups" - ); - } - Ok(Self { - total_num_groups, - block_size, - indices: Some(indices), - }) - } - - /// Returns the group count against which this selection was constructed. - pub fn total_num_groups(self) -> usize { - self.total_num_groups - } - - /// Ensures this selection is being applied to the same number of groups - /// against which it was constructed, see - /// `GroupSelection::validate_num_groups`. - pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> { - if actual_num_groups != self.total_num_groups { - return exec_err!( - "Group selection was constructed for {} groups but applied to {actual_num_groups} groups", - self.total_num_groups - ); - } - Ok(()) - } - - /// Returns the number of selected groups. - pub fn len(self) -> usize { - self.indices - .map_or(self.total_num_groups, |indices| indices.len()) - } - - /// Returns `true` if no groups are selected. - pub fn is_empty(self) -> bool { - self.len() == 0 - } - - /// Returns the selected groups in output order. - pub fn iter(self) -> impl Iterator + 'a { - let block_size = self.block_size; - let (all, indices): (_, &'a [BlocksIndex]) = match self.indices { - None => (0..self.total_num_groups, &[]), - Some(indices) => (0..0, indices), - }; - all.map(move |flat| BlocksIndex::from_flat(flat, block_size)) - .chain(indices.iter().copied()) - } -} - /// 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. @@ -281,46 +190,12 @@ pub trait BlockedGroupsAccumulator: Send + Any { /// * [`BlockedEmitTo::First`]: exactly one array fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result>; - /// Returns final values without changing the state or the group indices, - /// see `GroupsAccumulator::evaluate_preserving`. - /// - /// Rows are returned in the order of `selection`, in arrays of at most - /// [`Self::block_size`] rows. An empty selection returns no arrays. - fn evaluate_preserving( - &mut self, - _selection: BlockedGroupSelection<'_>, - ) -> Result> { - not_impl_err!("Preserving grouped evaluation is not implemented") - } - - /// Returns `true` if [`Self::evaluate_preserving`] is implemented. - fn supports_evaluate_preserving(&self) -> bool { - false - } - /// 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>>; - /// Returns intermediate state without changing the state or the group - /// indices, see `GroupsAccumulator::state_preserving`. - /// - /// Indexed `[chunk][state column]`, where chunks follow the same rules as - /// [`Self::evaluate_preserving`]. - fn state_preserving( - &mut self, - _selection: BlockedGroupSelection<'_>, - ) -> Result>> { - not_impl_err!("Preserving grouped state is not implemented") - } - - /// Returns `true` if [`Self::state_preserving`] is implemented. - fn supports_state_preserving(&self) -> bool { - false - } - /// Converts an input batch directly to intermediate state. /// /// Same contract as `GroupsAccumulator::convert_to_state`. diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 9c3cf0947a1c0..204088aab47bb 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -105,7 +105,7 @@ pub use datafusion_expr_common::accumulator::{ Accumulator, AggregateMetric, AggregateMetricRecorder, AggregateMetrics, }; pub use datafusion_expr_common::blocked_groups_accumulator::{ - BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex, + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, }; pub use datafusion_expr_common::columnar_value::ColumnarValue; pub use datafusion_expr_common::groups_accumulator::{ diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index b8a7ce85466c2..e18393fa1a292 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -35,10 +35,10 @@ use datafusion_common::{ stats::Precision, utils::expr::COUNT_STAR_EXPANSION, }; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, BlockedEmitTo, BlockedGroupSelection, - BlockedGroupsAccumulator, BlocksIndex, 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}, @@ -934,35 +934,6 @@ impl BlockedGroupsAccumulator for BlockedCountGroupsAccumulator { Ok(self.take(emit_to)) } - fn evaluate_preserving( - &mut self, - selection: BlockedGroupSelection<'_>, - ) -> Result> { - selection.validate_num_groups(self.counts.len())?; - // Build each output chunk in its own `Vec`, which becomes the array's - // buffer without copying - let block_size = self.counts.block_size(); - let mut remaining = selection.len(); - let mut arrays = Vec::with_capacity(remaining.div_ceil(block_size)); - let mut chunk = Vec::with_capacity(remaining.min(block_size)); - for index in selection.iter() { - chunk.push(self.counts.get(index)); - if chunk.len() == block_size { - remaining -= block_size; - let next = Vec::with_capacity(remaining.min(block_size)); - arrays.push(Self::counts_to_array(std::mem::replace(&mut chunk, next))); - } - } - if !chunk.is_empty() { - arrays.push(Self::counts_to_array(chunk)); - } - Ok(arrays) - } - - fn supports_evaluate_preserving(&self) -> bool { - true - } - fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { Ok(self .take(emit_to) @@ -971,21 +942,6 @@ impl BlockedGroupsAccumulator for BlockedCountGroupsAccumulator { .collect()) } - fn state_preserving( - &mut self, - selection: BlockedGroupSelection<'_>, - ) -> Result>> { - Ok(self - .evaluate_preserving(selection)? - .into_iter() - .map(|counts| vec![counts]) - .collect()) - } - - fn supports_state_preserving(&self) -> bool { - true - } - fn convert_to_state( &self, values: &[ArrayRef], @@ -1313,29 +1269,6 @@ mod tests { Ok(()) } - #[test] - fn blocked_count_preserving_reads_keep_state() -> Result<()> { - let mut acc = BlockedCountGroupsAccumulator::new(4); - let values: ArrayRef = Arc::new(Int32Array::from(vec![1; 7])); - acc.update_batch(&[values], &blocks_index(&[0, 1, 1, 5, 5, 5, 6], 4), None, 7)?; - - let all = acc.evaluate_preserving(BlockedGroupSelection::all(7, 4))?; - assert_eq!(blocked_values(&all), vec![vec![1, 2, 0, 0], vec![0, 3, 1]]); - - let indices = blocks_index(&[5, 0, 5], 4); - let selected = BlockedGroupSelection::try_from_indices(&indices, 7, 4)?; - let state = acc.state_preserving(selected)?; - assert_eq!(state.len(), 1); - assert_eq!(blocked_values(&state[0]), vec![vec![3, 1, 3]]); - - // state is unchanged - assert_eq!( - blocked_values(&acc.evaluate(BlockedEmitTo::All)?), - vec![vec![1, 2, 0, 0], vec![0, 3, 1]] - ); - Ok(()) - } - #[test] fn blocked_count_state_is_one_column_per_block() -> Result<()> { let mut acc = BlockedCountGroupsAccumulator::new(4); 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 d95c6405f8fea..fa0c1ff90a8f6 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 @@ -23,15 +23,16 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::Result; use datafusion_common::assert_or_internal_err; -use datafusion_common::{Result, internal_err}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{AggregateMetrics, EmitTo}; use crate::InputOrderMode; use crate::PhysicalExpr; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, + GroupByMetrics, GroupValues, new_group_values, }; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -42,11 +43,9 @@ use crate::aggregates::{ use super::AggregateTableMetrics; use super::common::{ AggregateAccumulator, AggregateBatchFn, AggregateHashTable, EvaluatedAggregateBatch, - HashAggregateAccumulator, MaterializeAccumulatorFn, create_accumulator_storage, -}; -use super::storage::{ - GroupKeys, MaterializedBatch, block_size, materialize_batches, use_blocked_keys, + MaterializeAccumulatorFn, create_group_accumulator, }; +use super::storage::{AccumulatorStorage, Emit, GroupIndices}; #[derive(Clone)] pub(in crate::aggregates) struct OrderedAggregateTableMetrics { @@ -166,9 +165,11 @@ pub(super) struct OrderedAggregateTableBuffer { /// Tracks how far ordered input allows this table to drain safely. pub(super) group_ordering: GroupOrdering, - /// Interned group keys, in the same group-id order used by accumulators, - /// and the scratch group id vector for the current input batch. - pub(super) keys: GroupKeys, + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec, /// One item per aggregate expression. /// @@ -201,43 +202,30 @@ impl OrderedAggregateTable { let group_ordering = GroupOrdering::try_new(input_order_mode)?; let group_schema = agg.group_by().group_schema(input_schema)?; - let block_size = block_size(batch_size); + let group_values = new_group_values(group_schema, &group_ordering)?; let aggregate_arguments = aggregate_expressions( agg.aggr_expr(), aggregate_mode, agg.group_by().num_group_exprs(), )?; - let accumulators: Vec = agg + let accumulators = agg .aggr_expr() .iter() .zip(aggregate_arguments) .zip(filters) .zip(metrics.submetrics.iter()) .map(|(((agg_expr, arguments), filter), submetrics)| { - let accumulator = create_accumulator_storage( - agg_expr, - Arc::clone(submetrics), - block_size, - )?; + let accumulator = + create_group_accumulator(agg_expr, Arc::clone(submetrics))?; Ok(AggregateAccumulator::new( Arc::clone(agg_expr), arguments, filter, - accumulator, + AccumulatorStorage::Flat(accumulator), Arc::clone(submetrics), )) }) .collect::>()?; - let keys = GroupKeys::try_new( - Arc::clone(&group_schema), - &group_ordering, - use_blocked_keys(agg, &group_schema, block_size), - block_size, - &accumulators - .iter() - .map(|acc| acc.storage()) - .collect::>(), - )?; Ok(Self { output_schema, @@ -250,7 +238,8 @@ impl OrderedAggregateTable { buffer: OrderedAggregateTableBuffer { group_by: Arc::clone(agg.group_by()), group_ordering, - keys, + group_values, + group_indices: vec![], accumulators, }, _mode: PhantomData, @@ -304,7 +293,7 @@ impl OrderedAggregateTable { /// Number of groups currently buffered. pub(in crate::aggregates) fn num_groups(&self) -> usize { - self.buffer.keys.len() + self.buffer.group_values.len() } /// Check if there is zero groups accumulated so far. @@ -319,8 +308,9 @@ impl OrderedAggregateTable { .iter() .map(|acc| acc.size()) .sum::() - + self.buffer.keys.size() + + self.buffer.group_values.size() + self.buffer.group_ordering.size() + + self.buffer.group_indices.allocated_size() } pub(in crate::aggregates) fn metrics(&self) -> OrderedAggregateTableMetrics { @@ -339,37 +329,41 @@ impl OrderedAggregateTable { /// active (incomplete) groups. Partial aggregation can pass those states to /// its final stage, while single and final aggregation sort and spill them /// before replay. - /// - /// 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( + pub(in crate::aggregates) fn take_state_batch( &mut self, - ) -> Result> { - if self.buffer.keys.is_empty() { - return Ok(vec![]); + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); } let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - let batches = self.group_by_metrics.time_emitting(|| { - materialize_batches( - &mut self.buffer.keys, - &mut self.buffer.accumulators, - EmitTo::All, - HashAggregateAccumulator::state, - AccumulatorPhase::State, - &accumulator_metrics, - &self.state_schema, - ) + 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(Emit::Flat(EmitTo::All)) + })? + .into_iter() + .flatten(), + ); + } + Ok::<_, datafusion_common::DataFusionError>(output) })?; - debug_assert!(batches.iter().all(|b| b.batch.num_rows() > 0)); - // Emitting all groups resets accumulator state. Explicitly shrink the + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released // before the batch is passed downstream or sorted for spilling. - self.buffer.keys.clear_shrink(0); + self.buffer.group_values.clear_shrink(0); + self.buffer.group_indices.clear(); + self.buffer.group_indices.shrink_to_fit(); self.buffer.group_ordering.reset(); - Ok(batches) + Ok(Some(batch)) } /// Aggregates one evaluated input batch after selecting the mode-specific @@ -389,13 +383,16 @@ impl OrderedAggregateTable { let group_by_metrics = self.group_by_metrics.clone(); for group_values in &evaluated_batch.grouping_set_args { let total_num_groups = group_by_metrics.time_group_key_preparation(|| { - let starting_num_groups = self.buffer.keys.len(); - self.buffer.keys.intern(group_values)?; - let total_num_groups = self.buffer.keys.len(); + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); if total_num_groups > starting_num_groups { - self.buffer.keys.record_new_groups( - &mut self.buffer.group_ordering, + self.buffer.group_ordering.new_groups( group_values, + &self.buffer.group_indices, + total_num_groups, )?; } Ok::<_, datafusion_common::DataFusionError>(total_num_groups) @@ -409,9 +406,13 @@ impl OrderedAggregateTable { .zip(evaluated_batch.accumulator_args.iter()) .enumerate() { - let group_indices = self.buffer.keys.indices_for(acc.storage()); accumulator_metrics.time(idx, accumulator_phase, || { - aggregate_fn(acc, values, group_indices, total_num_groups) + aggregate_fn( + acc, + values, + GroupIndices::Flat(&self.buffer.group_indices), + total_num_groups, + ) })?; } Ok::<(), datafusion_common::DataFusionError>(()) @@ -433,7 +434,7 @@ impl OrderedAggregateTable { materialize_accumulator_fn: MaterializeAccumulatorFn, accumulator_phase: AccumulatorPhase, ) -> Result> { - if self.buffer.keys.is_empty() { + if self.buffer.group_values.is_empty() { return Ok(None); } @@ -447,32 +448,21 @@ impl OrderedAggregateTable { } EmitTo::All => EmitTo::All, }; - // `emit_to` is capped at `batch_size`, at most one block, so this is - // one batch. - let mut batches = self.materialize_groups( - emit_to, - materialize_accumulator_fn, - accumulator_phase, - )?; - let Some(batch) = batches.pop().filter(|_| batches.is_empty()) else { - return internal_err!("{emit_to:?} must materialize exactly one batch"); - }; - Ok(Some(batch.batch)) + self.materialize_groups(emit_to, materialize_accumulator_fn, accumulator_phase) + .map(Some) } /// Removes the selected groups once and materializes their output columns. /// The caller chooses the completed prefix and any output-size limit. - /// - /// Returns one batch if every part is flat, otherwise one batch per - /// emitted block, see [`materialize_batches`]. pub(super) fn materialize_groups( &mut self, emit_to: EmitTo, materialize_accumulator_fn: MaterializeAccumulatorFn, accumulator_phase: AccumulatorPhase, - ) -> Result> { + ) -> Result { let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - let batches = self.group_by_metrics.time_emitting(|| { + let output = self.group_by_metrics.time_emitting(|| { + let mut output = self.buffer.group_values.emit(emit_to)?; // EOF can also emit a prefix when a caller limits its batch size, // but the completed ordering state no longer tracks group indexes. if let EmitTo::First(n) = emit_to @@ -481,18 +471,22 @@ impl OrderedAggregateTable { self.buffer.group_ordering.remove_groups(n); } - materialize_batches( - &mut self.buffer.keys, - &mut self.buffer.accumulators, - emit_to, - materialize_accumulator_fn, - accumulator_phase, - &accumulator_metrics, - &self.output_schema, - ) + for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { + output.extend( + accumulator_metrics + .time(idx, accumulator_phase, || { + materialize_accumulator_fn(acc, Emit::Flat(emit_to)) + })? + .into_iter() + .flatten(), + ); + } + Ok::<_, datafusion_common::DataFusionError>(output) })?; - debug_assert!(batches.iter().all(|b| b.batch.num_rows() > 0)); - Ok(batches) + let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + Ok(batch) } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index 621319750164c..f70fc0bf0d961 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -31,7 +31,6 @@ use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorP use super::common::HashAggregateAccumulator; use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; -use super::storage::MaterializedBatch; /// Implementation specific to final aggregation, where the table stores partial /// aggregate states and the input rows are also partial states. @@ -84,21 +83,21 @@ impl OrderedAggregateTable { /// Materializes final results for all groups proven complete by the input /// ordering, leaving the active ordered-key range in the table. /// - /// Returns no batch if there are no completed groups, and one batch per - /// emitted chunk for blocked storage. - pub(in crate::aggregates) fn take_completed_result_batches( + /// Returns None if there are no completed groups. + pub(in crate::aggregates) fn take_completed_result_batch( &mut self, - ) -> Result> { + ) -> Result> { if self.is_empty() { - return Ok(vec![]); + return Ok(None); } let Some(emit_to) = self.group_ordering().emit_to() else { - return Ok(vec![]); + return Ok(None); }; self.materialize_groups( emit_to, HashAggregateAccumulator::evaluate_to_columns, AccumulatorPhase::Evaluate, ) + .map(Some) } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index 011ddf518691e..104855c5f0d8f 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -42,7 +42,6 @@ use crate::aggregates::{ use super::common::HashAggregateAccumulator; use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; -use super::storage::MaterializedBatch; /// Implementation specific to partial aggregation, where the table stores /// partial aggregate states and the input rows are raw rows. @@ -92,19 +91,20 @@ impl OrderedAggregateTable { /// Materializes all groups proven complete by the input ordering, leaving /// the active ordered-key range in the table. - pub(in crate::aggregates) fn take_completed_state_batches( + pub(in crate::aggregates) fn take_completed_state_batch( &mut self, - ) -> Result> { + ) -> Result> { if self.is_empty() { - return Ok(vec![]); + return Ok(None); } let Some(emit_to) = self.group_ordering().emit_to() else { - return Ok(vec![]); + return Ok(None); }; self.materialize_groups( emit_to, HashAggregateAccumulator::state, AccumulatorPhase::State, ) + .map(Some) } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs index c6a58c8b57a95..a22b1ce679c96 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs @@ -188,26 +188,6 @@ impl GroupKeys { } } - /// Tells `group_ordering` about the groups added by the last interned - /// batch `cols`, see [`GroupOrdering::new_groups`]. - pub(super) fn record_new_groups( - &self, - group_ordering: &mut GroupOrdering, - cols: &[ArrayRef], - ) -> Result<()> { - match &self.storage { - KeyStorage::Flat(_) => { - group_ordering.new_groups(cols, &self.flat_indices, self.len()) - } - KeyStorage::Blocked(_) => group_ordering.new_blocked_groups( - cols, - &self.blocked_indices, - self.block_size, - self.len(), - ), - } - } - pub(super) fn len(&self) -> usize { match &self.storage { KeyStorage::Flat(group_values) => group_values.len(), diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index fa820725bd2de..d4202af0989a4 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -767,7 +767,7 @@ impl FinalHashAggregateStream { // groups so far. let result = hash_table .take_state_batches() - .and_then(|batch| spill_context.sort_and_spill(batch)); + .and_then(|batch| spill_context.sort_and_spill_batches(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -806,7 +806,7 @@ impl FinalHashAggregateStream { // Input was exhausted after spilling. Spill the last in-memory run hash_table .take_state_batches() - .and_then(|batch| spill_context.sort_and_spill(batch))?; + .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/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 8968f29d6fd5f..147e4e0d6f18d 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -19,7 +19,7 @@ use std::mem::size_of; use arrow::array::ArrayRef; use datafusion_common::Result; -use datafusion_expr::{BlocksIndex, EmitTo}; +use datafusion_expr::EmitTo; mod full; mod partial; @@ -147,32 +147,6 @@ impl GroupOrdering { Ok(()) } - /// [`Self::new_groups`] for blocked group indices, with `block_size` - /// groups per block. - pub(crate) fn new_blocked_groups( - &mut self, - batch_group_values: &[ArrayRef], - group_indices: &[BlocksIndex], - block_size: usize, - total_num_groups: usize, - ) -> Result<()> { - match self { - GroupOrdering::None => {} - GroupOrdering::Partial(partial) => { - partial.new_blocked_groups( - batch_group_values, - group_indices, - block_size, - total_num_groups, - )?; - } - GroupOrdering::Full(full) => { - full.new_groups(total_num_groups); - } - } - Ok(()) - } - /// Returns the size of memory used by the ordering state, in bytes. pub fn size(&self) -> usize { size_of::() diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 0038501aa5479..1603bb6d079be 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -25,7 +25,7 @@ use arrow_ord::partition::partition; use datafusion_common::utils::{compare_rows, get_row_at_idx}; use datafusion_common::{Result, ScalarValue}; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::{BlocksIndex, EmitTo}; +use datafusion_expr::EmitTo; /// Tracks grouping state when the data is ordered by some subset of /// the group keys. @@ -216,34 +216,6 @@ impl GroupOrderingPartial { batch_group_values: &[ArrayRef], group_indices: &[usize], total_num_groups: usize, - ) -> Result<()> { - self.new_groups_with(batch_group_values, total_num_groups, |row| { - group_indices[row] - }) - } - - /// [`Self::new_groups`] for blocked group indices, with `block_size` - /// groups per block. - pub(crate) fn new_blocked_groups( - &mut self, - batch_group_values: &[ArrayRef], - group_indices: &[BlocksIndex], - block_size: usize, - total_num_groups: usize, - ) -> Result<()> { - self.new_groups_with(batch_group_values, total_num_groups, |row| { - group_indices[row].flat(block_size) - }) - } - - /// Shared implementation of [`Self::new_groups`] and - /// [`Self::new_blocked_groups`]. `group_position(row)` is the position, - /// counting from the first group, of the group of `row`. - fn new_groups_with( - &mut self, - batch_group_values: &[ArrayRef], - total_num_groups: usize, - group_position: impl Fn(usize) -> usize, ) -> Result<()> { assert!(total_num_groups > 0); assert!(!batch_group_values.is_empty()); @@ -270,7 +242,7 @@ impl GroupOrderingPartial { let ranges = partition(&sort_keys)?.ranges(); let last_range = ranges.last().unwrap(); - let range_current_sort = group_position(last_range.start); + let range_current_sort = group_indices[last_range.start]; let range_sort_key = get_row_at_idx(&sort_keys, last_range.start)?; let (current_sort, sort_key) = if last_range.start == 0 { diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 1fcafc9102a27..4a70ef2df6909 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -28,7 +28,7 @@ use futures::stream::StreamExt; use super::AggregateExec; use super::aggregate_hash_table::{ - FinalMarker, MaterializedBatch, OrderedAggregateTable, OrderedAggregateTableMetrics, + FinalMarker, OrderedAggregateTable, OrderedAggregateTableMetrics, }; use super::spill::AggregateSpill; use crate::aggregates::AggregateMode; @@ -76,9 +76,8 @@ struct Aggregating { } struct Outputting { - /// Materialized final results: one batch for flat storage, one per chunk - /// for blocked storage. Each batch is emitted in slices of `batch_size` rows. - batches: Vec, + /// Materialized final results, emitted in slices of `batch_size` rows. + batch: RecordBatch, /// Aggregation stage to resume after output; `None` after EOF. resume: Option, } @@ -339,7 +338,7 @@ impl Aggregating { if self.table.is_empty() { return Err(oom); } - spill_context.sort_and_spill(self.table.take_state_batches()?)?; + spill_context.sort_and_spill(self.table.take_state_batch()?)?; reservation .try_resize(self.table.memory_size()) .map_err(|e| { @@ -357,13 +356,12 @@ impl Aggregating { // through replay before any more final results can be emitted. continue; } - let batches = self.table.take_completed_result_batches()?; - if batches.is_empty() { + let Some(batch) = self.table.take_completed_result_batch()? else { continue; - } + }; timer.done(); return Ok(Some(ExecutionStage::Outputting(Outputting { - batches, + batch, resume: Some(self), }))); } @@ -372,7 +370,7 @@ impl Aggregating { drop(self.input); let timer = elapsed_compute.timer(); if let Some(mut spill_context) = self.spill_context.filter(|s| s.has_spills()) { - spill_context.sort_and_spill(self.table.take_state_batches()?)?; + spill_context.sort_and_spill(self.table.take_state_batch()?)?; let metrics = self.table.metrics(); drop(self.table); reservation.try_resize(0)?; @@ -387,30 +385,29 @@ impl Aggregating { } self.table.input_done(); - let batches = self.table.take_completed_result_batches()?; + let output = self.table.take_completed_result_batch()?; drop(self.table); timer.done(); - if batches.is_empty() { + let Some(batch) = output else { reservation.try_resize(0)?; return Ok(None); - } + }; Ok(Some(ExecutionStage::Outputting(Outputting { - batches, + batch, resume: None, }))) } } impl Outputting { - /// Emits slices of the materialized batches without touching the aggregate - /// table. + /// Emits slices of one materialized batch without touching the aggregate table. async fn handle_stage( self, context: &OrderedFinalAggregateContext, reservation: &MemoryReservation, emitter: &mut TryEmitter, ) -> Result> { - let Self { batches, resume } = self; + let Self { mut batch, resume } = self; let elapsed_compute = context.baseline_metrics.elapsed_compute(); let mut timer = elapsed_compute.timer(); let (table_memory, next_stage) = match resume { @@ -420,43 +417,32 @@ impl Outputting { ), None => (0, None), }; - let batches_memory: usize = batches.iter().map(|b| b.memory_size).sum(); - match reservation.try_resize(table_memory + batches_memory) { + let batch_memory = batch.get_array_memory_size(); + match reservation.try_resize(table_memory + batch_memory) { Ok(()) => {} Err(DataFusionError::ResourcesExhausted(_)) => { - // If we cannot hold the batches while slicing, hand them off whole. + // If we cannot hold the batch while slicing, hand it off whole. reservation.try_resize(table_memory)?; timer.done(); - for MaterializedBatch { batch, .. } in batches { - emitter.emit(batch).await; - } + emitter.emit(batch).await; return Ok(next_stage); } Err(e) => return Err(e), } - for MaterializedBatch { - mut batch, - memory_size: batch_memory, - } in batches - { - while batch.num_rows() > context.batch_size { - let output = batch.slice(0, context.batch_size); - batch = batch - .slice(context.batch_size, batch.num_rows() - context.batch_size); - timer.done(); - emitter.emit(output).await; - timer = elapsed_compute.timer(); - } - - // The final slice transfers ownership of the buffers to the - // consumer, so this batch's memory is released batch by batch. - reservation.try_shrink(batch_memory)?; + while batch.num_rows() > context.batch_size { + let output = batch.slice(0, context.batch_size); + batch = + batch.slice(context.batch_size, batch.num_rows() - context.batch_size); timer.done(); - emitter.emit(batch).await; + emitter.emit(output).await; timer = elapsed_compute.timer(); } + + // The final slice transfers ownership of the buffers to the consumer. + reservation.try_shrink(batch_memory)?; timer.done(); + emitter.emit(batch).await; Ok(next_stage) } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 3f33364805d5a..dd845353c77eb 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -27,9 +27,7 @@ use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::StreamExt; use super::AggregateExec; -use super::aggregate_hash_table::{ - MaterializedBatch, OrderedAggregateTable, PartialMarker, -}; +use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; use crate::aggregates::order::GroupOrdering; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; @@ -124,10 +122,9 @@ struct Aggregating { } struct Outputting { - /// Materialized aggregate states: one batch for flat storage, one per - /// block or chunk for blocked storage. Each batch is emitted in slices of - /// `batch_size` rows. - batches: Vec, + /// Materialized aggregate states. Each iteration emits the first `batch_size` + /// rows and replaces this batch with the remaining slice. + batch: RecordBatch, /// Aggregation stage to resume after output; `None` after EOF. resume: Option, } @@ -313,32 +310,31 @@ impl Aggregating { let timer = elapsed_compute.timer(); self.table.aggregate_batch(&batch)?; - let batches = match reservation.try_resize(self.table.memory_size()) { - Ok(()) => self.table.take_completed_state_batches()?, + let output = match reservation.try_resize(self.table.memory_size()) { + Ok(()) => self.table.take_completed_state_batch()?, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { // Partial ordering may have an unbounded active key range. // The final stage can merge incomplete states emitted here. if matches!(self.table.group_ordering(), GroupOrdering::Full(_)) { return Err(oom); } - let batches = self.table.take_state_batches()?; - if batches.is_empty() { + let Some(batch) = self.table.take_state_batch()? else { return Err(oom); - } - batches + }; + Some(batch) } Err(e) => return Err(e), }; - if batches.is_empty() { + let Some(batch) = output else { continue; - } + }; timer.done(); // OOM, do early emit next, and go back to the current state to continue // aggregating return Ok(Some(ExecutionStage::Outputting(Outputting { - batches, + batch, resume: Some(self), }))); } @@ -347,23 +343,23 @@ impl Aggregating { drop(self.input); self.table.input_done(); let timer = elapsed_compute.timer(); - let batches = self.table.take_completed_state_batches()?; + let output = self.table.take_completed_state_batch()?; drop(self.table); timer.done(); - if batches.is_empty() { + let Some(batch) = output else { reservation.try_resize(0)?; return Ok(None); - } + }; Ok(Some(ExecutionStage::Outputting(Outputting { - batches, + batch, resume: None, }))) } } impl Outputting { - /// Emits slices of the materialized batches without touching the hash table. + /// Emits slices of one materialized batch without touching the hash table. /// /// See [`OrderedPartialAggregateStream::into_stream`] for stage transitions /// and output memory accounting. @@ -373,7 +369,7 @@ impl Outputting { reservation: &MemoryReservation, emitter: &mut TryEmitter, ) -> Result> { - let Self { batches, resume } = self; + let Self { mut batch, resume } = self; let elapsed_compute = context.baseline_metrics.elapsed_compute(); let mut timer = elapsed_compute.timer(); let (table_memory, next_stage) = match resume { @@ -383,49 +379,38 @@ impl Outputting { ), None => (0, None), }; - let batches_memory: usize = batches.iter().map(|b| b.memory_size).sum(); - match reservation.try_resize(table_memory + batches_memory) { + let batch_memory = batch.get_array_memory_size(); + match reservation.try_resize(table_memory + batch_memory) { Ok(()) => {} Err(DataFusionError::ResourcesExhausted(_)) => { - // If we cannot hold the batches while slicing, hand them off whole. + // If we cannot hold the batch while slicing, hand it off whole. // Only the retained table needs to remain reserved. reservation.try_resize(table_memory)?; + context.reduction_factor.add_part(batch.num_rows()); timer.done(); - for MaterializedBatch { batch, .. } in batches { - context.reduction_factor.add_part(batch.num_rows()); - emitter.emit(batch).await; - } + emitter.emit(batch).await; return Ok(next_stage); } Err(e) => return Err(e), } - for MaterializedBatch { - mut batch, - memory_size: batch_memory, - } in batches - { - while batch.num_rows() > context.batch_size { - // 1. Emit first `batch_size` rows from `batch` - // 2. Update `batch`` with the remaining tail - let output = batch.slice(0, context.batch_size); - batch = batch - .slice(context.batch_size, batch.num_rows() - context.batch_size); - context.reduction_factor.add_part(output.num_rows()); - timer.done(); - emitter.emit(output).await; - timer = elapsed_compute.timer(); - } - - // The final slice transfers ownership of the buffers to the - // consumer, so this batch's memory is released batch by batch. - reservation.try_shrink(batch_memory)?; - context.reduction_factor.add_part(batch.num_rows()); + while batch.num_rows() > context.batch_size { + // 1. Emit first `batch_size` rows from `batch` + // 2. Update `batch`` with the remaining tail + let output = batch.slice(0, context.batch_size); + batch = + batch.slice(context.batch_size, batch.num_rows() - context.batch_size); + context.reduction_factor.add_part(output.num_rows()); timer.done(); - emitter.emit(batch).await; + emitter.emit(output).await; timer = elapsed_compute.timer(); } + + // The final slice transfers ownership of the buffers to the consumer. + reservation.try_shrink(batch_memory)?; + context.reduction_factor.add_part(batch.num_rows()); timer.done(); + emitter.emit(batch).await; Ok(next_stage) } } diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 2c253e10f416e..97afab093a872 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -409,7 +409,7 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let mut result = table - .take_state_batches() + .take_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -456,7 +456,7 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let replay = match table - .take_state_batches() + .take_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index fee7e349dedf5..65e1c0d345241 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -441,7 +441,7 @@ impl SingleHashAggregateStream { let timer = elapsed_compute.timer(); let mut result = hash_table .take_state_batches() - .and_then(|batch| spill_context.sort_and_spill(batch)); + .and_then(|batch| spill_context.sort_and_spill_batches(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -489,7 +489,7 @@ impl SingleHashAggregateStream { let timer = elapsed_compute.timer(); let replay = match hash_table .take_state_batches() - .and_then(|batch| spill_context.sort_and_spill(batch)) + .and_then(|batch| spill_context.sort_and_spill_batches(batch)) { Ok(()) => { let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 7ad58acbf55ce..af9dab2d9c655 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -20,6 +20,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryReservation; @@ -40,7 +41,7 @@ use crate::{InputOrderMode, SendableRecordBatchStream}; /// /// Every aggregation stream that spills does so the same way. Each spill event /// drains all currently buffered groups as intermediate state (see -/// `take_state_batches` on the aggregate tables), sorts them by the full group +/// `take_state_batch` on the aggregate tables), sorts them by the full group /// key, and writes them to one spill file. After the original input ends, all /// files are merged and replayed through an [`OrderedFinalAggregateStream`], /// which merges the states and evaluates the final aggregate values. @@ -207,13 +208,45 @@ impl AggregateSpill { !self.spills.is_empty() } - /// Sorts `state_batches`, the intermediate state of all currently buffered - /// groups (empty if there are no groups), and writes each batch as one - /// spill file. Memory reservation should be updated by the caller. - /// - /// Flat aggregate storage produces one batch; blocked storage produces one - /// batch per block. + /// Sorts `state_batch`, the intermediate state of all currently buffered + /// groups (`None` if there are no groups), and writes it as one spill file. + /// Memory reservation should be updated by the caller. pub(super) fn sort_and_spill( + &mut self, + state_batch: Option, + ) -> Result<()> { + let Some(state_batch) = state_batch else { + return Ok(()); + }; + + let sorted_iter = IncrementalSortIterator::new( + state_batch, + self.spill_expr.clone(), + self.batch_size, + ); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + self.label, + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("{}: produced an empty spill", self.label); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// [`Self::sort_and_spill`] for the state batches of a hash aggregate + /// table: one batch for flat storage, one per block for blocked storage, + /// each written as its own spill file. + pub(super) fn sort_and_spill_batches( &mut self, state_batches: Vec, ) -> Result<()> { @@ -221,28 +254,8 @@ impl AggregateSpill { // sort.rs/rank.rs) so all blocks are written as one sorted run instead // of one run per block. for state_batch in state_batches { - let sorted_iter = IncrementalSortIterator::new( - state_batch.batch, - self.spill_expr.clone(), - self.batch_size, - ); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - self.label, - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("{}: produced an empty spill", self.label); - }; - - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); + self.sort_and_spill(Some(state_batch.batch))?; } - Ok(()) } diff --git a/datafusion/sqllogictest/test_files/aggregate_blocked.slt b/datafusion/sqllogictest/test_files/aggregate_blocked.slt deleted file mode 100644 index 86777be1c305a..0000000000000 --- a/datafusion/sqllogictest/test_files/aggregate_blocked.slt +++ /dev/null @@ -1,122 +0,0 @@ -# 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. - -# Blocked hash aggregation: `count` and `GROUP BY ` keys keep their -# state in blocks of 2^18 groups and emit block by block; other aggregates and -# key types keep flat storage in the same table. -# -# 600000 groups span 3 blocks. Keys are scrambled with `(v * 7) % 600000` -# (gcd(7, 600000) = 1, so still 600000 distinct keys) because generate_series -# is sorted, which would take the ordered path instead of the hash path. -# -# Each query summarizes the grouped result into one row. - -statement ok -SET datafusion.execution.target_partitions = 1 - -statement ok -CREATE TABLE t AS -SELECT (v * 7) % 600000 AS k, CASE WHEN v % 5 = 0 THEN NULL ELSE v END AS v -FROM generate_series(0, 1799999) AS s(v) - -# count per group across several blocks -query IIII -SELECT count(*), sum(c), min(c), max(c) -FROM (SELECT k, count(v) AS c FROM t GROUP BY k) ----- -600000 1440000 0 3 - -# count(*) per group -query II -SELECT count(*), sum(c) -FROM (SELECT k, count(*) AS c FROM t GROUP BY k) ----- -600000 1800000 - -# NULL group key, interned after the first block is full -query III -SELECT count(*), sum(c), sum(CASE WHEN k2 IS NULL THEN c END) -FROM ( - SELECT CASE WHEN k >= 599990 THEN NULL ELSE k END AS k2, count(v) AS c - FROM t GROUP BY k2 -) ----- -599991 1440000 24 - -# FILTER on count across blocks -query II -SELECT count(*), sum(c) -FROM (SELECT k, count(v) FILTER (WHERE v % 2 = 0) AS c FROM t GROUP BY k) ----- -600000 720000 - -# blocked count together with a flat aggregate (sum) -query III -SELECT count(*), sum(c), sum(s) -FROM (SELECT k, count(v) AS c, sum(v) AS s FROM t GROUP BY k) ----- -600000 1440000 1296000000000 - -# flat (string) keys with a blocked accumulator -query II -SELECT count(*), sum(c) -FROM (SELECT CAST(k AS VARCHAR) AS ks, count(v) AS c FROM t GROUP BY ks) ----- -600000 1440000 - -# multiple partitions: partial + final aggregation both blocked -statement ok -SET datafusion.execution.target_partitions = 4 - -query II -SELECT count(*), sum(c) -FROM (SELECT k, count(v) AS c FROM t GROUP BY k) ----- -600000 1440000 - -statement ok -SET datafusion.execution.target_partitions = 1 - -# ordered input (generate_series is sorted on v) with more groups than a block -query II -SELECT count(*), sum(c) -FROM (SELECT v, count(*) AS c FROM generate_series(0, 599999) AS s(v) GROUP BY v) ----- -600000 600000 - -# spill with more groups than a block -statement ok -SET datafusion.runtime.memory_limit = '8M' - -query II -SELECT count(*), sum(c) -FROM (SELECT k, count(v) AS c FROM t GROUP BY k) ----- -600000 1440000 - -statement ok -RESET datafusion.runtime.memory_limit - -statement ok -DROP TABLE t - -# Restore settings to slt runner defaults -statement ok -SET datafusion.execution.target_partitions = 4 - -statement ok -RESET datafusion.catalog.create_default_catalog_and_schema From c160407f7972e6087e8013eeba9ee760c398dcaf Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:42:30 +0300 Subject: [PATCH 04/10] docs(agg): describe per-aggregate blocked storage --- datafusion/expr-common/src/blocked_groups_accumulator.rs | 6 +++--- datafusion/expr/src/udaf.rs | 7 +++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/datafusion/expr-common/src/blocked_groups_accumulator.rs b/datafusion/expr-common/src/blocked_groups_accumulator.rs index 624ba73ca8bbb..0bc89b5e0746c 100644 --- a/datafusion/expr-common/src/blocked_groups_accumulator.rs +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -148,9 +148,9 @@ impl BlockedEmitTo { /// [`Self::block_size`] groups, and emitting returns one array per block /// instead of one large array. /// -/// Hash aggregation uses a `BlockedGroupsAccumulator` only when every -/// aggregate in the query (and the group keys) support blocked storage; -/// otherwise the whole aggregation uses [`GroupsAccumulator`]. +/// 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 { diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index d53ea58d02552..38ecf32fea47a 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -685,10 +685,9 @@ 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 every - /// aggregate in a hash aggregation returns `true` (and the group key is - /// supported), group state is stored in blocks and released block by - /// block. Defaults to `false`. + /// 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 } From 1d501ebdfa6d0034789b36fd6d0f1728a5cafa84 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:38:45 +0300 Subject: [PATCH 05/10] docs(agg): fix broken intra-doc links --- .../src/aggregates/aggregate_hash_table/common.rs | 4 ++++ .../physical-plan/src/aggregates/group_values/blocked/mod.rs | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) 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 3fad1c3bdf5f9..3480f96c129e6 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -121,6 +121,8 @@ pub(super) fn create_accumulator_storage( /// [`GroupsAccumulator`]. Both use columnar storage so aggregation can stay /// vectorized. /// +/// [`GroupValues`]: crate::aggregates::group_values::GroupValues +/// /// # Marker Type /// `AggrMode` selects the aggregate semantics. /// @@ -629,6 +631,8 @@ 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, diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs index 2371b0d78d7ab..75ad490d418ab 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked/mod.rs @@ -32,7 +32,7 @@ mod primitive; use primitive::BlockedGroupValuesPrimitive; /// Like [`GroupValues`], but the group keys are stored in blocks of -/// [`Self::block_size`] groups, matching the blocks of the +/// `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. /// From 1d12c0633175bf6f4089034f692661279c8e9e67 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:30:31 +0300 Subject: [PATCH 06/10] docs(agg): clarify that blocked keys mix with flat aggregates --- .../src/aggregates/aggregate_hash_table/storage.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs index a22b1ce679c96..bce3ef5ac8daf 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs @@ -23,8 +23,12 @@ //! 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, and a table with no blocked part runs exactly -//! the flat code. +//! 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 From 8fc173c180674fbb9b80e72cfa1743639f5f0c25 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:55:23 +0300 Subject: [PATCH 07/10] fix(agg): release early-emit reservation block by block --- .../src/aggregates/hash_stream.rs | 106 ++++++++++-------- .../src/aggregates/partial_reduce_stream.rs | 35 ++++-- 2 files changed, 80 insertions(+), 61 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index d4202af0989a4..cc9e5a59bff60 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -323,21 +323,33 @@ impl PartialHashAggregateStream { self.early_emit_count.add(1); timer.done(); + // 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 mut pending_memory: usize = materialized_group_states + 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(); - for batch in materialized_group_states { - pending_memory -= batch.memory_size; - self.emit_on_memory_pressure( - batch, - &mut emitter, - hash_table.memory_size() + pending_memory, - ) - .await?; + + // 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?; } } } @@ -445,64 +457,60 @@ impl PartialHashAggregateStream { // with batch slicing. remaining_groups: MaterializedBatch, emitter: &mut TryEmitter, - hash_table_mem_size: usize, ) -> Result<()> { let MaterializedBatch { batch: mut remaining_groups, memory_size: remaining_groups_memory, } = remaining_groups; - // 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)?; + // 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; diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d9001f4054751..421c52b22a545 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -30,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}; @@ -111,7 +113,7 @@ enum PartialReduceHashAggregateState { hash_table: AggregateHashTable, // After each incremental emitting step, the front of `remaining_groups` // is updated with batch slicing, and removed once fully emitted. - remaining_groups: VecDeque, + remaining_groups: VecDeque, }, ProducingOutput { hash_table: AggregateHashTable, @@ -295,8 +297,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 @@ -333,10 +335,7 @@ impl PartialReduceHashAggregateStream { ControlFlow::Continue( PartialReduceHashAggregateState::EmittingOnMemoryPressure { hash_table: original_state.into_hash_table(), - remaining_groups: remaining_groups - .into_iter() - .map(|b| b.batch) - .collect(), + remaining_groups: remaining_groups.into(), }, ) } @@ -362,17 +361,29 @@ impl PartialReduceHashAggregateStream { else { unreachable!("expected the EmittingOnMemoryPressure state") }; - let Some(batch) = remaining_groups.pop_front() else { + let Some(MaterializedBatch { batch, memory_size }) = remaining_groups.pop_front() + else { unreachable!("EmittingOnMemoryPressure always holds a batch") }; 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( - batch.slice(self.batch_size, batch.num_rows() - self.batch_size), - ); + 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() { From ce4e7c6779b0cfef1ac40a76c5f5cbbba6d4d9af Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:16:10 +0300 Subject: [PATCH 08/10] add unsafe, and always check in bounds also change the block size to be just batch size instead of the << 18 --- .../src/blocked_groups_accumulator.rs | 59 ++++++++++ .../groups_accumulator/blocked_vec.rs | 103 ++++++++++++++---- datafusion/functions-aggregate/src/count.rs | 22 ++-- .../aggregate_hash_table/storage.rs | 38 +++---- 4 files changed, 168 insertions(+), 54 deletions(-) diff --git a/datafusion/expr-common/src/blocked_groups_accumulator.rs b/datafusion/expr-common/src/blocked_groups_accumulator.rs index 0bc89b5e0746c..6c58a609c295c 100644 --- a/datafusion/expr-common/src/blocked_groups_accumulator.rs +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -78,6 +78,37 @@ impl BlocksIndex { 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) { @@ -213,6 +244,34 @@ pub trait BlockedGroupsAccumulator: Send + Any { 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!( 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 index a982d29f841ce..305c3d50610e5 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs @@ -232,11 +232,14 @@ impl BlockedVec { /// `group_indices` must only point to the first `total_num_groups` /// groups, which is the contract of /// [`BlockedGroupsAccumulator::update_batch`] and - /// [`BlockedGroupsAccumulator::merge_batch`]. It is only checked in debug - /// builds, like the flat accumulators' unchecked indexing. + /// [`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, @@ -244,15 +247,38 @@ impl BlockedVec { group_indices: &[BlocksIndex], nulls: Option<&NullBuffer>, opt_filter: Option<&BooleanArray>, - mut update_fn: F, + update_fn: F, ) where F: FnMut(&mut T), { self.grow_to(total_num_groups, starting_value); - self.debug_assert_in_bounds(group_indices); + 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, see the method docs + // SAFETY: indices are in bounds, guaranteed by the caller update_fn(unsafe { block.get_unchecked_mut(index.index_in_block()) }) }); return; @@ -263,7 +289,7 @@ impl BlockedVec { 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, see the method docs + // SAFETY: indices are in bounds, guaranteed by the caller *addr = unsafe { ptrs.ptr(index) }; } for &addr in &addrs[..chunk.len()] { @@ -274,7 +300,7 @@ impl BlockedVec { } } else { accumulate_blocked_indices(group_indices, nulls, opt_filter, |index| { - // SAFETY: indices are in bounds, see the method docs + // SAFETY: indices are in bounds, guaranteed by the caller update_fn(unsafe { &mut *ptrs.ptr(index) }) }); } @@ -289,31 +315,52 @@ impl BlockedVec { /// `group_indices` must only point to the first `total_num_groups` /// groups, which is the contract of /// [`BlockedGroupsAccumulator::update_batch`] and - /// [`BlockedGroupsAccumulator::merge_batch`]. It is only checked in debug - /// builds, like the flat accumulators' unchecked indexing. + /// [`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. + /// 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()); - self.grow_to(total_num_groups, starting_value); - self.debug_assert_in_bounds(group_indices); if let Some(block) = self.as_single_block_mut() { for (&index, &value) in group_indices.iter().zip(values) { - // SAFETY: indices are in bounds, see the method docs + // SAFETY: indices are in bounds, guaranteed by the caller update_fn( unsafe { block.get_unchecked_mut(index.index_in_block()) }, value, @@ -329,7 +376,7 @@ impl BlockedVec { .zip(values.chunks(RESOLVE_CHUNK)) { for (addr, &index) in addrs.iter_mut().zip(indices) { - // SAFETY: indices are in bounds, see the method docs + // SAFETY: indices are in bounds, guaranteed by the caller *addr = unsafe { ptrs.ptr(index) }; } for (&addr, &value) in addrs.iter().zip(values) { @@ -340,15 +387,13 @@ impl BlockedVec { } } - /// Checks, in debug builds only, that every index points to an element. + /// 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 debug_assert_in_bounds(&self, group_indices: &[BlocksIndex]) { - debug_assert!( - group_indices.iter().all(|index| { - self.blocks - .get(index.block_index()) - .is_some_and(|block| index.index_in_block() < block.len()) - }), + 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" ); } @@ -673,4 +718,18 @@ mod tests { 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 e18393fa1a292..d85c2fc07993d 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -892,16 +892,18 @@ impl BlockedGroupsAccumulator for BlockedCountGroupsAccumulator { let values = &values[0]; let nulls = values.logical_nulls().filter(|n| n.null_count() > 0); - // Add one to each group's counter for each non null, non - // filtered value - self.counts.update( - total_num_groups, - 0, - group_indices, - nulls.as_ref(), - opt_filter, - |count| *count += 1, - ); + 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(()) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs index bce3ef5ac8daf..676318a758eb3 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs @@ -60,15 +60,11 @@ use crate::aggregates::order::GroupOrdering; use super::common::{AggregateAccumulator, MaterializeAccumulatorFn}; -/// Smallest block, in groups. Tables with fewer groups stay in a single, -/// flat block and run exactly the flat hot loops. -const MIN_BLOCK_SIZE: usize = 1 << 18; - /// Number of groups per block for every blocked part of a table. /// -/// At least `batch_size`, so an output batch never spans two blocks. +/// `batch_size`, so every block is emitted as one output batch. pub(in crate::aggregates) fn block_size(batch_size: usize) -> usize { - batch_size.max(MIN_BLOCK_SIZE).next_power_of_two() + batch_size } /// Returns `true` if the group keys of `agg` can be stored in blocks: a single @@ -476,9 +472,7 @@ mod tests { 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), MIN_BLOCK_SIZE); - // The block always holds at least one output batch. - assert_eq!(block_size(MIN_BLOCK_SIZE + 1), MIN_BLOCK_SIZE * 2); + assert_eq!(block_size(8192), 8192); Ok(()) } @@ -486,13 +480,14 @@ mod tests { /// sliced at the block boundaries, and their memory is counted once. #[test] fn mixed_storage_output_is_aligned_and_counted_once() -> Result<()> { - let num_groups = MIN_BLOCK_SIZE + 10; + 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, 8192)?; + let mut table = single_table(&agg, batch_size)?; table.aggregate_batch(&input)?; let batches = table.take_state_batches()?; @@ -501,7 +496,7 @@ mod tests { .iter() .map(|b| b.batch.num_rows()) .collect::>(), - vec![MIN_BLOCK_SIZE, 10] + vec![batch_size, 10] ); for b in &batches { let keys = b.batch.column(0).as_primitive::().values(); @@ -527,8 +522,8 @@ mod tests { /// released once all its slices are handed out. #[test] fn output_releases_state_block_by_block() -> Result<()> { - let num_groups = MIN_BLOCK_SIZE + 10; - let batch_size = 8192; + 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))], @@ -563,15 +558,14 @@ mod tests { keys.sort_unstable(); assert_eq!(keys, (0..num_groups as i64).collect::>()); - // While the first block is sliced, the table still holds the second - // block; once the first block is handed out, only the small second - // block (10 groups) remains. - let first_block_batches = MIN_BLOCK_SIZE / batch_size; - assert!(memory_per_batch[0] <= memory_while_building); + // 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[first_block_batches] < memory_while_building / 4, - "memory after the first block: {}, while building: {memory_while_building}", - memory_per_batch[first_block_batches] + 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(()) From 51fca432f9bd5ad71a651b0b4935109cac97cd3a Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:50:58 +0300 Subject: [PATCH 09/10] fix(agg): spill all blocks as one sorted run --- .../physical-plan/src/aggregates/spill.rs | 85 +++++++++++++++---- 1 file changed, 69 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 20dfa54dcd0bc..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}; @@ -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( @@ -248,22 +317,6 @@ impl AggregateSpill { Ok(()) } - /// [`Self::sort_and_spill`] for the state batches of a hash aggregate - /// table: one batch for flat storage, one per block for blocked storage, - /// each written as its own spill file. - pub(super) fn sort_and_spill_batches( - &mut self, - state_batches: Vec, - ) -> Result<()> { - // TODO: sort across multiple arrays without concat (see #24928 - // sort.rs/rank.rs) so all blocks are written as one sorted run instead - // of one run per block. - for state_batch in state_batches { - self.sort_and_spill(Some(state_batch.batch))?; - } - Ok(()) - } - /// Merges every sorted run, and does the aggregate evaluation with /// [`OrderedFinalAggregateStream`]. pub(super) fn into_replay_stream( From c23a91fbadf6fec9c157a3dda56d837e42a89485 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 5 Oct 2026 14:28:06 +0300 Subject: [PATCH 10/10] fix compile --- datafusion/physical-plan/src/aggregates/single_stream.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index e4bde41ac585a..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_batches() - .and_then(|batch| spill_context.sort_and_spill_batches(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.