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

@@ -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

@@ -49,13 +49,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(),
         })
     }
@@ -232,8 +232,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 +271,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 +288,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;
@@ -479,4 +484,79 @@ 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::Utf8, 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_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", "three"]).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::Utf8, 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(())
+    }
 }