From 666ce238653ebbf6cca031b21c7ae8d05e64ffde Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 30 Jun 2026 13:11:19 -0700
Subject: [PATCH 01/11] allow sampled chunks to be spilled

---
 .../streaming/actor_graph/utils.py            | 38 ++++++++++++-------
 1 file changed, 25 insertions(+), 13 deletions(-)

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 02ae6b546a82..3a2e8e016cd6 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -52,6 +52,7 @@
         Coroutine,
         Generator,
         Iterator,
+        Mapping,
         Sequence,
     )
 
@@ -93,6 +94,17 @@ def __iter__(self) -> Generator[Message, None, None]:
             yield self._store.extract(mid=self._mids.popleft())
 
 
+def buffered_chunk_messages(
+    buffered_chunks: ChunkStore | Mapping[int, TableChunk],
+) -> Iterator[Message]:
+    """Yield buffered chunks as messages, consuming the buffer."""
+    if isinstance(buffered_chunks, ChunkStore):
+        yield from buffered_chunks
+    else:
+        for seq_num, chunk in buffered_chunks.items():
+            yield Message(seq_num, chunk)
+
+
 @contextlib.contextmanager
 def set_memory_resource(mr: rmm.mr.DeviceMemoryResource) -> Iterator[None]:
     """
@@ -731,8 +743,8 @@ def indices_to_names(indices: tuple[int, ...], schema: Schema) -> tuple[str, ...
 class TableSizeStats:
     """Sampled chunks and aggregate size/row stats for a table channel."""
 
-    chunks: dict[int, TableChunk] = field(default_factory=dict)
-    """The sampled chunks, keyed by sequence number."""
+    chunks: ChunkStore | dict[int, TableChunk] = field(default_factory=dict)
+    """The sampled chunks/messages, keyed or ordered by sequence number."""
     total_size: int = 0
     """The total estimated size of the table in bytes."""
     total_rows: int = 0
@@ -768,24 +780,24 @@ async def _sample_chunks(
     -------
     Sampled chunks and the extrapolated total size/rows for this rank.
     """
-    sampled_chunks: dict[int, TableChunk] = {}
+    sampled_chunks = ChunkStore(context)
+    sampled_count = 0
     total_size = 0
     total_rows = 0
     for _ in range(max_sample_chunks):
         msg = await ch.recv(context)
         if msg is None:
             break
-        chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill(
-            context.br(), allow_overbooking=True
-        )
-        sampled_chunks[msg.sequence_number] = chunk
+        chunk = TableChunk.from_message(msg, br=context.br())
         total_size += chunk.data_alloc_size()
         total_rows += chunk.shape[0]
+        sampled_count += 1
+        sampled_chunks.insert(Message(msg.sequence_number, chunk))
         if total_size >= max_sample_bytes:
             break
