ComputeFn to hold own registry (#3139)

Compute functions now hold their own kernel registry, allowing
third-parties to register kernels.

Author: gatesn

Cargo.lock

@@ -104,6 +104,7 @@ hashbrown = "0.15.1"
 homedir = "0.3.3"
 humansize = "2.1.3"
 indicatif = "0.17.8"
+inventory = "0.3.20"
 itertools = "0.14.0"
 jiff = "0.2.0"
 lending-iterator = "0.1.7"

Cargo.toml

@@ -10,7 +10,6 @@ use vortex::arrays::{
     VarBinViewArray,
 };
 use vortex::buffer::buffer;
-use vortex::compute::Filter;
 use vortex::dtype::datetime::{TIME_ID, TemporalMetadata, TimeUnit};
 use vortex::dtype::{DType, ExtDType, Nullability, PType};
 use vortex::encodings::alp::{ALPArray, Exponents, RDEncoder};
@@ -163,7 +162,6 @@ fn compute_funcs(encodings: &[ArrayRef]) {
         "compare",
         "fill_forward",
         "fill_null",
-        "filter",
         "scalar_at",
         "binary_numeric",
         "search_sorted",
@@ -180,7 +178,6 @@ fn compute_funcs(encodings: &[ArrayRef]) {
             encoding.compare_fn().is_some(),
             encoding.fill_forward_fn().is_some(),
             encoding.fill_null_fn().is_some(),
-            arr.find_kernel(&Filter).is_some(),
             encoding.scalar_at_fn().is_some(),
             encoding.binary_numeric_fn().is_some(),
             encoding.search_sorted_fn().is_some(),

bench-vortex/src/bin/notimplemented.rs

@@ -4,22 +4,19 @@ use std::fmt::Debug;
 
 use vortex_array::arrays::ConstantArray;
 use vortex_array::compute::{
-    BetweenFn, BetweenOptions, CompareFn, FilterKernel, FilterKernelAdapter, KernelRef, ScalarAtFn,
+    BetweenFn, BetweenOptions, CompareFn, FilterKernelAdapter, FilterKernelImpl, ScalarAtFn,
     SliceFn, StrictComparison, TakeFn, between, filter, scalar_at, slice, take,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::ComputeVTable;
-use vortex_array::{Array, ArrayComputeImpl, ArrayRef};
+use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_dtype::{NativePType, Nullability};
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_scalar::{Scalar, ScalarType};
 
 use crate::{ALPArray, ALPEncoding, ALPFloat, match_each_alp_float_ptype};
 
-impl ArrayComputeImpl for ALPArray {
-    const FILTER: Option<KernelRef> = FilterKernelAdapter(ALPEncoding).some();
-}
 impl ComputeVTable for ALPEncoding {
     fn between_fn(&self) -> Option<&dyn BetweenFn<&dyn Array>> {
         Some(self)
@@ -101,7 +98,7 @@ impl SliceFn<&ALPArray> for ALPEncoding {
     }
 }
 
-impl FilterKernel for ALPEncoding {
+impl FilterKernelImpl for ALPEncoding {
     fn filter(&self, array: &ALPArray, mask: &Mask) -> VortexResult<ArrayRef> {
         let patches = array
             .patches()
@@ -116,6 +113,8 @@ impl FilterKernel for ALPEncoding {
     }
 }
 
+register_kernel!(FilterKernelAdapter(ALPEncoding).lift());
+
 impl BetweenFn<&ALPArray> for ALPEncoding {
     fn between(
         &self,

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

@@ -1,11 +1,11 @@
-use vortex_array::compute::{FilterKernel, filter};
-use vortex_array::{Array, ArrayRef};
+use vortex_array::compute::{FilterKernelAdapter, FilterKernelImpl, filter};
+use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::{ALPRDArray, ALPRDEncoding};
 
-impl FilterKernel for ALPRDEncoding {
+impl FilterKernelImpl for ALPRDEncoding {
     fn filter(&self, array: &ALPRDArray, mask: &Mask) -> VortexResult<ArrayRef> {
         let left_parts_exceptions = array
             .left_parts_patches()
@@ -25,6 +25,8 @@ impl FilterKernel for ALPRDEncoding {
     }
 }
 
+register_kernel!(FilterKernelAdapter(ALPRDEncoding).lift());
+
 #[cfg(test)]
 mod test {
     use rstest::rstest;

encodings/alp/src/alp_rd/compute/filter.rs

@@ -1,19 +1,15 @@
-use vortex_array::compute::{FilterKernelAdapter, KernelRef, MaskFn, ScalarAtFn, SliceFn, TakeFn};
+use vortex_array::Array;
+use vortex_array::compute::{MaskFn, ScalarAtFn, SliceFn, TakeFn};
 use vortex_array::vtable::ComputeVTable;
-use vortex_array::{Array, ArrayComputeImpl};
 
-use crate::{ALPRDArray, ALPRDEncoding};
+use crate::ALPRDEncoding;
 
 mod filter;
 mod mask;
 mod scalar_at;
 mod slice;
 mod take;
 
-impl ArrayComputeImpl for ALPRDArray {
-    const FILTER: Option<KernelRef> = FilterKernelAdapter(ALPRDEncoding).some();
-}
-
 impl ComputeVTable for ALPRDEncoding {
     fn mask_fn(&self) -> Option<&dyn MaskFn<&dyn Array>> {
         Some(self)

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

@@ -3,16 +3,14 @@ use vortex_array::compute::{FillForwardFn, MaskFn, ScalarAtFn, SliceFn, TakeFn};
 use vortex_array::validity::Validity;
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::ComputeVTable;
-use vortex_array::{Array, ArrayComputeImpl, ArrayRef, ToCanonical};
+use vortex_array::{Array, ArrayRef, ToCanonical};
 use vortex_dtype::{Nullability, match_each_integer_ptype};
 use vortex_error::{VortexResult, vortex_err};
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
 
 use super::{ByteBoolArray, ByteBoolEncoding};
 
-impl ArrayComputeImpl for ByteBoolArray {}
-
 impl ComputeVTable for ByteBoolEncoding {
     fn fill_forward_fn(&self) -> Option<&dyn FillForwardFn<&dyn Array>> {
         None

encodings/bytebool/src/compute.rs

@@ -1,11 +1,11 @@
-use vortex_array::compute::{FilterKernel, filter};
-use vortex_array::{Array, ArrayRef};
+use vortex_array::compute::{FilterKernelAdapter, FilterKernelImpl, filter};
+use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::{DateTimePartsArray, DateTimePartsEncoding};
 
-impl FilterKernel for DateTimePartsEncoding {
+impl FilterKernelImpl for DateTimePartsEncoding {
     fn filter(&self, array: &DateTimePartsArray, mask: &Mask) -> VortexResult<ArrayRef> {
         Ok(DateTimePartsArray::try_new(
             array.dtype().clone(),
@@ -16,3 +16,4 @@ impl FilterKernel for DateTimePartsEncoding {
         .into_array())
     }
 }
+register_kernel!(FilterKernelAdapter(DateTimePartsEncoding).lift());

encodings/datetime-parts/src/compute/filter.rs

@@ -5,11 +5,10 @@ mod is_constant;
 mod take;
 
 use vortex_array::compute::{
-    CastFn, CompareFn, FilterKernelAdapter, IsConstantFn, KernelRef, ScalarAtFn, SliceFn, TakeFn,
-    scalar_at, slice,
+    CastFn, CompareFn, IsConstantFn, ScalarAtFn, SliceFn, TakeFn, scalar_at, slice,
 };
 use vortex_array::vtable::ComputeVTable;
-use vortex_array::{Array, ArrayComputeImpl, ArrayRef};
+use vortex_array::{Array, ArrayRef};
 use vortex_dtype::Nullability::{NonNullable, Nullable};
 use vortex_dtype::datetime::TemporalMetadata;
 use vortex_dtype::{DType, PType};
@@ -19,10 +18,6 @@ use vortex_scalar::Scalar;
 use crate::timestamp::{self, TimestampParts};
 use crate::{DateTimePartsArray, DateTimePartsEncoding};
 
-impl ArrayComputeImpl for DateTimePartsArray {
-    const FILTER: Option<KernelRef> = FilterKernelAdapter(DateTimePartsEncoding).some();
-}
-
 impl ComputeVTable for DateTimePartsEncoding {
     fn cast_fn(&self) -> Option<&dyn CastFn<&dyn Array>> {
         Some(self)

encodings/datetime-parts/src/compute/mod.rs

@@ -8,22 +8,18 @@ mod min_max;
 mod optimize;
 
 use vortex_array::compute::{
-    BinaryNumericFn, CompareFn, FillNullFn, FilterKernel, FilterKernelAdapter, IsConstantFn,
-    IsSortedFn, KernelRef, LikeFn, MinMaxFn, OptimizeFn, ScalarAtFn, SliceFn, TakeFn, filter,
-    scalar_at, slice, take,
+    BinaryNumericFn, CompareFn, FillNullFn, FilterKernelAdapter, FilterKernelImpl, IsConstantFn,
+    IsSortedFn, LikeFn, MinMaxFn, OptimizeFn, ScalarAtFn, SliceFn, TakeFn, filter, scalar_at,
+    slice, take,
 };
 use vortex_array::vtable::ComputeVTable;
-use vortex_array::{Array, ArrayComputeImpl, ArrayRef};
+use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
 
 use crate::{DictArray, DictEncoding};
 
-impl ArrayComputeImpl for DictArray {
-    const FILTER: Option<KernelRef> = FilterKernelAdapter(DictEncoding).some();
-}
-
 impl ComputeVTable for DictEncoding {
     fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
         Some(self)
@@ -84,13 +80,15 @@ impl TakeFn<&DictArray> for DictEncoding {
     }
 }
 
-impl FilterKernel for DictEncoding {
+impl FilterKernelImpl for DictEncoding {
     fn filter(&self, array: &DictArray, mask: &Mask) -> VortexResult<ArrayRef> {
         let codes = filter(array.codes(), mask)?;
         DictArray::try_new(codes, array.values().clone()).map(|a| a.into_array())
     }
 }
 
+register_kernel!(FilterKernelAdapter(DictEncoding).lift());
+
 impl SliceFn<&DictArray> for DictEncoding {
     // TODO(robert): Add function to trim the dictionary
     fn slice(&self, array: &DictArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {