Add Python CUDA bridge CI and buffer handoff ABI

Add explicit GPU-runner CI coverage for the Python CUDA bridge through
the vortex-data[cuda] optional-extra path.

Extend the private metadata bridge to carry host buffer-export capsules
instead of only a buffer count. The base Python package exports repr(C)
VortexBufferExport descriptors, and vortex-python-cuda imports them into
local BufferHandles before deserializing arrays through its own
VortexSession.

Tests now cover primitive, nullable, bool, and struct arrays across the
bridge, plus the existing CUDA Arrow Device smoke path.

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

Author: 0ax1

.github/workflows/ci.yml

@@ -132,6 +132,30 @@ jobs:
           uv run --all-packages make html
         working-directory: docs/
 
+  python-cuda-test:
+    name: "Python CUDA (test)"
+    if: github.repository == 'vortex-data/vortex'
+    runs-on: >-
+      ${{ format('runs-on={0}/runner=gpu/tag=python-cuda-test', github.run_id) }}
+    timeout-minutes: 30
+    env:
+      RUST_LOG: "info,maturin=off,uv=debug"
+      MATURIN_PEP517_ARGS: "--profile ci"
+    steps:
+      - uses: runs-on/action@v2
+        with:
+          sccache: s3
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0  # v7
+      - uses: ./.github/actions/setup-prebuild
+
+      - name: Pytest - PyVortex CUDA bridge
+        run: |
+          uv run --extra cuda \
+            --reinstall-package vortex-data \
+            --reinstall-package vortex-data-cuda \
+            pytest --benchmark-disable ../vortex-python-cuda/test/test_native_bridge.py
+        working-directory: vortex-python/
+
   rust-docs:
     name: "Rust (docs)"
     needs: duckdb-ready

Cargo.lock

@@ -28,6 +28,7 @@ extension-module = []
 
 [dependencies]
 arrow-schema = { workspace = true }
+bytes = { workspace = true }
 pyo3 = { workspace = true, features = ["abi3", "abi3-py311"] }
 vortex = { workspace = true }
 vortex-cuda = { workspace = true }

vortex-python-cuda/Cargo.toml

@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright the Vortex contributors
+# pyright: reportMissingModuleSource=false, reportPrivateUsage=false
 
-from ._lib import (  # pyright: ignore[reportMissingModuleSource]
-    _debug_array_metadata_dtype as _debug_array_metadata_dtype,  # pyright: ignore[reportPrivateUsage]
-)
-from ._lib import (  # pyright: ignore[reportMissingModuleSource]
-    cuda_available,
-    export_device_array,
-)
+from . import _lib
+
+_debug_array_metadata_dtype = _lib._debug_array_metadata_dtype
+_debug_array_metadata_scalar_values = _lib._debug_array_metadata_scalar_values
+cuda_available = _lib.cuda_available
+export_device_array = _lib.export_device_array
 
 __all__ = ["cuda_available", "export_device_array"]

vortex-python-cuda/python/vortex_cuda/__init__.py

@@ -2,6 +2,7 @@
 #  SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 def _debug_array_metadata_dtype(array: object) -> str: ...
+def _debug_array_metadata_scalar_values(array: object) -> list[str]: ...
 def cuda_available() -> bool: ...
 def export_device_array(
     array: object, requested_schema: object | None = None, **kwargs: object

vortex-python-cuda/python/vortex_cuda/_lib.pyi

@@ -35,7 +35,6 @@ use vortex::buffer::ByteBuffer;
 use vortex::dtype::DType;
 use vortex::error::VortexError;
 use vortex::error::VortexResult;
-use vortex::error::vortex_bail;
 use vortex::error::vortex_ensure;
 use vortex::error::vortex_err;
 use vortex::flatbuffers::FlatBuffer;
@@ -54,6 +53,117 @@ const ARROW_SCHEMA_CAPSULE_NAME: &CStr = c_str!("arrow_schema");
 const USED_ARROW_SCHEMA_CAPSULE_NAME: &CStr = c_str!("used_arrow_schema");
 const ARROW_DEVICE_ARRAY_CAPSULE_NAME: &CStr = c_str!("arrow_device_array");
 const USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME: &CStr = c_str!("used_arrow_device_array");
+const BUFFER_EXPORT_CAPSULE_NAME: &CStr = c_str!("vortex_buffer_export");
+
+/// Private C-ABI struct matching the layout in `vortex-python/src/arrays/mod.rs`.
+/// Both crates define this identically; the exact-version pin ensures layout agreement.
+#[repr(C)]
+struct VortexBufferExport {
+    version: u32,
+    kind: u32,
+    ptr: *const u8,
+    len: usize,
+    alignment: usize,
+    device_id: i32,
+    sync_event: *mut c_void,
+    private_data: *mut c_void,
+    release: Option<unsafe extern "C" fn(*mut VortexBufferExport)>,
+}
+
+const VORTEX_BUFFER_HOST: u32 = 0;
+
+struct BufferExportGuard {
+    export: NonNull<VortexBufferExport>,
+}
+
+impl BufferExportGuard {
+    fn export(&self) -> &VortexBufferExport {
+        unsafe { self.export.as_ref() }
+    }
+}
+
+impl AsRef<[u8]> for BufferExportGuard {
+    fn as_ref(&self) -> &[u8] {
+        let export = self.export();
+        if export.len == 0 {
+            &[]
+        } else {
+            unsafe { std::slice::from_raw_parts(export.ptr, export.len) }
+        }
+    }
+}
+
+impl Drop for BufferExportGuard {
+    fn drop(&mut self) {
+        let mut export = unsafe { Box::from_raw(self.export.as_ptr()) };
+        if let Some(release) = export.release.take() {
+            unsafe { release(&raw mut *export) };
+        }
+    }
+}
+
+// The guard is moved into `Bytes::from_owner`, which requires `Send + Sync`. After import we disable
+// the source capsule destructor and own the boxed C export until this guard is dropped.
+unsafe impl Send for BufferExportGuard {}
+unsafe impl Sync for BufferExportGuard {}
+
+fn import_buffer_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyResult<BufferHandle> {
+    let export_ptr = capsule
+        .pointer_checked(Some(BUFFER_EXPORT_CAPSULE_NAME))?
+        .cast::<VortexBufferExport>();
+    let export = unsafe { export_ptr.as_ref() };
+
+    if export.version != 1 {
+        return Err(PyValueError::new_err(format!(
+            "unsupported VortexBufferExport version {}",
+            export.version
+        )));
+    }
+    if export.kind != VORTEX_BUFFER_HOST {
+        return Err(PyValueError::new_err(format!(
+            "unsupported buffer kind {} (only host buffers are supported in metadata bridge)",
+            export.kind
+        )));
+    }
+
+    if export.len != 0 && export.ptr.is_null() {
+        return Err(PyValueError::new_err(
+            "non-empty VortexBufferExport has null data pointer",
+        ));
+    }
+
+    let len = export.len;
+    let alignment = vortex::buffer::Alignment::try_from(
+        u32::try_from(export.alignment)
+            .map_err(|_| PyValueError::new_err("buffer alignment exceeds u32"))?,
+    )
+    .map_err(|e| PyValueError::new_err(e.to_string()))?;
+
+    if len != 0 && !alignment.is_ptr_aligned(export.ptr) {
+        return Err(PyValueError::new_err(format!(
+            "buffer pointer is not aligned to requested alignment {alignment}"
+        )));
+    }
+
+    // Transfer ownership of the boxed VortexBufferExport from the producer capsule into the Bytes
+    // owner below. Otherwise the producer capsule could be dropped before the reconstructed
+    // BufferHandle, leaving the Bytes owner with a dangling export pointer.
+    unsafe { ffi::PyCapsule_SetDestructor(capsule.as_ptr(), None) };
+    if PyErr::occurred(capsule.py()) {
+        return Err(PyErr::fetch(capsule.py()));
+    }
+
+    let guard = BufferExportGuard { export: export_ptr };
+
+    let byte_buffer = if len == 0 {
+        drop(guard);
+        ByteBuffer::empty_aligned(alignment)
+    } else {
+        ByteBuffer::from(bytes::Bytes::from_owner(guard)).aligned(alignment)
+    };
+
+    Ok(BufferHandle::new_host(byte_buffer))
+}
 
 struct ExportedDeviceArray(ArrowDeviceArrayWithSchema);
 
@@ -101,7 +211,7 @@ struct ArrayMetadata {
     dtype: Vec<u8>,
     len: usize,
     metadata: Vec<u8>,
-    buffer_count: usize,
+    buffers: Vec<BufferHandle>,
     children: Vec<ArrayMetadata>,
 }
 
@@ -147,6 +257,16 @@ fn parse_array_metadata(value: &Bound<'_, PyAny>) -> PyResult<ArrayMetadata> {
         )));
     }
 
+    let buffers = tuple
+        .get_item(4)?
+        .cast::<PyList>()?
+        .iter()
+        .map(|item| {
+            let capsule: Bound<'_, PyCapsule> = item.extract()?;
+            import_buffer_from_capsule(&capsule)
+        })
+        .collect::<PyResult<Vec<_>>>()?;
+
     let children = tuple
         .get_item(5)?
         .cast::<PyList>()?
@@ -159,7 +279,7 @@ fn parse_array_metadata(value: &Bound<'_, PyAny>) -> PyResult<ArrayMetadata> {
         dtype: tuple.get_item(1)?.extract()?,
         len: tuple.get_item(2)?.extract()?,
         metadata: tuple.get_item(3)?.extract()?,
-        buffer_count: tuple.get_item(4)?.extract()?,
+        buffers,
         children,
     })
 }
@@ -173,14 +293,6 @@ fn deserialize_metadata_tree(
     metadata: &ArrayMetadata,
     session: &VortexSession,
 ) -> VortexResult<ArrayRef> {
-    if metadata.buffer_count != 0 {
-        vortex_bail!(
-            "metadata-only bridge cannot deserialize array {} with {} buffers yet",
-            metadata.encoding_id,
-            metadata.buffer_count
-        );
-    }
-
     let dtype = dtype_from_metadata(metadata, session)?;
     let children = metadata
         .children
@@ -194,12 +306,11 @@ fn deserialize_metadata_tree(
         .registry()
         .find(&encoding_id)
         .ok_or_else(|| vortex_err!("Unknown array encoding: {}", metadata.encoding_id))?;
-    let buffers: &[BufferHandle] = &[];
     let decoded = plugin.deserialize(
         &dtype,
         metadata.len,
         &metadata.metadata,
-        buffers,
+        &metadata.buffers,
         &children,
         session,
     )?;
@@ -246,6 +357,20 @@ fn _debug_array_metadata_dtype(array: Bound<'_, PyAny>) -> PyResult<String> {
     Ok(array.dtype().to_string())
 }
 
+/// Return scalar display strings after crossing the private vtable-metadata bridge.
+#[pyfunction]
+fn _debug_array_metadata_scalar_values(array: Bound<'_, PyAny>) -> PyResult<Vec<String>> {
+    let metadata = extract_array_metadata(&array)?;
+    let array = deserialize_metadata_tree(&metadata, &METADATA_SESSION).map_err(to_py_err)?;
+    (0..array.len())
+        .map(|index| {
+            #[expect(deprecated)]
+            array.scalar_at(index).map(|scalar| scalar.to_string())
+        })
+        .collect::<VortexResult<Vec<_>>>()
+        .map_err(to_py_err)
+}
+
 /// Export a PyVortex array as Arrow C Device schema and array PyCapsules.
 #[pyfunction]
 #[pyo3(signature = (array, requested_schema = None, **kwargs))]
@@ -461,6 +586,7 @@ unsafe extern "C" fn release_device_array_capsule(capsule: *mut ffi::PyObject) {
 fn _lib(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(cuda_available, m)?)?;
     m.add_function(wrap_pyfunction!(_debug_array_metadata_dtype, m)?)?;
+    m.add_function(wrap_pyfunction!(_debug_array_metadata_scalar_values, m)?)?;
     m.add_function(wrap_pyfunction!(export_device_array, m)?)?;
     Ok(())
 }

vortex-python-cuda/src/lib.rs

@@ -13,11 +13,58 @@ def test_debug_array_metadata_dtype_reads_base_vortex_array():
     assert vortex_cuda._debug_array_metadata_dtype(array) == str(array.dtype)  # pyright: ignore[reportPrivateUsage]
 
 
-def test_metadata_bridge_reports_arrays_that_need_buffer_handoff():
+def test_metadata_bridge_primitive_array():
     array = vortex.array([1, 2, 3])
 
-    with pytest.raises(RuntimeError, match="metadata-only bridge.*buffers"):
-        _ = vortex_cuda._debug_array_metadata_dtype(array)  # pyright: ignore[reportPrivateUsage]
+    assert vortex_cuda._debug_array_metadata_dtype(array) == str(array.dtype)  # pyright: ignore[reportPrivateUsage]
+    assert vortex_cuda._debug_array_metadata_scalar_values(array) == [  # pyright: ignore[reportPrivateUsage]
+        "1i64",
+        "2i64",
+        "3i64",
+    ]
+
+
+def test_metadata_bridge_nullable_array():
+    array = vortex.array([1, None, 3])
+
+    assert vortex_cuda._debug_array_metadata_dtype(array) == str(array.dtype)  # pyright: ignore[reportPrivateUsage]
+    assert vortex_cuda._debug_array_metadata_scalar_values(array) == [  # pyright: ignore[reportPrivateUsage]
+        "1i64",
+        "null",
+        "3i64",
+    ]
+
+
+def test_metadata_bridge_bool_array():
+    array = vortex.array([True, False, True])
+
+    assert vortex_cuda._debug_array_metadata_dtype(array) == str(array.dtype)  # pyright: ignore[reportPrivateUsage]
+    assert vortex_cuda._debug_array_metadata_scalar_values(array) == [  # pyright: ignore[reportPrivateUsage]
+        "true",
+        "false",
+        "true",
+    ]
+
+
+def test_metadata_bridge_struct_with_children():
+    import pyarrow as pa
+
+    arrow_table = pa.table({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
+    struct_array = vortex.Array.from_arrow(
+        pa.StructArray.from_arrays(  # pyright: ignore[reportUnknownMemberType]
+            [arrow_table.column("a").combine_chunks(), arrow_table.column("b").combine_chunks()],
+            names=["a", "b"],
+        )
+    )
+
+    assert vortex_cuda._debug_array_metadata_dtype(struct_array) == str(  # pyright: ignore[reportPrivateUsage]
+        struct_array.dtype
+    )
+    assert vortex_cuda._debug_array_metadata_scalar_values(struct_array) == [  # pyright: ignore[reportPrivateUsage]
+        "{a: 1i64, b: 4f64}",
+        "{a: 2i64, b: 5f64}",
+        "{a: 3i64, b: 6f64}",
+    ]
 
 
 def test_export_device_array_returns_capsules_or_clean_cuda_error():

vortex-python-cuda/test/test_native_bridge.py

@@ -26,7 +26,7 @@ class Array:
     @staticmethod
     def from_range(obj: range, *, dtype: DType | None = None) -> Array: ...
     def to_arrow_array(self) -> pa.Array[pa.Scalar[pa.DataType]]: ...
-    def __vortex_array_metadata__(self) -> tuple[str, bytes, int, bytes, int, list[object]]: ...
+    def __vortex_array_metadata__(self) -> tuple[str, bytes, int, bytes, list[object], list[object]]: ...
     @property
     def id(self) -> str: ...
     @property