Aggregates: ChunkedArray kernel (#6837)

First implementation of an aggregate push-down kernel. Requires making a few changes to the API. It also removes checked flag from the sum aggregate function since this is not supported in the old world so no point trying to re-create it yet.

Author: gatesn

vortex-array/public-api.lock

@@ -54,7 +54,7 @@ pub fn vortex_array::aggregate_fn::fns::sum::Sum::fmt(&self, f: &mut core::fmt::
 
 impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::sum::Sum
 
-pub type vortex_array::aggregate_fn::fns::sum::Sum::Options = vortex_array::aggregate_fn::fns::sum::SumOptions
+pub type vortex_array::aggregate_fn::fns::sum::Sum::Options = vortex_array::aggregate_fn::EmptyOptions
 
 pub type vortex_array::aggregate_fn::fns::sum::Sum::Partial = vortex_array::aggregate_fn::fns::sum::SumPartial
 
@@ -64,7 +64,7 @@ pub fn vortex_array::aggregate_fn::fns::sum::Sum::combine_partials(&self, partia
 
 pub fn vortex_array::aggregate_fn::fns::sum::Sum::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
 
-pub fn vortex_array::aggregate_fn::fns::sum::Sum::empty_partial(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
+pub fn vortex_array::aggregate_fn::fns::sum::Sum::empty_partial(&self, _options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
 
 pub fn vortex_array::aggregate_fn::fns::sum::Sum::finalize(&self, partials: vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
@@ -82,33 +82,19 @@ pub fn vortex_array::aggregate_fn::fns::sum::Sum::return_dtype(&self, _options:
 
 pub fn vortex_array::aggregate_fn::fns::sum::Sum::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
 
-pub struct vortex_array::aggregate_fn::fns::sum::SumOptions
-
-impl core::clone::Clone for vortex_array::aggregate_fn::fns::sum::SumOptions
-
-pub fn vortex_array::aggregate_fn::fns::sum::SumOptions::clone(&self) -> vortex_array::aggregate_fn::fns::sum::SumOptions
-
-impl core::cmp::Eq for vortex_array::aggregate_fn::fns::sum::SumOptions
-
-impl core::cmp::PartialEq for vortex_array::aggregate_fn::fns::sum::SumOptions
-
-pub fn vortex_array::aggregate_fn::fns::sum::SumOptions::eq(&self, other: &vortex_array::aggregate_fn::fns::sum::SumOptions) -> bool
-
-impl core::fmt::Debug for vortex_array::aggregate_fn::fns::sum::SumOptions
+pub struct vortex_array::aggregate_fn::fns::sum::SumPartial
 
-pub fn vortex_array::aggregate_fn::fns::sum::SumOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+pub mod vortex_array::aggregate_fn::kernels
 
-impl core::fmt::Display for vortex_array::aggregate_fn::fns::sum::SumOptions
+pub trait vortex_array::aggregate_fn::kernels::DynAggregateKernel: 'static + core::marker::Send + core::marker::Sync + core::fmt::Debug
 
-pub fn vortex_array::aggregate_fn::fns::sum::SumOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+pub fn vortex_array::aggregate_fn::kernels::DynAggregateKernel::aggregate(&self, aggregate_fn: &vortex_array::aggregate_fn::AggregateFnRef, batch: &vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::scalar::Scalar>>
 
-impl core::hash::Hash for vortex_array::aggregate_fn::fns::sum::SumOptions
+pub trait vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel: 'static + core::marker::Send + core::marker::Sync + core::fmt::Debug
 
-pub fn vortex_array::aggregate_fn::fns::sum::SumOptions::hash<__H: core::hash::Hasher>(&self, state: &mut __H)
+pub fn vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel::grouped_aggregate(&self, aggregate_fn: &vortex_array::aggregate_fn::AggregateFnRef, groups: &vortex_array::arrays::ListViewArray) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
 
-impl core::marker::StructuralPartialEq for vortex_array::aggregate_fn::fns::sum::SumOptions
-
-pub struct vortex_array::aggregate_fn::fns::sum::SumPartial
+pub fn vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel::grouped_aggregate_fixed_size(&self, aggregate_fn: &vortex_array::aggregate_fn::AggregateFnRef, groups: &vortex_array::arrays::FixedSizeListArray) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
 
 pub mod vortex_array::aggregate_fn::session
 
@@ -118,11 +104,13 @@ impl vortex_array::aggregate_fn::session::AggregateFnSession
 
 pub fn vortex_array::aggregate_fn::session::AggregateFnSession::register<V: vortex_array::aggregate_fn::AggregateFnVTable>(&self, vtable: V)
 
+pub fn vortex_array::aggregate_fn::session::AggregateFnSession::register_aggregate_kernel(&self, array_id: vortex_array::vtable::ArrayId, agg_fn_id: core::option::Option<vortex_array::aggregate_fn::AggregateFnId>, kernel: &'static dyn vortex_array::aggregate_fn::kernels::DynAggregateKernel)
+
 pub fn vortex_array::aggregate_fn::session::AggregateFnSession::registry(&self) -> &vortex_array::aggregate_fn::session::AggregateFnRegistry
 
 impl core::default::Default for vortex_array::aggregate_fn::session::AggregateFnSession
 
-pub fn vortex_array::aggregate_fn::session::AggregateFnSession::default() -> vortex_array::aggregate_fn::session::AggregateFnSession
+pub fn vortex_array::aggregate_fn::session::AggregateFnSession::default() -> Self
 
 impl core::fmt::Debug for vortex_array::aggregate_fn::session::AggregateFnSession
 
@@ -324,7 +312,7 @@ pub fn vortex_array::aggregate_fn::AggregateFnVTable::serialize(&self, options:
 
 impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::sum::Sum
 
-pub type vortex_array::aggregate_fn::fns::sum::Sum::Options = vortex_array::aggregate_fn::fns::sum::SumOptions
+pub type vortex_array::aggregate_fn::fns::sum::Sum::Options = vortex_array::aggregate_fn::EmptyOptions
 
 pub type vortex_array::aggregate_fn::fns::sum::Sum::Partial = vortex_array::aggregate_fn::fns::sum::SumPartial
 
@@ -334,7 +322,7 @@ pub fn vortex_array::aggregate_fn::fns::sum::Sum::combine_partials(&self, partia
 
 pub fn vortex_array::aggregate_fn::fns::sum::Sum::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
 
-pub fn vortex_array::aggregate_fn::fns::sum::Sum::empty_partial(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
+pub fn vortex_array::aggregate_fn::fns::sum::Sum::empty_partial(&self, _options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
 
 pub fn vortex_array::aggregate_fn::fns::sum::Sum::finalize(&self, partials: vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 

vortex-array/src/aggregate_fn/accumulator.rs

@@ -105,9 +105,17 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
                 break;
             }
 
-            let kernel_key = (self.vtable.id(), batch.encoding_id());
-            if let Some(kernel) = kernels.read().get(&kernel_key)
-                && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch)?
+            let kernels_r = kernels.read();
+            let batch_id = batch.encoding_id();
+            if let Some(result) = kernels_r
+                .get(&(batch_id.clone(), Some(self.aggregate_fn.id())))
+                .or_else(|| kernels_r.get(&(batch_id, None)))
+                .and_then(|kernel| {
+                    kernel
+                        .aggregate(&self.aggregate_fn, &batch, &mut ctx)
+                        .transpose()
+                })
+                .transpose()?
             {
                 vortex_ensure!(
                     result.dtype() == &self.partial_dtype,

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -164,21 +164,27 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
                 break;
             }
 
-            let kernel_key = (self.vtable.id(), elements.encoding_id());
-            if let Some(kernel) = kernels.read().get(&kernel_key) {
-                // 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(),
-                    )
-                };
-
-                if let Some(result) = kernel.grouped_aggregate(&self.aggregate_fn, &groups)? {
-                    return self.push_result(result);
-                }
+            let kernels_r = kernels.read();
+            if let Some(result) = kernels_r
+                .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
+                .or_else(|| kernels_r.get(&(elements.encoding_id(), None)))
+                .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()?
+            {
+                return self.push_result(result);
             }
 
             // Execute one step and try again
@@ -244,23 +250,28 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
                 break;
             }
 
-            let kernel_key = (self.vtable.id(), elements.encoding_id());
-            if let Some(kernel) = kernels.read().get(&kernel_key) {
-                // SAFETY: we assume that elements execution is safe
-                let groups = unsafe {
-                    FixedSizeListArray::new_unchecked(
-                        elements.clone(),
-                        groups.list_size(),
-                        groups.validity().clone(),
-                        groups.len(),
-                    )
-                };
-
-                if let Some(result) =
-                    kernel.grouped_aggregate_fixed_size(&self.aggregate_fn, &groups)?
-                {
-                    return self.push_result(result);
-                }
+            let kernels_r = kernels.read();
+            if let Some(result) = kernels_r
+                .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
+                .or_else(|| kernels_r.get(&(elements.encoding_id(), None)))
+                .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);
             }
 
             // Execute one step and try again