execute update

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

Author: joseph-isaacs

encodings/alp/src/alp/array.rs

@@ -479,6 +479,7 @@ mod tests {
     use std::sync::LazyLock;
 
     use rstest::rstest;
+    use vortex_array::ExecutionCtx;
     use vortex_array::ToCanonical;
     use vortex_array::VectorExecutor;
     use vortex_array::arrays::PrimitiveArray;
@@ -505,7 +506,10 @@ mod tests {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f32));
         let encoded = alp_encode(&values, None).unwrap();
 
-        let result_canonical = encoded.to_array().execute_session(&SESSION).unwrap();
+        let result_canonical = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            encoded.to_array().execute(&mut ctx).unwrap()
+        };
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -532,7 +536,10 @@ mod tests {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64));
         let encoded = alp_encode(&values, None).unwrap();
 
-        let result_canonical = encoded.to_array().execute_session(&SESSION).unwrap();
+        let result_canonical = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            encoded.to_array().execute(&mut ctx).unwrap()
+        };
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -565,7 +572,10 @@ mod tests {
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_canonical = encoded.to_array().execute_session(&SESSION).unwrap();
+        let result_canonical = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            encoded.to_array().execute(&mut ctx).unwrap()
+        };
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -596,7 +606,10 @@ mod tests {
         let array = PrimitiveArray::from_option_iter(values);
         let encoded = alp_encode(&array, None).unwrap();
 
-        let result_canonical = encoded.to_array().execute_session(&SESSION).unwrap();
+        let result_canonical = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            encoded.to_array().execute(&mut ctx).unwrap()
+        };
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -638,7 +651,10 @@ mod tests {
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_canonical = encoded.to_array().execute_session(&SESSION).unwrap();
+        let result_canonical = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            encoded.to_array().execute(&mut ctx).unwrap()
+        };
         // Compare against the traditional array-based decompress path
         let expected = decompress_into_array(encoded);
 
