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..82e89b2234c80 --- /dev/null +++ b/datafusion/expr-common/src/blocked_groups_accumulator.rs @@ -0,0 +1,398 @@ +// 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 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, +}; +use std::cmp::Ordering; + +/// 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!( + u32::try_from(value).is_ok(), + "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()) + } +} + +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_vec.rs b/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs new file mode 100644 index 0000000000000..edbe28d12e91b --- /dev/null +++ b/datafusion/expr-common/src/blocked_helpers/blocked_vec.rs @@ -0,0 +1,352 @@ +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}; + +/// 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 + } + + 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, |b| b.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, |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); + } + + 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(); + true + } else { + false + } + } + + 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: 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(); + true + } else { + false + } + } + + 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] + 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] + 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() } + } + + /// 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()) }; + unsafe { block.get_unchecked(index.index_in_block()) } + } + + /// 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()) }; + unsafe { block.get_unchecked_mut(index.index_in_block()) } + } + + 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`] + 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; + } + 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(); + + Some(block) + } + + 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 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..b6dca4fd7fa8e --- /dev/null +++ b/datafusion/expr-common/src/blocked_helpers/mod.rs @@ -0,0 +1,3 @@ +mod blocked_vec; + +pub use blocked_vec::BlockedVec; diff --git a/datafusion/expr-common/src/lib.rs b/datafusion/expr-common/src/lib.rs index c9a95fd294503..995425bc7cbf7 100644 --- a/datafusion/expr-common/src/lib.rs +++ b/datafusion/expr-common/src/lib.rs @@ -34,6 +34,8 @@ #![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; pub mod dyn_eq; diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs index 68865cbe1ca54..9816fb29434e4 100644 --- a/datafusion/expr/src/function.rs +++ b/datafusion/expr/src/function.rs @@ -24,7 +24,7 @@ use datafusion_common::Result; use std::sync::Arc; pub use datafusion_functions_aggregate_common::accumulator::{ - AccumulatorArgs, AccumulatorFactoryFunction, StateFieldsArgs, + AccumulatorArgs, AccumulatorFactoryFunction, BlockedAccumulatorArgs, StateFieldsArgs, }; use crate::expr::{AggregateFunction, WindowFunction}; diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index ed9b4df4ce2bd..491a8787203f8 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -104,10 +104,15 @@ pub use datafusion_doc::{ pub use datafusion_expr_common::accumulator::{ Accumulator, AggregateMetric, AggregateMetricRecorder, AggregateMetrics, }; +pub use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupSelection, BlockedGroupsAccumulator, BlocksIndex, +}; +pub use datafusion_expr_common::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::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; pub use datafusion_expr_common::signature::{ diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 7358568f0afe1..5041c2399b52c 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,7 +38,8 @@ use crate::expr::{ schema_name_from_sorts, }; use crate::function::{ - AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs, + AccumulatorArgs, AggregateFunctionSimplification, BlockedAccumulatorArgs, + StateFieldsArgs, }; use crate::groups_accumulator::GroupsAccumulator; use crate::udf_eq::UdfEq; @@ -270,6 +272,22 @@ 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, @@ -669,6 +687,35 @@ 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. /// @@ -1622,6 +1669,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 8db0ab4133dc0..2ec749e59f77f 100644 --- a/datafusion/functions-aggregate-common/src/accumulator.rs +++ b/datafusion/functions-aggregate-common/src/accumulator.rs @@ -20,6 +20,7 @@ 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 @@ -80,6 +81,41 @@ 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..5c5e94e431309 100644 --- a/datafusion/functions-aggregate-common/src/aggregate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate.rs @@ -16,6 +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; 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..c516ba7fc085c --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/blocked_groups_accumulator.rs @@ -0,0 +1,251 @@ +// 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`] + +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/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..2d0791930ef75 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -34,14 +34,17 @@ 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::blocked_helpers::BlockedVec; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupSelection, + Accumulator, AggregateUDFImpl, BlockedEmitTo, BlockedGroupSelection, + BlockedGroupsAccumulator, BlocksIndex, Documentation, EmitTo, Expr, GroupSelection, GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility, WindowFunctionDefinition, expr::WindowFunction, function::{AccumulatorArgs, StateFieldsArgs}, utils::{AggregateOrderSensitivity, format_state_name}, }; +use datafusion_functions_aggregate_common::accumulator::BlockedAccumulatorArgs; use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::{ count_distinct::Bitmap65536DistinctCountAccumulator, @@ -392,6 +395,19 @@ 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 } @@ -824,6 +840,213 @@ 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/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index 6d95d8ea12bd8..2a2e67ad944c7 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -54,10 +54,11 @@ use datafusion_expr::expr::{ use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{AggregateUDF, Expr, ReversedUDAF, SetMonotonicity}; use datafusion_expr_common::accumulator::{Accumulator, AggregateMetrics}; +use datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator; use datafusion_expr_common::groups_accumulator::GroupsAccumulator; use datafusion_expr_common::type_coercion::aggregates::check_arg_count; use datafusion_functions_aggregate_common::accumulator::{ - AccumulatorArgs, StateFieldsArgs, + AccumulatorArgs, BlockedAccumulatorArgs, StateFieldsArgs, }; use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -139,6 +140,7 @@ pub struct AggregateExprBuilder { is_distinct: bool, /// Whether the expression is reversed is_reversed: bool, + batch_size: Option, } impl AggregateExprBuilder { @@ -155,6 +157,7 @@ impl AggregateExprBuilder { ignore_nulls: false, is_distinct: false, is_reversed: false, + batch_size: None, } } @@ -261,6 +264,7 @@ impl AggregateExprBuilder { ignore_nulls, is_distinct, is_reversed, + batch_size, } = self; assert_or_internal_err!(!args.is_empty(), "args should not be empty"); @@ -325,6 +329,7 @@ impl AggregateExprBuilder { input_fields: input_exprs_fields, is_reversed, is_nullable, + batch_size, }) } @@ -374,6 +379,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 @@ -562,7 +572,13 @@ 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); @@ -670,6 +686,7 @@ pub struct AggregateFunctionExpr { is_reversed: bool, input_fields: Vec, is_nullable: bool, + batch_size: Option, } impl AggregateFunctionExpr { @@ -728,6 +745,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 @@ -827,7 +848,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); } @@ -943,6 +965,56 @@ 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) + } + /// Creates a groups accumulator and supplies optional aggregate-owned metrics. pub fn create_groups_accumulator_with_metrics( &self, @@ -1022,7 +1094,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 { @@ -1089,6 +1162,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 b25caf815eed1..7f56695fe545b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -26,9 +26,13 @@ use arrow::array::{ use arrow::compute::{filter_record_batch, prep_null_mask_filter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{Result, internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::{AggregateMetrics, EmitTo, GroupsAccumulator}; +use datafusion_expr::{AggregateMetrics, GroupsAccumulator}; +use datafusion_expr_common::blocked_groups_accumulator::{ + BlockedEmitTo, BlockedGroupsAccumulator, BlocksIndex, +}; +use datafusion_functions_aggregate_common::aggregate::blocked_groups_accumulator::BlockedGroupsAccumulatorAdapter; use datafusion_physical_expr::GroupsAccumulatorAdapter; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use log::debug; @@ -36,7 +40,7 @@ use log::debug; use crate::PhysicalExpr; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, GroupValues, new_group_values, + BlockedGroupValues, GroupByMetrics, new_blocked_group_values, }; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -80,6 +84,38 @@ pub(in crate::aggregates) 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(in crate::aggregates) fn create_blocked_group_accumulator( + agg_expr: &Arc, + block_size: usize, + metrics: Arc, +) -> 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" + ); + } + // TODO - native blocked accumulators do not receive `metrics` yet + 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, metrics)?; + Ok(Box::new(BlockedGroupsAccumulatorAdapter::new( + group_acc, block_size, + ))) + } +} + /// Grouped hash table shared by the partial and final paths. /// /// While building, it consumes input batches and updates group / accumulator @@ -167,8 +203,11 @@ impl AggregateHashTable { .zip(filters) .zip(metrics.submetrics.iter()) .map(|(((agg_expr, arguments), filter), submetrics)| { - let accumulator = - create_group_accumulator(agg_expr, Arc::clone(submetrics))?; + let accumulator = create_blocked_group_accumulator( + agg_expr, + batch_size, + Arc::clone(submetrics), + )?; Ok(HashAggregateAccumulator::new( Arc::clone(agg_expr), arguments, @@ -180,7 +219,8 @@ 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)?; Ok(Self { group_by_metrics: metrics.group_by, @@ -292,7 +332,6 @@ 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 = @@ -304,22 +343,39 @@ impl AggregateHashTable { // 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)?; + 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() { - columns.extend(accumulator_metrics.time( + let acc_cols = accumulator_metrics.time( idx, accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), - )?); + || 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>(columns) + Ok::<_, datafusion_common::DataFusionError>(blocked_columns) })?; - let batch = RecordBatch::try_new(output_schema, columns)?; - debug_assert!(batch.num_rows() > 0); - MaterializedAggregateOutput::new(batch) + let mut total_memory = 0; + + let batches = blocked_columns + .into_iter() + .map(|cols| { + let batch = + RecordBatch::try_new(Arc::clone(&output_schema), 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), @@ -330,13 +386,13 @@ impl AggregateHashTable { } }; - 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 { @@ -370,31 +426,45 @@ 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( - idx, - AccumulatorPhase::State, - || acc.state(EmitTo::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()); + }, + ); } - 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 @@ -403,7 +473,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 { @@ -535,7 +605,7 @@ pub(super) struct HashAggregateAccumulator { filter: Option>, /// Accumulator state for all groups for one aggregate expression. - accumulator: Box, + accumulator: Box, /// Optional internal metrics owned by this aggregate expression. submetrics: Arc, @@ -555,7 +625,7 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator; pub(super) type AggregateBatchFn = fn( &mut AggregateAccumulator, &CompactedAccumulatorArgs, - &[usize], + &[BlocksIndex], usize, ) -> Result<()>; @@ -566,7 +636,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 { @@ -610,13 +680,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. /// @@ -644,33 +714,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 } } @@ -680,9 +753,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() { @@ -715,7 +788,7 @@ impl HashAggregateAccumulator { aggregate_expr: Arc, arguments: Vec>, filter: Option>, - accumulator: Box, + accumulator: Box, submetrics: Arc, ) -> Self { Self { @@ -728,10 +801,13 @@ 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, Arc::clone(&self.submetrics))?; + let accumulator = create_blocked_group_accumulator( + &self.aggregate_expr, + self.accumulator.batch_size(), + Arc::clone(&self.submetrics), + )?; Ok(Self::new( Arc::clone(&self.aggregate_expr), self.arguments.clone(), @@ -835,7 +911,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 @@ -854,7 +930,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()); @@ -865,21 +941,26 @@ 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> { - Ok(vec![self.evaluate(emit_to)?]) + emit_to: BlockedEmitTo, + ) -> Result>> { + // 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 /// 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) } @@ -939,7 +1020,9 @@ 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); @@ -947,14 +1030,18 @@ 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()) ); } @@ -1048,8 +1135,11 @@ mod tests { .alias("SUM(value)") .build()?, ); - let accumulator = - create_group_accumulator(&aggregate_expr, Arc::clone(&submetrics))?; + let accumulator = create_blocked_group_accumulator( + &aggregate_expr, + 1, + Arc::clone(&submetrics), + )?; Ok(HashAggregateAccumulator::new( aggregate_expr, vec![argument], diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index fd9e92635fc65..1347f14826779 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,29 +21,30 @@ use std::marker::PhantomData; use std::sync::Arc; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion_common::Result; -use datafusion_common::assert_or_internal_err; -use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::{AggregateMetrics, EmitTo}; - use crate::InputOrderMode; use crate::PhysicalExpr; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, GroupValues, new_group_values, + BlockedGroupValues, GroupByMetrics, new_blocked_group_values, }; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, }; +use arrow::compute::concat_batches; +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::{AggregateMetrics, EmitTo}; +use datafusion_expr_common::blocked_groups_accumulator::{BlockedEmitTo, BlocksIndex}; use super::AggregateTableMetrics; use super::common::{ AggregateAccumulator, AggregateBatchFn, AggregateHashTable, EvaluatedAggregateBatch, - MaterializeAccumulatorFn, create_group_accumulator, + MaterializeAccumulatorFn, create_blocked_group_accumulator, }; #[derive(Clone)] @@ -165,10 +166,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. /// @@ -199,9 +200,10 @@ 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, @@ -214,8 +216,11 @@ impl OrderedAggregateTable { .zip(filters) .zip(metrics.submetrics.iter()) .map(|(((agg_expr, arguments), filter), submetrics)| { - let accumulator = - create_group_accumulator(agg_expr, Arc::clone(submetrics))?; + let accumulator = create_blocked_group_accumulator( + agg_expr, + batch_size, + Arc::clone(submetrics), + )?; Ok(AggregateAccumulator::new( Arc::clone(agg_expr), arguments, @@ -328,28 +333,43 @@ 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( - idx, - AccumulatorPhase::State, - || acc.state(EmitTo::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()); + }, + ); } - 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 @@ -359,7 +379,7 @@ impl OrderedAggregateTable { self.buffer.group_indices.shrink_to_fit(); self.buffer.group_ordering.reset(); - Ok(Some(batch)) + Ok(batches) } /// Aggregates one evaluated input batch after selecting the mode-specific @@ -391,7 +411,7 @@ impl OrderedAggregateTable { total_num_groups, )?; } - Ok::<_, datafusion_common::DataFusionError>(total_num_groups) + Ok::<_, DataFusionError>(total_num_groups) })?; group_by_metrics.time_aggregation(|| { @@ -411,7 +431,7 @@ impl OrderedAggregateTable { ) })?; } - Ok::<(), datafusion_common::DataFusionError>(()) + Ok::<(), DataFusionError>(()) })?; } @@ -456,9 +476,41 @@ impl OrderedAggregateTable { materialize_accumulator_fn: MaterializeAccumulatorFn, accumulator_phase: AccumulatorPhase, ) -> Result { + // Translate into blocked emits: full `NextBlock`s followed by a + // `First(rem)` tail, since `BlockedEmitTo::First` must be < block size. + let blocked_emits = match emit_to { + EmitTo::All => vec![BlockedEmitTo::All], + EmitTo::First(n) => { + let mut emits = vec![BlockedEmitTo::NextBlock; n / self.batch_size]; + if n % self.batch_size != 0 { + emits.push(BlockedEmitTo::First(n % self.batch_size)); + } + emits + } + }; + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); let output = self.group_by_metrics.time_emitting(|| { - let mut output = self.buffer.group_values.emit(emit_to)?; + let mut batches = vec![]; + for blocked_emit in blocked_emits { + let mut blocks = self.buffer.group_values.emit(blocked_emit)?; + for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { + let acc_blocks = + accumulator_metrics.time(idx, accumulator_phase, || { + materialize_accumulator_fn(acc, blocked_emit) + })?; + blocks.iter_mut().zip(acc_blocks).for_each(|(output, acc)| { + output.extend(acc); + }); + } + for cols in blocks { + batches.push(RecordBatch::try_new( + Arc::clone(&self.output_schema), + cols, + )?); + } + } + // EOF can also emit a prefix when a caller limits its batch size, // but the completed ordering state no longer tracks group indexes. if let EmitTo::First(n) = emit_to @@ -466,18 +518,12 @@ impl OrderedAggregateTable { { self.buffer.group_ordering.remove_groups(n); } - - for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { - output.extend(accumulator_metrics.time( - idx, - accumulator_phase, - || materialize_accumulator_fn(acc, emit_to), - )?); - } - Ok::<_, datafusion_common::DataFusionError>(output) + Ok::<_, DataFusionError>(batches) })?; - let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + // ponytail: concat undoes blocking for large completed prefixes; pass + // the blocks through to the output stage if that copy shows up. + let batch = concat_batches(&self.output_schema, &output)?; debug_assert!(batch.num_rows() > 0); Ok(batch) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 397b766f41697..b53e9e13a5a37 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -22,7 +22,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; -use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; +use crate::aggregates::group_values::{AccumulatorPhase, new_blocked_group_values}; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{AggregateExec, evaluate_group_by}; @@ -78,7 +78,11 @@ 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..f9e53612014c2 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/blocked.rs @@ -0,0 +1,290 @@ +use super::single_group_by::{ + blocked_primitive::BlockedGroupValuesPrimitive, +}; +use crate::aggregates::group_values::new_group_values; +use crate::aggregates::order::GroupOrdering; +use arrow::array::{ArrayRef, downcast_primitive}; +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; + +/// Stores the group values during hash aggregation. +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() + ); + + if self.len() == n { + self.inner.emit(EmitTo::All) + } else { + self.inner.emit(EmitTo::First(n)) + } + } + + 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(); + + macro_rules! downcast_helper { + ($t:ty, $d:ident) => { + return Ok(Box::new(BlockedGroupValuesPrimitive::<$t>::new( + $d.clone(), + block_size, + ))) + }; + } + + downcast_primitive! { + d => (downcast_helper, d), + _ => {} + } + } + + 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 80e6b335e3188..f4d444a29af31 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -46,9 +46,12 @@ use crate::aggregates::{ order::GroupOrdering, }; +mod blocked; mod metrics; mod null_builder; +pub use blocked::{BlockedGroupValues, new_blocked_group_values}; + pub(crate) use metrics::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, GroupByMetrics, aggregate_sub_metrics, 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..94fe32e4127f4 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/blocked_primitive.rs @@ -0,0 +1,295 @@ +// 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::{BlockedGroupValues, HashValue}; +use arrow::array::{ + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, + cast::AsArray, +}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::DataType; +use datafusion_common::Result; +use datafusion_common::hash_utils::RandomState; +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; + +/// 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, 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, 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>> { + self.map = HashTable::with_capacity(128); + + let all = self.values.take_all(); + let mut blocks = Vec::with_capacity(all.len()); + + 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); + }; + + // If nothing left + let null_group = if self.values.is_empty() { + 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..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,6 +17,7 @@ //! `GroupValues` implementations for single group by cases +pub(crate) mod blocked_primitive; pub(crate) mod boolean; pub(crate) mod bytes; pub(crate) mod bytes_view; diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index b9a6d2833ab7b..e92931c8ea21b 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -72,6 +72,7 @@ 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::BlocksIndex; use futures::ready; use futures::stream::{Stream, StreamExt}; use log::debug; @@ -528,7 +529,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, @@ -905,12 +906,16 @@ 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, )?; } @@ -1140,9 +1145,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 2ef7a6463f933..19964d4feac41 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}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; @@ -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, @@ -313,11 +309,14 @@ 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(); @@ -431,11 +430,16 @@ impl PartialHashAggregateStream { &mut self, // After each incremental emitting step, the `remaining_groups` will be updated // with batch slicing. - mut remaining_groups: RecordBatch, + 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. @@ -443,55 +447,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); + // 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; } - self.reduction_factor.add_part(remaining_groups.num_rows()); - debug_assert!(remaining_groups.num_rows() > 0); - - // We are no longer holding on the batch while slicing, so release the memory. - // The memory will now equal to the hash table size - self.reservation.try_shrink(remaining_groups_memory)?; - - emitter - .emit(remaining_groups.record_output(&self.baseline_metrics)) - .await; - Ok(()) } @@ -751,7 +732,7 @@ impl FinalHashAggregateStream { // Go to the next state to perform spilling the aggregated // groups so far. let result = hash_table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -790,7 +771,7 @@ impl FinalHashAggregateStream { // Input was exhausted after spilling. Spill the last in-memory run hash_table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch))?; // Construct the ordered input used to merge all spill files. diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 147e4e0d6f18d..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..88eec37e41049 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,8 @@ 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 +288,16 @@ 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 +320,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 +343,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 3a09dbd97d704..9a11f26bae617 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -444,7 +444,7 @@ impl OrderedFinalAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let mut result = table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -495,7 +495,7 @@ impl OrderedFinalAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let replay = match table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index dd845353c77eb..125c5510d1b7e 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -19,6 +19,7 @@ use std::sync::Arc; +use arrow::compute::concat_batches; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result}; @@ -318,10 +319,12 @@ impl Aggregating { if matches!(self.table.group_ordering(), GroupOrdering::Full(_)) { return Err(oom); } - let Some(batch) = self.table.take_state_batch()? else { + let batches = self.table.take_all_state_batch()?; + if batches.is_empty() { return Err(oom); - }; - Some(batch) + } + // ponytail: concat so the output stage can slice one batch + Some(concat_batches(&batches[0].schema(), &batches)?) } 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 97afab093a872..c045c6af50458 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -409,7 +409,7 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let mut result = table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -456,7 +456,7 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let replay = match table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d8ea2f711eb6c..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, @@ -107,9 +104,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, @@ -212,7 +207,6 @@ impl PartialReduceHashAggregateStream { Ok(Self { schema, input, - batch_size, baseline_metrics, reservation, early_emit_count, @@ -318,21 +312,24 @@ 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 +347,29 @@ 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 288bf025e4e39..84438716a8b59 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -440,7 +440,7 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let mut result = hash_table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated @@ -488,7 +488,7 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let replay = match hash_table - .take_state_batch() + .take_all_state_batch() .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index dea4fe1efd11d..6f1c65d6ed46e 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -41,7 +41,7 @@ use crate::{InputOrderMode, SendableRecordBatchStream}; /// /// Every aggregation stream that spills does so the same way. Each spill event /// drains all currently buffered groups as intermediate state (see -/// `take_state_batch` on the aggregate tables), sorts them by the full group +/// `take_all_state_batch` on the aggregate tables), sorts them by the full group /// key, and writes them to one spill file. After the original input ends, all /// files are merged and replayed through an [`OrderedFinalAggregateStream`], /// which merges the states and evaluates the final aggregate values. @@ -208,37 +208,37 @@ impl AggregateSpill { !self.spills.is_empty() } - /// Sorts `state_batch`, the intermediate state of all currently buffered - /// groups (`None` if there are no groups), and writes it as one spill file. - /// Memory reservation should be updated by the caller. + /// Sorts each of `state_batches`, the intermediate state of all currently + /// buffered groups (empty if there are no groups), and writes each as one + /// spill file. Memory reservation should be updated by the caller. pub(super) fn sort_and_spill( &mut self, - state_batch: Option, + state_batches: Vec, ) -> Result<()> { - let Some(state_batch) = state_batch else { - return Ok(()); - }; - - let sorted_iter = IncrementalSortIterator::new( - state_batch, - self.spill_expr.clone(), - self.batch_size, - ); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - self.label, - )?; + // TODO - this creates one (small) spill file per block; sort the + // blocks together instead of one by one. + for state_batch in state_batches { + let sorted_iter = IncrementalSortIterator::new( + state_batch, + self.spill_expr.clone(), + self.batch_size, + ); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + self.label, + )?; - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("{}: produced an empty spill", self.label); - }; + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("{}: produced an empty spill", self.label); + }; - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + } Ok(()) }