Arrow executor

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

encodings/alp/src/alp/array.rs

@@ -473,6 +473,7 @@ mod tests {
     use std::sync::LazyLock;
 
     use rstest::rstest;
+    use vortex_array::VectorExecutor;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::vtable::ValidityHelper;
     use vortex_dtype::PTypeDowncast;
@@ -497,7 +498,7 @@ mod tests {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f32));
         let encoded = alp_encode(&values, None).unwrap();
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -521,7 +522,7 @@ mod tests {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64));
         let encoded = alp_encode(&values, None).unwrap();
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -551,7 +552,7 @@ mod tests {
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -579,7 +580,7 @@ mod tests {
         let array = PrimitiveArray::from_option_iter(values);
         let encoded = alp_encode(&array, None).unwrap();
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -618,7 +619,7 @@ mod tests {
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_vector = encoded.to_array().execute(&SESSION).unwrap();
+        let result_vector = encoded.to_array().execute_vector(&SESSION).unwrap();
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -660,7 +661,7 @@ mod tests {
         let slice_len = slice_end - slice_start;
         let sliced_encoded = encoded.slice(slice_start..slice_end);
 
-        let result_vector = sliced_encoded.execute(&SESSION).unwrap();
+        let result_vector = sliced_encoded.execute_vector_optimized(&SESSION).unwrap();
         let result_primitive = result_vector.into_primitive().into_f64();
 
         for idx in 0..slice_len {

encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs

@@ -204,6 +204,7 @@ mod tests {
     use std::sync::LazyLock;
 
     use vortex_array::IntoArray;
+    use vortex_array::VectorExecutor;
     use vortex_array::assert_arrays_eq;
     use vortex_array::validity::Validity;
     use vortex_buffer::Buffer;
@@ -536,7 +537,7 @@ mod tests {
             let unpacked_array = unpack_array(&bitpacked);
 
             // Method 3: Using the execute() method (this is what would be used in production).
-            let executed = bitpacked.into_array().execute(&SESSION).unwrap();
+            let executed = bitpacked.into_array().execute_vector(&SESSION).unwrap();
 
             // All three should produce the same length.
             assert_eq!(vector_result.len(), array.len(), "vector length mismatch");
@@ -556,7 +557,10 @@ mod tests {
 
             // Verify that the execute() method works correctly by comparing with unpack_array.
             // We convert unpack_array result to a vector using execute() to compare.
-            let unpacked_executed = unpacked_array.into_array().execute(&SESSION).unwrap();
+            let unpacked_executed = unpacked_array
+                .into_array()
+                .execute_vector(&SESSION)
+                .unwrap();
             match (&executed, &unpacked_executed) {
                 (Vector::Primitive(exec_pv), Vector::Primitive(unpack_pv)) => {
                     assert_eq!(
@@ -593,7 +597,7 @@ mod tests {
         let sliced_bp = sliced.as_::<BitPackedVTable>();
         let vector_result = unpack_to_primitive_vector(sliced_bp);
         let unpacked_array = unpack_array(sliced_bp);
-        let executed = sliced.execute(&SESSION).unwrap();
+        let executed = sliced.execute_vector(&SESSION).unwrap();
 
         assert_eq!(
             vector_result.len(),

vortex-array/src/array/mod.rs

@@ -22,9 +22,6 @@ use vortex_error::vortex_ensure;
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
-use vortex_session::VortexSession;
-use vortex_vector::Vector;
-use vortex_vector::VectorOps;
 
 use crate::ArrayEq;
 use crate::ArrayHash;
@@ -372,35 +369,6 @@ impl dyn Array + '_ {
         }
         nbytes
     }
-
-    /// Execute the array and return the resulting vector.
-    ///
-    /// This entry-point function will choose an appropriate CPU-based execution strategy.
-    pub fn execute(&self, session: &VortexSession) -> VortexResult<Vector> {
-        let mut ctx = BindCtx::new(session.clone());
-
-        // NOTE(ngates): in the future we can choose a different mode of execution, or run
-        // optimization here, etc.
-        let kernel = self.bind_kernel(&mut ctx)?;
-        let result = kernel.execute()?;
-
-        vortex_ensure!(
-            result.len() == self.len(),
-            "Result length mismatch for {}",
-            self.encoding_id()
-        );
-
-        #[cfg(debug_assertions)]
-        {
-            vortex_ensure!(
-                vortex_vector::vector_matches_dtype(&result, self.dtype()),
-                "Executed vector dtype mismatch for {}",
-                self.encoding_id(),
-            );
-        }
-
-        Ok(result)
-    }
 }
 
 /// Trait for converting a type into a Vortex [`ArrayRef`].

vortex-array/src/arrays/scalar_fn/vtable/canonical.rs

@@ -3,13 +3,13 @@
 
 use itertools::Itertools;
 use vortex_error::VortexExpect;
-use vortex_vector::Datum;
 
 use crate::Array;
 use crate::Canonical;
 use crate::arrays::LEGACY_SESSION;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::executor::VectorExecutor;
 use crate::expr::ExecutionArgs;
 use crate::vectors::VectorIntoArray;
 use crate::vtable::CanonicalVTable;
@@ -20,8 +20,7 @@ impl CanonicalVTable<ScalarFnVTable> for ScalarFnVTable {
         let child_datums: Vec<_> = array
             .children()
             .iter()
-            // TODO(ngates): we could make all execution operate over datums
-            .map(|child| child.execute(&LEGACY_SESSION).map(Datum::Vector))
+            .map(|child| child.execute_datum_optimized(&LEGACY_SESSION))
             .try_collect()
             // FIXME(ngates): canonicalizing really ought to be fallible
             .vortex_expect(

vortex-array/src/arrays/scalar_fn/vtable/validity.rs

@@ -3,11 +3,15 @@
 
 use vortex_error::VortexExpect;
 use vortex_mask::Mask;
+use vortex_vector::Datum;
+use vortex_vector::ScalarOps;
+use vortex_vector::VectorOps;
 
 use crate::Array;
 use crate::arrays::LEGACY_SESSION;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::executor::VectorExecutor;
 use crate::vtable::ValidityVTable;
 
 impl ValidityVTable<ScalarFnVTable> for ScalarFnVTable {
@@ -46,9 +50,13 @@ impl ValidityVTable<ScalarFnVTable> for ScalarFnVTable {
     }
 
     fn validity_mask(array: &ScalarFnArray) -> Mask {
-        let vector = array
-            .execute(&LEGACY_SESSION)
+        let datum = array
+            .to_array()
+            .execute_datum(&LEGACY_SESSION)
             .vortex_expect("Validity mask computation should be fallible");
-        Mask::from_buffer(vector.into_bool().into_bits())
+        match datum {
+            Datum::Scalar(s) => Mask::new(array.len, s.is_valid()),
+            Datum::Vector(v) => v.validity().clone(),
+        }
     }
 }

vortex-array/src/executor.rs

@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_session::VortexSession;
+use vortex_vector::Datum;
+use vortex_vector::Vector;
+use vortex_vector::VectorOps;
+
+use crate::Array;
+use crate::ArrayRef;
+use crate::arrays::ConstantVTable;
+use crate::kernel::BindCtx;
+use crate::session::ArraySessionExt;
+
+/// Executor for exporting a Vortex [`Vector`] or [`Datum`] from an [`ArrayRef`].
+pub trait VectorExecutor {
+    /// Execute the array and return the resulting datum after running the optimizer.
+    fn execute_datum_optimized(&self, session: &VortexSession) -> VortexResult<Datum>;
+    /// Execute the array and return the resulting datum.
+    fn execute_datum(&self, session: &VortexSession) -> VortexResult<Datum>;
+    /// Execute the array and return the resulting vector after running the optimizer.
+    fn execute_vector_optimized(&self, session: &VortexSession) -> VortexResult<Vector>;
+    /// Execute the array and return the resulting vector.
+    fn execute_vector(&self, session: &VortexSession) -> VortexResult<Vector>;
+}
+
+impl VectorExecutor for ArrayRef {
+    fn execute_datum_optimized(&self, session: &VortexSession) -> VortexResult<Datum> {
+        session
+            .arrays()
+            .optimizer()
+            .optimize_array(self)?
+            .execute_datum(session)
+    }
+
+    fn execute_datum(&self, session: &VortexSession) -> VortexResult<Datum> {
+        // Attempt to short-circuit constant arrays.
+        if let Some(constant) = self.as_opt::<ConstantVTable>() {
+            return Ok(Datum::Scalar(constant.scalar().to_vector_scalar()));
+        }
+
+        let mut ctx = BindCtx::new(session.clone());
+
+        // NOTE(ngates): in the future we can choose a different mode of execution, or run
+        // optimization here, etc.
+        let kernel = self.bind_kernel(&mut ctx)?;
+        let result = kernel.execute()?;
+
+        vortex_ensure!(
+            result.len() == self.len(),
+            "Result length mismatch for {}",
+            self.encoding_id()
+        );
+
+        #[cfg(debug_assertions)]
+        {
+            vortex_ensure!(
+                vortex_vector::vector_matches_dtype(&result, self.dtype()),
+                "Executed vector dtype mismatch for {}",
+                self.encoding_id(),
+            );
+        }
+
+        Ok(Datum::Vector(result))
+    }
+
+    fn execute_vector_optimized(&self, session: &VortexSession) -> VortexResult<Vector> {
+        session
+            .arrays()
+            .optimizer()
+            .optimize_array(self)?
+            .execute_vector(session)
+    }
+
+    fn execute_vector(&self, session: &VortexSession) -> VortexResult<Vector> {
+        let len = self.len();
+        Ok(self.execute_datum(session)?.ensure_vector(len))
+    }
+}

vortex-array/src/expr/expression.rs

@@ -14,7 +14,6 @@ use vortex_dtype::DType;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
-use vortex_error::vortex_panic;
 
 use crate::ArrayRef;
 use crate::expr::Root;
@@ -124,12 +123,11 @@ impl Expression {
     }
 
     /// Evaluates the expression in the given scope, returning an array.
-    pub fn evaluate(&self, _scope: &ArrayRef) -> VortexResult<ArrayRef> {
-        vortex_panic!("DEPRECATED");
-        // if self.is::<Root>() {
-        //     return Ok(scope.clone());
-        // }
-        // self.scalar_fn.evaluate(self, scope)
+    pub fn evaluate(&self, scope: &ArrayRef) -> VortexResult<ArrayRef> {
+        if self.is::<Root>() {
+            return Ok(scope.clone());
+        }
+        self.scalar_fn.evaluate(self, scope)
     }
 
     /// An expression over zone-statistics which implies all records in the zone evaluate to false.

vortex-array/src/lib.rs

@@ -15,6 +15,7 @@
 pub use array::*;
 pub use canonical::*;
 pub use context::*;
+pub use executor::*;
 pub use hash::*;
 pub use mask_future::*;
 pub use metadata::*;
@@ -30,12 +31,13 @@ mod canonical;
 pub mod compute;
 mod context;
 pub mod display;
+mod executor;
 pub mod expr;
 mod expression;
 mod hash;
 pub mod iter;
 pub mod kernel;
-mod mask;
+pub mod mask;
 mod mask_future;
 mod metadata;
 pub mod optimizer;