diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py index b162ba6752c9..d1f6bfbf025e 100644 --- a/docs/cudf/source/conf.py +++ b/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"), ] diff --git a/python/cudf_polars/cudf_polars/containers/dataframe.py b/python/cudf_polars/cudf_polars/containers/dataframe.py index 3a095be3cfeb..6cb4012988af 100644 --- a/python/cudf_polars/cudf_polars/containers/dataframe.py +++ b/python/cudf_polars/cudf_polars/containers/dataframe.py @@ -55,7 +55,7 @@ def _create_polars_column_metadata( # This is also defined in pylibcudf.interop class _ObjectWithArrowMetadata: def __init__( - self, obj: plc.Table, metadata: list[plc.interop.ColumnMetadata] + self, obj: plc.Table | plc.Column, metadata: list[plc.interop.ColumnMetadata] ) -> None: self.obj = obj self.metadata = metadata diff --git a/python/pylibcudf/pylibcudf/_interop_helpers.pyx b/python/pylibcudf/pylibcudf/_interop_helpers.pyx index 3e474426898c..d51d46830acd 100644 --- a/python/pylibcudf/pylibcudf/_interop_helpers.pyx +++ b/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. diff --git a/python/pylibcudf/pylibcudf/column.pyi b/python/pylibcudf/pylibcudf/column.pyi index 42e2c7cccd1b..ff4945c123ef 100644 --- a/python/pylibcudf/pylibcudf/column.pyi +++ b/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 diff --git a/python/pylibcudf/pylibcudf/column.pyx b/python/pylibcudf/pylibcudf/column.pyx index a788b42c3940..06a488902b66 100644 --- a/python/pylibcudf/pylibcudf/column.pyx +++ b/python/pylibcudf/pylibcudf/column.pyx @@ -56,7 +56,7 @@ from ._interop_helpers cimport ( from .utils cimport _get_stream, _get_memory_resource 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 @@ -64,6 +64,14 @@ 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"] @@ -337,12 +345,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. diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index 0595bb4a777c..75edb59b093c 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/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,8 +104,13 @@ 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): @@ -114,70 +118,22 @@ if pa is not None: @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): diff --git a/python/pylibcudf/pylibcudf/scalar.pyi b/python/pylibcudf/pylibcudf/scalar.pyi index be84726ef18d..be38a979813a 100644 --- a/python/pylibcudf/pylibcudf/scalar.pyi +++ b/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, diff --git a/python/pylibcudf/pylibcudf/scalar.pyx b/python/pylibcudf/pylibcudf/scalar.pyx index 7c183b96bbba..46fb1faa9b08 100644 --- a/python/pylibcudf/pylibcudf/scalar.pyx +++ b/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, diff --git a/python/pylibcudf/pylibcudf/table.pyi b/python/pylibcudf/pylibcudf/table.pyi index a505d87d85e3..c8d11b99e953 100644 --- a/python/pylibcudf/pylibcudf/table.pyi +++ b/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 diff --git a/python/pylibcudf/pylibcudf/table.pyx b/python/pylibcudf/pylibcudf/table.pyx index afe8bc6ab4d2..38dcb0f0291e 100644 --- a/python/pylibcudf/pylibcudf/table.pyx +++ b/python/pylibcudf/pylibcudf/table.pyx @@ -37,7 +37,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"] @@ -62,6 +70,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.table(self) if + # no metadata is passed. + return pa.table(_ObjectWithArrowMetadata(self, metadata)) + @staticmethod def from_arrow(obj: ArrowLike, dtype: DataType | None = None) -> Table: """ diff --git a/python/pylibcudf/pylibcudf/types.pyi b/python/pylibcudf/pylibcudf/types.pyi index e3c0951da520..b56dc54dbcc3 100644 --- a/python/pylibcudf/pylibcudf/types.pyi +++ b/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: ... diff --git a/python/pylibcudf/pylibcudf/types.pyx b/python/pylibcudf/pylibcudf/types.pyx index a84171f991c5..2b4927f703b1 100644 --- a/python/pylibcudf/pylibcudf/types.pyx +++ b/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: """ diff --git a/python/pylibcudf/tests/test_column_to_arrow.py b/python/pylibcudf/tests/test_column_to_arrow.py new file mode 100644 index 000000000000..661210c75748 --- /dev/null +++ b/python/pylibcudf/tests/test_column_to_arrow.py @@ -0,0 +1,9 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. + +from utils import assert_column_eq + + +def test_column_to_arrow(table_data): + plc_tbl, _ = table_data + for col in plc_tbl.tbl.columns(): + assert_column_eq(col, col.to_arrow()) diff --git a/python/pylibcudf/tests/test_table.py b/python/pylibcudf/tests/test_table.py index ae39f80efef7..7c81f56b653c 100644 --- a/python/pylibcudf/tests/test_table.py +++ b/python/pylibcudf/tests/test_table.py @@ -2,6 +2,7 @@ import pyarrow as pa import pytest +from utils import assert_table_eq import pylibcudf as plc @@ -19,3 +20,14 @@ def test_table_shape(arrow_tbl): plc_tbl = plc.Table.from_arrow(arrow_tbl) assert plc_tbl.shape() == arrow_tbl.shape + + +def test_table_to_arrow(table_data): + plc_tbl, _ = table_data + expect = plc_tbl.tbl + got = expect.to_arrow() + # The order of `got` and `expect` is reversed here + # because in almost all pylibcudf tests the `expect` + # is a pyarrow object while `got` is a pylibcudf object, + # whereas in this case those types are reversed. + assert_table_eq(got, expect)