perf: specialize grouped sums for constant arrays

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<Option<ArrayRef>> {
+        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::<PrimitiveArray>(ctx)
+                .unwrap()
+        });
+}
+
+#[divan::bench(args = [2, 128], consts = [true, false])]
+fn sum_v2_constant_fixed<const SPECIALIZED: bool>(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<const SPECIALIZED: bool>(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);
+}

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
     }

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;
 

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<Option<ArrayRef>> {
+        if !aggregate_fn.is::<Sum>() && !aggregate_fn.is::<SumV2>() {
+            return Ok(None);
+        }
+        let Some(elements) = groups.elements().as_opt::<Constant>() 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::<PrimitiveArray>(&mut ctx)?
+            .into_array();
+        let make_groups = |elements: ArrayRef| -> VortexResult<GroupedArray> {
+            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::<i32>())]
+    #[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(())
+    }
+}