-    if sampled_chunks:
-        total_size = int((total_size / len(sampled_chunks)) * local_count)
-        total_rows = int((total_rows / len(sampled_chunks)) * local_count)
+    if sampled_count:
+        total_size = int((total_size / sampled_count) * local_count)
+        total_rows = int((total_rows / sampled_count) * local_count)
     return TableSizeStats(
         chunks=sampled_chunks,
         total_size=total_size,
@@ -798,7 +810,7 @@ async def replay_buffered_channel(
     context: Context,
     ch_out: Channel[TableChunk],
     ch_in: Channel[TableChunk],
-    buffered_chunks: dict[int, TableChunk],
+    buffered_chunks: ChunkStore | Mapping[int, TableChunk],
     metadata: ChannelMetadata,
     *,
     trace_ir: IR,
@@ -823,8 +835,8 @@ async def replay_buffered_channel(
     """
     async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir):
         await send_metadata(ch_out, context, metadata)
-        for seq_num, chunk in buffered_chunks.items():
-            await ch_out.send(context, Message(seq_num, chunk))
+        for msg in buffered_chunk_messages(buffered_chunks):
+            await ch_out.send(context, msg)
         while (msg := await ch_in.recv(context)) is not None:
             await ch_out.send(context, msg)
         await ch_out.drain(context)

From 6ad5b3ad9eff94ea9e781018be6e55f93825f21b Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 30 Jun 2026 13:21:29 -0700
Subject: [PATCH 02/11] require ChunkStore

---
 .../cudf_polars/streaming/actor_graph/join.py         | 11 +++++++++--
 .../cudf_polars/streaming/actor_graph/utils.py        |  6 +++---
 2 files changed, 12 insertions(+), 5 deletions(-)

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 031a7658f984..05ecbce4a25b 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
@@ -40,6 +40,7 @@
     CUDF_ROW_LIMIT,
     MAX_ROWS_PER_PARTITION,
     ChannelManager,
+    ChunkStore,
     NormalizedPartitioning,
     TableSizeStats,
     _sample_chunks,
@@ -1216,8 +1217,14 @@ async def _choose_strategy(
     if left_partitioning.is_aligned_with(right_partitioning, context.br()):
         # We can use a chunkwise join
         chunkwise = True
-        left_sample = TableSizeStats(total_chunks=left_metadata.local_count)
-        right_sample = TableSizeStats(total_chunks=right_metadata.local_count)
+        left_sample = TableSizeStats(
+            chunks=ChunkStore(context),
+            total_chunks=left_metadata.local_count,
+        )
+        right_sample = TableSizeStats(
+            chunks=ChunkStore(context),
+            total_chunks=right_metadata.local_count,
+        )
     else:
         # Need to shuffle or broadcast - Use sampled data to choose a strategy
         chunkwise = False
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 3a2e8e016cd6..f07132d743d0 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -12,7 +12,7 @@
 import time
 from collections import deque
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from functools import reduce
 from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast
 
@@ -743,8 +743,8 @@ def indices_to_names(indices: tuple[int, ...], schema: Schema) -> tuple[str, ...
 class TableSizeStats:
     """Sampled chunks and aggregate size/row stats for a table channel."""
 
-    chunks: ChunkStore | dict[int, TableChunk] = field(default_factory=dict)
-    """The sampled chunks/messages, keyed or ordered by sequence number."""
+    chunks: ChunkStore
+    """The sampled chunks/messages in replay order."""
     total_size: int = 0
     """The total estimated size of the table in bytes."""
     total_rows: int = 0

From c6243eef3b8466beb4d4a71cb5e5f934e454393b Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 30 Jun 2026 13:33:39 -0700
Subject: [PATCH 03/11] unify sort logic

---
 .../streaming/actor_graph/collectives/sort.py | 37 +++++++------------
 .../streaming/actor_graph/utils.py            | 16 +-------
 2 files changed, 15 insertions(+), 38 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
index c67f5d154555..7c5465cb7540 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
@@ -36,6 +36,7 @@
     ChannelManager,
     ChunkStore,
     NormalizedPartitioning,
+    _sample_chunks,
     allgather_reduce,
     chunk_to_frame,
     chunkwise_evaluate,
@@ -220,50 +221,38 @@ async def _sample_chunks_for_size_estimate(
     metadata_in: ChannelMetadata,
     executor: StreamingExecutor,
     collective_ids: list[int],
-) -> tuple[dict[int, TableChunk], int]:
+) -> tuple[ChunkStore, int]:
     """
     Sample chunks and estimate total data size to derive num_partitions dynamically.
 
-    The sampled chunks are returned keyed by sequence number. The caller is
+    The sampled chunks are returned in replay order. The caller is
     responsible for replaying them into a channel via replay_buffered_channel.
     """
     if executor.dynamic_planning is None:
-        return {}, num_partitions
+        return ChunkStore(context), num_partitions
 
     size_estimate_id = collective_ids.pop()
     target_partition_size = executor.target_partition_size
     sample_chunk_count = executor.dynamic_planning.sample_chunk_count
 
-    sampled_chunks: dict[int, TableChunk] = {}
-    sampled_bytes = 0
-    for _ in range(sample_chunk_count):
-        msg = await ch_in.recv(context)
-        if msg is None:
-            break
-        chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill(
-            context.br(), allow_overbooking=True
-        )
-        sampled_bytes += chunk.data_alloc_size()
-        sampled_chunks[msg.sequence_number] = chunk
-        if sampled_bytes >= target_partition_size:
-            break
-
-    # Extrapolate local size estimate from samples
-    local_count = metadata_in.local_count
-    local_size = (
-        int(sampled_bytes / len(sampled_chunks) * local_count) if sampled_chunks else 0
+    sample = await _sample_chunks(
+        context,
+        ch_in,
+        sample_chunk_count,
+        target_partition_size,
+        metadata_in.local_count,
     )
 
     # Allgather to get global size estimate across all ranks
     if comm.nranks > 1 and not metadata_in.duplicated:
         (global_size,) = await allgather_reduce(
-            context, comm, size_estimate_id, local_size
+            context, comm, size_estimate_id, sample.total_size
         )
     else:
-        global_size = local_size
+        global_size = sample.total_size
 
     num_partitions = max(1, global_size // target_partition_size)
-    return sampled_chunks, num_partitions
+    return sample.chunks, num_partitions
 
 
 async def _receive_and_buffer_chunks(
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 f07132d743d0..132ec4dc04b7 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -52,7 +52,6 @@
         Coroutine,
         Generator,
         Iterator,
-        Mapping,
         Sequence,
     )
 
@@ -94,17 +93,6 @@ def __iter__(self) -> Generator[Message, None, None]:
             yield self._store.extract(mid=self._mids.popleft())
 
 
-def buffered_chunk_messages(
-    buffered_chunks: ChunkStore | Mapping[int, TableChunk],
-) -> Iterator[Message]:
-    """Yield buffered chunks as messages, consuming the buffer."""
-    if isinstance(buffered_chunks, ChunkStore):
-        yield from buffered_chunks
-    else:
-        for seq_num, chunk in buffered_chunks.items():
-            yield Message(seq_num, chunk)
-
-
 @contextlib.contextmanager
 def set_memory_resource(mr: rmm.mr.DeviceMemoryResource) -> Iterator[None]:
     """
@@ -810,7 +798,7 @@ async def replay_buffered_channel(
     context: Context,
     ch_out: Channel[TableChunk],
     ch_in: Channel[TableChunk],
-    buffered_chunks: ChunkStore | Mapping[int, TableChunk],
+    buffered_chunks: ChunkStore,
     metadata: ChannelMetadata,
     *,
     trace_ir: IR,
@@ -835,7 +823,7 @@ async def replay_buffered_channel(
     """
     async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir):
         await send_metadata(ch_out, context, metadata)
-        for msg in buffered_chunk_messages(buffered_chunks):
+        for msg in buffered_chunks:
             await ch_out.send(context, msg)
         while (msg := await ch_in.recv(context)) is not None:
             await ch_out.send(context, msg)

From 199fe9043c76127079ba6e5825b46337b635c72a Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 30 Jun 2026 13:59:25 -0700
Subject: [PATCH 04/11] use sample_chunk_count as a multiplier

---
 .../streaming/actor_graph/utils.py            | 19 ++++++++++---------
 .../cudf_polars/cudf_polars/utils/config.py   |  6 ++++--
 2 files changed, 14 insertions(+), 11 deletions(-)

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 132ec4dc04b7..f9b6c27ff30d 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -744,8 +744,8 @@ class TableSizeStats:
 async def _sample_chunks(
     context: Context,
     ch: Channel[TableChunk],
-    max_sample_chunks: int,
-    max_sample_bytes: int,
+    sample_chunk_count: int,
+    target_partition_size: int,
     local_count: int,
 ) -> TableSizeStats:
     """
@@ -757,10 +757,12 @@ async def _sample_chunks(
         The context.
     ch
         The channel to sample from.
-    max_sample_chunks
-        The maximum number of chunks to sample.
-    max_sample_bytes
-        The maximum number of bytes to sample.
+    sample_chunk_count
+        Sampling budget multiplier. Sampling continues until
+        ``target_partition_size * sample_chunk_count`` bytes have been sampled
+        or the local input is exhausted.
+    target_partition_size
+        The target partition size used to derive the byte sampling budget.
     local_count
         The expected number of local chunks (used for extrapolation).
 
@@ -769,10 +771,11 @@ async def _sample_chunks(
     Sampled chunks and the extrapolated total size/rows for this rank.
     """
     sampled_chunks = ChunkStore(context)
+    sample_byte_limit = target_partition_size * sample_chunk_count
     sampled_count = 0
     total_size = 0
     total_rows = 0
-    for _ in range(max_sample_chunks):
+    while sampled_count < local_count and total_size < sample_byte_limit:
         msg = await ch.recv(context)
         if msg is None:
             break
@@ -781,8 +784,6 @@ async def _sample_chunks(
         total_rows += chunk.shape[0]
         sampled_count += 1
         sampled_chunks.insert(Message(msg.sequence_number, chunk))
-        if total_size >= max_sample_bytes:
-            break
     if sampled_count:
         total_size = int((total_size / sampled_count) * local_count)
         total_rows = int((total_rows / sampled_count) * local_count)
diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py
index 0a92d9223241..d161775631ec 100644
--- a/python/cudf_polars/cudf_polars/utils/config.py
+++ b/python/cudf_polars/cudf_polars/utils/config.py
@@ -311,8 +311,10 @@ class DynamicPlanningOptions:
     Parameters
     ----------
     sample_chunk_count
-        The maximum number of chunks to sample before deciding whether
-        to shuffle. Default is 2.
+        Multiplier used to derive the dynamic-planning sampling budget from
+        the target partition size. For example, the default value of 2 samples
+        up to roughly two target partitions of data before deciding whether to
+        shuffle. Default is 2.
     join_prefilter_threshold
         Row-count ratio (small / large) below which a join key prefilter is
         applied. Set to 0 to disable join prefiltering. Default is 0.5.

From 04fbec2ec56a1854220c1a0a33eac860cdc83967 Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Tue, 30 Jun 2026 14:14:34 -0700
Subject: [PATCH 05/11] accept coderabbit suggestion

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

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
index 7c5465cb7540..b29f9f6ad61b 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
@@ -251,7 +251,7 @@ async def _sample_chunks_for_size_estimate(
     else:
         global_size = sample.total_size
 
-    num_partitions = max(1, global_size // target_partition_size)
+    num_partitions = max(1, -(-global_size // target_partition_size))
     return sample.chunks, num_partitions
 
 

From a3aaf69439495cffa3a13fb13f442de472519a1f Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Wed, 1 Jul 2026 07:59:42 -0700
Subject: [PATCH 06/11] use Tom's suggestions

---
 .../cudf_polars/streaming/actor_graph/collectives/sort.py   | 3 ++-
 .../cudf_polars/cudf_polars/streaming/actor_graph/utils.py  | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
index b29f9f6ad61b..8777bb4a7cd0 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING
 
 import polars as pl
@@ -251,7 +252,7 @@ async def _sample_chunks_for_size_estimate(
     else:
         global_size = sample.total_size
 
-    num_partitions = max(1, -(-global_size // target_partition_size))
+    num_partitions = max(1, math.ceil(global_size / target_partition_size))
     return sample.chunks, num_partitions
 
 
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 f9b6c27ff30d..3408118ac617 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -734,11 +734,11 @@ class TableSizeStats:
     chunks: ChunkStore
     """The sampled chunks/messages in replay order."""
     total_size: int = 0
-    """The total estimated size of the table in bytes."""
+    """The estimated table size in bytes for the represented scope."""
     total_rows: int = 0
-    """The total estimated number of rows in the table."""
+    """The estimated number of rows for the represented scope."""
     total_chunks: int = 0
-    """The total estimated number of chunks in the table."""
+    """The estimated number of chunks for the represented scope."""
 
 
 async def _sample_chunks(

From 26ea4aefffc33336a10f3b576918bf9ebdfe79f6 Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Wed, 1 Jul 2026 09:05:54 -0700
Subject: [PATCH 07/11] fix existing bug exposed by this PR

---
 python/cudf_polars/cudf_polars/streaming/join.py | 14 +++++++++++++-
 python/cudf_polars/tests/streaming/test_join.py  |  8 +++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py
index ec31d27e79bb..1587515a9681 100644
--- a/python/cudf_polars/cudf_polars/streaming/join.py
+++ b/python/cudf_polars/cudf_polars/streaming/join.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Parallel Join Logic."""
 
@@ -8,6 +8,7 @@
 from functools import reduce
 from typing import TYPE_CHECKING
 
+from cudf_polars.dsl.expr import Col
 from cudf_polars.dsl.ir import ConditionalJoin, Join, Slice
 from cudf_polars.streaming.base import PartitionInfo
 from cudf_polars.streaming.dispatch import lower_ir_node
@@ -208,6 +209,17 @@ def _(
         )
         return rec(Slice(ir.schema, offset, length, new_join))
 
+    # Hash shuffle requires physical column keys. Computed join keys must
+    # fall back until they are materialized before shuffling.
+    for keys in [ir.left_on, ir.right_on]:
+        col_keys: list[Col] = [ne.value for ne in keys if isinstance(ne.value, Col)]
+        if len(col_keys) != len(keys):
+            return _lower_ir_fallback(
+                ir,
+                rec,
+                msg="Multi-partition Join not supported for keys with expressions.",
+            )
+
     # Lower children
     children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
     partition_info = reduce(operator.or_, _partition_info)
diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py
index 8ac799115032..069c0328a68f 100644
--- a/python/cudf_polars/tests/streaming/test_join.py
+++ b/python/cudf_polars/tests/streaming/test_join.py
@@ -593,6 +593,7 @@ def test_join_computed_expr_right_key(streaming_engine_factory) -> None:
             target_partition_size=1,
             max_rows_per_partition=4,
             broadcast_limit=1,  # Disable broadcast joins
+            fallback_mode="warn",
         ),
     )
     if engine.nranks < 2:
@@ -639,4 +640,9 @@ def test_join_computed_expr_right_key(streaming_engine_factory) -> None:
         left_on="zip_prefix",
         right_on=pl.col("full_zip").str.slice(0, 2),
     )
-    assert_gpu_result_equal(q, engine=engine, check_row_order=False)
+    with warns_on_spmd(
+        engine,
+        UserWarning,
+        match=r"Multi-partition Join not supported for keys with expressions\.",
+    ):
+        assert_gpu_result_equal(q, engine=engine, check_row_order=False)

From 51ec0fc94bd8cc840398f2df6ff8ab25af4a78ea Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Wed, 1 Jul 2026 09:17:02 -0700
Subject: [PATCH 08/11] take coderabbit suggestion

---
 .../cudf_polars/cudf_polars/streaming/join.py | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py
index 1587515a9681..47729b3822be 100644
--- a/python/cudf_polars/cudf_polars/streaming/join.py
+++ b/python/cudf_polars/cudf_polars/streaming/join.py
@@ -144,6 +144,15 @@ def _make_bcast_join(
     return new_node, partition_info
 
 
+def _has_expression_keys(ir: Join) -> bool:
+    """Return true if any join key is not a physical column reference."""
+    return any(
+        not isinstance(ne.value, Col)
+        for keys in (ir.left_on, ir.right_on)
+        for ne in keys
+    )
+
+
 @lower_ir_node.register(ConditionalJoin)
 def _(
     ir: ConditionalJoin, rec: LowerIRTransformer
@@ -209,17 +218,6 @@ def _(
         )
         return rec(Slice(ir.schema, offset, length, new_join))
 
-    # Hash shuffle requires physical column keys. Computed join keys must
-    # fall back until they are materialized before shuffling.
-    for keys in [ir.left_on, ir.right_on]:
-        col_keys: list[Col] = [ne.value for ne in keys if isinstance(ne.value, Col)]
-        if len(col_keys) != len(keys):
-            return _lower_ir_fallback(
-                ir,
-                rec,
-                msg="Multi-partition Join not supported for keys with expressions.",
-            )
-
     # Lower children
     children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
     partition_info = reduce(operator.or_, _partition_info)
@@ -230,6 +228,7 @@ def _(
 
     left, right = children
     output_count = max(partition_info[left].count, partition_info[right].count)
+    has_expression_keys = _has_expression_keys(ir)
     if output_count == 1 and not dynamic_planning:
         new_node = ir.reconstruct(children)
         partition_info[new_node] = PartitionInfo(count=1)
@@ -249,6 +248,12 @@ def _(
 
     # Check for dynamic planning - defer broadcast vs shuffle decision to runtime
     if dynamic_planning:  # pragma: no cover; Requires rapidsmpf runtime
+        if has_expression_keys:
+            return _lower_ir_fallback(
+                ir,
+                rec,
+                msg="Multi-partition Join not supported for keys with expressions.",
+            )
         new_node = ir.reconstruct(children)
         partition_info[new_node] = PartitionInfo(count=output_count)
         return new_node, partition_info
@@ -270,6 +275,14 @@ def _(
             left,
             right,
         )
+    elif has_expression_keys:
+        # Hash shuffle requires physical column keys. Computed join keys must
+        # fall back until they are materialized before shuffling.
+        return _lower_ir_fallback(
+            ir,
+            rec,
+            msg="Multi-partition Join not supported for keys with expressions.",
+        )
     else:
         # Create a hash join
         return _make_hash_join(

From d92399cb1144c843000870885e8b8423f206f850 Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Wed, 1 Jul 2026 11:21:15 -0700
Subject: [PATCH 09/11] partially roll-back sample_chunk_count meaning change

---
 .../streaming/actor_graph/utils.py            | 25 +++++++++++--------
 .../cudf_polars/cudf_polars/utils/config.py   |  6 ++---
 2 files changed, 16 insertions(+), 15 deletions(-)

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 3408118ac617..4506c333338f 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -744,8 +744,8 @@ class TableSizeStats:
 async def _sample_chunks(
     context: Context,
     ch: Channel[TableChunk],
-    sample_chunk_count: int,
-    target_partition_size: int,
+    max_sample_chunks: int,
+    max_sample_bytes: int,
     local_count: int,
 ) -> TableSizeStats:
     """
@@ -757,12 +757,10 @@ async def _sample_chunks(
         The context.
     ch
         The channel to sample from.
-    sample_chunk_count
-        Sampling budget multiplier. Sampling continues until
-        ``target_partition_size * sample_chunk_count`` bytes have been sampled
-        or the local input is exhausted.
-    target_partition_size
-        The target partition size used to derive the byte sampling budget.
+    max_sample_chunks
+        The maximum number of non-empty chunks to sample.
+    max_sample_bytes
+        The maximum number of bytes to sample.
     local_count
         The expected number of local chunks (used for extrapolation).
 
@@ -771,19 +769,24 @@ async def _sample_chunks(
     Sampled chunks and the extrapolated total size/rows for this rank.
     """
     sampled_chunks = ChunkStore(context)
-    sample_byte_limit = target_partition_size * sample_chunk_count
+    sampled_nonempty_count = 0
     sampled_count = 0
     total_size = 0
     total_rows = 0
-    while sampled_count < local_count and total_size < sample_byte_limit:
+    while sampled_nonempty_count < max_sample_chunks:
         msg = await ch.recv(context)
         if msg is None:
             break
         chunk = TableChunk.from_message(msg, br=context.br())
+        nrows = chunk.shape[0]
         total_size += chunk.data_alloc_size()
-        total_rows += chunk.shape[0]
+        total_rows += nrows
         sampled_count += 1
+        if nrows > 0:
+            sampled_nonempty_count += 1
         sampled_chunks.insert(Message(msg.sequence_number, chunk))
+        if total_size >= max_sample_bytes:
+            break
     if sampled_count:
         total_size = int((total_size / sampled_count) * local_count)
         total_rows = int((total_rows / sampled_count) * local_count)
diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py
index d161775631ec..d8cd5a64f7ec 100644
--- a/python/cudf_polars/cudf_polars/utils/config.py
+++ b/python/cudf_polars/cudf_polars/utils/config.py
@@ -311,10 +311,8 @@ class DynamicPlanningOptions:
     Parameters
     ----------
     sample_chunk_count
-        Multiplier used to derive the dynamic-planning sampling budget from
-        the target partition size. For example, the default value of 2 samples
-        up to roughly two target partitions of data before deciding whether to
-        shuffle. Default is 2.
+        Maximum number of non-empty chunks to sample before making
+        dynamic-planning decisions. Default is 2.
     join_prefilter_threshold
         Row-count ratio (small / large) below which a join key prefilter is
         applied. Set to 0 to disable join prefiltering. Default is 0.5.

From 98354b069f9fff2ceb013237da10f925fd823d10 Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Wed, 1 Jul 2026 15:58:38 -0700
Subject: [PATCH 10/11] further revert some changes

---
 .../cudf_polars/streaming/actor_graph/utils.py         | 10 +++-------
 python/cudf_polars/cudf_polars/utils/config.py         |  2 +-
 2 files changed, 4 insertions(+), 8 deletions(-)

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 4506c333338f..2459d94a347f 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
@@ -758,7 +758,7 @@ async def _sample_chunks(
     ch
         The channel to sample from.
     max_sample_chunks
-        The maximum number of non-empty chunks to sample.
+        The maximum number of chunks to sample.
     max_sample_bytes
         The maximum number of bytes to sample.
     local_count
@@ -769,21 +769,17 @@ async def _sample_chunks(
     Sampled chunks and the extrapolated total size/rows for this rank.
     """
     sampled_chunks = ChunkStore(context)
-    sampled_nonempty_count = 0
     sampled_count = 0
     total_size = 0
     total_rows = 0
-    while sampled_nonempty_count < max_sample_chunks:
+    for _ in range(max_sample_chunks):
         msg = await ch.recv(context)
         if msg is None:
             break
         chunk = TableChunk.from_message(msg, br=context.br())
-        nrows = chunk.shape[0]
         total_size += chunk.data_alloc_size()
-        total_rows += nrows
+        total_rows += chunk.shape[0]
         sampled_count += 1
-        if nrows > 0:
-            sampled_nonempty_count += 1
         sampled_chunks.insert(Message(msg.sequence_number, chunk))
         if total_size >= max_sample_bytes:
             break
diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py
index d8cd5a64f7ec..35100c5c38f2 100644
--- a/python/cudf_polars/cudf_polars/utils/config.py
+++ b/python/cudf_polars/cudf_polars/utils/config.py
@@ -311,7 +311,7 @@ class DynamicPlanningOptions:
     Parameters
     ----------
     sample_chunk_count
-        Maximum number of non-empty chunks to sample before making
+        The maximum number of chunks to sample before making
         dynamic-planning decisions. Default is 2.
     join_prefilter_threshold
         Row-count ratio (small / large) below which a join key prefilter is

From 41ce201e3b67cfd8717be8367cf45a58c50f6d3d Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Thu, 2 Jul 2026 06:43:41 -0700
Subject: [PATCH 11/11] update inject_gpu_engine.py

---
 .../cudf_polars/testing/inject_gpu_engine.py  | 45 -------------------
 1 file changed, 45 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py
index 6a1b5748bc5c..a9aff2a7ad0e 100644
--- a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py
+++ b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py
@@ -448,51 +448,6 @@ def pytest_report_header(config: pytest.Config) -> str:
     "tests/unit/operations/test_group_by.py::test_unique_head_tail_26429[1]": "https://github.com/rapidsai/cudf/issues/22075",
     "tests/unit/operations/test_group_by.py::test_unique_head_tail_26429[4]": "https://github.com/rapidsai/cudf/issues/22075",
     "tests/unit/operations/test_join.py::test_empty_outer_join_22206": "https://github.com/rapidsai/cudf/issues/22084",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes12]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes13]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes14]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes15]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes17]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes18]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes19]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes20]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes21]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes22]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes23]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes25]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes26]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes27]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes28]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes38]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes39]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes40]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes41]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes42]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes43]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[True-dtypes44]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes12]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes13]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes14]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes15]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes17]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes18]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes19]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes20]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes21]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes22]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes23]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes25]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes26]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes27]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes28]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes38]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes39]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes40]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes41]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes42]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes43]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes44]": "https://github.com/rapidsai/cudf/issues/22085",
-    "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_order": "https://github.com/rapidsai/cudf/issues/22085",
     "tests/unit/operations/test_window.py::test_over_literal_cum_sum_26800": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970",
     "tests/unit/sql/test_joins.py::test_cross_join_unnest_from_cte": "https://github.com/rapidsai/cudf/issues/22073",
     "tests/unit/sql/test_window_functions.py::test_over_with_cumulative_window_funcs": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970",