diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml index 77edaeff186..e20be195385 100644 --- a/conda/environments/all_cuda-118_arch-aarch64.yaml +++ b/conda/environments/all_cuda-118_arch-aarch64.yaml @@ -70,7 +70,7 @@ dependencies: - pyarrow>=14.0.0,<20.0.0a0 - pydata-sphinx-theme>=0.15.4 - pynvml>=12.0.0,<13.0.0a0 -- pytest-benchmark +- pytest-benchmark<5.1.0 - pytest-cases>=3.8.2 - pytest-cov - pytest-rerunfailures diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml index 0206564525e..040c3d0ba64 100644 --- a/conda/environments/all_cuda-118_arch-x86_64.yaml +++ b/conda/environments/all_cuda-118_arch-x86_64.yaml @@ -72,7 +72,7 @@ dependencies: - pyarrow>=14.0.0,<20.0.0a0 - pydata-sphinx-theme>=0.15.4 - pynvml>=12.0.0,<13.0.0a0 -- pytest-benchmark +- pytest-benchmark<5.1.0 - pytest-cases>=3.8.2 - pytest-cov - pytest-rerunfailures diff --git a/conda/environments/all_cuda-128_arch-aarch64.yaml b/conda/environments/all_cuda-128_arch-aarch64.yaml index 8921d72324d..264d60c2bd9 100644 --- a/conda/environments/all_cuda-128_arch-aarch64.yaml +++ b/conda/environments/all_cuda-128_arch-aarch64.yaml @@ -69,7 +69,7 @@ dependencies: - pydata-sphinx-theme>=0.15.4 - pynvjitlink>=0.0.0a0 - pynvml>=12.0.0,<13.0.0a0 -- pytest-benchmark +- pytest-benchmark<5.1.0 - pytest-cases>=3.8.2 - pytest-cov - pytest-rerunfailures diff --git a/conda/environments/all_cuda-128_arch-x86_64.yaml b/conda/environments/all_cuda-128_arch-x86_64.yaml index 0c8a43af017..ea2e83116fb 100644 --- a/conda/environments/all_cuda-128_arch-x86_64.yaml +++ b/conda/environments/all_cuda-128_arch-x86_64.yaml @@ -70,7 +70,7 @@ dependencies: - pydata-sphinx-theme>=0.15.4 - pynvjitlink>=0.0.0a0 - pynvml>=12.0.0,<13.0.0a0 -- pytest-benchmark +- pytest-benchmark<5.1.0 - pytest-cases>=3.8.2 - pytest-cov - pytest-rerunfailures diff --git a/cpp/include/cudf/interop.hpp b/cpp/include/cudf/interop.hpp index e60038419ff..7f01d3345f4 100644 --- a/cpp/include/cudf/interop.hpp +++ b/cpp/include/cudf/interop.hpp @@ -314,7 +314,7 @@ class arrow_column { * * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used for any allocations during conversion - * @return unique_column_view_t containing a view of the column data + * @return A view of the column data */ [[nodiscard]] column_view view() const; @@ -434,7 +434,7 @@ class arrow_table { * * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used for any allocations during conversion - * @return unique_table_view_t containing a view of the table data + * @return A view of the table data */ [[nodiscard]] table_view view() const; diff --git a/cpp/src/column/column_view.cpp b/cpp/src/column/column_view.cpp index a7718d19a94..d7214b6d06e 100644 --- a/cpp/src/column/column_view.cpp +++ b/cpp/src/column/column_view.cpp @@ -37,7 +37,10 @@ template void prefetch_col_data(ColumnView& col, void const* data_ptr, std::string_view key) noexcept { if (cudf::experimental::prefetch::detail::prefetch_config::instance().get(key)) { - if (cudf::is_fixed_width(col.type())) { + if (col.type().id() == cudf::type_id::EMPTY) { + // Skip prefetching for empty columns + return; + } else if (cudf::is_fixed_width(col.type())) { cudf::experimental::prefetch::detail::prefetch_noexcept( key, data_ptr, col.size() * size_of(col.type()), cudf::get_default_stream()); } else if (col.type().id() == type_id::STRING) { diff --git a/cpp/src/interop/from_arrow_device.cu b/cpp/src/interop/from_arrow_device.cu index 836da2987e2..74eb6d1e6b2 100644 --- a/cpp/src/interop/from_arrow_device.cu +++ b/cpp/src/interop/from_arrow_device.cu @@ -284,7 +284,7 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()( size_type const offset = input->offset; size_type const null_count = input->null_count; auto offsets_view = column_view{data_type(type_id::INT32), - offset + num_rows + 1, + (num_rows == 0) ? 0 : (offset + num_rows + 1), input->buffers[fixed_width_data_buffer_idx], nullptr, 0, @@ -300,8 +300,9 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()( // in the scenario where we were sliced and there are more elements in the child_view // than can be referenced by the sliced offsets, we need to slice the child_view // so that when `get_sliced_child` is called, we still produce the right result - auto max_child_offset = cudf::detail::get_value(offsets_view, offset + num_rows, stream); - child_view = cudf::slice(child_view, {0, max_child_offset}, stream).front(); + auto max_child_offset = + num_rows == 0 ? 0 : cudf::detail::get_value(offsets_view, offset + num_rows, stream); + child_view = cudf::slice(child_view, {0, max_child_offset}, stream).front(); return std::make_tuple( {type, diff --git a/cpp/src/interop/to_arrow_device.cu b/cpp/src/interop/to_arrow_device.cu index ececbc8ebdb..1f043420daf 100644 --- a/cpp/src/interop/to_arrow_device.cu +++ b/cpp/src/interop/to_arrow_device.cu @@ -91,6 +91,21 @@ int set_buffer(std::unique_ptr device_buf, int64_t i, ArrowArray* out) return NANOARROW_OK; } +int set_null_mask(column::contents& contents, ArrowArray* out) +{ + if (contents.null_mask) { + NANOARROW_RETURN_NOT_OK(set_buffer(std::move(contents.null_mask), validity_buffer_idx, out)); + } + return NANOARROW_OK; +} + +int set_contents(column::contents& contents, ArrowArray* out) +{ + NANOARROW_RETURN_NOT_OK(set_null_mask(contents, out)); + NANOARROW_RETURN_NOT_OK(set_buffer(std::move(contents.data), fixed_width_data_buffer_idx, out)); + return NANOARROW_OK; +} + struct dispatch_to_arrow_device { template () and not is_fixed_point())> @@ -117,23 +132,16 @@ struct dispatch_to_arrow_device { ArrowArrayMove(tmp.get(), out); return NANOARROW_OK; } - - int set_null_mask(column::contents& contents, ArrowArray* out) - { - if (contents.null_mask) { - NANOARROW_RETURN_NOT_OK(set_buffer(std::move(contents.null_mask), validity_buffer_idx, out)); - } - return NANOARROW_OK; - } - - int set_contents(column::contents& contents, ArrowArray* out) - { - NANOARROW_RETURN_NOT_OK(set_null_mask(contents, out)); - NANOARROW_RETURN_NOT_OK(set_buffer(std::move(contents.data), fixed_width_data_buffer_idx, out)); - return NANOARROW_OK; - } }; +int handle_empty_type_column(ArrowArray* array, cudf::column& column) +{ + NANOARROW_RETURN_NOT_OK(initialize_array(array, NANOARROW_TYPE_NA, column.view())); + auto child_contents = column.release(); + NANOARROW_RETURN_NOT_OK(set_contents(child_contents, array)); + return NANOARROW_OK; +} + template <> int dispatch_to_arrow_device::operator()(cudf::column&& column, rmm::cuda_stream_view stream, @@ -227,8 +235,12 @@ int dispatch_to_arrow_device::operator()(cudf::column&& colum for (size_t i = 0; i < size_t(tmp->n_children); ++i) { ArrowArray* child_ptr = tmp->children[i]; auto& child = contents.children[i]; - NANOARROW_RETURN_NOT_OK(cudf::type_dispatcher( - child->type(), dispatch_to_arrow_device{}, std::move(*child), stream, mr, child_ptr)); + if (child->type().id() == cudf::type_id::EMPTY) { + NANOARROW_RETURN_NOT_OK(handle_empty_type_column(child_ptr, *child)); + } else { + NANOARROW_RETURN_NOT_OK(cudf::type_dispatcher( + child->type(), dispatch_to_arrow_device{}, std::move(*child), stream, mr, child_ptr)); + } } ArrowArrayMove(tmp.get(), out); @@ -253,8 +265,12 @@ int dispatch_to_arrow_device::operator()(cudf::column&& column, NANOARROW_RETURN_NOT_OK(set_buffer(std::move(offsets_contents.data), 1, tmp.get())); auto& child = contents.children[cudf::lists_column_view::child_column_index]; - NANOARROW_RETURN_NOT_OK(cudf::type_dispatcher( - child->type(), dispatch_to_arrow_device{}, std::move(*child), stream, mr, tmp->children[0])); + if (child->type().id() == cudf::type_id::EMPTY) { + NANOARROW_RETURN_NOT_OK(handle_empty_type_column(tmp->children[0], *child)); + } else { + NANOARROW_RETURN_NOT_OK(cudf::type_dispatcher( + child->type(), dispatch_to_arrow_device{}, std::move(*child), stream, mr, tmp->children[0])); + } ArrowArrayMove(tmp.get(), out); return NANOARROW_OK; @@ -535,8 +551,12 @@ unique_device_array_t to_arrow_device(cudf::table&& table, for (size_t i = 0; i < cols.size(); ++i) { auto child = tmp->children[i]; auto col = cols[i].get(); - NANOARROW_THROW_NOT_OK(cudf::type_dispatcher( - col->type(), detail::dispatch_to_arrow_device{}, std::move(*col), stream, mr, child)); + if (col->type().id() == cudf::type_id::EMPTY) { + NANOARROW_THROW_NOT_OK(handle_empty_type_column(child, *col)); + } else { + NANOARROW_THROW_NOT_OK(cudf::type_dispatcher( + col->type(), detail::dispatch_to_arrow_device{}, std::move(*col), stream, mr, child)); + } } return create_device_array(std::move(tmp), stream); @@ -548,8 +568,12 @@ unique_device_array_t to_arrow_device(cudf::column&& col, { nanoarrow::UniqueArray tmp; - NANOARROW_THROW_NOT_OK(cudf::type_dispatcher( - col.type(), detail::dispatch_to_arrow_device{}, std::move(col), stream, mr, tmp.get())); + if (col.type().id() == cudf::type_id::EMPTY) { + NANOARROW_THROW_NOT_OK(handle_empty_type_column(tmp.get(), col)); + } else { + NANOARROW_THROW_NOT_OK(cudf::type_dispatcher( + col.type(), detail::dispatch_to_arrow_device{}, std::move(col), stream, mr, tmp.get())); + } return create_device_array(std::move(tmp), stream); } diff --git a/cpp/src/interop/to_arrow_schema.cpp b/cpp/src/interop/to_arrow_schema.cpp index aabba447ee2..5fbda0dea4c 100644 --- a/cpp/src/interop/to_arrow_schema.cpp +++ b/cpp/src/interop/to_arrow_schema.cpp @@ -152,8 +152,12 @@ int dispatch_to_arrow_type::operator()(column_view input, child->flags = col.has_nulls() ? ARROW_FLAG_NULLABLE : 0; - NANOARROW_RETURN_NOT_OK(cudf::type_dispatcher( - col.type(), detail::dispatch_to_arrow_type{}, col, metadata.children_meta[i], child)); + if (col.type().id() == cudf::type_id::EMPTY) { + NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(out->children[i], NANOARROW_TYPE_NA)); + } else { + NANOARROW_RETURN_NOT_OK(cudf::type_dispatcher( + col.type(), detail::dispatch_to_arrow_type{}, col, metadata.children_meta[i], child)); + } } return NANOARROW_OK; @@ -174,6 +178,9 @@ int dispatch_to_arrow_type::operator()(column_view input, out->flags = input.has_nulls() ? ARROW_FLAG_NULLABLE : 0; NANOARROW_RETURN_NOT_OK(ArrowSchemaSetName(out->children[0], child_meta.name.c_str())); out->children[0]->flags = child.has_nulls() ? ARROW_FLAG_NULLABLE : 0; + if (child.type().id() == cudf::type_id::EMPTY) { + return ArrowSchemaSetType(out->children[0], NANOARROW_TYPE_NA); + } return cudf::type_dispatcher( child.type(), detail::dispatch_to_arrow_type{}, child, child_meta, out->children[0]); } @@ -219,8 +226,12 @@ unique_schema_t to_arrow_schema(cudf::table_view const& input, NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(child, metadata[i].name.c_str())); child->flags = col.has_nulls() ? ARROW_FLAG_NULLABLE : 0; - NANOARROW_THROW_NOT_OK( - cudf::type_dispatcher(col.type(), detail::dispatch_to_arrow_type{}, col, metadata[i], child)); + if (col.type().id() == cudf::type_id::EMPTY) { + NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(child, NANOARROW_TYPE_NA)); + } else { + NANOARROW_THROW_NOT_OK(cudf::type_dispatcher( + col.type(), detail::dispatch_to_arrow_type{}, col, metadata[i], child)); + } } unique_schema_t out(new ArrowSchema, [](ArrowSchema* schema) { diff --git a/dependencies.yaml b/dependencies.yaml index 4af61084d1f..d17ec2e75ca 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -932,7 +932,14 @@ dependencies: - fastavro>=0.22.9 - hypothesis - mmh3 - - pytest-benchmark + # 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-cases>=3.8.2 - scipy - zstandard diff --git a/python/cudf/benchmarks/API/bench_dataframe.py b/python/cudf/benchmarks/API/bench_dataframe.py index ba243eb6a7c..5920aa9cb84 100644 --- a/python/cudf/benchmarks/API/bench_dataframe.py +++ b/python/cudf/benchmarks/API/bench_dataframe.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2024, NVIDIA CORPORATION. +# Copyright (c) 2022-2025, NVIDIA CORPORATION. """Benchmarks of DataFrame methods.""" @@ -6,12 +6,19 @@ import numba.cuda import numpy +import pyarrow as pa import pytest import pytest_cases from config import cudf, cupy from utils import benchmark_with_object +@pytest.mark.parametrize("N", [100, 1_000_000, 100_000_000]) +def bench_from_arrow(benchmark, N): + rng = numpy.random.default_rng(seed=10) + benchmark(cudf.DataFrame, {None: pa.array(rng.random(N))}) + + @pytest.mark.parametrize("N", [100, 1_000_000]) def bench_construction(benchmark, N): benchmark(cudf.DataFrame, {None: cupy.random.rand(N)}) diff --git a/python/cudf/cudf/core/column/column.py b/python/cudf/cudf/core/column/column.py index b5c1f1f7fd2..b0af1e5f784 100644 --- a/python/cudf/cudf/core/column/column.py +++ b/python/cudf/cudf/core/column/column.py @@ -329,7 +329,7 @@ def null_count(self) -> int: else: with acquire_spill_lock(): self._null_count = plc.null_mask.null_count( - self.base_mask.get_ptr(mode="read"), # type: ignore[union-attr] + plc.gpumemoryview(self.base_mask), # type: ignore[union-attr] self.offset, self.offset + self.size, ) @@ -750,7 +750,7 @@ def to_arrow(self) -> pa.Array: 4 ] """ - return plc.interop.to_arrow(self.to_pylibcudf(mode="read")).chunk(0) + return plc.interop.to_arrow(self.to_pylibcudf(mode="read")) @classmethod def from_arrow(cls, array: pa.Array) -> ColumnBase: diff --git a/python/cudf/pyproject.toml b/python/cudf/pyproject.toml index e43cbd05279..5080a578b91 100644 --- a/python/cudf/pyproject.toml +++ b/python/cudf/pyproject.toml @@ -56,7 +56,7 @@ test = [ "hypothesis", "mmh3", "msgpack", - "pytest-benchmark", + "pytest-benchmark<5.1.0", "pytest-cases>=3.8.2", "pytest-cov", "pytest-rerunfailures", diff --git a/python/pylibcudf/pylibcudf/CMakeLists.txt b/python/pylibcudf/pylibcudf/CMakeLists.txt index 147060f8747..0e96a7531fb 100644 --- a/python/pylibcudf/pylibcudf/CMakeLists.txt +++ b/python/pylibcudf/pylibcudf/CMakeLists.txt @@ -28,6 +28,7 @@ set(cython_sources groupby.pyx hashing.pyx interop.pyx + _interop_helpers.pyx jit.pyx join.pyx json.pyx @@ -61,11 +62,13 @@ rapids_cython_create_modules( LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX pylibcudf_ ASSOCIATED_TARGETS cudf ) -target_include_directories(pylibcudf_interop PUBLIC "$") - include(${rapids-cmake-dir}/export/find_package_root.cmake) include(../../../cpp/cmake/thirdparty/get_nanoarrow.cmake) -target_link_libraries(pylibcudf_interop PUBLIC nanoarrow) + +foreach(source interop _interop_helpers table column) + target_include_directories(pylibcudf_${source} PUBLIC "$") + target_link_libraries(pylibcudf_${source} PUBLIC nanoarrow) +endforeach() add_subdirectory(libcudf) add_subdirectory(strings) diff --git a/python/pylibcudf/pylibcudf/_interop_helpers.pxd b/python/pylibcudf/pylibcudf/_interop_helpers.pxd new file mode 100644 index 00000000000..72036be4e77 --- /dev/null +++ b/python/pylibcudf/pylibcudf/_interop_helpers.pxd @@ -0,0 +1,9 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. + +from pylibcudf.libcudf.interop cimport column_metadata + +cdef void _release_schema(object schema_capsule) noexcept + +cdef void _release_array(object array_capsule) noexcept + +cdef column_metadata _metadata_to_libcudf(metadata) diff --git a/python/pylibcudf/pylibcudf/_interop_helpers.pyx b/python/pylibcudf/pylibcudf/_interop_helpers.pyx new file mode 100644 index 00000000000..f2fa6ccf549 --- /dev/null +++ b/python/pylibcudf/pylibcudf/_interop_helpers.pyx @@ -0,0 +1,55 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. + +from cpython.pycapsule cimport PyCapsule_GetPointer + +from pylibcudf.libcudf.interop cimport ( + ArrowArray, + ArrowSchema, + column_metadata, + release_arrow_array_raw, + release_arrow_schema_raw, +) + +from dataclasses import dataclass, field + + +@dataclass +class ColumnMetadata: + """Metadata associated with a column. + + This is the Python representation of :cpp:class:`cudf::column_metadata`. + """ + name: str = "" + children_meta: list[ColumnMetadata] = field(default_factory=list) + + +cdef void _release_schema(object schema_capsule) noexcept: + """Release the ArrowSchema object stored in a PyCapsule.""" + cdef ArrowSchema* schema = PyCapsule_GetPointer( + schema_capsule, 'arrow_schema' + ) + release_arrow_schema_raw(schema) + + +cdef void _release_array(object array_capsule) noexcept: + """Release the ArrowArray object stored in a PyCapsule.""" + cdef ArrowArray* array = PyCapsule_GetPointer( + array_capsule, 'arrow_array' + ) + release_arrow_array_raw(array) + + +cdef column_metadata _metadata_to_libcudf(metadata): + """Convert a ColumnMetadata object to C++ column_metadata. + + Since this class is mutable and cheap, it is easier to create the C++ + object on the fly rather than have it directly backing the storage for + the Cython class. Additionally, this structure restricts the dependency + on C++ types to just within this module, allowing us to make the module a + pure Python module (from an import sense, i.e. no pxd declarations). + """ + cdef column_metadata c_metadata + c_metadata.name = metadata.name.encode() + for child_meta in metadata.children_meta: + c_metadata.children_meta.push_back(_metadata_to_libcudf(child_meta)) + return c_metadata diff --git a/python/pylibcudf/pylibcudf/column.pxd b/python/pylibcudf/pylibcudf/column.pxd index 14589eb3c8e..cdcb05a219a 100644 --- a/python/pylibcudf/pylibcudf/column.pxd +++ b/python/pylibcudf/pylibcudf/column.pxd @@ -15,6 +15,22 @@ from .gpumemoryview cimport gpumemoryview from .types cimport DataType +cdef class OwnerWithCAI: + cdef object owner + cdef dict cai + + @staticmethod + cdef create(column_view cv, object owner) + + +cdef class OwnerMaskWithCAI: + cdef object owner + cdef dict cai + + @staticmethod + cdef create(column_view cv, object owner) + + cdef class Column: # TODO: Should we document these attributes? Should we mark them readonly? cdef: diff --git a/python/pylibcudf/pylibcudf/column.pyx b/python/pylibcudf/pylibcudf/column.pyx index 26040f142a8..7e1fec675ce 100644 --- a/python/pylibcudf/pylibcudf/column.pyx +++ b/python/pylibcudf/pylibcudf/column.pyx @@ -1,30 +1,138 @@ # Copyright (c) 2023-2025, NVIDIA CORPORATION. from cython.operator cimport dereference + +from cpython.pycapsule cimport ( + PyCapsule_GetPointer, + PyCapsule_New, +) + from libc.stdint cimport uintptr_t + from libcpp.limits cimport numeric_limits from libcpp.memory cimport make_unique, unique_ptr from libcpp.utility cimport move -from rmm.pylibrmm.stream cimport Stream + from pylibcudf.libcudf.column.column cimport column, column_contents from pylibcudf.libcudf.column.column_factories cimport make_column_from_scalar -from pylibcudf.libcudf.scalar.scalar cimport scalar -from pylibcudf.libcudf.types cimport size_type, bitmask_type +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, + column_metadata, + to_arrow_host_raw, + to_arrow_schema_raw, +) from rmm.pylibrmm.device_buffer cimport DeviceBuffer +from rmm.pylibrmm.stream cimport Stream from .gpumemoryview cimport gpumemoryview from .filling cimport sequence from .scalar cimport Scalar from .types cimport DataType, size_of, type_id +from ._interop_helpers cimport ( + _release_schema, + _release_array, + _metadata_to_libcudf, +) +from .null_mask cimport bitmask_allocation_size_bytes from .utils cimport _get_stream -from functools import cache +from ._interop_helpers import ColumnMetadata + +import functools __all__ = ["Column", "ListColumnView", "is_c_contiguous"] +class _ArrowLikeMeta(type): + def __subclasscheck__(cls, other): + return hasattr(other, "__arrow_c_array__") + + +class _ArrowLike(metaclass=_ArrowLikeMeta): + pass + + +cdef class _ArrowColumnHolder: + """A holder for an Arrow column for gpumemoryview lifetime management.""" + cdef unique_ptr[arrow_column] col + + +cdef class OwnerWithCAI: + """An interface for column view's data with gpumemoryview via CAI.""" + @staticmethod + cdef create(column_view cv, object owner): + obj = OwnerWithCAI() + obj.owner = owner + cdef int size + cdef column_view offsets_column + cdef unique_ptr[scalar] last_offset + if cv.type().id() == type_id.EMPTY: + size = cv.size() + elif is_fixed_width(cv.type()): + size = cv.size() * cpp_size_of(cv.type()) + elif cv.type().id() == type_id.STRING: + # The size of the character array in the parent is the offsets size + num_children = cv.num_children() + size = 0 + # A strings column with no children is created for empty/all null + if num_children: + offsets_column = cv.child(0) + last_offset = get_element(offsets_column, offsets_column.size() - 1) + size = ( last_offset.get()).value() + else: + # All other types store data in the children, so the parent size is 0 + size = 0 + + obj.cai = { + "shape": (size,), + "strides": None, + # For the purposes in this function, just treat all of the types as byte + # streams of the appropriate size. This matches what we currently get from + # rmm.DeviceBuffer + "typestr": "|u1", + "data": ( cv.head[char](), False), + "version": 3, + } + return obj + + @property + def __cuda_array_interface__(self): + return self.cai + + +cdef class OwnerMaskWithCAI: + """An interface for column view's null mask with gpumemoryview via CAI.""" + @staticmethod + cdef create(column_view cv, object owner): + obj = OwnerMaskWithCAI() + obj.owner = owner + + obj.cai = { + "shape": (bitmask_allocation_size_bytes(cv.size()),), + "strides": None, + # For the purposes in this function, just treat all of the types as byte + # streams of the appropriate size. This matches what we currently get from + # rmm.DeviceBuffer + "typestr": "|u1", + "data": ( cv.null_mask(), False), + "version": 3, + } + return obj + + @property + def __cuda_array_interface__(self): + return self.cai + + class _Ravelled: def __init__(self, obj): self.obj = obj @@ -61,7 +169,26 @@ cdef class Column: children : list The children of this column if it is a compound column type. """ - def __init__( + def __init__(self, obj=None, *args, **kwargs): + self._init(obj, *args, **kwargs) + + __hash__ = None + + @functools.singledispatchmethod + def _init(self, obj, *args, **kwargs): + if obj is None: + if (data_type := kwargs.get("data_type")) is not None: + kwargs.pop("data_type") + self._init(data_type, *args, **kwargs) + return + elif (arrow_like := kwargs.get("arrow_like")) is not None: + kwargs.pop("arrow_like") + self._init(arrow_like, *args, **kwargs) + return + raise ValueError(f"Invalid input type {type(obj)}") + + @_init.register(DataType) + def _( self, DataType data_type not None, size_type size, gpumemoryview data, gpumemoryview mask, size_type null_count, size_type offset, list children @@ -77,7 +204,34 @@ cdef class Column: self._children = children self._num_children = len(children) - __hash__ = None + @_init.register(_ArrowLike) + def _(self, arrow_like): + schema, array = arrow_like.__arrow_c_array__() + cdef ArrowSchema* c_schema = ( + PyCapsule_GetPointer(schema, "arrow_schema") + ) + cdef ArrowArray* c_array = ( + PyCapsule_GetPointer(array, "arrow_array") + ) + + cdef _ArrowColumnHolder result = _ArrowColumnHolder() + cdef unique_ptr[arrow_column] c_result + 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(), + ) cdef column_view view(self) nogil: """Generate a libcudf column_view to pass to libcudf algorithms. @@ -269,12 +423,10 @@ cdef class Column: Column.from_column_view_of_arbitrary(cv.child(i), owner) ) - cdef gpumemoryview owning_data = gpumemoryview.from_pointer( - cv.head[char](), owner - ) - cdef gpumemoryview owning_mask = gpumemoryview.from_pointer( - cv.null_mask(), owner - ) + cdef gpumemoryview owning_data = gpumemoryview(OwnerWithCAI.create(cv, owner)) + cdef gpumemoryview owning_mask = None + if cv.null_count() > 0: + owning_mask = gpumemoryview(OwnerMaskWithCAI.create(cv, owner)) return Column( DataType.from_libcudf(cv.type()), @@ -529,6 +681,41 @@ cdef class Column: c_result = make_unique[column](self.view()) return Column.from_libcudf(move(c_result)) + def _create_nested_column_metadata(self): + return ColumnMetadata( + children_meta=[ + child._create_nested_column_metadata() for child in self.children() + ] + ) + + def _to_schema(self, metadata=None): + """Create an Arrow schema from this Column.""" + if metadata is None: + metadata = self._create_nested_column_metadata() + elif isinstance(metadata, str): + metadata = ColumnMetadata(metadata) + + cdef column_metadata c_metadata = _metadata_to_libcudf(metadata) + + cdef ArrowSchema* raw_schema_ptr + with nogil: + raw_schema_ptr = to_arrow_schema_raw(self.view(), c_metadata) + + return PyCapsule_New(raw_schema_ptr, 'arrow_schema', _release_schema) + + def _to_host_array(self): + cdef ArrowArray* raw_host_array_ptr + with nogil: + raw_host_array_ptr = to_arrow_host_raw(self.view()) + + return PyCapsule_New(raw_host_array_ptr, "arrow_array", _release_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() + cdef class ListColumnView: """Accessor for methods of a Column that are specific to lists.""" @@ -557,7 +744,7 @@ cdef class ListColumnView: return lists_column_view(self._column.view()) -@cache +@functools.cache def _datatype_from_dtype_desc(desc): mapping = { 'u1': type_id.UINT8, diff --git a/python/pylibcudf/pylibcudf/gpumemoryview.pxd b/python/pylibcudf/pylibcudf/gpumemoryview.pxd index 51fe5e4357e..c5cd32920dd 100644 --- a/python/pylibcudf/pylibcudf/gpumemoryview.pxd +++ b/python/pylibcudf/pylibcudf/gpumemoryview.pxd @@ -6,6 +6,4 @@ cdef class gpumemoryview: # to treat this object as something like a POD struct cdef readonly uintptr_t ptr cdef readonly object obj - - @staticmethod - cdef gpumemoryview from_pointer(uintptr_t ptr, object owner) + cdef readonly dict cai diff --git a/python/pylibcudf/pylibcudf/gpumemoryview.pyx b/python/pylibcudf/pylibcudf/gpumemoryview.pyx index a14a027c33a..6c13cac4f3f 100644 --- a/python/pylibcudf/pylibcudf/gpumemoryview.pyx +++ b/python/pylibcudf/pylibcudf/gpumemoryview.pyx @@ -2,7 +2,6 @@ import functools import operator -from libc.stdint cimport uintptr_t __all__ = ["gpumemoryview"] @@ -24,32 +23,13 @@ cdef class gpumemoryview: "the CUDA array interface" ) self.obj = obj + self.cai = cai # TODO: Need to respect readonly self.ptr = cai["data"][0] - @staticmethod - cdef gpumemoryview from_pointer(uintptr_t ptr, object owner): - """Create a gpumemoryview from a pointer and an owning object. - - Parameters - ---------- - ptr : uintptr_t - The pointer to the memory. - owner : object - The object that owns the data the pointer points to. - - Returns - ------- - gpumemoryview - """ - cdef gpumemoryview out = gpumemoryview.__new__(gpumemoryview) - out.obj = owner - out.ptr = ptr - return out - @property def __cuda_array_interface__(self): - return self.obj.__cuda_array_interface__ + return self.cai def __len__(self): return self.obj.__cuda_array_interface__["shape"][0] diff --git a/python/pylibcudf/pylibcudf/interop.pxd b/python/pylibcudf/pylibcudf/interop.pxd index 2a0a8c15fdd..7cf3be08e09 100644 --- a/python/pylibcudf/pylibcudf/interop.pxd +++ b/python/pylibcudf/pylibcudf/interop.pxd @@ -1,8 +1,7 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. from pylibcudf.table cimport Table - cpdef Table from_dlpack(object managed_tensor) cpdef object to_dlpack(Table input) diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index ec680a8e53e..b7b83c35586 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/python/pylibcudf/pylibcudf/interop.pyx @@ -8,27 +8,14 @@ from cpython.pycapsule cimport ( ) from libcpp.memory cimport unique_ptr from libcpp.utility cimport move -from libcpp.vector cimport vector -from dataclasses import dataclass, field from functools import singledispatch from pyarrow import lib as pa -from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.interop cimport ( - ArrowArray, - ArrowArrayStream, - ArrowSchema, DLManagedTensor, - column_metadata, - from_arrow_column as cpp_from_arrow_column, - from_arrow_stream as cpp_from_arrow_stream, from_dlpack as cpp_from_dlpack, - release_arrow_array_raw, - release_arrow_schema_raw, - to_arrow_host_raw, - to_arrow_schema_raw, to_dlpack as cpp_to_dlpack, ) from pylibcudf.libcudf.table.table cimport table @@ -38,6 +25,7 @@ from .column cimport Column from .scalar cimport Scalar from .table cimport Table from .types cimport DataType, type_id +from ._interop_helpers import ColumnMetadata __all__ = [ "ColumnMetadata", @@ -76,31 +64,6 @@ LIBCUDF_TO_ARROW_TYPES = { v: k for k, v in ARROW_TO_PYLIBCUDF_TYPES.items() } -cdef column_metadata _metadata_to_libcudf(metadata): - """Convert a ColumnMetadata object to C++ column_metadata. - - Since this class is mutable and cheap, it is easier to create the C++ - object on the fly rather than have it directly backing the storage for - the Cython class. Additionally, this structure restricts the dependency - on C++ types to just within this module, allowing us to make the module a - pure Python module (from an import sense, i.e. no pxd declarations). - """ - cdef column_metadata c_metadata - c_metadata.name = metadata.name.encode() - for child_meta in metadata.children_meta: - c_metadata.children_meta.push_back(_metadata_to_libcudf(child_meta)) - return c_metadata - - -@dataclass -class ColumnMetadata: - """Metadata associated with a column. - - This is the Python representation of :cpp:class:`cudf::column_metadata`. - """ - name: str = "" - children_meta: list[ColumnMetadata] = field(default_factory=list) - @singledispatch def from_arrow(pyarrow_object, *, DataType data_type=None): @@ -140,17 +103,7 @@ def _from_arrow_datatype(pyarrow_object): def _from_arrow_table(pyarrow_object, *, DataType data_type=None): if data_type is not None: raise ValueError("data_type may not be passed for tables") - stream = pyarrow_object.__arrow_c_stream__() - cdef ArrowArrayStream* c_stream = ( - PyCapsule_GetPointer(stream, "arrow_array_stream") - ) - - cdef unique_ptr[table] c_result - with nogil: - # The libcudf function here will release the stream. - c_result = cpp_from_arrow_stream(c_stream) - - return Table.from_libcudf(move(c_result)) + return Table(pyarrow_object) @from_arrow.register(pa.Scalar) @@ -173,33 +126,16 @@ def _from_arrow_column(pyarrow_object, *, DataType data_type=None): if data_type is not None: raise ValueError("data_type may not be passed for arrays") - schema, array = pyarrow_object.__arrow_c_array__() - cdef ArrowSchema* c_schema = ( - PyCapsule_GetPointer(schema, "arrow_schema") - ) - cdef ArrowArray* c_array = ( - PyCapsule_GetPointer(array, "arrow_array") - ) - - cdef unique_ptr[column] c_result - with nogil: - c_result = cpp_from_arrow_column(c_schema, c_array) - - # The capsule destructors should release automatically for us, but we - # choose to do it explicitly here for clarity. - c_schema.release(c_schema) - c_array.release(c_array) - - return Column.from_libcudf(move(c_result)) + return Column(pyarrow_object) @singledispatch -def to_arrow(cudf_object, metadata=None): +def to_arrow(plc_object, metadata=None): """Convert to a PyArrow object. Parameters ---------- - cudf_object : Union[Column, Table, Scalar] + plc_object : Union[Column, Table, Scalar] The cudf object to convert. metadata : list The metadata to attach to the columns of the table. @@ -209,11 +145,11 @@ def to_arrow(cudf_object, metadata=None): Union[pyarrow.Array, pyarrow.Table, pyarrow.Scalar] The converted object of type corresponding to the input type in PyArrow. """ - raise TypeError(f"Unsupported type {type(cudf_object)} for conversion to arrow") + raise TypeError(f"Unsupported type {type(plc_object)} for conversion to arrow") @to_arrow.register(DataType) -def _to_arrow_datatype(cudf_object, **kwargs): +def _to_arrow_datatype(plc_object, **kwargs): """ Convert a datatype to arrow. @@ -224,20 +160,20 @@ def _to_arrow_datatype(cudf_object, **kwargs): - When translating a struct type, provide ``fields`` - When translating a list type, provide the wrapped ``value_type`` """ - if cudf_object.id() in {type_id.DECIMAL32, type_id.DECIMAL64, type_id.DECIMAL128}: + 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, -cudf_object.scale()) - elif cudf_object.id() == type_id.STRUCT: + 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 cudf_object.id() == type_id.LIST: + 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" @@ -245,96 +181,39 @@ def _to_arrow_datatype(cudf_object, **kwargs): return pa.list_(value_type) else: try: - return LIBCUDF_TO_ARROW_TYPES[cudf_object.id()] + return LIBCUDF_TO_ARROW_TYPES[plc_object.id()] except KeyError: raise TypeError( - f"Unable to convert {cudf_object.id()} to arrow datatype" + f"Unable to convert {plc_object.id()} to arrow datatype" ) -cdef void _release_schema(object schema_capsule) noexcept: - """Release the ArrowSchema object stored in a PyCapsule.""" - cdef ArrowSchema* schema = PyCapsule_GetPointer( - schema_capsule, 'arrow_schema' - ) - release_arrow_schema_raw(schema) - - -cdef void _release_array(object array_capsule) noexcept: - """Release the ArrowArray object stored in a PyCapsule.""" - cdef ArrowArray* array = PyCapsule_GetPointer( - array_capsule, 'arrow_array' - ) - release_arrow_array_raw(array) - - -def _maybe_create_nested_column_metadata(Column col): - return ColumnMetadata( - children_meta=[ - _maybe_create_nested_column_metadata(child) for child in col.children() - ] - ) - - -def _table_to_schema(Table tbl, metadata): - if metadata is None: - metadata = [_maybe_create_nested_column_metadata(col) for col in tbl.columns()] - else: - metadata = [ColumnMetadata(m) if isinstance(m, str) else m for m in metadata] - - cdef vector[column_metadata] c_metadata - c_metadata.reserve(len(metadata)) - for meta in metadata: - c_metadata.push_back(_metadata_to_libcudf(meta)) - - cdef ArrowSchema* raw_schema_ptr - with nogil: - raw_schema_ptr = to_arrow_schema_raw(tbl.view(), c_metadata) - - return PyCapsule_New(raw_schema_ptr, 'arrow_schema', _release_schema) - - -def _table_to_host_array(Table tbl): - cdef ArrowArray* raw_host_array_ptr - with nogil: - raw_host_array_ptr = to_arrow_host_raw(tbl.view()) - - return PyCapsule_New(raw_host_array_ptr, "arrow_array", _release_array) - - -class _TableWithArrowMetadata: - def __init__(self, tbl, metadata=None): - self.tbl = tbl +class _ObjectWithArrowMetadata: + def __init__(self, obj, metadata=None): + self.obj = obj self.metadata = metadata def __arrow_c_array__(self, requested_schema=None): - return _table_to_schema(self.tbl, self.metadata), _table_to_host_array(self.tbl) + return self.obj._to_schema(self.metadata), self.obj._to_host_array() -# TODO: In the long run we should get rid of the `to_arrow` functions in favor of using -# the protocols directly via `pa.table(cudf_object, schema=...)` directly. We can do the -# same for columns. We cannot do this for scalars since there is no corresponding -# protocol. Since this will require broader changes throughout the codebase, the current -# approach is to leverage the protocol internally but to continue exposing `to_arrow`. @to_arrow.register(Table) -def _to_arrow_table(cudf_object, metadata=None): - test_table = _TableWithArrowMetadata(cudf_object, metadata) - return pa.table(test_table) +def _to_arrow_table(plc_object, metadata=None): + """Create a PyArrow table from a pylibcudf table.""" + return pa.table(_ObjectWithArrowMetadata(plc_object, metadata)) @to_arrow.register(Column) -def _to_arrow_array(cudf_object, metadata=None): +def _to_arrow_array(plc_object, metadata=None): """Create a PyArrow array from a pylibcudf column.""" - if metadata is not None: - metadata = [metadata] - return to_arrow(Table([cudf_object]), metadata)[0] + return pa.array(_ObjectWithArrowMetadata(plc_object, metadata)) @to_arrow.register(Scalar) -def _to_arrow_scalar(cudf_object, metadata=None): +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(cudf_object, 1), metadata=metadata)[0] + return to_arrow(Column.from_scalar(plc_object, 1), metadata=metadata)[0] cpdef Table from_dlpack(object managed_tensor): diff --git a/python/pylibcudf/pylibcudf/libcudf/interop.pxd b/python/pylibcudf/pylibcudf/libcudf/interop.pxd index b12990411b1..ab1d5bf5d5b 100644 --- a/python/pylibcudf/pylibcudf/libcudf/interop.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/interop.pxd @@ -46,13 +46,21 @@ cdef extern from "cudf/interop.hpp" namespace "cudf" \ string name vector[column_metadata] children_meta - cdef unique_ptr[table] from_arrow_stream( - ArrowArrayStream* input - ) except +libcudf_exception_handler - cdef unique_ptr[column] from_arrow_column( - const ArrowSchema* schema, - const ArrowArray* input - ) except +libcudf_exception_handler + +cdef extern from "cudf/interop.hpp" namespace "cudf::interop" \ + nogil: + cdef cppclass arrow_column: + arrow_column( + ArrowSchema&& schema, + ArrowArray&& array + ) except +libcudf_exception_handler + column_view view() except +libcudf_exception_handler + + cdef cppclass arrow_table: + arrow_table( + ArrowArrayStream&& stream, + ) except +libcudf_exception_handler + table_view view() except +libcudf_exception_handler cdef extern from *: @@ -72,6 +80,17 @@ cdef extern from *: return to_arrow_schema(input, metadata).release(); } + ArrowSchema* to_arrow_schema_raw( + cudf::column_view const& input, + cudf::column_metadata const& metadata) { + std::vector metadata_vec{metadata}; + cudf::table_view const& tbl = cudf::table_view({input}); + auto schema = cudf::to_arrow_schema(tbl, metadata_vec); + ArrowSchema *array_schema = new ArrowSchema(); + ArrowSchemaMove(schema->children[0], array_schema); + return array_schema; + } + void release_arrow_schema_raw(ArrowSchema *schema) { if (schema->release != nullptr) { schema->release(schema); @@ -90,6 +109,17 @@ cdef extern from *: return arr; } + ArrowArray* to_arrow_host_raw( + cudf::column_view const& col, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) { + // Assumes the sync event is null and the data is already on the host. + ArrowArray *arr = new ArrowArray(); + auto device_arr = cudf::to_arrow_host(col, stream, mr); + ArrowArrayMove(&device_arr->array, arr); + return arr; + } + void release_arrow_array_raw(ArrowArray *array) { if (array->release != nullptr) { array->release(array); @@ -97,16 +127,34 @@ cdef extern from *: delete array; } """ + # The `to_*_raw` functions are all defined in the above extern block as wrappers + # around libcudf functions that return unique_ptrs with non-default deleters, which + # are nontrivial to wrap in Cython. Since we need to manage them as raw pointers in + # Cython anyway, the inline C++ functions above are the simplest way to bridge the + # gap from a language syntax perspective. + # + # The corresponding `release_*_raw` functions are needed because while the arrow + # types are pure C structs, we allocate them with new in C++ and need to use delete + # to free them. Unfortunately, unless we lie to Cython and tell it that these types + # are cppclasses, Cython will not allow the usage of the del Python keyword to + # generate the delete call, so inline C++ is again the best option. cdef ArrowSchema *to_arrow_schema_raw( const table_view& tbl, const vector[column_metadata]& metadata, ) except +libcudf_exception_handler nogil + cdef ArrowSchema *to_arrow_schema_raw( + const column_view& tbl, + const column_metadata& metadata, + ) except +libcudf_exception_handler nogil cdef void release_arrow_schema_raw( ArrowSchema * ) except +libcudf_exception_handler nogil cdef ArrowArray* to_arrow_host_raw( const table_view& tbl ) except +libcudf_exception_handler nogil + cdef ArrowArray* to_arrow_host_raw( + const column_view& tbl + ) except +libcudf_exception_handler nogil cdef void release_arrow_array_raw( ArrowArray * ) except +libcudf_exception_handler nogil diff --git a/python/pylibcudf/pylibcudf/null_mask.pxd b/python/pylibcudf/pylibcudf/null_mask.pxd index 779a5aed306..bd6969cc415 100644 --- a/python/pylibcudf/pylibcudf/null_mask.pxd +++ b/python/pylibcudf/pylibcudf/null_mask.pxd @@ -2,6 +2,8 @@ from pylibcudf.libcudf.types cimport mask_state, size_type +from pylibcudf.gpumemoryview cimport gpumemoryview + from rmm.pylibrmm.device_buffer cimport DeviceBuffer from .column cimport Column @@ -17,4 +19,4 @@ cpdef tuple bitmask_and(list columns) cpdef tuple bitmask_or(list columns) -cpdef size_type null_count(Py_ssize_t bitmask, size_type start, size_type stop) +cpdef size_type null_count(gpumemoryview bitmask, size_type start, size_type stop) diff --git a/python/pylibcudf/pylibcudf/null_mask.pyi b/python/pylibcudf/pylibcudf/null_mask.pyi index ace18582bd1..524b3d432b2 100644 --- a/python/pylibcudf/pylibcudf/null_mask.pyi +++ b/python/pylibcudf/pylibcudf/null_mask.pyi @@ -3,6 +3,7 @@ from rmm.pylibrmm.device_buffer import DeviceBuffer from pylibcudf.column import Column +from pylibcudf.gpumemoryview import gpumemoryview from pylibcudf.types import MaskState def copy_bitmask(col: Column) -> DeviceBuffer: ... @@ -12,4 +13,4 @@ def create_null_mask( ) -> DeviceBuffer: ... def bitmask_and(columns: list[Column]) -> tuple[DeviceBuffer, int]: ... def bitmask_or(columns: list[Column]) -> tuple[DeviceBuffer, int]: ... -def null_count(bitmask: int, start: int, stop: int) -> int: ... +def null_count(bitmask: gpumemoryview, start: int, stop: int) -> int: ... diff --git a/python/pylibcudf/pylibcudf/null_mask.pyx b/python/pylibcudf/pylibcudf/null_mask.pyx index 045b7d726dc..8a38563481a 100644 --- a/python/pylibcudf/pylibcudf/null_mask.pyx +++ b/python/pylibcudf/pylibcudf/null_mask.pyx @@ -4,6 +4,7 @@ from libcpp.pair cimport pair from libcpp.utility cimport move from pylibcudf.libcudf cimport null_mask as cpp_null_mask from pylibcudf.libcudf.types cimport mask_state, size_type, bitmask_type +from pylibcudf.gpumemoryview cimport gpumemoryview from rmm.librmm.device_buffer cimport device_buffer from rmm.pylibrmm.device_buffer cimport DeviceBuffer @@ -150,7 +151,7 @@ cpdef tuple bitmask_or(list columns): return buffer_to_python(move(c_result.first)), c_result.second -cpdef size_type null_count(Py_ssize_t bitmask, size_type start, size_type stop): +cpdef size_type null_count(gpumemoryview bitmask, size_type start, size_type stop): """Given a validity bitmask, counts the number of null elements. For details, see :cpp:func:`null_count`. @@ -170,4 +171,4 @@ cpdef size_type null_count(Py_ssize_t bitmask, size_type start, size_type stop): The number of null elements in the specified range. """ with nogil: - return cpp_null_mask.null_count(bitmask, start, stop) + return cpp_null_mask.null_count((bitmask.ptr), start, stop) diff --git a/python/pylibcudf/pylibcudf/table.pyx b/python/pylibcudf/pylibcudf/table.pyx index 19d7b60dcbc..4992a02ab1c 100644 --- a/python/pylibcudf/pylibcudf/table.pyx +++ b/python/pylibcudf/pylibcudf/table.pyx @@ -1,19 +1,64 @@ # Copyright (c) 2023-2025, NVIDIA CORPORATION. from cython.operator cimport dereference -from libcpp.memory cimport unique_ptr + +from cpython.pycapsule cimport ( + PyCapsule_GetPointer, + PyCapsule_New, +) + +from libcpp.memory cimport unique_ptr, make_unique from libcpp.utility cimport move from libcpp.vector cimport vector + from rmm.pylibrmm.stream cimport Stream from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.table.table cimport table +from pylibcudf.libcudf.interop cimport ( + ArrowArray, + ArrowArrayStream, + ArrowSchema, + arrow_table, + column_metadata, + to_arrow_host_raw, + to_arrow_schema_raw, +) + from .column cimport Column from .utils cimport _get_stream +from pylibcudf._interop_helpers cimport ( + _release_schema, + _release_array, + _metadata_to_libcudf, +) +from ._interop_helpers import ColumnMetadata + +from functools import singledispatchmethod __all__ = ["Table"] + +class _ArrowLikeMeta(type): + # We cannot separate stream and array via singledispatch because the + # dispatch will often be ambiguous when objects expose both protocols. + def __subclasscheck__(cls, other): + return ( + hasattr(other, "__arrow_c_stream__") + or hasattr(other, "__arrow_c_array__") + ) + + +class _ArrowLike(metaclass=_ArrowLikeMeta): + pass + + +cdef class _ArrowTableHolder: + """A holder for an Arrow table for gpumemoryview lifetime management.""" + cdef unique_ptr[arrow_table] tbl + + cdef class Table: """A list of columns of the same size. @@ -22,12 +67,41 @@ cdef class Table: columns : list The columns in this table. """ - def __init__(self, list columns): + def __init__(self, obj): + self._init(obj) + + __hash__ = None + + @singledispatchmethod + def _init(self, obj): + raise ValueError(f"Invalid input type {type(obj)}") + + @_init.register(list) + def _(self, list columns): if not all(isinstance(c, Column) for c in columns): raise ValueError("All columns must be pylibcudf Column objects") self._columns = columns - __hash__ = None + @_init.register(_ArrowLike) + def _(self, arrow_like): + cdef ArrowArrayStream* c_stream + cdef _ArrowTableHolder result + cdef unique_ptr[arrow_table] c_result + if hasattr(arrow_like, "__arrow_c_stream__"): + stream = arrow_like.__arrow_c_stream__() + c_stream = ( + PyCapsule_GetPointer(stream, "arrow_array_stream") + ) + + result = _ArrowTableHolder() + with nogil: + c_result = make_unique[arrow_table](move(dereference(c_stream))) + result.tbl.swap(c_result) + + tmp = Table.from_table_view_of_arbitrary(result.tbl.get().view(), result) + self._columns = tmp.columns() + elif hasattr(arrow_like, "__arrow_c_array__"): + raise NotImplementedError("arrays not yet supported") cdef table_view view(self) nogil: """Generate a libcudf table_view to pass to libcudf algorithms. @@ -117,3 +191,38 @@ cdef class Table: cpdef list columns(self): """The columns in this table.""" return self._columns + + def _to_schema(self, metadata=None): + """Create an Arrow schema from this table.""" + if metadata is None: + metadata = [ + col._create_nested_column_metadata() for col in self.columns() + ] + else: + metadata = [ + ColumnMetadata(m) if isinstance(m, str) else m for m in metadata + ] + + cdef vector[column_metadata] c_metadata + c_metadata.reserve(len(metadata)) + for meta in metadata: + c_metadata.push_back(_metadata_to_libcudf(meta)) + + cdef ArrowSchema* raw_schema_ptr + with nogil: + raw_schema_ptr = to_arrow_schema_raw(self.view(), c_metadata) + + return PyCapsule_New(raw_schema_ptr, 'arrow_schema', _release_schema) + + def _to_host_array(self): + cdef ArrowArray* raw_host_array_ptr + with nogil: + raw_host_array_ptr = to_arrow_host_raw(self.view()) + + return PyCapsule_New(raw_host_array_ptr, "arrow_array", _release_array) + + def __arrow_c_array__(self, requested_schema=None): + if requested_schema is not None: + raise ValueError("pylibcudf.Table does not support alternative schema") + + return self._to_schema(), self._to_host_array() diff --git a/python/pylibcudf/pylibcudf/tests/test_labeling.py b/python/pylibcudf/pylibcudf/tests/test_labeling.py index 946d583d1cc..ec9a743a415 100644 --- a/python/pylibcudf/pylibcudf/tests/test_labeling.py +++ b/python/pylibcudf/pylibcudf/tests/test_labeling.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. import pyarrow as pa import pytest @@ -21,7 +21,7 @@ def test_label_bins(left_inclusive, right_inclusive): in_col, left_edges, left_inclusive, right_edges, right_inclusive ) ) - expected = pa.chunked_array([[0, 0, 0]], type=pa.int32()) + expected = pa.array([0, 0, 0], type=pa.int32()) assert result.equals(expected) diff --git a/python/pylibcudf/pylibcudf/tests/test_null_mask.py b/python/pylibcudf/pylibcudf/tests/test_null_mask.py index cd3da856de2..7f4f524d67c 100644 --- a/python/pylibcudf/pylibcudf/tests/test_null_mask.py +++ b/python/pylibcudf/pylibcudf/tests/test_null_mask.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. import pyarrow as pa import pytest @@ -24,11 +24,18 @@ def column(request, nullable): def test_copy_bitmask(column, nullable): - expected = column.null_mask().obj if nullable else rmm.DeviceBuffer() - got = plc.null_mask.copy_bitmask(column) + expected = ( + column.null_mask() + if nullable + else plc.gpumemoryview(rmm.DeviceBuffer()) + ) + got = plc.gpumemoryview(plc.null_mask.copy_bitmask(column)) - assert expected.size == got.size - assert expected.tobytes() == got.tobytes() + start = 0 + end = column.size() * 8 + assert plc.null_mask.null_count( + expected, start, end + ) == plc.null_mask.null_count(got, start, end) def test_bitmask_allocation_size_bytes(): diff --git a/python/pylibcudf/pylibcudf/tests/test_string_extract.py b/python/pylibcudf/pylibcudf/tests/test_string_extract.py index e70edf4fb33..2b395edc242 100644 --- a/python/pylibcudf/pylibcudf/tests/test_string_extract.py +++ b/python/pylibcudf/pylibcudf/tests/test_string_extract.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. import pyarrow as pa import pyarrow.compute as pc @@ -33,7 +33,5 @@ def test_extract_all_record(): ), ) result = plc.interop.to_arrow(plc_result) - expected = pa.chunked_array( - [pa.array([["a", "1"], ["b", "2"], None], type=result.type)] - ) + expected = pa.array([["a", "1"], ["b", "2"], None], type=result.type) assert result.equals(expected) diff --git a/python/pylibcudf/pylibcudf/tests/test_string_padding.py b/python/pylibcudf/pylibcudf/tests/test_string_padding.py index 79498132097..d2c69790bc6 100644 --- a/python/pylibcudf/pylibcudf/tests/test_string_padding.py +++ b/python/pylibcudf/pylibcudf/tests/test_string_padding.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. import pyarrow as pa import pyarrow.compute as pc @@ -15,7 +15,7 @@ def test_pad(): "!", ) result = plc.interop.to_arrow(plc_result) - expected = pa.chunked_array(pc.utf8_lpad(arr, 2, padding="!")) + expected = pa.array(pc.utf8_lpad(arr, 2, padding="!")) assert result.equals(expected) @@ -23,5 +23,5 @@ def test_zfill(): arr = pa.array(["a", "1", None]) plc_result = plc.strings.padding.zfill(plc.interop.from_arrow(arr), 2) result = plc.interop.to_arrow(plc_result) - expected = pa.chunked_array(pc.utf8_lpad(arr, 2, padding="0")) + expected = pa.array(pc.utf8_lpad(arr, 2, padding="0")) assert result.equals(expected) diff --git a/python/pylibcudf/pylibcudf/tests/test_string_repeat.py b/python/pylibcudf/pylibcudf/tests/test_string_repeat.py index c06c06be7c6..1c5f318fabd 100644 --- a/python/pylibcudf/pylibcudf/tests/test_string_repeat.py +++ b/python/pylibcudf/pylibcudf/tests/test_string_repeat.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. import pyarrow as pa import pyarrow.compute as pc @@ -17,5 +17,5 @@ def test_repeat_strings(repeats): else repeats, ) result = plc.interop.to_arrow(plc_result) - expected = pa.chunked_array(pc.binary_repeat(arr, repeats)) + expected = pa.array(pc.binary_repeat(arr, repeats)) assert result.equals(expected) diff --git a/python/pylibcudf/pylibcudf/tests/test_transform.py b/python/pylibcudf/pylibcudf/tests/test_transform.py index a63d94678ff..d6354508d9d 100644 --- a/python/pylibcudf/pylibcudf/tests/test_transform.py +++ b/python/pylibcudf/pylibcudf/tests/test_transform.py @@ -46,7 +46,7 @@ def test_bools_to_mask_roundtrip(): plc_output = plc.transform.mask_to_bools(mask.ptr, 0, len(pa_array)) result_pa = plc.interop.to_arrow(plc_output) - expected_pa = pa.chunked_array([[True, False, False]]) + expected_pa = pa.array([True, False, False]) assert result_pa.equals(expected_pa) @@ -68,7 +68,7 @@ def test_encode(): ) assert pa_table_result.equals(pa_table_expected) - pa_column_expected = pa.chunked_array([[0, 1, 2]], type=pa.int32()) + pa_column_expected = pa.array([0, 1, 2], type=pa.int32()) assert pa_column_result.equals(pa_column_expected)