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-python-cuda/python/vortex_cuda/__init__.py

@@ -2,11 +2,116 @@
 # SPDX-FileCopyrightText: Copyright the Vortex contributors
 # pyright: reportMissingModuleSource=false, reportPrivateUsage=false
 
+import importlib
+from typing import Protocol, cast, final
+
 from . import _lib
 
+# Private debug hooks used by CUDA bridge tests.
 _debug_array_metadata_dtype = _lib._debug_array_metadata_dtype
 _debug_array_metadata_display_values = _lib._debug_array_metadata_display_values
+_debug_arrow_device_array_capsule_summary = _lib._debug_arrow_device_array_capsule_summary
+_debug_consume_arrow_device_array_capsules = _lib._debug_consume_arrow_device_array_capsules
 cuda_available = _lib.cuda_available
 export_device_array = _lib.export_device_array
 
-__all__ = ["cuda_available", "export_device_array"]
+class _FromArrow(Protocol):
+    def from_arrow(self, obj: object) -> object: ...
+
+
+class _FromPylibcudf(Protocol):
+    def from_pylibcudf(self, obj: object) -> object: ...
+
+
+class _DataFrameFromPylibcudf(Protocol):
+    def from_pylibcudf(self, obj: object, metadata: dict[str, object] | None = None) -> object: ...
+
+
+class _PylibcudfModule(Protocol):
+    Column: _FromArrow
+    Table: _FromArrow
+
+
+class _CudfModule(Protocol):
+    Series: _FromPylibcudf
+    DataFrame: _DataFrameFromPylibcudf
+
+
+@final
+class _ArrowDeviceArrayExportable:
+    def __init__(self, array: object) -> None:
+        self._array: object = array
+
+    def __arrow_c_device_array__(
+        self,
+        requested_schema: object | None = None,
+        **kwargs: object,
+    ) -> tuple[object, object]:
+        return export_device_array(self._array, requested_schema, **kwargs)
+
+
+_SUPPORTED_FALLBACKS = frozenset({"error"})
+
+
+def _Array_to_cudf(self: object, *, fallback: str = "error") -> object:
+    return to_cudf(self, fallback=fallback)
+
+
+def _Array___arrow_c_device_array__(
+    self: object,
+    requested_schema: object | None = None,
+    **kwargs: object,
+) -> tuple[object, object]:
+    return export_device_array(self, requested_schema, **kwargs)
+
+
+def _install_vortex_array_methods() -> None:
+    import vortex
+
+    setattr(vortex.Array, "to_cudf", _Array_to_cudf)
+    if cuda_available():
+        setattr(vortex.Array, "__arrow_c_device_array__", _Array___arrow_c_device_array__)
+
+
+def _import_cudf_modules() -> tuple[_CudfModule, _PylibcudfModule]:
+    try:
+        cudf = importlib.import_module("cudf")
+        pylibcudf = importlib.import_module("pylibcudf")
+    except ImportError as err:
+        raise ImportError("vortex_cuda.to_cudf requires RAPIDS cuDF and pylibcudf to be installed") from err
+    return cast(_CudfModule, cast(object, cudf)), cast(_PylibcudfModule, cast(object, pylibcudf))
+
+
+def to_cudf(obj: object, *, fallback: str = "error") -> object:
+    """Convert a Vortex array to a cuDF object through the Arrow Device interface.
+
+    ``fallback`` is reserved for future policy choices. The initial implementation
+    supports only ``fallback="error"`` and never falls back to host Arrow conversion.
+    """
+    if fallback not in _SUPPORTED_FALLBACKS:
+        raise NotImplementedError("vortex_cuda.to_cudf currently supports only fallback='error'")
+
+    import vortex
+
+    if not isinstance(obj, vortex.Array):
+        raise TypeError(f"vortex_cuda.to_cudf expected a vortex.Array, got {type(obj).__name__}")
+
+    if not cuda_available():
+        raise RuntimeError("CUDA is not available; vortex_cuda.to_cudf requires a CUDA device")
+
+    cudf, pylibcudf = _import_cudf_modules()
+    exportable = _ArrowDeviceArrayExportable(obj)
+
+    dtype = obj.dtype
+    if isinstance(dtype, vortex.StructDType):
+        table = pylibcudf.Table.from_arrow(exportable)
+        return cudf.DataFrame.from_pylibcudf(table, metadata={"columns": dtype.names()})
+
+    column = pylibcudf.Column.from_arrow(exportable)
+    return cudf.Series.from_pylibcudf(column)
+
+
+_install_vortex_array_methods()
+
+
+__all__ = ["cuda_available", "export_device_array", "to_cudf"]

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

@@ -3,6 +3,10 @@
 
 def _debug_array_metadata_dtype(array: object) -> str: ...
 def _debug_array_metadata_display_values(array: object) -> str: ...
