[Experiment] Self-defined arrow-export

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

Cargo.lock

@@ -15,7 +15,9 @@ version = { workspace = true }
 
 [dependencies]
 arbitrary = { workspace = true, optional = true }
-arrow-array = { workspace = true, optional = true }
+arrow-array = { workspace = true }
+arrow-buffer = { workspace = true }
+arrow-schema = { workspace = true }
 itertools = { workspace = true }
 num-traits = { workspace = true }
 prost = { workspace = true }
@@ -39,7 +41,6 @@ vortex-array = { workspace = true, features = ["_test-harness"] }
 
 [features]
 arbitrary = ["dep:arbitrary", "vortex-array/arbitrary"]
-arrow = ["dep:arrow-array"]
 
 [[bench]]
 name = "run_end_null_count"

encodings/runend/Cargo.toml

@@ -88,12 +88,16 @@ pub fn vortex_runend::RunEnd::nbuffers(_array: &vortex_runend::RunEndArray) -> u
 
 pub fn vortex_runend::RunEnd::nchildren(_array: &vortex_runend::RunEndArray) -> usize
 
+pub fn vortex_runend::RunEnd::preferred_arrow_data_type(array: &vortex_runend::RunEndArray) -> core::option::Option<arrow_schema::datatype::DataType>
+
 pub fn vortex_runend::RunEnd::reduce_parent(array: &Self::Array, parent: &vortex_array::array::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
 pub fn vortex_runend::RunEnd::serialize(metadata: Self::Metadata) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
 
 pub fn vortex_runend::RunEnd::stats(array: &vortex_runend::RunEndArray) -> vortex_array::stats::array::StatsSetRef<'_>
 
+pub fn vortex_runend::RunEnd::to_arrow_array(array: &vortex_runend::RunEndArray, data_type: &arrow_schema::datatype::DataType, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<arrow_array::array::ArrayRef>>
+
 pub fn vortex_runend::RunEnd::with_children(array: &mut Self::Array, children: alloc::vec::Vec<vortex_array::array::ArrayRef>) -> vortex_error::VortexResult<()>
 
 impl vortex_array::vtable::operations::OperationsVTable<vortex_runend::RunEnd> for vortex_runend::RunEnd

encodings/runend/public-api.lock

@@ -3,7 +3,15 @@
 
 use std::fmt::Debug;
 use std::hash::Hash;
-
+use std::sync::Arc;
+
+use arrow_array::ArrayRef as ArrowArrayRef;
+use arrow_array::RunArray;
+use arrow_array::cast::AsArray;
+use arrow_array::types::*;
+use arrow_buffer::ArrowNativeType;
+use arrow_schema::DataType;
+use arrow_schema::Field;
 use vortex_array::ArrayEq;
 use vortex_array::ArrayHash;
 use vortex_array::ArrayRef;
@@ -17,6 +25,7 @@ use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
 use vortex_array::arrays::Primitive;
 use vortex_array::arrays::VarBinViewArray;
+use vortex_array::arrow::ArrowArrayExecutor;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
@@ -203,6 +212,41 @@ impl VTable for RunEnd {
         PARENT_KERNELS.execute(array, parent, child_idx, ctx)
     }
 
+    fn preferred_arrow_data_type(array: &RunEndArray) -> Option<DataType> {
+        let ends_dt = array.ends().dtype().to_arrow_dtype().ok()?;
+        let values_dt = array.values().dtype().to_arrow_dtype().ok()?;
+        let values_nullable = array.values().dtype().is_nullable();
+        Some(DataType::RunEndEncoded(
+            Arc::new(Field::new("run_ends", ends_dt, false)),
+            Arc::new(Field::new("values", values_dt, values_nullable)),
+        ))
+    }
+
+    fn to_arrow_array(
+        array: &RunEndArray,
+        data_type: &DataType,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrowArrayRef>> {
+        let DataType::RunEndEncoded(ends_field, values_field) = data_type else {
+            return Ok(None);
+        };
+        let arrow_ends = array
+            .ends()
+            .clone()
+            .execute_arrow(Some(ends_field.data_type()), ctx)?;
+        let arrow_values = array
+            .values()
+            .clone()
+            .execute_arrow(Some(values_field.data_type()), ctx)?;
+        Ok(Some(build_run_array(
+            &arrow_ends,
+            &arrow_values,
+            ends_field.data_type(),
+            array.offset(),
+            array.len(),
+        )?))
+    }
+
     fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
         run_end_canonicalize(array, ctx).map(ExecutionStep::Done)
     }
@@ -503,6 +547,59 @@ pub(super) fn run_end_canonicalize(
     })
 }
 
+fn build_run_array(
+    ends: &ArrowArrayRef,
+    values: &ArrowArrayRef,
+    ends_type: &DataType,
+    offset: usize,
+    length: usize,
+) -> VortexResult<ArrowArrayRef> {
+    match ends_type {
+        DataType::Int16 => build_run_array_typed::<Int16Type>(ends, values, offset, length),
+        DataType::Int32 => build_run_array_typed::<Int32Type>(ends, values, offset, length),
+        DataType::Int64 => build_run_array_typed::<Int64Type>(ends, values, offset, length),
+        _ => vortex_bail!("Unsupported run-end index type: {:?}", ends_type),
+    }
+}
+
+fn build_run_array_typed<R: RunEndIndexType>(
+    ends: &ArrowArrayRef,
+    values: &ArrowArrayRef,
+    offset: usize,
+    length: usize,
+) -> VortexResult<ArrowArrayRef>
+where
+    R::Native: std::ops::Sub<Output = R::Native> + Ord,
+{
+    let offset_native = R::Native::from_usize(offset).ok_or_else(|| {
+        vortex_error::vortex_err!("Offset {offset} exceeds run-end index capacity")
+    })?;
+    let length_native = R::Native::from_usize(length).ok_or_else(|| {
+        vortex_error::vortex_err!("Length {length} exceeds run-end index capacity")
+    })?;
+
+    let ends_prim = ends.as_primitive::<R>();
+    if offset == 0 && ends_prim.values().last() == Some(&length_native) {
+        return Ok(Arc::new(RunArray::<R>::try_new(ends_prim, values)?) as ArrowArrayRef);
+    }
+
+    // Trim to only include runs covering the [offset, offset+length) range.
+    let num_runs = (ends_prim
+        .values()
+        .partition_point(|&e| e - offset_native < length_native)
+        + 1)
+    .min(ends_prim.len());
+
+    let trimmed_ends = ends.slice(0, num_runs);
+    let trimmed_values = values.slice(0, num_runs);
+
+    let adjusted = trimmed_ends
+        .as_primitive::<R>()
+        .unary(|end| (end - offset_native).min(length_native));
+
+    Ok(Arc::new(RunArray::<R>::try_new(&adjusted, &trimmed_values)?) as ArrowArrayRef)
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::IntoArray;

encodings/runend/src/array.rs

@@ -9,7 +9,6 @@ pub use array::*;
 pub use iter::trimmed_ends_iter;
 
 mod array;
-#[cfg(feature = "arrow")]
 mod arrow;
 pub mod compress;
 mod compute;