Add support for the Arrow device capsule interfaces

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -47,6 +47,7 @@ dependencies:
 - moto>=4.0.8
 - msgpack-python
 - myst-nb
+- nanoarrow
 - nbconvert
 - nbformat
 - nbsphinx
@@ -70,7 +71,7 @@ dependencies:
 - pyarrow>=14.0.0,<20.0.0a0
 - pydata-sphinx-theme>=0.15.4
 - pynvml>=12.0.0,<13.0.0a0
-- pytest-benchmark<5.1.0
+- pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
 - pytest-rerunfailures

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -49,6 +49,7 @@ dependencies:
 - moto>=4.0.8
 - msgpack-python
 - myst-nb
+- nanoarrow
 - nbconvert
 - nbformat
 - nbsphinx
@@ -72,7 +73,7 @@ dependencies:
 - pyarrow>=14.0.0,<20.0.0a0
 - pydata-sphinx-theme>=0.15.4
 - pynvml>=12.0.0,<13.0.0a0
-- pytest-benchmark<5.1.0
+- pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
 - pytest-rerunfailures

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -47,6 +47,7 @@ dependencies:
 - moto>=4.0.8
 - msgpack-python
 - myst-nb
+- nanoarrow
 - nbconvert
 - nbformat
 - nbsphinx
@@ -69,7 +70,7 @@ dependencies:
 - pydata-sphinx-theme>=0.15.4
 - pynvjitlink>=0.0.0a0
 - pynvml>=12.0.0,<13.0.0a0
-- pytest-benchmark<5.1.0
+- pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
 - pytest-rerunfailures

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -48,6 +48,7 @@ dependencies:
 - moto>=4.0.8
 - msgpack-python
 - myst-nb
+- nanoarrow
 - nbconvert
 - nbformat
 - nbsphinx
@@ -70,7 +71,7 @@ dependencies:
 - pydata-sphinx-theme>=0.15.4
 - pynvjitlink>=0.0.0a0
 - pynvml>=12.0.0,<13.0.0a0
-- pytest-benchmark<5.1.0
+- pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
 - pytest-rerunfailures

dependencies.yaml

@@ -75,6 +75,7 @@ files:
       - test_python_common
       - test_python_cudf_common
       - test_python_cudf
+      - test_python_pylibcudf
       - depends_on_cudf
       - depends_on_pylibcudf
       - depends_on_libcudf
@@ -86,6 +87,7 @@ files:
       - test_python_common
       - test_python_cudf_common
       - test_python_dask_cudf
+      - test_python_pylibcudf
       - depends_on_cudf
       - depends_on_pylibcudf
       - depends_on_libcudf
@@ -920,28 +922,20 @@ dependencies:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
-          - fastavro>=0.22.9
-          # TODO: Remove hypothesis pinning once https://github.com/HypothesisWorks/hypothesis/issues/4365 is resolved
-          - hypothesis<6.131.1
+          - &fastavro fastavro>=0.22.9
+          - nanoarrow
           - *numpy
           - pandas
   test_python_cudf:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
           - cramjam
-          - fastavro>=0.22.9
+          - *fastavro
           # TODO: Remove hypothesis pinning once https://github.com/HypothesisWorks/hypothesis/issues/4365 is resolved
           - hypothesis<6.131.1
           - mmh3
-          # Version 5.1 is incompatible with pytest<8.2.
-          # https://github.com/ionelmc/pytest-benchmark/commit/8dfeeeca8a5c640a4dc4455904ace4f97f62e655
-          # Remove upper bound when we unbound pytest
-          # https://github.com/rapidsai/build-planning/issues/105
-          # or when a fixed version of pytest-benchmark 5.1 with the necessary bounds is released
-          # https://github.com/conda-forge/pytest-benchmark-feedstock/pull/27
-          # https://github.com/conda-forge/conda-forge-repodata-patches-feedstock/pull/990
-          - pytest-benchmark<5.1.0
+          - pytest-benchmark
           - pytest-cases>=3.8.2
           - scipy
           - zstandard

python/cudf/pyproject.toml