+def _debug_arrow_device_array_capsule_summary(schema: object, device_array: object) -> dict[str, object]: ...
+def _debug_consume_arrow_device_array_capsules(
+    schema: object, device_array: object
+) -> tuple[bool, bool, bool, bool, bool, bool]: ...
 def cuda_available() -> bool: ...
 def export_device_array(
     array: object, requested_schema: object | None = None, **kwargs: object

vortex-python-cuda/src/lib.rs

@@ -58,6 +58,9 @@ 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");
 
+// Owns a buffer export after we steal it from the source PyCapsule. `Bytes::from_owner`
+// keeps this guard alive for every cloned byte buffer and drops it to run the producer's
+// release callback once the reconstructed Vortex buffer is no longer referenced.
 struct BufferExportGuard {
     export: NonNull<VortexBufferExport>,
 }
@@ -89,8 +92,9 @@ impl Drop for BufferExportGuard {
     }
 }
 
-// The guard is moved into `Bytes::from_owner`, which requires `Send + Sync`. After import we disable
-// the source capsule destructor and own the C export until this guard is dropped.
+// The guard is moved into `Bytes::from_owner`, which requires `Send + Sync`.
+// After import we disable the source capsule destructor and own the C export
+// until this guard is dropped.
 unsafe impl Send for BufferExportGuard {}
 unsafe impl Sync for BufferExportGuard {}
 
@@ -461,6 +465,84 @@ fn release_exported(exported: &mut ArrowDeviceArrayWithSchema) {
     release_device_array(&mut exported.array);
 }
 
+/// Return non-owning details from Arrow Device capsules for Python-side smoke consumers.
+#[pyfunction]
+fn _debug_arrow_device_array_capsule_summary<'py>(
+    py: Python<'py>,
+    schema: Bound<'py, PyCapsule>,
+    device_array: Bound<'py, PyCapsule>,
+) -> PyResult<Bound<'py, PyDict>> {
+    let schema = unsafe {
+        schema
+            .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
+            .cast::<FFI_ArrowSchema>()
+            .as_ref()
+    };
+    let device_array = unsafe {
+        device_array
+            .pointer_checked(Some(ARROW_DEVICE_ARRAY_CAPSULE_NAME))?
+            .cast::<ArrowDeviceArray>()
+            .as_ref()
+    };
+
+    let summary = PyDict::new(py);
+    summary.set_item("schema_live", schema.release.is_some())?;
+    summary.set_item("array_live", device_array.array.release.is_some())?;
+    summary.set_item("is_cuda", device_array.device_type == ARROW_DEVICE_CUDA)?;
+    summary.set_item("device_type", device_array.device_type)?;
+    summary.set_item("device_id", device_array.device_id)?;
+    summary.set_item("length", device_array.array.length)?;
+    summary.set_item("null_count", device_array.array.null_count)?;
+    summary.set_item("n_buffers", device_array.array.n_buffers)?;
+    summary.set_item("n_children", device_array.array.n_children)?;
+    Ok(summary)
+}
+
+/// Simulate a Python Arrow Device consumer taking ownership from the returned capsules.
+#[pyfunction]
+fn _debug_consume_arrow_device_array_capsules(
+    schema: Bound<'_, PyCapsule>,
+    device_array: Bound<'_, PyCapsule>,
+) -> PyResult<(bool, bool, bool, bool, bool, bool)> {
+    let mut schema_ptr = schema
+        .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))?
+        .cast::<FFI_ArrowSchema>();
+    let mut device_array_ptr = device_array
+        .pointer_checked(Some(ARROW_DEVICE_ARRAY_CAPSULE_NAME))?
+        .cast::<ArrowDeviceArray>();
+
+    let schema_ref = unsafe { schema_ptr.as_mut() };
+    let device_array_ref = unsafe { device_array_ptr.as_mut() };
+    let schema_had_release = schema_ref.release.is_some();
+    let array_had_release = device_array_ref.array.release.is_some();
+
+    release_schema(schema_ref);
+    release_device_array(device_array_ref);
+
+    let schema_release_cleared = schema_ref.release.is_none();
+    let array_release_cleared = device_array_ref.array.release.is_none();
+
+    set_capsule_name(&schema, USED_ARROW_SCHEMA_CAPSULE_NAME)?;
+    set_capsule_name(&device_array, USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME)?;
+
+    Ok((
+        schema_had_release,
+        array_had_release,
+        schema_release_cleared,
+        array_release_cleared,
+        schema.is_valid_checked(Some(USED_ARROW_SCHEMA_CAPSULE_NAME)),
+        device_array.is_valid_checked(Some(USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME)),
+    ))
+}
+
+fn set_capsule_name(capsule: &Bound<'_, PyCapsule>, name: &CStr) -> PyResult<()> {
+    let result = unsafe { ffi::PyCapsule_SetName(capsule.as_ptr(), name.as_ptr()) };
+    if result != 0 {
+        return Err(PyErr::fetch(capsule.py()));
+    }
+    Ok(())
+}
+
 fn schema_capsule<'py>(
     py: Python<'py>,
     schema: FFI_ArrowSchema,
@@ -573,6 +655,14 @@ 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_display_values, m)?)?;
+    m.add_function(wrap_pyfunction!(
+        _debug_arrow_device_array_capsule_summary,
+        m
+    )?)?;
+    m.add_function(wrap_pyfunction!(
+        _debug_consume_arrow_device_array_capsules,
+        m
+    )?)?;
     m.add_function(wrap_pyfunction!(export_device_array, m)?)?;
     Ok(())
 }