From 03ef8ade6368af8c8a9aace1235f6c99b1486b2f Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 25 Jun 2026 10:04:35 +0000 Subject: [PATCH 01/23] Add dynamic join key prefilter planning Add generic dynamic-planning support for join key prefilters in the streaming actor graph. The planner evaluates join type, key compatibility, size estimates, and configured selectivity thresholds to decide when a small side can build a bloom/key prefilter for the larger side before shuffle. The implementation supports prefix key selection for multi-key joins, records structured trace metadata and skip reasons, preserves the original full join after the row-reduction stage, and exposes conservative dynamic-planning options for enabling, sizing, and tracing the prefilter path. --- .../actor_graph/collectives/common.py | 6 +- .../cudf_polars/streaming/actor_graph/join.py | 376 +++++++++++++++--- .../streaming/actor_graph/tracing.py | 8 +- .../streaming/actor_graph/utils.py | 1 + .../cudf_polars/cudf_polars/utils/config.py | 70 +++- .../cudf_polars/tests/streaming/test_join.py | 156 +++++++- .../tests/streaming/test_tracing.py | 6 + python/cudf_polars/tests/test_config.py | 83 ++++ 8 files changed, 648 insertions(+), 58 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py index d09c36aaa323..c2169daadb90 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.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 """Common utilities for collective operations.""" @@ -138,8 +138,8 @@ def __enter__(self) -> dict[IR, list[int]]: _get_new_collective_id_unsafe(), ] elif isinstance(node, Join) and self.dynamic_planning_enabled: - # Join needs 4 IDs: size allgather, left shuffle/bcast, - # right shuffle/bcast, bloom filter + # Join needs 4 IDs: size allgather, one strategy-specific + # allgather/bloom prefilter, left shuffle, and right shuffle. self.collective_id_map[node] = [ _get_new_collective_id_unsafe(), _get_new_collective_id_unsafe(), 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..1d464ab9ffa7 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -59,8 +59,7 @@ from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: - from collections.abc import Iterable, MutableMapping - from types import CoroutineType + from collections.abc import Coroutine, Iterable, MutableMapping from cudf_streaming.bloom_filter import BloomFilterChunk from rapidsmpf.communicator.communicator import Communicator @@ -94,6 +93,45 @@ class JoinStrategy: """The shuffle indices for the right side. Only used for shuffle joins.""" +@dataclass(frozen=True) +class JoinPrefilterDecision: + """Decision for an optional join-key prefilter stage.""" + + considered: bool + left_rows: int + right_rows: int + threshold: float + filter_side: Literal["left", "right"] | None = None + build_side: Literal["left", "right"] | None = None + build_indices: tuple[int, ...] = () + apply_indices: tuple[int, ...] = () + key_column_count: int = 0 + small_large_ratio: float | None = None + reason_skipped: str | None = None + + @property + def enabled(self) -> bool: + """Whether this decision applies a prefilter.""" + return self.reason_skipped is None and self.filter_side is not None + + def trace_dict(self) -> dict[str, Any]: + """Return structured trace metadata for this decision.""" + metadata: dict[str, Any] = { + "considered": self.considered, + "estimated_left_rows": self.left_rows, + "estimated_right_rows": self.right_rows, + "threshold": self.threshold, + "filtered_side": self.filter_side, + "build_side": self.build_side, + "key_column_count": self.key_column_count, + } + if self.small_large_ratio is not None: + metadata["small_large_ratio"] = self.small_large_ratio + if self.reason_skipped is not None: + metadata["reason_skipped"] = self.reason_skipped + return metadata + + @define_actor() async def broadcast_join_actor( context: Context, @@ -583,20 +621,193 @@ def use_bloom_filter( return large_rows > 0 and small_rows / large_rows < threshold +def _skipped_prefilter( + reason: str, + *, + left_rows: int, + right_rows: int, + threshold: float, + ratio: float | None = None, + key_column_count: int = 0, +) -> JoinPrefilterDecision: + return JoinPrefilterDecision( + considered=True, + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + key_column_count=key_column_count, + small_large_ratio=ratio, + reason_skipped=reason, + ) + + +def _select_join_prefilter( + join_type: Literal["Inner", "Left", "Right", "Full", "Semi", "Anti", "Cross"], + left_rows: int, + right_rows: int, + left_key_indices: tuple[int, ...], + right_key_indices: tuple[int, ...], + *, + threshold: float, + max_key_columns: int | None, +) -> JoinPrefilterDecision: + """ + Select a safe join-key prefilter. + + The prefilter only removes rows that cannot participate in the original + join. The full join still runs afterward with the complete key set. + """ + key_column_count = len(left_key_indices) + if threshold == 0.0: + return _skipped_prefilter( + "disabled", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + ) + if key_column_count == 0: + return _skipped_prefilter( + "no_join_keys", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + ) + if key_column_count != len(right_key_indices): + return _skipped_prefilter( + "mismatched_join_keys", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + ) + + if max_key_columns is not None: + key_column_count = min(key_column_count, max_key_columns) + + build_side: Literal["left", "right"] + filter_side: Literal["left", "right"] + small_rows: int + large_rows: int + + if join_type in ("Inner", "Semi"): + if left_rows <= right_rows: + build_side = "left" + filter_side = "right" + small_rows, large_rows = left_rows, right_rows + else: + build_side = "right" + filter_side = "left" + small_rows, large_rows = right_rows, left_rows + elif join_type in ("Left", "Anti"): + if left_rows >= right_rows: + ratio = right_rows / left_rows if left_rows > 0 else None + return _skipped_prefilter( + "no_legal_large_side", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + ratio=ratio, + key_column_count=key_column_count, + ) + build_side = "left" + filter_side = "right" + small_rows, large_rows = left_rows, right_rows + elif join_type == "Right": + if right_rows >= left_rows: + ratio = left_rows / right_rows if right_rows > 0 else None + return _skipped_prefilter( + "no_legal_large_side", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + ratio=ratio, + key_column_count=key_column_count, + ) + build_side = "right" + filter_side = "left" + small_rows, large_rows = right_rows, left_rows + else: + return _skipped_prefilter( + "unsupported_join_type", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + key_column_count=key_column_count, + ) + + if large_rows <= 0: + return _skipped_prefilter( + "no_large_side", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + key_column_count=key_column_count, + ) + + ratio = small_rows / large_rows + if ratio >= threshold: + return _skipped_prefilter( + "ratio_above_threshold", + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + ratio=ratio, + key_column_count=key_column_count, + ) + + if build_side == "left": + build_indices = left_key_indices[:key_column_count] + apply_indices = right_key_indices[:key_column_count] + else: + build_indices = right_key_indices[:key_column_count] + apply_indices = left_key_indices[:key_column_count] + + return JoinPrefilterDecision( + considered=True, + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + filter_side=filter_side, + build_side=build_side, + build_indices=build_indices, + apply_indices=apply_indices, + key_column_count=key_column_count, + small_large_ratio=ratio, + ) + + +async def trace_row_count_passthrough( + context: Context, + ch_in: Channel[TableChunk], + ch_out: Channel[TableChunk], + trace_stats: dict[str, Any], + *, + row_count_key: str, +) -> None: + """Forward a table-chunk channel while counting rows.""" + metadata = await recv_metadata(ch_in, context) + await send_metadata(ch_out, context, metadata) + row_count = 0 + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message(msg, br=context.br()) + row_count += chunk.shape[0] + await ch_out.send(context, Message(msg.sequence_number, chunk)) + trace_stats[row_count_key] = row_count + await ch_out.drain(context) + + def make_filter_tasks( context: Context, comm: Communicator, *, ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], - strategy: JoinStrategy, - left_rows: int, - right_rows: int, + decision: JoinPrefilterDecision, tag: int, + trace_stats: dict[str, Any] | None, ) -> tuple[ Channel[TableChunk], Channel[TableChunk], - list[CoroutineType[Any, Any, None]], + list[Coroutine[Any, Any, None]], list[Channel], ]: """ @@ -612,54 +823,79 @@ def make_filter_tasks( Left input channel ch_right Right input channel - strategy - Selected join strategy - left_rows - Estimate of number of rows in left table - right_rows - Estimate of number of rows in right table + decision + Selected prefilter decision tag Collective ID for combining partial filters across ranks + trace_stats + Mutable trace metadata to update with actual row counts, or None Returns ------- tuple Of new left and right channels, coroutines to await, and new channels to shutdown on error. """ + assert decision.enabled + assert decision.build_side in ("left", "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: + if decision.build_side == "left": passthrough_input = ch_left ch_left = passthrough_output - build_indices = strategy.left_indices + build_indices = decision.build_indices bloom_apply_input = ch_right - apply_indices = strategy.right_indices + apply_indices = decision.apply_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 + build_indices = decision.build_indices bloom_apply_input = ch_left - apply_indices = strategy.left_indices + apply_indices = decision.apply_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( - apply_meta, apply_indices, strategy.shuffle_modulus, comm.nranks - ): - # "large" side is already shuffled so no need to pre-filter - # TODO: Really we should pushdown the filter as far as possible, - # but the current implementation only prefilters "locally" in the - # query DAG. - return ch_left, ch_right, [], [] + # TODO: configure based on GPU L2 size nblocks = BloomFilter.fitting_num_blocks(32 * 1024 * 1024) filter = BloomFilter(context, comm, LIBCUDF_DEFAULT_HASH_SEED, nblocks) + filter_tasks: list[Coroutine[Any, Any, None]] = [] + chs_to_shutdown = [ + bloom_build_output, + bloom_build_input, + passthrough_output, + ] + + apply_input = bloom_apply_input + apply_output = bloom_apply_output + if trace_stats is not None: + counted_apply_input: Channel[TableChunk] = context.create_channel() + raw_apply_output: Channel[TableChunk] = context.create_channel() + filter_tasks.extend( + [ + trace_row_count_passthrough( + context, + bloom_apply_input, + counted_apply_input, + trace_stats, + row_count_key="input_rows", + ), + trace_row_count_passthrough( + context, + raw_apply_output, + bloom_apply_output, + trace_stats, + row_count_key="output_rows", + ), + ] + ) + chs_to_shutdown.extend([counted_apply_input, raw_apply_output]) + apply_input = counted_apply_input + apply_output = raw_apply_output + filter_tasks = [ + *filter_tasks, passthrough_split( context, passthrough_input, @@ -676,16 +912,11 @@ def make_filter_tasks( filter.apply( context, bloom_build_output, - bloom_apply_input, - bloom_apply_output, + apply_input, + apply_output, apply_indices, ), ] - chs_to_shutdown = [ - bloom_build_output, - bloom_build_input, - passthrough_output, - ] return ch_left, ch_right, filter_tasks, chs_to_shutdown @@ -702,7 +933,9 @@ async def _shuffle_join( *, row_counts: tuple[int, int], tracer: ActorTracer | None, - bloom_threshold: float, + prefilter_threshold: float, + prefilter_max_key_columns: int | None, + prefilter_trace: bool, ) -> None: """Execute a shuffle (hash) join.""" # Send output metadata @@ -722,18 +955,45 @@ async def _shuffle_join( await send_metadata(ch_out, context, metadata_out) left_rows, right_rows = row_counts bloom_tag = collective_ids.pop(0) - if use_bloom_filter(ir.options[0], left_rows, right_rows, bloom_threshold): + left_key_indices, right_key_indices, _ = _get_key_indices(ir, None) + prefilter_decision = _select_join_prefilter( + ir.options[0], + left_rows, + right_rows, + left_key_indices, + right_key_indices, + threshold=prefilter_threshold, + max_key_columns=prefilter_max_key_columns, + ) + prefilter_trace_stats = prefilter_decision.trace_dict() + if prefilter_decision.enabled: + apply_meta = ( + strategy.right_meta + if prefilter_decision.filter_side == "right" + else strategy.left_meta + ) + assert apply_meta is not None + prefilter_trace_stats["apply_side_prepartitioned"] = _is_already_partitioned( + apply_meta, + prefilter_decision.apply_indices, + strategy.shuffle_modulus, + comm.nranks, + ) + + if tracer is not None: + tracer.set_extra("join_prefilter", prefilter_trace_stats) + + if prefilter_decision.enabled: if tracer is not None: - tracer.decision = f"{tracer.decision or 'shuffle'}_filtered" + tracer.decision = f"{tracer.decision or 'shuffle'}_prefiltered" ch_left, ch_right, filter_tasks, chs_to_shutdown = make_filter_tasks( context, comm, ch_left=ch_left, ch_right=ch_right, - strategy=strategy, - left_rows=left_rows, - right_rows=right_rows, + decision=prefilter_decision, tag=bloom_tag, + trace_stats=prefilter_trace_stats if prefilter_trace else None, ) else: filter_tasks = [] @@ -1208,6 +1468,23 @@ async def join_actor( ) ) else: + dynamic_options = executor.dynamic_planning + prefilter_threshold = ( + dynamic_options.join_prefilter_threshold + if dynamic_options is not None + and dynamic_options.join_prefilter_threshold is not None + else 0.0 + ) + prefilter_max_key_columns = ( + dynamic_options.join_prefilter_max_key_columns + if dynamic_options is not None + else 1 + ) + prefilter_trace = ( + dynamic_options.join_prefilter_trace + if dynamic_options is not None + else False + ) actor_tasks.append( _shuffle_join( context, @@ -1224,11 +1501,9 @@ async def join_actor( right_sample.total_rows, ), tracer=tracer, - bloom_threshold=( - executor.dynamic_planning.bloom_filter_threshold - if executor.dynamic_planning is not None - else 0.0 - ), + prefilter_threshold=prefilter_threshold, + prefilter_max_key_columns=prefilter_max_key_columns, + prefilter_trace=prefilter_trace, ) ) await gather_in_task_group(*actor_tasks) @@ -1308,11 +1583,14 @@ def _( ): # Dynamic join - decide strategy at runtime collective_ids = list(rec.state["collective_id_map"].get(ir, [])) - # Join uses up to 3 collective IDs: 1 allgather + up to 2 (left/right shuffle) + # Join uses up to 4 collective IDs: size allgather, one + # strategy-specific allgather/bloom prefilter, left shuffle, and + # right shuffle. if len(collective_ids) < 4: raise ValueError( - "Dynamic join requires 3 reserved collective IDs " - "(allgather + left shuffle + right shuffle + bloom filter); got " + "Dynamic join requires 4 reserved collective IDs " + "(size allgather + strategy allgather/bloom prefilter " + "+ left shuffle + right shuffle); got " f"{len(collective_ids)} for this Join. " "Ensure ReserveOpIDs is run with dynamic_planning enabled." ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py index 5a39f4c18a60..ee5c75524d23 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py @@ -5,7 +5,7 @@ from __future__ import annotations import dataclasses -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from rapidsmpf.streaming.core.message import Message @@ -50,6 +50,7 @@ class ActorTracer: "chunk_count", "decision", "duplicated", + "extra", "ir_id", "ir_type", "row_count", @@ -62,6 +63,7 @@ def __init__(self, ir_id: int | None = None, ir_type: str | None = None) -> None self.chunk_count: int = 0 self.decision: str | None = None self.duplicated: bool = False + self.extra: dict[str, Any] = {} def add_chunk(self, *, chunk: TableChunk | None = None) -> None: """ @@ -83,6 +85,10 @@ def set_duplicated(self, *, duplicated: bool = True) -> None: """Mark output rows as duplicated across ranks.""" self.duplicated = duplicated + def set_extra(self, key: str, value: Any) -> None: + """Attach structured metadata to the actor trace event.""" + self.extra[key] = value + async def send_chunk( context: Context, 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 8e666310cdfa..02ae6b546a82 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -206,6 +206,7 @@ async def shutdown_on_error( record["row_count"] = tracer.row_count if tracer.decision is not None: record["decision"] = tracer.decision + record.update(tracer.extra) cudf_polars.dsl.tracing.log( "Streaming Actor", start=start, stop=stop, **record ) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 5786a5351cc4..a1d98a124dd1 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.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 """ @@ -144,6 +144,7 @@ class Cluster(enum.StrEnum): T = TypeVar("T") +_DEFAULT_JOIN_PREFILTER_MAX_KEY_COLUMNS: int | None = 1 def _make_default_factory( @@ -168,6 +169,18 @@ def _bool_converter(v: str) -> bool: raise ValueError(f"Invalid boolean value: '{v}'") +def _optional_float_converter(v: str) -> float | None: + if v.lower() in {"none", "null"}: + return None + return float(v) + + +def _optional_int_converter(v: str) -> int | None: + if v.lower() in {"none", "null"}: + return None + return int(v) + + @dataclasses.dataclass(frozen=True) class ParquetOptions: """ @@ -306,8 +319,19 @@ class DynamicPlanningOptions: to shuffle. Default is 2. bloom_filter_threshold Row-count ratio (small / large) below which a bloom filter is applied - to pre-filter the large side of an inner or semi shuffle join. - Set to 0 to disable bloom filtering. Default is 0.5. + to pre-filter a join side. This is retained as the legacy default for + ``join_prefilter_threshold``. Set to 0 to disable join prefiltering + when ``join_prefilter_threshold`` is unset. Default is 0.5. + join_prefilter_threshold + Row-count ratio (small / large) below which a join key prefilter is + applied. When unset, ``bloom_filter_threshold`` is used. Default is + unset. + join_prefilter_max_key_columns + Maximum number of join-key columns to use for the prefilter. Set to + ``None`` to use all join keys. Default is 1. + join_prefilter_trace + Whether to collect input/output row counts around applied join + prefilters. Default is False. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -322,6 +346,27 @@ class DynamicPlanningOptions: f"{_env_prefix}__BLOOM_FILTER_THRESHOLD", float, default=0.5 ) ) + join_prefilter_threshold: float | None = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_PREFILTER_THRESHOLD", + _optional_float_converter, + default=None, + ) + ) + join_prefilter_max_key_columns: int | None = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", + _optional_int_converter, + default=_DEFAULT_JOIN_PREFILTER_MAX_KEY_COLUMNS, + ) + ) + join_prefilter_trace: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_PREFILTER_TRACE", + _bool_converter, + default=False, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.sample_chunk_count, int): @@ -332,6 +377,25 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("bloom_filter_threshold must be a float") if not 0.0 <= self.bloom_filter_threshold <= 1.0: raise ValueError("bloom_filter_threshold must be between 0 and 1") + join_prefilter_threshold = self.join_prefilter_threshold + if join_prefilter_threshold is None: + join_prefilter_threshold = self.bloom_filter_threshold + object.__setattr__( + self, "join_prefilter_threshold", join_prefilter_threshold + ) + elif not isinstance(join_prefilter_threshold, float): + raise TypeError("join_prefilter_threshold must be a float or None") + if not 0.0 <= join_prefilter_threshold <= 1.0: + raise ValueError("join_prefilter_threshold must be between 0 and 1") + if self.join_prefilter_max_key_columns is not None: + if not isinstance(self.join_prefilter_max_key_columns, int): + raise TypeError("join_prefilter_max_key_columns must be an int or None") + if self.join_prefilter_max_key_columns < 1: + raise ValueError( + "join_prefilter_max_key_columns must be at least 1 or None" + ) + if not isinstance(self.join_prefilter_trace, bool): + raise TypeError("join_prefilter_trace must be a bool") @dataclasses.dataclass(frozen=True, eq=True) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index de6843aa2cc0..d08401c8a995 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_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 """Tests for dynamic join path in join_actor (including Right and Full joins).""" @@ -15,7 +15,10 @@ from cudf_polars.dsl.ir import Cache, Join from cudf_polars.dsl.traversal import traversal from cudf_polars.engine.options import StreamingOptions -from cudf_polars.streaming.actor_graph.join import _use_pwise_join +from cudf_polars.streaming.actor_graph.join import ( + _select_join_prefilter, + _use_pwise_join, +) from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.shuffle import Shuffle @@ -248,6 +251,155 @@ def test_bloom_filter_join(how, streaming_engine_factory): assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) +def test_multi_key_join_prefilter_preserves_full_join( + streaming_engine_factory, +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + max_rows_per_partition=2, + broadcast_limit=1, + target_partition_size=10, + dynamic_planning={ + "join_prefilter_threshold": 0.5, + "join_prefilter_max_key_columns": 1, + }, + ), + ) + fact = pl.LazyFrame( + { + "k1": range(200), + "k2": [i % 3 for i in range(200)], + "v": range(200), + } + ) + dim = pl.LazyFrame( + { + "k1": range(10), + "k2": [(i + 1) % 3 for i in range(10)], + "d": range(10), + } + ) + q = fact.join(dim, on=["k1", "k2"], how="inner") + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) + + +def test_join_prefilter_skips_when_sides_are_similar_size() -> None: + decision = _select_join_prefilter( + "Inner", + 100, + 120, + (0,), + (0,), + threshold=0.5, + max_key_columns=1, + ) + assert not decision.enabled + assert decision.reason_skipped == "ratio_above_threshold" + + +def test_join_prefilter_filters_large_side_with_key_prefix() -> None: + decision = _select_join_prefilter( + "Inner", + 10, + 1_000, + (0, 1), + (3, 4), + threshold=0.5, + max_key_columns=1, + ) + assert decision.enabled + assert decision.build_side == "left" + assert decision.filter_side == "right" + assert decision.build_indices == (0,) + assert decision.apply_indices == (3,) + assert decision.key_column_count == 1 + + +def test_join_prefilter_can_use_all_join_keys() -> None: + decision = _select_join_prefilter( + "Inner", + 10, + 1_000, + (0, 1), + (3, 4), + threshold=0.5, + max_key_columns=None, + ) + assert decision.enabled + assert decision.build_indices == (0, 1) + assert decision.apply_indices == (3, 4) + assert decision.key_column_count == 2 + + +@pytest.mark.parametrize("how", ["Left", "Anti"]) +def test_join_prefilter_outer_semantics_only_filter_right_side(how) -> None: + decision = _select_join_prefilter( + how, + 1_000, + 10, + (0,), + (0,), + threshold=0.5, + max_key_columns=1, + ) + assert not decision.enabled + assert decision.reason_skipped == "no_legal_large_side" + + decision = _select_join_prefilter( + how, + 10, + 1_000, + (0,), + (0,), + threshold=0.5, + max_key_columns=1, + ) + assert decision.enabled + assert decision.build_side == "left" + assert decision.filter_side == "right" + + +def test_join_prefilter_right_join_only_filters_left_side() -> None: + decision = _select_join_prefilter( + "Right", + 10, + 1_000, + (0,), + (0,), + threshold=0.5, + max_key_columns=1, + ) + assert not decision.enabled + assert decision.reason_skipped == "no_legal_large_side" + + decision = _select_join_prefilter( + "Right", + 1_000, + 10, + (0,), + (0,), + threshold=0.5, + max_key_columns=1, + ) + assert decision.enabled + assert decision.build_side == "right" + assert decision.filter_side == "left" + + +def test_join_prefilter_skips_unsupported_full_join() -> None: + decision = _select_join_prefilter( + "Full", + 10, + 1_000, + (0,), + (0,), + threshold=0.5, + max_key_columns=1, + ) + assert not decision.enabled + assert decision.reason_skipped == "unsupported_join_type" + + @pytest.mark.parametrize( "maintain_order", ["left_right", "right_left", "left", "right"] ) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index f584ceff6528..df83512e0c8a 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -42,6 +42,12 @@ def test_actor_tracer_counts_table_chunk_without_table_view(chunk: TableChunk) - assert tracer.row_count == 3 +def test_actor_tracer_records_extra_metadata() -> None: + tracer = ActorTracer() + tracer.set_extra("join_prefilter", {"considered": True}) + assert tracer.extra == {"join_prefilter": {"considered": True}} + + @pytest.mark.spmd def test_send_chunk_traces_and_sends_message( spmd_engine: SPMDEngine, chunk: TableChunk diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index f8c2a687bab7..8557bd919230 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -673,6 +673,9 @@ def test_dynamic_planning_defaults() -> None: assert config.executor.dynamic_planning is not None assert config.executor.dynamic_planning.sample_chunk_count == 2 assert config.executor.dynamic_planning.bloom_filter_threshold == 0.5 + assert config.executor.dynamic_planning.join_prefilter_threshold == 0.5 + assert config.executor.dynamic_planning.join_prefilter_max_key_columns == 1 + assert not config.executor.dynamic_planning.join_prefilter_trace def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -725,6 +728,86 @@ def test_bloom_filter_threshold_from_env(monkeypatch: pytest.MonkeyPatch) -> Non config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.executor.dynamic_planning is not None assert config.executor.dynamic_planning.bloom_filter_threshold == 0.3 + assert config.executor.dynamic_planning.join_prefilter_threshold == 0.3 + + +def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_THRESHOLD", "0.25" + ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_MAX_KEY_COLUMNS", + "none", + ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_TRACE", "1" + ) + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.join_prefilter_threshold == 0.25 + assert config.executor.dynamic_planning.join_prefilter_max_key_columns is None + assert config.executor.dynamic_planning.join_prefilter_trace + + +def test_join_prefilter_threshold_overrides_bloom_threshold() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": { + "bloom_filter_threshold": 0.2, + "join_prefilter_threshold": 0.4, + } + }, + ) + ) + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.bloom_filter_threshold == 0.2 + assert config.executor.dynamic_planning.join_prefilter_threshold == 0.4 + + +def test_validate_join_prefilter_threshold() -> None: + with pytest.raises(TypeError, match="join_prefilter_threshold must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_threshold": "bad"} + }, + ) + ) + with pytest.raises(ValueError, match="join_prefilter_threshold must be between"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_threshold": 1.5} + }, + ) + ) + + +def test_validate_join_prefilter_max_key_columns() -> None: + with pytest.raises(TypeError, match="join_prefilter_max_key_columns must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_max_key_columns": "bad"} + }, + ) + ) + with pytest.raises( + ValueError, match="join_prefilter_max_key_columns must be at least 1" + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_max_key_columns": 0} + }, + ) + ) def test_dynamic_planning_from_instance() -> None: From 09f04743d419647bcb734ec404d60e961a799574 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:14:58 +0000 Subject: [PATCH 02/23] Reject boolean join prefilter key limits Reject boolean values for join_prefilter_max_key_columns instead of accepting them through Python's bool-is-int relationship. Add explicit config validation coverage so only None and positive integer limits remain valid. --- python/cudf_polars/cudf_polars/utils/config.py | 4 +++- python/cudf_polars/tests/test_config.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index a1d98a124dd1..70be3dee9b8c 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -388,7 +388,9 @@ def __post_init__(self) -> None: # noqa: D105 if not 0.0 <= join_prefilter_threshold <= 1.0: raise ValueError("join_prefilter_threshold must be between 0 and 1") if self.join_prefilter_max_key_columns is not None: - if not isinstance(self.join_prefilter_max_key_columns, int): + if isinstance( + self.join_prefilter_max_key_columns, bool + ) or not isinstance(self.join_prefilter_max_key_columns, int): raise TypeError("join_prefilter_max_key_columns must be an int or None") if self.join_prefilter_max_key_columns < 1: raise ValueError( diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 8557bd919230..a9d19ecbdf36 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -797,6 +797,15 @@ def test_validate_join_prefilter_max_key_columns() -> None: }, ) ) + with pytest.raises(TypeError, match="join_prefilter_max_key_columns must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_max_key_columns": True} + }, + ) + ) with pytest.raises( ValueError, match="join_prefilter_max_key_columns must be at least 1" ): From 6f95d338f266e741a2372094e4f5de18203ff63e Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:15:23 +0000 Subject: [PATCH 03/23] Consolidate optional config converters Use a single typed helper for optional environment-value parsing and define the optional float and int converters in terms of it. This keeps the None/null handling in one place without changing accepted values. --- python/cudf_polars/cudf_polars/utils/config.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 70be3dee9b8c..aba314cd53e4 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -169,16 +169,18 @@ def _bool_converter(v: str) -> bool: raise ValueError(f"Invalid boolean value: '{v}'") -def _optional_float_converter(v: str) -> float | None: +def _optional_converter(v: str, parse: Callable[[str], T]) -> T | None: if v.lower() in {"none", "null"}: return None - return float(v) + return parse(v) + + +def _optional_float_converter(v: str) -> float | None: + return _optional_converter(v, float) def _optional_int_converter(v: str) -> int | None: - if v.lower() in {"none", "null"}: - return None - return int(v) + return _optional_converter(v, int) @dataclasses.dataclass(frozen=True) From 519cd1da10967a14751c4b2a13d636f4d74df9e6 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:16:34 +0000 Subject: [PATCH 04/23] Remove legacy bloom prefilter threshold path Make join_prefilter_threshold the single dynamic-planning threshold for join prefilters, preserving the existing 0.5 default. Remove the obsolete bloom_filter_threshold option, its fallback/override tests, and the unused use_bloom_filter helper now replaced by _select_join_prefilter. --- .../cudf_polars/streaming/actor_graph/join.py | 19 ------- .../cudf_polars/cudf_polars/utils/config.py | 36 +++---------- python/cudf_polars/tests/test_config.py | 50 ------------------- 3 files changed, 6 insertions(+), 99 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 1d464ab9ffa7..901969be0d50 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -603,24 +603,6 @@ async def passthrough_split( await ch_out.drain(context) -def use_bloom_filter( - join_type: Literal["Inner", "Left", "Right", "Full", "Semi", "Anti", "Cross"], - left_rows: int, - right_rows: int, - threshold: float, -) -> bool: - """Return True if bloom filter pre-filtering should be applied.""" - if ( - threshold == 0.0 - or join_type not in ("Inner", "Semi", "Left", "Right") - or (join_type == "Left" and right_rows <= left_rows) - or (join_type == "Right" and left_rows <= right_rows) - ): - return False - small_rows, large_rows = sorted([left_rows, right_rows]) - return large_rows > 0 and small_rows / large_rows < threshold - - def _skipped_prefilter( reason: str, *, @@ -1472,7 +1454,6 @@ async def join_actor( prefilter_threshold = ( dynamic_options.join_prefilter_threshold if dynamic_options is not None - and dynamic_options.join_prefilter_threshold is not None else 0.0 ) prefilter_max_key_columns = ( diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index aba314cd53e4..cf4b310861a8 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -175,10 +175,6 @@ def _optional_converter(v: str, parse: Callable[[str], T]) -> T | None: return parse(v) -def _optional_float_converter(v: str) -> float | None: - return _optional_converter(v, float) - - def _optional_int_converter(v: str) -> int | None: return _optional_converter(v, int) @@ -319,15 +315,9 @@ class DynamicPlanningOptions: sample_chunk_count The maximum number of chunks to sample before deciding whether to shuffle. Default is 2. - bloom_filter_threshold - Row-count ratio (small / large) below which a bloom filter is applied - to pre-filter a join side. This is retained as the legacy default for - ``join_prefilter_threshold``. Set to 0 to disable join prefiltering - when ``join_prefilter_threshold`` is unset. Default is 0.5. join_prefilter_threshold Row-count ratio (small / large) below which a join key prefilter is - applied. When unset, ``bloom_filter_threshold`` is used. Default is - unset. + applied. Set to 0 to disable join prefiltering. Default is 0.5. join_prefilter_max_key_columns Maximum number of join-key columns to use for the prefilter. Set to ``None`` to use all join keys. Default is 1. @@ -343,16 +333,11 @@ class DynamicPlanningOptions: f"{_env_prefix}__SAMPLE_CHUNK_COUNT", int, default=2 ) ) - bloom_filter_threshold: float = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__BLOOM_FILTER_THRESHOLD", float, default=0.5 - ) - ) - join_prefilter_threshold: float | None = dataclasses.field( + join_prefilter_threshold: float = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__JOIN_PREFILTER_THRESHOLD", - _optional_float_converter, - default=None, + float, + default=0.5, ) ) join_prefilter_max_key_columns: int | None = dataclasses.field( @@ -375,18 +360,9 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("sample_chunk_count must be an int") if self.sample_chunk_count < 1: raise ValueError("sample_chunk_count must be at least 1") - if not isinstance(self.bloom_filter_threshold, float): - raise TypeError("bloom_filter_threshold must be a float") - if not 0.0 <= self.bloom_filter_threshold <= 1.0: - raise ValueError("bloom_filter_threshold must be between 0 and 1") join_prefilter_threshold = self.join_prefilter_threshold - if join_prefilter_threshold is None: - join_prefilter_threshold = self.bloom_filter_threshold - object.__setattr__( - self, "join_prefilter_threshold", join_prefilter_threshold - ) - elif not isinstance(join_prefilter_threshold, float): - raise TypeError("join_prefilter_threshold must be a float or None") + if not isinstance(join_prefilter_threshold, float): + raise TypeError("join_prefilter_threshold must be a float") if not 0.0 <= join_prefilter_threshold <= 1.0: raise ValueError("join_prefilter_threshold must be between 0 and 1") if self.join_prefilter_max_key_columns is not None: diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index a9d19ecbdf36..ad1ffd6b8052 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -672,7 +672,6 @@ def test_dynamic_planning_defaults() -> None: # Dynamic planning is enabled by default assert config.executor.dynamic_planning is not None assert config.executor.dynamic_planning.sample_chunk_count == 2 - assert config.executor.dynamic_planning.bloom_filter_threshold == 0.5 assert config.executor.dynamic_planning.join_prefilter_threshold == 0.5 assert config.executor.dynamic_planning.join_prefilter_max_key_columns == 1 assert not config.executor.dynamic_planning.join_prefilter_trace @@ -699,38 +698,6 @@ def test_dynamic_planning_sample_chunk_count_from_env( assert config.executor.dynamic_planning.sample_chunk_count == 3 -def test_validate_bloom_filter_threshold_type() -> None: - with pytest.raises(TypeError, match="bloom_filter_threshold must be a float"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={ - "dynamic_planning": {"bloom_filter_threshold": "bad"} - }, - ) - ) - - -def test_validate_bloom_filter_threshold_range() -> None: - with pytest.raises(ValueError, match="bloom_filter_threshold must be between"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={"dynamic_planning": {"bloom_filter_threshold": 1.5}}, - ) - ) - - -def test_bloom_filter_threshold_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__BLOOM_FILTER_THRESHOLD", "0.3" - ) - config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.bloom_filter_threshold == 0.3 - assert config.executor.dynamic_planning.join_prefilter_threshold == 0.3 - - def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_THRESHOLD", "0.25" @@ -749,23 +716,6 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non assert config.executor.dynamic_planning.join_prefilter_trace -def test_join_prefilter_threshold_overrides_bloom_threshold() -> None: - config = ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={ - "dynamic_planning": { - "bloom_filter_threshold": 0.2, - "join_prefilter_threshold": 0.4, - } - }, - ) - ) - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.bloom_filter_threshold == 0.2 - assert config.executor.dynamic_planning.join_prefilter_threshold == 0.4 - - def test_validate_join_prefilter_threshold() -> None: with pytest.raises(TypeError, match="join_prefilter_threshold must be"): ConfigOptions.from_polars_engine( From 09a744ec02e83487a6515e9676a5d9e291b36d97 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 05:34:02 -0700 Subject: [PATCH 05/23] Fix linting --- python/cudf_polars/cudf_polars/utils/config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index cf4b310861a8..e8c4f3eb7846 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -366,9 +366,9 @@ def __post_init__(self) -> None: # noqa: D105 if not 0.0 <= join_prefilter_threshold <= 1.0: raise ValueError("join_prefilter_threshold must be between 0 and 1") if self.join_prefilter_max_key_columns is not None: - if isinstance( - self.join_prefilter_max_key_columns, bool - ) or not isinstance(self.join_prefilter_max_key_columns, int): + if isinstance(self.join_prefilter_max_key_columns, bool) or not isinstance( + self.join_prefilter_max_key_columns, int + ): raise TypeError("join_prefilter_max_key_columns must be an int or None") if self.join_prefilter_max_key_columns < 1: raise ValueError( From 7da769cf647acfce92fb4add3ca745cbe08b0bc2 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:38:36 +0000 Subject: [PATCH 06/23] Accept integer join prefilter thresholds Allow numeric integer values such as 0 for join_prefilter_threshold and normalize the stored value to float during DynamicPlanningOptions validation. Reject booleans explicitly and add config coverage for the documented disable value. --- python/cudf_polars/cudf_polars/utils/config.py | 8 ++++++-- python/cudf_polars/tests/test_config.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index e8c4f3eb7846..bb7320225e2a 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -361,8 +361,12 @@ def __post_init__(self) -> None: # noqa: D105 if self.sample_chunk_count < 1: raise ValueError("sample_chunk_count must be at least 1") join_prefilter_threshold = self.join_prefilter_threshold - if not isinstance(join_prefilter_threshold, float): - raise TypeError("join_prefilter_threshold must be a float") + if isinstance(join_prefilter_threshold, bool) or not isinstance( + join_prefilter_threshold, (int, float) + ): + raise TypeError("join_prefilter_threshold must be a float or int") + join_prefilter_threshold = float(join_prefilter_threshold) + object.__setattr__(self, "join_prefilter_threshold", join_prefilter_threshold) if not 0.0 <= join_prefilter_threshold <= 1.0: raise ValueError("join_prefilter_threshold must be between 0 and 1") if self.join_prefilter_max_key_columns is not None: diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index ad1ffd6b8052..f44d166f220d 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -717,6 +717,15 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non def test_validate_join_prefilter_threshold() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"dynamic_planning": {"join_prefilter_threshold": 0}}, + ) + ) + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.join_prefilter_threshold == 0.0 + with pytest.raises(TypeError, match="join_prefilter_threshold must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( @@ -726,6 +735,15 @@ def test_validate_join_prefilter_threshold() -> None: }, ) ) + with pytest.raises(TypeError, match="join_prefilter_threshold must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_threshold": True} + }, + ) + ) with pytest.raises(ValueError, match="join_prefilter_threshold must be between"): ConfigOptions.from_polars_engine( pl.GPUEngine( From fd6ba2303cab249a3b7b4ec40fe5669da1f62055 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 19:30:22 +0200 Subject: [PATCH 07/23] Implement default in-place without additional variable Co-authored-by: Lawrence Mitchell --- python/cudf_polars/cudf_polars/utils/config.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index bb7320225e2a..417464d73597 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -144,7 +144,6 @@ class Cluster(enum.StrEnum): T = TypeVar("T") -_DEFAULT_JOIN_PREFILTER_MAX_KEY_COLUMNS: int | None = 1 def _make_default_factory( @@ -344,7 +343,7 @@ class DynamicPlanningOptions: default_factory=_make_default_factory( f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", _optional_int_converter, - default=_DEFAULT_JOIN_PREFILTER_MAX_KEY_COLUMNS, + default=1, ) ) join_prefilter_trace: bool = dataclasses.field( From fbe886b803ab50acff133815057821d4d94c262f Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 17:29:37 +0000 Subject: [PATCH 08/23] Document join prefilter key prefix limit Clarify that join_prefilter_max_key_columns controls the size of the join-key prefix used by the prefilter, rather than selecting an arbitrary key subset. --- python/cudf_polars/cudf_polars/utils/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 417464d73597..16b3327905af 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -318,8 +318,8 @@ class DynamicPlanningOptions: Row-count ratio (small / large) below which a join key prefilter is applied. Set to 0 to disable join prefiltering. Default is 0.5. join_prefilter_max_key_columns - Maximum number of join-key columns to use for the prefilter. Set to - ``None`` to use all join keys. Default is 1. + Maximum number of columns from the join-key prefix to use for the + prefilter. Set to ``None`` to use the full join-key list. Default is 1. join_prefilter_trace Whether to collect input/output row counts around applied join prefilters. Default is False. From f01e62e5cace46af929124913341614a0a379e16 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 17:36:29 +0000 Subject: [PATCH 09/23] Use dataclass conversion for join prefilter trace Replace the hand-written JoinPrefilterDecision trace metadata mapping with dataclasses.asdict, so the trace output follows the dataclass fields without duplicating the field list. --- .../cudf_polars/streaming/actor_graph/join.py | 17 ++--------------- 1 file changed, 2 insertions(+), 15 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 901969be0d50..4125f4903124 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -4,7 +4,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Literal import pylibcudf as plc @@ -116,20 +116,7 @@ def enabled(self) -> bool: def trace_dict(self) -> dict[str, Any]: """Return structured trace metadata for this decision.""" - metadata: dict[str, Any] = { - "considered": self.considered, - "estimated_left_rows": self.left_rows, - "estimated_right_rows": self.right_rows, - "threshold": self.threshold, - "filtered_side": self.filter_side, - "build_side": self.build_side, - "key_column_count": self.key_column_count, - } - if self.small_large_ratio is not None: - metadata["small_large_ratio"] = self.small_large_ratio - if self.reason_skipped is not None: - metadata["reason_skipped"] = self.reason_skipped - return metadata + return asdict(self) @define_actor() From ffd5e8f36a1662e081ee42d6b109454de7b61026 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 17:42:01 +0000 Subject: [PATCH 10/23] Remove redundant join prefilter decision state Drop the always-true JoinPrefilterDecision.considered field since the presence of a decision already means the prefilter was considered. Inline skipped JoinPrefilterDecision construction at the return sites so the selector does not carry a helper that only forwards dataclass arguments. --- .../cudf_polars/streaming/actor_graph/join.py | 60 ++++++------------- .../tests/streaming/test_tracing.py | 4 +- 2 files changed, 21 insertions(+), 43 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 4125f4903124..6c41aa0cc5e1 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -97,7 +97,6 @@ class JoinStrategy: class JoinPrefilterDecision: """Decision for an optional join-key prefilter stage.""" - considered: bool left_rows: int right_rows: int threshold: float @@ -590,26 +589,6 @@ async def passthrough_split( await ch_out.drain(context) -def _skipped_prefilter( - reason: str, - *, - left_rows: int, - right_rows: int, - threshold: float, - ratio: float | None = None, - key_column_count: int = 0, -) -> JoinPrefilterDecision: - return JoinPrefilterDecision( - considered=True, - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - key_column_count=key_column_count, - small_large_ratio=ratio, - reason_skipped=reason, - ) - - def _select_join_prefilter( join_type: Literal["Inner", "Left", "Right", "Full", "Semi", "Anti", "Cross"], left_rows: int, @@ -628,25 +607,25 @@ def _select_join_prefilter( """ key_column_count = len(left_key_indices) if threshold == 0.0: - return _skipped_prefilter( - "disabled", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, + reason_skipped="disabled", ) if key_column_count == 0: - return _skipped_prefilter( - "no_join_keys", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, + reason_skipped="no_join_keys", ) if key_column_count != len(right_key_indices): - return _skipped_prefilter( - "mismatched_join_keys", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, + reason_skipped="mismatched_join_keys", ) if max_key_columns is not None: @@ -669,13 +648,13 @@ def _select_join_prefilter( elif join_type in ("Left", "Anti"): if left_rows >= right_rows: ratio = right_rows / left_rows if left_rows > 0 else None - return _skipped_prefilter( - "no_legal_large_side", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, - ratio=ratio, + small_large_ratio=ratio, key_column_count=key_column_count, + reason_skipped="no_legal_large_side", ) build_side = "left" filter_side = "right" @@ -683,44 +662,44 @@ def _select_join_prefilter( elif join_type == "Right": if right_rows >= left_rows: ratio = left_rows / right_rows if right_rows > 0 else None - return _skipped_prefilter( - "no_legal_large_side", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, - ratio=ratio, + small_large_ratio=ratio, key_column_count=key_column_count, + reason_skipped="no_legal_large_side", ) build_side = "right" filter_side = "left" small_rows, large_rows = right_rows, left_rows else: - return _skipped_prefilter( - "unsupported_join_type", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, key_column_count=key_column_count, + reason_skipped="unsupported_join_type", ) if large_rows <= 0: - return _skipped_prefilter( - "no_large_side", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, key_column_count=key_column_count, + reason_skipped="no_large_side", ) ratio = small_rows / large_rows if ratio >= threshold: - return _skipped_prefilter( - "ratio_above_threshold", + return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, threshold=threshold, - ratio=ratio, + small_large_ratio=ratio, key_column_count=key_column_count, + reason_skipped="ratio_above_threshold", ) if build_side == "left": @@ -731,7 +710,6 @@ def _select_join_prefilter( apply_indices = left_key_indices[:key_column_count] return JoinPrefilterDecision( - considered=True, left_rows=left_rows, right_rows=right_rows, threshold=threshold, diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index df83512e0c8a..fa7c508670bd 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -44,8 +44,8 @@ def test_actor_tracer_counts_table_chunk_without_table_view(chunk: TableChunk) - def test_actor_tracer_records_extra_metadata() -> None: tracer = ActorTracer() - tracer.set_extra("join_prefilter", {"considered": True}) - assert tracer.extra == {"join_prefilter": {"considered": True}} + tracer.set_extra("join_prefilter", {"enabled": True}) + assert tracer.extra == {"join_prefilter": {"enabled": True}} @pytest.mark.spmd From 6bee03814d4bb47b587ec21b5f428942907e82d7 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 18:28:52 +0000 Subject: [PATCH 11/23] Remove unused keyless prefilter skip Drop the no_join_keys branch from the join prefilter selector. Keyless cross joins are already unsupported by the selector, so cover that behavior directly instead of carrying a separate skip reason. --- .../cudf_polars/streaming/actor_graph/join.py | 7 ------- python/cudf_polars/tests/streaming/test_join.py | 14 ++++++++++++++ 2 files changed, 14 insertions(+), 7 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 6c41aa0cc5e1..648e241e3fed 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -613,13 +613,6 @@ def _select_join_prefilter( threshold=threshold, reason_skipped="disabled", ) - if key_column_count == 0: - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - reason_skipped="no_join_keys", - ) if key_column_count != len(right_key_indices): return JoinPrefilterDecision( left_rows=left_rows, diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index d08401c8a995..683535b22f5f 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -400,6 +400,20 @@ def test_join_prefilter_skips_unsupported_full_join() -> None: assert decision.reason_skipped == "unsupported_join_type" +def test_join_prefilter_skips_unsupported_cross_join() -> None: + decision = _select_join_prefilter( + "Cross", + 10, + 1_000, + (), + (), + threshold=0.5, + max_key_columns=1, + ) + assert not decision.enabled + assert decision.reason_skipped == "unsupported_join_type" + + @pytest.mark.parametrize( "maintain_order", ["left_right", "right_left", "left", "right"] ) From 47f2ce4668cfa612bc2711c398f7788a490580d9 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 18:31:37 +0000 Subject: [PATCH 12/23] Assert matching join prefilter key counts Replace the defensive mismatched-key skip reason with an assertion. A valid Join IR must provide the same number of left and right join keys, so reaching this state indicates malformed join metadata rather than a prefilter planning decision. --- .../cudf_polars/streaming/actor_graph/join.py | 10 +++------- python/cudf_polars/tests/streaming/test_join.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 7 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 648e241e3fed..b12e6c720aa7 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -606,6 +606,9 @@ def _select_join_prefilter( join. The full join still runs afterward with the complete key set. """ key_column_count = len(left_key_indices) + assert key_column_count == len(right_key_indices), ( + "left and right join key counts must match" + ) if threshold == 0.0: return JoinPrefilterDecision( left_rows=left_rows, @@ -613,13 +616,6 @@ def _select_join_prefilter( threshold=threshold, reason_skipped="disabled", ) - if key_column_count != len(right_key_indices): - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - reason_skipped="mismatched_join_keys", - ) if max_key_columns is not None: key_column_count = min(key_column_count, max_key_columns) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index 683535b22f5f..5fc7a5a1ec1f 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -414,6 +414,21 @@ def test_join_prefilter_skips_unsupported_cross_join() -> None: assert decision.reason_skipped == "unsupported_join_type" +def test_join_prefilter_asserts_mismatched_key_count() -> None: + with pytest.raises( + AssertionError, match="left and right join key counts must match" + ): + _select_join_prefilter( + "Inner", + 10, + 1_000, + (0,), + (0, 1), + threshold=0.5, + max_key_columns=1, + ) + + @pytest.mark.parametrize( "maintain_order", ["left_right", "right_left", "left", "right"] ) From 97466a313baa9f5f9938240479a7a11b30f1d0fe Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 18:33:40 +0000 Subject: [PATCH 13/23] Track only join prefilter apply side Remove the redundant build_side field from JoinPrefilterDecision. The bloom-filter build side is the inverse of filter_side, so task construction now derives that relationship from the side being filtered. --- .../cudf_polars/streaming/actor_graph/join.py | 13 +++---------- python/cudf_polars/tests/streaming/test_join.py | 3 --- 2 files changed, 3 insertions(+), 13 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 b12e6c720aa7..b1c7e0c90c39 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -101,7 +101,6 @@ class JoinPrefilterDecision: right_rows: int threshold: float filter_side: Literal["left", "right"] | None = None - build_side: Literal["left", "right"] | None = None build_indices: tuple[int, ...] = () apply_indices: tuple[int, ...] = () key_column_count: int = 0 @@ -620,18 +619,15 @@ def _select_join_prefilter( if max_key_columns is not None: key_column_count = min(key_column_count, max_key_columns) - build_side: Literal["left", "right"] filter_side: Literal["left", "right"] small_rows: int large_rows: int if join_type in ("Inner", "Semi"): if left_rows <= right_rows: - build_side = "left" filter_side = "right" small_rows, large_rows = left_rows, right_rows else: - build_side = "right" filter_side = "left" small_rows, large_rows = right_rows, left_rows elif join_type in ("Left", "Anti"): @@ -645,7 +641,6 @@ def _select_join_prefilter( key_column_count=key_column_count, reason_skipped="no_legal_large_side", ) - build_side = "left" filter_side = "right" small_rows, large_rows = left_rows, right_rows elif join_type == "Right": @@ -659,7 +654,6 @@ def _select_join_prefilter( key_column_count=key_column_count, reason_skipped="no_legal_large_side", ) - build_side = "right" filter_side = "left" small_rows, large_rows = right_rows, left_rows else: @@ -691,7 +685,7 @@ def _select_join_prefilter( reason_skipped="ratio_above_threshold", ) - if build_side == "left": + if filter_side == "right": build_indices = left_key_indices[:key_column_count] apply_indices = right_key_indices[:key_column_count] else: @@ -703,7 +697,6 @@ def _select_join_prefilter( right_rows=right_rows, threshold=threshold, filter_side=filter_side, - build_side=build_side, build_indices=build_indices, apply_indices=apply_indices, key_column_count=key_column_count, @@ -772,11 +765,11 @@ def make_filter_tasks( Of new left and right channels, coroutines to await, and new channels to shutdown on error. """ assert decision.enabled - assert decision.build_side in ("left", "right") + assert decision.filter_side in ("left", "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 decision.build_side == "left": + if decision.filter_side == "right": passthrough_input = ch_left ch_left = passthrough_output build_indices = decision.build_indices diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index 5fc7a5a1ec1f..8ac799115032 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -308,7 +308,6 @@ def test_join_prefilter_filters_large_side_with_key_prefix() -> None: max_key_columns=1, ) assert decision.enabled - assert decision.build_side == "left" assert decision.filter_side == "right" assert decision.build_indices == (0,) assert decision.apply_indices == (3,) @@ -355,7 +354,6 @@ def test_join_prefilter_outer_semantics_only_filter_right_side(how) -> None: max_key_columns=1, ) assert decision.enabled - assert decision.build_side == "left" assert decision.filter_side == "right" @@ -382,7 +380,6 @@ def test_join_prefilter_right_join_only_filters_left_side() -> None: max_key_columns=1, ) assert decision.enabled - assert decision.build_side == "right" assert decision.filter_side == "left" From 543a73d7c7991f39bf7a426d105d00d4b0d15b6e Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 18:35:54 +0000 Subject: [PATCH 14/23] Remove unused prefilter partition trace Drop apply_side_prepartitioned from join prefilter trace metadata. The prefilter does not make adaptive decisions from this value and the information is not consumed elsewhere, so keeping it in the filter trace adds noise without affecting planning. --- .../cudf_polars/streaming/actor_graph/join.py | 14 -------------- 1 file changed, 14 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 b1c7e0c90c39..2c47e2b71027 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -42,7 +42,6 @@ ChannelManager, NormalizedPartitioning, TableSizeStats, - _is_already_partitioned, _sample_chunks, allgather_reduce, chunk_to_frame, @@ -895,19 +894,6 @@ async def _shuffle_join( max_key_columns=prefilter_max_key_columns, ) prefilter_trace_stats = prefilter_decision.trace_dict() - if prefilter_decision.enabled: - apply_meta = ( - strategy.right_meta - if prefilter_decision.filter_side == "right" - else strategy.left_meta - ) - assert apply_meta is not None - prefilter_trace_stats["apply_side_prepartitioned"] = _is_already_partitioned( - apply_meta, - prefilter_decision.apply_indices, - strategy.shuffle_modulus, - comm.nranks, - ) if tracer is not None: tracer.set_extra("join_prefilter", prefilter_trace_stats) From 0dfc80da6e8defa90e2b50b964a4f49e5c4d9cce Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 18:45:28 +0000 Subject: [PATCH 15/23] Clarify dynamic join collective IDs Keep the dynamic-join reserved collective count at four, but remove the inaccurate strategy-allgather wording from the comment and error message. The four reserved IDs are the allgather, left shuffle, right shuffle, and bloom filter. --- .../streaming/actor_graph/collectives/common.py | 4 ++-- .../cudf_polars/cudf_polars/streaming/actor_graph/join.py | 8 +++----- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py index c2169daadb90..c528cf8e42e7 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py @@ -138,8 +138,8 @@ def __enter__(self) -> dict[IR, list[int]]: _get_new_collective_id_unsafe(), ] elif isinstance(node, Join) and self.dynamic_planning_enabled: - # Join needs 4 IDs: size allgather, one strategy-specific - # allgather/bloom prefilter, left shuffle, and right shuffle. + # Join needs 4 IDs: allgather, left shuffle, right shuffle, + # and bloom filter. self.collective_id_map[node] = [ _get_new_collective_id_unsafe(), _get_new_collective_id_unsafe(), 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 2c47e2b71027..b71de62b6d76 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -1497,14 +1497,12 @@ def _( ): # Dynamic join - decide strategy at runtime collective_ids = list(rec.state["collective_id_map"].get(ir, [])) - # Join uses up to 4 collective IDs: size allgather, one - # strategy-specific allgather/bloom prefilter, left shuffle, and - # right shuffle. + # Join uses up to 4 collective IDs: allgather, left shuffle, right + # shuffle, and bloom filter. if len(collective_ids) < 4: raise ValueError( "Dynamic join requires 4 reserved collective IDs " - "(size allgather + strategy allgather/bloom prefilter " - "+ left shuffle + right shuffle); got " + "(allgather + left shuffle + right shuffle + bloom filter); got " f"{len(collective_ids)} for this Join. " "Ensure ReserveOpIDs is run with dynamic_planning enabled." ) From 94e132cb17dec1c007d6fbb7b4e4d2bfe7dccdd8 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 18:49:25 +0000 Subject: [PATCH 16/23] Document actor trace extra metadata Clarify that ActorTracer extra metadata is for nested runtime decisions that do not have their own IR node but should still be logged with the parent actor trace. --- .../cudf_polars/streaming/actor_graph/tracing.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py index ee5c75524d23..d46c985de106 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py @@ -86,7 +86,11 @@ def set_duplicated(self, *, duplicated: bool = True) -> None: self.duplicated = duplicated def set_extra(self, key: str, value: Any) -> None: - """Attach structured metadata to the actor trace event.""" + """Attach structured metadata to the current actor trace event. + + This is useful for nested runtime decisions that do not have a + separate IR node, but should still be logged with their parent actor. + """ self.extra[key] = value From 918417d511f81914e295f0304370783e95f54fba Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:00:34 -0700 Subject: [PATCH 17/23] Fix docstring linting --- .../cudf_polars/cudf_polars/streaming/actor_graph/tracing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py index d46c985de106..c318affa4d2a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py @@ -86,7 +86,8 @@ def set_duplicated(self, *, duplicated: bool = True) -> None: self.duplicated = duplicated def set_extra(self, key: str, value: Any) -> None: - """Attach structured metadata to the current actor trace event. + """ + Attach structured metadata to the current actor trace event. This is useful for nested runtime decisions that do not have a separate IR node, but should still be logged with their parent actor. From c75249d5dab9b5d42035abc86b3823af6ae2ddc2 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 13:09:48 -0700 Subject: [PATCH 18/23] Fix optional join prefilter default typing Cast the join_prefilter_max_key_columns default to int | None so _make_default_factory infers the same optional type as _optional_int_converter. This avoids mypy narrowing the default to int and rejecting the converter. --- python/cudf_polars/cudf_polars/utils/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 16b3327905af..2afe86502d5c 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -27,7 +27,7 @@ import importlib.util import json import os -from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast from rmm.pylibrmm import CudaStreamFlags, CudaStreamPool @@ -343,7 +343,7 @@ class DynamicPlanningOptions: default_factory=_make_default_factory( f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", _optional_int_converter, - default=1, + default=cast("int | None", 1), ) ) join_prefilter_trace: bool = dataclasses.field( From 702c6adac3bb408ba4cc9e5993962ee9faabe33b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 12:18:45 +0000 Subject: [PATCH 19/23] Generalize default factory result typing Model the environment converter result and fallback default as separate type variables. This lets optional converters use concrete non-optional defaults without casts while preserving the factory's complete return type. --- python/cudf_polars/cudf_polars/utils/config.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 2afe86502d5c..0a74cade39b8 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -27,7 +27,7 @@ import importlib.util import json import os -from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar from rmm.pylibrmm import CudaStreamFlags, CudaStreamPool @@ -144,12 +144,13 @@ class Cluster(enum.StrEnum): T = TypeVar("T") +DefaultT = TypeVar("DefaultT") def _make_default_factory( - key: str, converter: Callable[[str], T], *, default: T -) -> Callable[[], T]: - def default_factory() -> T: + key: str, converter: Callable[[str], T], *, default: DefaultT +) -> Callable[[], T | DefaultT]: + def default_factory() -> T | DefaultT: v = os.environ.get(key) if v is None: return default @@ -343,7 +344,7 @@ class DynamicPlanningOptions: default_factory=_make_default_factory( f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", _optional_int_converter, - default=cast("int | None", 1), + default=1, ) ) join_prefilter_trace: bool = dataclasses.field( From a37608df617d0737342d848d8886e96f347e1b44 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 12:22:58 +0000 Subject: [PATCH 20/23] Document join prefilter selection inputs Describe the inputs that drive join prefilter planning and the decision returned by _select_join_prefilter. Remove execution details that are outside the selector's responsibility. --- .../cudf_polars/streaming/actor_graph/join.py | 26 ++++++++++++++++--- 1 file changed, 23 insertions(+), 3 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 b71de62b6d76..a2f0eaa0f9a3 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -598,10 +598,30 @@ def _select_join_prefilter( max_key_columns: int | None, ) -> JoinPrefilterDecision: """ - Select a safe join-key prefilter. + Determine whether to apply a prefilter to a join. - The prefilter only removes rows that cannot participate in the original - join. The full join still runs afterward with the complete key set. + Parameters + ---------- + join_type + Type of join. + left_rows + Estimated number of rows in the left table. + right_rows + Estimated number of rows in the right table. + left_key_indices + Column indices of the join keys in the left table. + right_key_indices + Column indices of the join keys in the right table. + threshold + Small-to-large row-count ratio at or above which filtering is disabled. + max_key_columns + Maximum number of columns to use from the key prefix. ``None`` uses all + join-key columns. + + Returns + ------- + JoinPrefilterDecision + The selected prefilter configuration, or the reason it was skipped. """ key_column_count = len(left_key_indices) assert key_column_count == len(right_key_indices), ( From 6bc75592e2e40cda4bfac18f2d69275da862b867 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 12:27:02 +0000 Subject: [PATCH 21/23] Simplify join prefilter decision flow Separate unsupported and disabled cases from supported join planning, then collect filter-side, ratio, and skip-reason state into one final JoinPrefilterDecision. Preserve the existing selection behavior and trace metadata while making the control flow easier to follow. --- .../cudf_polars/streaming/actor_graph/join.py | 88 +++++++------------ 1 file changed, 31 insertions(+), 57 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 a2f0eaa0f9a3..031a7658f984 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -638,44 +638,7 @@ def _select_join_prefilter( if max_key_columns is not None: key_column_count = min(key_column_count, max_key_columns) - filter_side: Literal["left", "right"] - small_rows: int - large_rows: int - - if join_type in ("Inner", "Semi"): - if left_rows <= right_rows: - filter_side = "right" - small_rows, large_rows = left_rows, right_rows - else: - filter_side = "left" - small_rows, large_rows = right_rows, left_rows - elif join_type in ("Left", "Anti"): - if left_rows >= right_rows: - ratio = right_rows / left_rows if left_rows > 0 else None - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - small_large_ratio=ratio, - key_column_count=key_column_count, - reason_skipped="no_legal_large_side", - ) - filter_side = "right" - small_rows, large_rows = left_rows, right_rows - elif join_type == "Right": - if right_rows >= left_rows: - ratio = left_rows / right_rows if right_rows > 0 else None - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - small_large_ratio=ratio, - key_column_count=key_column_count, - reason_skipped="no_legal_large_side", - ) - filter_side = "left" - small_rows, large_rows = right_rows, left_rows - else: + if join_type not in ("Inner", "Semi", "Left", "Anti", "Right"): return JoinPrefilterDecision( left_rows=left_rows, right_rows=right_rows, @@ -684,32 +647,42 @@ def _select_join_prefilter( reason_skipped="unsupported_join_type", ) - if large_rows <= 0: - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - key_column_count=key_column_count, - reason_skipped="no_large_side", - ) + small_rows, large_rows = sorted((left_rows, right_rows)) + ratio = small_rows / large_rows if large_rows > 0 else None + filter_side: Literal["left", "right"] | None = None + reason_skipped: str | None = None - ratio = small_rows / large_rows - if ratio >= threshold: - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - small_large_ratio=ratio, - key_column_count=key_column_count, - reason_skipped="ratio_above_threshold", - ) + if join_type in ("Inner", "Semi"): + filter_side = "right" if left_rows <= right_rows else "left" + elif join_type in ("Left", "Anti"): + if left_rows >= right_rows: + reason_skipped = "no_legal_large_side" + else: + filter_side = "right" + else: + if right_rows >= left_rows: + reason_skipped = "no_legal_large_side" + else: + filter_side = "left" + + if reason_skipped is None: + if ratio is None: + reason_skipped = "no_large_side" + elif ratio >= threshold: + reason_skipped = "ratio_above_threshold" + + if reason_skipped is not None: + filter_side = None if filter_side == "right": build_indices = left_key_indices[:key_column_count] apply_indices = right_key_indices[:key_column_count] - else: + elif filter_side == "left": build_indices = right_key_indices[:key_column_count] apply_indices = left_key_indices[:key_column_count] + else: + build_indices = () + apply_indices = () return JoinPrefilterDecision( left_rows=left_rows, @@ -720,6 +693,7 @@ def _select_join_prefilter( apply_indices=apply_indices, key_column_count=key_column_count, small_large_ratio=ratio, + reason_skipped=reason_skipped, ) From 903084df9c3c825e63e2106b9674cc94327c2762 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 06:36:36 -0700 Subject: [PATCH 22/23] Fix missing coverage --- python/cudf_polars/tests/test_config.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index f44d166f220d..9e4b4e605eb7 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -30,6 +30,7 @@ CUDAStreamPoolConfig, Cluster, ConfigOptions, + DynamicPlanningOptions, MemoryResourceConfig, StreamingExecutor, _default_cuda_stream_policy, @@ -716,6 +717,17 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non assert config.executor.dynamic_planning.join_prefilter_trace +@pytest.mark.parametrize("value", ["none", "null"]) +def test_join_prefilter_max_key_columns_none_from_env( + monkeypatch: pytest.MonkeyPatch, value: str +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_MAX_KEY_COLUMNS", + value, + ) + assert DynamicPlanningOptions().join_prefilter_max_key_columns is None + + def test_validate_join_prefilter_threshold() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( @@ -787,9 +799,17 @@ def test_validate_join_prefilter_max_key_columns() -> None: ) -def test_dynamic_planning_from_instance() -> None: - from cudf_polars.utils.config import DynamicPlanningOptions +def test_validate_join_prefilter_trace() -> None: + with pytest.raises(TypeError, match="join_prefilter_trace must be a bool"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"dynamic_planning": {"join_prefilter_trace": "bad"}}, + ) + ) + +def test_dynamic_planning_from_instance() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", From 4538aef834216dafbeff96edf532bea958a4bcc3 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 09:05:17 -0700 Subject: [PATCH 23/23] Fix one more missing coverage --- python/cudf_polars/tests/test_config.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 9e4b4e605eb7..ad00a5b69d59 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -717,15 +717,15 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non assert config.executor.dynamic_planning.join_prefilter_trace -@pytest.mark.parametrize("value", ["none", "null"]) -def test_join_prefilter_max_key_columns_none_from_env( - monkeypatch: pytest.MonkeyPatch, value: str +@pytest.mark.parametrize("value, expected", [("none", None), ("null", None), ("2", 2)]) +def test_join_prefilter_max_key_columns_from_env( + monkeypatch: pytest.MonkeyPatch, value: str, expected: int | None ) -> None: monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_MAX_KEY_COLUMNS", value, ) - assert DynamicPlanningOptions().join_prefilter_max_key_columns is None + assert DynamicPlanningOptions().join_prefilter_max_key_columns == expected def test_validate_join_prefilter_threshold() -> None: