NanCount AggregateFn (#7003)

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

encodings/alp/public-api.lock

@@ -24,10 +24,6 @@ impl vortex_array::arrays::slice::SliceKernel for vortex_alp::ALP
 
 pub fn vortex_alp::ALP::slice(array: &Self::Array, range: core::ops::range::Range<usize>, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
-impl vortex_array::compute::nan_count::NaNCountKernel for vortex_alp::ALP
-
-pub fn vortex_alp::ALP::nan_count(&self, array: &vortex_alp::ALPArray) -> vortex_error::VortexResult<usize>
-
 impl vortex_array::scalar_fn::fns::between::kernel::BetweenReduce for vortex_alp::ALP
 
 pub fn vortex_alp::ALP::between(array: &vortex_alp::ALPArray, lower: &vortex_array::array::ArrayRef, upper: &vortex_array::array::ArrayRef, options: &vortex_array::scalar_fn::fns::between::BetweenOptions) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
@@ -559,3 +555,5 @@ pub fn vortex_alp::alp_encode(parray: &vortex_array::arrays::primitive::array::P
 pub fn vortex_alp::alp_rd_decode<T: vortex_alp::ALPRDFloat>(left_parts: vortex_buffer::buffer::Buffer<u16>, left_parts_dict: &[u16], right_bit_width: u8, right_parts: vortex_buffer::buffer_mut::BufferMut<<T as vortex_alp::ALPRDFloat>::UINT>, left_parts_patches: core::option::Option<&vortex_array::patches::Patches>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_buffer::buffer::Buffer<T>>
 
 pub fn vortex_alp::decompress_into_array(array: vortex_alp::ALPArray, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrays::primitive::array::PrimitiveArray>
+
+pub fn vortex_alp::initialize(session: &mut vortex_session::VortexSession)

encodings/alp/src/alp/compute/mod.rs

@@ -6,7 +6,7 @@ mod cast;
 mod compare;
 mod filter;
 mod mask;
-mod nan_count;
+pub(crate) mod nan_count;
 mod slice;
 mod take;
 

encodings/alp/src/alp/compute/nan_count.rs

@@ -1,24 +1,45 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_array::compute::NaNCountKernel;
-use vortex_array::compute::NaNCountKernelAdapter;
-use vortex_array::compute::nan_count;
-use vortex_array::register_kernel;
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::aggregate_fn::AggregateFnRef;
+use vortex_array::aggregate_fn::fns::nan_count::NanCount;
+use vortex_array::aggregate_fn::fns::nan_count::nan_count;
+use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
+use vortex_array::scalar::Scalar;
 use vortex_error::VortexResult;
 
 use crate::ALP;
-use crate::ALPArray;
 
-impl NaNCountKernel for ALP {
-    fn nan_count(&self, array: &ALPArray) -> VortexResult<usize> {
-        // NANs can only be in patches
-        if let Some(patches) = array.patches() {
-            nan_count(patches.values())
-        } else {
-            Ok(0)
+/// ALP-specific NaN count kernel.
+///
+/// NaN values can only appear in the patches array of an ALP-encoded array, since the encoded
+/// integer values cannot represent NaN. This avoids decoding the entire array.
+#[derive(Debug)]
+pub(crate) struct ALPNanCountKernel;
+
+impl DynAggregateKernel for ALPNanCountKernel {
+    fn aggregate(
+        &self,
+        aggregate_fn: &AggregateFnRef,
+        batch: &ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Scalar>> {
+        if !aggregate_fn.is::<NanCount>() {
+            return Ok(None);
         }
+
+        let Some(alp) = batch.as_opt::<ALP>() else {
+            return Ok(None);
+        };
+
+        let count = if let Some(patches) = alp.patches() {
+            nan_count(patches.values(), ctx)?
+        } else {
+            0
+        };
+
+        Ok(Some(Scalar::from(count as u64)))
     }
 }
-
-register_kernel!(NaNCountKernelAdapter(ALP).lift());

encodings/alp/src/alp/mod.rs

@@ -15,7 +15,7 @@ use num_traits::ToPrimitive;
 
 mod array;
 mod compress;
-mod compute;
+pub(crate) mod compute;
 mod decompress;
 mod ops;
 mod rules;

encodings/alp/src/lib.rs

@@ -18,6 +18,24 @@
 
 pub use alp::*;
 pub use alp_rd::*;
+use vortex_array::aggregate_fn::AggregateFnVTable;
+use vortex_array::aggregate_fn::fns::nan_count::NanCount;
+use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
+use vortex_array::session::ArraySessionExt;
+use vortex_session::VortexSession;
 
 mod alp;
 mod alp_rd;
+
+/// Initialize ALP encoding in the given session.
+pub fn initialize(session: &mut VortexSession) {
+    session.arrays().register(ALP::ID, ALP);
+    session.arrays().register(ALPRD::ID, ALPRD);
+
+    // Register the ALP-specific NaN count aggregate kernel.
+    session.aggregate_fns().register_aggregate_kernel(
+        ALP::ID,
+        Some(NanCount.id()),
+        &compute::nan_count::ALPNanCountKernel,
+    );
+}

vortex-array/public-api.lock

@@ -30,6 +30,52 @@ pub mod vortex_array::aggregate_fn
 
 pub mod vortex_array::aggregate_fn::fns
 
+pub mod vortex_array::aggregate_fn::fns::nan_count
+
+pub struct vortex_array::aggregate_fn::fns::nan_count::NanCount
+
+impl core::clone::Clone for vortex_array::aggregate_fn::fns::nan_count::NanCount
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::clone(&self) -> vortex_array::aggregate_fn::fns::nan_count::NanCount
+
+impl core::fmt::Debug for vortex_array::aggregate_fn::fns::nan_count::NanCount
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::nan_count::NanCount
+
+pub type vortex_array::aggregate_fn::fns::nan_count::NanCount::Options = vortex_array::aggregate_fn::EmptyOptions
+
+pub type vortex_array::aggregate_fn::fns::nan_count::NanCount::Partial = u64
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::accumulate(&self, partial: &mut Self::Partial, batch: &vortex_array::Columnar, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::coerce_args(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::combine_partials(&self, partial: &mut Self::Partial, other: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::empty_partial(&self, _options: &Self::Options, _input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::finalize(&self, partials: vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::finalize_scalar(&self, partial: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::flush(&self, partial: &mut Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::is_saturated(&self, _partial: &Self::Partial) -> bool
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::partial_dtype(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::return_dtype(&self, _options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::nan_count(array: &vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<usize>
+
 pub mod vortex_array::aggregate_fn::fns::sum
 
 pub enum vortex_array::aggregate_fn::fns::sum::SumState
@@ -318,6 +364,38 @@ pub fn vortex_array::aggregate_fn::AggregateFnVTable::return_dtype(&self, option
 
 pub fn vortex_array::aggregate_fn::AggregateFnVTable::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
 
+impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::nan_count::NanCount
+
+pub type vortex_array::aggregate_fn::fns::nan_count::NanCount::Options = vortex_array::aggregate_fn::EmptyOptions
+
+pub type vortex_array::aggregate_fn::fns::nan_count::NanCount::Partial = u64
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::accumulate(&self, partial: &mut Self::Partial, batch: &vortex_array::Columnar, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::coerce_args(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::combine_partials(&self, partial: &mut Self::Partial, other: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::empty_partial(&self, _options: &Self::Options, _input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::finalize(&self, partials: vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::finalize_scalar(&self, partial: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::flush(&self, partial: &mut Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::is_saturated(&self, _partial: &Self::Partial) -> bool
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::partial_dtype(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::return_dtype(&self, _options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::nan_count::NanCount::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
 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::EmptyOptions
@@ -2954,10 +3032,6 @@ impl vortex_array::compute::MinMaxKernel for vortex_array::arrays::Primitive
 
 pub fn vortex_array::arrays::Primitive::min_max(&self, array: &vortex_array::arrays::PrimitiveArray) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::MinMaxResult>>
 
-impl vortex_array::compute::NaNCountKernel for vortex_array::arrays::Primitive
-
-pub fn vortex_array::arrays::Primitive::nan_count(&self, array: &vortex_array::arrays::PrimitiveArray) -> vortex_error::VortexResult<usize>
-
 impl vortex_array::optimizer::rules::ArrayParentReduceRule<vortex_array::arrays::Primitive> for vortex_array::arrays::primitive::PrimitiveMaskedValidityRule
 
 pub type vortex_array::arrays::primitive::PrimitiveMaskedValidityRule::Parent = vortex_array::arrays::Masked
@@ -6542,10 +6616,6 @@ impl vortex_array::compute::MinMaxKernel for vortex_array::arrays::Primitive
 
 pub fn vortex_array::arrays::Primitive::min_max(&self, array: &vortex_array::arrays::PrimitiveArray) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::MinMaxResult>>
 
-impl vortex_array::compute::NaNCountKernel for vortex_array::arrays::Primitive
-
-pub fn vortex_array::arrays::Primitive::nan_count(&self, array: &vortex_array::arrays::PrimitiveArray) -> vortex_error::VortexResult<usize>
-
 impl vortex_array::optimizer::rules::ArrayParentReduceRule<vortex_array::arrays::Primitive> for vortex_array::arrays::primitive::PrimitiveMaskedValidityRule
 
 pub type vortex_array::arrays::primitive::PrimitiveMaskedValidityRule::Parent = vortex_array::arrays::Masked
@@ -9634,24 +9704,6 @@ pub fn vortex_array::compute::MinMaxResult::fmt(&self, f: &mut core::fmt::Format
 
 impl core::marker::StructuralPartialEq for vortex_array::compute::MinMaxResult
 
-pub struct vortex_array::compute::NaNCountKernelAdapter<V: vortex_array::vtable::VTable>(pub V)
-
-impl<V: vortex_array::vtable::VTable + vortex_array::compute::NaNCountKernel> vortex_array::compute::NaNCountKernelAdapter<V>
-
-pub const fn vortex_array::compute::NaNCountKernelAdapter<V>::lift(&'static self) -> vortex_array::compute::NaNCountKernelRef
-
-impl<V: core::fmt::Debug + vortex_array::vtable::VTable> core::fmt::Debug for vortex_array::compute::NaNCountKernelAdapter<V>
-
-pub fn vortex_array::compute::NaNCountKernelAdapter<V>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
-
-impl<V: vortex_array::vtable::VTable + vortex_array::compute::NaNCountKernel> vortex_array::compute::Kernel for vortex_array::compute::NaNCountKernelAdapter<V>
-
-pub fn vortex_array::compute::NaNCountKernelAdapter<V>::invoke(&self, args: &vortex_array::compute::InvocationArgs<'_>) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::Output>>
-
-pub struct vortex_array::compute::NaNCountKernelRef(_)
-
-impl inventory::Collect for vortex_array::compute::NaNCountKernelRef
-
 pub struct vortex_array::compute::UnaryArgs<'a, O: vortex_array::compute::Options>
 
 pub vortex_array::compute::UnaryArgs::array: &'a dyn vortex_array::DynArray
@@ -9832,10 +9884,6 @@ impl<V: vortex_array::vtable::VTable + vortex_array::compute::MinMaxKernel> vort
 
 pub fn vortex_array::compute::MinMaxKernelAdapter<V>::invoke(&self, args: &vortex_array::compute::InvocationArgs<'_>) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::Output>>
 
-impl<V: vortex_array::vtable::VTable + vortex_array::compute::NaNCountKernel> vortex_array::compute::Kernel for vortex_array::compute::NaNCountKernelAdapter<V>
-
-pub fn vortex_array::compute::NaNCountKernelAdapter<V>::invoke(&self, args: &vortex_array::compute::InvocationArgs<'_>) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::Output>>
-
 pub trait vortex_array::compute::MinMaxKernel: vortex_array::vtable::VTable
 
 pub fn vortex_array::compute::MinMaxKernel::min_max(&self, array: &Self::Array) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::MinMaxResult>>
@@ -9896,14 +9944,6 @@ impl vortex_array::compute::MinMaxKernel for vortex_array::arrays::null::Null
 
 pub fn vortex_array::arrays::null::Null::min_max(&self, _array: &vortex_array::arrays::null::NullArray) -> vortex_error::VortexResult<core::option::Option<vortex_array::compute::MinMaxResult>>
 
-pub trait vortex_array::compute::NaNCountKernel: vortex_array::vtable::VTable
-
-pub fn vortex_array::compute::NaNCountKernel::nan_count(&self, array: &Self::Array) -> vortex_error::VortexResult<usize>
-
-impl vortex_array::compute::NaNCountKernel for vortex_array::arrays::Primitive
-
-pub fn vortex_array::arrays::Primitive::nan_count(&self, array: &vortex_array::arrays::PrimitiveArray) -> vortex_error::VortexResult<usize>
-
 pub trait vortex_array::compute::Options: 'static
 
 pub fn vortex_array::compute::Options::as_any(&self) -> &dyn core::any::Any

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

@@ -1,4 +1,5 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+pub mod nan_count;
 pub mod sum;