@@ -683,7 +699,10 @@ mod tests {
         let slice_len = slice_end - slice_start;
         let sliced_encoded = encoded.slice(slice_start..slice_end);
 
-        let result_canonical = sliced_encoded.execute_session(&SESSION).unwrap();
+        let result_canonical = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            sliced_encoded.execute(&mut ctx).unwrap()
+        };
         let result_primitive = result_canonical.into_primitive();
 
         for idx in 0..slice_len {

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

@@ -203,6 +203,7 @@ pub fn count_exceptions(bit_width: u8, bit_width_freq: &[usize]) -> usize {
 mod tests {
     use std::sync::LazyLock;
 
+    use vortex_array::ExecutionCtx;
     use vortex_array::IntoArray;
     use vortex_array::VectorExecutor;
     use vortex_array::assert_arrays_eq;
@@ -534,7 +535,10 @@ 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(&SESSION).unwrap();
+            let executed = {
+                let mut ctx = ExecutionCtx::new(SESSION.clone());
+                bitpacked.into_array().execute(&mut ctx).unwrap()
+            };
 
             // All three should produce the same length.
             assert_eq!(vector_result.len(), array.len(), "vector length mismatch");
@@ -554,11 +558,14 @@ mod tests {
 
             // Verify that the execute() method works correctly by comparing with unpack_array.
             // We convert unpack_array result to canonical to compare.
-            let unpacked_executed = unpacked_array
-                .into_array()
-                .execute_session(&SESSION)
-                .unwrap()
-                .into_primitive();
+            let unpacked_executed = {
+                let mut ctx = ExecutionCtx::new(SESSION.clone());
+                unpacked_array
+                    .into_array()
+                    .execute(&mut ctx)
+                    .unwrap()
+                    .into_primitive()
+            };
             assert_eq!(
                 executed_primitive.len(),
                 unpacked_executed.len(),
@@ -590,7 +597,10 @@ 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(&SESSION).unwrap();
+        let executed = {
+            let mut ctx = ExecutionCtx::new(SESSION.clone());
+            sliced.execute(&mut ctx).unwrap()
+        };
 
         assert_eq!(
             vector_result.len(),

vortex-array/src/arrow/compute/to_arrow/canonical.rs

@@ -57,6 +57,7 @@ use vortex_scalar::DecimalType;
 
 use crate::Array as _;
 use crate::Canonical;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::LEGACY_SESSION;
 use crate::VectorExecutor;
@@ -531,10 +532,11 @@ fn to_arrow_list<O: IntegerPType + OffsetSizeTrait>(
     };
 
     // Convert the child `offsets` and `validity` array to Arrow.
+    let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
     let offsets = list_array
         .offsets()
         .cast(DType::from(O::PTYPE))?
-        .execute_session(&LEGACY_SESSION)?
+        .execute(&mut ctx)?
         .to_vector_session(&LEGACY_SESSION)?
         .into_primitive()
         .downcast::<O>()
@@ -557,10 +559,11 @@ fn to_arrow_listview<O: IntegerPType + OffsetSizeTrait>(
 ) -> VortexResult<ArrowArrayRef> {
     // First we cast the offsets and sizes into the specified width (determined by `O::PTYPE`).
     let offsets_dtype = DType::Primitive(O::PTYPE, array.dtype().nullability());
+    let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
     let offsets = array
         .offsets()
         .cast(offsets_dtype.clone())?
-        .execute_session(&LEGACY_SESSION)?
+        .execute(&mut ctx)?
         .to_vector_session(&LEGACY_SESSION)?
         .into_primitive()
         .downcast::<O>()
@@ -569,7 +572,7 @@ fn to_arrow_listview<O: IntegerPType + OffsetSizeTrait>(
     let sizes = array
         .sizes()
         .cast(offsets_dtype)?
-        .execute_session(&LEGACY_SESSION)?
+        .execute(&mut ctx)?
         .to_vector_session(&LEGACY_SESSION)?
         .into_primitive()
         .downcast::<O>()

vortex-array/src/arrow/compute/to_arrow/list.rs

@@ -15,6 +15,7 @@ use vortex_dtype::PTypeDowncastExt;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::LEGACY_SESSION;
 use crate::VectorExecutor;
@@ -64,10 +65,11 @@ fn list_array_to_arrow_list<O: IntegerPType + OffsetSizeTrait>(
     element: Option<&FieldRef>,
 ) -> VortexResult<ArrowArrayRef> {
     // First we cast the offsets and sizes into the specified width (determined by `O::PTYPE`).
+    let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
     let offsets = array
         .offsets()
         .cast(DType::Primitive(O::PTYPE, array.dtype().nullability()))?
-        .execute_session(&LEGACY_SESSION)?
+        .execute(&mut ctx)?
         .to_vector_session(&LEGACY_SESSION)?
         .into_primitive()
         .downcast::<O>()

vortex-array/src/arrow/compute/to_arrow/temporal.rs

@@ -29,6 +29,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::LEGACY_SESSION;
 use crate::VectorExecutor;
@@ -128,10 +129,11 @@ where
     T::Native: NativePType,
 {
     let values_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability());
+    let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
     let values = array
         .temporal_values()
         .cast(values_dtype)?
-        .execute_session(&LEGACY_SESSION)?
+        .execute(&mut ctx)?
         .to_vector_session(&LEGACY_SESSION)?
         .into_primitive()
         .downcast::<T::Native>();

vortex-array/src/arrow/compute/to_arrow/varbin.rs

@@ -18,6 +18,7 @@ use vortex_error::vortex_bail;
 use vortex_error::vortex_panic;
 
 use crate::Array;
+use crate::ExecutionCtx;
 use crate::LEGACY_SESSION;
 use crate::VectorExecutor;
 use crate::arrays::VarBinArray;
@@ -80,11 +81,12 @@ impl ToArrowKernel for VarBinVTable {
 register_kernel!(ToArrowKernelAdapter(VarBinVTable).lift());
 
 fn to_arrow<O: IntegerPType + OffsetSizeTrait>(array: &VarBinArray) -> VortexResult<ArrowArrayRef> {
+    let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
     let offsets = cast(
         array.offsets(),
         &DType::Primitive(O::PTYPE, Nullability::NonNullable),
     )?
-    .execute_session(&LEGACY_SESSION)?
+    .execute(&mut ctx)?
     .to_vector_session(&LEGACY_SESSION)?
     .into_primitive()
     .downcast::<O>()

vortex-array/src/arrow/executor/bool.rs

@@ -9,6 +9,7 @@ use vortex_error::VortexResult;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::VectorExecutor;
 use crate::arrays::BoolArray;
 use crate::arrow::null_buffer::to_null_buffer;
@@ -25,6 +26,7 @@ pub(super) fn to_arrow_bool(
     array: ArrayRef,
     session: &VortexSession,
 ) -> VortexResult<ArrowArrayRef> {
-    let bool_array = array.execute_session(session)?.into_bool();
+    let mut ctx = ExecutionCtx::new(session.clone());
+    let bool_array = array.execute(&mut ctx)?.into_bool();
     Ok(canonical_bool_to_arrow(&bool_array))
 }

vortex-array/src/arrow/executor/byte.rs

@@ -16,6 +16,7 @@ use vortex_error::VortexResult;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::VectorExecutor;
 use crate::arrays::VarBinArray;
 use crate::arrays::VarBinVTable;
@@ -37,8 +38,9 @@ where
     }
 
     // Otherwise, we execute the array to a BinaryView vector and cast from there.
+    let mut ctx = ExecutionCtx::new(session.clone());
     let binary_view = array
-        .execute_session(session)?
+        .execute(&mut ctx)?
         .to_vector_session(session)?
         .into_arrow()?;
     arrow_cast::cast(&binary_view, &T::DATA_TYPE).map_err(VortexError::from)
@@ -53,10 +55,11 @@ where
     T::Offset: NativePType,
 {
     // We must cast the offsets to the required offset type.
+    let mut ctx = ExecutionCtx::new(session.clone());
     let offsets = array
         .offsets()
         .cast(DType::Primitive(T::Offset::PTYPE, Nullability::NonNullable))?
-        .execute_session(session)?
+        .execute(&mut ctx)?
         .to_vector_session(session)?
         .into_primitive()
         .downcast::<T::Offset>()

vortex-array/src/arrow/executor/byte_view.rs

@@ -14,6 +14,7 @@ use vortex_error::VortexResult;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::VectorExecutor;
 use crate::arrays::VarBinViewArray;
 use crate::arrow::null_buffer::to_null_buffer;
@@ -44,6 +45,7 @@ pub(super) fn to_arrow_byte_view<T: ByteViewType>(
     // flexible since there's no prescribed nullability in Arrow types.
     let array = array.cast(DType::from_arrow((&T::DATA_TYPE, Nullability::Nullable)))?;
 
-    let varbinview = array.execute_session(session)?.into_varbinview();
+    let mut ctx = ExecutionCtx::new(session.clone());
+    let varbinview = array.execute(&mut ctx)?.into_varbinview();
     Ok(canonical_varbinview_to_arrow::<T>(&varbinview))
 }

vortex-array/src/arrow/executor/decimal.rs

@@ -17,6 +17,7 @@ use vortex_error::vortex_ensure;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::VectorExecutor;
 use crate::arrow::null_buffer::to_null_buffer;
 use crate::builtins::ArrayBuiltins;
@@ -30,7 +31,7 @@ pub(super) fn to_arrow_decimal<D: DecimalType, N: NativeDecimalType>(
     session: &VortexSession,
 ) -> VortexResult<ArrowArrayRef> {
     // Since Vortex doesn't have physical types, our DecimalDType only contains precision and scale.
-    // When calling execute_vector, Vortex may use any physical type >= the smallest type that can
+    // When calling execute, Vortex may use any physical type >= the smallest type that can
     // hold the requested precision.
     //
     // We therefore create a fake precision that forces Vortex to use a native type that is at
@@ -45,8 +46,9 @@ pub(super) fn to_arrow_decimal<D: DecimalType, N: NativeDecimalType>(
     ))?;
 
     // Execute the array as a vector and downcast to our native type.
+    let mut ctx = ExecutionCtx::new(session.clone());
     let vector = array
-        .execute_session(session)?
+        .execute(&mut ctx)?
         .to_vector_session(session)?
         .into_decimal();
     vortex_ensure!(