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 1c60bde3562..10b8d673d06 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" 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 a511444ceee..5303f9a8d99 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) diff --git a/python/cudf_polars/cudf_polars/streaming/distinct.py b/python/cudf_polars/cudf_polars/streaming/distinct.py index 07904fcf357..b3e6ad8ab31 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 ef21cb33d17..887e9e9fecc 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 27b5454030d..fa8ce65a67e 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,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):