Add to_arrow method to pylibcudf core types

docs/cudf/source/conf.py

@@ -620,6 +620,7 @@ def on_missing_reference(app, env, node, contnode):
     ("py:class", "typing_extensions.Self"),
     ("py:class", "np.uint32"),
     ("py:class", "np.uint64"),
+    ("py:class", "ArrowLike"),
 ]
 
 

python/pylibcudf/pylibcudf/_interop_helpers.pyx

@@ -31,6 +31,15 @@ class ArrowLike(metaclass=_ArrowLikeMeta):
     pass
 
 
+class _ObjectWithArrowMetadata:
+    def __init__(self, obj, metadata=None):
+        self.obj = obj
+        self.metadata = metadata
+
+    def __arrow_c_array__(self, requested_schema=None):
+        return self.obj._to_schema(self.metadata), self.obj._to_host_array()
+
+
 @dataclass
 class ColumnMetadata:
     """Metadata associated with a column.

python/pylibcudf/pylibcudf/column.pyi

@@ -70,6 +70,7 @@ class Column:
     def from_rmm_buffer(
         buff: DeviceBuffer, dtype: DataType, size: int, children: list[Column]
     ) -> Column: ...
+    def to_arrow(self, metadata: list | str | None = None) -> ArrowLike: ...
     @staticmethod
     def from_arrow(
         obj: ArrowLike, dtype: DataType | None = None

python/pylibcudf/pylibcudf/column.pyx

@@ -55,14 +55,22 @@ from ._interop_helpers cimport (
 from .utils cimport _get_stream
 
 from .gpumemoryview import _datatype_from_dtype_desc
-from ._interop_helpers import ArrowLike, ColumnMetadata
+from ._interop_helpers import ArrowLike, ColumnMetadata, _ObjectWithArrowMetadata
 
 import array
 from itertools import accumulate
 import functools
 import operator
 from typing import Iterable
 
+try:
+    import pyarrow as pa
+    pa_err = None
+except ImportError as e:
+    pa = None
+    pa_err = e
+
+
 __all__ = ["Column", "ListColumnView", "is_c_contiguous"]
 
 
@@ -336,12 +344,38 @@ cdef class Column:
         self._children = children
         self._num_children = len(children)
 
+    def to_arrow(
+        self,
+        metadata: ColumnMetadata | str | None = None
+    ) -> ArrowLike:
+        """Create a PyArrow array from a pylibcudf column.
+
+        Parameters
+        ----------
+        metadata : list
+            The metadata to attach to the columns of the table.
+
+        Returns
+        -------
+        pyarrow.Array
+        """
+        if pa_err is not None:
+            raise RuntimeError(
+                "pyarrow was not found on your system. Please "
+                "pip install pylibcudf with the [pyarrow] extra for a "
+                "compatible pyarrow version."
+            ) from pa_err
+        # TODO: Once the arrow C device interface registers more
+        # types that it supports, we can call pa.array(self) if
+        # no metadata is passed.
+        return pa.array(_ObjectWithArrowMetadata(self, metadata))
+
     @staticmethod
     def from_arrow(
         obj: ArrowLike,
         dtype: DataType | None = None,
         Stream stream=None
-    ) -> Column:
+    ) -> ArrowLike:
         """
         Create a Column from an Arrow-like object using the Arrow C Data Interface.
 

python/pylibcudf/pylibcudf/interop.pyx

@@ -23,8 +23,7 @@ from rmm.pylibrmm.stream cimport Stream
 from .column cimport Column
 from .scalar cimport Scalar
 from .table cimport Table
-from .types cimport DataType, type_id
-from .types import LIBCUDF_TO_ARROW_TYPES
+from .types cimport DataType
 from .utils cimport _get_stream
 from ._interop_helpers import ColumnMetadata
 
@@ -71,7 +70,7 @@ def from_arrow(pyarrow_object, *, DataType data_type=None):
 
 
 @singledispatch
-def to_arrow(plc_object, metadata=None):
+def to_arrow(plc_object, **kwargs):
     """Convert to a PyArrow object.
 
     Parameters
@@ -105,79 +104,36 @@ if pa is not None:
         return Table.from_arrow(pyarrow_object, dtype=data_type)
 
     @from_arrow.register(pa.Scalar)
-    def _from_arrow_scalar(pyarrow_object, *, DataType data_type=None):
-        return Scalar.from_arrow(pyarrow_object, dtype=data_type)
+    def _from_arrow_scalar(
+        pyarrow_object,
+        *,
+        DataType data_type=None,
+        Stream stream = None
+    ):
+        return Scalar.from_arrow(pyarrow_object, dtype=data_type, stream=stream)
 
     @from_arrow.register(pa.Array)
     def _from_arrow_column(pyarrow_object, *, DataType data_type=None):
         return Column.from_arrow(pyarrow_object, dtype=data_type)
 
     @to_arrow.register(DataType)
     def _to_arrow_datatype(plc_object, **kwargs):
-        """
-        Convert a datatype to arrow.
-
-        Translation of some types requires extra information as a keyword
-        argument. Specifically:
-
-        - When translating a decimal type, provide ``precision``
-        - When translating a struct type, provide ``fields``
-        - When translating a list type, provide the wrapped ``value_type``
-        """
-        if plc_object.id() in {
-            type_id.DECIMAL32,
-            type_id.DECIMAL64,
-            type_id.DECIMAL128
-        }:
-            if not (precision := kwargs.get("precision")):
-                raise ValueError(
-                    "Precision must be provided for decimal types"
-                )
-                # no pa.decimal32 or pa.decimal64
-            return pa.decimal128(precision, -plc_object.scale())
-        elif plc_object.id() == type_id.STRUCT:
-            if not (fields := kwargs.get("fields")):
-                raise ValueError(
-                    "Fields must be provided for struct types"
-                )
-            return pa.struct(fields)
-        elif plc_object.id() == type_id.LIST:
-            if not (value_type := kwargs.get("value_type")):
-                raise ValueError(
-                    "Value type must be provided for list types"
-                )
-            return pa.list_(value_type)
-        else:
-            try:
-                return LIBCUDF_TO_ARROW_TYPES[plc_object.id()]
-            except KeyError:
-                raise TypeError(
-                    f"Unable to convert {plc_object.id()} to arrow datatype"
-                )
-
-    class _ObjectWithArrowMetadata:
-        def __init__(self, obj, metadata=None):
-            self.obj = obj
-            self.metadata = metadata
-
-        def __arrow_c_array__(self, requested_schema=None):
-            return self.obj._to_schema(self.metadata), self.obj._to_host_array()
+        """Convert a datatype to arrow."""
+        return plc_object.to_arrow(**kwargs)
 
     @to_arrow.register(Table)
     def _to_arrow_table(plc_object, metadata=None):
         """Create a PyArrow table from a pylibcudf table."""
-        return pa.table(_ObjectWithArrowMetadata(plc_object, metadata))
+        return plc_object.to_arrow(metadata=metadata)
 
     @to_arrow.register(Column)
     def _to_arrow_array(plc_object, metadata=None):
         """Create a PyArrow array from a pylibcudf column."""
-        return pa.array(_ObjectWithArrowMetadata(plc_object, metadata))
+        return plc_object.to_arrow(metadata=metadata)
 
     @to_arrow.register(Scalar)
     def _to_arrow_scalar(plc_object, metadata=None):
-        # Note that metadata for scalars is primarily important for preserving
-        # information on nested types since names are otherwise irrelevant.
-        return to_arrow(Column.from_scalar(plc_object, 1), metadata=metadata)[0]
+        return plc_object.to_arrow(metadata=metadata)
 
 
 cpdef Table from_dlpack(object managed_tensor, Stream stream=None):

python/pylibcudf/pylibcudf/scalar.pyi

@@ -9,12 +9,17 @@ from pylibcudf.types import DataType
 
 NpGeneric = type[Any]
 
+PaScalar = type[Any]
+
 class Scalar:
     def __init__(self): ...
     def type(self) -> DataType: ...
     def is_valid(self) -> bool: ...
     @staticmethod
     def empty_like(column: Column, stream: Stream | None = None) -> Scalar: ...
+    def to_arrow(
+        self, metadata: list | str | None = None, stream: Stream | None = None
+    ) -> PaScalar: ...
     @staticmethod
     def from_arrow(
         pa_val: Any,

python/pylibcudf/pylibcudf/scalar.pyx

@@ -59,6 +59,7 @@ from .traits cimport is_floating_point
 from .types cimport DataType
 from .utils cimport _get_stream
 from functools import singledispatch
+from ._interop_helpers import ArrowLike, ColumnMetadata
 
 try:
     import pyarrow as pa
@@ -150,6 +151,29 @@ cdef class Scalar:
         """True if the scalar is valid, false if not"""
         return self.get().is_valid()
 
+    def to_arrow(
+        self,
+        metadata: list[ColumnMetadata] | str | None = None,
+        stream: Stream = None,
+    ) -> ArrowLike:
+        """Create a PyArrow array from a pylibcudf scalar.
+
+        Parameters
+        ----------
+        metadata : list
+            The metadata to attach to the columns of the table.
+        stream : Stream | None
+            CUDA stream on which to perform the operation.
+
+        Returns
+        -------
+        pyarrow.Scalar
+        """
+        stream = _get_stream(stream)
+        # Note that metadata for scalars is primarily important for preserving
+        # information on nested types since names are otherwise irrelevant.
+        return Column.from_scalar(self, 1, stream).to_arrow(metadata=metadata)[0]
+
     @staticmethod
     def from_arrow(
         pa_val,

python/pylibcudf/pylibcudf/table.pyi

@@ -10,6 +10,7 @@ class Table:
     def num_rows(self) -> int: ...
     def shape(self) -> tuple[int, int]: ...
     def columns(self) -> list[Column]: ...
+    def to_arrow(self, metadata: list) -> ArrowLike: ...
     @staticmethod
     def from_arrow(
         arrow_like: ArrowLike, dtype: DataType | None = None

python/pylibcudf/pylibcudf/table.pyx

@@ -36,7 +36,15 @@ from pylibcudf._interop_helpers cimport (
     _release_device_array,
     _metadata_to_libcudf,
 )
-from ._interop_helpers import ArrowLike, ColumnMetadata
+from ._interop_helpers import ArrowLike, ColumnMetadata, _ObjectWithArrowMetadata
+
+try:
+    import pyarrow as pa
+    pa_err = None
+except ImportError as e:
+    pa = None
+    pa_err = e
+
 
 __all__ = ["Table"]
 
@@ -61,6 +69,22 @@ cdef class Table:
             raise ValueError("All columns must be pylibcudf Column objects")
         self._columns = columns
 
+    def to_arrow(
+        self,
+        metadata: list[ColumnMetadata | str] | None = None
+    ) -> ArrowLike:
+        """Create a PyArrow table from a pylibcudf table."""
+        if pa_err is not None:
+            raise RuntimeError(
+                "pyarrow was not found on your system. Please "
+                "pip install pylibcudf with the [pyarrow] extra for a "
+                "compatible pyarrow version."
+            ) from pa_err
+        # TODO: Once the arrow C device interface registers more
+        # types that it supports, we can call pa.array(self) if
+        # no metadata is passed.
+        return pa.table(_ObjectWithArrowMetadata(self, metadata))
+
     @staticmethod
     def from_arrow(obj: ArrowLike, dtype: DataType | None = None) -> Table:
         """

