diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ff2b5e218db6..5bd69b7cf36d 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,6 +2,9 @@ cpp/ @rapidsai/cudf-cpp-codeowners python/cudf/udf_cpp/ @rapidsai/cudf-cpp-codeowners +#libcudf_streaming code owners (order matters, this supersedes cpp/ for files in this directory) +cpp/libcudf_streaming/ @rapidsai/rapidsmpf-cpp-codeowners + #python code owners python/ @rapidsai/cudf-python-codeowners notebooks/ @rapidsai/cudf-python-codeowners @@ -12,6 +15,9 @@ python/cudf_polars/ @rapidsai/cudf-polars-codeowners CMakeLists.txt @rapidsai/cudf-cmake-codeowners **/cmake/ @rapidsai/cudf-cmake-codeowners *.cmake @rapidsai/cudf-cmake-codeowners +cpp/libcudf_streaming/CMakeLists.txt @rapidsai/rapidsmpf-cmake-codeowners +cpp/libcudf_streaming/**/cmake/ @rapidsai/rapidsmpf-cmake-codeowners +cpp/libcudf_streaming/*.cmake @rapidsai/rapidsmpf-cmake-codeowners #java code owners java/ @rapidsai/cudf-java-codeowners diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index c1d215601cec..c0bfeba287b6 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -177,6 +177,46 @@ jobs: date: ${{ inputs.date }} package-name: libcudf package-type: cpp + wheel-build-libcudf-streaming: + needs: [telemetry-setup, wheel-build-libcudf] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main + with: + # build for every combination of arch and CUDA version, but only for the latest Python + matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber))) + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + sha: ${{ inputs.sha }} + date: ${{ inputs.date }} + node_type: cpu16 + script: ci/build_wheel_libcudf_streaming.sh + package-name: libcudf_streaming + package-type: cpp + wheel-publish-libcudf-streaming: + needs: wheel-build-libcudf-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: libcudf_streaming + package-type: cpp wheel-build-pylibcudf: needs: [telemetry-setup, wheel-build-libcudf] permissions: diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index bcf12deeba83..80bf924f5be8 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -26,6 +26,7 @@ jobs: - conda-notebook-tests - docs-build - wheel-build-libcudf + - wheel-build-libcudf-streaming - wheel-build-pylibcudf - wheel-build-cudf - wheel-tests-cudf @@ -523,6 +524,24 @@ jobs: script: "ci/build_wheel_libcudf.sh" package-name: libcudf package-type: cpp + wheel-build-libcudf-streaming: + needs: [checks, wheel-build-libcudf] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main + with: + # build for every combination of arch and CUDA version, but only for the latest Python + matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber))) + build_type: pull-request + node_type: cpu16 + script: "ci/build_wheel_libcudf_streaming.sh" + package-name: libcudf_streaming + package-type: cpp wheel-build-pylibcudf: needs: [checks, wheel-build-libcudf] permissions: diff --git a/build.sh b/build.sh index e948e9dd3624..3482b0fa48e0 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 -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] [-v] [-g] [-n] [-h] [--cmake-args=\\\"\\\"] +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=\\\"\\\"] clean - remove all existing build artifacts and configuration (start over) libcudf - build the cudf C++ code only @@ -32,6 +32,7 @@ HELP="$0 [clean] [libcudf] [pylibcudf] [cudf] [cudf_polars] [dask_cudf] [benchma libcudf_kafka - build the libcudf_kafka C++ code only cudf_kafka - build the cudf_kafka Python package custreamz - build the custreamz Python package + libcudf_streaming - build the libcudf_streaming C++ code only -v - verbose build mode -g - build for debug -n - no install step (does not affect Python) @@ -51,6 +52,7 @@ HELP="$0 [clean] [libcudf] [pylibcudf] [cudf] [cudf_polars] [dask_cudf] [benchma " LIB_BUILD_DIR=${LIB_BUILD_DIR:=${REPODIR}/cpp/build} KAFKA_LIB_BUILD_DIR=${KAFKA_LIB_BUILD_DIR:=${REPODIR}/cpp/libcudf_kafka/build} +STREAMING_LIB_BUILD_DIR=${STREAMING_LIB_BUILD_DIR:=${REPODIR}/cpp/libcudf_streaming/build} CUDF_KAFKA_BUILD_DIR=${REPODIR}/python/cudf_kafka/build CUDF_BUILD_DIR=${REPODIR}/python/cudf/build DASK_CUDF_BUILD_DIR=${REPODIR}/python/dask_cudf/build @@ -58,7 +60,7 @@ PYLIBCUDF_BUILD_DIR=${REPODIR}/python/pylibcudf/build CUSTREAMZ_BUILD_DIR=${REPODIR}/python/custreamz/build CUDF_JAR_JAVA_BUILD_DIR="$REPODIR/java/target" -BUILD_DIRS="${LIB_BUILD_DIR} ${CUDF_BUILD_DIR} ${DASK_CUDF_BUILD_DIR} ${KAFKA_LIB_BUILD_DIR} ${CUDF_KAFKA_BUILD_DIR} ${CUSTREAMZ_BUILD_DIR} ${CUDF_JAR_JAVA_BUILD_DIR} ${PYLIBCUDF_BUILD_DIR}" +BUILD_DIRS="${LIB_BUILD_DIR} ${CUDF_BUILD_DIR} ${DASK_CUDF_BUILD_DIR} ${KAFKA_LIB_BUILD_DIR} ${CUDF_KAFKA_BUILD_DIR} ${CUSTREAMZ_BUILD_DIR} ${CUDF_JAR_JAVA_BUILD_DIR} ${PYLIBCUDF_BUILD_DIR} ${STREAMING_LIB_BUILD_DIR}" # Set defaults for vars modified by flags to this script VERBOSE_FLAG="" @@ -199,7 +201,7 @@ fi ################################################################################ # Configure, build, and install libcudf -if buildAll || hasArg libcudf || hasArg pylibcudf || hasArg cudf ; then +if buildAll || hasArg libcudf || hasArg libcudf_streaming || hasArg pylibcudf || hasArg cudf ; then if (( BUILD_ALL_GPU_ARCH == 0 )); then CUDF_CMAKE_CUDA_ARCHITECTURES="${CUDF_CMAKE_CUDA_ARCHITECTURES:-NATIVE}" if [[ "$CUDF_CMAKE_CUDA_ARCHITECTURES" == "NATIVE" ]]; then @@ -347,3 +349,21 @@ if hasArg custreamz; then cd "${REPODIR}/python/custreamz" python -m pip install "${PYTHON_ARGS_FOR_INSTALL[@]}" . fi + +# Build libcudf_streaming library +if hasArg libcudf_streaming; then + cmake -S "$REPODIR/cpp/libcudf_streaming" -B "${STREAMING_LIB_BUILD_DIR}" \ + -DCMAKE_INSTALL_PREFIX="${INSTALL_PREFIX}" \ + -DCMAKE_CUDA_ARCHITECTURES="${CUDF_CMAKE_CUDA_ARCHITECTURES}" \ + -DBUILD_TESTS=${BUILD_TESTS} \ + -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \ + "${EXTRA_CMAKE_ARGS[@]}" + + + cd "${STREAMING_LIB_BUILD_DIR}" + cmake --build . -j"${PARALLEL_LEVEL}" ${VERBOSE_FLAG} + + if [[ ${INSTALL_TARGET} != "" ]]; then + cmake --build . -j"${PARALLEL_LEVEL}" --target install ${VERBOSE_FLAG} + fi +fi diff --git a/ci/build_wheel_libcudf_streaming.sh b/ci/build_wheel_libcudf_streaming.sh new file mode 100755 index 000000000000..90df0e51a9e2 --- /dev/null +++ b/ci/build_wheel_libcudf_streaming.sh @@ -0,0 +1,44 @@ +#!/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="libcudf_streaming" +package_dir="python/libcudf_streaming" + +rapids-logger "Generating build requirements" + +rapids-dependency-file-generator \ + --output requirements \ + --file-key "py_build_${package_name}" \ + --file-key "py_rapids_build_${package_name}" \ + --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 + +./ci/build_wheel.sh "${package_name}" "${package_dir}" + +# repair wheels and write to the location that artifact-uploading code expects to find them +python -m auditwheel repair \ + --exclude libcudf.so \ + --exclude librapidsmpf.so \ + --exclude librapids_logger.so \ + --exclude librmm.so \ + --exclude libucxx.so \ + --exclude libucp.so.0 \ + -w "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" \ + ${package_dir}/dist/* + +./ci/validate_wheel.sh "${package_dir}" "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index 5e8b28c06c2d..4cc7ee84ccef 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -137,13 +137,14 @@ fi # Centralized version file update echo "${NEXT_FULL_TAG}" > VERSION +# The cudf version file must be a copy, see https://github.com/rapidsai/cudf/pull/18198 echo "${NEXT_FULL_TAG}" > python/cudf/cudf/VERSION echo "${RAPIDS_BRANCH_NAME}" > RAPIDS_BRANCH DEPENDENCIES=( cudf - cudf_kafka cudf-polars + cudf_kafka cugraph cuml custreamz @@ -152,8 +153,8 @@ DEPENDENCIES=( kvikio libcudf libcudf-example - libcudf_kafka libcudf-tests + libcudf_kafka libkvikio librmm pylibcudf diff --git a/ci/run_cudf_streaming_ctests.sh b/ci/run_cudf_streaming_ctests.sh new file mode 100755 index 000000000000..8ed8593ff6ed --- /dev/null +++ b/ci/run_cudf_streaming_ctests.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# Support customizing the ctests' install location +# First, try the installed location (CI/conda environments) +installed_test_location="${INSTALL_PREFIX:-${CONDA_PREFIX:-/usr}}/bin/gtests/libcudf_streaming/" +# Fall back to the build directory (devcontainer environments) +script_dir="$(dirname "$(realpath "${BASH_SOURCE[0]}")")" +devcontainers_test_location="${script_dir}/../cpp/libcudf_streaming/build/latest" +buildsh_test_location="${STREAMING_LIB_BUILD_DIR:-${script_dir}/../cpp/libcudf_streaming/build}" + +if [[ -d "${installed_test_location}" ]]; then + cd "${installed_test_location}" +elif [[ -d "${devcontainers_test_location}" ]]; then + cd "${devcontainers_test_location}" +elif [[ -d "${buildsh_test_location}" ]]; then + cd "${buildsh_test_location}" +else + echo "Error: Test location not found. Searched:" >&2 + echo " - ${installed_test_location}" >&2 + echo " - ${devcontainers_test_location}" >&2 + echo " - ${buildsh_test_location}" >&2 + exit 1 +fi + +ctest --output-on-failure --no-tests=error "$@" diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh index 4ae71ff3081c..3d723fe8be3f 100755 --- a/ci/test_cpp.sh +++ b/ci/test_cpp.sh @@ -32,5 +32,11 @@ if (( SUITEERROR == 0 )); then SUITEERROR=$? fi +if (( SUITEERROR == 0 )); then + rapids-logger "Run libcudf_streaming gtests" + timeout 30m ./ci/run_cudf_streaming_ctests.sh -j20 + SUITEERROR=$? +fi + rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index 14d7121260d2..cabf638a99f4 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -28,6 +28,10 @@ if [[ "${package_dir}" == "python/libcudf" ]]; then --max-allowed-size-compressed '350M' ) fi +elif [[ "${package_dir}" == "python/libcudf_streaming" ]]; then + PYDISTCHECK_ARGS+=( + --max-allowed-size-compressed '100M' + ) elif [[ "${package_dir}" != "python/cudf" ]] && \ [[ "${package_dir}" != "python/cudf_polars" ]] && \ [[ "${package_dir}" != "python/dask_cudf" ]] && \ diff --git a/conda/recipes/libcudf/recipe.yaml b/conda/recipes/libcudf/recipe.yaml index c2e3fff245b8..6077750b497f 100644 --- a/conda/recipes/libcudf/recipe.yaml +++ b/conda/recipes/libcudf/recipe.yaml @@ -38,7 +38,7 @@ cache: done ./build.sh -n -v \ - libcudf libcudf_kafka benchmarks tests \ + libcudf libcudf_kafka libcudf_streaming benchmarks tests \ --build_metrics --incl_cache_stats --allgpuarch \ --cmake-args="\"-DCUDF_ENABLE_ARROW_S3=ON\" $cmake_args" secrets: @@ -83,6 +83,7 @@ cache: - ${{ stdlib("c") }} host: - librmm =${{ minor_version }} + - librapidsmpf =${{ minor_version }} - libkvikio =${{ minor_version }} - cuda-cudart-dev - cuda-nvrtc-dev @@ -216,6 +217,53 @@ outputs: license: Apache-2.0 summary: libcudf_kafka library + - package: + name: libcudf-streaming + version: ${{ version }} + build: + script: + - cmake --install cpp/libcudf_streaming/build + string: cuda${{ cuda_major }}_${{ date_string }}_${{ head_rev }} + dynamic_linking: + overlinking_behavior: "error" + requirements: + build: + - cmake ${{ cmake_version }} + - ${{ stdlib("c") }} + host: + - ${{ pin_subpackage("libcudf", exact=True) }} + - librmm =${{ minor_version }} + - librapidsmpf =${{ minor_version }} + - cuda-version =${{ cuda_version }} + - cuda-cudart-dev + run: + - ${{ pin_subpackage("libcudf", exact=True) }} + - librmm =${{ minor_version }} + - librapidsmpf =${{ minor_version }} + - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} + ignore_run_exports: + by_name: + - cuda-cudart + - cuda-nvrtc + - cuda-nvtx + - cuda-version + - flatbuffers + - libcudf + - libcufile + - libcurand + - libkvikio + - librdkafka + - librmm + - libzlib + - libnvcomp + tests: + - script: + - test -f $PREFIX/lib/libcudf_streaming.so + about: + homepage: https://rapids.ai/ + license: Apache-2.0 + summary: libcudf-streaming library + - package: name: libcudf-example version: ${{ version }} @@ -301,6 +349,7 @@ outputs: script: - cmake --install cpp/build --component testing - cmake --install cpp/libcudf_kafka/build --component testing + - cmake --install cpp/libcudf_streaming/build --component testing string: cuda${{ cuda_major }}_${{ date_string }}_${{ head_rev }} dynamic_linking: overlinking_behavior: "error" @@ -313,6 +362,7 @@ outputs: host: - ${{ pin_subpackage("libcudf", exact=True) }} - ${{ pin_subpackage("libcudf_kafka", exact=True) }} + - ${{ pin_subpackage("libcudf-streaming", exact=True) }} - cuda-version =${{ cuda_version }} - libcurand-dev - cuda-cudart-dev @@ -321,6 +371,7 @@ outputs: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - ${{ pin_subpackage("libcudf", exact=True) }} - ${{ pin_subpackage("libcudf_kafka", exact=True) }} + - ${{ pin_subpackage("libcudf-streaming", exact=True) }} - libcurand - cuda-sanitizer-api ignore_run_exports: diff --git a/cpp/cmake/thirdparty/get_cudf.cmake b/cpp/cmake/thirdparty/get_cudf.cmake new file mode 100644 index 000000000000..02e9ce9e86d0 --- /dev/null +++ b/cpp/cmake/thirdparty/get_cudf.cmake @@ -0,0 +1,49 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +#[=======================================================================[.rst: +find_and_configure_cudf +----------------------- + +Find cudf and set any additional necessary environment variables. + +``VERSION`` + The cudf version to find. + +``EXPORT_SET`` + The CMake export set name to associate with this dependency. + +Result Variables +^^^^^^^^^^^^^^^^ + + :cmake:variable:`_SOURCE_DIR` is set to the path to the source directory of . + :cmake:variable:`_BINARY_DIR` is set to the path to the build directory of . + :cmake:variable:`_ADDED` is set to a true value if has not been added before. + +#]=======================================================================] +# cmake-lint: disable=C0112 +function(find_and_configure_cudf VERSION EXPORT_SET) + rapids_cmake_parse_version(MAJOR_MINOR ${VERSION} major_minor) + rapids_cpm_find( + cudf ${VERSION} + BUILD_EXPORT_SET ${EXPORT_SET} + INSTALL_EXPORT_SET ${EXPORT_SET} + CPM_ARGS + GIT_REPOSITORY https://github.com/rapidsai/cudf.git + GIT_TAG "${RAPIDS_BRANCH}" + GIT_SHALLOW TRUE SOURCE_SUBDIR cpp + OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" + ) + # If after loading cudf we now have the CMAKE_CUDA_COMPILER variable we know that we need to + # re-enable the cuda language. + if(CMAKE_CUDA_COMPILER) + set(cudf_REQUIRES_CUDA + TRUE + PARENT_SCOPE + ) + endif() +endfunction() diff --git a/cpp/libcudf_kafka/cmake/thirdparty/get_cudf.cmake b/cpp/libcudf_kafka/cmake/thirdparty/get_cudf.cmake index 2547238e3b4e..fc2e1b957b6b 100644 --- a/cpp/libcudf_kafka/cmake/thirdparty/get_cudf.cmake +++ b/cpp/libcudf_kafka/cmake/thirdparty/get_cudf.cmake @@ -5,33 +5,12 @@ # cmake-format: on # ============================================================================= -# This function finds cudf and sets any additional necessary environment variables. -function(find_and_configure_cudf VERSION) - rapids_cmake_parse_version(MAJOR_MINOR ${VERSION} major_minor) - rapids_cpm_find( - cudf ${VERSION} - BUILD_EXPORT_SET cudf_kafka-exports - INSTALL_EXPORT_SET cudf_kafka-exports - CPM_ARGS - GIT_REPOSITORY https://github.com/rapidsai/cudf.git - GIT_TAG "${RAPIDS_BRANCH}" - GIT_SHALLOW TRUE SOURCE_SUBDIR cpp - OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" - ) - # If after loading cudf we now have the CMAKE_CUDA_COMPILER variable we know that we need to - # re-enable the cuda language - if(CMAKE_CUDA_COMPILER) - set(cudf_REQUIRES_CUDA - TRUE - PARENT_SCOPE - ) - endif() -endfunction() +include("${CMAKE_CURRENT_LIST_DIR}/../../../cmake/thirdparty/get_cudf.cmake") set(CUDF_KAFKA_MIN_VERSION "${CUDF_KAFKA_VERSION_MAJOR}.${CUDF_KAFKA_VERSION_MINOR}.${CUDF_KAFKA_VERSION_PATCH}" ) -find_and_configure_cudf(${CUDF_KAFKA_MIN_VERSION}) +find_and_configure_cudf(${CUDF_KAFKA_MIN_VERSION} cudf_kafka-exports) if(cudf_REQUIRES_CUDA) rapids_cuda_init_architectures(CUDF_KAFKA) diff --git a/cpp/libcudf_streaming/CMakeLists.txt b/cpp/libcudf_streaming/CMakeLists.txt new file mode 100644 index 000000000000..385e691d6359 --- /dev/null +++ b/cpp/libcudf_streaming/CMakeLists.txt @@ -0,0 +1,140 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= +cmake_minimum_required(VERSION 4.0 FATAL_ERROR) + +include(../../cmake/rapids_config.cmake) +include(rapids-cmake) +include(rapids-cpm) +include(rapids-cuda) +include(rapids-export) +include(rapids-find) + +rapids_cuda_init_architectures(CUDF_STREAMING) + +project( + CUDF_STREAMING + VERSION "${RAPIDS_VERSION}" + LANGUAGES CXX CUDA +) + +# Set a default build type if none was specified +rapids_cmake_build_type(Release) + +# 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) + +# Include CUDA configuration (expt-extended-lambda, expt-relaxed-constexpr, etc.) +include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/Modules/ConfigureCUDA.cmake) + +# ################################################################################################## +# * conda environment ----------------------------------------------------------------------------- +rapids_cmake_support_conda_env(conda_env MODIFY_PREFIX_PATH) + +# ################################################################################################## +# * Build options +option(BUILD_TESTS "Build tests for libcudf_streaming" OFF) + +message(VERBOSE "CUDF_STREAMING: Build gtests: ${BUILD_TESTS}") + +# ################################################################################################## +# * Dependencies + +# add third party dependencies using CPM +rapids_cpm_init() +include(cmake/thirdparty/get_cudf.cmake) +include(cmake/thirdparty/get_rapidsmpf.cmake) +include(../cmake/thirdparty/get_cucollections.cmake) + +# ################################################################################################## +# * library target -------------------------------------------------------------------------------- +add_library( + cudf_streaming SHARED + src/integrations/bloom_filter.cu + src/integrations/partition.cpp + src/integrations/utils.cpp + src/streaming/bloom_filter.cpp + src/streaming/channel_metadata.cpp + src/streaming/parquet.cpp + src/streaming/partition.cpp + src/streaming/table_chunk.cpp +) + +# ################################################################################################## +# * include paths --------------------------------------------------------------------------------- +target_include_directories( + cudf_streaming PUBLIC "$" + "$" +) + +# ################################################################################################## +# * library properties ---------------------------------------------------------------------------- +set_target_properties( + cudf_streaming + PROPERTIES BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + POSITION_INDEPENDENT_CODE ON + INTERFACE_POSITION_INDEPENDENT_CODE ON +) + +target_compile_options( + cudf_streaming PRIVATE "$<$:${CUDF_CUDA_FLAGS};--expt-relaxed-constexpr>" +) + +# ################################################################################################## +# * library links --------------------------------------------------------------------------------- +target_link_libraries( + cudf_streaming + PUBLIC rapidsmpf::rapidsmpf cudf::cudf + PRIVATE cuco::cuco +) + +# Add Conda library, and include paths if specified +if(TARGET conda_env) + target_link_libraries(cudf_streaming PRIVATE conda_env) +endif() + +add_library(cudf_streaming::cudf_streaming ALIAS cudf_streaming) + +# # GTests if enabled +if(BUILD_TESTS) + # GoogleTest + include(../cmake/thirdparty/get_gtest.cmake) + + # include CTest module -- automatically calls enable_testing() + include(CTest) + add_subdirectory(tests) +endif() + +# ################################################################################################## +# * cudf_streaming Install ------------------------------------------------------------------------- +rapids_cmake_install_lib_dir(lib_dir) +install( + TARGETS cudf_streaming + DESTINATION ${lib_dir} + EXPORT cudf_streaming-exports +) +install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) + +rapids_export( + INSTALL cudf_streaming + EXPORT_SET cudf_streaming-exports + GLOBAL_TARGETS cudf_streaming + NAMESPACE cudf_streaming:: +) + +rapids_export( + BUILD cudf_streaming + EXPORT_SET cudf_streaming-exports + GLOBAL_TARGETS cudf_streaming + NAMESPACE cudf_streaming:: +) diff --git a/cpp/libcudf_streaming/cmake/thirdparty/get_cudf.cmake b/cpp/libcudf_streaming/cmake/thirdparty/get_cudf.cmake new file mode 100644 index 000000000000..d8c816fe1eb7 --- /dev/null +++ b/cpp/libcudf_streaming/cmake/thirdparty/get_cudf.cmake @@ -0,0 +1,27 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include("${CMAKE_CURRENT_LIST_DIR}/../../../cmake/thirdparty/get_cudf.cmake") + +set(CUDF_STREAMING_MIN_VERSION + "${CUDF_STREAMING_VERSION_MAJOR}.${CUDF_STREAMING_VERSION_MINOR}.${CUDF_STREAMING_VERSION_PATCH}" +) +find_and_configure_cudf(${CUDF_STREAMING_MIN_VERSION} cudf_streaming-exports) + +if(cudf_REQUIRES_CUDA) + rapids_cuda_init_architectures(CUDF_STREAMING) + + # Since we are building cudf as part of ourselves we need to enable the CUDA language in the + # top-most scope + enable_language(CUDA) + + # Since CUDF_STREAMING only enables CUDA optionally we need to manually include the file that + # rapids_cuda_init_architectures relies on `project` calling + if(DEFINED CMAKE_PROJECT_CUDF_STREAMING_INCLUDE) + include("${CMAKE_PROJECT_CUDF_STREAMING_INCLUDE}") + endif() +endif() diff --git a/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake b/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake new file mode 100644 index 000000000000..4d96e113970e --- /dev/null +++ b/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake @@ -0,0 +1,27 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +# This function finds rapidsmpf and sets any additional necessary environment variables. +function(find_and_configure_rapidsmpf VERSION) + rapids_cmake_parse_version(MAJOR_MINOR ${VERSION} major_minor) + rapids_cpm_find( + rapidsmpf ${VERSION} + BUILD_EXPORT_SET cudf_streaming-exports + INSTALL_EXPORT_SET cudf_streaming-exports + CPM_ARGS + GIT_REPOSITORY https://github.com/rapidsai/rapidsmpf.git + GIT_TAG "${RAPIDS_BRANCH}" + GIT_SHALLOW TRUE SOURCE_SUBDIR cpp + OPTIONS "BUILD_MPI_SUPPORT OFF" "BUILD_UCXX_SUPPORT OFF" "BUILD_SLURM_SUPPORT OFF" + "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF" + ) +endfunction() + +set(CUDF_STREAMING_MIN_VERSION + "${CUDF_STREAMING_VERSION_MAJOR}.${CUDF_STREAMING_VERSION_MINOR}.${CUDF_STREAMING_VERSION_PATCH}" +) +find_and_configure_rapidsmpf(${CUDF_STREAMING_MIN_VERSION}) diff --git a/cpp/libcudf_streaming/include/cudf_streaming/integrations/bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/integrations/bloom_filter.hpp new file mode 100644 index 000000000000..20759ace7689 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/integrations/bloom_filter.hpp @@ -0,0 +1,136 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace cudf_streaming::integrations { + +/** + * @brief A bloom filter, used for approximate set membership queries. + */ +struct BloomFilter { + /** + * @brief Create a filter. + * + * @param num_blocks Number of blocks in the filter. + * @param seed Seed used for hashing each value. + * @param storage Storage to view as a bloom filter, must be appropriately + * initialized. + * @param stream CUDA stream for device operations. + */ + BloomFilter(std::size_t num_blocks, + std::uint64_t seed, + void* storage, + rmm::cuda_stream_view stream); + + /** + * @brief Create a read-only filter. + * + * @param num_blocks Number of blocks in the filter. + * @param seed Seed used for hashing each value. + * @param storage View of storage, must be appropriately initialized. + * @param stream CUDA stream for device operations. + * + * @return A const-qualified bloom filter viewing the underlying storage. + */ + static const BloomFilter view(std::size_t num_blocks, + std::uint64_t seed, + void const* storage, + rmm::cuda_stream_view stream); + + /** + * @brief Create uninitialized storage for a filter. + * + * @param num_blocks Number of blocks in the filter. + * @param stream CUDA stream for device operations. + * @param mr Memory resource for allocations. + * + * @return Unique pointer to a device buffer containing storage for the requested + * number of filter blocks. + */ + static std::unique_ptr storage(std::size_t num_blocks, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + + /** + * @brief Add values to the filter. + * + * @param values_to_hash table of values to hash (with cudf::hashing::xxhash_64()) + * @param stream CUDA stream for allocations and device operations. + * @param mr Memory resource for allocations. + */ + void add(cudf::table_view const& values_to_hash, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + + /** + * @brief Merge two filters, computing their union. + * + * @param other Other filter to merge into this one. + * @param stream CUDA stream for device operations. + * + * @throws std::logic_error If `other` is not compatible with this filter. + */ + void merge(BloomFilter const& other, rmm::cuda_stream_view stream); + + /** + * @brief Return a mask of which rows are contained in the filter. + * + * @param values Value to check for set membership + * @param stream CUDA stream for allocations and device operations. + * @param mr Memory resource for allocations. + * + * @return Mask vector to be used for filtering the table. + */ + [[nodiscard]] rmm::device_uvector contains(cudf::table_view const& values, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + /** + * @brief @return The stream the underlying storage is valid on. + */ + [[nodiscard]] rmm::cuda_stream_view stream() const noexcept; + + /** + * @brief @return Pointer to the underlying storage. + */ + [[nodiscard]] void* data() noexcept; + + /** + * @brief @return Const Pointer to the underlying storage. + */ + [[nodiscard]] void const* data() const noexcept; + + /** + * @brief @return Size in bytes of the underlying storage. + */ + [[nodiscard]] std::size_t size() const noexcept; + + /** + * @brief @return Number of blocks to use if the filter should fit in a given L2 cache + * size. + * + * @param l2size Size of the L2 cache in bytes. + */ + [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept; + + private: + std::size_t num_blocks_; ///< Number of blocks used in the filter. + std::uint64_t seed_; ///< Seed used when hashing values. + void* storage_; ///< Backing storage. + rmm::cuda_stream_view stream_; ///< Stream storage is valid on. +}; + +} // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp b/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp new file mode 100644 index 000000000000..36e47d111d17 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp @@ -0,0 +1,189 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cudf_streaming::integrations { + +/** + * @brief Partitions rows from the input table into multiple output tables. + * + * @param table The table to partition. + * @param columns_to_hash Indices of input columns to hash. + * @param num_partitions The number of partitions. + * @param hash_function Hash function to use. + * @param seed Seed value to the hash function. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param br Buffer resource for memory allocations. + * @param allow_overbooking If true, allow overbooking (true by default) + * + * @return A vector of each partition and a table that owns the device memory. + * + * @throws std::out_of_range if index is `columns_to_hash` is invalid + * + * @see cudf::hash_partition + * @see cudf::split + */ +[[nodiscard]] std::pair, std::unique_ptr> +partition_and_split( + cudf::table_view const& table, + std::vector const& columns_to_hash, + int num_partitions, + cudf::hash_id hash_function, + std::uint32_t seed, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); + +/** + * @brief Partitions rows from the input table into multiple packed (serialized) tables. + * + * @param table The table to partition. + * @param columns_to_hash Indices of input columns to hash. + * @param num_partitions The number of partitions to use. + * @param hash_function Hash function to use. + * @param seed Seed value to the hash function. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param br Buffer resource for memory allocations. + * @param allow_overbooking If true, allow overbooking (true by default) + * // TODO: disable this by default https://github.com/rapidsmpf/rapidsmpf/issues/449 + * + * @return A map of partition IDs and their packed tables. + * + * @throws std::out_of_range if index is `columns_to_hash` is invalid + * + * @see unpack_and_concat + * @see cudf::hash_partition + * @see cudf::pack + */ +[[nodiscard]] std::unordered_map +partition_and_pack( + cudf::table_view const& table, + std::vector const& columns_to_hash, + int num_partitions, + cudf::hash_id hash_function, + std::uint32_t seed, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); + +/** + * @brief Splits rows from the input table into multiple packed (serialized) tables. + * + * @param table The table to split and pack into partitions. + * @param splits The split points, equivalent to cudf::split(), i.e. one less than + * the number of result partitions. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param br Buffer resource for memory allocations. + * @param allow_overbooking If true, allow overbooking (true by default) + * // TODO: disable this by default https://github.com/rapidsmpf/rapidsmpf/issues/449 + * + * @return A map of partition IDs and their packed tables. + * + * @throws std::out_of_range if the splits are invalid. + * + * @see unpack_and_concat + * @see cudf::split + * @see partition_and_pack + */ +[[nodiscard]] std::unordered_map split_and_pack( + cudf::table_view const& table, + std::vector const& splits, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); + +/** + * @brief Unpack (deserialize) 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 @p stream and synchronized with + * the unpacking. + * + * @param partitions Packed input tables (partitions). + * @param stream CUDA stream on which concatenation occurs and on which the resulting + * table is ordered. + * @param br Buffer resource used for memory allocations. + * @param allow_overbooking If true, allow overbooking (true by default). + * @return The concatenated table resulting from unpacking the input partitions. + * + * @throws rapidsmpf::reservation_error If the buffer resource cannot reserve enough + * memory to concatenate all partitions. + * @throws std::logic_error If the partitions are not in device memory. + * + * @see partition_and_pack + * @see cudf::unpack + * @see cudf::concatenate + */ +[[nodiscard]] std::unique_ptr unpack_and_concat( + std::vector&& partitions, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); + +/** + * @brief 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 that are + * already in host memory are passed through 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 thrown. Overbooking is not allowed. + * + * @param partitions The partitions to spill. + * @param br Buffer resource used to reserve host memory and perform the move. + * + * @return A vector of `PackedData`, where each buffer resides in host memory. + * + * @throws rapidsmpf::reservation_error If host memory reservation fails. + */ +std::vector spill_partitions(std::vector&& partitions, + rapidsmpf::BufferResource* br); + +/** + * @brief Move spilled partitions (i.e., packed tables in host memory) 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 the `allow_overbooking` flag. + * + * @param partitions The partitions to unspill, potentially containing host-resident data. + * @param br Buffer resource responsible for memory reservation and spills. + * @param allow_overbooking If false, ensures enough memory is freed to satisfy the + * reservation; otherwise, allows overbooking even if spilling was insufficient. + * + * @return A vector of `PackedData`, each with a buffer in device memory. + * + * @throws rapidsmpf::reservation_error If overbooking exceeds the amount spilled and + * `allow_overbooking` is false. + */ +std::vector unspill_partitions( + std::vector&& partitions, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking); + +} // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/include/cudf_streaming/integrations/utils.hpp b/cpp/libcudf_streaming/include/cudf_streaming/integrations/utils.hpp new file mode 100644 index 000000000000..18cb249c69f2 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/integrations/utils.hpp @@ -0,0 +1,73 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include + +#include +#include + +namespace cudf_streaming::integrations { + +/** + * @brief Converts the element at a specific index in a `cudf::column_view` to a string. + * + * @param col The column view containing the data. + * @param index The index of the element to convert. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param mr Memory resource for device memory allocation. + * @return A string representation of the element at the specified index. + */ +std::string str(cudf::column_view col, + cudf::size_type index, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + +/** + * @brief Converts all elements in a `cudf::column_view` to a string. + * + * @param col The column view containing the data. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param mr Memory resource for device memory allocation. + * @return A string representation of all elements in the column. + */ +std::string str(cudf::column_view col, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + +/** + * @brief Converts all rows in a `cudf::table_view` to a string. + * + * @param tbl The table view containing the data. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param mr Memory resource for device memory allocation. + * @return A string representation of all rows in the table. + */ +std::string str(cudf::table_view tbl, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + +/** + * @brief Estimate the memory usage of a column. + * + * @param col The column to estimate the memory usage of. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @return The estimated memory usage of the column. + */ +std::size_t estimated_memory_usage(cudf::column_view const& col, rmm::cuda_stream_view stream); + +/** + * @brief Estimate the memory usage of a table. + * + * @param tbl The table to estimate the memory usage of. + * @param stream CUDA stream used for device memory operations and kernel launches. + * @return The estimated memory usage of the table. + */ +std::size_t estimated_memory_usage(cudf::table_view const& tbl, rmm::cuda_stream_view stream); + +} // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp new file mode 100644 index 000000000000..e08661e83d28 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp @@ -0,0 +1,108 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include + +#include +#include +#include +#include + +#include +#include + +namespace cudf_streaming::streaming { + +/** + * @brief Utility managing construction and use of a bloom filter. + * + * This class provides methods to build a bloom filter from a stream of `TableChunk`s and + * then apply that filter to a different stream. + * + * A bloom filter is a fixed size probabilistic data structure that provides approximate + * set membership queries with no false negatives. That is, let `A` be some set and `f(A)` + * be the bloom filter representation of that set. Then, for all `a ∈ A` it holds that `a + * ∈ f(A)`. Conversely, there is a false positive rate that increases with the number of + * distinct values inserted into the bloom filter, and decreases with the number of filter + * blocks. That is, for any given bloom filter, there exists `a ∉ A` such that `a ∈ f(A)`. + * + * See https://arxiv.org/pdf/2512.15595 for details on the GPU implementation used. + * + * We use bloom filters to provide runtime pre-filtering of tables during shuffle-based + * joins. We gather the keys that will match from the build side and use those to + * pre-filter the probe side before shuffling. + */ +struct BloomFilter { + /** + * @brief Construct storage for a bloom filter. + * + * @param ctx Streaming context. + * @param comm Communicator for the collective operation. + * @param seed Hash seed used when hashing values into the filter. + * @param num_filter_blocks Number of blocks in the filter. + */ + explicit BloomFilter(std::shared_ptr ctx, + std::shared_ptr comm, + std::uint64_t seed, + std::size_t num_filter_blocks) noexcept + : ctx_{std::move(ctx)}, + comm_{std::move(comm)}, + seed_{seed}, + num_filter_blocks_{num_filter_blocks} + { + } + + /** + * @brief Gets the communicator associated with this BloomFilter. + * + * @return Shared pointer to communicator. + */ + [[nodiscard]] std::shared_ptr const& comm() const noexcept + { + return comm_; + } + + /** + * @brief Build a bloom filter from the input channel. + * + * @param ch_in Input channel of `TableChunk`s to build bloom filter for. + * @param ch_out Output channel receiving a single message containing the bloom + * filter. + * @param tag Disambiguating tag to combine filters across ranks. + * @return Coroutine representing the construction of the bloom filter. + */ + [[nodiscard]] rapidsmpf::streaming::Actor build( + std::shared_ptr ch_in, + std::shared_ptr ch_out, + rapidsmpf::OpID tag); + + /** + * @brief Apply a bloom filter to an input channel. + * + * @param bloom_filter Channel containing the bloom filter (a single message). + * @param ch_in Input channel of `TableChunk`s to apply bloom filter to. + * @param ch_out Output channel receiving filtered `TableChunk`s. + * @param keys Indices selecting the key columns for the hash fingerprint + * + * @note The application of the bloom filter expects _exactly one_ message to come + * through the `bloom_filter` channel, which must be drained after that message is + * sent. + * + * @return Coroutine representing the application of the bloom filter. + */ + [[nodiscard]] rapidsmpf::streaming::Actor apply( + std::shared_ptr bloom_filter, + std::shared_ptr ch_in, + std::shared_ptr ch_out, + std::vector keys); + + private: + std::shared_ptr ctx_{}; + std::shared_ptr comm_{}; + std::uint64_t seed_{}; + std::size_t num_filter_blocks_{}; +}; +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/include/cudf_streaming/streaming/channel_metadata.hpp b/cpp/libcudf_streaming/include/cudf_streaming/streaming/channel_metadata.hpp new file mode 100644 index 000000000000..aa22187b7ef4 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/streaming/channel_metadata.hpp @@ -0,0 +1,249 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace cudf_streaming::streaming { + +/** + * @brief Hash partitioning scheme. + * + * Rows are distributed by `hash(columns[column_indices]) % modulus`. + */ +struct HashScheme { + std::vector column_indices; ///< Column indices to hash on. + int modulus; ///< Hash modulus (number of partitions). + + /** + * @brief Equality comparison. + * @return True if both schemes are equal. + */ + bool operator==(HashScheme const&) const = default; +}; + +/** + * @brief A single sort key: column index, sort direction, and null placement. + */ +struct OrderKey { + cudf::size_type column_index; ///< Column to sort on. + cudf::order order; ///< ASCENDING or DESCENDING. + cudf::null_order null_order; ///< BEFORE or AFTER. + + /** + * @brief Equality comparison. + * @return True if all fields are equal. + */ + bool operator==(OrderKey const&) const = default; + + /** + * @brief Inequality comparison. + * @return True if any field is different. + */ + bool operator!=(OrderKey const&) const = default; +}; + +/** + * @brief 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: + * - Partition 0: values < boundaries[0] + * - Partition i (0 < i < N-1): boundaries[i-1] <= values < boundaries[i] + * - Partition N-1: values >= boundaries[N-2] + * + * `keys[i]` is the i-th sort column; ordering is lexicographic by `keys[0]`, + * then `keys[1]`, and so on. + * + * When `boundaries` is set, its columns must align with `keys` + * (same count and compatible dtypes). Mismatched dtypes are a usage error. + * + * `strict_boundaries`: when true, every row in a chunk belongs to a single partition's + * half-open key range (partition keys do not straddle chunk interiors). When false, + * a chunk may contain keys spanning multiple partitions. + */ +struct OrderScheme { + std::vector keys; ///< Sort keys (column, order, null_order per entry). + std::shared_ptr boundaries; ///< N-1 boundary rows for N partitions. + /// See struct-level note on `strict_boundaries` semantics. + bool strict_boundaries{false}; + + /// @brief Default constructor. Produces an invalid (empty) scheme. + OrderScheme() = default; + + /** + * @brief Construct a validated OrderScheme. + * + * @param keys Non-empty sort keys; size must equal `boundaries->shape().second`. + * @param boundaries Non-null, device-resident boundary table (N-1 rows for N + * partitions). Accepts a `unique_ptr` via implicit conversion. + * @param strict_boundaries See struct-level doc. Defaults to false. + * @throws std::invalid_argument if `keys` is empty, `boundaries` is null or not + * device-resident, or `keys.size() != boundaries->shape().second`. + */ + OrderScheme(std::vector keys, + std::shared_ptr boundaries, + bool strict_boundaries = false); + + /** + * @brief Return a new OrderScheme with updated key column indices, sharing + * boundaries. + * + * The new key count must match the existing boundary column count. + * + * @param new_keys Replacement sort keys; size must equal + * `boundaries->shape().second`. + * @return A new OrderScheme with `new_keys` and the same `boundaries` and + * `strict_boundaries`. + * @throws std::invalid_argument if `new_keys` is empty or size mismatches boundaries. + */ + [[nodiscard]] OrderScheme with_keys(std::vector new_keys) const; + + /** + * @brief Check whether boundary values are aligned with another scheme. + * + * @param other The OrderScheme to compare against. + * @param br Buffer resource used for temporary allocations during comparison. + * @return True when both schemes have matching boundary values and strict_boundaries + * attributes, and the schemes are otherwise compatible (same order and null_order). + */ + [[nodiscard]] bool boundaries_aligned_with(OrderScheme const& other, + rapidsmpf::BufferResource& br) const; +}; + +/** + * @brief Partitioning specification for a single hierarchical level. + * + * Represents how data is partitioned at one level of the hierarchy + * (e.g., inter-rank or local). Use the static factory methods to construct. + * + * - `none()`: No partitioning information at this level. + * - `inherit()`: Partitioning is inherited from the parent level unchanged. + * - `from_hash(h)`: Explicit hash partitioning with the given scheme. + * - `from_order(o)`: Explicit order/range partitioning with the given scheme. + */ +struct PartitioningSpec { + /** + * @brief Type tag for PartitioningSpec. + */ + enum class Type : std::uint8_t { + NONE, ///< No partitioning information at this level. + INHERIT, ///< Partitioning is inherited from parent level unchanged. + HASH, ///< Hash partitioning. + ORDER, ///< Order/range partitioning. + }; + + Type type = Type::NONE; ///< The type of partitioning. + std::optional hash; ///< Valid only when type == HASH. + std::optional order; ///< Valid only when type == ORDER. + + /** + * @brief Create a spec indicating no partitioning information. + * @return A PartitioningSpec with type NONE. + */ + static PartitioningSpec none() { return {}; } + + /** + * @brief Create a spec indicating partitioning passes through from parent. + * @return A PartitioningSpec with type INHERIT. + */ + static PartitioningSpec inherit() + { + return {.type = Type::INHERIT, .hash = std::nullopt, .order = std::nullopt}; + } + + /** + * @brief Create a spec for hash partitioning. + * @param h The hash scheme to use. + * @return A PartitioningSpec with type HASH. + */ + static PartitioningSpec from_hash(HashScheme h) + { + return {.type = Type::HASH, .hash = std::move(h), .order = std::nullopt}; + } + + /** + * @brief Create a spec for order/range partitioning. + * @param o The order scheme to use. `o.keys` must be non-empty; otherwise + * throws `std::invalid_argument`. + * @return A PartitioningSpec with type ORDER. + */ + static PartitioningSpec from_order(OrderScheme o); +}; + +/** + * @brief Hierarchical partitioning metadata for a data stream. + * + * Describes how data flowing through a channel is partitioned at multiple + * levels of the system hierarchy. Each level corresponds to a communicator + * used to shuffle data at that level: + * + * - `inter_rank`: Distribution across ranks, corresponding to the primary + * communicator (e.g., `Context::comm()`). Shuffle operations at this level + * move data between ranks. + * - `local`: Distribution within a rank, corresponding to a single-rank + * communicator. Operations at this level repartition data locally without + * network communication. + */ +struct Partitioning { + /// Distribution across ranks (corresponds to primary communicator). + PartitioningSpec inter_rank; + /// Distribution within a rank (corresponds to local/single communicator). + PartitioningSpec local; +}; + +/** + * @brief Channel-level metadata describing the data stream. + * + * Contains information about chunk counts, partitioning, and duplication + * status for the data flowing through a channel. + */ +struct ChannelMetadata { + std::uint64_t local_count{}; ///< Local chunk-count estimate for this rank. + Partitioning partitioning; ///< How the data is partitioned. + bool duplicated{}; ///< Whether data is duplicated on all workers. + + /// @brief Default constructor. + ChannelMetadata() = default; + + /** + * @brief Construct metadata with specified values. + * + * @param local_count Local chunk count. + * @param partitioning Partitioning metadata (default: no partitioning). + * @param duplicated Whether data is duplicated (default: false). + */ + ChannelMetadata(std::uint64_t local_count, + Partitioning partitioning = {}, + bool duplicated = false) + : local_count(local_count), partitioning(std::move(partitioning)), duplicated(duplicated) + { + } +}; + +/** + * @brief Wrap a `ChannelMetadata` into a `Message`. + * + * @param sequence_number Ordering identifier for the message. + * @param m The metadata to wrap. + * @return A `Message` encapsulating the metadata as its payload. + */ +rapidsmpf::streaming::Message to_message(std::uint64_t sequence_number, + std::unique_ptr m); + +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/include/cudf_streaming/streaming/parquet.hpp b/cpp/libcudf_streaming/include/cudf_streaming/streaming/parquet.hpp new file mode 100644 index 000000000000..481bbefa80e2 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/streaming/parquet.hpp @@ -0,0 +1,62 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +namespace cudf_streaming::streaming { + +/** + * @brief Filter ast expression with lifetime/stream management. + */ +struct Filter { + rmm::cuda_stream_view stream; ///< Stream the filter's scalars are valid on. + cudf::ast::expression& filter; ///< Filter expression. + rapidsmpf::OwningWrapper owner{}; ///< Owner of all objects in the filter. +}; + +namespace actor { +/** + * @brief Asynchronously read parquet files into an output channel. + * + * @note This is a collective operation, all ranks named by the execution context's + * communicator will participate. All ranks must specify the same set of options. + * Behaviour is undefined if a `read_parquet` actor appears only on a subset of the ranks + * named by the communicator, or the options differ between ranks. + * + * @param ctx The execution context to use. + * @param comm Communicator for distributing files across ranks. + * @param ch_out Channel to which `TableChunk`s are sent. + * @param num_producers Number of concurrent producer tasks. + * @param options Template reader options. The files within will be picked apart and used + * to reconstruct new options for each read chunk. The options should therefore specify + * the read options "as-if" one were reading the whole input in one go. + * @param num_rows_per_chunk Target (maximum) number of rows any sent `TableChunk` should + * have. + * @param filter Optional filter expression to apply to the read. + * + * @return Streaming actor representing the asynchronous read. + */ +rapidsmpf::streaming::Actor read_parquet(std::shared_ptr ctx, + std::shared_ptr comm, + std::shared_ptr ch_out, + std::size_t num_producers, + cudf::io::parquet_reader_options options, + // TODO: use byte count, not row count? + cudf::size_type num_rows_per_chunk, + std::unique_ptr filter = nullptr); +} // namespace actor +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/include/cudf_streaming/streaming/partition.hpp b/cpp/libcudf_streaming/include/cudf_streaming/streaming/partition.hpp new file mode 100644 index 000000000000..ae09017eb573 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/streaming/partition.hpp @@ -0,0 +1,81 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include + +#include +#include + +namespace cudf_streaming::streaming { + +namespace actor { + +/** + * @brief Asynchronously partitions input tables into multiple packed (serialized) tables. + * + * This is a streaming version of `rapidsmpf::partition_and_split` that operates on table + * chunks using channels. + * + * It receives tables from an input channel, partitions each row into one of + * `num_partitions` based on a hash of the selected columns, packs the resulting + * partitions, and sends them to an output channel. + * + * @param ctx The actor context to use. + * @param ch_in Input channel providing `TableChunk`s to partition. + * @param ch_out Output channel to which `PartitionMapChunk`s are sent. + * @param columns_to_hash Indices of input columns to hash. + * @param num_partitions The number of partitions to use. + * @param hash_function Hash function to use for partitioning. + * @param seed Seed value for the hash function. + * + * @return Streaming actor representing the asynchronous partitioning and packing + * operation. + * + * @throws std::out_of_range if any index in `columns_to_hash` is invalid. + * + * @see rapidsmpf::partition_and_split + */ +rapidsmpf::streaming::Actor partition_and_pack( + std::shared_ptr ctx, + std::shared_ptr ch_in, + std::shared_ptr ch_out, + std::vector columns_to_hash, + int num_partitions, + cudf::hash_id hash_function, + std::uint32_t seed); + +/** + * @brief Asynchronously unpacks and concatenates packed partitions. + * + * This is a streaming version of `rapidsmpf::unpack_and_concat` that operates on + * packed partition chunks using channels. + * + * It receives packed partitions from the input channel, deserializes and concatenates + * them, and sends the resulting tables to the output channel. Empty partitions are + * ignored. + * + * @param ctx The actor context to use. + * @param ch_in Input channel providing packed partitions as PartitionMapChunk or + * PartitionVectorChunk. + * @param ch_out Output channel to which unpacked and concatenated tables are sent. + * + * @return Streaming actor representing the asynchronous unpacking and concatenation + * operation. + * + * @see rapidsmpf::unpack_and_concat + */ +rapidsmpf::streaming::Actor unpack_and_concat( + std::shared_ptr ctx, + std::shared_ptr ch_in, + std::shared_ptr ch_out); + +} // namespace actor +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/include/cudf_streaming/streaming/table_chunk.hpp b/cpp/libcudf_streaming/include/cudf_streaming/streaming/table_chunk.hpp new file mode 100644 index 000000000000..a6106911f205 --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/streaming/table_chunk.hpp @@ -0,0 +1,337 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace cudf_streaming::streaming { + +/** + * @brief A unit of table data in a streaming pipeline. + * + * Represents either an unpacked `cudf::table` or a `PackedData`. + * + * TableChunks 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. + */ +class TableChunk { + public: + /** + * @brief Indicates whether the TableChunk holds an exclusive or shared view + * of the underlying table data. + * + * This boolean enum is used to explicitly express ownership semantics + * when constructing a TableChunk from a `cudf::table_view`. + * + * - `ExclusiveView::YES`: The TableChunk has exclusive ownership of + * the table's device memory and are considered spillable. + * + * - `ExclusiveView::NO`: The TableChunk is a non-owning view of data + * managed elsewhere. The memory may be shared or externally owned, + * and the chunk is therefore not spillable. + */ + enum class ExclusiveView : bool { + NO, + YES, + }; + + /** + * @brief Construct a TableChunk from a device table. + * + * @param table Device-resident table. + * @param stream The CUDA stream on which the table was created. + */ + TableChunk(std::unique_ptr table, rmm::cuda_stream_view stream); + + /** + * @brief Construct a TableChunk from a device table view. + * + * The TableChunk does not take ownership of the underlying data; instead, the + * provided @p owner object is kept alive for the lifetime of the TableChunk. + * The caller is responsible for ensuring that the underlying device memory + * referenced by @p table_view remains valid during this period. + * + * This constructor is typically used when creating a TableChunk from Python, + * where @p owner is used to keep the corresponding Python object alive until + * the TableChunk is destroyed. + * + * @param table_view Device-resident table view. + * @param stream CUDA stream on which the table was created. + * @param owner Object owning the memory backing @p table_view. This object will be + * destroyed last when the TableChunk is destroyed or spilled. + * @param exclusive_view Specifies whether this TableChunk has exclusive ownership + * semantics over the underlying table data: + * - When `ExclusiveView::YES`, the following guarantees must hold: + * - The @p table_view is the sole representation of the table. + * - The @p owner exclusively owns the table memory. + * These guarantees allow the TableChunk to be spillable and ensure that + * destroying @p owner will correctly free the associated device memory. + * - When `ExclusiveView::NO`, the chunk is considered a non-owning view and + * is therefore not spillable. + */ + TableChunk(cudf::table_view table_view, + rmm::cuda_stream_view stream, + rapidsmpf::OwningWrapper&& owner, + ExclusiveView exclusive_view); + + /** + * @brief Construct a TableChunk from a packed data blob. + * + * The packed data's CUDA stream will be associated the new table chunk. + * + * @param packed_data Serialized host/device data with metadata. + */ + TableChunk(std::unique_ptr packed_data); + + ~TableChunk() = default; + + /** + * @brief Move constructor + * + * @note After this call `other.is_available() == false`. + * @param other The TableChunk to move from. + */ + TableChunk(TableChunk&& other) noexcept; + + /** + * @brief Move assignment + * + * @note After this call `other.is_available() == false`. + * @param other The TableChunk to move from. + * @return Reference to this. + */ + TableChunk& operator=(TableChunk&& other) noexcept; + TableChunk(TableChunk const&) = delete; + TableChunk& operator=(TableChunk const&) = delete; + + /** + * @brief Returns the CUDA stream on which this table chunk was created. + * + * @return The CUDA stream view. + */ + [[nodiscard]] rmm::cuda_stream_view stream() const noexcept; + + /** + * @brief Number of bytes allocated for the data in the specified memory type. + * + * @param mem_type The memory type to query. + * @return Number of bytes allocated. + */ + [[nodiscard]] std::size_t data_alloc_size(rapidsmpf::MemoryType mem_type) const; + + /** + * @brief Indicates whether the underlying cudf table data is fully available in + * device memory. + * + * @return `true` if the table is already available; otherwise, `false`. + */ + [[nodiscard]] bool is_available() const noexcept; + + /** + * @brief Returns the estimated cost (in bytes) of making the table available. + * + * Currently, only device memory cost is tracked. + * + * @return The cost in bytes. + */ + [[nodiscard]] std::size_t make_available_cost() const noexcept; + + /** + * @brief Moves this table chunk into a new one with its cudf table made available. + * + * As part of the move, a copy or unpack may be performed, the associated CUDA + * stream is used. + * + * @param reservation Memory reservation for allocations if needed. + * @return A new TableChunk with data available on device. + * + * @note After this call, the current object is in a moved-from state; + * only reassignment, movement, or destruction are valid. + */ + [[nodiscard]] TableChunk make_available(rapidsmpf::MemoryReservation& reservation); + + /** + * @brief Moves this table chunk into a new one with its cudf table made available. + * + * Takes ownership of the memory reservation and consumes it entirely as part + * of making the data available on device. The full reservation is considered + * used, even if the actual allocation requires fewer bytes. + * + * @param reservation Memory reservation to be consumed for allocations. + * @return A new TableChunk with data available on device. + * + * @note After this call, the current object is in a moved-from state; only + * reassignment, movement, or destruction are valid. + */ + [[nodiscard]] TableChunk make_available(rapidsmpf::MemoryReservation&& reservation); + + /** + * @brief Move this table chunk into a new one with its cudf table made available. + * + * This variant of make_available() is a coroutine that may suspend if device + * memory is not immediately available. + * + * @note After this call, the current object is in a moved-from state; only + * reassignment, movement, or destruction are valid. + * + * @param ctx Streaming context used to access the memory reservation mechanism. + * @param net_memory_delta Estimated change in memory usage after the reservation + * is granted and all work using the returned `TableChunk` has completed. See + * `MemoryReserveOrWait::reserve_or_wait` for details. + * @return A new `TableChunk` that is available on device. + * + * @throws std::runtime_error If shutdown occurs before the reservation can be + * processed. + * @throws std::overflow_error If no progress is possible within the timeout and + * overbooking is disabled. + */ + [[nodiscard]] coro::task make_available( + std::shared_ptr ctx, + std::int64_t net_memory_delta = + rapidsmpf::streaming::MemoryReserveOrWait::missing_net_memory_delta); + + /** + * @brief Returns a view of the underlying table. + * + * The table must be available in device memory. + * + * @return cudf::table_view representing the table. + * + * @throws std::invalid_argument if `is_available() == false`. + */ + [[nodiscard]] cudf::table_view table_view() const; + + /** + * @brief Indicates whether this table chunk can be spilled from device to host memory. + * + * A table chunk is considered spillable if it owns its underlying memory. This is + * true when it was created from one of the following: + * - A device-owning source such as a `cudf::table`, `cudf::packed_columns`, or + * `PackedData`. + * - A `cudf::table_view` constructed with `is_exclusive_view == true`, indicating + * that the view is the sole representation of the underlying data and that its + * owner exclusively manages the table's memory. + * + * In contrast, chunks constructed from non-exclusive `cudf::table_view` instances are + * non-owning views of externally managed memory and therefore not spillable. + * + * To spill a table chunk from device to host memory, first call `copy()` to create a + * host-side copy, then delete or overwrite the original device chunk. If + * `is_spillable() == true`, destroying the original device chunk will release the + * associated device memory. + * + * @return `true` if the table chunk owns its memory and can be spilled; otherwise + * `false`. + */ + [[nodiscard]] bool is_spillable() const; + + /** + * @brief Create a deep copy of the table chunk. + * + * Allocates new memory for all buffers in the table using the specified + * `reservation`, which determines the target memory type (e.g., host or device). + * As a consequence, the `is_available()` status may differ in the new copy. For + * example, copying an available table chunk from device to host memory will result + * in an unavailable copy. + * + * @param reservation Memory reservation used to track and limit allocations. + * @return A new `TableChunk` instance containing copies of all buffers and metadata. + * + * @throws rapidsmpf::reservation_error If the total allocation size exceeds the + * available reservation. + */ + [[nodiscard]] TableChunk copy(rapidsmpf::MemoryReservation& reservation) const; + + /** + * @brief 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 `PackedData`), the packed data is moved out directly + * with no copy. Otherwise the table is serialized via `cudf::pack()`. + * + * @param br Buffer resource used for the device memory resource when packing + * is required. + * @return A unique pointer to the resulting `PackedData`. + * + * @throws std::invalid_argument If the data is not already packed and + * `is_available() == false`. + * + * @note After this call, this object is in a moved-from state; only reassignment, + * movement, or destruction are valid. + * + * @note No memory reservation is required. If the data is already in packed form, + * no allocation occurs. If packing is required, `cudf::pack()` allocates device + * memory that is not tracked via a reservation. + */ + [[nodiscard]] std::unique_ptr into_packed_data( + rapidsmpf::BufferResource* br) &&; + + /** + * @brief Return the shape of the table stored by the table chunk. + * + * @return Pair of number of rows and number of columns. + */ + [[nodiscard]] std::pair shape() const noexcept; + + private: + ///< @brief Optional owning object if the TableChunk was constructed from a + ///< table_view. + rapidsmpf::OwningWrapper owner_{}; + + // At most, one of the following unique pointers is non-null. If all of them are null, + // the TableChunk is a non-owning view. + // TODO: use a variant and drop the unique pointers? + std::unique_ptr table_; + std::unique_ptr packed_data_; + + // Has value iff this TableChunk is available. + std::optional table_view_; + + // Zero initialized data allocation size (one for each memory type). + std::array data_alloc_size_ = {}; + std::size_t make_available_cost_; // For now, only device memory cost is tracked. + + rmm::cuda_stream_view stream_; + bool is_spillable_; +}; + +/** + * @brief Generate a content description for a `TableChunk`. + * + * @param obj The object's content to describe. + * @return A new content description. + */ +rapidsmpf::ContentDescription get_content_description(TableChunk const& obj); + +/** + * @brief Wrap a `TableChunk` into a `Message`. + * + * @param sequence_number Ordering identifier for the message. + * @param chunk The chunk to wrap into a message. + * @return A `Message` encapsulating the provided chunk as its payload. + */ +rapidsmpf::streaming::Message to_message(std::uint64_t sequence_number, + std::unique_ptr chunk); + +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/src/integrations/bloom_filter.cu b/cpp/libcudf_streaming/src/integrations/bloom_filter.cu new file mode 100644 index 000000000000..3b2fb3195f4c --- /dev/null +++ b/cpp/libcudf_streaming/src/integrations/bloom_filter.cu @@ -0,0 +1,150 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include + +// cuco headers have sign-conversion issues; suppress for the host compiler +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wsign-conversion" +#elif defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wsign-conversion" +#endif +#include +#include +#include +#include +#ifdef __clang__ +#pragma clang diagnostic pop +#elif defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cudf_streaming::integrations { + +namespace { +using KeyType = std::uint64_t; + +using BloomFilterRefType = + cuco::bloom_filter_ref, + cuco::thread_scope_device, + cuco::arrow_filter_policy>; +using StorageType = BloomFilterRefType::filter_block_type; + +} // namespace + +BloomFilter::BloomFilter(std::size_t num_blocks, + std::uint64_t seed, + void* storage, + rmm::cuda_stream_view stream) + : num_blocks_{num_blocks}, seed_{seed}, storage_{storage}, stream_{stream} +{ + // TODO: use an aligned allocator adaptor to ensure this holds. + // Today all RMM device allocators guarantee at least 256 byte alignment, but that is + // an implementation detail. + RAPIDSMPF_EXPECTS( + reinterpret_cast(storage_) % std::alignment_of_v == 0, + "Allocation for bloom filter is not aligned."); +} + +BloomFilter const BloomFilter::view(std::size_t num_blocks, + std::uint64_t seed, + void const* storage, + rmm::cuda_stream_view stream) +{ + // const-cast is safe because the returned object is also const and therefore can't + // call methods that throw away constness. + return BloomFilter(num_blocks, seed, const_cast(storage), stream); +} + +std::unique_ptr BloomFilter::storage(std::size_t num_blocks, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + return std::make_unique(num_blocks * sizeof(StorageType), stream, mr); +} + +void BloomFilter::add(cudf::table_view const& values_to_hash, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + RAPIDSMPF_NVTX_FUNC_RANGE(); + auto filter_ref = BloomFilterRefType{ + static_cast(storage_), num_blocks_, cuco::thread_scope_device, {}}; + auto hashes = cudf::hashing::xxhash_64( + values_to_hash, seed_, stream, cudf::get_current_device_resource_ref()); + auto hash_view = hashes->view(); + RAPIDSMPF_EXPECTS(hash_view.type().id() == cudf::type_to_id(), + "Hash values do not have correct type"); + filter_ref.add_async(hash_view.begin(), hash_view.end(), stream); +} + +void BloomFilter::merge(BloomFilter const& other, rmm::cuda_stream_view stream) +{ + RAPIDSMPF_NVTX_FUNC_RANGE(); + RAPIDSMPF_EXPECTS(num_blocks_ == other.num_blocks_, "Mismatching number of blocks in filters"); + auto ref_this = BloomFilterRefType{ + static_cast(storage_), num_blocks_, cuco::thread_scope_device, {}}; + auto ref_other = BloomFilterRefType{ + static_cast(other.storage_), num_blocks_, cuco::thread_scope_device, {}}; + ref_this.merge_async(ref_other, stream); +} + +rmm::device_uvector BloomFilter::contains(cudf::table_view const& values, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + RAPIDSMPF_NVTX_FUNC_RANGE(); + auto filter_ref = BloomFilterRefType{ + static_cast(storage_), num_blocks_, cuco::thread_scope_device, {}}; + auto hashes = + cudf::hashing::xxhash_64(values, seed_, stream, cudf::get_current_device_resource_ref()); + auto view = hashes->view(); + rmm::device_uvector result{static_cast(view.size()), stream, mr}; + filter_ref.contains_async(view.begin(), view.end(), result.begin(), stream); + return result; +} + +std::size_t BloomFilter::fitting_num_blocks(std::size_t l2size) noexcept +{ + return (l2size * 2) / (3 * sizeof(StorageType)); +} + +rmm::cuda_stream_view BloomFilter::stream() const noexcept { return stream_; } + +void* BloomFilter::data() noexcept { return storage_; } + +void const* BloomFilter::data() const noexcept { return storage_; } + +std::size_t BloomFilter::size() const noexcept { return num_blocks_ * sizeof(StorageType); } + +} // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/src/integrations/partition.cpp b/cpp/libcudf_streaming/src/integrations/partition.cpp new file mode 100644 index 000000000000..f9af79f6fd76 --- /dev/null +++ b/cpp/libcudf_streaming/src/integrations/partition.cpp @@ -0,0 +1,226 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace cudf_streaming::integrations { + +std::pair, std::unique_ptr> partition_and_split( + cudf::table_view const& table, + std::vector const& columns_to_hash, + int num_partitions, + cudf::hash_id hash_function, + std::uint32_t seed, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking) +{ + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + if (table.num_rows() == 0) { + // Return views of a copy of the empty `table`. + auto owner = std::make_unique(table, stream, br->device_mr()); + return {std::vector(rapidsmpf::safe_cast(num_partitions), + owner->view()), + std::move(owner)}; + } + + // hash_partition does a deep-copy. Therefore, we need to reserve memory for + // at least the size of the table. + auto reservation = + br->reserve_device_memory_and_spill(estimated_memory_usage(table, stream), allow_overbooking); + auto [partition_table, offsets] = cudf::hash_partition( + table, columns_to_hash, num_partitions, hash_function, seed, stream, br->device_mr()); + reservation.clear(); + + // Notice, the offset argument for split() and hash_partition() doesn't align. + // hash_partition() returns the start offset of each partition thus we have to + // skip the first offset. See: . + auto partition_offsets = + cudf::host_span(offsets.data() + 1, offsets.size() - 2); + + // split does not make any copies. + auto tbl_partitioned = cudf::split(partition_table->view(), partition_offsets, stream); + + return {std::move(tbl_partitioned), std::move(partition_table)}; +} + +std::unordered_map partition_and_pack( + cudf::table_view const& table, + std::vector const& columns_to_hash, + int num_partitions, + cudf::hash_id hash_function, + std::uint32_t seed, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking) +{ + RAPIDSMPF_NVTX_FUNC_RANGE(); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + RAPIDSMPF_EXPECTS(num_partitions > 0, "Need to split to at least one partition"); + if (table.num_rows() == 0) { + auto splits = + std::vector(rapidsmpf::safe_cast(num_partitions - 1), 0); + return split_and_pack(table, splits, stream, br, allow_overbooking); + } + + // hash_partition does a deep-copy. Therefore, we need to reserve memory for + // at least the size of the table. + auto reservation = + br->reserve_device_memory_and_spill(estimated_memory_usage(table, stream), allow_overbooking); + auto [reordered, split_points] = cudf::hash_partition( + table, columns_to_hash, num_partitions, hash_function, seed, stream, br->device_mr()); + reservation.clear(); + std::vector splits(split_points.begin() + 1, split_points.end() - 1); + return split_and_pack(reordered->view(), splits, stream, br, allow_overbooking); +} + +std::unordered_map split_and_pack( + cudf::table_view const& table, + std::vector const& splits, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking) +{ + RAPIDSMPF_NVTX_FUNC_RANGE(); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + std::unordered_map ret; + + // contiguous split does a deep-copy. Therefore, we need to reserve memory for + // at least the size of the table. + auto reservation = br->reserve_device_memory_and_spill( + cudf::packed_size(table, stream, br->device_mr()), allow_overbooking); + auto packed = cudf::contiguous_split(table, splits, stream, br->device_mr()); + reservation.clear(); + ret.reserve(packed.size()); + for (rapidsmpf::shuffler::PartID i = 0; rapidsmpf::safe_cast(i) < packed.size(); + i++) { + auto pack = std::move(packed[i].data); + ret.emplace( + i, + rapidsmpf::PackedData(std::move(pack.metadata), br->move(std::move(pack.gpu_data), stream))); + } + return ret; +} + +std::unique_ptr unpack_and_concat(std::vector&& partitions, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking) +{ + RAPIDSMPF_NVTX_FUNC_RANGE(); + RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); + + // Let's find the total size of the partitions and how much of the packed data we + // need to move to device memory (unspill). + std::size_t total_size = 0; + std::size_t non_device_size = 0; + for (auto& packed_data : partitions) { + if (!packed_data.empty()) { + std::size_t size = packed_data.data->size; + total_size += size; + if (packed_data.data->mem_type() != rapidsmpf::MemoryType::DEVICE) { + non_device_size += size; + } + } + } + + std::vector unpacked; + std::vector references; + std::vector packed_data_streams; + unpacked.reserve(partitions.size()); + references.reserve(partitions.size()); + packed_data_streams.reserve(partitions.size()); + + // Reserve device memory for the unspill AND the cudf::unpack() calls. + auto reservation = + br->reserve_device_memory_and_spill(total_size + non_device_size, allow_overbooking); + for (auto& packed_data : partitions) { + if (!packed_data.empty()) { + if (packed_data.data->size > 0) { // No need to sync empty buffers. + packed_data_streams.push_back(packed_data.data->stream()); + } + unpacked.push_back(cudf::unpack(references.emplace_back( + std::move(packed_data.metadata), + br->move_to_device_buffer(std::move(packed_data.data), reservation)))); + } + } + reservation.clear(); + + // We need to synchronize `stream` with the packed_data and update their + // underlying device buffers to use `stream` going forward. This ensures + // the packed data are not deallocated before we have a chance to + // concatenate them on `stream`. + rapidsmpf::cuda_stream_join(std::views::single(stream), packed_data_streams); + for (cudf::packed_columns& packed_columns : references) { + packed_columns.gpu_data->set_stream(stream); + } + + reservation = br->reserve_device_memory_and_spill(total_size, allow_overbooking); + return cudf::concatenate(unpacked, stream, br->device_mr()); +} + +std::vector spill_partitions(std::vector&& partitions, + rapidsmpf::BufferResource* br) +{ + // Sum the total size of all packed data in device memory. + std::size_t device_size{0}; + for (auto& [_, data] : partitions) { + if (data->mem_type() == rapidsmpf::MemoryType::DEVICE) { device_size += data->size; } + } + // Spill each partition to host memory. + auto reservation = br->reserve_or_fail(device_size, rapidsmpf::SPILL_TARGET_MEMORY_TYPES); + std::vector ret; + ret.reserve(partitions.size()); + for (auto& [metadata, data] : partitions) { + ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation)); + } + return ret; +} + +std::vector unspill_partitions( + std::vector&& partitions, + rapidsmpf::BufferResource* br, + rapidsmpf::AllowOverbooking allow_overbooking) +{ + auto statistics = br->statistics(); + // Sum the total size of all packed data not in device memory already. + std::size_t non_device_size{0}; + for (auto& [_, data] : partitions) { + if (data->mem_type() != rapidsmpf::MemoryType::DEVICE) { non_device_size += data->size; } + } + + // Unspill each partition. + auto reservation = br->reserve_device_memory_and_spill(non_device_size, allow_overbooking); + std::vector ret; + ret.reserve(partitions.size()); + for (auto& [metadata, data] : partitions) { + ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation)); + } + + return ret; +} +} // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/src/integrations/utils.cpp b/cpp/libcudf_streaming/src/integrations/utils.cpp new file mode 100644 index 000000000000..94c5ba60e424 --- /dev/null +++ b/cpp/libcudf_streaming/src/integrations/utils.cpp @@ -0,0 +1,157 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cudf_streaming::integrations { + +namespace { +struct str_cudf_column_scalar_fn { + template + requires(cudf::is_numeric()) + std::string operator()(cudf::column_view col, + cudf::size_type index, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) + { + std::unique_ptr scalar = cudf::get_element(col, index, stream, mr); + if (!scalar->is_valid(stream)) { return "null"; } + auto typed_scalar = static_cast const*>(scalar.get()); + T val = typed_scalar->value(stream); + return std::to_string(val); + } + + template + requires(!cudf::is_numeric()) + std::string operator()(cudf::column_view /* col */, + cudf::size_type /* index */, + rmm::cuda_stream_view /* stream */, + rmm::device_async_resource_ref /* mr */ + ) + { + RAPIDSMPF_FAIL("not implemented"); + } +}; + +struct cudf_column_data_size_fn { + template + requires(cudf::is_fixed_width()) + std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) + { + return rapidsmpf::safe_cast(col.size()) * cudf::size_of(col.type()) + + bitmask_size(col); + } + + // string type specialization + template + requires(std::is_same_v) + std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view stream) + { + cudf::strings_column_view sv(col); + return rapidsmpf::safe_cast(sv.chars_size(stream)) + bitmask_size(col); + } + + // compound type specialization except string + template + requires(!std::is_same_v && cudf::is_compound()) + std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) + { + // compound types (except string) ie. list, dict, structs dont have a + // content::data buffer. Data is stored in children columns. So, just return the + // bitmask size. + return bitmask_size(col); + } + + template + std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) + { + RAPIDSMPF_FAIL("not implemented for type: " + cudf::type_to_name(col.type())); + } + + static std::size_t bitmask_size(cudf::column_view const& col) + { + return col.nullable() ? cudf::bitmask_allocation_size_bytes(col.size()) : 0; + } +}; + +} // namespace + +std::string str(cudf::column_view col, + cudf::size_type index, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + return cudf::type_dispatcher(col.type(), str_cudf_column_scalar_fn{}, col, index, stream, mr); +} + +std::string str(cudf::column_view col, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + std::stringstream ss; + ss << "Column(["; + for (cudf::size_type i = 0; i < col.size(); ++i) { + if (i != 0) { ss << ", "; } + ss << str(col, i, stream, mr); + } + ss << "])"; + return ss.str(); +} + +std::string str(cudf::table_view tbl, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + std::stringstream ss; + ss << "Table(["; + bool first = true; + for (auto col : tbl) { + if (!first) { ss << ", "; } + first = false; + ss << str(col, stream, mr); + } + ss << "])"; + return ss.str(); +} + +std::size_t estimated_memory_usage(cudf::column_view const& col, rmm::cuda_stream_view stream) +{ + return std::transform_reduce( + col.child_begin(), + col.child_end(), + cudf::type_dispatcher(col.type(), cudf_column_data_size_fn{}, col, stream), + std::plus{}, + [&stream](cudf::column_view const& child) { return estimated_memory_usage(child, stream); }); +} + +std::size_t estimated_memory_usage(cudf::table_view const& tbl, rmm::cuda_stream_view stream) +{ + return std::transform_reduce( + tbl.begin(), tbl.end(), std::size_t{0}, std::plus{}, [&stream](cudf::column_view const& col) { + return estimated_memory_usage(col, stream); + }); +} + +} // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/src/streaming/bloom_filter.cpp b/cpp/libcudf_streaming/src/streaming/bloom_filter.cpp new file mode 100644 index 000000000000..be7da4202571 --- /dev/null +++ b/cpp/libcudf_streaming/src/streaming/bloom_filter.cpp @@ -0,0 +1,138 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cudf_streaming::streaming { + +rapidsmpf::streaming::Actor BloomFilter::build( + std::shared_ptr ch_in, + std::shared_ptr ch_out, + rapidsmpf::OpID tag) +{ + rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; + co_await ctx_->executor()->schedule(); + co_await ch_in->shutdown_metadata(); + co_await ch_out->shutdown_metadata(); + auto const& br = ctx_->br(); + auto mr = br->device_mr(); + auto filter_stream = br->stream_pool().get_stream(); + rapidsmpf::CudaEvent event; + auto storage = + cudf_streaming::integrations::BloomFilter::storage(num_filter_blocks_, filter_stream, mr); + RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream)); + auto filter = cudf_streaming::integrations::BloomFilter( + num_filter_blocks_, seed_, storage->data(), filter_stream); + rapidsmpf::CudaEvent build_event; + build_event.record(filter_stream); + while (!ch_out->is_shutdown()) { + auto msg = co_await ch_in->receive(); + if (msg.empty()) { break; } + auto chunk = msg.release(); + chunk = co_await chunk.make_available( + ctx_, + -rapidsmpf::safe_cast(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE))); + // Filter is allocated on `filter_stream`, but we run the additions on the chunk's + // stream. The addition modifies global memory but we can safely launch two + // kernels doing that concurrently because the updates are atomic. + build_event.stream_wait(chunk.stream()); + filter.add(chunk.table_view(), chunk.stream(), mr); + rapidsmpf::cuda_stream_join(filter_stream, chunk.stream(), &event); + } + if (comm_->nranks() > 1) { + auto reducer = rapidsmpf::streaming::AllReduce( + ctx_, + comm_, + br->move(std::move(storage), filter_stream), + br->move( + cudf_streaming::integrations::BloomFilter::storage(num_filter_blocks_, filter_stream, mr), + filter_stream), + tag, + [num_blocks = num_filter_blocks_, seed = seed_](rapidsmpf::Buffer const* left, + rapidsmpf::Buffer* right) { + right->write_access([&](std::byte* out_bytes, rmm::cuda_stream_view stream) { + auto const in = + cudf_streaming::integrations::BloomFilter::view(num_blocks, seed, left->data(), stream); + cudf_streaming::integrations::BloomFilter(num_blocks, seed, out_bytes, stream) + .merge(in, stream); + }); + }); + auto result = co_await reducer.extract(); + auto [res, _] = br->reserve(rapidsmpf::MemoryType::DEVICE, 0, rapidsmpf::AllowOverbooking::YES); + storage = br->move_to_device_buffer(std::move(result.second), res); + } + co_await ch_out->send(rapidsmpf::streaming::Message{0, std::move(storage), {}, {}}); + co_await ch_out->drain(ctx_->executor()); +} + +rapidsmpf::streaming::Actor BloomFilter::apply( + std::shared_ptr bloom_filter, + std::shared_ptr ch_in, + std::shared_ptr ch_out, + std::vector keys) +{ + rapidsmpf::streaming::ShutdownAtExit c{bloom_filter, ch_in, ch_out}; + co_await ctx_->executor()->schedule(); + auto storage = (co_await bloom_filter->receive()).release(); + RAPIDSMPF_EXPECTS((co_await bloom_filter->receive()).empty(), + "Bloom filter channel contained more than one message"); + auto stream = storage.stream(); + rapidsmpf::CudaEvent event; + auto filter = + cudf_streaming::integrations::BloomFilter(num_filter_blocks_, seed_, storage.data(), stream); + auto meta = co_await ch_in->receive_metadata(); + if (!meta.empty()) { co_await ch_out->send_metadata(std::move(meta)); } + while (!ch_out->is_shutdown()) { + auto msg = co_await ch_in->receive(); + if (msg.empty()) { break; } + auto chunk = msg.release(); + chunk = co_await chunk.make_available( + ctx_, + -rapidsmpf::safe_cast(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE))); + auto chunk_stream = chunk.stream(); + rapidsmpf::cuda_stream_join(chunk_stream, stream, &event); + // Reservation for the mask construction and guess at output size. + auto res = co_await ctx_->memory(rapidsmpf::MemoryType::DEVICE) + ->reserve_or_wait(rapidsmpf::safe_cast(chunk.table_view().num_rows()) + // TODO: no magic numbers: the hashing algorithm in + // `contains` below returns an int64 column. + * (1 + sizeof(std::int64_t)) + // Guess at how selective the filter is. + + chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE) / 2, + 0); + auto mask = + filter.contains(chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr()); + rapidsmpf::cuda_stream_join(stream, chunk_stream, &event); + RAPIDSMPF_EXPECTS(mask.size() == static_cast(chunk.table_view().num_rows()), + "Invalid mask size"); + auto mask_view = cudf::column_view{cudf::data_type{cudf::type_id::BOOL8}, + rapidsmpf::safe_cast(mask.size()), + mask.data(), + {}, + 0}; + auto result = cudf::apply_boolean_mask( + chunk.table_view(), mask_view, chunk_stream, ctx_->br()->device_mr()); + std::ignore = std::move(chunk); + std::ignore = std::move(res); + co_await ch_out->send(to_message( + msg.sequence_number(), + std::make_unique(std::move(result), chunk_stream))); + } + co_await ch_out->drain(ctx_->executor()); +} +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/src/streaming/channel_metadata.cpp b/cpp/libcudf_streaming/src/streaming/channel_metadata.cpp new file mode 100644 index 000000000000..6e90a03ee97a --- /dev/null +++ b/cpp/libcudf_streaming/src/streaming/channel_metadata.cpp @@ -0,0 +1,102 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace cudf_streaming::streaming { + +OrderScheme::OrderScheme(std::vector keys, + std::shared_ptr boundaries, + bool strict_boundaries) + : keys{std::move(keys)}, boundaries{std::move(boundaries)}, strict_boundaries{strict_boundaries} +{ + RAPIDSMPF_EXPECTS( + !this->keys.empty(), "OrderScheme: keys must not be empty", std::invalid_argument); + RAPIDSMPF_EXPECTS( + this->boundaries != nullptr, "OrderScheme: boundaries must not be null", std::invalid_argument); + RAPIDSMPF_EXPECTS(this->boundaries->is_available(), + "OrderScheme: boundaries must be device-resident", + std::invalid_argument); + RAPIDSMPF_EXPECTS(this->keys.size() == static_cast(this->boundaries->shape().second), + "OrderScheme: number of keys must match number of boundary columns", + std::invalid_argument); +} + +PartitioningSpec PartitioningSpec::from_order(OrderScheme o) +{ + return {.type = Type::ORDER, .hash = std::nullopt, .order = std::move(o)}; +} + +OrderScheme OrderScheme::with_keys(std::vector new_keys) const +{ + return OrderScheme(std::move(new_keys), boundaries, strict_boundaries); +} + +bool OrderScheme::boundaries_aligned_with(OrderScheme const& other, + rapidsmpf::BufferResource& br) const +{ + if (strict_boundaries != other.strict_boundaries || + boundaries->shape() != other.boundaries->shape()) { + return false; + } + if (!std::equal( + keys.begin(), keys.end(), other.keys.begin(), [](OrderKey const& a, OrderKey const& b) { + return a.order == b.order && a.null_order == b.null_order; + })) { + return false; + } + if (boundaries->shape().first == 0) { return true; } + auto const lhs = boundaries->table_view(); + auto const rhs = other.boundaries->table_view(); + auto const stream = boundaries->stream(); + rapidsmpf::cuda_stream_join(stream, other.boundaries->stream()); + for (cudf::size_type i = 0; i < lhs.num_columns(); ++i) { + auto eq = cudf::binary_operation(lhs.column(i), + rhs.column(i), + cudf::binary_operator::NULL_EQUALS, + cudf::data_type{cudf::type_id::BOOL8}, + stream, + br.device_mr()); + auto result = cudf::reduce(eq->view(), + *cudf::make_all_aggregation(), + cudf::data_type{cudf::type_id::BOOL8}, + stream, + br.device_mr()); + auto& scalar = static_cast&>(*result); + if (!scalar.value(stream)) { return false; } + } + return true; +} + +rapidsmpf::streaming::Message to_message(std::uint64_t sequence_number, + std::unique_ptr m) +{ + return rapidsmpf::streaming::Message{ + sequence_number, + std::move(m), + {}, + [](rapidsmpf::streaming::Message const& msg, + rapidsmpf::MemoryReservation& /* reservation */) -> rapidsmpf::streaming::Message { + auto copy = std::make_unique(msg.get()); + return rapidsmpf::streaming::Message{ + msg.sequence_number(), std::move(copy), {}, msg.copy_cb()}; + }}; +} + +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/src/streaming/parquet.cpp b/cpp/libcudf_streaming/src/streaming/parquet.cpp new file mode 100644 index 000000000000..7c639b283434 --- /dev/null +++ b/cpp/libcudf_streaming/src/streaming/parquet.cpp @@ -0,0 +1,407 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cudf_streaming::streaming::actor { + +namespace { + +/** + * @brief Per-context cache for file-backed messages. + * + * FileCache caches file read results by storing message copies in the associated + * Context's SpillableMessages instance. By tying cached data to the Context, the + * lifetime of cached entries matches the lifetime of the Context itself. + * + * Each cache instance is scoped to a single Context and is shared across callers + * using that Context. + * + * The cache is thread-safe. + */ +class FileCache { + public: + struct Key { + std::vector filepaths; + std::int64_t skip_rows; + std::size_t skip_bytes; + std::optional num_rows; + std::optional num_bytes; + std::optional> column_names; + std::optional> column_indices; + std::vector> row_groups; + + // Lexicographical comparison of all data members. + auto operator<=>(Key const&) const = default; + }; + + /** + * @brief Construct a FileCache. + * + * @param mem_type Memory type used for cache storage. + */ + FileCache(rapidsmpf::MemoryType mem_type = rapidsmpf::MemoryType::HOST) : mem_type_{mem_type} {} + + /** + * @brief Insert a message into the cache. + * + * The message is copied into the memory type configured for this cache + * and stored in the associated Context's SpillableMessages instance. + * + * @param ctx Streaming context. + * @param key Cache key identifying the message. + * @param msg Message to cache. + * @return True if the message was inserted, false if the key already existed. + */ + bool insert(std::shared_ptr ctx, + Key key, + rapidsmpf::streaming::Message const& msg) + { + auto reservation = ctx->br()->reserve_or_fail(msg.copy_cost(), mem_type_); + auto msg_copy = msg.copy(reservation); + + std::lock_guard lock(mutex_); + if (cache_.contains(key)) { return false; } + cache_.emplace(std::move(key), ctx->spillable_messages()->insert(std::move(msg_copy))); + return true; + } + + /** + * @brief Retrieve a cached message. + * + * If the key exists, the cached message is copied out of spillable storage + * using newly reserved memory, prioritizing memory types in `MEMORY_TYPES` + * order. + * + * @param ctx Streaming context. + * @param key Cache key to look up. + * @return The cached message, or std::nullopt if the key is not present. + */ + std::optional get( + std::shared_ptr ctx, Key const& key) const + { + auto& stats = *ctx->statistics(); + + stats.add_report_entry("unbounded_file_read_cache hits", + {"unbounded_file_read_cache hits"}, + rapidsmpf::Statistics::Formatter::HitRate); + + rapidsmpf::streaming::SpillableMessages::MessageId mid; + { + std::lock_guard lock(mutex_); + auto it = cache_.find(key); + if (it == cache_.end()) { + stats.add_stat("unbounded_file_read_cache hits", 0); + return std::nullopt; + } + mid = it->second; + } + auto const size = ctx->spillable_messages()->get_content_description(mid).content_size(); + + stats.add_stat("unbounded_file_read_cache hits", 1); + stats.add_bytes_stat("unbounded_file_read_cache saved", size); + auto reservation = ctx->br()->reserve_or_fail(size, rapidsmpf::MEMORY_TYPES); + return ctx->spillable_messages()->copy(mid, reservation); + } + + /** + * @brief Get the FileCache instance for a Context. + * + * Each Context has exactly one FileCache instance for the lifetime of the + * process. If the `unbounded_file_read_cache` option is disabled, this + * function returns nullptr. + * + * @param ctx Context used to identify the cache instance. The same Context must + * be used for all subsequent insert and get operations. + * @return Shared pointer to the per-context FileCache, or nullptr if the cache + * is disabled. + */ + static std::shared_ptr instance(std::shared_ptr ctx) + { + static std::mutex mutex; + static std::unordered_map> instances; + + std::lock_guard lock(mutex); + auto const id = ctx->uid(); + auto it = instances.find(id); + if (it != instances.end()) { return it->second; } + + // Get the memory type of the file cache, if enabled. + auto const mem_type = ctx->options().get>( + "unbounded_file_read_cache", [](auto const& s) -> std::optional { + auto val = rapidsmpf::parse_optional(s); + if (!val.has_value() || val->empty()) { return std::nullopt; } + return rapidsmpf::parse_string(s); + }); + + if (mem_type.has_value()) { + auto ret = std::make_shared(*mem_type); + instances.emplace(id, ret); + return ret; + } + return nullptr; + } + + private: + mutable std::mutex mutex_; + std::map cache_; + rapidsmpf::MemoryType mem_type_; +}; + +/** + * @brief Read a single chunk from a parquet source. + * + * @param ctx The execution context to use. + * @param stream The stream on which to read the chunk. + * @param options The parquet reader options describing the data to read. + * @param sequence_number The ordered chunk id to reconstruct original ordering of the + * data. + * @return Message representing the read chunk. + */ +rapidsmpf::streaming::Message read_parquet_chunk(std::shared_ptr ctx, + rmm::cuda_stream_view stream, + cudf::io::parquet_reader_options options, + std::uint64_t sequence_number) +{ + auto do_read_parquet = [&]() -> rapidsmpf::streaming::Message { + return to_message( + sequence_number, + std::make_unique( + cudf::io::read_parquet(options, stream, ctx->br()->device_mr()).tbl, stream)); + }; + + auto file_cache = FileCache::instance(ctx); + if (file_cache == nullptr) { return do_read_parquet(); } + + FileCache::Key key{.filepaths = options.get_source().filepaths(), + .skip_rows = options.get_skip_rows(), + .skip_bytes = options.get_skip_bytes(), + .num_rows = options.get_num_rows(), + .num_bytes = options.get_num_bytes(), + .column_names = options.get_column_names(), + .column_indices = options.get_column_indices(), + .row_groups = options.get_row_groups()}; + + auto msg = file_cache->get(ctx, key); + if (msg.has_value()) { return std::move(*msg); } + + auto ret = do_read_parquet(); + file_cache->insert(ctx, key, ret); + return ret; +} + +struct ChunkDesc { + std::uint64_t sequence_number; + std::int64_t skip_rows; + std::int64_t num_rows; + cudf::io::source_info source; +}; + +/** + * @brief Read chunks and send them to an output channel. + * + * @param ctx Execution context to use. + * @param ch_out Channel to send output to. + * @param options Template reader options. + * @param chunks List of chunks from the input files to read. Processed in order. + * @param idx Index of the next chunk to process. + * + * @return Coroutine representing the processing of all chunks. + */ +rapidsmpf::streaming::Actor produce_chunks( + std::shared_ptr ctx, + std::shared_ptr ch_out, + std::vector& chunks, + cudf::io::parquet_reader_options options) +{ + // ShutdownAtExit c{ch_out}; + co_await ctx->executor()->schedule(); + for (auto& chunk : chunks) { + cudf::io::parquet_reader_options chunk_options{options}; + chunk_options.set_skip_rows(chunk.skip_rows); + chunk_options.set_num_rows(chunk.num_rows); + chunk_options.set_source(chunk.source); + auto stream = ctx->br()->stream_pool().get_stream(); + auto ticket = co_await ch_out->acquire(); + if (!ticket.has_value()) { + // Semaphore (and hence output channel) shutdown + break; + } + // Having acquire a ticket, let's move to a new thread. + co_await ctx->executor()->schedule(); + // TODO: This reads the metadata ntasks times. + // See https://github.com/rapidsai/cudf/issues/20311 + auto [msg, exception] = [&]() -> std::pair { + try { + return {read_parquet_chunk(ctx, stream, chunk_options, chunk.sequence_number), nullptr}; + } catch (...) { + return {rapidsmpf::streaming::Message{}, std::current_exception()}; + } + }(); + if (exception != nullptr) { + co_await ch_out->shutdown(); + std::rethrow_exception(exception); + } + auto sent = co_await ticket->send(std::move(msg)); + if (!sent) { + // Output channel is shutdown, no need for more reads. + break; + } + } + co_await ch_out->drain(ctx->executor()); +} +} // namespace + +rapidsmpf::streaming::Actor read_parquet(std::shared_ptr ctx, + std::shared_ptr comm, + std::shared_ptr ch_out, + std::size_t num_producers, + cudf::io::parquet_reader_options options, + cudf::size_type num_rows_per_chunk, + std::unique_ptr filter) +{ + rapidsmpf::streaming::ShutdownAtExit c{ch_out}; + co_await ctx->executor()->schedule(); + auto const size = rapidsmpf::safe_cast(comm->nranks()); + auto const rank = rapidsmpf::safe_cast(comm->rank()); + auto source = options.get_source(); + RAPIDSMPF_EXPECTS(source.type() == cudf::io::io_type::FILEPATH, + "Only implemented for file sources"); + // TODO: To handle this we need a prefix scan across all the ranks of the total + // number of rows that would be read by previous ranks. + RAPIDSMPF_EXPECTS(size == 1 || !options.get_num_rows().has_value(), + "Reading subset of rows not yet supported in multi-rank execution"); + // TODO: To handle this we need a prefix scan across all the ranks of the total + // number of rows that would be read by previous ranks. + RAPIDSMPF_EXPECTS(size == 1 || options.get_skip_rows() == 0, + "Skipping rows not yet supported in multi-rank execution"); + auto files = source.filepaths(); + RAPIDSMPF_EXPECTS(files.size() > 0, "Must have at least one file to read"); + RAPIDSMPF_EXPECTS(!options.get_filter().has_value(), + "Do not set filter on options, use the filter argument"); + if (filter != nullptr) { + options.set_filter(filter->filter); + // Let's just join all the possible streams here rather than inducing cross-stream + // deps in the tasks + rapidsmpf::cuda_stream_join( + std::ranges::transform_view( + std::ranges::iota_view(std::size_t{0}, ctx->br()->stream_pool().get_pool_size()), + [&](auto i) { return ctx->br()->stream_pool().get_stream(i); }), + std::ranges::single_view(filter->stream)); + } + // TODO: Handle case where multiple ranks are reading from a single file. + auto const files_per_rank = + rapidsmpf::safe_cast(files.size() / size + (rank < (files.size() % size))); + auto const file_offset = + rapidsmpf::safe_cast(rank * (files.size() / size) + std::min(rank, files.size() % size)); + auto local_files = + std::vector(files.begin() + file_offset, files.begin() + file_offset + files_per_rank); + std::uint64_t sequence_number = 0; + std::vector> chunks_per_producer(num_producers); + auto const num_files = local_files.size(); + // Estimate number of rows per file + std::size_t files_per_chunk = 1; + if (num_files > 1) { + auto nrows = cudf::io::read_parquet_metadata(cudf::io::source_info(local_files[0])).num_rows(); + files_per_chunk = + nrows > 0 + ? rapidsmpf::safe_cast(std::max(num_rows_per_chunk / nrows, 1)) + : 1; + } + auto to_skip = options.get_skip_rows(); + auto to_read = options.get_num_rows().value_or(std::numeric_limits::max()); + for (std::size_t file_offset = 0; file_offset < num_files; file_offset += files_per_chunk) { + std::vector chunk_files; + auto const nchunk_files = std::min(num_files - file_offset, files_per_chunk); + std::ranges::copy_n(local_files.begin() + rapidsmpf::safe_cast(file_offset), + rapidsmpf::safe_cast(nchunk_files), + std::back_inserter(chunk_files)); + auto source = cudf::io::source_info(chunk_files); + // Must read [skip_rows, skip_rows + num_rows) from full fileset + auto chunk_rows = cudf::io::read_parquet_metadata(source).num_rows() - to_skip; + auto chunk_skip_rows = to_skip; + // If the chunk is larger than the number rows we need to skip, on the next + // iteration we don't need to skip any more rows, otherwise we must skip the + // remainder. + to_skip = std::max(0l, -chunk_rows); + while (chunk_rows > 0 && to_read > 0) { + auto rows_read = + std::min({rapidsmpf::safe_cast(num_rows_per_chunk), chunk_rows, to_read}); + chunks_per_producer[sequence_number % num_producers].emplace_back( + sequence_number, chunk_skip_rows, rows_read, source); + sequence_number++; + to_read = std::max(0l, to_read - rows_read); + chunk_skip_rows += rows_read; + chunk_rows -= rows_read; + } + } + if (std::ranges::all_of(chunks_per_producer, [](auto&& v) { return v.empty(); })) { + if (local_files.size() > 0) { + // If we're on the hook to read some files, but the skip_rows/num_rows setup + // meant our slice was empty, send an empty table of correct shape. + // Anyone with no files will just immediately close their output channel. + auto empty_opts = options; + empty_opts.set_source(cudf::io::source_info(local_files[0])); + empty_opts.set_skip_rows(0); + empty_opts.set_num_rows(0); + co_await ctx->executor()->schedule(ch_out->send( + read_parquet_chunk(ctx, ctx->br()->stream_pool().get_stream(), std::move(empty_opts), 0))); + } + } else { + std::vector read_tasks; + read_tasks.reserve(1 + num_producers); + auto lineariser = rapidsmpf::streaming::Lineariser(ctx, ch_out, num_producers); + auto queues = lineariser.get_queues(); + for (std::size_t i = 0; i < num_producers; i++) { + read_tasks.push_back(produce_chunks(ctx, queues[i], chunks_per_producer[i], options)); + } + read_tasks.push_back(lineariser.drain()); + rapidsmpf::streaming::coro_results(co_await coro::when_all(std::move(read_tasks))); + } + co_await ch_out->drain(ctx->executor()); + if (filter != nullptr) { + // Let's just join all the possible streams here rather than inducing cross-stream + // deps in the tasks + rapidsmpf::cuda_stream_join( + std::ranges::single_view(filter->stream), + std::ranges::transform_view( + std::ranges::iota_view(std::size_t{0}, ctx->br()->stream_pool().get_pool_size()), + [&](auto i) { return ctx->br()->stream_pool().get_stream(i); })); + } +} +} // namespace cudf_streaming::streaming::actor diff --git a/cpp/libcudf_streaming/src/streaming/partition.cpp b/cpp/libcudf_streaming/src/streaming/partition.cpp new file mode 100644 index 000000000000..b39e39d68e66 --- /dev/null +++ b/cpp/libcudf_streaming/src/streaming/partition.cpp @@ -0,0 +1,88 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ +#include + +#include +#include +#include +#include +#include +#include + +#include + +namespace cudf_streaming::streaming::actor { + +rapidsmpf::streaming::Actor partition_and_pack( + std::shared_ptr ctx, + std::shared_ptr ch_in, + std::shared_ptr ch_out, + std::vector columns_to_hash, + int num_partitions, + cudf::hash_id hash_function, + std::uint32_t seed) +{ + rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; + + co_await ctx->executor()->schedule(); + while (true) { + auto msg = co_await ch_in->receive(); + if (msg.empty()) { break; } + auto table = msg.release(); + auto reservation = ctx->br()->reserve_device_memory_and_spill(table.make_available_cost(), + rapidsmpf::AllowOverbooking::NO); + auto tbl = table.make_available(reservation); + + rapidsmpf::streaming::PartitionMapChunk partition_map{ + .data = cudf_streaming::integrations::partition_and_pack(tbl.table_view(), + columns_to_hash, + num_partitions, + hash_function, + seed, + tbl.stream(), + ctx->br().get())}; + + co_await ch_out->send(to_message( + msg.sequence_number(), + std::make_unique(std::move(partition_map)))); + } + co_await ch_out->drain(ctx->executor()); +} + +rapidsmpf::streaming::Actor unpack_and_concat(std::shared_ptr ctx, + std::shared_ptr ch_in, + std::shared_ptr ch_out) +{ + rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; + co_await ctx->executor()->schedule(); + while (true) { + auto msg = co_await ch_in->receive(); + if (msg.empty()) { break; } + + // If receiving a partition map, we convert it to a vector and discard + // partition IDs. + std::uint64_t seq = msg.sequence_number(); + std::vector data; + if (msg.holds()) { + auto partition_map = msg.release(); + data = rapidsmpf::to_vector(std::move(partition_map.data)); + } else { + auto partition_vec = msg.release(); + data = std::move(partition_vec.data); + } + // Get a stream for the concatenated table chunk. + auto stream = ctx->br()->stream_pool().get_stream(); + + std::unique_ptr ret = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( + std::move(data), ctx->br().get(), rapidsmpf::AllowOverbooking::NO), + stream, + ctx->br().get()); + co_await ch_out->send(to_message(seq, std::make_unique(std::move(ret), stream))); + } + co_await ch_out->drain(ctx->executor()); +} + +} // namespace cudf_streaming::streaming::actor diff --git a/cpp/libcudf_streaming/src/streaming/table_chunk.cpp b/cpp/libcudf_streaming/src/streaming/table_chunk.cpp new file mode 100644 index 000000000000..779fd9098344 --- /dev/null +++ b/cpp/libcudf_streaming/src/streaming/table_chunk.cpp @@ -0,0 +1,295 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace cudf_streaming::streaming { + +TableChunk::TableChunk(std::unique_ptr table, rmm::cuda_stream_view stream) + : table_{std::move(table)}, stream_{stream}, is_spillable_{true} +{ + RAPIDSMPF_EXPECTS(table_ != nullptr, "table pointer cannot be null", std::invalid_argument); + table_view_ = table_->view(); + data_alloc_size_[static_cast(rapidsmpf::MemoryType::DEVICE)] = table_->alloc_size(); + make_available_cost_ = 0; +} + +TableChunk::TableChunk(cudf::table_view table_view, + rmm::cuda_stream_view stream, + rapidsmpf::OwningWrapper&& owner, + ExclusiveView exclusive_view) + : owner_{std::move(owner)}, + table_view_{table_view}, + stream_{stream}, + is_spillable_{static_cast(exclusive_view)} +{ + data_alloc_size_[static_cast(rapidsmpf::MemoryType::DEVICE)] = + cudf::packed_size(table_view, stream_, rmm::mr::get_current_device_resource_ref()); + make_available_cost_ = 0; +} + +TableChunk::TableChunk(std::unique_ptr packed_data) + : packed_data_{std::move(packed_data)}, is_spillable_{true} +{ + RAPIDSMPF_EXPECTS( + packed_data_ != nullptr, "packed data pointer cannot be null", std::invalid_argument); + RAPIDSMPF_EXPECTS(!packed_data_->empty(), "packed data cannot be empty", std::invalid_argument); + // Initialize stream_ here rather than in the member-initializer list to avoid + // dereferencing packed_data_ before the null check above. + stream_ = packed_data_->data->stream(); + data_alloc_size_[static_cast(packed_data_->data->mem_type())] = + packed_data_->data->size; + if (packed_data_->data->mem_type() != rapidsmpf::MemoryType::DEVICE) { + make_available_cost_ = packed_data_->data->size; + } else { + // table data is in device memory. We can trivially unpack it and make it + // available. + table_view_ = cudf::unpack(packed_data_->metadata->data(), + reinterpret_cast(packed_data_->data->data())); + make_available_cost_ = 0; + } +} + +TableChunk::TableChunk(TableChunk&& other) noexcept + : owner_(std::move(other.owner_)), + table_(std::move(other.table_)), + packed_data_(std::move(other.packed_data_)), + table_view_(std::exchange(other.table_view_, std::nullopt)), + data_alloc_size_(other.data_alloc_size_), + make_available_cost_(other.make_available_cost_), + stream_(other.stream_), + is_spillable_(other.is_spillable_) +{ +} + +TableChunk& TableChunk::operator=(TableChunk&& other) noexcept +{ + if (this != &other) { + owner_ = std::move(other.owner_); + table_ = std::move(other.table_); + packed_data_ = std::move(other.packed_data_); + table_view_ = std::exchange(other.table_view_, std::nullopt); + data_alloc_size_ = other.data_alloc_size_; + make_available_cost_ = other.make_available_cost_; + stream_ = other.stream_; + is_spillable_ = other.is_spillable_; + } + return *this; +} + +rmm::cuda_stream_view TableChunk::stream() const noexcept { return stream_; } + +std::size_t TableChunk::data_alloc_size(rapidsmpf::MemoryType mem_type) const +{ + return data_alloc_size_.at(static_cast(mem_type)); +} + +bool TableChunk::is_available() const noexcept { return table_view_.has_value(); } + +std::size_t TableChunk::make_available_cost() const noexcept { return make_available_cost_; } + +TableChunk TableChunk::make_available(rapidsmpf::MemoryReservation& reservation) +{ + if (is_available()) { return std::move(*this); } + // Table chunk is not available. This means that the table data is not in device + // memory. We need to move the table data to device memory using a device reservation. + RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE, + "device memory reservation is required"); + RAPIDSMPF_EXPECTS(packed_data_ != nullptr, "packed data pointer cannot be null"); + auto packed_data = std::move(packed_data_); + packed_data->data = reservation.br()->move(std::move(packed_data->data), reservation); + return TableChunk{std::move(packed_data)}; +} + +TableChunk TableChunk::make_available(rapidsmpf::MemoryReservation&& reservation) +{ + rapidsmpf::MemoryReservation& res = reservation; + return make_available(res); +} + +coro::task TableChunk::make_available( + std::shared_ptr ctx, std::int64_t net_memory_delta) +{ + co_return make_available(co_await reserve_memory(ctx, make_available_cost(), net_memory_delta)); +} + +cudf::table_view TableChunk::table_view() const +{ + RAPIDSMPF_EXPECTS(is_available(), + "the table view is unavailable, please make sure it is " + "unspilled and unpacked (see `make_available`).", + std::invalid_argument); + return table_view_.value(); +} + +bool TableChunk::is_spillable() const { return is_spillable_; } + +TableChunk TableChunk::copy(rapidsmpf::MemoryReservation& reservation) const +{ + // This method handles the two possible cases. Note that + // `!is_available() && packed_data_ == nullptr` is an invalid state, so the + // remaining valid combinations collapse into: + // + // 1. The chunk is available and not yet packed. The table is copied/packed + // into the reservation-specified memory type using libcudf: + // a. DEVICE - cudf-copy table_view() into device memory. + // b. PINNED_HOST - cudf::pack table_view() directly into pinned memory. + // c. HOST - cudf::pack table_view() into intermediate device + // memory and then copy to host memory. + // + // 2. The chunk data is already packed (packed_data_ != nullptr). + // Use buffer_copy() to copy the packed data into the reservation- + // specified memory type. The original memory type of the chunk does + // not matter. + rapidsmpf::BufferResource* br = reservation.br(); + + // If the table view is available and the table is not packed, we can use libcudf to + // copy the table in device memory, or pack it to pinned/ host memory. Else, fall + // through to case 2 (ie. use buffer_copy). + if (is_available() && packed_data_ == nullptr) { + switch (reservation.mem_type()) { + case rapidsmpf::MemoryType::DEVICE: // Case 1a. + { + // Use libcudf to copy the table_view(). + auto const nbytes = data_alloc_size(rapidsmpf::MemoryType::DEVICE); + auto statistics = br->statistics(); + rapidsmpf::StreamOrderedTiming timing{stream(), statistics}; + auto table = std::make_unique(table_view(), stream(), br->device_mr()); + statistics->record_copy( + rapidsmpf::MemoryType::DEVICE, rapidsmpf::MemoryType::DEVICE, nbytes, std::move(timing)); + // And update the provided `reservation`. + br->release(reservation, nbytes); + return TableChunk(std::move(table), stream()); + } + case rapidsmpf::MemoryType::PINNED_HOST: // Case 1b. + { + rapidsmpf::StreamOrderedTiming timing{stream(), br->statistics()}; + + // use cudf pack with pinned mr + auto packed_pinned = cudf::pack(table_view(), stream(), br->pinned_mr()); + auto nbytes = packed_pinned.gpu_data->size(); + + br->statistics()->record_copy(rapidsmpf::MemoryType::DEVICE, + rapidsmpf::MemoryType::PINNED_HOST, + nbytes, + std::move(timing)); + // update the provided `reservation` + br->release(reservation, nbytes); + auto host_buffer = br->move(std::move(packed_pinned.gpu_data), stream()); + return TableChunk(std::make_unique(std::move(packed_pinned.metadata), + std::move(host_buffer))); + } + case rapidsmpf::MemoryType::HOST: // Case 1c. + { + // We use libcudf's pack() to serialize `table_view()` into a + // packed_columns and then we move the packed_columns' gpu_data to a + // new host buffer. + // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, + // `cudf::pack()` allocates device memory we haven't reserved. + auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr()); + auto packed_data = std::make_unique( + std::move(packed_columns.metadata), + br->move(std::move(packed_columns.gpu_data), stream())); + + // Handle the case where `cudf::pack` allocates slightly more than the + // input size. This can occur because cudf uses aligned allocations, + // which may exceed the requested size. To accommodate this, we + // allow some wiggle room. + if (packed_data->data->size > reservation.size()) { + auto const wiggle_room = 1024 * static_cast(table_view().num_columns()); + if (packed_data->data->size <= reservation.size() + wiggle_room) { + reservation = + br->reserve( + reservation.mem_type(), packed_data->data->size, rapidsmpf::AllowOverbooking::YES) + .first; + } + } + packed_data->data = br->move(std::move(packed_data->data), reservation); + return TableChunk(std::move(packed_data)); + } + default: RAPIDSMPF_FAIL("MemoryType: unknown"); + } + } + // `!is_available() && packed_data_ == nullptr` is an invalid state, so + // reaching this point implies `packed_data_ != nullptr`. + RAPIDSMPF_EXPECTS(packed_data_ != nullptr, "something went wrong"); + + // Case 2. The chunk data is already packed (packed_data_ != nullptr). We need + // to copy the packed data into the reservation-specified memory type. + auto const nbytes = packed_data_->data->size; + auto metadata = std::make_unique>(*packed_data_->metadata); + auto data = br->make_buffer(nbytes, packed_data_->stream(), reservation); + rapidsmpf::buffer_copy(br->statistics(), *data, *packed_data_->data, nbytes); + return TableChunk(std::make_unique(std::move(metadata), std::move(data))); +} + +std::unique_ptr TableChunk::into_packed_data( + rapidsmpf::BufferResource* br) && +{ + if (packed_data_) { + table_view_ = std::nullopt; + return std::move(packed_data_); + } + RAPIDSMPF_EXPECTS(is_available(), "TableChunk must be available; call make_available() first"); + // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, + // `cudf::pack()` allocates device memory we haven't reserved. + auto packed_columns = cudf::pack(table_view_.value(), stream_, br->device_mr()); + table_view_ = std::nullopt; + return std::make_unique( + std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream_)); +} + +std::pair TableChunk::shape() const noexcept +{ + if (packed_data_ != nullptr) { + auto view = cudf::packed_metadata_view(*packed_data_->metadata); + return {view.num_rows(), view.num_columns()}; + } + assert(table_view_.has_value() && "shape() called on moved-from TableChunk"); + return {table_view_->num_rows(), table_view_->num_columns()}; +} + +rapidsmpf::ContentDescription get_content_description(TableChunk const& obj) +{ + rapidsmpf::ContentDescription ret{obj.is_spillable() + ? rapidsmpf::ContentDescription::Spillable::YES + : rapidsmpf::ContentDescription::Spillable::NO}; + for (auto mem_type : rapidsmpf::MEMORY_TYPES) { + ret.content_size(mem_type) = obj.data_alloc_size(mem_type); + } + return ret; +} + +rapidsmpf::streaming::Message to_message(std::uint64_t sequence_number, + std::unique_ptr chunk) +{ + auto cd = get_content_description(*chunk); + return rapidsmpf::streaming::Message{ + sequence_number, + std::move(chunk), + cd, + [](rapidsmpf::streaming::Message const& msg, + rapidsmpf::MemoryReservation& reservation) -> rapidsmpf::streaming::Message { + auto const& self = msg.get(); + auto chunk = std::make_unique(self.copy(reservation)); + auto cd = get_content_description(*chunk); + return rapidsmpf::streaming::Message{ + msg.sequence_number(), std::move(chunk), cd, msg.copy_cb()}; + }}; +} + +} // namespace cudf_streaming::streaming diff --git a/cpp/libcudf_streaming/tests/CMakeLists.txt b/cpp/libcudf_streaming/tests/CMakeLists.txt new file mode 100644 index 000000000000..600dbd51f74a --- /dev/null +++ b/cpp/libcudf_streaming/tests/CMakeLists.txt @@ -0,0 +1,81 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +# ################################################################################################## +# enable testing ----------------------------------------------------------------------------------- +# ################################################################################################## +enable_testing() + +include(rapids-test) +rapids_test_init() + +file(WRITE "${CUDF_STREAMING_BINARY_DIR}/CTestTestfile.cmake" "subdirs(\"tests\")\n") + +set(_cudf_streaming_gtests_link "${CMAKE_CURRENT_BINARY_DIR}/gtests") +if(NOT EXISTS "${_cudf_streaming_gtests_link}") + file(CREATE_LINK "${CUDF_STREAMING_BINARY_DIR}/gtests" "${_cudf_streaming_gtests_link}" SYMBOLIC) +endif() + +add_library(cudf_streaming_test_sources OBJECT) +set_target_properties( + cudf_streaming_test_sources + PROPERTIES CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON +) +target_include_directories( + cudf_streaming_test_sources PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}" +) +target_link_libraries( + cudf_streaming_test_sources + PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf::cudftestutil cudf::cudftestutil_impl + PUBLIC GTest::gmock GTest::gtest +) +target_sources( + cudf_streaming_test_sources + PRIVATE streaming/test_table_chunk.cpp + streaming/test_read_parquet.cpp + streaming/test_channel_metadata.cpp + streaming/test_partition.cpp + streaming/test_cudf_utils.cpp + test_partition.cpp +) + +add_executable(cudf_streaming_single_tests main/single.cpp) +set_target_properties( + cudf_streaming_single_tests + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/gtests" + INSTALL_RPATH "\$ORIGIN/../../../lib" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON +) +target_include_directories( + cudf_streaming_single_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}" +) +target_link_libraries( + cudf_streaming_single_tests + PRIVATE cudf_streaming rapidsmpf::rapidsmpf GTest::gmock GTest::gtest + $ cudf_streaming_test_sources +) +rapids_test_add( + NAME cudf_streaming_single_tests + COMMAND cudf_streaming_single_tests + GPUS 1 + PERCENT 100 + INSTALL_COMPONENT_SET testing +) + +rapids_test_install_relocatable( + INSTALL_COMPONENT_SET testing DESTINATION bin/gtests/libcudf_streaming +) diff --git a/cpp/libcudf_streaming/tests/environment.hpp b/cpp/libcudf_streaming/tests/environment.hpp new file mode 100644 index 000000000000..648d010ea11b --- /dev/null +++ b/cpp/libcudf_streaming/tests/environment.hpp @@ -0,0 +1,38 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +enum class TestEnvironmentType : int { + SINGLE, +}; + +class Environment : public ::testing::Environment { + public: + Environment(int argc, char** argv); + + void SetUp() override; + + void TearDown() override; + + void barrier(); + + [[nodiscard]] TestEnvironmentType type() const; + + constexpr rapidsmpf::config::Options& options() { return options_; } + + std::shared_ptr split_comm(); + + std::shared_ptr comm_; + + private: + std::shared_ptr split_comm_{nullptr}; + rapidsmpf::config::Options options_; +}; + +extern Environment* GlobalEnvironment; diff --git a/cpp/libcudf_streaming/tests/main/single.cpp b/cpp/libcudf_streaming/tests/main/single.cpp new file mode 100644 index 000000000000..22a4a9286790 --- /dev/null +++ b/cpp/libcudf_streaming/tests/main/single.cpp @@ -0,0 +1,45 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../environment.hpp" + +#include + +#include +#include + +#include + +Environment* GlobalEnvironment = nullptr; + +Environment::Environment(int, char**) {} + +TestEnvironmentType Environment::type() const { return TestEnvironmentType::SINGLE; } + +void Environment::SetUp() +{ + options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables()); + comm_ = + std::make_shared(options_, std::make_shared()); + split_comm_ = comm_; +} + +void Environment::TearDown() +{ + split_comm_ = nullptr; + comm_ = nullptr; +} + +void Environment::barrier() {} + +std::shared_ptr Environment::split_comm() { return split_comm_; } + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + GlobalEnvironment = new Environment(argc, argv); + ::testing::AddGlobalTestEnvironment(GlobalEnvironment); + return RUN_ALL_TESTS(); +} diff --git a/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp b/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp new file mode 100644 index 000000000000..415e34a6f046 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp @@ -0,0 +1,58 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include "../environment.hpp" + +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +extern Environment* GlobalEnvironment; + +class BaseStreamingFixture : public ::testing::Test { + protected: + void SetUp() override + { + SetUpWithThreads(1); // default number of streaming threads + } + + void TearDown() override + { + ctx.reset(); + br.reset(); + } + + void SetUpWithThreads(int num_streaming_threads, + std::unordered_map memory_limits = {}) + { + // create a new options object, since we can not modify values in the global + // options object + auto env_vars = rapidsmpf::config::get_environment_variables(); + env_vars["num_streaming_threads"] = std::to_string(num_streaming_threads); + rapidsmpf::config::Options options(std::move(env_vars)); + + stream = cudf::get_default_stream(); + br = rapidsmpf::BufferResource::create( + mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, std::move(memory_limits)); + ctx = std::make_shared( + std::move(options), GlobalEnvironment->comm_->logger(), br); + } + + rmm::cuda_stream_view stream; + rmm::mr::cuda_memory_resource mr_cuda; + std::shared_ptr br; + std::shared_ptr ctx; +}; diff --git a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp new file mode 100644 index 000000000000..95dec084ee1b --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp @@ -0,0 +1,227 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include + +using namespace cudf_streaming::streaming; + +class StreamingChannelMetadata : public ::testing::Test {}; + +TEST_F(StreamingChannelMetadata, HashScheme) +{ + HashScheme h{{0, 1}, 16}; + EXPECT_EQ(h.column_indices.size(), 2); + EXPECT_EQ(h.column_indices[0], 0); + EXPECT_EQ(h.column_indices[1], 1); + EXPECT_EQ(h.modulus, 16); + + // Equality + EXPECT_EQ(h, (HashScheme{{0, 1}, 16})); + EXPECT_NE(h, (HashScheme{{0, 1}, 32})); + EXPECT_NE(h, (HashScheme{{2}, 16})); +} + +TEST_F(StreamingChannelMetadata, OrderSchemeCtorRejectsEmptyKeys) +{ + EXPECT_THROW(static_cast(OrderScheme({}, nullptr)), std::invalid_argument); +} + +TEST_F(StreamingChannelMetadata, OrderSchemeCtorRejectsNullBoundaries) +{ + EXPECT_THROW(static_cast( + OrderScheme({{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}}, nullptr)), + std::invalid_argument); +} + +TEST_F(StreamingChannelMetadata, PartitioningSpec) +{ + // None + auto spec_none = PartitioningSpec::none(); + EXPECT_EQ(spec_none.type, PartitioningSpec::Type::NONE); + + // Inherit + auto spec_inherit = PartitioningSpec::inherit(); + EXPECT_EQ(spec_inherit.type, PartitioningSpec::Type::INHERIT); + + // Hash + auto spec_hash = PartitioningSpec::from_hash(HashScheme{{0}, 16}); + EXPECT_EQ(spec_hash.type, PartitioningSpec::Type::HASH); + EXPECT_EQ(spec_hash.hash->column_indices[0], 0); + EXPECT_EQ(spec_hash.hash->modulus, 16); + + // Type checks (operator== removed; use field comparisons) + EXPECT_EQ(spec_none.type, PartitioningSpec::Type::NONE); + EXPECT_EQ(spec_inherit.type, PartitioningSpec::Type::INHERIT); + EXPECT_EQ(spec_hash.type, PartitioningSpec::Type::HASH); + EXPECT_NE(spec_none.type, spec_inherit.type); + EXPECT_EQ(spec_hash.hash->modulus, 16); + EXPECT_NE((PartitioningSpec::from_hash(HashScheme{{0}, 32}).hash->modulus), 16); +} + +TEST_F(StreamingChannelMetadata, PartitioningScenarios) +{ + // Default construction + Partitioning p_default{}; + EXPECT_EQ(p_default.inter_rank.type, PartitioningSpec::Type::NONE); + EXPECT_EQ(p_default.local.type, PartitioningSpec::Type::NONE); + + // Direct global shuffle: inter_rank=Hash, local=Inherit + Partitioning p_global{PartitioningSpec::from_hash(HashScheme{{0}, 16}), + PartitioningSpec::inherit()}; + EXPECT_EQ(p_global.inter_rank.type, PartitioningSpec::Type::HASH); + EXPECT_EQ(p_global.local.type, PartitioningSpec::Type::INHERIT); + EXPECT_EQ(p_global.inter_rank.hash->modulus, 16); + + // Two-stage shuffle: inter_rank=Hash(nranks), local=Hash(N_l) + Partitioning p_twostage{PartitioningSpec::from_hash(HashScheme{{0}, 4}), + PartitioningSpec::from_hash(HashScheme{{0}, 8})}; + EXPECT_EQ(p_twostage.inter_rank.hash->modulus, 4); + EXPECT_EQ(p_twostage.local.hash->modulus, 8); + + // Field comparisons (Partitioning::operator== removed) + { + Partitioning p_same{PartitioningSpec::from_hash(HashScheme{{0}, 16}), + PartitioningSpec::inherit()}; + EXPECT_EQ(p_global.inter_rank.type, p_same.inter_rank.type); + EXPECT_EQ(p_global.inter_rank.hash->modulus, p_same.inter_rank.hash->modulus); + EXPECT_EQ(p_global.local.type, p_same.local.type); + } + EXPECT_NE(p_global.inter_rank.hash->modulus, p_twostage.inter_rank.hash->modulus); +} + +TEST_F(StreamingChannelMetadata, ChannelMetadata) +{ + // Full construction - use std::move to avoid GCC false positive on vector copy + Partitioning p{PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit()}; + ChannelMetadata m{4, std::move(p), true}; + EXPECT_EQ(m.local_count, 4); + EXPECT_EQ(m.partitioning.inter_rank.type, PartitioningSpec::Type::HASH); + EXPECT_EQ(m.partitioning.local.type, PartitioningSpec::Type::INHERIT); + EXPECT_TRUE(m.duplicated); + + // Minimal construction + ChannelMetadata m_minimal{4}; + EXPECT_EQ(m_minimal.local_count, 4); + EXPECT_FALSE(m_minimal.duplicated); + + // Equality - create fresh partitionings and move them + ChannelMetadata m_same{ + 4, + Partitioning{PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit()}, + true}; + ChannelMetadata m_diff{ + 8, + Partitioning{PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit()}, + true}; + // Field comparisons (ChannelMetadata::operator== removed) + EXPECT_EQ(m.local_count, m_same.local_count); + EXPECT_EQ(m.duplicated, m_same.duplicated); + EXPECT_EQ(m.partitioning.inter_rank.hash->modulus, m_same.partitioning.inter_rank.hash->modulus); + EXPECT_NE(m.local_count, m_diff.local_count); +} + +TEST_F(StreamingChannelMetadata, MessageRoundTrip) +{ + // ChannelMetadata round-trip + Partitioning part{PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit()}; + auto m = std::make_unique(4, std::move(part), false); + auto msg_m = to_message(99, std::move(m)); + EXPECT_EQ(msg_m.sequence_number(), 99); + EXPECT_TRUE(msg_m.holds()); + auto released = msg_m.release(); + EXPECT_EQ(released.local_count, 4); + EXPECT_FALSE(released.duplicated); + EXPECT_EQ(released.partitioning.inter_rank.hash->modulus, 16); + EXPECT_TRUE(msg_m.empty()); +} + +class StreamingChannelMetadataGPU : public ::testing::Test { + protected: + rmm::cuda_stream_view stream{cudf::get_default_stream()}; + std::shared_ptr br = + rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + + std::shared_ptr make_chunk(std::vector vals) + { + rmm::device_buffer buf(vals.data(), vals.size() * sizeof(int32_t), stream); + auto col = std::make_unique(cudf::data_type{cudf::type_id::INT32}, + static_cast(vals.size()), + std::move(buf), + rmm::device_buffer{}, + 0); + std::vector> cols; + cols.push_back(std::move(col)); + return std::make_shared(std::make_unique(std::move(cols)), stream); + } +}; + +TEST_F(StreamingChannelMetadataGPU, OrderSchemeReplaceKeys) +{ + OrderKey k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}; + OrderKey k5{5, cudf::order::DESCENDING, cudf::null_order::AFTER}; + + auto b = make_chunk({100, 200}); + OrderScheme o1({k0}, b); + auto o2 = o1.with_keys({k5}); + + EXPECT_EQ(o2.keys[0].column_index, 5); + EXPECT_EQ(o2.keys[0].order, cudf::order::DESCENDING); + EXPECT_EQ(o2.strict_boundaries, o1.strict_boundaries); + EXPECT_EQ(o2.boundaries->shape(), o1.boundaries->shape()); + EXPECT_EQ(o2.boundaries.get(), b.get()); + EXPECT_NE(o1.keys[0].column_index, o2.keys[0].column_index); + + EXPECT_THROW(static_cast(o1.with_keys({k0, k5})), std::invalid_argument); +} + +TEST_F(StreamingChannelMetadataGPU, OrderSchemeBoundariesAlignedWith) +{ + OrderKey k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}; + OrderKey k3{3, cudf::order::ASCENDING, cudf::null_order::BEFORE}; + + OrderScheme o1({k0}, make_chunk({100, 200})); + OrderScheme o2({k0}, make_chunk({100, 200})); + EXPECT_TRUE(o1.boundaries_aligned_with(o2, *br)); + + OrderScheme o_shifted({k3}, make_chunk({100, 200})); + EXPECT_TRUE(o1.boundaries_aligned_with(o_shifted, *br)); + + OrderScheme o_strict({k0}, make_chunk({100, 200}), /*strict=*/true); + EXPECT_FALSE(o1.boundaries_aligned_with(o_strict, *br)); + + OrderScheme o_diff({k0}, make_chunk({100, 300})); + EXPECT_FALSE(o1.boundaries_aligned_with(o_diff, *br)); +} + +TEST_F(StreamingChannelMetadataGPU, PartitioningSpecOrder) +{ + OrderKey k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}; + OrderScheme o({k0}, make_chunk({100, 200})); + + auto spec = PartitioningSpec::from_order(o); + EXPECT_EQ(spec.type, PartitioningSpec::Type::ORDER); + EXPECT_TRUE(spec.order.has_value()); + EXPECT_EQ(spec.order->keys[0].column_index, 0); + + // Type checks only (PartitioningSpec::operator== removed; ORDER value comparison + // requires boundaries_aligned_with on the OrderScheme directly) + EXPECT_EQ(spec.type, PartitioningSpec::Type::ORDER); + EXPECT_NE(spec.type, PartitioningSpec::from_hash(HashScheme{{0}, 16}).type); + EXPECT_NE(spec.type, PartitioningSpec::none().type); +} diff --git a/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp b/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp new file mode 100644 index 000000000000..a15559f54f36 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp @@ -0,0 +1,177 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include + +using namespace cudf_streaming::integrations; + +class BaseEstimatedMemoryUsageTest : public ::testing::Test { + protected: + void SetUp() override { stream = cudf::get_default_stream(); } + + rmm::cuda_stream_view stream; +}; + +/** + * @brief Templated test suite for testing estimated_memory_usage with different column + * types + */ +template +class EstimatedMemoryUsageTest : public BaseEstimatedMemoryUsageTest {}; + +// Define the types to test +using ColumnTypes = + ::testing::Types; + +TYPED_TEST_SUITE(EstimatedMemoryUsageTest, ColumnTypes); + +TYPED_TEST(EstimatedMemoryUsageTest, FixedWidthColumnMemoryUsage) +{ + using T = TypeParam; + // Test with different sizes + std::vector test_sizes = {0, 1, 10, 100, 1000, 1000000}; + + for (auto size : test_sizes) { + SCOPED_TRACE("test size: " + std::to_string(size)); + std::vector data(size); + + cudf::test::fixed_width_column_wrapper wrapper(data.begin(), data.end()); + auto column = wrapper.release(); + + std::size_t exp = column->alloc_size(); + std::size_t est = + cudf_streaming::integrations::estimated_memory_usage(column->view(), this->stream); + + EXPECT_EQ(exp, est); + } +} + +/** + * @brief Test suite for string column memory usage estimation + */ +TEST_F(BaseEstimatedMemoryUsageTest, StringType) +{ + // Test with different string data + std::vector> test_cases = { + {}, // Empty column + {"hello"}, // Single string + {"hello", "world", "test"}, // Multiple strings + {"café", "こんにちは", "gpu"}, // Multi-byte UTF-8 + {"", "a", "very long string that should take more memory", "short"}, // Mixed lengths + std::vector(100, "repeated string") // Many repeated strings + }; + + for (const auto& data : test_cases) { + // Create a string column + cudf::test::strings_column_wrapper wrapper(data.begin(), data.end()); + auto column = wrapper.release(); + + std::size_t exp = column->alloc_size(); + std::size_t est = cudf_streaming::integrations::estimated_memory_usage(column->view(), stream); + + EXPECT_EQ(exp, est); + } +} + +/** + * @brief Test suite for list column memory usage estimation + */ +TEST_F(BaseEstimatedMemoryUsageTest, ListType) +{ + // Test with different list data + std::vector> test_cases = { + {}, // Empty column + {1, 2, 3}, // Single list + {1, 2, 3, 4, 5, 6}, // Multiple values + {0, 1, 2, 3, 4, 5}, // Mixed values + std::vector(100, 42) // Many repeated values + }; + + for (const auto& data : test_cases) { + // Create a list column + cudf::test::lists_column_wrapper wrapper(data.begin(), data.end()); + auto column = wrapper.release(); + + std::size_t exp = column->alloc_size(); + std::size_t est = cudf_streaming::integrations::estimated_memory_usage(column->view(), stream); + + EXPECT_EQ(exp, est); + } +} + +/** + * @brief Test suite for struct column memory usage estimation + */ +TEST_F(BaseEstimatedMemoryUsageTest, StructType) +{ + // Test with different struct data configurations + std::vector>> test_cases = { + {}, // Empty struct column + {{std::make_pair(1, "hello")}}, // Single struct + {{std::make_pair(1, "hello"), std::make_pair(2, "world")}}, // Two structs + {{std::make_pair(0, ""), + std::make_pair(100, "very long string"), + std::make_pair(42, "short")}}, // Mixed data + std::vector>( + 50, std::make_pair(42, "repeated")) // Many repeated structs + }; + + for (const auto& data : test_cases) { + // Create struct columns for each field + std::vector int_data; + std::vector string_data; + + for (const auto& item : data) { + int_data.push_back(item.first); + string_data.push_back(item.second); + } + + cudf::test::fixed_width_column_wrapper int_wrapper(int_data.begin(), + int_data.end()); + cudf::test::strings_column_wrapper string_wrapper(string_data.begin(), string_data.end()); + + std::vector> children; + children.push_back(int_wrapper.release()); + children.push_back(string_wrapper.release()); + + cudf::test::structs_column_wrapper wrapper(std::move(children)); + auto column = wrapper.release(); + + std::size_t exp = column->alloc_size(); + std::size_t est = cudf_streaming::integrations::estimated_memory_usage(column->view(), stream); + + EXPECT_EQ(exp, est); + } +} + +/** + * @brief Test suite for dictionary column memory usage estimation + */ +TEST_F(BaseEstimatedMemoryUsageTest, DictionaryType) +{ + // Test with different dictionary data + std::vector> test_cases = { + {}, // Empty column + {"hello"}, // Single value + {"hello", "world", "test", "hello", "world"}, // Repeated values + {"", "a", "very long string", "short", "a", "very long string"}, // Mixed with repetition + std::vector(100, "repeated") // Many repeated values + }; + + for (const auto& data : test_cases) { + // Create a dictionary column + cudf::test::dictionary_column_wrapper wrapper(data.begin(), data.end()); + auto column = wrapper.release(); + + std::size_t exp = column->alloc_size(); + std::size_t est = cudf_streaming::integrations::estimated_memory_usage(column->view(), stream); + + EXPECT_EQ(exp, est); + } +} diff --git a/cpp/libcudf_streaming/tests/streaming/test_partition.cpp b/cpp/libcudf_streaming/tests/streaming/test_partition.cpp new file mode 100644 index 000000000000..7217f639e3c5 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_partition.cpp @@ -0,0 +1,141 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../utils.hpp" +#include "base_streaming_fixture.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace cudf_streaming::streaming; + +using StreamingPartition = BaseStreamingFixture; + +TEST_F(StreamingPartition, PackUnpackRoundTrip) +{ + int const num_partitions = 5; + int const num_rows = 100; + int const num_chunks = 10; + std::int64_t const seed = 42; + constexpr cudf::hash_id hash_function = cudf::hash_id::HASH_MURMUR3; + + std::vector expects; + for (int i = 0; i < num_chunks; ++i) { + expects.push_back(random_table_with_index(seed + i, num_rows, 0, 10)); + } + + std::vector inputs; + for (int i = 0; i < num_chunks; ++i) { + inputs.emplace_back(to_message( + i, + std::make_unique( + std::make_unique(expects[i], stream, ctx->br()->device_mr()), stream))); + } + + // Create and run the streaming pipeline. + std::vector outputs; + { + std::vector actors; + auto ch1 = ctx->create_channel(); + actors.push_back(rapidsmpf::streaming::actor::push_to_channel(ctx, ch1, std::move(inputs))); + + auto ch2 = ctx->create_channel(); + actors.push_back(cudf_streaming::streaming::actor::partition_and_pack( + ctx, ch1, ch2, {1}, num_partitions, hash_function, seed)); + + auto ch3 = ctx->create_channel(); + actors.push_back(cudf_streaming::streaming::actor::unpack_and_concat(ctx, ch2, ch3)); + + actors.push_back(rapidsmpf::streaming::actor::pull_from_channel(ctx, ch3, outputs)); + + rapidsmpf::streaming::run_actor_network(std::move(actors)); + } + + EXPECT_EQ(expects.size(), outputs.size()); + for (std::size_t i = 0; i < expects.size(); ++i) { + EXPECT_EQ(outputs[i].sequence_number(), i); + auto output = outputs[i].release(); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(sort_table(output.table_view()), + sort_table(expects[i].view())); + } +} + +TEST_F(StreamingPartition, PartitionMapChunkToMessage) +{ + constexpr std::uint64_t seq = 42; + std::unordered_map data; + data.emplace(0, generate_packed_data(10, 0, stream, *br)); + data.emplace(1, generate_packed_data(10, 10, stream, *br)); + auto chunk = std::make_unique(std::move(data)); + + rapidsmpf::streaming::Message m = to_message(seq, std::move(chunk)); + EXPECT_FALSE(m.empty()); + EXPECT_TRUE(m.holds()); + EXPECT_TRUE(m.content_description().spillable()); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 80); + EXPECT_EQ(m.sequence_number(), seq); + + auto res = br->reserve_or_fail(m.copy_cost(), rapidsmpf::MemoryType::DEVICE); + rapidsmpf::streaming::Message m2 = m.copy(res); + EXPECT_EQ(res.size(), 0); + EXPECT_FALSE(m2.empty()); + EXPECT_TRUE(m2.holds()); + EXPECT_TRUE(m2.content_description().spillable()); + EXPECT_EQ(m2.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m2.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 80); + + auto chunk2 = m2.release(); + validate_packed_data(std::move(chunk2.data.at(0)), 10, 0, stream, *br); + validate_packed_data(std::move(chunk2.data.at(1)), 10, 10, stream, *br); +} + +TEST_F(StreamingPartition, PartitionMapChunkContentDescription) +{ + // Create a packed data, one in device and one in host memory. + std::unordered_map data; + auto pack1 = generate_packed_data(5, 0, stream, *br); + auto pack1_size = pack1.data->size; + auto pack2_size = pack1.data->size * 2; + auto res = br->reserve_or_fail(pack2_size, rapidsmpf::MemoryType::HOST); + auto pack2 = generate_packed_data(10, 0, stream, *br).copy(res); + data.emplace(0, std::move(pack1)); + data.emplace(1, std::move(pack2)); + + auto chunk = std::make_unique(std::move(data)); + auto cd = get_content_description(*chunk); + EXPECT_TRUE(cd.spillable()); + EXPECT_EQ(cd.content_size(rapidsmpf::MemoryType::DEVICE), pack1_size); + EXPECT_EQ(cd.content_size(rapidsmpf::MemoryType::HOST), pack2_size); +} + +TEST_F(StreamingPartition, PartitionVectorChunkContentDescription) +{ + // Create a packed data, one in device and one in host memory. + std::vector data; + auto pack1 = generate_packed_data(5, 0, stream, *br); + auto pack1_size = pack1.data->size; + auto pack2_size = pack1.data->size * 2; + auto res = br->reserve_or_fail(pack2_size, rapidsmpf::MemoryType::HOST); + auto pack2 = generate_packed_data(10, 0, stream, *br).copy(res); + data.push_back(std::move(pack1)); + data.push_back(std::move(pack2)); + + auto chunk = std::make_unique(std::move(data)); + auto cd = get_content_description(*chunk); + EXPECT_TRUE(cd.spillable()); + EXPECT_EQ(cd.content_size(rapidsmpf::MemoryType::DEVICE), pack1_size); + EXPECT_EQ(cd.content_size(rapidsmpf::MemoryType::HOST), pack2_size); +} diff --git a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp new file mode 100644 index 000000000000..b0b4205b54b7 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp @@ -0,0 +1,245 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "base_streaming_fixture.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace cudf_streaming::streaming; + +class StreamingReadParquet : public BaseStreamingFixture { + protected: + void SetUp() override + { + BaseStreamingFixture::SetUp(); + constexpr int nfiles = 10; + constexpr int nrows = 10; + + temp_dir = std::filesystem::temp_directory_path() / "rapidsmpf_read_parquet_test"; + + for (int i = 0; i < nfiles; ++i) { + std::ostringstream filename_stream; + filename_stream << std::setw(3) << std::setfill('0') << i << ".pq"; + std::filesystem::path filepath = temp_dir / filename_stream.str(); + source_files.push_back(filepath.string()); + } + + if (GlobalEnvironment->comm_->rank() == 0) { + std::filesystem::create_directories(temp_dir); + + int start = 0; + for (auto& file : source_files) { + auto values = std::ranges::iota_view(start, start + nrows); + cudf::test::fixed_width_column_wrapper col(values.begin(), values.end()); + + std::vector> columns; + columns.push_back(col.release()); + auto table = std::make_unique(std::move(columns)); + + cudf::io::sink_info sink{file}; + auto options = cudf::io::parquet_writer_options::builder(sink, table->view()).build(); + cudf::io::write_parquet(options); + start += nrows + nrows / 2; + } + } + + GlobalEnvironment->barrier(); + } + + void TearDown() override + { + GlobalEnvironment->barrier(); + + if (GlobalEnvironment->comm_->rank() == 0 && std::filesystem::exists(temp_dir)) { + std::filesystem::remove_all(temp_dir); + } + + BaseStreamingFixture::TearDown(); + } + + [[nodiscard]] cudf::io::source_info get_source_info(bool truncate_file_list) const + { + if (truncate_file_list) { + std::vector files(source_files.begin(), source_files.begin() + 2); + return cudf::io::source_info(files); + } else { + return cudf::io::source_info(source_files); + } + } + + std::filesystem::path temp_dir; + std::vector source_files; +}; + +using ReadParquetParams = + std::tuple, std::optional, bool, bool>; + +class StreamingReadParquetParams : public StreamingReadParquet, + public ::testing::WithParamInterface {}; + +INSTANTIATE_TEST_SUITE_P(ReadParquetCombinations, + StreamingReadParquetParams, + ::testing::Combine( + // skip_rows + ::testing::Values(std::nullopt, + std::optional{7}, + std::optional{19}, + std::optional{113}), + // num_rows + ::testing::Values(std::nullopt, + std::optional{0}, + std::optional{3}, + std::optional{31}, + std::optional{83}), + // use_filter + ::testing::Values(false, true), + // truncate file list + ::testing::Values(false, true)), + [](const ::testing::TestParamInfo& info) { + auto const& skip_rows = std::get<0>(info.param); + auto const& num_rows = std::get<1>(info.param); + auto const& use_filter = std::get<2>(info.param); + auto const& truncate_file_list = std::get<3>(info.param); + std::string result = "skip_rows_"; + result += + skip_rows.has_value() ? std::to_string(skip_rows.value()) : "none"; + result += "_num_rows_"; + result += + num_rows.has_value() ? std::to_string(num_rows.value()) : "all"; + if (use_filter) { + result += "_with_filter"; + } else { + result += "_no_filter"; + } + if (truncate_file_list) { + result += "_one_file"; + } else { + result += "_all_files"; + } + return result; + }); + +TEST_P(StreamingReadParquetParams, ReadParquet) +{ + auto [skip_rows, num_rows, use_filter, truncate_file_list] = GetParam(); + auto source = get_source_info(truncate_file_list); + + auto options = cudf::io::parquet_reader_options::builder(source).build(); + if (skip_rows.has_value()) { options.set_skip_rows(skip_rows.value()); } + if (num_rows.has_value()) { options.set_num_rows(num_rows.value()); } + auto filter_expr = [&]() -> std::unique_ptr { + if (!use_filter) { return nullptr; } + auto stream = ctx->br()->stream_pool().get_stream(); + auto owner = new std::vector; + owner->push_back(std::make_shared>(15, true, stream)); + owner->push_back(std::make_shared( + *std::any_cast>>(owner->at(0)))); + owner->push_back(std::make_shared(0)); + owner->push_back(std::make_shared( + cudf::ast::ast_operator::LESS, + *std::any_cast>(owner->at(2)), + *std::any_cast>(owner->at(1)))); + return std::make_unique( + stream, + *std::any_cast>(owner->back()), + rapidsmpf::OwningWrapper(static_cast(owner), + [](void* p) { delete static_cast*>(p); })); + }(); + auto expected = [&]() { + if (filter_expr != nullptr) { + auto expected_options = options; + expected_options.set_filter(filter_expr->filter); + filter_expr->stream.synchronize(); + auto expected = cudf::io::read_parquet(expected_options).tbl; + filter_expr->stream.synchronize(); + return expected; + } else { + return cudf::io::read_parquet(options).tbl; + } + }(); + auto ch = ctx->create_channel(); + std::vector actors; + + actors.push_back(cudf_streaming::streaming::actor::read_parquet( + ctx, GlobalEnvironment->comm_, ch, 4, options, 3, std::move(filter_expr))); + + std::vector messages; + actors.push_back(rapidsmpf::streaming::actor::pull_from_channel(ctx, ch, messages)); + + if (GlobalEnvironment->comm_->nranks() > 1 && + (skip_rows.value_or(0) > 0 || num_rows.has_value())) { + // We don't yet implement skip_rows/num_rows in multi-rank mode + EXPECT_THROW(rapidsmpf::streaming::run_actor_network(std::move(actors)), std::logic_error); + return; + } + rapidsmpf::streaming::run_actor_network(std::move(actors)); + + rapidsmpf::coll::AllGather allgather(GlobalEnvironment->comm_, + /* op_id = */ 0, + br.get()); + + for (auto& msg : messages) { + auto chunk = msg.release(); + auto seq = msg.sequence_number(); + auto [reservation, _] = br->reserve( + rapidsmpf::MemoryType::DEVICE, chunk.make_available_cost(), rapidsmpf::AllowOverbooking::YES); + chunk = chunk.make_available(reservation); + auto packed_columns = cudf::pack(chunk.table_view(), chunk.stream(), br->device_mr()); + auto packed_data = + rapidsmpf::PackedData{std::move(packed_columns.metadata), + br->move(std::move(packed_columns.gpu_data), chunk.stream())}; + + allgather.insert(seq, std::move(packed_data)); + } + + allgather.insert_finished(); + + // May as well check on all ranks, so we also mildly exercise the allgather. + auto gathered_packed_data = allgather.wait_and_extract(rapidsmpf::coll::AllGather::Ordered::YES); + auto result = cudf_streaming::integrations::unpack_and_concat( + std::move(gathered_packed_data), rmm::cuda_stream_default, br.get()); + EXPECT_EQ(result->num_rows(), expected->num_rows()); + EXPECT_EQ(result->num_columns(), expected->num_columns()); + EXPECT_EQ(result->num_columns(), 1); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(result->view(), expected->view()); +} diff --git a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp new file mode 100644 index 000000000000..75c681889ef1 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp @@ -0,0 +1,507 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../utils.hpp" +#include "base_streaming_fixture.hpp" + +#include +#include +#include + +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include + +using namespace cudf_streaming::streaming; + +class StreamingTableChunk : public BaseStreamingFixture, + public ::testing::WithParamInterface { + protected: + void SetUp() override + { + rapidsmpf::config::Options options(rapidsmpf::config::get_environment_variables()); + + std::unordered_map memory_limits{}; + auto stream_pool = + std::make_shared(16, rmm::cuda_stream::flags::non_blocking); + stream = cudf::get_default_stream(); + br = rapidsmpf::BufferResource::create( + mr_cuda, // device_mr + rapidsmpf::PinnedMemoryResource::make_if_available(), // pinned_mr + memory_limits, // memory_limits + std::chrono::milliseconds{1}, // periodic_spill_check + stream_pool, // stream_pool + rapidsmpf::Statistics::disabled() // statistics + ); + ctx = std::make_shared( + options, GlobalEnvironment->comm_->logger(), br); + } + + rmm::cuda_stream_view stream; + rmm::mr::cuda_memory_resource mr_cuda; + std::shared_ptr br; + std::shared_ptr ctx; +}; + +TEST_F(StreamingTableChunk, FromTable) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + + TableChunk chunk{std::make_unique(expect), stream}; + EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_TRUE(chunk.is_available()); + EXPECT_TRUE(chunk.is_spillable()); + EXPECT_EQ(chunk.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect); + + auto chunk2 = chunk.make_available( + br->reserve_or_fail(chunk.make_available_cost(), rapidsmpf::MemoryType::DEVICE)); + EXPECT_FALSE(chunk.is_available()); + EXPECT_TRUE(chunk2.is_available()); + EXPECT_TRUE(chunk2.is_spillable()); + EXPECT_EQ(chunk2.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); +} + +TEST_F(StreamingTableChunk, TableChunkOwner) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + constexpr std::uint64_t seq = 42; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + // Static because the deleter function is a void(*)(void*) which precludes the use of + // a lambda with captures. + static std::size_t num_deletions{0}; + auto deleter = [](void* p) { + num_deletions++; + delete static_cast(p); + }; + auto make_chunk = [&](TableChunk::ExclusiveView exclusive_view) { + return TableChunk{expect, stream, rapidsmpf::OwningWrapper(new int, deleter), exclusive_view}; + }; + auto check_chunk = [&](TableChunk const& chunk, bool is_spillable) { + EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_TRUE(chunk.is_available()); + EXPECT_EQ(chunk.is_spillable(), is_spillable); + EXPECT_EQ(chunk.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect); + }; + { + auto chunk = make_chunk(TableChunk::ExclusiveView::NO); + check_chunk(chunk, false); + EXPECT_EQ(num_deletions, 0); + } + EXPECT_EQ(num_deletions, 1); + { + auto msg = + to_message(seq, std::make_unique(make_chunk(TableChunk::ExclusiveView::NO))); + EXPECT_EQ(num_deletions, 1); + } + EXPECT_EQ(num_deletions, 2); + { + auto msg = + to_message(seq, std::make_unique(make_chunk(TableChunk::ExclusiveView::YES))); + auto chunk = msg.release(); + check_chunk(chunk, true); + EXPECT_EQ(num_deletions, 2); + } + EXPECT_EQ(num_deletions, 3); + { + auto chunk = make_chunk(TableChunk::ExclusiveView::YES); + check_chunk(chunk, true); + auto res = br->reserve_or_fail(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE), + rapidsmpf::MemoryType::DEVICE); + // This is like spilling since the original `chunk` is ExclusiveView::YES and + // overwritten. + chunk = chunk.copy(res); + EXPECT_EQ(num_deletions, 4); + } +} + +TEST_F(StreamingTableChunk, FromPackedDataOnDevice) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + auto packed_columns = cudf::pack(expect, stream); + + auto packed_data = std::make_unique( + std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream)); + TableChunk chunk{std::move(packed_data)}; + + EXPECT_EQ(chunk.stream().value(), stream.value()); + // chunk was created from packed data on device, so it is available and make available + // cost is 0. + EXPECT_TRUE(chunk.is_available()); + EXPECT_TRUE(chunk.is_spillable()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, chunk.table_view()); + EXPECT_EQ(chunk.make_available_cost(), 0); + + auto chunk2 = chunk.make_available( + br->reserve_or_fail(chunk.make_available_cost(), rapidsmpf::MemoryType::DEVICE)); + EXPECT_FALSE(chunk.is_available()); + EXPECT_TRUE(chunk2.is_available()); + EXPECT_TRUE(chunk2.is_spillable()); + EXPECT_EQ(chunk2.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); +} + +INSTANTIATE_TEST_SUITE_P(StreamingTableChunkWithSpillTargets, + StreamingTableChunk, + ::testing::ValuesIn(rapidsmpf::SPILL_TARGET_MEMORY_TYPES), + [](testing::TestParamInfo const& info) { + return std::string{rapidsmpf::to_string(info.param)}; + }); + +TEST_P(StreamingTableChunk, FromPackedDataOn) +{ + auto const spill_mem_type = GetParam(); + if (spill_mem_type == rapidsmpf::MemoryType::PINNED_HOST && + !rapidsmpf::is_pinned_memory_resources_supported()) { + GTEST_SKIP() << "MemoryType::PINNED_HOST isn't supported on the system."; + } + + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + auto packed_columns = cudf::pack(expect, stream); + std::size_t const size = packed_columns.gpu_data->size(); + + // Move the gpu_data to a Buffer (still device memory). + auto gpu_data_on_device = br->move(std::move(packed_columns.gpu_data), stream); + + // Copy the GPU data to the current spill target memory type. + auto [res, _] = br->reserve(spill_mem_type, size, rapidsmpf::AllowOverbooking::YES); + auto gpu_data_in_spill_memory = br->move(std::move(gpu_data_on_device), res); + + auto packed_data = std::make_unique(std::move(packed_columns.metadata), + std::move(gpu_data_in_spill_memory)); + TableChunk chunk{std::move(packed_data)}; + + EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_FALSE(chunk.is_available()); + EXPECT_TRUE(chunk.is_spillable()); + EXPECT_THROW(std::ignore = chunk.table_view(), std::invalid_argument); + EXPECT_EQ(chunk.make_available_cost(), size); + + auto chunk2 = chunk.make_available( + br->reserve_or_fail(chunk.make_available_cost(), rapidsmpf::MemoryType::DEVICE)); + EXPECT_FALSE(chunk.is_available()); + EXPECT_TRUE(chunk2.is_available()); + EXPECT_TRUE(chunk2.is_spillable()); + EXPECT_EQ(chunk2.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); +} + +TEST_F(StreamingTableChunk, DeviceToDeviceCopy) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + + auto expect = random_table_with_index(seed, num_rows, 0, 10); + + cudf_streaming::streaming::TableChunk chunk{std::make_unique(expect), stream}; + EXPECT_TRUE(chunk.is_available()); + + auto res = br->reserve_or_fail(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE), + rapidsmpf::MemoryType::DEVICE); + auto chunk2 = chunk.copy(res); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); +} + +TEST_F(StreamingTableChunk, ShapeOnAvailableAndSpilledChunk) +{ + constexpr unsigned int num_rows = 64; + constexpr std::int64_t seed = 2025; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 5); + auto const expected_shape = + std::pair{expect.num_rows(), expect.num_columns()}; + + TableChunk device_chunk{std::make_unique(expect), stream}; + EXPECT_TRUE(device_chunk.is_available()); + EXPECT_EQ(device_chunk.shape(), expected_shape); + + auto [res, _] = br->reserve(rapidsmpf::MemoryType::HOST, + device_chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE), + rapidsmpf::AllowOverbooking::YES); + auto host_chunk = device_chunk.copy(res); + + EXPECT_FALSE(host_chunk.is_available()); + EXPECT_EQ(host_chunk.shape(), expected_shape); + + device_chunk = host_chunk.make_available( + br->reserve_or_fail(host_chunk.make_available_cost(), rapidsmpf::MemoryType::DEVICE)); + EXPECT_TRUE(device_chunk.is_available()); + EXPECT_EQ(device_chunk.shape(), expected_shape); +} + +TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy) +{ + auto const spill_mem_type = GetParam(); + if (spill_mem_type == rapidsmpf::MemoryType::PINNED_HOST && + !rapidsmpf::is_pinned_memory_resources_supported()) { + GTEST_SKIP() << "MemoryType::PINNED_HOST isn't supported on the system."; + } + + constexpr unsigned int num_rows = 64; + constexpr std::int64_t seed = 2025; + + auto expect = random_table_with_index(seed, num_rows, 0, 5); + + TableChunk dev_chunk{std::make_unique(expect), stream}; + EXPECT_TRUE(dev_chunk.is_available()); + EXPECT_TRUE(dev_chunk.is_spillable()); + EXPECT_EQ(dev_chunk.stream().value(), stream.value()); + EXPECT_EQ(dev_chunk.make_available_cost(), 0); + { + auto cd = get_content_description(dev_chunk); + EXPECT_EQ(cd.spillable(), dev_chunk.is_spillable()); + for (auto mem_type : rapidsmpf::MEMORY_TYPES) { + EXPECT_EQ(cd.content_size(mem_type), dev_chunk.data_alloc_size(mem_type)); + } + } + + // Copy to host memory -> new chunk should be unavailable. + auto host_res = + br->reserve_or_fail(dev_chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE), spill_mem_type); + auto host_copy = dev_chunk.copy(host_res); + EXPECT_FALSE(host_copy.is_available()); + EXPECT_TRUE(host_copy.is_spillable()); + EXPECT_EQ(host_copy.stream().value(), stream.value()); + EXPECT_GT(host_copy.make_available_cost(), 0); + { + auto cd = get_content_description(host_copy); + EXPECT_EQ(cd.spillable(), host_copy.is_spillable()); + for (auto mem_type : rapidsmpf::MEMORY_TYPES) { + EXPECT_EQ(cd.content_size(mem_type), host_copy.data_alloc_size(mem_type)); + } + } + + // Host to host copy. + auto host_res2 = br->reserve_or_fail(host_copy.data_alloc_size(spill_mem_type), spill_mem_type); + auto host_copy2 = host_copy.copy(host_res2); + EXPECT_FALSE(host_copy2.is_available()); + EXPECT_TRUE(host_copy2.is_spillable()); + EXPECT_EQ(host_copy2.stream().value(), stream.value()); + EXPECT_EQ(host_copy2.make_available_cost(), host_copy.make_available_cost()); + { + auto cd = get_content_description(host_copy2); + EXPECT_EQ(cd.spillable(), host_copy2.is_spillable()); + for (auto mem_type : rapidsmpf::MEMORY_TYPES) { + EXPECT_EQ(cd.content_size(mem_type), host_copy2.data_alloc_size(mem_type)); + } + } + + // Bring the new host copy back to device and verify equality. + auto dev_res = + br->reserve_or_fail(host_copy2.data_alloc_size(spill_mem_type), rapidsmpf::MemoryType::DEVICE); + auto dev_back = host_copy2.make_available(dev_res); + EXPECT_TRUE(dev_back.is_available()); + EXPECT_TRUE(dev_back.is_spillable()); + EXPECT_EQ(dev_back.stream().value(), stream.value()); + EXPECT_EQ(dev_back.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(dev_back.table_view(), expect); + { + auto cd = get_content_description(dev_back); + EXPECT_EQ(cd.spillable(), dev_back.is_spillable()); + for (auto mem_type : rapidsmpf::MEMORY_TYPES) { + EXPECT_EQ(cd.content_size(mem_type), dev_back.data_alloc_size(mem_type)); + } + } + + // Sanity check: a second device copy should also remain equivalent. + auto dev_res2 = br->reserve_or_fail(dev_back.data_alloc_size(rapidsmpf::MemoryType::DEVICE), + rapidsmpf::MemoryType::DEVICE); + auto dev_copy2 = dev_back.copy(dev_res2); + EXPECT_TRUE(dev_copy2.is_available()); + EXPECT_EQ(dev_copy2.make_available_cost(), 0); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(dev_copy2.table_view(), expect); + { + auto cd = get_content_description(dev_copy2); + EXPECT_EQ(cd.spillable(), dev_copy2.is_spillable()); + for (auto mem_type : rapidsmpf::MEMORY_TYPES) { + EXPECT_EQ(cd.content_size(mem_type), dev_copy2.data_alloc_size(mem_type)); + } + } +} + +TEST_F(StreamingTableChunk, ToMessageRoundTrip) +{ + constexpr unsigned int num_rows = 64; + constexpr std::int64_t seed = 2025; + constexpr std::uint64_t seq = 7; + + auto expect = random_table_with_index(seed, num_rows, 0, 5); + auto chunk = std::make_unique(std::make_unique(expect), stream); + + rapidsmpf::streaming::Message m = to_message(seq, std::move(chunk)); + EXPECT_FALSE(m.empty()); + EXPECT_TRUE(m.holds()); + EXPECT_TRUE(m.content_description().spillable()); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 1024); + EXPECT_EQ(m.sequence_number(), seq); + + // Deep-copy: device to host. + auto reservation = br->reserve_or_fail(m.copy_cost(), rapidsmpf::MemoryType::HOST); + rapidsmpf::streaming::Message m2 = m.copy(reservation); + EXPECT_EQ(reservation.size(), 0); + EXPECT_FALSE(m2.empty()); + EXPECT_TRUE(m2.holds()); + EXPECT_TRUE(m2.content_description().spillable()); + EXPECT_EQ(m2.content_description().content_size(rapidsmpf::MemoryType::HOST), 1024); + EXPECT_EQ(m2.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 0); + EXPECT_EQ(m2.sequence_number(), seq); + + // Deep-copy: host to host. + reservation = br->reserve_or_fail(m2.copy_cost(), rapidsmpf::MemoryType::HOST); + rapidsmpf::streaming::Message m3 = m.copy(reservation); + EXPECT_EQ(reservation.size(), 0); + EXPECT_FALSE(m3.empty()); + EXPECT_TRUE(m3.holds()); + EXPECT_TRUE(m3.content_description().spillable()); + EXPECT_EQ(m3.content_description().content_size(rapidsmpf::MemoryType::HOST), 1024); + EXPECT_EQ(m3.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 0); + EXPECT_EQ(m3.sequence_number(), seq); + + // Copy the chunk back to device and verify. + { + auto chunk = m3.release(); + auto res = br->reserve_or_fail(chunk.make_available_cost(), rapidsmpf::MemoryType::DEVICE); + chunk = chunk.make_available(res); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect); + } + + // Deep-copy: host to device. + reservation = br->reserve_or_fail(m2.copy_cost(), rapidsmpf::MemoryType::DEVICE); + rapidsmpf::streaming::Message m4 = m.copy(reservation); + EXPECT_EQ(reservation.size(), 0); + EXPECT_FALSE(m4.empty()); + EXPECT_TRUE(m4.holds()); + EXPECT_TRUE(m4.content_description().spillable()); + EXPECT_EQ(m4.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m4.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 1024); + EXPECT_EQ(m4.sequence_number(), seq); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(m4.get().table_view(), expect); + + // Deep-copy: device to device. + reservation = br->reserve_or_fail(m4.copy_cost(), rapidsmpf::MemoryType::DEVICE); + rapidsmpf::streaming::Message m5 = m.copy(reservation); + EXPECT_EQ(reservation.size(), 0); + EXPECT_FALSE(m5.empty()); + EXPECT_TRUE(m5.holds()); + EXPECT_TRUE(m5.content_description().spillable()); + EXPECT_EQ(m5.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m5.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 1024); + EXPECT_EQ(m5.sequence_number(), seq); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(m5.get().table_view(), expect); +} + +TEST_F(StreamingTableChunk, ToMessageNotSpillable) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + constexpr std::uint64_t seq = 42; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + + auto deleter = [](void* p) { delete static_cast(p); }; + auto chunk = std::make_unique( + expect, stream, rapidsmpf::OwningWrapper(new int, deleter), TableChunk::ExclusiveView::NO); + + rapidsmpf::streaming::Message m = to_message(seq, std::move(chunk)); + EXPECT_FALSE(m.empty()); + EXPECT_TRUE(m.holds()); + EXPECT_FALSE(m.content_description().spillable()); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::DEVICE), + cudf::packed_size(expect.view(), stream, rmm::mr::get_current_device_resource_ref())); + // packed size is greater than or equal to the alloc size due to buffer alignments. + EXPECT_GE(m.content_description().content_size(rapidsmpf::MemoryType::DEVICE), + expect.alloc_size()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(m.get().table_view(), expect); +} + +TEST_F(StreamingTableChunk, ToPackedDataFromPackedChunk) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + auto packed_columns = cudf::pack(expect, stream); + TableChunk chunk{std::make_unique( + std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream))}; + EXPECT_TRUE(chunk.is_available()); + + auto packed = std::move(chunk).into_packed_data(br.get()); + EXPECT_FALSE(chunk.is_available()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, TableChunk{std::move(packed)}.table_view()); +} + +TEST_F(StreamingTableChunk, ToPackedDataFromTable) +{ + constexpr unsigned int num_rows = 100; + constexpr std::int64_t seed = 1337; + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + TableChunk chunk{std::make_unique(expect), stream}; + EXPECT_TRUE(chunk.is_available()); + + auto packed = std::move(chunk).into_packed_data(br.get()); + EXPECT_FALSE(chunk.is_available()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, TableChunk{std::move(packed)}.table_view()); +} + +TEST_F(StreamingTableChunk, ToMessageUnalignedSize) +{ + constexpr unsigned int num_rows = 5; + constexpr std::int64_t seed = 2025; + constexpr std::uint64_t seq = 7; + + auto expect = random_table_with_index(seed, num_rows, 0, 5); + auto chunk = std::make_unique(std::make_unique(expect), stream); + + rapidsmpf::streaming::Message m = to_message(seq, std::move(chunk)); + EXPECT_EQ(m.sequence_number(), seq); + EXPECT_FALSE(m.empty()); + EXPECT_TRUE(m.holds()); + EXPECT_TRUE(m.content_description().spillable()); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::HOST), 0); + EXPECT_EQ(m.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 80); + EXPECT_EQ(m.copy_cost(), 80); + + // Deep copy: device → host. + // Note: `m.copy_cost() == 80`, but cudf performs 128-byte aligned allocations. + // This means `m.copy_cost()` is not always sufficient; however, TableChunk.copy() + // accounts for this alignment internally. + auto reservation = br->reserve_or_fail(m.copy_cost(), rapidsmpf::MemoryType::HOST); + rapidsmpf::streaming::Message m2 = m.copy(reservation); + EXPECT_EQ(reservation.size(), 0); + EXPECT_FALSE(m2.empty()); + EXPECT_TRUE(m2.holds()); + EXPECT_TRUE(m2.content_description().spillable()); + EXPECT_EQ(m2.copy_cost(), 128); + EXPECT_EQ(m2.content_description().content_size(rapidsmpf::MemoryType::HOST), 128); + EXPECT_EQ(m2.content_description().content_size(rapidsmpf::MemoryType::DEVICE), 0); + EXPECT_EQ(m2.sequence_number(), seq); +} diff --git a/cpp/libcudf_streaming/tests/test_partition.cpp b/cpp/libcudf_streaming/tests/test_partition.cpp new file mode 100644 index 000000000000..ede097a15c4f --- /dev/null +++ b/cpp/libcudf_streaming/tests/test_partition.cpp @@ -0,0 +1,132 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "utils.hpp" + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace cudf_streaming::integrations; + +class NumOfPartitions : public cudf::test::BaseFixtureWithParam> {}; + +// test different `num_partitions` and `num_rows`. +INSTANTIATE_TEST_SUITE_P(Partitions, + NumOfPartitions, + testing::Combine(testing::Range(1, 10), // num_partitions + testing::Range(1, 100, 9) // num_rows + )); + +TEST_P(NumOfPartitions, partition_and_pack) +{ + int const num_partitions = std::get<0>(GetParam()); + int const num_rows = std::get<1>(GetParam()); + std::int64_t const seed = 42; + cudf::hash_id const hash_fn = cudf::hash_id::HASH_MURMUR3; + auto stream = cudf::get_default_stream(); + auto br = rapidsmpf::BufferResource::create(mr()); + + cudf::table expect = random_table_with_index(seed, static_cast(num_rows), 0, 10); + + auto chunks = partition_and_pack(expect, {1}, num_partitions, hash_fn, seed, stream, br.get()); + + // Convert to a vector + std::vector chunks_vector; + for (auto& [_, chunk] : chunks) { + chunks_vector.push_back(std::move(chunk)); + } + EXPECT_EQ(chunks_vector.size(), num_partitions); + + auto result = unpack_and_concat(std::move(chunks_vector), stream, br.get()); + + // Compare the input table with the result. We ignore the row order by + // sorting by their index (first column). + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(sort_table(expect), sort_table(result)); +} + +TEST_P(NumOfPartitions, split_and_pack) +{ + int const num_partitions = std::get<0>(GetParam()); + int const num_rows = std::get<1>(GetParam()); + std::int64_t const seed = 42; + auto stream = cudf::get_default_stream(); + auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + + cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); + + std::vector splits; + for (int i = 1; i < num_partitions; ++i) { + splits.emplace_back(i * num_rows / num_partitions); + } + + auto chunks = split_and_pack(expect, splits, stream, br.get()); + + // Convert to a vector (restoring the original order). + std::vector chunks_vector; + for (int i = 0; i < num_partitions; ++i) { + chunks_vector.emplace_back(std::move(chunks.at(i))); + } + EXPECT_EQ(chunks_vector.size(), num_partitions); + + auto result = unpack_and_concat(std::move(chunks_vector), stream, br.get()); + + // Compare the input table with the result. + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result); +} + +class SpillingTest : public ::testing::Test { + protected: + void SetUp() override + { + br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + stream = cudf::get_default_stream(); + } + + std::shared_ptr br; + rmm::cuda_stream_view stream; +}; + +TEST_F(SpillingTest, SpillUnspillRoundtripPreservesDataAndMetadata) +{ + std::vector metadata{42, 99}; + std::vector payload{10, 20, 30}; + + // Create device input. + std::vector input; + input.push_back(create_packed_data(metadata, payload, stream, br.get())); + + // Device -> Device (moves data) + auto on_gpu = unspill_partitions(std::move(input), br.get(), rapidsmpf::AllowOverbooking::YES); + ASSERT_EQ(on_gpu.size(), 1); + EXPECT_EQ(on_gpu[0].data->mem_type(), rapidsmpf::MemoryType::DEVICE); + EXPECT_EQ(*on_gpu[0].metadata, metadata); + + // Device -> Host + auto back_on_host = spill_partitions(std::move(on_gpu), br.get()); + ASSERT_EQ(back_on_host.size(), 1); + EXPECT_EQ(back_on_host[0].data->mem_type(), rapidsmpf::MemoryType::HOST); + EXPECT_EQ(*back_on_host[0].metadata, metadata); + + // Check that contents match original + auto res = br->reserve_or_fail(back_on_host[0].data->size, rapidsmpf::MemoryType::HOST); + auto actual = br->move_to_host_buffer(std::move(back_on_host[0].data), res); + EXPECT_EQ(actual->copy_to_uint8_vector(), payload); +} diff --git a/cpp/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp new file mode 100644 index 000000000000..8c942db2f02e --- /dev/null +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -0,0 +1,294 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include +#include +#include + +#include +#include + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * @brief RAII temporary directory created under GTest's temp directory. + * + * The directory is created on construction and recursively removed on + * destruction. Removal errors are ignored. + */ +class TempDir { + public: + TempDir() : path_(unique_path()) + { + std::error_code ec; + if (!std::filesystem::create_directories(path_, ec) || ec) { + throw std::runtime_error("Failed to create temp directory: " + path_.string()); + } + } + + ~TempDir() noexcept + { + std::error_code ec; + std::filesystem::remove_all(path_, ec); + // Intentionally ignore errors in destructor. + } + + TempDir(TempDir const&) = delete; + TempDir& operator=(TempDir const&) = delete; + TempDir(TempDir&&) = delete; + TempDir& operator=(TempDir&&) = delete; + + /// @brief Returns the path to the temporary directory. + [[nodiscard]] std::filesystem::path const& path() const noexcept { return path_; } + + private: + static std::filesystem::path unique_path() + { + static std::atomic counter{0}; + return std::filesystem::path(testing::TempDir()) / + ("tmp_" + std::to_string(::getpid()) + "_" + + std::to_string(counter.fetch_add(1, std::memory_order_relaxed))); + } + + std::filesystem::path path_; +}; + +/// @brief User-defined literal for specifying memory sizes in KiB. +constexpr std::size_t operator"" _KiB(unsigned long long val) { return val * (1 << 10); } + +/// @brief User-defined literal for specifying memory sizes in MiB. +constexpr std::size_t operator"" _MiB(unsigned long long val) { return val * (1ull << 20); } + +/// @brief User-defined literal for specifying memory sizes in GiB. +constexpr std::size_t operator"" _GiB(unsigned long long val) { return val * (1 << 30); } + +template +[[nodiscard]] std::vector iota_vector(std::size_t nelem, T start = 0) +{ + std::vector ret(nelem); + std::iota(ret.begin(), ret.end(), start); + return ret; +} + +template +[[nodiscard]] inline std::unique_ptr iota_column(std::size_t nrows, T start = 0) +{ + std::vector vec = iota_vector(nrows, start); + cudf::test::fixed_width_column_wrapper ret(vec.begin(), vec.end()); + return ret.release(); +} + +template +[[nodiscard]] inline std::vector random_vector(std::int64_t seed, + std::size_t nelem, + T min = std::numeric_limits::min(), + T max = std::numeric_limits::max()) +{ + std::mt19937 rng(static_cast(seed)); + std::uniform_int_distribution dist(min, max); + std::vector ret(nelem); + std::generate(ret.begin(), ret.end(), [&]() { return dist(rng); }); + return ret; +} + +[[nodiscard]] inline std::unique_ptr random_column( + std::int64_t seed, + std::size_t nrows, + std::int64_t min = std::numeric_limits::min(), + std::int64_t max = std::numeric_limits::max()) +{ + std::vector vec = random_vector(seed, nrows, min, max); + cudf::test::fixed_width_column_wrapper ret(vec.begin(), vec.end()); + return ret.release(); +} + +[[nodiscard]] inline cudf::table random_table_with_index( + std::int64_t seed, + std::size_t nrows, + std::int64_t min = std::numeric_limits::min(), + std::int64_t max = std::numeric_limits::max()) +{ + std::vector> cols; + cols.push_back(iota_column(nrows)); + cols.push_back(random_column(seed, nrows, min, max)); + return cudf::table(std::move(cols)); +} + +[[nodiscard]] inline cudf::table sort_table( + cudf::table_view const& table, std::vector const& /* column_indices */ = {0}) +{ + if (table.num_columns() == 0) { return cudf::table(table); } + return cudf::gather(table, cudf::sorted_order(table.select({0}))->view())->release(); +} + +[[nodiscard]] inline cudf::table sort_table(std::unique_ptr const& table, + std::vector const& column_indices = { + 0}) +{ + return sort_table(table->view(), column_indices); +} + +/// @brief Create a PackedData object from a host buffer +[[nodiscard]] inline rapidsmpf::PackedData create_packed_data( + std::span metadata, + std::span data, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br) +{ + auto metadata_ptr = std::make_unique>(metadata.begin(), metadata.end()); + + auto reservation = + br->reserve(rapidsmpf::MemoryType::DEVICE, data.size(), rapidsmpf::AllowOverbooking::YES); + auto data_ptr = std::make_unique(data.data(), data.size(), stream); + return rapidsmpf::PackedData{std::move(metadata_ptr), br->move(std::move(data_ptr), stream)}; +} + +/** + * @brief Generate a packed data object with the given number of elements and offset. + * + * Both metadata and GPU data contain the same integer sequence. + * + * @param n_elements Number of elements in the sequence. + * @param offset Starting value of the sequence. + * @param stream CUDA stream for device allocation. + * @param br Buffer resource used for allocations. + * @return A packed data object containing metadata and GPU data. + */ +[[nodiscard]] inline rapidsmpf::PackedData generate_packed_data(int n_elements, + int offset, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource& br) +{ + auto values = iota_vector(n_elements, offset); + + auto metadata = std::make_unique>(n_elements * sizeof(int)); + std::memcpy(metadata->data(), values.data(), n_elements * sizeof(int)); + + auto data = std::make_unique( + values.data(), n_elements * sizeof(int), stream, br.device_mr()); + + return {std::move(metadata), br.move(std::move(data), stream)}; +} + +/** + * @brief Validate a packed data object by checking metadata and GPU data contents. + * + * @param packed_data Packed data object to validate. + * @param n_elements Expected number of elements. + * @param offset Expected starting value of the sequence. + * @param stream CUDA stream used for device-host transfers. + * @param br Buffer resource used for host allocation. + */ +inline void validate_packed_data(rapidsmpf::PackedData&& packed_data, + int n_elements, + int offset, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource& br) +{ + auto const& metadata = *packed_data.metadata; + EXPECT_EQ(n_elements * sizeof(int), metadata.size()); + + for (int i = 0; i < n_elements; i++) { + int val; + std::memcpy(&val, metadata.data() + i * sizeof(int), sizeof(int)); + EXPECT_EQ(offset + i, val); + } + + EXPECT_EQ(n_elements * sizeof(int), packed_data.data->size); + + auto res = br.reserve_or_fail(packed_data.data->size, rapidsmpf::MemoryType::HOST); + auto data_on_host = br.move_to_host_buffer(std::move(packed_data.data), res); + RAPIDSMPF_CUDA_TRY(cudaStreamSynchronize(stream)); + EXPECT_EQ(metadata, data_on_host->copy_to_uint8_vector()); +} + +/** + * @brief Device memory resource that can inject stream-ordered delays. + * + * When enabled, each allocation enqueues a host callback on the allocation + * stream that sleeps for a configurable duration. This blocks the CUDA stream + * (making `cudaEventQuery` return not-ready) without blocking the host thread, + * so the progress thread's event loop continues to run while data buffers + * appear unready. + */ +class DelayedMemoryResource { + public: + DelayedMemoryResource(rmm::device_async_resource_ref upstream, std::chrono::milliseconds delay) + : upstream_{upstream}, delay_{delay} + { + } + + void* allocate_sync(std::size_t, std::size_t) + { + RAPIDSMPF_FAIL("synchronous allocation not supported", std::invalid_argument); + } + + void deallocate_sync(void*, std::size_t, std::size_t) noexcept + { + RAPIDSMPF_FATAL("synchronous deallocation not supported"); + } + + void* allocate(rmm::cuda_stream_view stream, + std::size_t size, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + void* ptr = upstream_.allocate(stream, size, alignment); + if (size > 0) { + RAPIDSMPF_CUDA_TRY( + cudaLaunchHostFunc(stream.value(), sleep_on_stream, new std::chrono::milliseconds(delay_))); + } + return ptr; + } + + void deallocate(rmm::cuda_stream_view stream, + void* ptr, + std::size_t size, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + upstream_.deallocate(stream, ptr, size, alignment); + } + + bool operator==(DelayedMemoryResource const& other) const noexcept { return this == &other; } + + bool operator!=(DelayedMemoryResource const& other) const noexcept { return !(this == &other); } + + friend void get_property(DelayedMemoryResource const&, cuda::mr::device_accessible) noexcept {} + + private: + static void CUDART_CB sleep_on_stream(void* user_data) + { + auto* delay = static_cast(user_data); + std::this_thread::sleep_for(*delay); + delete delay; + } + + cuda::mr::any_resource upstream_; + std::chrono::milliseconds delay_; +}; + +static_assert(cuda::mr::resource); +static_assert(cuda::mr::resource_with); diff --git a/dependencies.yaml b/dependencies.yaml index a529b5d8ef7f..9ca0720be326 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -468,6 +468,34 @@ files: includes: - test_python_common - test_python_cudf_common + py_build_libcudf_streaming: + output: pyproject + pyproject_dir: python/libcudf_streaming + extras: + table: build-system + includes: + - rapids_build_skbuild + py_rapids_build_libcudf_streaming: + output: pyproject + pyproject_dir: python/libcudf_streaming + extras: + table: tool.rapids-build-backend + key: requires + includes: + - build_base + - build_cpp + - depends_on_libcudf + - depends_on_librmm + - depends_on_librapidsmpf + py_run_libcudf_streaming: + output: pyproject + pyproject_dir: python/libcudf_streaming + extras: + table: project + includes: + - depends_on_libcudf + - depends_on_librmm + - depends_on_librapidsmpf test_python_narwhals: output: none includes: @@ -1287,6 +1315,31 @@ dependencies: - matrix: packages: - *rapidsmpf_unsuffixed + depends_on_librapidsmpf: + common: + - output_types: conda + packages: + - &librapidsmpf_unsuffixed librapidsmpf==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: + - librapidsmpf-cu12==26.8.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - librapidsmpf-cu13==26.8.*,>=0.0.0a0 + - matrix: + packages: + - *librapidsmpf_unsuffixed depends_on_ray: common: - output_types: conda diff --git a/python/libcudf_streaming/CMakeLists.txt b/python/libcudf_streaming/CMakeLists.txt new file mode 100644 index 000000000000..c23fe0cecd92 --- /dev/null +++ b/python/libcudf_streaming/CMakeLists.txt @@ -0,0 +1,39 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# 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) +include(rapids-cuda) +rapids_cuda_init_architectures(libcudf-streaming-python) +rapids_cpm_init() + +project( + libcudf-streaming-python + VERSION "${RAPIDS_VERSION}" + LANGUAGES CXX CUDA +) + +# Check if cudf_streaming is already available. If so, it is the user's responsibility to ensure +# that the CMake package is also available at build time of the Python cudf_streaming package. +find_package(cudf_streaming "${RAPIDS_VERSION}") + +if(cudf_streaming_FOUND) + return() +endif() + +unset(cudf_streaming_FOUND) + +set(BUILD_TESTS OFF) +set(BUILD_BENCHMARKS OFF) +set(CUDF_BUILD_TESTUTIL OFF) +set(CUDF_BUILD_STREAMS_TEST_UTIL OFF) + +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib) + +add_subdirectory(../../cpp/libcudf_streaming cudf-streaming-cpp) diff --git a/python/libcudf_streaming/LICENSE b/python/libcudf_streaming/LICENSE new file mode 100644 index 000000000000..18bcb4316e6d --- /dev/null +++ b/python/libcudf_streaming/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2018 NVIDIA Corporation + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/python/libcudf_streaming/README.md b/python/libcudf_streaming/README.md new file mode 100644 index 000000000000..96815c10e73a --- /dev/null +++ b/python/libcudf_streaming/README.md @@ -0,0 +1,23 @@ +# libcudf_streaming + +libcudf_streaming is the C++ shared library wheel for cuDF Streaming, providing GPU-accelerated streaming data processing capabilities built on top of libcudf and librapidsmpf. + +Most users should not need to install this package directly; it is automatically pulled in as a dependency of `cudf-streaming`. + +This package provides `libcudf_streaming.so` and a Python `load_library()` helper to ensure the shared library and its dependencies are properly loaded at runtime. + +## Installation + +```bash +pip install libcudf-streaming-cu12 # For CUDA 12 +pip install libcudf-streaming-cu13 # For CUDA 13 +``` + +## Usage + +```python +import libcudf_streaming + +# Load the shared library and all dependencies +libcudf_streaming.load_library() +``` diff --git a/python/libcudf_streaming/libcudf_streaming/VERSION b/python/libcudf_streaming/libcudf_streaming/VERSION new file mode 120000 index 000000000000..d62dc733efdf --- /dev/null +++ b/python/libcudf_streaming/libcudf_streaming/VERSION @@ -0,0 +1 @@ +../../../VERSION \ No newline at end of file diff --git a/python/libcudf_streaming/libcudf_streaming/__init__.py b/python/libcudf_streaming/libcudf_streaming/__init__.py new file mode 100644 index 000000000000..6792d3404e98 --- /dev/null +++ b/python/libcudf_streaming/libcudf_streaming/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from libcudf_streaming._version import __git_commit__, __version__ +from libcudf_streaming.load import load_library + +__all__ = ["__git_commit__", "__version__", "load_library"] diff --git a/python/libcudf_streaming/libcudf_streaming/_version.py b/python/libcudf_streaming/libcudf_streaming/_version.py new file mode 100644 index 000000000000..f3ac607a490e --- /dev/null +++ b/python/libcudf_streaming/libcudf_streaming/_version.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +import importlib.resources + +__version__ = ( + importlib.resources.files(__package__) + .joinpath("VERSION") + .read_text() + .strip() +) +try: + __git_commit__ = ( + importlib.resources.files(__package__) + .joinpath("GIT_COMMIT") + .read_text() + .strip() + ) +except FileNotFoundError: + __git_commit__ = "" + +__all__ = ["__git_commit__", "__version__"] diff --git a/python/libcudf_streaming/libcudf_streaming/load.py b/python/libcudf_streaming/libcudf_streaming/load.py new file mode 100644 index 000000000000..82cebae62eac --- /dev/null +++ b/python/libcudf_streaming/libcudf_streaming/load.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +import ctypes +import os + +# Loading with RTLD_LOCAL adds the library itself to the loader's +# loaded library cache without loading any symbols into the global +# namespace. This allows libraries that express a dependency on +# this library to be loaded later and successfully satisfy this dependency +# without polluting the global symbol table with symbols from +# libcudf_streaming that could conflict with symbols from other DSOs. +PREFERRED_LOAD_FLAG = ctypes.RTLD_LOCAL + + +def _load_system_installation(soname: str): + """Try to dlopen() the library indicated by ``soname`` + Raises ``OSError`` if library cannot be loaded. + """ + return ctypes.CDLL(soname, PREFERRED_LOAD_FLAG) + + +def _load_wheel_installation(soname: str): + """Try to dlopen() the library indicated by ``soname`` + + Returns ``None`` if the library cannot be loaded. + """ + if os.path.isfile( + lib := os.path.join(os.path.dirname(__file__), "lib64", soname) + ): + return ctypes.CDLL(lib, PREFERRED_LOAD_FLAG) + return None + + +def load_library(): + """Dynamically load libcudf_streaming.so and its dependencies""" + try: + # These libraries must be loaded before libcudf_streaming because + # libcudf_streaming references their symbols. + import libcudf + import librapidsmpf + import librmm + + librmm.load_library() + libcudf.load_library() + librapidsmpf.load_library() + except ModuleNotFoundError: + # libcudf_streaming's runtime dependencies may be satisfied by + # natively installed libraries or conda packages, in which case + # the imports will fail and we assume the libraries are + # discoverable on system paths. + pass + + return _load_library("libcudf_streaming.so") + + +def _load_library(soname): + prefer_system_installation = ( + os.getenv( + "RAPIDS_LIBCUDF_STREAMING_PREFER_SYSTEM_LIBRARY", "false" + ).lower() + != "false" + ) + + found_lib = None + if prefer_system_installation: + # Prefer a system library if one is present to + # avoid clobbering symbols that other packages might expect, but if no + # other library is present use the one in the wheel. + try: + found_lib = _load_system_installation(soname) + except OSError: + found_lib = _load_wheel_installation(soname) + else: + # Prefer the libraries bundled in this package. If they aren't found + # (which might be the case in builds where the library was prebuilt + # before packaging the wheel), look for a system installation. + try: + found_lib = _load_wheel_installation(soname) + if found_lib is None: + found_lib = _load_system_installation(soname) + except OSError: + # If none of the searches above succeed, just silently return None + # and rely on other mechanisms (like RPATHs on other DSOs) to + # help the loader find the library. + pass + + # The caller almost never needs to do anything with this library, but no + # harm in offering the option since this object at least provides a handle + # to inspect where the library was loaded from. + return found_lib diff --git a/python/libcudf_streaming/pyproject.toml b/python/libcudf_streaming/pyproject.toml new file mode 100644 index 000000000000..21580a0de4c4 --- /dev/null +++ b/python/libcudf_streaming/pyproject.toml @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# 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 = "libcudf-streaming" +dynamic = ["version"] +description = "cuDF Streaming - GPU-accelerated streaming data processing (C++)" +readme = { file = "README.md", content-type = "text/markdown" } +authors = [ + { name = "NVIDIA Corporation" }, +] +license = "Apache-2.0" +requires-python = ">=3.11" +classifiers = [ + "Intended Audience :: Developers", + "Topic :: Database", + "Topic :: Scientific/Engineering", + "Programming Language :: C++", + "Environment :: GPU :: NVIDIA CUDA", +] +dependencies = [ + "libcudf==26.8.*,>=0.0.0a0", + "librapidsmpf==26.8.*,>=0.0.0a0", + "librmm==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.urls] +Homepage = "https://github.com/rapidsai/cudf" + +[project.entry-points."cmake.prefix"] +libcudf_streaming = "libcudf_streaming" + +[tool.pydistcheck] +select = [ + # note: max_allowed_size_compressed is set as a command-line arg + "distro-too-large-compressed", +] + +# Default precision is 2. Bump this a bit for better visibility. +output_file_size_precision = 4 + +[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.reproducible = true +wheel.packages = ["libcudf_streaming"] +wheel.install-dir = "libcudf_streaming" +wheel.py-api = "py3" + +[tool.scikit-build.metadata.version] +provider = "scikit_build_core.metadata.regex" +input = "libcudf_streaming/VERSION" +regex = "(?P.*)" + +[tool.rapids-build-backend] +build-backend = "scikit_build_core.build" +dependencies-file = "../../dependencies.yaml" +matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true" +requires = [ + "cmake>=4.0", + "libcudf==26.8.*,>=0.0.0a0", + "librapidsmpf==26.8.*,>=0.0.0a0", + "librmm==26.8.*,>=0.0.0a0", + "ninja", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.