diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 312c5c9c184..0aa9a4c3407 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -10,19 +10,19 @@ use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use rand::seq::SliceRandom; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; +use vortex_array::aggregate_fn::GroupIds; use vortex_array::aggregate_fn::GroupedAccumulator; use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; -use vortex_array::dtype::DType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -38,6 +38,22 @@ const GROUP_COUNT: usize = 128; const GROUP_SIZE_SEED: u64 = 42; const MIN_VALUES_PER_GROUP: usize = 1; const MAX_VALUES_PER_GROUP: usize = 15; +const CARDINALITY_ELEMENT_COUNT: usize = 1 << 16; + +#[derive(Clone, Copy, Debug)] +enum IdOrder { + Clustered, + Shuffled, +} + +const CARDINALITY_ARGS: &[(usize, IdOrder)] = &[ + (128, IdOrder::Clustered), + (128, IdOrder::Shuffled), + (1 << 12, IdOrder::Clustered), + (1 << 12, IdOrder::Shuffled), + (1 << 16, IdOrder::Clustered), + (1 << 16, IdOrder::Shuffled), +]; fn random_group_sizes() -> Vec { let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); @@ -50,44 +66,40 @@ fn total_element_count(group_sizes: &[usize]) -> usize { group_sizes.iter().sum() } -fn contiguous_list_view(elements: ArrayRef, group_sizes: &[usize]) -> ArrayRef { - let mut offset = 0usize; - let offsets: Buffer = group_sizes - .iter() - .map(|&size| { - let current_offset = offset; - offset += size; - current_offset as u32 - }) - .collect(); - let sizes: Buffer = group_sizes.iter().map(|&size| size as u32).collect(); +struct DenseGroupedInput { + values: ArrayRef, + group_ids: GroupIds, +} - assert_eq!(elements.len(), total_element_count(group_sizes)); +fn dense_grouped_input(values: ArrayRef, group_sizes: &[usize]) -> DenseGroupedInput { + assert_eq!(values.len(), total_element_count(group_sizes)); - ListViewArray::try_new( - elements, - offsets.into_array(), - sizes.into_array(), - Validity::NonNullable, + let group_ids = GroupIds::from_iter( + group_sizes + .iter() + .enumerate() + .flat_map(|(group_id, &size)| std::iter::repeat_n(group_id as u32, size)), + group_sizes.len(), ) - .unwrap() - .into_array() + .unwrap(); + + DenseGroupedInput { values, group_ids } } -fn i32_nullable_all_valid_input() -> ArrayRef { +fn i32_nullable_all_valid_input() -> DenseGroupedInput { let group_sizes = random_group_sizes(); let element_count = total_element_count(&group_sizes); let values: Buffer = (0..element_count) .map(|i| (i % 1024) as i32 - 512) .collect(); let validity = Validity::from_iter(std::iter::repeat_n(true, element_count)); - contiguous_list_view( + dense_grouped_input( PrimitiveArray::new(values, validity).into_array(), &group_sizes, ) } -fn i32_clustered_nulls_input() -> ArrayRef { +fn i32_clustered_nulls_input() -> DenseGroupedInput { let group_sizes = random_group_sizes(); let element_count = total_element_count(&group_sizes); let values = (0..element_count).map(|i| { @@ -97,26 +109,26 @@ fn i32_clustered_nulls_input() -> ArrayRef { Some((i % 1024) as i32 - 512) } }); - contiguous_list_view( + dense_grouped_input( PrimitiveArray::from_option_iter(values).into_array(), &group_sizes, ) } -fn f64_all_valid_input() -> ArrayRef { +fn f64_all_valid_input() -> DenseGroupedInput { let group_sizes = random_group_sizes(); let element_count = total_element_count(&group_sizes); let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); let values: Buffer = (0..element_count) .map(|_| rng.random_range(-1000.0..1000.0)) .collect(); - contiguous_list_view( + dense_grouped_input( PrimitiveArray::new(values, Validity::NonNullable).into_array(), &group_sizes, ) } -fn f64_clustered_nulls_input() -> ArrayRef { +fn f64_clustered_nulls_input() -> DenseGroupedInput { let group_sizes = random_group_sizes(); let element_count = total_element_count(&group_sizes); let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); @@ -127,42 +139,61 @@ fn f64_clustered_nulls_input() -> ArrayRef { Some(rng.random_range(-1000.0f64..1000.0)) } }); - contiguous_list_view( + dense_grouped_input( PrimitiveArray::from_option_iter(values).into_array(), &group_sizes, ) } -fn varbinview_input() -> ArrayRef { +fn varbinview_input() -> DenseGroupedInput { let group_sizes = random_group_sizes(); let element_count = total_element_count(&group_sizes); let values: Vec = (0..element_count) .map(|i| format!("value-{i:06}")) .collect(); - contiguous_list_view( + dense_grouped_input( VarBinViewArray::from_iter_str(values.iter().map(String::as_str)).into_array(), &group_sizes, ) } -fn list_element_dtype(list_view: &ArrayRef) -> DType { - match list_view.dtype() { - DType::List(element_dtype, _) => element_dtype.as_ref().clone(), - dtype => unreachable!("expected List dtype, got {dtype}"), +fn i32_cardinality_input(group_count: usize, order: IdOrder) -> DenseGroupedInput { + let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); + let values: Buffer = (0..CARDINALITY_ELEMENT_COUNT) + .map(|_| rng.random_range(-512..512)) + .collect(); + let mut group_ids: Vec = (0..CARDINALITY_ELEMENT_COUNT) + .map(|idx| (idx * group_count / CARDINALITY_ELEMENT_COUNT) as u32) + .collect(); + if matches!(order, IdOrder::Shuffled) { + group_ids.shuffle(&mut rng); + } + + DenseGroupedInput { + values: PrimitiveArray::new(values, Validity::NonNullable).into_array(), + group_ids: GroupIds::from_buffer(Buffer::from(group_ids), group_count).unwrap(), } } -fn grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef +fn grouped_accumulator(input: &DenseGroupedInput, vtable: V) -> ArrayRef where V: AggregateFnVTable + Clone, V::Options: Default, { let mut acc = - GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view)) + GroupedAccumulator::try_new(vtable, V::Options::default(), input.values.dtype().clone()) .unwrap(); - acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) + let num_groups = input.group_ids.num_groups(); + let mut ctx = SESSION.create_execution_ctx(); + acc.accumulate(&input.values, &input.group_ids, &mut ctx) .unwrap(); - divan::black_box(acc.finish().unwrap()) + let result = acc + .finish(num_groups) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + divan::black_box(result) } #[divan::bench] @@ -197,70 +228,33 @@ fn sum_f64_clustered_nulls(bencher: Bencher) { .bench_refs(|input| grouped_accumulator(input, Sum)); } -/// Like [`grouped_accumulator`], but executes the lazy finalize result to canonical so the -/// bench measures the full cost of producing usable sums. -fn grouped_accumulator_canonical(list_view: &ArrayRef, vtable: V) -> ArrayRef -where - V: AggregateFnVTable + Clone, - V::Options: Default, -{ - let mut acc = - GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view)) - .unwrap(); - let mut ctx = SESSION.create_execution_ctx(); - acc.accumulate_list(list_view, &mut ctx).unwrap(); - let result = acc - .finish() - .unwrap() - .execute::(&mut ctx) - .unwrap() - .into_array(); - divan::black_box(result) -} - #[divan::bench] -fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) { - let input = i32_nullable_all_valid_input(); - bencher - .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); -} - -#[divan::bench] -fn canonical_sum_i32_clustered_nulls(bencher: Bencher) { +fn count_i32_clustered_nulls(bencher: Bencher) { let input = i32_clustered_nulls_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); -} - -#[divan::bench] -fn canonical_sum_f64_all_valid(bencher: Bencher) { - let input = f64_all_valid_input(); - bencher - .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, Count)); } #[divan::bench] -fn canonical_sum_f64_clustered_nulls(bencher: Bencher) { - let input = f64_clustered_nulls_input(); +fn count_varbinview(bencher: Bencher) { + let input = varbinview_input(); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator_canonical(input, Sum)); + .bench_refs(|input| grouped_accumulator(input, Count)); } -#[divan::bench] -fn count_i32_clustered_nulls(bencher: Bencher) { - let input = i32_clustered_nulls_input(); +#[divan::bench(args = CARDINALITY_ARGS)] +fn sum_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) { + let input = i32_cardinality_input(group_count, order); bencher .with_inputs(|| &input) - .bench_refs(|input| grouped_accumulator(input, Count)); + .bench_refs(|input| grouped_accumulator(input, Sum)); } -#[divan::bench] -fn count_varbinview(bencher: Bencher) { - let input = varbinview_input(); +#[divan::bench(args = CARDINALITY_ARGS)] +fn count_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) { + let input = i32_cardinality_input(group_count, order); bencher .with_inputs(|| &input) .bench_refs(|input| grouped_accumulator(input, Count)); diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 52255481a43..654b5342776 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -181,7 +181,7 @@ impl DynAccumulator for Accumulator { } // 3. Iteratively check the registry against each intermediate encoding, executing one - // step between checks. Mirrors the loop in `GroupedAccumulator::accumulate_list_view`. + // step between checks. Mirrors the loop in `GroupedAccumulator::accumulate`. // Iteration 0 re-checks the initial encoding — a redundant HashMap miss, the price of // keeping the loop body uniform. Terminates on `AnyColumnar` (Canonical or Constant) // since the vtable's `accumulate(&Columnar)` handles both cases directly. diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 8bd26dad961..4571a4e4480 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -1,19 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use arrow_buffer::ArrowNativeType; +use std::any::Any; +use std::sync::Arc; +use std::sync::OnceLock; + +use num_traits::ToPrimitive; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; @@ -22,27 +22,29 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::session::AggregateFnSessionExt; -use crate::arrays::ChunkedArray; +use crate::array::ArrayId; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::columnar::AnyColumnar; +use crate::columnar::Columnar; use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; use crate::executor::max_iterations; use crate::match_each_integer_ptype; +use crate::scalar::Scalar; +use crate::validity::Validity; /// Reference-counted type-erased grouped accumulator. pub type GroupedAccumulatorRef = Box; -/// A batch of grouped values to aggregate. -/// -/// Each outer list value is one group, and the inner element array is shared by all groups. -/// Aggregate implementations can inspect the concrete grouped representation directly, or ask for -/// derived ranges when their algorithm is expressed in terms of `(offset, size)` pairs. +/// A canonical list representation used to adapt list-shaped groups to dense group IDs. pub enum GroupedArray { /// Groups represented as a list-view array with per-group offsets and sizes. ListView(ListViewArray), @@ -63,7 +65,7 @@ impl From for GroupedArray { } impl GroupedArray { - /// The inner element array shared by all groups. + /// Return the inner element array shared by all groups. pub fn elements(&self) -> &ArrayRef { match self { Self::ListView(groups) => groups.elements(), @@ -71,7 +73,7 @@ impl GroupedArray { } } - /// Return the `(offset, size)` ranges describing each group in `elements`. + /// Return the physical element ranges for each group. pub fn group_ranges(&self, ctx: &mut ExecutionCtx) -> VortexResult { match self { Self::ListView(groups) => list_view_group_ranges(groups, ctx), @@ -87,7 +89,7 @@ impl GroupedArray { } } - /// The number of groups in this batch. + /// Return the number of groups. pub fn len(&self) -> usize { match self { Self::ListView(groups) => groups.len(), @@ -95,41 +97,50 @@ impl GroupedArray { } } - /// Returns true when this batch contains no groups. + /// Return whether there are no groups. pub fn is_empty(&self) -> bool { self.len() == 0 } - /// Returns true when every group is valid. - pub fn all_groups_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(self.group_validity(ctx)?.all_true()) - } - - unsafe fn with_elements_unchecked(&self, elements: ArrayRef) -> VortexResult { - Ok(match self { - Self::ListView(groups) => unsafe { - ListViewArray::new_unchecked( - elements, - groups.offsets().clone(), - groups.sizes().clone(), - groups.validity()?, - ) + /// Convert list-shaped groups into a values array and parallel dense group IDs. + /// + /// Null groups contribute no values. Contiguous, all-valid groups reuse the original element + /// array; arbitrary list views are gathered into group order. + pub fn dense_input(&self, ctx: &mut ExecutionCtx) -> VortexResult<(ArrayRef, GroupIds)> { + let num_groups = self.len(); + validate_num_groups(num_groups)?; + let ranges = self.group_ranges(ctx)?; + let validity = self.group_validity(ctx)?; + let mut rows = Vec::new(); + let mut ids = Vec::new(); + let mut identity = true; + + for (group, ((offset, size), valid)) in ranges.iter().zip(validity.iter()).enumerate() { + if !valid { + identity = false; + continue; } - .into(), - Self::FixedSizeList(groups) => unsafe { - FixedSizeListArray::new_unchecked( - elements, - groups.list_size(), - groups.validity()?, - groups.len(), - ) + let group = u32::try_from(group)?; + for row in offset..offset + size { + identity &= row == rows.len(); + rows.push(u64::try_from(row)?); + ids.push(group); } - .into(), - }) + } + + identity &= rows.len() == self.elements().len(); + let values = if identity { + self.elements().clone() + } else { + self.elements() + .clone() + .take(Buffer::from_iter(rows).into_array())? + }; + Ok((values, GroupIds::from_iter(ids, num_groups)?)) } } -/// The physical ranges of a grouped array. +/// The physical element ranges of a canonical grouped list array. pub enum GroupRanges { /// Explicit ranges extracted from a list-view array. ListView { @@ -146,7 +157,7 @@ pub enum GroupRanges { } impl GroupRanges { - /// The number of groups described by these ranges. + /// Return the number of groups described by these ranges. pub fn len(&self) -> usize { match self { Self::ListView { ranges } => ranges.len(), @@ -154,32 +165,259 @@ impl GroupRanges { } } - /// Returns true when there are no groups. + /// Return whether no groups are described. pub fn is_empty(&self) -> bool { self.len() == 0 } - /// Return the `(offset, size)` range for the group at `index`. fn range(&self, index: usize) -> (usize, usize) { match self { Self::ListView { ranges } => ranges[index], Self::FixedSizeList { len, size } => { - assert!(index < *len, "range index out of bounds"); - (index * size, *size) + assert!(index < *len, "group range index out of bounds"); + (index * *size, *size) } } } - /// Iterate over all `(offset, size)` group ranges. + /// Iterate over `(offset, size)` ranges. pub fn iter(&self) -> impl Iterator + '_ { (0..self.len()).map(|index| self.range(index)) } } -/// An accumulator used for computing grouped aggregates. +/// Encoded group ids parallel to a grouped aggregate input batch. /// -/// Note that the groups must be processed in order, and the accumulator does not support random -/// access to groups. +/// The array must contain non-null `u32` ordinals. The ordinals are dense state slots in +/// `0..num_groups`, not raw group keys. Range validation may require executing the encoded array, +/// so kernels that can prove the invariant from encoded metadata should avoid materializing and +/// otherwise call [`Self::validated_ids`] before indexing group state. +#[derive(Clone, Debug)] +pub struct GroupIds { + ids: ArrayRef, + num_groups: usize, + validated: Arc>>, +} + +impl GroupIds { + /// Create group ids from an encoded non-null `u32` array. + pub fn new(ids: ArrayRef, num_groups: usize) -> VortexResult { + validate_num_groups(num_groups)?; + vortex_ensure!( + ids.dtype() == &DType::Primitive(PType::U32, Nullability::NonNullable), + "Group ids must be non-nullable u32, got {}", + ids.dtype() + ); + Ok(Self { + ids, + num_groups, + validated: Arc::new(OnceLock::new()), + }) + } + + /// Create group ids from a materialized buffer, validating the dense-id invariant once. + pub fn from_buffer(ids: Buffer, num_groups: usize) -> VortexResult { + validate_group_ids(ids.as_ref(), num_groups)?; + Ok(Self { + ids: PrimitiveArray::new(ids.clone(), Validity::NonNullable).into_array(), + num_groups, + validated: Arc::new(OnceLock::from(ids)), + }) + } + + /// Create group ids from materialized values. + pub fn from_iter(ids: impl IntoIterator, num_groups: usize) -> VortexResult { + Self::from_buffer(Buffer::from_iter(ids), num_groups) + } + + /// Return the encoded ids array. + pub fn ids(&self) -> &ArrayRef { + &self.ids + } + + /// Return the number of dense group state slots. + pub fn num_groups(&self) -> usize { + self.num_groups + } + + /// Return the number of ids. + pub fn len(&self) -> usize { + self.ids.len() + } + + /// Return whether there are no ids. + pub fn is_empty(&self) -> bool { + self.ids.is_empty() + } + + /// Return the encoding id for kernel dispatch. + pub fn encoding_id(&self) -> ArrayId { + self.ids.encoding_id() + } + + /// Execute the ids to a native buffer and validate every id is in range. + /// + /// The validated buffer is cached and shared by clones, so multiple aggregate functions over + /// the same group ids pay materialization and validation at most once. + pub fn validated_ids(&self, ctx: &mut ExecutionCtx) -> VortexResult> { + if let Some(ids) = self.validated.get() { + return Ok(ids.clone()); + } + + let ids = self.ids.clone().execute::>(ctx)?; + validate_group_ids(ids.as_ref(), self.num_groups)?; + drop(self.validated.set(ids.clone())); + Ok(ids) + } +} + +/// Aggregate-owned dense state used by grouped accumulation. +pub trait GroupedState: 'static + Send { + /// Expose the concrete state container to a typed grouped kernel. + fn as_any_mut(&mut self) -> &mut dyn Any; + + /// Return the number of allocated group slots. + fn len(&self) -> usize; + + /// Return whether no group slots are allocated. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Ensure that at least `num_groups` state slots exist. + fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()>; + + /// Return whether one group has reached a terminal state. + fn is_saturated(&self, group_id: usize) -> bool; + + /// Combine one scalar partial into a group. + fn combine_scalar(&mut self, group_id: usize, partial: Scalar) -> VortexResult<()>; + + /// Read one group's partial state. + fn partial_scalar(&self, group_id: usize) -> VortexResult; + + /// Fold an array of partial states into dense groups. + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + for (row_idx, &group_id) in group_ids.iter().enumerate() { + self.combine_scalar(group_id as usize, partials.execute_scalar(row_idx, ctx)?)?; + } + Ok(()) + } + + /// Flush all partial states into an array and reset the container. + fn flush_partials(&mut self, num_groups: usize) -> VortexResult; +} + +/// Default grouped state backed by one aggregate partial value per group. +pub(crate) struct DefaultGroupedState { + vtable: V, + options: V::Options, + input_dtype: DType, + partial_dtype: DType, + partials: Vec, +} + +impl DefaultGroupedState { + pub(crate) fn new( + vtable: V, + options: V::Options, + input_dtype: DType, + partial_dtype: DType, + ) -> Self { + Self { + vtable, + options, + input_dtype, + partial_dtype, + partials: Vec::new(), + } + } +} + +impl GroupedState for DefaultGroupedState { + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn len(&self) -> usize { + self.partials.len() + } + + fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { + self.partials + .reserve(num_groups.saturating_sub(self.partials.len())); + while self.partials.len() < num_groups { + self.partials.push( + self.vtable + .empty_partial(&self.options, &self.input_dtype)?, + ); + } + Ok(()) + } + + fn is_saturated(&self, group_id: usize) -> bool { + self.vtable.is_saturated(&self.partials[group_id]) + } + + fn combine_scalar(&mut self, group_id: usize, partial: Scalar) -> VortexResult<()> { + self.vtable + .combine_partials(&mut self.partials[group_id], partial) + } + + fn partial_scalar(&self, group_id: usize) -> VortexResult { + if let Some(partial) = self.partials.get(group_id) { + self.vtable.to_scalar(partial) + } else { + let partial = self + .vtable + .empty_partial(&self.options, &self.input_dtype)?; + self.vtable.to_scalar(&partial) + } + } + + fn flush_partials(&mut self, num_groups: usize) -> VortexResult { + vortex_ensure!( + num_groups >= self.partials.len(), + "Cannot flush {} groups after accumulating {} groups", + num_groups, + self.partials.len() + ); + self.ensure_groups(num_groups)?; + + if let Some(states) = self + .vtable + .partials_to_array(&self.partials, &self.partial_dtype)? + { + vortex_ensure!( + states.dtype() == &self.partial_dtype, + "Partial array DType mismatch: expected {}, got {}", + self.partial_dtype, + states.dtype() + ); + self.partials.clear(); + return Ok(states); + } + + let mut states = builder_with_capacity(&self.partial_dtype, num_groups); + for partial in &self.partials { + states.append_scalar(&self.vtable.to_scalar(partial)?)?; + } + self.partials.clear(); + Ok(states.finish()) + } +} + +/// An accumulator used for computing aggregates over group ids. +/// +/// Group ids are caller-assigned `u32` ordinals in the dense range `0..num_groups`. Input batches +/// may repeat, omit, and reorder those ids, but every id must identify a state slot rather than a +/// raw group key. The accumulator keeps one partial state per slot, so ordered and unordered +/// grouping only differ in how the caller assigns ids. pub struct GroupedAccumulator { /// The vtable of the aggregate function. vtable: V, @@ -193,8 +431,8 @@ pub struct GroupedAccumulator { return_dtype: DType, /// The DType of the partial accumulator state. partial_dtype: DType, - /// The accumulated state for prior batches of groups. - partials: Vec, + /// Aggregate-owned dense per-group state. + state: Box, } impl GroupedAccumulator { @@ -214,6 +452,7 @@ impl GroupedAccumulator { dtype ) })?; + let state = vtable.grouped_state(&options, &dtype, &partial_dtype)?; Ok(Self { vtable, @@ -222,180 +461,131 @@ impl GroupedAccumulator { dtype, return_dtype, partial_dtype, - partials: vec![], + state, }) } -} -/// A trait object for type-erased grouped accumulators, used for dynamic dispatch when the aggregate -/// function is not known at compile time. -pub trait DynGroupedAccumulator: 'static + Send { - /// Accumulate a list of groups into the accumulator. - fn accumulate_list(&mut self, groups: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()>; - - /// Finish the accumulation and return the partial aggregate results for all groups. - /// Resets the accumulator state for the next round of accumulation. - fn flush(&mut self) -> VortexResult; + fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { + validate_num_groups(num_groups)?; - /// Finish the accumulation and return the final aggregate results for all groups. - /// Resets the accumulator state for the next round of accumulation. - fn finish(&mut self) -> VortexResult; -} - -impl DynGroupedAccumulator for GroupedAccumulator { - fn accumulate_list(&mut self, groups: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { - let elements_dtype = match groups.dtype() { - DType::List(elem, _) => elem, - DType::FixedSizeList(elem, ..) => elem, - _ => vortex_bail!( - "Input DType mismatch: expected List or FixedSizeList, got {}", - groups.dtype() - ), - }; - vortex_ensure!( - elements_dtype.as_ref() == &self.dtype, - "Input DType mismatch: expected {}, got {}", - self.dtype, - elements_dtype - ); - - // We first execute the groups until it is a ListView or FixedSizeList, since we only - // dispatch the aggregate kernel over the elements of these arrays. - let canonical = match groups.clone().execute::(ctx)? { - Columnar::Canonical(c) => c, - Columnar::Constant(c) => c.into_array().execute::(ctx)?, - }; - match canonical { - Canonical::List(groups) => self.accumulate_grouped_array(groups.into(), ctx), - Canonical::FixedSizeList(groups) => self.accumulate_grouped_array(groups.into(), ctx), - _ => vortex_panic!("We checked the DType above, so this should never happen"), - } - } - - fn flush(&mut self) -> VortexResult { - let mut states = std::mem::take(&mut self.partials); - if states.len() == 1 { - return Ok(states.pop().vortex_expect("checked one partial")); - } - Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) + self.state.ensure_groups(num_groups) } - fn finish(&mut self) -> VortexResult { - let states = self.flush()?; - let results = self.vtable.finalize(states)?; + fn try_accumulate_kernel( + &mut self, + batch: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let session = ctx.session().clone(); - vortex_ensure!( - results.dtype() == &self.return_dtype, - "Return DType mismatch: expected {}, got {}", - self.return_dtype, - results.dtype() - ); + if let Some(kernel) = session.aggregate_fns().find_grouped_kernel( + self.aggregate_fn.id(), + batch.encoding_id(), + group_ids.encoding_id(), + ) && kernel.grouped_accumulate( + &self.aggregate_fn, + batch, + group_ids, + self.state.as_any_mut(), + ctx, + )? { + return Ok(true); + } - Ok(results) + Ok(false) } -} -impl GroupedAccumulator { - fn accumulate_grouped_array( + fn accumulate_fallback( &mut self, - groups: GroupedArray, + batch: &ArrayRef, + group_ids: &[u32], ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let mut elements = groups.elements().clone(); - let session = ctx.session().clone(); + let Some((&first, rest)) = group_ids.split_first() else { + return Ok(()); + }; + let mut first = first; + let mut last = first; + for &group_id in rest { + first = first.min(group_id); + last = last.max(group_id); + } - for _ in 0..max_iterations() { - // Try a registered grouped kernel for the current element encoding. - if let Some(kernel) = session - .aggregate_fns() - .find_grouped_encoding_kernel(elements.encoding_id(), self.aggregate_fn.id()) - { - // SAFETY: we assume that elements execution is safe - let kernel_groups = unsafe { groups.with_elements_unchecked(elements.clone())? }; - if let Some(result) = - kernel.grouped_aggregate(&self.aggregate_fn, &kernel_groups, ctx)? - { - return self.push_result(result); - } - } + let first = first as usize; + let span = last as usize - first + 1; + // Stable counting-sort the rows so every group becomes a slice of one gathered array. + let mut offsets = vec![0usize; span + 1]; + for &group_id in group_ids { + offsets[group_id as usize - first + 1] += 1; + } + for idx in 1..offsets.len() { + offsets[idx] += offsets[idx - 1]; + } + + let mut cursors = offsets.clone(); + let mut permutation = vec![0u64; group_ids.len()]; + for (row_idx, &group_id) in group_ids.iter().enumerate() { + let cursor = &mut cursors[group_id as usize - first]; + permutation[*cursor] = row_idx as u64; + *cursor += 1; + } - // Try a grouped kernel for the current aggregate regardless of element encoding. - if let Some(kernel) = session - .aggregate_fns() - .find_grouped_kernel(self.aggregate_fn.id()) - { - // SAFETY: we preserve the grouped shape and validity while replacing the - // elements with another representation of the same logical array. - let kernel_groups = unsafe { groups.with_elements_unchecked(elements.clone())? }; - if let Some(result) = - kernel.grouped_aggregate(&self.aggregate_fn, &kernel_groups, ctx)? - { - return self.push_result(result); - } + let batch = batch.clone().execute::(ctx)?.into_array(); + let gathered = batch.take(Buffer::from_iter(permutation).into_array())?; + for group_offset in 0..span { + let start = offsets[group_offset]; + let end = offsets[group_offset + 1]; + if start == end { + continue; } - if elements.is::() { - break; + let group = first + group_offset; + if self.state.is_saturated(group) { + continue; } - // Execute one step and try again - elements = elements.execute(ctx)?; + let mut accumulator = Accumulator::try_new( + self.vtable.clone(), + self.options.clone(), + self.dtype.clone(), + )?; + accumulator.accumulate(&gathered.slice(start..end)?, ctx)?; + let partial = accumulator.flush()?; + self.state.combine_scalar(group, partial)?; } - - let elements = elements.execute::(ctx)?.into_array(); - // SAFETY: we preserve the grouped shape and validity while replacing the elements with an - // executed form of the same logical array. - let grouped = unsafe { groups.with_elements_unchecked(elements)? }; - - // Otherwise, we iterate the offsets and sizes and accumulate each group one by one. - self.accumulate_grouped_fallback(&grouped, ctx) + Ok(()) } +} - fn accumulate_grouped_fallback( - &mut self, - grouped: &GroupedArray, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let mut accumulator = Accumulator::try_new( - self.vtable.clone(), - self.options.clone(), - self.dtype.clone(), - )?; - let mut states = builder_with_capacity(&self.partial_dtype, grouped.len()); - let group_ranges = grouped.group_ranges(ctx)?; - let group_validity = grouped.group_validity(ctx)?; - - for ((offset, size), valid) in group_ranges.iter().zip(group_validity.iter()) { - if valid { - let group = grouped.elements().slice(offset..offset + size)?; - accumulator.accumulate(&group, ctx)?; - states.append_scalar(&accumulator.flush()?)?; - } else { - states.append_null() - } - } - - self.push_result(states.finish()) - } +fn validate_num_groups(num_groups: usize) -> VortexResult<()> { + vortex_ensure!( + num_groups == 0 || u32::try_from(num_groups - 1).is_ok(), + "num_groups {} exceeds dense u32 group id capacity", + num_groups + ); + Ok(()) +} - fn push_result(&mut self, state: ArrayRef) -> VortexResult<()> { +fn validate_group_ids(group_ids: &[u32], num_groups: usize) -> VortexResult<()> { + validate_num_groups(num_groups)?; + for &group_id in group_ids { vortex_ensure!( - state.dtype() == &self.partial_dtype, - "State DType mismatch: expected {}, got {}", - self.partial_dtype, - state.dtype() + (group_id as usize) < num_groups, + "Group id {} out of range for {} groups", + group_id, + num_groups ); - self.partials.push(state); - Ok(()) } + Ok(()) } + fn list_view_group_ranges( groups: &ListViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult { let offsets = groups.offsets(); let sizes = groups.sizes().cast(offsets.dtype().clone())?; - let ranges = match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| { let offsets = offsets.clone().execute::>(ctx)?; let sizes = sizes.execute::>(ctx)?; @@ -411,7 +601,6 @@ fn list_view_group_ranges( }) .collect::>() }); - Ok(GroupRanges::ListView { ranges }) } @@ -421,3 +610,175 @@ fn fixed_size_list_group_ranges(groups: &FixedSizeListArray) -> GroupRanges { size: groups.list_size() as usize, } } + +/// A trait object for type-erased grouped accumulators, used for dynamic dispatch when the +/// aggregate function is not known at compile time. +pub trait DynGroupedAccumulator: 'static + Send { + /// Accumulate a values batch into dense group state. + /// + /// `group_ids` is parallel to `batch`. Each id must be a caller-assigned group ordinal in + /// `0..group_ids.num_groups()`; ids may repeat, appear out of order, or be absent from a + /// given batch. + fn accumulate( + &mut self, + batch: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()>; + + /// Fold columnar partial states into dense group state. + /// + /// `group_ids` is parallel to `partials` and follows the same dense ordinal contract as + /// [`Self::accumulate`]. + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()>; + + /// Merge one group from another grouped accumulator into this accumulator. + fn merge_group( + &mut self, + into: u32, + other: &dyn DynGroupedAccumulator, + from: u32, + ) -> VortexResult<()>; + + /// Return this accumulator's partial dtype. + fn partial_dtype(&self) -> &DType; + + /// Read one group's current partial state. + fn partial_scalar(&self, group_id: u32) -> VortexResult; + + /// Finish the accumulation and return partial aggregate results for all groups. + /// + /// Resets the accumulator state for the next round of accumulation. + fn flush_partials(&mut self, num_groups: usize) -> VortexResult; + + /// Finish the accumulation and return final aggregate results for all groups. + /// + /// Resets the accumulator state for the next round of accumulation. + fn finish(&mut self, num_groups: usize) -> VortexResult; +} + +impl DynGroupedAccumulator for GroupedAccumulator { + fn accumulate( + &mut self, + batch: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let num_groups = group_ids.num_groups(); + vortex_ensure!( + batch.dtype() == &self.dtype, + "Input DType mismatch: expected {}, got {}", + self.dtype, + batch.dtype() + ); + vortex_ensure!( + batch.len() == group_ids.len(), + "Grouped aggregate input length mismatch: {} values, {} group ids", + batch.len(), + group_ids.len() + ); + + self.ensure_groups(num_groups)?; + + if self.try_accumulate_kernel(batch, group_ids, ctx)? { + return Ok(()); + } + + let input = batch.clone(); + let mut batch = batch.clone(); + let mut tried_current = true; + for _ in 0..max_iterations() { + if batch.is::() { + break; + } + + if !tried_current && self.try_accumulate_kernel(&batch, group_ids, ctx)? { + return Ok(()); + } + + batch = batch.execute(ctx)?; + tried_current = false; + } + + if !tried_current && self.try_accumulate_kernel(&batch, group_ids, ctx)? { + return Ok(()); + } + + let group_ids = group_ids.validated_ids(ctx)?; + self.accumulate_fallback(&input, group_ids.as_ref(), ctx) + } + + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let num_groups = group_ids.num_groups(); + vortex_ensure!( + partials.dtype() == &self.partial_dtype, + "Partial DType mismatch: expected {}, got {}", + self.partial_dtype, + partials.dtype() + ); + vortex_ensure!( + partials.len() == group_ids.len(), + "Grouped aggregate partial length mismatch: {} partials, {} group ids", + partials.len(), + group_ids.len() + ); + + let group_ids = group_ids.validated_ids(ctx)?; + self.ensure_groups(num_groups)?; + self.state + .accumulate_partials(partials, group_ids.as_ref(), ctx) + } + + fn merge_group( + &mut self, + into: u32, + other: &dyn DynGroupedAccumulator, + from: u32, + ) -> VortexResult<()> { + vortex_ensure!( + other.partial_dtype() == &self.partial_dtype, + "Partial DType mismatch: expected {}, got {}", + self.partial_dtype, + other.partial_dtype() + ); + self.ensure_groups((into as usize) + 1)?; + self.state + .combine_scalar(into as usize, other.partial_scalar(from)?) + } + + fn partial_dtype(&self) -> &DType { + &self.partial_dtype + } + + fn partial_scalar(&self, group_id: u32) -> VortexResult { + self.state.partial_scalar(group_id as usize) + } + + fn flush_partials(&mut self, num_groups: usize) -> VortexResult { + self.state.flush_partials(num_groups) + } + + fn finish(&mut self, num_groups: usize) -> VortexResult { + let states = self.flush_partials(num_groups)?; + let results = self.vtable.finalize(states)?; + + vortex_ensure!( + results.dtype() == &self.return_dtype, + "Return DType mismatch: expected {}, got {}", + self.return_dtype, + results.dtype() + ); + + Ok(results) + } +} diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 39a957530bf..49e41fd0f18 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -1,202 +1,278 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; + +use num_traits::ToPrimitive; use vortex_buffer::Buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; use vortex_mask::Mask; use super::Count; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::GroupRanges; -use crate::aggregate_fn::GroupedArray; -use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::aggregate_fn::GroupIds; +use crate::aggregate_fn::GroupedState; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::kernels::GroupedAggregateKernel; +use crate::aggregate_fn::kernels::GroupedAggregateKernelAdapter; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::validity::Validity; +use crate::dtype::NativePType; +use crate::match_each_native_ptype; +use crate::scalar::Scalar; + +#[derive(Default)] +pub(crate) struct CountGroupedState { + counts: Vec, +} + +impl CountGroupedState { + fn counts_mut(&mut self) -> &mut [u64] { + &mut self.counts + } +} + +impl GroupedState for CountGroupedState { + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn len(&self) -> usize { + self.counts.len() + } + + fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { + self.counts.resize(num_groups.max(self.counts.len()), 0); + Ok(()) + } + + fn is_saturated(&self, _group_id: usize) -> bool { + false + } + + fn combine_scalar(&mut self, group_id: usize, partial: Scalar) -> VortexResult<()> { + self.counts[group_id] += partial + .as_primitive() + .typed_value::() + .vortex_expect("count partial should not be null"); + Ok(()) + } + + fn partial_scalar(&self, group_id: usize) -> VortexResult { + Ok(Scalar::primitive( + self.counts.get(group_id).copied().unwrap_or(0), + crate::dtype::Nullability::NonNullable, + )) + } + + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let partials = partials.clone().execute::>(ctx)?; + for (&partial, &group_id) in partials.iter().zip(group_ids) { + self.counts[group_id as usize] += partial; + } + Ok(()) + } + + fn flush_partials(&mut self, num_groups: usize) -> VortexResult { + vortex_ensure!( + num_groups >= self.len(), + "Cannot flush {} groups after accumulating {} groups", + num_groups, + self.len() + ); + self.ensure_groups(num_groups)?; + Ok(Buffer::from(std::mem::take(&mut self.counts)).into_array()) + } +} + +pub(crate) static COUNT_GROUPED_KERNEL: GroupedAggregateKernelAdapter = + GroupedAggregateKernelAdapter::new(CountGroupedKernel); -/// Encoding-independent grouped [`Count`] kernel. #[derive(Debug)] pub(crate) struct CountGroupedKernel; -impl DynGroupedAggregateKernel for CountGroupedKernel { - fn grouped_aggregate( +impl GroupedAggregateKernel for CountGroupedKernel { + type State = CountGroupedState; + + fn grouped_accumulate( &self, - aggregate_fn: &AggregateFnRef, - groups: &GroupedArray, + options: &NumericalAggregateOpts, + state: &mut Self::State, + batch: &ArrayRef, + group_ids: &GroupIds, ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let Some(options) = aggregate_fn.as_opt::() else { - return Ok(None); - }; - // NaN-skipping counts over floats must inspect the element values, which this - // validity-only kernel cannot do; fall back to the per-group accumulator path. - if options.skip_nans && groups.elements().dtype().is_float() { - return Ok(None); + ) -> VortexResult { + let states = state.counts_mut(); + if options.skip_nans && batch.dtype().is_float() { + let Some(primitive) = batch.as_opt::() else { + return Ok(false); + }; + let group_ids = group_ids.validated_ids(ctx)?; + accumulate_grouped_float_count( + states, + &primitive.into_owned(), + group_ids.as_ref(), + ctx, + )?; + return Ok(true); + } + + let group_ids = group_ids.validated_ids(ctx)?; + let validity = batch.validity()?.execute_mask(batch.len(), ctx)?; + if matches!(validity.indices(), AllOr::All) && has_long_group_runs(group_ids.as_ref()) { + for_each_group_run(group_ids.as_ref(), |group_id, start, end| { + states[group_id as usize] += + u64::try_from(end - start).vortex_expect("group run length must fit u64"); + }); + } else { + for_each_valid_idx(&validity, batch.len(), |idx| { + states[group_ids[idx] as usize] += 1; + }); } - try_grouped_count(groups, ctx) + Ok(true) } } -/// Count each valid group from the element validity mask. -/// -/// The [`Count`] partial dtype is non-nullable `U64`, so a null outer group cannot be represented -/// as a partial state. If any outer group is invalid, this returns `Ok(None)` and lets the caller -/// use the existing fallback behavior. -pub(super) fn try_grouped_count( - groups: &GroupedArray, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - if !groups.all_groups_valid(ctx)? { - return Ok(None); +fn has_long_group_runs(group_ids: &[u32]) -> bool { + let mut run_length = 1; + for ids in group_ids[..group_ids.len().min(256)].windows(2) { + if ids[0] == ids[1] { + run_length += 1; + if run_length >= 4 { + return true; + } + } else { + run_length = 1; + } + } + false +} + +fn for_each_group_run(group_ids: &[u32], mut f: impl FnMut(u32, usize, usize)) { + let Some(&first_group_id) = group_ids.first() else { + return; + }; + let mut group_id = first_group_id; + let mut start = 0; + for (idx, &next_group_id) in group_ids.iter().enumerate().skip(1) { + if next_group_id != group_id { + f(group_id, start, idx); + group_id = next_group_id; + start = idx; + } } - let group_ranges = groups.group_ranges(ctx)?; + f(group_id, start, group_ids.len()); +} - Ok(Some(grouped_count(groups.elements(), &group_ranges, ctx)?)) +fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { + match validity.indices() { + AllOr::All => (0..len).for_each(f), + AllOr::None => {} + AllOr::Some(indices) => indices.iter().copied().for_each(&mut f), + } } -/// Count the valid elements of each group described by `group_ranges` (element `(offset, size)` -/// pairs) into a non-nullable `U64` array, one entry per group. -fn grouped_count( - elements: &ArrayRef, - group_ranges: &GroupRanges, +fn accumulate_grouped_float_count( + states: &mut [u64], + primitive: &PrimitiveArray, + group_ids: &[u32], ctx: &mut ExecutionCtx, -) -> VortexResult { - let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?; - - let counts: Buffer = if elem_mask.all_true() { - group_ranges.iter().map(|(_, size)| size as u64).collect() - } else { - group_ranges - .iter() - .map(|(offset, size)| valid_count(&elem_mask, offset, size) as u64) - .collect() - }; +) -> VortexResult<()> { + let validity = primitive + .as_ref() + .validity()? + .execute_mask(primitive.as_ref().len(), ctx)?; - Ok(PrimitiveArray::new(counts, Validity::NonNullable).into_array()) + match_each_native_ptype!(primitive.ptype(), + unsigned: |_T| { unreachable!("float count received an unsigned primitive") }, + signed: |_T| { unreachable!("float count received a signed primitive") }, + floating: |T| { + let values = primitive.as_slice::(); + accumulate_valid_non_nan::(states, values, group_ids, &validity); + } + ); + Ok(()) } -/// Number of valid elements in the `[offset, offset + size)` range of the element mask. -fn valid_count(elem_mask: &Mask, offset: usize, size: usize) -> usize { - elem_mask.slice(offset..offset + size).true_count() +fn accumulate_valid_non_nan( + states: &mut [u64], + values: &[T], + group_ids: &[u32], + validity: &Mask, +) { + for_each_valid_idx(validity, values.len(), |idx| { + let value = ToPrimitive::to_f64(&values[idx]).vortex_expect("float to f64"); + if !value.is_nan() { + states[group_ids[idx] as usize] += 1; + } + }); } #[cfg(test)] mod tests { - #![allow(clippy::cast_possible_truncation)] - - use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; - use crate::ArrayRef; - use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::GroupIds; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::count::Count; use crate::array_session; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; + use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::Nullability::Nullable; + use crate::dtype::Nullability; use crate::dtype::PType; use crate::validity::Validity; - /// Run a grouped count through the accumulator. - fn grouped_count_actual( - groups: &ArrayRef, - elem_dtype: &DType, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut acc = GroupedAccumulator::try_new( - Count, - NumericalAggregateOpts::default(), - elem_dtype.clone(), - )?; - acc.accumulate_list(groups, ctx)?; - acc.finish() - } - - /// Reference valid-counts (non-nullable `U64`), one per group. - fn grouped_count_reference( - elements: &ArrayRef, - ranges: &[(usize, usize)], - ) -> VortexResult { + fn run_grouped_count( + values: &crate::ArrayRef, + ids: impl IntoIterator, + num_groups: usize, + options: NumericalAggregateOpts, + ) -> VortexResult { + let mut acc = GroupedAccumulator::try_new(Count, options, values.dtype().clone())?; + let group_ids = GroupIds::from_iter(ids, num_groups)?; let mut ctx = array_session().create_execution_ctx(); - let counts: Buffer = ranges - .iter() - .map(|&(offset, size)| { - Ok(elements - .slice(offset..offset + size)? - .valid_count(&mut ctx)? as u64) - }) - .collect::>()?; - Ok(PrimitiveArray::new(counts, Validity::NonNullable).into_array()) - } - - fn listview(elements: ArrayRef, ranges: &[(usize, usize)]) -> VortexResult { - let offsets = PrimitiveArray::from_iter(ranges.iter().map(|&(o, _)| o as i32)); - let sizes = PrimitiveArray::from_iter(ranges.iter().map(|&(_, s)| s as i32)); - Ok(ListViewArray::try_new( - elements, - offsets.into_array(), - sizes.into_array(), - Validity::NonNullable, - )? - .into_array()) + acc.accumulate(values, &group_ids, &mut ctx)?; + acc.finish(num_groups) } #[test] - fn listview_counts_all_valid() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let elements = - PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); - let elem_dtype = DType::Primitive(PType::I32, NonNullable); - let ranges = [(0, 2), (2, 1), (3, 3), (6, 0)]; - - let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; - let expected = grouped_count_reference(&elements, &ranges)?; - - let direct = - PrimitiveArray::new(buffer![2u64, 1, 3, 0], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut ctx); - assert_arrays_eq!(&actual, &expected, &mut ctx); - Ok(()) - } - - #[test] - fn listview_counts_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let elements = - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) + fn dense_ids_repeat_reorder_and_omit_groups() -> VortexResult<()> { + let values = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None, Some(6)]) .into_array(); - let elem_dtype = DType::Primitive(PType::I32, Nullable); - let ranges = [(0, 3), (3, 2), (5, 1)]; - - let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; - let expected = grouped_count_reference(&elements, &ranges)?; - - // Group 0: {1, null, 3} -> 2. Group 1: {null, null} -> 0. Group 2: {9} -> 1. - let direct = PrimitiveArray::new(buffer![2u64, 0, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut ctx); + let actual = run_grouped_count( + &values, + [2, 0, 2, 0, 2, 0], + 4, + NumericalAggregateOpts::default(), + )?; + let expected = PrimitiveArray::from_iter([2u64, 0, 2, 0]).into_array(); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] - fn listview_counts_varbinview_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let elements = VarBinViewArray::from_iter_nullable_str([ + fn varbinview_counts_nulls() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_str([ Some("a"), None, Some("bbb"), @@ -204,55 +280,96 @@ mod tests { Some("cc"), ]) .into_array(); - let elem_dtype = elements.dtype().clone(); - let ranges = [(0, 2), (2, 2), (4, 1)]; - - let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; - let expected = grouped_count_reference(&elements, &ranges)?; - - let direct = PrimitiveArray::new(buffer![1u64, 1, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut ctx); + let actual = run_grouped_count( + &values, + [0, 0, 1, 1, 2], + 3, + NumericalAggregateOpts::default(), + )?; + let expected = PrimitiveArray::from_iter([1u64, 1, 1]).into_array(); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] - fn fixed_size_counts_float_nans() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let elements = - PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(2.0)]) + fn float_nan_options_match_scalar_count() -> VortexResult<()> { + let values = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(3.0)]) .into_array(); - let elem_dtype = DType::Primitive(PType::F64, Nullable); - let groups = - FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array(); - - // NaNs are excluded by default and counted otherwise. - let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; - let expected = PrimitiveArray::new(buffer![1u64, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &expected, &mut ctx); + let skipped = + run_grouped_count(&values, [0, 0, 1, 1], 2, NumericalAggregateOpts::default())?; + let included = run_grouped_count( + &values, + [0, 0, 1, 1], + 2, + NumericalAggregateOpts::include_nans(), + )?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + &skipped, + &PrimitiveArray::from_iter([1u64, 1]).into_array(), + &mut ctx + ); + assert_arrays_eq!( + &included, + &PrimitiveArray::from_iter([2u64, 1]).into_array(), + &mut ctx + ); + Ok(()) + } - let mut acc = - GroupedAccumulator::try_new(Count, NumericalAggregateOpts::include_nans(), elem_dtype)?; - acc.accumulate_list(&groups, &mut ctx)?; - let actual = acc.finish()?; - let expected = PrimitiveArray::new(buffer![2u64, 1], Validity::NonNullable).into_array(); + #[test] + fn encoded_constant_group_ids() -> VortexResult<()> { + let values = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(); + let group_ids = GroupIds::new(ConstantArray::new(1u32, values.len()).into_array(), 3)?; + let mut ctx = array_session().create_execution_ctx(); + let mut acc = GroupedAccumulator::try_new( + Count, + NumericalAggregateOpts::default(), + values.dtype().clone(), + )?; + acc.accumulate(&values, &group_ids, &mut ctx)?; + let actual = acc.finish(3)?; + let expected = PrimitiveArray::from_iter([0u64, 3, 0]).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] - fn fixed_size_counts_with_nulls() -> VortexResult<()> { + fn rejects_out_of_range_group_id() -> VortexResult<()> { + assert!(GroupIds::from_iter([0u32, 2], 2).is_err()); + + let values = PrimitiveArray::new(buffer![1i32, 2], Validity::NonNullable).into_array(); + let group_ids = GroupIds::new( + PrimitiveArray::new(buffer![0u32, 2], Validity::NonNullable).into_array(), + 2, + )?; let mut ctx = array_session().create_execution_ctx(); - let elements = - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(); - let elem_dtype = DType::Primitive(PType::I32, Nullable); - let groups = - FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array(); + let mut acc = GroupedAccumulator::try_new( + Count, + NumericalAggregateOpts::default(), + values.dtype().clone(), + )?; + assert!(acc.accumulate(&values, &group_ids, &mut ctx).is_err()); + Ok(()) + } - let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; - let direct = PrimitiveArray::new(buffer![1u64, 2], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut ctx); + #[test] + fn accumulates_partials_and_merges_groups() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let partials = PrimitiveArray::from_iter([2u64, 3, 5]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let mut left = + GroupedAccumulator::try_new(Count, NumericalAggregateOpts::default(), dtype.clone())?; + left.accumulate_partials(&partials, &GroupIds::from_iter([0u32, 1, 1], 2)?, &mut ctx)?; + let mut right = + GroupedAccumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?; + right.merge_group(0, &left, 1)?; + let actual = right.finish(1)?; + let expected = PrimitiveArray::from_iter([8u64]).into_array(); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 8f7d68027bc..e8e1148fbd2 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod grouped; -pub(crate) use grouped::CountGroupedKernel; +pub(crate) use grouped::COUNT_GROUPED_KERNEL; +use grouped::CountGroupedState; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -10,10 +11,13 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::GroupedState; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::nan_count::nan_count; +use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -69,6 +73,15 @@ impl AggregateFnVTable for Count { }) } + fn grouped_state( + &self, + _options: &Self::Options, + _input_dtype: &DType, + _partial_dtype: &DType, + ) -> VortexResult> { + Ok(Box::new(CountGroupedState::default())) + } + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { let val = other .as_primitive() @@ -82,6 +95,16 @@ impl AggregateFnVTable for Count { Ok(Scalar::primitive(partial.count, Nullability::NonNullable)) } + fn partials_to_array( + &self, + partials: &[Self::Partial], + _partial_dtype: &DType, + ) -> VortexResult> { + Ok(Some( + PrimitiveArray::from_iter(partials.iter().map(|partial| partial.count)).into_array(), + )) + } + fn reset(&self, partial: &mut Self::Partial) { partial.count = 0; } diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 16a1b5431fe..ba59883b5f5 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -233,12 +233,12 @@ mod tests { use super::*; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::GroupIds; use crate::aggregate_fn::GroupedAccumulator; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::DecimalArray; - use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::dtype::DecimalDType; use crate::validity::Validity; @@ -466,11 +466,14 @@ mod tests { #[test] fn mean_grouped_finalize() -> VortexResult<()> { let cases = mean_nan_null(); - let elements = PrimitiveArray::from_option_iter( - cases.iter().flat_map(|(group, _)| group.iter().copied()), + let values = PrimitiveArray::from_option_iter( + (0..3).flat_map(|row| cases.iter().map(move |(group, _)| group[row])), ) .into_array(); - let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, cases.len())?; + let groups = (0..cases.len()) + .map(u32::try_from) + .collect::, _>>()?; + let group_ids = GroupIds::from_iter(std::iter::repeat_n(groups, 3).flatten(), cases.len())?; let mut acc = GroupedAccumulator::try_new( Mean::combined(), @@ -481,8 +484,8 @@ mod tests { DType::Primitive(PType::F64, Nullability::Nullable), )?; let mut ctx = array_session().create_execution_ctx(); - acc.accumulate_list(&groups.into_array(), &mut ctx)?; - let result = acc.finish()?; + acc.accumulate(&values, &group_ids, &mut ctx)?; + let result = acc.finish(cases.len())?; for (case, (_, expected)) in cases.into_iter().enumerate() { let actual = result.execute_scalar(case, &mut ctx)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 64b0ce81995..5ecf326a162 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -1,439 +1,729 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBuffer; -use vortex_buffer::BitBufferMut; +use num_traits::AsPrimitive; +use num_traits::CheckedAdd; +use num_traits::ToPrimitive; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_mask::AllOr; use vortex_mask::Mask; use super::Sum; +use super::SumAggregateOpts; +use super::checked_add_i64; +use super::checked_add_u64; +use super::grouped_state::SumGroupedState; +use super::grouped_state::SumGroupedValues; +use super::grouped_state::add_decimal; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; -use super::sum_partial_fields; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::GroupRanges; -use crate::aggregate_fn::GroupedArray; -use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::aggregate_fn::GroupIds; +use crate::aggregate_fn::kernels::GroupedAggregateKernel; +use crate::aggregate_fn::kernels::GroupedAggregateKernelAdapter; +use crate::arrays::Bool; use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::arrays::StructArray; +use crate::arrays::bool::BoolArrayExt; +use crate::dtype::NativeDecimalType; use crate::dtype::NativePType; -use crate::dtype::Nullability; +use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; -use crate::validity::Validity; -/// Encoding-specific grouped [`Sum`] kernel for primitive element arrays. +const MIN_GROUP_RUN_LENGTH: usize = 4; + +pub(crate) static SUM_GROUPED_KERNEL: GroupedAggregateKernelAdapter = + GroupedAggregateKernelAdapter::new(SumGroupedKernel); + #[derive(Debug)] -pub(crate) struct PrimitiveGroupedSumEncodingKernel; +pub(crate) struct SumGroupedKernel; + +impl GroupedAggregateKernel for SumGroupedKernel { + type State = SumGroupedState; -impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { - fn grouped_aggregate( + fn grouped_accumulate( &self, - aggregate_fn: &AggregateFnRef, - groups: &GroupedArray, + options: &SumAggregateOpts, + state: &mut Self::State, + batch: &ArrayRef, + group_ids: &GroupIds, ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let Some(options) = aggregate_fn.as_opt::() else { - return Ok(None); - }; - try_grouped_sum(groups, ctx, options.skip_nans) + ) -> VortexResult { + if let Some(primitive) = batch.as_opt::() { + let group_ids = group_ids.validated_ids(ctx)?; + accumulate_grouped_primitive( + state, + &primitive.into_owned(), + group_ids.as_ref(), + options.skip_nans, + ctx, + )?; + return Ok(true); + } + + if let Some(bools) = batch.as_opt::() { + let group_ids = group_ids.validated_ids(ctx)?; + accumulate_grouped_bool(state, &bools.into_owned(), group_ids.as_ref(), ctx)?; + return Ok(true); + } + + if let Some(decimals) = batch.as_opt::() { + let group_ids = group_ids.validated_ids(ctx)?; + accumulate_grouped_decimal(state, &decimals.into_owned(), group_ids.as_ref(), ctx)?; + return Ok(true); + } + + Ok(false) } } -/// Grouped [`Sum`] implementation for canonical primitive elements. -/// -/// Reuses the scalar primitive-sum reductions ([`sum_unsigned_all`]/[`sum_signed_all`]/ -/// [`sum_float_all`]) so the per-group semantics match scalar `sum` exactly (overflow saturates to -/// a null sum, NaNs are skipped). The element validity mask is materialized once and sliced per -/// group, rather than the per-group accumulator setup of the generic fallback path. -pub(super) fn try_grouped_sum( - groups: &GroupedArray, - ctx: &mut ExecutionCtx, +fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { + match validity.indices() { + AllOr::All => (0..len).for_each(f), + AllOr::None => {} + AllOr::Some(indices) => indices.iter().copied().for_each(&mut f), + } +} + +fn for_each_group_run(group_ids: &[u32], mut f: impl FnMut(u32, usize, usize)) { + let Some((&first, rest)) = group_ids.split_first() else { + return; + }; + let mut group_id = first; + let mut start = 0usize; + for (idx, &next_group_id) in rest.iter().enumerate() { + let idx = idx + 1; + if next_group_id != group_id { + f(group_id, start, idx); + group_id = next_group_id; + start = idx; + } + } + f(group_id, start, group_ids.len()); +} + +fn has_long_group_runs(group_ids: &[u32]) -> bool { + let mut run_length = 1; + for ids in group_ids[..group_ids.len().min(256)].windows(2) { + if ids[0] == ids[1] { + run_length += 1; + if run_length >= MIN_GROUP_RUN_LENGTH { + return true; + } + } else { + run_length = 1; + } + } + false +} + +fn accumulate_grouped_unsigned( + values: &mut [u64], + overflowed: &mut [u8], + empty: &mut [u8], + group_id: u32, + value: u64, +) { + let group = group_id as usize; + empty[group] = 0; + if checked_add_u64(&mut values[group], value) { + overflowed[group] = 1; + } +} + +fn accumulate_grouped_unsigned_run( + sums: &mut [u64], + overflowed: &mut [u8], + empty: &mut [u8], + group_id: u32, + values: &[T], +) where + T: NativePType + AsPrimitive, +{ + let group = group_id as usize; + empty[group] = 0; + if sum_unsigned_all(&mut sums[group], values) { + overflowed[group] = 1; + } +} + +fn accumulate_grouped_unsigned_all( + sums: &mut [u64], + overflowed: &mut [u8], + empty: &mut [u8], + values: &[T], + group_ids: &[u32], +) where + T: NativePType + AsPrimitive, +{ + if !has_long_group_runs(group_ids) { + for (&value, &group_id) in values.iter().zip(group_ids) { + accumulate_grouped_unsigned(sums, overflowed, empty, group_id, value.as_()); + } + return; + } + + for_each_group_run(group_ids, |group_id, start, end| { + empty[group_id as usize] = 0; + if end - start >= MIN_GROUP_RUN_LENGTH { + accumulate_grouped_unsigned_run(sums, overflowed, empty, group_id, &values[start..end]); + } else { + for &value in &values[start..end] { + accumulate_grouped_unsigned(sums, overflowed, empty, group_id, value.as_()); + } + } + }); +} + +fn accumulate_grouped_signed( + values: &mut [i64], + overflowed: &mut [u8], + empty: &mut [u8], + group_id: u32, + value: i64, +) { + let group = group_id as usize; + empty[group] = 0; + if checked_add_i64(&mut values[group], value) { + overflowed[group] = 1; + } +} + +fn accumulate_grouped_signed_run( + sums: &mut [i64], + overflowed: &mut [u8], + empty: &mut [u8], + group_id: u32, + values: &[T], +) where + T: NativePType + AsPrimitive, +{ + let group = group_id as usize; + empty[group] = 0; + if sum_signed_all(&mut sums[group], values) { + overflowed[group] = 1; + } +} + +fn accumulate_grouped_signed_all( + sums: &mut [i64], + overflowed: &mut [u8], + empty: &mut [u8], + values: &[T], + group_ids: &[u32], +) where + T: NativePType + AsPrimitive, +{ + if !has_long_group_runs(group_ids) { + for (&value, &group_id) in values.iter().zip(group_ids) { + accumulate_grouped_signed(sums, overflowed, empty, group_id, value.as_()); + } + return; + } + + for_each_group_run(group_ids, |group_id, start, end| { + if end - start >= MIN_GROUP_RUN_LENGTH { + accumulate_grouped_signed_run(sums, overflowed, empty, group_id, &values[start..end]); + } else { + for &value in &values[start..end] { + accumulate_grouped_signed(sums, overflowed, empty, group_id, value.as_()); + } + } + }); +} + +fn accumulate_grouped_float( + sums: &mut [f64], + empty: &mut [u8], + group_id: u32, + value: f64, skip_nans: bool, -) -> VortexResult> { - if !groups.elements().is::() { - return Ok(None); +) { + empty[group_id as usize] = 0; + if !skip_nans || !value.is_nan() { + sums[group_id as usize] += value; } - let elements = groups.elements().clone().downcast::(); - let group_ranges = groups.group_ranges(ctx)?; - let group_validity = groups.group_validity(ctx)?; - - Ok(Some(grouped_sum( - &elements, - &group_ranges, - &group_validity, - ctx, - skip_nans, - )?)) } -/// Sum each group described by `group_ranges` (element `(offset, size)` pairs), one sum per group. -fn grouped_sum( - elements: &PrimitiveArray, - group_ranges: &GroupRanges, - group_validity: &Mask, - ctx: &mut ExecutionCtx, +fn accumulate_grouped_float_all( + sums: &mut [f64], + empty: &mut [u8], + values: &[T], + group_ids: &[u32], + skip_nans: bool, +) { + if !has_long_group_runs(group_ids) { + for (value, &group_id) in values.iter().zip(group_ids) { + let value = ToPrimitive::to_f64(value).vortex_expect("float to f64"); + accumulate_grouped_float(sums, empty, group_id, value, skip_nans); + } + return; + } + + for_each_group_run(group_ids, |group_id, start, end| { + empty[group_id as usize] = 0; + if end - start >= MIN_GROUP_RUN_LENGTH { + sum_float_all(&mut sums[group_id as usize], &values[start..end], skip_nans); + } else { + for value in &values[start..end] { + let value = ToPrimitive::to_f64(value).vortex_expect("float to f64"); + accumulate_grouped_float(sums, empty, group_id, value, skip_nans); + } + } + }); +} + +fn accumulate_grouped_primitive( + state: &mut SumGroupedState, + primitive: &PrimitiveArray, + group_ids: &[u32], skip_nans: bool, -) -> VortexResult { - let elem_mask = elements + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let validity = primitive .as_ref() .validity()? - .execute_mask(elements.as_ref().len(), ctx)?; - let all_valid = elem_mask.all_true(); + .execute_mask(primitive.as_ref().len(), ctx)?; + let all_valid = matches!(validity.slices(), AllOr::All); + let (state, overflowed, empty) = state.parts_mut(); - let (sums, is_overflow, is_empty) = match_each_native_ptype!(elements.ptype(), + match_each_native_ptype!(primitive.ptype(), unsigned: |T| { - let values = elements.as_slice::(); - collect_sums::( - values, group_ranges, group_validity, &elem_mask, all_valid, sum_unsigned_all) + let SumGroupedValues::Unsigned(sums) = state else { + vortex_panic!("unsigned input with non-unsigned grouped sum state") + }; + let values = primitive.as_slice::(); + if all_valid { + accumulate_grouped_unsigned_all(sums, overflowed, empty, values, group_ids); + } else { + for_each_valid_idx(&validity, values.len(), |idx| { + accumulate_grouped_unsigned( + sums, + overflowed, + empty, + group_ids[idx], + values[idx].as_(), + ); + }); + } }, signed: |T| { - let values = elements.as_slice::(); - collect_sums::( - values, group_ranges, group_validity, &elem_mask, all_valid, sum_signed_all) + let SumGroupedValues::Signed(sums) = state else { + vortex_panic!("signed input with non-signed grouped sum state") + }; + let values = primitive.as_slice::(); + if all_valid { + accumulate_grouped_signed_all(sums, overflowed, empty, values, group_ids); + } else { + for_each_valid_idx(&validity, values.len(), |idx| { + accumulate_grouped_signed( + sums, + overflowed, + empty, + group_ids[idx], + values[idx].as_(), + ); + }); + } }, floating: |T| { - let values = elements.as_slice::(); - collect_sums::( - values, group_ranges, group_validity, &elem_mask, all_valid, - |acc, slice| { sum_float_all(acc, slice, skip_nans); false }) + let SumGroupedValues::Float(sums) = state else { + vortex_panic!("float input with non-float grouped sum state") + }; + let values = primitive.as_slice::(); + if all_valid { + accumulate_grouped_float_all(sums, empty, values, group_ids, skip_nans); + } else { + for_each_valid_idx(&validity, values.len(), |idx| { + let value = ToPrimitive::to_f64(&values[idx]).vortex_expect("float to f64"); + accumulate_grouped_float(sums, empty, group_ids[idx], value, skip_nans); + }); + } } ); + Ok(()) +} - let partial_fields = sum_partial_fields(sums.dtype().clone()); - - // SAFETY: all three children have one value per group and match `partial_fields`; the struct - // validity is derived from the same group count. - Ok(unsafe { - StructArray::new_unchecked( - vec![ - sums.into_array(), - BoolArray::new(is_overflow, Validity::NonNullable).into_array(), - BoolArray::new(is_empty, Validity::NonNullable).into_array(), - ], - partial_fields, - group_validity.len(), - Validity::from_mask(group_validity.clone(), Nullability::Nullable), - ) - } - .into_array()) +fn accumulate_grouped_bool( + state: &mut SumGroupedState, + bools: &BoolArray, + group_ids: &[u32], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let validity = bools + .as_ref() + .validity()? + .execute_mask(bools.as_ref().len(), ctx)?; + let values = bools.to_bit_buffer(); + let valid_true = match validity.bit_buffer() { + AllOr::All => values, + AllOr::None => return Ok(()), + AllOr::Some(validity) => &values & validity, + }; + let (state, overflowed, empty) = state.parts_mut(); + let SumGroupedValues::Unsigned(sums) = state else { + vortex_panic!("boolean input with non-unsigned grouped sum state") + }; + for_each_valid_idx(&validity, bools.as_ref().len(), |idx| { + empty[group_ids[idx] as usize] = 0; + }); + valid_true.for_each_set_index(|idx| { + accumulate_grouped_unsigned(sums, overflowed, empty, group_ids[idx], 1); + }); + Ok(()) } -/// Reduce each group's element slice into a non-null sum, overflow bitmap, and empty bitmap. -fn collect_sums( - values: &[T], - group_ranges: &GroupRanges, - group_validity: &Mask, - elem_mask: &Mask, - all_valid: bool, - sum_run: impl Fn(&mut A, &[T]) -> bool, -) -> (PrimitiveArray, BitBuffer, BitBuffer) { - let group_count = group_ranges.len(); - let mut is_overflow = BitBufferMut::new_unset(group_count); - let mut is_empty = BitBufferMut::new_unset(group_count); - let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| { - if !group_validity.value(i) { - return A::default(); - } - let mut acc = A::default(); - let (overflow, any_valid) = if all_valid { - (sum_run(&mut acc, &values[offset..offset + size]), size > 0) - } else { - sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run) - }; - if overflow { - // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per - // group. - unsafe { is_overflow.set_unchecked(i) }; - } - if !any_valid { - // SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per - // group. - unsafe { is_empty.set_unchecked(i) }; +fn accumulate_grouped_decimal( + state: &mut SumGroupedState, + decimals: &DecimalArray, + group_ids: &[u32], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let validity = decimals + .as_ref() + .validity()? + .execute_mask(decimals.as_ref().len(), ctx)?; + let output_dtype = state + .decimal_dtype() + .vortex_expect("decimal sum state dtype"); + let (state, overflowed, empty) = state.parts_mut(); + match_each_decimal_value_type!(decimals.values_type(), |T| { + let values = decimals.buffer::(); + match state { + SumGroupedValues::Decimal8(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + empty, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal16(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + empty, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal32(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + empty, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal64(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + empty, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal128(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + empty, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal256(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + empty, + values, + group_ids, + &validity, + output_dtype, + ), + _ => vortex_panic!("decimal input with non-decimal grouped sum state"), } - acc }); - let sums = PrimitiveArray::from_iter(sums); - (sums, is_overflow.freeze(), is_empty.freeze()) + Ok(()) } -/// Sum valid runs in one group, returning `(overflow, any_valid)`. -fn sum_masked_group( - acc: &mut A, - values: &[T], - offset: usize, - size: usize, - elem_mask: &Mask, - sum_run: &impl Fn(&mut A, &[T]) -> bool, -) -> (bool, bool) { - match elem_mask { - Mask::AllTrue(_) => (sum_run(acc, &values[offset..offset + size]), size > 0), - Mask::AllFalse(_) => (false, false), - Mask::Values(mask_values) => { - let validity = mask_values - .bit_buffer() - .as_view() - .slice(offset..offset + size); - let mut any_valid = false; - for (start, end) in validity.set_slices() { - any_valid = true; - if sum_run(acc, &values[offset + start..offset + end]) { - return (true, true); - } - } - (false, any_valid) - } - } +fn accumulate_grouped_decimal_values( + sums: &mut [I], + overflowed: &mut [u8], + empty: &mut [u8], + values: Buffer, + group_ids: &[u32], + validity: &Mask, + dtype: crate::dtype::DecimalDType, +) where + T: NativeDecimalType + AsPrimitive, + I: NativeDecimalType + CheckedAdd, +{ + for_each_valid_idx(validity, values.len(), |idx| { + empty[group_ids[idx] as usize] = 0; + add_decimal( + sums, + overflowed, + group_ids[idx] as usize, + values[idx].as_(), + dtype, + ); + }); } #[cfg(test)] mod tests { - #![allow(clippy::cast_possible_truncation)] - use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; - use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; + use crate::aggregate_fn::GroupIds; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumAggregateOpts; - use crate::aggregate_fn::fns::sum::sum; use crate::array_session; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; + use crate::arrays::BoolArray; + use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; use crate::assert_arrays_eq; - use crate::builders::builder_with_capacity; use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::Nullability::Nullable; + use crate::dtype::DecimalDType; + use crate::dtype::FieldName; + use crate::dtype::FieldNames; + use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::i256; + use crate::scalar::DecimalValue; use crate::validity::Validity; - /// Run a grouped sum through the accumulator. - fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = - GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; - acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; - acc.finish() + fn sum_partials( + sums: crate::ArrayRef, + overflowed: impl IntoIterator, + empty: impl IntoIterator, + ) -> VortexResult { + let len = sums.len(); + Ok(StructArray::try_new( + FieldNames::from_iter([ + FieldName::from("sum"), + FieldName::from("is_overflow"), + FieldName::from("is_empty"), + ]), + vec![ + sums, + BoolArray::from_iter(overflowed).into_array(), + BoolArray::from_iter(empty).into_array(), + ], + len, + Validity::AllValid, + )? + .into_array()) } - /// Reference sums computed exactly like the generic slow path: per-group scalar [`sum`] for - /// valid groups, a null sum for invalid groups. - fn grouped_sum_reference( - elements: &ArrayRef, - ranges: &[(usize, usize)], - group_valid: &[bool], - elem_dtype: &DType, - ) -> VortexResult { - use crate::aggregate_fn::AggregateFnVTable; - + fn run_grouped_sum( + values: &crate::ArrayRef, + ids: impl IntoIterator, + num_groups: usize, + options: SumAggregateOpts, + ) -> VortexResult { + let mut acc = GroupedAccumulator::try_new(Sum, options, values.dtype().clone())?; + let group_ids = GroupIds::from_iter(ids, num_groups)?; let mut ctx = array_session().create_execution_ctx(); - let sum_dtype = Sum - .return_dtype(&SumAggregateOpts::default(), elem_dtype) - .expect("sum return dtype"); - let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); - for (i, &(offset, size)) in ranges.iter().enumerate() { - if group_valid[i] { - let slice = elements.slice(offset..offset + size)?; - builder.append_scalar(&sum(&slice, &mut ctx)?)?; - } else { - builder.append_null(); - } - } - Ok(builder.finish()) - } - - fn offsets_sizes(ranges: &[(usize, usize)]) -> (ArrayRef, ArrayRef) { - let offsets = PrimitiveArray::from_iter(ranges.iter().map(|&(o, _)| o as i32)); - let sizes = PrimitiveArray::from_iter(ranges.iter().map(|&(_, s)| s as i32)); - (offsets.into_array(), sizes.into_array()) - } - - fn listview( - elements: ArrayRef, - ranges: &[(usize, usize)], - group_valid: &[bool], - ) -> VortexResult { - let (offsets, sizes) = offsets_sizes(ranges); - let validity = if group_valid.iter().all(|&v| v) { - Validity::NonNullable - } else { - Validity::from_iter(group_valid.iter().copied()) - }; - Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) + acc.accumulate(values, &group_ids, &mut ctx)?; + acc.finish(num_groups) } #[test] - fn listview_matches_reference_unsigned() -> VortexResult<()> { + fn dense_ids_repeat_reorder_and_omit_groups() -> VortexResult<()> { + let values = PrimitiveArray::from_option_iter([ + Some(1i32), + None, + Some(3), + Some(4), + Some(5), + Some(6), + ]) + .into_array(); + let actual = run_grouped_sum(&values, [2, 0, 2, 0, 2, 0], 4, SumAggregateOpts::default())?; + let expected = + PrimitiveArray::from_option_iter([Some(10i64), None, Some(9), None]).into_array(); let mut ctx = array_session().create_execution_ctx(); - let elements = - PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); - let elem_dtype = DType::Primitive(PType::U32, NonNullable); - let ranges = [(0, 2), (2, 1), (3, 3)]; - let valid = [true, true, true]; - - let groups = listview(elements.clone(), &ranges, &valid)?; - let actual = grouped_sum_actual(&groups, &elem_dtype)?; - let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - - // Unsigned input sums to U64. - let direct = PrimitiveArray::from_option_iter([Some(3u64), Some(3u64), Some(15u64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] - fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + fn bool_and_overflow_are_group_local() -> VortexResult<()> { + let bools: BoolArray = [true, false, true, true].into_iter().collect(); + let actual = run_grouped_sum( + &bools.into_array(), + [1, 0, 1, 0], + 2, + SumAggregateOpts::default(), + )?; let mut ctx = array_session().create_execution_ctx(); - // Offsets are not in group order and a group is null: the group validity must be indexed by - // group index, not by element offset. - let elements = - PrimitiveArray::new(buffer![10i32, 20, 30, 40, 50, 60], Validity::NonNullable) - .into_array(); - let elem_dtype = DType::Primitive(PType::I32, NonNullable); - let ranges = [(4, 2), (0, 2), (2, 2)]; - let valid = [true, false, true]; - - let groups = listview(elements.clone(), &ranges, &valid)?; - let actual = grouped_sum_actual(&groups, &elem_dtype)?; - let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; + assert_arrays_eq!( + &actual, + &PrimitiveArray::from_option_iter([Some(1u64), Some(2)]).into_array(), + &mut ctx + ); - let direct = PrimitiveArray::from_option_iter([Some(110i64), None, Some(70i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); - assert_arrays_eq!(&actual, &expected, &mut ctx); + let values = + PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); + let actual = run_grouped_sum(&values, [0, 0, 1, 1], 2, SumAggregateOpts::default())?; + assert_arrays_eq!( + &actual, + &PrimitiveArray::from_option_iter([None, Some(5i64)]).into_array(), + &mut ctx + ); Ok(()) } #[test] - fn listview_interior_and_full_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. - let elements = - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) + fn float_nan_options_match_scalar_sum() -> VortexResult<()> { + let values = + PrimitiveArray::new(buffer![1.0f64, f64::NAN, 2.0, 4.0], Validity::NonNullable) .into_array(); - let elem_dtype = DType::Primitive(PType::I32, Nullable); - let ranges = [(0, 3), (3, 2), (5, 0), (5, 1)]; - let valid = [true, true, true, true]; - - let groups = listview(elements.clone(), &ranges, &valid)?; - let actual = grouped_sum_actual(&groups, &elem_dtype)?; - let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - - let direct = PrimitiveArray::from_option_iter([Some(4i64), None, None, Some(9i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); - assert_arrays_eq!(&actual, &expected, &mut ctx); + let skipped = run_grouped_sum(&values, [0, 0, 1, 1], 2, SumAggregateOpts::default())?; + let included = run_grouped_sum(&values, [0, 0, 1, 1], 2, SumAggregateOpts::include_nans())?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + &skipped, + &PrimitiveArray::from_option_iter([Some(1.0f64), Some(6.0)]).into_array(), + &mut ctx + ); + let group_zero = included.execute_scalar(0, &mut ctx)?; + assert!( + group_zero + .as_primitive() + .typed_value::() + .vortex_expect("grouped float sum should be non-null") + .is_nan() + ); Ok(()) } #[test] - fn listview_overflow_group_is_null() -> VortexResult<()> { + fn exact_decimal_sum_with_reordered_ids_and_nulls() -> VortexResult<()> { + let input_dtype = DecimalDType::new(10, 2); + let values = DecimalArray::new( + buffer![100i64, 200, -50, 300, 400], + input_dtype, + Validity::from_iter([true, true, true, false, true]), + ) + .into_array(); + let actual = run_grouped_sum(&values, [2, 0, 2, 0, 2], 4, SumAggregateOpts::default())?; + let output_dtype = DecimalDType::new(20, 2); + let expected = DecimalArray::new( + buffer![200i64, 0, 450, 0], + output_dtype, + Validity::from_iter([true, false, true, false]), + ) + .into_array(); let mut ctx = array_session().create_execution_ctx(); - let elements = - PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); - let elem_dtype = DType::Primitive(PType::I64, NonNullable); - let ranges = [(0, 2), (2, 2)]; - let valid = [true, true]; - - let groups = listview(elements.clone(), &ranges, &valid)?; - let actual = grouped_sum_actual(&groups, &elem_dtype)?; - let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - - // First group overflows -> null sum; second group sums normally. - let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] - fn listview_float_nan_and_inf() -> VortexResult<()> { - let elements = PrimitiveArray::new( - buffer![1.0f64, f64::NAN, 2.0, f64::INFINITY, f64::NEG_INFINITY, 4.0], + fn exact_decimal_overflow_is_group_local() -> VortexResult<()> { + let one = i256::from_i128(1); + let large = i256::from_i128(10) + .checked_pow(76) + .vortex_expect("10^76 must fit in i256") + - one; + let dtype = DecimalDType::new(76, 0); + let values = DecimalArray::new( + buffer![large, i256::from_i128(7), large], + dtype, Validity::NonNullable, ) .into_array(); - let elem_dtype = DType::Primitive(PType::F64, NonNullable); - let ranges = [(0, 3), (3, 3)]; - let valid = [true, true]; - - let groups = listview(elements.clone(), &ranges, &valid)?; - let actual = grouped_sum_actual(&groups, &elem_dtype)?; - - // Group 0: NaN skipped -> 3.0. Group 1: INF + -INF = NaN. (Avoid array equality here since - // NaN != NaN; compare element scalars against the reference path instead.) + let actual = run_grouped_sum(&values, [0, 1, 0], 2, SumAggregateOpts::default())?; + let expected = DecimalArray::new( + buffer![i256::ZERO, i256::from_i128(7)], + dtype, + Validity::from_iter([false, true]), + ) + .into_array(); let mut ctx = array_session().create_execution_ctx(); - let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; - let g0 = actual.execute_scalar(0, &mut ctx)?; - assert_eq!(g0.as_primitive().typed_value::(), Some(3.0)); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + let group_one = actual.execute_scalar(1, &mut ctx)?; assert_eq!( - g0.as_primitive().typed_value::(), - expected - .execute_scalar(0, &mut ctx)? - .as_primitive() - .typed_value::() - ); - let g1 = actual.execute_scalar(1, &mut ctx)?; - assert!(g1.as_primitive().typed_value::().unwrap().is_nan()); - assert!( - expected - .execute_scalar(1, &mut ctx)? - .as_primitive() - .typed_value::() - .unwrap() - .is_nan() + group_one.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(7))) ); Ok(()) } #[test] - fn listview_float_nan_not_skipping() -> VortexResult<()> { - let elements = PrimitiveArray::new( - buffer![1.0f64, f64::NAN, 2.0, 3.0, 4.0], - Validity::NonNullable, - ) - .into_array(); - let elem_dtype = DType::Primitive(PType::F64, NonNullable); - let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; - - let mut acc = - GroupedAccumulator::try_new(Sum, SumAggregateOpts::include_nans(), elem_dtype)?; - acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; - let actual = acc.finish()?; - + fn accumulates_typed_primitive_partials() -> VortexResult<()> { + let input_dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let partials = sum_partials( + PrimitiveArray::new(buffer![2i64, 3, 5, 0], Validity::NonNullable).into_array(), + [false, false, false, true], + [false; 4], + )?; let mut ctx = array_session().create_execution_ctx(); - // Group 0 contains a NaN -> NaN sum; group 1 sums normally. - let g0 = actual.execute_scalar(0, &mut ctx)?; - assert!(g0.as_primitive().typed_value::().unwrap().is_nan()); - let g1 = actual.execute_scalar(1, &mut ctx)?; - assert_eq!(g1.as_primitive().typed_value::(), Some(7.0)); + let mut acc = GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), input_dtype)?; + acc.accumulate_partials( + &partials, + &GroupIds::from_iter([0u32, 1, 1, 0], 2)?, + &mut ctx, + )?; + let actual = acc.finish(2)?; + let expected = PrimitiveArray::from_option_iter([None, Some(8i64)]).into_array(); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] - fn fixed_size_overflow_and_nan() -> VortexResult<()> { + fn accumulates_typed_decimal_partials() -> VortexResult<()> { + let input_dtype = DecimalDType::new(10, 2); + let partial_dtype = DecimalDType::new(20, 2); + let partials = sum_partials( + DecimalArray::new( + buffer![200i64, 300, 500, 0], + partial_dtype, + Validity::NonNullable, + ) + .into_array(), + [false, false, false, true], + [false; 4], + )?; let mut ctx = array_session().create_execution_ctx(); - // FixedSize path: first group overflows -> null sum, second sums normally. - let elements = - PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); - let elem_dtype = DType::Primitive(PType::I64, NonNullable); - let groups = FixedSizeListArray::try_new(elements.clone(), 2, Validity::NonNullable, 2)? - .into_array(); - - let actual = grouped_sum_actual(&groups, &elem_dtype)?; - let expected = - grouped_sum_reference(&elements, &[(0, 2), (2, 2)], &[true, true], &elem_dtype)?; - let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + let mut acc = GroupedAccumulator::try_new( + Sum, + SumAggregateOpts::default(), + DType::Decimal(input_dtype, Nullability::Nullable), + )?; + acc.accumulate_partials( + &partials, + &GroupIds::from_iter([0u32, 1, 1, 0], 2)?, + &mut ctx, + )?; + let actual = acc.finish(2)?; + let expected = DecimalArray::new( + buffer![0i128, 800], + partial_dtype, + Validity::from_iter([false, true]), + ) + .into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped_state.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped_state.rs new file mode 100644 index 00000000000..bb9487cbc71 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped_state.rs @@ -0,0 +1,645 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; + +use num_traits::CheckedAdd; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use super::IS_EMPTY_FIELD; +use super::IS_OVERFLOW_FIELD; +use super::SUM_FIELD; +use super::checked_add_i64; +use super::checked_add_u64; +use super::decode_sum_partial_scalar; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::GroupedState; +use crate::arrays::BoolArray; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::struct_::StructArrayExt; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::validity::Validity; + +pub(super) enum SumGroupedValues { + Unsigned(Vec), + Signed(Vec), + Float(Vec), + Decimal8(Vec), + Decimal16(Vec), + Decimal32(Vec), + Decimal64(Vec), + Decimal128(Vec), + Decimal256(Vec), +} + +impl SumGroupedValues { + fn len(&self) -> usize { + match self { + Self::Unsigned(values) => values.len(), + Self::Signed(values) => values.len(), + Self::Float(values) => values.len(), + Self::Decimal8(values) => values.len(), + Self::Decimal16(values) => values.len(), + Self::Decimal32(values) => values.len(), + Self::Decimal64(values) => values.len(), + Self::Decimal128(values) => values.len(), + Self::Decimal256(values) => values.len(), + } + } + + fn resize(&mut self, len: usize) { + match self { + Self::Unsigned(values) => values.resize(len, 0), + Self::Signed(values) => values.resize(len, 0), + Self::Float(values) => values.resize(len, 0.0), + Self::Decimal8(values) => values.resize(len, 0), + Self::Decimal16(values) => values.resize(len, 0), + Self::Decimal32(values) => values.resize(len, 0), + Self::Decimal64(values) => values.resize(len, 0), + Self::Decimal128(values) => values.resize(len, 0), + Self::Decimal256(values) => values.resize(len, crate::dtype::i256::ZERO), + } + } +} + +pub(crate) struct SumGroupedState { + values: SumGroupedValues, + overflowed: Vec, + empty: Vec, + partial_dtype: DType, + return_dtype: DType, +} + +impl SumGroupedState { + pub(crate) fn try_new(partial_dtype: DType, return_dtype: DType) -> VortexResult { + let values = match &return_dtype { + DType::Primitive(PType::U64, _) => SumGroupedValues::Unsigned(Vec::new()), + DType::Primitive(PType::I64, _) => SumGroupedValues::Signed(Vec::new()), + DType::Primitive(PType::F64, _) => SumGroupedValues::Float(Vec::new()), + DType::Decimal(dtype, _) => match DecimalType::smallest_decimal_value_type(dtype) { + DecimalType::I8 => SumGroupedValues::Decimal8(Vec::new()), + DecimalType::I16 => SumGroupedValues::Decimal16(Vec::new()), + DecimalType::I32 => SumGroupedValues::Decimal32(Vec::new()), + DecimalType::I64 => SumGroupedValues::Decimal64(Vec::new()), + DecimalType::I128 => SumGroupedValues::Decimal128(Vec::new()), + DecimalType::I256 => SumGroupedValues::Decimal256(Vec::new()), + }, + dtype => vortex_bail!("Unsupported grouped sum return dtype: {dtype}"), + }; + Ok(Self { + values, + overflowed: Vec::new(), + empty: Vec::new(), + partial_dtype, + return_dtype, + }) + } + + pub(super) fn parts_mut(&mut self) -> (&mut SumGroupedValues, &mut [u8], &mut [u8]) { + (&mut self.values, &mut self.overflowed, &mut self.empty) + } + + pub(super) fn decimal_dtype(&self) -> Option { + self.return_dtype.as_decimal_opt().copied() + } +} + +impl GroupedState for SumGroupedState { + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn len(&self) -> usize { + self.values.len() + } + + fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { + let len = num_groups.max(self.len()); + self.values.resize(len); + self.overflowed.resize(len, 0); + self.empty.resize(len, 1); + Ok(()) + } + + fn is_saturated(&self, group_id: usize) -> bool { + if self.overflowed[group_id] != 0 { + return true; + } + matches!(&self.values, SumGroupedValues::Float(values) if values[group_id].is_nan()) + } + + fn combine_scalar(&mut self, group_id: usize, partial: Scalar) -> VortexResult<()> { + let (partial, partial_overflowed, partial_empty) = decode_sum_partial_scalar(partial)?; + if partial_empty { + return Ok(()); + } + self.empty[group_id] = 0; + if partial_overflowed { + self.overflowed[group_id] = 1; + return Ok(()); + } + if self.overflowed[group_id] != 0 { + return Ok(()); + } + + let decimal_dtype = self.decimal_dtype(); + let (values, overflowed, _) = self.parts_mut(); + match values { + SumGroupedValues::Unsigned(values) => { + let value = partial + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + overflowed[group_id] = u8::from(checked_add_u64(&mut values[group_id], value)); + } + SumGroupedValues::Signed(values) => { + let value = partial + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + overflowed[group_id] = u8::from(checked_add_i64(&mut values[group_id], value)); + } + SumGroupedValues::Float(values) => { + values[group_id] += partial + .as_primitive() + .typed_value::() + .vortex_expect("checked non-null"); + } + SumGroupedValues::Decimal8(values) => combine_decimal_scalar( + values, + overflowed, + group_id, + &partial, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal16(values) => combine_decimal_scalar( + values, + overflowed, + group_id, + &partial, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal32(values) => combine_decimal_scalar( + values, + overflowed, + group_id, + &partial, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal64(values) => combine_decimal_scalar( + values, + overflowed, + group_id, + &partial, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal128(values) => combine_decimal_scalar( + values, + overflowed, + group_id, + &partial, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal256(values) => combine_decimal_scalar( + values, + overflowed, + group_id, + &partial, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + } + Ok(()) + } + + fn partial_scalar(&self, group_id: usize) -> VortexResult { + let sum = match &self.values { + SumGroupedValues::Unsigned(values) => Scalar::primitive( + values.get(group_id).copied().unwrap_or(0), + Nullability::NonNullable, + ), + SumGroupedValues::Signed(values) => Scalar::primitive( + values.get(group_id).copied().unwrap_or(0), + Nullability::NonNullable, + ), + SumGroupedValues::Float(values) => Scalar::primitive( + values.get(group_id).copied().unwrap_or(0.0), + Nullability::NonNullable, + ), + SumGroupedValues::Decimal8(values) => decimal_scalar( + values.get(group_id).copied().unwrap_or(0), + self.decimal_dtype().vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal16(values) => decimal_scalar( + values.get(group_id).copied().unwrap_or(0), + self.decimal_dtype().vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal32(values) => decimal_scalar( + values.get(group_id).copied().unwrap_or(0), + self.decimal_dtype().vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal64(values) => decimal_scalar( + values.get(group_id).copied().unwrap_or(0), + self.decimal_dtype().vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal128(values) => decimal_scalar( + values.get(group_id).copied().unwrap_or(0), + self.decimal_dtype().vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal256(values) => decimal_scalar( + values + .get(group_id) + .copied() + .unwrap_or(crate::dtype::i256::ZERO), + self.decimal_dtype().vortex_expect("decimal state dtype"), + ), + }; + Ok(Scalar::struct_( + self.partial_dtype.clone(), + vec![ + sum, + Scalar::bool( + self.overflowed + .get(group_id) + .is_some_and(|&overflowed| overflowed != 0), + Nullability::NonNullable, + ), + Scalar::bool( + self.empty.get(group_id).is_none_or(|&empty| empty != 0), + Nullability::NonNullable, + ), + ], + )) + } + + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let partials = partials.clone().execute::(ctx)?; + let validity = partials + .as_ref() + .validity()? + .execute_mask(partials.as_ref().len(), ctx)?; + let sums = partials.unmasked_field_by_name(SUM_FIELD)?.clone(); + let partial_overflowed = partials + .unmasked_field_by_name(IS_OVERFLOW_FIELD)? + .clone() + .execute::(ctx)? + .into_bit_buffer(); + let partial_empty = partials + .unmasked_field_by_name(IS_EMPTY_FIELD)? + .clone() + .execute::(ctx)? + .into_bit_buffer(); + let rows = PartialRows { + group_ids, + validity: &validity, + overflowed: &partial_overflowed, + empty: &partial_empty, + }; + let decimal_dtype = self.decimal_dtype(); + let (values, overflowed, empty) = self.parts_mut(); + match values { + SumGroupedValues::Unsigned(values) => { + let sums = sums.execute::(ctx)?; + accumulate_primitive_partials( + PartialState { + values, + overflowed, + empty, + }, + sums.as_slice::(), + &rows, + checked_add_u64, + ); + } + SumGroupedValues::Signed(values) => { + let sums = sums.execute::(ctx)?; + accumulate_primitive_partials( + PartialState { + values, + overflowed, + empty, + }, + sums.as_slice::(), + &rows, + checked_add_i64, + ); + } + SumGroupedValues::Float(values) => { + let sums = sums.execute::(ctx)?; + accumulate_float_partials( + PartialState { + values, + overflowed, + empty, + }, + sums.as_slice::(), + &rows, + ); + } + SumGroupedValues::Decimal8(values) => accumulate_decimal_partials( + PartialState { + values, + overflowed, + empty, + }, + &sums.execute::(ctx)?, + &rows, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal16(values) => accumulate_decimal_partials( + PartialState { + values, + overflowed, + empty, + }, + &sums.execute::(ctx)?, + &rows, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal32(values) => accumulate_decimal_partials( + PartialState { + values, + overflowed, + empty, + }, + &sums.execute::(ctx)?, + &rows, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal64(values) => accumulate_decimal_partials( + PartialState { + values, + overflowed, + empty, + }, + &sums.execute::(ctx)?, + &rows, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal128(values) => accumulate_decimal_partials( + PartialState { + values, + overflowed, + empty, + }, + &sums.execute::(ctx)?, + &rows, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal256(values) => accumulate_decimal_partials( + PartialState { + values, + overflowed, + empty, + }, + &sums.execute::(ctx)?, + &rows, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + } + Ok(()) + } + + fn flush_partials(&mut self, num_groups: usize) -> VortexResult { + vortex_ensure!( + num_groups >= self.len(), + "Cannot flush {} groups after accumulating {} groups", + num_groups, + self.len() + ); + self.ensure_groups(num_groups)?; + let overflowed = std::mem::take(&mut self.overflowed); + let empty = std::mem::take(&mut self.empty); + let decimal_dtype = self.decimal_dtype(); + let sums = match &mut self.values { + SumGroupedValues::Unsigned(values) => { + PrimitiveArray::new(Buffer::from(std::mem::take(values)), Validity::NonNullable) + .into_array() + } + SumGroupedValues::Signed(values) => { + PrimitiveArray::new(Buffer::from(std::mem::take(values)), Validity::NonNullable) + .into_array() + } + SumGroupedValues::Float(values) => { + PrimitiveArray::new(Buffer::from(std::mem::take(values)), Validity::NonNullable) + .into_array() + } + SumGroupedValues::Decimal8(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + Validity::NonNullable, + ), + SumGroupedValues::Decimal16(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + Validity::NonNullable, + ), + SumGroupedValues::Decimal32(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + Validity::NonNullable, + ), + SumGroupedValues::Decimal64(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + Validity::NonNullable, + ), + SumGroupedValues::Decimal128(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + Validity::NonNullable, + ), + SumGroupedValues::Decimal256(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + Validity::NonNullable, + ), + }; + let names = self.partial_dtype.as_struct_fields().names().clone(); + Ok(StructArray::try_new( + names, + vec![ + sums, + BoolArray::from_iter(overflowed.into_iter().map(|value| value != 0)).into_array(), + BoolArray::from_iter(empty.into_iter().map(|value| value != 0)).into_array(), + ], + num_groups, + Validity::AllValid, + )? + .into_array()) + } +} + +struct PartialRows<'a> { + group_ids: &'a [u32], + validity: &'a Mask, + overflowed: &'a BitBuffer, + empty: &'a BitBuffer, +} + +struct PartialState<'a, T> { + values: &'a mut [T], + overflowed: &'a mut [u8], + empty: &'a mut [u8], +} + +fn for_each_nonempty_partial(rows: &PartialRows<'_>, mut f: impl FnMut(usize, bool)) { + for (idx, ((valid, overflowed), empty)) in rows + .validity + .iter() + .zip(rows.overflowed.iter()) + .zip(rows.empty.iter()) + .enumerate() + { + if valid && !empty { + f(idx, overflowed); + } + } +} + +fn accumulate_primitive_partials( + state: PartialState<'_, T>, + partials: &[T], + rows: &PartialRows<'_>, + checked_add: fn(&mut T, T) -> bool, +) { + for_each_nonempty_partial(rows, |idx, is_overflowed| { + let group = rows.group_ids[idx] as usize; + state.empty[group] = 0; + if is_overflowed || checked_add(&mut state.values[group], partials[idx]) { + state.overflowed[group] = 1; + } + }); +} + +fn accumulate_float_partials( + state: PartialState<'_, f64>, + partials: &[f64], + rows: &PartialRows<'_>, +) { + for_each_nonempty_partial(rows, |idx, is_overflowed| { + let group = rows.group_ids[idx] as usize; + state.empty[group] = 0; + if is_overflowed { + state.overflowed[group] = 1; + } else { + state.values[group] += partials[idx]; + } + }); +} + +fn accumulate_decimal_partials( + state: PartialState<'_, I>, + partials: &DecimalArray, + rows: &PartialRows<'_>, + dtype: DecimalDType, +) where + I: NativeDecimalType + CheckedAdd, +{ + match_each_decimal_value_type!(partials.values_type(), |T| { + accumulate_decimal_partial_values(state, &partials.buffer::(), rows, dtype); + }); +} + +fn accumulate_decimal_partial_values( + state: PartialState<'_, I>, + partials: &[T], + rows: &PartialRows<'_>, + dtype: DecimalDType, +) where + T: NativeDecimalType, + I: NativeDecimalType + CheckedAdd, +{ + for_each_nonempty_partial(rows, |idx, is_overflowed| { + let group = rows.group_ids[idx] as usize; + state.empty[group] = 0; + if is_overflowed { + state.overflowed[group] = 1; + } else { + let Some(value) = ::from(partials[idx]) else { + state.overflowed[group] = 1; + return; + }; + add_decimal(state.values, state.overflowed, group, value, dtype); + } + }); +} + +fn combine_decimal_scalar( + values: &mut [T], + overflowed: &mut [u8], + group_id: usize, + partial: &Scalar, + dtype: DecimalDType, +) where + T: NativeDecimalType + CheckedAdd, +{ + let value = partial + .as_decimal() + .decimal_value() + .vortex_expect("checked non-null") + .cast::() + .vortex_expect("decimal partial must use grouped state width"); + add_decimal(values, overflowed, group_id, value, dtype); +} + +pub(super) fn add_decimal( + values: &mut [T], + overflowed: &mut [u8], + group_id: usize, + value: T, + dtype: DecimalDType, +) where + T: NativeDecimalType + CheckedAdd, +{ + if overflowed[group_id] != 0 { + return; + } + let Some(result) = values[group_id].checked_add(&value) else { + overflowed[group_id] = 1; + return; + }; + let precision = usize::from(dtype.precision()); + if T::MIN_BY_PRECISION[precision] <= result && result <= T::MAX_BY_PRECISION[precision] { + values[group_id] = result; + } else { + overflowed[group_id] = 1; + } +} + +fn decimal_scalar(value: T, dtype: DecimalDType) -> Scalar +where + DecimalValue: From, +{ + Scalar::decimal(DecimalValue::from(value), dtype, Nullability::NonNullable) +} + +fn decimal_array( + values: Vec, + dtype: DecimalDType, + validity: Validity, +) -> ArrayRef { + DecimalArray::new(Buffer::from(values), dtype, validity).into_array() +} diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 89b19b5575f..44bf27bd481 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -5,12 +5,13 @@ mod bool; mod constant; mod decimal; mod grouped; +mod grouped_state; mod primitive; use std::fmt; use std::fmt::Display; use std::fmt::Formatter; -pub(crate) use grouped::PrimitiveGroupedSumEncodingKernel; +pub(crate) use grouped::SUM_GROUPED_KERNEL; use prost::Message; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -24,6 +25,7 @@ use vortex_session::registry::CachedId; use self::bool::accumulate_bool; use self::constant::multiply_constant; use self::decimal::accumulate_decimal; +use self::grouped_state::SumGroupedState; use self::primitive::accumulate_primitive; use crate::ArrayRef; use crate::Canonical; @@ -34,6 +36,7 @@ use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::GroupedState; use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::ConstantArray; use crate::arrays::StructArray; @@ -55,9 +58,9 @@ use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -const SUM_FIELD: &str = "sum"; -const IS_OVERFLOW_FIELD: &str = "is_overflow"; -const IS_EMPTY_FIELD: &str = "is_empty"; +pub(super) const SUM_FIELD: &str = "sum"; +pub(super) const IS_OVERFLOW_FIELD: &str = "is_overflow"; +pub(super) const IS_EMPTY_FIELD: &str = "is_empty"; /// Return the sum of an array. /// @@ -249,6 +252,21 @@ impl AggregateFnVTable for Sum { }) } + fn grouped_state( + &self, + options: &Self::Options, + input_dtype: &DType, + partial_dtype: &DType, + ) -> VortexResult> { + let return_dtype = self + .return_dtype(options, input_dtype) + .ok_or_else(|| vortex_err!("Unsupported sum dtype: {input_dtype}"))?; + Ok(Box::new(SumGroupedState::try_new( + partial_dtype.clone(), + return_dtype, + )?)) + } + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { let (other_sum, other_is_overflow, other_is_empty) = decode_sum_partial_scalar(other)?; validate_sum_field_dtype(&other_sum, &partial.return_dtype)?; @@ -427,7 +445,7 @@ pub enum SumState { }, } -fn decode_sum_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { +pub(super) fn decode_sum_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> { vortex_ensure!(!scalar.is_null(), "Sum partial must not be null"); let Some(fields) = scalar.as_struct_opt() else { @@ -657,6 +675,7 @@ mod arithmetic_tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; + use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::SumAggregateOpts; use crate::aggregate_fn::fns::sum::sum; @@ -688,6 +707,25 @@ mod arithmetic_tests { sum_result_partial_scalar(value, &return_dtype, false) } + fn accumulate_list( + acc: &mut GroupedAccumulator, + groups: &ArrayRef, + ctx: &mut crate::ExecutionCtx, + ) -> VortexResult { + let grouped: GroupedArray = + if let Some(groups) = groups.as_opt::() { + groups.into_owned().into() + } else if let Some(groups) = groups.as_opt::() { + groups.into_owned().into() + } else { + unreachable!("grouped sum test requires a canonical list array") + }; + let num_groups = grouped.len(); + let (values, group_ids) = grouped.dense_input(ctx)?; + acc.accumulate(&values, &group_ids, ctx)?; + Ok(num_groups) + } + /// Sum an array with an initial value (test-only helper). fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { let mut ctx = array_session().create_execution_ctx(); @@ -852,8 +890,9 @@ mod arithmetic_tests { fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { let mut acc = GroupedAccumulator::try_new(Sum, SumAggregateOpts::default(), elem_dtype.clone())?; - acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; - acc.finish() + let mut ctx = array_session().create_execution_ctx(); + let num_groups = accumulate_list(&mut acc, groups, &mut ctx)?; + acc.finish(num_groups) } #[test] @@ -944,16 +983,16 @@ mod arithmetic_tests { let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?; - acc.accumulate_list(&groups1.into_array(), &mut ctx)?; - let result1 = acc.finish()?; + let num_groups = accumulate_list(&mut acc, &groups1.into_array(), &mut ctx)?; + let result1 = acc.finish(num_groups)?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result1, &expected1, &mut ctx); let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; - acc.accumulate_list(&groups2.into_array(), &mut ctx)?; - let result2 = acc.finish()?; + let num_groups = accumulate_list(&mut acc, &groups2.into_array(), &mut ctx)?; + let result2 = acc.finish(num_groups)?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); assert_arrays_eq!(&result2, &expected2, &mut ctx); diff --git a/vortex-array/src/aggregate_fn/fns/sum/tests.rs b/vortex-array/src/aggregate_fn/fns/sum/tests.rs index c168b39e757..819552d8cd0 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/tests.rs @@ -22,6 +22,7 @@ use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::array_session; use crate::arrays::BoolArray; @@ -58,6 +59,25 @@ fn partial_with_value(value: Scalar) -> VortexResult { sum_result_partial_scalar(value, &return_dtype, false) } +fn accumulate_list( + acc: &mut GroupedAccumulator, + groups: &ArrayRef, + ctx: &mut crate::ExecutionCtx, +) -> VortexResult { + let grouped: GroupedArray = + if let Some(groups) = groups.as_opt::() { + groups.into_owned().into() + } else if let Some(groups) = groups.as_opt::() { + groups.into_owned().into() + } else { + unreachable!("grouped sum test requires a canonical list array") + }; + let num_groups = grouped.len(); + let (values, group_ids) = grouped.dense_input(ctx)?; + acc.accumulate(&values, &group_ids, ctx)?; + Ok(num_groups) +} + #[test] fn sum_uses_new_partial_shape_by_default() { let options = SumAggregateOpts::default(); @@ -754,8 +774,8 @@ fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult VortexRe SumAggregateOpts::default(), DType::Primitive(PType::I64, Nullable), )?; - acc.accumulate_list(&groups, &mut ctx)?; - let partials = acc.flush()?; + let num_groups = accumulate_list(&mut acc, &groups, &mut ctx)?; + let partials = acc.flush_partials(num_groups)?; let value = partials.execute_scalar(0, &mut ctx)?; let fields = value.as_struct(); @@ -799,7 +819,8 @@ fn grouped_sum_partial_distinguishes_empty_overflow_and_null_group() -> VortexRe Some(false) ); - for index in [1, 2] { + // Dense IDs represent both an empty list and a null/omitted list as an empty group state. + for index in [1, 2, 4] { let empty = partials.execute_scalar(index, &mut ctx)?; let fields = empty.as_struct(); assert_eq!( @@ -842,7 +863,6 @@ fn grouped_sum_partial_distinguishes_empty_overflow_and_null_group() -> VortexRe .and_then(|is_empty| is_empty.as_bool().value()), Some(false) ); - assert!(partials.execute_scalar(4, &mut ctx)?.is_null()); Ok(()) } @@ -955,16 +975,16 @@ fn grouped_sum_finish_resets() -> VortexResult<()> { let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); let groups1 = FixedSizeListArray::try_new(elements1, 2, Validity::NonNullable, 2)?; - acc.accumulate_list(&groups1.into_array(), &mut ctx)?; - let result1 = acc.finish()?; + let num_groups = accumulate_list(&mut acc, &groups1.into_array(), &mut ctx)?; + let result1 = acc.finish(num_groups)?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result1, &expected1, &mut ctx); let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; - acc.accumulate_list(&groups2.into_array(), &mut ctx)?; - let result2 = acc.finish()?; + let num_groups = accumulate_list(&mut acc, &groups2.into_array(), &mut ctx)?; + let result2 = acc.finish(num_groups)?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); assert_arrays_eq!(&result2, &expected2, &mut ctx); diff --git a/vortex-array/src/aggregate_fn/kernels.rs b/vortex-array/src/aggregate_fn/kernels.rs index c5af0902cbb..0651ec4ddff 100644 --- a/vortex-array/src/aggregate_fn/kernels.rs +++ b/vortex-array/src/aggregate_fn/kernels.rs @@ -4,14 +4,20 @@ //! Pluggable aggregate function kernels used to provide encoding-specific implementations of //! aggregate functions. +use std::any::Any; use std::fmt::Debug; +use std::marker::PhantomData; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::GroupIds; +use crate::aggregate_fn::GroupedState; use crate::scalar::Scalar; /// A pluggable kernel for an aggregate function. @@ -27,26 +33,108 @@ pub trait DynAggregateKernel: 'static + Send + Sync + Debug { ) -> VortexResult>; } +/// A typed grouped aggregate kernel. +/// +/// Implementations receive the concrete aggregate options and typed partial state. Return +/// `Ok(false)` when the kernel cannot handle the current values or group-id encodings. +pub trait GroupedAggregateKernel: 'static + Send + Sync + Debug { + /// Concrete aggregate-owned grouped state consumed by this kernel. + type State: GroupedState; + + /// Accumulate `batch` into `states` according to `group_ids`. + fn grouped_accumulate( + &self, + options: &V::Options, + state: &mut Self::State, + batch: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult; +} + +/// Bridges a typed [`GroupedAggregateKernel`] to type-erased grouped kernel dispatch. +pub struct GroupedAggregateKernelAdapter { + kernel: K, + _phantom: PhantomData V>, +} + +impl GroupedAggregateKernelAdapter { + /// Create a new adapter around `kernel`. + pub const fn new(kernel: K) -> Self { + Self { + kernel, + _phantom: PhantomData, + } + } +} + +impl Debug for GroupedAggregateKernelAdapter +where + V: AggregateFnVTable, + K: GroupedAggregateKernel, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GroupedAggregateKernelAdapter") + .field("kernel", &self.kernel) + .finish() + } +} + /// A pluggable kernel for batch aggregation of many groups. /// -/// A kernel can be registered either for an aggregate function regardless of the element encoding, -/// or for a specific aggregate function and element encoding. Element-encoding kernels are matched -/// on the inner array of the provided grouped array, not on the outer list encoding. This is more -/// pragmatic than having every kernel match on the outer list encoding and having to deal with the -/// possibility of multiple list encodings. +/// A grouped kernel can be registered for an aggregate function regardless of input encodings, or +/// for a specific aggregate function plus values and/or group-id encoding. /// -/// Each value in the grouped array represents a group and the result of the grouped aggregate -/// should be an array of the same length, where each element is the aggregate state of the -/// corresponding group. +/// Kernels receive the same dense group ordinals that the caller passed to the grouped accumulator +/// and may aggregate directly in the encoded domain. /// -/// Return `Ok(None)` if the kernel cannot be applied to the given aggregate function. +/// Return `Ok(false)` if the kernel cannot be applied to the given aggregate function or input +/// encodings. pub trait DynGroupedAggregateKernel: 'static + Send + Sync + Debug { - /// Aggregate each group in the provided grouped array and return an array of the aggregate - /// states. - fn grouped_aggregate( + /// Accumulate values into type-erased partial state. + fn grouped_accumulate( &self, aggregate_fn: &AggregateFnRef, - groups: &GroupedArray, + batch: &ArrayRef, + group_ids: &GroupIds, + states: &mut dyn Any, ctx: &mut ExecutionCtx, - ) -> VortexResult>; + ) -> VortexResult; +} + +impl DynGroupedAggregateKernel for GroupedAggregateKernelAdapter +where + V: AggregateFnVTable, + K: GroupedAggregateKernel, +{ + fn grouped_accumulate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + group_ids: &GroupIds, + states: &mut dyn Any, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(options) = aggregate_fn.as_opt::() else { + return Ok(false); + }; + + let Some(state) = states.downcast_mut::() else { + vortex_bail!( + "Grouped aggregate kernel for {} received incompatible partial state", + aggregate_fn.id() + ); + }; + + vortex_ensure!( + state.len() >= group_ids.num_groups(), + "Grouped aggregate kernel for {} received {} partial states for {} groups", + aggregate_fn.id(), + state.len(), + group_ids.num_groups() + ); + + self.kernel + .grouped_accumulate(options, state, batch, group_ids, ctx) + } } diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 0fd07afc9c4..0f7439b3682 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -20,8 +20,8 @@ use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; use crate::aggregate_fn::fns::bounded_max::BoundedMax; use crate::aggregate_fn::fns::bounded_min::BoundedMin; +use crate::aggregate_fn::fns::count::COUNT_GROUPED_KERNEL; use crate::aggregate_fn::fns::count::Count; -use crate::aggregate_fn::fns::count::CountGroupedKernel; use crate::aggregate_fn::fns::first::First; use crate::aggregate_fn::fns::is_constant::IsConstant; use crate::aggregate_fn::fns::is_sorted::IsSorted; @@ -31,7 +31,7 @@ use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; -use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel; +use crate::aggregate_fn::fns::sum::SUM_GROUPED_KERNEL; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; @@ -40,7 +40,6 @@ use crate::array::ArrayId; use crate::array::VTable; use crate::arrays::Chunked; use crate::arrays::Dict; -use crate::arrays::Primitive; use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; @@ -57,8 +56,7 @@ pub struct AggregateFnSession { registry: AggregateFnRegistry, kernels: AggregateKernelRegistry, - grouped_kernels: GroupedKernelRegistry, - grouped_encoding_kernels: GroupedEncodingKernelRegistry, + grouped_kernels: ArcSwapMap, } impl SessionVar for AggregateFnSession { @@ -72,25 +70,39 @@ impl SessionVar for AggregateFnSession { } type AggregateKernelKey = (ArrayId, Option); -type GroupedEncodingKernelKey = (ArrayId, AggregateFnId); + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +struct GroupedAggregateKernelKey { + aggregate_id: AggregateFnId, + values_id: Option, + group_ids_id: Option, +} + +impl GroupedAggregateKernelKey { + fn new( + aggregate_id: AggregateFnId, + values_id: Option, + group_ids_id: Option, + ) -> Self { + Self { + aggregate_id, + values_id, + group_ids_id, + } + } +} /// Registry of aggregate function plugins, keyed by aggregate function id. type AggregateFnRegistry = ArcSwapMap; /// Registry of aggregate kernels, keyed by encoding and optional aggregate function. type AggregateKernelRegistry = ArcSwapMap; -/// Registry of encoding-agnostic grouped aggregate kernels, keyed by aggregate function id. -type GroupedKernelRegistry = ArcSwapMap; -/// Registry of grouped aggregate kernels, keyed by encoding and aggregate function. -type GroupedEncodingKernelRegistry = - ArcSwapMap; impl Default for AggregateFnSession { fn default() -> Self { let this = Self { - registry: AggregateFnRegistry::default(), - kernels: AggregateKernelRegistry::default(), - grouped_kernels: GroupedKernelRegistry::default(), - grouped_encoding_kernels: GroupedEncodingKernelRegistry::default(), + registry: ArcSwapMap::default(), + kernels: ArcSwapMap::default(), + grouped_kernels: ArcSwapMap::default(), }; // Register the built-in aggregate functions @@ -118,14 +130,8 @@ impl Default for AggregateFnSession { this.register_aggregate_kernel(Dict.id(), Some(MinMax.id()), &DictMinMaxKernel); this.register_aggregate_kernel(Dict.id(), Some(IsConstant.id()), &DictIsConstantKernel); this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel); - - // Register the built-in grouped aggregate kernels. - this.register_grouped_kernel(Count.id(), &CountGroupedKernel); - this.register_grouped_encoding_kernel( - Primitive.id(), - Sum.id(), - &PrimitiveGroupedSumEncodingKernel, - ); + this.register_grouped_kernel(Count.id(), None, None, &COUNT_GROUPED_KERNEL); + this.register_grouped_kernel(Sum.id(), None, None, &SUM_GROUPED_KERNEL); this } @@ -195,54 +201,62 @@ impl AggregateFnSession { self.kernels.insert(id, kernel); } - /// Returns the grouped aggregate kernel registered for `agg_fn_id`, if any. + /// Returns the grouped aggregate kernel registered for this aggregate and pair of encodings. /// - /// These kernels are independent of the element encoding and are checked for each element - /// representation, after any kernel registered for the current element encoding. + /// Lookup first checks the exact `(aggregate, values encoding, group ids encoding)` key, then + /// falls back through `(aggregate, values encoding, any group ids)`, `(aggregate, any values, + /// group ids encoding)`, and finally `(aggregate, any values, any group ids)`. pub fn find_grouped_kernel( &self, agg_fn_id: impl Into, + values_id: impl Into, + group_ids_id: impl Into, ) -> Option<&'static dyn DynGroupedAggregateKernel> { let fn_id = agg_fn_id.into(); - self.grouped_kernels - .read(|kernels| kernels.get(&fn_id).copied()) - } - - /// Registers a grouped aggregate kernel for an aggregate function. - pub fn register_grouped_kernel( - &self, - agg_fn_id: impl Into, - kernel: &'static dyn DynGroupedAggregateKernel, - ) { - let fn_id = agg_fn_id.into(); - self.grouped_kernels.insert(fn_id, kernel) + let values_id = values_id.into(); + let group_ids_id = group_ids_id.into(); + self.grouped_kernels.read(|kernels| { + kernels + .get(&GroupedAggregateKernelKey::new( + fn_id, + Some(values_id), + Some(group_ids_id), + )) + .or_else(|| { + kernels.get(&GroupedAggregateKernelKey::new( + fn_id, + Some(values_id), + None, + )) + }) + .or_else(|| { + kernels.get(&GroupedAggregateKernelKey::new( + fn_id, + None, + Some(group_ids_id), + )) + }) + .or_else(|| kernels.get(&GroupedAggregateKernelKey::new(fn_id, None, None))) + .copied() + }) } - /// Returns the grouped aggregate kernel registered for `array_id` and `agg_fn_id`, if any. + /// Registers a grouped aggregate kernel. /// - /// These kernels are matched against each intermediate element encoding while the grouped - /// accumulator executes the element array. - pub fn find_grouped_encoding_kernel( - &self, - array_id: impl Into, - agg_fn_id: impl Into, - ) -> Option<&'static dyn DynGroupedAggregateKernel> { - let id = array_id.into(); - let fn_id = agg_fn_id.into(); - self.grouped_encoding_kernels - .read(|kernels| kernels.get(&(id, fn_id)).copied()) - } - - /// Registers a grouped aggregate kernel for a specific aggregate function and array encoding. - pub fn register_grouped_encoding_kernel( + /// `values_id` and `group_ids_id` are optional wildcards. Passing `None` for either dimension + /// makes the kernel a fallback for that encoding dimension. + pub fn register_grouped_kernel( &self, - array_id: impl Into, agg_fn_id: impl Into, + values_id: Option, + group_ids_id: Option, kernel: &'static dyn DynGroupedAggregateKernel, ) { - let id = array_id.into(); let fn_id = agg_fn_id.into(); - self.grouped_encoding_kernels.insert((id, fn_id), kernel) + self.grouped_kernels.insert( + GroupedAggregateKernelKey::new(fn_id, values_id, group_ids_id), + kernel, + ) } } @@ -254,3 +268,88 @@ pub trait AggregateFnSessionExt: SessionExt { } } impl AggregateFnSessionExt for S {} + +#[cfg(test)] +mod tests { + use std::any::Any; + + use vortex_error::VortexResult; + use vortex_session::registry::CachedId; + + use super::*; + use crate::ArrayRef; + use crate::ExecutionCtx; + use crate::aggregate_fn::AggregateFnRef; + use crate::aggregate_fn::GroupIds; + use crate::arrays::Constant; + use crate::arrays::Primitive; + + #[derive(Debug)] + struct TestGroupedKernel; + + impl DynGroupedAggregateKernel for TestGroupedKernel { + fn grouped_accumulate( + &self, + _aggregate_fn: &AggregateFnRef, + _batch: &ArrayRef, + _group_ids: &GroupIds, + _states: &mut dyn Any, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(false) + } + } + + static GENERIC_KERNEL: TestGroupedKernel = TestGroupedKernel; + static GROUP_IDS_KERNEL: TestGroupedKernel = TestGroupedKernel; + static VALUES_KERNEL: TestGroupedKernel = TestGroupedKernel; + static EXACT_KERNEL: TestGroupedKernel = TestGroupedKernel; + + fn assert_same_kernel( + actual: Option<&'static dyn DynGroupedAggregateKernel>, + expected: &'static dyn DynGroupedAggregateKernel, + ) { + assert!(std::ptr::eq( + actual.expect("expected registered grouped kernel"), + expected + )); + } + + #[test] + fn grouped_kernel_lookup_prefers_exact_then_value_then_group_ids() { + let session = AggregateFnSession::default(); + static AGGREGATE_ID: CachedId = CachedId::new("test.grouped_lookup"); + let aggregate_id = *AGGREGATE_ID; + let values_id = Primitive.id(); + let group_ids_id = Constant.id(); + + session.register_grouped_kernel(aggregate_id, None, None, &GENERIC_KERNEL); + assert_same_kernel( + session.find_grouped_kernel(aggregate_id, values_id, group_ids_id), + &GENERIC_KERNEL, + ); + + session.register_grouped_kernel(aggregate_id, None, Some(group_ids_id), &GROUP_IDS_KERNEL); + assert_same_kernel( + session.find_grouped_kernel(aggregate_id, values_id, group_ids_id), + &GROUP_IDS_KERNEL, + ); + + session.register_grouped_kernel(aggregate_id, Some(values_id), None, &VALUES_KERNEL); + assert_same_kernel( + session.find_grouped_kernel(aggregate_id, values_id, group_ids_id), + &VALUES_KERNEL, + ); + + session.register_grouped_kernel( + aggregate_id, + Some(values_id), + Some(group_ids_id), + &EXACT_KERNEL, + ); + assert_same_kernel( + session.find_grouped_kernel(aggregate_id, values_id, group_ids_id), + &EXACT_KERNEL, + ); + } +} diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 3f0a4cf567c..89305ad55ef 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -20,6 +20,8 @@ use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; +use crate::aggregate_fn::DefaultGroupedState; +use crate::aggregate_fn::GroupedState; use crate::dtype::DType; use crate::scalar::Scalar; @@ -114,6 +116,24 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { input_dtype: &DType, ) -> VortexResult; + /// Create the dense state container used by grouped accumulation. + /// + /// Aggregates may override this to keep monomorphic, cache-dense state. The default stores one + /// [`Self::Partial`] per group. + fn grouped_state( + &self, + options: &Self::Options, + input_dtype: &DType, + partial_dtype: &DType, + ) -> VortexResult> { + Ok(Box::new(DefaultGroupedState::new( + self.clone(), + options.clone(), + input_dtype.clone(), + partial_dtype.clone(), + ))) + } + /// Combine partial scalar state into the accumulator. fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()>; @@ -123,6 +143,17 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// options and input dtype used to construct the state. fn to_scalar(&self, partial: &Self::Partial) -> VortexResult; + /// Try to convert dense partial states directly into a partial-state array. + /// + /// Returning `Ok(None)` falls back to scalarizing each partial with [`Self::to_scalar`]. + fn partials_to_array( + &self, + _partials: &[Self::Partial], + _partial_dtype: &DType, + ) -> VortexResult> { + Ok(None) + } + /// Reset the state of the accumulator to an empty group. fn reset(&self, partial: &mut Self::Partial); diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs index 669f8344a96..2ad092d8515 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -15,9 +15,12 @@ use crate::IntoArray; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::GroupedArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::ListView; use crate::dtype::DType; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -131,9 +134,20 @@ fn list_sum_impl( options: &NumericalAggregateOpts, ctx: &mut ExecutionCtx, ) -> VortexResult { + let grouped: GroupedArray = if let Some(fsl) = canonical.as_opt::() { + fsl.into_owned().into() + } else if let Some(lv) = canonical.as_opt::() { + lv.into_owned().into() + } else { + let dtype = canonical.dtype(); + vortex_bail!("list_sum() requires List or FixedSizeList but got {dtype}") + }; + + let num_groups = grouped.len(); + let (values, group_ids) = grouped.dense_input(ctx)?; let mut acc = GroupedAccumulator::try_new(Sum, (*options).into(), elem_dtype)?; - acc.accumulate_list(&canonical, ctx)?; - acc.finish() + acc.accumulate(&values, &group_ids, ctx)?; + acc.finish(num_groups) } #[cfg(test)]