Use more execute

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

Author: gatesn

vortex-array/src/arrays/filter/vtable.rs

@@ -22,8 +22,8 @@ use crate::ArrayHash;
 use crate::ArrayRef;
 use crate::Canonical;
 use crate::IntoArray;
+use crate::LEGACY_SESSION;
 use crate::Precision;
-use crate::arrays::LEGACY_SESSION;
 use crate::arrays::filter::array::FilterArray;
 use crate::arrays::filter::kernel::FilterKernel;
 use crate::kernel::BindCtx;

vortex-array/src/arrays/mod.rs

@@ -6,8 +6,6 @@
 #[cfg(any(test, feature = "test-harness"))]
 mod assertions;
 
-use std::sync::LazyLock;
-
 #[cfg(any(test, feature = "test-harness"))]
 pub use assertions::format_indices;
 
@@ -59,12 +57,3 @@ pub use scalar_fn::*;
 pub use struct_::*;
 pub use varbin::*;
 pub use varbinview::*;
-use vortex_session::VortexSession;
-
-use crate::session::ArraySession;
-
-// TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute
-//  from the new array encodings to support back-compat for legacy encodings. So we hold a session
-//  here...
-static LEGACY_SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::empty().with::<ArraySession>());

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

@@ -5,7 +5,7 @@ use vortex_error::VortexExpect;
 
 use crate::Array;
 use crate::Canonical;
-use crate::arrays::LEGACY_SESSION;
+use crate::LEGACY_SESSION;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
 use crate::executor::VectorExecutor;

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

@@ -8,7 +8,7 @@ use vortex_vector::ScalarOps;
 use vortex_vector::VectorOps;
 
 use crate::Array;
-use crate::arrays::LEGACY_SESSION;
+use crate::LEGACY_SESSION;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
 use crate::executor::VectorExecutor;

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

@@ -53,6 +53,7 @@ use vortex_buffer::Buffer;
 use vortex_dtype::DType;
 use vortex_dtype::IntegerPType;
 use vortex_dtype::PType;
+use vortex_dtype::PTypeDowncastExt;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -62,7 +63,8 @@ use vortex_scalar::DecimalType;
 use crate::Array as _;
 use crate::Canonical;
 use crate::IntoArray;
-use crate::ToCanonical;
+use crate::LEGACY_SESSION;
+use crate::VectorExecutor;
 use crate::arrays::BoolArray;
 use crate::arrays::DecimalArray;
 use crate::arrays::FixedSizeListArray;
@@ -75,11 +77,15 @@ use crate::arrays::list_from_list_view;
 use crate::arrow::IntoArrowArray;
 use crate::arrow::array::ArrowArray;
 use crate::arrow::compute::ToArrowArgs;
+<<<<<<< Updated upstream
 use crate::arrow::null_buffer::to_null_buffer;
+=======
+use crate::arrow::compute::to_arrow::null_buffer::to_null_buffer;
+use crate::builtins::ArrayBuiltins;
+>>>>>>> Stashed changes
 use crate::compute::InvocationArgs;
 use crate::compute::Kernel;
 use crate::compute::Output;
-use crate::compute::cast;
 
 /// Implementation of `ToArrow` kernel for canonical Vortex arrays.
 #[derive(Debug)]
@@ -555,9 +561,13 @@ fn to_arrow_list<O: IntegerPType + OffsetSizeTrait>(
     };
 
     // Convert the child `offsets` and `validity` array to Arrow.
-    let offsets = cast(list_array.offsets().as_ref(), &DType::from(O::PTYPE))?
-        .to_primitive()
-        .buffer::<O>()
+    let offsets = list_array
+        .offsets()
+        .cast(DType::from(O::PTYPE))?
+        .execute_vector(&LEGACY_SESSION)?
+        .into_primitive()
+        .downcast::<O>()
+        .into_buffer()
         .into_arrow_offset_buffer();
     let nulls = to_null_buffer(list_array.validity_mask());
 
@@ -576,16 +586,24 @@ 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 offsets = cast(array.offsets(), &offsets_dtype)
-        .map_err(|err| err.with_context(format!("Failed to cast offsets to {offsets_dtype}")))?
-        .to_primitive();
-    let sizes = cast(array.sizes(), &offsets_dtype)
-        .map_err(|err| err.with_context(format!("Failed to cast sizes to {offsets_dtype}")))?
-        .to_primitive();
+    let offsets = array
+        .offsets()
+        .cast(offsets_dtype.clone())?
+        .execute_vector(&LEGACY_SESSION)?
+        .into_primitive()
+        .downcast::<O>()
+        .into_buffer()
+        .into_arrow_scalar_buffer();
+    let sizes = array
+        .sizes()
+        .cast(offsets_dtype)?
+        .execute_vector(&LEGACY_SESSION)?
+        .into_primitive()
+        .downcast::<O>()
+        .into_buffer()
+        .into_arrow_scalar_buffer();
 
     // Convert `offsets`, `sizes`, and `validity` to Arrow buffers.
