more test

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

Author: a10y

Cargo.lock

@@ -20,6 +20,13 @@ version = { workspace = true }
 workspace = true
 
 [dependencies]
+arrow-array = { workspace = true, features = ["ffi"] }
+arrow-schema = { workspace = true, features = ["ffi"] }
+futures = { workspace = true, features = ["executor"] }
+vortex-array = { workspace = true }
+vortex-buffer = { workspace = true }
+vortex-cuda = { path = "..", features = ["_test-harness"] }
+vortex-session = { workspace = true }
 
 [build-dependencies]
 bindgen = { workspace = true }

vortex-cuda/cudf-test/Cargo.toml

@@ -61,6 +61,7 @@ struct ArrowDeviceArray {
     int64_t device_id;
     ArrowDeviceType device_type;
     void* sync_event;
+    int64_t reserved[3];
 };
 
 // Error type: NULL on success, pointer to error string on failure.

vortex-cuda/cudf-test/cpp/cudf_arrow_ffi.h

@@ -195,6 +195,19 @@ impl Drop for CudfColumnView {
 
 #[cfg(test)]
 mod tests {
+    use arrow_array::ffi::FFI_ArrowSchema;
+    use arrow_schema::DataType;
+    use futures::executor::block_on;
+    use vortex_array::Canonical;
+    use vortex_array::IntoArray;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::validity::Validity;
+    use vortex_buffer::Buffer;
+    use vortex_cuda::CudaSession;
+    use vortex_cuda::arrow::CudaDeviceArrayExecute;
+    use vortex_cuda::executor::CudaArrayExt;
+    use vortex_session::VortexSession;
+
     use super::*;
 
     #[test]
@@ -210,4 +223,63 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_primitive_array_to_cudf_tableview() -> Result<()> {
+        // Create a PrimitiveArray with 100 i64 values
+        let data: Vec<i64> = (0..100).collect();
+        let expected_len = data.len();
+        let primitive_array =
+            PrimitiveArray::new(Buffer::from(data), Validity::NonNullable).into_array();
+
+        // Create CUDA execution context
+        let mut cuda_ctx = match CudaSession::create_execution_ctx(&VortexSession::empty()).unwrap();
+
+        // Export as ArrowDeviceArray using CudaDeviceArrayExecute
+        let device_array = block_on(Canonical::execute(
+            &primitive_array,
+            primitive_array.clone(),
+            &mut cuda_ctx,
+        ))
+        .unwrap();
+
+        // Synchronize the CUDA stream to ensure the data is ready
+        cuda_ctx.synchronize_stream().map_err(|e| CudfError {
+            message: e.to_string(),
+        })?;
+
+        // Create FFI_ArrowSchema from the data type
+        let mut ffi_schema =
+            FFI_ArrowSchema::try_from(&DataType::Int64).map_err(|e| CudfError {
+                message: format!("Failed to create FFI schema: {}", e),
+            })?;
+
+        // Create cudf context
+        let cudf_ctx = CudfContext::new()?;
+
+        // Import into cudf tableview
+        let tableview = unsafe {
+            cudf_ctx.tableview_from_device(
+                (&raw mut ffi_schema).cast::<ArrowSchema>(),
+                (&raw const device_array).cast::<ArrowDeviceArray>(),
+            )?
+        };
+
+        // Verify row count
+        let num_rows = tableview.num_rows()?;
+        assert_eq!(num_rows, expected_len as i64, "Row count mismatch");
+        println!(
+            "Successfully imported PrimitiveArray into cudf tableview with {} rows",
+            num_rows
+        );
+
+        // Verify column count (should be 1 for a primitive array)
+        let num_columns = tableview.num_columns()?;
+        assert_eq!(num_columns, 1, "Column count mismatch");
+        println!("Tableview has {} column(s)", num_columns);
+
+        // Tableview and cudf_ctx will be deallocated automatically via Drop
+
+        Ok(())
+    }
 }

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

@@ -3,6 +3,7 @@
 
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use cudarc::driver::sys;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
@@ -22,7 +23,7 @@ use crate::arrow::CudaPrivateData;
 use crate::arrow::DeviceType;
 use crate::executor::CudaArrayExt;
 
-// Impl it for the execution context instead here...I think this is right?
+#[async_trait]
 impl CudaDeviceArrayExecute for Canonical {
     async fn execute(
         &self,
@@ -32,7 +33,7 @@ impl CudaDeviceArrayExecute for Canonical {
         let cuda_array = array.execute_cuda(ctx).await?;
 
         let arrow_array = match cuda_array {
-            Canonical::Primitive(primitive) => export_primitive(primitive, ctx)?,
+            Canonical::Primitive(primitive) => export_primitive(primitive, ctx).await?,
             c => todo!("implement support for exporting {}", c.dtype()),
         };
 
@@ -46,19 +47,25 @@ impl CudaDeviceArrayExecute for Canonical {
     }
 }
 
-fn export_primitive(array: PrimitiveArray, ctx: &mut CudaExecutionCtx) -> VortexResult<ArrowArray> {
+async fn export_primitive(array: PrimitiveArray, ctx: &mut CudaExecutionCtx) -> VortexResult<ArrowArray> {
+    unsafe extern "C" fn release(array: *mut ArrowArray) {
+        // SAFETY: this is only safe if the caller provides a valid pointer to an `ArrowArray`.
+        drop(unsafe { Box::from_raw(array) });
+    }
+
     let len = array.len();
     let PrimitiveArrayParts {
         buffer,
-        ptype,
         validity,
         ..
     } = array.into_parts();
 
-    unsafe extern "C" fn release(array: *mut ArrowArray) {
-        // SAFETY: this is only safe if the caller provides a valid pointer to an `ArrowArray`.
-        drop(unsafe { Box::from_raw(array) });
-    }
+    let buffer = if buffer.is_on_device() {
+        buffer
+    } else {
+        // TODO(aduffy): I don't think this type parameter does anything
+        ctx.move_to_device::<u8>(buffer)?.await?
+    };
 
     let null_count = match validity {
         Validity::NonNullable | Validity::AllValid => 0,
@@ -106,4 +113,3 @@ fn export_primitive(array: PrimitiveArray, ctx: &mut CudaExecutionCtx) -> Vortex
         private_data: Box::into_raw(private_data).cast(),
     })
 }
-

vortex-cuda/src/arrow/canonical.rs

@@ -14,17 +14,15 @@ use std::ffi::c_void;
 use std::ptr::NonNull;
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use cudarc::driver::CudaStream;
 use cudarc::driver::sys;
 use cudarc::runtime::sys::cudaEvent_t;
 use vortex_array::ArrayRef;
-use vortex_array::Executable;
 use vortex_array::buffer::BufferHandle;
 use vortex_error::VortexResult;
 
 use crate::CudaExecutionCtx;
-use crate::executor::CudaArrayExt;
-use crate::executor::CudaExecute;
 
 #[derive(Debug, Copy, Clone)]
 #[repr(i32)]
@@ -53,7 +51,7 @@ pub type SyncEvent = Option<NonNull<cudaEvent_t>>;
 /// event that the client must wait on.
 #[repr(C)]
 #[derive(Debug)]
-pub(crate) struct ArrowDeviceArray {
+pub struct ArrowDeviceArray {
     array: ArrowArray,
     device_id: i64,
     device_type: DeviceType,
@@ -85,6 +83,7 @@ pub(crate) struct ArrowArray {
 }
 
 impl ArrowArray {
+    #[allow(unused)]
     pub fn empty() -> Self {
         Self {
             length: 0,
@@ -101,6 +100,7 @@ impl ArrowArray {
     }
 }
 
+#[expect(unused, reason = "cuda_stream and cuda_buffers need to have deferred drop")]
 pub(crate) struct CudaPrivateData {
     /// Hold a reference to the CudaStream so that it stays alive even after CudaExecutionCtx
     /// has been dropped.
@@ -113,7 +113,8 @@ pub(crate) struct CudaPrivateData {
 }
 
 /// Trait implemented for types that can be exported to [`ArrowDeviceArray`].
-pub(crate) trait CudaDeviceArrayExecute {
+#[async_trait]
+pub trait CudaDeviceArrayExecute {
     async fn execute(
         &self,
         array: ArrayRef,

vortex-cuda/src/arrow/mod.rs

@@ -12,7 +12,6 @@ use cudarc::driver::DevicePtr;
 use cudarc::driver::DeviceRepr;
 use cudarc::driver::sys;
 use futures::future::BoxFuture;
-use futures::future::ok;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::buffer::DeviceBuffer;
 use vortex_buffer::Alignment;

vortex-cuda/src/device_buffer.rs

@@ -5,7 +5,7 @@
 
 use std::process::Command;
 
-mod arrow;
+pub mod arrow;
 mod canonical;
 mod device_buffer;
 pub mod executor;