From beb6fad6f4d982bc299fc0507831d20620745837 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 24 Jun 2026 12:59:11 -0700 Subject: [PATCH 1/5] fix unique expr lowering to use dynamic planning --- .../cudf_polars/streaming/distinct.py | 67 ++++++++++++------- .../cudf_polars/streaming/expressions.py | 22 +++--- .../tests/streaming/test_unique.py | 22 +++++- 3 files changed, 75 insertions(+), 36 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/distinct.py b/python/cudf_polars/cudf_polars/streaming/distinct.py index 07904fcf357a..b3e6ad8ab31f 100644 --- a/python/cudf_polars/cudf_polars/streaming/distinct.py +++ b/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,6 +89,34 @@ 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 @@ -90,12 +124,7 @@ def _( # 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, diff --git a/python/cudf_polars/cudf_polars/streaming/expressions.py b/python/cudf_polars/cudf_polars/streaming/expressions.py index ef21cb33d17d..887e9e9fecce 100644 --- a/python/cudf_polars/cudf_polars/streaming/expressions.py +++ b/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, diff --git a/python/cudf_polars/tests/streaming/test_unique.py b/python/cudf_polars/tests/streaming/test_unique.py index 27b5454030d6..f2b05fdc210b 100644 --- a/python/cudf_polars/tests/streaming/test_unique.py +++ b/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,25 @@ 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(): + df = pl.LazyFrame({"y": list(range(10)) * 2}) + q = df.select(pl.col("y").unique()) + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={ + "dynamic_planning": {}, + "max_rows_per_partition": 5, + "min_device_size": 1 << 30, + }, + ) + + 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): From 1150fdbeef988f4fecfff4613c4659dee455374d Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 24 Jun 2026 13:23:48 -0700 Subject: [PATCH 2/5] use coderabbit suggestion --- python/cudf_polars/tests/streaming/test_unique.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cudf_polars/tests/streaming/test_unique.py b/python/cudf_polars/tests/streaming/test_unique.py index f2b05fdc210b..7fcdf1e16e39 100644 --- a/python/cudf_polars/tests/streaming/test_unique.py +++ b/python/cudf_polars/tests/streaming/test_unique.py @@ -67,6 +67,8 @@ def test_unique_select_dynamic_planning_uses_dynamic_distinct(): }, ) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + plan = explain_query(q, engine, physical=True) assert "DISTINCT" in plan From 7f44158e130ffc9b8b2278e018ad2c88c905f428 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 25 Jun 2026 09:30:08 -0700 Subject: [PATCH 3/5] update test --- python/cudf_polars/tests/streaming/test_unique.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_unique.py b/python/cudf_polars/tests/streaming/test_unique.py index 7fcdf1e16e39..fa8ce65a67e4 100644 --- a/python/cudf_polars/tests/streaming/test_unique.py +++ b/python/cudf_polars/tests/streaming/test_unique.py @@ -54,17 +54,14 @@ 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(): +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 = pl.GPUEngine( - executor="streaming", - raise_on_fail=True, - executor_options={ - "dynamic_planning": {}, - "max_rows_per_partition": 5, - "min_device_size": 1 << 30, - }, + engine = spmd_engine_factory( + StreamingOptions( + dynamic_planning={}, + max_rows_per_partition=5, + ) ) assert_gpu_result_equal(q, engine=engine, check_row_order=False) From 93dede18995ce05c840aa4ca73d9ef153d0f965b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 26 Jun 2026 10:27:15 -0700 Subject: [PATCH 4/5] minor safty tweak --- python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 1c60bde3562e..10b8d673d06d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/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" From bf563389e7033bb4a50b526fb73f9654f6855eeb Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 26 Jun 2026 12:10:09 -0700 Subject: [PATCH 5/5] fix bloom-filter bug --- .../cudf_polars/streaming/actor_graph/join.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 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 a511444ceeeb..5303f9a8d991 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/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)