From f185f493f4be4568e82d75fb8869f460e25e1c0f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 9 Sep 2026 14:32:06 -0400 Subject: [PATCH] perf: specialize grouped sums for constant arrays Signed-off-by: Connor Tsui --- vortex-array/benches/aggregate_grouped.rs | 82 ++++++- vortex-array/src/aggregate_fn/session.rs | 5 + .../src/arrays/constant/compute/mod.rs | 1 + .../src/arrays/constant/compute/sum.rs | 201 ++++++++++++++++++ 4 files changed, 288 insertions(+), 1 deletion(-) create mode 100644 vortex-array/src/arrays/constant/compute/sum.rs diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index 64820997778..ab95fcb4eab 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -11,22 +11,32 @@ use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::Canonical; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; use vortex_array::aggregate_fn::GroupedAccumulator; +use vortex_array::aggregate_fn::GroupedArray; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::count::Count; use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; +use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_error::VortexResult; use vortex_session::VortexSession; fn main() { @@ -149,7 +159,9 @@ fn varbinview_input() -> ArrayRef { fn list_element_dtype(list_view: &ArrayRef) -> DType { match list_view.dtype() { - DType::List(element_dtype, _) => element_dtype.as_ref().clone(), + DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { + element_dtype.as_ref().clone() + } dtype => unreachable!("expected List dtype, got {dtype}"), } } @@ -335,3 +347,71 @@ fn count_varbinview(bencher: Bencher) { .with_inputs(|| &input) .bench_refs(|input| grouped_accumulator(input, Count)); } + +/// Disable only the constant kernel so the same inputs exercise the existing grouped fallback. +#[derive(Debug)] +struct NoConstantSum; + +impl DynGroupedAggregateKernel for NoConstantSum { + fn grouped_aggregate( + &self, + _aggregate_fn: &AggregateFnRef, + _groups: &GroupedArray, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } +} + +fn bench_constant_sum(bencher: Bencher, groups: ArrayRef, specialized: bool) { + let session = vortex_array::array_session(); + if !specialized { + session.aggregate_fns().register_grouped_encoding_kernel( + Constant.id(), + SumV2.id(), + &NoConstantSum, + ); + } + let dtype = list_element_dtype(&groups); + bencher + .with_inputs(|| { + ( + GroupedAccumulator::try_new( + SumV2, + NumericalAggregateOpts::default(), + dtype.clone(), + ) + .unwrap(), + session.create_execution_ctx(), + ) + }) + .bench_refs(|(acc, ctx)| { + acc.accumulate_list(&groups, ctx).unwrap(); + acc.finish() + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench(args = [2, 128], consts = [true, false])] +fn sum_v2_constant_fixed(bencher: Bencher, size: u32) { + let len = 16_384; + let groups = FixedSizeListArray::try_new( + ConstantArray::new(3i32, len).into_array(), + size, + Validity::NonNullable, + len / size as usize, + ) + .unwrap() + .into_array(); + + bench_constant_sum(bencher, groups, SPECIALIZED); +} + +#[divan::bench(consts = [true, false])] +fn sum_v2_constant_list(bencher: Bencher) { + let sizes = random_group_sizes(); + let elements = ConstantArray::new(3i32, total_element_count(&sizes)).into_array(); + bench_constant_sum(bencher, contiguous_list_view(elements, &sizes), SPECIALIZED); +} diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 65b00d81031..3b9ad5d6e03 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -41,9 +41,11 @@ use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; use crate::array::ArrayId; use crate::array::VTable; use crate::arrays::Chunked; +use crate::arrays::Constant; use crate::arrays::Dict; use crate::arrays::Primitive; use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; +use crate::arrays::constant::compute::sum::ConstantGroupedSumKernel; use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; use crate::arrays::dict::compute::min_max::DictMinMaxKernel; @@ -134,6 +136,9 @@ impl Default for AggregateFnSession { SumV2.id(), &PrimitiveGroupedSumV2EncodingKernel, ); + for sum in [Sum.id(), SumV2.id()] { + this.register_grouped_encoding_kernel(Constant.id(), sum, &ConstantGroupedSumKernel); + } this } diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index a996bf79a0f..0d5f2449471 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -8,6 +8,7 @@ mod filter; mod not; pub(crate) mod rules; mod slice; +pub(crate) mod sum; mod take; pub(crate) mod uncompressed_size; diff --git a/vortex-array/src/arrays/constant/compute/sum.rs b/vortex-array/src/arrays/constant/compute/sum.rs new file mode 100644 index 00000000000..69eb520d5ee --- /dev/null +++ b/vortex-array/src/arrays/constant/compute/sum.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Grouped primitive sums over constant elements. +//! +//! Reuse the whole-array accumulator for each group size. Fixed-size groups share one +//! constant partial, and list-view groups avoid slicing the elements before accumulation. + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::GroupRanges; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::SumV2; +use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::builders::builder_with_capacity_in; +use crate::builtins::ArrayBuiltins; +use crate::validity::Validity; + +/// Grouped primitive sum kernel for constant elements. +#[derive(Debug)] +pub(crate) struct ConstantGroupedSumKernel; + +impl DynGroupedAggregateKernel for ConstantGroupedSumKernel { + fn grouped_aggregate( + &self, + aggregate_fn: &AggregateFnRef, + groups: &GroupedArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !aggregate_fn.is::() && !aggregate_fn.is::() { + return Ok(None); + } + let Some(elements) = groups.elements().as_opt::() else { + return Ok(None); + }; + if !elements.dtype().is_primitive() { + return Ok(None); + } + + let ranges = groups.group_ranges(ctx)?; + let validity = groups.group_validity(ctx)?; + let scalar = elements.scalar(); + let mut accumulator = aggregate_fn.accumulator(elements.dtype())?; + + if let GroupRanges::FixedSizeList { size, .. } = ranges { + let group = ConstantArray::new(scalar.clone(), size).into_array(); + accumulator.accumulate(&group, ctx)?; + let partials = ConstantArray::new(accumulator.flush()?, groups.len()).into_array(); + let mask = BoolArray::new(validity.to_bit_buffer(), Validity::NonNullable).into_array(); + return Ok(Some(partials.mask(mask)?)); + } + + let partial_dtype = aggregate_fn + .state_dtype(elements.dtype()) + .vortex_expect("The primitive sum accumulator has a partial dtype"); + let mut partials = builder_with_capacity_in(&partial_dtype, groups.len(), ctx.allocator()); + for ((_, size), valid) in ranges.iter().zip(validity.iter()) { + if !valid { + partials.append_null(); + continue; + } + + let group = ConstantArray::new(scalar.clone(), size).into_array(); + accumulator.accumulate(&group, ctx)?; + partials.append_scalar(&accumulator.flush()?)?; + } + + Ok(Some(partials.finish())) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::ConstantGroupedSumKernel; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::AggregateFnVTableExt; + use crate::aggregate_fn::GroupedArray; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::sum::Sum; + use crate::aggregate_fn::fns::sum_v2::SumV2; + use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; + use crate::array_session; + use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::ListViewArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability::NonNullable; + use crate::scalar::Scalar; + use crate::validity::Validity; + + fn check_groups( + scalar: Scalar, + fixed_size: bool, + options: NumericalAggregateOpts, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = ConstantArray::new(scalar, 8).into_array(); + let decoded = elements + .clone() + .execute::(&mut ctx)? + .into_array(); + let make_groups = |elements: ArrayRef| -> VortexResult { + if fixed_size { + Ok(FixedSizeListArray::try_new( + elements, + 2, + Validity::from_iter([true, false, true, true]), + 4, + )? + .into()) + } else { + Ok(ListViewArray::try_new( + elements, + buffer![0u32, 4, 2, 8, 1].into_array(), + buffer![3u32, 2, 4, 0, 0].into_array(), + Validity::from_iter([true, false, true, true, true]), + )? + .into()) + } + }; + let groups = make_groups(elements.clone())?; + let reference = make_groups(decoded.clone())?; + let as_array = |groups: &GroupedArray| match groups { + GroupedArray::ListView(array) => array.clone().into_array(), + GroupedArray::FixedSizeList(array) => array.clone().into_array(), + }; + let groups_array = as_array(&groups); + let reference_array = as_array(&reference); + + for aggregate in [Sum.bind(options), SumV2.bind(options)] { + assert!( + ConstantGroupedSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_some() + ); + let mut actual = aggregate.accumulator_grouped(elements.dtype())?; + actual.accumulate_list(&groups_array, &mut ctx)?; + let mut expected = aggregate.accumulator_grouped(decoded.dtype())?; + expected.accumulate_list(&reference_array, &mut ctx)?; + assert_arrays_eq!(actual.finish()?, expected.finish()?, &mut ctx); + } + + Ok(()) + } + + #[rstest] + #[case::i32((-3i32).into())] + #[case::null(Scalar::null_native::())] + #[case::signed_overflow(i64::MAX.into())] + #[case::negative_zero((-0.0f64).into())] + fn primitive_groups( + #[case] scalar: Scalar, + #[values(false, true)] fixed_size: bool, + ) -> VortexResult<()> { + check_groups(scalar, fixed_size, NumericalAggregateOpts::default()) + } + + #[rstest] + fn nan_groups(#[values(true, false)] skip_nans: bool) -> VortexResult<()> { + check_groups(f64::NAN.into(), false, NumericalAggregateOpts { skip_nans }) + } + + #[test] + fn decimal_kernel_declines() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(10, 2), NonNullable); + let elements = ConstantArray::new(Scalar::zero_value(&dtype), 4).into_array(); + let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into(); + + for aggregate in [ + Sum.bind(NumericalAggregateOpts::default()), + SumV2.bind(NumericalAggregateOpts::default()), + ] { + assert!( + ConstantGroupedSumKernel + .grouped_aggregate(&aggregate, &groups, &mut ctx)? + .is_none() + ); + } + + Ok(()) + } +}