feat[array]: refactor zip compute kernel to lazy zip via ScalarFnArray

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-array/src/arrays/chunked/compute/rules.rs

@@ -13,6 +13,7 @@ use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
 use crate::arrays::ScalarFnArray;
 use crate::compute::FillNullReduceAdaptor;
+use crate::compute::ZipReduceAdaptor;
 use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ParentRuleSet;
@@ -21,6 +22,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet<ChunkedVTable> = ParentRuleSet::new
     ParentRuleSet::lift(&ChunkedUnaryScalarFnPushDownRule),
     ParentRuleSet::lift(&ChunkedConstantScalarFnPushDownRule),
     ParentRuleSet::lift(&FillNullReduceAdaptor(ChunkedVTable)),
+    ParentRuleSet::lift(&ZipReduceAdaptor(ChunkedVTable)),
 ]);
 
 /// Push down any unary scalar function through chunked arrays.

vortex-array/src/arrays/chunked/compute/zip.rs

@@ -8,17 +8,14 @@ use crate::Array;
 use crate::ArrayRef;
 use crate::arrays::ChunkedArray;
 use crate::arrays::ChunkedVTable;
-use crate::compute::ZipKernel;
-use crate::compute::ZipKernelAdapter;
+use crate::compute::ZipReduce;
 use crate::compute::zip;
-use crate::register_kernel;
 
-// Push down the zip call to the chunks. Without this kernel
+// Push down the zip call to the chunks. Without this rule
 // the default implementation canonicalises the chunked array
 // then zips once.
