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 c0bfeba287b6..4e6046c6ae4b 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 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 }} + 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..031e26e4a048 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -27,6 +27,8 @@ jobs: - docs-build - wheel-build-libcudf - wheel-build-libcudf-streaming + - wheel-build-cudf-streaming + - wheel-tests-cudf-streaming - wheel-build-pylibcudf - wheel-build-cudf - wheel-tests-cudf @@ -542,6 +544,38 @@ 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-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/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/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/ci/build_wheel_cudf_streaming.sh b/ci/build_wheel_cudf_streaming.sh new file mode 100755 index 000000000000..d485f4741ac6 --- /dev/null +++ b/ci/build_wheel_cudf_streaming.sh @@ -0,0 +1,61 @@ +#!/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}" + +RAPIDS_PACKAGE_NAME="$(rapids-package-name wheel_python cudf_streaming --stable --cuda)" +export RAPIDS_PACKAGE_NAME diff --git a/ci/run_cudf_streaming_pytests.sh b/ci/run_cudf_streaming_pytests.sh new file mode 100755 index 000000000000..102ac539e6ba --- /dev/null +++ b/ci/run_cudf_streaming_pytests.sh @@ -0,0 +1,18 @@ +#!/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/tests + +# 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 9147ad1a8761..57404b76ebb0 100755 --- a/ci/test_python_other.sh +++ b/ci/test_python_other.sh @@ -54,5 +54,9 @@ 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" + 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..faa92e86b830 --- /dev/null +++ b/ci/test_wheel_cudf_streaming.sh @@ -0,0 +1,45 @@ +#!/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-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")") + +# 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/tests +EXITCODE=0 +timeout 30m python -m pytest \ + --cache-clear \ + --numprocesses=8 \ + --dist=worksteal \ + . || 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/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..2484991a82ea --- /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..ad6c95788d90 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: @@ -496,6 +497,49 @@ 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 + - depends_on_cupy + - numpy_run test_python_narwhals: output: none includes: @@ -1340,6 +1384,36 @@ 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_cudf_streaming: + common: + - output_types: conda + packages: + - cudf-streaming==26.8.*,>=0.0.0a0 depends_on_ray: common: - output_types: conda diff --git a/python/cudf_streaming/CMakeLists.txt b/python/cudf_streaming/CMakeLists.txt new file mode 100644 index 000000000000..024964265129 --- /dev/null +++ b/python/cudf_streaming/CMakeLists.txt @@ -0,0 +1,40 @@ +# ============================================================================= +# 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 +) + +# 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(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/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/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 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..c430608c26e8 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/CMakeLists.txt @@ -0,0 +1,17 @@ +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on + +set(cython_sources partition.pyx) +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_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/integrations/__init__.pxd b/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd new file mode 100644 index 000000000000..c60c8a61b990 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/__init__.pxd @@ -0,0 +1,4 @@ +# 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 new file mode 100644 index 000000000000..743c0b3ec454 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/__init__.py @@ -0,0 +1,7 @@ +# 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..cc607d841934 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx @@ -0,0 +1,371 @@ +# 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``. + """ + 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 + 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) 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..10d4362bead9 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/CMakeLists.txt @@ -0,0 +1,17 @@ +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# 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) + +rapids_cython_create_modules( + CXX ASSOCIATED_TARGETS cudf_streaming + SOURCE_FILES "${cython_sources}" + 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/cudf_streaming/streaming/__init__.pxd b/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd new file mode 100644 index 000000000000..2c2cb7e1959f --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/__init__.pxd @@ -0,0 +1,18 @@ +# 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 new file mode 100644 index 000000000000..4f5d020b6fac --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/__init__.py @@ -0,0 +1,36 @@ +# 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..789a409cd0b8 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pxd @@ -0,0 +1,32 @@ +# 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: + 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..8f9ed4bbd9d0 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/bloom_filter.pyx @@ -0,0 +1,251 @@ +# 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. + """ + # 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: + 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..8b3115790bf0 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/streaming/channel_metadata.pxd @@ -0,0 +1,133 @@ +# 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._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 + + +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 +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": + 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 +ex_handler + vector[cpp_OrderKey] keys + shared_ptr[cpp_TableChunk] boundaries + bool_t strict_boundaries + cpp_OrderScheme with_keys(vector[cpp_OrderKey]) except +ex_handler + bool_t boundaries_aligned_with( + const cpp_OrderScheme&, const cpp_BufferResource& + ) except +ex_handler + + 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 +ex_handler + cpp_Partitioning(const cpp_Partitioning&) except +ex_handler + + 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 +ex_handler + + cpp_Message cpp_to_message_channel_metadata \ + "cudf_streaming::streaming::to_message"( + uint64_t, unique_ptr[cpp_ChannelMetadata] + ) except +ex_handler + + +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..2f181e0970c9 --- /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, filter: 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..58a190a11ff1 --- /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_pack()`, + 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 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..6eb6cf7fc546 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/conftest.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING + +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 + + +@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]: + """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 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..cdf1d36a3a51 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py @@ -0,0 +1,174 @@ +# 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 + +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..ce526cc2801c --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -0,0 +1,417 @@ +# 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 + +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..f0948a7666c6 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_partition.py @@ -0,0 +1,92 @@ +# 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 + +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..df8b80af9e2c --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py @@ -0,0 +1,183 @@ +# 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 + +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..f48a83e84131 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_table_chunk.py @@ -0,0 +1,650 @@ +# 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 + +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 diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml new file mode 100644 index 000000000000..fc31a16dae1c --- /dev/null +++ b/python/cudf_streaming/pyproject.toml @@ -0,0 +1,103 @@ +# 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" +readme = { file = "README.md", content-type = "text/markdown" } +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 = [ + "cupy-cuda13x>=13.6.0,!=14.0.0,!=14.1.0", + "numpy>=1.26,<3.0", + "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.*)"