From 9081b83fb2bb36e9480e14c853cd772393edd4f8 Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 19 May 2026 09:06:40 -0700
Subject: [PATCH 1/2] move collectives into actor_graph

---
 python/cudf_polars/cudf_polars/engine/core.py          |  4 ++--
 python/cudf_polars/cudf_polars/engine/spmd.py          |  2 +-
 .../cudf_polars/streaming/actor_graph/__init__.py      |  7 ++++---
 .../{ => actor_graph}/collectives/__init__.py          |  2 +-
 .../{ => actor_graph}/collectives/allgather.py         |  0
 .../streaming/{ => actor_graph}/collectives/common.py  |  0
 .../streaming/{ => actor_graph}/collectives/shuffle.py |  0
 .../streaming/{ => actor_graph}/collectives/sort.py    |  4 ++--
 .../cudf_polars/streaming/actor_graph/groupby.py       |  2 +-
 .../cudf_polars/streaming/actor_graph/join.py          |  4 ++--
 .../cudf_polars/streaming/actor_graph/over.py          |  8 ++++----
 .../cudf_polars/streaming/actor_graph/repartition.py   |  2 +-
 .../cudf_polars/streaming/actor_graph/utils.py         |  2 +-
 .../cudf_polars/streaming/benchmarks/utils.py          |  2 +-
 python/cudf_polars/tests/streaming/test_allgather.py   |  2 +-
 python/cudf_polars/tests/streaming/test_groupby.py     |  2 +-
 python/cudf_polars/tests/streaming/test_metadata.py    |  8 ++++----
 python/cudf_polars/tests/streaming/test_shuffler.py    | 10 +++++-----
 python/cudf_polars/tests/streaming/test_spmd.py        |  2 +-
 19 files changed, 32 insertions(+), 31 deletions(-)
 rename python/cudf_polars/cudf_polars/streaming/{ => actor_graph}/collectives/__init__.py (82%)
 rename python/cudf_polars/cudf_polars/streaming/{ => actor_graph}/collectives/allgather.py (100%)
 rename python/cudf_polars/cudf_polars/streaming/{ => actor_graph}/collectives/common.py (100%)
 rename python/cudf_polars/cudf_polars/streaming/{ => actor_graph}/collectives/shuffle.py (100%)
 rename python/cudf_polars/cudf_polars/streaming/{ => actor_graph}/collectives/sort.py (99%)

diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py
index 23a295fad3a0..45aeaa8d270e 100644
--- a/python/cudf_polars/cudf_polars/engine/core.py
+++ b/python/cudf_polars/cudf_polars/engine/core.py
@@ -25,12 +25,12 @@
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IRExecutionContext
+from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs
+from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.streaming.actor_graph.core import generate_network
 from cudf_polars.streaming.actor_graph.tracing import log_query_plan
 from cudf_polars.streaming.actor_graph.utils import empty_table_chunk
 from cudf_polars.streaming.base import StatsCollector
-from cudf_polars.streaming.collectives import ReserveOpIDs
-from cudf_polars.streaming.collectives.common import reserve_op_id
 from cudf_polars.streaming.parallel import lower_ir_graph
 from cudf_polars.streaming.statistics import collect_statistics
 from cudf_polars.streaming.utils import _concat
diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py
index 10aae6ab1b3b..64e7b1893501 100644
--- a/python/cudf_polars/cudf_polars/engine/spmd.py
+++ b/python/cudf_polars/cudf_polars/engine/spmd.py
@@ -39,8 +39,8 @@
     HardwareBindingPolicy,
     bind_to_gpu,
 )
