Add more arrow execution

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

Author: gatesn

vortex-array/src/array/mod.rs

@@ -4,10 +4,12 @@
 mod visitor;
 
 use std::any::Any;
+use std::any::type_name;
 use std::fmt::Debug;
 use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
+use std::ops::Deref;
 use std::ops::Range;
 use std::sync::Arc;
 
@@ -18,6 +20,7 @@ use vortex_dtype::Nullability;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
@@ -73,6 +76,9 @@ pub trait Array:
     /// Returns the array as a reference to a generic [`Any`] trait object.
     fn as_any(&self) -> &dyn Any;
 
+    /// Returns the array as an `Arc<dyn Any + Send + Sync>`.
+    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
+
     /// Returns the array as an [`ArrayRef`].
     fn to_array(&self) -> ArrayRef;
 
@@ -211,6 +217,10 @@ impl Array for Arc<dyn Array> {
         self.as_ref().as_any()
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
+        self
+    }
+
     #[inline]
     fn to_array(&self) -> ArrayRef {
         self.clone()
@@ -350,6 +360,24 @@ impl dyn Array + '_ {
             .map(|array_adapter| &array_adapter.0)
     }
 
+    /// Returns the array downcast to the given `A` as an owned object.
+    pub fn try_into<V: VTable>(self: Arc<Self>) -> Result<V::Array, Arc<Self>> {
+        match self.is::<V>() {
+            true => {
+                let arc = self
+                    .as_any_arc()
+                    .downcast::<ArrayAdapter<V>>()
+                    .map_err(|_| vortex_err!("failed to downcast"))
+                    .vortex_expect("Failed to downcast");
+                Ok(match Arc::try_unwrap(arc) {
+                    Ok(array) => array.0,
+                    Err(arc) => arc.deref().0.clone(),
+                })
+            }
+            false => Err(self),
+        }
+    }
+
     /// Is self an array with encoding from vtable `V`.
     pub fn is<V: VTable>(&self) -> bool {
         self.as_opt::<V>().is_some()
@@ -443,6 +471,10 @@ impl<V: VTable> Array for ArrayAdapter<V> {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
+        self
+    }
+
     fn to_array(&self) -> ArrayRef {
         Arc::new(ArrayAdapter::<V>(self.0.clone()))
     }

vortex-array/src/arrays/dict/array.rs

@@ -114,6 +114,10 @@ impl DictArray {
         Ok(unsafe { Self::new_unchecked(codes, values) })
     }
 
+    pub fn into_parts(self) -> (ArrayRef, ArrayRef) {
+        (self.codes, self.values)
+    }
+
     #[inline]
     pub fn codes(&self) -> &ArrayRef {
         &self.codes

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

@@ -319,6 +319,10 @@ impl ListViewArray {
         .is_ok()
     }
 
+    pub fn into_parts(self) -> (ArrayRef, ArrayRef, ArrayRef, Validity) {
+        (self.elements, self.offsets, self.sizes, self.validity)
+    }
+
     /// Returns the offset at the given index.
     ///
     /// Note that it is possible the corresponding list view is null (which is only defined by the

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

@@ -14,7 +14,7 @@ use crate::VectorExecutor;
 use crate::arrow::null_buffer::to_null_buffer;
 
 pub(super) fn to_arrow_bool(
-    array: &ArrayRef,
+    array: ArrayRef,
     session: &VortexSession,
 ) -> VortexResult<ArrowArrayRef> {
     let bool_vector = array

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

@@ -0,0 +1,68 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use arrow_array::ArrayRef as ArrowArrayRef;
+use arrow_array::GenericBinaryArray;
+use arrow_array::types::ByteArrayType;
+use vortex_compute::arrow::IntoArrow;
+use vortex_dtype::DType;
+use vortex_dtype::NativePType;
+use vortex_dtype::Nullability;
+use vortex_dtype::PTypeDowncastExt;
+use vortex_error::VortexError;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+use crate::ArrayRef;
+use crate::VectorExecutor;
+use crate::arrays::VarBinArray;
+use crate::arrays::VarBinVTable;
+use crate::arrow::null_buffer::to_null_buffer;
+use crate::builtins::ArrayBuiltins;
+
+/// Convert a Vortex array into an Arrow GenericBinaryArray.
+pub(super) fn to_arrow_byte_array<T: ByteArrayType>(
+    array: ArrayRef,
+    session: &VortexSession,
+) -> VortexResult<ArrowArrayRef>
+where
+    T::Offset: NativePType,
+{
+    // If the Vortex array is already in VarBin format, we can directly convert it.
+    if let Some(array) = array.as_opt::<VarBinVTable>() {
+        return varbin_to_byte_array::<T>(array, session);
+    }
+
+    // Otherwise, we execute the array to a BinaryView vector and cast from there.
+    let binary_view = array.execute_vector(session)?.into_arrow()?;
+    arrow_cast::cast(&binary_view, &T::DATA_TYPE).map_err(VortexError::from)
+}
+
+/// Convert a Vortex VarBinArray into an Arrow GenericBinaryArray.
+fn varbin_to_byte_array<T: ByteArrayType>(
+    array: &VarBinArray,
+    session: &VortexSession,
+) -> VortexResult<ArrowArrayRef>
+where
+    T::Offset: NativePType,
+{
+    // We must cast the offsets to the required offset type.
+    let offsets = array
+        .offsets()
+        .cast(DType::Primitive(T::Offset::PTYPE, Nullability::NonNullable))?
+        .execute_vector(session)?
+        .into_primitive()
+        .downcast::<T::Offset>()
+        .into_buffer()
+        .into_arrow_offset_buffer();
+
+    let data = array.bytes().clone().into_arrow_buffer();
+
+    let null_buffer = to_null_buffer(array.validity_mask());
+
+    Ok(Arc::new(unsafe {
+        GenericBinaryArray::<T::Offset>::new_unchecked(offsets, data, null_buffer)
+    }))
+}

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

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use arrow_array::ArrayRef as ArrowArrayRef;
+use arrow_array::types::ByteViewType;
+use vortex_compute::arrow::IntoArrow;
+use vortex_dtype::DType;
+use vortex_dtype::Nullability;
+use vortex_dtype::arrow::FromArrowType;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+use crate::ArrayRef;
+use crate::VectorExecutor;
+use crate::builtins::ArrayBuiltins;
+
+pub(super) fn to_arrow_byte_view<T: ByteViewType>(
+    array: ArrayRef,
+    session: &VortexSession,
+) -> VortexResult<ArrowArrayRef> {
+    // First we cast the array into the desired ByteView type.
+    // We do this in case the vortex array is Utf8, and we want Binary or vice versa. By casting
+    // first, we may push this down through the Vortex array tree. We choose nullable to be most
+    // flexible since there's no prescribed nullability in Arrow types.
+    let array = array.cast(DType::from_arrow((&T::DATA_TYPE, Nullability::Nullable)))?;
+
+    // Perform a naive conversion via our VarBinView vector representation
+    array.execute_vector(session)?.into_arrow()
+}

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

@@ -19,15 +19,16 @@ use crate::arrays::DictVTable;
 use crate::arrow::ArrowArrayExecutor;
 
 pub(super) fn to_arrow_dictionary(
-    array: &ArrayRef,
+    array: ArrayRef,
     codes_type: &DataType,
     values_type: &DataType,
     session: &VortexSession,
 ) -> VortexResult<ArrowArrayRef> {
     // Check if we have a Vortex dictionary array
-    if let Some(array) = array.as_opt::<DictVTable>() {
-        return dict_to_dict(array, codes_type, values_type, session);
-    }
+    let array = match array.try_into::<DictVTable>() {
+        Ok(array) => return dict_to_dict(array, codes_type, values_type, session),
+        Err(a) => a,
+    };
 
     // Otherwise, we should try and build a dictionary.
     // Arrow hides this functionality inside the cast module!
@@ -41,13 +42,14 @@ pub(super) fn to_arrow_dictionary(
 
 /// Convert a Vortex dictionary array to an Arrow dictionary array.
 fn dict_to_dict(
-    array: &DictArray,
+    array: DictArray,
     codes_type: &DataType,
     values_type: &DataType,
     session: &VortexSession,
 ) -> VortexResult<ArrowArrayRef> {
-    let codes = array.codes().execute_arrow(codes_type, session)?;
-    let values = array.values().execute_arrow(values_type, session)?;
+    let (codes, values) = array.into_parts();
+    let codes = codes.execute_arrow(codes_type, session)?;
+    let values = values.execute_arrow(values_type, session)?;
 
     Ok(match codes_type {
         DataType::Int8 => Arc::new(unsafe {