diff --git a/datafusion/common/src/utils/blocked_vec.rs b/datafusion/common/src/utils/blocked_vec.rs new file mode 100644 index 0000000000000..792f8bccbc3a5 --- /dev/null +++ b/datafusion/common/src/utils/blocked_vec.rs @@ -0,0 +1,381 @@ +// 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. + +//! Per-group accumulator state, held in one allocation while it is small and in +//! fixed-size blocks once it is large. + +use std::mem::size_of; + +/// Groups per block once blocked: 2 MB for an eight byte state, and a whole number of +/// output batches, so a block can be handed to Arrow as a buffer of its own. +/// +/// Counted in groups rather than bytes so that two states of different element types +/// (`avg`'s sums and counts, say) always agree on where a block starts. +pub const BLOCK_LEN: usize = 1 << 18; + +/// Groups a state may reach before it is worth holding in blocks. +/// +/// Below this the state is one plain `Vec`: growing it copies at most this many elements +/// in total, which is not worth paying a second load on every group update to avoid. +/// +/// Equal to [`BLOCK_LEN`] so that the flat state is exactly one block when it switches, +/// and becomes block zero without being copied or split. +pub const THRESHOLD_LEN: usize = BLOCK_LEN; + +const BLOCK_SHIFT: u32 = BLOCK_LEN.trailing_zeros(); +const BLOCK_MASK: usize = BLOCK_LEN - 1; + +/// The block and offset a group index falls in. +#[inline] +pub fn block_offset(index: usize) -> (usize, usize) { + (index >> BLOCK_SHIFT, index & BLOCK_MASK) +} + +enum Storage { + /// One allocation. Growing copies, which is cheap while it is small. + Flat(Vec), + /// Fixed-size blocks. Growing appends a block and never moves what is already there. + Blocked(Vec>), +} + +/// A growable sequence of per-group state, addressed by group index. +/// +/// Growing one allocation copies everything already in it, and for a grouping with +/// millions of groups that copying is a large part of what the aggregate does. Reaching +/// through a block index costs a second load on *every* group update, though, so blocks +/// are only worth it once the copying they avoid outweighs the lookups they add: this +/// stays flat up to [`THRESHOLD_LEN`] groups and blocks above it. +/// +/// Callers should take [`Self::storage_mut`] once and loop inside the arm rather than +/// indexing through this type per group, so the representation is resolved once per +/// batch instead of once per update. +pub struct BlockedVec { + storage: Storage, + len: usize, +} + +impl Default for BlockedVec { + fn default() -> Self { + Self::new() + } +} + +impl std::fmt::Debug for BlockedVec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BlockedVec") + .field("len", &self.len) + .field("blocked", &self.is_blocked()) + .finish() + } +} + +/// A borrowed view of the state, taken once per batch so the update loop does not +/// re-check which representation is in use on every group. +pub enum StorageMut<'a, T> { + Flat(&'a mut [T]), + Blocked(&'a mut [Vec]), +} + +/// [`StorageMut`] for readers. +pub enum StorageRef<'a, T> { + Flat(&'a [T]), + Blocked(&'a [Vec]), +} + +impl BlockedVec { + pub fn new() -> Self { + Self { + storage: Storage::Flat(Vec::new()), + len: 0, + } + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Whether the state has moved to blocks. + pub fn is_blocked(&self) -> bool { + matches!(self.storage, Storage::Blocked(_)) + } + + /// Bytes held, counting space reserved but not filled. + pub fn capacity_bytes(&self) -> usize { + let elements = match &self.storage { + Storage::Flat(values) => values.capacity(), + Storage::Blocked(blocks) => blocks.iter().map(|b| b.capacity()).sum(), + }; + elements * size_of::() + } + + /// Resolves the representation once, so a reading loop is free of the check. + #[inline] + pub fn storage(&self) -> StorageRef<'_, T> { + match &self.storage { + Storage::Flat(values) => StorageRef::Flat(values.as_slice()), + Storage::Blocked(blocks) => StorageRef::Blocked(blocks.as_slice()), + } + } + + /// Reads one group without a bounds check. + /// + /// # Safety + /// `index` must be less than [`Self::len`]. + #[inline] + pub unsafe fn get_unchecked(&self, index: usize) -> &T { + match &self.storage { + Storage::Flat(values) => unsafe { values.get_unchecked(index) }, + Storage::Blocked(blocks) => { + let (block, offset) = block_offset(index); + unsafe { blocks.get_unchecked(block).get_unchecked(offset) } + } + } + } + + /// Appends one group. + pub fn push(&mut self, value: T) { + if self.len == THRESHOLD_LEN && !self.is_blocked() { + self.switch_to_blocked(); + } + match &mut self.storage { + Storage::Flat(values) => values.push(value), + Storage::Blocked(blocks) => { + let (block, _) = block_offset(self.len); + if block == blocks.len() { + blocks.push(Vec::with_capacity(BLOCK_LEN)); + } + blocks[block].push(value); + } + } + self.len += 1; + } + + /// Resolves the representation once, so the caller's loop is free of the check. + #[inline] + pub fn storage_mut(&mut self) -> StorageMut<'_, T> { + match &mut self.storage { + Storage::Flat(values) => StorageMut::Flat(values.as_mut_slice()), + Storage::Blocked(blocks) => StorageMut::Blocked(blocks.as_mut_slice()), + } + } + + /// Grows to `total_num_groups`, filling with `value`. + pub fn resize(&mut self, total_num_groups: usize, value: T) { + if total_num_groups <= self.len { + self.truncate(total_num_groups); + return; + } + if total_num_groups > THRESHOLD_LEN && !self.is_blocked() { + self.switch_to_blocked(); + } + match &mut self.storage { + Storage::Flat(values) => values.resize(total_num_groups, value), + Storage::Blocked(blocks) => { + let mut len = self.len; + while len < total_num_groups { + let (block, offset) = block_offset(len); + if block == blocks.len() { + blocks.push(Vec::with_capacity(BLOCK_LEN)); + } + let take = std::cmp::min(BLOCK_LEN - offset, total_num_groups - len); + blocks[block].resize(offset + take, value.clone()); + len += take; + } + } + } + self.len = total_num_groups; + } + + /// Hands the single allocation over as block zero. Paid once, at the threshold. + /// + /// The threshold is one block, so what is already stored fits a block exactly and + /// nothing has to be copied or split; it is only given room to fill that block out. + fn switch_to_blocked(&mut self) { + let Storage::Flat(values) = &mut self.storage else { + return; + }; + let mut rest = std::mem::take(values); + if rest.len() <= BLOCK_LEN { + // The usual case: growth switches at exactly one block, so the allocation + // is handed over as it stands and only given room to fill the block out. + rest.reserve_exact(BLOCK_LEN - rest.len()); + self.storage = Storage::Blocked(vec![rest]); + return; + } + // `take_first` can leave more than a block behind, which has to be split. + let mut blocks: Vec> = Vec::with_capacity(rest.len().div_ceil(BLOCK_LEN)); + while rest.len() > BLOCK_LEN { + let tail = rest.split_off(BLOCK_LEN); + blocks.push(rest); + rest = tail; + } + rest.reserve_exact(BLOCK_LEN - rest.len()); + blocks.push(rest); + self.storage = Storage::Blocked(blocks); + } + + fn truncate(&mut self, len: usize) { + if len >= self.len { + return; + } + match &mut self.storage { + Storage::Flat(values) => values.truncate(len), + Storage::Blocked(blocks) => { + blocks.truncate(len.div_ceil(BLOCK_LEN)); + if let Some(last) = blocks.last_mut() { + let offset = len & BLOCK_MASK; + if offset != 0 { + last.truncate(offset); + } + } + } + } + self.len = len; + } + + #[inline] + pub fn get(&self, index: usize) -> Option<&T> { + match &self.storage { + Storage::Flat(values) => values.get(index), + Storage::Blocked(blocks) => { + let (block, offset) = block_offset(index); + blocks.get(block)?.get(offset) + } + } + } + + /// Everything in one allocation, leaving this empty. + /// + /// While flat this hands over the allocation as it stands. Once blocked it copies, + /// which is what a per-block emit would avoid. + pub fn take_contiguous(&mut self) -> Vec { + let storage = std::mem::replace(&mut self.storage, Storage::Flat(Vec::new())); + self.len = 0; + match storage { + Storage::Flat(values) => values, + Storage::Blocked(mut blocks) => match blocks.len() { + 0 => Vec::new(), + 1 => blocks.pop().unwrap_or_default(), + _ => { + let mut out = + Vec::with_capacity(blocks.iter().map(|b| b.len()).sum()); + for block in &blocks { + out.extend(block.iter().cloned()); + } + out + } + }, + } + } + + /// Removes the first `n` elements, shifting the rest down. + pub fn take_first(&mut self, n: usize) -> Vec { + let mut all = self.take_contiguous(); + let rest = all.split_off(std::cmp::min(n, all.len())); + let len = rest.len(); + self.storage = Storage::Flat(rest); + self.len = len; + if len > THRESHOLD_LEN { + self.switch_to_blocked(); + } + all + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stays_flat_below_threshold() { + let mut v: BlockedVec = BlockedVec::new(); + v.resize(THRESHOLD_LEN, 0); + assert!(!v.is_blocked(), "should still be one allocation"); + v.resize(THRESHOLD_LEN + 1, 0); + assert!(v.is_blocked(), "should have switched to blocks"); + } + + #[test] + fn values_survive_the_switch() { + let mut v: BlockedVec = BlockedVec::new(); + let n = THRESHOLD_LEN + BLOCK_LEN + 7; + v.resize(1000, 0); + for i in 0..1000 { + unsafe { *v.get_unchecked_mut_for_test(i) = i as u64 }; + } + v.resize(n, 0); + for i in 1000..n { + unsafe { *v.get_unchecked_mut_for_test(i) = i as u64 }; + } + assert!(v.is_blocked()); + assert_eq!(v.len(), n); + for i in (0..n).step_by(997) { + assert_eq!(v.get(i).copied(), Some(i as u64), "at {i}"); + } + let flat = v.take_contiguous(); + assert_eq!(flat.len(), n); + assert!(flat.iter().enumerate().all(|(i, x)| *x == i as u64)); + } + + #[test] + fn resize_keeps_existing_values() { + let mut v: BlockedVec = BlockedVec::new(); + v.resize(10, 1); + v.resize(THRESHOLD_LEN + 5, 2); + assert_eq!(v.get(0).copied(), Some(1)); + assert_eq!(v.get(9).copied(), Some(1)); + assert_eq!(v.get(10).copied(), Some(2)); + assert_eq!(v.len(), THRESHOLD_LEN + 5); + } + + #[test] + fn take_first_shifts_down() { + let mut v: BlockedVec = BlockedVec::new(); + let n = THRESHOLD_LEN + 100; + v.resize(n, 0); + for i in 0..n { + unsafe { *v.get_unchecked_mut_for_test(i) = i as u64 }; + } + let taken = v.take_first(50); + assert_eq!(taken.len(), 50); + assert!(taken.iter().enumerate().all(|(i, x)| *x == i as u64)); + assert_eq!(v.len(), n - 50); + assert_eq!(v.get(0).copied(), Some(50)); + assert_eq!(v.get(v.len() - 1).copied(), Some(n as u64 - 1)); + } + + impl BlockedVec { + /// Indexing one group at a time, for tests only; the accumulators go through + /// [`BlockedVec::storage_mut`] instead. + /// + /// # Safety + /// `index` must be less than [`BlockedVec::len`]. + unsafe fn get_unchecked_mut_for_test(&mut self, index: usize) -> &mut T { + match &mut self.storage { + Storage::Flat(values) => unsafe { values.get_unchecked_mut(index) }, + Storage::Blocked(blocks) => { + let (block, offset) = block_offset(index); + unsafe { blocks.get_unchecked_mut(block).get_unchecked_mut(offset) } + } + } + } + } +} diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index e047db39a5740..f256a8e0ff4a3 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -18,6 +18,7 @@ //! This module provides the bisect function, which implements binary search. pub(crate) mod aggregate; +pub mod blocked_vec; pub mod expr; pub mod hex; pub mod memory; diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index b5610419166df..a32f3190e9922 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -19,6 +19,7 @@ //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`] pub mod accumulate; +pub mod blocked_vec; pub mod bool_op; pub mod nulls; pub mod prim_op; diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs new file mode 100644 index 0000000000000..5b18ceda9b2a2 --- /dev/null +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs @@ -0,0 +1,32 @@ +// 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. + +//! Applying [`EmitTo`] to blocked per-group state. + +pub use datafusion_common::utils::blocked_vec::{ + BLOCK_LEN, BlockedVec, StorageMut, StorageRef, THRESHOLD_LEN, block_offset, +}; + +use datafusion_expr_common::groups_accumulator::EmitTo; + +/// [`EmitTo`] applied to a [`BlockedVec`], returning one contiguous allocation. +pub fn take_blocked(values: &mut BlockedVec, emit_to: EmitTo) -> Vec { + match emit_to { + EmitTo::All => values.take_contiguous(), + EmitTo::First(n) => values.take_first(n), + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index c5d74978664c9..80402da5b44b6 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::mem::size_of; use std::sync::Arc; use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray}; @@ -26,6 +25,8 @@ use arrow::datatypes::DataType; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; +use super::blocked_vec::{BlockedVec, StorageMut, block_offset, take_blocked}; + use super::accumulate::NullState; /// An accumulator that implements a single operation over @@ -44,7 +45,7 @@ where F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static, { /// values per group, stored as the native type - values: Vec, + values: BlockedVec, /// The output type (needed for Decimal precision and scale) data_type: DataType, @@ -66,7 +67,7 @@ where { pub fn new(data_type: &DataType, prim_fn: F) -> Self { Self { - values: vec![], + values: BlockedVec::new(), data_type: data_type.clone(), null_state: NullState::new(), starting_value: T::default_value(), @@ -96,27 +97,54 @@ where assert_eq!(values.len(), 1, "single argument to update_batch"); let values = values[0].as_primitive::(); + let Self { + values: state, + null_state, + prim_fn, + starting_value, + .. + } = self; + // update values - self.values.resize(total_num_groups, self.starting_value); + state.resize(total_num_groups, *starting_value); + // Resolve how the state is stored once, so the loop below is not charged for + // re-checking it on every value. + // // NullState dispatches / handles tracking nulls and groups that saw no values - self.null_state.accumulate( - group_indices, - values, - opt_filter, - total_num_groups, - |group_index, new_value| { - // SAFETY: group_index is guaranteed to be in bounds - let value = unsafe { self.values.get_unchecked_mut(group_index) }; - (self.prim_fn)(value, new_value); - }, - ); + match state.storage_mut() { + StorageMut::Flat(state) => null_state.accumulate( + group_indices, + values, + opt_filter, + total_num_groups, + |group_index, new_value| { + // SAFETY: group_index is guaranteed to be in bounds + let value = unsafe { state.get_unchecked_mut(group_index) }; + prim_fn(value, new_value); + }, + ), + StorageMut::Blocked(blocks) => null_state.accumulate( + group_indices, + values, + opt_filter, + total_num_groups, + |group_index, new_value| { + let (block, offset) = block_offset(group_index); + // SAFETY: group_index is guaranteed to be in bounds + let value = unsafe { + blocks.get_unchecked_mut(block).get_unchecked_mut(offset) + }; + prim_fn(value, new_value); + }, + ), + } Ok(()) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - let values = emit_to.take_needed(&mut self.values); + let values = take_blocked(&mut self.values, emit_to); let nulls = self.null_state.build(emit_to); let values = PrimitiveArray::::new(values.into(), nulls) // no copy .with_data_type(self.data_type.clone()); @@ -190,6 +218,6 @@ where Ok(vec![Arc::new(state_values)]) } fn size(&self) -> usize { - self.values.capacity() * size_of::() + self.null_state.size() + self.values.capacity_bytes() + self.null_state.size() } } diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index 75bcec491d115..7ba584c14ee36 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -45,6 +45,9 @@ use datafusion_functions_aggregate_common::aggregate::avg_distinct::{ DecimalDistinctAvgAccumulator, Float64DistinctAvgAccumulator, }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::NullState; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::{ + BlockedVec, StorageMut, block_offset, take_blocked, +}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{ filtered_null_mask, set_nulls, }; @@ -990,10 +993,10 @@ where return_data_type: DataType, /// Count per group (use u64 to make UInt64Array) - counts: Vec, + counts: BlockedVec, /// Sums per group, stored as the native type - sums: Vec, + sums: BlockedVec, /// Track nulls in the input / filters null_state: NullState, @@ -1021,8 +1024,8 @@ where Self { return_data_type: return_data_type.clone(), sum_data_type: sum_data_type.clone(), - counts: vec![], - sums: vec![], + counts: BlockedVec::new(), + sums: BlockedVec::new(), null_state: NullState::new(), avg_fn, _phantom: PhantomData, @@ -1052,26 +1055,58 @@ where self.counts.resize(total_num_groups, 0); self.sums.resize(total_num_groups, S::default_value()); - self.null_state.accumulate( - group_indices, - values, - opt_filter, - total_num_groups, - |group_index, new_value| { - // SAFETY: group_index is guaranteed to be in bounds - let sum = unsafe { self.sums.get_unchecked_mut(group_index) }; - *sum = add_avg_sum::(*sum, new_value); - - self.counts[group_index] += 1; - }, - ); + let Self { + counts, + sums, + null_state, + .. + } = self; + + // Both states switch to blocks at the same group count, so they are always in + // the same representation; resolve it once rather than per value. + match (sums.storage_mut(), counts.storage_mut()) { + (StorageMut::Flat(sums), StorageMut::Flat(counts)) => null_state.accumulate( + group_indices, + values, + opt_filter, + total_num_groups, + |group_index, new_value| { + // SAFETY: group_index is guaranteed to be in bounds + let sum = unsafe { sums.get_unchecked_mut(group_index) }; + *sum = add_avg_sum::(*sum, new_value); + // SAFETY: group_index is in bounds after the resize above + unsafe { *counts.get_unchecked_mut(group_index) += 1 }; + }, + ), + (StorageMut::Blocked(sums), StorageMut::Blocked(counts)) => null_state + .accumulate( + group_indices, + values, + opt_filter, + total_num_groups, + |group_index, new_value| { + let (block, offset) = block_offset(group_index); + // SAFETY: group_index is guaranteed to be in bounds + let sum = unsafe { + sums.get_unchecked_mut(block).get_unchecked_mut(offset) + }; + *sum = add_avg_sum::(*sum, new_value); + // SAFETY: group_index is in bounds after the resize above + unsafe { + *counts.get_unchecked_mut(block).get_unchecked_mut(offset) += + 1 + }; + }, + ), + _ => unreachable!("counts and sums switch to blocks together"), + } Ok(()) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - let counts = emit_to.take_needed(&mut self.counts); - let sums = emit_to.take_needed(&mut self.sums); + let counts = take_blocked(&mut self.counts, emit_to); + let sums = take_blocked(&mut self.sums, emit_to); let nulls = self.null_state.build(emit_to); if let Some(nulls) = &nulls { @@ -1113,10 +1148,10 @@ where fn state(&mut self, emit_to: EmitTo) -> Result> { let nulls = self.null_state.build(emit_to); - let counts = emit_to.take_needed(&mut self.counts); + let counts = take_blocked(&mut self.counts, emit_to); let counts = UInt64Array::new(counts.into(), nulls.clone()); // zero copy - let sums = emit_to.take_needed(&mut self.sums); + let sums = take_blocked(&mut self.sums, emit_to); let sums = PrimitiveArray::::new(sums.into(), nulls) // zero copy .with_data_type(self.sum_data_type.clone()); @@ -1138,31 +1173,69 @@ where let partial_sums = values[1].as_primitive::(); // update counts with partial counts self.counts.resize(total_num_groups, 0); - self.null_state.accumulate( - group_indices, - partial_counts, - None, - total_num_groups, - |group_index, partial_count| { - // SAFETY: group_index is guaranteed to be in bounds - let count = unsafe { self.counts.get_unchecked_mut(group_index) }; - *count += partial_count; - }, - ); + let Self { + counts, null_state, .. + } = self; + match counts.storage_mut() { + StorageMut::Flat(counts) => null_state.accumulate( + group_indices, + partial_counts, + None, + total_num_groups, + |group_index, partial_count| { + // SAFETY: group_index is guaranteed to be in bounds + let count = unsafe { counts.get_unchecked_mut(group_index) }; + *count += partial_count; + }, + ), + StorageMut::Blocked(blocks) => null_state.accumulate( + group_indices, + partial_counts, + None, + total_num_groups, + |group_index, partial_count| { + let (block, offset) = block_offset(group_index); + // SAFETY: group_index is guaranteed to be in bounds + let count = unsafe { + blocks.get_unchecked_mut(block).get_unchecked_mut(offset) + }; + *count += partial_count; + }, + ), + } // update sums self.sums.resize(total_num_groups, S::default_value()); - self.null_state.accumulate( - group_indices, - partial_sums, - None, - total_num_groups, - |group_index, new_value: ::Native| { - // SAFETY: group_index is guaranteed to be in bounds - let sum = unsafe { self.sums.get_unchecked_mut(group_index) }; - *sum = add_avg_sum::(*sum, new_value); - }, - ); + let Self { + sums, null_state, .. + } = self; + match sums.storage_mut() { + StorageMut::Flat(sums) => null_state.accumulate( + group_indices, + partial_sums, + None, + total_num_groups, + |group_index, new_value: ::Native| { + // SAFETY: group_index is guaranteed to be in bounds + let sum = unsafe { sums.get_unchecked_mut(group_index) }; + *sum = add_avg_sum::(*sum, new_value); + }, + ), + StorageMut::Blocked(blocks) => null_state.accumulate( + group_indices, + partial_sums, + None, + total_num_groups, + |group_index, new_value: ::Native| { + let (block, offset) = block_offset(group_index); + // SAFETY: group_index is guaranteed to be in bounds + let sum = unsafe { + blocks.get_unchecked_mut(block).get_unchecked_mut(offset) + }; + *sum = add_avg_sum::(*sum, new_value); + }, + ), + } Ok(()) } @@ -1205,8 +1278,8 @@ where } fn size(&self) -> usize { // Heap buffers - self.counts.capacity() * size_of::() - + self.sums.capacity() * size_of::() + self.counts.capacity_bytes() + + self.sums.capacity_bytes() // Vec struct overhead (ptr, len, cap) for each field + size_of::>() + size_of::>() diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 87d3adeca27ab..f48c4c558a8ca 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -29,7 +29,6 @@ use arrow::{ }, }; use datafusion_common::hash_utils::RandomState; -use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::{ HashMap, Result, ScalarValue, downcast_value, exec_err, internal_err, not_impl_err, stats::Precision, utils::expr::COUNT_STAR_EXPANSION, @@ -43,6 +42,9 @@ use datafusion_expr::{ utils::format_state_name, }; use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::blocked_vec::{ + BlockedVec, StorageMut, block_offset, take_blocked, +}; use datafusion_functions_aggregate_common::aggregate::{ count_distinct::Bitmap65536DistinctCountAccumulator, count_distinct::Bitmap65536DistinctCountAccumulatorI16, @@ -635,12 +637,14 @@ struct CountGroupsAccumulator { /// output type of count is `DataType::Int64`. Thus by using `i64` /// for the counts, the output [`Int64Array`] can be created /// without copy. - counts: Vec, + counts: BlockedVec, } impl CountGroupsAccumulator { pub fn new() -> Self { - Self { counts: vec![] } + Self { + counts: BlockedVec::new(), + } } } @@ -658,16 +662,33 @@ impl GroupsAccumulator for CountGroupsAccumulator { // Add one to each group's counter for each non null, non // filtered value self.counts.resize(total_num_groups, 0); - 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; - }, - ); + let nulls = values.logical_nulls(); + // Resolve how the state is stored once, outside the counting loop. + match self.counts.storage_mut() { + StorageMut::Flat(counts) => accumulate_indices( + group_indices, + nulls.as_ref(), + opt_filter, + |group_index| { + // SAFETY: group_index is guaranteed to be in bounds + let count = unsafe { counts.get_unchecked_mut(group_index) }; + *count += 1; + }, + ), + StorageMut::Blocked(blocks) => accumulate_indices( + group_indices, + nulls.as_ref(), + opt_filter, + |group_index| { + let (block, offset) = block_offset(group_index); + // SAFETY: group_index is guaranteed to be in bounds + let count = unsafe { + blocks.get_unchecked_mut(block).get_unchecked_mut(offset) + }; + *count += 1; + }, + ), + } Ok(()) } @@ -688,17 +709,32 @@ impl GroupsAccumulator for CountGroupsAccumulator { // Adds the counts with the partial counts self.counts.resize(total_num_groups, 0); - group_indices.iter().zip(partial_counts.iter()).for_each( - |(&group_index, partial_count)| { - self.counts[group_index] += partial_count; - }, - ); + match self.counts.storage_mut() { + StorageMut::Flat(counts) => group_indices + .iter() + .zip(partial_counts.iter()) + .for_each(|(&group_index, partial_count)| { + // SAFETY: group_index is in bounds after the resize above + unsafe { *counts.get_unchecked_mut(group_index) += partial_count }; + }), + StorageMut::Blocked(blocks) => group_indices + .iter() + .zip(partial_counts.iter()) + .for_each(|(&group_index, partial_count)| { + let (block, offset) = block_offset(group_index); + // SAFETY: group_index is in bounds after the resize above + unsafe { + *blocks.get_unchecked_mut(block).get_unchecked_mut(offset) += + partial_count + }; + }), + } Ok(()) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - let counts = emit_to.take_needed(&mut self.counts); + let counts = take_blocked(&mut self.counts, emit_to); // Count is always non null (null inputs just don't contribute to the overall values) let nulls = None; @@ -709,7 +745,7 @@ impl GroupsAccumulator for CountGroupsAccumulator { // return arrays for counts fn state(&mut self, emit_to: EmitTo) -> Result> { - let counts = emit_to.take_needed(&mut self.counts); + let counts = take_blocked(&mut self.counts, emit_to); let counts: PrimitiveArray = Int64Array::from(counts); // zero copy, no nulls Ok(vec![Arc::new(counts) as ArrayRef]) } @@ -775,7 +811,7 @@ impl GroupsAccumulator for CountGroupsAccumulator { Ok(vec![state_array]) } fn size(&self) -> usize { - self.counts.heap_size(&mut DFHeapSizeCtx::default()) + self.counts.capacity_bytes() } } @@ -941,9 +977,8 @@ mod tests { let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); acc.update_batch(&[values], &[0, 1, 2], None, 3)?; - assert!(acc.counts.capacity() > 0); - let allocated_size = acc.counts.heap_size(&mut DFHeapSizeCtx::default()); - assert_eq!(allocated_size, acc.counts.capacity() * size_of::()); + let allocated_size = acc.counts.capacity_bytes(); + assert!(allocated_size > 0); assert_eq!(acc.size(), allocated_size); assert!(acc.size() > empty_size); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index b965ccc34f653..b24fd7d648787 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -29,9 +29,7 @@ use arrow::buffer::ScalarBuffer; use arrow::datatypes::DataType; use arrow::util::bit_util::apply_bitwise_binary_op; use datafusion_common::Result; -use datafusion_common::utils::split_vec_min_alloc; -use datafusion_execution::memory_pool::proxy::VecAllocExt; -use std::iter; +use datafusion_common::utils::blocked_vec::{BlockedVec, StorageRef, block_offset}; use std::sync::Arc; /// An implementation of [`GroupColumn`] for primitive values @@ -45,7 +43,7 @@ use std::sync::Arc; #[derive(Debug)] pub struct PrimitiveGroupValueBuilder { data_type: DataType, - group_values: Vec, + group_values: BlockedVec, nulls: NullBufferBuilder, } @@ -58,11 +56,37 @@ where pub fn new(data_type: DataType) -> Self { Self { data_type, - group_values: vec![], + group_values: BlockedVec::new(), nulls: NullBufferBuilder::empty(), } } + /// The comparison loop, taking how to read a stored group as a parameter so it is + /// compiled once per representation rather than branching per row. + #[inline] + fn compare_non_nullable( + lhs_rows: &[usize], + rhs_rows: &[usize], + array_values: &[T::Native], + equal_to_results: &BooleanBufferBuilder, + cmp_buf: &mut [u8], + left_at: impl Fn(usize) -> T::Native, + ) { + for (i, (&lhs_row, &rhs_row)) in lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(i) { + continue; + } + let left = left_at(lhs_row); + let right = unsafe { *array_values.get_unchecked(rhs_row) }; + // `left` was already canonicalized on append; canonicalize the + // input so ±0 (and any future equivalence class) compares equal. + if left.is_eq(right.canonicalize()) { + cmp_buf[i / 8] |= 1 << (i % 8); + } + } + } + fn vectorized_equal_to_non_nullable( &self, lhs_rows: &[usize], @@ -81,26 +105,28 @@ where let num_bytes = n.div_ceil(8); let mut cmp_buf = vec![0u8; num_bytes]; - for (i, (&lhs_row, &rhs_row)) in lhs_rows.iter().zip(rhs_rows.iter()).enumerate() - { - if !equal_to_results.get_bit(i) { - continue; - } - let left = if cfg!(debug_assertions) { - self.group_values[lhs_row] - } else { - unsafe { *self.group_values.get_unchecked(lhs_row) } - }; - let right = if cfg!(debug_assertions) { - array_values[rhs_row] - } else { - unsafe { *array_values.get_unchecked(rhs_row) } - }; - // `left` was already canonicalized on append; canonicalize the - // input so ±0 (and any future equivalence class) compares equal. - if left.is_eq(right.canonicalize()) { - cmp_buf[i / 8] |= 1 << (i % 8); - } + // Resolve how the group values are stored once, so the comparison loop is not + // charged for re-checking it on every row. + match self.group_values.storage() { + StorageRef::Flat(values) => Self::compare_non_nullable( + lhs_rows, + rhs_rows, + array_values, + equal_to_results, + &mut cmp_buf, + |row| unsafe { *values.get_unchecked(row) }, + ), + StorageRef::Blocked(blocks) => Self::compare_non_nullable( + lhs_rows, + rhs_rows, + array_values, + equal_to_results, + &mut cmp_buf, + |row| { + let (block, offset) = block_offset(row); + unsafe { *blocks.get_unchecked(block).get_unchecked(offset) } + }, + ), } // AND the comparison result into the existing equal_to_results bitmask @@ -124,24 +150,37 @@ where assert!(NULLABLE, "called with non-nullable input"); let array = array.as_primitive::(); - for (idx, (&lhs_row, &rhs_row)) in - lhs_rows.iter().zip(rhs_rows.iter()).enumerate() - { - if !equal_to_results.get_bit(idx) { - continue; - } - let exist_null = self.nulls.is_null(lhs_row); - let input_null = array.is_null(rhs_row); - if let Some(result) = nulls_equal_to(exist_null, input_null) { - if !result { + // Resolve how the group values are stored once, outside the loop. + let nulls = &self.nulls; + let mut compare = |left_at: &dyn Fn(usize) -> T::Native| { + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + let exist_null = nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + if !result { + equal_to_results.set_bit(idx, false); + } + continue; + } + + if !left_at(lhs_row).is_eq(array.value(rhs_row).canonicalize()) { equal_to_results.set_bit(idx, false); } - continue; } - - if !self.group_values[lhs_row].is_eq(array.value(rhs_row).canonicalize()) { - equal_to_results.set_bit(idx, false); + }; + match self.group_values.storage() { + StorageRef::Flat(values) => { + compare(&|row| unsafe { *values.get_unchecked(row) }) } + StorageRef::Blocked(blocks) => compare(&|row| { + let (block, offset) = block_offset(row); + unsafe { *blocks.get_unchecked(block).get_unchecked(offset) } + }), } } } @@ -162,7 +201,8 @@ where // Otherwise, we need to check their values } - self.group_values[lhs_row] + // SAFETY: `lhs_row` is an existing group index + unsafe { *self.group_values.get_unchecked(lhs_row) } .is_eq(array.as_primitive::().value(rhs_row).canonicalize()) } @@ -240,7 +280,7 @@ where (true, Nulls::All) => { self.nulls.append_n_nulls(rows.len()); self.group_values - .extend(iter::repeat_n(T::default_value(), rows.len())); + .resize(self.group_values.len() + rows.len(), T::default_value()); } (false, _) => { @@ -258,13 +298,13 @@ where } fn size(&self) -> usize { - self.group_values.allocated_size() + self.nulls.allocated_size() + self.group_values.capacity_bytes() + self.nulls.allocated_size() } fn build(self: Box) -> ArrayRef { let Self { data_type, - group_values, + mut group_values, nulls, } = *self; @@ -273,13 +313,16 @@ where assert!(nulls.is_none(), "unexpected nulls in non nullable input"); } - let arr = PrimitiveArray::::new(ScalarBuffer::from(group_values), nulls); + let arr = PrimitiveArray::::new( + ScalarBuffer::from(group_values.take_contiguous()), + nulls, + ); // Set timezone information for timestamp Arc::new(arr.with_data_type(data_type)) } fn take_n(&mut self, n: usize) -> ArrayRef { - let first_n = split_vec_min_alloc(&mut self.group_values, n); + let first_n = self.group_values.take_first(n); let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; Arc::new(