Add more arrow execution

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

Author: gatesn

vortex-array/Cargo.toml

@@ -25,6 +25,7 @@ arcref = { workspace = true }
 arrow-arith = { workspace = true }
 arrow-array = { workspace = true, features = ["ffi"] }
 arrow-buffer = { workspace = true }
+arrow-cast = { workspace = true }
 arrow-data = { workspace = true }
 arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }

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

@@ -0,0 +1,30 @@
+// 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::BooleanArray;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+use vortex_session::VortexSession;
+
+use crate::ArrayRef;
+use crate::VectorExecutor;
+use crate::arrow::null_buffer::to_null_buffer;
+
+pub(super) fn to_arrow_bool(
+    array: &ArrayRef,
+    session: &VortexSession,
+) -> VortexResult<ArrowArrayRef> {
+    let bool_vector = array
+        .execute_vector(session)?
+        .into_bool_opt()
+        .ok_or_else(|| vortex_err!("Failed to convert array to Bool vector"))?;
+
+    let (bits, validity) = bool_vector.into_parts();
+    Ok(Arc::new(BooleanArray::new(
+        bits.into(),
+        to_null_buffer(validity),
+    )))
+}

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

@@ -0,0 +1,79 @@
+// 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::DictionaryArray;
+use arrow_array::cast::AsArray;
+use arrow_array::types::*;
+use arrow_schema::DataType;
+use vortex_error::VortexError;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_session::VortexSession;
+
+use crate::ArrayRef;
+use crate::arrays::DictArray;
+use crate::arrays::DictVTable;
+use crate::arrow::ArrowArrayExecutor;
+
+pub(super) fn to_arrow_dictionary(
+    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);
+    }
+
+    // Otherwise, we should try and build a dictionary.
+    // Arrow hides this functionality inside the cast module!
+    let array = array.execute_arrow(values_type, session)?;
+    arrow_cast::cast(
+        &array,
+        &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())),
+    )
+    .map_err(VortexError::from)
+}
+
+/// Convert a Vortex dictionary array to an Arrow dictionary array.
+fn dict_to_dict(
+    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)?;
+
+    Ok(match codes_type {
+        DataType::Int8 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<Int8Type>().clone(), values)
+        }),
+        DataType::Int16 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<Int16Type>().clone(), values)
+        }),
+        DataType::Int32 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<Int32Type>().clone(), values)
+        }),
+        DataType::Int64 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<Int64Type>().clone(), values)
+        }),
+        DataType::UInt8 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<UInt8Type>().clone(), values)
+        }),
+        DataType::UInt16 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<UInt16Type>().clone(), values)
+        }),
+        DataType::UInt32 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<UInt32Type>().clone(), values)
+        }),
+        DataType::UInt64 => Arc::new(unsafe {
+            DictionaryArray::new_unchecked(codes.as_primitive::<UInt64Type>().clone(), values)
+        }),
+        _ => vortex_bail!("Unsupported dictionary codes type: {:?}", codes_type),
+    })
+}