Port BinaryBooleanFn to BooleanKernel (#3147)

gatesn · web-flow · commit f35381700e8f · 2025-04-28T19:55:32.000+01:00
diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs
@@ -63,7 +63,7 @@ mod primitive {
     use vortex_array::arrays::ConstantArray;
     use vortex_array::compute::StrictComparison::NonStrict;
     use vortex_array::compute::{
-        BetweenOptions, BinaryOperator, Operator, between, binary_boolean, compare,
+        BetweenOptions, BooleanOperator, Operator, between, binary_boolean, compare,
     };
     use vortex_dtype::NativePType;
     use vortex_error::VortexExpect;
@@ -89,7 +89,7 @@ mod primitive {
                 &compare(&arr, &ConstantArray::new(min, arr.len()), Operator::Gte)
                     .vortex_expect(""),
                 &compare(&arr, &ConstantArray::new(max, arr.len()), Operator::Lt).vortex_expect(""),
-                BinaryOperator::And,
+                BooleanOperator::And,
             )
             .vortex_expect("")
         })
@@ -132,7 +132,7 @@ mod bitpack {
     use vortex_array::arrays::ConstantArray;
     use vortex_array::compute::StrictComparison::NonStrict;
     use vortex_array::compute::{
-        BetweenOptions, BinaryOperator, Operator, between, binary_boolean, compare,
+        BetweenOptions, BooleanOperator, Operator, between, binary_boolean, compare,
     };
     use vortex_dtype::NativePType;
     use vortex_error::VortexExpect;
@@ -158,7 +158,7 @@ mod bitpack {
                 &compare(&arr, &ConstantArray::new(min, arr.len()), Operator::Gte)
                     .vortex_expect(""),
                 &compare(&arr, &ConstantArray::new(max, arr.len()), Operator::Lt).vortex_expect(""),
-                BinaryOperator::And,
+                BooleanOperator::And,
             )
         })
     }
@@ -200,7 +200,7 @@ mod alp {
     use vortex_array::arrays::ConstantArray;
     use vortex_array::compute::StrictComparison::NonStrict;
     use vortex_array::compute::{
-        BetweenOptions, BinaryOperator, Operator, between, binary_boolean, compare,
+        BetweenOptions, BooleanOperator, Operator, between, binary_boolean, compare,
     };
     use vortex_dtype::NativePType;
     use vortex_error::VortexExpect;
@@ -226,7 +226,7 @@ mod alp {
                 &compare(&arr, &ConstantArray::new(min, arr.len()), Operator::Gte)
                     .vortex_expect(""),
                 &compare(&arr, &ConstantArray::new(max, arr.len()), Operator::Lt).vortex_expect(""),
-                BinaryOperator::And,
+                BooleanOperator::And,
             )
         })
     }
diff --git a/vortex-array/src/arrays/bool/compute/mod.rs b/vortex-array/src/arrays/bool/compute/mod.rs
@@ -1,8 +1,8 @@
 use crate::Array;
 use crate::arrays::BoolEncoding;
 use crate::compute::{
-    BinaryBooleanFn, CastFn, FillForwardFn, FillNullFn, InvertFn, IsConstantFn, IsSortedFn, MaskFn,
-    MinMaxFn, ScalarAtFn, SliceFn, SumFn, TakeFn, ToArrowFn, UncompressedSizeFn,
+    CastFn, FillForwardFn, FillNullFn, InvertFn, IsConstantFn, IsSortedFn, MaskFn, MinMaxFn,
+    ScalarAtFn, SliceFn, SumFn, TakeFn, ToArrowFn, UncompressedSizeFn,
 };
 use crate::vtable::ComputeVTable;
 
