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

.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

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/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") }}

dependencies.yaml

@@ -335,6 +335,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 +1415,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,12 +14,12 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeVar
 
 import cuda.core
+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
 
@@ -42,10 +42,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

@@ -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

python/cudf_polars/cudf_polars/engine/ray.py

@@ -41,8 +41,8 @@
     import uuid
     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 ray.actor import ActorHandle
 
     from cudf_polars.dsl.ir import IR

python/cudf_polars/cudf_polars/engine/spmd.py

@@ -10,13 +10,13 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any, cast
 
+from cudf_streaming.integrations.partition import unpack_and_concat
 from rapidsmpf import bootstrap
 from rapidsmpf.coll import AllGather
 from rapidsmpf.communicator.single import (
     new_communicator as single_communicator,
 )
 from rapidsmpf.communicator.ucxx import barrier
-from rapidsmpf.integrations.cudf.partition import unpack_and_concat
 from rapidsmpf.memory.packed_data import PackedData
 from rapidsmpf.progress_thread import ProgressThread
 from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor
@@ -51,9 +51,9 @@
     import uuid
     from collections.abc import Callable
 
+    from cudf_streaming.streaming.channel_metadata import ChannelMetadata
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.config import Options
-    from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
 
     import polars as pl
 

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py

@@ -6,16 +6,16 @@
 
 from typing import TYPE_CHECKING
 
-from rapidsmpf.integrations.cudf.partition import unpack_and_concat
+from cudf_streaming.integrations.partition import unpack_and_concat
 from rapidsmpf.memory.packed_data import PackedData
 from rapidsmpf.streaming.coll.allgather import AllGather
 
 from pylibcudf.contiguous_split import pack
 
 if TYPE_CHECKING:
+    from cudf_streaming.streaming.table_chunk import TableChunk
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.context import Context
-    from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
     import pylibcudf as plc
     from rmm.pylibrmm.stream import Stream

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py

@@ -6,24 +6,24 @@
 
 from typing import TYPE_CHECKING, Any
 