@@ -56,7 +56,7 @@ test = [
     "hypothesis<6.131.1",
     "mmh3",
     "msgpack",
-    "pytest-benchmark<5.1.0",
+    "pytest-benchmark",
     "pytest-cases>=3.8.2",
     "pytest-cov",
     "pytest-rerunfailures",

python/pylibcudf/pylibcudf/_interop_helpers.pxd

@@ -6,4 +6,6 @@ cdef void _release_schema(object schema_capsule) noexcept
 
 cdef void _release_array(object array_capsule) noexcept
 
+cdef void _release_device_array(object array_capsule) noexcept
+
 cdef column_metadata _metadata_to_libcudf(metadata)

python/pylibcudf/pylibcudf/_interop_helpers.pyx

@@ -4,9 +4,11 @@ from cpython.pycapsule cimport PyCapsule_GetPointer
 
 from pylibcudf.libcudf.interop cimport (
     ArrowArray,
+    ArrowDeviceArray,
     ArrowSchema,
     column_metadata,
     release_arrow_array_raw,
+    release_arrow_device_array_raw,
     release_arrow_schema_raw,
 )
 
@@ -39,6 +41,14 @@ cdef void _release_array(object array_capsule) noexcept:
     release_arrow_array_raw(array)
 
 
+cdef void _release_device_array(object array_capsule) noexcept:
+    """Release the ArrowDeviceArray object stored in a PyCapsule."""
+    cdef ArrowDeviceArray* array = <ArrowDeviceArray*>PyCapsule_GetPointer(
+        array_capsule, 'arrow_device_array'
+    )
+    release_arrow_device_array_raw(array)
+
+
 cdef column_metadata _metadata_to_libcudf(metadata):
     """Convert a ColumnMetadata object to C++ column_metadata.
 

python/pylibcudf/pylibcudf/column.pyx

@@ -15,19 +15,21 @@ from libcpp.utility cimport move
 
 from pylibcudf.libcudf.column.column cimport column, column_contents
 from pylibcudf.libcudf.column.column_factories cimport make_column_from_scalar
-from pylibcudf.libcudf.interop cimport ArrowArray, ArrowSchema, arrow_column
-from pylibcudf.libcudf.scalar.scalar cimport scalar, numeric_scalar
-from pylibcudf.libcudf.types cimport size_type, size_of as cpp_size_of, bitmask_type
-from pylibcudf.libcudf.utilities.traits cimport is_fixed_width
-from pylibcudf.libcudf.copying cimport get_element
-
 from pylibcudf.libcudf.interop cimport (
     ArrowArray,
     ArrowSchema,
+    ArrowDeviceArray,
+    arrow_column,
     column_metadata,
     to_arrow_host_raw,
+    to_arrow_device_raw,
     to_arrow_schema_raw,
 )
+from pylibcudf.libcudf.scalar.scalar cimport scalar, numeric_scalar
+from pylibcudf.libcudf.types cimport size_type, size_of as cpp_size_of, bitmask_type
+from pylibcudf.libcudf.utilities.traits cimport is_fixed_width
+from pylibcudf.libcudf.copying cimport get_element
+
 
 from rmm.librmm.device_buffer cimport device_buffer
 from rmm.pylibrmm.device_buffer cimport DeviceBuffer
@@ -40,6 +42,7 @@ from .types cimport DataType, size_of, type_id
 from ._interop_helpers cimport (
     _release_schema,
     _release_array,
+    _release_device_array,
     _metadata_to_libcudf,
 )
 from .null_mask cimport bitmask_allocation_size_bytes
@@ -55,7 +58,12 @@ __all__ = ["Column", "ListColumnView", "is_c_contiguous"]
 
 class _ArrowLikeMeta(type):
     def __subclasscheck__(cls, other):
-        return hasattr(other, "__arrow_c_array__")
+        # We cannot separate these types via singledispatch because the dispatch
+        # will often be ambiguous when objects expose multiple protocols.
+        return (
+            hasattr(other, "__arrow_c_array__")
+            or hasattr(other, "__arrow_c_device_array__")
+        )
 
 
 class _ArrowLike(metaclass=_ArrowLikeMeta):
@@ -207,32 +215,59 @@ cdef class Column:
 
     @_init.register(_ArrowLike)
     def _(self, arrow_like):
-        schema, array = arrow_like.__arrow_c_array__()
-        cdef ArrowSchema* c_schema = (
-            <ArrowSchema*>PyCapsule_GetPointer(schema, "arrow_schema")
-        )
-        cdef ArrowArray* c_array = (
-            <ArrowArray*>PyCapsule_GetPointer(array, "arrow_array")
-        )
-
-        cdef _ArrowColumnHolder result = _ArrowColumnHolder()
+        cdef ArrowSchema* c_schema
+        cdef ArrowArray* c_array
+        cdef ArrowDeviceArray* c_device_array
+        cdef _ArrowColumnHolder result
         cdef unique_ptr[arrow_column] c_result
-        with nogil:
-            c_result = make_unique[arrow_column](
-                move(dereference(c_schema)), move(dereference(c_array))
+        if hasattr(arrow_like, "__arrow_c_device_array__"):
+            schema, array = arrow_like.__arrow_c_device_array__()
+            c_schema = <ArrowSchema*>PyCapsule_GetPointer(schema, "arrow_schema")
+            c_device_array = (
+                <ArrowDeviceArray*>PyCapsule_GetPointer(array, "arrow_device_array")
+            )
+
+            result = _ArrowColumnHolder()
+            with nogil:
+                c_result = make_unique[arrow_column](
+                    move(dereference(c_schema)), move(dereference(c_device_array))
+                )
+            result.col.swap(c_result)
+
+            tmp = Column.from_column_view_of_arbitrary(result.col.get().view(), result)
+            self._init(
+                tmp.type(),
+                tmp.size(),
+                tmp.data(),
+                tmp.null_mask(),
+                tmp.null_count(),
+                tmp.offset(),
+                tmp.children(),
             )
+        elif hasattr(arrow_like, "__arrow_c_array__"):
+            schema, array = arrow_like.__arrow_c_array__()
+            c_schema = <ArrowSchema*>PyCapsule_GetPointer(schema, "arrow_schema")
+            c_array = <ArrowArray*>PyCapsule_GetPointer(array, "arrow_array")
+
+            result = _ArrowColumnHolder()
+            with nogil:
+                c_result = make_unique[arrow_column](
+                    move(dereference(c_schema)), move(dereference(c_array))
+                )
             result.col.swap(c_result)
 
-        tmp = Column.from_column_view_of_arbitrary(result.col.get().view(), result)
-        self._init(
-            tmp.type(),
-            tmp.size(),
-            tmp.data(),
-            tmp.null_mask(),
-            tmp.null_count(),
-            tmp.offset(),
-            tmp.children(),
-        )
+            tmp = Column.from_column_view_of_arbitrary(result.col.get().view(), result)
+            self._init(
+                tmp.type(),
+                tmp.size(),
+                tmp.data(),
+                tmp.null_mask(),
+                tmp.null_count(),
+                tmp.offset(),
+                tmp.children(),
+            )
+        else:
+            raise ValueError("Invalid Arrow-like object")
 
     cdef column_view view(self) nogil:
         """Generate a libcudf column_view to pass to libcudf algorithms.
@@ -752,12 +787,37 @@ cdef class Column:
 
         return PyCapsule_New(<void*>raw_host_array_ptr, "arrow_array", _release_array)
 
+    def _to_device_array(self):
+        cdef ArrowDeviceArray* raw_device_array_ptr
+        with nogil:
+            raw_device_array_ptr = to_arrow_device_raw(self.view(), self)
+
+        return PyCapsule_New(
+            <void*>raw_device_array_ptr,
+            "arrow_device_array",
+            _release_device_array
+        )
+
     def __arrow_c_array__(self, requested_schema=None):
         if requested_schema is not None:
             raise ValueError("pylibcudf.Column does not support alternative schema")
 
         return self._to_schema(), self._to_host_array()
 
+    def __arrow_c_device_array__(self, requested_schema=None, **kwargs):
+        if requested_schema is not None:
+            raise ValueError("pylibcudf.Column does not support alternative schema")
+
+        non_default_kwargs = [
+            name for name, value in kwargs.items() if value is not None
+        ]
+        if non_default_kwargs:
+            raise NotImplementedError(
+                f"Received unsupported keyword argument(s): {non_default_kwargs}"
+            )
+
+        return self._to_schema(), self._to_device_array()
+
 
 cdef class ListColumnView:
     """Accessor for methods of a Column that are specific to lists."""