Add CUDA Arrow device schema export

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

.github/workflows/ci.yml

@@ -270,7 +270,7 @@ jobs:
         run: |
           git ls-files vortex-cuda vortex-cxx vortex-duckdb vortex-ffi \
             | grep -E '\.(cpp|hpp|cu|cuh|h)$' \
-            | grep -v 'arrow/reference/abi\.h$' \
+            | grep -v 'arrow/reference/arrow_c_device\.h$' \
             | grep -v 'kernels/src/bit_unpack_.*\.cu$' \
             | grep -v 'kernels/src/bit_unpack_.*_lanes\.cuh$' \
             | xargs clang-format --dry-run --Werror --style=file

Cargo.lock

@@ -46,3 +46,8 @@ path = ["vortex-cuda/kernels/src/bit_unpack_*"]
 precedence = "override"
 SPDX-FileCopyrightText = "Copyright the Vortex contributors"
 SPDX-License-Identifier = "Apache-2.0"
+
+[[annotations]]
+path = "vortex-cuda/src/arrow/reference/arrow_c_device.h"
+SPDX-FileCopyrightText = "2016-2025 Copyright The Apache Software Foundation"
+SPDX-License-Identifier = "Apache-2.0"

REUSE.toml

@@ -24,6 +24,7 @@ unstable_encodings = ["vortex/unstable_encodings"]
 
 [dependencies]
 arc-swap = { workspace = true }
+arrow-schema = { workspace = true, features = ["ffi"] }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 cudarc = { workspace = true, features = ["f16"] }

vortex-cuda/Cargo.toml

