refactor(cudf_polars): use cudf_streaming for cudf-specific streaming imports

.devcontainer/README.md

@@ -20,7 +20,6 @@ This container is a turnkey development environment for building and testing the
 By default, the following directories are bind-mounted into the devcontainer:
 
 * `${repo}:/home/coder/cudf`
-* `${repo}/../rapidsmpf:/home/coder/rapidsmpf`
 * `${repo}/../.aws:/home/coder/.aws`
 * `${repo}/../.local:/home/coder/.local`
 * `${repo}/../.cache:/home/coder/.cache`
@@ -29,13 +28,6 @@ By default, the following directories are bind-mounted into the devcontainer:
 
 This ensures caches, configurations, dependencies, and your commits are persisted on the host across container runs.
 
-The [rapidsmpf](https://github.com/rapidsai/rapidsmpf) repository is a required dependency of `cudf_polars` (that also requires `libcudf`) and must be cloned as a sibling directory to the cudf repo before launching the devcontainer:
-
-```
-# from the parent directory of your cudf clone
-git clone https://github.com/rapidsai/rapidsmpf.git
-```
-
 ## Launch a Dev Container
 
 To launch a devcontainer from VSCode, open the cuDF repo and select the "Reopen in Container" button in the bottom right:<br/><img src="https://user-images.githubusercontent.com/178183/221771999-97ab29d5-e718-4e5f-b32f-2cdd51bba25c.png"/>

.devcontainer/cuda12.9-conda/devcontainer.json

@@ -52,7 +52,7 @@
   "initializeCommand": [
     "/bin/bash",
     "-c",
-    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda12.9-envs} ${localWorkspaceFolder}/../rapidsmpf"
+    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda12.9-envs}"
   ],
   "postAttachCommand": [
     "/bin/bash",
@@ -66,8 +66,7 @@
     "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.conda/pkgs,target=/home/coder/.conda/pkgs,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda12.9-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent"
+    "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda12.9-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent"
   ],
   "customizations": {
     "vscode": {

.devcontainer/cuda12.9-pip/devcontainer.json

@@ -30,7 +30,7 @@
   "initializeCommand": [
     "/bin/bash",
     "-c",
-    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs} ${localWorkspaceFolder}/../rapidsmpf"
+    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs}"
   ],
   "postAttachCommand": [
     "/bin/bash",
@@ -43,8 +43,7 @@
     "source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent"
+    "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda12.9-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent"
   ],
   "customizations": {
     "vscode": {

.devcontainer/cuda13.2-conda/devcontainer.json

@@ -52,7 +52,7 @@
   "initializeCommand": [
     "/bin/bash",
     "-c",
-    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda13.2-envs} ${localWorkspaceFolder}/../rapidsmpf"
+    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda13.2-envs}"
   ],
   "postAttachCommand": [
     "/bin/bash",
@@ -66,8 +66,7 @@
     "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.conda/pkgs,target=/home/coder/.conda/pkgs,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda13.2-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent"
+    "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda13.2-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent"
   ],
   "customizations": {
     "vscode": {

.devcontainer/cuda13.2-pip/devcontainer.json

@@ -30,7 +30,7 @@
   "initializeCommand": [
     "/bin/bash",
     "-c",
-    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda13.2-venvs} ${localWorkspaceFolder}/../rapidsmpf"
+    "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda13.2-venvs}"
   ],
   "postAttachCommand": [
     "/bin/bash",
@@ -43,8 +43,7 @@
     "source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent",
     "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda13.2-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent",
-    "source=${localWorkspaceFolder}/../rapidsmpf,target=/home/coder/rapidsmpf,type=bind,consistency=consistent"
+    "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda13.2-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent"
   ],
   "customizations": {
     "vscode": {

.github/workflows/pr.yaml

@@ -646,7 +646,7 @@ jobs:
       package-type: python
       pure-wheel: true
   wheel-tests-cudf-polars:
-    needs: [wheel-build-cudf-polars, changed-files]
+    needs: [wheel-build-cudf-polars, wheel-build-cudf-streaming, changed-files]
     permissions:
       actions: read
       contents: read
@@ -663,7 +663,7 @@ jobs:
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
       script: "ci/test_wheel_cudf_polars.sh"
   cudf-polars-polars-tests:
-    needs: [wheel-build-cudf-polars, changed-files]
+    needs: [wheel-build-cudf-polars, wheel-build-cudf-streaming, changed-files]
     permissions:
       actions: read
       contents: read
@@ -732,14 +732,9 @@ jobs:
         SCCACHE_DIST_MAX_RETRIES=inf
         SCCACHE_SERVER_LOG=sccache=debug
         SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=false
-      # Call clone-rapidsmpf before creating the Python env to exclude cudf_polars' dependency on rapidsmpf
-      # librapidsmpf-cu13 wheels brings in a hardcoded libnuma-dev cmake target: https://github.com/NVIDIA/cuCascade/issues/118
-      # -DBUILD_TESTS=OFF to match rapidsmpf https://github.com/rapidsai/rapidsmpf/blob/main/.github/workflows/pr.yaml#L351 (leads to compilation errors)
-      # -DCUDF_BUILD_TESTUTIL=OFF to avoid IMPORTED_GLOBAL promotion errors when cuCascade's find_package(cudf) loads cudf-config.cmake from a CPM subdirectory
       build_command: |
         sccache --zero-stats;
-        clone-rapidsmpf -j$(nproc) -v -q --branch "$(cat ~/cudf/RAPIDS_BRANCH)" --clone-upstream --depth 1 --single-branch --shallow-submodules;
-        build-all -j0 -DBUILD_BENCHMARKS=OFF -DBUILD_NUMA_SUPPORT=OFF -DBUILD_TESTS=OFF -DCUDF_BUILD_TESTUTIL=OFF --verbose 2>&1 | tee telemetry-artifacts/build.log;
+        build-all -j0 -DBUILD_BENCHMARKS=ON --verbose 2>&1 | tee telemetry-artifacts/build.log;
         sccache --show-adv-stats | tee telemetry-artifacts/sccache-stats.txt;
   unit-tests-cudf-pandas:
     needs: [wheel-build-cudf, changed-files]

ci/test_cudf_polars_polars_tests.sh

@@ -12,6 +12,8 @@ RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
 CUDF_POLARS_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cudf_polars_${RAPIDS_PY_CUDA_SUFFIX}" RAPIDS_PY_WHEEL_PURE="1" rapids-download-wheels-from-github python)
 LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp)
 PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")")
+LIBCUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp)
+CUDF_STREAMING_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" cudf_streaming --stable --cuda "$RAPIDS_CUDA_VERSION")")
 
 # generate constraints (possibly pinning to oldest support versions of dependencies)
 rapids-generate-pip-constraints py_test_cudf_polars "${PIP_CONSTRAINT}"
@@ -29,7 +31,9 @@ rapids-pip-retry install \
     --constraint "${PIP_CONSTRAINT}" \
     "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test]" \
     "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
-    "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)"
+    "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
+    "$(echo "${LIBCUDF_STREAMING_WHEELHOUSE}"/libcudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
+    "$(echo "${CUDF_STREAMING_WHEELHOUSE}"/cudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)"
 
 TAG=$(python -c 'import polars; print(f"py-{polars.__version__}")')
 rapids-logger "Clone polars to ${TAG}"

ci/test_wheel_cudf_polars.sh

@@ -15,6 +15,10 @@ CUDF_POLARS_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="cudf_polars_${RAPIDS_PY_CUDA_SUFF
 LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp)
 PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")")
 
+# Download libcudf_streaming and cudf_streaming built in the previous step
+LIBCUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp)
+CUDF_STREAMING_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" cudf_streaming --stable --cuda "$RAPIDS_CUDA_VERSION")")
+
 # generate constraints (possibly pinning to oldest support versions of dependencies)
 rapids-generate-pip-constraints py_test_cudf_polars "${PIP_CONSTRAINT}"
 
@@ -65,6 +69,8 @@ for version in "${VERSIONS[@]}"; do
         "$(echo "${CUDF_POLARS_WHEELHOUSE}"/cudf_polars_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)[test,dask,ray]" \
         "$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
         "$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
