feat: add CUDA cuDF convenience API

Add vortex_cuda.to_cudf and install optional Vortex Array helpers for
cuDF conversion and Arrow C Device export. Keep the conversion path
CUDA-only by rejecting unsupported fallback policies and routing cuDF
ingestion through fresh Arrow C Device capsules.

Expand CUDA Python tests for the convenience API, installed Array
methods, Arrow Device export smoke coverage, and capsule ownership
paths.

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

Author: 0ax1

vortex-cuda/src/arrow/canonical.rs

@@ -764,7 +764,7 @@ fn gather_binary_values(
 ///
 /// Returns `None` for the buffer when Arrow can omit validity because all rows are valid.
 ///
-/// Returned buffers use zeroed 4-byte padding so cuDF's word-sized mask reads stay in bounds.
+/// Returned buffers use zeroed cuDF-sized padding so mask reads stay in bounds.
 /// Bits at positions `>= len + arrow_offset` within the final data byte are unspecified, as
 /// Arrow permits.
 pub(super) async fn export_arrow_validity_buffer(
@@ -798,21 +798,26 @@ pub(super) async fn export_arrow_validity_buffer(
             let bitmap = ctx.ensure_on_device(bits).await?;
             // ArrowDeviceArray uses ArrowArray layout with its buffers being device pointers.
             //
-            // Validity is one bit per row, addressed via the Arrow array offset. Reuse the bitmap
-            // when Vortex's validity offset already matches Arrow's; otherwise repack on the GPU
-            // so row i is at Arrow bit `arrow_offset + i`.
-            let bitmap = if meta.offset() == arrow_offset {
-                bitmap
-            } else {
-                repack_arrow_validity_buffer(&bitmap, meta.offset(), len, arrow_offset, ctx)?
-            };
+            // Validity is one bit per row, addressed via the Arrow array offset. Repack on the GPU
+            // so row i is at Arrow bit `arrow_offset + i` and the backing allocation has the
+            // zeroed cuDF-sized padding expected by Arrow Device consumers.
+            let bitmap =
+                repack_arrow_validity_buffer(&bitmap, meta.offset(), len, arrow_offset, ctx)?;
             // Keep nullable exports self-describing for consumers that require exact null counts.
             let null_count = count_arrow_validity_nulls(&bitmap, len, arrow_offset, ctx)?;
             Ok((Some(bitmap), null_count))
         }
     }
 }
 
+/// Minimum backing allocation quantum for Arrow validity buffers handed to cuDF.
+///
+/// Arrow exposes only the logical bitmap byte length, but cuDF imports null masks into 64-byte
+/// padded buffers and its kernels may read through that padded extent. Vortex therefore keeps the
+/// exported `BufferHandle` sliced to Arrow's logical length while zero-padding the underlying CUDA
+/// allocation to this boundary.
+const CUDF_VALIDITY_BUFFER_PADDING: usize = 64;
+
 /// Return the byte length needed for `len` validity bits at the given bit offset.
 fn validity_bitmap_byte_len(len: usize, arrow_offset: usize) -> VortexResult<usize> {
     Ok(len
@@ -821,12 +826,25 @@ fn validity_bitmap_byte_len(len: usize, arrow_offset: usize) -> VortexResult<usi
         .div_ceil(8))
 }
 
