add struct array for C export (#6294)

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

vortex-cuda/cudf-test/src/lib.rs

@@ -3,16 +3,19 @@
 
 //! This file is a simple C-compatible API that is called from the cudf-test-harness at CI time.
 
-#![allow(clippy::unwrap_used)]
+#![allow(clippy::unwrap_used, clippy::expect_used)]
 
 use std::sync::LazyLock;
 
-use arrow_schema::DataType;
 use arrow_schema::ffi::FFI_ArrowSchema;
 use futures::executor::block_on;
+use vortex::array::Array;
 use vortex::array::IntoArray;
 use vortex::array::arrays::PrimitiveArray;
+use vortex::array::arrays::StructArray;
 use vortex::array::session::ArraySession;
+use vortex::array::validity::Validity;
+use vortex::dtype::FieldNames;
 use vortex::expr::session::ExprSession;
 use vortex::io::session::RuntimeSession;
 use vortex::layout::session::LayoutSession;
@@ -42,9 +45,22 @@ pub extern "C" fn export_array(
 
     let primitive = PrimitiveArray::from_iter(0u32..1024);
 
-    *schema_ptr = FFI_ArrowSchema::try_from(DataType::UInt32).unwrap();
+    let array = StructArray::new(
+        FieldNames::from_iter(["a"]),
+        vec![primitive.into_array()],
+        1024,
+        Validity::NonNullable,
+    )
+    .into_array();
 
-    match block_on(primitive.into_array().export_device_array(&mut ctx)) {
+    let data_type = array
+        .dtype()
+        .to_arrow_dtype()
+        .expect("converting schema to Arrow DataType");
+
+    *schema_ptr = FFI_ArrowSchema::try_from(data_type).expect("data_type to FFI_ArrowSchema");
+
+    match block_on(array.export_device_array(&mut ctx)) {
         Ok(exported) => {
             *array_ptr = exported;
             0

vortex-cuda/src/arrow/canonical.rs

@@ -6,10 +6,13 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use cudarc::driver::sys;
+use futures::future::BoxFuture;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::PrimitiveArrayParts;
+use vortex_array::arrays::StructArray;
+use vortex_array::arrays::StructArrayParts;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::validity::Validity;
 use vortex_error::VortexResult;
@@ -41,10 +44,7 @@ impl ExportDeviceArray for CanonicalDeviceArrayExport {
     ) -> VortexResult<ArrowDeviceArray> {
         let cuda_array = array.execute_cuda(ctx).await?;
 
-        let (arrow_array, sync_event) = match cuda_array {
-            Canonical::Primitive(primitive) => export_primitive(primitive, ctx).await?,
-            c => todo!("implement support for exporting {}", c.dtype()),
-        };
+        let (arrow_array, sync_event) = export_canonical(cuda_array, ctx).await?;
 
         Ok(ArrowDeviceArray {
             array: arrow_array,
@@ -56,6 +56,85 @@ impl ExportDeviceArray for CanonicalDeviceArrayExport {
     }
 }
 
+fn export_canonical(
+    cuda_array: Canonical,
+    ctx: &mut CudaExecutionCtx,
+) -> BoxFuture<'_, VortexResult<(ArrowArray, SyncEvent)>> {
+    Box::pin(async {
+        match cuda_array {
+            Canonical::Struct(struct_array) => export_struct(struct_array, ctx).await,
+            Canonical::Primitive(primitive) => export_primitive(primitive, ctx).await,
+            c => todo!("support for exporting {} arrays", c.dtype()),
+        }
+    })
+}
+
+async fn export_struct(
+    array: StructArray,
+    ctx: &mut CudaExecutionCtx,
+) -> VortexResult<(ArrowArray, SyncEvent)> {
+    let len = array.len();
+    let StructArrayParts {
+        validity, fields, ..
+    } = array.into_parts();
+
+    let null_count = match validity {
+        Validity::NonNullable | Validity::AllValid => 0,
+        Validity::AllInvalid => len,
+        Validity::Array(_) => {
+            vortex_bail!("Exporting PrimitiveArray with non-trivial validity not supported yet")
+        }
+    };
+
+    // We need the children to be held across await points.
+    let mut children = Vec::with_capacity(fields.len());
+
+    for field in fields.iter() {
+        let cuda_field = field.clone().execute_cuda(ctx).await?;
+        let (arrow_field, _) = export_canonical(cuda_field, ctx).await?;
+        children.push(arrow_field);
+    }
+
+    let cuda_event = ctx
+        .stream()
+        .record_event(None)
+        .map_err(|_| vortex_err!("failed to create cudaEvent_t"))?;
+
+    let children = children
+        .into_iter()
+        .map(|array| Box::into_raw(Box::new(array)))
+        .collect::<Box<[_]>>();
+
+    let buffer_ptrs = vec![sys::CUdeviceptr::default()].into_boxed_slice();
+
+    let mut private_data = Box::new(PrivateData {
+        cuda_stream: Arc::clone(ctx.stream()),
+        buffers: Box::new([None]),
+        buffer_ptrs,
+        cuda_event_ptr: cuda_event.cu_event().cast(),
+        cuda_event,
+        children,
+    });
+
+    let sync_event: SyncEvent = NonNull::new(&raw mut private_data.cuda_event_ptr);
+
+    // Populate the ArrowArray with the child arrays.
+    let mut arrow_struct = ArrowArray::empty();
+    arrow_struct.length = len as i64;
+    arrow_struct.null_count = null_count as i64;
+    arrow_struct.n_children = fields.len() as i64;
+    arrow_struct.children = private_data.children.as_mut_ptr();
+
+    // StructArray _can_ contain a validity buffer. In this case, we just write a null pointer
+    // for it.
+    arrow_struct.n_buffers = 1;
+    arrow_struct.buffers = private_data.buffer_ptrs.as_mut_ptr();
+    arrow_struct.release = Some(release_array);
+    arrow_struct.private_data = Box::into_raw(private_data).cast();
+
+    Ok((arrow_struct, sync_event))
+}
+
 async fn export_primitive(
     array: PrimitiveArray,
     ctx: &mut CudaExecutionCtx,
@@ -109,6 +188,7 @@ async fn export_primitive(
 
     let mut private_data = Box::new(PrivateData {
         cuda_stream: Arc::clone(ctx.stream()),
+        children: Box::new([]),
         buffers,
         buffer_ptrs,
         cuda_event_ptr: cuda_event.cu_event().cast(),
@@ -145,7 +225,12 @@ unsafe extern "C" fn release_array(array: *mut ArrowArray) {
             std::ptr::replace(&raw mut (*array).private_data, std::ptr::null_mut());
 
         if !private_data_ptr.is_null() {
-            drop(Box::from_raw(private_data_ptr.cast::<PrivateData>()));
+            let mut private_data = Box::from_raw(private_data_ptr.cast::<PrivateData>());
+            let children = std::mem::take(&mut private_data.children);
+            for child in children {
+                release_array(child);
+            }
+            drop(private_data);
         }
 
         // update the release function to NULL to avoid any possibility of double-frees.

vortex-cuda/src/arrow/mod.rs

@@ -66,6 +66,9 @@ pub struct ArrowDeviceArray {
     _reserved: [i64; 3],
 }
 
+unsafe impl Send for ArrowDeviceArray {}
+unsafe impl Sync for ArrowDeviceArray {}
+
 /// An FFI-compatible version of the ArrowArray that holds pointers to device buffers.
 #[repr(C)]
 #[derive(Debug)]
@@ -77,7 +80,8 @@ pub(crate) struct ArrowArray {
     n_children: i64,
     buffers: *mut sys::CUdeviceptr,
     children: *mut *mut ArrowArray,
-    dictionary: *mut ArrowArray,
+    // NOTE: we don't support exporting dictionary arrays, so we leave this as an opaque pointer.
+    dictionary: *mut (),
     release: Option<unsafe extern "C" fn(arg1: *mut ArrowArray)>,
     // When exported, this MUST contain everything that is owned by this array.
     // for example, any buffer pointed to in `buffers` must be here, as well
@@ -105,6 +109,9 @@ impl ArrowArray {
     }
 }
 
+unsafe impl Send for ArrowArray {}
+unsafe impl Sync for ArrowArray {}
+
 #[expect(
     unused,
     reason = "cuda_stream and cuda_buffers need to have deferred drop"
@@ -120,6 +127,7 @@ pub(crate) struct PrivateData {
     pub(crate) buffer_ptrs: Box<[sys::CUdeviceptr]>,
     pub(crate) cuda_event: CudaEvent,
     pub(crate) cuda_event_ptr: cudaEvent_t,
+    pub(crate) children: Box<[*mut ArrowArray]>,
 }
 
 #[async_trait]