+        "$(echo "${LIBCUDF_STREAMING_WHEELHOUSE}"/libcudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
+        "$(echo "${CUDF_STREAMING_WHEELHOUSE}"/cudf_streaming_"${RAPIDS_PY_CUDA_SUFFIX}"*.whl)" \
         -r "polars-compat-${version}-requirements.txt"
 
     rapids-logger "Running tests for polars ${version}.*"

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -44,6 +44,7 @@ dependencies:
 - libkvikio==26.8.*,>=0.0.0a0
 - libnvcomp-dev==5.2.0.10
 - libnvjitlink-dev
+- librapidsmpf==26.8.*,>=0.0.0a0
 - librdkafka>=2.11.1,<2.12.0
 - librmm==26.8.*,>=0.0.0a0
 - make
@@ -87,6 +88,7 @@ dependencies:
 - rapids-build-backend>=0.4.0,<0.5.0
 - rapids-dask-dependency==26.8.*,>=0.0.0a0
 - rapids-logger==0.2.*,>=0.0.0a0
+- rapidsmpf==26.8.*,>=0.0.0a0
 - ray-default>=2.55.1
 - rich
 - rmm==26.8.*,>=0.0.0a0

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -44,6 +44,7 @@ dependencies:
 - libkvikio==26.8.*,>=0.0.0a0
 - libnvcomp-dev==5.2.0.10
 - libnvjitlink-dev
+- librapidsmpf==26.8.*,>=0.0.0a0
 - librdkafka>=2.11.1,<2.12.0
 - librmm==26.8.*,>=0.0.0a0
 - make
@@ -87,6 +88,7 @@ dependencies:
 - rapids-build-backend>=0.4.0,<0.5.0
 - rapids-dask-dependency==26.8.*,>=0.0.0a0
 - rapids-logger==0.2.*,>=0.0.0a0
+- rapidsmpf==26.8.*,>=0.0.0a0
 - ray-default>=2.55.1
 - rich
 - rmm==26.8.*,>=0.0.0a0

conda/environments/all_cuda-132_arch-aarch64.yaml

@@ -44,6 +44,7 @@ dependencies:
 - libkvikio==26.8.*,>=0.0.0a0
 - libnvcomp-dev==5.2.0.10
 - libnvjitlink-dev
+- librapidsmpf==26.8.*,>=0.0.0a0
 - librdkafka>=2.11.1,<2.12.0
 - librmm==26.8.*,>=0.0.0a0
 - make
@@ -87,6 +88,7 @@ dependencies:
 - rapids-build-backend>=0.4.0,<0.5.0
 - rapids-dask-dependency==26.8.*,>=0.0.0a0
 - rapids-logger==0.2.*,>=0.0.0a0
+- rapidsmpf==26.8.*,>=0.0.0a0
 - ray-default>=2.55.1
 - rich
 - rmm==26.8.*,>=0.0.0a0

conda/environments/all_cuda-132_arch-x86_64.yaml

@@ -44,6 +44,7 @@ dependencies:
 - libkvikio==26.8.*,>=0.0.0a0
 - libnvcomp-dev==5.2.0.10
 - libnvjitlink-dev
+- librapidsmpf==26.8.*,>=0.0.0a0
 - librdkafka>=2.11.1,<2.12.0
 - librmm==26.8.*,>=0.0.0a0
 - make
@@ -87,6 +88,7 @@ dependencies:
 - rapids-build-backend>=0.4.0,<0.5.0
 - rapids-dask-dependency==26.8.*,>=0.0.0a0
 - rapids-logger==0.2.*,>=0.0.0a0
+- rapidsmpf==26.8.*,>=0.0.0a0
 - ray-default>=2.55.1
 - rich
 - rmm==26.8.*,>=0.0.0a0

conda/recipes/cudf-polars/recipe.yaml

@@ -37,6 +37,7 @@ requirements:
     - python
     - pylibcudf =${{ version }}
     - rapidsmpf =${{ minor_version }}
+    - cudf-streaming =${{ version }}
     - polars>=1.35,<1.42
     - packaging
     - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}

cpp/CMakeLists.txt

@@ -290,6 +290,16 @@ endfunction()
 # * conda environment -----------------------------------------------------------------------------
 rapids_cmake_support_conda_env(conda_env MODIFY_PREFIX_PATH)
 