-    let arrow_offsets = offsets.buffer::<O>().into_arrow_scalar_buffer();
-    let arrow_sizes = sizes.buffer::<O>().into_arrow_scalar_buffer();
     let nulls = to_null_buffer(array.validity_mask());
 
     // Convert the child `elements` array to Arrow.
@@ -609,8 +627,8 @@ fn to_arrow_listview<O: IntegerPType + OffsetSizeTrait>(
 
     Ok(Arc::new(GenericListViewArray::new(
         element_field,
-        arrow_offsets,
-        arrow_sizes,
+        offsets,
+        sizes,
         elements,
         nulls,
     )))

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

@@ -11,19 +11,26 @@ use arrow_schema::Field;
 use arrow_schema::FieldRef;
 use vortex_dtype::DType;
 use vortex_dtype::IntegerPType;
+use vortex_dtype::PTypeDowncastExt;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
 use crate::IntoArray;
-use crate::ToCanonical;
+use crate::LEGACY_SESSION;
+use crate::VectorExecutor;
 use crate::arrays::ListArray;
 use crate::arrays::ListVTable;
 use crate::arrays::list_view_from_list;
 use crate::arrow::IntoArrowArray;
 use crate::arrow::compute::ToArrowKernel;
 use crate::arrow::compute::ToArrowKernelAdapter;
+<<<<<<< Updated upstream
 use crate::arrow::null_buffer::to_null_buffer;
 use crate::compute::cast;
+=======
+use crate::arrow::compute::to_arrow::null_buffer::to_null_buffer;
+use crate::builtins::ArrayBuiltins;
+>>>>>>> Stashed changes
 use crate::register_kernel;
 
 impl ToArrowKernel for ListVTable {
@@ -62,13 +69,16 @@ 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 offsets_dtype = DType::Primitive(O::PTYPE, array.dtype().nullability());
-    let offsets = cast(array.offsets(), &offsets_dtype)
-        .map_err(|err| err.with_context(format!("Failed to cast offsets to {offsets_dtype}")))?
-        .to_primitive();
+    let offsets = array
+        .offsets()
+        .cast(DType::Primitive(O::PTYPE, array.dtype().nullability()))?
+        .execute_vector(&LEGACY_SESSION)?
+        .into_primitive()
+        .downcast::<O>()
+        .into_buffer();
 
     // Convert `offsets` and `validity` to Arrow buffers.
-    let arrow_offsets = offsets.buffer::<O>().into_arrow_offset_buffer();
+    let arrow_offsets = offsets.into_arrow_offset_buffer();
     let nulls = to_null_buffer(array.validity_mask());
 
     // Convert the child `elements` array to Arrow.

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

@@ -12,6 +12,7 @@ use std::sync::LazyLock;
 use arcref::ArcRef;
 use arrow_array::ArrayRef as ArrowArrayRef;
 use arrow_schema::DataType;
+use vortex_compute::arrow::IntoArrow;
 use vortex_dtype::DType;
 use vortex_dtype::arrow::FromArrowType;
 use vortex_error::VortexError;
@@ -21,6 +22,9 @@ use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 
 use crate::Array;
+use crate::LEGACY_SESSION;
+use crate::USE_VORTEX_OPERATORS;
+use crate::VectorExecutor;
 use crate::arrow::array::ArrowArray;
 use crate::arrow::array::ArrowVTable;
 use crate::compute::ComputeFn;
@@ -124,15 +128,23 @@ impl ComputeFnVTable for ToArrow {
             return Ok(output);
         }
 
-        // Fall back to canonicalizing and then converting.
         if !array.is_canonical() {
-            let canonical_array = array.to_canonical();
-            let arrow_array = to_arrow_opts(
-                canonical_array.as_ref(),
-                &ToArrowOptions {
-                    arrow_type: arrow_type.cloned(),
-                },
-            )?;
+            let arrow_array = if *USE_VORTEX_OPERATORS {
+                array
+                    .to_array()
+                    .execute_vector(&LEGACY_SESSION)?
+                    .into_arrow()?
+            } else {
+                // Fall back to canonicalizing and then converting.
+                let canonical_array = array.to_canonical();
+                to_arrow_opts(
+                    canonical_array.as_ref(),
+                    &ToArrowOptions {
+                        arrow_type: arrow_type.cloned(),
+                    },
+                )?
+            };
+
             return Ok(ArrowArray::new(arrow_array, array.dtype().nullability())
                 .to_array()
                 .into());

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

@@ -21,25 +21,30 @@ use arrow_schema::DataType;
 use arrow_schema::TimeUnit as ArrowTimeUnit;
 use vortex_dtype::DType;
 use vortex_dtype::NativePType;
+use vortex_dtype::PTypeDowncastExt;
 use vortex_dtype::datetime::TemporalMetadata;
 use vortex_dtype::datetime::TimeUnit;
 use vortex_dtype::datetime::is_temporal_ext_type;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
-use crate::Array as _;
 use crate::IntoArray;
-use crate::ToCanonical;
+use crate::LEGACY_SESSION;
+use crate::VectorExecutor;
 use crate::arrays::ExtensionVTable;
 use crate::arrays::TemporalArray;
 use crate::arrow::array::ArrowArray;
 use crate::arrow::compute::to_arrow::ToArrowArgs;
+<<<<<<< Updated upstream
 use crate::arrow::null_buffer::to_null_buffer;
+=======
+use crate::arrow::compute::to_arrow::null_buffer::to_null_buffer;
+use crate::builtins::ArrayBuiltins;
+>>>>>>> Stashed changes
 use crate::compute::InvocationArgs;
 use crate::compute::Kernel;
 use crate::compute::Output;
-use crate::compute::cast;
 
 /// Implementation of `ToArrow` kernel for canonical Vortex arrays.
 #[derive(Debug)]
@@ -127,11 +132,17 @@ where
     T::Native: NativePType,
 {
     let values_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability());
-    let values = cast(array.temporal_values(), &values_dtype)?
-        .to_primitive()
-        .into_buffer()
-        .into_arrow_scalar_buffer();
-    let nulls = to_null_buffer(array.temporal_values().validity_mask());
+    let values = array
+        .temporal_values()
+        .cast(values_dtype)?
+        .execute_vector(&LEGACY_SESSION)?
+        .into_primitive()
+        .downcast::<T::Native>();
+
+    let (buffer, validity) = values.into_parts();
+    let values = buffer.into_arrow_scalar_buffer();
+    let nulls = to_null_buffer(validity);
+
     Ok(ArrowPrimitiveArray::<T>::new(values, nulls))
 }
 

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

@@ -12,12 +12,14 @@ use vortex_dtype::DType;
 use vortex_dtype::IntegerPType;
 use vortex_dtype::Nullability;
 use vortex_dtype::PType;
+use vortex_dtype::PTypeDowncastExt;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_panic;
 
 use crate::Array;
-use crate::ToCanonical;
+use crate::LEGACY_SESSION;
+use crate::VectorExecutor;
 use crate::arrays::VarBinArray;
 use crate::arrays::VarBinVTable;
 use crate::arrow::compute::ToArrowKernel;
@@ -82,7 +84,10 @@ fn to_arrow<O: IntegerPType + OffsetSizeTrait>(array: &VarBinArray) -> VortexRes
         array.offsets(),
         &DType::Primitive(O::PTYPE, Nullability::NonNullable),
     )?
-    .to_primitive();
+    .execute_vector(&LEGACY_SESSION)?
+    .into_primitive()
+    .downcast::<O>()
+    .into_buffer();
 
     let nulls = to_null_buffer(array.validity_mask());
     let data = array.bytes().clone();
@@ -91,14 +96,14 @@ fn to_arrow<O: IntegerPType + OffsetSizeTrait>(array: &VarBinArray) -> VortexRes
     Ok(match array.dtype() {
         DType::Binary(_) => Arc::new(unsafe {
             GenericBinaryArray::new_unchecked(
-                offsets.buffer::<O>().into_arrow_offset_buffer(),
+                offsets.into_arrow_offset_buffer(),
                 data.into_arrow_buffer(),
                 nulls,
             )
         }),
         DType::Utf8(_) => Arc::new(unsafe {
             GenericStringArray::new_unchecked(
-                offsets.buffer::<O>().into_arrow_offset_buffer(),
+                offsets.into_arrow_offset_buffer(),
                 data.into_arrow_buffer(),
                 nulls,
             )

vortex-array/src/lib.rs

@@ -12,13 +12,18 @@
 //! Every data type recognized by Vortex also has a canonical physical encoding format, which
 //! arrays can be [canonicalized](Canonical) into for ease of access in compute functions.
 
+use std::sync::LazyLock;
+
 pub use array::*;
 pub use canonical::*;
 pub use context::*;
 pub use executor::*;
 pub use hash::*;
 pub use mask_future::*;
 pub use metadata::*;
+use vortex_session::VortexSession;
+
+use crate::session::ArraySession;
 
 pub mod accessor;
 #[doc(hidden)]
@@ -60,3 +65,15 @@ pub mod flatbuffers {
     //! Re-exported autogenerated code from the core Vortex flatbuffer definitions.
     pub use vortex_flatbuffers::array::*;
 }
+
+static USE_VORTEX_OPERATORS: LazyLock<bool> = LazyLock::new(|| {
+    std::env::var("VORTEX_OPERATORS")
+        .map(|v| v == "1" || v.to_lowercase() == "true")
+        .unwrap_or(false)
+});
+
+// TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute
+//  from the new array encodings to support back-compat for legacy encodings. So we hold a session
+//  here...
+static LEGACY_SESSION: LazyLock<VortexSession> =
+    LazyLock::new(|| VortexSession::empty().with::<ArraySession>());