@@ -56,7 +56,7 @@ fn main() {
     generate_unpack::<u64>(&kernels_src, 16).expect("Failed to generate unpack for u64");
 
     let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set"));
-    generate_arrow_c_abi_bindings(Path::new(&manifest_dir), &out_dir);
+    generate_arrow_device_array_bindings(Path::new(&manifest_dir), &out_dir);
     generate_dynamic_dispatch_bindings(&kernels_src, &out_dir);
     generate_patches_bindings(&kernels_src, &out_dir);
 
@@ -197,9 +197,9 @@ fn nvcc_compile_ptx(
     Ok(())
 }
 
-/// Generate bindings for the vendored Arrow C ABI header.
-fn generate_arrow_c_abi_bindings(manifest_dir: &Path, out_dir: &Path) {
-    let header = manifest_dir.join("src/arrow/reference/abi.h");
+/// Generate bindings for the vendored Arrow C Device ABI header.
+fn generate_arrow_device_array_bindings(manifest_dir: &Path, out_dir: &Path) {
+    let header = manifest_dir.join("src/arrow/reference/arrow_c_device.h");
     println!("cargo:rerun-if-changed={}", header.display());
 
     let bindings = bindgen::Builder::default()
@@ -212,7 +212,7 @@ fn generate_arrow_c_abi_bindings(manifest_dir: &Path, out_dir: &Path) {
         .derive_debug(true)
         .layout_tests(false)
         .generate()
-        .expect("Failed to generate Arrow C ABI bindings");
+        .expect("Failed to generate Arrow C Device bindings");
 
     bindings
         .write_to_file(out_dir.join("arrow_c_abi.rs"))

vortex-cuda/build.rs

@@ -15,6 +15,7 @@ use vortex::array::arrays::decimal::DecimalDataParts;
 use vortex::array::arrays::extension::ExtensionArrayExt;
 use vortex::array::arrays::primitive::PrimitiveDataParts;
 use vortex::array::arrays::struct_::StructDataParts;
+use vortex::array::arrays::varbinview::VarBinViewDataParts;
 use vortex::array::buffer::BufferHandle;
 use vortex::dtype::DecimalType;
 use vortex::error::VortexResult;
@@ -30,8 +31,6 @@ use crate::arrow::ExportDeviceArray;
 use crate::arrow::PrivateData;
 use crate::arrow::SyncEvent;
 use crate::arrow::check_validity_empty;
-use crate::arrow::varbinview::BinaryParts;
-use crate::arrow::varbinview::copy_varbinview_to_varbin;
 use crate::executor::CudaArrayExt;
 
 /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by
@@ -49,13 +48,13 @@ impl ExportDeviceArray for CanonicalDeviceArrayExport {
     ) -> VortexResult<ArrowDeviceArray> {
         let cuda_array = array.execute_cuda(ctx).await?;
 
-        let (arrow_array, _) = export_canonical(cuda_array, ctx).await?;
+        let (arrow_array, sync_event) = export_canonical(cuda_array, ctx).await?;
 
         Ok(ArrowDeviceArray {
             array: arrow_array,
             device_id: ctx.stream().context().ordinal() as i64,
             device_type: ARROW_DEVICE_CUDA,
-            sync_event: ptr::null_mut(),
+            sync_event,
             reserved: Default::default(),
         })
     }
@@ -148,24 +147,32 @@ fn export_canonical(
             }
             Canonical::VarBinView(varbinview) => {
                 let len = varbinview.len();
-                check_validity_empty(&varbinview.validity()?)?;
+                let VarBinViewDataParts {
+                    views,
+                    buffers: data_buffers,
+                    validity,
+                    ..
+                } = varbinview.into_data_parts();
 
-                let BinaryParts { offsets, bytes } =
-                    copy_varbinview_to_varbin(varbinview, ctx).await?;
+                check_validity_empty(&validity)?;
 
-                let offsets = ctx.ensure_on_device(offsets).await?;
-                let bytes = ctx.ensure_on_device(bytes).await?;
+                let views = ctx.ensure_on_device(views).await?;
+                let mut buffers = Vec::with_capacity(data_buffers.len() + 2);
+                buffers.push(None);
+                buffers.push(Some(views));
+                for buffer in data_buffers.iter() {
+                    buffers.push(Some(ctx.ensure_on_device(buffer.clone()).await?));
+                }
 
-                let buffers = vec![None, Some(offsets), Some(bytes)];
+                let n_buffers = i64::try_from(buffers.len())?;
                 let mut private_data = PrivateData::new(buffers, vec![], ctx)?;
                 let sync_event = private_data.sync_event();
-                //
                 let arrow_array = ArrowArray {
                     length: len as i64,
                     null_count: 0,
                     offset: 0,
-                    // Arrow Utf8/Binary layout: optional null bitmap, offsets, and data bytes.
-                    n_buffers: 3,
+                    // Arrow Utf8View/BinaryView layout: optional null bitmap, views, and data buffers.
+                    n_buffers,
                     buffers: private_data.buffer_ptrs.as_mut_ptr(),
                     n_children: 0,
                     children: ptr::null_mut(),
@@ -232,8 +239,8 @@ fn export_fixed_size(
         "buffer must already be copied to device before calling"
     );
 
-    // TODO(aduffy): currently the null buffer is always None, in the future we will need
-    //  to pass it.
+    // Non-trivial validity is rejected before fixed-size export, so the Arrow null bitmap slot is
+    // always null for now. Future nullable export support should pass the validity bitmap here.
     let mut private_data = PrivateData::new(vec![None, Some(buffer)], vec![], ctx)?;
     let sync_event: SyncEvent = private_data.sync_event();
 
@@ -271,7 +278,9 @@ unsafe extern "C" fn release_array(array: *mut ArrowArray) {
             let children = mem::take(&mut private_data.children);
             for child in children {
                 if !child.is_null() {
-                    release_array(child);
+                    if let Some(release) = (*child).release {
+                        release(child);
+                    }
                     // Children are allocated with Box::into_raw in PrivateData::new, so the
                     // release callback must also reclaim the ArrowArray allocation itself.
                     drop(Box::from_raw(child));
@@ -286,6 +295,9 @@ unsafe extern "C" fn release_array(array: *mut ArrowArray) {
 
 #[cfg(test)]
 mod tests {
+    use arrow_schema::DataType;
+    use arrow_schema::Field;
+    use arrow_schema::Schema;
     use rstest::rstest;
     use vortex::array::ArrayRef;
     use vortex::array::IntoArray;
@@ -432,14 +444,10 @@ mod tests {
 
         assert_eq!(device_array.array.length, 3);
         assert_eq!(device_array.array.null_count, 0);
-        // VarBin export: null buffer + offsets + data
+        // VarBinView export: null buffer + views + data buffers
         assert_eq!(device_array.array.n_buffers, 3);
-        let buffers = unsafe {
-            std::slice::from_raw_parts(
-                device_array.array.buffers,
-                device_array.array.n_buffers as usize,
-            )
-        };
+        let n_buffers = usize::try_from(device_array.array.n_buffers)?;
+        let buffers = unsafe { std::slice::from_raw_parts(device_array.array.buffers, n_buffers) };
         assert!(buffers[0].is_null());
         assert!(!buffers[1].is_null());
         assert!(!buffers[2].is_null());
@@ -479,4 +487,84 @@ mod tests {
         unsafe { release_array(&raw mut device_array.array) };
         Ok(())
     }
+
+    #[crate::test]
+    async fn test_export_struct_with_schema() -> VortexResult<()> {
+        let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create execution context");
+
+        let array = StructArray::new(
+            FieldNames::from_iter(["a", "b", "c"]),
+            vec![
+                PrimitiveArray::from_iter(0u32..5).into_array(),
+                PrimitiveArray::from_iter(0i64..5).into_array(),
+                VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"])
+                    .into_array(),
+            ],
+            5,
+            Validity::NonNullable,
+        )
+        .into_array();
+        let mut exported = array.export_device_array_with_schema(&mut ctx).await?;
+
+        let schema = Schema::try_from(&exported.schema)?;
+        assert_eq!(
+            schema,
+            Schema::new(vec![
+                Field::new("a", DataType::UInt32, false),
+                Field::new("b", DataType::Int64, false),
+                Field::new("c", DataType::Utf8View, false),
+            ])
+        );
+        assert_eq!(exported.array.array.length, 5);
+        assert_eq!(exported.array.array.n_buffers, 1);
+        assert_eq!(exported.array.array.n_children, 3);
+        assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA);
+
+        unsafe { release_array(&raw mut exported.array.array) };
+        Ok(())
+    }
+
+    #[crate::test]
+    async fn test_export_primitive_with_schema_is_column_shaped() -> VortexResult<()> {
+        let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create execution context");
+
+        let array = PrimitiveArray::from_iter(0u32..5).into_array();
+        let mut exported = array.export_device_array_with_schema(&mut ctx).await?;
+
+        let field = Field::try_from(&exported.schema)?;
+        assert_eq!(field, Field::new("", DataType::UInt32, false));
+        assert_eq!(exported.array.array.length, 5);
+        assert_eq!(exported.array.array.n_buffers, 2);
+        assert_eq!(exported.array.array.n_children, 0);
+        assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA);
+
+        unsafe { release_array(&raw mut exported.array.array) };
+        Ok(())
+    }
+
+    #[crate::test]
+    async fn test_export_varbinview_with_schema_uses_utf8_view_layout() -> VortexResult<()> {
+        let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create execution context");
+
+        let array = VarBinViewArray::from_iter_str([
+            "one",
+            "two",
+            "this is a longer string for out-of-line storage",
+        ])
+        .into_array();
+        let mut exported = array.export_device_array_with_schema(&mut ctx).await?;
+
+        let field = Field::try_from(&exported.schema)?;
+        assert_eq!(field, Field::new("", DataType::Utf8View, false));
+        assert_eq!(exported.array.array.length, 3);
+        assert_eq!(exported.array.array.n_buffers, 3);
+        assert_eq!(exported.array.array.n_children, 0);
+        assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA);
+
+        unsafe { release_array(&raw mut exported.array.array) };
+        Ok(())
+    }
 }