comments

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -14,7 +14,6 @@ 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;
@@ -67,22 +66,6 @@ impl GroupIds {
         Self::from_buffer(Buffer::from_iter(ids), num_groups)
     }
 
-    /// Create group ids containing `0..num_groups`.
-    pub fn range(num_groups: usize) -> VortexResult<Self> {
-        validate_num_groups(num_groups)?;
-        if num_groups == 0 {
-            return Self::from_buffer(Buffer::<u32>::empty(), num_groups);
-        }
-
-        let last = u32::try_from(num_groups - 1).map_err(|_| {
-            vortex_err!(
-                "num_groups {} exceeds dense u32 group id capacity",
-                num_groups
-            )
-        })?;
-        Self::from_buffer((0..=last).collect(), num_groups)
-    }
-
     /// Return the encoded ids array.
     pub fn ids(&self) -> &ArrayRef {
         &self.ids
@@ -178,14 +161,6 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
         Ok(())
     }
 
-    fn accumulate_kernel_result(
-        &mut self,
-        result: GroupedAggregateKernelResult,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<()> {
-        self.accumulate_partials(result.partials(), result.group_ids(), ctx)
-    }
-
     fn try_accumulate_kernel(
         &mut self,
         batch: &ArrayRef,
@@ -198,10 +173,13 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
             self.aggregate_fn.id(),
             batch.encoding_id(),
             group_ids.encoding_id(),
-        ) && let Some(result) =
-            kernel.grouped_aggregate(&self.aggregate_fn, batch, group_ids, ctx)?
-        {
-            self.accumulate_kernel_result(result, ctx)?;
+        ) && kernel.grouped_accumulate(
+            &self.aggregate_fn,
+            batch,
+            group_ids,
+            &mut self.partials,
+            ctx,
+        )? {
             return Ok(true);
         }
 

vortex-array/src/aggregate_fn/fns/count/grouped.rs

@@ -6,48 +6,33 @@ use vortex_error::VortexResult;
 use super::Count;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
-use crate::IntoArray;
-use crate::aggregate_fn::AggregateFnRef;
+use crate::aggregate_fn::EmptyOptions;
 use crate::aggregate_fn::GroupIds;
-use crate::aggregate_fn::kernels::DynGroupedAggregateKernel;
-use crate::aggregate_fn::kernels::GroupedAggregateKernelResult;
-use crate::arrays::PrimitiveArray;
+use crate::aggregate_fn::kernels::GroupedAggregateKernel;
+use crate::aggregate_fn::kernels::GroupedAggregateKernelAdapter;
+
+pub(crate) static COUNT_GROUPED_KERNEL: GroupedAggregateKernelAdapter<Count, CountGroupedKernel> =
+    GroupedAggregateKernelAdapter::new(CountGroupedKernel);
 
 #[derive(Debug)]
 pub(crate) struct CountGroupedKernel;
 
-impl DynGroupedAggregateKernel for CountGroupedKernel {
-    fn grouped_aggregate(
+impl GroupedAggregateKernel<Count> for CountGroupedKernel {
+    fn grouped_accumulate(
         &self,
-        aggregate_fn: &AggregateFnRef,
+        _options: &EmptyOptions,
+        states: &mut [u64],
         batch: &ArrayRef,
         group_ids: &GroupIds,
         ctx: &mut ExecutionCtx,
-    ) -> VortexResult<Option<GroupedAggregateKernelResult>> {
-        if aggregate_fn.as_opt::<Count>().is_none() {
-            return Ok(None);
-        }
-
-        let partials = accumulate_grouped(batch, group_ids, ctx)?;
-        Ok(Some(GroupedAggregateKernelResult::dense(
-            PrimitiveArray::from_iter(partials).into_array(),
-            group_ids.num_groups(),
-        )?))
-    }
-}
-
-fn accumulate_grouped(
-    batch: &ArrayRef,
-    group_ids: &GroupIds,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<Vec<u64>> {
-    let ids = group_ids.validated_ids(ctx)?;
-    let mut partials = vec![0u64; group_ids.num_groups()];
-    let validity = batch.validity()?.execute_mask(batch.len(), ctx)?;
-    for (&group_id, valid) in ids.iter().zip(validity.iter()) {
-        if valid {
-            partials[group_id as usize] += 1;
+    ) -> VortexResult<bool> {
+        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(partials)
 }

vortex-array/src/aggregate_fn/fns/count/mod.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 mod grouped;
-pub(crate) use grouped::CountGroupedKernel;
+pub(crate) use grouped::COUNT_GROUPED_KERNEL;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 

vortex-array/src/aggregate_fn/fns/sum/grouped.rs

@@ -5,7 +5,6 @@ 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_panic;
 use vortex_mask::AllOr;
 use vortex_mask::Mask;
@@ -19,74 +18,51 @@ use super::primitive::sum_float_all;
 use super::primitive::sum_signed_all;
 use super::primitive::sum_unsigned_all;
 use crate::ArrayRef;
-use crate::Canonical;
-use crate::Columnar;
 use crate::ExecutionCtx;
-use crate::aggregate_fn::AggregateFnRef;
-use crate::aggregate_fn::AggregateFnVTable;
+use crate::aggregate_fn::EmptyOptions;
 use crate::aggregate_fn::GroupIds;
-use crate::aggregate_fn::kernels::DynGroupedAggregateKernel;
-use crate::aggregate_fn::kernels::GroupedAggregateKernelResult;
+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;
 use crate::match_each_native_ptype;
 
 const MIN_AVG_RUN_LENGTH_FOR_GROUPED_SUM_RUNS: usize = 4;
 
+pub(crate) static SUM_GROUPED_KERNEL: GroupedAggregateKernelAdapter<Sum, SumGroupedKernel> =
+    GroupedAggregateKernelAdapter::new(SumGroupedKernel);
+
 #[derive(Debug)]
 pub(crate) struct SumGroupedKernel;
 
-impl DynGroupedAggregateKernel for SumGroupedKernel {
-    fn grouped_aggregate(
+impl GroupedAggregateKernel<Sum> for SumGroupedKernel {
+    fn grouped_accumulate(
         &self,
-        aggregate_fn: &AggregateFnRef,
+        _options: &EmptyOptions,
+        partials: &mut [SumPartial],
         batch: &ArrayRef,
         group_ids: &GroupIds,
         ctx: &mut ExecutionCtx,
-    ) -> VortexResult<Option<GroupedAggregateKernelResult>> {
-        let Some(options) = aggregate_fn.as_opt::<Sum>() else {
-            return Ok(None);
-        };
-
-        let columnar = batch.clone().execute::<Columnar>(ctx)?;
-        match &columnar {
-            Columnar::Canonical(Canonical::Primitive(_))
-            | Columnar::Canonical(Canonical::Bool(_)) => {}
-            // Decimal and constants still use the universal grouped fallback.
-            Columnar::Canonical(Canonical::Decimal(_)) | Columnar::Constant(_) => return Ok(None),
-            Columnar::Canonical(_) => {
-                vortex_bail!("Unsupported canonical type for sum: {}", columnar.dtype())
-            }
+    ) -> VortexResult<bool> {
+        if let Some(primitive) = batch.as_opt::<Primitive>() {
+            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);
         }
 
-        let partial_dtype = Sum
-            .partial_dtype(options, batch.dtype())
-            .ok_or_else(|| vortex_error::vortex_err!("Unsupported sum dtype: {}", batch.dtype()))?;
-        let ids = group_ids.validated_ids(ctx)?;
-        let mut partials = (0..group_ids.num_groups())
-            .map(|_| Sum.empty_partial(options, batch.dtype()))
-            .collect::<VortexResult<Vec<_>>>()?;
-
-        match &columnar {
-            Columnar::Canonical(Canonical::Primitive(p)) => {
-                accumulate_grouped_primitive(&mut partials, p, ids.as_ref(), ctx)?;
-            }
-            Columnar::Canonical(Canonical::Bool(b)) => {
-                accumulate_grouped_bool(&mut partials, b, ids.as_ref(), ctx)?;
-            }
-            Columnar::Canonical(Canonical::Decimal(_)) | Columnar::Constant(_) => unreachable!(),
-            Columnar::Canonical(_) => unreachable!(),
+        if let Some(bools) = batch.as_opt::<Bool>() {
+            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);
         }
 
-        let Some(partials) = Sum.partials_to_array(&partials, &partial_dtype)? else {
-            return Ok(None);
-        };
-        Ok(Some(GroupedAggregateKernelResult::dense(
-            partials,
-            group_ids.num_groups(),
-        )?))
+        Ok(false)
     }
 }
 

vortex-array/src/aggregate_fn/fns/sum/mod.rs

@@ -7,7 +7,7 @@ mod decimal;
 mod grouped;
 mod primitive;
 
-pub(crate) use grouped::SumGroupedKernel;
+pub(crate) use grouped::SUM_GROUPED_KERNEL;
 use vortex_buffer::Buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;