python/pylibcudf/pylibcudf/types.pyi

@@ -85,6 +85,7 @@ class DataType:
     def __init__(self, type_id: TypeId, scale: int = 0): ...
     def id(self) -> TypeId: ...
     def scale(self) -> int: ...
+    def to_arrow(self, **kwargs) -> PyarrowDataType: ...
     @staticmethod
     def from_arrow(pa_typ: PyarrowDataType) -> DataType: ...
     def from_py(self, type: type) -> DataType: ...

python/pylibcudf/pylibcudf/types.pyx

@@ -24,6 +24,13 @@ from pylibcudf.libcudf.types import sorted as Sorted  # no-cython-lint, isort:sk
 
 from functools import cache
 
+try:
+    import pyarrow as pa
+    pa_err = None
+except ImportError as e:
+    pa = None
+    pa_err = e
+
 try:
     import pyarrow as pa
 
@@ -189,6 +196,60 @@ cdef class DataType:
         ret.c_obj = dt
         return ret
 
+    def to_arrow(self, **kwargs):
+        """
+        Convert a datatype to arrow.
+
+        Returns
+        -------
+        pyarrow.DataType
+
+        Notes
+        -----
+        Translation of some types requires extra information as a keyword
+        argument. Specifically:
+
+        - When translating a decimal type, provide ``precision``
+        - When translating a struct type, provide ``fields``
+        - When translating a list type, provide the wrapped ``value_type``
+        """
+        if pa_err is not None:
+            raise RuntimeError(
+                "pyarrow was not found on your system. Please "
+                "pip install pylibcudf with the [pyarrow] extra for a "
+                "compatible pyarrow version."
+            ) from pa_err
+        if self.id() in {
+            type_id.DECIMAL32,
+            type_id.DECIMAL64,
+            type_id.DECIMAL128
+        }:
+            if not (precision := kwargs.get("precision")):
+                raise ValueError(
+                    "Precision must be provided for decimal types"
+                )
+                # no pa.decimal32 or pa.decimal64
+            return pa.decimal128(precision, -self.scale())
+        elif self.id() == type_id.STRUCT:
+            if not (fields := kwargs.get("fields")):
+                raise ValueError(
+                    "Fields must be provided for struct types"
+                )
+            return pa.struct(fields)
+        elif self.id() == type_id.LIST:
+            if not (value_type := kwargs.get("value_type")):
+                raise ValueError(
+                    "Value type must be provided for list types"
+                )
+            return pa.list_(value_type)
+        else:
+            try:
+                return LIBCUDF_TO_ARROW_TYPES[self.id()]
+            except KeyError:
+                raise TypeError(
+                    f"Unable to convert {self.id()} to arrow datatype"
+                )
+
     @staticmethod
     def from_arrow(pa_typ) -> DataType:
         """