+from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.streaming.actor_graph.utils import set_memory_resource
-from cudf_polars.streaming.collectives.common import reserve_op_id
 from cudf_polars.utils.config import (
     MemoryResourceConfig,
     SPMDContext,
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py
index 9bc75910c2fd..a4dd049a0eae 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py
@@ -5,6 +5,9 @@
 
 from __future__ import annotations
 
+import cudf_polars.streaming.actor_graph.collectives.shuffle
+import cudf_polars.streaming.actor_graph.collectives.sort
+
 # Side-effect imports: each module registers
 # ``@generate_ir_sub_network.register(...)`` handlers at import time so the
 # dispatch table is populated before any query is evaluated.
@@ -13,8 +16,6 @@
 import cudf_polars.streaming.actor_graph.join
 import cudf_polars.streaming.actor_graph.over
 import cudf_polars.streaming.actor_graph.repartition
-import cudf_polars.streaming.actor_graph.union
-import cudf_polars.streaming.collectives.shuffle
-import cudf_polars.streaming.collectives.sort  # noqa: F401
+import cudf_polars.streaming.actor_graph.union  # noqa: F401
 
 __all__: list[str] = []
diff --git a/python/cudf_polars/cudf_polars/streaming/collectives/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
similarity index 82%
rename from python/cudf_polars/cudf_polars/streaming/collectives/__init__.py
rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
index 38ebc845b4e0..80126e9999e6 100644
--- a/python/cudf_polars/cudf_polars/streaming/collectives/__init__.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from cudf_polars.streaming.collectives.common import (
+from cudf_polars.streaming.actor_graph.collectives.common import (
     ReserveOpIDs,
     reserve_op_id,
 )
diff --git a/python/cudf_polars/cudf_polars/streaming/collectives/allgather.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py
similarity index 100%
rename from python/cudf_polars/cudf_polars/streaming/collectives/allgather.py
rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py
diff --git a/python/cudf_polars/cudf_polars/streaming/collectives/common.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py
similarity index 100%
rename from python/cudf_polars/cudf_polars/streaming/collectives/common.py
rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py
diff --git a/python/cudf_polars/cudf_polars/streaming/collectives/shuffle.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py
similarity index 100%
rename from python/cudf_polars/cudf_polars/streaming/collectives/shuffle.py
rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py
diff --git a/python/cudf_polars/cudf_polars/streaming/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
similarity index 99%
rename from python/cudf_polars/cudf_polars/streaming/collectives/sort.py
rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
index fa288b999da8..75087141c9b2 100644
--- a/python/cudf_polars/cudf_polars/streaming/collectives/sort.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
@@ -25,6 +25,8 @@
 from cudf_polars.dsl.expr import Col, NamedExpr
 from cudf_polars.dsl.ir import Empty, Sort
 from cudf_polars.dsl.utils.naming import names_to_indices, unique_names
+from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
+from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager
 from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network
 from cudf_polars.streaming.actor_graph.nodes import (
     default_node_single,
@@ -47,8 +49,6 @@
     replay_buffered_channel,
     send_metadata,
 )
-from cudf_polars.streaming.collectives.allgather import AllGatherManager
-from cudf_polars.streaming.collectives.shuffle import ShuffleManager
 from cudf_polars.streaming.repartition import Repartition
 from cudf_polars.streaming.sort import (
     _get_final_sort_boundaries,
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
index c0308d35de33..22be38dade4e 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
@@ -24,6 +24,7 @@
 from cudf_polars.dsl.expr import Col, NamedExpr
 from cudf_polars.dsl.ir import IR, Distinct, GroupBy, Select
 from cudf_polars.dsl.utils.naming import unique_names
+from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager
 from cudf_polars.streaming.actor_graph.dispatch import (
     generate_ir_sub_network,
 )
@@ -43,7 +44,6 @@
     send_metadata,
     shutdown_on_error,
 )
-from cudf_polars.streaming.collectives.shuffle import ShuffleManager
 from cudf_polars.streaming.groupby import combine, decompose
 from cudf_polars.streaming.repartition import Repartition
 
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
index d6f1828204a0..2c0efb5bf643 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
@@ -30,6 +30,8 @@
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IR, Join
 from cudf_polars.dsl.utils.naming import names_to_indices
+from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
+from cudf_polars.streaming.actor_graph.collectives.shuffle import _global_shuffle
 from cudf_polars.streaming.actor_graph.dispatch import (
     generate_ir_sub_network,
 )
@@ -51,8 +53,6 @@
     send_metadata,
     shutdown_on_error,
 )
-from cudf_polars.streaming.collectives.allgather import AllGatherManager
-from cudf_polars.streaming.collectives.shuffle import _global_shuffle
 from cudf_polars.streaming.repartition import Repartition
 from cudf_polars.streaming.utils import _concat
 
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
index b990e9ec5b62..d1569d561b21 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
@@ -59,6 +59,10 @@
 from cudf_polars.dsl.expressions.base import ExecutionContext
 from cudf_polars.dsl.utils.naming import unique_names
 from cudf_polars.dsl.utils.reshape import broadcast
+from cudf_polars.streaming.actor_graph.collectives.shuffle import (
+    LocalRepartitioner,
+    ShuffleManager,
+)
 from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network
 from cudf_polars.streaming.actor_graph.utils import (
     ChannelManager,
@@ -81,10 +85,6 @@
     send_metadata,
     shutdown_on_error,
 )
-from cudf_polars.streaming.collectives.shuffle import (
-    LocalRepartitioner,
-    ShuffleManager,
-)
 from cudf_polars.streaming.over import Over, _build_over_groupby_irs
 
 if TYPE_CHECKING:
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
index 5eda974db890..fcfc897e1728 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py
@@ -17,6 +17,7 @@
 )
 
 from cudf_polars.containers import DataFrame
+from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
 from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network
 from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error
 from cudf_polars.streaming.actor_graph.utils import (
@@ -25,7 +26,6 @@
     recv_metadata,
     send_metadata,
 )
-from cudf_polars.streaming.collectives.allgather import AllGatherManager
 from cudf_polars.streaming.repartition import Repartition
 from cudf_polars.streaming.utils import _concat
 
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
index 41bad69e5514..ac87b16d5872 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -41,8 +41,8 @@
 from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select
 from cudf_polars.dsl.tracing import Scope
 from cudf_polars.dsl.utils.naming import names_to_indices
+from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
 from cudf_polars.streaming.actor_graph.tracing import ActorTracer
