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/44] 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 6b71a9577d937201c7721ca9eec958892476da75 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 25 Jun 2026 10:04:47 +0000 Subject: [PATCH 02/44] Add generic derived join-domain prefilters Add a streaming optimizer pass that inserts generic derived key-domain semi joins before actor-graph lowering. The pass uses existing dynamic-planning scan statistics and join metadata to reduce large join inputs from selective domains while preserving the original full join for correctness. The optimizer handles simple selective-domain filters and constrained multi-key domains, including cases where a selective key on one side can narrow the domain used to prefilter a larger source. It skips unsupported join shapes, non-column keys, unselective simple domains, and non-inner joins. Wire the pass into streaming execution behind dynamic-planning options and add focused tests plus config coverage. --- .../streaming/join_domain_prefilter.py | 561 ++++++++++++++++++ .../cudf_polars/streaming/parallel.py | 9 +- .../cudf_polars/cudf_polars/utils/config.py | 59 ++ .../streaming/test_join_domain_prefilter.py | 209 +++++++ python/cudf_polars/tests/test_config.py | 68 +++ 5 files changed, 905 insertions(+), 1 deletion(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py create mode 100644 python/cudf_polars/tests/streaming/test_join_domain_prefilter.py diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py new file mode 100644 index 000000000000..f8499ce798a4 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -0,0 +1,561 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generic derived key-domain prefilters for streaming joins.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Cache, + DataFrameScan, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Scan, + Select, +) +from cudf_polars.dsl.tracing import Scope, log +from cudf_polars.dsl.traversal import traversal + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + + from cudf_polars.containers import DataType + from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.base import StatsCollector + from cudf_polars.utils.config import ConfigOptions, StreamingExecutor + + +@dataclass(frozen=True) +class _ColumnRef: + """A simple column join key.""" + + name: str + dtype: DataType + + +@dataclass(frozen=True) +class _Producer: + """A subtree that can provide a key domain.""" + + node: IR + column: str + rows: int + + +@dataclass(frozen=True) +class _Candidate: + """A derived key-domain prefilter candidate.""" + + mode: Literal["simple", "composite"] + target_side: Literal["left", "right"] + target: IR + target_key: _ColumnRef + domain: _Producer + domain_key: _ColumnRef + constraint_domain: _Producer | None = None + domain_constraint_key: _ColumnRef | None = None + target_constraint_key: _ColumnRef | None = None + + @property + def domain_rows(self) -> int: + """Estimated rows in the domain input.""" + return self.domain.rows + + @property + def target_rows(self) -> int: + """Estimated rows in the target input.""" + return _estimate_rows(self.target) or 0 + + @property + def score(self) -> tuple[int, int, int]: + """Prefer composite filters, then smaller constraint/domain inputs.""" + constraint_rows = ( + self.constraint_domain.rows + if self.constraint_domain is not None + else self.domain.rows + ) + return ( + 0 if self.mode == "composite" else 1, + constraint_rows, + self.domain.rows, + ) + + +_ROW_ESTIMATES: dict[IR, int | None] = {} +_SELECTIVE: dict[IR, bool] = {} +_STATS: StatsCollector | None = None + + +def optimize_join_domain_prefilters( + ir: IR, + stats: StatsCollector, + config_options: ConfigOptions[StreamingExecutor], +) -> IR: + """ + Insert generic semi-join key-domain prefilters before streaming lowering. + + The rewrite is intentionally conservative: only inner joins with simple + column equality keys are considered, and the original full join remains + after every inserted row-reduction semi join. + """ + dynamic_options = config_options.executor.dynamic_planning + if dynamic_options is None or not dynamic_options.join_domain_prefilter_enabled: + return ir + threshold = dynamic_options.join_domain_prefilter_threshold + trace = dynamic_options.join_domain_prefilter_trace + if threshold is None or threshold == 0 or trace is None: + return ir + + global _ROW_ESTIMATES, _SELECTIVE, _STATS + old_estimates, old_selective, old_stats = _ROW_ESTIMATES, _SELECTIVE, _STATS + _ROW_ESTIMATES, _SELECTIVE, _STATS = {}, {}, stats + try: + return _rewrite_node( + ir, + threshold=threshold, + trace=trace, + ) + finally: + _ROW_ESTIMATES, _SELECTIVE, _STATS = ( + old_estimates, + old_selective, + old_stats, + ) + + +def _rewrite_node(ir: IR, *, threshold: float, trace: bool) -> IR: + children = tuple( + _rewrite_node(child, threshold=threshold, trace=trace) for child in ir.children + ) + node = ir if children == ir.children else ir.reconstruct(children) + + if not isinstance(node, Join): + return node + + candidate, reason = _select_candidate(node, threshold) + if trace: + _trace_decision(node, threshold, candidate, reason) + if candidate is None: + return node + + left, right = node.children + target_filter = _make_target_filter(node, candidate) + if candidate.target_side == "left": + left = _replace_identity(left, candidate.target, target_filter) + else: + right = _replace_identity(right, candidate.target, target_filter) + return node.reconstruct((left, right)) + + +def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, str]: + if ir.options[0] != "Inner": + return None, "not_inner_join" + if ir.options[2] is not None: + return None, "sliced_join" + if ir.options[5] != "none": + return None, "maintain_order" + + left_keys = _simple_keys(ir.left_on, ir.children[0].schema) + right_keys = _simple_keys(ir.right_on, ir.children[1].schema) + if left_keys is None or right_keys is None: + return None, "non_column_join_key" + if len(left_keys) != len(right_keys): + return None, "key_count_mismatch" + + candidates: list[_Candidate] = [] + for target_side in ("left", "right"): + target_child, domain_child = ( + (ir.children[0], ir.children[1]) + if target_side == "left" + else (ir.children[1], ir.children[0]) + ) + target_keys, domain_keys = ( + (left_keys, right_keys) + if target_side == "left" + else (right_keys, left_keys) + ) + candidates.extend( + _composite_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + threshold, + ) + ) + candidates.extend( + _simple_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + threshold, + ) + ) + + if not candidates: + return None, "no_selective_domain" + return min(candidates, key=lambda c: c.score), "applied" + + +def _simple_keys( + keys: Sequence[expr.NamedExpr], schema: dict[str, DataType] +) -> tuple[_ColumnRef, ...] | None: + result: list[_ColumnRef] = [] + for key in keys: + if not isinstance(key.value, expr.Col): + return None + name = key.value.name + if name not in schema: + return None + result.append(_ColumnRef(name, schema[name])) + return tuple(result) + + +def _simple_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[_ColumnRef, ...], + domain_keys: tuple[_ColumnRef, ...], + threshold: float, +) -> Iterable[_Candidate]: + for target_key, domain_key in zip(target_keys, domain_keys, strict=True): + target = _largest_key_source(target_child, target_key.name) + if target is None: + continue + target_rows = _estimate_rows(target) + if target_rows is None or target_rows <= 0: + continue + domain = _smallest_key_producer( + domain_child, domain_key.name, require_selective=True + ) + if domain is None: + continue + if _contains_identity(target, domain.node): + continue + if domain.rows / target_rows > threshold: + continue + yield _Candidate( + mode="simple", + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=domain_key, + ) + + +def _composite_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[_ColumnRef, ...], + domain_keys: tuple[_ColumnRef, ...], + threshold: float, +) -> Iterable[_Candidate]: + if len(target_keys) < 2: + return + + for filter_index, (target_key, domain_key) in enumerate( + zip(target_keys, domain_keys, strict=True) + ): + target = _largest_key_source(target_child, target_key.name) + if target is None: + continue + target_rows = _estimate_rows(target) + if target_rows is None or target_rows <= 0: + continue + + for constraint_index, ( + target_constraint_key, + domain_constraint_key, + ) in enumerate(zip(target_keys, domain_keys, strict=True)): + if constraint_index == filter_index: + continue + domain = _smallest_node_containing_all( + domain_child, (domain_key.name, domain_constraint_key.name) + ) + if domain is None: + continue + constraint_domain = _smallest_key_producer( + target_child, + target_constraint_key.name, + require_selective=True, + exclude=target, + ) + if constraint_domain is None: + continue + if _contains_identity(target, domain.node) or _contains_identity( + target, constraint_domain.node + ): + continue + if domain.rows / target_rows > threshold: + continue + if constraint_domain.rows / domain.rows > threshold: + continue + yield _Candidate( + mode="composite", + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=domain_key, + constraint_domain=constraint_domain, + domain_constraint_key=domain_constraint_key, + target_constraint_key=target_constraint_key, + ) + + +def _make_target_filter(ir: Join, candidate: _Candidate) -> Join: + domain = _make_domain(candidate, ir) + return _make_semi_join( + candidate.target, + candidate.target_key, + domain, + _ColumnRef(candidate.domain_key.name, domain.schema[candidate.domain_key.name]), + nulls_equal=ir.options[1], + suffix=ir.options[3], + ) + + +def _make_domain(candidate: _Candidate, ir: Join) -> IR: + if candidate.mode == "simple": + return _select_key( + candidate.domain.node, + candidate.domain.column, + candidate.domain_key.name, + ) + + assert candidate.constraint_domain is not None + assert candidate.domain_constraint_key is not None + assert candidate.target_constraint_key is not None + + constraint_domain = _select_key( + candidate.constraint_domain.node, + candidate.constraint_domain.column, + candidate.target_constraint_key.name, + ) + constrained = _make_semi_join( + candidate.domain.node, + _ColumnRef( + candidate.domain_constraint_key.name, + candidate.domain.node.schema[candidate.domain_constraint_key.name], + ), + constraint_domain, + _ColumnRef( + candidate.target_constraint_key.name, + constraint_domain.schema[candidate.target_constraint_key.name], + ), + nulls_equal=ir.options[1], + suffix=ir.options[3], + ) + return _select_key(constrained, candidate.domain.column, candidate.domain_key.name) + + +def _select_key(source: IR, source_column: str, output_column: str) -> Select: + dtype = source.schema[source_column] + return Select( + {output_column: dtype}, + (expr.NamedExpr(output_column, expr.Col(dtype, source_column)),), + True, # noqa: FBT003 + source, + ) + + +def _make_semi_join( + target: IR, + target_key: _ColumnRef, + domain: IR, + domain_key: _ColumnRef, + *, + nulls_equal: bool, + suffix: str, +) -> Join: + return Join( + target.schema, + (expr.NamedExpr(target_key.name, expr.Col(target_key.dtype, target_key.name)),), + (expr.NamedExpr(domain_key.name, expr.Col(domain_key.dtype, domain_key.name)),), + ("Semi", nulls_equal, None, suffix, False, "none"), + target, + domain, + ) + + +def _smallest_key_producer( + root: IR, column: str, *, require_selective: bool, exclude: IR | None = None +) -> _Producer | None: + candidates = [] + for node in traversal([root]): + if node is exclude or column not in node.schema: + continue + rows = _estimate_rows(node) + if rows is None or rows <= 0: + continue + if require_selective and not _is_selective(node): + continue + candidates.append((rows, len(node.schema), _Producer(node, column, rows))) + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +def _smallest_node_containing_all(root: IR, columns: Sequence[str]) -> _Producer | None: + candidates = [] + needed = set(columns) + for node in traversal([root]): + if not needed.issubset(node.schema): + continue + rows = _estimate_rows(node) + if rows is None or rows <= 0: + continue + candidates.append((rows, len(node.schema), _Producer(node, columns[0], rows))) + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +def _largest_key_source(root: IR, column: str) -> IR | None: + source_candidates = [] + fallback_candidates = [] + for node in traversal([root]): + if column not in node.schema: + continue + rows = _estimate_rows(node) + if rows is None or rows <= 0: + continue + item = (rows, len(node.schema), node) + if isinstance(node, (Scan, DataFrameScan)): + source_candidates.append(item) + else: + fallback_candidates.append(item) + candidates = source_candidates or fallback_candidates + if not candidates: + return None + return max(candidates, key=lambda item: (item[0], -item[1]))[2] + + +def _estimate_rows(ir: IR) -> int | None: + try: + return _ROW_ESTIMATES[ir] + except KeyError: + pass + + rows: int | None + if isinstance(ir, (Scan, DataFrameScan)): + source = None if _STATS is None else _STATS.scan_stats.get(ir) + rows = None if source is None else source.row_count + if rows is None and isinstance(ir, DataFrameScan): + rows = ir.df.shape()[0] + elif isinstance(ir, (Select, Projection, HStack, Cache, Filter, Distinct, GroupBy)): + rows = _estimate_rows(ir.children[0]) + elif isinstance(ir, Join): + left_rows = _estimate_rows(ir.children[0]) + right_rows = _estimate_rows(ir.children[1]) + rows = _estimate_join_rows(ir.options[0], left_rows, right_rows) + else: + estimates = [ + estimate for child in ir.children if (estimate := _estimate_rows(child)) + ] + rows = max(estimates) if estimates else None + + _ROW_ESTIMATES[ir] = rows + return rows + + +def _estimate_join_rows( + how: str, left_rows: int | None, right_rows: int | None +) -> int | None: + if left_rows is None: + return right_rows + if right_rows is None: + return left_rows + if how in ("Inner", "Semi", "Anti"): + return min(left_rows, right_rows) + if how == "Left": + return left_rows + if how == "Right": + return right_rows + if how == "Full": + return max(left_rows, right_rows) + return None + + +def _is_selective(ir: IR) -> bool: + try: + return _SELECTIVE[ir] + except KeyError: + pass + + if isinstance(ir, Scan): + selective = ir.predicate is not None + elif isinstance(ir, Filter): + selective = True + else: + selective = any(_is_selective(child) for child in ir.children) + + _SELECTIVE[ir] = selective + return selective + + +def _contains_identity(root: IR, needle: IR) -> bool: + return any(node is needle for node in traversal([root])) + + +def _replace_identity(root: IR, old: IR, new: IR) -> IR: + if root is old: + return new + if not root.children: + return root + children = tuple(_replace_identity(child, old, new) for child in root.children) + if children == root.children: + return root + return root.reconstruct(children) + + +def _trace_decision( + ir: Join, threshold: float, candidate: _Candidate | None, reason: str +) -> None: + join_domain_prefilter: dict[str, Any] = { + "considered": True, + "threshold": threshold, + "reason": reason, + } + record = { + "scope": Scope.PLAN.value, + "join_domain_prefilter": join_domain_prefilter, + "actor_ir_id": ir.get_stable_id(), + "actor_ir_type": type(ir).__name__, + } + if candidate is not None: + join_domain_prefilter.update( + { + "mode": candidate.mode, + "target_side": candidate.target_side, + "target_key": candidate.target_key.name, + "domain_key": candidate.domain_key.name, + "estimated_target_rows": candidate.target_rows, + "estimated_domain_rows": candidate.domain_rows, + "target_node_type": type(candidate.target).__name__, + "domain_node_type": type(candidate.domain.node).__name__, + } + ) + if candidate.constraint_domain is not None: + join_domain_prefilter.update( + { + "constraint_key": candidate.target_constraint_key.name + if candidate.target_constraint_key is not None + else None, + "estimated_constraint_rows": candidate.constraint_domain.rows, + } + ) + log("Join Domain Prefilter", **record) diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 6f8734fd17b4..6cf8abfe6b44 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.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 evaluation.""" @@ -104,6 +104,13 @@ def lower_ir_graph( -------- lower_ir_node """ + if _dynamic_planning_on(config_options): + from cudf_polars.streaming.join_domain_prefilter import ( + optimize_join_domain_prefilters, + ) + + ir = optimize_join_domain_prefilters(ir, stats, config_options) + state: State = { "config_options": config_options, "stats": stats, diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index a1d98a124dd1..b52ec6b02a25 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -181,6 +181,12 @@ def _optional_int_converter(v: str) -> int | None: return int(v) +def _optional_bool_converter(v: str) -> bool | None: + if v.lower() in {"none", "null"}: + return None + return _bool_converter(v) + + @dataclasses.dataclass(frozen=True) class ParquetOptions: """ @@ -332,6 +338,16 @@ class DynamicPlanningOptions: join_prefilter_trace Whether to collect input/output row counts around applied join prefilters. Default is False. + join_domain_prefilter_enabled + Whether to insert generic derived key-domain semi-join filters before + lowering streaming joins. Default is True. + join_domain_prefilter_threshold + Row-count ratio (domain / target) below which a derived key-domain + semi-join filter is inserted. When unset, ``join_prefilter_threshold`` + is used. Default is unset. + join_domain_prefilter_trace + Whether to emit plan-time trace decisions for derived key-domain + prefilters. Default follows ``join_prefilter_trace``. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -367,6 +383,27 @@ class DynamicPlanningOptions: default=False, ) ) + join_domain_prefilter_enabled: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_DOMAIN_PREFILTER_ENABLED", + _bool_converter, + default=True, + ) + ) + join_domain_prefilter_threshold: float | None = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_DOMAIN_PREFILTER_THRESHOLD", + _optional_float_converter, + default=None, + ) + ) + join_domain_prefilter_trace: bool | None = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_DOMAIN_PREFILTER_TRACE", + _optional_bool_converter, + default=None, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.sample_chunk_count, int): @@ -396,6 +433,28 @@ def __post_init__(self) -> None: # noqa: D105 ) if not isinstance(self.join_prefilter_trace, bool): raise TypeError("join_prefilter_trace must be a bool") + if not isinstance(self.join_domain_prefilter_enabled, bool): + raise TypeError("join_domain_prefilter_enabled must be a bool") + join_domain_prefilter_threshold = self.join_domain_prefilter_threshold + if join_domain_prefilter_threshold is None: + join_domain_prefilter_threshold = join_prefilter_threshold + object.__setattr__( + self, + "join_domain_prefilter_threshold", + join_domain_prefilter_threshold, + ) + elif not isinstance(join_domain_prefilter_threshold, float): + raise TypeError("join_domain_prefilter_threshold must be a float or None") + if not 0.0 <= join_domain_prefilter_threshold <= 1.0: + raise ValueError("join_domain_prefilter_threshold must be between 0 and 1") + join_domain_prefilter_trace = self.join_domain_prefilter_trace + if join_domain_prefilter_trace is None: + join_domain_prefilter_trace = self.join_prefilter_trace + object.__setattr__( + self, "join_domain_prefilter_trace", join_domain_prefilter_trace + ) + elif not isinstance(join_domain_prefilter_trace, bool): + raise TypeError("join_domain_prefilter_trace must be a bool or None") @dataclasses.dataclass(frozen=True, eq=True) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py new file mode 100644 index 000000000000..44ebda768b8e --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal + +import polars as pl + +from cudf_polars.containers import DataType +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import Join, Scan +from cudf_polars.dsl.traversal import traversal +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.join_domain_prefilter import ( + optimize_join_domain_prefilters, +) +from cudf_polars.utils.config import ConfigOptions, ParquetOptions + +if TYPE_CHECKING: + from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.base import SerializedDataSourceInfo + +I64 = DataType(pl.Int64()) +BOOL = DataType(pl.Boolean()) + + +class _SourceInfo: + type: Literal["parquet"] = "parquet" + + def __init__(self, row_count: int | None) -> None: + self.row_count = row_count + + def column_storage_size(self, column: str) -> int | None: + del column + return None + + def serialize(self) -> SerializedDataSourceInfo: + return {"type": self.type, "row_count": self.row_count, "per_file_means": {}} + + @classmethod + def deserialize(cls, data: SerializedDataSourceInfo) -> _SourceInfo: + return cls(data["row_count"]) + + +def _scan(name: str, columns: tuple[str, ...], *, predicate: bool = False) -> Scan: + schema = dict.fromkeys(columns, I64) + mask = ( + expr.NamedExpr("__predicate", expr.Literal(BOOL, True)) # noqa: FBT003 + if predicate + else None + ) + return Scan( + schema, + "parquet", + {}, + None, + [f"/tmp/{name}.parquet"], + list(columns), + 0, + -1, + None, + None, + mask, + ParquetOptions(), + ) + + +def _key(node: IR, name: str) -> expr.NamedExpr: + return expr.NamedExpr(name, expr.Col(node.schema[name], name)) + + +def _join( + left: IR, + right: IR, + left_on: tuple[str, ...], + right_on: tuple[str, ...], + *, + how: str = "Inner", + maintain_order: str = "none", +) -> Join: + schema = dict(left.schema) + schema.update(right.schema) + return Join( + schema, + tuple(_key(left, name) for name in left_on), + tuple(_key(right, name) for name in right_on), + (how, False, None, "_right", False, maintain_order), + left, + right, + ) + + +def _stats(**row_counts: tuple[Scan, int]) -> StatsCollector: + stats = StatsCollector() + for scan, rows in row_counts.values(): + stats.scan_stats[scan] = _SourceInfo(rows) + return stats + + +def _config() -> ConfigOptions: + return ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": { + "join_domain_prefilter_enabled": True, + "join_domain_prefilter_trace": False, + } + }, + ) + ) + + +def _joins(ir: IR, how: str | None = None) -> list[Join]: + return [ + node + for node in traversal([ir]) + if isinstance(node, Join) and (how is None or node.options[0] == how) + ] + + +def test_simple_domain_prefilter_filters_large_side() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + lineitem = _scan("lineitem", ("l_partkey", "l_suppkey")) + root = _join(part, lineitem, ("p_partkey",), ("l_partkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(part=(part, 6), lineitem=(lineitem, 1_800)), + _config(), + ) + + assert isinstance(optimized, Join) + assert optimized.options[0] == "Inner" + assert isinstance(optimized.children[1], Join) + assert optimized.children[1].options[0] == "Semi" + assert optimized.children[1].children[0] is lineitem + assert optimized.children[0] is part + + +def test_no_simple_domain_prefilter_when_domain_is_not_selective() -> None: + supplier = _scan("supplier", ("s_suppkey",)) + lineitem = _scan("lineitem", ("l_suppkey",)) + root = _join(supplier, lineitem, ("s_suppkey",), ("l_suppkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(supplier=(supplier, 30), lineitem=(lineitem, 1_800)), + _config(), + ) + + assert optimized is root + assert not _joins(optimized, "Semi") + + +def test_composite_domain_prefilter_constrains_domain_first() -> None: + nation = _scan("nation", ("n_nationkey",), predicate=True) + orders = _scan("orders", ("o_orderkey", "n_nationkey")) + lineitem = _scan("lineitem", ("l_orderkey", "l_suppkey")) + supplier = _scan("supplier", ("s_suppkey", "s_nationkey")) + + nation_orders = _join(nation, orders, ("n_nationkey",), ("n_nationkey",)) + order_lineitem = _join( + nation_orders, + lineitem, + ("o_orderkey",), + ("l_orderkey",), + maintain_order="left", + ) + root = _join( + order_lineitem, + supplier, + ("l_suppkey", "n_nationkey"), + ("s_suppkey", "s_nationkey"), + ) + + optimized = optimize_join_domain_prefilters( + root, + _stats( + nation=(nation, 5), + orders=(orders, 900), + lineitem=(lineitem, 1_800), + supplier=(supplier, 30), + ), + _config(), + ) + + semis = _joins(optimized, "Semi") + assert isinstance(optimized, Join) + assert optimized.options[0] == "Inner" + assert optimized.children[1] is supplier + assert any(semi.children[0] is supplier for semi in semis) + assert any(semi.children[0] is lineitem for semi in semis) + + +def test_no_domain_prefilter_for_outer_join() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + lineitem = _scan("lineitem", ("l_partkey",)) + root = _join(part, lineitem, ("p_partkey",), ("l_partkey",), how="Left") + + optimized = optimize_join_domain_prefilters( + root, + _stats(part=(part, 6), lineitem=(lineitem, 1_800)), + _config(), + ) + + assert optimized is root + assert not _joins(optimized, "Semi") diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 8557bd919230..f8457380a9f9 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -676,6 +676,9 @@ def test_dynamic_planning_defaults() -> None: 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 + assert config.executor.dynamic_planning.join_domain_prefilter_enabled + assert config.executor.dynamic_planning.join_domain_prefilter_threshold == 0.5 + assert not config.executor.dynamic_planning.join_domain_prefilter_trace def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -747,6 +750,30 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non 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 + assert config.executor.dynamic_planning.join_domain_prefilter_threshold == 0.25 + assert config.executor.dynamic_planning.join_domain_prefilter_trace + + +def test_join_domain_prefilter_options_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_ENABLED", + "0", + ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_THRESHOLD", + "0.125", + ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_TRACE", + "1", + ) + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.dynamic_planning is not None + assert not config.executor.dynamic_planning.join_domain_prefilter_enabled + assert config.executor.dynamic_planning.join_domain_prefilter_threshold == 0.125 + assert config.executor.dynamic_planning.join_domain_prefilter_trace def test_join_prefilter_threshold_overrides_bloom_threshold() -> None: @@ -810,6 +837,47 @@ def test_validate_join_prefilter_max_key_columns() -> None: ) +def test_validate_join_domain_prefilter_options() -> None: + with pytest.raises(TypeError, match="join_domain_prefilter_enabled must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_domain_prefilter_enabled": "bad"} + }, + ) + ) + with pytest.raises(TypeError, match="join_domain_prefilter_threshold must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_domain_prefilter_threshold": "bad"} + }, + ) + ) + with pytest.raises( + ValueError, match="join_domain_prefilter_threshold must be between" + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_domain_prefilter_threshold": 1.5} + }, + ) + ) + with pytest.raises(TypeError, match="join_domain_prefilter_trace must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_domain_prefilter_trace": "bad"} + }, + ) + ) + + def test_dynamic_planning_from_instance() -> None: from cudf_polars.utils.config import DynamicPlanningOptions From b22929fe4dd3a9bb48099c8d15c05e58af916587 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 25 Jun 2026 14:12:21 +0000 Subject: [PATCH 03/44] Add profitability guards for derived join prefilters Teach the derived join-domain prefilter planner to estimate the source-scan cost needed to build candidate domains. Candidate selection now prefers cheaper domain producers before smaller estimated output rows, and rejects prefilters whose domain-build cost is too large relative to the target being reduced. Also prevent a single-source domain from being stacked onto a target source that is already below the filtered side of a semi join. This keeps the profitable Q9 domain choices while avoiding the Q8-shaped memory regression caused by adding a second high-cardinality source-only prefilter to an already-reduced lineitem scan. Extend trace metadata with estimated target, domain, and constraint costs, and add regression tests for both the Q9 profitability case and the Q8 stacked-prefilter case. --- .../streaming/join_domain_prefilter.py | 143 ++++++++++++++++-- .../streaming/test_join_domain_prefilter.py | 76 ++++++++++ 2 files changed, 204 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index f8499ce798a4..9df83b2eb161 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -47,6 +47,7 @@ class _Producer: node: IR column: str rows: int + cost: int @dataclass(frozen=True) @@ -68,6 +69,11 @@ def domain_rows(self) -> int: """Estimated rows in the domain input.""" return self.domain.rows + @property + def domain_cost(self) -> int: + """Estimated scan-row cost to build the domain input.""" + return self.domain.cost + @property def target_rows(self) -> int: """Estimated rows in the target input.""" @@ -75,20 +81,22 @@ def target_rows(self) -> int: @property def score(self) -> tuple[int, int, int]: - """Prefer composite filters, then smaller constraint/domain inputs.""" - constraint_rows = ( - self.constraint_domain.rows + """Prefer composite filters, then cheaper constraint/domain inputs.""" + constraint_cost = ( + self.constraint_domain.cost if self.constraint_domain is not None - else self.domain.rows + else self.domain.cost ) return ( 0 if self.mode == "composite" else 1, - constraint_rows, - self.domain.rows, + constraint_cost, + self.domain.cost, ) _ROW_ESTIMATES: dict[IR, int | None] = {} +_SOURCE_COSTS: dict[IR, int | None] = {} +_SOURCE_COUNTS: dict[IR, int] = {} _SELECTIVE: dict[IR, bool] = {} _STATS: StatsCollector | None = None @@ -113,9 +121,21 @@ def optimize_join_domain_prefilters( if threshold is None or threshold == 0 or trace is None: return ir - global _ROW_ESTIMATES, _SELECTIVE, _STATS - old_estimates, old_selective, old_stats = _ROW_ESTIMATES, _SELECTIVE, _STATS - _ROW_ESTIMATES, _SELECTIVE, _STATS = {}, {}, stats + global _ROW_ESTIMATES, _SOURCE_COSTS, _SOURCE_COUNTS, _SELECTIVE, _STATS + old_estimates, old_costs, old_counts, old_selective, old_stats = ( + _ROW_ESTIMATES, + _SOURCE_COSTS, + _SOURCE_COUNTS, + _SELECTIVE, + _STATS, + ) + _ROW_ESTIMATES, _SOURCE_COSTS, _SOURCE_COUNTS, _SELECTIVE, _STATS = ( + {}, + {}, + {}, + {}, + stats, + ) try: return _rewrite_node( ir, @@ -123,8 +143,10 @@ def optimize_join_domain_prefilters( trace=trace, ) finally: - _ROW_ESTIMATES, _SELECTIVE, _STATS = ( + _ROW_ESTIMATES, _SOURCE_COSTS, _SOURCE_COUNTS, _SELECTIVE, _STATS = ( old_estimates, + old_costs, + old_counts, old_selective, old_stats, ) @@ -203,7 +225,7 @@ def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, st ) if not candidates: - return None, "no_selective_domain" + return None, "no_profitable_domain" return min(candidates, key=lambda c: c.score), "applied" @@ -243,8 +265,14 @@ def _simple_candidates( continue if _contains_identity(target, domain.node): continue + if _is_source_only_domain(domain.node) and _has_filtering_semi_ancestor( + target_child, target + ): + continue if domain.rows / target_rows > threshold: continue + if not _domain_cost_is_small(domain.node, target, threshold): + continue yield _Candidate( mode="simple", target_side=target_side, @@ -303,6 +331,12 @@ def _composite_candidates( continue if constraint_domain.rows / domain.rows > threshold: continue + if not _domain_cost_is_small(domain.node, target, threshold): + continue + if not _domain_cost_is_small( + constraint_domain.node, domain.node, threshold + ): + continue yield _Candidate( mode="composite", target_side=target_side, @@ -403,10 +437,14 @@ def _smallest_key_producer( continue if require_selective and not _is_selective(node): continue - candidates.append((rows, len(node.schema), _Producer(node, column, rows))) + cost = _source_cost(node) + if cost is None: + continue + producer = _Producer(node, column, rows, cost) + candidates.append((cost, rows, len(node.schema), producer)) if not candidates: return None - return min(candidates, key=lambda item: (item[0], item[1]))[2] + return min(candidates, key=lambda item: (item[0], item[1], item[2]))[3] def _smallest_node_containing_all(root: IR, columns: Sequence[str]) -> _Producer | None: @@ -418,10 +456,14 @@ def _smallest_node_containing_all(root: IR, columns: Sequence[str]) -> _Producer rows = _estimate_rows(node) if rows is None or rows <= 0: continue - candidates.append((rows, len(node.schema), _Producer(node, columns[0], rows))) + cost = _source_cost(node) + if cost is None: + continue + producer = _Producer(node, columns[0], rows, cost) + candidates.append((cost, rows, len(node.schema), producer)) if not candidates: return None - return min(candidates, key=lambda item: (item[0], item[1]))[2] + return min(candidates, key=lambda item: (item[0], item[1], item[2]))[3] def _largest_key_source(root: IR, column: str) -> IR | None: @@ -490,6 +532,74 @@ def _estimate_join_rows( return None +def _domain_cost_is_small(domain: IR, target: IR, threshold: float) -> bool: + """Return whether building a domain is cheap enough for the target it reduces.""" + domain_cost = _source_cost(domain) + target_rows = _estimate_rows(target) + if domain_cost is None or target_rows is None or target_rows <= 0: + return False + return domain_cost / target_rows <= threshold + + +def _source_cost(ir: IR) -> int | None: + """Estimate source rows that must be scanned to materialize an IR subtree.""" + try: + return _SOURCE_COSTS[ir] + except KeyError: + pass + + sources: list[int] = [] + seen: set[IR] = set() + for node in traversal([ir]): + if node in seen: + continue + seen.add(node) + if not isinstance(node, (Scan, DataFrameScan)): + continue + rows = _estimate_rows(node) + if rows is not None and rows > 0: + sources.append(rows) + + cost = sum(sources) if sources else _estimate_rows(ir) + _SOURCE_COSTS[ir] = cost + return cost + + +def _source_count(ir: IR) -> int: + """Count unique source scans in an IR subtree.""" + try: + return _SOURCE_COUNTS[ir] + except KeyError: + pass + + sources = { + node for node in traversal([ir]) if isinstance(node, (Scan, DataFrameScan)) + } + count = len(sources) + _SOURCE_COUNTS[ir] = count + return count + + +def _is_source_only_domain(ir: IR) -> bool: + """Return whether a domain is derived from just one source scan.""" + return _source_count(ir) == 1 + + +def _has_filtering_semi_ancestor(root: IR, target: IR) -> bool: + """Return whether target is already below a semi join's filtered side.""" + if root is target: + return False + + for index, child in enumerate(root.children): + if not _contains_identity(child, target): + continue + if isinstance(root, Join) and root.options[0] == "Semi" and index == 0: + return True + if _has_filtering_semi_ancestor(child, target): + return True + return False + + def _is_selective(ir: IR) -> bool: try: return _SELECTIVE[ir] @@ -545,6 +655,8 @@ def _trace_decision( "domain_key": candidate.domain_key.name, "estimated_target_rows": candidate.target_rows, "estimated_domain_rows": candidate.domain_rows, + "estimated_target_cost": _source_cost(candidate.target), + "estimated_domain_cost": candidate.domain_cost, "target_node_type": type(candidate.target).__name__, "domain_node_type": type(candidate.domain.node).__name__, } @@ -556,6 +668,7 @@ def _trace_decision( if candidate.target_constraint_key is not None else None, "estimated_constraint_rows": candidate.constraint_domain.rows, + "estimated_constraint_cost": candidate.constraint_domain.cost, } ) log("Join Domain Prefilter", **record) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 44ebda768b8e..1d6fd4f2487f 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -120,6 +120,18 @@ def _joins(ir: IR, how: str | None = None) -> list[Join]: ] +def _contains_node(ir: IR, needle: IR) -> bool: + return any(node is needle for node in traversal([ir])) + + +def _join_key_names(keys: tuple[expr.NamedExpr, ...]) -> tuple[str, ...]: + names = [] + for key in keys: + assert isinstance(key.value, expr.Col) + names.append(key.value.name) + return tuple(names) + + def test_simple_domain_prefilter_filters_large_side() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey", "l_suppkey")) @@ -194,6 +206,70 @@ def test_composite_domain_prefilter_constrains_domain_first() -> None: assert any(semi.children[0] is lineitem for semi in semis) +def test_prefilter_uses_cheaper_source_domain_and_skips_expensive_domain() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + partsupp = _scan("partsupp", ("ps_partkey", "ps_suppkey")) + supplier = _scan("supplier", ("s_suppkey",)) + lineitem = _scan("lineitem", ("l_partkey", "l_suppkey", "l_orderkey")) + orders = _scan("orders", ("o_orderkey",)) + + part_partsupp = _join(part, partsupp, ("p_partkey",), ("ps_partkey",)) + part_partsupp_supplier = _join( + part_partsupp, supplier, ("ps_suppkey",), ("s_suppkey",) + ) + q9_left = _join( + part_partsupp_supplier, + lineitem, + ("p_partkey", "ps_suppkey"), + ("l_partkey", "l_suppkey"), + ) + root = _join(q9_left, orders, ("l_orderkey",), ("o_orderkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats( + part=(part, 60), + partsupp=(partsupp, 120), + supplier=(supplier, 30), + lineitem=(lineitem, 1_800), + orders=(orders, 900), + ), + _config(), + ) + + semis = _joins(optimized, "Semi") + lineitem_semis = [semi for semi in semis if semi.children[0] is lineitem] + assert lineitem_semis + assert not any(semi.children[0] is orders for semi in semis) + assert _contains_node(lineitem_semis[0].children[1], part) + assert not _contains_node(lineitem_semis[0].children[1], supplier) + + +def test_source_only_domain_does_not_stack_on_prefiltered_source() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + lineitem = _scan("lineitem", ("l_partkey", "l_orderkey")) + orders = _scan("orders", ("o_orderkey",), predicate=True) + + part_lineitem = _join(part, lineitem, ("p_partkey",), ("l_partkey",)) + root = _join(part_lineitem, orders, ("l_orderkey",), ("o_orderkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(part=(part, 60), lineitem=(lineitem, 1_800), orders=(orders, 150)), + _config(), + ) + + lineitem_semis = [ + semi for semi in _joins(optimized, "Semi") if semi.children[0] is lineitem + ] + assert any( + _join_key_names(semi.left_on) == ("l_partkey",) for semi in lineitem_semis + ) + assert not any( + _join_key_names(semi.left_on) == ("l_orderkey",) for semi in lineitem_semis + ) + + def test_no_domain_prefilter_for_outer_join() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",)) From 09f04743d419647bcb734ec404d60e961a799574 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:14:58 +0000 Subject: [PATCH 04/44] 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 05/44] 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 06/44] 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 07/44] 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 08/44] 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 09/44] 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 10/44] 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 11/44] 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 12/44] 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 13/44] 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 14/44] 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 15/44] 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 16/44] 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 17/44] 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 18/44] 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 19/44] 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 20/44] 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 4cc6213ccefc8203c3c1fbe7c0b12b7d23db72f1 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 08:09:35 +0000 Subject: [PATCH 21/44] Use column expressions for join domain keys Represent simple join keys directly as expr.Col nodes instead of duplicating their names and dtypes in a private wrapper. Rely on valid Join IR to bind each column to its input schema, while retaining explicit handling for non-column key expressions. --- .../streaming/join_domain_prefilter.py | 63 +++++++------------ 1 file changed, 22 insertions(+), 41 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index f8499ce798a4..9289eda1ea34 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -26,20 +26,11 @@ if TYPE_CHECKING: from collections.abc import Iterable, Sequence - from cudf_polars.containers import DataType from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector from cudf_polars.utils.config import ConfigOptions, StreamingExecutor -@dataclass(frozen=True) -class _ColumnRef: - """A simple column join key.""" - - name: str - dtype: DataType - - @dataclass(frozen=True) class _Producer: """A subtree that can provide a key domain.""" @@ -56,12 +47,12 @@ class _Candidate: mode: Literal["simple", "composite"] target_side: Literal["left", "right"] target: IR - target_key: _ColumnRef + target_key: expr.Col domain: _Producer - domain_key: _ColumnRef + domain_key: expr.Col constraint_domain: _Producer | None = None - domain_constraint_key: _ColumnRef | None = None - target_constraint_key: _ColumnRef | None = None + domain_constraint_key: expr.Col | None = None + target_constraint_key: expr.Col | None = None @property def domain_rows(self) -> int: @@ -162,9 +153,9 @@ def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, st if ir.options[5] != "none": return None, "maintain_order" - left_keys = _simple_keys(ir.left_on, ir.children[0].schema) - right_keys = _simple_keys(ir.right_on, ir.children[1].schema) - if left_keys is None or right_keys is None: + left_keys = _simple_keys(ir.left_on) + right_keys = _simple_keys(ir.right_on) + if len(left_keys) != len(ir.left_on) or len(right_keys) != len(ir.right_on): return None, "non_column_join_key" if len(left_keys) != len(right_keys): return None, "key_count_mismatch" @@ -207,26 +198,16 @@ def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, st return min(candidates, key=lambda c: c.score), "applied" -def _simple_keys( - keys: Sequence[expr.NamedExpr], schema: dict[str, DataType] -) -> tuple[_ColumnRef, ...] | None: - result: list[_ColumnRef] = [] - for key in keys: - if not isinstance(key.value, expr.Col): - return None - name = key.value.name - if name not in schema: - return None - result.append(_ColumnRef(name, schema[name])) - return tuple(result) +def _simple_keys(keys: Sequence[expr.NamedExpr]) -> tuple[expr.Col, ...]: + return tuple(key.value for key in keys if isinstance(key.value, expr.Col)) def _simple_candidates( target_side: Literal["left", "right"], target_child: IR, domain_child: IR, - target_keys: tuple[_ColumnRef, ...], - domain_keys: tuple[_ColumnRef, ...], + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], threshold: float, ) -> Iterable[_Candidate]: for target_key, domain_key in zip(target_keys, domain_keys, strict=True): @@ -259,8 +240,8 @@ def _composite_candidates( target_side: Literal["left", "right"], target_child: IR, domain_child: IR, - target_keys: tuple[_ColumnRef, ...], - domain_keys: tuple[_ColumnRef, ...], + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], threshold: float, ) -> Iterable[_Candidate]: if len(target_keys) < 2: @@ -322,7 +303,7 @@ def _make_target_filter(ir: Join, candidate: _Candidate) -> Join: candidate.target, candidate.target_key, domain, - _ColumnRef(candidate.domain_key.name, domain.schema[candidate.domain_key.name]), + expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), nulls_equal=ir.options[1], suffix=ir.options[3], ) @@ -347,14 +328,14 @@ def _make_domain(candidate: _Candidate, ir: Join) -> IR: ) constrained = _make_semi_join( candidate.domain.node, - _ColumnRef( - candidate.domain_constraint_key.name, + expr.Col( candidate.domain.node.schema[candidate.domain_constraint_key.name], + candidate.domain_constraint_key.name, ), constraint_domain, - _ColumnRef( - candidate.target_constraint_key.name, + expr.Col( constraint_domain.schema[candidate.target_constraint_key.name], + candidate.target_constraint_key.name, ), nulls_equal=ir.options[1], suffix=ir.options[3], @@ -374,17 +355,17 @@ def _select_key(source: IR, source_column: str, output_column: str) -> Select: def _make_semi_join( target: IR, - target_key: _ColumnRef, + target_key: expr.Col, domain: IR, - domain_key: _ColumnRef, + domain_key: expr.Col, *, nulls_equal: bool, suffix: str, ) -> Join: return Join( target.schema, - (expr.NamedExpr(target_key.name, expr.Col(target_key.dtype, target_key.name)),), - (expr.NamedExpr(domain_key.name, expr.Col(domain_key.dtype, domain_key.name)),), + (expr.NamedExpr(target_key.name, target_key),), + (expr.NamedExpr(domain_key.name, domain_key),), ("Semi", nulls_equal, None, suffix, False, "none"), target, domain, From c148219194d561300fa1358623fa24696ae66166 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 08:11:51 +0000 Subject: [PATCH 22/44] Use shared DAG utilities for domain prefilter planning Rewrite join nodes with CachingVisitor and singledispatch, carry analysis through explicit visitor state, and compute row estimates and selectivity in post-order traversals. Replace target nodes with the shared DAG replacement helper, eliminating module-global caches and hand-written recursive traversal. --- .../streaming/join_domain_prefilter.py | 238 ++++++++++-------- 1 file changed, 135 insertions(+), 103 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 9289eda1ea34..d506bd3f6162 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -5,10 +5,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal +from functools import singledispatch +from typing import TYPE_CHECKING, Any, Literal, TypedDict from cudf_polars.dsl import expr from cudf_polars.dsl.ir import ( + IR, Cache, DataFrameScan, Distinct, @@ -21,13 +23,19 @@ Select, ) from cudf_polars.dsl.tracing import Scope, log -from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.traversal import ( + CachingVisitor, + post_traversal, + reuse_if_unchanged, + traversal, +) +from cudf_polars.dsl.utils.replace import replace if TYPE_CHECKING: from collections.abc import Iterable, Sequence - from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector + from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -50,6 +58,7 @@ class _Candidate: target_key: expr.Col domain: _Producer domain_key: expr.Col + target_rows: int constraint_domain: _Producer | None = None domain_constraint_key: expr.Col | None = None target_constraint_key: expr.Col | None = None @@ -59,11 +68,6 @@ def domain_rows(self) -> int: """Estimated rows in the domain input.""" return self.domain.rows - @property - def target_rows(self) -> int: - """Estimated rows in the target input.""" - return _estimate_rows(self.target) or 0 - @property def score(self) -> tuple[int, int, int]: """Prefer composite filters, then smaller constraint/domain inputs.""" @@ -79,9 +83,13 @@ def score(self) -> tuple[int, int, int]: ) -_ROW_ESTIMATES: dict[IR, int | None] = {} -_SELECTIVE: dict[IR, bool] = {} -_STATS: StatsCollector | None = None +class _RewriteState(TypedDict): + """State shared by the join-domain prefilter DAG rewrite.""" + + threshold: float + trace: bool + row_estimates: dict[IR, int | None] + selective_nodes: set[IR] def optimize_join_domain_prefilters( @@ -104,48 +112,59 @@ def optimize_join_domain_prefilters( if threshold is None or threshold == 0 or trace is None: return ir - global _ROW_ESTIMATES, _SELECTIVE, _STATS - old_estimates, old_selective, old_stats = _ROW_ESTIMATES, _SELECTIVE, _STATS - _ROW_ESTIMATES, _SELECTIVE, _STATS = {}, {}, stats - try: - return _rewrite_node( - ir, - threshold=threshold, - trace=trace, - ) - finally: - _ROW_ESTIMATES, _SELECTIVE, _STATS = ( - old_estimates, - old_selective, - old_stats, - ) + state = _RewriteState( + threshold=threshold, + trace=trace, + row_estimates=_estimate_row_counts(ir, stats), + selective_nodes=_collect_selective_nodes(ir), + ) + mapper: GenericTransformer[IR, IR, _RewriteState] = CachingVisitor( + _rewrite, state=state + ) + return mapper(ir) -def _rewrite_node(ir: IR, *, threshold: float, trace: bool) -> IR: - children = tuple( - _rewrite_node(child, threshold=threshold, trace=trace) for child in ir.children - ) - node = ir if children == ir.children else ir.reconstruct(children) +@singledispatch +def _rewrite(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + raise AssertionError - if not isinstance(node, Join): - return node - candidate, reason = _select_candidate(node, threshold) - if trace: - _trace_decision(node, threshold, candidate, reason) +@_rewrite.register(IR) +def _(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + return reuse_if_unchanged(node, rec) + + +@_rewrite.register(Join) +def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + rewritten = reuse_if_unchanged(node, rec) + assert isinstance(rewritten, Join) + node = rewritten + candidate, reason = _select_candidate( + node, + rec.state["threshold"], + rec.state["row_estimates"], + rec.state["selective_nodes"], + ) + if rec.state["trace"]: + _trace_decision(node, rec.state["threshold"], candidate, reason) if candidate is None: return node left, right = node.children target_filter = _make_target_filter(node, candidate) if candidate.target_side == "left": - left = _replace_identity(left, candidate.target, target_filter) + (left,) = replace([left], {candidate.target: target_filter}) else: - right = _replace_identity(right, candidate.target, target_filter) + (right,) = replace([right], {candidate.target: target_filter}) return node.reconstruct((left, right)) -def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, str]: +def _select_candidate( + ir: Join, + threshold: float, + row_estimates: dict[IR, int | None], + selective_nodes: set[IR], +) -> tuple[_Candidate | None, str]: if ir.options[0] != "Inner": return None, "not_inner_join" if ir.options[2] is not None: @@ -180,6 +199,8 @@ def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, st target_keys, domain_keys, threshold, + row_estimates, + selective_nodes, ) ) candidates.extend( @@ -190,6 +211,8 @@ def _select_candidate(ir: Join, threshold: float) -> tuple[_Candidate | None, st target_keys, domain_keys, threshold, + row_estimates, + selective_nodes, ) ) @@ -209,16 +232,22 @@ def _simple_candidates( target_keys: tuple[expr.Col, ...], domain_keys: tuple[expr.Col, ...], threshold: float, + row_estimates: dict[IR, int | None], + selective_nodes: set[IR], ) -> Iterable[_Candidate]: for target_key, domain_key in zip(target_keys, domain_keys, strict=True): - target = _largest_key_source(target_child, target_key.name) + target = _largest_key_source(target_child, target_key.name, row_estimates) if target is None: continue - target_rows = _estimate_rows(target) + target_rows = row_estimates.get(target) if target_rows is None or target_rows <= 0: continue domain = _smallest_key_producer( - domain_child, domain_key.name, require_selective=True + domain_child, + domain_key.name, + row_estimates, + selective_nodes, + require_selective=True, ) if domain is None: continue @@ -233,6 +262,7 @@ def _simple_candidates( target_key=target_key, domain=domain, domain_key=domain_key, + target_rows=target_rows, ) @@ -243,6 +273,8 @@ def _composite_candidates( target_keys: tuple[expr.Col, ...], domain_keys: tuple[expr.Col, ...], threshold: float, + row_estimates: dict[IR, int | None], + selective_nodes: set[IR], ) -> Iterable[_Candidate]: if len(target_keys) < 2: return @@ -250,10 +282,10 @@ def _composite_candidates( for filter_index, (target_key, domain_key) in enumerate( zip(target_keys, domain_keys, strict=True) ): - target = _largest_key_source(target_child, target_key.name) + target = _largest_key_source(target_child, target_key.name, row_estimates) if target is None: continue - target_rows = _estimate_rows(target) + target_rows = row_estimates.get(target) if target_rows is None or target_rows <= 0: continue @@ -264,13 +296,17 @@ def _composite_candidates( if constraint_index == filter_index: continue domain = _smallest_node_containing_all( - domain_child, (domain_key.name, domain_constraint_key.name) + domain_child, + (domain_key.name, domain_constraint_key.name), + row_estimates, ) if domain is None: continue constraint_domain = _smallest_key_producer( target_child, target_constraint_key.name, + row_estimates, + selective_nodes, require_selective=True, exclude=target, ) @@ -291,6 +327,7 @@ def _composite_candidates( target_key=target_key, domain=domain, domain_key=domain_key, + target_rows=target_rows, constraint_domain=constraint_domain, domain_constraint_key=domain_constraint_key, target_constraint_key=target_constraint_key, @@ -373,16 +410,22 @@ def _make_semi_join( def _smallest_key_producer( - root: IR, column: str, *, require_selective: bool, exclude: IR | None = None + root: IR, + column: str, + row_estimates: dict[IR, int | None], + selective_nodes: set[IR], + *, + require_selective: bool, + exclude: IR | None = None, ) -> _Producer | None: candidates = [] for node in traversal([root]): if node is exclude or column not in node.schema: continue - rows = _estimate_rows(node) + rows = row_estimates.get(node) if rows is None or rows <= 0: continue - if require_selective and not _is_selective(node): + if require_selective and node not in selective_nodes: continue candidates.append((rows, len(node.schema), _Producer(node, column, rows))) if not candidates: @@ -390,13 +433,15 @@ def _smallest_key_producer( return min(candidates, key=lambda item: (item[0], item[1]))[2] -def _smallest_node_containing_all(root: IR, columns: Sequence[str]) -> _Producer | None: +def _smallest_node_containing_all( + root: IR, columns: Sequence[str], row_estimates: dict[IR, int | None] +) -> _Producer | None: candidates = [] needed = set(columns) for node in traversal([root]): if not needed.issubset(node.schema): continue - rows = _estimate_rows(node) + rows = row_estimates.get(node) if rows is None or rows <= 0: continue candidates.append((rows, len(node.schema), _Producer(node, columns[0], rows))) @@ -405,13 +450,15 @@ def _smallest_node_containing_all(root: IR, columns: Sequence[str]) -> _Producer return min(candidates, key=lambda item: (item[0], item[1]))[2] -def _largest_key_source(root: IR, column: str) -> IR | None: +def _largest_key_source( + root: IR, column: str, row_estimates: dict[IR, int | None] +) -> IR | None: source_candidates = [] fallback_candidates = [] for node in traversal([root]): if column not in node.schema: continue - rows = _estimate_rows(node) + rows = row_estimates.get(node) if rows is None or rows <= 0: continue item = (rows, len(node.schema), node) @@ -425,32 +472,33 @@ def _largest_key_source(root: IR, column: str) -> IR | None: return max(candidates, key=lambda item: (item[0], -item[1]))[2] -def _estimate_rows(ir: IR) -> int | None: - try: - return _ROW_ESTIMATES[ir] - except KeyError: - pass - - rows: int | None - if isinstance(ir, (Scan, DataFrameScan)): - source = None if _STATS is None else _STATS.scan_stats.get(ir) - rows = None if source is None else source.row_count - if rows is None and isinstance(ir, DataFrameScan): - rows = ir.df.shape()[0] - elif isinstance(ir, (Select, Projection, HStack, Cache, Filter, Distinct, GroupBy)): - rows = _estimate_rows(ir.children[0]) - elif isinstance(ir, Join): - left_rows = _estimate_rows(ir.children[0]) - right_rows = _estimate_rows(ir.children[1]) - rows = _estimate_join_rows(ir.options[0], left_rows, right_rows) - else: - estimates = [ - estimate for child in ir.children if (estimate := _estimate_rows(child)) - ] - rows = max(estimates) if estimates else None - - _ROW_ESTIMATES[ir] = rows - return rows +def _estimate_row_counts(ir: IR, stats: StatsCollector) -> dict[IR, int | None]: + estimates: dict[IR, int | None] = {} + for node in post_traversal([ir]): + if isinstance(node, (Scan, DataFrameScan)): + source = stats.scan_stats.get(node) + rows = None if source is None else source.row_count + if rows is None and isinstance(node, DataFrameScan): + rows = node.df.shape()[0] + elif isinstance( + node, (Select, Projection, HStack, Cache, Filter, Distinct, GroupBy) + ): + rows = estimates[node.children[0]] + elif isinstance(node, Join): + rows = _estimate_join_rows( + node.options[0], + estimates[node.children[0]], + estimates[node.children[1]], + ) + else: + child_estimates = [ + estimate + for child in node.children + if (estimate := estimates[child]) is not None + ] + rows = max(child_estimates) if child_estimates else None + estimates[node] = rows + return estimates def _estimate_join_rows( @@ -471,20 +519,15 @@ def _estimate_join_rows( return None -def _is_selective(ir: IR) -> bool: - try: - return _SELECTIVE[ir] - except KeyError: - pass - - if isinstance(ir, Scan): - selective = ir.predicate is not None - elif isinstance(ir, Filter): - selective = True - else: - selective = any(_is_selective(child) for child in ir.children) - - _SELECTIVE[ir] = selective +def _collect_selective_nodes(ir: IR) -> set[IR]: + selective: set[IR] = set() + for node in post_traversal([ir]): + if ( + (isinstance(node, Scan) and node.predicate is not None) + or isinstance(node, Filter) + or any(child in selective for child in node.children) + ): + selective.add(node) return selective @@ -492,17 +535,6 @@ def _contains_identity(root: IR, needle: IR) -> bool: return any(node is needle for node in traversal([root])) -def _replace_identity(root: IR, old: IR, new: IR) -> IR: - if root is old: - return new - if not root.children: - return root - children = tuple(_replace_identity(child, old, new) for child in root.children) - if children == root.children: - return root - return root.reconstruct(children) - - def _trace_decision( ir: Join, threshold: float, candidate: _Candidate | None, reason: str ) -> None: From fe17507fb00c4e10077a61fd74542aa7ee0caffb Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 08:12:15 +0000 Subject: [PATCH 23/44] Restore optional float config conversion Restore the optional float environment converter used by the join-domain prefilter threshold after the join-prefilter branch merge removed the legacy definition. This keeps the documented numeric and null environment values valid and restores static-checking correctness. --- python/cudf_polars/cudf_polars/utils/config.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 991fe018f731..b987fa08f036 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -174,6 +174,10 @@ 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) From 1c0e3df13b20999bf9989a761bcb135ed1d7c1e0 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 10:04:08 +0000 Subject: [PATCH 24/44] Reanalyze rewritten join-domain subtrees Refresh row-count and selectivity analysis when bottom-up rewriting reconstructs a join subtree. This lets parent joins rank domains using newly inserted semi joins, preserving derived-filter propagation for Q5-like plans and avoiding harmful stacked filters for Q8-like plans. Add regression coverage for both behaviors. --- .../streaming/join_domain_prefilter.py | 16 +++++- .../streaming/test_join_domain_prefilter.py | 51 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index d506bd3f6162..e1dc0637453a 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -88,6 +88,7 @@ class _RewriteState(TypedDict): threshold: float trace: bool + stats: StatsCollector row_estimates: dict[IR, int | None] selective_nodes: set[IR] @@ -115,6 +116,7 @@ def optimize_join_domain_prefilters( state = _RewriteState( threshold=threshold, trace=trace, + stats=stats, row_estimates=_estimate_row_counts(ir, stats), selective_nodes=_collect_selective_nodes(ir), ) @@ -136,14 +138,24 @@ def _(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: @_rewrite.register(Join) def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + original = node rewritten = reuse_if_unchanged(node, rec) assert isinstance(rewritten, Join) node = rewritten + if node is original: + row_estimates = rec.state["row_estimates"] + selective_nodes = rec.state["selective_nodes"] + else: + # Child rewrites introduce new semi joins and reconstructed ancestors. + # Re-analyze that current subtree so parent joins can use the derived + # selectivity and cardinality when ranking their own candidates. + row_estimates = _estimate_row_counts(node, rec.state["stats"]) + selective_nodes = _collect_selective_nodes(node) candidate, reason = _select_candidate( node, rec.state["threshold"], - rec.state["row_estimates"], - rec.state["selective_nodes"], + row_estimates, + selective_nodes, ) if rec.state["trace"]: _trace_decision(node, rec.state["threshold"], candidate, reason) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 44ebda768b8e..155199c26632 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -194,6 +194,57 @@ def test_composite_domain_prefilter_constrains_domain_first() -> None: assert any(semi.children[0] is lineitem for semi in semis) +def test_derived_selectivity_propagates_through_rewritten_children() -> None: + region = _scan("region", ("r_regionkey",), predicate=True) + nation = _scan("nation", ("n_nationkey", "n_regionkey")) + customer = _scan("customer", ("c_custkey", "c_nationkey")) + orders = _scan("orders", ("o_orderkey", "o_custkey")) + + region_nation = _join(region, nation, ("r_regionkey",), ("n_regionkey",)) + nation_customer = _join(region_nation, customer, ("n_nationkey",), ("c_nationkey",)) + root = _join(nation_customer, orders, ("c_custkey",), ("o_custkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats( + region=(region, 1), + nation=(nation, 25), + customer=(customer, 150), + orders=(orders, 1_500), + ), + _config(), + ) + + filtered = {semi.children[0] for semi in _joins(optimized, "Semi")} + assert {nation, customer, orders} <= filtered + + +def test_rewritten_domain_filters_other_side_instead_of_stacking() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + lineitem = _scan("lineitem", ("l_orderkey", "l_partkey", "l_suppkey")) + supplier = _scan("supplier", ("s_suppkey",)) + orders = _scan("orders", ("o_orderkey",), predicate=True) + + part_lineitem = _join(part, lineitem, ("p_partkey",), ("l_partkey",)) + line_supplier = _join(part_lineitem, supplier, ("l_suppkey",), ("s_suppkey",)) + root = _join(line_supplier, orders, ("l_orderkey",), ("o_orderkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats( + part=(part, 60), + lineitem=(lineitem, 1_800), + supplier=(supplier, 30), + orders=(orders, 150), + ), + _config(), + ) + + semis = _joins(optimized, "Semi") + assert sum(semi.children[0] is lineitem for semi in semis) == 1 + assert any(semi.children[0] is orders for semi in semis) + + def test_no_domain_prefilter_for_outer_join() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",)) From 702c6adac3bb408ba4cc9e5993962ee9faabe33b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 12:18:45 +0000 Subject: [PATCH 25/44] 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 26/44] 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 27/44] 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 28/44] 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 29/44] 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: From b3049c29b99fa0ddee6733fbdac5edf5bf0cfdd4 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 29 Jun 2026 13:13:26 -0700 Subject: [PATCH 30/44] Add missing coverage for JOIN_DOMAIN_PREFILTER_TRACE --- python/cudf_polars/tests/test_config.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 246b8877f117..08f836a4573e 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -624,6 +624,20 @@ def test_join_domain_prefilter_options_from_env( assert config.executor.dynamic_planning.join_domain_prefilter_trace +@pytest.mark.parametrize("value", ["none", "null"]) +def test_join_domain_prefilter_trace_inherits_from_env( + monkeypatch: pytest.MonkeyPatch, value: str +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_TRACE", "1" + ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_TRACE", + value, + ) + assert DynamicPlanningOptions().join_domain_prefilter_trace + + @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 From 0b93f55e4b2e88512c315abe05b681b881355e15 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 30 Jun 2026 20:10:10 +0000 Subject: [PATCH 31/44] Simplify join-domain candidate selection Remove an unreachable key-count skip, express row-estimate fallback through max's default, and iterate over explicit left/right side descriptors when collecting prefilter candidates. --- .../streaming/join_domain_prefilter.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index e1dc0637453a..b78b1f18c80e 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -188,21 +188,23 @@ def _select_candidate( right_keys = _simple_keys(ir.right_on) if len(left_keys) != len(ir.left_on) or len(right_keys) != len(ir.right_on): return None, "non_column_join_key" - if len(left_keys) != len(right_keys): - return None, "key_count_mismatch" candidates: list[_Candidate] = [] - for target_side in ("left", "right"): - target_child, domain_child = ( - (ir.children[0], ir.children[1]) - if target_side == "left" - else (ir.children[1], ir.children[0]) - ) - target_keys, domain_keys = ( - (left_keys, right_keys) - if target_side == "left" - else (right_keys, left_keys) - ) + left: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( + "left", + ir.children[0], + left_keys, + ) + right: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( + "right", + ir.children[1], + right_keys, + ) + for (target_side, target_child, target_keys), ( + _, + domain_child, + domain_keys, + ) in ((left, right), (right, left)): candidates.extend( _composite_candidates( target_side, @@ -508,7 +510,7 @@ def _estimate_row_counts(ir: IR, stats: StatsCollector) -> dict[IR, int | None]: for child in node.children if (estimate := estimates[child]) is not None ] - rows = max(child_estimates) if child_estimates else None + rows = max(child_estimates, default=None) estimates[node] = rows return estimates From 0226a831c3898a2280b8d7a7ac07844d22c2eeb5 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 30 Jun 2026 20:13:48 +0000 Subject: [PATCH 32/44] Track column bindings through domain subplans Resolve join keys through proven output-to-input bindings instead of searching descendant schemas by name. Carry the bound source names into simple and composite prefilters, stop at ambiguous transformations, and cover target, domain, and composite renames with regression tests. --- .../streaming/join_domain_prefilter.py | 181 ++++++++++++++---- .../streaming/test_join_domain_prefilter.py | 90 ++++++++- 2 files changed, 230 insertions(+), 41 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index b78b1f18c80e..7476f72f0611 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -21,6 +21,7 @@ Projection, Scan, Select, + Sort, ) from cudf_polars.dsl.tracing import Scope, log from cudf_polars.dsl.traversal import ( @@ -41,12 +42,17 @@ @dataclass(frozen=True) class _Producer: - """A subtree that can provide a key domain.""" + """A subtree and its bound column names at an insertion point.""" node: IR - column: str + columns: tuple[str, ...] rows: int + @property + def column(self) -> str: + """First bound column in the producer.""" + return self.columns[0] + @dataclass(frozen=True) class _Candidate: @@ -54,7 +60,7 @@ class _Candidate: mode: Literal["simple", "composite"] target_side: Literal["left", "right"] - target: IR + target: _Producer target_key: expr.Col domain: _Producer domain_key: expr.Col @@ -165,9 +171,9 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: left, right = node.children target_filter = _make_target_filter(node, candidate) if candidate.target_side == "left": - (left,) = replace([left], {candidate.target: target_filter}) + (left,) = replace([left], {candidate.target.node: target_filter}) else: - (right,) = replace([right], {candidate.target: target_filter}) + (right,) = replace([right], {candidate.target.node: target_filter}) return node.reconstruct((left, right)) @@ -253,9 +259,6 @@ def _simple_candidates( target = _largest_key_source(target_child, target_key.name, row_estimates) if target is None: continue - target_rows = row_estimates.get(target) - if target_rows is None or target_rows <= 0: - continue domain = _smallest_key_producer( domain_child, domain_key.name, @@ -265,9 +268,9 @@ def _simple_candidates( ) if domain is None: continue - if _contains_identity(target, domain.node): + if _contains_identity(target.node, domain.node): continue - if domain.rows / target_rows > threshold: + if domain.rows / target.rows > threshold: continue yield _Candidate( mode="simple", @@ -276,7 +279,7 @@ def _simple_candidates( target_key=target_key, domain=domain, domain_key=domain_key, - target_rows=target_rows, + target_rows=target.rows, ) @@ -299,9 +302,6 @@ def _composite_candidates( target = _largest_key_source(target_child, target_key.name, row_estimates) if target is None: continue - target_rows = row_estimates.get(target) - if target_rows is None or target_rows <= 0: - continue for constraint_index, ( target_constraint_key, @@ -322,15 +322,15 @@ def _composite_candidates( row_estimates, selective_nodes, require_selective=True, - exclude=target, + exclude=target.node, ) if constraint_domain is None: continue - if _contains_identity(target, domain.node) or _contains_identity( - target, constraint_domain.node + if _contains_identity(target.node, domain.node) or _contains_identity( + target.node, constraint_domain.node ): continue - if domain.rows / target_rows > threshold: + if domain.rows / target.rows > threshold: continue if constraint_domain.rows / domain.rows > threshold: continue @@ -341,7 +341,7 @@ def _composite_candidates( target_key=target_key, domain=domain, domain_key=domain_key, - target_rows=target_rows, + target_rows=target.rows, constraint_domain=constraint_domain, domain_constraint_key=domain_constraint_key, target_constraint_key=target_constraint_key, @@ -350,9 +350,10 @@ def _composite_candidates( def _make_target_filter(ir: Join, candidate: _Candidate) -> Join: domain = _make_domain(candidate, ir) + target = candidate.target return _make_semi_join( - candidate.target, - candidate.target_key, + target.node, + expr.Col(target.node.schema[target.column], target.column), domain, expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), nulls_equal=ir.options[1], @@ -380,8 +381,8 @@ def _make_domain(candidate: _Candidate, ir: Join) -> IR: constrained = _make_semi_join( candidate.domain.node, expr.Col( - candidate.domain.node.schema[candidate.domain_constraint_key.name], - candidate.domain_constraint_key.name, + candidate.domain.node.schema[candidate.domain.columns[1]], + candidate.domain.columns[1], ), constraint_domain, expr.Col( @@ -433,15 +434,17 @@ def _smallest_key_producer( exclude: IR | None = None, ) -> _Producer | None: candidates = [] - for node in traversal([root]): - if node is exclude or column not in node.schema: + for node, bound_column in _column_bindings(root, column): + if node is exclude: continue rows = row_estimates.get(node) if rows is None or rows <= 0: continue if require_selective and node not in selective_nodes: continue - candidates.append((rows, len(node.schema), _Producer(node, column, rows))) + candidates.append( + (rows, len(node.schema), _Producer(node, (bound_column,), rows)) + ) if not candidates: return None return min(candidates, key=lambda item: (item[0], item[1]))[2] @@ -451,14 +454,30 @@ def _smallest_node_containing_all( root: IR, columns: Sequence[str], row_estimates: dict[IR, int | None] ) -> _Producer | None: candidates = [] - needed = set(columns) - for node in traversal([root]): - if not needed.issubset(node.schema): - continue - rows = row_estimates.get(node) - if rows is None or rows <= 0: - continue - candidates.append((rows, len(node.schema), _Producer(node, columns[0], rows))) + lineages = [tuple(_column_bindings(root, column)) for column in columns] + if not lineages or any(not lineage for lineage in lineages): + return None + for node, first_column in lineages[0]: + bound_columns = [first_column] + for lineage in lineages[1:]: + match = next( + (bound_column for candidate, bound_column in lineage if candidate is node), + None, + ) + if match is None: + break + bound_columns.append(match) + else: + rows = row_estimates.get(node) + if rows is None or rows <= 0: + continue + candidates.append( + ( + rows, + len(node.schema), + _Producer(node, tuple(bound_columns), rows), + ) + ) if not candidates: return None return min(candidates, key=lambda item: (item[0], item[1]))[2] @@ -466,16 +485,14 @@ def _smallest_node_containing_all( def _largest_key_source( root: IR, column: str, row_estimates: dict[IR, int | None] -) -> IR | None: +) -> _Producer | None: source_candidates = [] fallback_candidates = [] - for node in traversal([root]): - if column not in node.schema: - continue + for node, bound_column in _column_bindings(root, column): rows = row_estimates.get(node) if rows is None or rows <= 0: continue - item = (rows, len(node.schema), node) + item = (rows, len(node.schema), _Producer(node, (bound_column,), rows)) if isinstance(node, (Scan, DataFrameScan)): source_candidates.append(item) else: @@ -486,6 +503,90 @@ def _largest_key_source( return max(candidates, key=lambda item: (item[0], -item[1]))[2] +def _column_bindings(root: IR, column: str) -> Iterable[tuple[IR, str]]: + """Yield exact output-to-input bindings for a column through a subplan.""" + node = root + while column in node.schema: + yield node, column + binding = _input_binding(node, column) + if binding is None: + return + node, column = binding + + +def _input_binding(node: IR, column: str) -> tuple[IR, str] | None: + """Return a proven direct input binding, stopping at ambiguous operations.""" + child = node.children[0] if len(node.children) == 1 else None + if isinstance(node, Select): + selected = next((item for item in node.exprs if item.name == column), None) + return _column_expression_binding(child, selected) + if isinstance(node, HStack): + stacked = next((item for item in node.columns if item.name == column), None) + if stacked is not None: + return _column_expression_binding(child, stacked) + return _passthrough_binding(child, column) + if isinstance(node, GroupBy): + if node.zlice is not None: + return None + key = next((item for item in node.keys if item.name == column), None) + return _column_expression_binding(child, key) + if isinstance(node, Join): + return _join_input_binding(node, column) + if isinstance(node, Distinct): + return ( + None + if node.zlice is not None + else _passthrough_binding(child, column) + ) + if isinstance(node, Sort): + return ( + None + if node.zlice is not None + else _passthrough_binding(child, column) + ) + if isinstance(node, (Cache, Filter, Projection)): + return _passthrough_binding(child, column) + return None + + +def _column_expression_binding( + child: IR | None, expression: expr.NamedExpr | None +) -> tuple[IR, str] | None: + if ( + child is not None + and expression is not None + and isinstance(expression.value, expr.Col) + and expression.value.name in child.schema + ): + return child, expression.value.name + return None + + +def _passthrough_binding(child: IR | None, column: str) -> tuple[IR, str] | None: + if child is not None and column in child.schema: + return child, column + return None + + +def _join_input_binding(node: Join, column: str) -> tuple[IR, str] | None: + if node.options[0] != "Inner" or node.options[2] is not None: + return None + left, right = node.children + bindings = [] + if column in left.schema: + bindings.append((left, column)) + suffix = node.options[3] + for right_column in right.schema: + output_column = ( + f"{right_column}{suffix}" if right_column in left.schema else right_column + ) + if output_column == column and output_column in node.schema: + bindings.append((right, right_column)) + if len(bindings) == 1: + return bindings[0] + return None + + def _estimate_row_counts(ir: IR, stats: StatsCollector) -> dict[IR, int | None]: estimates: dict[IR, int | None] = {} for node in post_traversal([ir]): @@ -572,7 +673,7 @@ def _trace_decision( "domain_key": candidate.domain_key.name, "estimated_target_rows": candidate.target_rows, "estimated_domain_rows": candidate.domain_rows, - "target_node_type": type(candidate.target).__name__, + "target_node_type": type(candidate.target.node).__name__, "domain_node_type": type(candidate.domain.node).__name__, } ) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 155199c26632..ead347910af9 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -9,10 +9,11 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr -from cudf_polars.dsl.ir import Join, Scan +from cudf_polars.dsl.ir import Join, Scan, Select from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.join_domain_prefilter import ( + _smallest_node_containing_all, optimize_join_domain_prefilters, ) from cudf_polars.utils.config import ConfigOptions, ParquetOptions @@ -70,6 +71,19 @@ def _key(node: IR, name: str) -> expr.NamedExpr: return expr.NamedExpr(name, expr.Col(node.schema[name], name)) +def _select(node: IR, **columns: str) -> Select: + schema = {output: node.schema[source] for output, source in columns.items()} + return Select( + schema, + tuple( + expr.NamedExpr(output, expr.Col(schema[output], source)) + for output, source in columns.items() + ), + True, # noqa: FBT003 + node, + ) + + def _join( left: IR, right: IR, @@ -245,6 +259,80 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking() -> None: assert any(semi.children[0] is orders for semi in semis) +def test_target_source_follows_join_key_through_rename() -> None: + big = _scan("big", ("left_key", "other")) + renamed_big = _select(big, foo="left_key", other="other") + small = _scan("small", ("left_key", "other2")) + joined = _join( + renamed_big, + small, + ("other",), + ("other2",), + maintain_order="left", + ) + domain = _scan("domain", ("domain_key",), predicate=True) + root = _join(joined, domain, ("left_key",), ("domain_key",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(big=(big, 1_000), small=(small, 100), domain=(domain, 5)), + _config(), + ) + + semis = _joins(optimized, "Semi") + assert any(semi.children[0] is small for semi in semis) + assert not any(semi.children[0] is big for semi in semis) + + +def test_domain_source_follows_join_key_through_rename() -> None: + target = _scan("target", ("target_key",)) + unrelated = _scan("unrelated", ("domain_key", "other"), predicate=True) + renamed_unrelated = _select(unrelated, foo="domain_key", other="other") + domain_source = _scan("domain_source", ("domain_key", "other2"), predicate=True) + domain = _join( + renamed_unrelated, + domain_source, + ("other",), + ("other2",), + maintain_order="left", + ) + root = _join(target, domain, ("target_key",), ("domain_key",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats( + target=(target, 1_000), + unrelated=(unrelated, 1), + domain_source=(domain_source, 5), + ), + _config(), + ) + + semi = next(semi for semi in _joins(optimized, "Semi") if semi.children[0] is target) + selected_domain = semi.children[1] + assert isinstance(selected_domain, Select) + assert selected_domain.children[0] is domain + + +def test_composite_domain_columns_follow_renames() -> None: + source = _scan("source", ("raw_key", "raw_constraint")) + renamed = _select( + source, + domain_key="raw_key", + domain_constraint="raw_constraint", + ) + + producer = _smallest_node_containing_all( + renamed, + ("domain_key", "domain_constraint"), + {renamed: 20, source: 10}, + ) + + assert producer is not None + assert producer.node is source + assert producer.columns == ("raw_key", "raw_constraint") + + def test_no_domain_prefilter_for_outer_join() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",)) From 8bd90237041bfca4146d18876029ea56c171dfb1 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 30 Jun 2026 20:14:28 +0000 Subject: [PATCH 33/44] Keep domain prefilter replacement side-scoped Document why target replacement starts from the selected join child rather than the join root, and cover a DAG-shared target to ensure the domain side remains unchanged. --- .../streaming/join_domain_prefilter.py | 2 ++ .../streaming/test_join_domain_prefilter.py | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 7476f72f0611..947ae0fa7d46 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -170,6 +170,8 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: left, right = node.children target_filter = _make_target_filter(node, candidate) + # A DAG may share the target with the domain side, so only rewrite the + # side for which this candidate was selected. if candidate.target_side == "left": (left,) = replace([left], {candidate.target.node: target_filter}) else: diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index ead347910af9..6d9cc029e677 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -333,6 +333,34 @@ def test_composite_domain_columns_follow_renames() -> None: assert producer.columns == ("raw_key", "raw_constraint") +def test_target_replacement_does_not_rewrite_shared_domain_side() -> None: + shared = _scan("shared", ("target_key", "other")) + domain_source = _scan( + "domain_source", ("domain_key", "other2"), predicate=True + ) + domain = _join( + shared, + domain_source, + ("other",), + ("other2",), + maintain_order="left", + ) + root = _join(shared, domain, ("target_key",), ("domain_key",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(shared=(shared, 1_000), domain_source=(domain_source, 5)), + _config(), + ) + + assert isinstance(optimized, Join) + assert isinstance(optimized.children[0], Join) + assert optimized.children[0].options[0] == "Semi" + assert optimized.children[0].children[0] is shared + assert optimized.children[1] is domain + assert domain.children[0] is shared + + def test_no_domain_prefilter_for_outer_join() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",)) From 9511c94107fb1472f486d3dbd35403cc24017453 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 30 Jun 2026 20:15:23 +0000 Subject: [PATCH 34/44] Format join-domain review updates --- .../streaming/join_domain_prefilter.py | 18 +++++++----------- .../streaming/test_join_domain_prefilter.py | 8 ++++---- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 947ae0fa7d46..fed92717f31f 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -463,7 +463,11 @@ def _smallest_node_containing_all( bound_columns = [first_column] for lineage in lineages[1:]: match = next( - (bound_column for candidate, bound_column in lineage if candidate is node), + ( + bound_column + for candidate, bound_column in lineage + if candidate is node + ), None, ) if match is None: @@ -535,17 +539,9 @@ def _input_binding(node: IR, column: str) -> tuple[IR, str] | None: if isinstance(node, Join): return _join_input_binding(node, column) if isinstance(node, Distinct): - return ( - None - if node.zlice is not None - else _passthrough_binding(child, column) - ) + return None if node.zlice is not None else _passthrough_binding(child, column) if isinstance(node, Sort): - return ( - None - if node.zlice is not None - else _passthrough_binding(child, column) - ) + return None if node.zlice is not None else _passthrough_binding(child, column) if isinstance(node, (Cache, Filter, Projection)): return _passthrough_binding(child, column) return None diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 6d9cc029e677..01f870651658 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -308,7 +308,9 @@ def test_domain_source_follows_join_key_through_rename() -> None: _config(), ) - semi = next(semi for semi in _joins(optimized, "Semi") if semi.children[0] is target) + semi = next( + semi for semi in _joins(optimized, "Semi") if semi.children[0] is target + ) selected_domain = semi.children[1] assert isinstance(selected_domain, Select) assert selected_domain.children[0] is domain @@ -335,9 +337,7 @@ def test_composite_domain_columns_follow_renames() -> None: def test_target_replacement_does_not_rewrite_shared_domain_side() -> None: shared = _scan("shared", ("target_key", "other")) - domain_source = _scan( - "domain_source", ("domain_key", "other2"), predicate=True - ) + domain_source = _scan("domain_source", ("domain_key", "other2"), predicate=True) domain = _join( shared, domain_source, From ae443dd986d18bba5731ffe689b50afe92b484db Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 30 Jun 2026 20:17:12 +0000 Subject: [PATCH 35/44] Trace bindings through filtering joins Follow output columns through the left input of semi and anti joins so later domain-prefilter decisions can still reach original producers. Strengthen the derived-domain regression to reject stacked semi filters. --- .../cudf_polars/streaming/join_domain_prefilter.py | 6 +++++- .../tests/streaming/test_join_domain_prefilter.py | 4 ++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index fed92717f31f..175d58e94455 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -567,9 +567,13 @@ def _passthrough_binding(child: IR | None, column: str) -> tuple[IR, str] | None def _join_input_binding(node: Join, column: str) -> tuple[IR, str] | None: - if node.options[0] != "Inner" or node.options[2] is not None: + if node.options[2] is not None: return None left, right = node.children + if node.options[0] in ("Semi", "Anti"): + return _passthrough_binding(left, column) + if node.options[0] != "Inner": + return None bindings = [] if column in left.schema: bindings.append((left, column)) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 01f870651658..2d339613f5b5 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -257,6 +257,10 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking() -> None: semis = _joins(optimized, "Semi") assert sum(semi.children[0] is lineitem for semi in semis) == 1 assert any(semi.children[0] is orders for semi in semis) + assert not any( + isinstance(semi.children[0], Join) and semi.children[0].options[0] == "Semi" + for semi in semis + ) def test_target_source_follows_join_key_through_rename() -> None: From 571181bc3ef3127c5eaefa706d51de936214734a Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 Jul 2026 18:51:18 +0000 Subject: [PATCH 36/44] Check cheap domain guards first Evaluate local row-ratio guards before identity checks and additional producer searches so unprofitable candidates avoid unnecessary subgraph traversals. --- .../cudf_polars/streaming/join_domain_prefilter.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 175d58e94455..9bbbfd060cd6 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -270,10 +270,10 @@ def _simple_candidates( ) if domain is None: continue - if _contains_identity(target.node, domain.node): - continue if domain.rows / target.rows > threshold: continue + if _contains_identity(target.node, domain.node): + continue yield _Candidate( mode="simple", target_side=target_side, @@ -318,6 +318,8 @@ def _composite_candidates( ) if domain is None: continue + if domain.rows / target.rows > threshold: + continue constraint_domain = _smallest_key_producer( target_child, target_constraint_key.name, @@ -328,14 +330,12 @@ def _composite_candidates( ) if constraint_domain is None: continue + if constraint_domain.rows / domain.rows > threshold: + continue if _contains_identity(target.node, domain.node) or _contains_identity( target.node, constraint_domain.node ): continue - if domain.rows / target.rows > threshold: - continue - if constraint_domain.rows / domain.rows > threshold: - continue yield _Candidate( mode="composite", target_side=target_side, From cf59505b7b6b335cd586bb6959dce3e47913b818 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 Jul 2026 18:51:56 +0000 Subject: [PATCH 37/44] Inline target prefilter construction Build the selected domain and target semi join directly in the join rewrite, removing the single-use _make_target_filter wrapper. --- .../streaming/join_domain_prefilter.py | 24 ++++++++----------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 9bbbfd060cd6..6da3a9b149d3 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -169,7 +169,16 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: return node left, right = node.children - target_filter = _make_target_filter(node, candidate) + domain = _make_domain(candidate, node) + target = candidate.target + target_filter = _make_semi_join( + target.node, + expr.Col(target.node.schema[target.column], target.column), + domain, + expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), + nulls_equal=node.options[1], + suffix=node.options[3], + ) # A DAG may share the target with the domain side, so only rewrite the # side for which this candidate was selected. if candidate.target_side == "left": @@ -350,19 +359,6 @@ def _composite_candidates( ) -def _make_target_filter(ir: Join, candidate: _Candidate) -> Join: - domain = _make_domain(candidate, ir) - target = candidate.target - return _make_semi_join( - target.node, - expr.Col(target.node.schema[target.column], target.column), - domain, - expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), - nulls_equal=ir.options[1], - suffix=ir.options[3], - ) - - def _make_domain(candidate: _Candidate, ir: Join) -> IR: if candidate.mode == "simple": return _select_key( From 3a57f1435b86c7cde836ce21770e3f1a561c5a41 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 Jul 2026 18:53:14 +0000 Subject: [PATCH 38/44] Clarify bound key projection Rename the key projection helper around its binding-aware purpose, pass the join-visible column expression explicitly, and assert that its dtype matches the producer-bound source column. --- .../streaming/join_domain_prefilter.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 6da3a9b149d3..cdce3ccd2eb8 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -361,20 +361,20 @@ def _composite_candidates( def _make_domain(candidate: _Candidate, ir: Join) -> IR: if candidate.mode == "simple": - return _select_key( + return _project_bound_key( candidate.domain.node, candidate.domain.column, - candidate.domain_key.name, + candidate.domain_key, ) assert candidate.constraint_domain is not None assert candidate.domain_constraint_key is not None assert candidate.target_constraint_key is not None - constraint_domain = _select_key( + constraint_domain = _project_bound_key( candidate.constraint_domain.node, candidate.constraint_domain.column, - candidate.target_constraint_key.name, + candidate.target_constraint_key, ) constrained = _make_semi_join( candidate.domain.node, @@ -390,14 +390,18 @@ def _make_domain(candidate: _Candidate, ir: Join) -> IR: nulls_equal=ir.options[1], suffix=ir.options[3], ) - return _select_key(constrained, candidate.domain.column, candidate.domain_key.name) + return _project_bound_key( + constrained, candidate.domain.column, candidate.domain_key + ) -def _select_key(source: IR, source_column: str, output_column: str) -> Select: - dtype = source.schema[source_column] +def _project_bound_key(source: IR, bound_column: str, output_key: expr.Col) -> Select: + """Project a bound source column under its join-visible key name.""" + dtype = source.schema[bound_column] + assert dtype == output_key.dtype return Select( - {output_column: dtype}, - (expr.NamedExpr(output_column, expr.Col(dtype, source_column)),), + {output_key.name: dtype}, + (expr.NamedExpr(output_key.name, expr.Col(dtype, bound_column)),), True, # noqa: FBT003 source, ) From 4261279ddad21baf286a5b9527df5d6cb46b17aa Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 Jul 2026 19:08:04 +0000 Subject: [PATCH 39/44] Decouple join-domain prefilter configuration Move the logical domain-prefilter controls into a dedicated executor option group and run the rewrite independently of dynamic shuffle planning. This keeps static planning eligible for the same logical row reduction and gives the rewrite its own environment-variable namespace. --- .../cudf_polars/cudf_polars/engine/options.py | 10 ++ .../streaming/join_domain_prefilter.py | 10 +- .../cudf_polars/streaming/parallel.py | 9 +- .../cudf_polars/cudf_polars/utils/config.py | 131 +++++++++--------- .../streaming/test_join_domain_prefilter.py | 28 +++- python/cudf_polars/tests/test_config.py | 82 +++++------ 6 files changed, 139 insertions(+), 131 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 559c05be8c94..35d838c42efa 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -25,6 +25,7 @@ from cudf_polars.utils.config import ( DynamicPlanningOptions, + JoinDomainPrefilterOptions, ParquetOptions, ) @@ -247,6 +248,12 @@ class StreamingOptions: Env: ``CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING``. Default: enabled. Category: executor. + join_domain_prefilter + Join-domain prefilter config, dict or + :class:`~cudf_polars.utils.config.JoinDomainPrefilterOptions`. + Env: ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__*``. + Default: enabled. + Category: executor. sink_to_directory Whether multi-partition sink operations should write to a directory rather than a single file. The ``spmd``/``ray``/``dask`` engines @@ -341,6 +348,9 @@ class StreamingOptions: dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = ( _opt("executor") ) + join_domain_prefilter: dict[str, Any] | JoinDomainPrefilterOptions | Unspecified = ( + _opt("executor") + ) sink_to_directory: bool | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__SINK_TO_DIRECTORY", parse_boolean ) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index cdce3ccd2eb8..310dcbcde071 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -111,12 +111,12 @@ def optimize_join_domain_prefilters( column equality keys are considered, and the original full join remains after every inserted row-reduction semi join. """ - dynamic_options = config_options.executor.dynamic_planning - if dynamic_options is None or not dynamic_options.join_domain_prefilter_enabled: + options = config_options.executor.join_domain_prefilter + if not options.enabled: return ir - threshold = dynamic_options.join_domain_prefilter_threshold - trace = dynamic_options.join_domain_prefilter_trace - if threshold is None or threshold == 0 or trace is None: + threshold = options.threshold + trace = options.trace + if threshold == 0: return ir state = _RewriteState( diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 6cf8abfe6b44..391a99928547 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -104,12 +104,11 @@ def lower_ir_graph( -------- lower_ir_node """ - if _dynamic_planning_on(config_options): - from cudf_polars.streaming.join_domain_prefilter import ( - optimize_join_domain_prefilters, - ) + from cudf_polars.streaming.join_domain_prefilter import ( + optimize_join_domain_prefilters, + ) - ir = optimize_join_domain_prefilters(ir, stats, config_options) + ir = optimize_join_domain_prefilters(ir, stats, config_options) state: State = { "config_options": config_options, diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 4dcb3234c6e7..5374ba0bb8e1 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -51,6 +51,7 @@ "DaskContext", "DynamicPlanningOptions", "InMemoryExecutor", + "JoinDomainPrefilterOptions", "ParquetOptions", "RayContext", "SPMDContext", @@ -173,20 +174,10 @@ 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) -def _optional_bool_converter(v: str) -> bool | None: - if v.lower() in {"none", "null"}: - return None - return _bool_converter(v) - - @dataclasses.dataclass(frozen=True) class ParquetOptions: """ @@ -332,16 +323,6 @@ class DynamicPlanningOptions: join_prefilter_trace Whether to collect input/output row counts around applied join prefilters. Default is False. - join_domain_prefilter_enabled - Whether to insert generic derived key-domain semi-join filters before - lowering streaming joins. Default is True. - join_domain_prefilter_threshold - Row-count ratio (domain / target) below which a derived key-domain - semi-join filter is inserted. When unset, ``join_prefilter_threshold`` - is used. Default is unset. - join_domain_prefilter_trace - Whether to emit plan-time trace decisions for derived key-domain - prefilters. Default follows ``join_prefilter_trace``. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -372,27 +353,6 @@ class DynamicPlanningOptions: default=False, ) ) - join_domain_prefilter_enabled: bool = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__JOIN_DOMAIN_PREFILTER_ENABLED", - _bool_converter, - default=True, - ) - ) - join_domain_prefilter_threshold: float | None = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__JOIN_DOMAIN_PREFILTER_THRESHOLD", - _optional_float_converter, - default=None, - ) - ) - join_domain_prefilter_trace: bool | None = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__JOIN_DOMAIN_PREFILTER_TRACE", - _optional_bool_converter, - default=None, - ) - ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.sample_chunk_count, int): @@ -419,28 +379,59 @@ def __post_init__(self) -> None: # noqa: D105 ) if not isinstance(self.join_prefilter_trace, bool): raise TypeError("join_prefilter_trace must be a bool") - if not isinstance(self.join_domain_prefilter_enabled, bool): - raise TypeError("join_domain_prefilter_enabled must be a bool") - join_domain_prefilter_threshold = self.join_domain_prefilter_threshold - if join_domain_prefilter_threshold is None: - join_domain_prefilter_threshold = join_prefilter_threshold - object.__setattr__( - self, - "join_domain_prefilter_threshold", - join_domain_prefilter_threshold, - ) - elif not isinstance(join_domain_prefilter_threshold, float): - raise TypeError("join_domain_prefilter_threshold must be a float or None") - if not 0.0 <= join_domain_prefilter_threshold <= 1.0: - raise ValueError("join_domain_prefilter_threshold must be between 0 and 1") - join_domain_prefilter_trace = self.join_domain_prefilter_trace - if join_domain_prefilter_trace is None: - join_domain_prefilter_trace = self.join_prefilter_trace - object.__setattr__( - self, "join_domain_prefilter_trace", join_domain_prefilter_trace - ) - elif not isinstance(join_domain_prefilter_trace, bool): - raise TypeError("join_domain_prefilter_trace must be a bool or None") + + +@dataclasses.dataclass(frozen=True) +class JoinDomainPrefilterOptions: + """ + Configuration for the logical join-domain prefilter rewrite. + + These options can be configured via environment variables with the prefix + ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__``. + + Parameters + ---------- + enabled + Whether to insert generic derived key-domain semi-join filters before + lowering streaming joins. Default is True. + threshold + Row-count ratio (domain / target) below which a derived key-domain + semi-join filter is inserted. Default is 0.5. + trace + Whether to emit plan-time trace decisions for derived key-domain + prefilters. Default is False. + """ + + _env_prefix = "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER" + + enabled: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__ENABLED", _bool_converter, default=True + ) + ) + threshold: float = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__THRESHOLD", float, default=0.5 + ) + ) + trace: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__TRACE", _bool_converter, default=False + ) + ) + + def __post_init__(self) -> None: # noqa: D105 + if not isinstance(self.enabled, bool): + raise TypeError("enabled must be a bool") + threshold = self.threshold + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)): + raise TypeError("threshold must be a float or int") + threshold = float(threshold) + object.__setattr__(self, "threshold", threshold) + if not 0.0 <= threshold <= 1.0: + raise ValueError("threshold must be between 0 and 1") + if not isinstance(self.trace, bool): + raise TypeError("trace must be a bool") @dataclasses.dataclass(frozen=True, eq=True) @@ -703,6 +694,9 @@ class StreamingExecutor: dynamic_planning Options controlling dynamic shuffle planning. See :class:`~cudf_polars.utils.config.DynamicPlanningOptions` for more. + join_domain_prefilter + Options controlling the logical join-domain prefilter rewrite. See + :class:`~cudf_polars.utils.config.JoinDomainPrefilterOptions` for more. max_io_threads Maximum number of IO threads. Default is 4. This controls the parallelism of IO operations when reading data. @@ -766,6 +760,9 @@ class StreamingExecutor: dynamic_planning: DynamicPlanningOptions | None = dataclasses.field( default_factory=DynamicPlanningOptions ) + join_domain_prefilter: JoinDomainPrefilterOptions = dataclasses.field( + default_factory=JoinDomainPrefilterOptions + ) max_io_threads: int = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__MAX_IO_THREADS", int, default=4 @@ -821,6 +818,13 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) + if isinstance(self.join_domain_prefilter, dict): + object.__setattr__( + self, + "join_domain_prefilter", + JoinDomainPrefilterOptions(**self.join_domain_prefilter), + ) + if self.cluster in ("spmd", "ray", "dask"): if self.sink_to_directory is False: raise ValueError( @@ -853,6 +857,7 @@ def __hash__(self) -> int: # noqa: D105 # to json and hash that. d = dataclasses.asdict(self) d["dynamic_planning"] = json.dumps(d["dynamic_planning"]) + d["join_domain_prefilter"] = json.dumps(d["join_domain_prefilter"]) return hash(tuple(sorted(d.items()))) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 2d339613f5b5..b8a375b08db9 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -112,16 +112,16 @@ def _stats(**row_counts: tuple[Scan, int]) -> StatsCollector: return stats -def _config() -> ConfigOptions: +def _config(*, dynamic_planning: bool = True) -> ConfigOptions: + executor_options: dict[str, object] = { + "join_domain_prefilter": {"enabled": True, "trace": False} + } + if not dynamic_planning: + executor_options["dynamic_planning"] = None return ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={ - "dynamic_planning": { - "join_domain_prefilter_enabled": True, - "join_domain_prefilter_trace": False, - } - }, + executor_options=executor_options, ) ) @@ -153,6 +153,20 @@ def test_simple_domain_prefilter_filters_large_side() -> None: assert optimized.children[0] is part +def test_domain_prefilter_is_independent_of_dynamic_planning() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + lineitem = _scan("lineitem", ("l_partkey",)) + root = _join(part, lineitem, ("p_partkey",), ("l_partkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(part=(part, 6), lineitem=(lineitem, 1_800)), + _config(dynamic_planning=False), + ) + + assert _joins(optimized, "Semi") + + def test_no_simple_domain_prefilter_when_domain_is_not_selective() -> None: supplier = _scan("supplier", ("s_suppkey",)) lineitem = _scan("lineitem", ("l_suppkey",)) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 08f836a4573e..4170fb2f3d98 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -29,6 +29,7 @@ Cluster, ConfigOptions, DynamicPlanningOptions, + JoinDomainPrefilterOptions, MemoryResourceConfig, StreamingExecutor, ) @@ -556,9 +557,9 @@ def test_dynamic_planning_defaults() -> None: 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 - assert config.executor.dynamic_planning.join_domain_prefilter_enabled - assert config.executor.dynamic_planning.join_domain_prefilter_threshold == 0.5 - assert not config.executor.dynamic_planning.join_domain_prefilter_trace + assert config.executor.join_domain_prefilter.enabled + assert config.executor.join_domain_prefilter.threshold == 0.5 + assert not config.executor.join_domain_prefilter.trace def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -598,44 +599,22 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non 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 - assert config.executor.dynamic_planning.join_domain_prefilter_threshold == 0.25 - assert config.executor.dynamic_planning.join_domain_prefilter_trace + assert config.executor.join_domain_prefilter.threshold == 0.5 + assert not config.executor.join_domain_prefilter.trace def test_join_domain_prefilter_options_from_env( monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__ENABLED", "0") monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_ENABLED", - "0", - ) - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_THRESHOLD", - "0.125", - ) - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_TRACE", - "1", + "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__THRESHOLD", "0.125" ) + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__TRACE", "1") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - assert config.executor.dynamic_planning is not None - assert not config.executor.dynamic_planning.join_domain_prefilter_enabled - assert config.executor.dynamic_planning.join_domain_prefilter_threshold == 0.125 - assert config.executor.dynamic_planning.join_domain_prefilter_trace - - -@pytest.mark.parametrize("value", ["none", "null"]) -def test_join_domain_prefilter_trace_inherits_from_env( - monkeypatch: pytest.MonkeyPatch, value: str -) -> None: - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_PREFILTER_TRACE", "1" - ) - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__JOIN_DOMAIN_PREFILTER_TRACE", - value, - ) - assert DynamicPlanningOptions().join_domain_prefilter_trace + assert not config.executor.join_domain_prefilter.enabled + assert config.executor.join_domain_prefilter.threshold == 0.125 + assert config.executor.join_domain_prefilter.trace @pytest.mark.parametrize("value, expected", [("none", None), ("null", None), ("2", 2)]) @@ -731,46 +710,47 @@ def test_validate_join_prefilter_trace() -> None: def test_validate_join_domain_prefilter_options() -> None: - with pytest.raises(TypeError, match="join_domain_prefilter_enabled must be"): + with pytest.raises(TypeError, match="enabled must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={ - "dynamic_planning": {"join_domain_prefilter_enabled": "bad"} - }, + executor_options={"join_domain_prefilter": {"enabled": "bad"}}, ) ) - with pytest.raises(TypeError, match="join_domain_prefilter_threshold must be"): + with pytest.raises(TypeError, match="threshold must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={ - "dynamic_planning": {"join_domain_prefilter_threshold": "bad"} - }, + executor_options={"join_domain_prefilter": {"threshold": "bad"}}, ) ) - with pytest.raises( - ValueError, match="join_domain_prefilter_threshold must be between" - ): + with pytest.raises(ValueError, match="threshold must be between"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={ - "dynamic_planning": {"join_domain_prefilter_threshold": 1.5} - }, + executor_options={"join_domain_prefilter": {"threshold": 1.5}}, ) ) - with pytest.raises(TypeError, match="join_domain_prefilter_trace must be"): + with pytest.raises(TypeError, match="trace must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={ - "dynamic_planning": {"join_domain_prefilter_trace": "bad"} - }, + executor_options={"join_domain_prefilter": {"trace": "bad"}}, ) ) +def test_join_domain_prefilter_from_instance() -> None: + options = JoinDomainPrefilterOptions(enabled=False, threshold=0.25, trace=True) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_domain_prefilter": options}, + ) + ) + assert config.executor.join_domain_prefilter is options + + def test_dynamic_planning_from_instance() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( From ff32f43ff883866099de3e6271ae103fa851b21b Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 Jul 2026 19:09:05 +0000 Subject: [PATCH 40/44] Test nullable join-domain prefilters Exercise the domain-prefilter rewrite with nullable join keys for both null-equality modes. The regression verifies that inserted semi joins inherit the original join semantics and that optimized GPU execution matches Polars CPU results. --- .../streaming/test_join_domain_prefilter.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index b8a375b08db9..d5dc56339adb 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -5,8 +5,11 @@ from typing import TYPE_CHECKING, Literal +import pytest + import polars as pl +from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl import expr from cudf_polars.dsl.ir import Join, Scan, Select @@ -16,9 +19,13 @@ _smallest_node_containing_all, optimize_join_domain_prefilters, ) +from cudf_polars.streaming.statistics import collect_statistics +from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import ConfigOptions, ParquetOptions if TYPE_CHECKING: + import concurrent.futures + from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import SerializedDataSourceInfo @@ -167,6 +174,46 @@ def test_domain_prefilter_is_independent_of_dynamic_planning() -> None: assert _joins(optimized, "Semi") +@pytest.mark.parametrize( + "nulls_equal", [False, True], ids=["nulls_not_equal", "nulls_equal"] +) +def test_nullable_join_keys_preserve_results( + nulls_equal: bool, # noqa: FBT001 + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + domain = pl.LazyFrame( + { + "key": [None, 1, 2, 9], + "active": [True, True, True, False], + } + ).filter("active") + target = pl.LazyFrame( + { + "key": [None, 1, 2, 3] * 10, + "value": range(40), + } + ) + query = domain.join(target, on="key", nulls_equal=nulls_equal) + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={"join_domain_prefilter": {"enabled": True, "threshold": 0.5}}, + ) + + ir = Translator(query._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_join_domain_prefilters( + ir, + collect_statistics(ir, config, parquet_stats_executor), + config, + ) + + semi_joins = _joins(optimized, "Semi") + assert semi_joins + assert all(join.options[1] is nulls_equal for join in semi_joins) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + def test_no_simple_domain_prefilter_when_domain_is_not_selective() -> None: supplier = _scan("supplier", ("s_suppkey",)) lineitem = _scan("lineitem", ("l_suppkey",)) From 011ba8746e2002f7f2004ec32c7ce7217e0420ce Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 Jul 2026 20:02:03 +0000 Subject: [PATCH 41/44] Validate join-domain prefilter configuration Reject unsupported join-domain prefilter option values during streaming executor construction so invalid configuration cannot fail later in logical optimization. --- python/cudf_polars/cudf_polars/utils/config.py | 5 +++++ python/cudf_polars/tests/test_config.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 5374ba0bb8e1..3f4e827c4901 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -824,6 +824,11 @@ def __post_init__(self) -> None: # noqa: D105 "join_domain_prefilter", JoinDomainPrefilterOptions(**self.join_domain_prefilter), ) + if not isinstance(self.join_domain_prefilter, JoinDomainPrefilterOptions): + raise TypeError( + "join_domain_prefilter must be a JoinDomainPrefilterOptions " + "instance or dict" + ) if self.cluster in ("spmd", "ray", "dask"): if self.sink_to_directory is False: diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 4170fb2f3d98..778dcbefcc75 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -740,6 +740,20 @@ def test_validate_join_domain_prefilter_options() -> None: ) +@pytest.mark.parametrize("value", [None, object()]) +def test_validate_join_domain_prefilter_type(value: object) -> None: + with pytest.raises( + TypeError, + match="join_domain_prefilter must be a JoinDomainPrefilterOptions instance", + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_domain_prefilter": value}, + ) + ) + + def test_join_domain_prefilter_from_instance() -> None: options = JoinDomainPrefilterOptions(enabled=False, threshold=0.25, trace=True) config = ConfigOptions.from_polars_engine( From 7188a61d94cf18454317ba8b81653b53a30c6604 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 2 Jul 2026 16:05:02 +0000 Subject: [PATCH 42/44] Use optional join-domain prefilter options Use None as the join-domain prefilter disable sentinel, matching dynamic planning. Preserve enabled defaults, add top-level environment disabling, and cover option propagation and rewrite bypass behavior. --- .../cudf_polars/cudf_polars/engine/options.py | 12 +++--- .../streaming/join_domain_prefilter.py | 2 +- .../cudf_polars/cudf_polars/utils/config.py | 33 +++++++++------ .../streaming/test_join_domain_prefilter.py | 22 ++++++++-- .../tests/streaming/test_options.py | 5 +++ python/cudf_polars/tests/test_config.py | 40 ++++++++++++------- 6 files changed, 78 insertions(+), 36 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 35d838c42efa..1ad8bdf4ee13 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -250,8 +250,10 @@ class StreamingOptions: Category: executor. join_domain_prefilter Join-domain prefilter config, dict or - :class:`~cudf_polars.utils.config.JoinDomainPrefilterOptions`. - Env: ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__*``. + :class:`~cudf_polars.utils.config.JoinDomainPrefilterOptions`. ``None`` + disables the rewrite. + Env: ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER`` and + ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__*``. Default: enabled. Category: executor. sink_to_directory @@ -348,9 +350,9 @@ class StreamingOptions: dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = ( _opt("executor") ) - join_domain_prefilter: dict[str, Any] | JoinDomainPrefilterOptions | Unspecified = ( - _opt("executor") - ) + join_domain_prefilter: ( + dict[str, Any] | JoinDomainPrefilterOptions | None | Unspecified + ) = _opt("executor") sink_to_directory: bool | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__SINK_TO_DIRECTORY", parse_boolean ) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index 310dcbcde071..fd9c75f120e2 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -112,7 +112,7 @@ def optimize_join_domain_prefilters( after every inserted row-reduction semi join. """ options = config_options.executor.join_domain_prefilter - if not options.enabled: + if options is None: return ir threshold = options.threshold trace = options.trace diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 3f4e827c4901..7d47e893eb3f 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -386,14 +386,14 @@ class JoinDomainPrefilterOptions: """ Configuration for the logical join-domain prefilter rewrite. + Pass ``None`` to ``StreamingExecutor(join_domain_prefilter=...)`` to + disable the rewrite. + These options can be configured via environment variables with the prefix ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__``. Parameters ---------- - enabled - Whether to insert generic derived key-domain semi-join filters before - lowering streaming joins. Default is True. threshold Row-count ratio (domain / target) below which a derived key-domain semi-join filter is inserted. Default is 0.5. @@ -404,11 +404,6 @@ class JoinDomainPrefilterOptions: _env_prefix = "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER" - enabled: bool = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__ENABLED", _bool_converter, default=True - ) - ) threshold: float = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__THRESHOLD", float, default=0.5 @@ -421,8 +416,6 @@ class JoinDomainPrefilterOptions: ) def __post_init__(self) -> None: # noqa: D105 - if not isinstance(self.enabled, bool): - raise TypeError("enabled must be a bool") threshold = self.threshold if isinstance(threshold, bool) or not isinstance(threshold, (int, float)): raise TypeError("threshold must be a float or int") @@ -697,6 +690,7 @@ class StreamingExecutor: join_domain_prefilter Options controlling the logical join-domain prefilter rewrite. See :class:`~cudf_polars.utils.config.JoinDomainPrefilterOptions` for more. + ``None`` disables the rewrite. max_io_threads Maximum number of IO threads. Default is 4. This controls the parallelism of IO operations when reading data. @@ -760,7 +754,7 @@ class StreamingExecutor: dynamic_planning: DynamicPlanningOptions | None = dataclasses.field( default_factory=DynamicPlanningOptions ) - join_domain_prefilter: JoinDomainPrefilterOptions = dataclasses.field( + join_domain_prefilter: JoinDomainPrefilterOptions | None = dataclasses.field( default_factory=JoinDomainPrefilterOptions ) max_io_threads: int = dataclasses.field( @@ -824,10 +818,12 @@ def __post_init__(self) -> None: # noqa: D105 "join_domain_prefilter", JoinDomainPrefilterOptions(**self.join_domain_prefilter), ) - if not isinstance(self.join_domain_prefilter, JoinDomainPrefilterOptions): + if self.join_domain_prefilter is not None and not isinstance( + self.join_domain_prefilter, JoinDomainPrefilterOptions + ): raise TypeError( "join_domain_prefilter must be a JoinDomainPrefilterOptions " - "instance or dict" + "instance, dict, or None" ) if self.cluster in ("spmd", "ray", "dask"): @@ -982,6 +978,17 @@ def from_polars_engine( if not _bool_converter(env_dynamic_planning): user_executor_options["dynamic_planning"] = None + # Handle join_domain_prefilter: check user config, then env var + user_join_domain_prefilter = user_executor_options.get( + "join_domain_prefilter", None + ) + if user_join_domain_prefilter is None: + env_join_domain_prefilter = os.environ.get( + "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER", "1" + ) + if not _bool_converter(env_join_domain_prefilter): + user_executor_options["join_domain_prefilter"] = None + executor = StreamingExecutor(**user_executor_options) case _: # pragma: no cover; Unreachable raise ValueError(f"Unsupported executor: {user_executor}") diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index d5dc56339adb..315a27968f90 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -119,9 +119,11 @@ def _stats(**row_counts: tuple[Scan, int]) -> StatsCollector: return stats -def _config(*, dynamic_planning: bool = True) -> ConfigOptions: +def _config( + *, dynamic_planning: bool = True, join_domain_prefilter: bool = True +) -> ConfigOptions: executor_options: dict[str, object] = { - "join_domain_prefilter": {"enabled": True, "trace": False} + "join_domain_prefilter": {"trace": False} if join_domain_prefilter else None } if not dynamic_planning: executor_options["dynamic_planning"] = None @@ -174,6 +176,20 @@ def test_domain_prefilter_is_independent_of_dynamic_planning() -> None: assert _joins(optimized, "Semi") +def test_domain_prefilter_can_be_disabled() -> None: + part = _scan("part", ("p_partkey",), predicate=True) + lineitem = _scan("lineitem", ("l_partkey",)) + root = _join(part, lineitem, ("p_partkey",), ("l_partkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(part=(part, 6), lineitem=(lineitem, 1_800)), + _config(join_domain_prefilter=False), + ) + + assert optimized is root + + @pytest.mark.parametrize( "nulls_equal", [False, True], ids=["nulls_not_equal", "nulls_equal"] ) @@ -197,7 +213,7 @@ def test_nullable_join_keys_preserve_results( engine = pl.GPUEngine( executor="streaming", raise_on_fail=True, - executor_options={"join_domain_prefilter": {"enabled": True, "threshold": 0.5}}, + executor_options={"join_domain_prefilter": {"threshold": 0.5}}, ) ir = Translator(query._ldf.visit(), engine).translate_ir() diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index c5a42062c9a7..c9d07f687ea4 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -83,6 +83,11 @@ def test_executor_options_sink_to_directory_absent_when_unspecified() -> None: assert "sink_to_directory" not in StreamingOptions().to_executor_options() +def test_executor_options_join_domain_prefilter_disabled() -> None: + result = StreamingOptions(join_domain_prefilter=None).to_executor_options() + assert result["join_domain_prefilter"] is None + + # --------------------------------------------------------------------------- # to_engine_options # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 778dcbefcc75..fcfd7771f9dc 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -557,7 +557,7 @@ def test_dynamic_planning_defaults() -> None: 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 - assert config.executor.join_domain_prefilter.enabled + assert config.executor.join_domain_prefilter is not None assert config.executor.join_domain_prefilter.threshold == 0.5 assert not config.executor.join_domain_prefilter.trace @@ -599,6 +599,7 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non 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 + assert config.executor.join_domain_prefilter is not None assert config.executor.join_domain_prefilter.threshold == 0.5 assert not config.executor.join_domain_prefilter.trace @@ -606,17 +607,25 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non def test_join_domain_prefilter_options_from_env( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__ENABLED", "0") monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__THRESHOLD", "0.125" ) monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__TRACE", "1") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - assert not config.executor.join_domain_prefilter.enabled + assert config.executor.join_domain_prefilter is not None assert config.executor.join_domain_prefilter.threshold == 0.125 assert config.executor.join_domain_prefilter.trace +def test_join_domain_prefilter_disabled_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER", "0") + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__TRACE", "1") + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.join_domain_prefilter is None + + @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 @@ -710,13 +719,6 @@ def test_validate_join_prefilter_trace() -> None: def test_validate_join_domain_prefilter_options() -> None: - with pytest.raises(TypeError, match="enabled must be"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={"join_domain_prefilter": {"enabled": "bad"}}, - ) - ) with pytest.raises(TypeError, match="threshold must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( @@ -740,8 +742,7 @@ def test_validate_join_domain_prefilter_options() -> None: ) -@pytest.mark.parametrize("value", [None, object()]) -def test_validate_join_domain_prefilter_type(value: object) -> None: +def test_validate_join_domain_prefilter_type() -> None: with pytest.raises( TypeError, match="join_domain_prefilter must be a JoinDomainPrefilterOptions instance", @@ -749,13 +750,13 @@ def test_validate_join_domain_prefilter_type(value: object) -> None: ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": value}, + executor_options={"join_domain_prefilter": object()}, ) ) def test_join_domain_prefilter_from_instance() -> None: - options = JoinDomainPrefilterOptions(enabled=False, threshold=0.25, trace=True) + options = JoinDomainPrefilterOptions(threshold=0.25, trace=True) config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", @@ -765,6 +766,17 @@ def test_join_domain_prefilter_from_instance() -> None: assert config.executor.join_domain_prefilter is options +def test_join_domain_prefilter_disabled_from_options() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_domain_prefilter": None}, + ) + ) + assert config.executor.join_domain_prefilter is None + assert hash(config) == hash(config) + + def test_dynamic_planning_from_instance() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( From 9cf3029a6c9354d0ca5cad2ceb72f9041928cb8f Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 2 Jul 2026 12:53:56 -0700 Subject: [PATCH 43/44] Fix join-domain prefilter CI regressions Shut down the default singleton created by the nullable-key execution test so later explicit engine fixtures can initialize. Register the join-domain options class in the Sphinx API and options references. --- docs/cudf/source/cudf_polars/api.md | 1 + docs/cudf/source/cudf_polars/options.md | 1 + .../tests/streaming/test_join_domain_prefilter.py | 6 +++++- 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 6acff73d4f03..3c2b0d88dc32 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -67,6 +67,7 @@ Most users interact with them through `StreamingOptions` fields rather than dire .. automodule:: cudf_polars.utils.config :members: DynamicPlanningOptions, + JoinDomainPrefilterOptions, MemoryResourceConfig, ParquetOptions, StreamingExecutor, diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index 5d813e74bbe1..94ca7d52e1bf 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/docs/cudf/source/cudf_polars/options.md @@ -108,6 +108,7 @@ Environment variables follow these patterns: | `broadcast_limit` | Maximum number of bytes for broadcast joins. | auto | | `target_partition_size` | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto. | auto | | `dynamic_planning` | Dynamic planning configuration, dict or {class}`~cudf_polars.utils.config.DynamicPlanningOptions`. `None` disables. | enabled | +| `join_domain_prefilter` | Join-domain prefilter configuration, dict or {class}`~cudf_polars.utils.config.JoinDomainPrefilterOptions`. `None` disables. | enabled | | `sink_to_directory` | Whether `.sink_*()` writes its output as a directory. The `spmd`, `ray`, and `dask` engines always use `True`; passing `False` raises `ValueError`. | `True` | ### Category: `engine` diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index 315a27968f90..0aead0061d49 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -14,6 +14,7 @@ from cudf_polars.dsl import expr from cudf_polars.dsl.ir import Join, Scan, Select from cudf_polars.dsl.traversal import traversal +from cudf_polars.engine.default_singleton_engine import DefaultSingletonEngine from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.join_domain_prefilter import ( _smallest_node_containing_all, @@ -227,7 +228,10 @@ def test_nullable_join_keys_preserve_results( semi_joins = _joins(optimized, "Semi") assert semi_joins assert all(join.options[1] is nulls_equal for join in semi_joins) - assert_gpu_result_equal(query, engine=engine, check_row_order=False) + try: + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + finally: + DefaultSingletonEngine.shutdown() def test_no_simple_domain_prefilter_when_domain_is_not_selective() -> None: From 7f9d2b2bb9f15c70e78ab0bb1d9205f4638fc051 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 3 Jul 2026 18:58:57 +0000 Subject: [PATCH 44/44] Reuse existing semi-join domains Discover selective key domains that are already attached to a semi join and reuse them to prefilter another large join input. Follow exact column bindings and shared source provenance, reject uncached source duplication, and retain the original full join for final correctness. For SF30K Q18, this inserts an early key-domain reduction on the second lineitem input. The traced six-node run improved from 18.35s/10.80s to 11.51s/11.89s, and the untraced first iteration improved to 9.82s. This materially reduces memory pressure, although the raw detail payload remains large enough that a later hot iteration can still exhaust six nodes. Q18 plan before, relevant subtree: JOIN Inner (o_orderkey) (l_orderkey) JOIN Semi (o_orderkey) (l_orderkey) STREAMINGSCAN orders PROJECTION selected_l_orderkey FILTER sum_quantity > 300 GROUPBY lineitem by l_orderkey CACHE lineitem (l_orderkey, l_quantity) Q18 plan after, relevant subtree: JOIN Inner (o_orderkey) (l_orderkey) JOIN Semi (o_orderkey) (l_orderkey) STREAMINGSCAN orders PROJECTION selected_l_orderkey FILTER sum_quantity > 300 GROUPBY lineitem by l_orderkey JOIN Semi (l_orderkey) (l_orderkey) CACHE lineitem (l_orderkey, l_quantity) PROJECTION selected_l_orderkey FILTER sum_quantity > 300 GROUPBY lineitem by l_orderkey --- .../streaming/join_domain_prefilter.py | 138 ++++++++++++++++- .../streaming/test_join_domain_prefilter.py | 142 +++++++++++++++++- 2 files changed, 275 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py index cc76ebb97be9..c9e22f7003f2 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -59,7 +59,7 @@ def column(self) -> str: class _Candidate: """A derived key-domain prefilter candidate.""" - mode: Literal["simple", "composite"] + mode: Literal["simple", "composite", "reused"] target_side: Literal["left", "right"] target: _Producer target_key: expr.Col @@ -83,13 +83,14 @@ def domain_cost(self) -> int: @property def score(self) -> tuple[int, int, int]: """Prefer composite filters, then cheaper constraint/domain inputs.""" + mode_score = {"composite": 0, "simple": 1, "reused": 2}[self.mode] constraint_cost = ( self.constraint_domain.cost if self.constraint_domain is not None else self.domain.cost ) return ( - 0 if self.mode == "composite" else 1, + mode_score, constraint_cost, self.domain.cost, ) @@ -268,6 +269,19 @@ def _select_candidate( selective_nodes, ) ) + candidates.extend( + _reused_semi_domain_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + row_estimates, + source_costs, + source_counts, + selective_nodes, + ) + ) if not candidates: return None, "no_profitable_domain" @@ -329,6 +343,49 @@ def _simple_candidates( ) +def _reused_semi_domain_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], + row_estimates: dict[IR, int | None], + source_costs: dict[IR, int | None], + source_counts: dict[IR, int], + selective_nodes: set[IR], +) -> Iterable[_Candidate]: + for target_key, domain_key in zip(target_keys, domain_keys, strict=True): + for domain, semi_domain_key in _filtering_semi_domains( + domain_child, + domain_key, + row_estimates, + source_costs, + selective_nodes, + ): + target = _largest_key_reusable_target( + target_child, target_key.name, domain.node, row_estimates, source_costs + ) + if target is None: + continue + if domain.rows > target.rows: + continue + if _contains_identity(target.node, domain.node): + continue + if source_counts.get(domain.node) == 1 and _has_filtering_semi_ancestor( + target_child, target.node + ): + continue + yield _Candidate( + mode="reused", + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=semi_domain_key, + target_rows=target.rows, + ) + + def _composite_candidates( target_side: Literal["left", "right"], target_child: IR, @@ -412,7 +469,7 @@ def _composite_candidates( def _make_domain(candidate: _Candidate, ir: Join) -> IR: - if candidate.mode == "simple": + if candidate.mode in ("simple", "reused"): return _project_bound_key( candidate.domain.node, candidate.domain.column, @@ -507,6 +564,37 @@ def _smallest_key_producer( return min(candidates, key=lambda item: (item[0], item[1], item[2]))[3] +def _filtering_semi_domains( + root: IR, + filtered_key: expr.Col, + row_estimates: dict[IR, int | None], + source_costs: dict[IR, int | None], + selective_nodes: set[IR], +) -> Iterable[tuple[_Producer, expr.Col]]: + """Yield domains from semi joins on the exact filtered-key lineage.""" + for node, bound_key in _column_bindings(root, filtered_key.name): + if not isinstance(node, Join) or node.options[0] != "Semi": + continue + left_keys = _simple_keys(node.left_on) + right_keys = _simple_keys(node.right_on) + if len(left_keys) != len(node.left_on) or len(right_keys) != len(node.right_on): + continue + assert len(left_keys) == len(right_keys) + for left_key, right_key in zip(left_keys, right_keys, strict=True): + if left_key.name != bound_key: + continue + domain = _smallest_key_producer( + node.children[1], + right_key.name, + row_estimates, + source_costs, + selective_nodes, + require_selective=True, + ) + if domain is not None: + yield domain, right_key + + def _smallest_node_containing_all( root: IR, columns: Sequence[str], @@ -571,6 +659,50 @@ def _largest_key_source( return max(candidates, key=lambda item: (item[0], -item[1]))[2] +def _largest_key_reusable_target( + root: IR, + column: str, + domain: IR, + row_estimates: dict[IR, int | None], + source_costs: dict[IR, int | None], +) -> _Producer | None: + """Select a target without duplicating a source used by the domain.""" + shared_cache_candidates = [] + source_candidates = [] + fallback_candidates = [] + domain_sources = _scan_sources(domain) + for order, (node, bound_column) in enumerate(_column_bindings(root, column)): + rows = row_estimates.get(node) + if rows is None or rows <= 0: + continue + cost = source_costs.get(node) + if cost is None: + continue + producer = _Producer(node, (bound_column,), rows, cost) + item = (rows, len(node.schema), -order, producer) + if _contains_identity(domain, node): + if isinstance(node, Cache): + shared_cache_candidates.append(item) + continue + if not _scan_sources(node).isdisjoint(domain_sources): + continue + if isinstance(node, (Scan, DataFrameScan)): + source_candidates.append(item) + else: + fallback_candidates.append(item) + candidates = shared_cache_candidates or source_candidates or fallback_candidates + if not candidates: + return None + return max(candidates, key=lambda item: (item[0], -item[1], item[2]))[3] + + +def _scan_sources(root: IR) -> frozenset[IR]: + """Return physical scan nodes reachable from a subtree.""" + return frozenset( + node for node in traversal([root]) if isinstance(node, (Scan, DataFrameScan)) + ) + + def _column_bindings(root: IR, column: str) -> Iterable[tuple[IR, str]]: """Yield exact output-to-input bindings for a column through a subplan.""" node = root diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py index bd92e52b509e..b978938239ae 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -12,7 +12,7 @@ from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl import expr -from cudf_polars.dsl.ir import Join, Scan, Select +from cudf_polars.dsl.ir import Cache, Filter, GroupBy, Join, Scan, Select from cudf_polars.dsl.traversal import traversal from cudf_polars.engine.default_singleton_engine import DefaultSingletonEngine from cudf_polars.streaming.base import StatsCollector @@ -156,6 +156,28 @@ def _join_key_names(keys: tuple[expr.NamedExpr, ...]) -> tuple[str, ...]: return tuple(names) +def _filtered_groupby_domain(source: IR, key: str) -> Select: + grouped = GroupBy( + {key: source.schema[key]}, + (_key(source, key),), + (), + False, # noqa: FBT003 + None, + source, + ) + filtered = Filter( + grouped.schema, + expr.NamedExpr("__predicate", expr.Literal(BOOL, True)), # noqa: FBT003 + grouped, + ) + return Select( + {key: filtered.schema[key]}, + (_key(filtered, key),), + True, # noqa: FBT003 + filtered, + ) + + def test_simple_domain_prefilter_filters_large_side() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey", "l_suppkey")) @@ -474,7 +496,7 @@ def test_domain_source_follows_join_key_through_rename() -> None: ) selected_domain = semi.children[1] assert isinstance(selected_domain, Select) - assert selected_domain.children[0] is domain + assert selected_domain.children[0] is domain_source def test_composite_domain_columns_follow_renames() -> None: @@ -523,6 +545,122 @@ def test_target_replacement_does_not_rewrite_shared_domain_side() -> None: assert domain.children[0] is shared +def test_reuses_existing_semi_domain_to_filter_shared_target() -> None: + lineitem = _scan("lineitem", ("l_orderkey", "l_quantity")) + cached_lineitem = Cache(lineitem.schema, 0, 2, lineitem) + selected_orders = _filtered_groupby_domain(cached_lineitem, "l_orderkey") + orders = _scan("orders", ("o_orderkey", "o_custkey")) + + filtered_orders = _join( + orders, selected_orders, ("o_orderkey",), ("l_orderkey",), how="Semi" + ) + root = _join(filtered_orders, cached_lineitem, ("o_orderkey",), ("l_orderkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(lineitem=(lineitem, 1_800), orders=(orders, 450)), + _config(), + ) + + assert isinstance(optimized, Join) + assert optimized.options[0] == "Inner" + assert optimized.children[0] is filtered_orders + assert isinstance(optimized.children[1], Join) + assert optimized.children[1].options[0] == "Semi" + assert optimized.children[1].children[0] is cached_lineitem + assert _contains_node(optimized.children[1].children[1], selected_orders) + + +def test_reused_semi_domain_follows_filtered_key_rename() -> None: + lineitem = _scan("lineitem", ("l_orderkey", "l_quantity")) + cached_lineitem = Cache(lineitem.schema, 0, 2, lineitem) + selected_orders = _filtered_groupby_domain(cached_lineitem, "l_orderkey") + orders = _scan("orders", ("o_orderkey", "o_custkey")) + + filtered_orders = _join( + orders, selected_orders, ("o_orderkey",), ("l_orderkey",), how="Semi" + ) + renamed_orders = _select( + filtered_orders, + joined_orderkey="o_orderkey", + o_custkey="o_custkey", + ) + root = _join( + renamed_orders, + cached_lineitem, + ("joined_orderkey",), + ("l_orderkey",), + ) + + optimized = optimize_join_domain_prefilters( + root, + _stats(lineitem=(lineitem, 1_800), orders=(orders, 450)), + _config(), + ) + + assert isinstance(optimized, Join) + assert isinstance(optimized.children[1], Join) + assert optimized.children[1].options[0] == "Semi" + assert optimized.children[1].children[0] is cached_lineitem + assert _contains_node(optimized.children[1].children[1], selected_orders) + + +def test_reused_semi_domain_does_not_duplicate_uncached_source() -> None: + lineitem = _scan("lineitem", ("l_orderkey", "l_quantity")) + selected_orders = _filtered_groupby_domain(lineitem, "l_orderkey") + orders = _scan("orders", ("o_orderkey", "o_custkey")) + + filtered_orders = _join( + orders, selected_orders, ("o_orderkey",), ("l_orderkey",), how="Semi" + ) + root = _join(filtered_orders, lineitem, ("o_orderkey",), ("l_orderkey",)) + + optimized = optimize_join_domain_prefilters( + root, + _stats(lineitem=(lineitem, 1_800), orders=(orders, 450)), + _config(), + ) + + lineitem_semis = [ + semi for semi in _joins(optimized, "Semi") if semi.children[0] is lineitem + ] + assert not lineitem_semis + + +def test_reused_semi_domain_does_not_duplicate_wrapped_uncached_source() -> None: + lineitem = _scan("lineitem", ("l_orderkey", "l_quantity")) + projected_lineitem = _select( + lineitem, + l_orderkey="l_orderkey", + l_quantity="l_quantity", + ) + selected_orders = _filtered_groupby_domain(projected_lineitem, "l_orderkey") + orders = _scan("orders", ("o_orderkey", "o_custkey")) + + filtered_orders = _join( + orders, selected_orders, ("o_orderkey",), ("l_orderkey",), how="Semi" + ) + root = _join( + filtered_orders, + projected_lineitem, + ("o_orderkey",), + ("l_orderkey",), + ) + + optimized = optimize_join_domain_prefilters( + root, + _stats(lineitem=(lineitem, 1_800), orders=(orders, 450)), + _config(), + ) + + lineitem_semis = [ + semi + for semi in _joins(optimized, "Semi") + if semi.children[0] is projected_lineitem + ] + assert not lineitem_semis + + def test_no_domain_prefilter_for_outer_join() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",))