@@ -24,14 +24,6 @@ mod to_arrow;
 mod uncompressed_size;
 
 impl ComputeVTable for BoolEncoding {
-    fn binary_boolean_fn(&self) -> Option<&dyn BinaryBooleanFn<&dyn Array>> {
-        // We only implement this when other is a constant value, otherwise we fall back to the
-        // default implementation that canonicalizes to Arrow.
-        // TODO(ngates): implement this for constants.
-        // other.is_constant().then_some(self)
-        None
-    }
-
     fn cast_fn(&self) -> Option<&dyn CastFn<&dyn Array>> {
         Some(self)
     }
diff --git a/vortex-array/src/arrays/chunked/compute/boolean.rs b/vortex-array/src/arrays/chunked/compute/boolean.rs
@@ -2,15 +2,15 @@ use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
 use crate::arrays::{ChunkedArray, ChunkedEncoding};
-use crate::compute::{BinaryBooleanFn, BinaryOperator, binary_boolean, slice};
-use crate::{Array, ArrayRef};
+use crate::compute::{BooleanKernel, BooleanKernelAdapter, BooleanOperator, binary_boolean, slice};
+use crate::{Array, ArrayRef, register_kernel};
 
-impl BinaryBooleanFn<&ChunkedArray> for ChunkedEncoding {
+impl BooleanKernel for ChunkedEncoding {
     fn binary_boolean(
         &self,
         lhs: &ChunkedArray,
         rhs: &dyn Array,
-        op: BinaryOperator,
+        op: BooleanOperator,
     ) -> VortexResult<Option<ArrayRef>> {
         let mut idx = 0;
         let mut chunks = Vec::with_capacity(lhs.nchunks());
@@ -28,14 +28,16 @@ impl BinaryBooleanFn<&ChunkedArray> for ChunkedEncoding {
     }
 }
 
+register_kernel!(BooleanKernelAdapter(ChunkedEncoding).lift());
+
 #[cfg(test)]
 mod tests {
     use vortex_dtype::{DType, Nullability};
 
     use crate::array::Array;
     use crate::arrays::{BoolArray, BooleanBuffer, ChunkedArray};
     use crate::canonical::ToCanonical;
-    use crate::compute::{BinaryOperator, binary_boolean};
+    use crate::compute::{BooleanOperator, binary_boolean};
 
     #[test]
     fn test_bin_bool_chunked() {
@@ -49,7 +51,7 @@ mod tests {
         let chunked2 =
             ChunkedArray::try_new(vec![arr2, arr3], DType::Bool(Nullability::Nullable)).unwrap();
 
-        let result = binary_boolean(&chunked1, &chunked2, BinaryOperator::Or)
+        let result = binary_boolean(&chunked1, &chunked2, BooleanOperator::Or)
             .unwrap()
             .to_bool()
             .unwrap();
diff --git a/vortex-array/src/arrays/chunked/compute/mod.rs b/vortex-array/src/arrays/chunked/compute/mod.rs
@@ -4,8 +4,8 @@ use vortex_error::VortexResult;
 use crate::arrays::ChunkedEncoding;
 use crate::arrays::chunked::ChunkedArray;
 use crate::compute::{
-    BinaryBooleanFn, BinaryNumericFn, CastFn, FillNullFn, InvertFn, IsConstantFn, IsSortedFn,
-    MaskFn, MinMaxFn, ScalarAtFn, SliceFn, SumFn, TakeFn, UncompressedSizeFn, try_cast,
+    BinaryNumericFn, CastFn, FillNullFn, InvertFn, IsConstantFn, IsSortedFn, MaskFn, MinMaxFn,
+    ScalarAtFn, SliceFn, SumFn, TakeFn, UncompressedSizeFn, try_cast,
 };
 use crate::vtable::ComputeVTable;
 use crate::{Array, ArrayRef};
@@ -27,10 +27,6 @@ mod take;
 mod uncompressed_size;
 
 impl ComputeVTable for ChunkedEncoding {
-    fn binary_boolean_fn(&self) -> Option<&dyn BinaryBooleanFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
         Some(self)
     }
diff --git a/vortex-array/src/arrays/constant/compute/boolean.rs b/vortex-array/src/arrays/constant/compute/boolean.rs
@@ -3,15 +3,15 @@ use vortex_error::{VortexResult, vortex_bail, vortex_err};
 use vortex_scalar::Scalar;
 
 use crate::arrays::{ConstantArray, ConstantEncoding};
-use crate::compute::{BinaryBooleanFn, BinaryOperator};
-use crate::{Array, ArrayRef};
+use crate::compute::{BooleanKernel, BooleanKernelAdapter, BooleanOperator};
+use crate::{Array, ArrayRef, register_kernel};
 
-impl BinaryBooleanFn<&ConstantArray> for ConstantEncoding {
+impl BooleanKernel for ConstantEncoding {
     fn binary_boolean(
         &self,
         lhs: &ConstantArray,
         rhs: &dyn Array,
-        op: BinaryOperator,
+        op: BooleanOperator,
     ) -> VortexResult<Option<ArrayRef>> {
         // We only implement this for constant <-> constant arrays, otherwise we allow fall back
         // to the Arrow implementation.
@@ -31,10 +31,10 @@ impl BinaryBooleanFn<&ConstantArray> for ConstantEncoding {
             .value();
 
         let result = match op {
-            BinaryOperator::And => and(lhs, rhs),
-            BinaryOperator::AndKleene => kleene_and(lhs, rhs),
-            BinaryOperator::Or => or(lhs, rhs),
-            BinaryOperator::OrKleene => kleene_or(lhs, rhs),
+            BooleanOperator::And => and(lhs, rhs),
+            BooleanOperator::AndKleene => kleene_and(lhs, rhs),
+            BooleanOperator::Or => or(lhs, rhs),
+            BooleanOperator::OrKleene => kleene_or(lhs, rhs),
         };
 
         let scalar = result
@@ -45,6 +45,8 @@ impl BinaryBooleanFn<&ConstantArray> for ConstantEncoding {
     }
 }
 
+register_kernel!(BooleanKernelAdapter(ConstantEncoding).lift());
+
 fn and(left: Option<bool>, right: Option<bool>) -> Option<bool> {
     left.zip(right).map(|(l, r)| l & r)
 }
diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs
@@ -15,18 +15,14 @@ use vortex_scalar::{FromPrimitiveOrF16, PrimitiveScalar, Scalar};
 use crate::arrays::ConstantEncoding;
 use crate::arrays::constant::ConstantArray;
 use crate::compute::{
-    BinaryBooleanFn, BinaryNumericFn, CastFn, InvertFn, ScalarAtFn, SearchSortedFn, SliceFn, SumFn,
-    TakeFn, UncompressedSizeFn,
+    BinaryNumericFn, CastFn, InvertFn, ScalarAtFn, SearchSortedFn, SliceFn, SumFn, TakeFn,
+    UncompressedSizeFn,
 };
 use crate::stats::Stat;
 use crate::vtable::ComputeVTable;
 use crate::{Array, ArrayRef};
 
 impl ComputeVTable for ConstantEncoding {
-    fn binary_boolean_fn(&self) -> Option<&dyn BinaryBooleanFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
         Some(self)
     }
diff --git a/vortex-array/src/compute/between.rs b/vortex-array/src/compute/between.rs
@@ -8,7 +8,7 @@ use vortex_scalar::Scalar;
 use crate::arcref::ArcRef;
 use crate::arrays::ConstantArray;
 use crate::compute::{
-    BinaryOperator, ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Options, Output,
+    BooleanOperator, ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Options, Output,
     binary_boolean, compare,
 };
 use crate::{Array, ArrayRef, Canonical, Encoding, IntoArray};
@@ -19,7 +19,7 @@ use crate::{Array, ArrayRef, Canonical, Encoding, IntoArray};
 /// This semantics is equivalent to:
 /// ```
 /// use vortex_array::{Array, ArrayRef};
-/// use vortex_array::compute::{binary_boolean, compare, BetweenOptions, BinaryOperator, Operator};///
+/// use vortex_array::compute::{binary_boolean, compare, BetweenOptions, BooleanOperator, Operator};///
 /// use vortex_error::VortexResult;
 ///
 /// fn between(
@@ -31,7 +31,7 @@ use crate::{Array, ArrayRef, Canonical, Encoding, IntoArray};
 ///     binary_boolean(
 ///         &compare(lower, arr, options.lower_strict.to_operator())?,
 ///         &compare(arr, upper,  options.upper_strict.to_operator())?,
-///         BinaryOperator::And
+///         BooleanOperator::And
 ///     )
 /// }
 ///  ```
@@ -142,7 +142,7 @@ impl ComputeFnVTable for Between {
         Ok(binary_boolean(
             &compare(lower, array, options.lower_strict.to_operator())?,
             &compare(array, upper, options.upper_strict.to_operator())?,
-            BinaryOperator::And,
+            BooleanOperator::And,
         )?
         .into())
     }
diff --git a/vortex-array/src/compute/boolean.rs b/vortex-array/src/compute/boolean.rs
diff --git a/vortex-array/src/compute/mod.rs b/vortex-array/src/compute/mod.rs
diff --git a/vortex-array/src/vtable/compute.rs b/vortex-array/src/vtable/compute.rs