-from cudf_polars.streaming.collectives.allgather import AllGatherManager
 from cudf_polars.streaming.utils import _concat
 from cudf_polars.utils.dtypes import make_empty_column
 
diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py
index 1abc833a3183..502ecda5abe6 100644
--- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py
@@ -1187,7 +1187,7 @@ def run_polars_spmd(
         from cudf_polars.engine.spmd import (
             allgather_polars_dataframe,
         )
-        from cudf_polars.streaming.collectives.common import reserve_op_id
+        from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 
         def _allgather_result(df: pl.DataFrame) -> pl.DataFrame:
             with reserve_op_id() as op_id:
diff --git a/python/cudf_polars/tests/streaming/test_allgather.py b/python/cudf_polars/tests/streaming/test_allgather.py
index af97089b129d..48d84afdb00f 100644
--- a/python/cudf_polars/tests/streaming/test_allgather.py
+++ b/python/cudf_polars/tests/streaming/test_allgather.py
@@ -15,8 +15,8 @@
 import pylibcudf as plc
 
 from cudf_polars.dsl.ir import IRExecutionContext
+from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
 from cudf_polars.streaming.actor_graph.utils import allgather_reduce
-from cudf_polars.streaming.collectives.allgather import AllGatherManager
 
 
 async def _test_allgather(engine) -> None:
diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py
index fbc18e3fa8e6..99835e125fd9 100644
--- a/python/cudf_polars/tests/streaming/test_groupby.py
+++ b/python/cudf_polars/tests/streaming/test_groupby.py
@@ -12,7 +12,7 @@
 import polars as pl
 
 from cudf_polars.engine.options import StreamingOptions
-from cudf_polars.streaming.collectives.shuffle import ShuffleManager
+from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager
 from cudf_polars.testing.asserts import assert_gpu_result_equal
 
 
diff --git a/python/cudf_polars/tests/streaming/test_metadata.py b/python/cudf_polars/tests/streaming/test_metadata.py
index b226e51c2fbe..5c555e3593ba 100644
--- a/python/cudf_polars/tests/streaming/test_metadata.py
+++ b/python/cudf_polars/tests/streaming/test_metadata.py
@@ -24,15 +24,15 @@
 from cudf_polars.dsl import expr
 from cudf_polars.dsl.ir import GroupBy, HStack, Projection, Select, Sort
 from cudf_polars.engine.options import StreamingOptions
+from cudf_polars.streaming.actor_graph.collectives.sort import (
+    _is_already_sorted,
+    _sort_to_order_keys,
+)
 from cudf_polars.streaming.actor_graph.core import evaluate_logical_plan
 from cudf_polars.streaming.actor_graph.utils import (
     NormalizedPartitioning,
     maybe_remap_partitioning,
 )
-from cudf_polars.streaming.collectives.sort import (
-    _is_already_sorted,
-    _sort_to_order_keys,
-)
 from cudf_polars.utils.config import ConfigOptions
 
 
diff --git a/python/cudf_polars/tests/streaming/test_shuffler.py b/python/cudf_polars/tests/streaming/test_shuffler.py
index df71e7b07cdc..7f4fd46d15e1 100644
--- a/python/cudf_polars/tests/streaming/test_shuffler.py
+++ b/python/cudf_polars/tests/streaming/test_shuffler.py
@@ -18,14 +18,14 @@
 from cudf_polars.containers import DataFrame, DataType
 from cudf_polars.engine.options import StreamingOptions
 from cudf_polars.engine.spmd import allgather_polars_dataframe
-from cudf_polars.streaming.actor_graph.utils import (
-    _is_already_partitioned,
-)
-from cudf_polars.streaming.collectives.common import reserve_op_id
-from cudf_polars.streaming.collectives.shuffle import (
+from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
+from cudf_polars.streaming.actor_graph.collectives.shuffle import (
     LocalRepartitioner,
     ShuffleManager,
 )
+from cudf_polars.streaming.actor_graph.utils import (
+    _is_already_partitioned,
+)
 from cudf_polars.testing.asserts import assert_gpu_result_equal
 
 
diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py
index 6b95c91c178d..75ca4888d083 100644
--- a/python/cudf_polars/tests/streaming/test_spmd.py
+++ b/python/cudf_polars/tests/streaming/test_spmd.py
@@ -22,7 +22,7 @@
     SPMDEngine,
     allgather_polars_dataframe,
 )
-from cudf_polars.streaming.collectives.common import reserve_op_id
+from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.testing.asserts import assert_gpu_result_equal
 from cudf_polars.utils.config import MemoryResourceConfig
 

From e417b31ba950d7263d22874a20060c9b66ef7b5b Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 19 May 2026 10:53:34 -0700
Subject: [PATCH 2/2] arbitrary change to trigger CI

---
 .../cudf_polars/streaming/actor_graph/collectives/__init__.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
index 80126e9999e6..5cbf011ca353 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py
@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
-"""Collective operations for the RapidsMPF streaming runtime."""
+"""Collective operations for building a RapidsMPF actor graph."""
 
 from __future__ import annotations