From 6bdcd32b62d3852247efe48a7a3a2af7d7fd922c Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 11 Jun 2026 17:35:52 -0400 Subject: [PATCH 1/9] Support dense grouped aggregate accumulation Signed-off-by: "Nicholas Gates" --- vortex-array/benches/aggregate_grouped.rs | 80 ++- vortex-array/src/aggregate_fn/accumulator.rs | 2 +- .../src/aggregate_fn/accumulator_grouped.rs | 478 ++++++++++-------- .../src/aggregate_fn/fns/count/mod.rs | 56 ++ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 267 +++++++--- vortex-array/src/aggregate_fn/kernels.rs | 64 ++- vortex-array/src/aggregate_fn/vtable.rs | 29 ++ 7 files changed, 637 insertions(+), 339 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index b067314c1d9..2d46a5cce8a 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -18,10 +18,8 @@ use vortex_array::aggregate_fn::EmptyOptions; 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; @@ -45,44 +43,42 @@ 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 +struct DenseGroupedInput { + values: ArrayRef, + group_ids: Vec, + num_groups: usize, +} + +fn dense_grouped_input(values: ArrayRef, group_sizes: &[usize]) -> DenseGroupedInput { + assert_eq!(values.len(), total_element_count(group_sizes)); + + let group_ids = group_sizes .iter() - .map(|&size| { - let current_offset = offset; - offset += size; - current_offset as u32 - }) + .enumerate() + .flat_map(|(group_id, &size)| std::iter::repeat_n(group_id as u32, size)) .collect(); - let sizes: Buffer = group_sizes.iter().map(|&size| size as u32).collect(); - assert_eq!(elements.len(), total_element_count(group_sizes)); - - ListViewArray::try_new( - elements, - offsets.into_array(), - sizes.into_array(), - Validity::NonNullable, - ) - .unwrap() - .into_array() + DenseGroupedInput { + values, + group_ids, + num_groups: group_sizes.len(), + } } -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| { @@ -92,26 +88,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); @@ -122,40 +118,38 @@ 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 grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef +fn grouped_accumulator(input: &DenseGroupedInput, vtable: V) -> ArrayRef where V: AggregateFnVTable + Clone, { let mut acc = - GroupedAccumulator::try_new(vtable, EmptyOptions, list_element_dtype(list_view)).unwrap(); - acc.accumulate_list(list_view, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); - divan::black_box(acc.finish().unwrap()) + GroupedAccumulator::try_new(vtable, EmptyOptions, input.values.dtype().clone()).unwrap(); + acc.accumulate( + &input.values, + &input.group_ids, + input.num_groups, + &mut LEGACY_SESSION.create_execution_ctx(), + ) + .unwrap(); + divan::black_box(acc.finish(input.num_groups).unwrap()) } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index c89418e67a6..ab4e0ee26ba 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -172,7 +172,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 4b94159127b..66da32b4085 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -1,19 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use arrow_buffer::ArrowNativeType; 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::AnyCanonical; use crate::ArrayRef; -use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; @@ -22,26 +15,21 @@ use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; +use crate::aggregate_fn::kernels::GroupedAggregateKernelResult; use crate::aggregate_fn::session::AggregateFnSessionExt; -use crate::arrays::ChunkedArray; -use crate::arrays::FixedSizeListArray; -use crate::arrays::ListViewArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; use crate::builders::builder_with_capacity; -use crate::builtins::ArrayBuiltins; +use crate::columnar::AnyColumnar; use crate::dtype::DType; -use crate::dtype::IntegerPType; use crate::executor::max_iterations; -use crate::match_each_integer_ptype; +use crate::scalar::Scalar; /// Reference-counted type-erased grouped accumulator. pub type GroupedAccumulatorRef = Box; -/// An accumulator used for computing grouped aggregates. +/// An accumulator used for computing aggregates over dense group ids. /// -/// Note that the groups must be processed in order, and the accumulator does not support random -/// access to groups. +/// Group ids are dense `u32` slots in the range `0..num_groups`. 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, @@ -55,8 +43,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, + /// Dense per-group partial state. + partials: Vec, } impl GroupedAccumulator { @@ -84,249 +72,315 @@ impl GroupedAccumulator { dtype, return_dtype, partial_dtype, - partials: vec![], + partials: Vec::new(), }) } -} - -/// 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<()> { + vortex_ensure!( + num_groups <= (u32::MAX as usize) + 1, + "num_groups {} exceeds dense u32 group id capacity", + 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; -} + while self.partials.len() < num_groups { + self.partials + .push(self.vtable.empty_partial(&self.options, &self.dtype)?); + } + Ok(()) + } -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() - ), - }; + fn validate_group_ids(&self, group_ids: &[u32], num_groups: usize) -> VortexResult<()> { vortex_ensure!( - elements_dtype.as_ref() == &self.dtype, - "Input DType mismatch: expected {}, got {}", - self.dtype, - elements_dtype + num_groups <= (u32::MAX as usize) + 1, + "num_groups {} exceeds dense u32 group id capacity", + num_groups ); - - // 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_list_view(&groups, ctx), - Canonical::FixedSizeList(groups) => self.accumulate_fixed_size_list(&groups, ctx), - _ => vortex_panic!("We checked the DType above, so this should never happen"), + for &group_id in group_ids { + vortex_ensure!( + (group_id as usize) < num_groups, + "Group id {} out of range for {} groups", + group_id, + num_groups + ); } + Ok(()) } - fn flush(&mut self) -> VortexResult { - let states = std::mem::take(&mut self.partials); - Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array()) + fn accumulate_kernel_result( + &mut self, + result: GroupedAggregateKernelResult, + num_groups: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + self.accumulate_partials(result.partials(), result.group_ids(), num_groups, ctx) } - fn finish(&mut self) -> VortexResult { - let states = self.flush()?; - let results = self.vtable.finalize(states)?; + fn accumulate_fallback( + &mut self, + batch: &ArrayRef, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + 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); + } - vortex_ensure!( - results.dtype() == &self.return_dtype, - "Return DType mismatch: expected {}, got {}", - self.return_dtype, - results.dtype() - ); + let first = first as usize; + let mut buckets = vec![Vec::new(); last as usize - first + 1]; + for (row_idx, &group_id) in group_ids.iter().enumerate() { + buckets[group_id as usize - first].push(row_idx as u64); + } - Ok(results) + for (offset, rows) in buckets.into_iter().enumerate() { + if rows.is_empty() { + continue; + } + + let group = first + offset; + if self.vtable.is_saturated(&self.partials[group]) { + continue; + } + + let taken = batch.clone().take(Buffer::from_iter(rows).into_array())?; + let mut accumulator = Accumulator::try_new( + self.vtable.clone(), + self.options.clone(), + self.dtype.clone(), + )?; + accumulator.accumulate(&taken, ctx)?; + let partial = accumulator.flush()?; + self.vtable + .combine_partials(&mut self.partials[group], partial)?; + } + Ok(()) } } -impl GroupedAccumulator { - fn accumulate_list_view( +/// 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. + fn accumulate( &mut self, - groups: &ListViewArray, + batch: &ArrayRef, + group_ids: &[u32], + num_groups: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()>; + + /// Fold columnar partial states into dense group state. + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &[u32], + num_groups: usize, + 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: &[u32], + num_groups: usize, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let mut elements = groups.elements().clone(); - let groups_validity = groups.validity()?; + 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.validate_group_ids(group_ids, num_groups)?; + self.ensure_groups(num_groups)?; + let session = ctx.session().clone(); + if let Some(kernel) = session + .aggregate_fns() + .find_grouped_kernel(batch.encoding_id(), self.aggregate_fn.id()) + && let Some(result) = + kernel.grouped_aggregate(&self.aggregate_fn, batch, group_ids, num_groups, ctx)? + { + return self.accumulate_kernel_result(result, num_groups, ctx); + } + + if self.vtable.try_accumulate_grouped( + &mut self.partials[..num_groups], + batch, + group_ids, + ctx, + )? { + return Ok(()); + } + + let input = batch.clone(); + let mut batch = batch.clone(); for _ in 0..max_iterations() { - if elements.is::() { + if batch.is::() { break; } - if let Some(result) = session + if let Some(kernel) = session .aggregate_fns() - .find_grouped_kernel(elements.encoding_id(), self.aggregate_fn.id()) - .and_then(|kernel| { - // SAFETY: we assume that elements execution is safe - let groups = unsafe { - ListViewArray::new_unchecked( - elements.clone(), - groups.offsets().clone(), - groups.sizes().clone(), - groups_validity.clone(), - ) - }; - kernel - .grouped_aggregate(&self.aggregate_fn, &groups) - .transpose() - }) - .transpose()? + .find_grouped_kernel(batch.encoding_id(), self.aggregate_fn.id()) + && let Some(result) = kernel.grouped_aggregate( + &self.aggregate_fn, + &batch, + group_ids, + num_groups, + ctx, + )? { - return self.push_result(result); + return self.accumulate_kernel_result(result, num_groups, ctx); } - // Execute one step and try again - elements = elements.execute(ctx)?; + batch = batch.execute(ctx)?; } - // Otherwise, we iterate the offsets and sizes and accumulate each group one by one. - let elements = elements.execute::(ctx)?.into_array(); - let offsets = groups.offsets(); - let sizes = groups.sizes().cast(offsets.dtype().clone())?; - let validity = groups_validity.execute_mask(offsets.len(), ctx)?; - - match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| { - let offsets = offsets.clone().execute::>(ctx)?; - let sizes = sizes.execute::>(ctx)?; - self.accumulate_list_view_typed( - &elements, - offsets.as_ref(), - sizes.as_ref(), - &validity, - ctx, - ) - }) + let columnar = batch.clone().execute::(ctx)?; + if self.vtable.accumulate_grouped( + &mut self.partials[..num_groups], + &columnar, + group_ids, + ctx, + )? { + return Ok(()); + } + + self.accumulate_fallback(&input, group_ids, ctx) } - fn accumulate_list_view_typed( + fn accumulate_partials( &mut self, - elements: &ArrayRef, - offsets: &[O], - sizes: &[O], - validity: &Mask, + partials: &ArrayRef, + group_ids: &[u32], + num_groups: usize, 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, offsets.len()); - - // `validity` is the per-group list-view validity, so it is zipped element-wise with the - // offsets and sizes (one entry per group). - for ((offset, size), valid) in offsets.iter().zip(sizes.iter()).zip(validity.iter()) { - let offset = offset.to_usize().vortex_expect("Offset value is not usize"); - let size = size.to_usize().vortex_expect("Size value is not usize"); - - if valid { - let group = elements.slice(offset..offset + size)?; - accumulator.accumulate(&group, ctx)?; - states.append_scalar(&accumulator.flush()?)?; - } else { - states.append_null() - } - } + 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() + ); - self.push_result(states.finish()) + self.validate_group_ids(group_ids, num_groups)?; + self.ensure_groups(num_groups)?; + + for (row_idx, &group_id) in group_ids.iter().enumerate() { + let partial = partials.execute_scalar(row_idx, ctx)?; + self.vtable + .combine_partials(&mut self.partials[group_id as usize], partial)?; + } + Ok(()) } - fn accumulate_fixed_size_list( + fn merge_group( &mut self, - groups: &FixedSizeListArray, - ctx: &mut ExecutionCtx, + into: u32, + other: &dyn DynGroupedAccumulator, + from: u32, ) -> VortexResult<()> { - let mut elements = groups.elements().clone(); - let groups_validity = groups.validity()?; - let session = ctx.session().clone(); - - for _ in 0..64 { - if elements.is::() { - break; - } + 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)?; + let partial = other.partial_scalar(from)?; + self.vtable + .combine_partials(&mut self.partials[into as usize], partial) + } - if let Some(result) = session - .aggregate_fns() - .find_grouped_kernel(elements.encoding_id(), self.aggregate_fn.id()) - .and_then(|kernel| { - // SAFETY: we assume that elements execution is safe - let groups = unsafe { - FixedSizeListArray::new_unchecked( - elements.clone(), - groups.list_size(), - groups_validity.clone(), - groups.len(), - ) - }; - - kernel - .grouped_aggregate_fixed_size(&self.aggregate_fn, &groups) - .transpose() - }) - .transpose()? - { - return self.push_result(result); - } + fn partial_dtype(&self) -> &DType { + &self.partial_dtype + } - // Execute one step and try again - elements = elements.execute(ctx)?; + fn partial_scalar(&self, group_id: u32) -> VortexResult { + if let Some(partial) = self.partials.get(group_id as usize) { + self.vtable.to_scalar(partial) + } else { + let partial = self.vtable.empty_partial(&self.options, &self.dtype)?; + self.vtable.to_scalar(&partial) } + } - // Otherwise, we iterate the offsets and sizes and accumulate each group one by one. - let elements = elements.execute::(ctx)?.into_array(); - let validity = groups_validity.execute_mask(groups.len(), ctx)?; - - let mut accumulator = Accumulator::try_new( - self.vtable.clone(), - self.options.clone(), - self.dtype.clone(), - )?; - let mut states = builder_with_capacity(&self.partial_dtype, groups.len()); - - let mut offset = 0; - let size = groups - .list_size() - .to_usize() - .vortex_expect("List size is not usize"); - - for valid in validity.iter() { - if valid { - let group = elements.slice(offset..offset + size)?; - accumulator.accumulate(&group, ctx)?; - states.append_scalar(&accumulator.flush()?)?; - } else { - states.append_null() - } - offset += size; + 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)?; + + 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(); - self.push_result(states.finish()) + Ok(states.finish()) } - fn push_result(&mut self, state: ArrayRef) -> VortexResult<()> { + fn finish(&mut self, num_groups: usize) -> VortexResult { + let states = self.flush_partials(num_groups)?; + let results = self.vtable.finalize(states)?; + vortex_ensure!( - state.dtype() == &self.partial_dtype, - "State DType mismatch: expected {}, got {}", - self.partial_dtype, - state.dtype() + results.dtype() == &self.return_dtype, + "Return DType mismatch: expected {}, got {}", + self.return_dtype, + results.dtype() ); - self.partials.push(state); - Ok(()) + + Ok(results) } } diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index e25c42e0845..53afa28d912 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -82,6 +82,22 @@ impl AggregateFnVTable for Count { Ok(true) } + fn try_accumulate_grouped( + &self, + states: &mut [Self::Partial], + batch: &ArrayRef, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let validity = batch.validity()?.execute_mask(batch.len(), ctx)?; + for (&group_id, valid) in group_ids.iter().zip(validity.iter()) { + if valid { + states[group_id as usize] += 1; + } + } + Ok(true) + } + fn accumulate( &self, _partial: &mut Self::Partial, @@ -114,11 +130,14 @@ mod tests { use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; + use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::fns::count::Count; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -225,6 +244,43 @@ mod tests { Ok(()) } + #[test] + fn grouped_count_dense_ids() -> VortexResult<()> { + let values = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None, Some(6)]) + .into_array(); + let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?; + acc.accumulate( + &values, + &[0, 0, 1, 1, 2, 2], + 3, + &mut LEGACY_SESSION.create_execution_ctx(), + )?; + + let actual = acc.finish(3)?; + let expected = PrimitiveArray::from_iter([1u64, 2, 1]).into_array(); + assert_arrays_eq!(&actual, &expected); + Ok(()) + } + + #[test] + fn grouped_count_accumulate_partials_and_merge_group() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let partials = PrimitiveArray::from_iter([2u64, 3, 5]).into_array(); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + + let mut left = GroupedAccumulator::try_new(Count, EmptyOptions, dtype.clone())?; + left.accumulate_partials(&partials, &[0, 1, 1], 2, &mut ctx)?; + + let mut right = GroupedAccumulator::try_new(Count, EmptyOptions, 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); + Ok(()) + } + #[test] fn count_constant_non_null() -> VortexResult<()> { let array = ConstantArray::new(42i32, 10); diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 24799570ff7..4e75a1390f3 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -6,11 +6,15 @@ mod constant; mod decimal; mod primitive; +use num_traits::AsPrimitive; +use num_traits::ToPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::AllOr; +use vortex_mask::Mask; use self::bool::accumulate_bool; use self::constant::multiply_constant; @@ -25,14 +29,19 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::arrays::BoolArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::bool::BoolArrayExt; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::MAX_PRECISION; +use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; +use crate::match_each_native_ptype; use crate::scalar::DecimalValue; use crate::scalar::Scalar; @@ -253,6 +262,30 @@ impl AggregateFnVTable for Sum { Ok(()) } + fn accumulate_grouped( + &self, + partials: &mut [Self::Partial], + batch: &Columnar, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + match batch { + Columnar::Canonical(Canonical::Primitive(p)) => { + accumulate_grouped_primitive(partials, p, group_ids, ctx)?; + Ok(true) + } + Columnar::Canonical(Canonical::Bool(b)) => { + accumulate_grouped_bool(partials, b, group_ids, ctx)?; + Ok(true) + } + // Decimal and constants still use the universal grouped fallback. + Columnar::Canonical(Canonical::Decimal(_)) | Columnar::Constant(_) => Ok(false), + Columnar::Canonical(_) => { + vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()) + } + } + } + fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } @@ -299,6 +332,146 @@ fn make_zero_state(return_dtype: &DType) -> SumState { } } +fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { + match validity.indices() { + AllOr::All => { + for idx in 0..len { + f(idx); + } + } + AllOr::None => {} + AllOr::Some(indices) => { + for &idx in indices { + f(idx); + } + } + } +} + +fn accumulate_grouped_unsigned(partials: &mut [SumPartial], group_id: u32, value: u64) { + let partial = &mut partials[group_id as usize]; + let saturated = match partial.current.as_mut() { + None => return, + Some(SumState::Unsigned(acc)) => checked_add_u64(acc, value), + Some(_) => vortex_panic!("unsigned sum state with non-unsigned input"), + }; + if saturated { + partial.current = None; + } +} + +fn accumulate_grouped_signed(partials: &mut [SumPartial], group_id: u32, value: i64) { + let partial = &mut partials[group_id as usize]; + let saturated = match partial.current.as_mut() { + None => return, + Some(SumState::Signed(acc)) => checked_add_i64(acc, value), + Some(_) => vortex_panic!("signed sum state with non-signed input"), + }; + if saturated { + partial.current = None; + } +} + +fn accumulate_grouped_float(partials: &mut [SumPartial], group_id: u32, value: f64) { + if value.is_nan() { + return; + } + + match partials[group_id as usize].current.as_mut() { + None => {} + Some(SumState::Float(acc)) => *acc += value, + Some(_) => vortex_panic!("float sum state with non-float input"), + } +} + +fn accumulate_grouped_primitive( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let validity = primitive + .as_ref() + .validity()? + .execute_mask(primitive.as_ref().len(), ctx)?; + match_each_native_ptype!(primitive.ptype(), + unsigned: |T| { + accumulate_grouped_primitive_unsigned::(partials, primitive, group_ids, &validity); + Ok(()) + }, + signed: |T| { + accumulate_grouped_primitive_signed::(partials, primitive, group_ids, &validity); + Ok(()) + }, + floating: |T| { + accumulate_grouped_primitive_float::(partials, primitive, group_ids, &validity); + Ok(()) + } + ) +} + +fn accumulate_grouped_primitive_unsigned( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], + validity: &Mask, +) where + T: NativePType + AsPrimitive, +{ + let values = primitive.as_slice::(); + for_each_valid_idx(validity, values.len(), |idx| { + accumulate_grouped_unsigned(partials, group_ids[idx], values[idx].as_()); + }); +} + +fn accumulate_grouped_primitive_signed( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], + validity: &Mask, +) where + T: NativePType + AsPrimitive, +{ + let values = primitive.as_slice::(); + for_each_valid_idx(validity, values.len(), |idx| { + accumulate_grouped_signed(partials, group_ids[idx], values[idx].as_()); + }); +} + +fn accumulate_grouped_primitive_float( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], + validity: &Mask, +) where + T: NativePType + ToPrimitive, +{ + let values = primitive.as_slice::(); + for_each_valid_idx(validity, values.len(), |idx| { + let value = values[idx].to_f64().vortex_expect("float to f64"); + accumulate_grouped_float(partials, group_ids[idx], value); + }); +} + +fn accumulate_grouped_bool( + partials: &mut [SumPartial], + 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(); + for_each_valid_idx(&validity, values.len(), |idx| { + if values.value(idx) { + accumulate_grouped_unsigned(partials, group_ids[idx], 1); + } + }); + Ok(()) +} + /// Checked add for u64, returning true if overflow occurred. #[inline(always)] fn checked_add_u64(acc: &mut u64, val: u64) -> bool { @@ -346,8 +519,6 @@ mod tests { use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -512,20 +683,26 @@ mod tests { // Grouped sum tests - fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, elem_dtype.clone())?; - acc.accumulate_list(groups, &mut LEGACY_SESSION.create_execution_ctx())?; - acc.finish() + fn run_grouped_sum( + values: &ArrayRef, + group_ids: &[u32], + num_groups: usize, + ) -> VortexResult { + let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, values.dtype().clone())?; + acc.accumulate( + values, + group_ids, + num_groups, + &mut LEGACY_SESSION.create_execution_ctx(), + )?; + acc.finish(num_groups) } #[test] - fn grouped_sum_fixed_size_list() -> VortexResult<()> { - let elements = + fn grouped_sum_dense_ids() -> VortexResult<()> { + let values = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); - let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; - - let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + let result = run_grouped_sum(&values, &[0, 0, 0, 1, 1, 1], 2)?; let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array(); assert_arrays_eq!(&result, &expected); @@ -534,13 +711,10 @@ mod tests { #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { - let elements = + let values = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); - let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; - - let elem_dtype = DType::Primitive(PType::I32, Nullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + let result = run_grouped_sum(&values, &[0, 0, 0, 1, 1, 1], 2)?; let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array(); assert_arrays_eq!(&result, &expected); @@ -548,30 +722,22 @@ mod tests { } #[test] - fn grouped_sum_with_null_group() -> VortexResult<()> { - let elements = - PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) - .into_array(); - let validity = Validity::from_iter([true, false, true]); - let groups = FixedSizeListArray::try_new(elements, 3, validity, 3)?; - - let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + fn grouped_sum_empty_group() -> VortexResult<()> { + let values = + PrimitiveArray::new(buffer![1i32, 2, 3, 7, 8, 9], Validity::NonNullable).into_array(); + let result = run_grouped_sum(&values, &[0, 0, 0, 2, 2, 2], 3)?; let expected = - PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array(); + PrimitiveArray::from_option_iter([Some(6i64), Some(0i64), Some(24i64)]).into_array(); assert_arrays_eq!(&result, &expected); Ok(()) } #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { - let elements = + let values = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); - let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; - - let elem_dtype = DType::Primitive(PType::I32, Nullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + let result = run_grouped_sum(&values, &[0, 0, 1, 1], 2)?; let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result, &expected); @@ -580,12 +746,8 @@ mod tests { #[test] fn grouped_sum_bool() -> VortexResult<()> { - let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); - let groups = - FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; - - let elem_dtype = DType::Bool(Nullability::NonNullable); - let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; + let values: BoolArray = [true, false, true, true, true, true].into_iter().collect(); + let result = run_grouped_sum(&values.into_array(), &[0, 0, 0, 1, 1, 1], 2)?; let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).into_array(); assert_arrays_eq!(&result, &expected); @@ -598,19 +760,17 @@ mod tests { let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, elem_dtype)?; - let elements1 = + let values1 = 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()?; + acc.accumulate(&values1, &[0, 0, 1, 1], 2, &mut ctx)?; + let result1 = acc.finish(2)?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result1, &expected1); - 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 values2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); + acc.accumulate(&values2, &[0, 0], 1, &mut ctx)?; + let result2 = acc.finish(1)?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); assert_arrays_eq!(&result2, &expected2); @@ -618,20 +778,13 @@ mod tests { } #[test] - fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let elements = + fn grouped_sum_out_of_order_group_ids() -> VortexResult<()> { + let values = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); - let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); - let sizes = PrimitiveArray::new(buffer![1i32, 1, 1], Validity::NonNullable).into_array(); - let validity = Validity::from_iter([true, false, true]); - let groups = ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array(); - - let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let result = run_grouped_sum(&groups, &elem_dtype)?; + let result = run_grouped_sum(&values, &[2, 0, 1], 3)?; - // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. let expected = - PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); + PrimitiveArray::from_option_iter([Some(200i64), Some(300), Some(100)]).into_array(); assert_arrays_eq!(&result, &expected); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/kernels.rs b/vortex-array/src/aggregate_fn/kernels.rs index d806b18d84d..23ad6a934e5 100644 --- a/vortex-array/src/aggregate_fn/kernels.rs +++ b/vortex-array/src/aggregate_fn/kernels.rs @@ -6,13 +6,12 @@ use std::fmt::Debug; +use vortex_buffer::Buffer; use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFnRef; -use crate::arrays::FixedSizeListArray; -use crate::arrays::ListViewArray; use crate::scalar::Scalar; /// A pluggable kernel for an aggregate function. @@ -28,36 +27,49 @@ pub trait DynAggregateKernel: 'static + Send + Sync + Debug { ) -> VortexResult>; } -/// A pluggable kernel for batch aggregation of many groups. +/// Partial grouped aggregate output produced by an encoding-specific grouped kernel. /// -/// The kernel is matched on the encoding of the _elements_ array, which is the inner array of the -/// provided `ListViewArray`. 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. +/// `group_ids` is parallel to `partials`: each row in `partials` is a partial state for the +/// corresponding dense group id. The grouped accumulator merges this batch through +/// `accumulate_partials`. +#[derive(Clone, Debug)] +pub struct GroupedAggregateKernelResult { + group_ids: Buffer, + partials: ArrayRef, +} + +impl GroupedAggregateKernelResult { + pub fn new(group_ids: Buffer, partials: ArrayRef) -> Self { + Self { + group_ids, + partials, + } + } + + pub fn group_ids(&self) -> &[u32] { + self.group_ids.as_ref() + } + + pub fn partials(&self) -> &ArrayRef { + &self.partials + } +} + +/// A pluggable kernel for batch aggregation of many groups. /// -/// Each element of the list 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. +/// The kernel is matched on the encoding of the values array. It receives the same dense group ids +/// 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. pub trait DynGroupedAggregateKernel: 'static + Send + Sync + Debug { - /// Aggregate each group in the provided `ListViewArray` and return an array of the - /// aggregate states. + /// Aggregate values into a partial-state batch keyed by dense group id. fn grouped_aggregate( &self, aggregate_fn: &AggregateFnRef, - groups: &ListViewArray, - ) -> VortexResult>; - - /// Aggregate each group in the provided `FixedSizeListArray` and return an array of the - /// aggregate states. - fn grouped_aggregate_fixed_size( - &self, - aggregate_fn: &AggregateFnRef, - groups: &FixedSizeListArray, - ) -> VortexResult> { - // TODO(ngates): we could automatically delegate to `grouped_aggregate` if SequenceArray - // was in the vortex-array crate - let _ = (aggregate_fn, groups); - Ok(None) - } + batch: &ArrayRef, + group_ids: &[u32], + num_groups: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 28b91d45166..b6c0915f2e7 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -146,6 +146,35 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { ctx: &mut ExecutionCtx, ) -> VortexResult<()>; + /// Try to accumulate a raw values batch into dense per-group states before decompression. + /// + /// `group_ids` is parallel to `batch` and contains dense ids in `0..states.len()`. Returns + /// `true` when the batch was fully handled. + fn try_accumulate_grouped( + &self, + _states: &mut [Self::Partial], + _batch: &ArrayRef, + _group_ids: &[u32], + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(false) + } + + /// Accumulate a canonical values batch into dense per-group states. + /// + /// `group_ids` is parallel to `batch` and contains dense ids in `0..states.len()`. Returns + /// `true` when the batch was fully handled. The provided default preserves universal + /// correctness through [`GroupedAccumulator`]'s fallback. + fn accumulate_grouped( + &self, + _states: &mut [Self::Partial], + _batch: &Columnar, + _group_ids: &[u32], + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(false) + } + /// Finalize an array of accumulator states into an array of aggregate results. /// /// The provides `states` array has dtype as specified by `state_dtype`, the result array From 9bd157fe5670f24592c01f7cb9fe35b9058a1df6 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 11 Jun 2026 17:42:21 -0400 Subject: [PATCH 2/9] Clarify dense grouped aggregate ids Signed-off-by: "Nicholas Gates" --- .../src/aggregate_fn/accumulator_grouped.rs | 12 ++++++++++-- vortex-array/src/aggregate_fn/fns/count/mod.rs | 10 ++++++++++ vortex-array/src/aggregate_fn/kernels.rs | 13 +++++++------ vortex-array/src/aggregate_fn/vtable.rs | 10 ++++++---- 4 files changed, 33 insertions(+), 12 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 66da32b4085..990a420eecb 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -28,8 +28,10 @@ pub type GroupedAccumulatorRef = Box; /// An accumulator used for computing aggregates over dense group ids. /// -/// Group ids are dense `u32` slots in the range `0..num_groups`. The accumulator keeps one partial -/// state per slot, so ordered and unordered grouping only differ in how the caller assigns 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, @@ -167,6 +169,9 @@ impl GroupedAccumulator { /// 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..num_groups`; ids may repeat, appear out of order, or be absent from a given batch. fn accumulate( &mut self, batch: &ArrayRef, @@ -176,6 +181,9 @@ pub trait DynGroupedAccumulator: 'static + Send { ) -> 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, diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 53afa28d912..07395211ca8 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -263,6 +263,16 @@ mod tests { Ok(()) } + #[test] + fn grouped_count_rejects_out_of_range_group_id() -> VortexResult<()> { + let values = PrimitiveArray::new(buffer![1i32, 2], Validity::NonNullable).into_array(); + let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + + assert!(acc.accumulate(&values, &[0, 2], 2, &mut ctx).is_err()); + Ok(()) + } + #[test] fn grouped_count_accumulate_partials_and_merge_group() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::Nullable); diff --git a/vortex-array/src/aggregate_fn/kernels.rs b/vortex-array/src/aggregate_fn/kernels.rs index 23ad6a934e5..91248091437 100644 --- a/vortex-array/src/aggregate_fn/kernels.rs +++ b/vortex-array/src/aggregate_fn/kernels.rs @@ -30,8 +30,9 @@ pub trait DynAggregateKernel: 'static + Send + Sync + Debug { /// Partial grouped aggregate output produced by an encoding-specific grouped kernel. /// /// `group_ids` is parallel to `partials`: each row in `partials` is a partial state for the -/// corresponding dense group id. The grouped accumulator merges this batch through -/// `accumulate_partials`. +/// corresponding dense group ordinal. The ids may repeat, omit, and reorder groups, but must be +/// valid slots in the accumulator's `0..num_groups` range. The grouped accumulator merges this +/// batch through `accumulate_partials`. #[derive(Clone, Debug)] pub struct GroupedAggregateKernelResult { group_ids: Buffer, @@ -57,13 +58,13 @@ impl GroupedAggregateKernelResult { /// A pluggable kernel for batch aggregation of many groups. /// -/// The kernel is matched on the encoding of the values array. It receives the same dense group ids -/// that the caller passed to the grouped accumulator and may aggregate directly in the encoded -/// domain. +/// The kernel is matched on the encoding of the values array. It receives 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. pub trait DynGroupedAggregateKernel: 'static + Send + Sync + Debug { - /// Aggregate values into a partial-state batch keyed by dense group id. + /// Aggregate values into a partial-state batch keyed by dense group ordinal. fn grouped_aggregate( &self, aggregate_fn: &AggregateFnRef, diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index b6c0915f2e7..e30f41f012e 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -148,8 +148,9 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Try to accumulate a raw values batch into dense per-group states before decompression. /// - /// `group_ids` is parallel to `batch` and contains dense ids in `0..states.len()`. Returns - /// `true` when the batch was fully handled. + /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in + /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch. + /// Returns `true` when the batch was fully handled. fn try_accumulate_grouped( &self, _states: &mut [Self::Partial], @@ -162,8 +163,9 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Accumulate a canonical values batch into dense per-group states. /// - /// `group_ids` is parallel to `batch` and contains dense ids in `0..states.len()`. Returns - /// `true` when the batch was fully handled. The provided default preserves universal + /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in + /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch. + /// Returns `true` when the batch was fully handled. The provided default preserves universal /// correctness through [`GroupedAccumulator`]'s fallback. fn accumulate_grouped( &self, From 50701b29f24915144a50e4fe6d88942b5d0808cc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 11 Jun 2026 18:04:41 -0400 Subject: [PATCH 3/9] Fix grouped aggregate CI checks Signed-off-by: "Nicholas Gates" --- AGENTS.md | 8 +++++++ .../src/aggregate_fn/accumulator_grouped.rs | 21 ++++++++++--------- vortex-array/src/aggregate_fn/vtable.rs | 2 +- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e5c3d0cc13b..2a1ad73df22 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -113,6 +113,14 @@ cargo +nightly fmt --all cargo clippy --all-targets --all-features ``` +Do not push Rust code changes before running the applicable lint command above. If the change adds +or edits Rustdoc on public APIs, also run the CI docs command so broken intra-doc links are caught +locally: + +```bash +RUSTDOCFLAGS="-D warnings" cargo doc --profile ci --no-deps +``` + Notes: - For `.github/` changes, follow `.github/AGENTS.md` and run diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 990a420eecb..938ee2b0dd6 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -79,11 +79,7 @@ impl GroupedAccumulator { } fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { - vortex_ensure!( - num_groups <= (u32::MAX as usize) + 1, - "num_groups {} exceeds dense u32 group id capacity", - num_groups - ); + validate_num_groups(num_groups)?; while self.partials.len() < num_groups { self.partials @@ -93,11 +89,7 @@ impl GroupedAccumulator { } fn validate_group_ids(&self, group_ids: &[u32], num_groups: usize) -> VortexResult<()> { - vortex_ensure!( - num_groups <= (u32::MAX as usize) + 1, - "num_groups {} exceeds dense u32 group id capacity", - num_groups - ); + validate_num_groups(num_groups)?; for &group_id in group_ids { vortex_ensure!( (group_id as usize) < num_groups, @@ -165,6 +157,15 @@ impl GroupedAccumulator { } } +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(()) +} + /// 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 { diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index e30f41f012e..24c2113e64a 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -166,7 +166,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch. /// Returns `true` when the batch was fully handled. The provided default preserves universal - /// correctness through [`GroupedAccumulator`]'s fallback. + /// correctness through [`crate::aggregate_fn::GroupedAccumulator`]'s fallback. fn accumulate_grouped( &self, _states: &mut [Self::Partial], From adae76e3122b5c71387eaef1de187b07694bb6e1 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 11 Jun 2026 18:05:45 -0400 Subject: [PATCH 4/9] DCO Remediation Commit for Nicholas Gates I, Nicholas Gates , hereby add my Signed-off-by to this commit: 6bdcd32b62d3852247efe48a7a3a2af7d7fd922c I, Nicholas Gates , hereby add my Signed-off-by to this commit: 9bd157fe5670f24592c01f7cb9fe35b9058a1df6 I, Nicholas Gates , hereby add my Signed-off-by to this commit: 50701b29f24915144a50e4fe6d88942b5d0808cc Signed-off-by: Nicholas Gates --- AGENTS.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 2a1ad73df22..759008d730b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -198,5 +198,8 @@ you ran and call out any checks you could not run. All commits must be signed off by the committers in this form: ```text -Signed-off-by: "COMMITTER" +Signed-off-by: COMMITTER ``` + +Do not wrap the committer name in quotes; the DCO check expects the exact unquoted name/email +pair from the commit author. From c61b7bbabb015e7518a4138d85c34623d4200435 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 12 Jun 2026 14:26:49 -0400 Subject: [PATCH 5/9] Restore dense grouped aggregate test coverage Signed-off-by: Nicholas Gates --- .../src/aggregate_fn/fns/count/mod.rs | 53 ++++++++++++++++--- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 33 ++++++++++++ 2 files changed, 78 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 1ce3588a6ef..60896476e91 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -132,6 +132,7 @@ mod tests { use crate::arrays::ChunkedArray; 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; @@ -239,25 +240,61 @@ mod tests { Ok(()) } + fn run_grouped_count( + values: &ArrayRef, + group_ids: &[u32], + num_groups: usize, + ) -> VortexResult { + let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?; + acc.accumulate( + values, + group_ids, + num_groups, + &mut LEGACY_SESSION.create_execution_ctx(), + )?; + acc.finish(num_groups) + } + #[test] fn grouped_count_dense_ids() -> VortexResult<()> { let values = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None, Some(6)]) .into_array(); - let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?; - acc.accumulate( - &values, - &[0, 0, 1, 1, 2, 2], - 3, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let actual = run_grouped_count(&values, &[0, 0, 1, 1, 2, 2], 3)?; - let actual = acc.finish(3)?; let expected = PrimitiveArray::from_iter([1u64, 2, 1]).into_array(); assert_arrays_eq!(&actual, &expected); Ok(()) } + #[test] + fn grouped_count_omitted_group() -> VortexResult<()> { + let values = + PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); + let actual = run_grouped_count(&values, &[0, 0, 1, 2, 2, 2], 4)?; + + let expected = PrimitiveArray::from_iter([2u64, 1, 3, 0]).into_array(); + assert_arrays_eq!(&actual, &expected); + Ok(()) + } + + #[test] + fn grouped_count_varbinview_with_nulls() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_str([ + Some("a"), + None, + Some("bbb"), + None, + Some("cc"), + ]) + .into_array(); + let actual = run_grouped_count(&values, &[0, 0, 1, 1, 2], 3)?; + + let expected = PrimitiveArray::from_iter([1u64, 1, 1]).into_array(); + assert_arrays_eq!(&actual, &expected); + Ok(()) + } + #[test] fn grouped_count_rejects_out_of_range_group_id() -> VortexResult<()> { let values = PrimitiveArray::new(buffer![1i32, 2], Validity::NonNullable).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 e26c07c1f1a..8626e810399 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -641,6 +641,39 @@ mod tests { Ok(()) } + #[test] + fn grouped_sum_overflow_group_is_null() -> VortexResult<()> { + let values = + PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); + let result = run_grouped_sum(&values, &[0, 0, 1, 1], 2)?; + + let expected = PrimitiveArray::from_option_iter([None, Some(5i64)]).into_array(); + assert_arrays_eq!(&result, &expected); + Ok(()) + } + + #[test] + fn grouped_sum_float_nan_and_inf() -> VortexResult<()> { + let values = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, f64::INFINITY, f64::NEG_INFINITY, 4.0], + Validity::NonNullable, + ) + .into_array(); + let actual = run_grouped_sum(&values, &[0, 0, 0, 1, 1, 1], 2)?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + + let g0 = actual.execute_scalar(0, &mut ctx)?; + assert_eq!(g0.as_primitive().typed_value::(), Some(3.0)); + + let g1 = actual.execute_scalar(1, &mut ctx)?; + let g1_value = g1 + .as_primitive() + .typed_value::() + .vortex_expect("group sum should be non-null"); + assert!(g1_value.is_nan()); + Ok(()) + } + // Chunked array tests #[test] From 2ef64b2397d508460bda2f6d27a22b4835f20a3a Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 12 Jun 2026 16:46:40 -0400 Subject: [PATCH 6/9] Optimize dense grouped aggregates Signed-off-by: Nicholas Gates --- .../src/aggregate_fn/accumulator_grouped.rs | 14 ++ .../src/aggregate_fn/fns/count/mod.rs | 12 ++ .../src/aggregate_fn/fns/sum/grouped.rs | 142 +++++++++++++++++- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 89 +++++++++++ vortex-array/src/aggregate_fn/vtable.rs | 11 ++ 5 files changed, 265 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index bf771609ac8..7a614ceed63 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -385,6 +385,20 @@ impl DynGroupedAccumulator for GroupedAccumulator { ); 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)?)?; diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 60896476e91..e53a378b5a9 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -8,9 +8,11 @@ use vortex_error::VortexResult; 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::EmptyOptions; +use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -64,6 +66,16 @@ impl AggregateFnVTable for Count { Ok(Scalar::primitive(*partial, Nullability::NonNullable)) } + fn partials_to_array( + &self, + partials: &[Self::Partial], + _partial_dtype: &DType, + ) -> VortexResult> { + Ok(Some( + PrimitiveArray::from_iter(partials.iter().copied()).into_array(), + )) + } + fn reset(&self, partial: &mut Self::Partial) { *partial = 0; } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 432c48b4cf4..81304f1eb9f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -13,6 +13,9 @@ use super::SumPartial; use super::SumState; use super::checked_add_i64; use super::checked_add_u64; +use super::primitive::sum_float_all; +use super::primitive::sum_signed_all; +use super::primitive::sum_unsigned_all; use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; @@ -20,6 +23,8 @@ use crate::arrays::bool::BoolArrayExt; use crate::dtype::NativePType; use crate::match_each_native_ptype; +const MIN_AVG_RUN_LENGTH_FOR_GROUPED_SUM_RUNS: usize = 4; + fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { match validity.indices() { AllOr::All => { @@ -36,6 +41,41 @@ fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { } } +fn should_accumulate_group_runs(group_ids: &[u32]) -> bool { + let Some((&first, rest)) = group_ids.split_first() else { + return false; + }; + + let mut run_count = 1usize; + let mut group_id = first; + for &next_group_id in rest { + if next_group_id != group_id { + run_count += 1; + group_id = next_group_id; + } + } + + run_count * MIN_AVG_RUN_LENGTH_FOR_GROUPED_SUM_RUNS <= group_ids.len() +} + +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 accumulate_grouped_unsigned(partials: &mut [SumPartial], group_id: u32, value: u64) { let partial = &mut partials[group_id as usize]; let saturated = match partial.current.as_mut() { @@ -48,6 +88,21 @@ fn accumulate_grouped_unsigned(partials: &mut [SumPartial], group_id: u32, value } } +fn accumulate_grouped_unsigned_run(partials: &mut [SumPartial], group_id: u32, values: &[T]) +where + T: NativePType + AsPrimitive, +{ + let partial = &mut partials[group_id as usize]; + let saturated = match partial.current.as_mut() { + None => return, + Some(SumState::Unsigned(acc)) => sum_unsigned_all(acc, values), + Some(_) => vortex_panic!("unsigned sum state with non-unsigned input"), + }; + if saturated { + partial.current = None; + } +} + fn accumulate_grouped_signed(partials: &mut [SumPartial], group_id: u32, value: i64) { let partial = &mut partials[group_id as usize]; let saturated = match partial.current.as_mut() { @@ -60,6 +115,21 @@ fn accumulate_grouped_signed(partials: &mut [SumPartial], group_id: u32, value: } } +fn accumulate_grouped_signed_run(partials: &mut [SumPartial], group_id: u32, values: &[T]) +where + T: NativePType + AsPrimitive, +{ + let partial = &mut partials[group_id as usize]; + let saturated = match partial.current.as_mut() { + None => return, + Some(SumState::Signed(acc)) => sum_signed_all(acc, values), + Some(_) => vortex_panic!("signed sum state with non-signed input"), + }; + if saturated { + partial.current = None; + } +} + fn accumulate_grouped_float(partials: &mut [SumPartial], group_id: u32, value: f64) { if value.is_nan() { return; @@ -72,6 +142,18 @@ fn accumulate_grouped_float(partials: &mut [SumPartial], group_id: u32, value: f } } +fn accumulate_grouped_float_run( + partials: &mut [SumPartial], + group_id: u32, + values: &[T], +) { + match partials[group_id as usize].current.as_mut() { + None => {} + Some(SumState::Float(acc)) => sum_float_all(acc, values), + Some(_) => vortex_panic!("float sum state with non-float input"), + } +} + pub(super) fn accumulate_grouped_primitive( partials: &mut [SumPartial], primitive: &PrimitiveArray, @@ -82,17 +164,32 @@ pub(super) fn accumulate_grouped_primitive( .as_ref() .validity()? .execute_mask(primitive.as_ref().len(), ctx)?; + let use_runs = + matches!(validity.slices(), AllOr::All) && should_accumulate_group_runs(group_ids); + match_each_native_ptype!(primitive.ptype(), unsigned: |T| { - accumulate_grouped_primitive_unsigned::(partials, primitive, group_ids, &validity); + if use_runs { + accumulate_grouped_primitive_unsigned_runs::(partials, primitive, group_ids); + } else { + accumulate_grouped_primitive_unsigned::(partials, primitive, group_ids, &validity); + } Ok(()) }, signed: |T| { - accumulate_grouped_primitive_signed::(partials, primitive, group_ids, &validity); + if use_runs { + accumulate_grouped_primitive_signed_runs::(partials, primitive, group_ids); + } else { + accumulate_grouped_primitive_signed::(partials, primitive, group_ids, &validity); + } Ok(()) }, floating: |T| { - accumulate_grouped_primitive_float::(partials, primitive, group_ids, &validity); + if use_runs { + accumulate_grouped_primitive_float_runs::(partials, primitive, group_ids); + } else { + accumulate_grouped_primitive_float::(partials, primitive, group_ids, &validity); + } Ok(()) } ) @@ -112,6 +209,19 @@ fn accumulate_grouped_primitive_unsigned( }); } +fn accumulate_grouped_primitive_unsigned_runs( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], +) where + T: NativePType + AsPrimitive, +{ + let values = primitive.as_slice::(); + for_each_group_run(group_ids, |group_id, start, end| { + accumulate_grouped_unsigned_run(partials, group_id, &values[start..end]); + }); +} + fn accumulate_grouped_primitive_signed( partials: &mut [SumPartial], primitive: &PrimitiveArray, @@ -126,6 +236,19 @@ fn accumulate_grouped_primitive_signed( }); } +fn accumulate_grouped_primitive_signed_runs( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], +) where + T: NativePType + AsPrimitive, +{ + let values = primitive.as_slice::(); + for_each_group_run(group_ids, |group_id, start, end| { + accumulate_grouped_signed_run(partials, group_id, &values[start..end]); + }); +} + fn accumulate_grouped_primitive_float( partials: &mut [SumPartial], primitive: &PrimitiveArray, @@ -141,6 +264,19 @@ fn accumulate_grouped_primitive_float( }); } +fn accumulate_grouped_primitive_float_runs( + partials: &mut [SumPartial], + primitive: &PrimitiveArray, + group_ids: &[u32], +) where + T: NativePType, +{ + let values = primitive.as_slice::(); + for_each_group_run(group_ids, |group_id, start, end| { + accumulate_grouped_float_run(partials, group_id, &values[start..end]); + }); +} + pub(super) fn accumulate_grouped_bool( partials: &mut [SumPartial], bools: &BoolArray, diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 8626e810399..eff487d55e8 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -7,6 +7,7 @@ mod decimal; mod grouped; mod primitive; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -21,11 +22,13 @@ use crate::ArrayRef; use crate::Canonical; use crate::Columnar; use crate::ExecutionCtx; +use crate::IntoArray; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::MAX_PRECISION; @@ -36,6 +39,7 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::validity::Validity; /// Return the sum of an array. /// @@ -201,6 +205,29 @@ impl AggregateFnVTable for Sum { }) } + fn partials_to_array( + &self, + partials: &[Self::Partial], + partial_dtype: &DType, + ) -> VortexResult> { + Ok(match partial_dtype { + DType::Primitive(PType::U64, _) => Some(sum_primitive_partials_to_array( + partials, + unsigned_sum_state_value, + )), + DType::Primitive(PType::I64, _) => Some(sum_primitive_partials_to_array( + partials, + signed_sum_state_value, + )), + DType::Primitive(PType::F64, _) => Some(sum_primitive_partials_to_array( + partials, + float_sum_state_value, + )), + DType::Decimal(..) => None, + _ => vortex_bail!("Unsupported sum partial dtype: {}", partial_dtype), + }) + } + fn reset(&self, partial: &mut Self::Partial) { partial.current = Some(make_zero_state(&partial.return_dtype)); } @@ -309,6 +336,54 @@ pub enum SumState { }, } +fn sum_primitive_partials_to_array( + partials: &[SumPartial], + value_from_state: fn(&SumState) -> T, +) -> ArrayRef +where + T: crate::dtype::NativePType, +{ + if partials.iter().all(|partial| partial.current.is_some()) { + let values = Buffer::from_iter(partials.iter().map(|partial| { + value_from_state( + partial + .current + .as_ref() + .vortex_expect("checked non-null partial"), + ) + })); + return PrimitiveArray::new(values, Validity::AllValid).into_array(); + } + + PrimitiveArray::from_option_iter( + partials + .iter() + .map(|partial| partial.current.as_ref().map(value_from_state)), + ) + .into_array() +} + +fn unsigned_sum_state_value(state: &SumState) -> u64 { + match state { + SumState::Unsigned(v) => *v, + _ => vortex_panic!("unsigned sum state with non-unsigned partial dtype"), + } +} + +fn signed_sum_state_value(state: &SumState) -> i64 { + match state { + SumState::Signed(v) => *v, + _ => vortex_panic!("signed sum state with non-signed partial dtype"), + } +} + +fn float_sum_state_value(state: &SumState) -> f64 { + match state { + SumState::Float(v) => *v, + _ => vortex_panic!("float sum state with non-float partial dtype"), + } +} + fn make_zero_state(return_dtype: &DType) -> SumState { match return_dtype { DType::Primitive(ptype, _) => match ptype { @@ -641,6 +716,20 @@ mod tests { Ok(()) } + #[test] + fn grouped_sum_contiguous_group_runs() -> VortexResult<()> { + let values = PrimitiveArray::new( + buffer![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + Validity::NonNullable, + ) + .into_array(); + let result = run_grouped_sum(&values, &[0, 0, 0, 0, 1, 1, 1, 1], 2)?; + + let expected = PrimitiveArray::from_option_iter([Some(10.0f64), Some(26.0)]).into_array(); + assert_arrays_eq!(&result, &expected); + Ok(()) + } + #[test] fn grouped_sum_overflow_group_is_null() -> VortexResult<()> { let values = diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 24c2113e64a..09eab6c5a9c 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -115,6 +115,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); From 4ba3988ee06ca430758c5c9715e4d71d3a17897e Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Mon, 29 Jun 2026 12:28:35 +0100 Subject: [PATCH 7/9] group ids as array ref, multi encoding kernel lookup (#8550) --- vortex-array/benches/aggregate_grouped.rs | 27 ++- .../src/aggregate_fn/accumulator_grouped.rs | 191 +++++++++++------- .../src/aggregate_fn/fns/count/grouped.rs | 38 +++- .../src/aggregate_fn/fns/count/mod.rs | 38 ++-- .../src/aggregate_fn/fns/sum/grouped.rs | 41 ++++ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 36 +--- vortex-array/src/aggregate_fn/kernels.rs | 120 ++++++++--- vortex-array/src/aggregate_fn/session.rs | 191 ++++++++++++++---- vortex-array/src/aggregate_fn/vtable.rs | 31 --- 9 files changed, 473 insertions(+), 240 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 2d46a5cce8a..e99e4619143 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -15,6 +15,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::EmptyOptions; +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; @@ -45,24 +46,22 @@ fn total_element_count(group_sizes: &[usize]) -> usize { struct DenseGroupedInput { values: ArrayRef, - group_ids: Vec, - num_groups: usize, + group_ids: GroupIds, } fn dense_grouped_input(values: ArrayRef, group_sizes: &[usize]) -> DenseGroupedInput { assert_eq!(values.len(), total_element_count(group_sizes)); - let group_ids = group_sizes - .iter() - .enumerate() - .flat_map(|(group_id, &size)| std::iter::repeat_n(group_id as u32, size)) - .collect(); + 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(); - DenseGroupedInput { - values, - group_ids, - num_groups: group_sizes.len(), - } + DenseGroupedInput { values, group_ids } } fn i32_nullable_all_valid_input() -> DenseGroupedInput { @@ -142,14 +141,14 @@ where { let mut acc = GroupedAccumulator::try_new(vtable, EmptyOptions, input.values.dtype().clone()).unwrap(); + let num_groups = input.group_ids.num_groups(); acc.accumulate( &input.values, &input.group_ids, - input.num_groups, &mut LEGACY_SESSION.create_execution_ctx(), ) .unwrap(); - divan::black_box(acc.finish(input.num_groups).unwrap()) + divan::black_box(acc.finish(num_groups).unwrap()) } #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 7a614ceed63..46064e3b000 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -7,7 +7,6 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::ArrayRef; -use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::Accumulator; @@ -15,18 +14,92 @@ use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::kernels::GroupedAggregateKernelResult; use crate::aggregate_fn::session::AggregateFnSessionExt; +use crate::array::ArrayId; +use crate::arrays::PrimitiveArray; use crate::builders::builder_with_capacity; use crate::columnar::AnyColumnar; use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; use crate::executor::max_iterations; use crate::scalar::Scalar; +use crate::validity::Validity; /// Reference-counted type-erased grouped accumulator. pub type GroupedAccumulatorRef = Box; -/// An accumulator used for computing aggregates over dense group ids. +/// Encoded group ids parallel to a grouped aggregate input batch. +/// +/// 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, +} + +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 }) + } + + /// Create group ids from a materialized buffer. + pub fn from_buffer(ids: Buffer, num_groups: usize) -> VortexResult { + Self::new( + PrimitiveArray::new(ids, Validity::NonNullable).into_array(), + num_groups, + ) + } + + /// 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. + pub fn validated_ids(&self, ctx: &mut ExecutionCtx) -> VortexResult> { + let ids = self.ids.clone().execute::>(ctx)?; + validate_group_ids(ids.as_ref(), self.num_groups)?; + Ok(ids) + } +} + +/// 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 @@ -88,54 +161,25 @@ impl GroupedAccumulator { Ok(()) } - fn validate_group_ids(&self, group_ids: &[u32], num_groups: usize) -> VortexResult<()> { - validate_num_groups(num_groups)?; - for &group_id in group_ids { - vortex_ensure!( - (group_id as usize) < num_groups, - "Group id {} out of range for {} groups", - group_id, - num_groups - ); - } - Ok(()) - } - - fn accumulate_kernel_result( - &mut self, - result: GroupedAggregateKernelResult, - num_groups: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - self.accumulate_partials(result.partials(), result.group_ids(), num_groups, ctx) - } - fn try_accumulate_kernel( &mut self, batch: &ArrayRef, - group_ids: &[u32], - num_groups: usize, + group_ids: &GroupIds, ctx: &mut ExecutionCtx, ) -> VortexResult { let session = ctx.session().clone(); - if let Some(kernel) = session - .aggregate_fns() - .find_grouped_encoding_kernel(batch.encoding_id(), self.aggregate_fn.id()) - && let Some(result) = - kernel.grouped_aggregate(&self.aggregate_fn, batch, group_ids, num_groups, ctx)? - { - self.accumulate_kernel_result(result, num_groups, ctx)?; - return Ok(true); - } - - if let Some(kernel) = session - .aggregate_fns() - .find_grouped_kernel(self.aggregate_fn.id()) - && let Some(result) = - kernel.grouped_aggregate(&self.aggregate_fn, batch, group_ids, num_groups, ctx)? - { - self.accumulate_kernel_result(result, num_groups, ctx)?; + 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, + &mut self.partials, + ctx, + )? { return Ok(true); } @@ -198,18 +242,31 @@ fn validate_num_groups(num_groups: usize) -> VortexResult<()> { Ok(()) } +fn validate_group_ids(group_ids: &[u32], num_groups: usize) -> VortexResult<()> { + validate_num_groups(num_groups)?; + for &group_id in group_ids { + vortex_ensure!( + (group_id as usize) < num_groups, + "Group id {} out of range for {} groups", + group_id, + num_groups + ); + } + Ok(()) +} + /// 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..num_groups`; ids may repeat, appear out of order, or be absent from a given batch. + /// `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: &[u32], - num_groups: usize, + group_ids: &GroupIds, ctx: &mut ExecutionCtx, ) -> VortexResult<()>; @@ -220,8 +277,7 @@ pub trait DynGroupedAccumulator: 'static + Send { fn accumulate_partials( &mut self, partials: &ArrayRef, - group_ids: &[u32], - num_groups: usize, + group_ids: &GroupIds, ctx: &mut ExecutionCtx, ) -> VortexResult<()>; @@ -254,10 +310,10 @@ impl DynGroupedAccumulator for GroupedAccumulator { fn accumulate( &mut self, batch: &ArrayRef, - group_ids: &[u32], - num_groups: usize, + 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 {}", @@ -271,56 +327,43 @@ impl DynGroupedAccumulator for GroupedAccumulator { group_ids.len() ); - self.validate_group_ids(group_ids, num_groups)?; self.ensure_groups(num_groups)?; - if self.try_accumulate_kernel(batch, group_ids, num_groups, ctx)? { - return Ok(()); - } - - if self.vtable.try_accumulate_grouped( - &mut self.partials[..num_groups], - batch, - group_ids, - ctx, - )? { + 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 self.try_accumulate_kernel(&batch, group_ids, num_groups, ctx)? { + if !tried_current && self.try_accumulate_kernel(&batch, group_ids, ctx)? { return Ok(()); } batch = batch.execute(ctx)?; + tried_current = false; } - let columnar = batch.clone().execute::(ctx)?; - if self.vtable.accumulate_grouped( - &mut self.partials[..num_groups], - &columnar, - group_ids, - ctx, - )? { + if !tried_current && self.try_accumulate_kernel(&batch, group_ids, ctx)? { return Ok(()); } - self.accumulate_fallback(&input, group_ids, ctx) + 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: &[u32], - num_groups: usize, + 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 {}", @@ -334,7 +377,7 @@ impl DynGroupedAccumulator for GroupedAccumulator { group_ids.len() ); - self.validate_group_ids(group_ids, num_groups)?; + let group_ids = group_ids.validated_ids(ctx)?; self.ensure_groups(num_groups)?; for (row_idx, &group_id) in group_ids.iter().enumerate() { diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 03e2b1b49ae..68ea4e05d26 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -3,20 +3,36 @@ use vortex_error::VortexResult; +use super::Count; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::GroupIds; +use crate::aggregate_fn::kernels::GroupedAggregateKernel; +use crate::aggregate_fn::kernels::GroupedAggregateKernelAdapter; -pub(super) fn try_accumulate_grouped( - states: &mut [u64], - batch: &ArrayRef, - group_ids: &[u32], - ctx: &mut ExecutionCtx, -) -> VortexResult { - let validity = batch.validity()?.execute_mask(batch.len(), ctx)?; - for (&group_id, valid) in group_ids.iter().zip(validity.iter()) { - if valid { - states[group_id as usize] += 1; +pub(crate) static COUNT_GROUPED_KERNEL: GroupedAggregateKernelAdapter = + GroupedAggregateKernelAdapter::new(CountGroupedKernel); + +#[derive(Debug)] +pub(crate) struct CountGroupedKernel; + +impl GroupedAggregateKernel for CountGroupedKernel { + fn grouped_accumulate( + &self, + _options: &EmptyOptions, + states: &mut [u64], + batch: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let group_ids = group_ids.validated_ids(ctx)?; + let validity = batch.validity()?.execute_mask(batch.len(), ctx)?; + for (&group_id, valid) in group_ids.iter().zip(validity.iter()) { + if valid { + states[group_id as usize] += 1; + } } + Ok(true) } - Ok(true) } diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index e53a378b5a9..c6a9c27d52f 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod grouped; +pub(crate) use grouped::COUNT_GROUPED_KERNEL; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -95,16 +96,6 @@ impl AggregateFnVTable for Count { Ok(true) } - fn try_accumulate_grouped( - &self, - states: &mut [Self::Partial], - batch: &ArrayRef, - group_ids: &[u32], - ctx: &mut ExecutionCtx, - ) -> VortexResult { - grouped::try_accumulate_grouped(states, batch, group_ids, ctx) - } - fn accumulate( &self, _partial: &mut Self::Partial, @@ -139,6 +130,7 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::GroupIds; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::fns::count::Count; use crate::arrays::ChunkedArray; @@ -258,10 +250,10 @@ mod tests { num_groups: usize, ) -> VortexResult { let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?; + let group_ids = GroupIds::from_iter(group_ids.iter().copied(), num_groups)?; acc.accumulate( values, - group_ids, - num_groups, + &group_ids, &mut LEGACY_SESSION.create_execution_ctx(), )?; acc.finish(num_groups) @@ -307,13 +299,30 @@ mod tests { Ok(()) } + #[test] + fn grouped_count_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 = LEGACY_SESSION.create_execution_ctx(); + let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, 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); + Ok(()) + } + #[test] fn grouped_count_rejects_out_of_range_group_id() -> VortexResult<()> { let values = PrimitiveArray::new(buffer![1i32, 2], Validity::NonNullable).into_array(); let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?; let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let group_ids = GroupIds::from_iter([0u32, 2], 2)?; - assert!(acc.accumulate(&values, &[0, 2], 2, &mut ctx).is_err()); + assert!(acc.accumulate(&values, &group_ids, &mut ctx).is_err()); Ok(()) } @@ -324,7 +333,8 @@ mod tests { let mut ctx = LEGACY_SESSION.create_execution_ctx(); let mut left = GroupedAccumulator::try_new(Count, EmptyOptions, dtype.clone())?; - left.accumulate_partials(&partials, &[0, 1, 1], 2, &mut ctx)?; + let group_ids = GroupIds::from_iter([0u32, 1, 1], 2)?; + left.accumulate_partials(&partials, &group_ids, &mut ctx)?; let mut right = GroupedAccumulator::try_new(Count, EmptyOptions, dtype)?; right.merge_group(0, &left, 1)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 81304f1eb9f..e7a73059fc3 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; +use super::Sum; use super::SumPartial; use super::SumState; use super::checked_add_i64; @@ -16,8 +17,15 @@ use super::checked_add_u64; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; +use crate::ArrayRef; use crate::ExecutionCtx; +use crate::aggregate_fn::EmptyOptions; +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::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::dtype::NativePType; @@ -25,6 +33,39 @@ use crate::match_each_native_ptype; const MIN_AVG_RUN_LENGTH_FOR_GROUPED_SUM_RUNS: usize = 4; +pub(crate) static SUM_GROUPED_KERNEL: GroupedAggregateKernelAdapter = + GroupedAggregateKernelAdapter::new(SumGroupedKernel); + +#[derive(Debug)] +pub(crate) struct SumGroupedKernel; + +impl GroupedAggregateKernel for SumGroupedKernel { + fn grouped_accumulate( + &self, + _options: &EmptyOptions, + partials: &mut [SumPartial], + batch: &ArrayRef, + group_ids: &GroupIds, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if let Some(primitive) = batch.as_opt::() { + let group_ids = group_ids.validated_ids(ctx)?; + let primitive = primitive.into_owned(); + accumulate_grouped_primitive(partials, &primitive, group_ids.as_ref(), ctx)?; + return Ok(true); + } + + if let Some(bools) = batch.as_opt::() { + let group_ids = group_ids.validated_ids(ctx)?; + let bools = bools.into_owned(); + accumulate_grouped_bool(partials, &bools, group_ids.as_ref(), ctx)?; + return Ok(true); + } + + Ok(false) + } +} + fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { match validity.indices() { AllOr::All => { diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index eff487d55e8..207b8140922 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -7,6 +7,7 @@ mod decimal; mod grouped; mod primitive; +pub(crate) use grouped::SUM_GROUPED_KERNEL; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -281,30 +282,6 @@ impl AggregateFnVTable for Sum { Ok(()) } - fn accumulate_grouped( - &self, - partials: &mut [Self::Partial], - batch: &Columnar, - group_ids: &[u32], - ctx: &mut ExecutionCtx, - ) -> VortexResult { - match batch { - Columnar::Canonical(Canonical::Primitive(p)) => { - grouped::accumulate_grouped_primitive(partials, p, group_ids, ctx)?; - Ok(true) - } - Columnar::Canonical(Canonical::Bool(b)) => { - grouped::accumulate_grouped_bool(partials, b, group_ids, ctx)?; - Ok(true) - } - // Decimal and constants still use the universal grouped fallback. - Columnar::Canonical(Canonical::Decimal(_)) | Columnar::Constant(_) => Ok(false), - Columnar::Canonical(_) => { - vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()) - } - } - } - fn finalize(&self, partials: ArrayRef) -> VortexResult { Ok(partials) } @@ -439,6 +416,7 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::GroupIds; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; @@ -616,10 +594,10 @@ mod tests { num_groups: usize, ) -> VortexResult { let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, values.dtype().clone())?; + let group_ids = GroupIds::from_iter(group_ids.iter().copied(), num_groups)?; acc.accumulate( values, - group_ids, - num_groups, + &group_ids, &mut LEGACY_SESSION.create_execution_ctx(), )?; acc.finish(num_groups) @@ -689,14 +667,16 @@ mod tests { let values1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - acc.accumulate(&values1, &[0, 0, 1, 1], 2, &mut ctx)?; + let group_ids1 = GroupIds::from_iter([0u32, 0, 1, 1], 2)?; + acc.accumulate(&values1, &group_ids1, &mut ctx)?; let result1 = acc.finish(2)?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); assert_arrays_eq!(&result1, &expected1); let values2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); - acc.accumulate(&values2, &[0, 0], 1, &mut ctx)?; + let group_ids2 = GroupIds::from_iter([0u32, 0], 1)?; + acc.accumulate(&values2, &group_ids2, &mut ctx)?; let result2 = acc.finish(1)?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); diff --git a/vortex-array/src/aggregate_fn/kernels.rs b/vortex-array/src/aggregate_fn/kernels.rs index e0b1d42e41e..51d47c33a2e 100644 --- a/vortex-array/src/aggregate_fn/kernels.rs +++ b/vortex-array/src/aggregate_fn/kernels.rs @@ -4,14 +4,19 @@ //! 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_buffer::Buffer; 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::AggregateFnVTable; +use crate::aggregate_fn::GroupIds; use crate::scalar::Scalar; /// A pluggable kernel for an aggregate function. @@ -27,53 +32,110 @@ pub trait DynAggregateKernel: 'static + Send + Sync + Debug { ) -> VortexResult>; } -/// Partial grouped aggregate output produced by an encoding-specific grouped kernel. +/// A typed grouped aggregate kernel. /// -/// `group_ids` is parallel to `partials`: each row in `partials` is a partial state for the -/// corresponding dense group ordinal. The ids may repeat, omit, and reorder groups, but must be -/// valid slots in the accumulator's `0..num_groups` range. The grouped accumulator merges this -/// batch through `accumulate_partials`. -#[derive(Clone, Debug)] -pub struct GroupedAggregateKernelResult { - group_ids: Buffer, - partials: ArrayRef, +/// 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 { + /// Accumulate `batch` into `states` according to `group_ids`. + fn grouped_accumulate( + &self, + options: &V::Options, + states: &mut [V::Partial], + 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 GroupedAggregateKernelResult { - pub fn new(group_ids: Buffer, partials: ArrayRef) -> Self { +impl GroupedAggregateKernelAdapter { + /// Create a new adapter around `kernel`. + pub const fn new(kernel: K) -> Self { Self { - group_ids, - partials, + kernel, + _phantom: PhantomData, } } +} - pub fn group_ids(&self) -> &[u32] { - self.group_ids.as_ref() - } - - pub fn partials(&self) -> &ArrayRef { - &self.partials +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 grouped kernel can be registered for an aggregate function regardless of input encoding, or -/// for a specific aggregate function and array encoding. Encoding-specific kernels are matched on -/// the values array, not on a pre-grouped list wrapper. +/// 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. /// /// 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 values into a partial-state batch keyed by dense group ordinal. - fn grouped_aggregate( + /// Accumulate values into type-erased partial state. + fn grouped_accumulate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + group_ids: &GroupIds, + states: &mut dyn Any, + ctx: &mut ExecutionCtx, + ) -> VortexResult; +} + +impl DynGroupedAggregateKernel for GroupedAggregateKernelAdapter +where + V: AggregateFnVTable, + K: GroupedAggregateKernel, +{ + fn grouped_accumulate( &self, aggregate_fn: &AggregateFnRef, batch: &ArrayRef, - group_ids: &[u32], - num_groups: usize, + group_ids: &GroupIds, + states: &mut dyn Any, ctx: &mut ExecutionCtx, - ) -> VortexResult>; + ) -> VortexResult { + let Some(options) = aggregate_fn.as_opt::() else { + return Ok(false); + }; + + let Some(states) = states.downcast_mut::>() else { + vortex_bail!( + "Grouped aggregate kernel for {} received incompatible partial state", + aggregate_fn.id() + ); + }; + + vortex_ensure!( + states.len() >= group_ids.num_groups(), + "Grouped aggregate kernel for {} received {} partial states for {} groups", + aggregate_fn.id(), + states.len(), + group_ids.num_groups() + ); + + self.kernel.grouped_accumulate( + options, + &mut states[..group_ids.num_groups()], + batch, + group_ids, + ctx, + ) + } } diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 78b139bf36f..14a5ccb261d 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -18,6 +18,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::first::First; use crate::aggregate_fn::fns::is_constant::IsConstant; use crate::aggregate_fn::fns::is_sorted::IsSorted; @@ -27,6 +29,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::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; @@ -51,9 +54,7 @@ pub struct AggregateFnSession { registry: ArcSwapMap, kernels: ArcSwapMap, - grouped_kernels: ArcSwapMap, - grouped_encoding_kernels: - ArcSwapMap, + grouped_kernels: ArcSwapMap, } impl SessionVar for AggregateFnSession { @@ -67,7 +68,27 @@ 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, + } + } +} impl Default for AggregateFnSession { fn default() -> Self { @@ -75,7 +96,6 @@ impl Default for AggregateFnSession { registry: ArcSwapMap::default(), kernels: ArcSwapMap::default(), grouped_kernels: ArcSwapMap::default(), - grouped_encoding_kernels: ArcSwapMap::default(), }; // Register the built-in aggregate functions @@ -103,6 +123,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); + this.register_grouped_kernel(Count.id(), None, None, &COUNT_GROUPED_KERNEL); + this.register_grouped_kernel(Sum.id(), None, None, &SUM_GROUPED_KERNEL); this } @@ -156,54 +178,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, + ) } } @@ -215,3 +245,86 @@ pub trait AggregateFnSessionExt: SessionExt { } } impl AggregateFnSessionExt for S {} + +#[cfg(test)] +mod tests { + use std::any::Any; + + use vortex_error::VortexResult; + + 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(); + let aggregate_id = AggregateFnId::new("test.grouped_lookup"); + 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 09eab6c5a9c..ab9edae5862 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -157,37 +157,6 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { ctx: &mut ExecutionCtx, ) -> VortexResult<()>; - /// Try to accumulate a raw values batch into dense per-group states before decompression. - /// - /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in - /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch. - /// Returns `true` when the batch was fully handled. - fn try_accumulate_grouped( - &self, - _states: &mut [Self::Partial], - _batch: &ArrayRef, - _group_ids: &[u32], - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(false) - } - - /// Accumulate a canonical values batch into dense per-group states. - /// - /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in - /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch. - /// Returns `true` when the batch was fully handled. The provided default preserves universal - /// correctness through [`crate::aggregate_fn::GroupedAccumulator`]'s fallback. - fn accumulate_grouped( - &self, - _states: &mut [Self::Partial], - _batch: &Columnar, - _group_ids: &[u32], - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(false) - } - /// Finalize an array of accumulator states into an array of aggregate results. /// /// The provides `states` array has dtype as specified by `state_dtype`, the result array From 35606668b0387e55b3b2e297da84e1fd06419d7c Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 7 Aug 2026 19:15:54 -0400 Subject: [PATCH 8/9] Optimize grouped aggregate hot paths Signed-off-by: Nicholas Gates --- vortex-array/benches/aggregate_grouped.rs | 33 ++++++ .../src/aggregate_fn/accumulator_grouped.rs | 61 ++++++++--- .../src/aggregate_fn/fns/count/grouped.rs | 7 +- vortex-array/src/aggregate_fn/fns/mean/mod.rs | 9 +- .../src/aggregate_fn/fns/sum/grouped.rs | 102 ++++++++++-------- 5 files changed, 148 insertions(+), 64 deletions(-) diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 0cd36ffa978..985339feda3 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -37,6 +37,8 @@ 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 SHUFFLED_ELEMENT_COUNT: usize = 1 << 16; +const SHUFFLED_GROUP_COUNT: usize = 1 << 12; fn random_group_sizes() -> Vec { let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); @@ -140,6 +142,21 @@ fn varbinview_input() -> DenseGroupedInput { ) } +fn i32_shuffled_input() -> DenseGroupedInput { + let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); + let values: Buffer = (0..SHUFFLED_ELEMENT_COUNT) + .map(|_| rng.random_range(-512..512)) + .collect(); + let group_ids: Buffer = (0..SHUFFLED_ELEMENT_COUNT) + .map(|_| rng.random_range(0..SHUFFLED_GROUP_COUNT as u32)) + .collect(); + + DenseGroupedInput { + values: PrimitiveArray::new(values, Validity::NonNullable).into_array(), + group_ids: GroupIds::from_buffer(group_ids, SHUFFLED_GROUP_COUNT).unwrap(), + } +} + fn grouped_accumulator(input: &DenseGroupedInput, vtable: V) -> ArrayRef where V: AggregateFnVTable + Clone, @@ -207,3 +224,19 @@ fn count_varbinview(bencher: Bencher) { .with_inputs(|| &input) .bench_refs(|input| grouped_accumulator(input, Count)); } + +#[divan::bench] +fn sum_i32_shuffled_4k_groups(bencher: Bencher) { + let input = i32_shuffled_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, Sum)); +} + +#[divan::bench] +fn count_i32_shuffled_4k_groups(bencher: Bencher) { + let input = i32_shuffled_input(); + bencher + .with_inputs(|| &input) + .bench_refs(|input| grouped_accumulator(input, Count)); +} diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 77bd17e483d..1d586a25429 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -1,6 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; +use std::sync::OnceLock; + use num_traits::ToPrimitive; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -28,6 +31,7 @@ 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; @@ -191,6 +195,7 @@ impl GroupRanges { pub struct GroupIds { ids: ArrayRef, num_groups: usize, + validated: Arc>>, } impl GroupIds { @@ -202,15 +207,21 @@ impl GroupIds { "Group ids must be non-nullable u32, got {}", ids.dtype() ); - Ok(Self { ids, num_groups }) + Ok(Self { + ids, + num_groups, + validated: Arc::new(OnceLock::new()), + }) } - /// Create group ids from a materialized buffer. + /// Create group ids from a materialized buffer, validating the dense-id invariant once. pub fn from_buffer(ids: Buffer, num_groups: usize) -> VortexResult { - Self::new( - PrimitiveArray::new(ids, Validity::NonNullable).into_array(), + 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. @@ -244,9 +255,17 @@ impl GroupIds { } /// 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) } } @@ -306,6 +325,8 @@ impl GroupedAccumulator { fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { validate_num_groups(num_groups)?; + 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.dtype)?); @@ -355,28 +376,44 @@ impl GroupedAccumulator { } let first = first as usize; - let mut buckets = vec![Vec::new(); last as usize - first + 1]; + 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() { - buckets[group_id as usize - first].push(row_idx as u64); + let cursor = &mut cursors[group_id as usize - first]; + permutation[*cursor] = row_idx as u64; + *cursor += 1; } - for (offset, rows) in buckets.into_iter().enumerate() { - if rows.is_empty() { + 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; } - let group = first + offset; + let group = first + group_offset; if self.vtable.is_saturated(&self.partials[group]) { continue; } - let taken = batch.clone().take(Buffer::from_iter(rows).into_array())?; let mut accumulator = Accumulator::try_new( self.vtable.clone(), self.options.clone(), self.dtype.clone(), )?; - accumulator.accumulate(&taken, ctx)?; + accumulator.accumulate(&gathered.slice(start..end)?, ctx)?; let partial = accumulator.flush()?; self.vtable .combine_partials(&mut self.partials[group], partial)?; diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 58e21754d6d..371b615251e 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -223,8 +223,13 @@ mod tests { #[test] 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::from_iter([0u32, 2], 2)?; + let group_ids = GroupIds::new( + PrimitiveArray::new(buffer![0u32, 2], Validity::NonNullable).into_array(), + 2, + )?; let mut ctx = array_session().create_execution_ctx(); let mut acc = GroupedAccumulator::try_new( Count, diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index cf8341c121b..8083e71441c 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -466,18 +466,13 @@ mod tests { fn mean_grouped_finalize() -> VortexResult<()> { let cases = mean_nan_null(); let values = PrimitiveArray::from_option_iter( - cases.iter().flat_map(|(group, _)| group.iter().copied()), + (0..3).flat_map(|row| cases.iter().map(move |(group, _)| group[row])), ) .into_array(); let groups = (0..cases.len()) .map(u32::try_from) .collect::, _>>()?; - let group_ids = GroupIds::from_iter( - groups - .into_iter() - .flat_map(|group| std::iter::repeat_n(group, 3)), - cases.len(), - )?; + let group_ids = GroupIds::from_iter(std::iter::repeat_n(groups, 3).flatten(), cases.len())?; let mut acc = GroupedAccumulator::try_new( Mean::combined(), diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index aa669500209..7ca85d83273 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -18,6 +18,7 @@ use super::checked_add_u64; use super::primitive::sum_float_all; use super::primitive::sum_signed_all; use super::primitive::sum_unsigned_all; +use super::sum_decimal_dtype; use crate::ArrayRef; use crate::ExecutionCtx; use crate::aggregate_fn::GroupIds; @@ -38,7 +39,7 @@ use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; use crate::scalar::DecimalValue; -const MIN_AVG_RUN_LENGTH_FOR_GROUPED_SUM_RUNS: usize = 4; +const MIN_GROUP_RUN_LENGTH: usize = 4; pub(crate) static SUM_GROUPED_KERNEL: GroupedAggregateKernelAdapter = GroupedAggregateKernelAdapter::new(SumGroupedKernel); @@ -91,21 +92,6 @@ fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { } } -fn should_accumulate_group_runs(group_ids: &[u32]) -> bool { - let Some((&first, rest)) = group_ids.split_first() else { - return false; - }; - let mut run_count = 1usize; - let mut group_id = first; - for &next_group_id in rest { - if next_group_id != group_id { - run_count += 1; - group_id = next_group_id; - } - } - run_count * MIN_AVG_RUN_LENGTH_FOR_GROUPED_SUM_RUNS <= group_ids.len() -} - 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; @@ -150,6 +136,21 @@ where } } +fn accumulate_grouped_unsigned_all(partials: &mut [SumPartial], values: &[T], group_ids: &[u32]) +where + T: NativePType + AsPrimitive, +{ + for_each_group_run(group_ids, |group_id, start, end| { + if end - start >= MIN_GROUP_RUN_LENGTH { + accumulate_grouped_unsigned_run(partials, group_id, &values[start..end]); + } else { + for &value in &values[start..end] { + accumulate_grouped_unsigned(partials, group_id, value.as_()); + } + } + }); +} + fn accumulate_grouped_signed(partials: &mut [SumPartial], group_id: u32, value: i64) { let partial = &mut partials[group_id as usize]; let saturated = match partial.current.as_mut() { @@ -177,6 +178,21 @@ where } } +fn accumulate_grouped_signed_all(partials: &mut [SumPartial], values: &[T], group_ids: &[u32]) +where + T: NativePType + AsPrimitive, +{ + for_each_group_run(group_ids, |group_id, start, end| { + if end - start >= MIN_GROUP_RUN_LENGTH { + accumulate_grouped_signed_run(partials, group_id, &values[start..end]); + } else { + for &value in &values[start..end] { + accumulate_grouped_signed(partials, group_id, value.as_()); + } + } + }); +} + fn accumulate_grouped_float( partials: &mut [SumPartial], group_id: u32, @@ -206,6 +222,24 @@ fn accumulate_grouped_float_run( } } +fn accumulate_grouped_float_all( + partials: &mut [SumPartial], + values: &[T], + group_ids: &[u32], + skip_nans: bool, +) { + for_each_group_run(group_ids, |group_id, start, end| { + if end - start >= MIN_GROUP_RUN_LENGTH { + accumulate_grouped_float_run(partials, group_id, &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(partials, group_id, value, skip_nans); + } + } + }); +} + fn accumulate_grouped_primitive( partials: &mut [SumPartial], primitive: &PrimitiveArray, @@ -217,16 +251,13 @@ fn accumulate_grouped_primitive( .as_ref() .validity()? .execute_mask(primitive.as_ref().len(), ctx)?; - let use_runs = - matches!(validity.slices(), AllOr::All) && should_accumulate_group_runs(group_ids); + let all_valid = matches!(validity.slices(), AllOr::All); match_each_native_ptype!(primitive.ptype(), unsigned: |T| { let values = primitive.as_slice::(); - if use_runs { - for_each_group_run(group_ids, |group_id, start, end| { - accumulate_grouped_unsigned_run(partials, group_id, &values[start..end]); - }); + if all_valid { + accumulate_grouped_unsigned_all(partials, values, group_ids); } else { for_each_valid_idx(&validity, values.len(), |idx| { accumulate_grouped_unsigned(partials, group_ids[idx], values[idx].as_()); @@ -235,10 +266,8 @@ fn accumulate_grouped_primitive( }, signed: |T| { let values = primitive.as_slice::(); - if use_runs { - for_each_group_run(group_ids, |group_id, start, end| { - accumulate_grouped_signed_run(partials, group_id, &values[start..end]); - }); + if all_valid { + accumulate_grouped_signed_all(partials, values, group_ids); } else { for_each_valid_idx(&validity, values.len(), |idx| { accumulate_grouped_signed(partials, group_ids[idx], values[idx].as_()); @@ -247,15 +276,8 @@ fn accumulate_grouped_primitive( }, floating: |T| { let values = primitive.as_slice::(); - if use_runs { - for_each_group_run(group_ids, |group_id, start, end| { - accumulate_grouped_float_run( - partials, - group_id, - &values[start..end], - skip_nans, - ); - }); + if all_valid { + accumulate_grouped_float_all(partials, 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"); @@ -299,15 +321,7 @@ fn accumulate_grouped_decimal( .as_ref() .validity()? .execute_mask(decimals.as_ref().len(), ctx)?; - let Some(output_dtype) = partials - .iter() - .find_map(|partial| match partial.current.as_ref() { - Some(SumState::Decimal { dtype, .. }) => Some(*dtype), - _ => None, - }) - else { - return Ok(()); - }; + let output_dtype = sum_decimal_dtype(&decimals.decimal_dtype()); let output_type = DecimalType::smallest_decimal_value_type(&output_dtype); match_each_decimal_value_type!(decimals.values_type(), |T| { match_each_decimal_value_type!(output_type, |I| { From f52c15d1976128ba864a50754fadfbbd6df522e5 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 7 Aug 2026 20:08:40 -0400 Subject: [PATCH 9/9] Optimize grouped aggregate state storage Signed-off-by: Nicholas Gates --- vortex-array/benches/aggregate_grouped.rs | 44 +- .../src/aggregate_fn/accumulator_grouped.rs | 215 +++++-- .../src/aggregate_fn/fns/count/grouped.rs | 132 +++- .../src/aggregate_fn/fns/count/mod.rs | 11 + .../src/aggregate_fn/fns/sum/grouped.rs | 362 +++++++---- .../src/aggregate_fn/fns/sum/grouped_state.rs | 573 ++++++++++++++++++ vortex-array/src/aggregate_fn/fns/sum/mod.rs | 12 + vortex-array/src/aggregate_fn/kernels.rs | 21 +- vortex-array/src/aggregate_fn/vtable.rs | 20 + 9 files changed, 1180 insertions(+), 210 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/sum/grouped_state.rs diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 985339feda3..d214854de80 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -10,6 +10,7 @@ use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use rand::seq::SliceRandom; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -37,8 +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 SHUFFLED_ELEMENT_COUNT: usize = 1 << 16; -const SHUFFLED_GROUP_COUNT: usize = 1 << 12; +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); @@ -142,18 +157,21 @@ fn varbinview_input() -> DenseGroupedInput { ) } -fn i32_shuffled_input() -> DenseGroupedInput { +fn i32_cardinality_input(group_count: usize, order: IdOrder) -> DenseGroupedInput { let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED); - let values: Buffer = (0..SHUFFLED_ELEMENT_COUNT) + let values: Buffer = (0..CARDINALITY_ELEMENT_COUNT) .map(|_| rng.random_range(-512..512)) .collect(); - let group_ids: Buffer = (0..SHUFFLED_ELEMENT_COUNT) - .map(|_| rng.random_range(0..SHUFFLED_GROUP_COUNT as u32)) + 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(group_ids, SHUFFLED_GROUP_COUNT).unwrap(), + group_ids: GroupIds::from_buffer(Buffer::from(group_ids), group_count).unwrap(), } } @@ -225,17 +243,17 @@ fn count_varbinview(bencher: Bencher) { .bench_refs(|input| grouped_accumulator(input, Count)); } -#[divan::bench] -fn sum_i32_shuffled_4k_groups(bencher: Bencher) { - let input = i32_shuffled_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, Sum)); } -#[divan::bench] -fn count_i32_shuffled_4k_groups(bencher: Bencher) { - let input = i32_shuffled_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_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 1d586a25429..4571a4e4480 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; use std::sync::Arc; use std::sync::OnceLock; @@ -270,6 +271,147 @@ impl GroupIds { } } +/// 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 @@ -289,8 +431,8 @@ pub struct GroupedAccumulator { return_dtype: DType, /// The DType of the partial accumulator state. partial_dtype: DType, - /// Dense per-group partial state. - partials: Vec, + /// Aggregate-owned dense per-group state. + state: Box, } impl GroupedAccumulator { @@ -310,6 +452,7 @@ impl GroupedAccumulator { dtype ) })?; + let state = vtable.grouped_state(&options, &dtype, &partial_dtype)?; Ok(Self { vtable, @@ -318,20 +461,14 @@ impl GroupedAccumulator { dtype, return_dtype, partial_dtype, - partials: Vec::new(), + state, }) } fn ensure_groups(&mut self, num_groups: usize) -> VortexResult<()> { validate_num_groups(num_groups)?; - 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.dtype)?); - } - Ok(()) + self.state.ensure_groups(num_groups) } fn try_accumulate_kernel( @@ -350,7 +487,7 @@ impl GroupedAccumulator { &self.aggregate_fn, batch, group_ids, - &mut self.partials, + self.state.as_any_mut(), ctx, )? { return Ok(true); @@ -404,7 +541,7 @@ impl GroupedAccumulator { } let group = first + group_offset; - if self.vtable.is_saturated(&self.partials[group]) { + if self.state.is_saturated(group) { continue; } @@ -415,8 +552,7 @@ impl GroupedAccumulator { )?; accumulator.accumulate(&gathered.slice(start..end)?, ctx)?; let partial = accumulator.flush()?; - self.vtable - .combine_partials(&mut self.partials[group], partial)?; + self.state.combine_scalar(group, partial)?; } Ok(()) } @@ -599,13 +735,8 @@ impl DynGroupedAccumulator for GroupedAccumulator { let group_ids = group_ids.validated_ids(ctx)?; self.ensure_groups(num_groups)?; - - for (row_idx, &group_id) in group_ids.iter().enumerate() { - let partial = partials.execute_scalar(row_idx, ctx)?; - self.vtable - .combine_partials(&mut self.partials[group_id as usize], partial)?; - } - Ok(()) + self.state + .accumulate_partials(partials, group_ids.as_ref(), ctx) } fn merge_group( @@ -621,9 +752,8 @@ impl DynGroupedAccumulator for GroupedAccumulator { other.partial_dtype() ); self.ensure_groups((into as usize) + 1)?; - let partial = other.partial_scalar(from)?; - self.vtable - .combine_partials(&mut self.partials[into as usize], partial) + self.state + .combine_scalar(into as usize, other.partial_scalar(from)?) } fn partial_dtype(&self) -> &DType { @@ -631,44 +761,11 @@ impl DynGroupedAccumulator for GroupedAccumulator { } fn partial_scalar(&self, group_id: u32) -> VortexResult { - if let Some(partial) = self.partials.get(group_id as usize) { - self.vtable.to_scalar(partial) - } else { - let partial = self.vtable.empty_partial(&self.options, &self.dtype)?; - self.vtable.to_scalar(&partial) - } + self.state.partial_scalar(group_id as usize) } 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()) + self.state.flush_partials(num_groups) } fn finish(&mut self, num_groups: usize) -> VortexResult { diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 371b615251e..49e41fd0f18 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -1,17 +1,22 @@ // 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 super::CountPartial; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; 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; @@ -19,6 +24,76 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; 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); @@ -27,14 +102,17 @@ pub(crate) static COUNT_GROUPED_KERNEL: GroupedAggregateKernelAdapter for CountGroupedKernel { + type State = CountGroupedState; + fn grouped_accumulate( &self, options: &NumericalAggregateOpts, - states: &mut [CountPartial], + state: &mut Self::State, batch: &ArrayRef, group_ids: &GroupIds, ctx: &mut ExecutionCtx, ) -> VortexResult { + let states = state.counts_mut(); if options.skip_nans && batch.dtype().is_float() { let Some(primitive) = batch.as_opt::() else { return Ok(false); @@ -51,13 +129,51 @@ impl GroupedAggregateKernel for CountGroupedKernel { let group_ids = group_ids.validated_ids(ctx)?; let validity = batch.validity()?.execute_mask(batch.len(), ctx)?; - for_each_valid_idx(&validity, batch.len(), |idx| { - states[group_ids[idx] as usize].count += 1; - }); + 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; + }); + } Ok(true) } } +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; + } + } + f(group_id, start, group_ids.len()); +} + fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { match validity.indices() { AllOr::All => (0..len).for_each(f), @@ -67,7 +183,7 @@ fn for_each_valid_idx(validity: &Mask, len: usize, mut f: impl FnMut(usize)) { } fn accumulate_grouped_float_count( - states: &mut [CountPartial], + states: &mut [u64], primitive: &PrimitiveArray, group_ids: &[u32], ctx: &mut ExecutionCtx, @@ -89,7 +205,7 @@ fn accumulate_grouped_float_count( } fn accumulate_valid_non_nan( - states: &mut [CountPartial], + states: &mut [u64], values: &[T], group_ids: &[u32], validity: &Mask, @@ -97,7 +213,7 @@ fn accumulate_valid_non_nan( 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].count += 1; + states[group_ids[idx] as usize] += 1; } }); } diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 04550f53453..e8e1148fbd2 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -3,6 +3,7 @@ mod grouped; pub(crate) use grouped::COUNT_GROUPED_KERNEL; +use grouped::CountGroupedState; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -13,6 +14,7 @@ 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; @@ -71,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() diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 7ca85d83273..bf3e79140b2 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use num_traits::AsPrimitive; +use num_traits::CheckedAdd; use num_traits::ToPrimitive; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -11,14 +12,14 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use super::Sum; -use super::SumPartial; -use super::SumState; 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_decimal_dtype; use crate::ArrayRef; use crate::ExecutionCtx; use crate::aggregate_fn::GroupIds; @@ -32,12 +33,10 @@ use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; -use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; use crate::dtype::NativePType; use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; -use crate::scalar::DecimalValue; const MIN_GROUP_RUN_LENGTH: usize = 4; @@ -48,10 +47,12 @@ pub(crate) static SUM_GROUPED_KERNEL: GroupedAggregateKernelAdapter for SumGroupedKernel { + type State = SumGroupedState; + fn grouped_accumulate( &self, options: &NumericalAggregateOpts, - partials: &mut [SumPartial], + state: &mut Self::State, batch: &ArrayRef, group_ids: &GroupIds, ctx: &mut ExecutionCtx, @@ -59,7 +60,7 @@ impl GroupedAggregateKernel for SumGroupedKernel { if let Some(primitive) = batch.as_opt::() { let group_ids = group_ids.validated_ids(ctx)?; accumulate_grouped_primitive( - partials, + state, &primitive.into_owned(), group_ids.as_ref(), options.skip_nans, @@ -70,13 +71,13 @@ impl GroupedAggregateKernel for SumGroupedKernel { if let Some(bools) = batch.as_opt::() { let group_ids = group_ids.validated_ids(ctx)?; - accumulate_grouped_bool(partials, &bools.into_owned(), group_ids.as_ref(), 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(partials, &decimals.into_owned(), group_ids.as_ref(), ctx)?; + accumulate_grouped_decimal(state, &decimals.into_owned(), group_ids.as_ref(), ctx)?; return Ok(true); } @@ -109,139 +110,154 @@ fn for_each_group_run(group_ids: &[u32], mut f: impl FnMut(u32, usize, usize)) { f(group_id, start, group_ids.len()); } -fn accumulate_grouped_unsigned(partials: &mut [SumPartial], group_id: u32, value: u64) { - let partial = &mut partials[group_id as usize]; - let saturated = match partial.current.as_mut() { - None => return, - Some(SumState::Unsigned(acc)) => checked_add_u64(acc, value), - Some(_) => vortex_panic!("unsigned sum state with non-unsigned input"), - }; - if saturated { - partial.current = 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 >= MIN_GROUP_RUN_LENGTH { + return true; + } + } else { + run_length = 1; + } } + false } -fn accumulate_grouped_unsigned_run(partials: &mut [SumPartial], group_id: u32, values: &[T]) -where +fn accumulate_grouped_unsigned( + values: &mut [u64], + overflowed: &mut [u8], + group_id: u32, + value: u64, +) { + let group = group_id as usize; + if checked_add_u64(&mut values[group], value) { + overflowed[group] = 1; + } +} + +fn accumulate_grouped_unsigned_run( + sums: &mut [u64], + overflowed: &mut [u8], + group_id: u32, + values: &[T], +) where T: NativePType + AsPrimitive, { - let partial = &mut partials[group_id as usize]; - let saturated = match partial.current.as_mut() { - None => return, - Some(SumState::Unsigned(acc)) => sum_unsigned_all(acc, values), - Some(_) => vortex_panic!("unsigned sum state with non-unsigned input"), - }; - if saturated { - partial.current = None; + let group = group_id as usize; + if sum_unsigned_all(&mut sums[group], values) { + overflowed[group] = 1; } } -fn accumulate_grouped_unsigned_all(partials: &mut [SumPartial], values: &[T], group_ids: &[u32]) -where +fn accumulate_grouped_unsigned_all( + sums: &mut [u64], + overflowed: &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, group_id, value.as_()); + } + return; + } + for_each_group_run(group_ids, |group_id, start, end| { if end - start >= MIN_GROUP_RUN_LENGTH { - accumulate_grouped_unsigned_run(partials, group_id, &values[start..end]); + accumulate_grouped_unsigned_run(sums, overflowed, group_id, &values[start..end]); } else { for &value in &values[start..end] { - accumulate_grouped_unsigned(partials, group_id, value.as_()); + accumulate_grouped_unsigned(sums, overflowed, group_id, value.as_()); } } }); } -fn accumulate_grouped_signed(partials: &mut [SumPartial], group_id: u32, value: i64) { - let partial = &mut partials[group_id as usize]; - let saturated = match partial.current.as_mut() { - None => return, - Some(SumState::Signed(acc)) => checked_add_i64(acc, value), - Some(_) => vortex_panic!("signed sum state with non-signed input"), - }; - if saturated { - partial.current = None; +fn accumulate_grouped_signed(values: &mut [i64], overflowed: &mut [u8], group_id: u32, value: i64) { + let group = group_id as usize; + if checked_add_i64(&mut values[group], value) { + overflowed[group] = 1; } } -fn accumulate_grouped_signed_run(partials: &mut [SumPartial], group_id: u32, values: &[T]) -where +fn accumulate_grouped_signed_run( + sums: &mut [i64], + overflowed: &mut [u8], + group_id: u32, + values: &[T], +) where T: NativePType + AsPrimitive, { - let partial = &mut partials[group_id as usize]; - let saturated = match partial.current.as_mut() { - None => return, - Some(SumState::Signed(acc)) => sum_signed_all(acc, values), - Some(_) => vortex_panic!("signed sum state with non-signed input"), - }; - if saturated { - partial.current = None; + let group = group_id as usize; + if sum_signed_all(&mut sums[group], values) { + overflowed[group] = 1; } } -fn accumulate_grouped_signed_all(partials: &mut [SumPartial], values: &[T], group_ids: &[u32]) -where +fn accumulate_grouped_signed_all( + sums: &mut [i64], + overflowed: &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, 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(partials, group_id, &values[start..end]); + accumulate_grouped_signed_run(sums, overflowed, group_id, &values[start..end]); } else { for &value in &values[start..end] { - accumulate_grouped_signed(partials, group_id, value.as_()); + accumulate_grouped_signed(sums, overflowed, group_id, value.as_()); } } }); } -fn accumulate_grouped_float( - partials: &mut [SumPartial], - group_id: u32, - value: f64, - skip_nans: bool, -) { - if skip_nans && value.is_nan() { - return; - } - match partials[group_id as usize].current.as_mut() { - None => {} - Some(SumState::Float(acc)) => *acc += value, - Some(_) => vortex_panic!("float sum state with non-float input"), - } -} - -fn accumulate_grouped_float_run( - partials: &mut [SumPartial], - group_id: u32, - values: &[T], - skip_nans: bool, -) { - match partials[group_id as usize].current.as_mut() { - None => {} - Some(SumState::Float(acc)) => sum_float_all(acc, values, skip_nans), - Some(_) => vortex_panic!("float sum state with non-float input"), +fn accumulate_grouped_float(sums: &mut [f64], group_id: u32, value: f64, skip_nans: bool) { + if !skip_nans || !value.is_nan() { + sums[group_id as usize] += value; } } fn accumulate_grouped_float_all( - partials: &mut [SumPartial], + sums: &mut [f64], 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, group_id, value, skip_nans); + } + return; + } + for_each_group_run(group_ids, |group_id, start, end| { if end - start >= MIN_GROUP_RUN_LENGTH { - accumulate_grouped_float_run(partials, group_id, &values[start..end], skip_nans); + 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(partials, group_id, value, skip_nans); + accumulate_grouped_float(sums, group_id, value, skip_nans); } } }); } fn accumulate_grouped_primitive( - partials: &mut [SumPartial], + state: &mut SumGroupedState, primitive: &PrimitiveArray, group_ids: &[u32], skip_nans: bool, @@ -252,36 +268,46 @@ fn accumulate_grouped_primitive( .validity()? .execute_mask(primitive.as_ref().len(), ctx)?; let all_valid = matches!(validity.slices(), AllOr::All); + let (state, overflowed) = state.parts_mut(); match_each_native_ptype!(primitive.ptype(), unsigned: |T| { + 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(partials, values, group_ids); + accumulate_grouped_unsigned_all(sums, overflowed, values, group_ids); } else { for_each_valid_idx(&validity, values.len(), |idx| { - accumulate_grouped_unsigned(partials, group_ids[idx], values[idx].as_()); + accumulate_grouped_unsigned(sums, overflowed, group_ids[idx], values[idx].as_()); }); } }, signed: |T| { + 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(partials, values, group_ids); + accumulate_grouped_signed_all(sums, overflowed, values, group_ids); } else { for_each_valid_idx(&validity, values.len(), |idx| { - accumulate_grouped_signed(partials, group_ids[idx], values[idx].as_()); + accumulate_grouped_signed(sums, overflowed, group_ids[idx], values[idx].as_()); }); } }, floating: |T| { + 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(partials, values, group_ids, skip_nans); + accumulate_grouped_float_all(sums, 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(partials, group_ids[idx], value, skip_nans); + accumulate_grouped_float(sums, group_ids[idx], value, skip_nans); }); } } @@ -290,7 +316,7 @@ fn accumulate_grouped_primitive( } fn accumulate_grouped_bool( - partials: &mut [SumPartial], + state: &mut SumGroupedState, bools: &BoolArray, group_ids: &[u32], ctx: &mut ExecutionCtx, @@ -305,14 +331,18 @@ fn accumulate_grouped_bool( AllOr::None => return Ok(()), AllOr::Some(validity) => &values & validity, }; + let (state, overflowed) = state.parts_mut(); + let SumGroupedValues::Unsigned(sums) = state else { + vortex_panic!("boolean input with non-unsigned grouped sum state") + }; valid_true.for_each_set_index(|idx| { - accumulate_grouped_unsigned(partials, group_ids[idx], 1); + accumulate_grouped_unsigned(sums, overflowed, group_ids[idx], 1); }); Ok(()) } fn accumulate_grouped_decimal( - partials: &mut [SumPartial], + state: &mut SumGroupedState, decimals: &DecimalArray, group_ids: &[u32], ctx: &mut ExecutionCtx, @@ -321,46 +351,86 @@ fn accumulate_grouped_decimal( .as_ref() .validity()? .execute_mask(decimals.as_ref().len(), ctx)?; - let output_dtype = sum_decimal_dtype(&decimals.decimal_dtype()); - let output_type = DecimalType::smallest_decimal_value_type(&output_dtype); + let output_dtype = state + .decimal_dtype() + .vortex_expect("decimal sum state dtype"); + let (state, overflowed) = state.parts_mut(); match_each_decimal_value_type!(decimals.values_type(), |T| { - match_each_decimal_value_type!(output_type, |I| { - accumulate_grouped_decimal_values::( - partials, - decimals.buffer::(), + let values = decimals.buffer::(); + match state { + SumGroupedValues::Decimal8(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal16(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal32(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal64(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + values, + group_ids, + &validity, + output_dtype, + ), + SumGroupedValues::Decimal128(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + values, group_ids, &validity, - ); - }) + output_dtype, + ), + SumGroupedValues::Decimal256(sums) => accumulate_grouped_decimal_values( + sums, + overflowed, + values, + group_ids, + &validity, + output_dtype, + ), + _ => vortex_panic!("decimal input with non-decimal grouped sum state"), + } }); Ok(()) } fn accumulate_grouped_decimal_values( - partials: &mut [SumPartial], + sums: &mut [I], + overflowed: &mut [u8], values: Buffer, group_ids: &[u32], validity: &Mask, + dtype: crate::dtype::DecimalDType, ) where T: NativeDecimalType + AsPrimitive, - I: NativeDecimalType, - DecimalValue: From, + I: NativeDecimalType + CheckedAdd, { for_each_valid_idx(validity, values.len(), |idx| { - let partial = &mut partials[group_ids[idx] as usize]; - let Some(SumState::Decimal { value, dtype }) = partial.current.as_mut() else { - return; - }; - let operand = DecimalValue::from(values[idx].as_()); - let Some(result) = value.checked_add(&operand) else { - partial.current = None; - return; - }; - if result.fits_in_precision(*dtype) { - *value = result; - } else { - partial.current = None; - } + add_decimal( + sums, + overflowed, + group_ids[idx] as usize, + values[idx].as_(), + dtype, + ); }); } @@ -382,7 +452,10 @@ mod tests { use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::dtype::DType; use crate::dtype::DecimalDType; + use crate::dtype::Nullability; + use crate::dtype::PType; use crate::dtype::i256; use crate::scalar::DecimalValue; use crate::validity::Validity; @@ -535,4 +608,55 @@ mod tests { ); Ok(()) } + + #[test] + fn accumulates_typed_primitive_partials() -> VortexResult<()> { + let input_dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let partials = + PrimitiveArray::from_option_iter([Some(2i64), Some(3), Some(5), None]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let mut acc = + GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::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 accumulates_typed_decimal_partials() -> VortexResult<()> { + let input_dtype = DecimalDType::new(10, 2); + let partial_dtype = DecimalDType::new(20, 2); + let partials = DecimalArray::new( + buffer![200i64, 300, 500, 700], + partial_dtype, + Validity::from_iter([true, true, true, false]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + let mut acc = GroupedAccumulator::try_new( + Sum, + NumericalAggregateOpts::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..561591a3108 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped_state.rs @@ -0,0 +1,573 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; + +use num_traits::CheckedAdd; +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::AllOr; +use vortex_mask::Mask; + +use super::checked_add_i64; +use super::checked_add_u64; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::GroupedState; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +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, + partial_dtype: DType, +} + +impl SumGroupedState { + pub(crate) fn try_new(partial_dtype: DType) -> VortexResult { + let values = match &partial_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 partial dtype: {dtype}"), + }; + Ok(Self { + values, + overflowed: Vec::new(), + partial_dtype, + }) + } + + pub(super) fn parts_mut(&mut self) -> (&mut SumGroupedValues, &mut [u8]) { + (&mut self.values, &mut self.overflowed) + } + + pub(super) fn decimal_dtype(&self) -> Option { + self.partial_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); + 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<()> { + if partial.is_null() { + 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 { + if self + .overflowed + .get(group_id) + .is_some_and(|&overflowed| overflowed != 0) + { + return Ok(Scalar::null(self.partial_dtype.clone())); + } + + Ok(match &self.values { + SumGroupedValues::Unsigned(values) => Scalar::primitive( + values.get(group_id).copied().unwrap_or(0), + Nullability::Nullable, + ), + SumGroupedValues::Signed(values) => Scalar::primitive( + values.get(group_id).copied().unwrap_or(0), + Nullability::Nullable, + ), + SumGroupedValues::Float(values) => Scalar::primitive( + values.get(group_id).copied().unwrap_or(0.0), + Nullability::Nullable, + ), + 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"), + ), + }) + } + + fn accumulate_partials( + &mut self, + partials: &ArrayRef, + group_ids: &[u32], + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let validity = partials.validity()?.execute_mask(partials.len(), ctx)?; + let decimal_dtype = self.decimal_dtype(); + let (values, overflowed) = self.parts_mut(); + match values { + SumGroupedValues::Unsigned(values) => { + let partials = partials.clone().execute::(ctx)?; + accumulate_primitive_partials( + values, + overflowed, + partials.as_slice::(), + group_ids, + &validity, + checked_add_u64, + ); + } + SumGroupedValues::Signed(values) => { + let partials = partials.clone().execute::(ctx)?; + accumulate_primitive_partials( + values, + overflowed, + partials.as_slice::(), + group_ids, + &validity, + checked_add_i64, + ); + } + SumGroupedValues::Float(values) => { + let partials = partials.clone().execute::(ctx)?; + accumulate_float_partials( + values, + overflowed, + partials.as_slice::(), + group_ids, + &validity, + ); + } + SumGroupedValues::Decimal8(values) => accumulate_decimal_partials( + values, + overflowed, + &partials.clone().execute::(ctx)?, + group_ids, + &validity, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal16(values) => accumulate_decimal_partials( + values, + overflowed, + &partials.clone().execute::(ctx)?, + group_ids, + &validity, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal32(values) => accumulate_decimal_partials( + values, + overflowed, + &partials.clone().execute::(ctx)?, + group_ids, + &validity, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal64(values) => accumulate_decimal_partials( + values, + overflowed, + &partials.clone().execute::(ctx)?, + group_ids, + &validity, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal128(values) => accumulate_decimal_partials( + values, + overflowed, + &partials.clone().execute::(ctx)?, + group_ids, + &validity, + decimal_dtype.vortex_expect("decimal state dtype"), + ), + SumGroupedValues::Decimal256(values) => accumulate_decimal_partials( + values, + overflowed, + &partials.clone().execute::(ctx)?, + group_ids, + &validity, + 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 validity = validity_from_overflow(&self.overflowed); + self.overflowed.clear(); + let decimal_dtype = self.decimal_dtype(); + Ok(match &mut self.values { + SumGroupedValues::Unsigned(values) => { + PrimitiveArray::new(Buffer::from(std::mem::take(values)), validity).into_array() + } + SumGroupedValues::Signed(values) => { + PrimitiveArray::new(Buffer::from(std::mem::take(values)), validity).into_array() + } + SumGroupedValues::Float(values) => { + PrimitiveArray::new(Buffer::from(std::mem::take(values)), validity).into_array() + } + SumGroupedValues::Decimal8(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + validity, + ), + SumGroupedValues::Decimal16(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + validity, + ), + SumGroupedValues::Decimal32(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + validity, + ), + SumGroupedValues::Decimal64(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + validity, + ), + SumGroupedValues::Decimal128(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + validity, + ), + SumGroupedValues::Decimal256(values) => decimal_array( + std::mem::take(values), + decimal_dtype.vortex_expect("decimal state dtype"), + validity, + ), + }) + } +} + +fn for_each_partial_row(validity: &Mask, len: usize, mut f: impl FnMut(usize, bool)) { + match validity.indices() { + AllOr::All => (0..len).for_each(|idx| f(idx, true)), + AllOr::None => (0..len).for_each(|idx| f(idx, false)), + AllOr::Some(valid_indices) => { + let mut valid_indices = valid_indices.iter().copied().peekable(); + for idx in 0..len { + let valid = valid_indices.next_if_eq(&idx).is_some(); + f(idx, valid); + } + } + } +} + +fn accumulate_primitive_partials( + values: &mut [T], + overflowed: &mut [u8], + partials: &[T], + group_ids: &[u32], + validity: &Mask, + checked_add: fn(&mut T, T) -> bool, +) { + for_each_partial_row(validity, partials.len(), |idx, valid| { + let group = group_ids[idx] as usize; + if !valid || checked_add(&mut values[group], partials[idx]) { + overflowed[group] = 1; + } + }); +} + +fn accumulate_float_partials( + values: &mut [f64], + overflowed: &mut [u8], + partials: &[f64], + group_ids: &[u32], + validity: &Mask, +) { + for_each_partial_row(validity, partials.len(), |idx, valid| { + let group = group_ids[idx] as usize; + if !valid { + overflowed[group] = 1; + } else { + values[group] += partials[idx]; + } + }); +} + +fn accumulate_decimal_partials( + values: &mut [I], + overflowed: &mut [u8], + partials: &DecimalArray, + group_ids: &[u32], + validity: &Mask, + dtype: DecimalDType, +) where + I: NativeDecimalType + CheckedAdd, +{ + match_each_decimal_value_type!(partials.values_type(), |T| { + accumulate_decimal_partial_values( + values, + overflowed, + &partials.buffer::(), + group_ids, + validity, + dtype, + ); + }); +} + +fn accumulate_decimal_partial_values( + values: &mut [I], + overflowed: &mut [u8], + partials: &[T], + group_ids: &[u32], + validity: &Mask, + dtype: DecimalDType, +) where + T: NativeDecimalType, + I: NativeDecimalType + CheckedAdd, +{ + for_each_partial_row(validity, partials.len(), |idx, valid| { + let group = group_ids[idx] as usize; + if !valid { + overflowed[group] = 1; + } else { + let Some(value) = ::from(partials[idx]) else { + overflowed[group] = 1; + return; + }; + add_decimal(values, 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 validity_from_overflow(overflowed: &[u8]) -> Validity { + if overflowed.iter().all(|&overflowed| overflowed == 0) { + Validity::AllValid + } else { + Validity::from_iter(overflowed.iter().map(|&overflowed| overflowed == 0)) + } +} + +fn decimal_scalar(value: T, dtype: DecimalDType) -> Scalar +where + DecimalValue: From, +{ + Scalar::decimal(DecimalValue::from(value), dtype, Nullability::Nullable) +} + +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 515178585ee..1d41a40b583 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -5,6 +5,7 @@ mod bool; mod constant; mod decimal; mod grouped; +mod grouped_state; mod primitive; pub(crate) use grouped::SUM_GROUPED_KERNEL; @@ -20,6 +21,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; @@ -30,6 +32,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::DecimalArray; use crate::arrays::PrimitiveArray; @@ -164,6 +167,15 @@ impl AggregateFnVTable for Sum { }) } + fn grouped_state( + &self, + _options: &Self::Options, + _input_dtype: &DType, + partial_dtype: &DType, + ) -> VortexResult> { + Ok(Box::new(SumGroupedState::try_new(partial_dtype.clone())?)) + } + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { if other.is_null() { // A null partial means the sub-accumulator saturated (overflow). diff --git a/vortex-array/src/aggregate_fn/kernels.rs b/vortex-array/src/aggregate_fn/kernels.rs index 51d47c33a2e..0651ec4ddff 100644 --- a/vortex-array/src/aggregate_fn/kernels.rs +++ b/vortex-array/src/aggregate_fn/kernels.rs @@ -17,6 +17,7 @@ use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFnRef; 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. @@ -37,11 +38,14 @@ pub trait DynAggregateKernel: 'static + Send + Sync + Debug { /// 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, - states: &mut [V::Partial], + state: &mut Self::State, batch: &ArrayRef, group_ids: &GroupIds, ctx: &mut ExecutionCtx, @@ -115,7 +119,7 @@ where return Ok(false); }; - let Some(states) = states.downcast_mut::>() else { + let Some(state) = states.downcast_mut::() else { vortex_bail!( "Grouped aggregate kernel for {} received incompatible partial state", aggregate_fn.id() @@ -123,19 +127,14 @@ where }; vortex_ensure!( - states.len() >= group_ids.num_groups(), + state.len() >= group_ids.num_groups(), "Grouped aggregate kernel for {} received {} partial states for {} groups", aggregate_fn.id(), - states.len(), + state.len(), group_ids.num_groups() ); - self.kernel.grouped_accumulate( - options, - &mut states[..group_ids.num_groups()], - batch, - group_ids, - ctx, - ) + self.kernel + .grouped_accumulate(options, state, batch, group_ids, ctx) } } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 68bb0e7ec0c..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<()>;