+# ##################################################################################################
+# * find C++ libraries from RAPIDS wheels in in site-packages --------------------------------------
+find_package(Python COMPONENTS Interpreter)
+if(Python_FOUND)
+  include(${rapids-cmake-dir}/cython-core/find_prefix_paths.cmake)
+  rapids_cython_find_prefix_paths("${Python_EXECUTABLE}" wheel_prefix_paths)
+  set_property(GLOBAL PROPERTY FIND_LIBRARY_USE_LIB64_PATHS TRUE)
+  list(APPEND CMAKE_PREFIX_PATH ${wheel_prefix_paths})
+endif()
+
 # ##################################################################################################
 # * compiler options ------------------------------------------------------------------------------
 set(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Modules ${CMAKE_MODULE_PATH})

dependencies.yaml

@@ -23,8 +23,10 @@ files:
       - depends_on_libkvikio
       - depends_on_librmm
       - depends_on_libnvcomp
+      - depends_on_librapidsmpf
       - depends_on_numba_cuda
       - depends_on_rapids_logger
+      - depends_on_rapidsmpf
       - depends_on_ray
       - depends_on_rmm
       - develop
@@ -335,6 +337,7 @@ files:
     includes:
       - run_cudf_polars
       - depends_on_rapidsmpf
+      - depends_on_cudf_streaming
       - depends_on_pylibcudf
       - depends_on_cuda_python
   # TODO(26.10): Remove this alias in favor of dask
@@ -1414,7 +1417,27 @@ dependencies:
     common:
       - output_types: conda
         packages:
-          - cudf-streaming==26.8.*,>=0.0.0a0
+          - &cudf_streaming_unsuffixed cudf-streaming==26.8.*,>=0.0.0a0
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+              cuda_suffixed: "true"
+            packages:
+              - cudf-streaming-cu12==26.8.*,>=0.0.0a0
+          - matrix:
+              cuda: "13.*"
+              cuda_suffixed: "true"
+            packages:
+              - cudf-streaming-cu13==26.8.*,>=0.0.0a0
+          - matrix:
+            packages:
+              - *cudf_streaming_unsuffixed
   depends_on_ray:
     common:
       - output_types: conda

python/cudf_polars/cudf_polars/engine/core.py

@@ -14,14 +14,15 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeVar
 
 import cuda.core
+
+import polars as pl
+
+from cudf_streaming.streaming.table_chunk import TableChunk
 from rapidsmpf.coll import AllGather
 from rapidsmpf.config import Options, get_environment_variables
 from rapidsmpf.memory.packed_data import PackedData
 from rapidsmpf.statistics import Statistics
 from rapidsmpf.streaming.core.actor import run_actor_network
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
-
-import polars as pl
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IRExecutionContext
@@ -42,10 +43,10 @@
     from concurrent.futures import Executor, ThreadPoolExecutor
 
     import rapidsmpf.config
+    from cudf_streaming.streaming.channel_metadata import ChannelMetadata
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.memory.buffer_resource import BufferResource
     from rapidsmpf.streaming.core.context import Context
-    from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
 
     from cudf_polars.dsl.ir import IR
     from cudf_polars.streaming.base import PartitionInfo

python/cudf_polars/cudf_polars/engine/dask.py

@@ -16,6 +16,10 @@
 import distributed.system
 import pynvml
 import ucxx._lib.libucxx as ucx_api
+
+import polars as pl
+
+import rmm.mr
 from rapidsmpf import bootstrap
 from rapidsmpf.communicator.ucxx import barrier, get_root_ucxx_address, new_communicator
 from rapidsmpf.config import Options
@@ -24,10 +28,6 @@
 from rapidsmpf.statistics import Statistics
 from rapidsmpf.streaming.core.context import Context
 
-import polars as pl
-
-import rmm.mr
-
 from cudf_polars.engine.core import (
     ClusterInfo,
     StreamingEngine,
@@ -44,8 +44,8 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from cudf_streaming.streaming.channel_metadata import ChannelMetadata
     from rapidsmpf.communicator.communicator import Communicator
-    from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
 
     from cudf_polars.dsl.ir import IR
     from cudf_polars.engine.core import T