Fix Distinct expression lowering

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

@@ -601,7 +601,7 @@ async def _choose_strategy(
     output_count_limit = local_count if skip_global_comm else total_chunk_count
     output_count = min(ideal_count, output_count_limit)
     if not use_tree:
-        output_count = max(output_count, min_row_limit_count)
+        output_count = max(2, output_count, min_row_limit_count)
     if tracer is not None:
         tracer.decision = (
             "tree_local"

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

@@ -626,26 +626,17 @@ def make_filter_tasks(
     tuple
        Of new left and right channels, coroutines to await, and new channels to shutdown on error.
     """
-    bloom_build_output: Channel[BloomFilterChunk] = context.create_channel()
-    bloom_build_input: Channel[TableChunk] = context.create_channel()
-    passthrough_output: Channel[TableChunk] = context.create_channel()
     if left_rows < right_rows:
         passthrough_input = ch_left
-        ch_left = passthrough_output
         build_indices = strategy.left_indices
         bloom_apply_input = ch_right
         apply_indices = strategy.right_indices
-        ch_right = context.create_channel()
-        bloom_apply_output = ch_right
         apply_meta = strategy.right_meta
     else:
         passthrough_input = ch_right
-        ch_right = passthrough_output
         build_indices = strategy.right_indices
         bloom_apply_input = ch_left
         apply_indices = strategy.left_indices
-        ch_left = context.create_channel()
-        bloom_apply_output = ch_left
         apply_meta = strategy.left_meta
     assert apply_meta is not None
     if _is_already_partitioned(
@@ -656,6 +647,18 @@ def make_filter_tasks(
         # but the current implementation only prefilters "locally" in the
         # query DAG.
         return ch_left, ch_right, [], []
+
+    bloom_build_output: Channel[BloomFilterChunk] = context.create_channel()
+    bloom_build_input: Channel[TableChunk] = context.create_channel()
+    passthrough_output: Channel[TableChunk] = context.create_channel()
+    if left_rows < right_rows:
+        ch_left = passthrough_output
+        ch_right = context.create_channel()
+        bloom_apply_output = ch_right
+    else:
+        ch_right = passthrough_output
+        ch_left = context.create_channel()
+        bloom_apply_output = ch_left
     # TODO: configure based on GPU L2 size
     nblocks = BloomFilter.fitting_num_blocks(32 * 1024 * 1024)
     filter = BloomFilter(context, comm, LIBCUDF_DEFAULT_HASH_SEED, nblocks)

python/cudf_polars/cudf_polars/streaming/distinct.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Multi-partition Distinct logic."""
 
@@ -28,14 +28,22 @@
     from cudf_polars.utils.config import ConfigOptions, StreamingExecutor
 
 
-def lower_distinct(
+def _distinct_keys(ir: Distinct) -> tuple[NamedExpr, ...]:
+    subset: frozenset[str] = ir.subset or frozenset(ir.schema)
+    return tuple(
+        NamedExpr(name, Col(ir.schema[name], name))
+        for name in ir.schema
+        if name in subset
+    )
+
+
+def _lower_distinct_static(
     ir: Distinct,
     child: IR,
     partition_info: MutableMapping[IR, PartitionInfo],
-    config_options: ConfigOptions[StreamingExecutor],
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     """
-    Lower a Distinct IR into partition-wise stages.
+    Lower a Distinct IR into partition-wise stages with static planning.
 
     Note: Edge cases (KEEP_NONE + ordering, complex slice, pre-shuffle)
     must be handled by the caller before calling this function.
@@ -50,8 +58,6 @@ def lower_distinct(
     partition_info
         A mapping from all unique IR nodes to the
         associated partitioning information.
-    config_options
-        GPUEngine configuration options.
 
     Returns
     -------
@@ -83,19 +89,42 @@ def lower_distinct(
     return new_node, partition_info
 
 
+def lower_distinct(
+    ir: Distinct,
+    child: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    config_options: ConfigOptions[StreamingExecutor],
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """
+    Lower a Distinct IR with an already-lowered child.
+
+    Note: Edge cases (KEEP_NONE + ordering, complex slice, pre-shuffle)
+    must be handled by the caller before calling this function.
+    """
+    if _dynamic_planning_on(
+        config_options
+    ):  # pragma: no cover; Requires rapidsmpf runtime
+        # Dynamic planning: Reconstruct the Distinct.
+        # The runtime distinct_node will handle strategy selection,
+        # respecting ir.stable, ir.keep, and ir.zlice attributes.
+        dynamic_node = ir.reconstruct([child])
+        partition_info[dynamic_node] = PartitionInfo(
+            count=partition_info[child].count,
+            partitioned_on=_distinct_keys(ir),
+        )
+        return dynamic_node, partition_info
+
+    return _lower_distinct_static(ir, child, partition_info)
+
+
 @lower_ir_node.register(Distinct)
 def _(
     ir: Distinct, rec: LowerIRTransformer
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     # Extract child partitioning
     child, partition_info = rec(ir.children[0])
     child_count = partition_info[child].count
-    subset: frozenset[str] = ir.subset or frozenset(ir.schema)
-    distinct_keys = tuple(
-        NamedExpr(name, Col(ir.schema[name], name))
-        for name in ir.schema
-        if name in subset
-    )
+    distinct_keys = _distinct_keys(ir)
 
     config_options = rec.state["config_options"]
 
@@ -135,20 +164,6 @@ def _(
             msg="Complex slice not supported for multiple partitions.",
         )
 
-    # Branch based on dynamic planning
-    if _dynamic_planning_on(
-        config_options
-    ):  # pragma: no cover; Requires rapidsmpf runtime
-        # Dynamic planning: Reconstruct the Distinct.
-        # The runtime distinct_node will handle strategy selection,
-        # respecting ir.stable, ir.keep, and ir.zlice attributes.
-        dynamic_node = ir.reconstruct([child])
-        partition_info[dynamic_node] = PartitionInfo(
-            count=child_count,
-            partitioned_on=distinct_keys,
-        )
-        return dynamic_node, partition_info
-
     return lower_distinct(
         ir,
         child,

python/cudf_polars/cudf_polars/streaming/expressions.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
 """
 Multi-partition Expr classes and utilities.
@@ -192,15 +192,19 @@ def _decompose_unique(
     )
     (column,) = columns
 
+    distinct_ir = Distinct(
+        {column.name: column.dtype},
+        plc.stream_compaction.DuplicateKeepOption.KEEP_ANY,
+        None,
+        None,
+        maintain_order,
+        input_ir,
+    )
+    # Expr.unique() always lowers to KEEP_ANY with no subset or slice,
+    # so the Distinct fallback cases are not reachable here. We can call
+    # lower_distinct directly.
     input_ir, partition_info = lower_distinct(
-        Distinct(
-            {column.name: column.dtype},
-            plc.stream_compaction.DuplicateKeepOption.KEEP_ANY,
-            None,
-            None,
-            maintain_order,
-            input_ir,
-        ),
+        distinct_ir,
         input_ir,
         partition_info,
         config_options,

python/cudf_polars/tests/streaming/test_unique.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
 
 from __future__ import annotations
@@ -9,6 +9,7 @@
 from polars.testing import assert_frame_equal
 
 from cudf_polars.engine.options import StreamingOptions
+from cudf_polars.streaming.explain import explain_query
 from cudf_polars.testing.asserts import assert_gpu_result_equal
 from cudf_polars.testing.engine_utils import warns_on_spmd
 
@@ -53,6 +54,24 @@ def test_unique_select(df, streaming_engine_factory, maintain_order):
     assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
+def test_unique_select_dynamic_planning_uses_dynamic_distinct(spmd_engine_factory):
+    df = pl.LazyFrame({"y": list(range(10)) * 2})
+    q = df.select(pl.col("y").unique())
+    engine = spmd_engine_factory(
+        StreamingOptions(
+            dynamic_planning={},
+            max_rows_per_partition=5,
+        )
+    )
+
+    assert_gpu_result_equal(q, engine=engine, check_row_order=False)
+
+    plan = explain_query(q, engine, physical=True)
+
+    assert "DISTINCT" in plan
+    assert "REPARTITION" not in plan
+
+
 @pytest.mark.parametrize("keep", ["first", "last", "any"])
 @pytest.mark.parametrize("zlice", ["head", "tail"])
 def test_unique_head_tail(keep, zlice, streaming_engine_factory):