Improve chunk sampling in dynamic actors

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
@@ -36,6 +37,7 @@
     ChannelManager,
     ChunkStore,
     NormalizedPartitioning,
+    _sample_chunks,
     allgather_reduce,
     chunk_to_frame,
     chunkwise_evaluate,
@@ -220,50 +222,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
+    num_partitions = max(1, math.ceil(global_size / target_partition_size))
+    return sample.chunks, num_partitions
 
 
 async def _receive_and_buffer_chunks(

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

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
 
@@ -731,14 +731,14 @@ 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
+    """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(
@@ -768,24 +768,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 +798,7 @@ async def replay_buffered_channel(
     context: Context,
     ch_out: Channel[TableChunk],
     ch_in: Channel[TableChunk],
-    buffered_chunks: dict[int, TableChunk],
+    buffered_chunks: ChunkStore,
     metadata: ChannelMetadata,
     *,
     trace_ir: IR,
@@ -823,8 +823,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_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)

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
@@ -143,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
@@ -218,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)
@@ -237,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
@@ -258,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(

python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py

@@ -451,51 +451,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",

python/cudf_polars/cudf_polars/utils/config.py

@@ -311,8 +311,8 @@ class DynamicPlanningOptions:
     Parameters
     ----------
     sample_chunk_count
-        The maximum number of chunks to sample before deciding whether
-        to shuffle. Default is 2.
+        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
         applied. Set to 0 to disable join prefiltering. Default is 0.5.

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)