From 4e83deee1ee030759686b197709ce76706a66b63 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 12:34:10 +0300 Subject: [PATCH 1/8] feat(agg): add initial Blocked aggregate api --- .../src/blocked_groups_accumulator.rs | 441 ++++++++++++++ .../blocked_boolean_buffer_builder.rs | 485 +++++++++++++++ .../src/blocked_helpers/blocked_vec.rs | 565 ++++++++++++++++++ .../expr-common/src/blocked_helpers/mod.rs | 8 + datafusion/expr-common/src/lib.rs | 2 + datafusion/expr/src/function.rs | 1 + datafusion/expr/src/lib.rs | 6 + datafusion/expr/src/udaf.rs | 44 +- .../src/accumulator.rs | 33 + .../src/aggregate.rs | 1 + .../aggregate/blocked_groups_accumulator.rs | 208 +++++++ .../blocked_groups_accumulator/accumulate.rs | 391 ++++++++++++ .../blocked_prim_op.rs | 239 ++++++++ .../groups_accumulator/accumulate.rs | 12 +- datafusion/functions-aggregate/src/count.rs | 219 ++++++- datafusion/functions-aggregate/src/sum.rs | 28 +- datafusion/physical-expr/src/aggregate.rs | 79 ++- .../aggregates/aggregate_hash_table/common.rs | 196 +++--- .../aggregate_hash_table/common_ordered.rs | 83 ++- .../aggregate_hash_table/partial_table.rs | 4 +- .../src/aggregates/group_values/blocked.rs | 278 +++++++++ .../src/aggregates/group_values/mod.rs | 3 + .../single_group_by/blocked_boolean.rs | 282 +++++++++ .../single_group_by/blocked_primitive.rs | 291 +++++++++ .../group_values/single_group_by/mod.rs | 2 + .../src/aggregates/grouped_hash_stream.rs | 39 +- .../src/aggregates/hash_stream.rs | 106 ++-- .../physical-plan/src/aggregates/order/mod.rs | 12 +- .../src/aggregates/order/partial.rs | 19 +- .../src/aggregates/ordered_final_stream.rs | 36 +- .../src/aggregates/ordered_partial_stream.rs | 55 +- .../src/aggregates/ordered_single_stream.rs | 38 +- .../src/aggregates/partial_reduce_stream.rs | 32 +- .../src/aggregates/single_stream.rs | 42 +- 34 files changed, 3961 insertions(+), 319 deletions(-) create mode 100644 datafusion/expr-common/src/blocked_groups_accumulator.rs create mode 100644 datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs create mode 100644 datafusion/expr-common/src/blocked_helpers/blocked_vec.rs create mode 100644 datafusion/expr-common/src/blocked_helpers/mod.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs create mode 100644 datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs create mode 100644 datafusion/physical-plan/src/aggregates/group_values/blocked.rs create mode 100644 datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs create mode 100644 datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs 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..50fb30d4706ca --- /dev/null +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -0,0 +1,441 @@ +// 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 [`GroupsAccumulator`] + +use arrow::array::{ArrayRef, BooleanArray}; +use datafusion_common::{ + Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err, not_impl_err, + utils::split_vec_min_alloc, +}; +use std::cmp::Ordering; +use crate::groups_accumulator::EmitTo; + +/// Selects groups for a non-destructive grouped aggregation read. +/// +/// Unlike [`BlockedEmitTo`], this selection does not remove groups or change their +/// indices. Selections created by [`Self::try_from_indices`] preserve the +/// requested order and support duplicate indices. +/// +/// A selection is validated once when it is constructed and can then be reused +/// for the group values and accumulators participating in the same snapshot. +/// Construct a new selection if the number or indexing of those groups changes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlockedGroupSelection<'a> { + block_size: usize, + total_num_groups: usize, + indices: Option<&'a [BlocksIndex]>, +} + +impl<'a> BlockedGroupSelection<'a> { + /// Selects all `total_num_groups` groups in group-index order. + pub fn all(total_num_groups: usize, block_size: usize) -> Self { + Self { + total_num_groups, + indices: None, + block_size, + } + } + + /// Selects groups in the order specified by `indices`. + /// + /// Returns an error if an index is not less than `total_num_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 { + let total_num_groups_parsed = + BlocksIndex::from_index_in_fixed_block_size(total_num_groups, block_size); + if let Some(index) = indices + .iter() + .find(|&&index| index >= total_num_groups_parsed) + { + return exec_err!( + "Group index {index:?} is out of bounds for {total_num_groups_parsed:?} ({total_num_groups}) groups" + ); + } + Ok(Self { + total_num_groups, + indices: Some(indices), + block_size, + }) + } + + /// 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. + /// + /// Preserving-read implementations should call this method with their + /// stored group count before using [`Self::iter`]. This check is `O(1)`; + /// the selected indices were already checked by [`Self::try_from_indices`]. + 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 group indices in output order. + pub fn iter(self) -> impl Iterator + 'a { + let (all, indices): (_, &'a [BlocksIndex]) = match self.indices { + None => (0..self.total_num_groups, &[]), + Some(indices) => (0..0, indices), + }; + let block_size = self.block_size; + all.map(move |index| { + BlocksIndex::from_index_in_fixed_block_size(index, block_size) + }) + .chain(indices.iter().copied()) + } + + #[doc(hidden)] + pub fn indices(&self) -> Option<&[BlocksIndex]> { + self.indices + } +} + +/// Two `u32`s (8 bytes) rather than two `usize`s: an index is stored per group in every blocked +/// hash table entry, so its size decides the size of the tables. Both halves fit `u32` by a wide +/// margin, `index_in_block < block_size` (the batch size) and the number of blocks is +/// `groups / block_size`. The API keeps taking and returning `usize`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)] +pub struct BlocksIndex { + block_index: u32, + index_in_block: u32, +} + +impl BlocksIndex { + pub const ZERO: Self = Self { + // flat_index: 0, + block_index: 0, + index_in_block: 0, + }; + pub const MAX: Self = Self { + // flat_index: usize::MAX, + block_index: u32::MAX, + index_in_block: u32::MAX, + }; + + pub fn new(block_index: usize, index_in_block: usize) -> Self { + Self { + block_index: Self::narrow(block_index), + index_in_block: Self::narrow(index_in_block), + } + } + + #[inline(always)] + fn narrow(value: usize) -> u32 { + debug_assert!(value <= u32::MAX as usize, "block index part {value} does not fit in u32"); + value as u32 + } + + #[inline(always)] + pub fn new_in_first_block(index_in_block: usize) -> Self { + // Implementation note: + // not having From that will do this instead even when it will be more convenient + // so we can later change the layout to be a single usize with bit shifts + Self::new(0, index_in_block) + // Self { + // flat_index: index_in_block, + // } + } + + #[inline(always)] + pub fn from_index_in_fixed_block_size(index: usize, block_size: usize) -> Self { + Self::new(index / block_size, index % block_size) + } + + #[inline(always)] + pub fn into_index_in_fixed_block_size(self, block_size: usize) -> usize { + self.block_index as usize * block_size + self.index_in_block as usize + } + + #[inline(always)] + pub fn block_index(&self) -> usize { + self.block_index as usize + } + + #[inline(always)] + pub fn index_in_block(&self) -> usize { + self.index_in_block as usize + } + + pub fn sub_flat_checked(self, rhs_flat: usize, block_size: usize) -> Option { + let self_flat = self.into_index_in_fixed_block_size(block_size); + self_flat.checked_sub(rhs_flat).map(|v| Self::from_index_in_fixed_block_size(v, block_size)) + } + + pub fn sub_flat(self, rhs_flat: usize, block_size: usize) -> Self { + Self::from_index_in_fixed_block_size(self.into_index_in_fixed_block_size(block_size) - rhs_flat, block_size) + } + + pub fn gte_flat(self, rhs_flat: usize, block_size: usize) -> bool { + self.into_index_in_fixed_block_size(block_size) >= rhs_flat + } + + pub fn prev_block(mut self) -> Self { + self.block_index -= 1; + + self + } + + pub fn prev_block_checked(mut self) -> Option { + self.block_index.checked_sub(1).map(|v| { + self.block_index = v; + self + }) + } + + pub fn add(self, n: usize, block_size: usize) -> BlocksIndex { + if self.index_in_block + n as u32 >= block_size as u32 { + let new_flat = self.into_index_in_fixed_block_size(block_size) + n; + Self::from_index_in_fixed_block_size(new_flat, block_size) + } else { + Self::new(self.block_index as usize, self.index_in_block as usize + n) + } + } + + pub fn increment(mut self, block_size: usize) -> Self { + if self.index_in_block + 1 >= block_size as u32 { + self.block_index += 1; + self.index_in_block = 0; + } else { + self.index_in_block += 1; + } + + self + } +} + +impl PartialOrd for BlocksIndex { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for BlocksIndex { + fn cmp(&self, other: &Self) -> Ordering { + self.block_index.cmp(&other.block_index).then(self.index_in_block.cmp(&other.index_in_block)) + } +} + +/// Describes how many rows should be emitted during grouping. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BlockedEmitTo { + All, + /// Emit next group + NextBlock, + /// Emit only the first `n` groups and shift all existing group + /// indexes down by `n`. + /// + /// For example, if `n=10`, group_index `0, 1, ... 9` are emitted + /// and group indexes `10, 11, 12, ...` become `0, 1, 2, ...`. + /// + /// Requirements: + /// 1. `n` must be smaller than block_size + /// 2. `n` is not 0 + First(usize), +} + +impl BlockedEmitTo { + pub fn into_emit_to( + self, + len: usize, + block_size: usize, + ) -> Result + Clone> { + let mut to_emit = match self { + Self::All => len, + Self::NextBlock => len.min(block_size), + Self::First(n) => { + assert_ne_or_internal_err!(n, 0); + assert_or_internal_err!( + n <= len, + "n ({n}) must be less than or equal current length ({})", + len + ); + assert_or_internal_err!( + n < block_size, + "n ({n}) must be less than current block size ({})", + block_size + ); + n + } + }; + + let mut blocks = vec![]; + + let emitting_all = len == to_emit; + + while to_emit > block_size { + blocks.push(EmitTo::First(block_size)); + to_emit -= block_size + } + + if to_emit > 0 && emitting_all { + blocks.push(EmitTo::All); + } else if to_emit > 0 { + blocks.push(EmitTo::First(to_emit)); + } + + Ok(blocks.into_iter()) + } +} + +/// `GroupsAccumulator` implements a single aggregate (e.g. AVG) and +/// stores the state for *all* groups internally. +/// +/// Logically, a [`GroupsAccumulator`] stores a mapping from each group index to +/// the state of the aggregate for that group. For example an implementation for +/// `min` might look like +/// +/// ```text +/// ┌─────┐ +/// │ 0 │───────────▶ 100 +/// ├─────┤ +/// │ 1 │───────────▶ 200 +/// └─────┘ +/// ... ... +/// ┌─────┐ +/// │ N-2 │───────────▶ 50 +/// ├─────┤ +/// │ N-1 │───────────▶ 200 +/// └─────┘ +/// +/// +/// Logical group Current Min +/// number value for that +/// group +/// ``` +/// +/// # Notes on Implementing `GroupsAccumulator` +/// +/// All aggregates must first implement the simpler [`Accumulator`] trait, which +/// handles state for a single group. Implementing `GroupsAccumulator` is +/// optional and is harder to implement than `Accumulator`, but can be much +/// faster for queries with many group values. See the [Aggregating Millions of +/// Groups Fast blog] for more background. +/// For more background, please also see the [Aggregating Millions of Groups Fast in Apache Arrow DataFusion 28.0.0 blog] +/// +/// [Aggregating Millions of Groups Fast in Apache Arrow DataFusion 28.0.0 blog]: https://datafusion.apache.org/blog/2023/08/05/datafusion_fast_grouping +/// +/// [`NullState`] can help keep the state for groups that have not seen any +/// values and produce the correct output for those groups. +/// +/// [`NullState`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/struct.NullState.html +/// +/// # Details +/// Each group is assigned a `group_index` by the hash table and each +/// accumulator manages the specific state, one per `group_index`. +/// +/// `group_index`es are contiguous (there aren't gaps), and thus it is +/// expected that each `GroupsAccumulator` will use something like `Vec<..>` +/// to store the group states. +/// +/// [`Accumulator`]: crate::accumulator::Accumulator +/// [Aggregating Millions of Groups Fast blog]: https://arrow.apache.org/blog/2023/08/05/datafusion_fast_grouping/ +pub trait BlockedGroupsAccumulator: Send + std::any::Any { + fn batch_size(&self) -> usize; + + fn update_batch( + &mut self, + values: &[ArrayRef], + // Get blocks index rather than flat index to avoid extra computation when inserting and searching (BlocksIndex can be implemented as flat usize) + group_indices: &[BlocksIndex], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()>; + + /// Same as `GroupsAccumulator::evaluate` but returning vec of blocks and different EmitTo + /// For + /// - `BlockedEmitTo::All` it should return `Vec` + /// - `BlockedEmitTo::NextBlock` it should return single item vector with the block or empty vec in case of no blocks + /// - `BlockedEmitTo::First(n)` it should return single item vector with the first n rows in the first block. n must be smaller than block size and length + /// + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result>; + + // Same as `GroupsAccumulator::evaluate_preserving` but with blocked indices + 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 + } + + /// Same as `GroupsAccumulator::state` but returning vec of blocks and different EmitTo + /// For + /// - `BlockedEmitTo::All` it should return `Vec` + /// - `BlockedEmitTo::NextBlock` it should return single item vector with the block or empty vec in case of no blocks + /// - `BlockedEmitTo::First(n)` it should return single item vector with the first n rows in the first block. n must be smaller than block size and length + /// + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>>; + + /// Same as `GroupsAccumulator::state_preserving` but with blocked indices + 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 + } + fn merge_batch( + &mut self, + values: &[ArrayRef], + + // Get blocks index rather than flat index to avoid extra computation when inserting and searching (BlocksIndex can be implemented as flat usize) + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()>; + + /// Same as `GroupsAccumulator::convert_to_state` + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result>; + + /// Same as `GroupsAccumulator::size` + fn size(&self) -> usize; +} diff --git a/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs b/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs new file mode 100644 index 0000000000000..73c6425c0d7fe --- /dev/null +++ b/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs @@ -0,0 +1,485 @@ +use datafusion_common::utils::proxy::{VecAllocExt}; +use std::collections::VecDeque; +use std::fmt::Debug; +use std::ops::{Index, IndexMut}; +use arrow::array::BooleanBufferBuilder; +use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow::datatypes::ArrowNativeType; +use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; + + +/// Blocked Vec +/// +/// # Implementation Notes +/// +/// ## Why `T: Copy`? +/// +/// +/// 1. So the [`BlockedVec::allocated_size`] will be accurate since the size of T is known, and not include heap allocations (like `String` or `Vec`) that are not part of the allocated size of the `BlockedVec` +/// 2. So we can provide mutable access to the items (e.g. `IndexMut`) since if `T` is not `Copy` (like when `T` is a `Vec`) we could change the size of it without the [`BlockedVec::allocated_size`] changing, which would be confusing and lead to bugs +/// +#[derive(Debug)] +pub struct BlockedBooleanBufferBuilder { + /// Using `VecDeque` so we can remove the first block and reclaim memory + blocks: VecDeque, + + /// The size of each block + block_size: usize, + + /// The total number of items, not the number of offset since in each block there is the initial offset + len: usize, + + /// The index of the current block + current_block_index: usize, + should_count_current_block: bool, + + finished_blocks_allocated_memory: usize, +} + +impl BlockedBooleanBufferBuilder +{ + // TODO - some want to preallocate the blocks and some don't, + // there should be a way while avoiding having a lot of memory used if all are prealocatting + pub fn new(block_size: usize) -> Self { + assert_ne!(block_size, 0, "block size must be greater than 0"); + + let blocks = VecDeque::from(vec![BooleanBufferBuilder::new(block_size)]); + Self { + blocks, + block_size, + should_count_current_block: false, + len: 0, + current_block_index: 0, + finished_blocks_allocated_memory: 0, + } + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn num_blocks(&self) -> usize { + self.current_block_index + (self.should_count_current_block() as usize) + } + + fn should_count_current_block(&self) -> bool { + self.should_count_current_block || !self.blocks[self.current_block_index].is_empty() + } + + pub fn block(&self, block_index: usize) -> &BooleanBufferBuilder { + &self.blocks[block_index] + } + + pub fn block_size(&self) -> usize { + self.block_size + } + + pub fn allocated_size(&self) -> usize { + self.finished_blocks_allocated_memory + + self.blocks.capacity() * size_of::() + + self.blocks.back().map_or(0, |b| builder_allocated_size(b)) + } + + /// Get the number of elements in the current block + pub fn current_block_len(&self) -> usize { + self.blocks[self.current_block_index].len() + } + + pub fn start_new_block(&mut self) { + // a block that was only pre-opened becomes the started block instead of an empty + // block of its own + if self.should_count_current_block() { + self.end_current_block(); + } + self.should_count_current_block = true; + } + + pub(crate) fn mark_current_block_counted(&mut self) { + self.should_count_current_block = true; + } + + pub fn end_current_block(&mut self) { + self.end_current_block_inner(); + } + + fn end_current_block_inner(&mut self) { + // Don't add to number of blocks since we might not insert into it + self.current_block_index += 1; + self.finished_blocks_allocated_memory += + self.blocks.back().map_or(0, |b| builder_allocated_size(b)); + let new_block = BooleanBufferBuilder::new(self.block_size); + + // Don't count current block since we might not insert into it + self.should_count_current_block = false; + self.blocks.push_back(new_block); + } + + pub(crate) fn reserve_blocks(&mut self, n: usize) { + self.blocks.reserve(n); + } + + /// Push length and return if the current block is now full + pub fn push(&mut self, value: bool) -> bool { + let block = &mut self.blocks[self.current_block_index]; + + block.append(value); + self.len += 1; + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Extends iterator of lengths within current block + /// Returns if the current block has finished + /// + /// # Panics + /// Panics if the iterator length exceeds the remaining size of the current block + pub(super) fn extend_in_block( + &mut self, + iter: impl Iterator, + ) -> bool { + let block = &mut self.blocks[self.current_block_index]; + + let prev_block_len = block.len(); + + // TODO - add to arrow extend support for iterators of bools, but for now just append one by one + for value in iter { + block.append(value); + } + + assert!( + block.len() <= self.block_size, + "overflow from block new block length: {}, block size: {}", + block.len(), + self.block_size + ); + + let added_items = block.len() - prev_block_len; + self.len += added_items; + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Extends from slice within current block + /// Returns if the current block has finished + /// + /// # Panics + /// Panics if the iterator length exceeds the remaining size of the current block + pub(super) fn extend_from_slice_in_block( + &mut self, + slice: &[bool], + ) -> bool + { + let block = &mut self.blocks[self.current_block_index]; + + let prev_block_len = block.len(); + block.append_slice(slice); + + assert!( + block.len() <= self.block_size, + "overflow from block new block length: {}, block size: {}", + block.len(), + self.block_size + ); + + let added_items = block.len() - prev_block_len; + self.len += added_items; + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Extend the length from the current offsets + pub fn extend_from_slice( + &mut self, + mut buffer: &[bool], + ) { + let number_of_blocks_to_reserve = buffer + .len() + .saturating_sub(self.current_block_remaining_len()) + .div_ceil(self.block_size); + self.reserve_blocks(number_of_blocks_to_reserve); + + while !buffer.is_empty() { + let remaining_in_current_block = self.current_block_remaining_len(); + let to_add = remaining_in_current_block.min(buffer.len()); + + let (to_copy, rest) = buffer.split_at(to_add); + buffer = rest; + + self.extend_from_slice_in_block(to_copy); + } + } + + pub fn current_block_remaining_len(&self) -> usize { + self.block_size - self.blocks[self.current_block_index].len() + } + + pub fn push_value_n_within_block( + &mut self, + value: bool, + n: usize, + ) -> bool + { + self.len += n; + let block = &mut self.blocks[self.current_block_index]; + + let new_len = block.len() + n; + assert!( + new_len <= self.block_size, + "overflow from block new block length: {new_len}, block size: {}", + self.block_size + ); + block.append_n(n, value); + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + pub fn push_value_n( + &mut self, + value: bool, + mut n: usize, + ) { + let number_of_blocks_to_reserve = n + .saturating_sub(self.current_block_remaining_len()) + .div_ceil(self.block_size); + self.reserve_blocks(number_of_blocks_to_reserve); + + while n > 0 { + let remaining_in_current_block = self.current_block_remaining_len(); + + let to_add = remaining_in_current_block.min(n); + n -= to_add; + + self.push_value_n_within_block(value, to_add); + } + } + + pub fn push_value_n_to_len(&mut self, value: bool, new_len: usize) { + assert!(new_len >= self.len, "new_len must be greater than or equal to current len"); + let n = new_len - self.len; + self.push_value_n(value, n); + } + + pub fn get_bit(&self, index: BlocksIndex) -> bool { + let block_index = index.block_index(); + let index_in_block = index.index_in_block(); + + self.blocks[block_index].get_bit(index_in_block) + } + + pub fn set_bit(&mut self, index: BlocksIndex, value: bool) { + let block_index = index.block_index(); + let index_in_block = index.index_in_block(); + + self.blocks[block_index].set_bit(index_in_block, value) + } + + /// Block `block_index` without bounds checking + /// + /// # Safety + /// `block_index < self.num_blocks()` + #[inline] + pub unsafe fn block_unchecked(&self, block_index: usize) -> &BooleanBufferBuilder { + debug_assert!(block_index < self.blocks.len()); + unsafe { self.blocks.get(block_index).unwrap_unchecked() } + } + + /// Mutable block `block_index` without bounds checking + /// + /// # Safety + /// `block_index < self.num_blocks()` + #[inline] + pub unsafe fn block_unchecked_mut( + &mut self, + block_index: usize, + ) -> &mut BooleanBufferBuilder { + debug_assert!(block_index < self.blocks.len()); + unsafe { self.blocks.get_mut(block_index).unwrap_unchecked() } + } + + pub fn blocks_mut(&mut self) -> impl Iterator { + let num_blocks = self.num_blocks(); + self.blocks.iter_mut().take(num_blocks) + } + + pub fn current_block_mut(&mut self) -> &mut BooleanBufferBuilder { + &mut self.blocks[self.current_block_index] + } + + pub fn current_or_block_mut(&mut self, block_index: usize) -> &mut BooleanBufferBuilder { + &mut self.blocks[block_index] + } + + pub fn emit(&mut self, emit_to: BlockedEmitTo) -> Vec { + match emit_to { + BlockedEmitTo::All => { + let counts = self.take_all(); + + counts + } + BlockedEmitTo::NextBlock => { + self.take_block().map_or(vec![], |next| vec![next]) + } + BlockedEmitTo::First(n) => { + vec![self.take_n(n)] + } + } + } + + /// Take every block that counts, see [`Self::num_blocks`] + pub fn take_all(&mut self) -> Vec { + let num_blocks = self.num_blocks(); + let mut blocks = std::mem::take(&mut self.blocks); + blocks.truncate(num_blocks); + self.reset(); + + blocks.into_iter().map(BooleanBuffer::from).collect() + } + + /// Take the first block, `None` once there are no more items + pub fn take_block(&mut self) -> Option { + if self.num_blocks() == 0 { + return None; + } + Some(self.take_first_block()) + } + + /// Take the first block even when it is empty, for callers that know from + /// elsewhere that the block holds items + pub fn take_first_block(&mut self) -> BooleanBuffer { + let block = self.blocks.pop_front().expect("always at least one block"); + + if self.blocks.is_empty() { + self.current_block_index = 0; + self.should_count_current_block = false; + self.blocks.push_back(BooleanBufferBuilder::new(self.block_size)); + } else { + self.current_block_index -= 1; + + // Only reduce memory if not the last one since the last block is calculated separately + self.finished_blocks_allocated_memory -= builder_allocated_size(&block); + } + + self.len -= block.len(); + + block.build() + } + + pub fn take_n( + &mut self, + n: usize, + ) -> BooleanBuffer { + assert_ne!(n, 0, "n must be greater than 0"); + assert!(n <= self.len, "n ({n}) must be <= len ({})", self.len); + assert!(n < self.block_size, "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", self.block_size); + + if n == self.len { + let blocks = self.take_all(); + return blocks.into_iter().next().unwrap_or_else(|| BooleanBuffer::new_unset(0)) + } + + if n == self.blocks[0].len() { + return self.take_block().expect("must have at least one block since n < len and n != 0"); + } + + let prev_len = self.len; + + + debug_assert!( + self.blocks.back().is_some_and(|block| block.len() < self.block_size), + "the last block must not be full (since it should be the writable tail)" + ); + + // The emitted items are always fully contained in the first block + let mut taken = BooleanBufferBuilder::new(n); + taken.append_packed_range(0..n, self.blocks[0].as_slice()); + + { + let first = &mut self.blocks[0]; + let first_len = first.len(); + + // TODO - reuse the same allocation by adding copy_within + let mut new_first = BooleanBufferBuilder::new(first_len - n); + new_first.append_packed_range(n..first_len, first.as_slice()); + *first = new_first; + } + + // Every block gives its first `n` items to the end of the previous block and + // shifts the rest down, so every finished block stays full and keeps its allocation + for i in 1..self.blocks.len() { + let mut block = std::mem::replace(&mut self.blocks[i], BooleanBufferBuilder::new(self.block_size)); + let block_len = block.len(); + let moved = n.min(block_len); + + self.blocks[i - 1].append_packed_range(0..moved, block.as_slice()); + + // Copy to the new block + self.blocks[i].append_packed_range(moved..block.len(), block.as_slice()); + } + + // The old last block is the only one that can end up empty. + // If the block before it is full it is the writable tail as is and keeps its allocation, + // otherwise the block before it is the writable tail and the empty one is not part of the layout + let tail_is_empty = self.blocks.back().is_some_and(|b| b.is_empty()); + let prev_is_full = self.blocks.len() >= 2 && self.blocks[self.blocks.len() - 2].len() == self.block_size; + + if tail_is_empty && !prev_is_full { + self.blocks.pop_back(); + } + + self.len = prev_len - n; + self.current_block_index = self.blocks.len() - 1; + self.should_count_current_block = !tail_is_empty || !prev_is_full; + self.finished_blocks_allocated_memory = self + .blocks + .iter() + .take(self.current_block_index) + .map(|block| builder_allocated_size(block)) + .sum(); + + taken.build() + } + + pub fn reset(&mut self) { + self.blocks = VecDeque::from(vec![BooleanBufferBuilder::new(self.block_size)]); + self.len = 0; + self.current_block_index = 0; + self.finished_blocks_allocated_memory = 0; + self.should_count_current_block = false; + } +} + +fn builder_allocated_size(block: &BooleanBufferBuilder) -> usize { + block.capacity() / 8 // capacity in bits +} diff --git a/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs b/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs new file mode 100644 index 0000000000000..beb3feccf8f17 --- /dev/null +++ b/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs @@ -0,0 +1,565 @@ +use datafusion_common::utils::proxy::{VecAllocExt}; +use std::collections::VecDeque; +use std::fmt::Debug; +use std::ops::{Index, IndexMut}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::ArrowNativeType; +use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; + + +/// Blocked Vec +/// +/// # Implementation Notes +/// +/// ## Why `T: Copy`? +/// +/// +/// 1. So the [`BlockedVec::allocated_size`] will be accurate since the size of T is known, and not include heap allocations (like `String` or `Vec`) that are not part of the allocated size of the `BlockedVec` +/// 2. So we can provide mutable access to the items (e.g. `IndexMut`) since if `T` is not `Copy` (like when `T` is a `Vec`) we could change the size of it without the [`BlockedVec::allocated_size`] changing, which would be confusing and lead to bugs +/// +#[derive(Debug)] +pub struct BlockedVec { + /// Using `VecDeque` so we can remove the first block and reclaim memory + blocks: VecDeque>, + + /// The size of each block + block_size: usize, + + /// The total number of items, not the number of offset since in each block there is the initial offset + len: usize, + + /// The index of the current block + current_block_index: usize, + should_count_current_block: bool, + + finished_blocks_allocated_memory: usize, +} + +impl BlockedVec +{ + // TODO - some want to preallocate the blocks and some don't, + // there should be a way while avoiding having a lot of memory used if all are prealocatting + pub fn new(block_size: usize) -> Self { + assert_ne!(block_size, 0, "block size must be greater than 0"); + + let blocks = VecDeque::from(vec![vec![]]); + BlockedVec { + blocks, + block_size, + should_count_current_block: false, + len: 0, + current_block_index: 0, + finished_blocks_allocated_memory: 0, + } + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn num_blocks(&self) -> usize { + self.current_block_index + (self.should_count_current_block() as usize) + } + + fn should_count_current_block(&self) -> bool { + self.should_count_current_block || !self.blocks[self.current_block_index].is_empty() + } + + pub fn block(&self, block_index: usize) -> &Vec { + &self.blocks[block_index] + } + + pub fn block_size(&self) -> usize { + self.block_size + } + + pub fn allocated_size(&self) -> usize { + self.finished_blocks_allocated_memory + + self.blocks.capacity() * size_of::>() + + self.blocks.back().map_or(0, |b| b.allocated_size()) + } + + /// Get the number of elements in the current block + pub fn current_block_len(&self) -> usize { + self.blocks[self.current_block_index].len() + } + + pub fn start_new_block(&mut self) { + // a block that was only pre-opened becomes the started block instead of an empty + // block of its own + if self.should_count_current_block() { + self.end_current_block(); + } + self.should_count_current_block = true; + } + + pub(crate) fn mark_current_block_counted(&mut self) { + self.should_count_current_block = true; + } + + pub fn end_current_block(&mut self) { + self.end_current_block_inner(); + } + + fn end_current_block_inner(&mut self) { + // Don't add to number of blocks since we might not insert into it + self.current_block_index += 1; + self.finished_blocks_allocated_memory += + self.blocks.back().map_or(0, |b| b.allocated_size()); + let new_block = vec![]; + + // Don't count current block since we might not insert into it + self.should_count_current_block = false; + self.blocks.push_back(new_block); + } + + pub(crate) fn reserve_blocks(&mut self, n: usize) { + self.blocks.reserve(n); + } + + /// Push length and return if the current block is now full + pub fn push(&mut self, value: T) -> bool { + let block = &mut self.blocks[self.current_block_index]; + + block.push(value); + self.len += 1; + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Extends iterator of lengths within current block + /// Returns if the current block has finished + /// + /// # Panics + /// Panics if the iterator length exceeds the remaining size of the current block + pub(super) fn extend_in_block( + &mut self, + iter: impl Iterator, + ) -> bool { + let block = &mut self.blocks[self.current_block_index]; + + let prev_block_len = block.len(); + block.extend(iter); + + assert!( + block.len() <= self.block_size, + "overflow from block new block length: {}, block size: {}", + block.len(), + self.block_size + ); + + let added_items = block.len() - prev_block_len; + self.len += added_items; + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Extends from slice within current block + /// Returns if the current block has finished + /// + /// # Panics + /// Panics if the iterator length exceeds the remaining size of the current block + pub(super) fn extend_from_slice_in_block( + &mut self, + slice: &[T], + ) -> bool + { + let block = &mut self.blocks[self.current_block_index]; + + let prev_block_len = block.len(); + block.copy_from_slice(slice); + + assert!( + block.len() <= self.block_size, + "overflow from block new block length: {}, block size: {}", + block.len(), + self.block_size + ); + + let added_items = block.len() - prev_block_len; + self.len += added_items; + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Extend the length from the current offsets + pub fn extend_from_slice( + &mut self, + mut buffer: &[T], + ) { + let number_of_blocks_to_reserve = buffer + .len() + .saturating_sub(self.current_block_remaining_len()) + .div_ceil(self.block_size); + self.reserve_blocks(number_of_blocks_to_reserve); + + while !buffer.is_empty() { + let remaining_in_current_block = self.current_block_remaining_len(); + let to_add = remaining_in_current_block.min(buffer.len()); + + let (to_copy, rest) = buffer.split_at(to_add); + buffer = rest; + + self.extend_from_slice_in_block(to_copy); + } + } + + pub fn current_block_remaining_len(&self) -> usize { + self.block_size - self.blocks[self.current_block_index].len() + } + + pub fn push_value_n_within_block( + &mut self, + value: T, + n: usize, + ) -> bool + { + self.len += n; + let block = &mut self.blocks[self.current_block_index]; + + let new_len = block.len() + n; + assert!( + new_len <= self.block_size, + "overflow from block new block length: {new_len}, block size: {}", + self.block_size + ); + block.resize(new_len, value); + + let finished_block = block.len() == self.block_size; + + if finished_block { + self.end_current_block_inner(); + true + } else { + false + } + } + + /// Push default + pub fn push_default_n(&mut self, n: usize) + where + T: Default, + { + self.push_value_n(T::default(), n); + } + + pub fn push_value_n( + &mut self, + value: T, + mut n: usize, + ) { + let number_of_blocks_to_reserve = n + .saturating_sub(self.current_block_remaining_len()) + .div_ceil(self.block_size); + self.reserve_blocks(number_of_blocks_to_reserve); + + while n > 0 { + let remaining_in_current_block = self.current_block_remaining_len(); + + let to_add = remaining_in_current_block.min(n); + n -= to_add; + + self.push_value_n_within_block(value, to_add); + } + } + + pub fn push_value_n_to_len(&mut self, value: T, new_len: usize) { + assert!(new_len >= self.len, "new_len must be greater than or equal to current len"); + let n = new_len - self.len; + self.push_value_n(value, n); + } + + /// Block `block_index` without bounds checking + /// + /// # Safety + /// `block_index < self.num_blocks()` + #[inline] + pub unsafe fn block_unchecked(&self, block_index: usize) -> &Vec { + debug_assert!(block_index < self.blocks.len()); + unsafe { self.blocks.get(block_index).unwrap_unchecked() } + } + + /// Mutable block `block_index` without bounds checking + /// + /// # Safety + /// `block_index < self.num_blocks()` + #[inline] + pub unsafe fn block_unchecked_mut( + &mut self, + block_index: usize, + ) -> &mut Vec { + debug_assert!(block_index < self.blocks.len()); + unsafe { self.blocks.get_mut(block_index).unwrap_unchecked() } + } + + pub fn blocks_mut(&mut self) -> impl Iterator> { + let num_blocks = self.num_blocks(); + self.blocks.iter_mut().take(num_blocks) + } + + pub fn current_block_mut(&mut self) -> &mut Vec { + &mut self.blocks[self.current_block_index] + } + + pub fn current_or_block_mut(&mut self, block_index: usize) -> &mut Vec { + &mut self.blocks[block_index] + } + + pub fn get_unchecked(&self, index: BlocksIndex) -> &T { + let block = unsafe { self.block_unchecked(index.block_index()) }; + let item = unsafe { block.get_unchecked(index.index_in_block()) }; + item + } + + pub fn get_unchecked_mut(&mut self, index: BlocksIndex) -> &mut T { + let block = unsafe { self.block_unchecked_mut(index.block_index()) }; + let item = unsafe { block.get_unchecked_mut(index.index_in_block()) }; + item + } + + pub fn emit(&mut self, emit_to: BlockedEmitTo) -> Vec> { + match emit_to { + BlockedEmitTo::All => { + let counts = self.take_all(); + + counts + } + BlockedEmitTo::NextBlock => { + self.take_block().map_or(vec![], |next| vec![next]) + } + BlockedEmitTo::First(n) => { + vec![self.take_n(n)] + } + } + } + + /// Take every block that counts, see [`Self::num_blocks`] + pub fn take_all(&mut self) -> Vec> { + let num_blocks = self.num_blocks(); + let mut blocks = std::mem::take(&mut self.blocks); + blocks.truncate(num_blocks); + self.reset(); + + blocks.into() + } + + /// Take the first block, `None` once there are no more items + pub fn take_block(&mut self) -> Option> { + if self.num_blocks() == 0 { + return None; + } + Some(self.take_first_block()) + } + + /// Take the first block even when it is empty, for callers that know from + /// elsewhere that the block holds items + pub fn take_first_block(&mut self) -> Vec { + let block = self.blocks.pop_front().expect("always at least one block"); + + if self.blocks.is_empty() { + self.current_block_index = 0; + self.should_count_current_block = false; + self.blocks.push_back(vec![]); + } else { + self.current_block_index -= 1; + + // Only reduce memory if not the last one since the last block is calculated separately + self.finished_blocks_allocated_memory -= block.allocated_size(); + } + + self.len -= block.len(); + + block + } + + pub fn take_block_finished(&mut self) -> Option> + where T: ArrowNativeType + { + let block = self.take_block()?; + + let finished = ScalarBuffer::from(block); + + Some(finished) + } + + pub fn take_n( + &mut self, + n: usize, + ) -> Vec { + assert_ne!(n, 0, "n must be greater than 0"); + assert!(n <= self.len, "n ({n}) must be <= len ({})", self.len); + assert!(n < self.block_size, "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", self.block_size); + + if n == self.len { + let blocks = self.take_all(); + return blocks.into_iter().next().unwrap_or_default(); + } + + if n == self.blocks[0].len() { + return self.take_block().expect("must have at least one block since n < len and n != 0"); + } + + let prev_len = self.len; + + + debug_assert!( + self.blocks.back().is_some_and(|block| block.len() < self.block_size), + "the last block must not be full (since it should be the writable tail)" + ); + + // The emitted items are always fully contained in the first block + let mut taken = Vec::with_capacity(n); + taken.extend_from_slice(&self.blocks[0][..n]); + + { + let first = &mut self.blocks[0]; + let first_len = first.len(); + first.copy_within(n.., 0); + first.truncate(first_len - n); + } + + // Every block gives its first `n` items to the end of the previous block and + // shifts the rest down, so every finished block stays full and keeps its allocation + for i in 1..self.blocks.len() { + let mut block = std::mem::take(&mut self.blocks[i]); + let block_len = block.len(); + let moved = n.min(block_len); + + self.blocks[i - 1].extend_from_slice(&block[..moved]); + + block.copy_within(moved.., 0); + block.truncate(block_len - moved); + + self.blocks[i] = block; + } + + // The old last block is the only one that can end up empty. + // If the block before it is full it is the writable tail as is and keeps its allocation, + // otherwise the block before it is the writable tail and the empty one is not part of the layout + let tail_is_empty = self.blocks.back().is_some_and(Vec::is_empty); + let prev_is_full = self.blocks.len() >= 2 && self.blocks[self.blocks.len() - 2].len() == self.block_size; + + if tail_is_empty && !prev_is_full { + self.blocks.pop_back(); + } + + self.len = prev_len - n; + self.current_block_index = self.blocks.len() - 1; + self.should_count_current_block = !tail_is_empty || !prev_is_full; + self.finished_blocks_allocated_memory = self + .blocks + .iter() + .take(self.current_block_index) + .map(|block| block.allocated_size()) + .sum(); + + taken + } + + pub fn reset(&mut self) { + self.blocks = VecDeque::from(vec![vec![]]); + self.len = 0; + self.current_block_index = 0; + self.finished_blocks_allocated_memory = 0; + self.should_count_current_block = false; + } +} + +impl + Extend + for BlockedVec +{ + fn extend>( + &mut self, + iter: U, + ) { + let mut iter = iter.into_iter(); + + loop { + let remaining_in_current_block = self.current_block_remaining_len(); + let block_finished = + self.extend_in_block(iter.by_ref().take(remaining_in_current_block)); + + if !block_finished { + break; + } + } + } +} + +impl Index + for BlockedVec +where + T: Copy, +{ + type Output = T; + + fn index(&self, index: usize) -> &Self::Output { + self.index(BlocksIndex::from_index_in_fixed_block_size( + index, + self.block_size, + )) + } +} + +impl IndexMut + for BlockedVec +where + T: Copy, +{ + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + self.index_mut(BlocksIndex::from_index_in_fixed_block_size( + index, + self.block_size, + )) + } +} + +impl Index + for BlockedVec +where + T: Copy, +{ + type Output = T; + + fn index(&self, index: BlocksIndex) -> &Self::Output { + &self.blocks[index.block_index()] + [index.index_in_block()] + } +} + +impl IndexMut + for BlockedVec +where + T: Copy, +{ + fn index_mut(&mut self, index: BlocksIndex) -> &mut Self::Output { + &mut self.blocks[index.block_index()] + [index.index_in_block()] + } +} diff --git a/datafusion/expr-common/src/blocked_helpers/mod.rs b/datafusion/expr-common/src/blocked_helpers/mod.rs new file mode 100644 index 0000000000000..e7389dcfd32c6 --- /dev/null +++ b/datafusion/expr-common/src/blocked_helpers/mod.rs @@ -0,0 +1,8 @@ + +mod blocked_vec; +mod blocked_boolean_buffer_builder; + +pub use blocked_vec::{ + BlockedVec, +}; +pub use blocked_boolean_buffer_builder::BlockedBooleanBufferBuilder; diff --git a/datafusion/expr-common/src/lib.rs b/datafusion/expr-common/src/lib.rs index c9a95fd294503..6d9d470f39a36 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_helpers; pub mod casts; pub mod columnar_value; pub mod dyn_eq; @@ -45,5 +46,6 @@ pub mod signature; pub mod sort_properties; pub mod statistics; pub mod type_coercion; +pub mod blocked_groups_accumulator; pub use placement::ExpressionPlacement; diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs index 68865cbe1ca54..160a9faf05cf9 100644 --- a/datafusion/expr/src/function.rs +++ b/datafusion/expr/src/function.rs @@ -25,6 +25,7 @@ use std::sync::Arc; pub use datafusion_functions_aggregate_common::accumulator::{ AccumulatorArgs, AccumulatorFactoryFunction, StateFieldsArgs, + BlockedAccumulatorArgs, }; use crate::expr::{AggregateFunction, WindowFunction}; diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 649f3de2520ac..6926d49c43057 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -106,6 +106,11 @@ pub use datafusion_expr_common::columnar_value::ColumnarValue; pub use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, }; +pub use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlocksIndex, BlockedGroupSelection, BlockedGroupsAccumulator, +}; +pub use datafusion_expr_common::blocked_helpers; + pub use datafusion_expr_common::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; pub use datafusion_expr_common::signature::{ @@ -153,6 +158,7 @@ pub use udf::{ pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl}; pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits}; + #[cfg(test)] #[ctor::ctor(unsafe)] fn init() { diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index de65702ad9aa6..cb7045fced32f 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -27,6 +27,7 @@ use std::vec; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::{Result, ScalarValue, Statistics, exec_err, not_impl_err}; +use datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator; use datafusion_expr_common::dyn_eq::{DynEq, DynHash}; use datafusion_expr_common::operator::Operator; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -37,9 +38,9 @@ use crate::expr::{ schema_name_from_sorts, }; use crate::function::{ - AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs, + AccumulatorArgs, BlockedAccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs, }; -use crate::groups_accumulator::GroupsAccumulator; +use crate::groups_accumulator::{GroupsAccumulator}; use crate::udf_eq::UdfEq; use crate::utils::AggregateOrderSensitivity; use crate::utils::format_state_name; @@ -271,6 +272,19 @@ 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: BlockedAccumulatorArgs) -> 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: BlockedAccumulatorArgs, + ) -> Result> { + self.inner.create_blocked_groups_accumulator(args) + } + pub fn create_sliding_accumulator( &self, args: AccumulatorArgs, @@ -670,6 +684,32 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { not_impl_err!("GroupsAccumulator hasn't been implemented for {self:?} yet") } + /// If the aggregate expression has a specialized + /// [`BlockedGroupsAccumulator`] implementation. If this returns true, + /// `[Self::create_blocked_groups_accumulator]` will be called. + /// + /// # Notes + /// + /// Even if this function returns true, DataFusion will still use + /// [`Self::accumulator`] for certain queries, such as when this aggregate is + /// used as a window function or when there no GROUP BY columns in the + /// query. + fn blocked_groups_accumulator_supported(&self, _args: BlockedAccumulatorArgs) -> bool { + false + } + + /// Return a specialized [`BlockedGroupsAccumulator`] that manages state + /// for all groups. + /// + /// For maximum performance, a [`BlockedGroupsAccumulator`] should be + /// implemented in addition to [`Accumulator`]. + fn create_blocked_groups_accumulator( + &self, + _args: BlockedAccumulatorArgs, + ) -> 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. /// diff --git a/datafusion/functions-aggregate-common/src/accumulator.rs b/datafusion/functions-aggregate-common/src/accumulator.rs index 8db0ab4133dc0..9f3eb58589fe6 100644 --- a/datafusion/functions-aggregate-common/src/accumulator.rs +++ b/datafusion/functions-aggregate-common/src/accumulator.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::ops::{Deref, DerefMut}; use arrow::datatypes::{DataType, FieldRef, Schema}; use datafusion_common::Result; use datafusion_expr_common::accumulator::Accumulator; @@ -80,6 +81,38 @@ impl AccumulatorArgs<'_> { } } +/// [`AccumulatorArgs`] contains information about how an aggregate +/// function was called, including the types of its arguments and any optional +/// ordering expressions. +#[derive(Debug, Clone)] +pub struct BlockedAccumulatorArgs<'a> { + pub accumulator_args: AccumulatorArgs<'a>, + + pub batch_size: usize, +} +impl<'a> Deref for BlockedAccumulatorArgs<'a> { + type Target = AccumulatorArgs<'a>; + + fn deref(&self) -> &Self::Target { + &self.accumulator_args + } +} +impl DerefMut for BlockedAccumulatorArgs<'_> { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.accumulator_args + } +} + +impl<'a> BlockedAccumulatorArgs<'a> { + /// Returns the return type of the aggregate function. + pub fn from_accumulator_args(accumulator_args: AccumulatorArgs<'a>, batch_size: usize) -> Self { + Self { + accumulator_args, + batch_size + } + } +} + /// Factory that returns an accumulator for the given aggregate function. pub type AccumulatorFactoryFunction = Arc Result> + Send + Sync>; diff --git a/datafusion/functions-aggregate-common/src/aggregate.rs b/datafusion/functions-aggregate-common/src/aggregate.rs index aadce907e7cc3..e22bd1d524f9b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate.rs @@ -19,3 +19,4 @@ pub mod avg_distinct; pub mod count_distinct; pub mod groups_accumulator; pub mod sum_distinct; +pub mod blocked_groups_accumulator; diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs new file mode 100644 index 0000000000000..92446c068d2b8 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs @@ -0,0 +1,208 @@ +// 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. + +//! Utilities for implementing GroupsAccumulator +//! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] + +pub mod blocked_prim_op; +pub mod accumulate; + +use arrow::{ + array::{ArrayRef, BooleanArray}, +}; +use datafusion_common::{Result}; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex}; +use datafusion_expr_common::groups_accumulator::{GroupSelection, GroupsAccumulator}; + +/// An adapter that implements [`GroupsAccumulator`] for any [`Accumulator`] +/// +/// While [`Accumulator`] are simpler to implement and can support +/// more general calculations (like retractable window functions), +/// they are not as fast as a specialized `GroupsAccumulator`. This +/// interface bridges the gap so the group by operator only operates +/// in terms of [`Accumulator`]. +/// +/// Internally, this adapter creates a new [`Accumulator`] for each group which +/// stores the state for that group. This both requires an allocation for each +/// Accumulator, internal indices, as well as whatever internal allocations the +/// Accumulator itself requires. +/// +/// For example, a `MinAccumulator` that computes the minimum string value with +/// a [`ScalarValue::Utf8`]. That will require at least two allocations per group +/// (one for the `MinAccumulator` and one for the `ScalarValue::Utf8`). +/// +/// ```text +/// ┌─────────────────────────────────┐ +/// │MinAccumulator { │ +/// ┌─────▶│ min: ScalarValue::Utf8("A") │───────┐ +/// │ │} │ │ +/// │ └─────────────────────────────────┘ └───────▶ "A" +/// ┌─────┐ │ ┌─────────────────────────────────┐ +/// │ 0 │─────┘ │MinAccumulator { │ +/// ├─────┤ ┌─────▶│ min: ScalarValue::Utf8("Z") │───────────────▶ "Z" +/// │ 1 │─────┘ │} │ +/// └─────┘ └─────────────────────────────────┘ ... +/// ... ... +/// ┌─────┐ ┌────────────────────────────────┐ +/// │ N-2 │ │MinAccumulator { │ +/// ├─────┤ │ min: ScalarValue::Utf8("A") │────────────────▶ "A" +/// │ N-1 │─────┐ │} │ +/// └─────┘ │ └────────────────────────────────┘ +/// │ ┌────────────────────────────────┐ ┌───────▶ "Q" +/// │ │MinAccumulator { │ │ +/// └─────▶│ min: ScalarValue::Utf8("Q") │────────┘ +/// │} │ +/// └────────────────────────────────┘ +/// +/// +/// Logical group Current Min/Max value for that group stored +/// number as a ScalarValue which points to an +/// individually allocated String +/// ``` +/// +/// # Optimizations +/// +/// The adapter minimizes the number of calls to [`Accumulator::update_batch`] +/// by first collecting the input rows for each group into a contiguous array +/// using [`compute::take`] +pub struct BlockedGroupsAccumulatorAdapter { + inner: Box, + batch_size: usize, + number_of_groups: usize, +} + +impl BlockedGroupsAccumulatorAdapter { + /// Create a new adapter that will create a new [`Accumulator`] + /// for each group, using the specified factory function + pub fn new(inner: Box, batch_size: usize) -> Self { + Self { + batch_size, + inner, + number_of_groups: 0, + } + } + + fn update_number_of_groups_after_emit(&mut self, emit_to: BlockedEmitTo) { + match emit_to { + BlockedEmitTo::All => { + self.number_of_groups = 0; + } + BlockedEmitTo::NextBlock => { + self.number_of_groups = self.number_of_groups.saturating_sub(self.batch_size); + } + BlockedEmitTo::First(n) => { + self.number_of_groups = self.number_of_groups.saturating_sub(n); + } + } + } +} + +impl BlockedGroupsAccumulator for BlockedGroupsAccumulatorAdapter { + fn batch_size(&self) -> usize { + self.batch_size + } + + fn update_batch(&mut self, values: &[ArrayRef], group_indices: &[BlocksIndex], opt_filter: Option<&BooleanArray>, total_num_groups: usize) -> Result<()> { + let group_indices_flatten = group_indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + self.number_of_groups = total_num_groups; + self.inner.update_batch(values, &group_indices_flatten, opt_filter, total_num_groups) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + let iter = emit_to.into_emit_to(self.number_of_groups, self.batch_size)?; + + let mut output = vec![]; + for mapped_emit_to in iter { + let item = self.inner.evaluate(mapped_emit_to)?; + output.push(item); + } + + self.update_number_of_groups_after_emit(emit_to); + + Ok(output) + } + + fn supports_evaluate_preserving(&self) -> bool { + self.inner.supports_evaluate_preserving() + } + + fn evaluate_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result { + match selection.indices() { + None => { + let selection = GroupSelection::all(selection.total_num_groups()); + + self.inner.evaluate_preserving(selection) + } + Some(indices) => { + let indices_flatten = indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + + let selection = GroupSelection::try_from_indices(&indices_flatten, selection.total_num_groups())?; + + self.inner.evaluate_preserving(selection) + } + } + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + let iter = emit_to.into_emit_to(self.number_of_groups, self.batch_size)?; + + let mut output = vec![]; + for mapped_emit_to in iter { + let item = self.inner.state(mapped_emit_to)?; + output.push(item); + } + + self.update_number_of_groups_after_emit(emit_to); + + Ok(output) + } + + fn supports_state_preserving(&self) -> bool { + self.inner.supports_state_preserving() + } + + fn state_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result> { + match selection.indices() { + None => { + let selection = GroupSelection::all(selection.total_num_groups()); + + self.inner.state_preserving(selection) + } + Some(indices) => { + let indices_flatten = indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + + let selection = GroupSelection::try_from_indices(&indices_flatten, selection.total_num_groups())?; + + self.inner.state_preserving(selection) + } + } + } + + fn merge_batch(&mut self, values: &[ArrayRef], group_indices: &[BlocksIndex], total_num_groups: usize) -> Result<()> { + let group_indices_flatten = group_indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + self.number_of_groups = total_num_groups; + self.inner.merge_batch(values, &group_indices_flatten, total_num_groups) + } + + fn convert_to_state(&self, values: &[ArrayRef], opt_filter: Option<&BooleanArray>) -> Result> { + self.inner.convert_to_state(values, opt_filter) + } + + fn size(&self) -> usize { + self.inner.size() + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs new file mode 100644 index 0000000000000..5b5dbd1026381 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs @@ -0,0 +1,391 @@ +// 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. + +//! [`GroupsAccumulator`] helpers: [`BlockedNullState`] and [`accumulate_indices`] +//! +//! [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator + +use arrow::array::{Array, BooleanArray, BooleanBufferBuilder, PrimitiveArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::ArrowPrimitiveType; + +use datafusion_common::Result; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlocksIndex}; +use datafusion_expr_common::blocked_helpers::{BlockedBooleanBufferBuilder, BlockedVec}; +use crate::aggregate::groups_accumulator::accumulate::accumulate; + +/// If the input has nulls, then the accumulator must potentially +/// handle each input null value specially (e.g. for `SUM` to mark the +/// corresponding sum as null) +/// +/// `NullState` tracks if it has seen *any* value for each group when filters or +/// sparse group indices may omit input for a registered group. +#[derive(Debug)] +pub enum BlockedSeenValues { + /// All groups seen so far have seen at least one non-null value + All { + num_values: usize, + block_size: usize, + }, + // Some groups have not yet seen a non-null value + Some { + values: BlockedBooleanBufferBuilder, + }, +} + +impl BlockedSeenValues { + + pub fn new(block_size: usize) -> Self { + BlockedSeenValues::All { num_values: 0, block_size } + } + + pub fn block_size(&self) -> usize { + match self { + BlockedSeenValues::All { block_size, .. } => *block_size, + BlockedSeenValues::Some { values } => values.block_size(), + } + } + + /// Return a mutable reference to the `BooleanBufferBuilder` in `BlockedSeenValues::Some`. + /// + /// If `self` is `BlockedSeenValues::All`, it is transitioned to `BlockedSeenValues::Some` + /// by creating a new `BooleanBufferBuilder` where the first `num_values` are true. + /// + /// The builder is then ensured to have at least `total_num_groups` length, + /// with any new entries initialized to false. + fn get_builder(&mut self, total_num_groups: usize) -> &mut BlockedBooleanBufferBuilder { + match self { + BlockedSeenValues::All { num_values, block_size } => { + let mut builder = BlockedBooleanBufferBuilder::new(*block_size); + builder.push_value_n(true, *num_values); + if total_num_groups > *num_values { + builder.push_value_n(false, total_num_groups - *num_values); + } + *self = BlockedSeenValues::Some { values: builder }; + match self { + BlockedSeenValues::Some { values } => values, + _ => unreachable!(), + } + } + BlockedSeenValues::Some { values } => { + if values.len() < total_num_groups { + values.push_value_n(false, total_num_groups - values.len()); + } + values + } + } + } +} + +/// Returns true when all newly registered groups are present in `group_indices`. +/// +/// Group indices are assigned in first-seen order, so an unfiltered batch visits +/// new groups in ascending order. Pre-filtered input can omit a new group, making +/// the indices sparse even though the accumulator no longer receives a filter. +fn new_groups_are_dense( + group_indices: &[BlocksIndex], + block_size: usize, + first_new_group: usize, + total_num_groups: usize, +) -> bool { + let first_new_group = BlocksIndex::from_index_in_fixed_block_size(first_new_group, block_size); + let total_num_groups = BlocksIndex::from_index_in_fixed_block_size(total_num_groups, block_size); + + if first_new_group == total_num_groups { + return true; + } + + let mut next_new_group = first_new_group; + for &group_index in group_indices { + if group_index == next_new_group { + next_new_group = next_new_group.increment(block_size); + } else if group_index > next_new_group { + return false; + } + } + next_new_group == total_num_groups +} + +/// Track the accumulator null state per row: if any values for that +/// group were null and if any values have been seen at all for that group. +/// +/// This is part of the inner loop for many [`GroupsAccumulator`]s, +/// and thus the performance is critical and so there are multiple +/// specialized implementations, invoked depending on the specific +/// combinations of the input. +/// +/// Typically there are 4 potential combinations of inputs must be +/// special cased for performance: +/// +/// * With / Without filter +/// * With / Without nulls in the input +/// +/// If the input has nulls, then the accumulator must potentially +/// handle each input null value specially (e.g. for `SUM` to mark the +/// corresponding sum as null) +/// +/// `NullState` tracks if it has seen *any* value for each group when filters or +/// sparse group indices may omit input for a registered group. +/// +/// [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator +#[derive(Debug)] +pub struct BlockedNullState { + /// Have we seen any non-filtered input values for `group_index`? + /// + /// If `seen_values` is `SeenValues::Some(buffer)` and buffer\[i\] is true, have seen at least one non null + /// value for group `i` + /// + /// If `seen_values` is `SeenValues::Some(buffer)` and buffer\[i\] is false, have not seen any values that + /// pass the filter yet for group `i` + /// + /// If `seen_values` is `SeenValues::All`, all groups have seen at least one non null value + seen_values: BlockedSeenValues, +} + +impl BlockedNullState { + pub fn new(block_size: usize) -> Self { + Self { + seen_values: BlockedSeenValues::All { num_values: 0, block_size }, + } + } + + /// return the size of all buffers allocated by this null state, not including self + pub fn size(&self) -> usize { + match &self.seen_values { + BlockedSeenValues::All { .. } => 0, + BlockedSeenValues::Some { values } => { + values.allocated_size() + }, + } + } + + /// Invokes `value_fn(group_index, value)` for each non null, non + /// filtered value of `value`, while tracking which groups have + /// seen null inputs and which groups have seen any inputs if necessary + // + /// # Arguments: + /// + /// * `values`: the input arguments to the accumulator + /// * `group_indices`: To which groups do the rows in `values` belong, (aka group_index) + /// * `opt_filter`: if present, only rows for which is Some(true) are included + /// * `value_fn`: function invoked for (group_index, value) where value is non null + /// + /// See [`accumulate`], for more details on how value_fn is called + /// + /// When value_fn is called it also sets + /// + /// 1. `self.seen_values[group_index]` to true for all rows that had a non null value + pub fn accumulate( + &mut self, + group_indices: &[BlocksIndex], + values: &PrimitiveArray, + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + mut value_fn: F, + ) where + T: ArrowPrimitiveType + Send, + F: FnMut(BlocksIndex, T::Native) + Send, + { + // Skip per-value null handling when every input value is valid and all + // newly registered groups are represented. Pre-filtered inputs can have + // sparse group indices despite not passing a filter to the accumulator. + if opt_filter.is_none() + && values.null_count() == 0 + && let BlockedSeenValues::All { num_values, block_size } = &mut self.seen_values + && new_groups_are_dense(group_indices, *block_size, *num_values, total_num_groups) + { + accumulate(group_indices, values, None, value_fn); + *num_values = total_num_groups; + return; + } + + let seen_values = self.seen_values.get_builder(total_num_groups); + accumulate(group_indices, values, opt_filter, |group_index, value| { + seen_values.set_bit(group_index, true); + value_fn(group_index, value); + }); + } + + /// Invokes `value_fn(group_index, value)` for each non null, non + /// filtered value in `values`, while tracking which groups have + /// seen null inputs and which groups have seen any inputs, for + /// [`BooleanArray`]s. + /// + /// Since `BooleanArray` is not a [`PrimitiveArray`] it must be + /// handled specially. + /// + /// See [`Self::accumulate`], which handles `PrimitiveArray`s, for + /// more details on other arguments. + pub fn accumulate_boolean( + &mut self, + group_indices: &[BlocksIndex], + values: &BooleanArray, + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + mut value_fn: F, + ) where + F: FnMut(BlocksIndex, bool) + Send, + { + let data = values.values(); + assert_eq!(data.len(), group_indices.len()); + + // Skip per-value null handling when every input value is valid and all + // newly registered groups are represented. Pre-filtered inputs can have + // sparse group indices despite not passing a filter to the accumulator. + if opt_filter.is_none() + && values.null_count() == 0 + && let BlockedSeenValues::All { num_values, block_size } = &mut self.seen_values + && new_groups_are_dense(group_indices, *block_size, *num_values, total_num_groups) + { + group_indices + .iter() + .zip(data.iter()) + .for_each(|(&group_index, new_value)| value_fn(group_index, new_value)); + *num_values = total_num_groups; + + return; + } + + let seen_values = self.seen_values.get_builder(total_num_groups); + + // These could be made more performant by iterating in chunks of 64 bits at a time + match (values.null_count() > 0, opt_filter) { + // no nulls, no filter, + (false, None) => { + // if we have previously seen nulls, ensure the null + // buffer is big enough (start everything at valid) + group_indices.iter().zip(data.iter()).for_each( + |(&group_index, new_value)| { + seen_values.set_bit(group_index, true); + value_fn(group_index, new_value) + }, + ) + } + // nulls, no filter + (true, None) => { + let nulls = values.nulls().unwrap(); + group_indices + .iter() + .zip(data.iter()) + .zip(nulls.iter()) + .for_each(|((&group_index, new_value), is_valid)| { + if is_valid { + seen_values.set_bit(group_index, true); + value_fn(group_index, new_value); + } + }) + } + // no nulls, but a filter + (false, Some(filter)) => { + assert_eq!(filter.len(), group_indices.len()); + + group_indices + .iter() + .zip(data.iter()) + .zip(filter.iter()) + .for_each(|((&group_index, new_value), filter_value)| { + if filter_value == Some(true) { + seen_values.set_bit(group_index, true); + value_fn(group_index, new_value); + } + }) + } + // both null values and filters + (true, Some(filter)) => { + assert_eq!(filter.len(), group_indices.len()); + filter + .iter() + .zip(group_indices.iter()) + .zip(values.iter()) + .for_each(|((filter_value, &group_index), new_value)| { + if filter_value == Some(true) + && let Some(new_value) = new_value + { + seen_values.set_bit(group_index, true); + value_fn(group_index, new_value) + } + }) + } + } + } + + /// Creates a [`NullBuffer`] for `selection` without changing this state. + pub fn build_preserving( + &self, + selection: BlockedGroupSelection<'_>, + ) -> Result> { + let selected_len = selection.len(); + match &self.seen_values { + BlockedSeenValues::All { num_values, block_size: _ } => { + selection.validate_num_groups(*num_values)?; + Ok(None) + } + BlockedSeenValues::Some { values } => { + selection.validate_num_groups(values.len())?; + let mut selected = BooleanBufferBuilder::new(selected_len); + for index in selection.iter() { + selected.append(values.get_bit(index)); + } + Ok(Some(NullBuffer::new(selected.finish()))) + } + } + } + + /// Creates the a [`NullBuffer`] representing which group_indices + /// should have null values (because they never saw any values) + /// for the `emit_to` rows. + /// + /// resets the internal state appropriately + pub fn build(&mut self, emit_to: BlockedEmitTo) -> Vec> { + let block_size = self.seen_values.block_size(); + let old_seen = std::mem::replace(&mut self.seen_values, BlockedSeenValues::new(block_size)); + + match old_seen { + BlockedSeenValues::All { num_values, block_size } => { + match emit_to { + BlockedEmitTo::All => vec![None; num_values.div_ceil(block_size)], + BlockedEmitTo::NextBlock => { + self.seen_values = BlockedSeenValues::All { num_values: num_values.saturating_sub(block_size), block_size }; + + if num_values == 0 { + vec![] + } else { + vec![None] + } + } + BlockedEmitTo::First(n) => { + assert!(n < block_size, "emit_to First(n) must be less than block_size"); + self.seen_values = BlockedSeenValues::All { num_values: num_values.saturating_sub(n), block_size }; + if num_values == 0 { + vec![] + } else { + vec![None] + } + } + } + } + BlockedSeenValues::Some { mut values } => { + let emitted = values.emit(emit_to); + + self.seen_values = BlockedSeenValues::Some { values }; + + emitted.into_iter().map(|block| Some(NullBuffer::new(block))).collect() + } + } + } +} + diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs new file mode 100644 index 0000000000000..1bd88415287a7 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs @@ -0,0 +1,239 @@ +// 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, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::NullBuffer; +use arrow::compute; +use arrow::datatypes::ArrowPrimitiveType; +use arrow::datatypes::DataType; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex}; +use datafusion_expr_common::blocked_helpers::BlockedVec; +use datafusion_expr_common::groups_accumulator::{ + EmitTo, GroupSelection, GroupsAccumulator, +}; + +use super::accumulate::BlockedNullState; + +/// An accumulator that implements a single operation over +/// [`ArrowPrimitiveType`] where the accumulated 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). +/// +/// [`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 + } +} + +impl BlockedGroupsAccumulator for BlockedPrimitiveGroupsAccumulator +where + T: ArrowPrimitiveType + Send, + F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, +{ + fn batch_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::(); + + // update values + self.values.push_value_n_to_len(self.starting_value, total_num_groups); + + // NullState dispatches / handles tracking nulls and groups that saw no values + self.null_state.accumulate( + group_indices, + values, + opt_filter, + total_num_groups, + |group_index, new_value| { + // SAFETY: group_index is guaranteed to be in bounds + let value = unsafe { self.values.get_unchecked_mut(group_index) }; + (self.prim_fn)(value, new_value); + }, + ); + + Ok(()) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + let values = self.values.emit(emit_to); + let nulls = self.null_state.build(emit_to); + + let blocks = values.into_iter().zip(nulls).map(|(values, nulls)| { + let values = PrimitiveArray::::new(values.into(), nulls) // no copy + .with_data_type(self.data_type.clone()); + Arc::new(values) as ArrayRef + }).collect::>(); + + Ok(blocks) + } + + fn evaluate_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result { + selection.validate_num_groups(self.values.len())?; + let values: Vec = + selection.iter().map(|index| self.values[index]).collect(); + let nulls = self.null_state.build_preserving(selection)?; + let values = PrimitiveArray::::new(values.into(), nulls) + .with_data_type(self.data_type.clone()); + Ok(Arc::new(values)) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + let blocks = self.evaluate(emit_to)?; + + Ok(blocks.into_iter() + // Convert each block into a Vec to match the return type + .map(|arr| vec![arr]) + .collect::>()) + } + + fn state_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|arr| vec![arr]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + + 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) + } + + /// Converts an input batch directly to a state batch + /// + /// The state is: + /// - self.prim_fn for all non null, non filtered values + /// - null otherwise + fn convert_to_state( + &self, + 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)]) + } + + fn size(&self) -> usize { + self.values.allocated_size() + self.null_state.size() + } +} 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..130711554dd94 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -418,14 +418,14 @@ impl NullState { /// value_fn(0, 200) /// value_fn(0, 300) /// ``` -pub fn accumulate( - group_indices: &[usize], +pub fn accumulate( + group_indices: &[I], values: &PrimitiveArray, opt_filter: Option<&BooleanArray>, mut value_fn: F, ) where T: ArrowPrimitiveType + Send, - F: FnMut(usize, T::Native) + Send, + F: FnMut(I, T::Native) + Send, { let data: &[T::Native] = values.values(); assert_eq!(data.len(), group_indices.len()); @@ -588,13 +588,13 @@ pub fn accumulate_multiple( /// /// See [`NullState::accumulate`], for more details on other /// arguments. -pub fn accumulate_indices( - group_indices: &[usize], +pub fn accumulate_indices( + group_indices: &[I], nulls: Option<&NullBuffer>, opt_filter: Option<&BooleanArray>, mut index_fn: F, ) where - F: FnMut(usize) + Send, + F: FnMut(I) + Send, { match (nulls, opt_filter) { (None, None) => { diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index ed2abe88644ae..c13bd7de2aaf3 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -34,14 +34,7 @@ use datafusion_common::{ HashMap, Result, ScalarValue, downcast_value, exec_err, internal_err, not_impl_err, stats::Precision, utils::expr::COUNT_STAR_EXPANSION, }; -use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupSelection, - GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, - TypeSignature, Volatility, WindowFunctionDefinition, - expr::WindowFunction, - function::{AccumulatorArgs, StateFieldsArgs}, - utils::{AggregateOrderSensitivity, format_state_name}, -}; +use datafusion_expr::{expr::WindowFunction, function::{AccumulatorArgs, StateFieldsArgs}, utils::{AggregateOrderSensitivity, format_state_name}, Accumulator, AggregateUDFImpl, BlockedGroupsAccumulator, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility, WindowFunctionDefinition, BlocksIndex, BlockedEmitTo, BlockedGroupSelection}; use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::{ count_distinct::Bitmap65536DistinctCountAccumulator, @@ -65,6 +58,7 @@ use std::{ ops::BitAnd, sync::Arc, }; +use datafusion_expr::blocked_helpers::{BlockedVec, CopyItemBlockedVecBuilder}; make_udaf_expr_and_func!( Count, @@ -731,9 +725,9 @@ impl GroupsAccumulator for CountGroupsAccumulator { fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { selection.validate_num_groups(self.counts.len())?; let counts = selection - .iter() - .map(|index| self.counts[index]) - .collect::>(); + .iter() + .map(|index| self.counts[index]) + .collect::>(); Ok(Arc::new(Int64Array::from(counts))) } @@ -824,6 +818,209 @@ impl GroupsAccumulator for CountGroupsAccumulator { } } +/// An accumulator to compute the counts of [`PrimitiveArray`]. +/// Stores values as native types, and does overflow checking +/// +/// Unlike most other accumulators, COUNT never produces NULLs. If no +/// non-null values are seen in any group the output is 0. Thus, this +/// accumulator has no additional null or seen filter tracking. +#[derive(Debug)] +struct CountBlockedGroupsAccumulator { + /// Count per group. + /// + /// Note this is an i64 and not a u64 (or usize) because the + /// output type of count is `DataType::Int64`. Thus by using `i64` + /// for the counts, the output [`Int64Array`] can be created + /// without copy. + counts: BlockedVec, +} + +impl CountBlockedGroupsAccumulator { + pub fn new(block_size: usize) -> Self { + Self { counts: BlockedVec::new(block_size) } + } +} + +impl BlockedGroupsAccumulator for CountBlockedGroupsAccumulator { + + fn batch_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]; + + self.counts.push_value_n_to_len(0, total_num_groups); + + accumulate_indices( + group_indices, + values.logical_nulls().as_ref(), + opt_filter, + |group_index| { + // SAFETY: group_index is guaranteed to be in bounds + let count = unsafe { self.counts.get_unchecked_mut(group_index) }; + *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"); + // first batch is counts, second is partial sums + 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.push_value_n_to_len(0, total_num_groups); + group_indices.iter().zip(partial_counts.iter()).for_each( + |(&group_index, partial_count)| { + self.counts[group_index] += partial_count; + }, + ); + + Ok(()) + } + + fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { + let blocks = self.counts.emit(emit_to); + + let blocks_ready = blocks + .into_iter() + .map(|block| { + // Count is always non null (null inputs just don't contribute to the overall values) + let array = PrimitiveArray::::new(block.into(), None); + Arc::new(array) as ArrayRef + }) + .collect::>(); + + Ok(blocks_ready) + } + + fn evaluate_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result { + selection.validate_num_groups(self.counts.len())?; + let counts = selection + .iter() + .map(|index| self.counts[index]) + .collect::>(); + Ok(Arc::new(Int64Array::from(counts))) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + + // return arrays for counts + fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { + let blocks = self.counts.emit(emit_to); + + let blocks_ready = blocks + .into_iter() + .map(|block| { + // Count is always non null (null inputs just don't contribute to the overall values) + let array = PrimitiveArray::::new(block.into(), None); + + // Each state only have 1 column + vec![Arc::new(array) as ArrayRef] + }) + .collect::>(); + + Ok(blocks_ready) + } + + fn state_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|array| vec![array]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + + /// Converts an input batch directly to a state batch + /// + /// The state of `COUNT` is always a single Int64Array: + /// * `1` (for non-null, non filtered values) + /// * `0` (for null values) + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + let values = &values[0]; + + let state_array = match (values.logical_nulls(), opt_filter) { + (None, None) => { + // In case there is no nulls in input and no filter, returning array of 1 + Arc::new(Int64Array::from_value(1, values.len())) + } + (Some(nulls), None) => { + // If there are any nulls in input values -- casting `nulls` (true for values, false for nulls) + // of input array to Int64 + let nulls = BooleanArray::new(nulls.into_inner(), None); + compute::cast(&nulls, &DataType::Int64)? + } + (None, Some(filter)) => { + // If there is only filter + // - applying filter null mask to filter values by bitand filter values and nulls buffers + // (using buffers guarantees absence of nulls in result) + // - casting result of bitand to Int64 array + let (filter_values, filter_nulls) = filter.clone().into_parts(); + + let state_buf = match filter_nulls { + Some(filter_nulls) => &filter_values & filter_nulls.inner(), + None => filter_values, + }; + + let boolean_state = BooleanArray::new(state_buf, None); + compute::cast(&boolean_state, &DataType::Int64)? + } + (Some(nulls), Some(filter)) => { + // For both input nulls and filter + // - applying filter null mask to filter values by bitand filter values and nulls buffers + // (using buffers guarantees absence of nulls in result) + // - applying values null mask to filter buffer by another bitand on filter result and + // nulls from input values + // - casting result to Int64 array + let (filter_values, filter_nulls) = filter.clone().into_parts(); + + let filter_buf = match filter_nulls { + Some(filter_nulls) => &filter_values & filter_nulls.inner(), + None => filter_values, + }; + let state_buf = &filter_buf & nulls.inner(); + + let boolean_state = BooleanArray::new(state_buf, None); + compute::cast(&boolean_state, &DataType::Int64)? + } + }; + + Ok(vec![state_array]) + } + + 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 { diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs index 1228b458a82b9..0f49bfc35655a 100644 --- a/datafusion/functions-aggregate/src/sum.rs +++ b/datafusion/functions-aggregate/src/sum.rs @@ -39,16 +39,14 @@ use datafusion_expr::expr::AggregateFunction; 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, - Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, - TypeSignatureClass, Volatility, -}; +use datafusion_expr::{Accumulator, AggregateUDFImpl, BlockedGroupsAccumulator, Coercion, Documentation, Expr, GroupsAccumulator, Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; +use datafusion_functions_aggregate_common::aggregate::blocked_groups_accumulator::blocked_prim_op::BlockedPrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator; use datafusion_macros::user_doc; use datafusion_physical_expr::expressions::{CastExpr, Column}; use std::mem::{size_of, size_of_val}; +use datafusion_functions_aggregate_common::accumulator::BlockedAccumulatorArgs; make_udaf_expr_and_func!( Sum, @@ -317,6 +315,26 @@ impl AggregateUDFImpl for Sum { downcast_sum!(args, helper) } + fn blocked_groups_accumulator_supported(&self, args: BlockedAccumulatorArgs) -> bool { + !args.is_distinct + } + + fn create_blocked_groups_accumulator( + &self, + args: BlockedAccumulatorArgs, + ) -> Result> { + macro_rules! helper { + ($t:ty, $dt:expr) => { + Ok(Box::new(BlockedPrimitiveGroupsAccumulator::<$t, _>::new( + &$dt, + args.batch_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 df22bc69d8706..2f1c4a22133d2 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -52,13 +52,12 @@ use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, NullTreatment, physical_name, }; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; -use datafusion_expr::{AggregateUDF, Expr, ReversedUDAF, SetMonotonicity}; +use datafusion_expr::{AggregateUDF, AggregateUDFImpl, Expr, ReversedUDAF, SetMonotonicity}; use datafusion_expr_common::accumulator::Accumulator; -use datafusion_expr_common::groups_accumulator::GroupsAccumulator; +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::{ - AccumulatorArgs, StateFieldsArgs, -}; +use datafusion_functions_aggregate_common::accumulator::{AccumulatorArgs, BlockedAccumulatorArgs, StateFieldsArgs}; use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -139,6 +138,7 @@ pub struct AggregateExprBuilder { is_distinct: bool, /// Whether the expression is reversed is_reversed: bool, + batch_size: Option, } impl AggregateExprBuilder { @@ -155,6 +155,7 @@ impl AggregateExprBuilder { ignore_nulls: false, is_distinct: false, is_reversed: false, + batch_size: None, } } @@ -261,6 +262,7 @@ impl AggregateExprBuilder { ignore_nulls, is_distinct, is_reversed, + batch_size, } = self; assert_or_internal_err!(!args.is_empty(), "args should not be empty"); @@ -316,6 +318,7 @@ impl AggregateExprBuilder { input_fields: input_exprs_fields, is_reversed, is_nullable, + batch_size }) } @@ -365,6 +368,11 @@ impl AggregateExprBuilder { self } + pub fn with_batch_size(mut self, batch_size: Option) -> Self { + self.batch_size = batch_size; + self + } + pub fn distinct(mut self) -> Self { self.is_distinct = true; self @@ -553,7 +561,8 @@ impl<'a> LoweredAggregateBuilder<'a> { .alias(name) .output_metadata(output_metadata) .with_ignore_nulls(ignore_nulls) - .with_distinct(*distinct); + .with_distinct(*distinct) + .with_batch_size(execution_props.config_options.as_ref().map(|x| x.execution.batch_size.get())); if let Some(human_display) = human_display { builder = builder.human_display(human_display.expression); @@ -661,6 +670,7 @@ pub struct AggregateFunctionExpr { is_reversed: bool, input_fields: Vec, is_nullable: bool, + batch_size: Option, } impl AggregateFunctionExpr { @@ -719,6 +729,10 @@ impl AggregateFunctionExpr { self.is_nullable } + pub fn batch_size(&self) -> Option { + self.batch_size + } + /// the field of the final result of this aggregation. pub fn field(&self) -> FieldRef { self.return_field @@ -812,7 +826,8 @@ impl AggregateFunctionExpr { .output_metadata(self.return_field_metadata()) .with_ignore_nulls(self.ignore_nulls) .with_distinct(self.is_distinct) - .with_reversed(self.is_reversed); + .with_reversed(self.is_reversed) + .with_batch_size(self.batch_size); if let Some(human_display) = self.human_display() { builder = builder.human_display(human_display); } @@ -928,6 +943,52 @@ impl AggregateFunctionExpr { self.fun.create_groups_accumulator(args) } + /// If the aggregate expression has a specialized + /// [`GroupsAccumulator`] implementation. If this returns true, + /// `[Self::create_groups_accumulator`] will be called. + pub fn blocked_groups_accumulator_supported(&self) -> bool { + let Some(batch_size) = self.batch_size else { + return false + }; + 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, + }; + let blocked = BlockedAccumulatorArgs::from_accumulator_args(args, batch_size); + self.fun.blocked_groups_accumulator_supported(blocked) + } + + /// Return a specialized [`GroupsAccumulator`] that manages state + /// for all groups. + /// + /// For maximum performance, a [`GroupsAccumulator`] should be + /// implemented in addition to [`Accumulator`]. + pub fn create_blocked_groups_accumulator(&self) -> Result> { + let Some(batch_size) = self.batch_size else { + return internal_err!("batch size must be set if blocked group accumulator is supported"); + }; + 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, + }; + let blocked = BlockedAccumulatorArgs::from_accumulator_args(args, batch_size); + self.fun.create_blocked_groups_accumulator(blocked) + } + /// 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, @@ -997,7 +1058,8 @@ impl AggregateFunctionExpr { .output_metadata(self.return_field_metadata()) .with_ignore_nulls(self.ignore_nulls) .with_distinct(self.is_distinct) - .with_reversed(!self.is_reversed); + .with_reversed(!self.is_reversed) + .with_batch_size(self.batch_size); if let Some(human_display) = human_display { builder = builder.human_display(human_display.expression); if let Some(alias) = human_display.alias { @@ -1067,6 +1129,7 @@ impl AggregateFunctionExpr { is_reversed: false, input_fields: self.input_fields.clone(), is_nullable: self.is_nullable, + batch_size: self.batch_size, }) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 5f101950870d2..916eccd0ada9e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -29,14 +29,12 @@ use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; -use crate::aggregates::group_values::{ - AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, GroupValues, new_group_values, -}; -use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::group_values::{new_blocked_group_values, new_group_values, AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, BlockedGroupValues}; +use crate::aggregates::grouped_hash_stream::{create_blocked_group_accumulator, create_group_accumulator}; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, @@ -138,7 +136,7 @@ impl AggregateHashTable { .zip(aggregate_arguments) .zip(filters) .map(|((agg_expr, arguments), filter)| { - let accumulator = create_group_accumulator(agg_expr)?; + let accumulator = create_blocked_group_accumulator(agg_expr, batch_size)?; Ok(HashAggregateAccumulator::new( Arc::clone(agg_expr), arguments, @@ -149,7 +147,7 @@ impl AggregateHashTable { .collect::>()?; let group_schema = agg.group_by.group_schema(&input_schema)?; - let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let group_values = new_blocked_group_values(group_schema, &GroupOrdering::None, batch_size)?; let metrics = AggregateTableMetrics::new(agg, partition); @@ -262,51 +260,60 @@ impl AggregateHashTable { accumulator_phase: AccumulatorPhase, ) -> Result> { let output_schema = Arc::clone(&self.output_schema); - let batch_size = self.batch_size; let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); let mut output = - match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { - AggregateHashTableState::Outputting(mut state) => { - if state.group_values.is_empty() { - return Ok(None); - } - - // Accumulator output consumes internal state. Materialize all - // groups once, then slice the materialized batch on later polls. - let emit_to = EmitTo::All; - let columns = self.group_by_metrics.time_emitting(|| { - let mut columns = state.group_values.emit(emit_to)?; - for (idx, acc) in state.accumulators.iter_mut().enumerate() { - columns.extend(accumulator_metrics.time( - idx, - accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), - )?); - } - Ok::<_, datafusion_common::DataFusionError>(columns) - })?; - - let batch = RecordBatch::try_new(output_schema, columns)?; - debug_assert!(batch.num_rows() > 0); - MaterializedAggregateOutput::new(batch) - } - AggregateHashTableState::OutputtingMaterialized(output) => output, - AggregateHashTableState::Done => return Ok(None), - AggregateHashTableState::Building(_) => { - return internal_err!( + match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { + AggregateHashTableState::Outputting(mut state) => { + if state.group_values.is_empty() { + return Ok(None); + } + + // Accumulator output consumes internal state. Materialize all + // groups once, then slice the materialized batch on later polls. + let blocked_columns = self.group_by_metrics.time_emitting(|| { + let mut blocked_columns = state.group_values.emit(BlockedEmitTo::All)?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + let acc_cols = accumulator_metrics.time( + idx, + accumulator_phase, + || materialize_accumulator_fn(acc, BlockedEmitTo::All), + )?; + + blocked_columns.iter_mut().zip(acc_cols).for_each(|(output, acc_cols)| { + output.extend_from_slice(acc_cols.as_ref()); + }); + } + Ok::<_, datafusion_common::DataFusionError>(blocked_columns) + })?; + + let mut total_memory = 0; + + let batches = blocked_columns.into_iter().map(|cols| { + let batch = RecordBatch::try_new(output_schema.clone(), cols)?; + debug_assert!(batch.num_rows() > 0); + total_memory += batch.get_array_memory_size(); + Ok(batch) + }).collect::>(); + + MaterializedAggregateOutput::new(batches, total_memory) + } + AggregateHashTableState::OutputtingMaterialized(output) => output, + AggregateHashTableState::Done => return Ok(None), + AggregateHashTableState::Building(_) => { + return internal_err!( "next_output_batch must be called in the outputting state" ); - } - }; + } + }; - let batch = output.next_batch(batch_size); + let batch = output.next_batch().transpose(); if output.is_exhausted() { self.state = AggregateHashTableState::Done; } else { self.state = AggregateHashTableState::OutputtingMaterialized(output); } - Ok(batch) + batch } pub(in crate::aggregates) fn memory_size(&self) -> usize { @@ -340,31 +347,43 @@ 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( + pub(in crate::aggregates) fn take_all_state_batch( &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); + return Ok(vec![]); } let output = group_by_metrics.time_emitting(|| { - let mut output = state.group_values.emit(EmitTo::All)?; + let mut blocked_output = state.group_values.emit(BlockedEmitTo::All)?; + + for (idx, acc) in state.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( + let blocked_state = accumulator_metrics.time( idx, AccumulatorPhase::State, - || acc.state(EmitTo::All), - )?); + || acc.state(BlockedEmitTo::All), + )?; + + blocked_output.iter_mut().zip(blocked_state).for_each(|(output, states)| { + output.extend_from_slice(states.as_ref()); + }); } - Ok::<_, datafusion_common::DataFusionError>(output) + + Ok::<_, datafusion_common::DataFusionError>(blocked_output) })?; - let batch = RecordBatch::try_new(state_schema, output)?; - debug_assert!(batch.num_rows() > 0); + let batches = output.into_iter().map(|cols| { + let batch = RecordBatch::try_new(Arc::clone(&state_schema), cols)?; + debug_assert!(batch.num_rows() > 0); + + Ok(batch) + }).collect::>>()?; + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released @@ -373,7 +392,7 @@ impl AggregateHashTable { state.batch_group_indices.clear(); state.batch_group_indices.shrink_to_fit(); - Ok(Some(batch)) + Ok(batches) } pub(in crate::aggregates) fn is_building(&self) -> bool { @@ -505,7 +524,7 @@ pub(super) struct HashAggregateAccumulator { filter: Option>, /// Accumulator state for all groups for one aggregate expression. - accumulator: Box, + accumulator: Box, } pub(super) type AggregateAccumulator = HashAggregateAccumulator; @@ -522,7 +541,7 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator; pub(super) type AggregateBatchFn = fn( &mut AggregateAccumulator, &CompactedAccumulatorArgs, - &[usize], + &[BlocksIndex], usize, ) -> Result<()>; @@ -533,7 +552,7 @@ pub(super) type AggregateBatchFn = fn( /// * accumulator to materialize. /// * group range to emit from the accumulator. pub(super) type MaterializeAccumulatorFn = - fn(&mut AggregateAccumulator, EmitTo) -> Result>; + fn(&mut AggregateAccumulator, BlockedEmitTo) -> Result>>; /// Aggregate arguments compacted according to one aggregate's `FILTER`. pub(super) struct CompactedAccumulatorArgs { @@ -577,13 +596,13 @@ pub(super) struct AggregateHashTableBuffer { pub(super) group_by: Arc, /// Interned group keys. Accumulator state is stored separately by group index. - pub(super) group_values: Box, + 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, + pub(super) batch_group_indices: Vec, /// One item per aggregate expression. /// @@ -611,33 +630,36 @@ pub(super) enum AggregateHashTableState { /// output should not repeatedly renumber group values with `EmitTo::First`. /// Materialize once and then slice to honor `batch_size` across output polls. pub(super) struct MaterializedAggregateOutput { - batch: RecordBatch, - offset: usize, + batches: std::vec::IntoIter>, + total_memory: usize, } impl MaterializedAggregateOutput { - pub(super) fn new(batch: RecordBatch) -> Self { - Self { batch, offset: 0 } + pub(super) fn new(batches: Vec>, total_memory: usize) -> Self { + Self { + batches: batches.into_iter(), + total_memory, + } } - pub(super) fn next_batch(&mut self, batch_size: usize) -> Option { - debug_assert!(batch_size > 0); - if self.is_exhausted() { - return None; + pub(super) fn next_batch(&mut self) -> Option> { + let next = self.batches.next(); + match next { + Some(Ok(batch)) => { + debug_assert!(batch.num_rows() > 0); + self.total_memory -= batch.get_array_memory_size(); + Some(Ok(batch)) + } + res => res } - - let length = batch_size.min(self.batch.num_rows() - self.offset); - let batch = self.batch.slice(self.offset, length); - self.offset += length; - Some(batch) } pub(super) fn is_exhausted(&self) -> bool { - self.offset >= self.batch.num_rows() + self.batches.len() == 0 } pub(super) fn memory_size(&self) -> usize { - self.batch.get_array_memory_size() + self.total_memory } } @@ -647,9 +669,9 @@ impl MaterializedAggregateOutput { /// slice without allocating. At high selectivity, copying contiguous selected /// ranges avoids branching once per row. fn compact_group_indices( - group_indices: &[usize], + group_indices: &[BlocksIndex], filter: &BooleanArray, -) -> Option> { +) -> Option> { debug_assert_eq!(group_indices.len(), filter.len()); let filter = match filter.null_count() { @@ -682,7 +704,7 @@ impl HashAggregateAccumulator { aggregate_expr: Arc, arguments: Vec>, filter: Option>, - accumulator: Box, + accumulator: Box, ) -> Self { Self { aggregate_expr, @@ -693,9 +715,10 @@ impl HashAggregateAccumulator { } /// Construct a new accumulator with the same definition, but with empty internal - /// state buffers (empty [`GroupsAccumulator`]). + /// state buffers (empty [`BlockedGroupsAccumulator`]). pub(super) fn empty_like(&self) -> Result { - let accumulator = create_group_accumulator(&self.aggregate_expr)?; + let batch_size = self.accumulator.batch_size(); + let accumulator = create_blocked_group_accumulator(&self.aggregate_expr, batch_size)?; Ok(Self::new( Arc::clone(&self.aggregate_expr), self.arguments.clone(), @@ -798,7 +821,7 @@ impl HashAggregateAccumulator { pub(super) fn update_batch( &mut self, values: &CompactedAccumulatorArgs, - group_indices: &[usize], + group_indices: &[BlocksIndex], total_num_groups: usize, ) -> Result<()> { let filtered_group_indices = values @@ -817,7 +840,7 @@ impl HashAggregateAccumulator { pub(super) fn merge_batch( &mut self, values: &CompactedAccumulatorArgs, - group_indices: &[usize], + group_indices: &[BlocksIndex], total_num_groups: usize, ) -> Result<()> { debug_assert!(values.selection.is_none()); @@ -828,21 +851,21 @@ impl HashAggregateAccumulator { /// Evaluating final aggregate results according to `EmitTo`, 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 { + pub(super) fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { self.accumulator.evaluate(emit_to) } pub(super) fn evaluate_to_columns( &mut self, - emit_to: EmitTo, - ) -> Result> { + emit_to: BlockedEmitTo, + ) -> Result>> { Ok(vec![self.evaluate(emit_to)?]) } /// Evaluating partial aggregate results according to `EmitTo`, 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> { + pub(super) fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { self.accumulator.state(emit_to) } @@ -901,7 +924,7 @@ mod tests { #[test] fn compact_group_indices_uses_filter_bitmap() { - let group_indices = (0..10).collect::>(); + let group_indices = (0..10).map(BlocksIndex::new_in_first_block).collect::>(); let all_true = BooleanArray::from(vec![true; 10]); assert_eq!(compact_group_indices(&group_indices, &all_true), None); @@ -909,14 +932,14 @@ mod tests { BooleanArray::from((0..10).map(|index| index != 4).collect::>()); assert_eq!( compact_group_indices(&group_indices, &high_selectivity), - Some(vec![0, 1, 2, 3, 5, 6, 7, 8, 9]) + Some([0, 1, 2, 3, 5, 6, 7, 8, 9].map(BlocksIndex::new_in_first_block).to_vec()) ); let with_nulls = BooleanArray::from(vec![Some(true), None, Some(false), Some(true), None]); assert_eq!( compact_group_indices(&group_indices[..5], &with_nulls), - Some(vec![0, 3]) + Some(vec![0, 3].map(BlocksIndex::new_in_first_block).to_vec()) ); } @@ -1012,6 +1035,7 @@ mod tests { vec![argument], Some(Arc::new(Column::new(filter_name, filter_index))), accumulator, + 8192 )) } 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 97ef898b51f4d..ee3e3aaf25019 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,18 +23,15 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_common::assert_or_internal_err; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::EmitTo; - +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; use crate::InputOrderMode; use crate::PhysicalExpr; -use crate::aggregates::group_values::{ - AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, GroupValues, new_group_values, -}; -use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::group_values::{new_blocked_group_values, new_group_values, AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, BlockedGroupValues}; +use crate::aggregates::grouped_hash_stream::{create_blocked_group_accumulator, create_group_accumulator}; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, @@ -160,10 +157,10 @@ pub(super) struct OrderedAggregateTableBuffer { pub(super) group_ordering: GroupOrdering, /// Interned group keys, in the same group-id order used by accumulators. - pub(super) group_values: Box, + pub(super) group_values: Box, /// Scratch group id vector for the current input batch. - pub(super) group_indices: Vec, + pub(super) group_indices: Vec, /// One item per aggregate expression. /// @@ -194,9 +191,9 @@ impl OrderedAggregateTable { "OrderedAggregateTable requires config batch_size >= 1" ); - let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_ordering = GroupOrdering::try_new(input_order_mode, batch_size)?; let group_schema = agg.group_by.group_schema(input_schema)?; - let group_values = new_group_values(group_schema, &group_ordering)?; + let group_values = new_blocked_group_values(group_schema, &group_ordering, batch_size)?; let aggregate_arguments = aggregate_expressions( &agg.aggr_expr, aggregate_mode, @@ -208,7 +205,7 @@ impl OrderedAggregateTable { .zip(aggregate_arguments) .zip(filters) .map(|((agg_expr, arguments), filter)| { - let accumulator = create_group_accumulator(agg_expr)?; + let accumulator = create_blocked_group_accumulator(agg_expr, batch_size)?; Ok(AggregateAccumulator::new( Arc::clone(agg_expr), arguments, @@ -318,28 +315,40 @@ 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( + pub(in crate::aggregates) fn take_all_state_batch( &mut self, - ) -> Result> { + ) -> Result> { if self.buffer.group_values.is_empty() { - return Ok(None); + 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)?; + let mut blocked_output = self.buffer.group_values.emit(BlockedEmitTo::All)?; + + for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( + let blocked_state = accumulator_metrics.time( idx, AccumulatorPhase::State, - || acc.state(EmitTo::All), - )?); + || acc.state(BlockedEmitTo::All), + )?; + + blocked_output.iter_mut().zip(blocked_state).for_each(|(output, states)| { + output.extend_from_slice(states.as_ref()); + }); } - Ok::<_, datafusion_common::DataFusionError>(output) + + Ok::<_, DataFusionError>(blocked_output) })?; - let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; - debug_assert!(batch.num_rows() > 0); + let batches = output.into_iter().map(|cols| { + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), cols)?; + debug_assert!(batch.num_rows() > 0); + + Ok(batch) + }).collect::>>()?; // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released @@ -349,7 +358,7 @@ impl OrderedAggregateTable { self.buffer.group_indices.shrink_to_fit(); self.buffer.group_ordering.reset(); - Ok(Some(batch)) + Ok(batches) } /// Returns the [`EmitTo`], clamped to the specified batch size @@ -362,11 +371,12 @@ impl OrderedAggregateTable { &self, group_count: usize, emit_to: EmitTo, - ) -> (EmitTo, bool) { + ) -> (BlockedEmitTo, bool) { match emit_to { - EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), - EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), - EmitTo::All => (EmitTo::First(self.batch_size), false), + EmitTo::First(n) if n < self.batch_size => (BlockedEmitTo::First(n), true), + EmitTo::First(_) => (BlockedEmitTo::NextBlock, true), + EmitTo::All if group_count <= self.batch_size => (BlockedEmitTo::All, false), + EmitTo::All => (BlockedEmitTo::NextBlock, true), } } @@ -399,7 +409,7 @@ impl OrderedAggregateTable { total_num_groups, )?; } - Ok::<_, datafusion_common::DataFusionError>(total_num_groups) + Ok::<_, DataFusionError>(total_num_groups) })?; group_by_metrics.time_aggregation(|| { @@ -419,7 +429,7 @@ impl OrderedAggregateTable { ) })?; } - Ok::<(), datafusion_common::DataFusionError>(()) + Ok::<(), DataFusionError>(()) })?; } @@ -451,13 +461,16 @@ impl OrderedAggregateTable { 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)?; + assert_eq!(output.len(), 1, "must only get single group key batch per emit"); + let mut output = output.into_iter().next().unwrap(); if should_remove_groups { match emit_to { - EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), + BlockedEmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), + BlockedEmitTo::NextBlock => self.buffer.group_ordering.remove_groups(self.batch_size), // `EmitTo::All` is only used after `input_done`, when all // buffered groups are known complete and the ordering state is // no longer needed. - EmitTo::All => {} + BlockedEmitTo::All => {} } } @@ -465,10 +478,14 @@ impl OrderedAggregateTable { output.extend(accumulator_metrics.time( idx, accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), + || { + let mut acc = materialize_accumulator_fn(acc, emit_to)?; + assert_eq!(acc.len(), 1, "must emit single batch"); + Ok::<_, DataFusionError>(acc.into_iter().next().unwrap()) + }, )?); } - Ok::<_, datafusion_common::DataFusionError>(output) + Ok::<_, DataFusionError>(output) })?; let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; 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 377abdae71f4c..0c091d340436b 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::{new_blocked_group_values, new_group_values, AccumulatorPhase}; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{AggregateExec, evaluate_group_by}; @@ -78,7 +78,7 @@ 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)?; + let group_values = new_blocked_group_values(group_schema, &GroupOrdering::None, self.batch_size)?; let accumulators = state .accumulators .iter() diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs new file mode 100644 index 0000000000000..e33d6e367f006 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs @@ -0,0 +1,278 @@ +use arrow::array::{downcast_primitive, ArrayRef}; +use arrow_schema::{DataType, SchemaRef, TimeUnit}; +use datafusion_common::{assert_eq_or_internal_err, assert_ne_or_internal_err, assert_or_internal_err, not_impl_err, unwrap_or_internal_err}; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlocksIndex}; +use datafusion_expr_common::groups_accumulator::EmitTo; +use datafusion_physical_expr_common::binary_map::OutputType; +use crate::aggregates::group_values::{multi_group_by, new_group_values}; +use crate::aggregates::group_values::single_group_by::blocked_boolean::BlockedGroupValuesBoolean; +use crate::aggregates::group_values::single_group_by::boolean::GroupValuesBoolean; +use crate::aggregates::group_values::single_group_by::bytes::GroupValuesBytes; +use crate::aggregates::group_values::single_group_by::bytes_view::GroupValuesBytesView; +use crate::aggregates::order::GroupOrdering; + +/// Stores the group values during hash aggregation. +/// +/// # Background +/// +/// In a query such as `SELECT a, b, count(*) FROM t GROUP BY a, b`, the group values +/// identify each group, and correspond to all the distinct values of `(a,b)`. +/// +/// ```sql +/// -- Input has 4 rows with 3 distinct combinations of (a,b) ("groups") +/// create table t(a int, b varchar) +/// as values (1, 'a'), (2, 'b'), (1, 'a'), (3, 'c'); +/// +/// select a, b, count(*) from t group by a, b; +/// ---- +/// 1 a 2 +/// 2 b 1 +/// 3 c 1 +/// ``` +/// +/// # Design +/// +/// Managing group values is a performance critical operation in hash +/// aggregation. The major operations are: +/// +/// 1. Intern: Quickly finding existing and adding new group values +/// 2. Emit: Returning the group values as an array +/// +/// There are multiple specialized implementations of this trait optimized for +/// different data types and number of columns, optimized for these operations. +/// See [`crate::aggregates::group_values::new_group_values`] for details. +/// +/// # Group Ids +/// +/// Each distinct group in a hash aggregation is identified by a unique group id +/// (usize) which is assigned by instances of this trait. Group ids are +/// continuous without gaps, starting from 0. +pub trait BlockedGroupValues: Send { + fn block_size(&self) -> usize; + + /// Calculates the group id for each input row of `cols`, assigning new + /// group ids as necessary. + /// + /// When the function returns, `groups` must contain the group id for each + /// row in `cols`. + /// + /// If a row has the same value as a previous row, the same group id is + /// assigned. If a row has a new value, the next available group id is + /// assigned. + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> datafusion_common::Result<()>; + + /// Returns the number of bytes of memory used by this [`BlockedGroupValues`]. + /// + /// May be expensive; check the implementation before calling on hot paths. + fn size(&self) -> usize; + + /// Returns true if this [`BlockedGroupValues`] is empty + fn is_empty(&self) -> bool; + + /// The number of values (distinct group values) stored in this [`BlockedGroupValues`] + fn len(&self) -> usize; + + /// Materializes selected group values without changing the stored values or + /// their group indices. + /// + /// Rows are returned in the order specified by `selection`. An empty + /// selection returns one correctly typed empty array per group-value column. + /// + /// This method requires exclusive access because implementations may mutate + /// internal caches or builders, even though stored values are unchanged. + fn values_preserving( + &mut self, + _selection: BlockedGroupSelection<'_>, + ) -> datafusion_common::Result> { + not_impl_err!("Preserving group values are not implemented") + } + + /// Returns `true` if [`Self::values_preserving`] is implemented. + fn supports_values_preserving(&self) -> bool { + false + } + + /// Emits the group values + fn emit(&mut self, emit_to: BlockedEmitTo) -> datafusion_common::Result>> { + match emit_to { + BlockedEmitTo::All => { + self.emit_all() + } + BlockedEmitTo::NextBlock => { + let len = self.len(); + let block_size = self.block_size(); + + if len == 0 { + return Ok(vec![]); + } + + if len <= block_size { + return self.emit_all(); + } + + let block = self.emit_block()?; + let block = unwrap_or_internal_err!(block); + + // Assert that all arrays length equal block size since length is greater than block size + for arr in &block { + assert_eq_or_internal_err!(arr.len(), block_size); + } + + Ok(vec![block]) + } + BlockedEmitTo::First(n) => { + assert_ne_or_internal_err!(n, 0); + assert_or_internal_err!(n <= self.len(), "n ({n}) must be less than or equal current length ({})", self.len()); + assert_or_internal_err!(n < self.block_size(), "n ({n}) must be less than current block size ({})", self.block_size()); + + if n == self.len() { + self.emit_all() + } else { + self.emit_first_n(n).map(|first_n| vec![first_n]) + } + } + } + } + + /// Emit all group values + fn emit_all(&mut self) -> datafusion_common::Result>>; + + /// Emit the next block + /// returns Ok(None) when there are no blocks + fn emit_block(&mut self) -> datafusion_common::Result>>; + + /// Emit first `n` values and shift all values to fit into blocks + /// + /// `n` must be smaller than [`Self::block_size`] and larger than `0` + /// `n` must be smaller or equal to [`Self::len`] + fn emit_first_n(&mut self, n: usize) -> datafusion_common::Result>; + + // TODO - add into iterator which move the entry state into an iterator that will output ready batches + // this is so it won't need to update the underlying hash map whenever calling emit block + // this is for the case when we need to emit all, but we don't want to materialize all right away + // but we want to skip the internal hash map updates, so the iterator will avoid that while clearing memory + // the iterator should expose `allocated_size()` function + + /// 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); +} + + +// This is just an adapter until all is implemented +pub struct BlockedGroupValuesAdapter { + block_size: usize, + inner: Box, +} + +impl BlockedGroupValuesAdapter { + pub fn new(block_size: usize, inner: Box) -> Self { + Self { block_size, inner } + } +} + +impl BlockedGroupValues for BlockedGroupValuesAdapter { + fn block_size(&self) -> usize { + self.block_size + } + + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> datafusion_common::Result<()> { + let block_size = self.block_size; + let mut group_indices_flattened = groups.iter().map(|i| i.into_index_in_fixed_block_size(block_size)).collect::>(); + self.inner.intern(cols, &mut group_indices_flattened)?; + *groups = group_indices_flattened.iter().map(|index| BlocksIndex::from_index_in_fixed_block_size(*index, block_size)).collect::>(); + + Ok(()) + } + + fn size(&self) -> usize { + self.inner.size() + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + fn len(&self) -> usize { + self.inner.len() + } + + fn emit_all(&mut self) -> datafusion_common::Result>> { + let mut blocks = vec![]; + + while self.len() > self.block_size { + blocks.push(self.inner.emit(EmitTo::First(self.block_size))?); + } + + if self.len() > 0 { + blocks.push(self.inner.emit(EmitTo::All)?); + } + + Ok(blocks) + } + + fn emit_block(&mut self) -> datafusion_common::Result>> { + if self.len() == 0 { + return Ok(None); + } + + let output = if self.len() <= self.block_size { + self.inner.emit(EmitTo::All) + } else { + self.inner.emit(EmitTo::First(self.block_size)) + }; + + Ok(Some(output?)) + } + + fn emit_first_n(&mut self, n: usize) -> datafusion_common::Result> { + assert_ne_or_internal_err!(n, 0); + assert_or_internal_err!(n <= self.len(), "n ({n}) must be less than or equal current length ({})", self.len()); + assert_or_internal_err!(n < self.block_size(), "n ({n}) must be less than current block size ({})", self.block_size()); + + let output = if self.len() == n { + self.inner.emit(EmitTo::All) + } else { + self.inner.emit(EmitTo::First(n)) + }; + + Ok(output?) + } + + fn clear_shrink(&mut self, num_rows: usize) { + self.inner.clear_shrink(num_rows) + } +} + +/// Return a specialized implementation of [`BlockedGroupValues`] for the given schema. +/// +/// [`BlockedGroupValues`] implementations choosing logic: +/// +/// - If group by single column, and type of this column has +/// the specific [`BlockedGroupValues`] implementation, such implementation +/// will be chosen. +/// +/// - If group by multiple columns, and all column types have the specific +/// `GroupColumn` implementations, `GroupValuesColumn` will be chosen. +/// +/// - Otherwise, the general implementation `GroupValuesRows` will be chosen. +/// +/// `GroupColumn`: crate::aggregates_blocked::group_values::multi_group_by::GroupColumn +/// `GroupValuesColumn`: crate::aggregates_blocked::group_values::multi_group_by::GroupValuesColumn +/// `GroupValuesRows`: crate::aggregates_blocked::group_values::GroupValuesRows +pub fn new_blocked_group_values( + schema: SchemaRef, + group_ordering: &GroupOrdering, + block_size: usize, +) -> datafusion_common::Result> { + if schema.fields.len() == 1 { + let d = schema.fields[0].data_type(); + + if matches!(d, DataType::Boolean) { + return Ok(Box::new(BlockedGroupValuesBoolean::new(block_size))); + } + } + + let group_values = new_group_values(schema, group_ordering)?; + + Ok(Box::new(BlockedGroupValuesAdapter::new(block_size, group_values))) +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index b764fd7792b39..d11f9683c3074 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -48,6 +48,9 @@ use crate::aggregates::{ mod metrics; mod null_builder; +mod blocked; + +pub use blocked::{BlockedGroupValues, new_blocked_group_values}; pub(crate) use metrics::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs new file mode 100644 index 0000000000000..a0b7af8413de2 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs @@ -0,0 +1,282 @@ +// 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 crate::aggregates::group_values::GroupValues; + +use arrow::array::{ + ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, +}; +use datafusion_common::Result; +use datafusion_expr::EmitTo; +use std::{mem::size_of, sync::Arc}; +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; +use crate::aggregates::group_values::blocked::BlockedGroupValues; + +#[derive(Debug)] +pub struct BlockedGroupValuesBoolean { + false_group: Option, + true_group: Option, + null_group: Option, + block_size: usize, +} + +impl BlockedGroupValuesBoolean { + pub fn new(block_size: usize) -> Self { + assert_ne!(block_size, 0); + Self { + block_size, + false_group: None, + true_group: None, + null_group: None, + } + } +} + +impl BlockedGroupValues for BlockedGroupValuesBoolean { + fn block_size(&self) -> usize { + self.block_size + } + + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + let array = cols[0].as_boolean(); + groups.clear(); + + for value in array.iter() { + let index = match value { + Some(false) => { + if let Some(index) = self.false_group { + index + } else { + let index = self.len(); + self.false_group = Some(index); + index + } + } + Some(true) => { + if let Some(index) = self.true_group { + index + } else { + let index = self.len(); + self.true_group = Some(index); + index + } + } + None => { + if let Some(index) = self.null_group { + index + } else { + let index = self.len(); + self.null_group = Some(index); + index + } + } + }; + + groups.push(BlocksIndex::from_index_in_fixed_block_size(index, self.block_size)); + } + + Ok(()) + } + + fn size(&self) -> usize { + size_of::() + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { + self.false_group.is_some() as usize + + self.true_group.is_some() as usize + + self.null_group.is_some() as usize + } + + fn emit_all(&mut self) -> Result>> { + let mut blocks = vec![]; + + while let Some(block) = self.emit_block()? { + blocks.push(block); + } + + Ok(blocks) + } + + fn emit_block(&mut self) -> Result>> { + let len = self.len(); + if len == 0 { + return Ok(None); + } + // Group indices are always contiguous from 0, so the last block holds the remaining groups + let emit_count = self.block_size.min(len); + let mut builder = BooleanBufferBuilder::new(emit_count); + builder.append_n(emit_count, false); + if let Some(idx) = self.true_group.as_mut() { + if *idx < emit_count { + builder.set_bit(*idx, true); + self.true_group = None; + } else { + *idx -= emit_count; + } + } + + if let Some(idx) = self.false_group.as_mut() { + if *idx < emit_count { + // already false, no need to set + self.false_group = None; + } else { + *idx -= emit_count; + } + } + + let values = builder.finish(); + + let nulls = if let Some(idx) = self.null_group.as_mut() { + if *idx < emit_count { + let mut buffer = NullBufferBuilder::new(len); + buffer.append_n_non_nulls(*idx); + buffer.append_null(); + buffer.append_n_non_nulls(emit_count - *idx - 1); + + self.null_group = None; + Some(buffer.finish().unwrap()) + } else { + *idx -= emit_count; + None + } + } else { + None + }; + + Ok(Some(vec![Arc::new(BooleanArray::new(values, nulls)) as _])) + } + + fn emit_first_n(&mut self, n: usize) -> Result> { + assert!(n < self.block_size, "{n} must be less than block_size {}", self.block_size); + + { + let len = self.len(); + assert!(n < len, "{n} must be less than length {len}"); + } + + let mut builder = BooleanBufferBuilder::new(n); + builder.append_n(n, false); + if let Some(idx) = self.true_group.as_mut() { + if *idx < n { + builder.set_bit(*idx, true); + self.true_group = None; + } else { + *idx -= n; + } + } + + if let Some(idx) = self.false_group.as_mut() { + if *idx < n { + // already false, no need to set + self.false_group = None; + } else { + *idx -= n; + } + } + + let values = builder.finish(); + + let nulls = if let Some(idx) = self.null_group.as_mut() { + if *idx < n { + let mut buffer = NullBufferBuilder::new(builder.len()); + buffer.append_n_non_nulls(*idx); + buffer.append_null(); + buffer.append_n_non_nulls(n - *idx - 1); + + self.null_group = None; + Some(buffer.finish().unwrap()) + } else { + *idx -= n; + None + } + } else { + None + }; + + Ok(vec![Arc::new(BooleanArray::new(values, nulls)) as _]) + } + + fn clear_shrink(&mut self, _num_rows: usize) { + self.false_group = None; + self.true_group = None; + self.null_group = None; + } +} + +#[cfg(test)] +mod tests { + use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; + use super::*; + + fn intern(group_values: &mut BlockedGroupValuesBoolean, values: &[Option]) -> Vec { + let mut groups = vec![]; + let array: ArrayRef = Arc::new(BooleanArray::from(values.to_vec())); + group_values.intern(&[array], &mut groups).unwrap(); + groups + } + + fn emitted(blocks: Vec>) -> Vec>> { + blocks + .into_iter() + .map(|block| block[0].as_boolean().iter().collect()) + .collect() + } + + #[test] + fn emit_all_only_emits_existing_groups() { + let mut group_values = BlockedGroupValuesBoolean::new(531); + let groups = intern(&mut group_values, &[Some(true), None, Some(true), Some(false)]); + assert_eq!( + groups, + [0, 1, 0, 2].map(BlocksIndex::new_in_first_block).to_vec() + ); + assert_eq!(group_values.len(), 3); + + let blocks = emitted(group_values.emit_all().unwrap()); + assert_eq!(blocks, vec![vec![Some(true), None, Some(false)]]); + assert!(group_values.is_empty()); + assert!(group_values.emit_block().unwrap().is_none()); + } + + #[test] + fn emit_all_splits_into_blocks() { + let mut group_values = BlockedGroupValuesBoolean::new(2); + intern(&mut group_values, &[Some(false), Some(true), None]); + + let blocks = emitted(group_values.emit_all().unwrap()); + assert_eq!(blocks, vec![vec![Some(false), Some(true)], vec![None]]); + assert!(group_values.is_empty()); + } + + #[test] + fn emit_first_n_shifts_remaining() { + let mut group_values = BlockedGroupValuesBoolean::new(4); + intern(&mut group_values, &[None, Some(false), Some(true)]); + + let first = group_values.emit_first_n(1).unwrap(); + assert_eq!(first[0].as_boolean().iter().collect::>(), vec![None]); + assert_eq!(group_values.len(), 2); + + let blocks = emitted(group_values.emit_all().unwrap()); + assert_eq!(blocks, vec![vec![Some(false), Some(true)]]); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs new file mode 100644 index 0000000000000..84d276c5ac6ac --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs @@ -0,0 +1,291 @@ +// 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 arrow::buffer::ScalarBuffer; +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_execution::memory_pool::proxy::VecAllocExt; +use hashbrown::hash_table::HashTable; +use std::mem::size_of; +use std::sync::Arc; +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; +use datafusion_expr_common::blocked_helpers::BlockedVec; +use crate::aggregates::group_values::{BlockedGroupValues, HashValue}; + +/// A [`GroupValues`] storing a single column of primitive values +/// +/// This specialization is significantly faster than using the more general +/// purpose `Row`s format +pub struct BlockedGroupValuesPrimitive { + block_size: usize, + /// The data type of the output array + data_type: DataType, + /// Stores the `(group_index, hash)` based on the hash of its value + /// + /// We also store `hash` is for reducing cost of rehashing. Such cost + /// is obvious in high cardinality group by situation. + /// More details can see: + /// + map: HashTable<(BlocksIndex, u64)>, + /// 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 fn new(data_type: DataType, block_size: usize) -> Self { + assert!(PrimitiveArray::::is_compatible(&data_type)); + Self { + block_size, + data_type, + map: HashTable::with_capacity(128), + values: BlockedVec::new(block_size), + null_group: None, + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + } + } + + fn build_primitive( + values: impl Into>, + null_idx: Option, + ) -> PrimitiveArray { + let values = values.into(); + let nulls = null_idx.map(|null_idx| { + let mut buffer = NullBufferBuilder::new(values.len()); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(values.len() - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }); + PrimitiveArray::::new(values.into(), nulls) + } + + fn build_no_nulls_primitive_arc( + values: impl Into>, + data_type: DataType, + ) -> ArrayRef { + Arc::new(PrimitiveArray::::new(values.into(), None).with_data_type(data_type)) + } + + fn build_with_nulls_primitive_arc( + values: impl Into>, + null_idx: usize, + data_type: DataType, + ) -> ArrayRef { + let values = values.into(); + let nulls = { + let mut buffer = NullBufferBuilder::new(values.len()); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(values.len() - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }; + Arc::new( + PrimitiveArray::::new(values.into(), Some(nulls)) + .with_data_type(data_type), + ) + } +} + +impl BlockedGroupValues for BlockedGroupValuesPrimitive +where + T::Native: HashValue, +{ + fn block_size(&self) -> usize { + self.block_size + } + + 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(|| { + let group_id = self.values.len(); + self.values.push(Default::default()); + BlocksIndex::from_index_in_fixed_block_size(group_id, self.block_size) + }), + 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, + |&(g, h)| { + // SAFETY: every group in the map was appended to `values` + hash == h && unsafe { self.values.get_unchecked(g) }.is_eq(key) + }, + |&(_, h)| h, + ); + + match insert { + hashbrown::hash_table::Entry::Occupied(o) => o.get().0, + hashbrown::hash_table::Entry::Vacant(v) => { + let g = BlocksIndex::from_index_in_fixed_block_size(self.values.len(), self.block_size); + v.insert((g, hash)); + self.values.push(key); + g + } + } + } + }; + groups.push(group_id) + } + Ok(()) + } + + fn size(&self) -> usize { + self.map.capacity() * size_of::<(BlocksIndex, u64)>() + self.values.allocated_size() + } + + fn is_empty(&self) -> bool { + self.values.is_empty() + } + + fn len(&self) -> usize { + self.values.len() + } + + fn emit_all(&mut self) -> Result>> { + let mut blocks = Vec::with_capacity(self.values.num_blocks()); + self.map = HashTable::with_capacity(128); + + let all = self.values.take_all(); + let mut into_iter = all.into_iter(); + + if let Some(null_idx) = self.null_group.take() { + for values in into_iter.by_ref().take(null_idx.block_index()) { + blocks.push(vec![Self::build_no_nulls_primitive_arc( + values, + self.data_type.clone(), + )]); + } + + let block_with_nulls = into_iter.next().expect("must have block for nulls"); + blocks.push(vec![Self::build_with_nulls_primitive_arc( + block_with_nulls, + null_idx.index_in_block(), + self.data_type.clone(), + )]); + } + + for values in into_iter { + blocks.push(vec![Self::build_no_nulls_primitive_arc( + values, + self.data_type.clone(), + )]); + } + + Ok(blocks) + } + + fn emit_block(&mut self) -> Result>> { + let Some(values) = self.values.take_block() else { + return Ok(None); + }; + let block_size = self.block_size; + + // If nothing left + let null_group = if self.values.len() == 0 { + self.map.clear(); + + assert_eq!(self.null_group.map_or(0, |index| index.block_index()), 0); + + self.null_group.take() + } else { + self.map.retain(|entry| { + // Decrement group index by n + let group_idx = entry.0; + match group_idx.prev_block_checked() { + Some(new_block_index) => { + entry.0 = new_block_index; + true + } + None => false, + } + }); + + match &mut self.null_group { + Some(v) if v.block_index() != 0 => { + *v = v.prev_block(); + None + } + Some(_) => self.null_group.take(), + None => None, + } + }; + + let array = Self::build_primitive(values, null_group.map(|i| i.index_in_block())); + + Ok(Some(vec![Arc::new( + array.with_data_type(self.data_type.clone()), + )])) + } + + fn emit_first_n(&mut self, n: usize) -> Result> { + let block_size = self.block_size; + let array: PrimitiveArray = { + self.map.retain(|entry| { + // Decrement group index by n + let group_idx = entry.0; + match group_idx.sub_flat_checked(n, block_size) { + // Group index was >= n, shift value down + Some(sub) => { + entry.0 = sub; + true + } + // Group index was < n, so remove from table + None => false, + } + }); + let null_group = match &mut self.null_group { + Some(v) if v.gte_flat(n, block_size) => { + *v = v.sub_flat(n, block_size); + None + } + Some(_) => self.null_group.take().map(|g| g.index_in_block()), + None => None, + }; + let first_values = self.values.take_n(n); + Self::build_primitive(first_values, null_group) + }; + + Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) + } + + fn clear_shrink(&mut self, num_rows: usize) { + // self.values.clear(); + // self.values.shrink_to(num_rows); + self.values.reset(); + self.map.clear(); + self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs index 89c6b624e8e0a..dc06b42124a5e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs @@ -21,3 +21,5 @@ pub(crate) mod boolean; pub(crate) mod bytes; pub(crate) mod bytes_view; pub(crate) mod primitive; +pub(crate) mod blocked_boolean; +mod blocked_primitive; diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index d8f2543996f0e..03a2ffd11d35b 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -76,6 +76,8 @@ use datafusion_common::utils::memory::get_record_batch_memory_size; use futures::ready; use futures::stream::{Stream, StreamExt}; use log::debug; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedGroupsAccumulator, BlocksIndex}; +use datafusion_functions_aggregate_common::aggregate::blocked_groups_accumulator::BlockedGroupsAccumulatorAdapter; #[derive(Debug, Clone)] /// This object tracks the aggregation phase (input/output) @@ -523,7 +525,7 @@ impl GroupedHashAggregateStream { .collect::>() .join(", "); let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})"); - let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?; + let group_ordering = GroupOrdering::try_new(&agg.input_order_mode, batch_size)?; let oom_mode = match (agg.mode, &group_ordering) { // In partial aggregation mode, always prefer to emit incomplete results early. (AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly, @@ -666,6 +668,30 @@ pub(crate) fn create_group_accumulator( } } +/// Create an accumulator for `agg_expr` -- a [`BlockedGroupsAccumulator`] if +/// that is supported by the aggregate, or a +/// [`BlockedGroupsAccumulatorAdapter`] if not. +pub(crate) fn create_blocked_group_accumulator( + agg_expr: &Arc, + block_size: usize +) -> Result> { + // TODO - WHAT ABOUT DYNAMIC BLOCK SIZE? + if let Some(agg_batch_size) = agg_expr.batch_size() { + assert_eq_or_internal_err!(agg_batch_size, block_size, "Block size mismatch for blocked groups accumulator"); + } + if agg_expr.blocked_groups_accumulator_supported() { + agg_expr.create_blocked_groups_accumulator() + } else { + // Note in the log when the slow path is used + debug!( + "Creating BlockedGroupsAccumulatorAdapter for {}: {agg_expr:?}", + agg_expr.name() + ); + let group_acc = create_group_accumulator(agg_expr)?; + Ok(Box::new(BlockedGroupsAccumulatorAdapter::new(group_acc, block_size))) + } +} + impl Stream for GroupedHashAggregateStream { type Item = Result; @@ -920,12 +946,12 @@ impl GroupedHashAggregateStream { let starting_num_groups = self.group_values.len(); self.group_values .intern(group_values, &mut self.current_group_indices)?; - let group_indices = &self.current_group_indices; + let group_indices_as_blocks = self.current_group_indices.iter().map(|&idx| BlocksIndex::new_in_first_block(idx)).collect::>(); let total_num_groups = self.group_values.len(); if total_num_groups > starting_num_groups { self.group_ordering.new_groups( group_values, - group_indices, + &group_indices_as_blocks, total_num_groups, )?; } @@ -1155,9 +1181,14 @@ impl GroupedHashAggregateStream { .intern(&cols, &mut self.current_group_indices)?; let total_groups = self.group_values.len(); if total_groups > starting_groups { + let current_group_indices_as_blocks = self + .current_group_indices + .iter() + .map(|&idx| BlocksIndex::new_in_first_block(idx)) + .collect::>(); self.group_ordering.new_groups( &cols, - &self.current_group_indices, + ¤t_group_indices_as_blocks, total_groups, )?; } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 152082d88dd90..dc555b75d5bbe 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -292,27 +292,29 @@ impl FinalSpillContext { &mut self, hash_table: &mut AggregateHashTable, ) -> Result<()> { - let Some(batch) = hash_table.take_state_batch()? else { - return Ok(()); - }; + for batch in hash_table.take_all_state_batch()? { + // TODO - avoid this - it will create MANY small spill files as oppose to before, + // CHANGE TO sort iterator that will sort multiple batches instead of one batches without concat + // and will output multiple batches + let sorted_iter = + IncrementalSortIterator::new(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, + "FinalHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Final hash aggregation produced an empty spill"); + }; - let sorted_iter = - IncrementalSortIterator::new(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, - "FinalHashAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Final hash aggregation produced an empty spill"); - }; + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + } - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); Ok(()) } @@ -476,11 +478,13 @@ 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_all_state_batch()?; + + if materialized_group_states.is_empty() { + return internal_err!( "Partial hash aggregate ran out of memory with no aggregated groups" - ) - })?; + ); + } self.early_emit_count.add(1); timer.done(); @@ -594,11 +598,13 @@ impl PartialHashAggregateStream { &mut self, // After each incremental emitting step, the `remaining_groups` will be updated // with batch slicing. - mut remaining_groups: RecordBatch, + mut remaining_groups: Vec, emitter: &mut TryEmitter, hash_table_mem_size: usize, ) -> Result<()> { - let remaining_groups_memory = remaining_groups.get_array_memory_size(); + let sizes = remaining_groups.iter().map(|b| b.get_array_memory_size()).collect::>(); + // Skip first since we are emitting it right away + let total_size: usize = sizes.iter().skip(1).sum(); // Emitting clears the aggregate table and releases its // accumulated memory. Update the reservation accordingly. @@ -606,54 +612,32 @@ impl PartialHashAggregateStream { // if there is not enough memory, fallback to emit large batch match self .reservation - .try_resize(hash_table_mem_size + remaining_groups_memory) + // TODO - also reserve for the vec we are holding between emit + .try_resize(hash_table_mem_size + total_size) { 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)?; - - self.reduction_factor.add_part(remaining_groups.num_rows()); - emitter - .emit(remaining_groups.record_output(&self.baseline_metrics)) - .await; - - return Ok(()); - } + // TODO - even in out-of-memory we shouldn't concat and emit large batch since it will have huge tmp memory and may lead to machine OOM 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); - - remaining_groups = remaining_groups.slice( - self.batch_size, - remaining_groups.num_rows() - self.batch_size, - ); - + for (i, (output, size)) in remaining_groups.into_iter().zip(sizes).enumerate() { self.reduction_factor.add_part(output.num_rows()); debug_assert!(output.num_rows() > 0); - emitter - .emit(output.record_output(&self.baseline_metrics)) - .await; - } - - self.reduction_factor.add_part(remaining_groups.num_rows()); - debug_assert!(remaining_groups.num_rows() > 0); + // Do not shrink the first batch since we did not count it + if i > 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)?; + // 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(size)?; + } - emitter - .emit(remaining_groups.record_output(&self.baseline_metrics)) - .await; + emitter + .emit(output.record_output(&self.baseline_metrics)) + .await; + } Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 147e4e0d6f18d..1c3e170d39303 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -20,6 +20,7 @@ use std::mem::size_of; use arrow::array::ArrayRef; use datafusion_common::Result; use datafusion_expr::EmitTo; +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; mod full; mod partial; @@ -41,11 +42,11 @@ pub enum GroupOrdering { impl GroupOrdering { /// Create a `GroupOrdering` for the specified ordering - pub fn try_new(mode: &InputOrderMode) -> Result { + pub fn try_new(mode: &InputOrderMode, block_size: usize) -> Result { match mode { InputOrderMode::Linear => Ok(GroupOrdering::None), InputOrderMode::PartiallySorted(order_indices) => { - GroupOrderingPartial::try_new(order_indices.clone()) + GroupOrderingPartial::try_new(order_indices.clone(), block_size) .map(GroupOrdering::Partial) } InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())), @@ -128,7 +129,7 @@ impl GroupOrdering { pub fn new_groups( &mut self, batch_group_values: &[ArrayRef], - group_indices: &[usize], + group_indices: &[BlocksIndex], total_num_groups: usize, ) -> Result<()> { match self { @@ -180,14 +181,15 @@ mod tests { /// distinct values such as `[1, 2, 3]` create boundaries, while repeated /// values such as `[1, 1, 1]` do not. fn partial_ordering(sort_key_values: Vec) -> Result { + let block_size = 8192; let mut group_ordering = - GroupOrdering::Partial(GroupOrderingPartial::try_new(vec![0])?); + GroupOrdering::Partial(GroupOrderingPartial::try_new(vec![0], block_size)?); let batch_group_values: Vec = vec![ Arc::new(Int32Array::from(sort_key_values)), Arc::new(Int32Array::from(vec![10, 20, 30])), ]; - let group_indices = vec![0, 1, 2]; + let group_indices = [0, 1, 2].map(BlocksIndex::new_in_first_block).to_vec(); group_ordering.new_groups(&batch_group_values, &group_indices, 3)?; diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 1603bb6d079be..eeaa4417f59f0 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -26,6 +26,7 @@ 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_common::blocked_groups_accumulator::BlocksIndex; /// Tracks grouping state when the data is ordered by some subset of /// the group keys. @@ -67,6 +68,8 @@ pub struct GroupOrderingPartial { /// State machine state: State, + block_size: usize, + /// The indexes of the group by columns that form the sort key. /// For example if grouping by `id, state` and ordered by `state` /// this would be `[1]`. @@ -117,11 +120,12 @@ impl State { impl GroupOrderingPartial { /// TODO: Remove unnecessary `input_schema` parameter. - pub fn try_new(order_indices: Vec) -> Result { + pub fn try_new(order_indices: Vec, block_size: usize) -> Result { debug_assert!(!order_indices.is_empty()); Ok(Self { state: State::Start, order_indices, + block_size, }) } @@ -214,7 +218,7 @@ impl GroupOrderingPartial { pub fn new_groups( &mut self, batch_group_values: &[ArrayRef], - group_indices: &[usize], + group_indices: &[BlocksIndex], total_num_groups: usize, ) -> Result<()> { assert!(total_num_groups > 0); @@ -242,7 +246,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_indices[last_range.start].into_index_in_fixed_block_size(self.block_size); let range_sort_key = get_row_at_idx(&sort_keys, last_range.start)?; let (current_sort, sort_key) = if last_range.start == 0 { @@ -283,14 +287,15 @@ mod tests { fn test_group_ordering_partial() -> Result<()> { // Ordered on column a let order_indices = vec![0]; - let mut group_ordering = GroupOrderingPartial::try_new(order_indices)?; + let block_size = 8192; + let mut group_ordering = GroupOrderingPartial::try_new(order_indices, block_size)?; let batch_group_values: Vec = vec![ Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(Int32Array::from(vec![2, 1, 3])), ]; - let group_indices = vec![0, 1, 2]; + let group_indices = [0, 1, 2].map(BlocksIndex::new_in_first_block).to_vec(); let total_num_groups = 3; group_ordering.new_groups( @@ -313,7 +318,7 @@ mod tests { Arc::new(Int32Array::from(vec![3, 3, 3])), Arc::new(Int32Array::from(vec![2, 1, 7])), ]; - let group_indices = vec![3, 4, 5]; + let group_indices = [3, 4, 5].map(BlocksIndex::new_in_first_block).to_vec(); let total_num_groups = 6; group_ordering.new_groups( @@ -336,7 +341,7 @@ mod tests { Arc::new(Int32Array::from(vec![4, 4, 4])), Arc::new(Int32Array::from(vec![1, 1, 1])), ]; - let group_indices = vec![6, 7, 8]; + let group_indices = [6, 7, 8].map(BlocksIndex::new_in_first_block).to_vec(); let total_num_groups = 9; group_ordering.new_groups( diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index ac6f2491b481e..4e5a32461b94b 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -188,27 +188,25 @@ impl OrderedFinalSpillContext { &mut self, table: &mut OrderedAggregateTable, ) -> Result<()> { - let Some(batch) = table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(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, - "OrderedFinalAggregateSpill", - )?; + for batch in table.take_all_state_batch()? { + let sorted_iter = + IncrementalSortIterator::new(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, + "OrderedFinalAggregateSpill", + )?; - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Ordered final aggregation produced an empty spill"); - }; + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Ordered final aggregation produced an empty spill"); + }; - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + } Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 8c2315588ad7a..308e59d1fc5ff 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -25,7 +25,7 @@ use datafusion_common::{DataFusionError, Result}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::{Stream, StreamExt}; - +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::VecAllocExt; use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; @@ -239,16 +239,37 @@ impl OrderedPartialAggregateStream { let input_rows = batch.num_rows(); self.reduction_factor.add_total(input_rows); - let timer = elapsed_compute.timer(); + let mut timer = elapsed_compute.timer(); table.aggregate_batch(&batch)?; // Check memory reservation. See function comments for details. - if let Some(batch) = self.resize_or_take_state_batch(table)? { - self.reduction_factor.add_part(batch.num_rows()); - drop(timer); - emitter.emit(batch).await; - continue; + { + let (batches, mut size) = self.resize_or_take_state_batch(table)?; + // Exclude the first batch since we are emitting it right away + size -= batches.first().map_or(0, |(_, size)| *size); + + let table_size = table.memory_size(); + self.reservation.try_resize(table_size + size)?; + + if !batches.is_empty() { + for (index, (batch, batch_size)) in batches.into_iter().enumerate() { + self.reduction_factor.add_part(batch.num_rows()); + + // Exclude the first batch since size does not include it + if index > 0 { + size -= batch_size; + self.reservation.try_resize(table_size + size)?; + } + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + self.reservation.try_resize(table_size)?; + + continue; + } } let Some(batch) = table.next_output_batch()? else { @@ -284,9 +305,9 @@ impl OrderedPartialAggregateStream { fn resize_or_take_state_batch( &mut self, table: &mut OrderedAggregateTable, - ) -> Result> { + ) -> Result<(Vec<(RecordBatch, usize)>, usize)> { let oom = match self.reservation.try_resize(table.memory_size()) { - Ok(()) => return Ok(None), + Ok(()) => return Ok((vec![], 0)), Err(e @ DataFusionError::ResourcesExhausted(_)) => e, Err(e) => return Err(e), }; @@ -295,11 +316,19 @@ impl OrderedPartialAggregateStream { return Err(oom); } - let Some(batch) = table.take_state_batch()? else { - return Err(oom); + let mut total_size = 0; + + let batches = match table.take_all_state_batch() { + Ok(batches) if batches.is_empty() => return Err(oom), + Ok(batches) => batches.into_iter().map(|b| { + let size = b.get_array_memory_size(); + total_size += size; + (b, size) + }).collect::>(), + Err(e) => return Err(e), }; - self.reservation.try_resize(table.memory_size())?; - Ok(Some(batch)) + + Ok((batches, total_size)) } /// Emits one batch after input is exhausted. diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 701dee4e8146d..50c0609b8ceac 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -248,27 +248,25 @@ impl OrderedSingleSpillContext { &mut self, table: &mut OrderedAggregateTable, ) -> Result<()> { - let Some(batch) = table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(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, - "OrderedSingleAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Ordered single aggregation produced an empty spill"); - }; + for batch in table.take_all_state_batch()? { + let sorted_iter = + IncrementalSortIterator::new(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, + "OrderedSingleAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Ordered single aggregation produced an empty spill"); + }; - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + } 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..e242b10f8f0d7 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -107,9 +107,7 @@ enum PartialReduceHashAggregateState { /// 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, + remaining_groups: std::vec::IntoIter, }, ProducingOutput { hash_table: AggregateHashTable, @@ -318,21 +316,22 @@ 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_all_state_batch(); 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(), }, ) } - // 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,28 +349,27 @@ impl PartialReduceHashAggregateStream { ) -> PartialReduceHashAggregateStateTransition { let PartialReduceHashAggregateState::EmittingOnMemoryPressure { hash_table, - remaining_groups: batch, + remaining_groups: mut batches, } = original_state else { unreachable!("expected the EmittingOnMemoryPressure state") }; - let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + let next_batch = batches.next().expect("must have at least one batch to emit"); + + let (output_batch, next_state) = if batches.len() == 0 { // Go back to `ReadingInput` ( - batch, + next_batch, 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, + next_batch, PartialReduceHashAggregateState::EmittingOnMemoryPressure { hash_table, - remaining_groups: remaining, + remaining_groups: batches, }, ) }; diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index cfa708c51129b..0ca88a616f686 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -273,27 +273,29 @@ impl SingleSpillContext { &mut self, hash_table: &mut AggregateHashTable, ) -> Result<()> { - let Some(batch) = hash_table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(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, - "SingleHashAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Single hash aggregation produced an empty spill"); - }; + for batch in hash_table.take_all_state_batch()? { + // TODO - avoid this - it will create MANY small spill files as oppose to before, + // CHANGE TO sort iterator that will sort multiple batches instead of one batches without concat + // and will output multiple batches + + let sorted_iter = + IncrementalSortIterator::new(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, + "SingleHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Single hash aggregation produced an empty spill"); + }; - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + } Ok(()) } From baa07d177a4baae5992e600a726897748295a698 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:06:20 +0300 Subject: [PATCH 2/8] add primitive --- .../src/blocked_groups_accumulator.rs | 29 +- .../blocked_boolean_buffer_builder.rs | 273 ++--------- .../src/blocked_helpers/blocked_vec.rs | 317 ++---------- .../expr-common/src/blocked_helpers/mod.rs | 7 +- datafusion/expr-common/src/lib.rs | 2 +- datafusion/expr/src/function.rs | 3 +- datafusion/expr/src/lib.rs | 9 +- datafusion/expr/src/udaf.rs | 26 +- .../src/accumulator.rs | 9 +- .../src/aggregate.rs | 2 +- .../aggregate/blocked_groups_accumulator.rs | 84 +++- .../blocked_groups_accumulator/accumulate.rs | 121 +++-- .../blocked_prim_op.rs | 39 +- datafusion/functions-aggregate/src/count.rs | 75 +-- datafusion/physical-expr/src/aggregate.rs | 29 +- .../aggregates/aggregate_hash_table/common.rs | 157 +++--- .../aggregate_hash_table/common_ordered.rs | 74 +-- .../aggregate_hash_table/partial_table.rs | 10 +- .../src/aggregates/group_values/blocked.rs | 408 +++++++++------- .../src/aggregates/group_values/mod.rs | 2 +- .../single_group_by/blocked_boolean.rs | 38 +- .../single_group_by/blocked_primitive.rs | 451 +++++++++--------- .../group_values/single_group_by/mod.rs | 4 +- .../src/aggregates/grouped_hash_stream.rs | 24 +- .../src/aggregates/hash_stream.rs | 37 +- .../src/aggregates/order/partial.rs | 6 +- .../src/aggregates/ordered_final_stream.rs | 21 +- .../src/aggregates/ordered_partial_stream.rs | 26 +- .../src/aggregates/ordered_single_stream.rs | 21 +- .../src/aggregates/partial_reduce_stream.rs | 8 +- .../src/aggregates/single_stream.rs | 17 +- 31 files changed, 1116 insertions(+), 1213 deletions(-) diff --git a/datafusion/expr-common/src/blocked_groups_accumulator.rs b/datafusion/expr-common/src/blocked_groups_accumulator.rs index 50fb30d4706ca..7ac34ee834b49 100644 --- a/datafusion/expr-common/src/blocked_groups_accumulator.rs +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -17,13 +17,12 @@ //! Vectorized [`GroupsAccumulator`] +use crate::groups_accumulator::EmitTo; use arrow::array::{ArrayRef, BooleanArray}; use datafusion_common::{ Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err, not_impl_err, - utils::split_vec_min_alloc, }; use std::cmp::Ordering; -use crate::groups_accumulator::EmitTo; /// Selects groups for a non-destructive grouped aggregation read. /// @@ -151,15 +150,18 @@ impl BlocksIndex { }; pub fn new(block_index: usize, index_in_block: usize) -> Self { - Self { - block_index: Self::narrow(block_index), - index_in_block: Self::narrow(index_in_block), - } + Self { + block_index: Self::narrow(block_index), + index_in_block: Self::narrow(index_in_block), + } } #[inline(always)] fn narrow(value: usize) -> u32 { - debug_assert!(value <= u32::MAX as usize, "block index part {value} does not fit in u32"); + debug_assert!( + u32::try_from(value).is_ok(), + "block index part {value} does not fit in u32" + ); value as u32 } @@ -196,11 +198,16 @@ impl BlocksIndex { pub fn sub_flat_checked(self, rhs_flat: usize, block_size: usize) -> Option { let self_flat = self.into_index_in_fixed_block_size(block_size); - self_flat.checked_sub(rhs_flat).map(|v| Self::from_index_in_fixed_block_size(v, block_size)) + self_flat + .checked_sub(rhs_flat) + .map(|v| Self::from_index_in_fixed_block_size(v, block_size)) } pub fn sub_flat(self, rhs_flat: usize, block_size: usize) -> Self { - Self::from_index_in_fixed_block_size(self.into_index_in_fixed_block_size(block_size) - rhs_flat, block_size) + Self::from_index_in_fixed_block_size( + self.into_index_in_fixed_block_size(block_size) - rhs_flat, + block_size, + ) } pub fn gte_flat(self, rhs_flat: usize, block_size: usize) -> bool { @@ -249,7 +256,9 @@ impl PartialOrd for BlocksIndex { impl Ord for BlocksIndex { fn cmp(&self, other: &Self) -> Ordering { - self.block_index.cmp(&other.block_index).then(self.index_in_block.cmp(&other.index_in_block)) + self.block_index + .cmp(&other.block_index) + .then(self.index_in_block.cmp(&other.index_in_block)) } } diff --git a/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs b/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs index 73c6425c0d7fe..e1783050d2bdf 100644 --- a/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs +++ b/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs @@ -1,23 +1,10 @@ -use datafusion_common::utils::proxy::{VecAllocExt}; +use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; +use arrow::array::BooleanBufferBuilder; +use arrow::buffer::BooleanBuffer; use std::collections::VecDeque; use std::fmt::Debug; -use std::ops::{Index, IndexMut}; -use arrow::array::BooleanBufferBuilder; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; -use arrow::datatypes::ArrowNativeType; -use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; - -/// Blocked Vec -/// -/// # Implementation Notes -/// -/// ## Why `T: Copy`? -/// -/// -/// 1. So the [`BlockedVec::allocated_size`] will be accurate since the size of T is known, and not include heap allocations (like `String` or `Vec`) that are not part of the allocated size of the `BlockedVec` -/// 2. So we can provide mutable access to the items (e.g. `IndexMut`) since if `T` is not `Copy` (like when `T` is a `Vec`) we could change the size of it without the [`BlockedVec::allocated_size`] changing, which would be confusing and lead to bugs -/// +/// Blocked [`BooleanBufferBuilder`] #[derive(Debug)] pub struct BlockedBooleanBufferBuilder { /// Using `VecDeque` so we can remove the first block and reclaim memory @@ -36,10 +23,7 @@ pub struct BlockedBooleanBufferBuilder { finished_blocks_allocated_memory: usize, } -impl BlockedBooleanBufferBuilder -{ - // TODO - some want to preallocate the blocks and some don't, - // there should be a way while avoiding having a lot of memory used if all are prealocatting +impl BlockedBooleanBufferBuilder { pub fn new(block_size: usize) -> Self { assert_ne!(block_size, 0, "block size must be greater than 0"); @@ -62,16 +46,13 @@ impl BlockedBooleanBufferBuilder self.len } - pub fn num_blocks(&self) -> usize { + fn num_blocks(&self) -> usize { self.current_block_index + (self.should_count_current_block() as usize) } fn should_count_current_block(&self) -> bool { - self.should_count_current_block || !self.blocks[self.current_block_index].is_empty() - } - - pub fn block(&self, block_index: usize) -> &BooleanBufferBuilder { - &self.blocks[block_index] + self.should_count_current_block + || !self.blocks[self.current_block_index].is_empty() } pub fn block_size(&self) -> usize { @@ -81,36 +62,14 @@ impl BlockedBooleanBufferBuilder pub fn allocated_size(&self) -> usize { self.finished_blocks_allocated_memory + self.blocks.capacity() * size_of::() - + self.blocks.back().map_or(0, |b| builder_allocated_size(b)) - } - - /// Get the number of elements in the current block - pub fn current_block_len(&self) -> usize { - self.blocks[self.current_block_index].len() - } - - pub fn start_new_block(&mut self) { - // a block that was only pre-opened becomes the started block instead of an empty - // block of its own - if self.should_count_current_block() { - self.end_current_block(); - } - self.should_count_current_block = true; - } - - pub(crate) fn mark_current_block_counted(&mut self) { - self.should_count_current_block = true; - } - - pub fn end_current_block(&mut self) { - self.end_current_block_inner(); + + self.blocks.back().map_or(0, builder_allocated_size) } - fn end_current_block_inner(&mut self) { + fn end_current_block(&mut self) { // Don't add to number of blocks since we might not insert into it self.current_block_index += 1; self.finished_blocks_allocated_memory += - self.blocks.back().map_or(0, |b| builder_allocated_size(b)); + self.blocks.back().map_or(0, builder_allocated_size); let new_block = BooleanBufferBuilder::new(self.block_size); // Don't count current block since we might not insert into it @@ -118,7 +77,7 @@ impl BlockedBooleanBufferBuilder self.blocks.push_back(new_block); } - pub(crate) fn reserve_blocks(&mut self, n: usize) { + fn reserve_blocks(&mut self, n: usize) { self.blocks.reserve(n); } @@ -132,118 +91,18 @@ impl BlockedBooleanBufferBuilder let finished_block = block.len() == self.block_size; if finished_block { - self.end_current_block_inner(); - true - } else { - false - } - } - - /// Extends iterator of lengths within current block - /// Returns if the current block has finished - /// - /// # Panics - /// Panics if the iterator length exceeds the remaining size of the current block - pub(super) fn extend_in_block( - &mut self, - iter: impl Iterator, - ) -> bool { - let block = &mut self.blocks[self.current_block_index]; - - let prev_block_len = block.len(); - - // TODO - add to arrow extend support for iterators of bools, but for now just append one by one - for value in iter { - block.append(value); - } - - assert!( - block.len() <= self.block_size, - "overflow from block new block length: {}, block size: {}", - block.len(), - self.block_size - ); - - let added_items = block.len() - prev_block_len; - self.len += added_items; - - let finished_block = block.len() == self.block_size; - - if finished_block { - self.end_current_block_inner(); - true - } else { - false - } - } - - /// Extends from slice within current block - /// Returns if the current block has finished - /// - /// # Panics - /// Panics if the iterator length exceeds the remaining size of the current block - pub(super) fn extend_from_slice_in_block( - &mut self, - slice: &[bool], - ) -> bool - { - let block = &mut self.blocks[self.current_block_index]; - - let prev_block_len = block.len(); - block.append_slice(slice); - - assert!( - block.len() <= self.block_size, - "overflow from block new block length: {}, block size: {}", - block.len(), - self.block_size - ); - - let added_items = block.len() - prev_block_len; - self.len += added_items; - - let finished_block = block.len() == self.block_size; - - if finished_block { - self.end_current_block_inner(); + self.end_current_block(); true } else { false } } - /// Extend the length from the current offsets - pub fn extend_from_slice( - &mut self, - mut buffer: &[bool], - ) { - let number_of_blocks_to_reserve = buffer - .len() - .saturating_sub(self.current_block_remaining_len()) - .div_ceil(self.block_size); - self.reserve_blocks(number_of_blocks_to_reserve); - - while !buffer.is_empty() { - let remaining_in_current_block = self.current_block_remaining_len(); - let to_add = remaining_in_current_block.min(buffer.len()); - - let (to_copy, rest) = buffer.split_at(to_add); - buffer = rest; - - self.extend_from_slice_in_block(to_copy); - } - } - pub fn current_block_remaining_len(&self) -> usize { self.block_size - self.blocks[self.current_block_index].len() } - pub fn push_value_n_within_block( - &mut self, - value: bool, - n: usize, - ) -> bool - { + fn push_value_n_within_block(&mut self, value: bool, n: usize) -> bool { self.len += n; let block = &mut self.blocks[self.current_block_index]; @@ -258,18 +117,14 @@ impl BlockedBooleanBufferBuilder let finished_block = block.len() == self.block_size; if finished_block { - self.end_current_block_inner(); + self.end_current_block(); true } else { false } } - pub fn push_value_n( - &mut self, - value: bool, - mut n: usize, - ) { + pub fn push_value_n(&mut self, value: bool, mut n: usize) { let number_of_blocks_to_reserve = n .saturating_sub(self.current_block_remaining_len()) .div_ceil(self.block_size); @@ -285,12 +140,6 @@ impl BlockedBooleanBufferBuilder } } - pub fn push_value_n_to_len(&mut self, value: bool, new_len: usize) { - assert!(new_len >= self.len, "new_len must be greater than or equal to current len"); - let n = new_len - self.len; - self.push_value_n(value, n); - } - pub fn get_bit(&self, index: BlocksIndex) -> bool { let block_index = index.block_index(); let index_in_block = index.index_in_block(); @@ -305,60 +154,18 @@ impl BlockedBooleanBufferBuilder self.blocks[block_index].set_bit(index_in_block, value) } - /// Block `block_index` without bounds checking - /// - /// # Safety - /// `block_index < self.num_blocks()` - #[inline] - pub unsafe fn block_unchecked(&self, block_index: usize) -> &BooleanBufferBuilder { - debug_assert!(block_index < self.blocks.len()); - unsafe { self.blocks.get(block_index).unwrap_unchecked() } - } - - /// Mutable block `block_index` without bounds checking - /// - /// # Safety - /// `block_index < self.num_blocks()` - #[inline] - pub unsafe fn block_unchecked_mut( - &mut self, - block_index: usize, - ) -> &mut BooleanBufferBuilder { - debug_assert!(block_index < self.blocks.len()); - unsafe { self.blocks.get_mut(block_index).unwrap_unchecked() } - } - - pub fn blocks_mut(&mut self) -> impl Iterator { - let num_blocks = self.num_blocks(); - self.blocks.iter_mut().take(num_blocks) - } - - pub fn current_block_mut(&mut self) -> &mut BooleanBufferBuilder { - &mut self.blocks[self.current_block_index] - } - - pub fn current_or_block_mut(&mut self, block_index: usize) -> &mut BooleanBufferBuilder { - &mut self.blocks[block_index] - } - pub fn emit(&mut self, emit_to: BlockedEmitTo) -> Vec { match emit_to { - BlockedEmitTo::All => { - let counts = self.take_all(); - - counts - } + BlockedEmitTo::All => self.take_all(), BlockedEmitTo::NextBlock => { self.take_block().map_or(vec![], |next| vec![next]) } - BlockedEmitTo::First(n) => { - vec![self.take_n(n)] - } + BlockedEmitTo::First(n) => vec![self.take_n(n)], } } /// Take every block that counts, see [`Self::num_blocks`] - pub fn take_all(&mut self) -> Vec { + fn take_all(&mut self) -> Vec { let num_blocks = self.num_blocks(); let mut blocks = std::mem::take(&mut self.blocks); blocks.truncate(num_blocks); @@ -368,7 +175,7 @@ impl BlockedBooleanBufferBuilder } /// Take the first block, `None` once there are no more items - pub fn take_block(&mut self) -> Option { + fn take_block(&mut self) -> Option { if self.num_blocks() == 0 { return None; } @@ -377,13 +184,14 @@ impl BlockedBooleanBufferBuilder /// Take the first block even when it is empty, for callers that know from /// elsewhere that the block holds items - pub fn take_first_block(&mut self) -> BooleanBuffer { + fn take_first_block(&mut self) -> BooleanBuffer { let block = self.blocks.pop_front().expect("always at least one block"); if self.blocks.is_empty() { self.current_block_index = 0; self.should_count_current_block = false; - self.blocks.push_back(BooleanBufferBuilder::new(self.block_size)); + self.blocks + .push_back(BooleanBufferBuilder::new(self.block_size)); } else { self.current_block_index -= 1; @@ -396,28 +204,35 @@ impl BlockedBooleanBufferBuilder block.build() } - pub fn take_n( - &mut self, - n: usize, - ) -> BooleanBuffer { + fn take_n(&mut self, n: usize) -> BooleanBuffer { assert_ne!(n, 0, "n must be greater than 0"); assert!(n <= self.len, "n ({n}) must be <= len ({})", self.len); - assert!(n < self.block_size, "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", self.block_size); + assert!( + n < self.block_size, + "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", + self.block_size + ); if n == self.len { let blocks = self.take_all(); - return blocks.into_iter().next().unwrap_or_else(|| BooleanBuffer::new_unset(0)) + return blocks + .into_iter() + .next() + .unwrap_or_else(|| BooleanBuffer::new_unset(0)); } if n == self.blocks[0].len() { - return self.take_block().expect("must have at least one block since n < len and n != 0"); + return self + .take_block() + .expect("must have at least one block since n < len and n != 0"); } let prev_len = self.len; - debug_assert!( - self.blocks.back().is_some_and(|block| block.len() < self.block_size), + self.blocks + .back() + .is_some_and(|block| block.len() < self.block_size), "the last block must not be full (since it should be the writable tail)" ); @@ -438,7 +253,10 @@ impl BlockedBooleanBufferBuilder // Every block gives its first `n` items to the end of the previous block and // shifts the rest down, so every finished block stays full and keeps its allocation for i in 1..self.blocks.len() { - let mut block = std::mem::replace(&mut self.blocks[i], BooleanBufferBuilder::new(self.block_size)); + let block = std::mem::replace( + &mut self.blocks[i], + BooleanBufferBuilder::new(self.block_size), + ); let block_len = block.len(); let moved = n.min(block_len); @@ -452,7 +270,8 @@ impl BlockedBooleanBufferBuilder // If the block before it is full it is the writable tail as is and keeps its allocation, // otherwise the block before it is the writable tail and the empty one is not part of the layout let tail_is_empty = self.blocks.back().is_some_and(|b| b.is_empty()); - let prev_is_full = self.blocks.len() >= 2 && self.blocks[self.blocks.len() - 2].len() == self.block_size; + let prev_is_full = self.blocks.len() >= 2 + && self.blocks[self.blocks.len() - 2].len() == self.block_size; if tail_is_empty && !prev_is_full { self.blocks.pop_back(); @@ -465,7 +284,7 @@ impl BlockedBooleanBufferBuilder .blocks .iter() .take(self.current_block_index) - .map(|block| builder_allocated_size(block)) + .map(builder_allocated_size) .sum(); taken.build() diff --git a/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs b/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs index beb3feccf8f17..edbe28d12e91b 100644 --- a/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs +++ b/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs @@ -1,11 +1,8 @@ -use datafusion_common::utils::proxy::{VecAllocExt}; +use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; +use datafusion_common::utils::proxy::VecAllocExt; use std::collections::VecDeque; use std::fmt::Debug; use std::ops::{Index, IndexMut}; -use arrow::buffer::ScalarBuffer; -use arrow::datatypes::ArrowNativeType; -use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; - /// Blocked Vec /// @@ -35,8 +32,7 @@ pub struct BlockedVec { finished_blocks_allocated_memory: usize, } -impl BlockedVec -{ +impl BlockedVec { // TODO - some want to preallocate the blocks and some don't, // there should be a way while avoiding having a lot of memory used if all are prealocatting pub fn new(block_size: usize) -> Self { @@ -61,16 +57,13 @@ impl BlockedVec self.len } - pub fn num_blocks(&self) -> usize { + fn num_blocks(&self) -> usize { self.current_block_index + (self.should_count_current_block() as usize) } fn should_count_current_block(&self) -> bool { - self.should_count_current_block || !self.blocks[self.current_block_index].is_empty() - } - - pub fn block(&self, block_index: usize) -> &Vec { - &self.blocks[block_index] + self.should_count_current_block + || !self.blocks[self.current_block_index].is_empty() } pub fn block_size(&self) -> usize { @@ -83,33 +76,11 @@ impl BlockedVec + self.blocks.back().map_or(0, |b| b.allocated_size()) } - /// Get the number of elements in the current block - pub fn current_block_len(&self) -> usize { - self.blocks[self.current_block_index].len() - } - - pub fn start_new_block(&mut self) { - // a block that was only pre-opened becomes the started block instead of an empty - // block of its own - if self.should_count_current_block() { - self.end_current_block(); - } - self.should_count_current_block = true; - } - - pub(crate) fn mark_current_block_counted(&mut self) { - self.should_count_current_block = true; - } - - pub fn end_current_block(&mut self) { - self.end_current_block_inner(); - } - - fn end_current_block_inner(&mut self) { + fn end_current_block(&mut self) { // Don't add to number of blocks since we might not insert into it self.current_block_index += 1; self.finished_blocks_allocated_memory += - self.blocks.back().map_or(0, |b| b.allocated_size()); + self.blocks.back().map_or(0, |b| b.allocated_size()); let new_block = vec![]; // Don't count current block since we might not insert into it @@ -117,7 +88,7 @@ impl BlockedVec self.blocks.push_back(new_block); } - pub(crate) fn reserve_blocks(&mut self, n: usize) { + fn reserve_blocks(&mut self, n: usize) { self.blocks.reserve(n); } @@ -131,114 +102,18 @@ impl BlockedVec let finished_block = block.len() == self.block_size; if finished_block { - self.end_current_block_inner(); - true - } else { - false - } - } - - /// Extends iterator of lengths within current block - /// Returns if the current block has finished - /// - /// # Panics - /// Panics if the iterator length exceeds the remaining size of the current block - pub(super) fn extend_in_block( - &mut self, - iter: impl Iterator, - ) -> bool { - let block = &mut self.blocks[self.current_block_index]; - - let prev_block_len = block.len(); - block.extend(iter); - - assert!( - block.len() <= self.block_size, - "overflow from block new block length: {}, block size: {}", - block.len(), - self.block_size - ); - - let added_items = block.len() - prev_block_len; - self.len += added_items; - - let finished_block = block.len() == self.block_size; - - if finished_block { - self.end_current_block_inner(); - true - } else { - false - } - } - - /// Extends from slice within current block - /// Returns if the current block has finished - /// - /// # Panics - /// Panics if the iterator length exceeds the remaining size of the current block - pub(super) fn extend_from_slice_in_block( - &mut self, - slice: &[T], - ) -> bool - { - let block = &mut self.blocks[self.current_block_index]; - - let prev_block_len = block.len(); - block.copy_from_slice(slice); - - assert!( - block.len() <= self.block_size, - "overflow from block new block length: {}, block size: {}", - block.len(), - self.block_size - ); - - let added_items = block.len() - prev_block_len; - self.len += added_items; - - let finished_block = block.len() == self.block_size; - - if finished_block { - self.end_current_block_inner(); + self.end_current_block(); true } else { false } } - /// Extend the length from the current offsets - pub fn extend_from_slice( - &mut self, - mut buffer: &[T], - ) { - let number_of_blocks_to_reserve = buffer - .len() - .saturating_sub(self.current_block_remaining_len()) - .div_ceil(self.block_size); - self.reserve_blocks(number_of_blocks_to_reserve); - - while !buffer.is_empty() { - let remaining_in_current_block = self.current_block_remaining_len(); - let to_add = remaining_in_current_block.min(buffer.len()); - - let (to_copy, rest) = buffer.split_at(to_add); - buffer = rest; - - self.extend_from_slice_in_block(to_copy); - } - } - - pub fn current_block_remaining_len(&self) -> usize { + fn current_block_remaining_len(&self) -> usize { self.block_size - self.blocks[self.current_block_index].len() } - pub fn push_value_n_within_block( - &mut self, - value: T, - n: usize, - ) -> bool - { + fn push_value_n_within_block(&mut self, value: T, n: usize) -> bool { self.len += n; let block = &mut self.blocks[self.current_block_index]; @@ -253,26 +128,14 @@ impl BlockedVec let finished_block = block.len() == self.block_size; if finished_block { - self.end_current_block_inner(); + self.end_current_block(); true } else { false } } - /// Push default - pub fn push_default_n(&mut self, n: usize) - where - T: Default, - { - self.push_value_n(T::default(), n); - } - - pub fn push_value_n( - &mut self, - value: T, - mut n: usize, - ) { + pub fn push_value_n(&mut self, value: T, mut n: usize) { let number_of_blocks_to_reserve = n .saturating_sub(self.current_block_remaining_len()) .div_ceil(self.block_size); @@ -289,7 +152,10 @@ impl BlockedVec } pub fn push_value_n_to_len(&mut self, value: T, new_len: usize) { - assert!(new_len >= self.len, "new_len must be greater than or equal to current len"); + assert!( + new_len >= self.len, + "new_len must be greater than or equal to current len" + ); let n = new_len - self.len; self.push_value_n(value, n); } @@ -299,7 +165,7 @@ impl BlockedVec /// # Safety /// `block_index < self.num_blocks()` #[inline] - pub unsafe fn block_unchecked(&self, block_index: usize) -> &Vec { + unsafe fn block_unchecked(&self, block_index: usize) -> &Vec { debug_assert!(block_index < self.blocks.len()); unsafe { self.blocks.get(block_index).unwrap_unchecked() } } @@ -309,46 +175,32 @@ impl BlockedVec /// # Safety /// `block_index < self.num_blocks()` #[inline] - pub unsafe fn block_unchecked_mut( - &mut self, - block_index: usize, - ) -> &mut Vec { + unsafe fn block_unchecked_mut(&mut self, block_index: usize) -> &mut Vec { debug_assert!(block_index < self.blocks.len()); unsafe { self.blocks.get_mut(block_index).unwrap_unchecked() } } - pub fn blocks_mut(&mut self) -> impl Iterator> { - let num_blocks = self.num_blocks(); - self.blocks.iter_mut().take(num_blocks) - } - - pub fn current_block_mut(&mut self) -> &mut Vec { - &mut self.blocks[self.current_block_index] - } - - pub fn current_or_block_mut(&mut self, block_index: usize) -> &mut Vec { - &mut self.blocks[block_index] - } - - pub fn get_unchecked(&self, index: BlocksIndex) -> &T { + /// Get reference to the item at `index` without bounds checking + /// + /// # Safety + /// index must be in bound + pub unsafe fn get_unchecked(&self, index: BlocksIndex) -> &T { let block = unsafe { self.block_unchecked(index.block_index()) }; - let item = unsafe { block.get_unchecked(index.index_in_block()) }; - item + unsafe { block.get_unchecked(index.index_in_block()) } } - pub fn get_unchecked_mut(&mut self, index: BlocksIndex) -> &mut T { + /// Get mutable reference to the item at `index` without bounds checking + /// + /// # Safety + /// index must be in bound + pub unsafe fn get_unchecked_mut(&mut self, index: BlocksIndex) -> &mut T { let block = unsafe { self.block_unchecked_mut(index.block_index()) }; - let item = unsafe { block.get_unchecked_mut(index.index_in_block()) }; - item + unsafe { block.get_unchecked_mut(index.index_in_block()) } } pub fn emit(&mut self, emit_to: BlockedEmitTo) -> Vec> { match emit_to { - BlockedEmitTo::All => { - let counts = self.take_all(); - - counts - } + BlockedEmitTo::All => self.take_all(), BlockedEmitTo::NextBlock => { self.take_block().map_or(vec![], |next| vec![next]) } @@ -373,12 +225,6 @@ impl BlockedVec if self.num_blocks() == 0 { return None; } - Some(self.take_first_block()) - } - - /// Take the first block even when it is empty, for callers that know from - /// elsewhere that the block holds items - pub fn take_first_block(&mut self) -> Vec { let block = self.blocks.pop_front().expect("always at least one block"); if self.blocks.is_empty() { @@ -394,26 +240,17 @@ impl BlockedVec self.len -= block.len(); - block + Some(block) } - pub fn take_block_finished(&mut self) -> Option> - where T: ArrowNativeType - { - let block = self.take_block()?; - - let finished = ScalarBuffer::from(block); - - Some(finished) - } - - pub fn take_n( - &mut self, - n: usize, - ) -> Vec { + pub fn take_n(&mut self, n: usize) -> Vec { assert_ne!(n, 0, "n must be greater than 0"); assert!(n <= self.len, "n ({n}) must be <= len ({})", self.len); - assert!(n < self.block_size, "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", self.block_size); + assert!( + n < self.block_size, + "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", + self.block_size + ); if n == self.len { let blocks = self.take_all(); @@ -421,14 +258,17 @@ impl BlockedVec } if n == self.blocks[0].len() { - return self.take_block().expect("must have at least one block since n < len and n != 0"); + return self + .take_block() + .expect("must have at least one block since n < len and n != 0"); } let prev_len = self.len; - debug_assert!( - self.blocks.back().is_some_and(|block| block.len() < self.block_size), + self.blocks + .back() + .is_some_and(|block| block.len() < self.block_size), "the last block must not be full (since it should be the writable tail)" ); @@ -462,7 +302,8 @@ impl BlockedVec // If the block before it is full it is the writable tail as is and keeps its allocation, // otherwise the block before it is the writable tail and the empty one is not part of the layout let tail_is_empty = self.blocks.back().is_some_and(Vec::is_empty); - let prev_is_full = self.blocks.len() >= 2 && self.blocks[self.blocks.len() - 2].len() == self.block_size; + let prev_is_full = self.blocks.len() >= 2 + && self.blocks[self.blocks.len() - 2].len() == self.block_size; if tail_is_empty && !prev_is_full { self.blocks.pop_back(); @@ -490,76 +331,22 @@ impl BlockedVec } } -impl - Extend - for BlockedVec -{ - fn extend>( - &mut self, - iter: U, - ) { - let mut iter = iter.into_iter(); - - loop { - let remaining_in_current_block = self.current_block_remaining_len(); - let block_finished = - self.extend_in_block(iter.by_ref().take(remaining_in_current_block)); - - if !block_finished { - break; - } - } - } -} - -impl Index - for BlockedVec -where - T: Copy, -{ - type Output = T; - - fn index(&self, index: usize) -> &Self::Output { - self.index(BlocksIndex::from_index_in_fixed_block_size( - index, - self.block_size, - )) - } -} - -impl IndexMut - for BlockedVec -where - T: Copy, -{ - fn index_mut(&mut self, index: usize) -> &mut Self::Output { - self.index_mut(BlocksIndex::from_index_in_fixed_block_size( - index, - self.block_size, - )) - } -} - -impl Index - for BlockedVec +impl Index for BlockedVec where T: Copy, { type Output = T; fn index(&self, index: BlocksIndex) -> &Self::Output { - &self.blocks[index.block_index()] - [index.index_in_block()] + &self.blocks[index.block_index()][index.index_in_block()] } } -impl IndexMut - for BlockedVec +impl IndexMut for BlockedVec where T: Copy, { fn index_mut(&mut self, index: BlocksIndex) -> &mut Self::Output { - &mut self.blocks[index.block_index()] - [index.index_in_block()] + &mut self.blocks[index.block_index()][index.index_in_block()] } } diff --git a/datafusion/expr-common/src/blocked_helpers/mod.rs b/datafusion/expr-common/src/blocked_helpers/mod.rs index e7389dcfd32c6..2c1ee4414c333 100644 --- a/datafusion/expr-common/src/blocked_helpers/mod.rs +++ b/datafusion/expr-common/src/blocked_helpers/mod.rs @@ -1,8 +1,5 @@ - -mod blocked_vec; mod blocked_boolean_buffer_builder; +mod blocked_vec; -pub use blocked_vec::{ - BlockedVec, -}; pub use blocked_boolean_buffer_builder::BlockedBooleanBufferBuilder; +pub use blocked_vec::BlockedVec; diff --git a/datafusion/expr-common/src/lib.rs b/datafusion/expr-common/src/lib.rs index 6d9d470f39a36..995425bc7cbf7 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 blocked_helpers; pub mod casts; pub mod columnar_value; @@ -46,6 +47,5 @@ pub mod signature; pub mod sort_properties; pub mod statistics; pub mod type_coercion; -pub mod blocked_groups_accumulator; pub use placement::ExpressionPlacement; diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs index 160a9faf05cf9..9816fb29434e4 100644 --- a/datafusion/expr/src/function.rs +++ b/datafusion/expr/src/function.rs @@ -24,8 +24,7 @@ use datafusion_common::Result; use std::sync::Arc; pub use datafusion_functions_aggregate_common::accumulator::{ - AccumulatorArgs, AccumulatorFactoryFunction, StateFieldsArgs, - BlockedAccumulatorArgs, + AccumulatorArgs, AccumulatorFactoryFunction, BlockedAccumulatorArgs, StateFieldsArgs, }; use crate::expr::{AggregateFunction, WindowFunction}; diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 6926d49c43057..29c5d79aa5c66 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -102,14 +102,14 @@ pub use datafusion_doc::{ scalar_doc_sections, window_doc_sections, }; pub use datafusion_expr_common::accumulator::Accumulator; +pub use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex, +}; +pub use datafusion_expr_common::blocked_helpers; pub use datafusion_expr_common::columnar_value::ColumnarValue; pub use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, }; -pub use datafusion_expr_common::blocked_groups_accumulator::{ - BlockedEmitTo, BlocksIndex, BlockedGroupSelection, BlockedGroupsAccumulator, -}; -pub use datafusion_expr_common::blocked_helpers; pub use datafusion_expr_common::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; @@ -158,7 +158,6 @@ pub use udf::{ pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl}; pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits}; - #[cfg(test)] #[ctor::ctor(unsafe)] fn init() { diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index cb7045fced32f..a6e6cd00d0996 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -38,9 +38,10 @@ use crate::expr::{ schema_name_from_sorts, }; use crate::function::{ - AccumulatorArgs, BlockedAccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs, + AccumulatorArgs, AggregateFunctionSimplification, BlockedAccumulatorArgs, + StateFieldsArgs, }; -use crate::groups_accumulator::{GroupsAccumulator}; +use crate::groups_accumulator::GroupsAccumulator; use crate::udf_eq::UdfEq; use crate::utils::AggregateOrderSensitivity; use crate::utils::format_state_name; @@ -273,7 +274,10 @@ impl AggregateUDF { } /// See [`AggregateUDFImpl::blocked_groups_accumulator_supported`] for more details. - pub fn blocked_groups_accumulator_supported(&self, args: BlockedAccumulatorArgs) -> bool { + pub fn blocked_groups_accumulator_supported( + &self, + args: BlockedAccumulatorArgs, + ) -> bool { self.inner.blocked_groups_accumulator_supported(args) } @@ -694,7 +698,10 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { /// [`Self::accumulator`] for certain queries, such as when this aggregate is /// used as a window function or when there no GROUP BY columns in the /// query. - fn blocked_groups_accumulator_supported(&self, _args: BlockedAccumulatorArgs) -> bool { + fn blocked_groups_accumulator_supported( + &self, + _args: BlockedAccumulatorArgs, + ) -> bool { false } @@ -1663,6 +1670,17 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl { self.inner.create_groups_accumulator(args) } + fn blocked_groups_accumulator_supported(&self, args: BlockedAccumulatorArgs) -> bool { + self.inner.blocked_groups_accumulator_supported(args) + } + + fn create_blocked_groups_accumulator( + &self, + args: BlockedAccumulatorArgs, + ) -> Result> { + self.inner.create_blocked_groups_accumulator(args) + } + fn create_sliding_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/functions-aggregate-common/src/accumulator.rs b/datafusion/functions-aggregate-common/src/accumulator.rs index 9f3eb58589fe6..2ec749e59f77f 100644 --- a/datafusion/functions-aggregate-common/src/accumulator.rs +++ b/datafusion/functions-aggregate-common/src/accumulator.rs @@ -15,12 +15,12 @@ // specific language governing permissions and limitations // under the License. -use std::ops::{Deref, DerefMut}; use arrow::datatypes::{DataType, FieldRef, Schema}; use datafusion_common::Result; use datafusion_expr_common::accumulator::Accumulator; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use std::ops::{Deref, DerefMut}; use std::sync::Arc; /// [`AccumulatorArgs`] contains information about how an aggregate @@ -105,10 +105,13 @@ impl DerefMut for BlockedAccumulatorArgs<'_> { impl<'a> BlockedAccumulatorArgs<'a> { /// Returns the return type of the aggregate function. - pub fn from_accumulator_args(accumulator_args: AccumulatorArgs<'a>, batch_size: usize) -> Self { + pub fn from_accumulator_args( + accumulator_args: AccumulatorArgs<'a>, + batch_size: usize, + ) -> Self { Self { accumulator_args, - batch_size + batch_size, } } } diff --git a/datafusion/functions-aggregate-common/src/aggregate.rs b/datafusion/functions-aggregate-common/src/aggregate.rs index e22bd1d524f9b..5c5e94e431309 100644 --- a/datafusion/functions-aggregate-common/src/aggregate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate.rs @@ -16,7 +16,7 @@ // under the License. pub mod avg_distinct; +pub mod blocked_groups_accumulator; pub mod count_distinct; pub mod groups_accumulator; pub mod sum_distinct; -pub mod blocked_groups_accumulator; diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs index 92446c068d2b8..7769e2e97ba01 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs @@ -18,14 +18,14 @@ //! Utilities for implementing GroupsAccumulator //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] -pub mod blocked_prim_op; pub mod accumulate; +pub mod blocked_prim_op; -use arrow::{ - array::{ArrayRef, BooleanArray}, +use arrow::array::{ArrayRef, BooleanArray}; +use datafusion_common::Result; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex, }; -use datafusion_common::{Result}; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex}; use datafusion_expr_common::groups_accumulator::{GroupSelection, GroupsAccumulator}; /// An adapter that implements [`GroupsAccumulator`] for any [`Accumulator`] @@ -102,7 +102,8 @@ impl BlockedGroupsAccumulatorAdapter { self.number_of_groups = 0; } BlockedEmitTo::NextBlock => { - self.number_of_groups = self.number_of_groups.saturating_sub(self.batch_size); + self.number_of_groups = + self.number_of_groups.saturating_sub(self.batch_size); } BlockedEmitTo::First(n) => { self.number_of_groups = self.number_of_groups.saturating_sub(n); @@ -116,10 +117,24 @@ impl BlockedGroupsAccumulator for BlockedGroupsAccumulatorAdapter { self.batch_size } - fn update_batch(&mut self, values: &[ArrayRef], group_indices: &[BlocksIndex], opt_filter: Option<&BooleanArray>, total_num_groups: usize) -> Result<()> { - let group_indices_flatten = group_indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + let group_indices_flatten = group_indices + .iter() + .map(|index| index.into_index_in_fixed_block_size(self.batch_size)) + .collect::>(); self.number_of_groups = total_num_groups; - self.inner.update_batch(values, &group_indices_flatten, opt_filter, total_num_groups) + self.inner.update_batch( + values, + &group_indices_flatten, + opt_filter, + total_num_groups, + ) } fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { @@ -140,7 +155,10 @@ impl BlockedGroupsAccumulator for BlockedGroupsAccumulatorAdapter { self.inner.supports_evaluate_preserving() } - fn evaluate_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result { + fn evaluate_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result { match selection.indices() { None => { let selection = GroupSelection::all(selection.total_num_groups()); @@ -148,9 +166,15 @@ impl BlockedGroupsAccumulator for BlockedGroupsAccumulatorAdapter { self.inner.evaluate_preserving(selection) } Some(indices) => { - let indices_flatten = indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + let indices_flatten = indices + .iter() + .map(|index| index.into_index_in_fixed_block_size(self.batch_size)) + .collect::>(); - let selection = GroupSelection::try_from_indices(&indices_flatten, selection.total_num_groups())?; + let selection = GroupSelection::try_from_indices( + &indices_flatten, + selection.total_num_groups(), + )?; self.inner.evaluate_preserving(selection) } @@ -175,7 +199,10 @@ impl BlockedGroupsAccumulator for BlockedGroupsAccumulatorAdapter { self.inner.supports_state_preserving() } - fn state_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result> { + fn state_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result> { match selection.indices() { None => { let selection = GroupSelection::all(selection.total_num_groups()); @@ -183,22 +210,41 @@ impl BlockedGroupsAccumulator for BlockedGroupsAccumulatorAdapter { self.inner.state_preserving(selection) } Some(indices) => { - let indices_flatten = indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + let indices_flatten = indices + .iter() + .map(|index| index.into_index_in_fixed_block_size(self.batch_size)) + .collect::>(); - let selection = GroupSelection::try_from_indices(&indices_flatten, selection.total_num_groups())?; + let selection = GroupSelection::try_from_indices( + &indices_flatten, + selection.total_num_groups(), + )?; self.inner.state_preserving(selection) } } } - fn merge_batch(&mut self, values: &[ArrayRef], group_indices: &[BlocksIndex], total_num_groups: usize) -> Result<()> { - let group_indices_flatten = group_indices.iter().map(|index| index.into_index_in_fixed_block_size(self.batch_size)).collect::>(); + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[BlocksIndex], + total_num_groups: usize, + ) -> Result<()> { + let group_indices_flatten = group_indices + .iter() + .map(|index| index.into_index_in_fixed_block_size(self.batch_size)) + .collect::>(); self.number_of_groups = total_num_groups; - self.inner.merge_batch(values, &group_indices_flatten, total_num_groups) + self.inner + .merge_batch(values, &group_indices_flatten, total_num_groups) } - fn convert_to_state(&self, values: &[ArrayRef], opt_filter: Option<&BooleanArray>) -> Result> { + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { self.inner.convert_to_state(values, opt_filter) } diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs index 5b5dbd1026381..4aff483a39610 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs @@ -23,10 +23,12 @@ use arrow::array::{Array, BooleanArray, BooleanBufferBuilder, PrimitiveArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::ArrowPrimitiveType; -use datafusion_common::Result; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlocksIndex}; -use datafusion_expr_common::blocked_helpers::{BlockedBooleanBufferBuilder, BlockedVec}; use crate::aggregate::groups_accumulator::accumulate::accumulate; +use datafusion_common::Result; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, BlocksIndex, +}; +use datafusion_expr_common::blocked_helpers::BlockedBooleanBufferBuilder; /// If the input has nulls, then the accumulator must potentially /// handle each input null value specially (e.g. for `SUM` to mark the @@ -48,9 +50,11 @@ pub enum BlockedSeenValues { } impl BlockedSeenValues { - pub fn new(block_size: usize) -> Self { - BlockedSeenValues::All { num_values: 0, block_size } + BlockedSeenValues::All { + num_values: 0, + block_size, + } } pub fn block_size(&self) -> usize { @@ -67,9 +71,15 @@ impl BlockedSeenValues { /// /// The builder is then ensured to have at least `total_num_groups` length, /// with any new entries initialized to false. - fn get_builder(&mut self, total_num_groups: usize) -> &mut BlockedBooleanBufferBuilder { + fn get_builder( + &mut self, + total_num_groups: usize, + ) -> &mut BlockedBooleanBufferBuilder { match self { - BlockedSeenValues::All { num_values, block_size } => { + BlockedSeenValues::All { + num_values, + block_size, + } => { let mut builder = BlockedBooleanBufferBuilder::new(*block_size); builder.push_value_n(true, *num_values); if total_num_groups > *num_values { @@ -102,8 +112,10 @@ fn new_groups_are_dense( first_new_group: usize, total_num_groups: usize, ) -> bool { - let first_new_group = BlocksIndex::from_index_in_fixed_block_size(first_new_group, block_size); - let total_num_groups = BlocksIndex::from_index_in_fixed_block_size(total_num_groups, block_size); + let first_new_group = + BlocksIndex::from_index_in_fixed_block_size(first_new_group, block_size); + let total_num_groups = + BlocksIndex::from_index_in_fixed_block_size(total_num_groups, block_size); if first_new_group == total_num_groups { return true; @@ -159,7 +171,10 @@ pub struct BlockedNullState { impl BlockedNullState { pub fn new(block_size: usize) -> Self { Self { - seen_values: BlockedSeenValues::All { num_values: 0, block_size }, + seen_values: BlockedSeenValues::All { + num_values: 0, + block_size, + }, } } @@ -167,9 +182,7 @@ impl BlockedNullState { pub fn size(&self) -> usize { match &self.seen_values { BlockedSeenValues::All { .. } => 0, - BlockedSeenValues::Some { values } => { - values.allocated_size() - }, + BlockedSeenValues::Some { values } => values.allocated_size(), } } @@ -205,8 +218,16 @@ impl BlockedNullState { // sparse group indices despite not passing a filter to the accumulator. if opt_filter.is_none() && values.null_count() == 0 - && let BlockedSeenValues::All { num_values, block_size } = &mut self.seen_values - && new_groups_are_dense(group_indices, *block_size, *num_values, total_num_groups) + && let BlockedSeenValues::All { + num_values, + block_size, + } = &mut self.seen_values + && new_groups_are_dense( + group_indices, + *block_size, + *num_values, + total_num_groups, + ) { accumulate(group_indices, values, None, value_fn); *num_values = total_num_groups; @@ -248,8 +269,16 @@ impl BlockedNullState { // sparse group indices despite not passing a filter to the accumulator. if opt_filter.is_none() && values.null_count() == 0 - && let BlockedSeenValues::All { num_values, block_size } = &mut self.seen_values - && new_groups_are_dense(group_indices, *block_size, *num_values, total_num_groups) + && let BlockedSeenValues::All { + num_values, + block_size, + } = &mut self.seen_values + && new_groups_are_dense( + group_indices, + *block_size, + *num_values, + total_num_groups, + ) { group_indices .iter() @@ -330,7 +359,10 @@ impl BlockedNullState { ) -> Result> { let selected_len = selection.len(); match &self.seen_values { - BlockedSeenValues::All { num_values, block_size: _ } => { + BlockedSeenValues::All { + num_values, + block_size: _, + } => { selection.validate_num_groups(*num_values)?; Ok(None) } @@ -352,40 +384,45 @@ impl BlockedNullState { /// resets the internal state appropriately pub fn build(&mut self, emit_to: BlockedEmitTo) -> Vec> { let block_size = self.seen_values.block_size(); - let old_seen = std::mem::replace(&mut self.seen_values, BlockedSeenValues::new(block_size)); + let old_seen = + std::mem::replace(&mut self.seen_values, BlockedSeenValues::new(block_size)); match old_seen { - BlockedSeenValues::All { num_values, block_size } => { - match emit_to { - BlockedEmitTo::All => vec![None; num_values.div_ceil(block_size)], - BlockedEmitTo::NextBlock => { - self.seen_values = BlockedSeenValues::All { num_values: num_values.saturating_sub(block_size), block_size }; + BlockedSeenValues::All { + num_values, + block_size, + } => match emit_to { + BlockedEmitTo::All => vec![None; num_values.div_ceil(block_size)], + BlockedEmitTo::NextBlock => { + self.seen_values = BlockedSeenValues::All { + num_values: num_values.saturating_sub(block_size), + block_size, + }; - if num_values == 0 { - vec![] - } else { - vec![None] - } - } - BlockedEmitTo::First(n) => { - assert!(n < block_size, "emit_to First(n) must be less than block_size"); - self.seen_values = BlockedSeenValues::All { num_values: num_values.saturating_sub(n), block_size }; - if num_values == 0 { - vec![] - } else { - vec![None] - } - } + if num_values == 0 { vec![] } else { vec![None] } } - } + BlockedEmitTo::First(n) => { + assert!( + n < block_size, + "emit_to First(n) must be less than block_size" + ); + self.seen_values = BlockedSeenValues::All { + num_values: num_values.saturating_sub(n), + block_size, + }; + if num_values == 0 { vec![] } else { vec![None] } + } + }, BlockedSeenValues::Some { mut values } => { let emitted = values.emit(emit_to); self.seen_values = BlockedSeenValues::Some { values }; - emitted.into_iter().map(|block| Some(NullBuffer::new(block))).collect() + emitted + .into_iter() + .map(|block| Some(NullBuffer::new(block))) + .collect() } } } } - diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs index 1bd88415287a7..3b9d08aabfa2b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::mem::size_of; use std::sync::Arc; use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray}; @@ -24,11 +23,10 @@ use arrow::compute; use arrow::datatypes::ArrowPrimitiveType; use arrow::datatypes::DataType; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex}; -use datafusion_expr_common::blocked_helpers::BlockedVec; -use datafusion_expr_common::groups_accumulator::{ - EmitTo, GroupSelection, GroupsAccumulator, +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex, }; +use datafusion_expr_common::blocked_helpers::BlockedVec; use super::accumulate::BlockedNullState; @@ -105,7 +103,8 @@ where let values = values[0].as_primitive::(); // update values - self.values.push_value_n_to_len(self.starting_value, total_num_groups); + self.values + .push_value_n_to_len(self.starting_value, total_num_groups); // NullState dispatches / handles tracking nulls and groups that saw no values self.null_state.accumulate( @@ -127,16 +126,23 @@ where let values = self.values.emit(emit_to); let nulls = self.null_state.build(emit_to); - let blocks = values.into_iter().zip(nulls).map(|(values, nulls)| { - let values = PrimitiveArray::::new(values.into(), nulls) // no copy - .with_data_type(self.data_type.clone()); - Arc::new(values) as ArrayRef - }).collect::>(); + let blocks = values + .into_iter() + .zip(nulls) + .map(|(values, nulls)| { + let values = PrimitiveArray::::new(values.into(), nulls) // no copy + .with_data_type(self.data_type.clone()); + Arc::new(values) as ArrayRef + }) + .collect::>(); Ok(blocks) } - fn evaluate_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result { + fn evaluate_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result { selection.validate_num_groups(self.values.len())?; let values: Vec = selection.iter().map(|index| self.values[index]).collect(); @@ -153,10 +159,11 @@ where fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { let blocks = self.evaluate(emit_to)?; - Ok(blocks.into_iter() - // Convert each block into a Vec to match the return type - .map(|arr| vec![arr]) - .collect::>()) + Ok(blocks + .into_iter() + // Convert each block into a Vec to match the return type + .map(|arr| vec![arr]) + .collect::>()) } fn state_preserving( diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index c13bd7de2aaf3..d714d734f09fe 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -34,7 +34,16 @@ use datafusion_common::{ HashMap, Result, ScalarValue, downcast_value, exec_err, internal_err, not_impl_err, stats::Precision, utils::expr::COUNT_STAR_EXPANSION, }; -use datafusion_expr::{expr::WindowFunction, function::{AccumulatorArgs, StateFieldsArgs}, utils::{AggregateOrderSensitivity, format_state_name}, Accumulator, AggregateUDFImpl, BlockedGroupsAccumulator, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility, WindowFunctionDefinition, BlocksIndex, BlockedEmitTo, BlockedGroupSelection}; +use datafusion_expr::blocked_helpers::BlockedVec; +use datafusion_expr::{ + 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::Bitmap65536DistinctCountAccumulator, @@ -58,7 +67,7 @@ use std::{ ops::BitAnd, sync::Arc, }; -use datafusion_expr::blocked_helpers::{BlockedVec, CopyItemBlockedVecBuilder}; +use datafusion_functions_aggregate_common::accumulator::BlockedAccumulatorArgs; make_udaf_expr_and_func!( Count, @@ -386,6 +395,14 @@ impl AggregateUDFImpl for Count { create_distinct_count_groups_accumulator(&args) } + fn blocked_groups_accumulator_supported(&self, args: BlockedAccumulatorArgs) -> bool { + args.expr_fields.len() == 1 && !args.is_distinct + } + + fn create_blocked_groups_accumulator(&self, args: BlockedAccumulatorArgs) -> Result> { + Ok(Box::new(CountBlockedGroupsAccumulator::new(args.batch_size))) + } + fn reverse_expr(&self) -> ReversedUDAF { ReversedUDAF::Identical } @@ -725,9 +742,9 @@ impl GroupsAccumulator for CountGroupsAccumulator { fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { selection.validate_num_groups(self.counts.len())?; let counts = selection - .iter() - .map(|index| self.counts[index]) - .collect::>(); + .iter() + .map(|index| self.counts[index]) + .collect::>(); Ok(Arc::new(Int64Array::from(counts))) } @@ -837,12 +854,13 @@ struct CountBlockedGroupsAccumulator { impl CountBlockedGroupsAccumulator { pub fn new(block_size: usize) -> Self { - Self { counts: BlockedVec::new(block_size) } + Self { + counts: BlockedVec::new(block_size), + } } } impl BlockedGroupsAccumulator for CountBlockedGroupsAccumulator { - fn batch_size(&self) -> usize { self.counts.block_size() } @@ -902,23 +920,26 @@ impl BlockedGroupsAccumulator for CountBlockedGroupsAccumulator { let blocks = self.counts.emit(emit_to); let blocks_ready = blocks - .into_iter() - .map(|block| { - // Count is always non null (null inputs just don't contribute to the overall values) - let array = PrimitiveArray::::new(block.into(), None); - Arc::new(array) as ArrayRef - }) - .collect::>(); + .into_iter() + .map(|block| { + // Count is always non null (null inputs just don't contribute to the overall values) + let array = PrimitiveArray::::new(block.into(), None); + Arc::new(array) as ArrayRef + }) + .collect::>(); Ok(blocks_ready) } - fn evaluate_preserving(&mut self, selection: BlockedGroupSelection<'_>) -> Result { + fn evaluate_preserving( + &mut self, + selection: BlockedGroupSelection<'_>, + ) -> Result { selection.validate_num_groups(self.counts.len())?; let counts = selection - .iter() - .map(|index| self.counts[index]) - .collect::>(); + .iter() + .map(|index| self.counts[index]) + .collect::>(); Ok(Arc::new(Int64Array::from(counts))) } @@ -931,15 +952,15 @@ impl BlockedGroupsAccumulator for CountBlockedGroupsAccumulator { let blocks = self.counts.emit(emit_to); let blocks_ready = blocks - .into_iter() - .map(|block| { - // Count is always non null (null inputs just don't contribute to the overall values) - let array = PrimitiveArray::::new(block.into(), None); - - // Each state only have 1 column - vec![Arc::new(array) as ArrayRef] - }) - .collect::>(); + .into_iter() + .map(|block| { + // Count is always non null (null inputs just don't contribute to the overall values) + let array = PrimitiveArray::::new(block.into(), None); + + // Each state only have 1 column + vec![Arc::new(array) as ArrayRef] + }) + .collect::>(); Ok(blocks_ready) } diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index 2f1c4a22133d2..2cd7ee036575e 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -52,12 +52,16 @@ use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, NullTreatment, physical_name, }; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; -use datafusion_expr::{AggregateUDF, AggregateUDFImpl, Expr, ReversedUDAF, SetMonotonicity}; +use datafusion_expr::{ + AggregateUDF, Expr, ReversedUDAF, SetMonotonicity, +}; use datafusion_expr_common::accumulator::Accumulator; use datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator; -use datafusion_expr_common::groups_accumulator::{GroupsAccumulator}; +use datafusion_expr_common::groups_accumulator::GroupsAccumulator; use datafusion_expr_common::type_coercion::aggregates::check_arg_count; -use datafusion_functions_aggregate_common::accumulator::{AccumulatorArgs, BlockedAccumulatorArgs, StateFieldsArgs}; +use datafusion_functions_aggregate_common::accumulator::{ + AccumulatorArgs, BlockedAccumulatorArgs, StateFieldsArgs, +}; use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -318,7 +322,7 @@ impl AggregateExprBuilder { input_fields: input_exprs_fields, is_reversed, is_nullable, - batch_size + batch_size, }) } @@ -562,7 +566,12 @@ impl<'a> LoweredAggregateBuilder<'a> { .output_metadata(output_metadata) .with_ignore_nulls(ignore_nulls) .with_distinct(*distinct) - .with_batch_size(execution_props.config_options.as_ref().map(|x| x.execution.batch_size.get())); + .with_batch_size( + execution_props + .config_options + .as_ref() + .map(|x| x.execution.batch_size.get()), + ); if let Some(human_display) = human_display { builder = builder.human_display(human_display.expression); @@ -948,7 +957,7 @@ impl AggregateFunctionExpr { /// `[Self::create_groups_accumulator`] will be called. pub fn blocked_groups_accumulator_supported(&self) -> bool { let Some(batch_size) = self.batch_size else { - return false + return false; }; let args = AccumulatorArgs { return_field: Arc::clone(&self.return_field), @@ -970,9 +979,13 @@ impl AggregateFunctionExpr { /// /// For maximum performance, a [`GroupsAccumulator`] should be /// implemented in addition to [`Accumulator`]. - pub fn create_blocked_groups_accumulator(&self) -> Result> { + pub fn create_blocked_groups_accumulator( + &self, + ) -> Result> { let Some(batch_size) = self.batch_size else { - return internal_err!("batch size must be set if blocked group accumulator is supported"); + return internal_err!( + "batch size must be set if blocked group accumulator is supported" + ); }; let args = AccumulatorArgs { return_field: Arc::clone(&self.return_field), 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 916eccd0ada9e..0102e091b5a9b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -28,13 +28,18 @@ 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::{EmitTo, GroupsAccumulator}; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex}; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, +}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; -use crate::aggregates::group_values::{new_blocked_group_values, new_group_values, AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, BlockedGroupValues}; -use crate::aggregates::grouped_hash_stream::{create_blocked_group_accumulator, create_group_accumulator}; +use crate::aggregates::group_values::{ + AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + BlockedGroupValues, GroupByMetrics, new_blocked_group_values, +}; +use crate::aggregates::grouped_hash_stream:: +create_blocked_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, @@ -147,7 +152,8 @@ impl AggregateHashTable { .collect::>()?; let group_schema = agg.group_by.group_schema(&input_schema)?; - let group_values = new_blocked_group_values(group_schema, &GroupOrdering::None, batch_size)?; + let group_values = + new_blocked_group_values(group_schema, &GroupOrdering::None, batch_size)?; let metrics = AggregateTableMetrics::new(agg, partition); @@ -263,49 +269,56 @@ impl AggregateHashTable { let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); let mut output = - match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { - AggregateHashTableState::Outputting(mut state) => { - if state.group_values.is_empty() { - return Ok(None); - } - - // Accumulator output consumes internal state. Materialize all - // groups once, then slice the materialized batch on later polls. - let blocked_columns = self.group_by_metrics.time_emitting(|| { - let mut blocked_columns = state.group_values.emit(BlockedEmitTo::All)?; - for (idx, acc) in state.accumulators.iter_mut().enumerate() { - let acc_cols = accumulator_metrics.time( - idx, - accumulator_phase, - || materialize_accumulator_fn(acc, BlockedEmitTo::All), - )?; - - blocked_columns.iter_mut().zip(acc_cols).for_each(|(output, acc_cols)| { - output.extend_from_slice(acc_cols.as_ref()); - }); - } - Ok::<_, datafusion_common::DataFusionError>(blocked_columns) - })?; - - let mut total_memory = 0; - - let batches = blocked_columns.into_iter().map(|cols| { - let batch = RecordBatch::try_new(output_schema.clone(), cols)?; - debug_assert!(batch.num_rows() > 0); - total_memory += batch.get_array_memory_size(); - Ok(batch) - }).collect::>(); - - MaterializedAggregateOutput::new(batches, total_memory) - } - AggregateHashTableState::OutputtingMaterialized(output) => output, - AggregateHashTableState::Done => return Ok(None), - AggregateHashTableState::Building(_) => { - return internal_err!( + match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { + AggregateHashTableState::Outputting(mut state) => { + if state.group_values.is_empty() { + return Ok(None); + } + + // Accumulator output consumes internal state. Materialize all + // groups once, then slice the materialized batch on later polls. + let blocked_columns = self.group_by_metrics.time_emitting(|| { + let mut blocked_columns = + state.group_values.emit(BlockedEmitTo::All)?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + let acc_cols = accumulator_metrics.time( + idx, + accumulator_phase, + || materialize_accumulator_fn(acc, BlockedEmitTo::All), + )?; + + blocked_columns.iter_mut().zip(acc_cols).for_each( + |(output, acc_cols)| { + output.extend_from_slice(acc_cols.as_ref()); + }, + ); + } + Ok::<_, datafusion_common::DataFusionError>(blocked_columns) + })?; + + let mut total_memory = 0; + + let batches = blocked_columns + .into_iter() + .map(|cols| { + let batch = + RecordBatch::try_new(output_schema.clone(), cols)?; + debug_assert!(batch.num_rows() > 0); + total_memory += batch.get_array_memory_size(); + Ok(batch) + }) + .collect::>(); + + MaterializedAggregateOutput::new(batches, total_memory) + } + AggregateHashTableState::OutputtingMaterialized(output) => output, + AggregateHashTableState::Done => return Ok(None), + AggregateHashTableState::Building(_) => { + return internal_err!( "next_output_batch must be called in the outputting state" ); - } - }; + } + }; let batch = output.next_batch().transpose(); if output.is_exhausted() { @@ -361,29 +374,31 @@ impl AggregateHashTable { let output = group_by_metrics.time_emitting(|| { let mut blocked_output = state.group_values.emit(BlockedEmitTo::All)?; - for (idx, acc) in state.accumulators.iter_mut().enumerate() { - let blocked_state = accumulator_metrics.time( - idx, - AccumulatorPhase::State, - || acc.state(BlockedEmitTo::All), - )?; - - blocked_output.iter_mut().zip(blocked_state).for_each(|(output, states)| { - output.extend_from_slice(states.as_ref()); - }); + let blocked_state = + accumulator_metrics.time(idx, AccumulatorPhase::State, || { + acc.state(BlockedEmitTo::All) + })?; + + blocked_output.iter_mut().zip(blocked_state).for_each( + |(output, states)| { + output.extend_from_slice(states.as_ref()); + }, + ); } Ok::<_, datafusion_common::DataFusionError>(blocked_output) })?; - let batches = output.into_iter().map(|cols| { - let batch = RecordBatch::try_new(Arc::clone(&state_schema), cols)?; - debug_assert!(batch.num_rows() > 0); - - Ok(batch) - }).collect::>>()?; + let batches = output + .into_iter() + .map(|cols| { + let batch = RecordBatch::try_new(Arc::clone(&state_schema), cols)?; + debug_assert!(batch.num_rows() > 0); + Ok(batch) + }) + .collect::>>()?; // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released @@ -650,7 +665,7 @@ impl MaterializedAggregateOutput { self.total_memory -= batch.get_array_memory_size(); Some(Ok(batch)) } - res => res + res => res, } } @@ -718,7 +733,8 @@ impl HashAggregateAccumulator { /// state buffers (empty [`BlockedGroupsAccumulator`]). pub(super) fn empty_like(&self) -> Result { let batch_size = self.accumulator.batch_size(); - let accumulator = create_blocked_group_accumulator(&self.aggregate_expr, batch_size)?; + let accumulator = + create_blocked_group_accumulator(&self.aggregate_expr, batch_size)?; Ok(Self::new( Arc::clone(&self.aggregate_expr), self.arguments.clone(), @@ -924,7 +940,9 @@ mod tests { #[test] fn compact_group_indices_uses_filter_bitmap() { - let group_indices = (0..10).map(BlocksIndex::new_in_first_block).collect::>(); + let group_indices = (0..10) + .map(BlocksIndex::new_in_first_block) + .collect::>(); let all_true = BooleanArray::from(vec![true; 10]); assert_eq!(compact_group_indices(&group_indices, &all_true), None); @@ -932,7 +950,11 @@ mod tests { BooleanArray::from((0..10).map(|index| index != 4).collect::>()); assert_eq!( compact_group_indices(&group_indices, &high_selectivity), - Some([0, 1, 2, 3, 5, 6, 7, 8, 9].map(BlocksIndex::new_in_first_block).to_vec()) + Some( + [0, 1, 2, 3, 5, 6, 7, 8, 9] + .map(BlocksIndex::new_in_first_block) + .to_vec() + ) ); let with_nulls = @@ -1029,13 +1051,12 @@ mod tests { .alias("SUM(value)") .build()?, ); - let accumulator = create_group_accumulator(&aggregate_expr)?; + let accumulator = create_blocked_group_accumulator(&aggregate_expr)?; Ok(HashAggregateAccumulator::new( aggregate_expr, vec![argument], Some(Arc::new(Column::new(filter_name, filter_index))), accumulator, - 8192 )) } 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 ee3e3aaf25019..bed4f5a14fb56 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 @@ -21,22 +21,26 @@ use std::marker::PhantomData; use std::sync::Arc; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result}; -use datafusion_common::assert_or_internal_err; -use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::EmitTo; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; use crate::InputOrderMode; use crate::PhysicalExpr; -use crate::aggregates::group_values::{new_blocked_group_values, new_group_values, AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, GroupValues, BlockedGroupValues}; -use crate::aggregates::grouped_hash_stream::{create_blocked_group_accumulator, create_group_accumulator}; +use crate::aggregates::group_values::{ + AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + BlockedGroupValues, GroupByMetrics, GroupValues, new_blocked_group_values, +}; +use crate::aggregates::grouped_hash_stream:: +create_blocked_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, }; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::assert_or_internal_err; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; use super::AggregateTableMetrics; use super::common::{ @@ -193,7 +197,8 @@ impl OrderedAggregateTable { let group_ordering = GroupOrdering::try_new(input_order_mode, batch_size)?; let group_schema = agg.group_by.group_schema(input_schema)?; - let group_values = new_blocked_group_values(group_schema, &group_ordering, batch_size)?; + let group_values = + new_blocked_group_values(group_schema, &group_ordering, batch_size)?; let aggregate_arguments = aggregate_expressions( &agg.aggr_expr, aggregate_mode, @@ -327,28 +332,31 @@ impl OrderedAggregateTable { let output = self.group_by_metrics.time_emitting(|| { let mut blocked_output = self.buffer.group_values.emit(BlockedEmitTo::All)?; - for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { - let blocked_state = accumulator_metrics.time( - idx, - AccumulatorPhase::State, - || acc.state(BlockedEmitTo::All), - )?; + let blocked_state = + accumulator_metrics.time(idx, AccumulatorPhase::State, || { + acc.state(BlockedEmitTo::All) + })?; - blocked_output.iter_mut().zip(blocked_state).for_each(|(output, states)| { - output.extend_from_slice(states.as_ref()); - }); + blocked_output.iter_mut().zip(blocked_state).for_each( + |(output, states)| { + output.extend_from_slice(states.as_ref()); + }, + ); } Ok::<_, DataFusionError>(blocked_output) })?; - let batches = output.into_iter().map(|cols| { - let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), cols)?; - debug_assert!(batch.num_rows() > 0); + let batches = output + .into_iter() + .map(|cols| { + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), cols)?; + debug_assert!(batch.num_rows() > 0); - Ok(batch) - }).collect::>>()?; + Ok(batch) + }) + .collect::>>()?; // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released @@ -460,13 +468,21 @@ impl OrderedAggregateTable { 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)?; - assert_eq!(output.len(), 1, "must only get single group key batch per emit"); + let output = self.buffer.group_values.emit(emit_to)?; + assert_eq!( + output.len(), + 1, + "must only get single group key batch per emit" + ); let mut output = output.into_iter().next().unwrap(); if should_remove_groups { match emit_to { - BlockedEmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), - BlockedEmitTo::NextBlock => self.buffer.group_ordering.remove_groups(self.batch_size), + BlockedEmitTo::First(n) => { + self.buffer.group_ordering.remove_groups(n) + } + BlockedEmitTo::NextBlock => { + self.buffer.group_ordering.remove_groups(self.batch_size) + } // `EmitTo::All` is only used after `input_done`, when all // buffered groups are known complete and the ordering state is // no longer needed. @@ -479,7 +495,7 @@ impl OrderedAggregateTable { idx, accumulator_phase, || { - let mut acc = materialize_accumulator_fn(acc, emit_to)?; + let acc = materialize_accumulator_fn(acc, emit_to)?; assert_eq!(acc.len(), 1, "must emit single batch"); Ok::<_, DataFusionError>(acc.into_iter().next().unwrap()) }, 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 0c091d340436b..8666e008ddac1 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,9 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; -use crate::aggregates::group_values::{new_blocked_group_values, new_group_values, AccumulatorPhase}; +use crate::aggregates::group_values::{ + AccumulatorPhase, new_blocked_group_values, +}; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{AggregateExec, evaluate_group_by}; @@ -78,7 +80,11 @@ impl AggregateHashTable { ) -> Result> { let state = self.state.building(); let group_schema = state.group_by.group_schema(&self.input_schema)?; - let group_values = new_blocked_group_values(group_schema, &GroupOrdering::None, self.batch_size)?; + let group_values = new_blocked_group_values( + group_schema, + &GroupOrdering::None, + self.batch_size, + )?; let accumulators = state .accumulators .iter() diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs index e33d6e367f006..f93dc842e6db1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs @@ -1,15 +1,19 @@ +use crate::aggregates::group_values::new_group_values; +use super::single_group_by::{ + blocked_boolean::BlockedGroupValuesBoolean, + blocked_primitive::BlockedGroupValuesPrimitive, +}; +use crate::aggregates::order::GroupOrdering; use arrow::array::{downcast_primitive, ArrayRef}; -use arrow_schema::{DataType, SchemaRef, TimeUnit}; -use datafusion_common::{assert_eq_or_internal_err, assert_ne_or_internal_err, assert_or_internal_err, not_impl_err, unwrap_or_internal_err}; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlockedGroupSelection, BlocksIndex}; +use arrow_schema::{DataType, SchemaRef}; +use datafusion_common::{ + assert_eq_or_internal_err, assert_ne_or_internal_err, assert_or_internal_err, + not_impl_err, unwrap_or_internal_err, +}; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, BlocksIndex, +}; use datafusion_expr_common::groups_accumulator::EmitTo; -use datafusion_physical_expr_common::binary_map::OutputType; -use crate::aggregates::group_values::{multi_group_by, new_group_values}; -use crate::aggregates::group_values::single_group_by::blocked_boolean::BlockedGroupValuesBoolean; -use crate::aggregates::group_values::single_group_by::boolean::GroupValuesBoolean; -use crate::aggregates::group_values::single_group_by::bytes::GroupValuesBytes; -use crate::aggregates::group_values::single_group_by::bytes_view::GroupValuesBytesView; -use crate::aggregates::order::GroupOrdering; /// Stores the group values during hash aggregation. /// @@ -48,199 +52,232 @@ use crate::aggregates::order::GroupOrdering; /// (usize) which is assigned by instances of this trait. Group ids are /// continuous without gaps, starting from 0. pub trait BlockedGroupValues: Send { - fn block_size(&self) -> usize; - - /// Calculates the group id for each input row of `cols`, assigning new - /// group ids as necessary. - /// - /// When the function returns, `groups` must contain the group id for each - /// row in `cols`. - /// - /// If a row has the same value as a previous row, the same group id is - /// assigned. If a row has a new value, the next available group id is - /// assigned. - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> datafusion_common::Result<()>; - - /// Returns the number of bytes of memory used by this [`BlockedGroupValues`]. - /// - /// May be expensive; check the implementation before calling on hot paths. - fn size(&self) -> usize; - - /// Returns true if this [`BlockedGroupValues`] is empty - fn is_empty(&self) -> bool; - - /// The number of values (distinct group values) stored in this [`BlockedGroupValues`] - fn len(&self) -> usize; - - /// Materializes selected group values without changing the stored values or - /// their group indices. - /// - /// Rows are returned in the order specified by `selection`. An empty - /// selection returns one correctly typed empty array per group-value column. - /// - /// This method requires exclusive access because implementations may mutate - /// internal caches or builders, even though stored values are unchanged. - fn values_preserving( - &mut self, - _selection: BlockedGroupSelection<'_>, - ) -> datafusion_common::Result> { - not_impl_err!("Preserving group values are not implemented") - } - - /// Returns `true` if [`Self::values_preserving`] is implemented. - fn supports_values_preserving(&self) -> bool { - false - } - - /// Emits the group values - fn emit(&mut self, emit_to: BlockedEmitTo) -> datafusion_common::Result>> { - match emit_to { - BlockedEmitTo::All => { - self.emit_all() - } - BlockedEmitTo::NextBlock => { - let len = self.len(); - let block_size = self.block_size(); - - if len == 0 { - return Ok(vec![]); - } - - if len <= block_size { - return self.emit_all(); - } - - let block = self.emit_block()?; - let block = unwrap_or_internal_err!(block); - - // Assert that all arrays length equal block size since length is greater than block size - for arr in &block { - assert_eq_or_internal_err!(arr.len(), block_size); - } + fn block_size(&self) -> usize; + + /// Calculates the group id for each input row of `cols`, assigning new + /// group ids as necessary. + /// + /// When the function returns, `groups` must contain the group id for each + /// row in `cols`. + /// + /// If a row has the same value as a previous row, the same group id is + /// assigned. If a row has a new value, the next available group id is + /// assigned. + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + ) -> datafusion_common::Result<()>; + + /// Returns the number of bytes of memory used by this [`BlockedGroupValues`]. + /// + /// May be expensive; check the implementation before calling on hot paths. + fn size(&self) -> usize; + + /// Returns true if this [`BlockedGroupValues`] is empty + fn is_empty(&self) -> bool; + + /// The number of values (distinct group values) stored in this [`BlockedGroupValues`] + fn len(&self) -> usize; + + /// Materializes selected group values without changing the stored values or + /// their group indices. + /// + /// Rows are returned in the order specified by `selection`. An empty + /// selection returns one correctly typed empty array per group-value column. + /// + /// This method requires exclusive access because implementations may mutate + /// internal caches or builders, even though stored values are unchanged. + fn values_preserving( + &mut self, + _selection: BlockedGroupSelection<'_>, + ) -> datafusion_common::Result> { + not_impl_err!("Preserving group values are not implemented") + } - Ok(vec![block]) - } - BlockedEmitTo::First(n) => { - assert_ne_or_internal_err!(n, 0); - assert_or_internal_err!(n <= self.len(), "n ({n}) must be less than or equal current length ({})", self.len()); - assert_or_internal_err!(n < self.block_size(), "n ({n}) must be less than current block size ({})", self.block_size()); + /// Returns `true` if [`Self::values_preserving`] is implemented. + fn supports_values_preserving(&self) -> bool { + false + } - if n == self.len() { - self.emit_all() - } else { - self.emit_first_n(n).map(|first_n| vec![first_n]) + /// Emits the group values + fn emit( + &mut self, + emit_to: BlockedEmitTo, + ) -> datafusion_common::Result>> { + match emit_to { + BlockedEmitTo::All => self.emit_all(), + BlockedEmitTo::NextBlock => { + let len = self.len(); + let block_size = self.block_size(); + + if len == 0 { + return Ok(vec![]); + } + + if len <= block_size { + return self.emit_all(); + } + + let block = self.emit_block()?; + let block = unwrap_or_internal_err!(block); + + // Assert that all arrays length equal block size since length is greater than block size + for arr in &block { + assert_eq_or_internal_err!(arr.len(), block_size); + } + + Ok(vec![block]) + } + BlockedEmitTo::First(n) => { + assert_ne_or_internal_err!(n, 0); + assert_or_internal_err!( + n <= self.len(), + "n ({n}) must be less than or equal current length ({})", + self.len() + ); + assert_or_internal_err!( + n < self.block_size(), + "n ({n}) must be less than current block size ({})", + self.block_size() + ); + + if n == self.len() { + self.emit_all() + } else { + self.emit_first_n(n).map(|first_n| vec![first_n]) + } + } } - } } - } - /// Emit all group values - fn emit_all(&mut self) -> datafusion_common::Result>>; + /// Emit all group values + fn emit_all(&mut self) -> datafusion_common::Result>>; - /// Emit the next block - /// returns Ok(None) when there are no blocks - fn emit_block(&mut self) -> datafusion_common::Result>>; + /// Emit the next block + /// returns Ok(None) when there are no blocks + fn emit_block(&mut self) -> datafusion_common::Result>>; - /// Emit first `n` values and shift all values to fit into blocks - /// - /// `n` must be smaller than [`Self::block_size`] and larger than `0` - /// `n` must be smaller or equal to [`Self::len`] - fn emit_first_n(&mut self, n: usize) -> datafusion_common::Result>; + /// Emit first `n` values and shift all values to fit into blocks + /// + /// `n` must be smaller than [`Self::block_size`] and larger than `0` + /// `n` must be smaller or equal to [`Self::len`] + fn emit_first_n(&mut self, n: usize) -> datafusion_common::Result>; - // TODO - add into iterator which move the entry state into an iterator that will output ready batches - // this is so it won't need to update the underlying hash map whenever calling emit block - // this is for the case when we need to emit all, but we don't want to materialize all right away - // but we want to skip the internal hash map updates, so the iterator will avoid that while clearing memory - // the iterator should expose `allocated_size()` function + // TODO - add into iterator which move the entry state into an iterator that will output ready batches + // this is so it won't need to update the underlying hash map whenever calling emit block + // this is for the case when we need to emit all, but we don't want to materialize all right away + // but we want to skip the internal hash map updates, so the iterator will avoid that while clearing memory + // the iterator should expose `allocated_size()` function - /// 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); + /// 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); } - // This is just an adapter until all is implemented pub struct BlockedGroupValuesAdapter { - block_size: usize, - inner: Box, + block_size: usize, + inner: Box, } impl BlockedGroupValuesAdapter { - pub fn new(block_size: usize, inner: Box) -> Self { - Self { block_size, inner } - } + pub fn new( + block_size: usize, + inner: Box, + ) -> Self { + Self { block_size, inner } + } } impl BlockedGroupValues for BlockedGroupValuesAdapter { - fn block_size(&self) -> usize { - self.block_size - } - - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> datafusion_common::Result<()> { - let block_size = self.block_size; - let mut group_indices_flattened = groups.iter().map(|i| i.into_index_in_fixed_block_size(block_size)).collect::>(); - self.inner.intern(cols, &mut group_indices_flattened)?; - *groups = group_indices_flattened.iter().map(|index| BlocksIndex::from_index_in_fixed_block_size(*index, block_size)).collect::>(); - - Ok(()) - } - - fn size(&self) -> usize { - self.inner.size() - } - - fn is_empty(&self) -> bool { - self.inner.is_empty() - } + fn block_size(&self) -> usize { + self.block_size + } - fn len(&self) -> usize { - self.inner.len() - } + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + ) -> datafusion_common::Result<()> { + let block_size = self.block_size; + let mut group_indices_flattened = groups + .iter() + .map(|i| i.into_index_in_fixed_block_size(block_size)) + .collect::>(); + self.inner.intern(cols, &mut group_indices_flattened)?; + *groups = group_indices_flattened + .iter() + .map(|index| BlocksIndex::from_index_in_fixed_block_size(*index, block_size)) + .collect::>(); + + Ok(()) + } - fn emit_all(&mut self) -> datafusion_common::Result>> { - let mut blocks = vec![]; + fn size(&self) -> usize { + self.inner.size() + } - while self.len() > self.block_size { - blocks.push(self.inner.emit(EmitTo::First(self.block_size))?); + fn is_empty(&self) -> bool { + self.inner.is_empty() } - if self.len() > 0 { - blocks.push(self.inner.emit(EmitTo::All)?); + fn len(&self) -> usize { + self.inner.len() } - Ok(blocks) - } + fn emit_all(&mut self) -> datafusion_common::Result>> { + let mut blocks = vec![]; + + while self.len() > self.block_size { + blocks.push(self.inner.emit(EmitTo::First(self.block_size))?); + } - fn emit_block(&mut self) -> datafusion_common::Result>> { - if self.len() == 0 { - return Ok(None); + if self.len() > 0 { + blocks.push(self.inner.emit(EmitTo::All)?); + } + + Ok(blocks) } - let output = if self.len() <= self.block_size { - self.inner.emit(EmitTo::All) - } else { - self.inner.emit(EmitTo::First(self.block_size)) - }; + fn emit_block(&mut self) -> datafusion_common::Result>> { + if self.len() == 0 { + return Ok(None); + } - Ok(Some(output?)) - } + let output = if self.len() <= self.block_size { + self.inner.emit(EmitTo::All) + } else { + self.inner.emit(EmitTo::First(self.block_size)) + }; - fn emit_first_n(&mut self, n: usize) -> datafusion_common::Result> { - assert_ne_or_internal_err!(n, 0); - assert_or_internal_err!(n <= self.len(), "n ({n}) must be less than or equal current length ({})", self.len()); - assert_or_internal_err!(n < self.block_size(), "n ({n}) must be less than current block size ({})", self.block_size()); + Ok(Some(output?)) + } - let output = if self.len() == n { - self.inner.emit(EmitTo::All) - } else { - self.inner.emit(EmitTo::First(n)) - }; + fn emit_first_n(&mut self, n: usize) -> datafusion_common::Result> { + assert_ne_or_internal_err!(n, 0); + assert_or_internal_err!( + n <= self.len(), + "n ({n}) must be less than or equal current length ({})", + self.len() + ); + assert_or_internal_err!( + n < self.block_size(), + "n ({n}) must be less than current block size ({})", + self.block_size() + ); + + let output = if self.len() == n { + self.inner.emit(EmitTo::All) + } else { + self.inner.emit(EmitTo::First(n)) + }; - Ok(output?) - } + Ok(output?) + } - fn clear_shrink(&mut self, num_rows: usize) { - self.inner.clear_shrink(num_rows) - } + fn clear_shrink(&mut self, num_rows: usize) { + self.inner.clear_shrink(num_rows) + } } /// Return a specialized implementation of [`BlockedGroupValues`] for the given schema. @@ -260,19 +297,32 @@ impl BlockedGroupValues for BlockedGroupValuesAdapter { /// `GroupValuesColumn`: crate::aggregates_blocked::group_values::multi_group_by::GroupValuesColumn /// `GroupValuesRows`: crate::aggregates_blocked::group_values::GroupValuesRows pub fn new_blocked_group_values( - schema: SchemaRef, - group_ordering: &GroupOrdering, - block_size: usize, + schema: SchemaRef, + group_ordering: &GroupOrdering, + block_size: usize, ) -> datafusion_common::Result> { - if schema.fields.len() == 1 { - let d = schema.fields[0].data_type(); + if schema.fields.len() == 1 { + let d = schema.fields[0].data_type(); + + macro_rules! downcast_helper { + ($t:ty, $d:ident) => { + return Ok(Box::new(BlockedGroupValuesPrimitive::<$t>::new($d.clone(), block_size))) + }; + } - if matches!(d, DataType::Boolean) { - return Ok(Box::new(BlockedGroupValuesBoolean::new(block_size))); + downcast_primitive! { + d => (downcast_helper, d), + DataType::Boolean => { + return Ok(Box::new(BlockedGroupValuesBoolean::new(block_size))); + }, + _ => {} + } } - } - let group_values = new_group_values(schema, group_ordering)?; + let group_values = new_group_values(schema, group_ordering)?; - Ok(Box::new(BlockedGroupValuesAdapter::new(block_size, group_values))) + Ok(Box::new(BlockedGroupValuesAdapter::new( + block_size, + group_values, + ))) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index d11f9683c3074..1483526e13db1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -46,9 +46,9 @@ use crate::aggregates::{ order::GroupOrdering, }; +mod blocked; mod metrics; mod null_builder; -mod blocked; pub use blocked::{BlockedGroupValues, new_blocked_group_values}; diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs index a0b7af8413de2..4e35dd69d5012 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs @@ -17,14 +17,13 @@ use crate::aggregates::group_values::GroupValues; +use crate::aggregates::group_values::blocked::BlockedGroupValues; use arrow::array::{ ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, }; use datafusion_common::Result; -use datafusion_expr::EmitTo; -use std::{mem::size_of, sync::Arc}; use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; -use crate::aggregates::group_values::blocked::BlockedGroupValues; +use std::{mem::size_of, sync::Arc}; #[derive(Debug)] pub struct BlockedGroupValuesBoolean { @@ -86,7 +85,10 @@ impl BlockedGroupValues for BlockedGroupValuesBoolean { } }; - groups.push(BlocksIndex::from_index_in_fixed_block_size(index, self.block_size)); + groups.push(BlocksIndex::from_index_in_fixed_block_size( + index, + self.block_size, + )); } Ok(()) @@ -102,8 +104,8 @@ impl BlockedGroupValues for BlockedGroupValuesBoolean { fn len(&self) -> usize { self.false_group.is_some() as usize - + self.true_group.is_some() as usize - + self.null_group.is_some() as usize + + self.true_group.is_some() as usize + + self.null_group.is_some() as usize } fn emit_all(&mut self) -> Result>> { @@ -166,7 +168,11 @@ impl BlockedGroupValues for BlockedGroupValuesBoolean { } fn emit_first_n(&mut self, n: usize) -> Result> { - assert!(n < self.block_size, "{n} must be less than block_size {}", self.block_size); + assert!( + n < self.block_size, + "{n} must be less than block_size {}", + self.block_size + ); { let len = self.len(); @@ -224,10 +230,13 @@ impl BlockedGroupValues for BlockedGroupValuesBoolean { #[cfg(test)] mod tests { - use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; use super::*; + use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; - fn intern(group_values: &mut BlockedGroupValuesBoolean, values: &[Option]) -> Vec { + fn intern( + group_values: &mut BlockedGroupValuesBoolean, + values: &[Option], + ) -> Vec { let mut groups = vec![]; let array: ArrayRef = Arc::new(BooleanArray::from(values.to_vec())); group_values.intern(&[array], &mut groups).unwrap(); @@ -236,15 +245,18 @@ mod tests { fn emitted(blocks: Vec>) -> Vec>> { blocks - .into_iter() - .map(|block| block[0].as_boolean().iter().collect()) - .collect() + .into_iter() + .map(|block| block[0].as_boolean().iter().collect()) + .collect() } #[test] fn emit_all_only_emits_existing_groups() { let mut group_values = BlockedGroupValuesBoolean::new(531); - let groups = intern(&mut group_values, &[Some(true), None, Some(true), Some(false)]); + let groups = intern( + &mut group_values, + &[Some(true), None, Some(true), Some(false)], + ); assert_eq!( groups, [0, 1, 0, 2].map(BlocksIndex::new_in_first_block).to_vec() diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs index 84d276c5ac6ac..0e4f40ab9d0b2 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs @@ -15,277 +15,282 @@ // specific language governing permissions and limitations // under the License. -use arrow::buffer::ScalarBuffer; +use crate::aggregates::group_values::{BlockedGroupValues, HashValue}; use arrow::array::{ - ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, - cast::AsArray, + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, + cast::AsArray, }; -use arrow::datatypes::{DataType}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; +use datafusion_expr_common::blocked_helpers::BlockedVec; use hashbrown::hash_table::HashTable; use std::mem::size_of; use std::sync::Arc; -use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; -use datafusion_expr_common::blocked_helpers::BlockedVec; -use crate::aggregates::group_values::{BlockedGroupValues, HashValue}; /// A [`GroupValues`] storing a single column of primitive values /// /// This specialization is significantly faster than using the more general /// purpose `Row`s format pub struct BlockedGroupValuesPrimitive { - block_size: usize, - /// The data type of the output array - data_type: DataType, - /// Stores the `(group_index, hash)` based on the hash of its value - /// - /// We also store `hash` is for reducing cost of rehashing. Such cost - /// is obvious in high cardinality group by situation. - /// More details can see: - /// - map: HashTable<(BlocksIndex, u64)>, - /// 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, + block_size: usize, + /// The data type of the output array + data_type: DataType, + /// Stores the `(group_index, hash)` based on the hash of its value + /// + /// We also store `hash` is for reducing cost of rehashing. Such cost + /// is obvious in high cardinality group by situation. + /// More details can see: + /// + map: HashTable<(BlocksIndex, u64)>, + /// 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 fn new(data_type: DataType, block_size: usize) -> Self { - assert!(PrimitiveArray::::is_compatible(&data_type)); - Self { - block_size, - data_type, - map: HashTable::with_capacity(128), - values: BlockedVec::new(block_size), - null_group: None, - random_state: crate::aggregates::AGGREGATION_HASH_SEED, + pub fn new(data_type: DataType, block_size: usize) -> Self { + assert!(PrimitiveArray::::is_compatible(&data_type)); + Self { + block_size, + data_type, + map: HashTable::with_capacity(128), + values: BlockedVec::new(block_size), + null_group: None, + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + } } - } - fn build_primitive( - values: impl Into>, - null_idx: Option, - ) -> PrimitiveArray { - let values = values.into(); - let nulls = null_idx.map(|null_idx| { - let mut buffer = NullBufferBuilder::new(values.len()); - buffer.append_n_non_nulls(null_idx); - buffer.append_null(); - buffer.append_n_non_nulls(values.len() - null_idx - 1); - // NOTE: The inner builder must be constructed as there is at least one null - buffer.finish().unwrap() - }); - PrimitiveArray::::new(values.into(), nulls) - } + fn build_primitive( + values: impl Into>, + null_idx: Option, + ) -> PrimitiveArray { + let values = values.into(); + let nulls = null_idx.map(|null_idx| { + let mut buffer = NullBufferBuilder::new(values.len()); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(values.len() - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }); + PrimitiveArray::::new(values.into(), nulls) + } - fn build_no_nulls_primitive_arc( - values: impl Into>, - data_type: DataType, - ) -> ArrayRef { - Arc::new(PrimitiveArray::::new(values.into(), None).with_data_type(data_type)) - } + fn build_no_nulls_primitive_arc( + values: impl Into>, + data_type: DataType, + ) -> ArrayRef { + Arc::new(PrimitiveArray::::new(values.into(), None).with_data_type(data_type)) + } - fn build_with_nulls_primitive_arc( - values: impl Into>, - null_idx: usize, - data_type: DataType, - ) -> ArrayRef { - let values = values.into(); - let nulls = { - let mut buffer = NullBufferBuilder::new(values.len()); - buffer.append_n_non_nulls(null_idx); - buffer.append_null(); - buffer.append_n_non_nulls(values.len() - null_idx - 1); - // NOTE: The inner builder must be constructed as there is at least one null - buffer.finish().unwrap() - }; - Arc::new( - PrimitiveArray::::new(values.into(), Some(nulls)) - .with_data_type(data_type), - ) - } + fn build_with_nulls_primitive_arc( + values: impl Into>, + null_idx: usize, + data_type: DataType, + ) -> ArrayRef { + let values = values.into(); + let nulls = { + let mut buffer = NullBufferBuilder::new(values.len()); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(values.len() - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }; + Arc::new( + PrimitiveArray::::new(values.into(), Some(nulls)) + .with_data_type(data_type), + ) + } } impl BlockedGroupValues for BlockedGroupValuesPrimitive where - T::Native: HashValue, + T::Native: HashValue, { - fn block_size(&self) -> usize { - self.block_size - } + fn block_size(&self) -> usize { + self.block_size + } - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { - assert_eq!(cols.len(), 1); - groups.clear(); + 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(|| { - let group_id = self.values.len(); - self.values.push(Default::default()); - BlocksIndex::from_index_in_fixed_block_size(group_id, self.block_size) - }), - 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, - |&(g, h)| { - // SAFETY: every group in the map was appended to `values` - hash == h && unsafe { self.values.get_unchecked(g) }.is_eq(key) - }, - |&(_, h)| h, - ); + for v in cols[0].as_primitive::() { + let group_id = match v { + None => *self.null_group.get_or_insert_with(|| { + let group_id = self.values.len(); + self.values.push(Default::default()); + BlocksIndex::from_index_in_fixed_block_size(group_id, self.block_size) + }), + 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, + |&(g, h)| { + // SAFETY: every group in the map was appended to `values` + hash == h + && unsafe { self.values.get_unchecked(g) }.is_eq(key) + }, + |&(_, h)| h, + ); - match insert { - hashbrown::hash_table::Entry::Occupied(o) => o.get().0, - hashbrown::hash_table::Entry::Vacant(v) => { - let g = BlocksIndex::from_index_in_fixed_block_size(self.values.len(), self.block_size); - v.insert((g, hash)); - self.values.push(key); - g - } - } + match insert { + hashbrown::hash_table::Entry::Occupied(o) => o.get().0, + hashbrown::hash_table::Entry::Vacant(v) => { + let g = BlocksIndex::from_index_in_fixed_block_size( + self.values.len(), + self.block_size, + ); + v.insert((g, hash)); + self.values.push(key); + g + } + } + } + }; + groups.push(group_id) } - }; - groups.push(group_id) + Ok(()) } - Ok(()) - } - fn size(&self) -> usize { - self.map.capacity() * size_of::<(BlocksIndex, u64)>() + self.values.allocated_size() - } + fn size(&self) -> usize { + self.map.capacity() * size_of::<(BlocksIndex, u64)>() + + self.values.allocated_size() + } - fn is_empty(&self) -> bool { - self.values.is_empty() - } + fn is_empty(&self) -> bool { + self.values.is_empty() + } - fn len(&self) -> usize { - self.values.len() - } + fn len(&self) -> usize { + self.values.len() + } - fn emit_all(&mut self) -> Result>> { - let mut blocks = Vec::with_capacity(self.values.num_blocks()); - self.map = HashTable::with_capacity(128); + fn emit_all(&mut self) -> Result>> { + self.map = HashTable::with_capacity(128); - let all = self.values.take_all(); - let mut into_iter = all.into_iter(); + let all = self.values.take_all(); + let mut blocks = Vec::with_capacity(all.len()); - if let Some(null_idx) = self.null_group.take() { - for values in into_iter.by_ref().take(null_idx.block_index()) { - blocks.push(vec![Self::build_no_nulls_primitive_arc( - values, - self.data_type.clone(), - )]); - } + let mut into_iter = all.into_iter(); - let block_with_nulls = into_iter.next().expect("must have block for nulls"); - blocks.push(vec![Self::build_with_nulls_primitive_arc( - block_with_nulls, - null_idx.index_in_block(), - self.data_type.clone(), - )]); - } + if let Some(null_idx) = self.null_group.take() { + for values in into_iter.by_ref().take(null_idx.block_index()) { + blocks.push(vec![Self::build_no_nulls_primitive_arc( + values, + self.data_type.clone(), + )]); + } - for values in into_iter { - blocks.push(vec![Self::build_no_nulls_primitive_arc( - values, - self.data_type.clone(), - )]); - } + let block_with_nulls = into_iter.next().expect("must have block for nulls"); + blocks.push(vec![Self::build_with_nulls_primitive_arc( + block_with_nulls, + null_idx.index_in_block(), + self.data_type.clone(), + )]); + } - Ok(blocks) - } + for values in into_iter { + blocks.push(vec![Self::build_no_nulls_primitive_arc( + values, + self.data_type.clone(), + )]); + } - fn emit_block(&mut self) -> Result>> { - let Some(values) = self.values.take_block() else { - return Ok(None); - }; - let block_size = self.block_size; + Ok(blocks) + } - // If nothing left - let null_group = if self.values.len() == 0 { - self.map.clear(); + fn emit_block(&mut self) -> Result>> { + let Some(values) = self.values.take_block() else { + return Ok(None); + }; - assert_eq!(self.null_group.map_or(0, |index| index.block_index()), 0); + // If nothing left + let null_group = if self.values.len() == 0 { + self.map.clear(); - self.null_group.take() - } else { - self.map.retain(|entry| { - // Decrement group index by n - let group_idx = entry.0; - match group_idx.prev_block_checked() { - Some(new_block_index) => { - entry.0 = new_block_index; - true - } - None => false, - } - }); + assert_eq!(self.null_group.map_or(0, |index| index.block_index()), 0); - match &mut self.null_group { - Some(v) if v.block_index() != 0 => { - *v = v.prev_block(); - None - } - Some(_) => self.null_group.take(), - None => None, - } - }; + self.null_group.take() + } else { + self.map.retain(|entry| { + // Decrement group index by n + let group_idx = entry.0; + match group_idx.prev_block_checked() { + Some(new_block_index) => { + entry.0 = new_block_index; + true + } + None => false, + } + }); - let array = Self::build_primitive(values, null_group.map(|i| i.index_in_block())); + match &mut self.null_group { + Some(v) if v.block_index() != 0 => { + *v = v.prev_block(); + None + } + Some(_) => self.null_group.take(), + None => None, + } + }; - Ok(Some(vec![Arc::new( - array.with_data_type(self.data_type.clone()), - )])) - } + let array = Self::build_primitive(values, null_group.map(|i| i.index_in_block())); - fn emit_first_n(&mut self, n: usize) -> Result> { - let block_size = self.block_size; - let array: PrimitiveArray = { - self.map.retain(|entry| { - // Decrement group index by n - let group_idx = entry.0; - match group_idx.sub_flat_checked(n, block_size) { - // Group index was >= n, shift value down - Some(sub) => { - entry.0 = sub; - true - } - // Group index was < n, so remove from table - None => false, - } - }); - let null_group = match &mut self.null_group { - Some(v) if v.gte_flat(n, block_size) => { - *v = v.sub_flat(n, block_size); - None - } - Some(_) => self.null_group.take().map(|g| g.index_in_block()), - None => None, - }; - let first_values = self.values.take_n(n); - Self::build_primitive(first_values, null_group) - }; + Ok(Some(vec![Arc::new( + array.with_data_type(self.data_type.clone()), + )])) + } + + fn emit_first_n(&mut self, n: usize) -> Result> { + let block_size = self.block_size; + let array: PrimitiveArray = { + self.map.retain(|entry| { + // Decrement group index by n + let group_idx = entry.0; + match group_idx.sub_flat_checked(n, block_size) { + // Group index was >= n, shift value down + Some(sub) => { + entry.0 = sub; + true + } + // Group index was < n, so remove from table + None => false, + } + }); + let null_group = match &mut self.null_group { + Some(v) if v.gte_flat(n, block_size) => { + *v = v.sub_flat(n, block_size); + None + } + Some(_) => self.null_group.take().map(|g| g.index_in_block()), + None => None, + }; + let first_values = self.values.take_n(n); + Self::build_primitive(first_values, null_group) + }; - Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) - } + Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) + } - fn clear_shrink(&mut self, num_rows: usize) { - // self.values.clear(); - // self.values.shrink_to(num_rows); - self.values.reset(); - self.map.clear(); - self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared - } + fn clear_shrink(&mut self, num_rows: usize) { + // self.values.clear(); + // self.values.shrink_to(num_rows); + self.values.reset(); + self.map.clear(); + self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs index dc06b42124a5e..28721678e43bc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs @@ -17,9 +17,9 @@ //! `GroupValues` implementations for single group by cases +pub(crate) mod blocked_boolean; +pub(crate) mod blocked_primitive; pub(crate) mod boolean; pub(crate) mod bytes; pub(crate) mod bytes_view; pub(crate) mod primitive; -pub(crate) mod blocked_boolean; -mod blocked_primitive; diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 03a2ffd11d35b..a3400f0725acf 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -73,11 +73,13 @@ use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use crate::sorts::IncrementalSortIterator; use datafusion_common::utils::memory::get_record_batch_memory_size; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedGroupsAccumulator, BlocksIndex, +}; +use datafusion_functions_aggregate_common::aggregate::blocked_groups_accumulator::BlockedGroupsAccumulatorAdapter; use futures::ready; use futures::stream::{Stream, StreamExt}; use log::debug; -use datafusion_expr_common::blocked_groups_accumulator::{BlockedGroupsAccumulator, BlocksIndex}; -use datafusion_functions_aggregate_common::aggregate::blocked_groups_accumulator::BlockedGroupsAccumulatorAdapter; #[derive(Debug, Clone)] /// This object tracks the aggregation phase (input/output) @@ -673,11 +675,15 @@ pub(crate) fn create_group_accumulator( /// [`BlockedGroupsAccumulatorAdapter`] if not. pub(crate) fn create_blocked_group_accumulator( agg_expr: &Arc, - block_size: usize + block_size: usize, ) -> Result> { // TODO - WHAT ABOUT DYNAMIC BLOCK SIZE? if let Some(agg_batch_size) = agg_expr.batch_size() { - assert_eq_or_internal_err!(agg_batch_size, block_size, "Block size mismatch for blocked groups accumulator"); + assert_eq_or_internal_err!( + agg_batch_size, + block_size, + "Block size mismatch for blocked groups accumulator" + ); } if agg_expr.blocked_groups_accumulator_supported() { agg_expr.create_blocked_groups_accumulator() @@ -688,7 +694,9 @@ pub(crate) fn create_blocked_group_accumulator( agg_expr.name() ); let group_acc = create_group_accumulator(agg_expr)?; - Ok(Box::new(BlockedGroupsAccumulatorAdapter::new(group_acc, block_size))) + Ok(Box::new(BlockedGroupsAccumulatorAdapter::new( + group_acc, block_size, + ))) } } @@ -946,7 +954,11 @@ impl GroupedHashAggregateStream { let starting_num_groups = self.group_values.len(); self.group_values .intern(group_values, &mut self.current_group_indices)?; - let group_indices_as_blocks = self.current_group_indices.iter().map(|&idx| BlocksIndex::new_in_first_block(idx)).collect::>(); + let group_indices_as_blocks = self + .current_group_indices + .iter() + .map(|&idx| BlocksIndex::new_in_first_block(idx)) + .collect::>(); let total_num_groups = self.group_values.len(); if total_num_groups > starting_num_groups { self.group_ordering.new_groups( diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index dc555b75d5bbe..87d5a1b1f0ac3 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -26,7 +26,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{ - DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, + DataFusionError, Result, assert_ne_or_internal_err, internal_err, }; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -292,18 +292,21 @@ impl FinalSpillContext { &mut self, hash_table: &mut AggregateHashTable, ) -> Result<()> { - for batch in hash_table.take_all_state_batch()? { + for batch in hash_table.take_all_state_batch()? { // TODO - avoid this - it will create MANY small spill files as oppose to before, // CHANGE TO sort iterator that will sort multiple batches instead of one batches without concat // and will output multiple batches - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let sorted_iter = IncrementalSortIterator::new( + 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, - "FinalHashAggregateSpill", - )?; + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "FinalHashAggregateSpill", + )?; let Some((file, max_record_batch_memory)) = spill_file else { return internal_err!("Final hash aggregation produced an empty spill"); @@ -315,7 +318,6 @@ impl FinalSpillContext { }); } - Ok(()) } @@ -478,7 +480,8 @@ impl PartialHashAggregateStream { break; } HandleInputResult::OOM => { - let materialized_group_states = hash_table.take_all_state_batch()?; + let materialized_group_states = + hash_table.take_all_state_batch()?; if materialized_group_states.is_empty() { return internal_err!( @@ -598,11 +601,14 @@ impl PartialHashAggregateStream { &mut self, // After each incremental emitting step, the `remaining_groups` will be updated // with batch slicing. - mut remaining_groups: Vec, + remaining_groups: Vec, emitter: &mut TryEmitter, hash_table_mem_size: usize, ) -> Result<()> { - let sizes = remaining_groups.iter().map(|b| b.get_array_memory_size()).collect::>(); + let sizes = remaining_groups + .iter() + .map(|b| b.get_array_memory_size()) + .collect::>(); // Skip first since we are emitting it right away let total_size: usize = sizes.iter().skip(1).sum(); @@ -628,15 +634,14 @@ impl PartialHashAggregateStream { // Do not shrink the first batch since we did not count it if i > 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(size)?; } emitter - .emit(output.record_output(&self.baseline_metrics)) - .await; + .emit(output.record_output(&self.baseline_metrics)) + .await; } Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index eeaa4417f59f0..88eec37e41049 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -246,7 +246,8 @@ impl GroupOrderingPartial { let ranges = partition(&sort_keys)?.ranges(); let last_range = ranges.last().unwrap(); - let range_current_sort = group_indices[last_range.start].into_index_in_fixed_block_size(self.block_size); + let range_current_sort = group_indices[last_range.start] + .into_index_in_fixed_block_size(self.block_size); let range_sort_key = get_row_at_idx(&sort_keys, last_range.start)?; let (current_sort, sort_key) = if last_range.start == 0 { @@ -288,7 +289,8 @@ mod tests { // Ordered on column a let order_indices = vec![0]; let block_size = 8192; - let mut group_ordering = GroupOrderingPartial::try_new(order_indices, block_size)?; + let mut group_ordering = + GroupOrderingPartial::try_new(order_indices, block_size)?; let batch_group_values: Vec = vec![ Arc::new(Int32Array::from(vec![1, 2, 3])), diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 4e5a32461b94b..00ceb9e6407ea 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -189,17 +189,22 @@ impl OrderedFinalSpillContext { table: &mut OrderedAggregateTable, ) -> Result<()> { for batch in table.take_all_state_batch()? { - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let sorted_iter = IncrementalSortIterator::new( + 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, - "OrderedFinalAggregateSpill", - )?; + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "OrderedFinalAggregateSpill", + )?; let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Ordered final aggregation produced an empty spill"); + return internal_err!( + "Ordered final aggregation produced an empty spill" + ); }; self.spills.push(SortedSpillFile { diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 308e59d1fc5ff..b43a56eb3f63c 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -19,13 +19,6 @@ use std::sync::Arc; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result}; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; -use futures::stream::{Stream, StreamExt}; -use datafusion_functions_aggregate_common::aggregate::groups_accumulator::VecAllocExt; use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; @@ -33,6 +26,12 @@ use crate::aggregates::order::GroupOrdering; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; +use futures::stream::{Stream, StreamExt}; /// Partial aggregate stream for `InputOrderMode::Sorted` and /// `InputOrderMode::PartiallySorted`. @@ -320,11 +319,14 @@ impl OrderedPartialAggregateStream { let batches = match table.take_all_state_batch() { Ok(batches) if batches.is_empty() => return Err(oom), - Ok(batches) => batches.into_iter().map(|b| { - let size = b.get_array_memory_size(); - total_size += size; - (b, size) - }).collect::>(), + Ok(batches) => batches + .into_iter() + .map(|b| { + let size = b.get_array_memory_size(); + total_size += size; + (b, size) + }) + .collect::>(), Err(e) => return Err(e), }; diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 50c0609b8ceac..25584f87db1ff 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -249,17 +249,22 @@ impl OrderedSingleSpillContext { table: &mut OrderedAggregateTable, ) -> Result<()> { for batch in table.take_all_state_batch()? { - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let sorted_iter = IncrementalSortIterator::new( + 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, - "OrderedSingleAggregateSpill", - )?; + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "OrderedSingleAggregateSpill", + )?; let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Ordered single aggregation produced an empty spill"); + return internal_err!( + "Ordered single aggregation produced an empty spill" + ); }; self.spills.push(SortedSpillFile { diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index e242b10f8f0d7..48d0abc0f4973 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -321,7 +321,9 @@ impl PartialReduceHashAggregateStream { match state_batch_result { // 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) if remaining_groups.is_empty() => { + Self::break_with_err(oom) + } Ok(remaining_groups) => { self.early_emit_count.add(1); @@ -355,7 +357,9 @@ impl PartialReduceHashAggregateStream { unreachable!("expected the EmittingOnMemoryPressure state") }; - let next_batch = batches.next().expect("must have at least one batch to emit"); + let next_batch = batches + .next() + .expect("must have at least one batch to emit"); let (output_batch, next_state) = if batches.len() == 0 { // Go back to `ReadingInput` diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 0ca88a616f686..ace286b4f2b7d 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -278,14 +278,17 @@ impl SingleSpillContext { // CHANGE TO sort iterator that will sort multiple batches instead of one batches without concat // and will output multiple batches - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let sorted_iter = IncrementalSortIterator::new( + 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, - "SingleHashAggregateSpill", - )?; + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "SingleHashAggregateSpill", + )?; let Some((file, max_record_batch_memory)) = spill_file else { return internal_err!("Single hash aggregation produced an empty spill"); From 92fc7ce7d84d12831a0e9c8d7c9bb28588f0e68c Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:41:16 +0300 Subject: [PATCH 3/8] fix(agg): return one column per block from evaluate_to_columns evaluate_to_columns wrapped all evaluated blocks as a single entry, producing batches with one column per block once there was more than one block. --- .../src/aggregates/aggregate_hash_table/common.rs | 7 ++++++- 1 file changed, 6 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 c4ef97ba2fe6a..b25cbedf56323 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -949,7 +949,12 @@ impl HashAggregateAccumulator { &mut self, emit_to: BlockedEmitTo, ) -> Result>> { - Ok(vec![self.evaluate(emit_to)?]) + // One output column per block + Ok(self + .evaluate(emit_to)? + .into_iter() + .map(|block| vec![block]) + .collect()) } /// Evaluating partial aggregate results according to `EmitTo`, and reset inner From 55d5b398fbe0be29c4ce2b2d3bc07e341388b4d5 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:47:26 +0300 Subject: [PATCH 4/8] fix clippy --- .../physical-plan/src/aggregates/aggregate_hash_table/common.rs | 2 +- 1 file changed, 1 insertion(+), 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 b25cbedf56323..7f56695fe545b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -368,7 +368,7 @@ impl AggregateHashTable { .into_iter() .map(|cols| { let batch = - RecordBatch::try_new(output_schema.clone(), cols)?; + RecordBatch::try_new(Arc::clone(&output_schema), cols)?; debug_assert!(batch.num_rows() > 0); total_memory += batch.get_array_memory_size(); Ok(batch) From e4d25aac899bf9c1e0b5c47739b7eb989f881bc9 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:10:31 +0300 Subject: [PATCH 5/8] remove some impl --- .../blocked_boolean_buffer_builder.rs | 304 ------------- .../expr-common/src/blocked_helpers/mod.rs | 2 - .../aggregate/blocked_groups_accumulator.rs | 3 - .../blocked_groups_accumulator/accumulate.rs | 428 ------------------ .../blocked_prim_op.rs | 246 ---------- datafusion/functions-aggregate/src/sum.rs | 34 +- .../src/aggregates/group_values/blocked.rs | 6 +- .../single_group_by/blocked_boolean.rs | 2 - .../single_group_by/blocked_primitive.rs | 7 +- .../src/aggregates/hash_stream.rs | 4 - .../src/aggregates/partial_reduce_stream.rs | 4 - 11 files changed, 11 insertions(+), 1029 deletions(-) delete mode 100644 datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs delete mode 100644 datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs delete mode 100644 datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs diff --git a/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs b/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs deleted file mode 100644 index e1783050d2bdf..0000000000000 --- a/datafusion/expr-common/src/blocked_helpers/blocked_boolean_buffer_builder.rs +++ /dev/null @@ -1,304 +0,0 @@ -use crate::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; -use arrow::array::BooleanBufferBuilder; -use arrow::buffer::BooleanBuffer; -use std::collections::VecDeque; -use std::fmt::Debug; - -/// Blocked [`BooleanBufferBuilder`] -#[derive(Debug)] -pub struct BlockedBooleanBufferBuilder { - /// Using `VecDeque` so we can remove the first block and reclaim memory - blocks: VecDeque, - - /// The size of each block - block_size: usize, - - /// The total number of items, not the number of offset since in each block there is the initial offset - len: usize, - - /// The index of the current block - current_block_index: usize, - should_count_current_block: bool, - - finished_blocks_allocated_memory: usize, -} - -impl BlockedBooleanBufferBuilder { - pub fn new(block_size: usize) -> Self { - assert_ne!(block_size, 0, "block size must be greater than 0"); - - let blocks = VecDeque::from(vec![BooleanBufferBuilder::new(block_size)]); - Self { - blocks, - block_size, - should_count_current_block: false, - len: 0, - current_block_index: 0, - finished_blocks_allocated_memory: 0, - } - } - - pub fn is_empty(&self) -> bool { - self.len == 0 - } - - pub fn len(&self) -> usize { - self.len - } - - fn num_blocks(&self) -> usize { - self.current_block_index + (self.should_count_current_block() as usize) - } - - fn should_count_current_block(&self) -> bool { - self.should_count_current_block - || !self.blocks[self.current_block_index].is_empty() - } - - pub fn block_size(&self) -> usize { - self.block_size - } - - pub fn allocated_size(&self) -> usize { - self.finished_blocks_allocated_memory - + self.blocks.capacity() * size_of::() - + self.blocks.back().map_or(0, builder_allocated_size) - } - - fn end_current_block(&mut self) { - // Don't add to number of blocks since we might not insert into it - self.current_block_index += 1; - self.finished_blocks_allocated_memory += - self.blocks.back().map_or(0, builder_allocated_size); - let new_block = BooleanBufferBuilder::new(self.block_size); - - // Don't count current block since we might not insert into it - self.should_count_current_block = false; - self.blocks.push_back(new_block); - } - - fn reserve_blocks(&mut self, n: usize) { - self.blocks.reserve(n); - } - - /// Push length and return if the current block is now full - pub fn push(&mut self, value: bool) -> bool { - let block = &mut self.blocks[self.current_block_index]; - - block.append(value); - self.len += 1; - - let finished_block = block.len() == self.block_size; - - if finished_block { - self.end_current_block(); - true - } else { - false - } - } - - pub fn current_block_remaining_len(&self) -> usize { - self.block_size - self.blocks[self.current_block_index].len() - } - - fn push_value_n_within_block(&mut self, value: bool, n: usize) -> bool { - self.len += n; - let block = &mut self.blocks[self.current_block_index]; - - let new_len = block.len() + n; - assert!( - new_len <= self.block_size, - "overflow from block new block length: {new_len}, block size: {}", - self.block_size - ); - block.append_n(n, value); - - let finished_block = block.len() == self.block_size; - - if finished_block { - self.end_current_block(); - true - } else { - false - } - } - - pub fn push_value_n(&mut self, value: bool, mut n: usize) { - let number_of_blocks_to_reserve = n - .saturating_sub(self.current_block_remaining_len()) - .div_ceil(self.block_size); - self.reserve_blocks(number_of_blocks_to_reserve); - - while n > 0 { - let remaining_in_current_block = self.current_block_remaining_len(); - - let to_add = remaining_in_current_block.min(n); - n -= to_add; - - self.push_value_n_within_block(value, to_add); - } - } - - pub fn get_bit(&self, index: BlocksIndex) -> bool { - let block_index = index.block_index(); - let index_in_block = index.index_in_block(); - - self.blocks[block_index].get_bit(index_in_block) - } - - pub fn set_bit(&mut self, index: BlocksIndex, value: bool) { - let block_index = index.block_index(); - let index_in_block = index.index_in_block(); - - self.blocks[block_index].set_bit(index_in_block, value) - } - - pub fn emit(&mut self, emit_to: BlockedEmitTo) -> Vec { - match emit_to { - BlockedEmitTo::All => self.take_all(), - BlockedEmitTo::NextBlock => { - self.take_block().map_or(vec![], |next| vec![next]) - } - BlockedEmitTo::First(n) => vec![self.take_n(n)], - } - } - - /// Take every block that counts, see [`Self::num_blocks`] - fn take_all(&mut self) -> Vec { - let num_blocks = self.num_blocks(); - let mut blocks = std::mem::take(&mut self.blocks); - blocks.truncate(num_blocks); - self.reset(); - - blocks.into_iter().map(BooleanBuffer::from).collect() - } - - /// Take the first block, `None` once there are no more items - fn take_block(&mut self) -> Option { - if self.num_blocks() == 0 { - return None; - } - Some(self.take_first_block()) - } - - /// Take the first block even when it is empty, for callers that know from - /// elsewhere that the block holds items - fn take_first_block(&mut self) -> BooleanBuffer { - let block = self.blocks.pop_front().expect("always at least one block"); - - if self.blocks.is_empty() { - self.current_block_index = 0; - self.should_count_current_block = false; - self.blocks - .push_back(BooleanBufferBuilder::new(self.block_size)); - } else { - self.current_block_index -= 1; - - // Only reduce memory if not the last one since the last block is calculated separately - self.finished_blocks_allocated_memory -= builder_allocated_size(&block); - } - - self.len -= block.len(); - - block.build() - } - - fn take_n(&mut self, n: usize) -> BooleanBuffer { - assert_ne!(n, 0, "n must be greater than 0"); - assert!(n <= self.len, "n ({n}) must be <= len ({})", self.len); - assert!( - n < self.block_size, - "n ({n}) must be lower than the block size ({}), instead use `take_block` and take_n with the remainder", - self.block_size - ); - - if n == self.len { - let blocks = self.take_all(); - return blocks - .into_iter() - .next() - .unwrap_or_else(|| BooleanBuffer::new_unset(0)); - } - - if n == self.blocks[0].len() { - return self - .take_block() - .expect("must have at least one block since n < len and n != 0"); - } - - let prev_len = self.len; - - debug_assert!( - self.blocks - .back() - .is_some_and(|block| block.len() < self.block_size), - "the last block must not be full (since it should be the writable tail)" - ); - - // The emitted items are always fully contained in the first block - let mut taken = BooleanBufferBuilder::new(n); - taken.append_packed_range(0..n, self.blocks[0].as_slice()); - - { - let first = &mut self.blocks[0]; - let first_len = first.len(); - - // TODO - reuse the same allocation by adding copy_within - let mut new_first = BooleanBufferBuilder::new(first_len - n); - new_first.append_packed_range(n..first_len, first.as_slice()); - *first = new_first; - } - - // Every block gives its first `n` items to the end of the previous block and - // shifts the rest down, so every finished block stays full and keeps its allocation - for i in 1..self.blocks.len() { - let block = std::mem::replace( - &mut self.blocks[i], - BooleanBufferBuilder::new(self.block_size), - ); - let block_len = block.len(); - let moved = n.min(block_len); - - self.blocks[i - 1].append_packed_range(0..moved, block.as_slice()); - - // Copy to the new block - self.blocks[i].append_packed_range(moved..block.len(), block.as_slice()); - } - - // The old last block is the only one that can end up empty. - // If the block before it is full it is the writable tail as is and keeps its allocation, - // otherwise the block before it is the writable tail and the empty one is not part of the layout - let tail_is_empty = self.blocks.back().is_some_and(|b| b.is_empty()); - let prev_is_full = self.blocks.len() >= 2 - && self.blocks[self.blocks.len() - 2].len() == self.block_size; - - if tail_is_empty && !prev_is_full { - self.blocks.pop_back(); - } - - self.len = prev_len - n; - self.current_block_index = self.blocks.len() - 1; - self.should_count_current_block = !tail_is_empty || !prev_is_full; - self.finished_blocks_allocated_memory = self - .blocks - .iter() - .take(self.current_block_index) - .map(builder_allocated_size) - .sum(); - - taken.build() - } - - pub fn reset(&mut self) { - self.blocks = VecDeque::from(vec![BooleanBufferBuilder::new(self.block_size)]); - self.len = 0; - self.current_block_index = 0; - self.finished_blocks_allocated_memory = 0; - self.should_count_current_block = false; - } -} - -fn builder_allocated_size(block: &BooleanBufferBuilder) -> usize { - block.capacity() / 8 // capacity in bits -} diff --git a/datafusion/expr-common/src/blocked_helpers/mod.rs b/datafusion/expr-common/src/blocked_helpers/mod.rs index 2c1ee4414c333..b6dca4fd7fa8e 100644 --- a/datafusion/expr-common/src/blocked_helpers/mod.rs +++ b/datafusion/expr-common/src/blocked_helpers/mod.rs @@ -1,5 +1,3 @@ -mod blocked_boolean_buffer_builder; mod blocked_vec; -pub use blocked_boolean_buffer_builder::BlockedBooleanBufferBuilder; pub use blocked_vec::BlockedVec; diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs index 7769e2e97ba01..c516ba7fc085c 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs @@ -18,9 +18,6 @@ //! Utilities for implementing GroupsAccumulator //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] -pub mod accumulate; -pub mod blocked_prim_op; - use arrow::array::{ArrayRef, BooleanArray}; use datafusion_common::Result; use datafusion_expr_common::blocked_groups_accumulator::{ diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs deleted file mode 100644 index 4aff483a39610..0000000000000 --- a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/accumulate.rs +++ /dev/null @@ -1,428 +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. - -//! [`GroupsAccumulator`] helpers: [`BlockedNullState`] and [`accumulate_indices`] -//! -//! [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator - -use arrow::array::{Array, BooleanArray, BooleanBufferBuilder, PrimitiveArray}; -use arrow::buffer::NullBuffer; -use arrow::datatypes::ArrowPrimitiveType; - -use crate::aggregate::groups_accumulator::accumulate::accumulate; -use datafusion_common::Result; -use datafusion_expr_common::blocked_groups_accumulator::{ - BlockedEmitTo, BlockedGroupSelection, BlocksIndex, -}; -use datafusion_expr_common::blocked_helpers::BlockedBooleanBufferBuilder; - -/// If the input has nulls, then the accumulator must potentially -/// handle each input null value specially (e.g. for `SUM` to mark the -/// corresponding sum as null) -/// -/// `NullState` tracks if it has seen *any* value for each group when filters or -/// sparse group indices may omit input for a registered group. -#[derive(Debug)] -pub enum BlockedSeenValues { - /// All groups seen so far have seen at least one non-null value - All { - num_values: usize, - block_size: usize, - }, - // Some groups have not yet seen a non-null value - Some { - values: BlockedBooleanBufferBuilder, - }, -} - -impl BlockedSeenValues { - pub fn new(block_size: usize) -> Self { - BlockedSeenValues::All { - num_values: 0, - block_size, - } - } - - pub fn block_size(&self) -> usize { - match self { - BlockedSeenValues::All { block_size, .. } => *block_size, - BlockedSeenValues::Some { values } => values.block_size(), - } - } - - /// Return a mutable reference to the `BooleanBufferBuilder` in `BlockedSeenValues::Some`. - /// - /// If `self` is `BlockedSeenValues::All`, it is transitioned to `BlockedSeenValues::Some` - /// by creating a new `BooleanBufferBuilder` where the first `num_values` are true. - /// - /// The builder is then ensured to have at least `total_num_groups` length, - /// with any new entries initialized to false. - fn get_builder( - &mut self, - total_num_groups: usize, - ) -> &mut BlockedBooleanBufferBuilder { - match self { - BlockedSeenValues::All { - num_values, - block_size, - } => { - let mut builder = BlockedBooleanBufferBuilder::new(*block_size); - builder.push_value_n(true, *num_values); - if total_num_groups > *num_values { - builder.push_value_n(false, total_num_groups - *num_values); - } - *self = BlockedSeenValues::Some { values: builder }; - match self { - BlockedSeenValues::Some { values } => values, - _ => unreachable!(), - } - } - BlockedSeenValues::Some { values } => { - if values.len() < total_num_groups { - values.push_value_n(false, total_num_groups - values.len()); - } - values - } - } - } -} - -/// Returns true when all newly registered groups are present in `group_indices`. -/// -/// Group indices are assigned in first-seen order, so an unfiltered batch visits -/// new groups in ascending order. Pre-filtered input can omit a new group, making -/// the indices sparse even though the accumulator no longer receives a filter. -fn new_groups_are_dense( - group_indices: &[BlocksIndex], - block_size: usize, - first_new_group: usize, - total_num_groups: usize, -) -> bool { - let first_new_group = - BlocksIndex::from_index_in_fixed_block_size(first_new_group, block_size); - let total_num_groups = - BlocksIndex::from_index_in_fixed_block_size(total_num_groups, block_size); - - if first_new_group == total_num_groups { - return true; - } - - let mut next_new_group = first_new_group; - for &group_index in group_indices { - if group_index == next_new_group { - next_new_group = next_new_group.increment(block_size); - } else if group_index > next_new_group { - return false; - } - } - next_new_group == total_num_groups -} - -/// Track the accumulator null state per row: if any values for that -/// group were null and if any values have been seen at all for that group. -/// -/// This is part of the inner loop for many [`GroupsAccumulator`]s, -/// and thus the performance is critical and so there are multiple -/// specialized implementations, invoked depending on the specific -/// combinations of the input. -/// -/// Typically there are 4 potential combinations of inputs must be -/// special cased for performance: -/// -/// * With / Without filter -/// * With / Without nulls in the input -/// -/// If the input has nulls, then the accumulator must potentially -/// handle each input null value specially (e.g. for `SUM` to mark the -/// corresponding sum as null) -/// -/// `NullState` tracks if it has seen *any* value for each group when filters or -/// sparse group indices may omit input for a registered group. -/// -/// [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator -#[derive(Debug)] -pub struct BlockedNullState { - /// Have we seen any non-filtered input values for `group_index`? - /// - /// If `seen_values` is `SeenValues::Some(buffer)` and buffer\[i\] is true, have seen at least one non null - /// value for group `i` - /// - /// If `seen_values` is `SeenValues::Some(buffer)` and buffer\[i\] is false, have not seen any values that - /// pass the filter yet for group `i` - /// - /// If `seen_values` is `SeenValues::All`, all groups have seen at least one non null value - seen_values: BlockedSeenValues, -} - -impl BlockedNullState { - pub fn new(block_size: usize) -> Self { - Self { - seen_values: BlockedSeenValues::All { - num_values: 0, - block_size, - }, - } - } - - /// return the size of all buffers allocated by this null state, not including self - pub fn size(&self) -> usize { - match &self.seen_values { - BlockedSeenValues::All { .. } => 0, - BlockedSeenValues::Some { values } => values.allocated_size(), - } - } - - /// Invokes `value_fn(group_index, value)` for each non null, non - /// filtered value of `value`, while tracking which groups have - /// seen null inputs and which groups have seen any inputs if necessary - // - /// # Arguments: - /// - /// * `values`: the input arguments to the accumulator - /// * `group_indices`: To which groups do the rows in `values` belong, (aka group_index) - /// * `opt_filter`: if present, only rows for which is Some(true) are included - /// * `value_fn`: function invoked for (group_index, value) where value is non null - /// - /// See [`accumulate`], for more details on how value_fn is called - /// - /// When value_fn is called it also sets - /// - /// 1. `self.seen_values[group_index]` to true for all rows that had a non null value - pub fn accumulate( - &mut self, - group_indices: &[BlocksIndex], - values: &PrimitiveArray, - opt_filter: Option<&BooleanArray>, - total_num_groups: usize, - mut value_fn: F, - ) where - T: ArrowPrimitiveType + Send, - F: FnMut(BlocksIndex, T::Native) + Send, - { - // Skip per-value null handling when every input value is valid and all - // newly registered groups are represented. Pre-filtered inputs can have - // sparse group indices despite not passing a filter to the accumulator. - if opt_filter.is_none() - && values.null_count() == 0 - && let BlockedSeenValues::All { - num_values, - block_size, - } = &mut self.seen_values - && new_groups_are_dense( - group_indices, - *block_size, - *num_values, - total_num_groups, - ) - { - accumulate(group_indices, values, None, value_fn); - *num_values = total_num_groups; - return; - } - - let seen_values = self.seen_values.get_builder(total_num_groups); - accumulate(group_indices, values, opt_filter, |group_index, value| { - seen_values.set_bit(group_index, true); - value_fn(group_index, value); - }); - } - - /// Invokes `value_fn(group_index, value)` for each non null, non - /// filtered value in `values`, while tracking which groups have - /// seen null inputs and which groups have seen any inputs, for - /// [`BooleanArray`]s. - /// - /// Since `BooleanArray` is not a [`PrimitiveArray`] it must be - /// handled specially. - /// - /// See [`Self::accumulate`], which handles `PrimitiveArray`s, for - /// more details on other arguments. - pub fn accumulate_boolean( - &mut self, - group_indices: &[BlocksIndex], - values: &BooleanArray, - opt_filter: Option<&BooleanArray>, - total_num_groups: usize, - mut value_fn: F, - ) where - F: FnMut(BlocksIndex, bool) + Send, - { - let data = values.values(); - assert_eq!(data.len(), group_indices.len()); - - // Skip per-value null handling when every input value is valid and all - // newly registered groups are represented. Pre-filtered inputs can have - // sparse group indices despite not passing a filter to the accumulator. - if opt_filter.is_none() - && values.null_count() == 0 - && let BlockedSeenValues::All { - num_values, - block_size, - } = &mut self.seen_values - && new_groups_are_dense( - group_indices, - *block_size, - *num_values, - total_num_groups, - ) - { - group_indices - .iter() - .zip(data.iter()) - .for_each(|(&group_index, new_value)| value_fn(group_index, new_value)); - *num_values = total_num_groups; - - return; - } - - let seen_values = self.seen_values.get_builder(total_num_groups); - - // These could be made more performant by iterating in chunks of 64 bits at a time - match (values.null_count() > 0, opt_filter) { - // no nulls, no filter, - (false, None) => { - // if we have previously seen nulls, ensure the null - // buffer is big enough (start everything at valid) - group_indices.iter().zip(data.iter()).for_each( - |(&group_index, new_value)| { - seen_values.set_bit(group_index, true); - value_fn(group_index, new_value) - }, - ) - } - // nulls, no filter - (true, None) => { - let nulls = values.nulls().unwrap(); - group_indices - .iter() - .zip(data.iter()) - .zip(nulls.iter()) - .for_each(|((&group_index, new_value), is_valid)| { - if is_valid { - seen_values.set_bit(group_index, true); - value_fn(group_index, new_value); - } - }) - } - // no nulls, but a filter - (false, Some(filter)) => { - assert_eq!(filter.len(), group_indices.len()); - - group_indices - .iter() - .zip(data.iter()) - .zip(filter.iter()) - .for_each(|((&group_index, new_value), filter_value)| { - if filter_value == Some(true) { - seen_values.set_bit(group_index, true); - value_fn(group_index, new_value); - } - }) - } - // both null values and filters - (true, Some(filter)) => { - assert_eq!(filter.len(), group_indices.len()); - filter - .iter() - .zip(group_indices.iter()) - .zip(values.iter()) - .for_each(|((filter_value, &group_index), new_value)| { - if filter_value == Some(true) - && let Some(new_value) = new_value - { - seen_values.set_bit(group_index, true); - value_fn(group_index, new_value) - } - }) - } - } - } - - /// Creates a [`NullBuffer`] for `selection` without changing this state. - pub fn build_preserving( - &self, - selection: BlockedGroupSelection<'_>, - ) -> Result> { - let selected_len = selection.len(); - match &self.seen_values { - BlockedSeenValues::All { - num_values, - block_size: _, - } => { - selection.validate_num_groups(*num_values)?; - Ok(None) - } - BlockedSeenValues::Some { values } => { - selection.validate_num_groups(values.len())?; - let mut selected = BooleanBufferBuilder::new(selected_len); - for index in selection.iter() { - selected.append(values.get_bit(index)); - } - Ok(Some(NullBuffer::new(selected.finish()))) - } - } - } - - /// Creates the a [`NullBuffer`] representing which group_indices - /// should have null values (because they never saw any values) - /// for the `emit_to` rows. - /// - /// resets the internal state appropriately - pub fn build(&mut self, emit_to: BlockedEmitTo) -> Vec> { - let block_size = self.seen_values.block_size(); - let old_seen = - std::mem::replace(&mut self.seen_values, BlockedSeenValues::new(block_size)); - - match old_seen { - BlockedSeenValues::All { - num_values, - block_size, - } => match emit_to { - BlockedEmitTo::All => vec![None; num_values.div_ceil(block_size)], - BlockedEmitTo::NextBlock => { - self.seen_values = BlockedSeenValues::All { - num_values: num_values.saturating_sub(block_size), - block_size, - }; - - if num_values == 0 { vec![] } else { vec![None] } - } - BlockedEmitTo::First(n) => { - assert!( - n < block_size, - "emit_to First(n) must be less than block_size" - ); - self.seen_values = BlockedSeenValues::All { - num_values: num_values.saturating_sub(n), - block_size, - }; - if num_values == 0 { vec![] } else { vec![None] } - } - }, - BlockedSeenValues::Some { mut values } => { - let emitted = values.emit(emit_to); - - self.seen_values = BlockedSeenValues::Some { values }; - - emitted - .into_iter() - .map(|block| Some(NullBuffer::new(block))) - .collect() - } - } - } -} diff --git a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs deleted file mode 100644 index 3b9d08aabfa2b..0000000000000 --- a/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator/blocked_prim_op.rs +++ /dev/null @@ -1,246 +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. - -use std::sync::Arc; - -use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray}; -use arrow::buffer::NullBuffer; -use arrow::compute; -use arrow::datatypes::ArrowPrimitiveType; -use arrow::datatypes::DataType; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use datafusion_expr_common::blocked_groups_accumulator::{ - BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex, -}; -use datafusion_expr_common::blocked_helpers::BlockedVec; - -use super::accumulate::BlockedNullState; - -/// An accumulator that implements a single operation over -/// [`ArrowPrimitiveType`] where the accumulated 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). -/// -/// [`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 - } -} - -impl BlockedGroupsAccumulator for BlockedPrimitiveGroupsAccumulator -where - T: ArrowPrimitiveType + Send, - F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, -{ - fn batch_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::(); - - // update values - self.values - .push_value_n_to_len(self.starting_value, total_num_groups); - - // NullState dispatches / handles tracking nulls and groups that saw no values - self.null_state.accumulate( - group_indices, - values, - opt_filter, - total_num_groups, - |group_index, new_value| { - // SAFETY: group_index is guaranteed to be in bounds - let value = unsafe { self.values.get_unchecked_mut(group_index) }; - (self.prim_fn)(value, new_value); - }, - ); - - Ok(()) - } - - fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result> { - let values = self.values.emit(emit_to); - let nulls = self.null_state.build(emit_to); - - let blocks = values - .into_iter() - .zip(nulls) - .map(|(values, nulls)| { - let values = PrimitiveArray::::new(values.into(), nulls) // no copy - .with_data_type(self.data_type.clone()); - Arc::new(values) as ArrayRef - }) - .collect::>(); - - Ok(blocks) - } - - fn evaluate_preserving( - &mut self, - selection: BlockedGroupSelection<'_>, - ) -> Result { - selection.validate_num_groups(self.values.len())?; - let values: Vec = - selection.iter().map(|index| self.values[index]).collect(); - let nulls = self.null_state.build_preserving(selection)?; - let values = PrimitiveArray::::new(values.into(), nulls) - .with_data_type(self.data_type.clone()); - Ok(Arc::new(values)) - } - - fn supports_evaluate_preserving(&self) -> bool { - true - } - - fn state(&mut self, emit_to: BlockedEmitTo) -> Result>> { - let blocks = self.evaluate(emit_to)?; - - Ok(blocks - .into_iter() - // Convert each block into a Vec to match the return type - .map(|arr| vec![arr]) - .collect::>()) - } - - fn state_preserving( - &mut self, - selection: BlockedGroupSelection<'_>, - ) -> Result> { - self.evaluate_preserving(selection).map(|arr| vec![arr]) - } - - fn supports_state_preserving(&self) -> bool { - true - } - - 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) - } - - /// Converts an input batch directly to a state batch - /// - /// The state is: - /// - self.prim_fn for all non null, non filtered values - /// - null otherwise - fn convert_to_state( - &self, - 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)]) - } - - fn size(&self) -> usize { - self.values.allocated_size() + self.null_state.size() - } -} diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs index 0f49bfc35655a..9ca78f2c25c8e 100644 --- a/datafusion/functions-aggregate/src/sum.rs +++ b/datafusion/functions-aggregate/src/sum.rs @@ -20,9 +20,9 @@ use arrow::array::{Array, ArrayRef, ArrowNativeTypeOp, ArrowNumericType, AsArray}; use arrow::datatypes::Field; use arrow::datatypes::{ - ArrowNativeType, DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, - DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DataType, Decimal32Type, - Decimal64Type, Decimal128Type, Decimal256Type, DurationMicrosecondType, + ArrowNativeType, + DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DataType, Decimal128Type, Decimal256Type, Decimal32Type, + Decimal64Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, DurationSecondType, FieldRef, Float64Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType, TimeUnit, UInt64Type, @@ -31,22 +31,20 @@ use datafusion_common::hash_utils::RandomState; use datafusion_common::internal_err; use datafusion_common::stats::Precision; use datafusion_common::types::{ - NativeType, logical_float64, logical_int8, logical_int16, logical_int32, - logical_int64, logical_uint8, logical_uint16, logical_uint32, logical_uint64, + NativeType, logical_float64, logical_int16, logical_int32, + logical_int64, logical_int8, logical_uint16, logical_uint32, logical_uint64, logical_uint8, }; use datafusion_common::{HashMap, Result, ScalarValue, exec_err, not_impl_err}; use datafusion_expr::expr::AggregateFunction; 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, Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility}; +use datafusion_expr::{Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, GroupsAccumulator, Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; -use datafusion_functions_aggregate_common::aggregate::blocked_groups_accumulator::blocked_prim_op::BlockedPrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator; use datafusion_macros::user_doc; use datafusion_physical_expr::expressions::{CastExpr, Column}; use std::mem::{size_of, size_of_val}; -use datafusion_functions_aggregate_common::accumulator::BlockedAccumulatorArgs; make_udaf_expr_and_func!( Sum, @@ -315,26 +313,6 @@ impl AggregateUDFImpl for Sum { downcast_sum!(args, helper) } - fn blocked_groups_accumulator_supported(&self, args: BlockedAccumulatorArgs) -> bool { - !args.is_distinct - } - - fn create_blocked_groups_accumulator( - &self, - args: BlockedAccumulatorArgs, - ) -> Result> { - macro_rules! helper { - ($t:ty, $dt:expr) => { - Ok(Box::new(BlockedPrimitiveGroupsAccumulator::<$t, _>::new( - &$dt, - args.batch_size, - |x, y| *x = x.add_wrapping(y), - ))) - }; - } - downcast_sum!(args, helper) - } - fn create_sliding_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs index 2f4369888cdc8..4bbf784e49193 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs @@ -266,13 +266,11 @@ impl BlockedGroupValues for BlockedGroupValuesAdapter { self.block_size() ); - let output = if self.len() == n { + if self.len() == n { self.inner.emit(EmitTo::All) } else { self.inner.emit(EmitTo::First(n)) - }; - - Ok(output?) + } } fn clear_shrink(&mut self, num_rows: usize) { diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs index 4e35dd69d5012..ddb84a80cd8b7 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs @@ -15,8 +15,6 @@ // specific language governing permissions and limitations // under the License. -use crate::aggregates::group_values::GroupValues; - use crate::aggregates::group_values::blocked::BlockedGroupValues; use arrow::array::{ ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs index 0e4f40ab9d0b2..94fe32e4127f4 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs @@ -24,7 +24,6 @@ use arrow::buffer::ScalarBuffer; use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; -use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; use datafusion_expr_common::blocked_helpers::BlockedVec; use hashbrown::hash_table::HashTable; @@ -80,7 +79,7 @@ impl BlockedGroupValuesPrimitive { // NOTE: The inner builder must be constructed as there is at least one null buffer.finish().unwrap() }); - PrimitiveArray::::new(values.into(), nulls) + PrimitiveArray::::new(values, nulls) } fn build_no_nulls_primitive_arc( @@ -105,7 +104,7 @@ impl BlockedGroupValuesPrimitive { buffer.finish().unwrap() }; Arc::new( - PrimitiveArray::::new(values.into(), Some(nulls)) + PrimitiveArray::::new(values, Some(nulls)) .with_data_type(data_type), ) } @@ -219,7 +218,7 @@ where }; // If nothing left - let null_group = if self.values.len() == 0 { + let null_group = if self.values.is_empty() { self.map.clear(); assert_eq!(self.null_group.map_or(0, |index| index.block_index()), 0); diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 94ce7f09ccbbf..19964d4feac41 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -144,9 +144,6 @@ pub(crate) struct PartialHashAggregateStream { /// Input batches containing raw rows, not partial aggregate state. input: SendableRecordBatchStream, - /// Target output batch size from configuration. - batch_size: usize, - /// Memory reservation for group keys and accumulators. reservation: MemoryReservation, @@ -271,7 +268,6 @@ impl PartialHashAggregateStream { Ok(Self { schema, input, - batch_size, baseline_metrics, reservation, reduction_factor, diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 48d0abc0f4973..a65ff19d7e91e 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -79,9 +79,6 @@ pub(crate) struct PartialReduceHashAggregateStream { /// Input batches containing partial aggregate state rows. input: SendableRecordBatchStream, - /// Target output batch size from configuration. - batch_size: usize, - /// Execution metrics shared with the aggregate plan node. baseline_metrics: BaselineMetrics, @@ -210,7 +207,6 @@ impl PartialReduceHashAggregateStream { Ok(Self { schema, input, - batch_size, baseline_metrics, reservation, early_emit_count, From 57bb439c4aa767dde0ccf44ccb763a97dceda81f Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:14:00 +0300 Subject: [PATCH 6/8] remove some --- datafusion/functions-aggregate/src/sum.rs | 16 +- .../src/aggregates/group_values/blocked.rs | 4 - .../single_group_by/blocked_boolean.rs | 292 ------------------ .../group_values/single_group_by/mod.rs | 1 - 4 files changed, 10 insertions(+), 303 deletions(-) delete mode 100644 datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs index 9ca78f2c25c8e..1228b458a82b9 100644 --- a/datafusion/functions-aggregate/src/sum.rs +++ b/datafusion/functions-aggregate/src/sum.rs @@ -20,9 +20,9 @@ use arrow::array::{Array, ArrayRef, ArrowNativeTypeOp, ArrowNumericType, AsArray}; use arrow::datatypes::Field; use arrow::datatypes::{ - ArrowNativeType, - DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DataType, Decimal128Type, Decimal256Type, Decimal32Type, - Decimal64Type, DurationMicrosecondType, + ArrowNativeType, DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, + DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DataType, Decimal32Type, + Decimal64Type, Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, DurationSecondType, FieldRef, Float64Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType, TimeUnit, UInt64Type, @@ -31,15 +31,19 @@ use datafusion_common::hash_utils::RandomState; use datafusion_common::internal_err; use datafusion_common::stats::Precision; use datafusion_common::types::{ - NativeType, logical_float64, logical_int16, logical_int32, - logical_int64, logical_int8, logical_uint16, logical_uint32, logical_uint64, logical_uint8, + NativeType, logical_float64, logical_int8, logical_int16, logical_int32, + logical_int64, logical_uint8, logical_uint16, logical_uint32, logical_uint64, }; use datafusion_common::{HashMap, Result, ScalarValue, exec_err, not_impl_err}; use datafusion_expr::expr::AggregateFunction; 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, Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility}; +use datafusion_expr::{ + Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, GroupsAccumulator, + Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, + TypeSignatureClass, Volatility, +}; 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; diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs index 4bbf784e49193..f3d6562ddc4a3 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs @@ -1,5 +1,4 @@ use super::single_group_by::{ - blocked_boolean::BlockedGroupValuesBoolean, blocked_primitive::BlockedGroupValuesPrimitive, }; use crate::aggregates::group_values::new_group_values; @@ -313,9 +312,6 @@ pub fn new_blocked_group_values( downcast_primitive! { d => (downcast_helper, d), - DataType::Boolean => { - return Ok(Box::new(BlockedGroupValuesBoolean::new(block_size))); - }, _ => {} } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs deleted file mode 100644 index ddb84a80cd8b7..0000000000000 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_boolean.rs +++ /dev/null @@ -1,292 +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. - -use crate::aggregates::group_values::blocked::BlockedGroupValues; -use arrow::array::{ - ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, -}; -use datafusion_common::Result; -use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; -use std::{mem::size_of, sync::Arc}; - -#[derive(Debug)] -pub struct BlockedGroupValuesBoolean { - false_group: Option, - true_group: Option, - null_group: Option, - block_size: usize, -} - -impl BlockedGroupValuesBoolean { - pub fn new(block_size: usize) -> Self { - assert_ne!(block_size, 0); - Self { - block_size, - false_group: None, - true_group: None, - null_group: None, - } - } -} - -impl BlockedGroupValues for BlockedGroupValuesBoolean { - fn block_size(&self) -> usize { - self.block_size - } - - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { - let array = cols[0].as_boolean(); - groups.clear(); - - for value in array.iter() { - let index = match value { - Some(false) => { - if let Some(index) = self.false_group { - index - } else { - let index = self.len(); - self.false_group = Some(index); - index - } - } - Some(true) => { - if let Some(index) = self.true_group { - index - } else { - let index = self.len(); - self.true_group = Some(index); - index - } - } - None => { - if let Some(index) = self.null_group { - index - } else { - let index = self.len(); - self.null_group = Some(index); - index - } - } - }; - - groups.push(BlocksIndex::from_index_in_fixed_block_size( - index, - self.block_size, - )); - } - - Ok(()) - } - - fn size(&self) -> usize { - size_of::() - } - - fn is_empty(&self) -> bool { - self.len() == 0 - } - - fn len(&self) -> usize { - self.false_group.is_some() as usize - + self.true_group.is_some() as usize - + self.null_group.is_some() as usize - } - - fn emit_all(&mut self) -> Result>> { - let mut blocks = vec![]; - - while let Some(block) = self.emit_block()? { - blocks.push(block); - } - - Ok(blocks) - } - - fn emit_block(&mut self) -> Result>> { - let len = self.len(); - if len == 0 { - return Ok(None); - } - // Group indices are always contiguous from 0, so the last block holds the remaining groups - let emit_count = self.block_size.min(len); - let mut builder = BooleanBufferBuilder::new(emit_count); - builder.append_n(emit_count, false); - if let Some(idx) = self.true_group.as_mut() { - if *idx < emit_count { - builder.set_bit(*idx, true); - self.true_group = None; - } else { - *idx -= emit_count; - } - } - - if let Some(idx) = self.false_group.as_mut() { - if *idx < emit_count { - // already false, no need to set - self.false_group = None; - } else { - *idx -= emit_count; - } - } - - let values = builder.finish(); - - let nulls = if let Some(idx) = self.null_group.as_mut() { - if *idx < emit_count { - let mut buffer = NullBufferBuilder::new(len); - buffer.append_n_non_nulls(*idx); - buffer.append_null(); - buffer.append_n_non_nulls(emit_count - *idx - 1); - - self.null_group = None; - Some(buffer.finish().unwrap()) - } else { - *idx -= emit_count; - None - } - } else { - None - }; - - Ok(Some(vec![Arc::new(BooleanArray::new(values, nulls)) as _])) - } - - fn emit_first_n(&mut self, n: usize) -> Result> { - assert!( - n < self.block_size, - "{n} must be less than block_size {}", - self.block_size - ); - - { - let len = self.len(); - assert!(n < len, "{n} must be less than length {len}"); - } - - let mut builder = BooleanBufferBuilder::new(n); - builder.append_n(n, false); - if let Some(idx) = self.true_group.as_mut() { - if *idx < n { - builder.set_bit(*idx, true); - self.true_group = None; - } else { - *idx -= n; - } - } - - if let Some(idx) = self.false_group.as_mut() { - if *idx < n { - // already false, no need to set - self.false_group = None; - } else { - *idx -= n; - } - } - - let values = builder.finish(); - - let nulls = if let Some(idx) = self.null_group.as_mut() { - if *idx < n { - let mut buffer = NullBufferBuilder::new(builder.len()); - buffer.append_n_non_nulls(*idx); - buffer.append_null(); - buffer.append_n_non_nulls(n - *idx - 1); - - self.null_group = None; - Some(buffer.finish().unwrap()) - } else { - *idx -= n; - None - } - } else { - None - }; - - Ok(vec![Arc::new(BooleanArray::new(values, nulls)) as _]) - } - - fn clear_shrink(&mut self, _num_rows: usize) { - self.false_group = None; - self.true_group = None; - self.null_group = None; - } -} - -#[cfg(test)] -mod tests { - use super::*; - use datafusion_expr_common::blocked_groups_accumulator::BlocksIndex; - - fn intern( - group_values: &mut BlockedGroupValuesBoolean, - values: &[Option], - ) -> Vec { - let mut groups = vec![]; - let array: ArrayRef = Arc::new(BooleanArray::from(values.to_vec())); - group_values.intern(&[array], &mut groups).unwrap(); - groups - } - - fn emitted(blocks: Vec>) -> Vec>> { - blocks - .into_iter() - .map(|block| block[0].as_boolean().iter().collect()) - .collect() - } - - #[test] - fn emit_all_only_emits_existing_groups() { - let mut group_values = BlockedGroupValuesBoolean::new(531); - let groups = intern( - &mut group_values, - &[Some(true), None, Some(true), Some(false)], - ); - assert_eq!( - groups, - [0, 1, 0, 2].map(BlocksIndex::new_in_first_block).to_vec() - ); - assert_eq!(group_values.len(), 3); - - let blocks = emitted(group_values.emit_all().unwrap()); - assert_eq!(blocks, vec![vec![Some(true), None, Some(false)]]); - assert!(group_values.is_empty()); - assert!(group_values.emit_block().unwrap().is_none()); - } - - #[test] - fn emit_all_splits_into_blocks() { - let mut group_values = BlockedGroupValuesBoolean::new(2); - intern(&mut group_values, &[Some(false), Some(true), None]); - - let blocks = emitted(group_values.emit_all().unwrap()); - assert_eq!(blocks, vec![vec![Some(false), Some(true)], vec![None]]); - assert!(group_values.is_empty()); - } - - #[test] - fn emit_first_n_shifts_remaining() { - let mut group_values = BlockedGroupValuesBoolean::new(4); - intern(&mut group_values, &[None, Some(false), Some(true)]); - - let first = group_values.emit_first_n(1).unwrap(); - assert_eq!(first[0].as_boolean().iter().collect::>(), vec![None]); - assert_eq!(group_values.len(), 2); - - let blocks = emitted(group_values.emit_all().unwrap()); - assert_eq!(blocks, vec![vec![Some(false), Some(true)]]); - } -} diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs index 28721678e43bc..6381fc50cb5bf 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/mod.rs @@ -17,7 +17,6 @@ //! `GroupValues` implementations for single group by cases -pub(crate) mod blocked_boolean; pub(crate) mod blocked_primitive; pub(crate) mod boolean; pub(crate) mod bytes; From 725baf83a37b674f841655a3ebcdc2b8d15717a5 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:15:25 +0300 Subject: [PATCH 7/8] chore: align datafusion-testing submodule with main --- datafusion-testing | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion-testing b/datafusion-testing index 13bbae38776c2..81c7440e08b9f 160000 --- a/datafusion-testing +++ b/datafusion-testing @@ -1 +1 @@ -Subproject commit 13bbae38776c2bfbc1fab1be7e7220222d4284bf +Subproject commit 81c7440e08b9fe81a982e6b0b7e0fe3a6e82fb78 From 1b2a8f15433ce491ce9a3f8db0df69365e420e6b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:29:33 +0300 Subject: [PATCH 8/8] remove comments --- .../src/blocked_groups_accumulator.rs | 52 ------------------- .../src/aggregates/group_values/blocked.rs | 35 ------------- 2 files changed, 87 deletions(-) diff --git a/datafusion/expr-common/src/blocked_groups_accumulator.rs b/datafusion/expr-common/src/blocked_groups_accumulator.rs index 7ac34ee834b49..82e89b2234c80 100644 --- a/datafusion/expr-common/src/blocked_groups_accumulator.rs +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -324,58 +324,6 @@ impl BlockedEmitTo { } } -/// `GroupsAccumulator` implements a single aggregate (e.g. AVG) and -/// stores the state for *all* groups internally. -/// -/// Logically, a [`GroupsAccumulator`] stores a mapping from each group index to -/// the state of the aggregate for that group. For example an implementation for -/// `min` might look like -/// -/// ```text -/// ┌─────┐ -/// │ 0 │───────────▶ 100 -/// ├─────┤ -/// │ 1 │───────────▶ 200 -/// └─────┘ -/// ... ... -/// ┌─────┐ -/// │ N-2 │───────────▶ 50 -/// ├─────┤ -/// │ N-1 │───────────▶ 200 -/// └─────┘ -/// -/// -/// Logical group Current Min -/// number value for that -/// group -/// ``` -/// -/// # Notes on Implementing `GroupsAccumulator` -/// -/// All aggregates must first implement the simpler [`Accumulator`] trait, which -/// handles state for a single group. Implementing `GroupsAccumulator` is -/// optional and is harder to implement than `Accumulator`, but can be much -/// faster for queries with many group values. See the [Aggregating Millions of -/// Groups Fast blog] for more background. -/// For more background, please also see the [Aggregating Millions of Groups Fast in Apache Arrow DataFusion 28.0.0 blog] -/// -/// [Aggregating Millions of Groups Fast in Apache Arrow DataFusion 28.0.0 blog]: https://datafusion.apache.org/blog/2023/08/05/datafusion_fast_grouping -/// -/// [`NullState`] can help keep the state for groups that have not seen any -/// values and produce the correct output for those groups. -/// -/// [`NullState`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/struct.NullState.html -/// -/// # Details -/// Each group is assigned a `group_index` by the hash table and each -/// accumulator manages the specific state, one per `group_index`. -/// -/// `group_index`es are contiguous (there aren't gaps), and thus it is -/// expected that each `GroupsAccumulator` will use something like `Vec<..>` -/// to store the group states. -/// -/// [`Accumulator`]: crate::accumulator::Accumulator -/// [Aggregating Millions of Groups Fast blog]: https://arrow.apache.org/blog/2023/08/05/datafusion_fast_grouping/ pub trait BlockedGroupsAccumulator: Send + std::any::Any { fn batch_size(&self) -> usize; diff --git a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs index f3d6562ddc4a3..f9e53612014c2 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/blocked.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs @@ -15,41 +15,6 @@ use datafusion_expr_common::blocked_groups_accumulator::{ use datafusion_expr_common::groups_accumulator::EmitTo; /// Stores the group values during hash aggregation. -/// -/// # Background -/// -/// In a query such as `SELECT a, b, count(*) FROM t GROUP BY a, b`, the group values -/// identify each group, and correspond to all the distinct values of `(a,b)`. -/// -/// ```sql -/// -- Input has 4 rows with 3 distinct combinations of (a,b) ("groups") -/// create table t(a int, b varchar) -/// as values (1, 'a'), (2, 'b'), (1, 'a'), (3, 'c'); -/// -/// select a, b, count(*) from t group by a, b; -/// ---- -/// 1 a 2 -/// 2 b 1 -/// 3 c 1 -/// ``` -/// -/// # Design -/// -/// Managing group values is a performance critical operation in hash -/// aggregation. The major operations are: -/// -/// 1. Intern: Quickly finding existing and adding new group values -/// 2. Emit: Returning the group values as an array -/// -/// There are multiple specialized implementations of this trait optimized for -/// different data types and number of columns, optimized for these operations. -/// See [`crate::aggregates::group_values::new_group_values`] for details. -/// -/// # Group Ids -/// -/// Each distinct group in a hash aggregation is identified by a unique group id -/// (usize) which is assigned by instances of this trait. Group ids are -/// continuous without gaps, starting from 0. pub trait BlockedGroupValues: Send { fn block_size(&self) -> usize;