-from rapidsmpf.communicator.single import new_communicator as single_comm
-from rapidsmpf.config import Options, get_environment_variables
-from rapidsmpf.integrations.cudf.partition import (
+from cudf_streaming.integrations.partition import (
     partition_and_pack as py_partition_and_pack,
     split_and_pack as py_split_and_pack,
     unpack_and_concat as py_unpack_and_concat,
 )
+from cudf_streaming.streaming.channel_metadata import (
+    ChannelMetadata,
+    HashScheme,
+    Partitioning,
+)
+from cudf_streaming.streaming.table_chunk import TableChunk
+from rapidsmpf.communicator.single import new_communicator as single_comm
+from rapidsmpf.config import Options, get_environment_variables
 from rapidsmpf.shuffler import PartitionAssignment
 from rapidsmpf.streaming.coll.shuffler import ShufflerAsync
 from rapidsmpf.streaming.core.actor import define_actor
 from rapidsmpf.streaming.core.context import Context
 from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.channel_metadata import (
-    ChannelMetadata,
-    HashScheme,
-    Partitioning,
-)
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
 import pylibcudf as plc
 import pylibcudf.partitioning
@@ -445,7 +445,7 @@ async def shuffle_actor(
 
     This node combines partition_and_pack, shuffler, and unpack_and_concat
     into a single Python node using rapidsmpf.shuffler.Shuffler and utilities
-    from rapidsmpf.integrations.cudf.partition.
+    from cudf_streaming.integrations.partition.
 
     Parameters
     ----------

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py

@@ -6,16 +6,16 @@
 
 from typing import TYPE_CHECKING
 
-from rapidsmpf.shuffler import PartitionAssignment
-from rapidsmpf.streaming.core.actor import define_actor
-from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.channel_metadata import (
+from cudf_streaming.streaming.channel_metadata import (
     ChannelMetadata,
     OrderKey,
     OrderScheme,
     Partitioning,
 )
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
+from cudf_streaming.streaming.table_chunk import TableChunk
+from rapidsmpf.shuffler import PartitionAssignment
+from rapidsmpf.streaming.core.actor import define_actor
+from rapidsmpf.streaming.core.message import Message
 
 import polars as pl
 

python/cudf_polars/cudf_polars/streaming/actor_graph/core.py

@@ -30,12 +30,12 @@
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
 
+    from cudf_streaming.streaming.channel_metadata import ChannelMetadata
+    from cudf_streaming.streaming.table_chunk import TableChunk
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
     from rapidsmpf.streaming.core.leaf_actor import DeferredMessages
-    from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
-    from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
     import polars as pl
 

python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py

@@ -7,16 +7,16 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+from cudf_streaming.streaming.channel_metadata import (
+    ChannelMetadata,
+    HashScheme,
+)
+from cudf_streaming.streaming.table_chunk import TableChunk
 from rapidsmpf.communicator.single import new_communicator as single_comm
 from rapidsmpf.config import Options, get_environment_variables
 from rapidsmpf.streaming.core.actor import define_actor
 from rapidsmpf.streaming.core.context import Context
 from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.channel_metadata import (
-    ChannelMetadata,
-    HashScheme,
-)
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
 import pylibcudf as plc
 

python/cudf_polars/cudf_polars/streaming/actor_graph/io.py

@@ -8,13 +8,13 @@
 import math
 from typing import TYPE_CHECKING, Any
 
+from cudf_streaming.streaming.channel_metadata import ChannelMetadata
+from cudf_streaming.streaming.table_chunk import TableChunk
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.streaming.core.memory_reserve_or_wait import (
     reserve_memory,
 )
 from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
 import pylibcudf as plc
 
@@ -481,7 +481,7 @@ def make_rapidsmpf_read_parquet_node(
     The RapidsMPF read parquet node, or None if the predicate cannot be
     converted to a parquet filter (caller should fall back to scan_node).
     """
-    from rapidsmpf.streaming.cudf.parquet import Filter, read_parquet
+    from cudf_streaming.streaming.parquet import Filter, read_parquet
 
     # Build ParquetReaderOptions
     try:

python/cudf_polars/cudf_polars/streaming/actor_graph/join.py

@@ -7,22 +7,22 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal
 
-from rapidsmpf.memory.memory_reservation import opaque_memory_usage
-from rapidsmpf.streaming.core.actor import define_actor
-from rapidsmpf.streaming.core.memory_reserve_or_wait import (
-    reserve_memory,
-)
-from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.bloom_filter import BloomFilter
-from rapidsmpf.streaming.cudf.channel_metadata import (
+from cudf_streaming.streaming.bloom_filter import BloomFilter
+from cudf_streaming.streaming.channel_metadata import (
     ChannelMetadata,
     HashScheme,
     Partitioning,
 )
-from rapidsmpf.streaming.cudf.table_chunk import (
+from cudf_streaming.streaming.table_chunk import (
     TableChunk,
     make_table_chunks_available_or_wait,
 )
+from rapidsmpf.memory.memory_reservation import opaque_memory_usage
+from rapidsmpf.streaming.core.actor import define_actor
+from rapidsmpf.streaming.core.memory_reserve_or_wait import (
+    reserve_memory,
+)
+from rapidsmpf.streaming.core.message import Message
 
 import pylibcudf as plc
 from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED
@@ -60,10 +60,10 @@
     from collections.abc import Iterable, MutableMapping
     from types import CoroutineType
 
+    from cudf_streaming.streaming.bloom_filter import BloomFilterChunk
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
-    from rapidsmpf.streaming.cudf.bloom_filter import BloomFilterChunk
 
     from cudf_polars.dsl.ir import IR, IRExecutionContext
     from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator

python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py

@@ -7,16 +7,16 @@
 import asyncio
 from typing import TYPE_CHECKING, Any, cast
 
+from cudf_streaming.streaming.channel_metadata import ChannelMetadata
+from cudf_streaming.streaming.table_chunk import (
+    TableChunk,
+    make_table_chunks_available_or_wait,
+)
 from rapidsmpf.memory.buffer import MemoryType
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.streaming.core.actor import define_actor
 from rapidsmpf.streaming.core.message import Message
 from rapidsmpf.streaming.core.spillable_messages import SpillableMessages
-from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata
-from rapidsmpf.streaming.cudf.table_chunk import (
-    TableChunk,
-    make_table_chunks_available_or_wait,
-)
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IR, Empty