+/// Return the CUDA allocation size for a logical Arrow validity bitmap byte length.
+///
+/// The returned allocation length may be larger than `byte_len`; callers slice the exported
+/// `BufferHandle` back to `byte_len` while retaining the padded backing allocation. A zero-length
+/// bitmap still gets one byte so we never request a zero-sized CUDA allocation.
+fn validity_bitmap_allocation_byte_len(byte_len: usize) -> usize {
+    if byte_len == 0 {
+        1
+    } else {
+        byte_len.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING)
+    }
+}
+
 /// Allocate a zeroed device buffer with cuDF-safe padding for Arrow validity masks.
 fn device_zeroed_byte_buffer(
     byte_len: usize,
     ctx: &mut CudaExecutionCtx,
 ) -> VortexResult<BufferHandle> {
-    let allocation_len = byte_len.next_multiple_of(size_of::<u32>()).max(1);
+    let allocation_len = validity_bitmap_allocation_byte_len(byte_len);
     let mut buffer = ctx.device_alloc::<u8>(allocation_len)?;
     ctx.stream()
         .memset_zeros(&mut buffer)
@@ -894,8 +912,8 @@ pub fn count_arrow_validity_nulls(
 ///
 /// Vortex bitmaps may start at any bit offset. Arrow exposes only a byte-addressed validity buffer
 /// plus an array offset, so sliced compact exports need a GPU rewrite when either side has a
-/// bit-level offset. The kernel writes the output one 64-bit word at a time, funnel-shifting two
-/// adjacent input words, so the allocation is padded to whole words (zeroed by the edge masks).
+/// bit-level offset. The output handle keeps Arrow's logical byte length, while the backing
+/// allocation is zero-padded to cuDF's mask allocation size for consumers that read full masks.
 pub fn repack_arrow_validity_buffer(
     input_buffer: &BufferHandle,
     input_offset: usize,
@@ -904,7 +922,13 @@ pub fn repack_arrow_validity_buffer(
     ctx: &mut CudaExecutionCtx,
 ) -> VortexResult<BufferHandle> {
     let output_bytes = validity_bitmap_byte_len(len, arrow_offset)?;
+    // The CUDA kernel writes the bitmap as u64 words, so round the logical byte length up to the
+    // number of words that cover the exported Arrow bytes.
     let output_words = output_bytes.div_ceil(size_of::<u64>());
+    // `device_alloc::<u64>` takes a word count, while the padding policy is expressed in bytes.
+    // Round up so the padded byte allocation is fully represented by whole u64 words.
+    let allocation_words =
+        validity_bitmap_allocation_byte_len(output_bytes).div_ceil(size_of::<u64>());
 
     // The kernel loads the input bitmap as 64-bit words.
     if !input_buffer
@@ -914,7 +938,12 @@ pub fn repack_arrow_validity_buffer(
         vortex_bail!("Arrow validity repack requires an 8-byte aligned device buffer");
     }
 
-    let output = ctx.device_alloc::<u64>(output_words.max(1))?;
+    let mut output = ctx.device_alloc::<u64>(allocation_words.max(1))?;
+    // The repack kernel writes only the logical bitmap words. Zero the whole backing allocation so
+    // cuDF's padded mask reads see invalid rows, not uninitialized CUDA memory.
+    ctx.stream()
+        .memset_zeros(&mut output)
+        .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer padding: {err}"))?;
     let output_device = CudaDeviceBuffer::new(output);
 
     if output_words > 0 {
@@ -1337,6 +1366,7 @@ mod tests {
     use crate::arrow::ArrowDeviceArray;
     use crate::arrow::DeviceArrayExt;
     use crate::arrow::PrivateData;
+    use crate::arrow::canonical::CUDF_VALIDITY_BUFFER_PADDING;
     use crate::arrow::canonical::export_arrow_validity_buffer;
     use crate::arrow::canonical::repack_arrow_validity_buffer;
     use crate::device_buffer::cuda_backing_allocation;
@@ -2955,7 +2985,43 @@ mod tests {
         let backing_bytes = backing.to_host_sync();
         assert_eq!(
             backing_bytes.len(),
-            output_bytes.next_multiple_of(size_of::<u64>())
+            output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING)
+        );
+        assert!(backing_bytes[output_bytes..].iter().all(|byte| *byte == 0));
+
+        Ok(())
+    }
+
+    #[crate::test]
+    async fn test_export_validity_buffer_pads_matching_offset() -> VortexResult<()> {
+        let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())
+            .vortex_expect("failed to create execution context");
+
+        let len = 3;
+        let arrow_offset = 0;
+        let (buffer, null_count) = export_arrow_validity_buffer(
+            Validity::from(BitBuffer::from_iter([true, false, true])),
+            len,
+            arrow_offset,
+            &mut ctx,
+        )
+        .await?;
+        ctx.synchronize_stream()?;
+
+        assert_eq!(null_count, 1);
+        let buffer = buffer.vortex_expect("nullable validity should export a null buffer");
+        let output_bytes = (len + arrow_offset).div_ceil(8);
+        assert_eq!(buffer.len(), output_bytes);
+        let actual = BitBuffer::new(buffer.to_host_sync(), len + arrow_offset)
+            .iter()
+            .collect::<Vec<_>>();
+        assert_eq!(actual, [true, false, true]);
+
+        let backing = cuda_backing_allocation(&buffer)?;
+        let backing_bytes = backing.to_host_sync();
+        assert_eq!(
+            backing_bytes.len(),
+            output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING)
         );
         assert!(backing_bytes[output_bytes..].iter().all(|byte| *byte == 0));
 
@@ -2983,6 +3049,11 @@ mod tests {
         let bytes = buffer.to_host_sync();
         assert_eq!(bytes.len(), (len + arrow_offset).div_ceil(8));
         assert!(bytes.iter().all(|byte| *byte == 0));
+        let backing = cuda_backing_allocation(&buffer)?;
+        assert_eq!(
+            backing.len(),
+            bytes.len().next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING)
+        );
 
         Ok(())
     }

vortex-cuda/src/arrow/mod.rs

@@ -140,6 +140,16 @@ pub(crate) struct PrivateData {
 }
 
 impl PrivateData {
+    fn retain_buffers(&self, buffers: &mut Vec<BufferHandle>) -> VortexResult<()> {
+        // `BufferHandle::clone` is shallow for device buffers; this retains the Arc-backed CUDA
+        // allocations for cuDF without copying device memory.
+        buffers.extend(self.buffers.iter().flatten().cloned());
+        for child in &self.children {
+            retain_arrow_array_buffers(unsafe { child.as_ref() }, buffers)?;
+        }
+        retain_arrow_array_buffers(unsafe { self.dictionary.as_ref() }, buffers)
+    }
+
     /// Create private data for arrays that own buffers and child arrays but no dictionary.
     pub(crate) fn new(
         buffers: Vec<Option<BufferHandle>>,
@@ -204,6 +214,49 @@ impl PrivateData {
     }
 }
 
+/// Cloned device buffers retained outside an Arrow Device release callback.
+///
+/// This is useful for Arrow Device consumers that consume and release the Arrow C structs during
+/// import but keep zero-copy views over the producer's device buffers in their returned object.
+#[derive(Debug)]
+pub struct ArrowDeviceArrayBufferKeepAlive {
+    #[allow(
+        dead_code,
+        reason = "buffers are retained until this keepalive is dropped"
+    )]
+    buffers: Vec<BufferHandle>,
+}
+
+/// Clone the device buffers owned by a Vortex-created Arrow Device array.
+///
+/// The returned owner must be dropped only after consumers are finished with any zero-copy views over
+/// `array`'s buffers. Call this before the consumer invokes the Arrow release callback.
+///
+/// # Safety
+///
+/// `array` must be a live Arrow Device array exported by Vortex. Passing a foreign Arrow array, or a
+/// Vortex array whose release callback has already run, may dereference invalid private data.
+pub unsafe fn retain_arrow_device_array_buffers(
+    array: &ArrowDeviceArray,
+) -> VortexResult<ArrowDeviceArrayBufferKeepAlive> {
+    let mut buffers = Vec::new();
+    retain_arrow_array_buffers(Some(&array.array), &mut buffers)?;
+    Ok(ArrowDeviceArrayBufferKeepAlive { buffers })
+}
+
+fn retain_arrow_array_buffers(
+    array: Option<&ArrowArray>,
+    buffers: &mut Vec<BufferHandle>,
+) -> VortexResult<()> {
+    let Some(array) = array else {
+        return Ok(());
+    };
+    let Some(private_data) = (unsafe { array.private_data.cast::<PrivateData>().as_ref() }) else {
+        return Ok(());
+    };
+    private_data.retain_buffers(buffers)
+}
+
 /// A Vortex array exported as an Arrow schema and Arrow Device array pair.
 #[derive(Debug)]
 pub struct ArrowDeviceArrayWithSchema {