Batch execution API

encodings/sequence/src/operator.rs

@@ -188,9 +188,9 @@ mod tests {
             .unwrap()
             .into_array();
 
-        let selection = bitbuffer![1 0 1 0 1].into_array();
+        let selection = bitbuffer![1 0 1 0 1].into();
         let result = seq
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_primitive()
             .into_i32();
@@ -208,9 +208,9 @@ mod tests {
             .unwrap()
             .into_array();
 
-        let selection = bitbuffer![1 1 0 0 0].into_array();
+        let selection = bitbuffer![1 1 0 0 0].into();
         let result = seq
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_primitive()
             .into_i64();
@@ -225,9 +225,9 @@ mod tests {
             .unwrap()
             .into_array();
 
-        let selection = bitbuffer![0 0 1 1].into_array();
+        let selection = bitbuffer![0 0 1 1].into();
         let result = seq
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_primitive()
             .into_u64();
@@ -245,8 +245,8 @@ mod tests {
             .unwrap()
             .into_array();
 
-        let selection = bitbuffer![0 0 0 0].into_array();
-        let result = seq.execute_with_selection(Some(&selection)).unwrap();
+        let selection = bitbuffer![0 0 0 0].into();
+        let result = seq.execute_with_selection(&selection).unwrap();
         assert!(result.is_empty())
     }
 
@@ -257,9 +257,9 @@ mod tests {
             .unwrap()
             .into_array();
 
-        let selection = bitbuffer![1 1 1 1].into_array();
+        let selection = bitbuffer![1 1 1 1].into();
         let result = seq
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_primitive()
             .into_i16();
@@ -277,9 +277,9 @@ mod tests {
             .unwrap()
             .into_array();
 
-        let selection = bitbuffer![1 0 0 1 0 1].into_array();
+        let selection = bitbuffer![1 0 0 1 0 1].into();
         let result = seq
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_primitive()
             .into_i32();

vortex-array/src/array/operator.rs

@@ -3,11 +3,11 @@
 
 use std::sync::Arc;
 
-use vortex_dtype::DType;
-use vortex_error::{VortexResult, vortex_bail};
-use vortex_vector::Vector;
+use vortex_error::{VortexResult, vortex_panic};
+use vortex_mask::Mask;
+use vortex_vector::{Vector, VectorOps, vector_matches_dtype};
 
-use crate::execution::{BatchKernelRef, BindCtx};
+use crate::execution::{BatchKernelRef, BindCtx, DummyExecutionCtx, ExecutionCtx};
 use crate::vtable::{OperatorVTable, VTable};
 use crate::{Array, ArrayAdapter, ArrayRef};
 
@@ -16,13 +16,13 @@ use crate::{Array, ArrayAdapter, ArrayRef};
 /// Note: the public functions such as "execute" should move onto the main `Array` trait when
 /// operators is stabilized. The other functions should remain on a `pub(crate)` trait.
 pub trait ArrayOperator: 'static + Send + Sync {
-    /// Execute the array producing a canonical vector.
-    fn execute(&self) -> VortexResult<Vector> {
-        self.execute_with_selection(None)
-    }
-
-    /// Execute the array with a selection mask, producing a canonical vector.
-    fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector>;
+    /// Execute the array's batch kernel with the given selection mask.
+    ///
+    /// # Panics
+    ///
+    /// If the mask length does not match the array length.
+    /// If the array's implementation returns an invalid vector (wrong length, wrong type, etc).
+    fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector>;
 
     /// Optimize the array by running the optimization rules.
     fn reduce_children(&self) -> VortexResult<Option<ArrayRef>>;
@@ -39,8 +39,8 @@ pub trait ArrayOperator: 'static + Send + Sync {
 }
 
 impl ArrayOperator for Arc<dyn Array> {
-    fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
-        self.as_ref().execute_with_selection(selection)
+    fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        self.as_ref().execute_batch(selection, ctx)
     }
 
     fn reduce_children(&self) -> VortexResult<Option<ArrayRef>> {
@@ -61,23 +61,31 @@ impl ArrayOperator for Arc<dyn Array> {
 }
 
 impl<V: VTable> ArrayOperator for ArrayAdapter<V> {
-    fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
-        if let Some(selection) = selection.as_ref() {
-            if !matches!(selection.dtype(), DType::Bool(_)) {
-                vortex_bail!(
-                    "Selection array must be of boolean type, got {}",
-                    selection.dtype()
-                );
-            }
-            if selection.len() != self.len() {
-                vortex_bail!(
-                    "Selection array length {} does not match array length {}",
-                    selection.len(),
-                    self.len()
+    fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        let vector =
+            <V::OperatorVTable as OperatorVTable<V>>::execute_batch(&self.0, selection, ctx)?;
+
+        // Such a cheap check that we run it always. More expensive DType checks live in
+        // debug_assertions.
+        assert_eq!(
+            vector.len(),
+            selection.true_count(),
+            "Batch execution returned vector of incorrect length"
+        );
+
+        #[cfg(debug_assertions)]
+        {
+            // Checks for correct type and nullability.
+            if !vector_matches_dtype(&vector, self.dtype()) {
+                vortex_panic!(
+                    "Returned vector {:?} does not match expected dtype {}",
+                    vector,
+                    self.dtype()
                 );
             }
         }
-        self.bind(selection, &mut ())?.execute()
+
+        Ok(vector)
     }
 
     fn reduce_children(&self) -> VortexResult<Option<ArrayRef>> {
@@ -107,3 +115,14 @@ impl BindCtx for () {
         array.bind(selection, self)
     }
 }
+
+impl dyn Array + '_ {
+    pub fn execute(&self) -> VortexResult<Vector> {
+        self.execute_batch(&Mask::new_true(self.len()), &mut DummyExecutionCtx)
+    }
+
+    pub fn execute_with_selection(&self, mask: &Mask) -> VortexResult<Vector> {
+        assert_eq!(self.len(), mask.len());
+        self.execute_batch(mask, &mut DummyExecutionCtx)
+    }
+}

vortex-array/src/arrays/listview/vtable/operator.rs

@@ -43,8 +43,6 @@ impl OperatorVTable<ListViewVTable> for ListViewVTable {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
-
     use vortex_dtype::PTypeDowncast;
     use vortex_mask::Mask;
     use vortex_vector::VectorOps;
@@ -53,7 +51,7 @@ mod tests {
     use crate::arrays::listview::tests::common::{
         create_basic_listview, create_nullable_listview, create_overlapping_listview,
     };
-    use crate::arrays::{BoolArray, ListViewArray, PrimitiveArray};
+    use crate::arrays::{ListViewArray, PrimitiveArray};
     use crate::validity::Validity;
 
     #[test]
@@ -99,12 +97,10 @@ mod tests {
         let listview = ListViewArray::new(elements, offsets, sizes, Validity::AllValid);
 
         // Create selection mask: [true, false, true, false, true, false].
-        let selection = BoolArray::from_iter([true, false, true, false, true, false]).into_array();
+        let selection = Mask::from_iter([true, false, true, false, true, false]);
 
         // Execute with selection.
-        let result = listview
-            .execute_with_selection(Some(&Arc::new(selection)))
-            .unwrap();
+        let result = listview.execute_with_selection(&selection).unwrap();
 
         // Verify filtered length (3 lists selected).
         assert_eq!(result.len(), 3);
@@ -133,12 +129,10 @@ mod tests {
         let listview = create_nullable_listview();
 
         // Create selection mask: [true, true, false].
-        let selection = BoolArray::from_iter([true, true, false]).into_array();
+        let selection = Mask::from_iter([true, true, false]);
 
         // Execute with selection.
-        let result = listview
-            .execute_with_selection(Some(&Arc::new(selection)))
-            .unwrap();
+        let result = listview.execute_with_selection(&selection).unwrap();
 
         // Verify filtered length (2 lists selected, including the null).
         assert_eq!(result.len(), 2);
@@ -168,12 +162,10 @@ mod tests {
         let listview = create_overlapping_listview();
 
         // Create selection mask: [true, false, true, true, false].
-        let selection = BoolArray::from_iter([true, false, true, true, false]).into_array();
+        let selection = Mask::from_iter([true, false, true, true, false]);
 
         // Execute with selection.
-        let result = listview
-            .execute_with_selection(Some(&Arc::new(selection)))
-            .unwrap();
+        let result = listview.execute_with_selection(&selection).unwrap();
 
         // Verify filtered length (3 lists selected).
         assert_eq!(result.len(), 3);

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

@@ -41,9 +41,8 @@ impl OperatorVTable<StructVTable> for StructVTable {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
-
     use vortex_dtype::{FieldNames, PTypeDowncast};
+    use vortex_mask::Mask;
     use vortex_vector::VectorOps;
 
     use crate::IntoArray;
@@ -98,12 +97,10 @@ mod tests {
         .unwrap();
 
         // Create a selection mask that selects indices 0, 2, 4 (alternating pattern).
-        let selection = BoolArray::from_iter([true, false, true, false, true, false]).into_array();
+        let selection = Mask::from_iter([true, false, true, false, true, false]);
 
         // Execute with selection mask.
-        let result = struct_array
-            .execute_with_selection(Some(&Arc::new(selection)))
-            .unwrap();
+        let result = struct_array.execute_with_selection(&selection).unwrap();
 
         // Verify the result has the filtered length.
         assert_eq!(result.len(), 3);
@@ -152,12 +149,10 @@ mod tests {
         .unwrap();
 
         // Create a selection mask that selects indices 0, 1, 2, 4, 5.
-        let selection = BoolArray::from_iter([true, true, true, false, true, true]).into_array();
+        let selection = Mask::from_iter([true, true, true, false, true, true]);
 
         // Execute with selection mask.
-        let result = struct_array
-            .execute_with_selection(Some(&Arc::new(selection)))
-            .unwrap();
+        let result = struct_array.execute_with_selection(&selection).unwrap();
 
         assert_eq!(result.len(), 5);
 

vortex-array/src/compute/arrays/arithmetic.rs

@@ -347,7 +347,7 @@ mod tests {
 
     use crate::arrays::PrimitiveArray;
     use crate::compute::arrays::arithmetic::{ArithmeticArray, ArithmeticOperator};
-    use crate::{ArrayOperator, ArrayRef, IntoArray};
+    use crate::{ArrayRef, IntoArray};
 
     fn add(lhs: ArrayRef, rhs: ArrayRef) -> ArrayRef {
         ArithmeticArray::new(lhs, rhs, ArithmeticOperator::Add).into_array()
@@ -418,10 +418,8 @@ mod tests {
         let lhs = PrimitiveArray::from_iter([1u32, 2, 3]).into_array();
         let rhs = PrimitiveArray::from_iter([10u32, 20, 30]).into_array();
 
-        let selection = bitbuffer![1 0 1].into_array();
-
         let result = add(lhs, rhs)
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&bitbuffer![1 0 1].into())
             .unwrap()
             .into_primitive()
             .downcast::<u32>();

vortex-array/src/compute/arrays/get_item.rs

@@ -147,10 +147,10 @@ mod tests {
     use vortex_dtype::{FieldNames, Nullability, PTypeDowncast};
     use vortex_vector::VectorOps;
 
+    use crate::IntoArray;
     use crate::arrays::{BoolArray, PrimitiveArray, StructArray};
     use crate::compute::arrays::get_item::GetItemArray;
     use crate::validity::Validity;
-    use crate::{ArrayOperator, IntoArray};
 
     #[test]
     fn test_get_item_basic() {
@@ -233,9 +233,9 @@ mod tests {
             .into_array();
 
         // Apply selection mask [1 0 1 0 1 0] => select indices 0, 2, 4
-        let selection = bitbuffer![1 0 1 0 1 0].into_array();
+        let selection = bitbuffer![1 0 1 0 1 0].into();
         let result = get_item
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_primitive()
             .into_i32();

vortex-array/src/compute/arrays/logical.rs

@@ -213,7 +213,7 @@ mod tests {
     use vortex_buffer::bitbuffer;
 
     use crate::compute::arrays::logical::{LogicalArray, LogicalOperator};
-    use crate::{ArrayOperator, ArrayRef, IntoArray};
+    use crate::{ArrayRef, IntoArray};
 
     fn and_(lhs: ArrayRef, rhs: ArrayRef) -> ArrayRef {
         LogicalArray::new(lhs, rhs, LogicalOperator::And).into_array()
@@ -232,10 +232,10 @@ mod tests {
         let lhs = bitbuffer![0 1 0].into_array();
         let rhs = bitbuffer![0 1 1].into_array();
 
-        let selection = bitbuffer![0 1 1].into_array();
+        let selection = bitbuffer![0 1 1].into();
 
         let result = and_(lhs, rhs)
-            .execute_with_selection(Some(&selection))
+            .execute_with_selection(&selection)
             .unwrap()
             .into_bool();
         assert_eq!(result.bits(), &bitbuffer![1 0]);

vortex-array/src/execution/mod.rs

@@ -7,3 +7,18 @@ mod validity;
 
 pub use batch::*;
 pub use mask::*;
+
+/// Execution context for batch array compute.
+// NOTE(ngates): This context will eventually hold cached resources for execution, such as CSE
+//  nodes, and may well eventually support a type-map interface for arrays to stash arbitrary
+//  execution-related data.
+pub trait ExecutionCtx: private::Sealed {}
+
+/// A crate-internal dummy execution context.
+pub(crate) struct DummyExecutionCtx;
+impl ExecutionCtx for DummyExecutionCtx {}
+
+mod private {
+    pub trait Sealed {}
+    impl Sealed for super::DummyExecutionCtx {}
+}

vortex-array/src/optimizer.rs

@@ -63,9 +63,9 @@ mod tests {
     use vortex_dtype::PTypeDowncast;
     use vortex_vector::VectorOps;
 
+    use crate::IntoArray;
     use crate::arrays::{BoolArray, MaskedArray, PrimitiveArray};
     use crate::validity::Validity;
-    use crate::{ArrayOperator, IntoArray};
 
     #[test]
     fn test_masked_pushdown() {