-impl ZipKernel for ChunkedVTable {
+impl ZipReduce for ChunkedVTable {
     fn zip(
-        &self,
         if_true: &ChunkedArray,
         if_false: &dyn Array,
         mask: &Mask,
@@ -72,8 +69,6 @@ impl ZipKernel for ChunkedVTable {
     }
 }
 
-register_kernel!(ZipKernelAdapter(ChunkedVTable).lift());
-
 #[cfg(test)]
 mod tests {
     use vortex_buffer::buffer;

vortex-array/src/arrays/struct_/compute/zip.rs

@@ -10,21 +10,20 @@ use vortex_mask::Mask;
 
 use crate::Array;
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::arrays::StructArray;
 use crate::arrays::StructVTable;
 use crate::compute::ZipKernel;
-use crate::compute::ZipKernelAdapter;
 use crate::compute::zip;
-use crate::register_kernel;
 use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 
 impl ZipKernel for StructVTable {
     fn zip(
-        &self,
         if_true: &StructArray,
         if_false: &dyn Array,
         mask: &Mask,
+        _ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         let Some(if_false) = if_false.as_opt::<StructVTable>() else {
             return Ok(None);
@@ -66,8 +65,6 @@ impl ZipKernel for StructVTable {
     }
 }
 
-register_kernel!(ZipKernelAdapter(StructVTable).lift());
-
 #[cfg(test)]
 mod tests {
     use vortex_dtype::FieldNames;

vortex-array/src/arrays/struct_/vtable/kernel.rs

@@ -3,7 +3,10 @@
 
 use crate::arrays::StructVTable;
 use crate::arrays::TakeExecuteAdaptor;
+use crate::compute::ZipExecuteAdaptor;
 use crate::kernel::ParentKernelSet;
 
-pub(super) const PARENT_KERNELS: ParentKernelSet<StructVTable> =
-    ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(StructVTable))]);
+pub(super) const PARENT_KERNELS: ParentKernelSet<StructVTable> = ParentKernelSet::new(&[
+    ParentKernelSet::lift(&TakeExecuteAdaptor(StructVTable)),
+    ParentKernelSet::lift(&ZipExecuteAdaptor(StructVTable)),
+]);

vortex-array/src/arrays/varbinview/compute/zip.rs

@@ -12,22 +12,21 @@ use vortex_vector::binaryview::BinaryView;
 
 use crate::Array;
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::VarBinViewVTable;
 use crate::builders::DeduplicatedBuffers;
 use crate::builders::LazyBitBufferBuilder;
 use crate::compute::ZipKernel;
-use crate::compute::ZipKernelAdapter;
-use crate::register_kernel;
 
 // A dedicated VarBinView zip kernel that builds the result directly by adjusting views and validity,
 // instead of routing through the generic builder (which would redo buffer lookups per mask slice).
 impl ZipKernel for VarBinViewVTable {
     fn zip(
-        &self,
         if_true: &VarBinViewArray,
         if_false: &dyn Array,
         mask: &Mask,
+        _ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         let Some(if_false) = if_false.as_opt::<VarBinViewVTable>() else {
             return Ok(None);
@@ -37,7 +36,6 @@ impl ZipKernel for VarBinViewVTable {
             vortex_bail!("input arrays to zip must have the same dtype");
         }
 
-        // compute fn already asserts if_true.len() == if_false.len()
         let len = if_true.len();
         let dtype = if_true
             .dtype()
@@ -205,8 +203,6 @@ fn push_view(
     validity_builder.append_non_null();
 }
 
-register_kernel!(ZipKernelAdapter(VarBinViewVTable).lift());
-
 #[cfg(test)]
 mod tests {
     use vortex_dtype::DType;

vortex-array/src/arrays/varbinview/vtable/kernel.rs

@@ -3,7 +3,10 @@
 
 use crate::arrays::TakeExecuteAdaptor;
 use crate::arrays::VarBinViewVTable;
+use crate::compute::ZipExecuteAdaptor;
 use crate::kernel::ParentKernelSet;
 
-pub(super) const PARENT_KERNELS: ParentKernelSet<VarBinViewVTable> =
-    ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(VarBinViewVTable))]);
+pub(super) const PARENT_KERNELS: ParentKernelSet<VarBinViewVTable> = ParentKernelSet::new(&[
+    ParentKernelSet::lift(&TakeExecuteAdaptor(VarBinViewVTable)),
+    ParentKernelSet::lift(&ZipExecuteAdaptor(VarBinViewVTable)),
+]);

vortex-array/src/builtins.rs

@@ -13,9 +13,11 @@ use vortex_dtype::DType;
 use vortex_dtype::FieldName;
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
+use vortex_session::VortexSession;
 
 use crate::Array;
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::ScalarFnArrayExt;
@@ -28,6 +30,7 @@ use crate::expr::IsNull;
 use crate::expr::Mask;
 use crate::expr::Not;
 use crate::expr::VTableExt;
+use crate::expr::Zip;
 use crate::optimizer::ArrayOptimizer;
 
 /// A collection of built-in scalar functions that can be applied to expressions or arrays.
@@ -51,6 +54,9 @@ pub trait ExprBuiltins: Sized {
 
     /// Boolean negation.
     fn not(&self) -> VortexResult<Expression>;
+
+    /// Conditional selection: `result[i] = if mask[i] then self[i] else if_false[i]`.
+    fn zip(&self, if_false: Expression, mask: Expression) -> VortexResult<Expression>;
 }
 
 impl ExprBuiltins for Expression {
@@ -77,6 +83,10 @@ impl ExprBuiltins for Expression {
     fn not(&self) -> VortexResult<Expression> {
         Not.try_new_expr(EmptyOptions, [self.clone()])
     }
+
+    fn zip(&self, if_false: Expression, mask: Expression) -> VortexResult<Expression> {
+        Zip.try_new_expr(EmptyOptions, [self.clone(), if_false, mask])
+    }
 }
 
 pub trait ArrayBuiltins: Sized {
@@ -99,6 +109,9 @@ pub trait ArrayBuiltins: Sized {
 
     /// Boolean negation.
     fn not(&self) -> VortexResult<ArrayRef>;
+
+    /// Conditional selection: `result[i] = if mask[i] then self[i] else if_false[i]`.
+    fn zip(&self, if_false: ArrayRef, mask: ArrayRef) -> VortexResult<ArrayRef>;
 }
 
 impl ArrayBuiltins for ArrayRef {
@@ -141,4 +154,11 @@ impl ArrayBuiltins for ArrayRef {
         Not.try_new_array(self.len(), EmptyOptions, [self.clone()])?
             .optimize()
     }
+
+    fn zip(&self, if_false: ArrayRef, mask: ArrayRef) -> VortexResult<ArrayRef> {
+        let scalar_fn =
+            Zip.try_new_array(self.len(), EmptyOptions, [self.clone(), if_false, mask])?;
+        let mut ctx = ExecutionCtx::new(VortexSession::empty());
+        scalar_fn.execute::<ArrayRef>(&mut ctx)
+    }
 }

vortex-array/src/compute/mod.rs

@@ -49,6 +49,10 @@ pub use crate::expr::FillNullExecuteAdaptor;
 pub use crate::expr::FillNullKernel;
 pub use crate::expr::FillNullReduce;
 pub use crate::expr::FillNullReduceAdaptor;
+pub use crate::expr::ZipExecuteAdaptor;
+pub use crate::expr::ZipKernel;
+pub use crate::expr::ZipReduce;
+pub use crate::expr::ZipReduceAdaptor;
 
 #[cfg(feature = "arbitrary")]
 mod arbitrary;
@@ -99,7 +103,6 @@ pub fn warm_up_vtables() {
     nan_count::warm_up_vtable();
     numeric::warm_up_vtable();
     sum::warm_up_vtable();
-    zip::warm_up_vtable();
 }
 
 impl ComputeFn {