Compare kernel (#3144)

Migrate the CompareFn to the new kernels

Author: gatesn

bench-vortex/src/bin/notimplemented.rs

@@ -159,7 +159,6 @@ fn compute_funcs(encodings: &[ArrayRef]) {
     table_builder.push_record(vec![
         "Encoding",
         "cast",
-        "compare",
         "fill_forward",
         "fill_null",
         "scalar_at",
@@ -175,7 +174,6 @@ fn compute_funcs(encodings: &[ArrayRef]) {
         let mut impls = vec![id.as_ref()];
         for (j, func) in [
             encoding.cast_fn().is_some(),
-            encoding.compare_fn().is_some(),
             encoding.fill_forward_fn().is_some(),
             encoding.fill_null_fn().is_some(),
             encoding.scalar_at_fn().is_some(),

encodings/alp/src/alp/array.rs

@@ -24,6 +24,7 @@ pub struct ALPArray {
     stats_set: ArrayStats,
 }
 
+#[derive(Debug)]
 pub struct ALPEncoding;
 impl Encoding for ALPEncoding {
     type Array = ALPArray;

encodings/alp/src/alp/compute/compare.rs

@@ -1,8 +1,8 @@
 use std::fmt::Debug;
 
 use vortex_array::arrays::ConstantArray;
-use vortex_array::compute::{CompareFn, Operator, compare};
-use vortex_array::{Array, ArrayRef};
+use vortex_array::compute::{CompareKernel, CompareKernelAdapter, Operator, compare};
+use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_dtype::NativePType;
 use vortex_error::{VortexResult, vortex_bail};
 use vortex_scalar::{PrimitiveScalar, Scalar};
@@ -11,7 +11,7 @@ use crate::{ALPArray, ALPEncoding, ALPFloat, match_each_alp_float_ptype};
 
 // TODO(joe): add fuzzing.
 
-impl CompareFn<&ALPArray> for ALPEncoding {
+impl CompareKernel for ALPEncoding {
     fn compare(
         &self,
         lhs: &ALPArray,
@@ -42,6 +42,8 @@ impl CompareFn<&ALPArray> for ALPEncoding {
     }
 }
 
+register_kernel!(CompareKernelAdapter(ALPEncoding).lift());
+
 // We can compare a scalar to an ALPArray by encoding the scalar into the ALP domain and comparing
 // the encoded value to the encoded values in the ALPArray. There are fixups when the value doesn't
 // encode into the ALP domain.

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

@@ -0,0 +1,23 @@
+use vortex_array::compute::{FilterKernel, FilterKernelAdapter, filter};
+use vortex_array::{Array, ArrayRef, register_kernel};
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use crate::{ALPArray, ALPEncoding};
+
+impl FilterKernel for ALPEncoding {
+    fn filter(&self, array: &ALPArray, mask: &Mask) -> VortexResult<ArrayRef> {
+        let patches = array
+            .patches()
+            .map(|p| p.filter(mask))
+            .transpose()?
+            .flatten();
+
+        Ok(
+            ALPArray::try_new(filter(array.encoded(), mask)?, array.exponents(), patches)?
+                .into_array(),
+        )
+    }
+}
+
+register_kernel!(FilterKernelAdapter(ALPEncoding).lift());

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

@@ -1,18 +1,18 @@
 mod compare;
+mod filter;
 
 use std::fmt::Debug;
 
 use vortex_array::arrays::ConstantArray;
 use vortex_array::compute::{
-    BetweenFn, BetweenOptions, CompareFn, FilterKernelAdapter, FilterKernelImpl, ScalarAtFn,
-    SliceFn, StrictComparison, TakeFn, between, filter, scalar_at, slice, take,
+    BetweenFn, BetweenOptions, ScalarAtFn, SliceFn, StrictComparison, TakeFn, between, scalar_at,
+    slice, take,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::ComputeVTable;
-use vortex_array::{Array, ArrayRef, register_kernel};
+use vortex_array::{Array, ArrayRef};
 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};
@@ -22,10 +22,6 @@ impl ComputeVTable for ALPEncoding {
         Some(self)
     }
 
-    fn compare_fn(&self) -> Option<&dyn CompareFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
         Some(self)
     }
@@ -98,23 +94,6 @@ impl SliceFn<&ALPArray> for ALPEncoding {
     }
 }
 
-impl FilterKernelImpl for ALPEncoding {
-    fn filter(&self, array: &ALPArray, mask: &Mask) -> VortexResult<ArrayRef> {
-        let patches = array
-            .patches()
-            .map(|p| p.filter(mask))
-            .transpose()?
-            .flatten();
-
-        Ok(
-            ALPArray::try_new(filter(array.encoded(), mask)?, array.exponents(), patches)?
-                .into_array(),
-        )
-    }
-}
-
-register_kernel!(FilterKernelAdapter(ALPEncoding).lift());
-
 impl BetweenFn<&ALPArray> for ALPEncoding {
     fn between(
         &self,

encodings/alp/src/alp_rd/array.rs

@@ -28,6 +28,7 @@ pub struct ALPRDArray {
     stats_set: ArrayStats,
 }
 
+#[derive(Debug)]
 pub struct ALPRDEncoding;
 impl Encoding for ALPRDEncoding {
     type Array = ALPRDArray;

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

@@ -1,11 +1,11 @@
-use vortex_array::compute::{FilterKernelAdapter, FilterKernelImpl, filter};
+use vortex_array::compute::{FilterKernel, FilterKernelAdapter, filter};
 use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::{ALPRDArray, ALPRDEncoding};
 
-impl FilterKernelImpl for ALPRDEncoding {
+impl FilterKernel for ALPRDEncoding {
     fn filter(&self, array: &ALPRDArray, mask: &Mask) -> VortexResult<ArrayRef> {
         let left_parts_exceptions = array
             .left_parts_patches()

encodings/bytebool/src/array.rs

@@ -25,6 +25,7 @@ pub struct ByteBoolArray {
 
 try_from_array_ref!(ByteBoolArray);
 
+#[derive(Debug)]
 pub struct ByteBoolEncoding;
 impl Encoding for ByteBoolEncoding {
     type Array = ByteBoolArray;

encodings/datetime-parts/src/array.rs

@@ -25,6 +25,7 @@ pub struct DateTimePartsArray {
     stats_set: ArrayStats,
 }
 
+#[derive(Debug)]
 pub struct DateTimePartsEncoding;
 impl Encoding for DateTimePartsEncoding {
     type Array = DateTimePartsArray;

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

@@ -1,14 +1,16 @@
 use vortex_array::arrays::ConstantArray;
-use vortex_array::compute::{CompareFn, Operator, and, compare, or, try_cast};
-use vortex_array::{Array, ArrayRef};
+use vortex_array::compute::{
+    CompareKernel, CompareKernelAdapter, Operator, and, compare, or, try_cast,
+};
+use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_dtype::DType;
 use vortex_dtype::datetime::TemporalMetadata;
 use vortex_error::{VortexExpect as _, VortexResult};
 
 use crate::array::{DateTimePartsArray, DateTimePartsEncoding};
 use crate::timestamp;
 
-impl CompareFn<&DateTimePartsArray> for DateTimePartsEncoding {
+impl CompareKernel for DateTimePartsEncoding {
     /// Compares two arrays and returns a new boolean array with the result of the comparison.
     /// Or, returns None if comparison is not supported.
     fn compare(
@@ -54,6 +56,8 @@ impl CompareFn<&DateTimePartsArray> for DateTimePartsEncoding {
     }
 }
 
+register_kernel!(CompareKernelAdapter(DateTimePartsEncoding).lift());
+
 fn compare_eq(
     lhs: &DateTimePartsArray,
     ts_parts: &timestamp::TimestampParts,