From 688f83a0b9494751555b9bb0428533745f8556e1 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 2 Jun 2026 00:43:03 +0000 Subject: [PATCH 01/19] feat(cudf_streaming): add Cython/Python package scaffolding --- python/cudf_streaming/CMakeLists.txt | 44 ++++++++ python/cudf_streaming/cudf_streaming/VERSION | 1 + .../cudf_streaming/cudf_streaming/__init__.py | 13 +++ .../integrations/CMakeLists.txt | 6 ++ .../cudf_streaming/integrations/__init__.pxd | 2 + .../cudf_streaming/integrations/__init__.py | 3 + .../cudf_streaming/streaming/CMakeLists.txt | 6 ++ .../cudf_streaming/streaming/__init__.pxd | 2 + .../cudf_streaming/streaming/__init__.py | 3 + python/cudf_streaming/pyproject.toml | 100 ++++++++++++++++++ 10 files changed, 180 insertions(+) create mode 100644 python/cudf_streaming/CMakeLists.txt create mode 100644 python/cudf_streaming/cudf_streaming/VERSION create mode 100644 python/cudf_streaming/cudf_streaming/__init__.py create mode 100644 python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt create mode 100644 python/cudf_streaming/cudf_streaming/integrations/__init__.pxd create mode 100644 python/cudf_streaming/cudf_streaming/integrations/__init__.py create mode 100644 python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt create mode 100644 python/cudf_streaming/cudf_streaming/streaming/__init__.pxd create mode 100644 python/cudf_streaming/cudf_streaming/streaming/__init__.py create mode 100644 python/cudf_streaming/pyproject.toml diff --git a/python/cudf_streaming/CMakeLists.txt b/python/cudf_streaming/CMakeLists.txt new file mode 100644 index 000000000000..793e64b33bdc --- /dev/null +++ b/python/cudf_streaming/CMakeLists.txt @@ -0,0 +1,44 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +cmake_minimum_required(VERSION 4.0 FATAL_ERROR) + +include(../../cmake/rapids_config.cmake) +include(rapids-cpm) +rapids_cpm_init() +include(rapids-cuda) +include(rapids-find) + +rapids_cuda_init_architectures(cudf-streaming-python) + +project( + cudf-streaming-python + VERSION "${RAPIDS_VERSION}" + LANGUAGES CXX CUDA +) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# For now, disable CMake's automatic module scanning for C++ files. There is an sccache bug in the +# version RAPIDS uses in CI that causes it to handle the resulting -M* flags incorrectly with +# gcc>=14. We can remove this once we upgrade to a newer sccache version. +set(CMAKE_CXX_SCAN_FOR_MODULES OFF) + +# ################################################################################################## +# * find C++ libraries from RAPIDS wheels in site-packages ----------------------------------- +set_property(GLOBAL PROPERTY FIND_LIBRARY_USE_LIB64_PATHS TRUE) + +find_package(cudf_streaming "${RAPIDS_VERSION}" REQUIRED) +find_package(rapidsmpf REQUIRED) +find_package(CUDAToolkit REQUIRED) + +include(rapids-cython-core) +rapids_cython_init() + +add_subdirectory(cudf_streaming/integrations) +add_subdirectory(cudf_streaming/streaming) diff --git a/python/cudf_streaming/cudf_streaming/VERSION b/python/cudf_streaming/cudf_streaming/VERSION new file mode 100644 index 000000000000..4e6864b4ca93 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/VERSION @@ -0,0 +1 @@ +26.08.00 diff --git a/python/cudf_streaming/cudf_streaming/__init__.py b/python/cudf_streaming/cudf_streaming/__init__.py new file mode 100644 index 000000000000..56b13692fab0 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/__init__.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""cuDF Streaming library.""" + +# If libcudf_streaming was installed as a wheel, request it to load the library +# symbols. Otherwise, assume the library is on a system path that ld can find. +try: + import libcudf_streaming +except ModuleNotFoundError: + pass +else: + libcudf_streaming.load_library() + del libcudf_streaming diff --git a/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt new file mode 100644 index 000000000000..93535df35ef7 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt @@ -0,0 +1,6 @@ +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on + +# Placeholder for integrations Cython extensions. diff --git a/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd b/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd new file mode 100644 index 000000000000..23a377ecc291 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 diff --git a/python/cudf_streaming/cudf_streaming/integrations/__init__.py b/python/cudf_streaming/cudf_streaming/integrations/__init__.py new file mode 100644 index 000000000000..0cd9cb2345d6 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Collection of cuDF specific functions.""" diff --git a/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt new file mode 100644 index 000000000000..7d8d987a7798 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt @@ -0,0 +1,6 @@ +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on + +# Placeholder for streaming Cython extensions. diff --git a/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd b/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd new file mode 100644 index 000000000000..23a377ecc291 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 diff --git a/python/cudf_streaming/cudf_streaming/streaming/__init__.py b/python/cudf_streaming/cudf_streaming/streaming/__init__.py new file mode 100644 index 000000000000..10f4b88ea5b7 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Submodule for streaming cudf operations.""" diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml new file mode 100644 index 000000000000..c73b4cee7f7f --- /dev/null +++ b/python/cudf_streaming/pyproject.toml @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +build-backend = "rapids_build_backend.build" +requires = [ + "rapids-build-backend>=0.4.0,<0.5.0", + "scikit-build-core[pyproject]>=0.11.0", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[project] +name = "cudf_streaming" +dynamic = ["version"] +description = "cuDF Streaming library" +authors = [ + { name = "NVIDIA Corporation" }, +] +license = "Apache-2.0" +requires-python = ">=3.11" +dependencies = [ + "libcudf_streaming==26.8.*,>=0.0.0a0", + "pylibcudf==26.8.*,>=0.0.0a0", + "rapidsmpf==26.8.*,>=0.0.0a0", + "rmm==26.8.*,>=0.0.0a0", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[project.optional-dependencies] +test = [ + "pytest", + "pytest-cov", + "pytest-xdist", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[project.urls] +Homepage = "https://github.com/rapidsai/cudf" +Documentation = "https://docs.rapids.ai/api/cudf/stable/" + +[tool.ruff] +extend = "../../pyproject.toml" + +[tool.ruff.lint.isort] +combine-as-imports = true +known-first-party = ["cudf_streaming"] +section-order = ["future", "standard-library", "third-party", "dask", "rapids", "first-party", "local-folder"] + +[tool.ruff.lint.isort.sections] +dask = ["dask", "distributed", "dask_cuda", "streamz"] +rapids = ["rmm", "cudf", "dask_cudf", "cudf_streaming", "rapidsmpf"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["E402", "F401"] + +[tool.pydistcheck] +select = [ + "distro-too-large-compressed", +] + +# PyPI hard limit is 1GiB, but try to keep these as small as possible +max_allowed_size_compressed = '75M' + +[tool.pytest.ini_options] +addopts = "--tb=native --strict-config --strict-markers" +empty_parameter_set_mark = "fail_at_collect" +filterwarnings = [ + "error", +] +xfail_strict = true + +[tool.rapids-build-backend] +build-backend = "scikit_build_core.build" +requires-python = ">=3.11" +dependencies-file = "../../dependencies.yaml" +matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true" +requires = [ + "cmake>=4.0", + "cython>=3.2.2", + "libcudf_streaming==26.8.*,>=0.0.0a0", + "librapidsmpf==26.8.*,>=0.0.0a0", + "librmm==26.8.*,>=0.0.0a0", + "ninja", + "pylibcudf==26.8.*,>=0.0.0a0", + "rapidsmpf==26.8.*,>=0.0.0a0", + "rmm==26.8.*,>=0.0.0a0", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[tool.scikit-build] +build-dir = "build/{wheel_tag}" +cmake.build-type = "Release" +cmake.version = "CMakeLists.txt" +minimum-version = "build-system.requires" +ninja.make-fallback = false +sdist.exclude = ["*tests*"] +sdist.reproducible = true +wheel.packages = ["cudf_streaming"] +wheel.exclude = ["*.pyx", "CMakeLists.txt"] + +[tool.scikit-build.metadata.version] +provider = "scikit_build_core.metadata.regex" +input = "cudf_streaming/VERSION" +regex = "(?P.*)" From b3d08e78f07f0b5daa6f3df06748fcd3b42bc68b Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 2 Jun 2026 01:11:48 +0000 Subject: [PATCH 02/19] feat(cudf_streaming): add integrations/ Cython bindings --- .../integrations/CMakeLists.txt | 12 +- .../cudf_streaming/integrations/__init__.pxd | 4 +- .../cudf_streaming/integrations/__init__.py | 6 +- .../cudf_streaming/integrations/partition.pxd | 29 ++ .../cudf_streaming/integrations/partition.pyi | 41 ++ .../cudf_streaming/integrations/partition.pyx | 369 ++++++++++++++++++ 6 files changed, 457 insertions(+), 4 deletions(-) create mode 100644 python/cudf_streaming/cudf_streaming/integrations/partition.pxd create mode 100644 python/cudf_streaming/cudf_streaming/integrations/partition.pyi create mode 100644 python/cudf_streaming/cudf_streaming/integrations/partition.pyx diff --git a/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt index 93535df35ef7..c6c1ba807667 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt +++ b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt @@ -1,6 +1,14 @@ # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on -# Placeholder for integrations Cython extensions. +set(cython_sources partition.pyx) +set(linked_libraries cudf_streaming::cudf_streaming rapidsmpf::rapidsmpf) + +rapids_cython_create_modules( + CXX ASSOCIATED_TARGETS cudf_streaming rapidsmpf + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" + MODULE_PREFIX cudf_streaming_integrations_ +) diff --git a/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd b/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd index 23a377ecc291..c60c8a61b990 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd +++ b/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd @@ -1,2 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 + +from cudf_streaming.integrations.partition cimport pack, unpack_and_concat diff --git a/python/cudf_streaming/cudf_streaming/integrations/__init__.py b/python/cudf_streaming/cudf_streaming/integrations/__init__.py index 0cd9cb2345d6..743c0b3ec454 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/__init__.py +++ b/python/cudf_streaming/cudf_streaming/integrations/__init__.py @@ -1,3 +1,7 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Collection of cuDF specific functions.""" + +from cudf_streaming.integrations.partition import unpack_and_concat + +__all__ = ["unpack_and_concat"] diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pxd b/python/cudf_streaming/cudf_streaming/integrations/partition.pxd new file mode 100644 index 000000000000..774ccb338c10 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pxd @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from pylibcudf.table cimport Table +from rmm.pylibrmm.stream cimport Stream + +from rapidsmpf.memory.buffer_resource cimport BufferResource + + +cpdef object partition_and_pack( + Table table, + object columns_to_hash, + int num_partitions, + Stream stream, + BufferResource br, +) +cpdef object split_and_pack( + Table table, + object splits, + Stream stream, + BufferResource br, +) +cpdef object unpack_and_concat(object partitions, Stream stream, BufferResource br) +cpdef object spill_partitions(object partitions, BufferResource br) +cpdef object unspill_partitions( + object partitions, + BufferResource br, + object allow_overbooking, +) diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pyi b/python/cudf_streaming/cudf_streaming/integrations/partition.pyi new file mode 100644 index 000000000000..1133c65ab1c9 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pyi @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Iterable + +from pylibcudf.table import Table + +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.packed_data import PackedData +from rmm.pylibrmm.stream import Stream + +def partition_and_pack( + table: Table, + columns_to_hash: Iterable[int], + num_partitions: int, + stream: Stream, + br: BufferResource, +) -> dict[int, PackedData]: ... +def split_and_pack( + table: Table, + splits: Iterable[int], + stream: Stream, + br: BufferResource, +) -> dict[int, PackedData]: ... +def unpack_and_concat( + partitions: Iterable[PackedData], + stream: Stream, + br: BufferResource, +) -> Table: ... +def spill_partitions( + partitions: Iterable[PackedData], + *, + br: BufferResource, +) -> list[PackedData]: ... +def unspill_partitions( + partitions: Iterable[PackedData], + *, + br: BufferResource, + allow_overbooking: bool, +) -> list[PackedData]: ... diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pyx b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx new file mode 100644 index 000000000000..bdf5cd67be3c --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Partitioning of cuDF tables.""" + +from cython.operator cimport dereference as deref +from cython.operator cimport postincrement +from libc.stdint cimport uint32_t +from libcpp.memory cimport make_unique, unique_ptr +from libcpp.unordered_map cimport unordered_map +from libcpp.utility cimport move +from libcpp.vector cimport vector +from pylibcudf.libcudf.table.table cimport table as cpp_table +from pylibcudf.libcudf.table.table_view cimport table_view +from pylibcudf.libcudf.types cimport size_type +from pylibcudf.table cimport Table +from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.pylibrmm.stream cimport Stream + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.memory.buffer_resource cimport (AllowOverbooking, + BufferResource, + cpp_BufferResource) +from rapidsmpf.memory.packed_data cimport (PackedData, cpp_PackedData, + packed_data_vector_to_list) + + +cdef extern from "" nogil: + int cpp_HASH_MURMUR3"cudf::hash_id::HASH_MURMUR3" + uint32_t cpp_DEFAULT_HASH_SEED"cudf::DEFAULT_HASH_SEED", + + cdef unordered_map[uint32_t, cpp_PackedData] cpp_partition_and_pack \ + "cudf_streaming::integrations::partition_and_pack"( + const table_view& table, + const vector[size_type] &columns_to_hash, + int num_partitions, + int hash_function, + uint32_t seed, + cuda_stream_view stream, + cpp_BufferResource* br, + ) except +ex_handler + + cdef unordered_map[uint32_t, cpp_PackedData] cpp_split_and_pack \ + "cudf_streaming::integrations::split_and_pack"( + const table_view& table, + const vector[size_type] &splits, + cuda_stream_view stream, + cpp_BufferResource* br, + ) except +ex_handler + + +cpdef object partition_and_pack( + Table table, + object columns_to_hash, + int num_partitions, + Stream stream, + BufferResource br, +): + """ + Partition rows from the input table into multiple packed tables. + + Parameters + ---------- + table + The input table to partition. + columns_to_hash + Indices of the input columns to use for hashing. + num_partitions + The number of partitions to create. + stream + The CUDA stream used for memory operations. + br + Buffer resource for memory allocations. + + Returns + ------- + A dictionary where the keys are partition IDs and the values are packed tables. + + Raises + ------ + IndexError + If any index in ``columns_to_hash`` is invalid. + + See Also + -------- + cudf_streaming.integrations.partition.unpack_and_concat + pylibcudf.partitioning.hash_partition + pylibcudf.contiguous_split.pack + cudf_streaming.integrations.partition.split_and_pack + """ + cdef cuda_stream_view _stream = stream.view() + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[size_type] _columns_to_hash = tuple(columns_to_hash) + cdef unordered_map[uint32_t, cpp_PackedData] _ret + cdef table_view tbl = table.view() + with nogil: + _ret = cpp_partition_and_pack( + tbl, + _columns_to_hash, + num_partitions, + cpp_HASH_MURMUR3, + cpp_DEFAULT_HASH_SEED, + _stream, + _br, + ) + ret = {} + cdef unordered_map[uint32_t, cpp_PackedData].iterator it = _ret.begin() + while(it != _ret.end()): + ret[deref(it).first] = PackedData.from_librapidsmpf( + make_unique[cpp_PackedData](move(deref(it).second)), + br, + ) + postincrement(it) + return ret + + +cpdef object split_and_pack( + Table table, + object splits, + Stream stream, + BufferResource br, +): + """ + Split rows from the input table into multiple packed tables. + + Parameters + ---------- + table + The input table to split and pack. The table cannot be empty (the + split points would not be valid). + splits + The split points, one less than the number of result partitions. + stream + The CUDA stream used for memory operations. + br + Buffer resource for memory allocations. + + Returns + ------- + A map of partition IDs and their packed tables. + + Raises + ------ + IndexError + If the splits are out of range for ``[0, len(table)]``. + + See Also + -------- + cudf_streaming.integrations.partition.unpack_and_concat + pylibcudf.copying.split + cudf_streaming.integrations.partition.partition_and_pack + """ + cdef cuda_stream_view _stream = stream.view() + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[size_type] _splits = tuple(splits) + cdef unordered_map[uint32_t, cpp_PackedData] _ret + cdef table_view tbl = table.view() + with nogil: + _ret = cpp_split_and_pack( + tbl, + _splits, + _stream, + _br, + ) + ret = {} + cdef unordered_map[uint32_t, cpp_PackedData].iterator it = _ret.begin() + while(it != _ret.end()): + ret[deref(it).first] = PackedData.from_librapidsmpf( + make_unique[cpp_PackedData](move(deref(it).second)), + br, + ) + postincrement(it) + return ret + + +cdef object pack( + Table table, + Stream stream, + BufferResource br, +): + """Pack a table into a single ``PackedData`` instance.""" + packed = split_and_pack(table, (), stream, br) + return packed[0] + + +cdef extern from "" nogil: + cdef unique_ptr[cpp_table] cpp_unpack_and_concat \ + "cudf_streaming::integrations::unpack_and_concat"( + vector[cpp_PackedData] partition, + cuda_stream_view stream, + cpp_BufferResource* br, + ) except +ex_handler + + +cdef vector[cpp_PackedData] _partitions_py_to_cpp(partitions): + cdef vector[cpp_PackedData] ret + for part in partitions: + if not (part).c_obj: + raise ValueError("PackedData was empty") + ret.push_back(move(deref((part).c_obj))) + return move(ret) + + +cpdef object unpack_and_concat( + object partitions, + Stream stream, + BufferResource br, +): + """ + Unpack input partitions and concatenate them into a single table. + + Empty partitions are ignored. + + The unpacking of each partition is stream-ordered on that partition's own CUDA + stream. The returned table is stream-ordered on the provided ``stream`` and + synchronized with the unpacking. + + Notes + ----- + The input partitions are released and left empty on return. + + Parameters + ---------- + partitions + Packed input tables (partitions). + stream + CUDA stream on which concatenation occurs and on which the resulting + table is ordered. + br + Buffer resource used for memory allocations. + + Returns + ------- + The concatenated table resulting from unpacking the input partitions. + + Raises + ------ + ReservationError + If the buffer resource cannot reserve enough memory to concatenate all + partitions. + + See Also + -------- + cudf_streaming.integrations.partition.partition_and_pack + """ + cdef cuda_stream_view _stream = stream.view() + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) + cdef unique_ptr[cpp_table] _ret + with nogil: + _ret = cpp_unpack_and_concat( + move(_partitions), + _stream, + _br, + ) + return Table.from_libcudf(move(_ret), stream, br._device_mr) + + +cdef extern from "" nogil: + cdef vector[cpp_PackedData] cpp_spill_partitions \ + "cudf_streaming::integrations::spill_partitions"( + vector[cpp_PackedData] partitions, + cpp_BufferResource* br, + ) except +ex_handler + + +cpdef object spill_partitions( + object partitions, + BufferResource br, +): + """ + Spill partitions from device memory to host memory. + + Moves the buffer of each ``PackedData`` from device memory to host memory using + the provided buffer resource and the buffer's CUDA stream. Partitions already + in host memory are returned unchanged. + + For device-resident partitions, a host memory reservation is made before moving + the buffer. If the reservation fails due to insufficient host memory, an + exception is raised. Overbooking is not allowed. + + The input partitions are released and are left empty on return. + + Parameters + ---------- + partitions + The partitions to spill. + br + Buffer resource used to reserve host memory and perform the move. + + Returns + ------- + A list of partitions whose buffers reside in host memory. + + Raises + ------ + ReservationError + If host memory reservation fails. + """ + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) + cdef vector[cpp_PackedData] _ret + with nogil: + _ret = cpp_spill_partitions( + move(_partitions), + _br, + ) + return packed_data_vector_to_list(move(_ret), br) + + +cdef extern from "" nogil: + cdef vector[cpp_PackedData] cpp_unspill_partitions \ + "cudf_streaming::integrations::unspill_partitions"( + vector[cpp_PackedData] partitions, + cpp_BufferResource* br, + AllowOverbooking allow_overbooking, + ) except +ex_handler + + +cpdef object unspill_partitions( + object partitions, + BufferResource br, + object allow_overbooking, +): + """ + Move spilled partitions back to device memory. + + Each partition is inspected to determine whether its buffer resides in device + memory. Buffers already in device memory are left untouched. Host-resident buffers + are moved to device memory using the provided buffer resource and the buffer's CUDA + stream. + + If insufficient device memory is available, the buffer resource's spill manager is + invoked to free memory. If overbooking occurs and spilling fails to reclaim enough + memory, behavior depends on ``allow_overbooking``. + + The input partitions are released and are left empty on return. + + Parameters + ---------- + partitions + The partitions to unspill, potentially containing host-resident data. + br + Buffer resource responsible for memory reservation and spills. + allow_overbooking + If False, ensures enough memory is freed to satisfy the reservation; + otherwise, allows overbooking even if spilling was insufficient. + + Returns + ------- + A list of partitions whose buffers reside in device memory. + + Raises + ------ + ReservationError + If overbooking exceeds the amount spilled and ``allow_overbooking is False``. + """ + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) + cdef vector[cpp_PackedData] _ret + cdef AllowOverbooking ab = ( + AllowOverbooking.YES if allow_overbooking else AllowOverbooking.NO + ) + with nogil: + _ret = cpp_unspill_partitions( + move(_partitions), + _br, + ab, + ) + return packed_data_vector_to_list(move(_ret), br) From 4bb32e507d7061edff4a30f749f321eabe9cba1e Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 2 Jun 2026 01:13:20 +0000 Subject: [PATCH 03/19] feat(cudf_streaming): add streaming/ Cython bindings --- .../cudf_streaming/streaming/CMakeLists.txt | 12 +- .../cudf_streaming/streaming/__init__.pxd | 18 +- .../cudf_streaming/streaming/__init__.py | 35 +- .../cudf_streaming/streaming/bloom_filter.pxd | 46 ++ .../cudf_streaming/streaming/bloom_filter.pyi | 49 ++ .../cudf_streaming/streaming/bloom_filter.pyx | 247 +++++++ .../streaming/channel_metadata.pxd | 131 ++++ .../streaming/channel_metadata.pyi | 94 +++ .../streaming/channel_metadata.pyx | 405 +++++++++++ .../cudf_streaming/streaming/parquet.pyi | 25 + .../cudf_streaming/streaming/parquet.pyx | 145 ++++ .../cudf_streaming/streaming/partition.pyi | 28 + .../cudf_streaming/streaming/partition.pyx | 150 ++++ .../cudf_streaming/streaming/table_chunk.pxd | 46 ++ .../cudf_streaming/streaming/table_chunk.pyi | 72 ++ .../cudf_streaming/streaming/table_chunk.pyx | 664 ++++++++++++++++++ 16 files changed, 2163 insertions(+), 4 deletions(-) create mode 100644 python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd create mode 100644 python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyi create mode 100644 python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx create mode 100644 python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd create mode 100644 python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyi create mode 100644 python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyx create mode 100644 python/cudf_streaming/cudf_streaming/streaming/parquet.pyi create mode 100644 python/cudf_streaming/cudf_streaming/streaming/parquet.pyx create mode 100644 python/cudf_streaming/cudf_streaming/streaming/partition.pyi create mode 100644 python/cudf_streaming/cudf_streaming/streaming/partition.pyx create mode 100644 python/cudf_streaming/cudf_streaming/streaming/table_chunk.pxd create mode 100644 python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyi create mode 100644 python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyx diff --git a/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt index 7d8d987a7798..1428c8f3305e 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt +++ b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt @@ -1,6 +1,14 @@ # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on -# Placeholder for streaming Cython extensions. +set(cython_sources bloom_filter.pyx channel_metadata.pyx parquet.pyx partition.pyx table_chunk.pyx) +set(linked_libraries cudf_streaming::cudf_streaming rapidsmpf::rapidsmpf) + +rapids_cython_create_modules( + CXX ASSOCIATED_TARGETS cudf_streaming + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" + MODULE_PREFIX cudf_streaming_streaming_ +) diff --git a/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd b/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd index 23a377ecc291..2c2cb7e1959f 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd +++ b/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd @@ -1,2 +1,18 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 + +from cudf_streaming.streaming.bloom_filter cimport BloomFilter, cpp_BloomFilter +from cudf_streaming.streaming.channel_metadata cimport ( + ChannelMetadata, + HashScheme, + OrderKey, + OrderScheme, + Partitioning, + cpp_ChannelMetadata, + cpp_HashScheme, + cpp_OrderKey, + cpp_OrderScheme, + cpp_Partitioning, + cpp_PartitioningSpec, +) +from cudf_streaming.streaming.table_chunk cimport TableChunk, cpp_TableChunk diff --git a/python/cudf_streaming/cudf_streaming/streaming/__init__.py b/python/cudf_streaming/cudf_streaming/streaming/__init__.py index 10f4b88ea5b7..4f5d020b6fac 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/__init__.py +++ b/python/cudf_streaming/cudf_streaming/streaming/__init__.py @@ -1,3 +1,36 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Submodule for streaming cudf operations.""" + +from cudf_streaming.streaming.bloom_filter import BloomFilter +from cudf_streaming.streaming.channel_metadata import ( + ChannelMetadata, + HashScheme, + OrderKey, + OrderScheme, + Partitioning, +) +from cudf_streaming.streaming.parquet import Filter, read_parquet +from cudf_streaming.streaming.partition import ( + partition_and_pack, + unpack_and_concat, +) +from cudf_streaming.streaming.table_chunk import ( + TableChunk, + make_table_chunks_available_or_wait, +) + +__all__ = [ + "BloomFilter", + "ChannelMetadata", + "Filter", + "HashScheme", + "OrderKey", + "OrderScheme", + "Partitioning", + "TableChunk", + "make_table_chunks_available_or_wait", + "partition_and_pack", + "read_parquet", + "unpack_and_concat", +] diff --git a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd new file mode 100644 index 000000000000..3e4e09654828 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stddef cimport size_t +from libc.stdint cimport uint64_t +from libcpp.memory cimport shared_ptr, unique_ptr + +from rapidsmpf.communicator.communicator cimport Communicator, cpp_Communicator +from rapidsmpf.streaming.core.context cimport cpp_Context + + +cdef extern from "" nogil: + cdef cppclass cpp_BloomFilter "cudf_streaming::streaming::BloomFilter": + cpp_BloomFilter( + shared_ptr[cpp_Context] ctx, + shared_ptr[cpp_Communicator] comm, + uint64_t seed, + size_t num_filter_blocks, + ) noexcept + const shared_ptr[cpp_Communicator]& comm() noexcept + + +cdef extern from "" nogil: + size_t cpp_fitting_num_blocks \ + "cudf_streaming::integrations::BloomFilter::fitting_num_blocks"( + size_t l2size + ) noexcept + + +cdef class BloomFilter: + """ + Streaming bloom filter construction and application. + + Parameters + ---------- + ctx + Streaming context. + comm + The communicator the bloom filter construction is collective over. + seed + Seed used for hashing values into the bloom filter. + num_filter_blocks + Number of blocks used to size the filter. + """ + cdef unique_ptr[cpp_BloomFilter] _handle + cdef Communicator _comm diff --git a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyi b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyi new file mode 100644 index 000000000000..15148005007a --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyi @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Iterable +from typing import Self + +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.communicator.communicator import Communicator +from rapidsmpf.streaming.core.channel import Channel +from rapidsmpf.streaming.core.context import Context +from rapidsmpf.streaming.core.message import Message + +class BloomFilterChunk: + # Note: if you go looking for this type in the cython bindings, you + # won't find it. This is purely to provide for better type-checking of + # the generic Channel argument to BloomFilter.build/apply below. + @classmethod + def from_message(cls: type[Self], message: Message[Self]) -> Self: ... + def into_message( + self, sequence_number: int, message: Message[Self] + ) -> None: ... + +class BloomFilter: + def __init__( + self, + ctx: Context, + comm: Communicator, + seed: int, + num_filter_blocks: int, + ) -> None: ... + @property + def comm(self) -> Communicator: ... + @staticmethod + def fitting_num_blocks(l2size: int) -> int: ... + async def build( + self, + ctx: Context, + ch_in: Channel[TableChunk], + ch_out: Channel[BloomFilterChunk], + tag: int, + ) -> None: ... + async def apply( + self, + ctx: Context, + bloom_filter: Channel[BloomFilterChunk], + ch_in: Channel[TableChunk], + ch_out: Channel[TableChunk], + keys: Iterable[int], + ) -> None: ... diff --git a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx new file mode 100644 index 000000000000..b43613f9d83d --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from cpython.object cimport PyObject +from cpython.ref cimport Py_INCREF +from cython.operator cimport dereference as deref +from libc.stddef cimport size_t +from libc.stdint cimport int32_t, uint64_t +from libcpp.memory cimport make_unique, shared_ptr +from libcpp.utility cimport move +from libcpp.vector cimport vector +from pylibcudf.libcudf.types cimport size_type + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.communicator.communicator cimport Communicator +from rapidsmpf.owning_wrapper cimport cpp_OwningWrapper +from rapidsmpf.streaming._detail.libcoro_spawn_task cimport cpp_set_py_future +from rapidsmpf.streaming.chunks.utils cimport py_deleter +from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel +from rapidsmpf.streaming.core.context cimport Context, cpp_Context + +import asyncio + + +cdef extern from * nogil: + """ + namespace { + void cpp_bloom_filter_build( + std::shared_ptr ctx, + cudf_streaming::streaming::BloomFilter& bloom_filter, + std::shared_ptr ch_in, + std::shared_ptr ch_out, + int32_t tag, + void (*cpp_set_py_future)(void*, const char *), + rapidsmpf::OwningWrapper py_future + ) { + RAPIDSMPF_EXPECTS( + ctx->executor()->spawn_detached( + cython_libcoro_task_wrapper( + cpp_set_py_future, + std::move(py_future), + bloom_filter.build( + std::move(ch_in), + std::move(ch_out), + tag + ) + ) + ), + "libcoro's spawn_detached() failed to spawn task" + ); + } + } // namespace + """ + void cpp_bloom_filter_build( + shared_ptr[cpp_Context] ctx, + cpp_BloomFilter& bloom_filter, + shared_ptr[cpp_Channel] ch_in, + shared_ptr[cpp_Channel] ch_out, + int32_t tag, + void (*cpp_set_py_future)(void*, const char *), + cpp_OwningWrapper py_future + ) except +ex_handler + + +cdef extern from * nogil: + """ + namespace { + void cpp_bloom_filter_apply( + std::shared_ptr ctx, + cudf_streaming::streaming::BloomFilter& bloom_filter, + std::shared_ptr bloom_filter_ch, + std::shared_ptr ch_in, + std::shared_ptr ch_out, + std::vector keys, + void (*cpp_set_py_future)(void*, const char *), + rapidsmpf::OwningWrapper py_future + ) { + RAPIDSMPF_EXPECTS( + ctx->executor()->spawn_detached( + cython_libcoro_task_wrapper( + cpp_set_py_future, + std::move(py_future), + bloom_filter.apply( + std::move(bloom_filter_ch), + std::move(ch_in), + std::move(ch_out), + std::move(keys) + ) + ) + ), + "libcoro's spawn_detached() failed to spawn task" + ); + } + } // namespace + """ + void cpp_bloom_filter_apply( + shared_ptr[cpp_Context] ctx, + cpp_BloomFilter& bloom_filter, + shared_ptr[cpp_Channel] bloom_filter_ch, + shared_ptr[cpp_Channel] ch_in, + shared_ptr[cpp_Channel] ch_out, + vector[size_type] keys, + void (*cpp_set_py_future)(void*, const char *), + cpp_OwningWrapper py_future + ) except +ex_handler + + +cdef class BloomFilter: + """ + Streaming bloom filter construction and application. + + Parameters + ---------- + ctx + Streaming context. + comm + The communicator the bloom filter construction is collective over. + seed + Seed used for hashing values into the bloom filter. + num_filter_blocks + Number of blocks used to size the filter. + """ + + def __init__( + self, + Context ctx not None, + Communicator comm not None, + uint64_t seed, + size_t num_filter_blocks, + ): + self._comm = comm + with nogil: + self._handle = make_unique[cpp_BloomFilter]( + ctx._handle, + comm._handle, + seed, + num_filter_blocks, + ) + + def __dealloc__(self): + with nogil: + self._handle.reset() + + @property + def comm(self): + """ + Get the communicator used by the bloom filter. + + Returns + ------- + The communicator. + """ + return self._comm + + @staticmethod + def fitting_num_blocks(size_t l2size): + """ + Return the number of blocks needed to fit within an L2 cache size. + + Parameters + ---------- + l2size + Size of the L2 cache in bytes. + + Returns + ------- + Number of blocks to use in the filter. + """ + cdef size_t ret + with nogil: + ret = cpp_fitting_num_blocks(l2size) + return ret + + async def build( + self, + Context ctx not None, + Channel ch_in not None, + Channel ch_out not None, + int32_t tag, + ): + """ + Build a bloom filter from input table chunks. + + Parameters + ---------- + ctx + The current streaming context. + ch_in + Input channel of ``TableChunk`` objects. + ch_out + Output channel receiving a single bloom filter message. + tag + Disambiguating tag to combine filters across ranks. + """ + ret = asyncio.get_running_loop().create_future() + Py_INCREF(ret) + with nogil: + cpp_bloom_filter_build( + ctx._handle, + deref(self._handle), + ch_in._handle, + ch_out._handle, + tag, + cpp_set_py_future, + move(cpp_OwningWrapper(ret, py_deleter)), + ) + await ret + + async def apply( + self, + Context ctx not None, + Channel bloom_filter not None, + Channel ch_in not None, + Channel ch_out not None, + keys, + ): + """ + Apply a bloom filter to incoming table chunks. + + Parameters + ---------- + ctx + The current streaming context. + bloom_filter + Channel containing the bloom filter (a single message). + ch_in + Input channel of ``TableChunk`` objects to filter. + ch_out + Output channel receiving filtered ``TableChunk`` objects. + keys + Indices selecting the key columns for hash fingerprints. + """ + cdef vector[size_type] c_keys = tuple(keys) + ret = asyncio.get_running_loop().create_future() + Py_INCREF(ret) + with nogil: + cpp_bloom_filter_apply( + ctx._handle, + deref(self._handle), + bloom_filter._handle, + ch_in._handle, + ch_out._handle, + move(c_keys), + cpp_set_py_future, + move(cpp_OwningWrapper(ret, py_deleter)), + ) + await ret diff --git a/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd new file mode 100644 index 000000000000..f5118f8bbc34 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stdint cimport int32_t, uint64_t +from libcpp cimport bool as bool_t +from libcpp.memory cimport shared_ptr, unique_ptr +from libcpp.optional cimport optional +from libcpp.utility cimport pair +from libcpp.vector cimport vector +from pylibcudf.libcudf.types cimport null_order as cpp_null_order +from pylibcudf.libcudf.types cimport order as cpp_order +from rmm.librmm.cuda_stream_view cimport cuda_stream_view + +from rapidsmpf.memory.buffer_resource cimport cpp_BufferResource +from rapidsmpf.streaming.core.message cimport cpp_Message +from cudf_streaming.streaming.table_chunk cimport TableChunk, cpp_TableChunk + + +cdef extern from "" \ + namespace "cudf_streaming::streaming" nogil: + + cdef cppclass cpp_HashScheme "cudf_streaming::streaming::HashScheme": + vector[int32_t] column_indices + int modulus + cpp_HashScheme() except + + cpp_HashScheme(vector[int32_t], int) except + + bool_t operator==(const cpp_HashScheme&) + + cdef cppclass cpp_OrderKey "cudf_streaming::streaming::OrderKey": + cpp_OrderKey() noexcept + cpp_OrderKey(int32_t, cpp_order, cpp_null_order) noexcept + int32_t column_index + cpp_order order + cpp_null_order null_order + bool_t operator==(const cpp_OrderKey&) noexcept + + cdef cppclass cpp_OrderScheme "cudf_streaming::streaming::OrderScheme": + cpp_OrderScheme() noexcept + cpp_OrderScheme( + vector[cpp_OrderKey], unique_ptr[cpp_TableChunk], bool_t + ) except + + vector[cpp_OrderKey] keys + shared_ptr[cpp_TableChunk] boundaries + bool_t strict_boundaries + cpp_OrderScheme with_keys(vector[cpp_OrderKey]) except + + bool_t boundaries_aligned_with( + const cpp_OrderScheme&, const cpp_BufferResource& + ) except + + + cdef cppclass cpp_PartitioningSpec "cudf_streaming::streaming::PartitioningSpec": + enum cpp_Type "cudf_streaming::streaming::PartitioningSpec::Type": + NONE "cudf_streaming::streaming::PartitioningSpec::Type::NONE" + INHERIT "cudf_streaming::streaming::PartitioningSpec::Type::INHERIT" + HASH "cudf_streaming::streaming::PartitioningSpec::Type::HASH" + ORDER "cudf_streaming::streaming::PartitioningSpec::Type::ORDER" + + cpp_Type type + optional[cpp_HashScheme] hash + optional[cpp_OrderScheme] order + + @staticmethod + cpp_PartitioningSpec none() + + @staticmethod + cpp_PartitioningSpec inherit() + + @staticmethod + cpp_PartitioningSpec from_hash(cpp_HashScheme) + + @staticmethod + cpp_PartitioningSpec from_order(cpp_OrderScheme) + + cdef cppclass cpp_Partitioning "cudf_streaming::streaming::Partitioning": + cpp_PartitioningSpec inter_rank + cpp_PartitioningSpec local + cpp_Partitioning() except + + cpp_Partitioning(const cpp_Partitioning&) except + + + cdef cppclass cpp_ChannelMetadata "cudf_streaming::streaming::ChannelMetadata": + uint64_t local_count + cpp_Partitioning partitioning + bool_t duplicated + cpp_ChannelMetadata( + uint64_t, + cpp_Partitioning, + bool_t + ) except + + + cpp_Message cpp_to_message_channel_metadata \ + "cudf_streaming::streaming::to_message"( + uint64_t, unique_ptr[cpp_ChannelMetadata] + ) except + + + +cdef class HashScheme: + cdef cpp_HashScheme _handle + + @staticmethod + cdef HashScheme from_cpp(cpp_HashScheme scheme) + + +cdef class OrderKey: + cdef cpp_OrderKey _handle + + @staticmethod + cdef OrderKey from_cpp(cpp_OrderKey key) + + +cdef class OrderScheme: + cdef cpp_OrderScheme _handle + + @staticmethod + cdef OrderScheme from_cpp(cpp_OrderScheme scheme) + + +cdef class Partitioning: + cdef cpp_Partitioning _handle + + @staticmethod + cdef Partitioning from_cpp(cpp_Partitioning data) + + +cdef class ChannelMetadata: + cdef unique_ptr[cpp_ChannelMetadata] _handle + + @staticmethod + cdef ChannelMetadata from_handle(unique_ptr[cpp_ChannelMetadata] handle) + + cdef const cpp_ChannelMetadata* handle_ptr(self) except NULL + + cdef unique_ptr[cpp_ChannelMetadata] release_handle(self) diff --git a/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyi b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyi new file mode 100644 index 000000000000..02b039af7ab1 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyi @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Type stubs for channel_metadata module.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Literal, Self + +import pylibcudf as plc + +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.streaming.core.message import Message + +class HashScheme: + def __init__( + self, column_indices: Sequence[int], modulus: int + ) -> None: ... + @property + def column_indices(self) -> tuple[int, ...]: ... + @property + def modulus(self) -> int: ... + def __eq__(self, other: object) -> bool: ... + def __repr__(self) -> str: ... + +@dataclass(frozen=True, slots=True) +class OrderKey: + """Sort key: column index, direction, and null ordering.""" + + column_index: int + order: plc.types.Order + null_order: plc.types.NullOrder + +class OrderScheme: + def __init__( + self, + keys: Sequence[OrderKey], + boundaries: TableChunk, + *, + strict_boundaries: bool = False, + ) -> None: ... + @property + def keys(self) -> tuple[OrderKey, ...]: ... + @property + def strict_boundaries(self) -> bool: ... + @property + def num_boundaries(self) -> int: ... + def get_boundaries(self, br: BufferResource) -> TableChunk: ... + def with_keys(self, new_keys: Sequence[OrderKey]) -> OrderScheme: ... + def boundaries_aligned_with( + self, other: OrderScheme, br: BufferResource + ) -> bool: ... + def __repr__(self) -> str: ... + +PartitioningSpecValue = HashScheme | OrderScheme | None | Literal["inherit"] + +class Partitioning: + def __init__( + self, + inter_rank: PartitioningSpecValue = None, + local: PartitioningSpecValue = None, + ) -> None: ... + @property + def inter_rank(self) -> PartitioningSpecValue: ... + @property + def local(self) -> PartitioningSpecValue: ... + def __eq__(self, other: object) -> bool: ... + def __repr__(self) -> str: ... + +class ChannelMetadata: + def __init__( + self, + local_count: int, + *, + partitioning: Partitioning | None = None, + duplicated: bool = False, + ) -> None: ... + @classmethod + def from_message( + cls: type[Self], message: Message[Self] + ) -> ChannelMetadata: ... + def into_message( + self, sequence_number: int, message: Message[Self] + ) -> None: ... + @property + def local_count(self) -> int: ... + @property + def partitioning(self) -> Partitioning: ... + @property + def duplicated(self) -> bool: ... + def __eq__(self, other: object) -> bool: ... + def __repr__(self) -> str: ... diff --git a/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyx b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyx new file mode 100644 index 000000000000..2dd0d3ad741f --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pyx @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Channel metadata types for streaming pipelines.""" + +from cython.operator cimport dereference as deref +from libc.stdint cimport int32_t, uint64_t +from libcpp.memory cimport make_unique, unique_ptr +from libcpp.utility cimport move +from libcpp.vector cimport vector +from pylibcudf.libcudf.types cimport null_order as cpp_null_order +from pylibcudf.libcudf.types cimport order as cpp_order +from pylibcudf.table cimport Table +from rmm.pylibrmm.stream cimport Stream + +from rapidsmpf.memory.buffer_resource cimport BufferResource +from rapidsmpf.streaming.core.message cimport Message +from cudf_streaming.streaming.table_chunk cimport TableChunk, cpp_TableChunk + + +cdef extern from * nogil: + """ + #include + #include + + static std::unique_ptr + cpp_channel_metadata_from_message(rapidsmpf::streaming::Message msg) { + return std::make_unique( + msg.release() + ); + } + """ + unique_ptr[cpp_ChannelMetadata] cpp_channel_metadata_from_message( + cpp_Message + ) except + + + +cdef class HashScheme: + """Hash partitioning scheme: rows distributed by hash(column_indices) % modulus.""" + + def __init__(self, object column_indices, int modulus): + """ + Parameters + ---------- + column_indices + Column indices to hash on. + modulus + Number of hash partitions. + """ + cdef vector[int32_t] cols + for c in column_indices: + cols.push_back(c) + self._handle = cpp_HashScheme(cols, modulus) + + @staticmethod + cdef HashScheme from_cpp(cpp_HashScheme scheme): + cdef HashScheme ret = HashScheme.__new__(HashScheme) + ret._handle = move(scheme) + return ret + + @property + def column_indices(self) -> tuple: + """Column indices used for hashing.""" + return tuple(self._handle.column_indices) + + @property + def modulus(self) -> int: + """Number of hash partitions.""" + return self._handle.modulus + + def __eq__(self, other): + if not isinstance(other, HashScheme): + return NotImplemented + return self._handle == (other)._handle + + def __repr__(self): + return f"HashScheme({self.column_indices!r}, {self.modulus})" + + +cdef class OrderKey: + """A single sort key: column index, sort direction, and null placement.""" + + def __init__( + self, + int column_index, + cpp_order order, + cpp_null_order null_order, + ): + """ + Parameters + ---------- + column_index + Zero-based index of the sort column. + order + Sort direction (ascending or descending). + null_order + Null placement (before or after non-null values). + """ + self._handle = cpp_OrderKey(column_index, order, null_order) + + @staticmethod + cdef OrderKey from_cpp(cpp_OrderKey key): + cdef OrderKey ret = OrderKey.__new__(OrderKey) + ret._handle = key + return ret + + @property + def column_index(self) -> int: + """Zero-based index of the sort column.""" + return self._handle.column_index + + @property + def order(self): + """Sort direction (ascending or descending).""" + return self._handle.order + + @property + def null_order(self): + """Null placement (before or after non-null values).""" + return self._handle.null_order + + def __eq__(self, other): + if not isinstance(other, OrderKey): + return NotImplemented + return self._handle == (other)._handle + + def __repr__(self): + return f"OrderKey({self.column_index}, {self.order!r}, {self.null_order!r})" + + +cdef class OrderScheme: + """Order-based partitioning scheme for sorted/range-partitioned data. + + Data is partitioned by value ranges based on predetermined boundaries. + For N partitions, there are N-1 boundary rows. + + Parameters + ---------- + keys + Sequence of ``OrderKey`` objects (one per sort column). + boundaries + Optional ``TableChunk`` of N-1 boundary rows for N partitions. + strict_boundaries + When true, every row in a chunk falls in a single partition's half-open key + range (keys do not straddle chunk interiors). See the C++ ``OrderScheme`` docs. + Default false. + """ + + def __init__( + self, + object keys, + TableChunk boundaries not None, + *, + bint strict_boundaries = False, + ): + cdef vector[cpp_OrderKey] cpp_keys + for key in keys: + cpp_keys.push_back((key)._handle) + if cpp_keys.empty(): + raise ValueError("OrderScheme: keys must not be empty") + self._handle = cpp_OrderScheme( + move(cpp_keys), move(boundaries.release_handle()), strict_boundaries + ) + + @staticmethod + cdef OrderScheme from_cpp(cpp_OrderScheme scheme): + cdef OrderScheme ret = OrderScheme.__new__(OrderScheme) + ret._handle = move(scheme) + return ret + + @property + def keys(self) -> tuple: + """Sort keys, one per sort column.""" + cdef int i + cdef int n = self._handle.keys.size() + return tuple(OrderKey.from_cpp(self._handle.keys[i]) for i in range(n)) + + @property + def strict_boundaries(self) -> bool: + """Same semantics as the C++ ``OrderScheme::strict_boundaries`` field.""" + return self._handle.strict_boundaries + + @property + def num_boundaries(self) -> int: + """Number of boundary rows (N-1 for N partitions).""" + return self._handle.boundaries.get().shape().first + + def get_boundaries(self, BufferResource br not None) -> TableChunk: + """ + Return the boundary rows. + + Parameters + ---------- + br + Buffer resource to associate with the returned table chunk. + + Returns + ------- + TableChunk + A non-exclusive view of the boundary rows owned by this scheme. + """ + cdef const cpp_TableChunk* chunk = self._handle.boundaries.get() + cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().value()) + tbl = Table.from_table_view_of_arbitrary( + chunk.table_view(), owner=self, stream=stream + ) + return TableChunk.from_pylibcudf_table( + tbl, stream, exclusive_view=False, br=br + ) + + def with_keys(self, object new_keys) -> OrderScheme: + """Return a new ``OrderScheme`` with updated key column indices.""" + cdef vector[cpp_OrderKey] cpp_keys + for key in new_keys: + cpp_keys.push_back((key)._handle) + return OrderScheme.from_cpp(self._handle.with_keys(move(cpp_keys))) + + def boundaries_aligned_with( + self, OrderScheme other not None, BufferResource br not None + ) -> bool: + """ + Check whether boundary values are aligned with another scheme. + + Parameters + ---------- + other + The scheme to compare against. + br + Buffer resource for temporary allocations during comparison. + """ + return self._handle.boundaries_aligned_with(other._handle, deref(br.ptr())) + + def __repr__(self): + return ( + f"OrderScheme({self.keys!r}, " + f"strict_boundaries={self.strict_boundaries})" + ) + + +cdef void _apply_spec(cpp_PartitioningSpec& spec, obj) except *: + """Set *spec* in-place from a Python value.""" + if obj is None: + spec = cpp_PartitioningSpec.none() + elif obj == "inherit": + spec = cpp_PartitioningSpec.inherit() + elif isinstance(obj, HashScheme): + spec = cpp_PartitioningSpec.from_hash((obj)._handle) + elif isinstance(obj, OrderScheme): + spec = cpp_PartitioningSpec.from_order((obj)._handle) + else: + raise TypeError( + f"Expected HashScheme, OrderScheme, None, or 'inherit', " + f"got {type(obj).__name__}" + ) + + +cdef object _from_spec(const cpp_PartitioningSpec& spec): + """Convert PartitioningSpec (by reference) to a Python object.""" + if spec.type == cpp_PartitioningSpec.cpp_Type.NONE: + return None + elif spec.type == cpp_PartitioningSpec.cpp_Type.INHERIT: + return "inherit" + elif spec.type == cpp_PartitioningSpec.cpp_Type.HASH: + return HashScheme.from_cpp(deref(spec.hash)) + elif spec.type == cpp_PartitioningSpec.cpp_Type.ORDER: + return OrderScheme.from_cpp(deref(spec.order)) # copies out of optional + else: + raise ValueError("Unknown PartitioningSpec.Type") + + +cdef class Partitioning: + """ + Hierarchical partitioning metadata for a data stream. + + Parameters + ---------- + inter_rank + Distribution across ranks. Can be a HashScheme, OrderScheme, None, + or 'inherit'. + local + Distribution within a rank. Can be a HashScheme, OrderScheme, None, + or 'inherit'. + """ + + def __init__(self, inter_rank=None, local=None): + _apply_spec(self._handle.inter_rank, inter_rank) + _apply_spec(self._handle.local, local) + + @staticmethod + cdef Partitioning from_cpp(cpp_Partitioning data): + cdef Partitioning ret = Partitioning.__new__(Partitioning) + ret._handle = move(data) + return ret + + @property + def inter_rank(self): + """Inter-rank partitioning spec.""" + return _from_spec(self._handle.inter_rank) + + @property + def local(self): + """Intra-rank (local) partitioning spec.""" + return _from_spec(self._handle.local) + + def __repr__(self): + return f"Partitioning(inter_rank={self.inter_rank!r}, local={self.local!r})" + + +cdef class ChannelMetadata: + """ + Channel-level metadata describing a data stream. + + Parameters + ---------- + local_count + Estimated number of chunks for this rank. + partitioning + How the data is partitioned (default: no partitioning). + duplicated + Whether data is duplicated on all workers (default: False). + """ + + def __init__( + self, + int local_count, + *, + partitioning: Partitioning | None = None, + bint duplicated = False, + ): + if local_count < 0: + raise ValueError(f"local_count must be non-negative, got {local_count}") + + cdef cpp_Partitioning part + if partitioning is not None: + part = (partitioning)._handle + + self._handle = make_unique[cpp_ChannelMetadata]( + local_count, part, duplicated + ) + + def __dealloc__(self): + with nogil: + self._handle.reset() + + @staticmethod + cdef ChannelMetadata from_handle(unique_ptr[cpp_ChannelMetadata] handle): + cdef ChannelMetadata ret = ChannelMetadata.__new__(ChannelMetadata) + ret._handle = move(handle) + return ret + + @staticmethod + def from_message(Message message not None): + """Construct by consuming a Message (message becomes empty).""" + return ChannelMetadata.from_handle( + cpp_channel_metadata_from_message(move(message._handle)) + ) + + def into_message(self, uint64_t sequence_number, Message message not None): + """ + Move this ChannelMetadata into a Message. + + Parameters + ---------- + sequence_number + Ordering identifier for the message. + message + Empty message that will take ownership of this metadata. + """ + if not message.empty(): + raise ValueError("cannot move into a non-empty message") + message._handle = cpp_to_message_channel_metadata( + sequence_number, move(self.release_handle()) + ) + + cdef const cpp_ChannelMetadata* handle_ptr(self) except NULL: + """Return pointer to underlying handle, raising if released.""" + if not self._handle: + raise ValueError("ChannelMetadata is uninitialized, has it been released?") + return self._handle.get() + + @property + def local_count(self) -> int: + """Estimated number of chunks for this rank.""" + return self.handle_ptr().local_count + + @property + def partitioning(self) -> Partitioning: + """How the data is partitioned.""" + return Partitioning.from_cpp(self.handle_ptr().partitioning) + + @property + def duplicated(self) -> bool: + """Whether data is duplicated on all workers.""" + return self.handle_ptr().duplicated + + def __repr__(self): + return ( + f"ChannelMetadata(local_count={self.local_count}, " + f"partitioning={self.partitioning!r}, " + f"duplicated={self.duplicated})" + ) + + cdef unique_ptr[cpp_ChannelMetadata] release_handle(self): + if not self._handle: + raise ValueError("is uninitialized, has it been released?") + return move(self._handle) diff --git a/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi b/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi new file mode 100644 index 000000000000..7a393e1650e6 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from pylibcudf.expressions import Expression +from pylibcudf.io.parquet import ParquetReaderOptions + +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.communicator.communicator import Communicator +from rapidsmpf.streaming.core.actor import CppActor +from rapidsmpf.streaming.core.channel import Channel +from rapidsmpf.streaming.core.context import Context +from rmm.pylibrmm.stream import Stream + +class Filter: + def __init__(self, stream: Stream, expression: Expression) -> None: ... + +def read_parquet( + ctx: Context, + comm: Communicator, + ch_out: Channel[TableChunk], + num_producers: int, + options: ParquetReaderOptions, + num_rows_per_chunk: int, + filter: Filter | None = None, +) -> CppActor: ... diff --git a/python/cudf_streaming/cudf_streaming/streaming/parquet.pyx b/python/cudf_streaming/cudf_streaming/streaming/parquet.pyx new file mode 100644 index 000000000000..13d7a6821852 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/parquet.pyx @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from cpython.object cimport PyObject +from cpython.ref cimport Py_INCREF +from cython.operator cimport dereference as deref +from libc.stddef cimport size_t +from libcpp.memory cimport make_unique, shared_ptr, unique_ptr +from libcpp.utility cimport move +from pylibcudf.expressions cimport Expression +from pylibcudf.io.parquet cimport ParquetReaderOptions +from pylibcudf.libcudf.expressions cimport expression +from pylibcudf.libcudf.io.parquet cimport parquet_reader_options +from pylibcudf.libcudf.types cimport size_type +from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.pylibrmm.stream cimport Stream + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.communicator.communicator cimport Communicator, cpp_Communicator +from rapidsmpf.streaming.chunks.arbitrary cimport cpp_OwningWrapper +from rapidsmpf.streaming.chunks.utils cimport py_deleter +from rapidsmpf.streaming.core.actor cimport CppActor, cpp_Actor +from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel +from rapidsmpf.streaming.core.context cimport Context, cpp_Context + + +cdef extern from "" nogil: + cdef cppclass cpp_Filter "cudf_streaming::streaming::Filter": + cpp_Filter(cuda_stream_view, expression, cpp_OwningWrapper) + + cdef cpp_Actor cpp_read_parquet \ + "cudf_streaming::streaming::actor::read_parquet"( + shared_ptr[cpp_Context] ctx, + shared_ptr[cpp_Communicator] comm, + shared_ptr[cpp_Channel] ch_out, + size_t num_producers, + parquet_reader_options options, + size_type num_rows_per_chunk, + unique_ptr[cpp_Filter], + ) except +ex_handler + + +cdef class Filter: + """ + A filter expression for parquet reads. + + Parameters + ---------- + stream + The stream any scalars in the expression are valid on. + expression + The filter expression + + Notes + ----- + The object safely manages the lifetime of the expressions when called + from C++ coroutines, so it is safe to drop the expression passed in on + the python side. + """ + cdef unique_ptr[cpp_Filter] _handle + + def __init__(self, Stream stream not None, Expression filter not None): + Py_INCREF(filter) + self._handle = make_unique[cpp_Filter]( + stream.view(), + deref(filter.c_obj), + cpp_OwningWrapper( + filter, py_deleter + ) + ) + + cdef unique_ptr[cpp_Filter] release_handle(self): + """ + Move the owning C++ handle out of the object. + + Returns + ------- + unique_ptr to the C++ Filter object. + + Raises + ------ + ValueError + If this Filter has already been used and the handle is already released. + """ + if not self._handle: + raise ValueError("Filter is uninitialized, has it been released?") + return move(self._handle) + + def __dealloc__(self): + with nogil: + self._handle.reset() + + +def read_parquet( + Context ctx not None, + Communicator comm not None, + Channel ch_out not None, + size_t num_producers, + ParquetReaderOptions options not None, + size_type num_rows_per_chunk, + Filter filter = None, +): + """ + Create a streaming actor to read from parquet. + + Parameters + ---------- + ctx + Streaming execution context. + comm + The communicator. + ch_out + Output channel to receive the TableChunks. + num_producers + Number of concurrent producers of output chunks. + options + Reader options. + num_rows_per_chunk + Target (maximum) number of rows per output chunk. + filter + Optional filter object. If provided, is consumed by this function + and not subsequently usable. + + Notes + ----- + This is a collective operation, all ranks participating via the + communicator must call it with the same options. + """ + cdef cpp_Actor _ret + cdef unique_ptr[cpp_Filter] c_filter + if filter is not None: + c_filter = move(filter.release_handle()) + with nogil: + _ret = cpp_read_parquet( + ctx._handle, + comm._handle, + ch_out._handle, + num_producers, + options.c_obj, + num_rows_per_chunk, + move(c_filter) + ) + return CppActor.from_handle( + make_unique[cpp_Actor](move(_ret)), owner=None + ) diff --git a/python/cudf_streaming/cudf_streaming/streaming/partition.pyi b/python/cudf_streaming/cudf_streaming/streaming/partition.pyi new file mode 100644 index 000000000000..940e953ea3ac --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/partition.pyi @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Iterable + +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.streaming.chunks.partition import ( + PartitionMapChunk, + PartitionVectorChunk, +) +from rapidsmpf.streaming.core.actor import CppActor +from rapidsmpf.streaming.core.channel import Channel +from rapidsmpf.streaming.core.context import Context + +def partition_and_pack( + ctx: Context, + ch_in: Channel[TableChunk], + ch_out: Channel[PartitionMapChunk], + columns_to_hash: Iterable[int], + num_partitions: int, +) -> CppActor: ... +def unpack_and_concat( + ctx: Context, + ch_in: Channel[PartitionMapChunk] + | Channel[PartitionVectorChunk] + | Channel[PartitionMapChunk | PartitionVectorChunk], + ch_out: Channel[TableChunk], +) -> CppActor: ... diff --git a/python/cudf_streaming/cudf_streaming/streaming/partition.pyx b/python/cudf_streaming/cudf_streaming/streaming/partition.pyx new file mode 100644 index 000000000000..1df59390a2b5 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/partition.pyx @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + + +from libc.stdint cimport uint32_t +from libcpp.memory cimport make_unique, shared_ptr +from libcpp.utility cimport move +from libcpp.vector cimport vector +from pylibcudf.libcudf.types cimport size_type + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.streaming.core.actor cimport CppActor, cpp_Actor +from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel +from rapidsmpf.streaming.core.context cimport Context, cpp_Context + + +cdef extern from "" nogil: + int cpp_HASH_MURMUR3"cudf::hash_id::HASH_MURMUR3" + uint32_t cpp_DEFAULT_HASH_SEED"cudf::DEFAULT_HASH_SEED", + cdef cpp_Actor cpp_partition_and_pack \ + "cudf_streaming::streaming::actor::partition_and_pack"( + shared_ptr[cpp_Context] ctx, + shared_ptr[cpp_Channel] ch_in, + shared_ptr[cpp_Channel] ch_out, + vector[size_type] columns_to_hash, + int num_partitions, + int hash_function, + uint32_t seed, + ) except +ex_handler + cdef cpp_Actor cpp_unpack_and_concat \ + "cudf_streaming::streaming::actor::unpack_and_concat"( + shared_ptr[cpp_Context] ctx, + shared_ptr[cpp_Channel] ch_in, + shared_ptr[cpp_Channel] ch_out, + ) except +ex_handler + + +def partition_and_pack( + Context ctx not None, + Channel ch_in not None, + Channel ch_out not None, + object columns_to_hash not None, + int num_partitions, +): + """ + Asynchronously partition and pack table chunks. + + This is the streaming equivalent of + :func:`cudf_streaming.integrations.partition.partition_and_split()`, + operating on incoming table chunks via channels. + + Each incoming table from `ch_in` is partitioned into `num_partitions` outputs + based on a hash of the specified columns. Each partition is then serialized + (packed) and sent to the output channel `ch_out`. + + Parameters + ---------- + ctx + The streaming actor context used to create and manage the asynchronous task. + ch_in + Input channel that provides ``TableChunk`` objects to partition. + ch_out + Output channel to which packed partitions (``PartitionMapChunk`` objects) + are sent. + columns_to_hash + Indices of input columns to hash when computing partition assignments. + num_partitions + Number of output partitions to create. + + Returns + ------- + A streaming actor representing the asynchronous partitioning and packing operation. + + Raises + ------ + ValueError + If any index in ``columns_to_hash`` is invalid. + + See Also + -------- + cudf_streaming.integrations.partition.partition_and_pack + Non-streaming variant operating on static tables. + cudf_streaming.streaming.partition.unpack_and_concat + The inverse operation that unpacks and concatenates packed partitions. + """ + cdef vector[size_type] _columns_to_hash = tuple(columns_to_hash) + cdef cpp_Actor _ret + with nogil: + _ret = cpp_partition_and_pack( + ctx._handle, + ch_in._handle, + ch_out._handle, + _columns_to_hash, + num_partitions, + cpp_HASH_MURMUR3, + cpp_DEFAULT_HASH_SEED, + ) + return CppActor.from_handle( + make_unique[cpp_Actor](move(_ret)), owner = None + ) + + +def unpack_and_concat( + Context ctx not None, + Channel ch_in not None, + Channel ch_out not None, +): + """ + Asynchronously unpack and concatenate packed partitions. + + This is the streaming equivalent of + :func:`cudf_streaming.integrations.partition.unpack_and_concat()`, + operating on packed partition chunks via channels. + + The function receives packed partitions from `ch_in`, deserializes them, + concatenates the partitions belonging to the same logical table, and sends the + resulting tables to `ch_out`. Empty partitions are automatically ignored. + + Parameters + ---------- + ctx + The streaming actor context used to manage asynchronous execution. + ch_in + Input channel providing packed partitions (``PartitionMapChunk`` or + ``PartitionVectorChunk``). + ch_out + Output channel receiving the unpacked and concatenated ``TableChunk`` objects. + + Returns + ------- + A streaming actor representing the asynchronous unpacking and concatenation + operation. + + See Also + -------- + cudf_streaming.integrations.partition.unpack_and_concat + Non-streaming version. + cudf_streaming.streaming.partition.partition_and_pack + The inverse operation that partitions and packs tables into partitions. + """ + cdef cpp_Actor _ret + with nogil: + _ret = cpp_unpack_and_concat( + ctx._handle, + ch_in._handle, + ch_out._handle, + ) + return CppActor.from_handle( + make_unique[cpp_Actor](move(_ret)), owner = None + ) diff --git a/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pxd b/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pxd new file mode 100644 index 000000000000..9f7a09785e79 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pxd @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stddef cimport size_t +from libc.stdint cimport uint64_t +from libcpp cimport bool as bool_t +from libcpp.memory cimport unique_ptr +from libcpp.utility cimport pair +from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view +from pylibcudf.libcudf.types cimport size_type +from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.pylibrmm.stream cimport Stream + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.memory.buffer cimport MemoryType +from rapidsmpf.memory.buffer_resource cimport (BufferResource, + cpp_BufferResource) +from rapidsmpf.memory.memory_reservation cimport cpp_MemoryReservation +from rapidsmpf.memory.packed_data cimport cpp_PackedData + + +cdef extern from "" nogil: + cdef cppclass cpp_TableChunk "cudf_streaming::streaming::TableChunk": + cpp_TableChunk(unique_ptr[cpp_PackedData]) except +ex_handler + cuda_stream_view stream() noexcept + size_t data_alloc_size(MemoryType mem_type) except +ex_handler + bool_t is_available() noexcept + size_t make_available_cost() noexcept + cpp_table_view table_view() except +ex_handler + bool_t is_spillable() noexcept + cpp_TableChunk copy(cpp_MemoryReservation& reservation) except +ex_handler + pair[size_type, size_type] shape() noexcept + unique_ptr[cpp_PackedData] into_packed_data( + cpp_BufferResource* br + ) except +ex_handler + +cdef class TableChunk: + cdef unique_ptr[cpp_TableChunk] _handle + # Keep the BufferResource alive as long as this object is so that when this + # object is deallocated the associated stream and memory resource are still alive. + cdef BufferResource _br + + @staticmethod + cdef TableChunk from_handle(unique_ptr[cpp_TableChunk] handle, BufferResource br) + cdef const cpp_TableChunk* handle_ptr(self) + cdef unique_ptr[cpp_TableChunk] release_handle(self) diff --git a/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyi b/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyi new file mode 100644 index 000000000000..2228f93a9729 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyi @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Iterable +from typing import Self, overload + +from pylibcudf.table import Table + +from rapidsmpf.memory.buffer import MemoryType +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.memory_reservation import MemoryReservation +from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.streaming.core.context import Context +from rapidsmpf.streaming.core.message import Message +from rmm.pylibrmm.stream import Stream + +class TableChunk: + @staticmethod + def from_pylibcudf_table( + table: Table, + stream: Stream, + *, + exclusive_view: bool, + br: BufferResource, + ) -> TableChunk: ... + @staticmethod + def from_packed_data(pd: PackedData, br: BufferResource) -> TableChunk: ... + @classmethod + def from_message( + cls: type[Self], message: Message[Self], br: BufferResource + ) -> Self: ... + def into_message( + self, sequence_number: int, message: Message[Self] + ) -> None: ... + @property + def stream(self) -> Stream: ... + def data_alloc_size(self, mem_type: MemoryType | None = None) -> int: ... + def is_available(self) -> bool: ... + def make_available_cost(self) -> int: ... + def make_available(self, reservation: MemoryReservation) -> TableChunk: ... + async def make_available_or_wait( + self, ctx: Context, *, net_memory_delta: int + ) -> TableChunk: ... + def make_available_and_spill( + self, br: BufferResource, *, allow_overbooking: bool + ) -> TableChunk: ... + def table_view(self) -> Table: ... + def is_spillable(self) -> bool: ... + def copy(self, reservation: MemoryReservation) -> TableChunk: ... + def into_packed_data(self, br: BufferResource) -> PackedData: ... + @property + def shape(self) -> tuple[int, int]: ... + +@overload +async def make_table_chunks_available_or_wait( + context: Context, + chunks: TableChunk, + *, + reserve_extra: int, + net_memory_delta: int, + allow_overbooking: bool | None = None, +) -> tuple[TableChunk, MemoryReservation]: ... +@overload +async def make_table_chunks_available_or_wait( + context: Context, + chunks: Iterable[TableChunk], + *, + reserve_extra: int, + net_memory_delta: int, + allow_overbooking: bool | None = None, +) -> tuple[list[TableChunk], MemoryReservation]: ... diff --git a/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyx b/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyx new file mode 100644 index 000000000000..fb12b3573d24 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/table_chunk.pyx @@ -0,0 +1,664 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from cpython.object cimport PyObject +from cython cimport no_gc_clear +from cython.operator cimport dereference as deref +from libc.stdint cimport int64_t, uint64_t +from libcpp.memory cimport make_unique, unique_ptr +from libcpp.utility cimport move +from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view +from pylibcudf.table cimport Table + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.memory.buffer_resource cimport (BufferResource, + cpp_BufferResource) +from rapidsmpf.memory.memory_reservation cimport (MemoryReservation, + cpp_MemoryReservation) +from rapidsmpf.memory.packed_data cimport PackedData +# Need the header include for inline C++ code +from rapidsmpf.owning_wrapper cimport cpp_OwningWrapper # no-cython-lint +from rapidsmpf.streaming.chunks.utils cimport py_deleter +from rapidsmpf.streaming.core.context cimport Context +from rapidsmpf.streaming.core.message cimport Message, cpp_Message + +from rapidsmpf.memory.buffer import MemoryType as py_MemoryType +from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory + + +cdef extern from "" nogil: + cpp_Message cpp_to_message"cudf_streaming::streaming::to_message"\ + (uint64_t sequence_number, unique_ptr[cpp_TableChunk]) except +ex_handler + + +cdef extern from * nogil: + """ + namespace { + std::unique_ptr + cpp_release_table_chunk_from_message( + rapidsmpf::streaming::Message &&msg + ) { + return std::make_unique( + msg.release() + ); + } + + std::unique_ptr + cpp_from_table_view_with_owner( + cudf::table_view view, + rmm::cuda_stream_view stream, + PyObject *owner, + void(*py_deleter)(void *), + bool exclusive_view + ) { + // Called holding the gil. + // Decref is done by the deleter. + Py_XINCREF(owner); + return std::make_unique( + view, + stream, + rapidsmpf::OwningWrapper(owner, py_deleter), + exclusive_view ? + cudf_streaming::streaming::TableChunk::ExclusiveView::YES + : cudf_streaming::streaming::TableChunk::ExclusiveView::NO + ); + } + + std::unique_ptr cpp_table_make_available( + std::unique_ptr &&table, + rapidsmpf::MemoryReservation* reservation + ) { + return std::make_unique( + table->make_available(*reservation) + ); + } + + std::unique_ptr cpp_table_copy( + std::unique_ptr const& table, + rapidsmpf::MemoryReservation* reservation + ) { + return std::make_unique( + table->copy(*reservation) + ); + } + } // namespace + """ + unique_ptr[cpp_TableChunk] cpp_release_table_chunk_from_message( + cpp_Message + ) except +ex_handler + unique_ptr[cpp_TableChunk] cpp_from_table_view_with_owner(...) except +ex_handler + unique_ptr[cpp_TableChunk] cpp_table_make_available( + unique_ptr[cpp_TableChunk], cpp_MemoryReservation* + ) except +ex_handler + unique_ptr[cpp_TableChunk] cpp_table_copy( + unique_ptr[cpp_TableChunk], cpp_MemoryReservation* + ) except +ex_handler + + +@no_gc_clear +cdef class TableChunk: + """ + A unit of table data in a streaming pipeline. + + Represents either an unpacked pylibcudf table, a packed (serialized) table, + or `rapidsmpf.memory.packed_data.PackedData`. + + A TableChunk may be initially unavailable (e.g., if the data is packed or + spilled), and can be made available (i.e., materialized to device memory) + on demand. + + Use the factory functions `from_pylibcudf_table` and `from_message` to + create a new table chunk. + """ + def __init__(self): + raise ValueError("use the `from_*` factory functions") + + def __dealloc__(self): + with nogil: + self._handle.reset() + + @staticmethod + cdef TableChunk from_handle( + unique_ptr[cpp_TableChunk] handle, BufferResource br, + ): + """ + Construct a TableChunk from an existing C++ handle. + + Parameters + ---------- + handle + A unique pointer to a C++ TableChunk. + br + A BufferResource to keep alive. + + Returns + ------- + A new TableChunk wrapping the given handle. + """ + cdef TableChunk ret = TableChunk.__new__(TableChunk) + ret._handle = move(handle) + ret._br = br + return ret + + @staticmethod + def from_pylibcudf_table( + Table table not None, + Stream stream not None, + *, + bool_t exclusive_view, + BufferResource br not None, + ): + """ + Construct a TableChunk from a pylibcudf Table. + + Parameters + ---------- + table + A pylibcudf Table to wrap as a TableChunk. + stream + The CUDA stream on which this chunk was created. + exclusive_view + Indicates that this TableChunk has exclusive ownership semantics for the + underlying table view. + + When ``True``, the following guarantees must hold: + - The pylibcudf Table is the sole representation of the table data, + i.e. no views exist. + - The Table object exclusively owns the table's device memory. + + These guarantees allow the TableChunk to be spillable and ensure that + when the owner is destroyed, the underlying device memory is correctly + freed. + + Returns + ------- + A new TableChunk wrapping the given pylibcudf Table. + + Notes + ----- + The returned TableChunk maintains a reference to ``table`` to ensure + its underlying buffers remain valid for the lifetime of the chunk. + This reference is managed by the underlying C++ object, so it + persists even when the chunk is transferred through Channels. + + """ + cdef cuda_stream_view _stream = stream.view() + cdef cpp_table_view view = table.view() + return TableChunk.from_handle( + cpp_from_table_view_with_owner( + view, + _stream, + table, + py_deleter, + exclusive_view, + ), + br, + ) + + @staticmethod + def from_packed_data(PackedData pd not None, BufferResource br not None): + """ + Construct a TableChunk from packed data. + + Parameters + ---------- + pd + The PackedData object + + Returns + ------- + A new TableChunk owning the packed data. + + Notes + ----- + This takes ownership of the data in the PackedData object, which is left empty. + """ + return TableChunk.from_handle(make_unique[cpp_TableChunk](move(pd.c_obj)), br) + + @staticmethod + def from_message(Message message not None, BufferResource br not None): + """ + Construct a TableChunk by consuming a Message. + + Parameters + ---------- + message + Message containing a TableChunk. The message is released and is empty + after this call. + + Returns + ------- + A new TableChunk extracted from the given message. + """ + return TableChunk.from_handle( + cpp_release_table_chunk_from_message(move(message._handle)), + br, + ) + + def into_message(self, uint64_t sequence_number, Message message not None): + """ + Move this TableChunk into an empty Message. + + This method is not typically called directly. Instead, it is invoked by + `Message.__init__()` when creating a new Message with this TableChunk as + its payload. + + Parameters + ---------- + sequence_number + Ordering identifier for the message. + message + Message object that will take ownership of this TableChunk. + + Raises + ------ + ValueError + If the provided message is not empty. + + Warnings + -------- + The original table chunk is released and must not be used after this call. + """ + if not message.empty(): + raise ValueError("cannot move into a non-empty message") + message._handle = cpp_to_message( + sequence_number, move(self.release_handle()) + ) + + cdef const cpp_TableChunk* handle_ptr(self): + """ + Return a pointer to the underlying C++ TableChunk. + + Returns + ------- + Raw pointer to the underlying C++ object. + + Raises + ------ + ValueError + If the TableChunk is uninitialized. + """ + if not self._handle: + raise ValueError("TableChunk is uninitialized, has it been released?") + return self._handle.get() + + cdef unique_ptr[cpp_TableChunk] release_handle(self): + """ + Release ownership of the underlying C++ TableChunk. + + After this call, the current object is in a moved-from state and + must not be accessed. + + Returns + ------- + Unique pointer to the underlying C++ object. + + Raises + ------ + ValueError + If the TableChunk is uninitialized. + """ + if not self._handle: + raise ValueError("TableChunk is uninitialized, has it been released?") + return move(self._handle) + + @property + def stream(self): + """ + Return the CUDA stream on which this chunk was created. + + Returns + ------- + Stream + The CUDA stream. + """ + return Stream._from_cudaStream_t( + deref(self.handle_ptr()).stream().value() + ) + + def data_alloc_size(self, mem_type=None): + """ + Number of bytes allocated for the data in the specified memory type. + + Parameters + ---------- + mem_type + The memory type to query. If None, returns the total size across + all memory types. + + Returns + ------- + Number of bytes allocated. + """ + if mem_type is None: + return sum(self.data_alloc_size(m) for m in py_MemoryType) + return deref(self.handle_ptr()).data_alloc_size(mem_type) + + def is_available(self): + """ + Indicates whether the underlying table data is fully available in + device memory. + + Returns + ------- + True if the table is already available; otherwise, False. + """ + return deref(self.handle_ptr()).is_available() + + def make_available_cost(self): + """ + Return the estimated cost (in bytes) of making the table available. + + Currently, only device memory usage is accounted for in this estimate. + + Returns + ------- + The estimated cost in bytes. + """ + return deref(self.handle_ptr()).make_available_cost() + + def make_available(self, MemoryReservation reservation not None): + """ + Move this table chunk into a new one with its data made available. + + As part of the move, a copy or unpack operation may be performed, + using the associated CUDA stream for execution. + + Parameters + ---------- + reservation + Memory reservation used for allocations, if making data available + is needed. + + Returns + ------- + A new table chunk with its data available on device. + + Warnings + -------- + The original table chunk is released and must not be used after this call. + """ + cdef cpp_MemoryReservation* res = reservation._handle.get() + cdef unique_ptr[cpp_TableChunk] handle = self.release_handle() + cdef unique_ptr[cpp_TableChunk] ret + with nogil: + ret = cpp_table_make_available(move(handle), res) + return TableChunk.from_handle(move(ret), self._br) + + async def make_available_or_wait( + self, Context ctx not None, *, int64_t net_memory_delta + ): + """ + Move this table chunk into a new one with its data made available. + + This is an asynchronous variant of :meth:`make_available`. The coroutine may + suspend if the required device memory is not immediately available and + resumes once a memory reservation has been granted or an error condition is + reached. + + Parameters + ---------- + ctx + Streaming context used to access the memory reservation mechanism. + net_memory_delta + Estimated change in memory usage after the reservation is granted and + all work using the returned table chunk has completed. + + Returns + ------- + A new table chunk with its data available on device. + + Raises + ------ + RuntimeError + If shutdown occurs before the reservation can be processed. + OverflowError + If no progress is possible within the timeout and overbooking is + disabled. + + Warnings + -------- + The original table chunk is released and must not be used after this call. + """ + return self.make_available( + await reserve_memory( + ctx, self.make_available_cost(), net_memory_delta=net_memory_delta + ) + ) + + def make_available_and_spill( + self, BufferResource br not None, *, allow_overbooking + ): + """ + Make this table chunk available on device, spilling other data if necessary. + + Ensures that the data backing this table chunk is made available in device + memory. If there is insufficient free memory to complete the operation, the + buffer resource may spill other data until enough space has been freed. + + Parameters + ---------- + br + Buffer resource used for allocations and spill management. + allow_overbooking + Whether the memory reservation may temporarily exceed the current + allocation limit. + + Returns + ------- + A new table chunk with its data made available on device. + + Raises + ------ + ReservationError + If the allocation or spilling process fails to free enough memory. + + Warnings + -------- + The original table chunk is released and must not be used after this call. + + Examples + -------- + >>> # Make the data of an existing chunk available on device + >>> chunk = chunk.make_available_and_spill(br, allow_overbooking=False) + >>> chunk.table_view() + """ + + cdef MemoryReservation res = br.reserve_device_memory_and_spill( + self.make_available_cost(), + allow_overbooking=allow_overbooking + ) + return self.make_available(res) + + def table_view(self): + """ + Returns a view of the underlying pylibcudf table. + + The table must be available in device memory. + + Returns + ------- + A view of the underlying table. The view holds a reference to this + `TableChunk` to ensure it remains alive. + + Raises + ------ + ValueError + If ``self.is_available() is False``. + """ + cdef const cpp_TableChunk* handle = self.handle_ptr() + cdef cpp_table_view ret + with nogil: + ret = deref(handle).table_view() + return Table.from_table_view_of_arbitrary(ret, owner=self, stream=self.stream) + + def is_spillable(self): + """ + Indicates whether this chunk can be spilled. + + A chunk is considered spillable if it was created from one of the following: + - A message (via ``.from_message()``). + - An exclusive pylibcudf table (via + ``.from_pylibcudf_table(..., exclusive_view=True)``). + + Both of these creation paths imply device-owning semantics, meaning the + TableChunk owns its underlying memory and can safely be spilled to host memory. + + Returns + ------- + True if the table chunk can be spilled, otherwise, False. + """ + return deref(self.handle_ptr()).is_spillable() + + def copy(self, MemoryReservation reservation not None): + """ + Create a deep copy of this table chunk. + + All buffers are allocated for the new table chunk using the provided + memory reservation, which also determines the target memory type of + the copy. + + Parameters + ---------- + reservation + Memory reservation to consume for allocating the buffers of the + new table chunk. + + Returns + ------- + TableChunk + A new table chunk containing a deep copy of this chunk's data and + metadata. + """ + cdef unique_ptr[cpp_TableChunk] ret + cdef cpp_MemoryReservation* res = reservation._handle.get() + with nogil: + ret = cpp_table_copy(self._handle, res) + return TableChunk.from_handle(move(ret), self._br) + + def into_packed_data(self, BufferResource br not None): + """ + Convert this table chunk to a PackedData, avoiding unnecessary copies. + + If the chunk's data is already in packed form (e.g., it arrived over the + network or was constructed from a :class:`PackedData`), the packed data is + moved out directly with no copy. Otherwise the table is serialized via + ``cudf.pack()``. + + Parameters + ---------- + br + Buffer resource used when packing is required. + + Returns + ------- + PackedData + The resulting packed data. + + Raises + ------ + ValueError + If the data is not already packed and the table is not available + (i.e., ``is_available() == False``). + + Warnings + -------- + The original table chunk is released and must not be used after this call. + """ + cdef unique_ptr[cpp_PackedData] result + cdef cpp_BufferResource* _br = br.ptr() + cdef unique_ptr[cpp_TableChunk] handle = self.release_handle() + with nogil: + result = move(deref(handle)).into_packed_data(_br) + return PackedData.from_librapidsmpf(move(result), br) + + @property + def shape(self): + """Return the shape of the table in this TableChunk. + + Returns + ------- + Tuple of shape ``(num_rows, num_columns)```. + """ + return deref(self._handle).shape() + + +async def make_table_chunks_available_or_wait( + Context ctx not None, + chunks, + *, + size_t reserve_extra, + int64_t net_memory_delta, + allow_overbooking=None, +): + """ + Make one or more table chunks available, waiting on a memory reservation if needed. + + This helper combines :meth:`TableChunk.make_available` with :func:`reserve_memory`. + It computes the device-memory cost of making the provided table chunks available, + reserves that amount (plus ``reserve_extra``), and then returns new table chunks + whose data are available on device. + + The coroutine may suspend if the required device memory is not immediately + available and resumes once a memory reservation has been granted or an error + condition is reached. The behavior when the progress timeout expires depends + on whether overbooking is allowed. + + Parameters + ---------- + ctx + Streaming context used to access the memory reservation mechanism. + chunks + A table chunk or an iterable of table chunks to make available on device. + reserve_extra + Additional bytes to include in the reservation beyond the aggregated + availability cost of ``chunks``. + net_memory_delta + Estimated change in memory usage after the reservation is granted and all + work using the returned table chunks has completed. This value is used as + a heuristic to prioritize eligible requests. + allow_overbooking + Whether to allow overbooking if no progress is possible. + - If ``True``, the reservation may overbook memory when no further + progress can be made. If ``False``, the call fails when no progress + is possible. + - If ``None``, the behavior is determined by the configuration option + ``"allow_overbooking_by_default"``, which is read via ``ctx.options()``. + + Returns + ------- + A tuple containing: + - If a chunk is provided: the new table chunk with its data available on device. + - If multiple chunks are provided: a list of new table chunks with their data + available on device. + - The memory reservation used with a remaining size of @p reserve_extra. + + Raises + ------ + RuntimeError + If shutdown occurs before the reservation can be processed. + OverflowError + If no progress is possible within the timeout and overbooking is disabled. + + Warnings + -------- + The original table chunks are released and must not be used after this call. + """ + # Handle both single chunk and iterable of chunks + input_chunks = chunks + if isinstance(input_chunks, TableChunk): + chunks = (input_chunks,) + + size = sum(chunk.make_available_cost() for chunk in chunks) + res = await reserve_memory( + ctx, + size + reserve_extra, + net_memory_delta=net_memory_delta, + mem_type=MemoryType.DEVICE, + allow_overbooking=allow_overbooking, + ) + available_chunks = [chunk.make_available(res) for chunk in chunks] + if isinstance(input_chunks, TableChunk): + return available_chunks[0], res + else: + return available_chunks, res From 006a715dbafc623b93517c98863e4b795386a9c6 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 2 Jun 2026 01:20:51 +0000 Subject: [PATCH 04/19] feat(cudf_streaming): add Python tests --- .../cudf_streaming/tests/conftest.py | 35 + .../cudf_streaming/tests/test_bloom_filter.py | 176 +++++ .../tests/test_channel_metadata.py | 419 +++++++++++ .../cudf_streaming/tests/test_partition.py | 94 +++ .../cudf_streaming/tests/test_read_parquet.py | 185 +++++ .../cudf_streaming/tests/test_table_chunk.py | 652 ++++++++++++++++++ 6 files changed, 1561 insertions(+) create mode 100644 python/cudf_streaming/cudf_streaming/tests/conftest.py create mode 100644 python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py create mode 100644 python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py create mode 100644 python/cudf_streaming/cudf_streaming/tests/test_partition.py create mode 100644 python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py create mode 100644 python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py diff --git a/python/cudf_streaming/cudf_streaming/tests/conftest.py b/python/cudf_streaming/cudf_streaming/tests/conftest.py new file mode 100644 index 000000000000..8fc1023afa2c --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/conftest.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +import pytest + +import rmm.mr +from rapidsmpf.config import Options, get_environment_variables +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +from rapidsmpf.streaming.core.context import Context + +sys.path.insert(0, str(Path(__file__).parents[2])) + +if TYPE_CHECKING: + from collections.abc import Generator + + from rapidsmpf.communicator.communicator import Communicator + + +@pytest.fixture +def context(comm: Communicator) -> Generator[Context, None, None]: + """ + Fixture to get a streaming context. + """ + options = Options(get_environment_variables()) + mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + br = BufferResource(mr) + + with Context(comm.logger, br, options) as ctx: + yield ctx diff --git a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py new file mode 100644 index 000000000000..c0d7f718aa6b --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +import numpy as np +import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming.streaming") + +from cudf_streaming.streaming import ChannelMetadata +from cudf_streaming.streaming.bloom_filter import BloomFilter +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.streaming.core.actor import define_actor, run_actor_network +from rapidsmpf.streaming.core.leaf_actor import ( + pull_from_channel, + push_to_channel, +) +from rapidsmpf.streaming.core.message import Message +from rapidsmpf.testing import assert_eq + +if TYPE_CHECKING: + from collections.abc import Awaitable + + from cudf_streaming.streaming.bloom_filter import BloomFilterChunk + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.memory.buffer_resource import BufferResource + from rapidsmpf.streaming.core.actor import CppActor + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +def make_table( + values: np.ndarray, stream: Stream, br: BufferResource +) -> TableChunk: + table = plc.Table([plc.Column.from_array(values, stream=stream)]) + return TableChunk.from_pylibcudf_table( + table, stream, exclusive_view=True, br=br + ) + + +@define_actor() +async def add_metadata( + ctx: Context, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk] +) -> None: + await ch_out.send_metadata(ctx, Message(0, ChannelMetadata(1))) + await ch_out.drain_metadata(ctx) + while (msg := await ch_in.recv(ctx)) is not None: + await ch_out.send(ctx, msg) + await ch_out.drain(ctx) + + +@define_actor() +async def receive_metadata( + ctx: Context, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk] +) -> None: + m = await ch_in.recv_metadata(ctx) + assert m is not None + meta = ChannelMetadata.from_message(m) + assert meta.local_count == 1 + while (msg := await ch_in.recv(ctx)) is not None: + await ch_out.send(ctx, msg) + await ch_out.drain(ctx) + + +@define_actor() +async def bloom_pipeline( + ctx: Context, + bloom: BloomFilter, + ch_build: Channel[TableChunk], + ch_probe: Channel[TableChunk], + ch_out: Channel[TableChunk], +) -> None: + ch_filter: Channel[BloomFilterChunk] = ctx.create_channel() + await asyncio.gather( + bloom.build(ctx, ch_in=ch_build, ch_out=ch_filter, tag=0), + bloom.apply( + ctx, + bloom_filter=ch_filter, + ch_in=ch_probe, + ch_out=ch_out, + keys=(0,), + ), + ) + + +def run_bloom_filter_pipeline( + context: Context, + comm: Communicator, + build_table: TableChunk, + probe_table: TableChunk, + *, + seed: int = 42, + l2size: int = 1 << 20, +) -> list[Message]: + bloom = BloomFilter( + context, + comm, + seed=seed, + num_filter_blocks=BloomFilter.fitting_num_blocks(l2size), + ) + + build_msg = Message(0, build_table) + probe_msg = Message(0, probe_table) + + ch_build: Channel[TableChunk] = context.create_channel() + ch_probe: Channel[TableChunk] = context.create_channel() + ch_probe_meta: Channel[TableChunk] = context.create_channel() + ch_out_meta: Channel[TableChunk] = context.create_channel() + ch_out: Channel[TableChunk] = context.create_channel() + + actors: list[CppActor | Awaitable[None]] = [ + push_to_channel(context, ch_build, [build_msg]), + push_to_channel(context, ch_probe, [probe_msg]), + add_metadata(context, ch_probe, ch_probe_meta), + bloom_pipeline(context, bloom, ch_build, ch_probe_meta, ch_out_meta), + receive_metadata(context, ch_out_meta, ch_out), + ] + pull_actor, deferred = pull_from_channel(context, ch_out) + actors.append(pull_actor) + run_actor_network(context, actors=actors) + return deferred.release() + + +def test_bloom_filter_roundtrip(context: Context, comm: Communicator) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + + stream = context.get_stream_from_pool() + values = np.arange(10, dtype=np.int32) + build_table = make_table(values, stream=stream, br=context.br()) + probe_table = make_table(values, stream=stream, br=context.br()) + messages = run_bloom_filter_pipeline( + context, comm, build_table, probe_table + ) + assert len(messages) == 1 + + result = TableChunk.from_message(messages[0], br=context.br()) + expected = plc.Table([plc.Column.from_array(values, stream=result.stream)]) + result.stream.synchronize() + assert_eq(result.table_view(), expected) + + +def test_bloom_filter_empty_build_filters_all( + context: Context, comm: Communicator +) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + + stream = context.get_stream_from_pool() + build_table = make_table( + np.array([], dtype=np.int32), stream=stream, br=context.br() + ) + probe_table = make_table( + np.arange(5, dtype=np.int32), stream=stream, br=context.br() + ) + messages = run_bloom_filter_pipeline( + context, comm, build_table, probe_table + ) + assert len(messages) == 1 + + result = TableChunk.from_message(messages[0], br=context.br()) + expected = plc.Table( + [ + plc.Column.from_array( + np.array([], dtype=np.int32), stream=result.stream + ) + ] + ) + result.stream.synchronize() + assert_eq(result.table_view(), expected) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py new file mode 100644 index 000000000000..ce5c456db706 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -0,0 +1,419 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for streaming metadata types (Partitioning and ChannelMetadata).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming.streaming") + +from cudf_streaming.streaming import ( + ChannelMetadata, + HashScheme, + OrderKey, + OrderScheme, + Partitioning, + TableChunk, +) +from rapidsmpf.streaming.core.message import Message + +if TYPE_CHECKING: + from rapidsmpf.streaming.core.context import Context + + +def _make_boundaries(context: Context, table: plc.Table) -> TableChunk: + stream = context.get_stream_from_pool() + return TableChunk.from_pylibcudf_table( + table, + stream, + exclusive_view=False, + br=context.br(), + ) + + +def _two_key_order_scheme( + context: Context, *, strict_boundaries: bool = False +) -> OrderScheme: + """Two-key OrderScheme with a 1-row boundary table (2 partitions).""" + boundaries = _make_boundaries( + context, + plc.Table( + [ + plc.Column.from_iterable_of_py( + [100], plc.DataType(plc.TypeId.INT64) + ), + plc.Column.from_iterable_of_py( + ["abc"], plc.DataType(plc.TypeId.STRING) + ), + ] + ), + ) + return OrderScheme( + [ + OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ], + boundaries, + strict_boundaries=strict_boundaries, + ) + + +def test_hash_scheme() -> None: + """Test HashScheme construction, properties, equality, and repr.""" + h1 = HashScheme((0, 1), 16) + assert h1.column_indices == (0, 1) + assert h1.modulus == 16 + assert repr(h1) == "HashScheme((0, 1), 16)" + + # Equality + assert h1 == HashScheme((0, 1), 16) + assert h1 != HashScheme((0, 1), 32) + assert h1 != HashScheme((2,), 16) + + +def test_order_key() -> None: + k = OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) + assert k.column_index == 0 + assert k.order == plc.types.Order.ASCENDING + assert k.null_order == plc.types.NullOrder.BEFORE + assert k == OrderKey( + 0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE + ) + assert k != OrderKey( + 1, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE + ) + assert "OrderKey" in repr(k) + + +def test_order_scheme(context: Context) -> None: + """Test OrderScheme construction, properties, equality, and repr.""" + o1 = _two_key_order_scheme(context) + assert o1.keys == ( + OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ) + assert not o1.strict_boundaries + assert o1.num_boundaries == 1 + assert "OrderScheme" in repr(o1) + + assert o1.boundaries_aligned_with( + _two_key_order_scheme(context), context.br() + ) + + o_strict = _two_key_order_scheme(context, strict_boundaries=True) + assert o_strict.strict_boundaries + assert not o1.boundaries_aligned_with(o_strict, context.br()) + assert o_strict.boundaries_aligned_with( + _two_key_order_scheme(context, strict_boundaries=True), context.br() + ) + + with pytest.raises(TypeError, match="OrderKey"): + OrderScheme( + [ + (0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) + ], # ty: ignore[invalid-argument-type] + _make_boundaries( + context, + plc.Table( + [ + plc.Column.from_iterable_of_py( + [0], plc.DataType(plc.TypeId.INT64) + ) + ] + ), + ), + ) + + with pytest.raises(ValueError, match="empty"): + OrderScheme( + [], + _make_boundaries( + context, + plc.Table( + [ + plc.Column.from_iterable_of_py( + [0], plc.DataType(plc.TypeId.INT64) + ) + ] + ), + ), + ) + + +def test_order_scheme_get_boundaries(context: Context) -> None: + scheme = _two_key_order_scheme(context) + chunk = scheme.get_boundaries(context.br()) + assert chunk.table_view().num_columns() == 2 + assert chunk.table_view().num_rows() == 1 + scheme2 = OrderScheme( + scheme.keys, + chunk, + strict_boundaries=scheme.strict_boundaries, + ) + assert scheme2.boundaries_aligned_with(scheme, context.br()) + + +def test_order_scheme_with_keys(context: Context) -> None: + """with_keys shares boundaries and updates column indices.""" + o1 = _two_key_order_scheme(context) + new_keys = [ + OrderKey(5, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(3, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ] + o2 = o1.with_keys(new_keys) + assert o2.keys[0].column_index == 5 + assert o2.keys[1].column_index == 3 + assert o2.num_boundaries == o1.num_boundaries + assert o2.strict_boundaries == o1.strict_boundaries + # Schemes with different key indices but shared boundaries are boundary-aligned + assert o1.boundaries_aligned_with(o2, context.br()) + + +def test_order_scheme_boundaries_aligned_with(context: Context) -> None: + """boundaries_aligned_with performs value-level boundary comparison, ignoring key indices.""" + df = plc.Table( + [ + plc.Column.from_iterable_of_py( + [100, 200], plc.DataType(plc.TypeId.INT64) + ), + plc.Column.from_iterable_of_py( + ["abc", "xyz"], plc.DataType(plc.TypeId.STRING) + ), + ] + ) + keys = [ + OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ] + o1 = OrderScheme(keys, _make_boundaries(context, df)) + o2 = OrderScheme(keys, _make_boundaries(context, df)) + assert o1.boundaries_aligned_with(o2, context.br()) + + # Different key column indices but same boundary values → still aligned + shifted_keys = [ + OrderKey(2, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(3, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ] + o_shifted = OrderScheme(shifted_keys, _make_boundaries(context, df)) + assert o1.boundaries_aligned_with(o_shifted, context.br()) + + # Different boundary values → not aligned (shape matches, values differ) + df_diff = plc.Table( + [ + plc.Column.from_iterable_of_py( + [100, 300], plc.DataType(plc.TypeId.INT64) + ), + plc.Column.from_iterable_of_py( + ["abc", "xyz"], plc.DataType(plc.TypeId.STRING) + ), + ] + ) + o3 = OrderScheme(keys, _make_boundaries(context, df_diff)) + assert not o1.boundaries_aligned_with(o3, context.br()) + + # Different strict_boundaries → not aligned + o_strict = OrderScheme( + keys, _make_boundaries(context, df), strict_boundaries=True + ) + assert not o1.boundaries_aligned_with(o_strict, context.br()) + + +def test_order_scheme_key_column_mismatch(context: Context) -> None: + """OrderScheme rejects key/column count mismatch.""" + boundaries = _make_boundaries( + context, + plc.Table( + [ + plc.Column.from_iterable_of_py( + [0], plc.DataType(plc.TypeId.INT64) + ) + ] + ), + ) + with pytest.raises(ValueError, match="keys must match"): + OrderScheme( + [ + OrderKey( + 0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE + ), + OrderKey( + 1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER + ), + ], + boundaries, # 1 column, but 2 keys + ) + + +def test_partitioning_scenarios(context: Context) -> None: + """Test various partitioning configurations.""" + # Default / None + p_default = Partitioning() + assert p_default.inter_rank is None + assert p_default.local is None + assert Partitioning(None, None).inter_rank is None + assert Partitioning(None, None).local is None + + # Direct global shuffle: inter_rank=Hash, local=Aligned + p_global = Partitioning(HashScheme((0,), 16), "inherit") + assert p_global.inter_rank == HashScheme((0,), 16) + assert p_global.local == "inherit" + + # Two-stage shuffle: inter_rank=Hash(nranks), local=Hash(N_l) + p_twostage = Partitioning(HashScheme((0,), 4), HashScheme((0,), 8)) + assert p_twostage.inter_rank == HashScheme((0,), 4) + assert p_twostage.local == HashScheme((0,), 8) + + # Order-based partitioning (range partitioned / sorted) + order_scheme = _two_key_order_scheme(context) + p_ordered = Partitioning(order_scheme, "inherit") + assert isinstance(p_ordered.inter_rank, OrderScheme) + assert p_ordered.inter_rank.boundaries_aligned_with( + order_scheme, context.br() + ) + assert p_ordered.local == "inherit" + + # Mixed: inter_rank=Order, local=Hash + p_mixed = Partitioning( + _two_key_order_scheme(context), + HashScheme((1,), 8), + ) + assert isinstance(p_mixed.inter_rank, OrderScheme) + assert isinstance(p_mixed.local, HashScheme) + + # Repr + assert "Partitioning" in repr(p_global) + assert "inter_rank" in repr(p_global) + + # Invalid type + with pytest.raises(TypeError): + Partitioning("invalid", None) # ty: ignore[invalid-argument-type] + + +def test_channel_metadata() -> None: + """Test ChannelMetadata construction and properties.""" + # Basic construction + m = ChannelMetadata(local_count=4) + assert m.local_count == 4 + assert not m.duplicated + + # With partitioning and duplicated + p = Partitioning(HashScheme((0,), 16), "inherit") + m_full = ChannelMetadata(local_count=4, partitioning=p, duplicated=True) + assert m_full.partitioning.inter_rank == HashScheme((0,), 16) + assert m_full.partitioning.local == "inherit" + assert m_full.duplicated + + # Field comparisons (ChannelMetadata.__eq__ removed) + m2 = ChannelMetadata(local_count=4) + assert m.local_count == m2.local_count + assert m.duplicated == m2.duplicated + assert ChannelMetadata(local_count=8).local_count != m.local_count + assert "local_count=4" in repr(m) + + # Validation + with pytest.raises(ValueError, match="local_count must be non-negative"): + ChannelMetadata(local_count=-1) + + +def test_message_roundtrip() -> None: + """Test ChannelMetadata can round-trip through Message.""" + m = ChannelMetadata( + local_count=4, + partitioning=Partitioning(HashScheme((0,), 16), "inherit"), + duplicated=True, + ) + msg_m = Message(99, m) + assert msg_m.sequence_number == 99 + got_m = ChannelMetadata.from_message(msg_m) + assert got_m.local_count == 4 + assert got_m.duplicated + assert got_m.partitioning.inter_rank == HashScheme((0,), 16) + assert msg_m.empty() + + +def test_message_roundtrip_with_order_scheme(context: Context) -> None: + """Test ChannelMetadata with OrderScheme can round-trip through Message.""" + table = plc.Table( + [ + plc.Column.from_iterable_of_py( + [100, 200], plc.DataType(plc.TypeId.INT64) + ), + plc.Column.from_iterable_of_py( + ["abc", "xyz"], plc.DataType(plc.TypeId.STRING) + ), + ] + ) + boundaries = _make_boundaries(context, table) + order_scheme = OrderScheme( + [ + OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ], + boundaries, + strict_boundaries=True, + ) + m = ChannelMetadata( + local_count=8, + partitioning=Partitioning(order_scheme, "inherit"), + duplicated=True, + ) + msg_m = Message(42, m) + assert msg_m.sequence_number == 42 + got_m = ChannelMetadata.from_message(msg_m) + assert got_m.local_count == 8 + assert got_m.duplicated + assert isinstance(got_m.partitioning.inter_rank, OrderScheme) + assert got_m.partitioning.inter_rank.keys == ( + OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), + OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), + ) + assert got_m.partitioning.local == "inherit" + assert got_m.partitioning.inter_rank.strict_boundaries + assert got_m.partitioning.inter_rank.num_boundaries == 2 + assert got_m.partitioning.inter_rank.boundaries_aligned_with( + order_scheme, context.br() + ) + assert msg_m.empty() + + +def test_order_scheme_roundtrip_from_metadata(context: Context) -> None: + """An OrderScheme read back from ChannelMetadata can be re-used in a new Partitioning.""" + src = ChannelMetadata( + local_count=1, + partitioning=Partitioning(_two_key_order_scheme(context), "inherit"), + ) + scheme = src.partitioning.inter_rank + assert isinstance(scheme, OrderScheme) + + p2 = Partitioning(scheme, None) + assert isinstance(p2.inter_rank, OrderScheme) + assert p2.inter_rank.boundaries_aligned_with( + _two_key_order_scheme(context), context.br() + ) + + +def test_access_after_move_raises() -> None: + """Test that accessing a released ChannelMetadata raises ValueError.""" + m = ChannelMetadata( + local_count=4, + partitioning=Partitioning(HashScheme((0,), 16), "inherit"), + ) + # Move into a message (releases the handle) + _ = Message(0, m) + + # Accessing any property should raise ValueError + with pytest.raises(ValueError, match="uninitialized"): + _ = m.local_count + + with pytest.raises(ValueError, match="uninitialized"): + _ = m.partitioning + + with pytest.raises(ValueError, match="uninitialized"): + _ = m.duplicated + + with pytest.raises(ValueError, match="uninitialized"): + repr(m) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_partition.py b/python/cudf_streaming/cudf_streaming/tests/test_partition.py new file mode 100644 index 000000000000..15dc5562684e --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_partition.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming.streaming") + +from cudf_streaming.streaming.partition import ( + partition_and_pack, + unpack_and_concat, +) +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.streaming.core.actor import run_actor_network +from rapidsmpf.streaming.core.leaf_actor import ( + pull_from_channel, + push_to_channel, +) +from rapidsmpf.streaming.core.message import Message +from rapidsmpf.testing import assert_eq + +if TYPE_CHECKING: + from rapidsmpf.streaming.chunks.partition import PartitionMapChunk + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) +def test_partition_and_pack_unpack( + context: Context, stream: Stream, num_partitions: int +) -> None: + expects = [ + plc.Table( + [ + plc.Column.from_iterable_of_py( + [1, 2, 3], plc.DataType(plc.TypeId.INT64) + ), + plc.Column.from_iterable_of_py( + [2, 1, 1], plc.DataType(plc.TypeId.INT64) + ), + ] + ), + plc.Table( + [ + plc.Column.from_iterable_of_py( + [], plc.DataType(plc.TypeId.INT64) + ), + plc.Column.from_iterable_of_py( + [], plc.DataType(plc.TypeId.INT64) + ), + ] + ), + ] + table_chunks = [ + Message( + seq, + TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=False, br=context.br() + ), + ) + for seq, expect in enumerate(expects) + ] + ch1: Channel[TableChunk] = context.create_channel() + actor1 = push_to_channel(context, ch_out=ch1, messages=table_chunks) + + ch2: Channel[PartitionMapChunk] = context.create_channel() + actor2 = partition_and_pack( + context, + ch_in=ch1, + ch_out=ch2, + columns_to_hash=(1,), + num_partitions=num_partitions, + ) + + ch3: Channel[TableChunk] = context.create_channel() + actor3 = unpack_and_concat( + context, + ch_in=ch2, + ch_out=ch3, + ) + + actor4, output = pull_from_channel(context, ch_in=ch3) + run_actor_network(context, actors=(actor1, actor2, actor3, actor4)) + + results = output.release() + for seq, (result, expect) in enumerate(zip(results, expects, strict=True)): + assert result.sequence_number == seq + tbl = TableChunk.from_message(result, br=context.br()) + assert_eq(tbl.table_view(), expect, sort_rows=0) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py b/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py new file mode 100644 index 000000000000..2eb95e681128 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import numpy as np +import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming.streaming") + +from cudf_streaming.streaming.parquet import Filter, read_parquet +from cudf_streaming.streaming.table_chunk import TableChunk +from rapidsmpf.streaming.core.actor import run_actor_network +from rapidsmpf.streaming.core.leaf_actor import pull_from_channel + +if TYPE_CHECKING: + from typing import Literal + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.actor import CppActor + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +@pytest.fixture(scope="module") +def source( + tmp_path_factory: pytest.TempPathFactory, +) -> plc.io.SourceInfo: + path = tmp_path_factory.mktemp("read_parquet") + + nrows = 10 + start = 0 + sources = [] + for i in range(10): + table = plc.Table( + [ + plc.Column.from_array( + np.arange(start, start + nrows, dtype="int32") + ) + ] + ) + # gaps in the column numbering we produce + start += nrows + nrows // 2 + filename = path / f"{i:3d}.pq" + sink = plc.io.SinkInfo([filename]) + options = plc.io.parquet.ParquetWriterOptions.builder( + sink, table + ).build() + plc.io.parquet.write_parquet(options) + sources.append(filename) + return plc.io.SourceInfo(sources) + + +def make_filter(stream: Stream) -> plc.expressions.Expression: + return plc.expressions.Operation( + plc.expressions.ASTOperator.LESS, + plc.expressions.ColumnReference(0), + plc.expressions.Literal( + plc.Scalar.from_py( + 15, dtype=plc.DataType(plc.TypeId.INT32), stream=stream + ) + ), + ) + + +def make_producer( + context: Context, + comm: Communicator, + ch: Channel[TableChunk], + options: plc.io.parquet.ParquetReaderOptions, + *, + use_filter: bool, +) -> CppActor: + if use_filter: + fstream = context.get_stream_from_pool() + return read_parquet( + context, + comm, + ch, + 4, + options, + 3, + Filter(fstream, make_filter(fstream)), + ) + else: + return read_parquet(context, comm, ch, 4, options, 3) + + +def get_expected( + ctx: Context, + source: plc.io.SourceInfo, + skip_rows: int | Literal["none"], + num_rows: int | Literal["all"], + *, + use_filter: bool, +) -> plc.Table: + options = plc.io.parquet.ParquetReaderOptions.builder(source).build() + + if skip_rows != "none": + options.set_skip_rows(skip_rows) + if num_rows != "all": + options.set_num_rows(num_rows) + if use_filter: + fstream = ctx.get_stream_from_pool() + filter = make_filter(fstream) + fstream.synchronize() + options.set_filter(filter) + + expected = plc.io.parquet.read_parquet(options).tbl + + if use_filter: + fstream.synchronize() + return expected + + +@pytest.mark.parametrize( + "skip_rows", ["none", 7, 19, 113], ids=lambda s: f"skip_rows_{s}" +) +@pytest.mark.parametrize( + "num_rows", ["all", 0, 3, 31, 83], ids=lambda s: f"nrows_{s}" +) +@pytest.mark.parametrize("use_filter", [False, True]) +def test_read_parquet( + context: Context, + comm: Communicator, + source: plc.io.SourceInfo, + skip_rows: int | Literal["none"], + num_rows: int | Literal["all"], + use_filter: bool, +) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + + ch: Channel[TableChunk] = context.create_channel() + + options = plc.io.parquet.ParquetReaderOptions.builder(source).build() + + if skip_rows != "none": + options.set_skip_rows(skip_rows) + if num_rows != "all": + options.set_num_rows(num_rows) + + producer = make_producer(context, comm, ch, options, use_filter=use_filter) + + consumer, deferred_messages = pull_from_channel(context, ch) + + run_actor_network(context, actors=[producer, consumer]) + + messages = deferred_messages.release() + assert all( + m1.sequence_number < m2.sequence_number + for m1, m2 in itertools.pairwise(messages) + ) + chunks = [TableChunk.from_message(m, br=context.br()) for m in messages] + for chunk in chunks: + chunk.stream.synchronize() + + got = plc.concatenate.concatenate([chunk.table_view() for chunk in chunks]) + for chunk in chunks: + chunk.stream.synchronize() + + expected = get_expected( + context, source, skip_rows, num_rows, use_filter=use_filter + ) + + assert got.num_rows() == expected.num_rows() + assert got.num_columns() == expected.num_columns() + assert got.num_columns() == 1 + + all_equal = plc.reduce.reduce( + plc.binaryop.binary_operation( + got.columns()[0], + expected.columns()[0], + plc.binaryop.BinaryOperator.EQUAL, + plc.DataType(plc.TypeId.BOOL8), + ), + plc.aggregation.all(), + plc.DataType(plc.TypeId.BOOL8), + ) + assert all_equal.to_py() diff --git a/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py b/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py new file mode 100644 index 000000000000..537a774463cb --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py @@ -0,0 +1,652 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import cupy +import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming.streaming") + +from cudf_streaming.streaming.table_chunk import ( + TableChunk, + make_table_chunks_available_or_wait, +) +from rapidsmpf.cuda_stream import is_equal_streams +from rapidsmpf.memory.buffer import MemoryType +from rapidsmpf.memory.content_description import ContentDescription +from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.streaming.core.actor import define_actor, run_actor_network +from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.core.spillable_messages import SpillableMessages +from rapidsmpf.testing import assert_eq + +if TYPE_CHECKING: + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +def random_table(nbytes: int) -> plc.Table: + assert nbytes % 4 == 0 + return plc.Table( + [ + plc.Column.from_array( + cupy.random.random(nbytes // 4, dtype=cupy.float32) + ) + ] + ) + + +@pytest.mark.parametrize( + "exclusive_view", + [True, False], +) +def test_roundtrip( + context: Context, stream: Stream, *, exclusive_view: bool +) -> None: + seq = 42 + expect = random_table(1024) + table_chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=exclusive_view, br=context.br() + ) + assert is_equal_streams(table_chunk.stream, stream) + assert table_chunk.is_available() + assert table_chunk.make_available_cost() == 0 + assert table_chunk.is_spillable() == exclusive_view + assert_eq(expect, table_chunk.table_view()) + + # Message roundtrip check. + msg1 = Message(seq, table_chunk) + assert msg1.sequence_number == seq + assert msg1.get_content_description() == ContentDescription( + content_sizes={ + MemoryType.DEVICE: 1024, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=exclusive_view, + ) + + # Make a copy of msg1 in host memory. + assert msg1.copy_cost() == 1024 + res, _ = context.br().reserve( + MemoryType.HOST, 1024, allow_overbooking=True + ) + msg2 = msg1.copy(res) + assert res.size == 0 + + # msg1 is availabe + table_chunk2 = TableChunk.from_message(msg1, br=context.br()) + assert is_equal_streams(table_chunk2.stream, stream) + assert table_chunk2.is_available() + assert table_chunk2.make_available_cost() == 0 + assert_eq(expect, table_chunk2.table_view()) + + # Make a copy of msg2 back to device memory. + assert msg2.copy_cost() == 1024 + res, _ = context.br().reserve( + MemoryType.DEVICE, 1024, allow_overbooking=True + ) + msg3 = msg2.copy(res) + assert res.size == 0 + + # msg2 is on host and is not availabe + table_chunk3 = TableChunk.from_message(msg2, br=context.br()) + assert is_equal_streams(table_chunk3.stream, stream) + assert not table_chunk3.is_available() + assert table_chunk3.make_available_cost() == 1024 + # but we can make its table available using `make_available()`. + res, _ = context.br().reserve( + MemoryType.DEVICE, 1024, allow_overbooking=True + ) + table_chunk4 = table_chunk3.make_available(res) + assert is_equal_streams(table_chunk4.stream, stream) + assert table_chunk4.is_available() + assert table_chunk4.make_available_cost() == 0 + assert_eq(expect, table_chunk4.table_view()) + + # msg3 is on device (was created by copying the host msg2). During the copy this + # is made available trivially. + table_chunk5 = TableChunk.from_message(msg3, br=context.br()) + assert is_equal_streams(table_chunk5.stream, stream) + assert table_chunk5.is_available() + # and it cost no device memory to make available. + assert table_chunk5.make_available_cost() == 0 + res, _ = context.br().reserve(MemoryType.DEVICE, 0, allow_overbooking=True) + table_chunk6 = table_chunk5.make_available(res) + assert table_chunk6.is_available() + assert table_chunk6.make_available_cost() == 0 + assert_eq(expect, table_chunk6.table_view()) + + +def test_copy_roundtrip(context: Context, stream: Stream) -> None: + for nrows, ncols in [(1, 1), (1000, 100), (1, 1000)]: + expect = plc.Table( + [ + plc.Column.from_array( + cupy.random.random(nrows, dtype=cupy.float32) + ) + for _ in range(ncols) + ] + ) + + tbl1 = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + res, _ = context.br().reserve( + MemoryType.HOST, + tbl1.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + tbl2 = tbl1.copy(res) + res, _ = context.br().reserve( + MemoryType.DEVICE, + tbl2.make_available_cost(), + allow_overbooking=True, + ) + tbl3 = tbl2.make_available(res) + assert_eq(expect, tbl3.table_view()) + + +def test_spillable_messages(context: Context, stream: Stream) -> None: + seq = 42 + df1 = random_table(1024) + df2 = random_table(2048) + + sm = SpillableMessages(context.br()) + sm.insert( + Message( + seq, + TableChunk.from_pylibcudf_table( + df1, stream, exclusive_view=True, br=context.br() + ), + ) + ) + assert sm.get_content_descriptions() == { + 0: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 1024, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=True, + ) + } + sm.insert( + Message( + seq, + TableChunk.from_pylibcudf_table( + df2, stream, exclusive_view=False, br=context.br() + ), + ) + ) + assert sm.get_content_descriptions() == { + 0: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 1024, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=True, + ), + 1: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 2048, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=False, + ), + } + assert sm.spill(mid=0, br=context.br()) == 1024 + assert sm.get_content_descriptions() == { + 0: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 0, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 1024, + }, + spillable=True, + ), + 1: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 2048, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=False, + ), + } + assert sm.spill(mid=1, br=context.br()) == 0 + assert sm.get_content_descriptions() == { + 0: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 0, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 1024, + }, + spillable=True, + ), + 1: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 2048, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=False, + ), + } + + # Extract, make available, and check table chunk 1. + df1_got = TableChunk.from_message(sm.extract(mid=0), br=context.br()) + res, _ = context.br().reserve( + MemoryType.DEVICE, + df1_got.make_available_cost(), + allow_overbooking=True, + ) + df1_got = df1_got.make_available(res) + assert_eq(df1, df1_got.table_view()) + + with pytest.raises(IndexError, match="Invalid key"): + sm.extract(mid=0) + + df2_got = TableChunk.from_message(sm.extract(mid=1), br=context.br()) + df2_got = df2_got.make_available_and_spill( + context.br(), allow_overbooking=True + ) + assert_eq(df2, df2_got.table_view()) + assert sm.get_content_descriptions() == {} + + +def test_spillable_messages_by_context( + context: Context, stream: Stream +) -> None: + seq = 42 + expect = random_table(1024) + + mid = context.spillable_messages().insert( + Message( + seq, + TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ), + ) + ) + assert context.spillable_messages().get_content_descriptions() == { + 0: ContentDescription( + content_sizes={ + MemoryType.DEVICE: 1024, + MemoryType.PINNED_HOST: 0, + MemoryType.HOST: 0, + }, + spillable=True, + ) + } + got = TableChunk.from_message( + context.spillable_messages().extract(mid=mid), br=context.br() + ) + assert_eq(expect, got.table_view()) + + +def test_make_available_or_wait_already_available( + context: Context, stream: Stream +) -> None: + expect = random_table(1024) + chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + result_holder: list[TableChunk] = [] + + @define_actor() + async def test_actor(ctx: Context) -> None: + result = await chunk.make_available_or_wait(ctx, net_memory_delta=0) + result_holder.append(result) + + run_actor_network(context, actors=[test_actor(context)]) + assert_eq(expect, result_holder[0].table_view()) + + +@pytest.mark.parametrize("net_memory_delta", [0, 512]) +def test_make_available_or_wait_from_host( + context: Context, + stream: Stream, + *, + net_memory_delta: int, +) -> None: + expect = random_table(1024) + device_chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + res, _ = context.br().reserve( + MemoryType.HOST, + device_chunk.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + host_chunk = device_chunk.copy(res) + result_holder: list[TableChunk] = [] + + @define_actor() + async def test_actor(ctx: Context) -> None: + result = await host_chunk.make_available_or_wait( + ctx, net_memory_delta=net_memory_delta + ) + result_holder.append(result) + + run_actor_network(context, actors=[test_actor(context)]) + assert_eq(expect, result_holder[0].table_view()) + + +def test_data_alloc_size(context: Context, stream: Stream) -> None: + # Create a table chunk on device memory. + expect = random_table(1024) + device_chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + + # Check device memory size. + assert device_chunk.data_alloc_size(MemoryType.DEVICE) == 1024 + assert device_chunk.data_alloc_size(MemoryType.HOST) == 0 + assert device_chunk.data_alloc_size(MemoryType.PINNED_HOST) == 0 + + # Check that None returns the total across all memory types. + total_size = device_chunk.data_alloc_size(None) + assert total_size == 1024 + assert total_size == ( + device_chunk.data_alloc_size(MemoryType.DEVICE) + + device_chunk.data_alloc_size(MemoryType.HOST) + + device_chunk.data_alloc_size(MemoryType.PINNED_HOST) + ) + + # Check that calling without arguments (default None) works the same. + assert device_chunk.data_alloc_size() == 1024 + assert device_chunk.data_alloc_size() == device_chunk.data_alloc_size(None) + + # Copy to host memory and verify memory distribution. + res, _ = context.br().reserve( + MemoryType.HOST, + device_chunk.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + host_chunk = device_chunk.copy(res) + + assert host_chunk.data_alloc_size(MemoryType.DEVICE) == 0 + assert host_chunk.data_alloc_size(MemoryType.HOST) == 1024 + assert host_chunk.data_alloc_size(MemoryType.PINNED_HOST) == 0 + + # Check that None still returns the correct total. + total_size = host_chunk.data_alloc_size(None) + assert total_size == 1024 + assert total_size == ( + host_chunk.data_alloc_size(MemoryType.DEVICE) + + host_chunk.data_alloc_size(MemoryType.HOST) + + host_chunk.data_alloc_size(MemoryType.PINNED_HOST) + ) + + # Verify default parameter works after copy too. + assert host_chunk.data_alloc_size() == 1024 + assert host_chunk.data_alloc_size() == host_chunk.data_alloc_size(None) + + +@pytest.mark.parametrize( + "from_pack", [False, True], ids=["from_table", "from_pack"] +) +def test_shape_accessor( + context: Context, stream: Stream, from_pack: bool +) -> None: + nrows = 64 + expect = plc.Table( + [ + plc.Column.from_iterable_of_py( + ("abc" for _ in range(nrows)), stream=stream + ), + plc.Column.from_iterable_of_py(range(nrows), stream=stream), + ] + ) + expected_shape = (expect.num_rows(), expect.num_columns()) + + if from_pack: + pd = PackedData.from_cudf_packed_columns( + plc.contiguous_split.pack(expect, stream), stream, context.br() + ) + device_chunk = TableChunk.from_packed_data(pd, br=context.br()) + else: + device_chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + assert device_chunk.is_available() + assert device_chunk.shape == expected_shape + + res, _ = context.br().reserve( + MemoryType.HOST, + device_chunk.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + host_chunk = device_chunk.copy(res) + assert not host_chunk.is_available() + assert host_chunk.shape == expected_shape + + res, _ = context.br().reserve( + MemoryType.DEVICE, + host_chunk.make_available_cost(), + allow_overbooking=True, + ) + device_chunk = host_chunk.make_available(res) + assert device_chunk.is_available() + assert device_chunk.shape == expected_shape + + +@pytest.mark.parametrize( + "from_pack", [False, True], ids=["from_table", "from_pack"] +) +def test_into_packed_data( + context: Context, stream: Stream, from_pack: bool +) -> None: + expect = random_table(1024) + if from_pack: + pd = PackedData.from_cudf_packed_columns( + plc.contiguous_split.pack(expect, stream), stream, context.br() + ) + chunk = TableChunk.from_packed_data(pd, br=context.br()) + else: + chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + assert chunk.is_available() + + result = chunk.into_packed_data(context.br()) + assert isinstance(result, PackedData) + + # Wrap the PackedData back into a TableChunk and verify contents. + result_chunk = TableChunk.from_packed_data(result, br=context.br()) + assert result_chunk.is_available() + assert_eq(expect, result_chunk.table_view()) + + +@pytest.mark.parametrize("chunk_location", ["device", "host"]) +def test_make_table_chunks_available_or_wait_single_chunk( + context: Context, + stream: Stream, + *, + chunk_location: str, +) -> None: + expect = random_table(1024) + device_chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + + if chunk_location == "host": + res_holder, _ = context.br().reserve( + MemoryType.HOST, + device_chunk.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + chunk = device_chunk.copy(res_holder) + else: + chunk = device_chunk + + result_holder: list[tuple] = [] + + @define_actor() + async def test_actor(ctx: Context) -> None: + result_chunk, res = await make_table_chunks_available_or_wait( + ctx, chunk, reserve_extra=0, net_memory_delta=0 + ) + result_holder.append((result_chunk, res)) + + run_actor_network(context, actors=[test_actor(context)]) + chunk, res = result_holder[0] + assert chunk.is_available() + assert_eq(expect, chunk.table_view()) + # Reservation should be consumed by making the chunk available. + assert res.size == 0 + + +@pytest.mark.parametrize("num_chunks", [1, 2, 3, 5]) +def test_make_table_chunks_available_or_wait_multiple_chunks( + context: Context, + stream: Stream, + *, + num_chunks: int, +) -> None: + # Create multiple chunks with different sizes. + sizes = [1024, 2048, 512, 768, 1536][:num_chunks] + expects = [random_table(size) for size in sizes] + + # Create host chunks. + device_chunks = [ + TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + for expect in expects + ] + + host_chunks = [] + for device_chunk in device_chunks: + res, _ = context.br().reserve( + MemoryType.HOST, + device_chunk.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + host_chunks.append(device_chunk.copy(res)) + + result_holder: list[tuple] = [] + + @define_actor() + async def test_actor(ctx: Context) -> None: + chunks, res = await make_table_chunks_available_or_wait( + ctx, + host_chunks, + reserve_extra=0, + net_memory_delta=0, + ) + result_holder.append((chunks, res)) + + run_actor_network(context, actors=[test_actor(context)]) + chunks, res = result_holder[0] + assert len(chunks) == num_chunks + assert all(chunk.is_available() for chunk in chunks) + for i, expect in enumerate(expects): + assert_eq(expect, chunks[i].table_view()) + # Reservation should be consumed. + assert res.size == 0 + + +@pytest.mark.parametrize( + "reserve_extra,net_memory_delta,allow_overbooking", + [ + # Test reserve_extra variations. + (0, 0, None), + (512, 0, None), + (1024, 0, None), + # Test net_memory_delta variations. + (0, -1024, None), + (0, 512, None), + (0, 2048, None), + # Test allow_overbooking variations. + (0, 0, True), + (0, 0, False), + ], +) +def test_make_table_chunks_available_or_wait( + context: Context, + stream: Stream, + *, + reserve_extra: int, + net_memory_delta: int, + allow_overbooking: bool | None, +) -> None: + expect = random_table(1024) + device_chunk = TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=True, br=context.br() + ) + res_holder, _ = context.br().reserve( + MemoryType.HOST, + device_chunk.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + host_chunk = device_chunk.copy(res_holder) + result_holder: list[tuple] = [] + + @define_actor() + async def test_actor(ctx: Context) -> None: + chunk, res = await make_table_chunks_available_or_wait( + ctx, + host_chunk, + reserve_extra=reserve_extra, + net_memory_delta=net_memory_delta, + allow_overbooking=allow_overbooking, + ) + result_holder.append((chunk, res)) + + run_actor_network(context, actors=[test_actor(context)]) + chunk, res = result_holder[0] + assert chunk.is_available() + assert_eq(expect, chunk.table_view()) + # Reservation should have reserve_extra bytes remaining. + assert res.size == reserve_extra + + +def test_make_table_chunks_available_or_wait_mixed_availability( + context: Context, stream: Stream +) -> None: + expect1 = random_table(1024) + expect2 = random_table(2048) + + # First chunk is already available on device. + available_chunk = TableChunk.from_pylibcudf_table( + expect1, stream, exclusive_view=True, br=context.br() + ) + + # Second chunk is on host memory. + device_chunk2 = TableChunk.from_pylibcudf_table( + expect2, stream, exclusive_view=True, br=context.br() + ) + res2, _ = context.br().reserve( + MemoryType.HOST, + device_chunk2.data_alloc_size(MemoryType.DEVICE), + allow_overbooking=True, + ) + host_chunk = device_chunk2.copy(res2) + result_holder: list[tuple] = [] + + @define_actor() + async def test_actor(ctx: Context) -> None: + chunks, res = await make_table_chunks_available_or_wait( + ctx, + [available_chunk, host_chunk], + reserve_extra=0, + net_memory_delta=0, + ) + result_holder.append((chunks, res)) + + run_actor_network(context, actors=[test_actor(context)]) + chunks, res = result_holder[0] + assert len(chunks) == 2 + assert all(chunk.is_available() for chunk in chunks) + assert_eq(expect1, chunks[0].table_view()) + assert_eq(expect2, chunks[1].table_view()) + # Only the host chunk required device memory. + assert res.size == 0 From e7cb26f290d97ecc3e4d9d75b2516fb470980623 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 2 Jun 2026 05:50:42 +0000 Subject: [PATCH 05/19] feat(cudf_streaming): add packaging --- .github/workflows/build.yaml | 40 +++++++ .github/workflows/pr.yaml | 19 ++++ ci/build_wheel_cudf_streaming.sh | 58 ++++++++++ ci/run_cudf_streaming_pytests.sh | 10 ++ ci/validate_wheel.sh | 4 + .../cudf_streaming/conda_build_config.yaml | 17 +++ conda/recipes/cudf_streaming/recipe.yaml | 106 ++++++++++++++++++ dependencies.yaml | 66 +++++++++++ python/cudf_streaming/CMakeLists.txt | 4 - .../integrations/CMakeLists.txt | 11 +- .../cudf_streaming/streaming/CMakeLists.txt | 9 +- python/cudf_streaming/pyproject.toml | 2 +- 12 files changed, 334 insertions(+), 12 deletions(-) create mode 100755 ci/build_wheel_cudf_streaming.sh create mode 100755 ci/run_cudf_streaming_pytests.sh create mode 100644 conda/recipes/cudf_streaming/conda_build_config.yaml create mode 100644 conda/recipes/cudf_streaming/recipe.yaml diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index c0bfeba287b6..cf6a8c8995c1 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -217,6 +217,46 @@ jobs: date: ${{ inputs.date }} package-name: libcudf_streaming package-type: cpp + wheel-build-cudf-streaming: + needs: [telemetry-setup, wheel-build-libcudf-streaming, wheel-build-pylibcudf] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main + with: + # build for every combination of arch and CUDA version, but only for the latest Python + matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber))) + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + sha: ${{ inputs.sha }} + date: ${{ inputs.date }} + node_type: cpu16 + script: ci/build_wheel_cudf_streaming.sh + package-name: cudf_streaming + package-type: python + wheel-publish-cudf-streaming: + needs: wheel-build-cudf-streaming + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@main + secrets: + CONDA_RAPIDSAI_WHEELS_NIGHTLY_TOKEN: ${{ secrets.CONDA_RAPIDSAI_WHEELS_NIGHTLY_TOKEN }} + RAPIDSAI_PYPI_TOKEN: ${{ secrets.RAPIDSAI_PYPI_TOKEN }} + with: + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + sha: ${{ inputs.sha }} + date: ${{ inputs.date }} + package-name: cudf_streaming + package-type: python wheel-build-pylibcudf: needs: [telemetry-setup, wheel-build-libcudf] permissions: diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 80bf924f5be8..cb196ebd9bc8 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -27,6 +27,7 @@ jobs: - docs-build - wheel-build-libcudf - wheel-build-libcudf-streaming + - wheel-build-cudf-streaming - wheel-build-pylibcudf - wheel-build-cudf - wheel-tests-cudf @@ -542,6 +543,24 @@ jobs: script: "ci/build_wheel_libcudf_streaming.sh" package-name: libcudf_streaming package-type: cpp + wheel-build-cudf-streaming: + needs: [checks, wheel-build-libcudf-streaming, wheel-build-pylibcudf] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main + with: + # Build a wheel for each CUDA x ARCH x minimum supported Python version + matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) + build_type: pull-request + node_type: cpu16 + script: "ci/build_wheel_cudf_streaming.sh" + package-name: cudf_streaming + package-type: python wheel-build-pylibcudf: needs: [checks, wheel-build-libcudf] permissions: diff --git a/ci/build_wheel_cudf_streaming.sh b/ci/build_wheel_cudf_streaming.sh new file mode 100755 index 000000000000..5df50084feea --- /dev/null +++ b/ci/build_wheel_cudf_streaming.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +source rapids-init-pip + +package_name="cudf-streaming" +package_dir="python/cudf_streaming" +dependency_file_key_suffix="cudf_streaming" + +RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" + +# Downloads libcudf_streaming wheel from this current build, +# then ensures 'cudf_streaming' wheel builds always use the 'libcudf_streaming' just built in the same CI run. +LIBCUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) +echo "libcudf-streaming-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo ${LIBCUDF_STREAMING_WHEELHOUSE}/libcudf_streaming_*.whl)" >> "${PIP_CONSTRAINT}" + +rapids-logger "Generating build requirements" + +rapids-dependency-file-generator \ + --output requirements \ + --file-key "py_build_${dependency_file_key_suffix}" \ + --file-key "py_rapids_build_${dependency_file_key_suffix}" \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION};cuda_suffixed=true" \ + | tee /tmp/requirements-build.txt + +rapids-logger "Installing build requirements" +rapids-pip-retry install \ + -v \ + --prefer-binary \ + -r /tmp/requirements-build.txt + +# build with '--no-build-isolation', for better sccache hit rate +# 0 really means "add --no-build-isolation" (ref: https://github.com/pypa/pip/issues/5735) +export PIP_NO_BUILD_ISOLATION=0 + +# TODO: move this variable into `ci-wheel` +# Format Python limited API version string +RAPIDS_PY_API="cp${RAPIDS_PY_VERSION//./}" +export RAPIDS_PY_API + +./ci/build_wheel.sh "${package_name}" "${package_dir}" --stable + +# repair wheels and write to the location that artifact-uploading code expects to find them +python -m auditwheel repair \ + --exclude libcudf.so \ + --exclude libcudf_streaming.so \ + --exclude librapidsmpf.so \ + --exclude librapids_logger.so \ + --exclude librmm.so \ + --exclude libucxx.so \ + --exclude libucp.so.0 \ + -w "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" \ + ${package_dir}/dist/* + +./ci/validate_wheel.sh "${package_dir}" "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" diff --git a/ci/run_cudf_streaming_pytests.sh b/ci/run_cudf_streaming_pytests.sh new file mode 100755 index 000000000000..0c2e4b495ab2 --- /dev/null +++ b/ci/run_cudf_streaming_pytests.sh @@ -0,0 +1,10 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# Support invoking run_cudf_streaming_pytests.sh outside the script directory +cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_streaming/cudf_streaming + +pytest --cache-clear "$@" tests diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index cabf638a99f4..7100d3f2b978 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -32,6 +32,10 @@ elif [[ "${package_dir}" == "python/libcudf_streaming" ]]; then PYDISTCHECK_ARGS+=( --max-allowed-size-compressed '100M' ) +elif [[ "${package_dir}" == "python/cudf_streaming" ]]; then + PYDISTCHECK_ARGS+=( + --max-allowed-size-compressed '75M' + ) elif [[ "${package_dir}" != "python/cudf" ]] && \ [[ "${package_dir}" != "python/cudf_polars" ]] && \ [[ "${package_dir}" != "python/dask_cudf" ]] && \ diff --git a/conda/recipes/cudf_streaming/conda_build_config.yaml b/conda/recipes/cudf_streaming/conda_build_config.yaml new file mode 100644 index 000000000000..8ca00d2b4cc3 --- /dev/null +++ b/conda/recipes/cudf_streaming/conda_build_config.yaml @@ -0,0 +1,17 @@ +c_compiler_version: + - 14 + +cxx_compiler_version: + - 14 + +c_stdlib: + - sysroot + +c_stdlib_version: + - "2.28" + +cmake_version: + - ">=4.0" + +cuda_compiler: + - cuda-nvcc diff --git a/conda/recipes/cudf_streaming/recipe.yaml b/conda/recipes/cudf_streaming/recipe.yaml new file mode 100644 index 000000000000..f5859351c4f9 --- /dev/null +++ b/conda/recipes/cudf_streaming/recipe.yaml @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +schema_version: 1 + +context: + version: ${{ env.get("RAPIDS_PACKAGE_VERSION") }} + minor_version: ${{ (version | split("."))[:2] | join(".") }} + cuda_version: ${{ (env.get("RAPIDS_CUDA_VERSION") | split("."))[:2] | join(".") }} + cuda_major: '${{ (env.get("RAPIDS_CUDA_VERSION") | split("."))[0] }}' + date_string: '${{ env.get("RAPIDS_DATE_STRING") }}' + head_rev: '${{ git.head_rev(".")[:8] }}' + py_abi_min: ${{ env.get("RAPIDS_PY_VERSION") }} + py_buildstring : ${{ py_abi_min | version_to_buildstring }} + py_runtime_latest: "3.14" + +package: + name: cudf_streaming + version: ${{ version }} + +source: + path: ../../.. + +build: + python: + version_independent: true + string: cuda${{ cuda_major }}_cp${{ py_buildstring }}_abi3_${{ date_string }}_${{ head_rev }} + script: + content: | + ./build.sh cudf_streaming + secrets: + - AWS_ACCESS_KEY_ID + - AWS_SECRET_ACCESS_KEY + - AWS_SESSION_TOKEN + - SCCACHE_DIST_AUTH_TOKEN + env: + CMAKE_C_COMPILER_LAUNCHER: ${{ env.get("CMAKE_C_COMPILER_LAUNCHER") }} + CMAKE_CUDA_COMPILER_LAUNCHER: ${{ env.get("CMAKE_CUDA_COMPILER_LAUNCHER") }} + CMAKE_CXX_COMPILER_LAUNCHER: ${{ env.get("CMAKE_CXX_COMPILER_LAUNCHER") }} + CMAKE_GENERATOR: ${{ env.get("CMAKE_GENERATOR") }} + NVCC_APPEND_FLAGS: ${{ env.get("NVCC_APPEND_FLAGS", default="") }} + PARALLEL_LEVEL: ${{ env.get("PARALLEL_LEVEL", default="8") }} + RAPIDS_ARTIFACTS_DIR: ${{ env.get("RAPIDS_ARTIFACTS_DIR", default="") }} + RAPIDS_PY_VERSION: ${{ env.get("RAPIDS_PY_VERSION", default="") }} + SCCACHE_BUCKET: ${{ env.get("SCCACHE_BUCKET", default="") }} + SCCACHE_DIST_AUTH_TYPE: ${{ env.get("SCCACHE_DIST_AUTH_TYPE", default="token") }} + SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE: ${{ env.get("SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE", default="false") }} + SCCACHE_DIST_MAX_RETRIES: ${{ env.get("SCCACHE_DIST_MAX_RETRIES", default="inf") }} + SCCACHE_DIST_REQUEST_TIMEOUT: ${{ env.get("SCCACHE_DIST_REQUEST_TIMEOUT", default="7140") }} + SCCACHE_DIST_SCHEDULER_URL: ${{ env.get("SCCACHE_DIST_SCHEDULER_URL", default="") }} + SCCACHE_ERROR_LOG: ${{ env.get("SCCACHE_ERROR_LOG", default="/tmp/sccache.log") }} + SCCACHE_IDLE_TIMEOUT: ${{ env.get("SCCACHE_IDLE_TIMEOUT", default="0") }} + SCCACHE_NO_CACHE: ${{ env.get("SCCACHE_NO_CACHE", default="") }} + SCCACHE_RECACHE: ${{ env.get("SCCACHE_RECACHE", default="") }} + SCCACHE_REGION: ${{ env.get("SCCACHE_REGION", default="") }} + SCCACHE_S3_KEY_PREFIX: cudf-streaming-${{ env.get("RAPIDS_CONDA_ARCH") }} + SCCACHE_S3_NO_CREDENTIALS: ${{ env.get("SCCACHE_S3_NO_CREDENTIALS", default="false") }} + SCCACHE_S3_PREPROCESSOR_CACHE_KEY_PREFIX: cudf-streaming-${{ env.get("RAPIDS_CONDA_ARCH") }}-cuda${{ cuda_major }}-conda-preprocessor-cache + SCCACHE_S3_USE_PREPROCESSOR_CACHE_MODE: ${{ env.get("SCCACHE_S3_USE_PREPROCESSOR_CACHE_MODE", default="true") }} + SCCACHE_S3_USE_SSL: ${{ env.get("SCCACHE_S3_USE_SSL", default="true") }} + SCCACHE_SERVER_LOG: ${{ env.get("SCCACHE_SERVER_LOG", default="sccache=debug") }} + +requirements: + build: + - cmake ${{ cmake_version }} + - ninja + - ${{ compiler("c") }} + - ${{ compiler("cxx") }} + - ${{ compiler("cuda") }} + - cuda-version =${{ cuda_version }} + - ${{ stdlib("c") }} + host: + - cython >=3.2.2 + - pip + - python =${{ py_abi_min }} + - python-abi3 ${{ py_abi_min }}.* + - cuda-version =${{ cuda_version }} + - pylibcudf =${{ version }} + - libcudf_streaming =${{ minor_version }} + - rapidsmpf =${{ minor_version }} + - rapids-build-backend >=0.4.0,<0.5.0 + - scikit-build-core>=0.11.0 + - cuda-cudart-dev + run: + - python + - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} + - libcudf_streaming =${{ minor_version }} + - pylibcudf =${{ version }} + - rapidsmpf =${{ minor_version }} + - cuda-cudart + ignore_run_exports: + from_package: + - cuda-cudart-dev + by_name: + - cuda-version + +tests: + - python: + imports: + - cudf_streaming + python_version: ${{ py_runtime_latest }}.* + pip_check: false + +about: + homepage: ${{ load_from_file("python/cudf_streaming/pyproject.toml").project.urls.Homepage }} + license: ${{ load_from_file("python/cudf_streaming/pyproject.toml").project.license }} + summary: ${{ load_from_file("python/cudf_streaming/pyproject.toml").project.description }} diff --git a/dependencies.yaml b/dependencies.yaml index c81fd85595de..3172e871857b 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -496,6 +496,47 @@ files: - depends_on_libcudf - depends_on_librmm - depends_on_librapidsmpf + py_build_cudf_streaming: + output: pyproject + pyproject_dir: python/cudf_streaming + extras: + table: build-system + includes: + - rapids_build_skbuild + py_rapids_build_cudf_streaming: + output: pyproject + pyproject_dir: python/cudf_streaming + extras: + table: tool.rapids-build-backend + key: requires + includes: + - build_base + - build_cpp + - build_python_common + - depends_on_libcudf_streaming + - depends_on_librapidsmpf + - depends_on_librmm + - depends_on_pylibcudf + - depends_on_rapidsmpf + - depends_on_rmm + py_run_cudf_streaming: + output: pyproject + pyproject_dir: python/cudf_streaming + extras: + table: project + includes: + - depends_on_libcudf_streaming + - depends_on_pylibcudf + - depends_on_rapidsmpf + - depends_on_rmm + py_test_cudf_streaming: + output: pyproject + pyproject_dir: python/cudf_streaming + extras: + table: project.optional-dependencies + key: test + includes: + - test_python_common test_python_narwhals: output: none includes: @@ -1340,6 +1381,31 @@ dependencies: - matrix: packages: - *librapidsmpf_unsuffixed + depends_on_libcudf_streaming: + common: + - output_types: conda + packages: + - &libcudf_streaming_unsuffixed libcudf_streaming==26.8.*,>=0.0.0a0 + - output_types: requirements + packages: + # pip recognizes the index as a global option for the requirements.txt file + - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple + specific: + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + cuda_suffixed: "true" + packages: + - libcudf-streaming-cu12==26.8.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - libcudf-streaming-cu13==26.8.*,>=0.0.0a0 + - matrix: + packages: + - *libcudf_streaming_unsuffixed depends_on_ray: common: - output_types: conda diff --git a/python/cudf_streaming/CMakeLists.txt b/python/cudf_streaming/CMakeLists.txt index 793e64b33bdc..024964265129 100644 --- a/python/cudf_streaming/CMakeLists.txt +++ b/python/cudf_streaming/CMakeLists.txt @@ -21,9 +21,6 @@ project( LANGUAGES CXX CUDA ) -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED ON) - # For now, disable CMake's automatic module scanning for C++ files. There is an sccache bug in the # version RAPIDS uses in CI that causes it to handle the resulting -M* flags incorrectly with # gcc>=14. We can remove this once we upgrade to a newer sccache version. @@ -34,7 +31,6 @@ set(CMAKE_CXX_SCAN_FOR_MODULES OFF) set_property(GLOBAL PROPERTY FIND_LIBRARY_USE_LIB64_PATHS TRUE) find_package(cudf_streaming "${RAPIDS_VERSION}" REQUIRED) -find_package(rapidsmpf REQUIRED) find_package(CUDAToolkit REQUIRED) include(rapids-cython-core) diff --git a/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt index c6c1ba807667..c430608c26e8 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt +++ b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt @@ -4,11 +4,14 @@ # cmake-format: on set(cython_sources partition.pyx) -set(linked_libraries cudf_streaming::cudf_streaming rapidsmpf::rapidsmpf) +set(linked_libraries cudf_streaming::cudf_streaming) rapids_cython_create_modules( - CXX ASSOCIATED_TARGETS cudf_streaming rapidsmpf + CXX ASSOCIATED_TARGETS cudf_streaming SOURCE_FILES "${cython_sources}" - LINKED_LIBRARIES "${linked_libraries}" - MODULE_PREFIX cudf_streaming_integrations_ + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX cudf_streaming_integrations_ ) + +foreach(target IN LISTS RAPIDS_CYTHON_CREATED_TARGETS) + set_target_properties(${target} PROPERTIES CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON) +endforeach() diff --git a/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt index 1428c8f3305e..10d4362bead9 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt +++ b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt @@ -4,11 +4,14 @@ # cmake-format: on set(cython_sources bloom_filter.pyx channel_metadata.pyx parquet.pyx partition.pyx table_chunk.pyx) -set(linked_libraries cudf_streaming::cudf_streaming rapidsmpf::rapidsmpf) +set(linked_libraries cudf_streaming::cudf_streaming) rapids_cython_create_modules( CXX ASSOCIATED_TARGETS cudf_streaming SOURCE_FILES "${cython_sources}" - LINKED_LIBRARIES "${linked_libraries}" - MODULE_PREFIX cudf_streaming_streaming_ + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX cudf_streaming_streaming_ ) + +foreach(target IN LISTS RAPIDS_CYTHON_CREATED_TARGETS) + set_target_properties(${target} PROPERTIES CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON) +endforeach() diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index c73b4cee7f7f..81764d22216e 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -9,7 +9,7 @@ requires = [ ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. [project] -name = "cudf_streaming" +name = "cudf-streaming" dynamic = ["version"] description = "cuDF Streaming library" authors = [ From b8379bbda83e1666733045bc6893e423771b0b29 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 3 Jun 2026 03:32:16 +0000 Subject: [PATCH 06/19] fix(cudf_streaming): add README.md for twine validation twine requires long_description metadata. Add readme field to pyproject.toml and a README.md file. --- python/cudf_streaming/README.md | 10 ++++++++++ python/cudf_streaming/pyproject.toml | 1 + 2 files changed, 11 insertions(+) create mode 100644 python/cudf_streaming/README.md diff --git a/python/cudf_streaming/README.md b/python/cudf_streaming/README.md new file mode 100644 index 000000000000..1af9635ab96f --- /dev/null +++ b/python/cudf_streaming/README.md @@ -0,0 +1,10 @@ +# cudf_streaming + +cudf_streaming provides Python/Cython bindings for libcudf_streaming, enabling GPU-accelerated streaming data processing built on top of libcudf and rapidsmpf. + +## Installation + +```bash +pip install cudf-streaming-cu12 # For CUDA 12 +pip install cudf-streaming-cu13 # For CUDA 13 +``` diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index 81764d22216e..bab6499f29c0 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -12,6 +12,7 @@ requires = [ name = "cudf-streaming" dynamic = ["version"] description = "cuDF Streaming library" +readme = { file = "README.md", content-type = "text/markdown" } authors = [ { name = "NVIDIA Corporation" }, ] From 65037309c2df6f626ec7a85e030cb2e51d4495e1 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 3 Jun 2026 06:00:04 +0000 Subject: [PATCH 07/19] feat(cudf_streaming): add wheel and conda test jobs - Add ci/test_wheel_cudf_streaming.sh for wheel-based pytest - Add wheel-tests-cudf-streaming job to pr.yaml - Add cudf_streaming pytest to ci/test_python_other.sh (conda tests) --- .github/workflows/pr.yaml | 15 +++++++++++ ci/test_python_other.sh | 6 +++++ ci/test_wheel_cudf_streaming.sh | 38 ++++++++++++++++++++++++++++ dependencies.yaml | 1 + python/cudf_streaming/pyproject.toml | 1 + 5 files changed, 61 insertions(+) create mode 100755 ci/test_wheel_cudf_streaming.sh diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index cb196ebd9bc8..031e26e4a048 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -28,6 +28,7 @@ jobs: - wheel-build-libcudf - wheel-build-libcudf-streaming - wheel-build-cudf-streaming + - wheel-tests-cudf-streaming - wheel-build-pylibcudf - wheel-build-cudf - wheel-tests-cudf @@ -561,6 +562,20 @@ jobs: script: "ci/build_wheel_cudf_streaming.sh" package-name: cudf_streaming package-type: python + wheel-tests-cudf-streaming: + needs: [wheel-build-cudf-streaming, changed-files] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main + if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels + with: + build_type: pull-request + script: ci/test_wheel_cudf_streaming.sh wheel-build-pylibcudf: needs: [checks, wheel-build-libcudf] permissions: diff --git a/ci/test_python_other.sh b/ci/test_python_other.sh index 9147ad1a8761..4e38e539da6e 100755 --- a/ci/test_python_other.sh +++ b/ci/test_python_other.sh @@ -54,5 +54,11 @@ rapids-logger "pytest cudf-polars" --durations=10 --durations-min=10 \ -ra +rapids-logger "pytest cudf_streaming" +timeout 30m ./ci/run_cudf_streaming_pytests.sh \ + --junitxml="${RAPIDS_TESTS_DIR}/junit-cudf-streaming.xml" \ + --numprocesses=8 \ + --dist=worksteal + rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/ci/test_wheel_cudf_streaming.sh b/ci/test_wheel_cudf_streaming.sh new file mode 100755 index 000000000000..e62acabe8933 --- /dev/null +++ b/ci/test_wheel_cudf_streaming.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +source rapids-init-pip + +RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" + +# Download cudf_streaming, libcudf_streaming, and pylibcudf built in previous steps +CUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github python) +LIBCUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) +LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) +PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")") + +# generate constraints (possibly pinning to oldest support versions of dependencies) +rapids-generate-pip-constraints py_test_cudf_streaming "${PIP_CONSTRAINT}" + +rapids-logger "Install cudf_streaming and its dependencies" + +rapids-pip-retry install \ + -v \ + --prefer-binary \ + --constraint "${PIP_CONSTRAINT}" \ + "$(echo "${CUDF_STREAMING_WHEELHOUSE}"/cudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]" \ + "$(echo "${LIBCUDF_STREAMING_WHEELHOUSE}"/libcudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ + "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \ + "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" + +rapids-logger "pytest cudf_streaming" +pushd python/cudf_streaming/cudf_streaming +timeout 30m python -m pytest \ + --cache-clear \ + --numprocesses=8 \ + --dist=worksteal \ + tests +popd diff --git a/dependencies.yaml b/dependencies.yaml index 3172e871857b..e58b02a671ff 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -537,6 +537,7 @@ files: key: test includes: - test_python_common + - depends_on_cupy test_python_narwhals: output: none includes: diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index bab6499f29c0..e477988e9e84 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -27,6 +27,7 @@ dependencies = [ [project.optional-dependencies] test = [ + "cupy-cuda13x>=13.6.0,!=14.0.0,!=14.1.0", "pytest", "pytest-cov", "pytest-xdist", From 748cb08977541735f823ff19631d809ad90b16eb Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 3 Jun 2026 17:35:40 +0000 Subject: [PATCH 08/19] fix(cudf_streaming): fix test import path and fixture resolution - Use importmode=importlib in pytest config to prevent source tree shadowing the installed package in site-packages. - Change test scripts to cd into tests/ directory and run pytest on '.' (consistent with pylibcudf wheel test pattern). - Conditionally import rapidsmpf.tests.conftest fixtures (comm, stream) when available (conda), with fallback skip-fixtures for wheel envs. - Handle pytest exit code 5 (all tests skipped) in wheel test script since communicator support may not be available. --- ci/run_cudf_streaming_pytests.sh | 4 +-- ci/test_wheel_cudf_streaming.sh | 11 ++++++-- .../cudf_streaming/tests/conftest.py | 26 ++++++++++++++++--- python/cudf_streaming/pyproject.toml | 2 +- 4 files changed, 34 insertions(+), 9 deletions(-) diff --git a/ci/run_cudf_streaming_pytests.sh b/ci/run_cudf_streaming_pytests.sh index 0c2e4b495ab2..8005c02e58fc 100755 --- a/ci/run_cudf_streaming_pytests.sh +++ b/ci/run_cudf_streaming_pytests.sh @@ -5,6 +5,6 @@ set -euo pipefail # Support invoking run_cudf_streaming_pytests.sh outside the script directory -cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_streaming/cudf_streaming +cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_streaming/cudf_streaming/tests -pytest --cache-clear "$@" tests +pytest --cache-clear "$@" . diff --git a/ci/test_wheel_cudf_streaming.sh b/ci/test_wheel_cudf_streaming.sh index e62acabe8933..62766eb43080 100755 --- a/ci/test_wheel_cudf_streaming.sh +++ b/ci/test_wheel_cudf_streaming.sh @@ -29,10 +29,17 @@ rapids-pip-retry install \ "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" rapids-logger "pytest cudf_streaming" -pushd python/cudf_streaming/cudf_streaming +pushd python/cudf_streaming/cudf_streaming/tests +EXITCODE=0 timeout 30m python -m pytest \ --cache-clear \ --numprocesses=8 \ --dist=worksteal \ - tests + . || EXITCODE=$? + +# Exit code 5 means no tests were collected (all skipped); acceptable when +# communicator support (MPI/UCXX) is unavailable in the wheel test environment. +if [ ${EXITCODE} -ne 0 ] && [ ${EXITCODE} -ne 5 ]; then + exit ${EXITCODE} +fi popd diff --git a/python/cudf_streaming/cudf_streaming/tests/conftest.py b/python/cudf_streaming/cudf_streaming/tests/conftest.py index 8fc1023afa2c..b8646628e562 100644 --- a/python/cudf_streaming/cudf_streaming/tests/conftest.py +++ b/python/cudf_streaming/cudf_streaming/tests/conftest.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import sys -from pathlib import Path +import importlib.util from typing import TYPE_CHECKING import pytest @@ -14,13 +13,19 @@ from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.streaming.core.context import Context -sys.path.insert(0, str(Path(__file__).parents[2])) - if TYPE_CHECKING: from collections.abc import Generator from rapidsmpf.communicator.communicator import Communicator +# Import fixtures (comm, device_mr, etc.) from rapidsmpf's test conftest +# if available (conda installs include tests; wheel installs may not). +_HAS_RAPIDSMPF_TEST_FIXTURES = ( + importlib.util.find_spec("rapidsmpf.tests.conftest") is not None +) +if _HAS_RAPIDSMPF_TEST_FIXTURES: + pytest_plugins = ["rapidsmpf.tests.conftest"] + @pytest.fixture def context(comm: Communicator) -> Generator[Context, None, None]: @@ -33,3 +38,16 @@ def context(comm: Communicator) -> Generator[Context, None, None]: with Context(comm.logger, br, options) as ctx: yield ctx + + +if not _HAS_RAPIDSMPF_TEST_FIXTURES: + + @pytest.fixture(params=["mpi", "ucxx"]) + def comm(request): + """Fallback comm fixture that skips when rapidsmpf test infra is unavailable.""" + pytest.skip("rapidsmpf test fixtures not installed") + + @pytest.fixture + def stream(): + """Fallback stream fixture that skips when rapidsmpf test infra is unavailable.""" + pytest.skip("rapidsmpf test fixtures not installed") diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index e477988e9e84..d8ef2661817f 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -61,7 +61,7 @@ select = [ max_allowed_size_compressed = '75M' [tool.pytest.ini_options] -addopts = "--tb=native --strict-config --strict-markers" +addopts = "--tb=native --strict-config --strict-markers --import-mode=importlib" empty_parameter_set_mark = "fail_at_collect" filterwarnings = [ "error", From 2a9be05edb24e0e8b7d2aa5283747c3a79edaa7b Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 3 Jun 2026 18:36:57 +0000 Subject: [PATCH 09/19] fix(cudf_streaming): add MPI support for conda tests - Add depends_on_cudf_streaming to test_python_other so the conda test environment installs cudf_streaming and its transitive deps (rapidsmpf, mpi4py, openmpi, ucxx). - Use mpirun in run_cudf_streaming_pytests.sh to initialize MPI, which is required by the rapidsmpf comm fixture (MPI.COMM_WORLD.barrier()). - Set OMPI_ALLOW_RUN_AS_ROOT and OMPI_MCA_opal_cuda_support env vars matching rapidsmpf's CI configuration. - Remove pytest-xdist flags from cudf_streaming in test_python_other.sh since xdist is incompatible with MPI process management. --- ci/run_cudf_streaming_pytests.sh | 10 +++++++++- ci/test_python_other.sh | 4 +--- dependencies.yaml | 6 ++++++ 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/ci/run_cudf_streaming_pytests.sh b/ci/run_cudf_streaming_pytests.sh index 8005c02e58fc..102ac539e6ba 100755 --- a/ci/run_cudf_streaming_pytests.sh +++ b/ci/run_cudf_streaming_pytests.sh @@ -7,4 +7,12 @@ set -euo pipefail # Support invoking run_cudf_streaming_pytests.sh outside the script directory cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_streaming/cudf_streaming/tests -pytest --cache-clear "$@" . +# OpenMPI specific options (CI runs as root) +export OMPI_ALLOW_RUN_AS_ROOT=1 +export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 +export OMPI_MCA_opal_cuda_support=1 + +# cudf_streaming tests require MPI for the communicator fixtures. +# Run with mpirun; currently single-rank only tests exist. +mpirun --map-by node --bind-to none -np 1 \ + python -m pytest --cache-clear "$@" . diff --git a/ci/test_python_other.sh b/ci/test_python_other.sh index 4e38e539da6e..57404b76ebb0 100755 --- a/ci/test_python_other.sh +++ b/ci/test_python_other.sh @@ -56,9 +56,7 @@ rapids-logger "pytest cudf-polars" rapids-logger "pytest cudf_streaming" timeout 30m ./ci/run_cudf_streaming_pytests.sh \ - --junitxml="${RAPIDS_TESTS_DIR}/junit-cudf-streaming.xml" \ - --numprocesses=8 \ - --dist=worksteal + --junitxml="${RAPIDS_TESTS_DIR}/junit-cudf-streaming.xml" rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/dependencies.yaml b/dependencies.yaml index e58b02a671ff..4756dc05fd82 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -118,6 +118,7 @@ files: - depends_on_custreamz - depends_on_cudf_polars - depends_on_ray + - depends_on_cudf_streaming test_java: output: none includes: @@ -1407,6 +1408,11 @@ dependencies: - matrix: packages: - *libcudf_streaming_unsuffixed + depends_on_cudf_streaming: + common: + - output_types: conda + packages: + - cudf_streaming==26.8.*,>=0.0.0a0 depends_on_ray: common: - output_types: conda From 5b66b8d1bc61249f7cdf714405780399e6590d57 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 3 Jun 2026 22:13:51 +0000 Subject: [PATCH 10/19] fix(cudf_streaming): build conda package in CI and fix conftest import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add cudf_streaming to ci/build_python.sh so the conda package is available in the test_python_other environment. - Use hyphenated conda package names (cudf-streaming, libcudf-streaming) consistent with the cpp branch changes. - Replace importlib.util.find_spec with try/except import in conftest.py to handle the case where rapidsmpf.tests is partially present in the wheel but not fully importable (exit code 4 → graceful fallback). --- ci/build_python.sh | 10 ++++++++++ conda/recipes/cudf_streaming/recipe.yaml | 6 +++--- dependencies.yaml | 2 +- .../cudf_streaming/cudf_streaming/tests/conftest.py | 11 ++++++----- 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/ci/build_python.sh b/ci/build_python.sh index 8e66d73df95b..75e00f5d3f14 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -63,6 +63,16 @@ rapids-telemetry-record build-cudf_kafka.log \ rapids-telemetry-record sccache-stats-cudf_kafka.txt sccache --show-adv-stats sccache --stop-server >/dev/null 2>&1 || true +rapids-logger "Building cudf_streaming" + +rapids-telemetry-record build-cudf_streaming.log \ + rattler-build build --recipe conda/recipes/cudf_streaming \ + "${RATTLER_ARGS[@]}" \ + "${RATTLER_CHANNELS[@]}" + +rapids-telemetry-record sccache-stats-cudf_streaming.txt sccache --show-adv-stats +sccache --stop-server >/dev/null 2>&1 || true + # remove build_cache directory rm -rf "$RAPIDS_CONDA_BLD_OUTPUT_DIR"/build_cache diff --git a/conda/recipes/cudf_streaming/recipe.yaml b/conda/recipes/cudf_streaming/recipe.yaml index f5859351c4f9..2484991a82ea 100644 --- a/conda/recipes/cudf_streaming/recipe.yaml +++ b/conda/recipes/cudf_streaming/recipe.yaml @@ -14,7 +14,7 @@ context: py_runtime_latest: "3.14" package: - name: cudf_streaming + name: cudf-streaming version: ${{ version }} source: @@ -75,7 +75,7 @@ requirements: - python-abi3 ${{ py_abi_min }}.* - cuda-version =${{ cuda_version }} - pylibcudf =${{ version }} - - libcudf_streaming =${{ minor_version }} + - libcudf-streaming =${{ minor_version }} - rapidsmpf =${{ minor_version }} - rapids-build-backend >=0.4.0,<0.5.0 - scikit-build-core>=0.11.0 @@ -83,7 +83,7 @@ requirements: run: - python - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - - libcudf_streaming =${{ minor_version }} + - libcudf-streaming =${{ minor_version }} - pylibcudf =${{ version }} - rapidsmpf =${{ minor_version }} - cuda-cudart diff --git a/dependencies.yaml b/dependencies.yaml index 4756dc05fd82..306b8a6dd950 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -1412,7 +1412,7 @@ dependencies: common: - output_types: conda packages: - - cudf_streaming==26.8.*,>=0.0.0a0 + - cudf-streaming==26.8.*,>=0.0.0a0 depends_on_ray: common: - output_types: conda diff --git a/python/cudf_streaming/cudf_streaming/tests/conftest.py b/python/cudf_streaming/cudf_streaming/tests/conftest.py index b8646628e562..da7af13ba80b 100644 --- a/python/cudf_streaming/cudf_streaming/tests/conftest.py +++ b/python/cudf_streaming/cudf_streaming/tests/conftest.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import importlib.util from typing import TYPE_CHECKING import pytest @@ -20,11 +19,13 @@ # Import fixtures (comm, device_mr, etc.) from rapidsmpf's test conftest # if available (conda installs include tests; wheel installs may not). -_HAS_RAPIDSMPF_TEST_FIXTURES = ( - importlib.util.find_spec("rapidsmpf.tests.conftest") is not None -) -if _HAS_RAPIDSMPF_TEST_FIXTURES: +try: + import rapidsmpf.tests.conftest # noqa: F401 + + _HAS_RAPIDSMPF_TEST_FIXTURES = True pytest_plugins = ["rapidsmpf.tests.conftest"] +except (ImportError, ModuleNotFoundError): + _HAS_RAPIDSMPF_TEST_FIXTURES = False @pytest.fixture From 9cffff7697637339142f48a37eb588f7ec96a053 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 16:06:21 +0000 Subject: [PATCH 11/19] fix(cudf_streaming): add build.sh target and fix wheel artifact resolution - Add 'cudf_streaming' to build.sh VALIDARGS and build section so conda recipe's './build.sh cudf_streaming' invocation succeeds. - Switch ci/test_wheel_cudf_streaming.sh to use rapids-download-from-github with rapids-package-name --stable for abi3 artifact resolution, matching how cudf handles Python-version-independent wheel downloads. --- build.sh | 15 +++++++++++++-- ci/test_wheel_cudf_streaming.sh | 2 +- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/build.sh b/build.sh index 3482b0fa48e0..50d7037039ec 100755 --- a/build.sh +++ b/build.sh @@ -18,8 +18,8 @@ ARGS=$* # script, and that this script resides in the repo dir! REPODIR=$(cd "$(dirname "$0")"; pwd) -VALIDARGS="clean libcudf pylibcudf cudf cudf_polars dask_cudf benchmarks tests libcudf_kafka cudf_kafka custreamz libcudf_streaming -v -g -n --pydevelop -l --allgpuarch --disable_nvtx --opensource_nvcomp --show_depr_warn --ptds -h --build_metrics --incl_cache_stats --disable_large_strings" -HELP="$0 [clean] [libcudf] [pylibcudf] [cudf] [cudf_polars] [dask_cudf] [benchmarks] [tests] [libcudf_kafka] [cudf_kafka] [custreamz] [libcudf_streaming] [-v] [-g] [-n] [-h] [--cmake-args=\\\"\\\"] +VALIDARGS="clean libcudf pylibcudf cudf cudf_polars dask_cudf benchmarks tests libcudf_kafka cudf_kafka custreamz libcudf_streaming cudf_streaming -v -g -n --pydevelop -l --allgpuarch --disable_nvtx --opensource_nvcomp --show_depr_warn --ptds -h --build_metrics --incl_cache_stats --disable_large_strings" +HELP="$0 [clean] [libcudf] [pylibcudf] [cudf] [cudf_polars] [dask_cudf] [benchmarks] [tests] [libcudf_kafka] [cudf_kafka] [custreamz] [libcudf_streaming] [cudf_streaming] [-v] [-g] [-n] [-h] [--cmake-args=\\\"\\\"] clean - remove all existing build artifacts and configuration (start over) libcudf - build the cudf C++ code only @@ -33,6 +33,7 @@ HELP="$0 [clean] [libcudf] [pylibcudf] [cudf] [cudf_polars] [dask_cudf] [benchma cudf_kafka - build the cudf_kafka Python package custreamz - build the custreamz Python package libcudf_streaming - build the libcudf_streaming C++ code only + cudf_streaming - build the cudf_streaming Python package -v - verbose build mode -g - build for debug -n - no install step (does not affect Python) @@ -367,3 +368,13 @@ if hasArg libcudf_streaming; then cmake --build . -j"${PARALLEL_LEVEL}" --target install ${VERBOSE_FLAG} fi fi + +# build cudf_streaming Python package +if hasArg cudf_streaming; then + cd "${REPODIR}/python/cudf_streaming" + SKBUILD_CMAKE_ARGS="-DCMAKE_PREFIX_PATH=${INSTALL_PREFIX};-DCMAKE_LIBRARY_PATH=${STREAMING_LIB_BUILD_DIR};${EXTRA_CMAKE_ARGS[*]}" \ + python -m pip install \ + "${PYTHON_ARGS_FOR_INSTALL[@]}" \ + "${PY_API_ARGS[@]}" \ + . +fi diff --git a/ci/test_wheel_cudf_streaming.sh b/ci/test_wheel_cudf_streaming.sh index 62766eb43080..faa92e86b830 100755 --- a/ci/test_wheel_cudf_streaming.sh +++ b/ci/test_wheel_cudf_streaming.sh @@ -9,7 +9,7 @@ source rapids-init-pip RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" # Download cudf_streaming, libcudf_streaming, and pylibcudf built in previous steps -CUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github python) +CUDF_STREAMING_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" cudf_streaming --stable --cuda "$RAPIDS_CUDA_VERSION")") LIBCUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")") From 7a01161d6437e25c97d1602cc4a66674310746a4 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 17:31:01 +0000 Subject: [PATCH 12/19] Set wheel upload name correctly --- ci/build_wheel_cudf_streaming.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ci/build_wheel_cudf_streaming.sh b/ci/build_wheel_cudf_streaming.sh index 5df50084feea..d485f4741ac6 100755 --- a/ci/build_wheel_cudf_streaming.sh +++ b/ci/build_wheel_cudf_streaming.sh @@ -56,3 +56,6 @@ python -m auditwheel repair \ ${package_dir}/dist/* ./ci/validate_wheel.sh "${package_dir}" "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" + +RAPIDS_PACKAGE_NAME="$(rapids-package-name wheel_python cudf_streaming --stable --cuda)" +export RAPIDS_PACKAGE_NAME From 2ddc2a837a4a8cdc262c85fe7b0d5415081ca57b Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 17:46:41 +0000 Subject: [PATCH 13/19] Remove incorrect importlib setting --- python/cudf_streaming/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index d8ef2661817f..e477988e9e84 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -61,7 +61,7 @@ select = [ max_allowed_size_compressed = '75M' [tool.pytest.ini_options] -addopts = "--tb=native --strict-config --strict-markers --import-mode=importlib" +addopts = "--tb=native --strict-config --strict-markers" empty_parameter_set_mark = "fail_at_collect" filterwarnings = [ "error", From 74811d39a72712132dd631500fe642e1c9d48212 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 17:53:45 +0000 Subject: [PATCH 14/19] fix(cudf_streaming): inline test fixtures from rapidsmpf Copy comm, stream, and supporting fixtures directly into our conftest instead of importing rapidsmpf.tests.conftest. The rapidsmpf test module is not installed in wheel environments and the try/import/fallback logic added unnecessary complexity. The fixtures are simple and self-contained. --- .../cudf_streaming/tests/conftest.py | 82 +++++++++++++------ 1 file changed, 58 insertions(+), 24 deletions(-) diff --git a/python/cudf_streaming/cudf_streaming/tests/conftest.py b/python/cudf_streaming/cudf_streaming/tests/conftest.py index da7af13ba80b..6eb6cf7fc546 100644 --- a/python/cudf_streaming/cudf_streaming/tests/conftest.py +++ b/python/cudf_streaming/cudf_streaming/tests/conftest.py @@ -7,48 +7,82 @@ import pytest import rmm.mr +from rapidsmpf.communicator import COMMUNICATORS from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.progress_thread import ProgressThread from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.streaming.core.context import Context +from rmm.pylibrmm.stream import DEFAULT_STREAM if TYPE_CHECKING: from collections.abc import Generator from rapidsmpf.communicator.communicator import Communicator + from rmm.pylibrmm.stream import Stream -# Import fixtures (comm, device_mr, etc.) from rapidsmpf's test conftest -# if available (conda installs include tests; wheel installs may not). -try: - import rapidsmpf.tests.conftest # noqa: F401 - _HAS_RAPIDSMPF_TEST_FIXTURES = True - pytest_plugins = ["rapidsmpf.tests.conftest"] -except (ImportError, ModuleNotFoundError): - _HAS_RAPIDSMPF_TEST_FIXTURES = False +@pytest.fixture(scope="session") +def _mpi_comm() -> Communicator: + """Session-wide MPI communicator.""" + if "mpi" not in COMMUNICATORS: + pytest.skip("RapidsMPF not built with MPI support") + + from mpi4py import MPI + + from rapidsmpf.communicator.mpi import new_communicator + + return new_communicator( + MPI.COMM_WORLD, Options(get_environment_variables()), ProgressThread() + ) + + +@pytest.fixture(scope="session") +def _ucxx_comm() -> Communicator: + """Session-wide UCXX communicator.""" + if "ucxx" not in COMMUNICATORS: + pytest.skip("RapidsMPF not built with UCXX support") + if "mpi" not in COMMUNICATORS: + pytest.skip("MPI required to bootstrap UCXX test communicator") + + from rapidsmpf.communicator.testing import ucxx_mpi_setup + + return ucxx_mpi_setup( + None, Options(get_environment_variables()), ProgressThread() + ) + + +@pytest.fixture(params=["mpi", "ucxx"]) +def comm( + request: pytest.FixtureRequest, +) -> Generator[Communicator, None, None]: + """Communicator fixture parametrized over MPI and UCXX transports.""" + comm_name = request.param + + if "mpi" not in COMMUNICATORS: + pytest.skip("RapidsMPF not built with MPI support") + if "ucxx" not in COMMUNICATORS: + pytest.skip("RapidsMPF not built with UCXX support") + + from mpi4py import MPI + + MPI.COMM_WORLD.barrier() + yield request.getfixturevalue(f"_{comm_name}_comm") + MPI.COMM_WORLD.barrier() + + +@pytest.fixture +def stream() -> Stream: + """CUDA stream for test operations.""" + return DEFAULT_STREAM @pytest.fixture def context(comm: Communicator) -> Generator[Context, None, None]: - """ - Fixture to get a streaming context. - """ + """Streaming context backed by a fresh memory resource.""" options = Options(get_environment_variables()) mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) br = BufferResource(mr) with Context(comm.logger, br, options) as ctx: yield ctx - - -if not _HAS_RAPIDSMPF_TEST_FIXTURES: - - @pytest.fixture(params=["mpi", "ucxx"]) - def comm(request): - """Fallback comm fixture that skips when rapidsmpf test infra is unavailable.""" - pytest.skip("rapidsmpf test fixtures not installed") - - @pytest.fixture - def stream(): - """Fallback stream fixture that skips when rapidsmpf test infra is unavailable.""" - pytest.skip("rapidsmpf test fixtures not installed") From 19a6e9befee6225f118d653f9e45ffdfa9f1cd77 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 18:25:59 +0000 Subject: [PATCH 15/19] fix(cudf_streaming): use hyphenated conda package name for libcudf-streaming Use 'libcudf-streaming' instead of 'libcudf_streaming' in conda package specs in dependencies.yaml for precision. The dependency file generator propagates the fix to pyproject.toml as well. --- dependencies.yaml | 2 +- python/cudf_streaming/pyproject.toml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dependencies.yaml b/dependencies.yaml index 306b8a6dd950..3f68f56d5b7f 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -1387,7 +1387,7 @@ dependencies: common: - output_types: conda packages: - - &libcudf_streaming_unsuffixed libcudf_streaming==26.8.*,>=0.0.0a0 + - &libcudf_streaming_unsuffixed libcudf-streaming==26.8.*,>=0.0.0a0 - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index e477988e9e84..acee2abbbdc9 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -19,7 +19,7 @@ authors = [ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ - "libcudf_streaming==26.8.*,>=0.0.0a0", + "libcudf-streaming==26.8.*,>=0.0.0a0", "pylibcudf==26.8.*,>=0.0.0a0", "rapidsmpf==26.8.*,>=0.0.0a0", "rmm==26.8.*,>=0.0.0a0", @@ -76,7 +76,7 @@ matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true" requires = [ "cmake>=4.0", "cython>=3.2.2", - "libcudf_streaming==26.8.*,>=0.0.0a0", + "libcudf-streaming==26.8.*,>=0.0.0a0", "librapidsmpf==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "ninja", From e3adb1d2bf7527e1b792fcd9780a8e88dbcff5c0 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 18:26:36 +0000 Subject: [PATCH 16/19] fix(cudf_streaming): address review feedback - build.yaml: use min_by for stable-ABI wheel builds (match pr.yaml) - parquet.pyi: fix Filter.__init__ parameter name (expression -> filter) - partition.pyx: fix docstring reference to partition_and_pack - channel_metadata.pxd: use except +ex_handler consistently --- .github/CODEOWNERS | 1 + .github/workflows/build.yaml | 4 ++-- .../cudf_streaming/integrations/partition.pyx | 2 ++ .../streaming/channel_metadata.pxd | 20 ++++++++++--------- .../cudf_streaming/streaming/parquet.pyi | 2 +- .../cudf_streaming/streaming/partition.pyx | 2 +- 6 files changed, 18 insertions(+), 13 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 5bd69b7cf36d..489b2b065280 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -10,6 +10,7 @@ python/ @rapidsai/cudf-python-codeowners notebooks/ @rapidsai/cudf-python-codeowners python/dask_cudf/ @rapidsai/cudf-dask-codeowners python/cudf_polars/ @rapidsai/cudf-polars-codeowners +python/cudf_streaming/ @rapidsai/rapidsmpf-python-codeowners #cmake code owners CMakeLists.txt @rapidsai/cudf-cmake-codeowners diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index cf6a8c8995c1..4e6046c6ae4b 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -228,8 +228,8 @@ jobs: secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main with: - # build for every combination of arch and CUDA version, but only for the latest Python - matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber))) + # Build a wheel for each CUDA x ARCH x minimum supported Python version + matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} sha: ${{ inputs.sha }} diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pyx b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx index bdf5cd67be3c..cc607d841934 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/partition.pyx +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx @@ -354,6 +354,8 @@ cpdef object unspill_partitions( ReservationError If overbooking exceeds the amount spilled and ``allow_overbooking is False``. """ + if not isinstance(allow_overbooking, bool): + raise TypeError("allow_overbooking must be a bool") cdef cpp_BufferResource* _br = br.ptr() cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) cdef vector[cpp_PackedData] _ret diff --git a/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd index f5118f8bbc34..8b3115790bf0 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd +++ b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd @@ -11,6 +11,8 @@ from pylibcudf.libcudf.types cimport null_order as cpp_null_order from pylibcudf.libcudf.types cimport order as cpp_order from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rapidsmpf._detail.exception_handling cimport ex_handler + from rapidsmpf.memory.buffer_resource cimport cpp_BufferResource from rapidsmpf.streaming.core.message cimport cpp_Message from cudf_streaming.streaming.table_chunk cimport TableChunk, cpp_TableChunk @@ -22,8 +24,8 @@ cdef extern from "" \ cdef cppclass cpp_HashScheme "cudf_streaming::streaming::HashScheme": vector[int32_t] column_indices int modulus - cpp_HashScheme() except + - cpp_HashScheme(vector[int32_t], int) except + + cpp_HashScheme() except +ex_handler + cpp_HashScheme(vector[int32_t], int) except +ex_handler bool_t operator==(const cpp_HashScheme&) cdef cppclass cpp_OrderKey "cudf_streaming::streaming::OrderKey": @@ -38,14 +40,14 @@ cdef extern from "" \ cpp_OrderScheme() noexcept cpp_OrderScheme( vector[cpp_OrderKey], unique_ptr[cpp_TableChunk], bool_t - ) except + + ) except +ex_handler vector[cpp_OrderKey] keys shared_ptr[cpp_TableChunk] boundaries bool_t strict_boundaries - cpp_OrderScheme with_keys(vector[cpp_OrderKey]) except + + cpp_OrderScheme with_keys(vector[cpp_OrderKey]) except +ex_handler bool_t boundaries_aligned_with( const cpp_OrderScheme&, const cpp_BufferResource& - ) except + + ) except +ex_handler cdef cppclass cpp_PartitioningSpec "cudf_streaming::streaming::PartitioningSpec": enum cpp_Type "cudf_streaming::streaming::PartitioningSpec::Type": @@ -73,8 +75,8 @@ cdef extern from "" \ cdef cppclass cpp_Partitioning "cudf_streaming::streaming::Partitioning": cpp_PartitioningSpec inter_rank cpp_PartitioningSpec local - cpp_Partitioning() except + - cpp_Partitioning(const cpp_Partitioning&) except + + cpp_Partitioning() except +ex_handler + cpp_Partitioning(const cpp_Partitioning&) except +ex_handler cdef cppclass cpp_ChannelMetadata "cudf_streaming::streaming::ChannelMetadata": uint64_t local_count @@ -84,12 +86,12 @@ cdef extern from "" \ uint64_t, cpp_Partitioning, bool_t - ) except + + ) except +ex_handler cpp_Message cpp_to_message_channel_metadata \ "cudf_streaming::streaming::to_message"( uint64_t, unique_ptr[cpp_ChannelMetadata] - ) except + + ) except +ex_handler cdef class HashScheme: diff --git a/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi b/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi index 7a393e1650e6..2f181e0970c9 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi +++ b/python/cudf_streaming/cudf_streaming/streaming/parquet.pyi @@ -12,7 +12,7 @@ from rapidsmpf.streaming.core.context import Context from rmm.pylibrmm.stream import Stream class Filter: - def __init__(self, stream: Stream, expression: Expression) -> None: ... + def __init__(self, stream: Stream, filter: Expression) -> None: ... def read_parquet( ctx: Context, diff --git a/python/cudf_streaming/cudf_streaming/streaming/partition.pyx b/python/cudf_streaming/cudf_streaming/streaming/partition.pyx index 1df59390a2b5..58a190a11ff1 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/partition.pyx +++ b/python/cudf_streaming/cudf_streaming/streaming/partition.pyx @@ -46,7 +46,7 @@ def partition_and_pack( Asynchronously partition and pack table chunks. This is the streaming equivalent of - :func:`cudf_streaming.integrations.partition.partition_and_split()`, + :func:`cudf_streaming.integrations.partition.partition_and_pack()`, operating on incoming table chunks via channels. Each incoming table from `ch_in` is partitioned into `num_partitions` outputs From 80d31c8b08324c86ab9c12a0106bf67f5cd930ea Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 18:40:24 +0000 Subject: [PATCH 17/19] fix(cudf_streaming): address reviewer feedback - Remove unnecessary pytest.importorskip for cudf_streaming.streaming (always available when running these tests) - Add numpy to test dependencies since tests use it directly --- dependencies.yaml | 1 + python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py | 2 -- .../cudf_streaming/tests/test_channel_metadata.py | 2 -- python/cudf_streaming/cudf_streaming/tests/test_partition.py | 2 -- python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py | 2 -- python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py | 2 -- python/cudf_streaming/pyproject.toml | 1 + 7 files changed, 2 insertions(+), 10 deletions(-) diff --git a/dependencies.yaml b/dependencies.yaml index 3f68f56d5b7f..ad6c95788d90 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -539,6 +539,7 @@ files: includes: - test_python_common - depends_on_cupy + - numpy_run test_python_narwhals: output: none includes: diff --git a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py index c0d7f718aa6b..cdf1d36a3a51 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py @@ -10,8 +10,6 @@ import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming.streaming") - from cudf_streaming.streaming import ChannelMetadata from cudf_streaming.streaming.bloom_filter import BloomFilter from cudf_streaming.streaming.table_chunk import TableChunk diff --git a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py index ce5c456db706..ce526cc2801c 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -9,8 +9,6 @@ import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming.streaming") - from cudf_streaming.streaming import ( ChannelMetadata, HashScheme, diff --git a/python/cudf_streaming/cudf_streaming/tests/test_partition.py b/python/cudf_streaming/cudf_streaming/tests/test_partition.py index 15dc5562684e..f0948a7666c6 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_partition.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_partition.py @@ -8,8 +8,6 @@ import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming.streaming") - from cudf_streaming.streaming.partition import ( partition_and_pack, unpack_and_concat, diff --git a/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py b/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py index 2eb95e681128..df8b80af9e2c 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py @@ -10,8 +10,6 @@ import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming.streaming") - from cudf_streaming.streaming.parquet import Filter, read_parquet from cudf_streaming.streaming.table_chunk import TableChunk from rapidsmpf.streaming.core.actor import run_actor_network diff --git a/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py b/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py index 537a774463cb..f48a83e84131 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py @@ -9,8 +9,6 @@ import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming.streaming") - from cudf_streaming.streaming.table_chunk import ( TableChunk, make_table_chunks_available_or_wait, diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index acee2abbbdc9..fc31a16dae1c 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ [project.optional-dependencies] test = [ "cupy-cuda13x>=13.6.0,!=14.0.0,!=14.1.0", + "numpy>=1.26,<3.0", "pytest", "pytest-cov", "pytest-xdist", From 4efd65d6248252e61bd30679515abc084a42f49a Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 20:41:29 +0000 Subject: [PATCH 18/19] Make version a symlink --- python/cudf_streaming/cudf_streaming/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) mode change 100644 => 120000 python/cudf_streaming/cudf_streaming/VERSION diff --git a/python/cudf_streaming/cudf_streaming/VERSION b/python/cudf_streaming/cudf_streaming/VERSION deleted file mode 100644 index 4e6864b4ca93..000000000000 --- a/python/cudf_streaming/cudf_streaming/VERSION +++ /dev/null @@ -1 +0,0 @@ -26.08.00 diff --git a/python/cudf_streaming/cudf_streaming/VERSION b/python/cudf_streaming/cudf_streaming/VERSION new file mode 120000 index 000000000000..d62dc733efdf --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/VERSION @@ -0,0 +1 @@ +../../../VERSION \ No newline at end of file From 32ca1673193c37a8e22120e460c028fed5fc87a6 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 4 Jun 2026 20:56:43 +0000 Subject: [PATCH 19/19] Address review: remove pxd docstring, add coroutine bridging comment --- .../cudf_streaming/streaming/bloom_filter.pxd | 14 -------------- .../cudf_streaming/streaming/bloom_filter.pyx | 4 ++++ 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd index 3e4e09654828..789a409cd0b8 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd @@ -28,19 +28,5 @@ cdef extern from "" nogil: cdef class BloomFilter: - """ - Streaming bloom filter construction and application. - - Parameters - ---------- - ctx - Streaming context. - comm - The communicator the bloom filter construction is collective over. - seed - Seed used for hashing values into the bloom filter. - num_filter_blocks - Number of blocks used to size the filter. - """ cdef unique_ptr[cpp_BloomFilter] _handle cdef Communicator _comm diff --git a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx index b43613f9d83d..8f9ed4bbd9d0 100644 --- a/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx @@ -192,6 +192,10 @@ cdef class BloomFilter: tag Disambiguating tag to combine filters across ranks. """ + # Coroutine bridging pattern: create a Python future, transfer + # ownership to C++ via OwningWrapper (Py_INCREF here, py_deleter + # calls Py_DECREF when the C++ task completes), then spawn the + # C++ coroutine which resolves the future via cpp_set_py_future. ret = asyncio.get_running_loop().create_future() Py_INCREF(ret) with nogil: