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/65] 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/65] 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 09f04743d419647bcb734ec404d60e961a799574 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Fri, 26 Jun 2026 12:14:58 +0000 Subject: [PATCH 03/65] 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 04/65] 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 05/65] 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 06/65] 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 07/65] 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 08/65] 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 09/65] 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 10/65] 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 11/65] 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 12/65] 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 13/65] 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 14/65] 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 15/65] 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 16/65] 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 17/65] 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 18/65] 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 19/65] 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 20/65] 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 21/65] 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 22/65] 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 23/65] 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 24/65] 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 25/65] 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 26/65] 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 27/65] 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 28/65] 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 29/65] 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 30/65] 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 31/65] 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 32/65] 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 33/65] 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 34/65] 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 35/65] 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 36/65] 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 37/65] 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 38/65] 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 39/65] 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 40/65] 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 41/65] 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 42/65] 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 a5f919be3b6d7fd486bcbcc4ae369b8a48a6d049 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 13 Jul 2026 15:18:44 +0100 Subject: [PATCH 43/65] New optimize_with_stats stage in lowering Lowering now returns info on the optimized IR DAG along with the lowered (after optimization) DAG, and partition info mapping. The optimized DAG is used as the physical plan in quent traces. --- python/cudf_polars/cudf_polars/engine/core.py | 14 ++-- .../cudf_polars/streaming/explain.py | 14 ++-- .../cudf_polars/streaming/parallel.py | 68 ++++++++++++++----- python/cudf_polars/tests/quent/test_quent.py | 5 +- .../tests/streaming/test_dataframescan.py | 12 ++-- .../tests/streaming/test_hstack.py | 8 ++- .../cudf_polars/tests/streaming/test_join.py | 23 ++++--- .../cudf_polars/tests/streaming/test_scan.py | 4 +- 8 files changed, 99 insertions(+), 49 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index e87728b4d01f..a13a90a177a3 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -732,11 +732,17 @@ def evaluate_on_rank( """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) + lowering, node_map = lower_ir_graph_with_node_map( + ir, config_options, stats, rank=comm.rank, nranks=comm.nranks + ) + optimized = lowering.optimized + ir = lowering.lowered + partition_info = lowering.partition_info if config_options.executor.quent_context is not None: assert local_quent_context is not None - logical_plan_id = ir.get_stable_plan_id() + logical_plan_id = optimized.get_stable_plan_id() plan, ops, ports, logical_op_by_id = build_plan( - ir, + optimized, config_options, query=local_quent_context.context.query, plan_id=logical_plan_id, @@ -750,10 +756,6 @@ def evaluate_on_rank( local_quent_context.logger, plan, ops, ports ) - ir, partition_info, node_map = lower_ir_graph_with_node_map( - ir, config_options, stats, rank=comm.rank, nranks=comm.nranks - ) - if comm.rank == 0: log_query_plan(ir, config_options) diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 8e62e161cd7c..b585138d583a 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -123,8 +123,10 @@ def explain_query( if physical: with cm: stats = collect_statistics(ir, config, executor) - lowered_ir, partition_info = lower_ir_graph(ir, config, stats) - return _repr_ir_tree(lowered_ir, partition_info, stats=stats, config=config) + lowered = lower_ir_graph(ir, config, stats) + return _repr_ir_tree( + lowered.lowered, lowered.partition_info, stats=stats, config=config + ) else: if config.executor.name == "streaming": # Include row-count statistics for the logical plan @@ -150,7 +152,9 @@ def collect_partition_plan( with concurrent.futures.ThreadPoolExecutor() as executor: stats = collect_statistics(ir, config, executor) - lowered_ir, partition_info = lower_ir_graph(ir, config, stats) + lowered = lower_ir_graph(ir, config, stats) + lowered_ir = lowered.lowered + partition_info = lowered.partition_info seen: set[tuple] = set() rows: list[PartitionPlanRow] = [] @@ -755,7 +759,9 @@ def from_ir( if lowered: with cm: stats = collect_statistics(ir, config_options, executor) - ir, partition_info_d = lower_ir_graph(ir, config_options, stats) + lowering = lower_ir_graph(ir, config_options, stats) + ir = lowering.lowered + partition_info_d = lowering.partition_info partition_info_dict = {} nodes: dict[str, SerializableIRNode] = {} diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 83863d41f7d3..de88570440c9 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -4,6 +4,7 @@ from __future__ import annotations +import dataclasses import operator from functools import partial, reduce from typing import TYPE_CHECKING @@ -65,6 +66,44 @@ def _( ) +@dataclasses.dataclass +class LoweringInfo: + """Information produced by optimizing and lowering an IR graph.""" + + optimized: IR # IR after optimization + lowered: IR # optimized IR after lowering + partition_info: MutableMapping[ + IR, PartitionInfo + ] # Partition mapping for nodes in the lowered IR. + + +def optimize_with_stats( + ir: IR, config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector +) -> IR: + """ + Optimize an IR graph given some statistics. + + Parameters + ---------- + ir + Root of the graph to optimize. + config_options + GPUEngine configuration options. + stats + Pre-computed statistics. + + Returns + ------- + IR + The optimized IR graph. + """ + from cudf_polars.streaming.join_domain_prefilter import ( + optimize_join_domain_prefilters, + ) + + return optimize_join_domain_prefilters(ir, stats, config_options) + + def _lower_ir_graph_impl( ir: IR, config_options: ConfigOptions[StreamingExecutor], @@ -72,20 +111,19 @@ def _lower_ir_graph_impl( *, rank: int = 0, nranks: int = 1, -) -> tuple[tuple[IR, MutableMapping[IR, PartitionInfo]], LowerIRTransformer]: - from cudf_polars.streaming.join_domain_prefilter import ( - optimize_join_domain_prefilters, - ) - - ir = optimize_join_domain_prefilters(ir, stats, config_options) +) -> tuple[LoweringInfo, LowerIRTransformer]: state: State = { "config_options": config_options, "stats": stats, "rank": rank, "nranks": nranks, } + optimized = optimize_with_stats(ir, config_options, stats) mapper: LowerIRTransformer = CachingVisitor(lower_ir_node, state=state) - return mapper(ir), mapper + lowered, partition_info = mapper(optimized) + return LoweringInfo( + optimized=optimized, lowered=lowered, partition_info=partition_info + ), mapper def lower_ir_graph( @@ -95,7 +133,7 @@ def lower_ir_graph( *, rank: int = 0, nranks: int = 1, -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: +) -> LoweringInfo: """ Rewrite an IR graph and extract partitioning information. @@ -114,9 +152,7 @@ def lower_ir_graph( Returns ------- - new_ir, partition_info - The rewritten graph and a mapping from unique nodes - in the new graph to associated partitioning information. + LoweringInfo Notes ----- @@ -137,7 +173,7 @@ def lower_ir_graph_with_node_map( *, rank: int = 0, nranks: int = 1, -) -> tuple[IR, MutableMapping[IR, PartitionInfo], dict[str, list[str]]]: +) -> tuple[LoweringInfo, dict[str, list[str]]]: """ Lower an IR graph and return a mapping from physical to logical stable IDs. @@ -160,10 +196,8 @@ def lower_ir_graph_with_node_map( Returns ------- - new_ir - The rewritten IR graph. - partition_info - Mapping from unique nodes in the new graph to partitioning info. + LoweringInfo + Information about the lowered IR graph. node_map Mapping ``{physical_stable_id: [logical_stable_id, ...]}`` built from the internal :class:`CachingVisitor` cache. Nodes inserted @@ -178,7 +212,7 @@ def lower_ir_graph_with_node_map( old_key = str(old_node.get_stable_id()) node_map.setdefault(new_key, []).append(old_key) - return *result, node_map + return result, node_map def evaluate_streaming( diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index be1aedc770b8..2b7217d0689c 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -449,10 +449,9 @@ def test_lower_ir_graph_with_node_map() -> None: ir, config_options, concurrent.futures.ThreadPoolExecutor() ) - _lowered_ir, _partition_info, node_map = lower_ir_graph_with_node_map( - ir, config_options, stats - ) + lowering, node_map = lower_ir_graph_with_node_map(ir, config_options, stats) + assert lowering.optimized is ir assert len(node_map) > 0 for physical_sid, logical_sids in node_map.items(): assert isinstance(physical_sid, str) diff --git a/python/cudf_polars/tests/streaming/test_dataframescan.py b/python/cudf_polars/tests/streaming/test_dataframescan.py index 014851f689e7..0fd27f2b57e8 100644 --- a/python/cudf_polars/tests/streaming/test_dataframescan.py +++ b/python/cudf_polars/tests/streaming/test_dataframescan.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 from __future__ import annotations @@ -59,7 +59,7 @@ def test_parallel_dataframescan( ) qir = Translator(df._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, info = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -68,6 +68,8 @@ def test_parallel_dataframescan( parquet_stats_executor, ), ) + ir = lowering.lowered + info = lowering.partition_info count = info[ir].count if max_rows_per_partition < total_row_count: assert count > 1 @@ -106,7 +108,7 @@ def test_join_in_memory_lazy_stable_id_pickle( right = pl.LazyFrame({"k": [2, 3, 4], "y": [1, 2, 3]}).collect(engine=engine).lazy() qir = Translator(left.join(right, on="k")._ldf.visit(), engine).translate_ir() config_options = ConfigOptions.from_polars_engine(engine) - ir, _ = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -115,6 +117,7 @@ def test_join_in_memory_lazy_stable_id_pickle( parquet_stats_executor, ), ) + ir = lowering.lowered _assert_stable_ids_match(ir, pickle.loads(pickle.dumps(ir))) @@ -128,7 +131,7 @@ def test_dataframescan_pickle( ) qir = Translator(df._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, _ = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -137,6 +140,7 @@ def test_dataframescan_pickle( parquet_stats_executor, ), ) + ir = lowering.lowered # Pickle and unpickle the IR (which contains DataFrameScan) pickled = pickle.dumps(ir) diff --git a/python/cudf_polars/tests/streaming/test_hstack.py b/python/cudf_polars/tests/streaming/test_hstack.py index f7a3c6213340..232ad33e255e 100644 --- a/python/cudf_polars/tests/streaming/test_hstack.py +++ b/python/cudf_polars/tests/streaming/test_hstack.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 CSE HStack handling in the streaming executor.""" @@ -150,14 +150,16 @@ def test_cse_agg_shared_decomposition( assert len(inner_hstacks) == (1 if comm_subexpr_elim else 0) config_options = ConfigOptions.from_polars_engine(engine) - lowered, _ = lower_ir_graph( + lowering = lower_ir_graph( ir, config_options, collect_statistics(ir, config_options, parquet_stats_executor), ) # Both paths must lower to a single Repartition computing one aggregation. - repartitions = [n for n in traversal([lowered]) if isinstance(n, Repartition)] + repartitions = [ + n for n in traversal([lowering.lowered]) if isinstance(n, Repartition) + ] assert len(repartitions) == 1 assert len(repartitions[0].children[0].exprs) == 1 # type: ignore[attr-defined] assert_gpu_result_equal(q, engine=engine, collect_kwargs={"optimizations": opts}) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index 069c0328a68f..51e27f15155e 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -492,18 +492,17 @@ def test_broadcast_limit( q = left.join(right, on="y", how="inner") ir = Translator(q._ldf.visit(), engine).translate_ir() config_options = ConfigOptions.from_polars_engine(engine) - shuffle_nodes = [ - type(node) - for node in lower_ir_graph( + lowering = lower_ir_graph( + ir, + config_options, + collect_statistics( ir, config_options, - collect_statistics( - ir, - config_options, - parquet_stats_executor, - ), - )[1] - if isinstance(node, Shuffle) + parquet_stats_executor, + ), + ) + shuffle_nodes = [ + type(node) for node in lowering.partition_info if isinstance(node, Shuffle) ] # NOTE: Expect small table to have 3 partitions (9 / 3). @@ -543,11 +542,13 @@ def test_cache_preserves_partitioning_join( config_options = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() - lowered_ir, partition_info = lower_ir_graph( + lowering = lower_ir_graph( ir, config_options, collect_statistics(ir, config_options, parquet_stats_executor), ) + lowered_ir = lowering.lowered + partition_info = lowering.partition_info # Cache should preserve partitioning on 'key' cache_partitioning = [ diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 13e88ead7731..ed672e7677aa 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -205,7 +205,7 @@ def test_target_partition_size( ) qir = Translator(q._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, info = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -214,6 +214,8 @@ def test_target_partition_size( parquet_stats_executor, ), ) + ir = lowering.lowered + info = lowering.partition_info count = info[ir].count if blocksize <= 12_000: assert count > n_files From 3c6adf8eb2cc8e10cdc726d31f2c45eaf45caecd Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 13 Jul 2026 17:34:42 +0100 Subject: [PATCH 44/65] Introduce utility to map value domains of columns --- .../cudf_polars/dsl/utils/column_domain.py | 124 +++++++++++ .../tests/dsl/test_column_domain.py | 209 ++++++++++++++++++ 2 files changed, 333 insertions(+) create mode 100644 python/cudf_polars/cudf_polars/dsl/utils/column_domain.py create mode 100644 python/cudf_polars/tests/dsl/test_column_domain.py diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py new file mode 100644 index 000000000000..b2ad4a7d7aeb --- /dev/null +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for tracking column value domains between IR nodes.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import singledispatch +from typing import TYPE_CHECKING + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Cache, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, + Slice, + Sort, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + + from cudf_polars.dsl.ir import IR + +__all__ = ["ColumnRef", "column_domain_bindings"] + + +@dataclass(frozen=True) +class ColumnRef: + """A named column produced by an IR node.""" + + node: IR + name: str + + +@singledispatch +def column_domain_bindings(node: IR) -> Mapping[str, ColumnRef]: + """ + Map output columns to child columns containing their value domains. + + For every ``output_name -> ColumnRef(child, input_name)`` binding, every + value appearing in ``node[output_name]`` is guaranteed to appear in + ``child[input_name]``. Row order, multiplicity, and cardinality are not + preserved. + + If a name in ``node.schema`` does not appear in the mapping it means + that it was not possible to derive a relationship between the domain of + the output and input values for that column. + """ + return {} + + +@column_domain_bindings.register(Select) +def _(node: Select) -> Mapping[str, ColumnRef]: + child = node.children[0] + return { + item.name: ColumnRef(child, item.value.name) + for item in node.exprs + if isinstance(item.value, expr.Col) + } + + +@column_domain_bindings.register(HStack) +def _(node: HStack) -> Mapping[str, ColumnRef]: + child = node.children[0] + replaced = {item.name for item in node.columns} + return { + name: ColumnRef(child, name) for name in child.schema if name not in replaced + } | { + item.name: ColumnRef(child, item.value.name) + for item in node.columns + if isinstance(item.value, expr.Col) + } + + +@column_domain_bindings.register(GroupBy) +def _(node: GroupBy) -> Mapping[str, ColumnRef]: + child = node.children[0] + return { + key.name: ColumnRef(child, key.value.name) + for key in node.keys + if isinstance(key.value, expr.Col) + } + + +@column_domain_bindings.register(Join) +def _(node: Join) -> Mapping[str, ColumnRef]: + left, right = node.children + how = node.options[0] + if how in ("Semi", "Anti"): + return { + name: ColumnRef(left, name) for name in node.schema if name in left.schema + } + if how != "Inner": + return {} + + bindings = {name: ColumnRef(left, name) for name in left.schema} + suffix = node.options[3] + for name in right.schema: + output_name = f"{name}{suffix}" if name in left.schema else name + if output_name in node.schema: + bindings[output_name] = ColumnRef(right, name) + return bindings + + +@column_domain_bindings.register(Cache) +@column_domain_bindings.register(Distinct) +@column_domain_bindings.register(Filter) +@column_domain_bindings.register(Projection) +@column_domain_bindings.register(Slice) +@column_domain_bindings.register(Sort) +def _( + node: Cache | Distinct | Filter | Projection | Slice | Sort, +) -> Mapping[str, ColumnRef]: + child = node.children[0] + return { + name: ColumnRef(child, name) for name in node.schema if name in child.schema + } diff --git a/python/cudf_polars/tests/dsl/test_column_domain.py b/python/cudf_polars/tests/dsl/test_column_domain.py new file mode 100644 index 000000000000..3d1ec6b347b5 --- /dev/null +++ b/python/cudf_polars/tests/dsl/test_column_domain.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 + +import polars as pl + +import pylibcudf as plc + +from cudf_polars.containers import DataType +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Cache, + DataFrameScan, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, + Slice, + Sort, +) +from cudf_polars.dsl.utils.column_domain import ( + ColumnRef, + column_domain_bindings, +) + +I64 = DataType(pl.Int64()) +BOOL = DataType(pl.Boolean()) + + +def make_scan(*names: str) -> DataFrameScan: + frame = pl.DataFrame({name: [1] for name in names}) + return DataFrameScan(dict.fromkeys(names, I64), frame._df, None) + + +def col(name: str) -> expr.Col: + return expr.Col(I64, name) + + +def named_col(output: str, source: str) -> expr.NamedExpr: + return expr.NamedExpr(output, col(source)) + + +def test_source_has_no_column_domain_bindings() -> None: + assert column_domain_bindings(make_scan("a")) == {} + + +def test_select_binds_aliases_and_omits_derived_columns() -> None: + child = make_scan("a", "b") + node = Select( + {"renamed": I64, "derived": I64}, + ( + named_col("renamed", "a"), + expr.NamedExpr("derived", expr.Literal(I64, 1)), + ), + True, # noqa: FBT003 + child, + ) + + assert column_domain_bindings(node) == { + "renamed": ColumnRef(child, "a"), + } + + +def test_hstack_binds_passthrough_alias_and_override() -> None: + child = make_scan("a", "b") + node = HStack( + {"a": I64, "b": I64, "alias": I64}, + ( + expr.NamedExpr("a", expr.Literal(I64, 1)), + named_col("alias", "b"), + ), + True, # noqa: FBT003 + child, + ) + + assert column_domain_bindings(node) == { + "b": ColumnRef(child, "b"), + "alias": ColumnRef(child, "b"), + } + + +def test_groupby_binds_only_direct_keys() -> None: + child = make_scan("a", "b") + node = GroupBy( + {"key": I64, "value": I64}, + (named_col("key", "a"),), + (named_col("value", "b"),), + False, # noqa: FBT003 + None, + child, + ) + + assert column_domain_bindings(node) == { + "key": ColumnRef(child, "a"), + } + + +def test_inner_join_binds_left_right_and_suffixed_columns() -> None: + left = make_scan("key", "left_value") + right = make_scan("key", "right_value") + node = Join( + { + "key": I64, + "left_value": I64, + "key_right": I64, + "right_value": I64, + }, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Inner", False, (0, 1), "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnRef(left, "key"), + "left_value": ColumnRef(left, "left_value"), + "key_right": ColumnRef(right, "key"), + "right_value": ColumnRef(right, "right_value"), + } + + +def test_inner_join_omits_coalesced_right_key() -> None: + left = make_scan("key", "left_value") + right = make_scan("key", "right_value") + node = Join( + {"key": I64, "left_value": I64, "right_value": I64}, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Inner", False, None, "_right", True, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnRef(left, "key"), + "left_value": ColumnRef(left, "left_value"), + "right_value": ColumnRef(right, "right_value"), + } + + +def test_semi_join_binds_only_left_columns() -> None: + left = make_scan("key", "value") + right = make_scan("key") + node = Join( + left.schema, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Semi", False, None, "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnRef(left, "key"), + "value": ColumnRef(left, "value"), + } + + +def test_outer_join_has_no_column_domain_bindings() -> None: + left = make_scan("key") + right = make_scan("other") + node = Join( + {**left.schema, **right.schema}, + (named_col("key", "key"),), + (named_col("other", "other"),), + ("Left", False, None, "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == {} + + +def test_passthrough_nodes_bind_same_named_columns() -> None: + child = make_scan("a", "b") + mask = expr.NamedExpr("mask", expr.Literal(BOOL, True)) # noqa: FBT003 + nodes = ( + Cache(child.schema, 1, None, child), + Filter(child.schema, mask, child), + Projection({"b": I64}, child), + Slice(child.schema, 0, 1, child), + Distinct( + child.schema, + plc.stream_compaction.DuplicateKeepOption.KEEP_ANY, + None, + (0, 1), + False, # noqa: FBT003 + child, + ), + Sort( + child.schema, + (named_col("a", "a"),), + (plc.types.Order.ASCENDING,), + (plc.types.NullOrder.AFTER,), + False, # noqa: FBT003 + (0, 1), + child, + ), + ) + + for node in nodes: + assert column_domain_bindings(node) == { + name: ColumnRef(child, name) for name in node.schema + } From 135af451be067286c3721396d77cd6bcdf2c955d Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 13 Jul 2026 17:46:23 +0100 Subject: [PATCH 45/65] Use column_domain_bindings in select_column_targets --- .../cudf_polars/cudf_polars/streaming/actor_graph/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 8e25c8166013..a8f1e2b62af7 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -40,6 +40,7 @@ from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select from cudf_polars.dsl.tracing import Scope +from cudf_polars.dsl.utils.column_domain import column_domain_bindings from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk @@ -430,9 +431,8 @@ def _derived_ordering( def _select_column_targets(select: Select) -> dict[str, dict[str, None]]: old_to_new_names: defaultdict[str, dict[str, None]] = defaultdict(dict) - for ne in select.exprs: - if isinstance(ne.value, Col): - old_to_new_names[ne.value.name][ne.name] = None + for output_name, source in column_domain_bindings(select).items(): + old_to_new_names[source.name][output_name] = None return dict(old_to_new_names) From 7e583208a92a90a54de558cafd83009510ef7e81 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 13 Jul 2026 19:38:54 +0100 Subject: [PATCH 46/65] WIP: gather plan facts for join prefilter in one place --- .../cudf_polars/dsl/utils/column_domain.py | 19 +- .../streaming/join_domain_prefilter.py | 318 ++++++++---------- .../streaming/test_join_domain_prefilter.py | 64 +++- 3 files changed, 210 insertions(+), 191 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py index b2ad4a7d7aeb..e9cb02014994 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -24,11 +24,11 @@ ) if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Iterator, Mapping from cudf_polars.dsl.ir import IR -__all__ = ["ColumnRef", "column_domain_bindings"] +__all__ = ["ColumnLineage", "ColumnRef", "column_domain_bindings"] @dataclass(frozen=True) @@ -39,6 +39,21 @@ class ColumnRef: name: str +@dataclass(frozen=True) +class ColumnLineage: + """Persistent value-domain lineage, sharing suffixes across DAG branches.""" + + column: ColumnRef + source: ColumnLineage | None = None + + def __iter__(self) -> Iterator[ColumnRef]: + """Iterate from the output column towards its furthest known source.""" + lineage: ColumnLineage | None = self + while lineage is not None: + yield lineage.column + lineage = lineage.source + + @singledispatch def column_domain_bindings(node: IR) -> Mapping[str, ColumnRef]: """ 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 fd9c75f120e2..a8f5ea2688cc 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, + Slice, Sort, ) from cudf_polars.dsl.tracing import Scope, log @@ -30,10 +31,15 @@ reuse_if_unchanged, traversal, ) +from cudf_polars.dsl.utils.column_domain import ( + ColumnLineage, + ColumnRef, + column_domain_bindings, +) from cudf_polars.dsl.utils.replace import replace if TYPE_CHECKING: - from collections.abc import Iterable, Sequence + from collections.abc import Iterable, Iterator, Mapping, Sequence from cudf_polars.streaming.base import StatsCollector from cudf_polars.typing import GenericTransformer @@ -89,14 +95,90 @@ def score(self) -> tuple[int, int, int]: ) +@dataclass(frozen=True) +class PlanFacts: + """Facts derived in one bottom-up traversal of an IR DAG.""" + + row_estimates: Mapping[IR, int | None] + selective_nodes: frozenset[IR] + column_lineages: Mapping[ColumnRef, ColumnLineage] + + class _RewriteState(TypedDict): """State shared by the join-domain prefilter DAG rewrite.""" threshold: float trace: bool stats: StatsCollector - row_estimates: dict[IR, int | None] - selective_nodes: set[IR] + facts: PlanFacts + + +def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: + """Derive row, selectivity, and column-domain facts in post-order.""" + row_estimates: dict[IR, int | None] = {} + selective_nodes: set[IR] = set() + column_lineages: dict[ColumnRef, ColumnLineage] = {} + + for node in post_traversal([ir]): + if isinstance(node, (Scan, DataFrameScan)): + source_info = stats.scan_stats.get(node) + rows = None if source_info is None else source_info.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 = row_estimates[node.children[0]] + elif isinstance(node, Join): + rows = _estimate_join_rows( + node.options[0], + row_estimates[node.children[0]], + row_estimates[node.children[1]], + ) + else: + child_estimates = [ + estimate + for child in node.children + if (estimate := row_estimates[child]) is not None + ] + rows = max(child_estimates, default=None) + row_estimates[node] = rows + + if ( + (isinstance(node, Scan) and node.predicate is not None) + or isinstance(node, Filter) + or any(child in selective_nodes for child in node.children) + ): + selective_nodes.add(node) + + bindings = column_domain_bindings(node) + for name in node.schema: + column = ColumnRef(node, name) + source = bindings.get(name) + source_lineage = None if source is None else column_lineages[source] + column_lineages[column] = ColumnLineage(column, source_lineage) + + return PlanFacts( + row_estimates=row_estimates, + selective_nodes=frozenset(selective_nodes), + column_lineages=column_lineages, + ) + + +def prefilter_lineage(facts: PlanFacts, root: IR, column: str) -> Iterator[ColumnRef]: + """Iterate over column-domain lineage valid for prefilter insertion.""" + lineage = facts.column_lineages.get(ColumnRef(root, column)) + if lineage is None: + return + for reference in lineage: + yield reference + node = reference.node + if isinstance(node, Slice): + return + if isinstance(node, (Distinct, GroupBy, Sort)) and node.zlice is not None: + return + if isinstance(node, Join) and node.options[2] is not None: + return def optimize_join_domain_prefilters( @@ -123,8 +205,7 @@ def optimize_join_domain_prefilters( threshold=threshold, trace=trace, stats=stats, - row_estimates=_estimate_row_counts(ir, stats), - selective_nodes=_collect_selective_nodes(ir), + facts=analyze_plan(ir, stats), ) mapper: GenericTransformer[IR, IR, _RewriteState] = CachingVisitor( _rewrite, state=state @@ -149,19 +230,16 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: assert isinstance(rewritten, Join) node = rewritten if node is original: - row_estimates = rec.state["row_estimates"] - selective_nodes = rec.state["selective_nodes"] + facts = rec.state["facts"] 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) + facts = analyze_plan(node, rec.state["stats"]) candidate, reason = _select_candidate( node, rec.state["threshold"], - row_estimates, - selective_nodes, + facts, ) if rec.state["trace"]: _trace_decision(node, rec.state["threshold"], candidate, reason) @@ -191,8 +269,7 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: def _select_candidate( ir: Join, threshold: float, - row_estimates: dict[IR, int | None], - selective_nodes: set[IR], + facts: PlanFacts, ) -> tuple[_Candidate | None, str]: if ir.options[0] != "Inner": return None, "not_inner_join" @@ -230,8 +307,7 @@ def _select_candidate( target_keys, domain_keys, threshold, - row_estimates, - selective_nodes, + facts, ) ) candidates.extend( @@ -242,8 +318,7 @@ def _select_candidate( target_keys, domain_keys, threshold, - row_estimates, - selective_nodes, + facts, ) ) @@ -263,18 +338,16 @@ 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], + facts: PlanFacts, ) -> Iterable[_Candidate]: for target_key, domain_key in zip(target_keys, domain_keys, strict=True): - target = _largest_key_source(target_child, target_key.name, row_estimates) + target = _largest_key_source(target_child, target_key.name, facts) if target is None: continue domain = _smallest_key_producer( domain_child, domain_key.name, - row_estimates, - selective_nodes, + facts, require_selective=True, ) if domain is None: @@ -301,8 +374,7 @@ 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], + facts: PlanFacts, ) -> Iterable[_Candidate]: if len(target_keys) < 2: return @@ -310,7 +382,7 @@ 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, row_estimates) + target = _largest_key_source(target_child, target_key.name, facts) if target is None: continue @@ -323,7 +395,7 @@ def _composite_candidates( domain = _smallest_node_containing_all( domain_child, (domain_key.name, domain_constraint_key.name), - row_estimates, + facts, ) if domain is None: continue @@ -332,8 +404,7 @@ def _composite_candidates( constraint_domain = _smallest_key_producer( target_child, target_constraint_key.name, - row_estimates, - selective_nodes, + facts, require_selective=True, exclude=target.node, ) @@ -429,20 +500,20 @@ def _make_semi_join( def _smallest_key_producer( root: IR, column: str, - row_estimates: dict[IR, int | None], - selective_nodes: set[IR], + facts: PlanFacts, *, require_selective: bool, exclude: IR | None = None, ) -> _Producer | None: candidates = [] - for node, bound_column in _column_bindings(root, column): + for reference in prefilter_lineage(facts, root, column): + node, bound_column = reference.node, reference.name if node is exclude: continue - rows = row_estimates.get(node) + rows = facts.row_estimates.get(node) if rows is None or rows <= 0: continue - if require_selective and node not in selective_nodes: + if require_selective and node not in facts.selective_nodes: continue candidates.append( (rows, len(node.schema), _Producer(node, (bound_column,), rows)) @@ -453,49 +524,47 @@ def _smallest_key_producer( def _smallest_node_containing_all( - root: IR, columns: Sequence[str], row_estimates: dict[IR, int | None] + root: IR, columns: Sequence[str], facts: PlanFacts ) -> _Producer | None: candidates = [] - lineages = [tuple(_column_bindings(root, column)) for column in columns] + lineages = [ + { + id(reference.node): reference + for reference in prefilter_lineage(facts, 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), - ) + # A common producer must be the same node occurrence in the IR DAG, not + # merely an equal node, so index each lineage by object identity. + for node_id, first_reference in lineages[0].items(): + node = first_reference.node + try: + bound_columns = tuple(lineage[node_id].name for lineage in lineages) + except KeyError: + continue + rows = facts.row_estimates.get(node) + if rows is None or rows <= 0: + continue + candidates.append( + ( + rows, + len(node.schema), + _Producer(node, bound_columns, 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, row_estimates: dict[IR, int | None] -) -> _Producer | None: +def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: source_candidates = [] fallback_candidates = [] - for node, bound_column in _column_bindings(root, column): - rows = row_estimates.get(node) + for reference in prefilter_lineage(facts, root, column): + node, bound_column = reference.node, reference.name + rows = facts.row_estimates.get(node) if rows is None or rows <= 0: continue item = (rows, len(node.schema), _Producer(node, (bound_column,), rows)) @@ -509,115 +578,6 @@ 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[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)) - 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]): - 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, default=None) - estimates[node] = rows - return estimates - - def _estimate_join_rows( how: str, left_rows: int | None, right_rows: int | None ) -> int | None: @@ -636,18 +596,6 @@ def _estimate_join_rows( return None -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 - - def _contains_identity(root: IR, needle: IR) -> bool: return any(node is needle for node in traversal([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 0aead0061d49..8e611056558f 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -12,13 +12,17 @@ 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 Join, Scan, Select, Slice from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.utils.column_domain import ColumnRef from cudf_polars.engine.default_singleton_engine import DefaultSingletonEngine from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.join_domain_prefilter import ( + PlanFacts, _smallest_node_containing_all, + analyze_plan, optimize_join_domain_prefilters, + prefilter_lineage, ) from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -409,10 +413,14 @@ def test_composite_domain_columns_follow_renames() -> None: domain_constraint="raw_constraint", ) + analyzed = analyze_plan(renamed, _stats(source=(source, 10))) + facts = PlanFacts( + row_estimates={renamed: 20, source: 10}, + selective_nodes=analyzed.selective_nodes, + column_lineages=analyzed.column_lineages, + ) producer = _smallest_node_containing_all( - renamed, - ("domain_key", "domain_constraint"), - {renamed: 20, source: 10}, + renamed, ("domain_key", "domain_constraint"), facts ) assert producer is not None @@ -420,6 +428,54 @@ def test_composite_domain_columns_follow_renames() -> None: assert producer.columns == ("raw_key", "raw_constraint") +def test_plan_facts_share_lineage_suffixes_across_shared_dag() -> None: + source = _scan("source", ("raw_key",)) + left = _select(source, left_key="raw_key") + right = _select(source, right_key="raw_key") + root = _join(left, right, ("left_key",), ("right_key",)) + + facts = analyze_plan(root, _stats(source=(source, 10))) + left_lineage = facts.column_lineages[ColumnRef(left, "left_key")] + right_lineage = facts.column_lineages[ColumnRef(right, "right_key")] + + assert left_lineage.source is right_lineage.source + assert tuple(left_lineage) == ( + ColumnRef(left, "left_key"), + ColumnRef(source, "raw_key"), + ) + assert tuple(right_lineage) == ( + ColumnRef(right, "right_key"), + ColumnRef(source, "raw_key"), + ) + + +def test_target_prefilter_does_not_move_below_slice() -> None: + target = _scan("target", ("target_key",)) + sliced = Slice(target.schema, 0, 100, target) + domain = _scan("domain", ("domain_key",), predicate=True) + root = _join(sliced, domain, ("target_key",), ("domain_key",)) + + stats = _stats(target=(target, 1_000), domain=(domain, 5)) + facts = analyze_plan(root, stats) + assert tuple(facts.column_lineages[ColumnRef(sliced, "target_key")]) == ( + ColumnRef(sliced, "target_key"), + ColumnRef(target, "target_key"), + ) + assert tuple(prefilter_lineage(facts, sliced, "target_key")) == ( + ColumnRef(sliced, "target_key"), + ) + + optimized = optimize_join_domain_prefilters( + root, + stats, + _config(), + ) + + semis = _joins(optimized, "Semi") + assert any(semi.children[0] is sliced for semi in semis) + assert not any(semi.children[0] is target for semi in semis) + + 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) From 4ae2b152a6bed03fed7aa843f04fcca904d786e6 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 11:25:23 +0100 Subject: [PATCH 47/65] Refactor semijoin pushdown candidates --- .../streaming/join_domain_prefilter.py | 73 +++++++++++++++---- .../streaming/test_join_domain_prefilter.py | 62 ++++++++++++---- 2 files changed, 109 insertions(+), 26 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 a8f5ea2688cc..dc2585fcc762 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -12,6 +12,7 @@ from cudf_polars.dsl.ir import ( IR, Cache, + ConditionalJoin, DataFrameScan, Distinct, Filter, @@ -19,10 +20,12 @@ HStack, Join, Projection, + Rolling, Scan, Select, Slice, Sort, + Union, ) from cudf_polars.dsl.tracing import Scope, log from cudf_polars.dsl.traversal import ( @@ -165,19 +168,63 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: ) -def prefilter_lineage(facts: PlanFacts, root: IR, column: str) -> Iterator[ColumnRef]: - """Iterate over column-domain lineage valid for prefilter insertion.""" - lineage = facts.column_lineages.get(ColumnRef(root, column)) - if lineage is None: +def blocks_pushdown(node: IR) -> bool: + """ + Return whether a node blocks filter pushdown. + + Parameters + ---------- + node + Node to check + + Returns + ------- + bool + True if a semijoin cannot be pushed past this node, otherwise False. + """ + return ( + # TODO: Distinct and Rolling only block pushdown in some + # circumstances, but we'd need to make the logic more complicated: + # - We can push through distinct if the filter applies to the columns + # that are being used to determine distinct rows + # - We can push through rolling if the filter applies to the + # groupby keys. + # TODO: We can push through an unsliced Union, but need to + # distribute the filter onto every child. + isinstance(node, (Distinct, Rolling, Slice, Union)) + # Can't push through anything that is sliced. + or (isinstance(node, (GroupBy, Sort)) and node.zlice is not None) + or (isinstance(node, (ConditionalJoin, Join)) and node.options[2] is not None) + ) + + +def semijoin_pushdown_candidates( + facts: PlanFacts, root: IR, column: str +) -> Iterator[ColumnRef]: + """ + Yield column domain lineage providing valid locations for semijoin pushdown. + + Parameters + ---------- + facts + Gathered facts about the plan + root + Root node to search from + column + Name of column we're finding the lineage of. + + Returns + ------- + Iterator + Of valid insertion points for a semijoin filter on the given column name. + """ + try: + lineage = facts.column_lineages[ColumnRef(root, column)] + except KeyError: return for reference in lineage: yield reference - node = reference.node - if isinstance(node, Slice): - return - if isinstance(node, (Distinct, GroupBy, Sort)) and node.zlice is not None: - return - if isinstance(node, Join) and node.options[2] is not None: + if blocks_pushdown(reference.node): return @@ -506,7 +553,7 @@ def _smallest_key_producer( exclude: IR | None = None, ) -> _Producer | None: candidates = [] - for reference in prefilter_lineage(facts, root, column): + for reference in semijoin_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name if node is exclude: continue @@ -530,7 +577,7 @@ def _smallest_node_containing_all( lineages = [ { id(reference.node): reference - for reference in prefilter_lineage(facts, root, column) + for reference in semijoin_pushdown_candidates(facts, root, column) } for column in columns ] @@ -562,7 +609,7 @@ def _smallest_node_containing_all( def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: source_candidates = [] fallback_candidates = [] - for reference in prefilter_lineage(facts, root, column): + for reference in semijoin_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name rows = facts.row_estimates.get(node) if rows is None or rows <= 0: 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 8e611056558f..b2e8903d62b5 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -12,17 +12,17 @@ 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, Slice +from cudf_polars.dsl.ir import Distinct, Join, Scan, Select, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.column_domain import ColumnRef -from cudf_polars.engine.default_singleton_engine import DefaultSingletonEngine +from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.join_domain_prefilter import ( PlanFacts, _smallest_node_containing_all, analyze_plan, optimize_join_domain_prefilters, - prefilter_lineage, + semijoin_pushdown_candidates, ) from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -32,12 +32,24 @@ import concurrent.futures from cudf_polars.dsl.ir import IR + from cudf_polars.engine.spmd import SPMDEngine from cudf_polars.streaming.base import SerializedDataSourceInfo I64 = DataType(pl.Int64()) BOOL = DataType(pl.Boolean()) +@pytest.fixture +def engine(spmd_engine_factory) -> SPMDEngine: + """Return an SPMD engine configured for join-domain prefilter tests.""" + return spmd_engine_factory( + StreamingOptions( + join_domain_prefilter={"threshold": 0.5}, + raise_on_fail=True, + ) + ) + + class _SourceInfo: type: Literal["parquet"] = "parquet" @@ -200,6 +212,7 @@ def test_domain_prefilter_can_be_disabled() -> None: ) def test_nullable_join_keys_preserve_results( nulls_equal: bool, # noqa: FBT001 + engine: SPMDEngine, parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ) -> None: domain = pl.LazyFrame( @@ -215,11 +228,6 @@ def test_nullable_join_keys_preserve_results( } ) query = domain.join(target, on="key", nulls_equal=nulls_equal) - engine = pl.GPUEngine( - executor="streaming", - raise_on_fail=True, - executor_options={"join_domain_prefilter": {"threshold": 0.5}}, - ) ir = Translator(query._ldf.visit(), engine).translate_ir() config = ConfigOptions.from_polars_engine(engine) @@ -232,10 +240,38 @@ 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) - try: - assert_gpu_result_equal(query, engine=engine, check_row_order=False) - finally: - DefaultSingletonEngine.shutdown() + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_prefilter_does_not_move_below_distinct_on_non_subset_column( + engine: SPMDEngine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + target = pl.LazyFrame( + { + "group": [1, 1] * 100, + "key": [0, 1] * 100, + } + ).unique(subset="group", keep="first", maintain_order=True) + domain = pl.LazyFrame( + { + "key": [1, 2], + "active": [True, False], + } + ).filter("active") + query = target.join(domain, on="key") + + 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, + ) + + semis = _joins(optimized, "Semi") + assert any(isinstance(semi.children[0], Distinct) for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_no_simple_domain_prefilter_when_domain_is_not_selective() -> None: @@ -461,7 +497,7 @@ def test_target_prefilter_does_not_move_below_slice() -> None: ColumnRef(sliced, "target_key"), ColumnRef(target, "target_key"), ) - assert tuple(prefilter_lineage(facts, sliced, "target_key")) == ( + assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( ColumnRef(sliced, "target_key"), ) From 4c2e988587c849dd1c32d6d54c0bb95ddd280485 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 12:09:13 +0100 Subject: [PATCH 48/65] docstring --- .../streaming/join_domain_prefilter.py | 15 ++++++++++++++- 1 file changed, 14 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 dc2585fcc762..7edc7cdccca5 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -117,7 +117,20 @@ class _RewriteState(TypedDict): def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: - """Derive row, selectivity, and column-domain facts in post-order.""" + """ + Derive row, selectivity, and column-domain facts for an IR DAG + + Parameters + ---------- + ir + Root node to gather facts for. + stats + Pre-populated statistics + + Returns + ------- + Gather facts about the plan. + """ row_estimates: dict[IR, int | None] = {} selective_nodes: set[IR] = set() column_lineages: dict[ColumnRef, ColumnLineage] = {} From 6da9b87e04f2dd320ce8fbc83b2b1e4179ead87a Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 12:09:55 +0100 Subject: [PATCH 49/65] fixup --- .../cudf_polars/cudf_polars/streaming/join_domain_prefilter.py | 2 +- 1 file changed, 1 insertion(+), 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 7edc7cdccca5..940979b48ad7 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -118,7 +118,7 @@ class _RewriteState(TypedDict): def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: """ - Derive row, selectivity, and column-domain facts for an IR DAG + Derive row, selectivity, and column-domain facts for an IR DAG. Parameters ---------- From ccc95034c7684b8fb9b40a7283e535cf566a3d00 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 12:33:05 +0100 Subject: [PATCH 50/65] More refactoring --- .../streaming/join_domain_prefilter.py | 142 ++++++++++-------- .../streaming/test_join_domain_prefilter.py | 38 +++-- 2 files changed, 103 insertions(+), 77 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 940979b48ad7..d303c746abb6 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -6,7 +6,7 @@ from dataclasses import dataclass from functools import singledispatch -from typing import TYPE_CHECKING, Any, Literal, TypedDict +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict from cudf_polars.dsl import expr from cudf_polars.dsl.ir import ( @@ -64,38 +64,59 @@ def column(self) -> str: @dataclass(frozen=True) -class _Candidate: - """A derived key-domain prefilter candidate.""" +class SimpleCandidate: + """A direct key-domain prefilter candidate.""" - mode: Literal["simple", "composite"] + mode = "simple" target_side: Literal["left", "right"] target: _Producer 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 @property - def domain_rows(self) -> int: - """Estimated rows in the domain input.""" - return self.domain.rows + def score(self) -> tuple[int, int, int]: + """Rank after composite candidates, then by domain size.""" + return (1, self.domain.rows, self.domain.rows) + + +@dataclass(frozen=True) +class CompositeCandidate: + """A key-domain prefilter constrained by another join key.""" + + mode = "composite" + target_side: Literal["left", "right"] + target: _Producer + target_key: expr.Col + domain: _Producer + domain_key: expr.Col + constraint_domain: _Producer + domain_constraint_key: expr.Col + target_constraint_key: expr.Col @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, - ) + """Prefer smaller constraint and domain inputs.""" + return (0, self.constraint_domain.rows, self.domain.rows) + + +Candidate: TypeAlias = SimpleCandidate | CompositeCandidate +DecisionReason: TypeAlias = Literal[ + "applied", + "maintain_order", + "no_selective_domain", + "non_column_join_key", + "not_inner_join", + "sliced_join", +] + + +@dataclass(frozen=True) +class Decision: + """Result of considering a join for a domain prefilter.""" + + reason: DecisionReason + candidate: Candidate | None = None @dataclass(frozen=True) @@ -250,8 +271,7 @@ def optimize_join_domain_prefilters( 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. + column equality keys are considered. """ options = config_options.executor.join_domain_prefilter if options is None: @@ -296,26 +316,30 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: # Re-analyze that current subtree so parent joins can use the derived # selectivity and cardinality when ranking their own candidates. facts = analyze_plan(node, rec.state["stats"]) - candidate, reason = _select_candidate( + decision = _select_candidate( node, rec.state["threshold"], facts, ) if rec.state["trace"]: - _trace_decision(node, rec.state["threshold"], candidate, reason) - if candidate is None: + _trace_decision(node, rec.state["threshold"], decision) + if decision.candidate is None: return node + return apply_candidate(node, decision.candidate) + - left, right = node.children - domain = _make_domain(candidate, node) +def apply_candidate(ir: Join, candidate: Candidate) -> IR: + """Apply a selected join-domain prefilter candidate to a join.""" + left, right = ir.children + domain = _make_domain(candidate, ir) 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], + nulls_equal=ir.options[1], + suffix=ir.options[3], ) # A DAG may share the target with the domain side, so only rewrite the # side for which this candidate was selected. @@ -323,27 +347,27 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: (left,) = replace([left], {candidate.target.node: target_filter}) else: (right,) = replace([right], {candidate.target.node: target_filter}) - return node.reconstruct((left, right)) + return ir.reconstruct((left, right)) def _select_candidate( ir: Join, threshold: float, facts: PlanFacts, -) -> tuple[_Candidate | None, str]: +) -> Decision: if ir.options[0] != "Inner": - return None, "not_inner_join" + return Decision(reason="not_inner_join") if ir.options[2] is not None: - return None, "sliced_join" + return Decision(reason="sliced_join") if ir.options[5] != "none": - return None, "maintain_order" + return Decision(reason="maintain_order") 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" + return Decision(reason="non_column_join_key") - candidates: list[_Candidate] = [] + candidates: list[Candidate] = [] left: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( "left", ir.children[0], @@ -383,8 +407,8 @@ def _select_candidate( ) if not candidates: - return None, "no_selective_domain" - return min(candidates, key=lambda c: c.score), "applied" + return Decision(reason="no_selective_domain") + return Decision(reason="applied", candidate=min(candidates, key=lambda c: c.score)) def _simple_keys(keys: Sequence[expr.NamedExpr]) -> tuple[expr.Col, ...]: @@ -399,7 +423,7 @@ def _simple_candidates( domain_keys: tuple[expr.Col, ...], threshold: float, facts: PlanFacts, -) -> Iterable[_Candidate]: +) -> Iterable[SimpleCandidate]: for target_key, domain_key in zip(target_keys, domain_keys, strict=True): target = _largest_key_source(target_child, target_key.name, facts) if target is None: @@ -416,14 +440,12 @@ def _simple_candidates( continue if _contains_identity(target.node, domain.node): continue - yield _Candidate( - mode="simple", + yield SimpleCandidate( target_side=target_side, target=target, target_key=target_key, domain=domain, domain_key=domain_key, - target_rows=target.rows, ) @@ -435,7 +457,7 @@ def _composite_candidates( domain_keys: tuple[expr.Col, ...], threshold: float, facts: PlanFacts, -) -> Iterable[_Candidate]: +) -> Iterable[CompositeCandidate]: if len(target_keys) < 2: return @@ -476,32 +498,26 @@ def _composite_candidates( target.node, constraint_domain.node ): continue - yield _Candidate( - mode="composite", + yield CompositeCandidate( target_side=target_side, target=target, 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, ) -def _make_domain(candidate: _Candidate, ir: Join) -> IR: - if candidate.mode == "simple": +def _make_domain(candidate: Candidate, ir: Join) -> IR: + if isinstance(candidate, SimpleCandidate): return _project_bound_key( candidate.domain.node, candidate.domain.column, 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 = _project_bound_key( candidate.constraint_domain.node, candidate.constraint_domain.column, @@ -660,13 +676,11 @@ def _contains_identity(root: IR, needle: IR) -> bool: return any(node is needle for node in traversal([root])) -def _trace_decision( - ir: Join, threshold: float, candidate: _Candidate | None, reason: str -) -> None: +def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: join_domain_prefilter: dict[str, Any] = { "considered": True, "threshold": threshold, - "reason": reason, + "reason": decision.reason, } record = { "scope": Scope.PLAN.value, @@ -674,25 +688,23 @@ def _trace_decision( "actor_ir_id": ir.get_stable_id(), "actor_ir_type": type(ir).__name__, } - if candidate is not None: + if (candidate := decision.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, + "estimated_target_rows": candidate.target.rows, + "estimated_domain_rows": candidate.domain.rows, "target_node_type": type(candidate.target.node).__name__, "domain_node_type": type(candidate.domain.node).__name__, } ) - if candidate.constraint_domain is not None: + if isinstance(candidate, CompositeCandidate): join_domain_prefilter.update( { - "constraint_key": candidate.target_constraint_key.name - if candidate.target_constraint_key is not None - else None, + "constraint_key": candidate.target_constraint_key.name, "estimated_constraint_rows": candidate.constraint_domain.rows, } ) 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 b2e8903d62b5..196453bf79ef 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -18,9 +18,14 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.join_domain_prefilter import ( + CompositeCandidate, + Decision, PlanFacts, + SimpleCandidate, + _select_candidate, _smallest_node_containing_all, analyze_plan, + apply_candidate, optimize_join_domain_prefilters, semijoin_pushdown_candidates, ) @@ -165,11 +170,12 @@ def test_simple_domain_prefilter_filters_large_side() -> None: 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(), - ) + facts = analyze_plan(root, _stats(part=(part, 6), lineitem=(lineitem, 1_800))) + decision = _select_candidate(root, 0.5, facts) + + assert decision.reason == "applied" + assert isinstance(decision.candidate, SimpleCandidate) + optimized = apply_candidate(root, decision.candidate) assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" @@ -278,13 +284,17 @@ 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",)) + stats = _stats(supplier=(supplier, 30), lineitem=(lineitem, 1_800)) + + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) optimized = optimize_join_domain_prefilters( root, - _stats(supplier=(supplier, 30), lineitem=(lineitem, 1_800)), + stats, _config(), ) + assert decision == Decision(reason="no_selective_domain") assert optimized is root assert not _joins(optimized, "Semi") @@ -310,17 +320,21 @@ def test_composite_domain_prefilter_constrains_domain_first() -> None: ("s_suppkey", "s_nationkey"), ) + stats = _stats( + nation=(nation, 5), + orders=(orders, 900), + lineitem=(lineitem, 1_800), + supplier=(supplier, 30), + ) + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) optimized = optimize_join_domain_prefilters( root, - _stats( - nation=(nation, 5), - orders=(orders, 900), - lineitem=(lineitem, 1_800), - supplier=(supplier, 30), - ), + stats, _config(), ) + assert decision.reason == "applied" + assert isinstance(decision.candidate, CompositeCandidate) semis = _joins(optimized, "Semi") assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" From 7c8393aea9bc50fe3d5f14eeb224b28d9dcdfeb9 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 13:57:46 +0100 Subject: [PATCH 51/65] Disambiguate column lineage with edges --- .../cudf_polars/dsl/utils/column_domain.py | 2 + .../streaming/join_domain_prefilter.py | 90 ++++++++++++------- .../streaming/test_join_domain_prefilter.py | 35 +++++++- 3 files changed, 93 insertions(+), 34 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py index e9cb02014994..7959a18a1e9e 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -45,6 +45,8 @@ class ColumnLineage: column: ColumnRef source: ColumnLineage | None = None + source_child_index: int | None = None + """Unique child edge leading to ``source``, or None if absent or ambiguous.""" def __iter__(self) -> Iterator[ColumnRef]: """Iterate from the output column towards its furthest known source.""" 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 d303c746abb6..866f9e842353 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -192,8 +192,22 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: for name in node.schema: column = ColumnRef(node, name) source = bindings.get(name) - source_lineage = None if source is None else column_lineages[source] - column_lineages[column] = ColumnLineage(column, source_lineage) + if source is None: + source_lineage = None + source_child_index = None + else: + source_lineage = column_lineages[source] + source_children = tuple( + index + for index, child in enumerate(node.children) + if child == source.node + ) + source_child_index = ( + source_children[0] if len(source_children) == 1 else None + ) + column_lineages[column] = ColumnLineage( + column, source_lineage, source_child_index + ) return PlanFacts( row_estimates=row_estimates, @@ -256,10 +270,15 @@ def semijoin_pushdown_candidates( lineage = facts.column_lineages[ColumnRef(root, column)] except KeyError: return - for reference in lineage: - yield reference - if blocks_pushdown(reference.node): + while True: + yield lineage.column + if ( + blocks_pushdown(lineage.column.node) + or lineage.source is None + or lineage.source_child_index is None + ): return + lineage = lineage.source def optimize_join_domain_prefilters( @@ -438,7 +457,7 @@ def _simple_candidates( continue if domain.rows / target.rows > threshold: continue - if _contains_identity(target.node, domain.node): + if contains_node(target.node, domain.node): continue yield SimpleCandidate( target_side=target_side, @@ -494,7 +513,7 @@ def _composite_candidates( continue if constraint_domain.rows / domain.rows > threshold: continue - if _contains_identity(target.node, domain.node) or _contains_identity( + if contains_node(target.node, domain.node) or contains_node( target.node, constraint_domain.node ): continue @@ -603,33 +622,37 @@ def _smallest_node_containing_all( root: IR, columns: Sequence[str], facts: PlanFacts ) -> _Producer | None: candidates = [] - lineages = [ - { - id(reference.node): reference - for reference in semijoin_pushdown_candidates(facts, root, column) - } - for column in columns - ] - if not lineages or any(not lineage for lineage in lineages): + lineages: list[ColumnLineage] = [] + for column in columns: + lineage = facts.column_lineages.get(ColumnRef(root, column)) + if lineage is None: + return None + lineages.append(lineage) + if not lineages: return None - # A common producer must be the same node occurrence in the IR DAG, not - # merely an equal node, so index each lineage by object identity. - for node_id, first_reference in lineages[0].items(): - node = first_reference.node - try: - bound_columns = tuple(lineage[node_id].name for lineage in lineages) - except KeyError: - continue + while True: + node = lineages[0].column.node + if any(lineage.column.node != node for lineage in lineages[1:]): + break + bound_columns = tuple(lineage.column.name for lineage in lineages) rows = facts.row_estimates.get(node) - if rows is None or rows <= 0: - continue - candidates.append( - ( - rows, - len(node.schema), - _Producer(node, bound_columns, rows), + if rows is not None and rows > 0: + candidates.append( + ( + rows, + len(node.schema), + _Producer(node, bound_columns, rows), + ) ) - ) + if blocks_pushdown(node): + break + source_child_indices = {lineage.source_child_index for lineage in lineages} + if len(source_child_indices) != 1 or None in source_child_indices: + break + sources = [lineage.source for lineage in lineages] + if any(source is None for source in sources): + break + lineages = [source for source in sources if source is not None] if not candidates: return None return min(candidates, key=lambda item: (item[0], item[1]))[2] @@ -672,8 +695,9 @@ def _estimate_join_rows( return None -def _contains_identity(root: IR, needle: IR) -> bool: - return any(node is needle for node in traversal([root])) +def contains_node(root: IR, needle: IR) -> bool: + """Return whether an equal node occurs in a DAG rooted at ``root``.""" + return needle in traversal([root]) def _trace_decision(ir: Join, threshold: float, decision: Decision) -> 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 196453bf79ef..7c3a8e700262 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 Distinct, Join, Scan, Select, Slice +from cudf_polars.dsl.ir import Cache, Distinct, Join, Scan, Select, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.column_domain import ColumnRef from cudf_polars.engine.options import StreamingOptions @@ -26,6 +26,7 @@ _smallest_node_containing_all, analyze_plan, apply_candidate, + contains_node, optimize_join_domain_prefilters, semijoin_pushdown_candidates, ) @@ -478,6 +479,38 @@ def test_composite_domain_columns_follow_renames() -> None: assert producer.columns == ("raw_key", "raw_constraint") +def test_composite_domain_columns_do_not_reconverge_after_join( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame({"key": [1, 1, 2], "value": [10, 20, 30]}) + query = source.join(source, on="key", suffix="_right") + joined = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(joined, Join) + assert isinstance(joined.children[0], Cache) + assert joined.children[0] is joined.children[1] + + facts = analyze_plan(joined, StatsCollector()) + producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) + + assert tuple(semijoin_pushdown_candidates(facts, joined, "value")) == ( + ColumnRef(joined, "value"), + ) + assert producer is not None + assert producer.node is joined + assert producer.columns == ("value", "value_right") + + +def test_contains_node_uses_dag_equality() -> None: + source = _scan("source", ("key",)) + equal_source = _scan("source", ("key",)) + root = _select(source, key="key") + + assert source is not equal_source + assert source == equal_source + assert contains_node(root, equal_source) + + def test_plan_facts_share_lineage_suffixes_across_shared_dag() -> None: source = _scan("source", ("raw_key",)) left = _select(source, left_key="raw_key") From f9926f8448cf8f2d2be002112351b003e8877fbf Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 14:45:20 +0100 Subject: [PATCH 52/65] Remove Cache nodes when lowering to streaming engine The refcount-based tracking that inserts fanouts doesn't need the Cache at all, and removing them will simplify join prefilter lineage determination. --- .../cudf_polars/streaming/parallel.py | 16 ++++++++++- .../cudf_polars/tests/streaming/test_join.py | 14 ++++++---- .../tests/streaming/test_parallel.py | 28 +++++++++++++++++++ 3 files changed, 51 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index de88570440c9..5cbe812e62e0 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -36,7 +36,7 @@ Slice, Union, ) -from cudf_polars.dsl.traversal import CachingVisitor, traversal +from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged, traversal from cudf_polars.dsl.utils.naming import unique_names from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node @@ -53,6 +53,7 @@ from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.dispatch import LowerIRTransformer, State + from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -77,6 +78,18 @@ class LoweringInfo: ] # Partition mapping for nodes in the lowered IR. +def remove_cache_nodes(ir: IR) -> IR: + """Remove logical cache nodes while preserving shared DAG structure.""" + + def rewrite(node: IR, rec: GenericTransformer[IR, IR, None]) -> IR: + if isinstance(node, Cache): + return rec(node.children[0]) + return reuse_if_unchanged(node, rec) + + mapper: GenericTransformer[IR, IR, None] = CachingVisitor(rewrite, state=None) + return mapper(ir) + + def optimize_with_stats( ir: IR, config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector ) -> IR: @@ -101,6 +114,7 @@ def optimize_with_stats( optimize_join_domain_prefilters, ) + ir = remove_cache_nodes(ir) return optimize_join_domain_prefilters(ir, stats, config_options) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index 51e27f15155e..ef317dc004b3 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -516,7 +516,7 @@ def test_broadcast_limit( assert len(shuffle_nodes) == 0 -def test_cache_preserves_partitioning_join( +def test_shared_join_preserves_partitioning( parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ): engine = pl.GPUEngine( @@ -550,14 +550,16 @@ def test_cache_preserves_partitioning_join( lowered_ir = lowering.lowered partition_info = lowering.partition_info - # Cache should preserve partitioning on 'key' - cache_partitioning = [ + assert not any(isinstance(node, Cache) for node in traversal([lowered_ir])) + + # Removing Cache should preserve the shared join's partitioning on 'key'. + join_partitioning = [ [ne.name for ne in partition_info[node].partitioned_on] for node in traversal([lowered_ir]) - if isinstance(node, Cache) + if isinstance(node, Join) ] - assert cache_partitioning == [["key"]], ( - f"Cache should preserve partitioning on 'key', got {cache_partitioning}" + assert join_partitioning == [["key"]], ( + f"Shared join should be partitioned on 'key', got {join_partitioning}" ) # Only 2 shuffles needed (for join sides, not for groupby) diff --git a/python/cudf_polars/tests/streaming/test_parallel.py b/python/cudf_polars/tests/streaming/test_parallel.py index 44d1b4ce0757..742daadd6cdb 100644 --- a/python/cudf_polars/tests/streaming/test_parallel.py +++ b/python/cudf_polars/tests/streaming/test_parallel.py @@ -12,9 +12,13 @@ from polars.testing import assert_frame_equal from cudf_polars import Translator +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.base import StatsCollector +from cudf_polars.streaming.parallel import optimize_with_stats from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.utils.config import ConfigOptions @pytest.mark.parametrize("column", ["a", "b"]) @@ -88,6 +92,30 @@ def test_evaluate_streaming(streaming_engine): assert_frame_equal(expected, got_streaming) +def test_optimize_removes_cache_nodes() -> None: + source = pl.LazyFrame({"key": [1, 1, 2], "value": [10, 20, 30]}) + query = source.join(source, on="key", suffix="_right") + engine = GPUEngine( + executor="streaming", + executor_options={"join_domain_prefilter": None}, + ) + ir = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(ir, Join) + assert isinstance(ir.children[0], Cache) + assert ir.children[0] is ir.children[1] + + optimized = optimize_with_stats( + ir, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + assert optimized.children[0] is optimized.children[1] + assert not any(isinstance(node, Cache) for node in traversal([optimized])) + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_parallel.py (round 3) # --------------------------------------------------------------------------- From e26e1ac3663e949b0ae6662957dcd0d69d1de752 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 16:37:42 +0100 Subject: [PATCH 53/65] Remove treatment of Cache nodes in lowering and raise if we see one --- .../cudf_polars/streaming/actor_graph/utils.py | 4 ++-- python/cudf_polars/cudf_polars/streaming/join.py | 2 +- python/cudf_polars/cudf_polars/streaming/parallel.py | 11 ++++++++--- 3 files changed, 11 insertions(+), 6 deletions(-) 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 a8f1e2b62af7..db528b709624 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -38,7 +38,7 @@ import cudf_polars.dsl.tracing from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction -from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select +from cudf_polars.dsl.ir import Filter, GroupBy, HStack, Join, Projection, Select from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.utils.column_domain import column_domain_bindings from cudf_polars.dsl.utils.naming import names_to_indices @@ -587,7 +587,7 @@ def maybe_remap_partitioning( ), local=_remap_scheme_simple(ir, partitioning.local, ir.children[0]), ) - if isinstance(ir, (Cache, Join, Projection, Filter)): + if isinstance(ir, (Join, Projection, Filter)): child = child_ir if child_ir is not None else ir.children[0] return Partitioning( inter_rank=_remap_scheme_simple(ir, partitioning.inter_rank, child), diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 47729b3822be..e2117ee3e0da 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -75,7 +75,7 @@ def _make_hash_join( partition_info, output_count, ) - # Always reconstruct in case children contain Cache nodes + # Reconstruct with the lowered and possibly shuffled children. ir = ir.reconstruct([left, right]) # Record new partitioning info diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 5cbe812e62e0..2dcd87bdccc8 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -6,7 +6,7 @@ import dataclasses import operator -from functools import partial, reduce +from functools import reduce from typing import TYPE_CHECKING import polars as pl @@ -67,6 +67,13 @@ def _( ) +@lower_ir_node.register(Cache) +def _( + ir: Cache, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: # pragma: no cover + raise AssertionError("Cache nodes should have been removed before lowering") + + @dataclasses.dataclass class LoweringInfo: """Information produced by optimizing and lowering an IR graph.""" @@ -335,8 +342,6 @@ def _lower_ir_pwise( return new_node, partition_info -_lower_ir_pwise_preserve = partial(_lower_ir_pwise, preserve_partitioning=True) -lower_ir_node.register(Cache, _lower_ir_pwise_preserve) lower_ir_node.register(HConcat, _lower_ir_pwise) From 97e693a8e2e0c34a555f2f84ca88e7fc9695b400 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 16:59:28 +0100 Subject: [PATCH 54/65] Correctly handle pushdown through self-joins If a node is reachable through both sides of a join, we only want to push a filter onto one side. To do this, build a path (rather than just a single node key) for column lineages, and do targeted replacement of a child in that path. --- .../cudf_polars/dsl/utils/column_domain.py | 68 ++++++------ .../streaming/join_domain_prefilter.py | 105 ++++++++++++------ .../tests/dsl/test_column_domain.py | 32 +++--- .../streaming/test_join_domain_prefilter.py | 89 ++++++++++++--- 4 files changed, 194 insertions(+), 100 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py index 7959a18a1e9e..382d0027f649 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -11,7 +11,6 @@ from cudf_polars.dsl import expr from cudf_polars.dsl.ir import ( - Cache, Distinct, Filter, GroupBy, @@ -24,11 +23,24 @@ ) if TYPE_CHECKING: - from collections.abc import Iterator, Mapping + from collections.abc import Mapping from cudf_polars.dsl.ir import IR -__all__ = ["ColumnLineage", "ColumnRef", "column_domain_bindings"] +__all__ = [ + "ColumnBinding", + "ColumnLineage", + "ColumnRef", + "column_domain_bindings", +] + + +@dataclass(frozen=True) +class ColumnBinding: + """A direct binding to a named column on a specific child edge.""" + + child_index: int + name: str @dataclass(frozen=True) @@ -46,25 +58,18 @@ class ColumnLineage: column: ColumnRef source: ColumnLineage | None = None source_child_index: int | None = None - """Unique child edge leading to ``source``, or None if absent or ambiguous.""" - - def __iter__(self) -> Iterator[ColumnRef]: - """Iterate from the output column towards its furthest known source.""" - lineage: ColumnLineage | None = self - while lineage is not None: - yield lineage.column - lineage = lineage.source + """Child edge leading to ``source``, or None if there is no source.""" @singledispatch -def column_domain_bindings(node: IR) -> Mapping[str, ColumnRef]: +def column_domain_bindings(node: IR) -> Mapping[str, ColumnBinding]: """ Map output columns to child columns containing their value domains. - For every ``output_name -> ColumnRef(child, input_name)`` binding, every - value appearing in ``node[output_name]`` is guaranteed to appear in - ``child[input_name]``. Row order, multiplicity, and cardinality are not - preserved. + For every ``output_name -> ColumnBinding(child_index, input_name)`` binding, + every value appearing in ``node[output_name]`` is guaranteed to appear in + ``node.children[child_index][input_name]``. Row order, multiplicity, and + cardinality are not preserved. If a name in ``node.schema`` does not appear in the mapping it means that it was not possible to derive a relationship between the domain of @@ -74,68 +79,65 @@ def column_domain_bindings(node: IR) -> Mapping[str, ColumnRef]: @column_domain_bindings.register(Select) -def _(node: Select) -> Mapping[str, ColumnRef]: - child = node.children[0] +def _(node: Select) -> Mapping[str, ColumnBinding]: return { - item.name: ColumnRef(child, item.value.name) + item.name: ColumnBinding(0, item.value.name) for item in node.exprs if isinstance(item.value, expr.Col) } @column_domain_bindings.register(HStack) -def _(node: HStack) -> Mapping[str, ColumnRef]: +def _(node: HStack) -> Mapping[str, ColumnBinding]: child = node.children[0] replaced = {item.name for item in node.columns} return { - name: ColumnRef(child, name) for name in child.schema if name not in replaced + name: ColumnBinding(0, name) for name in child.schema if name not in replaced } | { - item.name: ColumnRef(child, item.value.name) + item.name: ColumnBinding(0, item.value.name) for item in node.columns if isinstance(item.value, expr.Col) } @column_domain_bindings.register(GroupBy) -def _(node: GroupBy) -> Mapping[str, ColumnRef]: - child = node.children[0] +def _(node: GroupBy) -> Mapping[str, ColumnBinding]: return { - key.name: ColumnRef(child, key.value.name) + key.name: ColumnBinding(0, key.value.name) for key in node.keys if isinstance(key.value, expr.Col) } @column_domain_bindings.register(Join) -def _(node: Join) -> Mapping[str, ColumnRef]: +def _(node: Join) -> Mapping[str, ColumnBinding]: left, right = node.children how = node.options[0] if how in ("Semi", "Anti"): return { - name: ColumnRef(left, name) for name in node.schema if name in left.schema + name: ColumnBinding(0, name) for name in node.schema if name in left.schema } if how != "Inner": return {} - bindings = {name: ColumnRef(left, name) for name in left.schema} + bindings = {name: ColumnBinding(0, name) for name in left.schema} suffix = node.options[3] for name in right.schema: output_name = f"{name}{suffix}" if name in left.schema else name if output_name in node.schema: - bindings[output_name] = ColumnRef(right, name) + bindings[output_name] = ColumnBinding(1, name) return bindings -@column_domain_bindings.register(Cache) @column_domain_bindings.register(Distinct) @column_domain_bindings.register(Filter) @column_domain_bindings.register(Projection) @column_domain_bindings.register(Slice) @column_domain_bindings.register(Sort) def _( - node: Cache | Distinct | Filter | Projection | Slice | Sort, -) -> Mapping[str, ColumnRef]: + node: Distinct | Filter | Projection | Slice | Sort, +) -> Mapping[str, ColumnBinding]: child = node.children[0] return { - name: ColumnRef(child, name) for name in node.schema if name in child.schema + name: ColumnBinding(0, name) for name in node.schema if name in child.schema } 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 866f9e842353..c2afe100d960 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -11,7 +11,6 @@ from cudf_polars.dsl import expr from cudf_polars.dsl.ir import ( IR, - Cache, ConditionalJoin, DataFrameScan, Distinct, @@ -39,7 +38,6 @@ ColumnRef, column_domain_bindings, ) -from cudf_polars.dsl.utils.replace import replace if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping, Sequence @@ -56,6 +54,8 @@ class _Producer: node: IR columns: tuple[str, ...] rows: int + path: tuple[int, ...] = () + """Child-edge path from the candidate root to ``node``.""" @property def column(self) -> str: @@ -162,9 +162,7 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: rows = None if source_info is None else source_info.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) - ): + elif isinstance(node, (Select, Projection, HStack, Filter, Distinct, GroupBy)): rows = row_estimates[node.children[0]] elif isinstance(node, Join): rows = _estimate_join_rows( @@ -191,20 +189,17 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: bindings = column_domain_bindings(node) for name in node.schema: column = ColumnRef(node, name) - source = bindings.get(name) - if source is None: + binding = bindings.get(name) + if binding is None: source_lineage = None source_child_index = None else: - source_lineage = column_lineages[source] - source_children = tuple( - index - for index, child in enumerate(node.children) - if child == source.node - ) - source_child_index = ( - source_children[0] if len(source_children) == 1 else None + source_child_index = binding.child_index + source = ColumnRef( + node.children[source_child_index], + binding.name, ) + source_lineage = column_lineages[source] column_lineages[column] = ColumnLineage( column, source_lineage, source_child_index ) @@ -248,7 +243,7 @@ def blocks_pushdown(node: IR) -> bool: def semijoin_pushdown_candidates( facts: PlanFacts, root: IR, column: str -) -> Iterator[ColumnRef]: +) -> Iterator[tuple[ColumnRef, tuple[int, ...]]]: """ Yield column domain lineage providing valid locations for semijoin pushdown. @@ -264,21 +259,22 @@ def semijoin_pushdown_candidates( Returns ------- Iterator - Of valid insertion points for a semijoin filter on the given column name. + Of valid insertion points and their child-edge paths from ``root``. """ try: lineage = facts.column_lineages[ColumnRef(root, column)] except KeyError: return + path: tuple[int, ...] = () while True: - yield lineage.column - if ( - blocks_pushdown(lineage.column.node) - or lineage.source is None - or lineage.source_child_index is None - ): + yield lineage.column, path + source = lineage.source + source_child_index = lineage.source_child_index + if blocks_pushdown(lineage.column.node) or source is None: return - lineage = lineage.source + assert source_child_index is not None + path = (*path, source_child_index) + lineage = source def optimize_join_domain_prefilters( @@ -360,15 +356,46 @@ def apply_candidate(ir: Join, candidate: Candidate) -> IR: nulls_equal=ir.options[1], suffix=ir.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": - (left,) = replace([left], {candidate.target.node: target_filter}) + left = replace_at_path(left, target.path, target_filter) else: - (right,) = replace([right], {candidate.target.node: target_filter}) + right = replace_at_path(right, target.path, target_filter) return ir.reconstruct((left, right)) +def replace_at_path(root: IR, path: Sequence[int], replacement: IR) -> IR: + """ + Replace a specific child in a DAG starting at root. + + Parameters + ---------- + root + Root of DAG to carry out replacement. + path + Breadcrumb trail selecting which child at every level to recurse + into. + replacement + Replacement node to return when the path becomes empty. + + Returns + ------- + IR + New DAG with the selected child replaced with replacement. + + Notes + ----- + This specifically does not use replacement by equality so that we can + disambiguate between shared children in the DAG where we only want to + replace one. + """ + if not path: + return replacement + index, *path = path + children = list(root.children) + children[index] = replace_at_path(children[index], path, replacement) + return root.reconstruct(children) + + def _select_candidate( ir: Join, threshold: float, @@ -601,7 +628,7 @@ def _smallest_key_producer( exclude: IR | None = None, ) -> _Producer | None: candidates = [] - for reference in semijoin_pushdown_candidates(facts, root, column): + for reference, path in semijoin_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name if node is exclude: continue @@ -611,7 +638,7 @@ def _smallest_key_producer( if require_selective and node not in facts.selective_nodes: continue candidates.append( - (rows, len(node.schema), _Producer(node, (bound_column,), rows)) + (rows, len(node.schema), _Producer(node, (bound_column,), rows, path)) ) if not candidates: return None @@ -630,6 +657,7 @@ def _smallest_node_containing_all( lineages.append(lineage) if not lineages: return None + path: tuple[int, ...] = () while True: node = lineages[0].column.node if any(lineage.column.node != node for lineage in lineages[1:]): @@ -641,17 +669,20 @@ def _smallest_node_containing_all( ( rows, len(node.schema), - _Producer(node, bound_columns, rows), + _Producer(node, bound_columns, rows, path), ) ) if blocks_pushdown(node): break - source_child_indices = {lineage.source_child_index for lineage in lineages} - if len(source_child_indices) != 1 or None in source_child_indices: + source_child_index = lineages[0].source_child_index + if source_child_index is None or any( + lineage.source_child_index != source_child_index for lineage in lineages[1:] + ): break sources = [lineage.source for lineage in lineages] if any(source is None for source in sources): break + path = (*path, source_child_index) lineages = [source for source in sources if source is not None] if not candidates: return None @@ -661,12 +692,16 @@ def _smallest_node_containing_all( def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: source_candidates = [] fallback_candidates = [] - for reference in semijoin_pushdown_candidates(facts, root, column): + for reference, path in semijoin_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name rows = facts.row_estimates.get(node) if rows is None or rows <= 0: continue - item = (rows, len(node.schema), _Producer(node, (bound_column,), rows)) + item = ( + rows, + len(node.schema), + _Producer(node, (bound_column,), rows, path), + ) if isinstance(node, (Scan, DataFrameScan)): source_candidates.append(item) else: diff --git a/python/cudf_polars/tests/dsl/test_column_domain.py b/python/cudf_polars/tests/dsl/test_column_domain.py index 3d1ec6b347b5..999f3f2463d2 100644 --- a/python/cudf_polars/tests/dsl/test_column_domain.py +++ b/python/cudf_polars/tests/dsl/test_column_domain.py @@ -10,7 +10,6 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr from cudf_polars.dsl.ir import ( - Cache, DataFrameScan, Distinct, Filter, @@ -23,7 +22,7 @@ Sort, ) from cudf_polars.dsl.utils.column_domain import ( - ColumnRef, + ColumnBinding, column_domain_bindings, ) @@ -61,7 +60,7 @@ def test_select_binds_aliases_and_omits_derived_columns() -> None: ) assert column_domain_bindings(node) == { - "renamed": ColumnRef(child, "a"), + "renamed": ColumnBinding(0, "a"), } @@ -78,8 +77,8 @@ def test_hstack_binds_passthrough_alias_and_override() -> None: ) assert column_domain_bindings(node) == { - "b": ColumnRef(child, "b"), - "alias": ColumnRef(child, "b"), + "b": ColumnBinding(0, "b"), + "alias": ColumnBinding(0, "b"), } @@ -95,7 +94,7 @@ def test_groupby_binds_only_direct_keys() -> None: ) assert column_domain_bindings(node) == { - "key": ColumnRef(child, "a"), + "key": ColumnBinding(0, "a"), } @@ -117,10 +116,10 @@ def test_inner_join_binds_left_right_and_suffixed_columns() -> None: ) assert column_domain_bindings(node) == { - "key": ColumnRef(left, "key"), - "left_value": ColumnRef(left, "left_value"), - "key_right": ColumnRef(right, "key"), - "right_value": ColumnRef(right, "right_value"), + "key": ColumnBinding(0, "key"), + "left_value": ColumnBinding(0, "left_value"), + "key_right": ColumnBinding(1, "key"), + "right_value": ColumnBinding(1, "right_value"), } @@ -137,9 +136,9 @@ def test_inner_join_omits_coalesced_right_key() -> None: ) assert column_domain_bindings(node) == { - "key": ColumnRef(left, "key"), - "left_value": ColumnRef(left, "left_value"), - "right_value": ColumnRef(right, "right_value"), + "key": ColumnBinding(0, "key"), + "left_value": ColumnBinding(0, "left_value"), + "right_value": ColumnBinding(1, "right_value"), } @@ -156,8 +155,8 @@ def test_semi_join_binds_only_left_columns() -> None: ) assert column_domain_bindings(node) == { - "key": ColumnRef(left, "key"), - "value": ColumnRef(left, "value"), + "key": ColumnBinding(0, "key"), + "value": ColumnBinding(0, "value"), } @@ -180,7 +179,6 @@ def test_passthrough_nodes_bind_same_named_columns() -> None: child = make_scan("a", "b") mask = expr.NamedExpr("mask", expr.Literal(BOOL, True)) # noqa: FBT003 nodes = ( - Cache(child.schema, 1, None, child), Filter(child.schema, mask, child), Projection({"b": I64}, child), Slice(child.schema, 0, 1, child), @@ -205,5 +203,5 @@ def test_passthrough_nodes_bind_same_named_columns() -> None: for node in nodes: assert column_domain_bindings(node) == { - name: ColumnRef(child, name) for name in node.schema + name: ColumnBinding(0, name) for name in node.schema } 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 7c3a8e700262..0dfbd9393a94 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -30,6 +30,7 @@ optimize_join_domain_prefilters, semijoin_pushdown_candidates, ) +from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes 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 @@ -489,12 +490,16 @@ def test_composite_domain_columns_do_not_reconverge_after_join( assert isinstance(joined, Join) assert isinstance(joined.children[0], Cache) assert joined.children[0] is joined.children[1] + joined = remove_cache_nodes(joined) + assert isinstance(joined, Join) + assert joined.children[0] is joined.children[1] facts = analyze_plan(joined, StatsCollector()) producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) assert tuple(semijoin_pushdown_candidates(facts, joined, "value")) == ( - ColumnRef(joined, "value"), + (ColumnRef(joined, "value"), ()), + (ColumnRef(joined.children[0], "value"), (0,)), ) assert producer is not None assert producer.node is joined @@ -520,16 +525,13 @@ def test_plan_facts_share_lineage_suffixes_across_shared_dag() -> None: facts = analyze_plan(root, _stats(source=(source, 10))) left_lineage = facts.column_lineages[ColumnRef(left, "left_key")] right_lineage = facts.column_lineages[ColumnRef(right, "right_key")] + source_lineage = facts.column_lineages[ColumnRef(source, "raw_key")] - assert left_lineage.source is right_lineage.source - assert tuple(left_lineage) == ( - ColumnRef(left, "left_key"), - ColumnRef(source, "raw_key"), - ) - assert tuple(right_lineage) == ( - ColumnRef(right, "right_key"), - ColumnRef(source, "raw_key"), - ) + assert left_lineage.column == ColumnRef(left, "left_key") + assert right_lineage.column == ColumnRef(right, "right_key") + assert left_lineage.source is source_lineage + assert right_lineage.source is source_lineage + assert source_lineage.source is None def test_target_prefilter_does_not_move_below_slice() -> None: @@ -540,12 +542,12 @@ def test_target_prefilter_does_not_move_below_slice() -> None: stats = _stats(target=(target, 1_000), domain=(domain, 5)) facts = analyze_plan(root, stats) - assert tuple(facts.column_lineages[ColumnRef(sliced, "target_key")]) == ( - ColumnRef(sliced, "target_key"), - ColumnRef(target, "target_key"), - ) + lineage = facts.column_lineages[ColumnRef(sliced, "target_key")] + assert lineage.column == ColumnRef(sliced, "target_key") + assert lineage.source is facts.column_lineages[ColumnRef(target, "target_key")] + assert lineage.source.source is None assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( - ColumnRef(sliced, "target_key"), + (ColumnRef(sliced, "target_key"), ()), ) optimized = optimize_join_domain_prefilters( @@ -585,6 +587,63 @@ def test_target_replacement_does_not_rewrite_shared_domain_side() -> None: assert domain.children[0] is shared +def test_target_prefilter_rewrites_only_selected_self_join_edge( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame( + { + "key": [1, 1, 2, 2], + "value": [10, 20, 30, 40], + } + ) + domain = ( + pl.LazyFrame( + { + "domain_value": [10, 999], + "active": [True, False], + } + ) + .filter("active") + .select("domain_value") + ) + query = source.join(source, on="key", suffix="_right").join( + domain, + left_on="value", + right_on="domain_value", + ) + translated = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(translated, Join) + translated_self_join = translated.children[0] + assert isinstance(translated_self_join, Join) + shared_cache = translated_self_join.children[0] + assert isinstance(shared_cache, Cache) + assert translated_self_join.children[1] is shared_cache + source_ir = shared_cache.children[0] + + optimized = optimize_with_stats( + translated, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + rewritten_self_join = optimized.children[0] + assert isinstance(rewritten_self_join, Join) + filtered, unfiltered = rewritten_self_join.children + assert isinstance(filtered, Join) + assert filtered.options[0] == "Semi" + assert filtered.children[0] is source_ir + assert unfiltered is source_ir + + expected = query.collect() + assert sorted(expected.select("value", "value_right").rows()) == [ + (10, 10), + (10, 20), + ] + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + def test_no_domain_prefilter_for_outer_join() -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",)) From 3a46fcf0db1e3c73dd7e02a31ec59f35756d2755 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 17:45:29 +0100 Subject: [PATCH 55/65] Add high-level description of what is going on --- .../streaming/join_domain_prefilter.py | 65 ++++++++++++++++++- 1 file changed, 64 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 c2afe100d960..1a546b73675f 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -1,6 +1,69 @@ # 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.""" +""" +Insert derived key-domain prefilters for streaming joins. + +For a supported inner equijoin, this optimization tries to use the join-key +values produced by one input to reduce the other input before the original +join. In relational notation, a simple rewrite is:: + + left join[left.key = right.key] right + + -> + + (left semijoin[left.key = right.key] project(right.key)) + join[left.key = right.key] right + +In this example, the right hand table is selected to pre-filter the left +table before performing the inner join. The inserted semi-join is therefore +an exact filter: it only removes target rows that could not match the +domain input. + +The implementation uses the following terms: + +``column lineage`` + A chain from a named output column towards columns in its input subplan. + Each step guarantees that every value in the output column also appears in + the referenced child column, although row order and multiplicity are not + preserved and the child may contain additional values. Candidate traversal + does not cross an operator with which a semi-join cannot safely commute. +``child edge`` + One particular parent-to-child position in the IR DAG. The same child node + may occur on more than one edge, so a lineage records child indices and a + rewrite follows the resulting edge path to change only the chosen + occurrence. +``producer`` + A node on a column lineage, together with the column name at that node and + its edge path from the join input. Producers are possible locations for + inserting a target semi-join or projecting a domain key. +``target`` + The join input to reduce. A semi-join is inserted at a producer in this + input's column lineage and replaces only the selected child-edge + occurrence. +``domain`` + The other join input, whose join-key values provide the semi-join domain. + The domain producer may be below projections, renames, or other operators + through which the semi-join can safely be pushed. +``constraint domain`` + Selective values of another join key from the target input, used to reduce + the domain before deriving the values that will filter the target. +``simple candidate`` + A rewrite that projects one domain join key and uses it to filter the + corresponding target key directly. +``composite candidate`` + For a multi-key join, a rewrite that first semi-joins the domain using the + constraint domain, then projects the reduced domain's key used to filter + the target. + +Plan rewrite has three stages. ``analyze_plan`` gathers row estimates, +selective nodes, and column value-domain lineages. Candidate selection +consumes those facts and returns a decision. ``apply_candidate`` then +constructs the selected semi-join rewrite. + +Row estimates, selectivity propagation, thresholds, and candidate scores are +only heuristics for deciding whether a safe rewrite is likely to improve +execution. Poor estimates can choose an unprofitable rewrite. +""" from __future__ import annotations From 7904b204063a5582f237c5a1b151cbfe98de0479 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 14 Jul 2026 18:00:19 +0100 Subject: [PATCH 56/65] Actually test we get the right answers --- .../streaming/test_join_domain_prefilter.py | 206 ++++++++++++------ 1 file changed, 142 insertions(+), 64 deletions(-) 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 0dfbd9393a94..8e638ae05316 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -301,48 +301,71 @@ def test_no_simple_domain_prefilter_when_domain_is_not_selective() -> None: 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", +def test_composite_domain_prefilter_constrains_domain_first( + engine: SPMDEngine, +) -> None: + nation = ( + pl.LazyFrame( + { + "n_nationkey": range(10), + "active": [True] * 5 + [False] * 5, + } + ) + .filter("active") + .select("n_nationkey") ) - root = _join( - order_lineitem, - supplier, - ("l_suppkey", "n_nationkey"), - ("s_suppkey", "s_nationkey"), + orders = pl.LazyFrame( + { + "o_orderkey": range(90), + "n_nationkey": [i % 10 for i in range(90)], + } ) - - stats = _stats( - nation=(nation, 5), - orders=(orders, 900), - lineitem=(lineitem, 1_800), - supplier=(supplier, 30), + lineitem = pl.LazyFrame( + { + "l_orderkey": [i % 90 for i in range(180)], + "l_suppkey": [i % 30 for i in range(180)], + } ) - decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) - optimized = optimize_join_domain_prefilters( - root, - stats, - _config(), + supplier = pl.LazyFrame( + { + "s_suppkey": range(30), + "s_nationkey": [i % 10 for i in range(30)], + } ) + query = ( + nation.join(orders, on="n_nationkey") + .join( + lineitem, + left_on="o_orderkey", + right_on="l_orderkey", + maintain_order="left", + ) + .join( + supplier, + left_on=("l_suppkey", "n_nationkey"), + right_on=("s_suppkey", "s_nationkey"), + ) + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + config = ConfigOptions.from_polars_engine(engine) + stats = StatsCollector() + + assert isinstance(root, Join) + order_lineitem, supplier_ir = root.children + assert isinstance(order_lineitem, Join) + lineitem_ir = order_lineitem.children[1] + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) + optimized = optimize_join_domain_prefilters(root, stats, config) assert decision.reason == "applied" assert isinstance(decision.candidate, CompositeCandidate) 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) + assert optimized.children[1] is supplier_ir + assert any(semi.children[0] is supplier_ir for semi in semis) + assert any(semi.children[0] is lineitem_ir for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_derived_selectivity_propagates_through_rewritten_children() -> None: @@ -400,61 +423,116 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking() -> None: ) -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, +def test_target_source_follows_join_key_through_rename( + engine: SPMDEngine, +) -> None: + big = pl.LazyFrame( + { + "left_key": range(20), + "other": [i % 5 for i in range(20)], + } + ) + renamed_big = big.select(pl.col("left_key").alias("foo"), "other") + small = pl.LazyFrame( + { + "left_key": range(10), + "other2": [i % 5 for i in range(10)], + } + ) + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + query = renamed_big.join( small, - ("other",), - ("other2",), + left_on="other", + right_on="other2", maintain_order="left", + ).join( + domain, + left_on="left_key", + right_on="domain_key", ) - domain = _scan("domain", ("domain_key",), predicate=True) - root = _join(joined, domain, ("left_key",), ("domain_key",)) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + assert isinstance(root, Join) + joined = root.children[0] + assert isinstance(joined, Join) + renamed_big_ir, small_ir = joined.children + assert isinstance(renamed_big_ir, Select) + big_ir = renamed_big_ir.children[0] optimized = optimize_join_domain_prefilters( root, - _stats(big=(big, 1_000), small=(small, 100), domain=(domain, 5)), - _config(), + StatsCollector(), + ConfigOptions.from_polars_engine(engine), ) 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) + assert any(semi.children[0] is small_ir for semi in semis) + assert not any(semi.children[0] is big_ir for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) -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, +def test_domain_source_follows_join_key_through_rename( + engine: SPMDEngine, +) -> None: + target = pl.LazyFrame({"target_key": range(20)}) + unrelated = pl.LazyFrame( + { + "domain_key": [100, 101], + "other": [0, 1], + "active": [True, False], + } + ) + renamed_unrelated = unrelated.filter("active").select( + pl.col("domain_key").alias("foo"), "other" + ) + domain_source = pl.LazyFrame( + { + "domain_key": range(1, 6), + "other2": range(5), + } + ) + domain = renamed_unrelated.join( domain_source, - ("other",), - ("other2",), + left_on="other", + right_on="other2", maintain_order="left", ) - root = _join(target, domain, ("target_key",), ("domain_key",)) + query = target.join( + domain, + left_on="target_key", + right_on="domain_key", + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + assert isinstance(root, Join) + target_ir, domain_ir = root.children + assert isinstance(domain_ir, Join) + renamed_unrelated_ir, domain_source_ir = domain_ir.children optimized = optimize_join_domain_prefilters( root, - _stats( - target=(target, 1_000), - unrelated=(unrelated, 1), - domain_source=(domain_source, 5), - ), - _config(), + StatsCollector(), + ConfigOptions.from_polars_engine(engine), ) semi = next( - semi for semi in _joins(optimized, "Semi") if semi.children[0] is target + semi for semi in _joins(optimized, "Semi") if semi.children[0] is target_ir ) selected_domain = semi.children[1] assert isinstance(selected_domain, Select) - assert selected_domain.children[0] is domain + rewritten_domain_source = selected_domain.children[0] + assert isinstance(rewritten_domain_source, Join) + assert rewritten_domain_source.options[0] == "Semi" + assert rewritten_domain_source.children[0] is domain_source_ir + assert rewritten_domain_source.children[0] is not renamed_unrelated_ir + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_composite_domain_columns_follow_renames() -> None: From 881718153803c4b875eeaa4f24b1bad5a3eb25aa Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 10:01:13 +0100 Subject: [PATCH 57/65] Assert optimised query in way independent of Polars version --- .../streaming/test_join_domain_prefilter.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) 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 8e638ae05316..83b68f729a76 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -575,13 +575,14 @@ def test_composite_domain_columns_do_not_reconverge_after_join( facts = analyze_plan(joined, StatsCollector()) producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) - assert tuple(semijoin_pushdown_candidates(facts, joined, "value")) == ( - (ColumnRef(joined, "value"), ()), - (ColumnRef(joined.children[0], "value"), (0,)), - ) + candidates = tuple(semijoin_pushdown_candidates(facts, joined, "value")) + assert candidates[0] == (ColumnRef(joined, "value"), ()) + assert len(candidates) >= 2 + assert all(path == (0,) * len(path) for _, path in candidates[1:]) assert producer is not None assert producer.node is joined assert producer.columns == ("value", "value_right") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_contains_node_uses_dag_equality() -> None: @@ -709,10 +710,11 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( rewritten_self_join = optimized.children[0] assert isinstance(rewritten_self_join, Join) filtered, unfiltered = rewritten_self_join.children - assert isinstance(filtered, Join) - assert filtered.options[0] == "Semi" - assert filtered.children[0] is source_ir assert unfiltered is source_ir + filtered_semis = _joins(filtered, "Semi") + assert len(filtered_semis) == 1 + assert not _joins(unfiltered, "Semi") + assert any(filtered_semis[0].children[0] is node for node in traversal([source_ir])) expected = query.collect() assert sorted(expected.select("value", "value_right").rows()) == [ From 6f16f9f5815a70f7676fa6fd288bc90acc23e41c Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 10:23:57 +0100 Subject: [PATCH 58/65] Assert more correctness in tests --- .../streaming/test_join_domain_prefilter.py | 284 +++++++++++++----- 1 file changed, 204 insertions(+), 80 deletions(-) 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 83b68f729a76..3d67286480ac 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 Cache, Distinct, Join, Scan, Select, Slice +from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Scan, Select, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.column_domain import ColumnRef from cudf_polars.engine.options import StreamingOptions @@ -167,12 +167,45 @@ def _joins(ir: IR, how: str | None = None) -> list[Join]: ] -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",)) +def translate_query(query: pl.LazyFrame, engine: SPMDEngine) -> IR: + """Translate a public Polars query and remove logical Cache nodes.""" + return remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + + +def dataframe_scan(ir: IR, column: str) -> DataFrameScan: + """Return the unique in-memory scan containing ``column``.""" + (match,) = ( + node + for node in traversal([ir]) + if isinstance(node, DataFrameScan) and column in node.schema + ) + return match + + +def test_simple_domain_prefilter_filters_large_side(engine: SPMDEngine) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 10 for i in range(20)], + "l_suppkey": range(20), + } + ) + query = part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + root = translate_query(query, engine) - facts = analyze_plan(root, _stats(part=(part, 6), lineitem=(lineitem, 1_800))) + assert isinstance(root, Join) + part_ir, _ = root.children + lineitem_ir = dataframe_scan(root, "l_partkey") + facts = analyze_plan(root, StatsCollector()) decision = _select_candidate(root, 0.5, facts) assert decision.reason == "applied" @@ -181,10 +214,11 @@ def test_simple_domain_prefilter_filters_large_side() -> None: 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 + semis = _joins(optimized, "Semi") + assert len(semis) == 1 + assert semis[0].children[0] is lineitem_ir + assert not _joins(part_ir, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_domain_prefilter_is_independent_of_dynamic_planning() -> None: @@ -368,59 +402,114 @@ def test_composite_domain_prefilter_constrains_domain_first( assert_gpu_result_equal(query, engine=engine, check_row_order=False) -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",)) +def test_derived_selectivity_propagates_through_rewritten_children( + engine: SPMDEngine, +) -> None: + region = ( + pl.LazyFrame( + { + "r_regionkey": [0, 1], + "active": [True, False], + } + ) + .filter("active") + .select("r_regionkey") + ) + nation = pl.LazyFrame( + { + "n_nationkey": range(10), + "n_regionkey": [i % 2 for i in range(10)], + } + ) + customer = pl.LazyFrame( + { + "c_custkey": range(40), + "c_nationkey": [i % 10 for i in range(40)], + } + ) + orders = pl.LazyFrame( + { + "o_orderkey": range(200), + "o_custkey": [i % 40 for i in range(200)], + } + ) + query = ( + region.join(nation, left_on="r_regionkey", right_on="n_regionkey") + .join(customer, left_on="n_nationkey", right_on="c_nationkey") + .join(orders, left_on="c_custkey", right_on="o_custkey") + ) + root = translate_query(query, engine) optimized = optimize_join_domain_prefilters( root, - _stats( - region=(region, 1), - nation=(nation, 25), - customer=(customer, 150), - orders=(orders, 1_500), - ), - _config(), + StatsCollector(), + ConfigOptions.from_polars_engine(engine), ) - filtered = {semi.children[0] for semi in _joins(optimized, "Semi")} - assert {nation, customer, orders} <= filtered - + semis = _joins(optimized, "Semi") + expected_targets = { + dataframe_scan(root, "n_nationkey"), + dataframe_scan(root, "c_custkey"), + dataframe_scan(root, "o_orderkey"), + } + assert expected_targets <= {semi.children[0] for semi in semis} + assert_gpu_result_equal(query, engine=engine, check_row_order=False) -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",)) +def test_rewritten_domain_filters_other_side_instead_of_stacking( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(6), + "part_active": [True] * 3 + [False] * 3, + } + ) + .filter("part_active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_orderkey": [i % 15 for i in range(180)], + "l_partkey": [i % 6 for i in range(180)], + "l_suppkey": [i % 3 for i in range(180)], + } + ) + supplier = pl.LazyFrame({"s_suppkey": range(3)}) + orders = ( + pl.LazyFrame( + { + "o_orderkey": range(15), + "order_active": [True] * 8 + [False] * 7, + } + ) + .filter("order_active") + .select("o_orderkey") + ) + query = ( + part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + .join(supplier, left_on="l_suppkey", right_on="s_suppkey") + .join(orders, left_on="l_orderkey", right_on="o_orderkey") + ) + root = translate_query(query, engine) optimized = optimize_join_domain_prefilters( root, - _stats( - part=(part, 60), - lineitem=(lineitem, 1_800), - supplier=(supplier, 30), - orders=(orders, 150), - ), - _config(), + StatsCollector(), + ConfigOptions.from_polars_engine(engine), ) + lineitem_ir = dataframe_scan(root, "l_orderkey") + orders_ir = dataframe_scan(root, "o_orderkey") 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 sum(semi.children[0] is lineitem_ir for semi in semis) == 1 + assert any(semi.children[0] is orders_ir for semi in semis) assert not any( isinstance(semi.children[0], Join) and semi.children[0].options[0] == "Semi" for semi in semis ) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_target_source_follows_join_key_through_rename( @@ -613,18 +702,33 @@ def test_plan_facts_share_lineage_suffixes_across_shared_dag() -> None: assert source_lineage.source is None -def test_target_prefilter_does_not_move_below_slice() -> None: - target = _scan("target", ("target_key",)) - sliced = Slice(target.schema, 0, 100, target) - domain = _scan("domain", ("domain_key",), predicate=True) - root = _join(sliced, domain, ("target_key",), ("domain_key",)) +def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: + target = ( + pl.LazyFrame({"target_key": range(20)}) + .filter(pl.col("target_key") >= 0) + .slice(0, 10) + ) + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + query = target.join(domain, left_on="target_key", right_on="domain_key") + root = translate_query(query, engine) - stats = _stats(target=(target, 1_000), domain=(domain, 5)) + assert isinstance(root, Join) + sliced = root.children[0] + assert isinstance(sliced, Slice) + target_ir = dataframe_scan(root, "target_key") + stats = StatsCollector() facts = analyze_plan(root, stats) lineage = facts.column_lineages[ColumnRef(sliced, "target_key")] assert lineage.column == ColumnRef(sliced, "target_key") - assert lineage.source is facts.column_lineages[ColumnRef(target, "target_key")] - assert lineage.source.source is None assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( (ColumnRef(sliced, "target_key"), ()), ) @@ -632,38 +736,63 @@ def test_target_prefilter_does_not_move_below_slice() -> None: optimized = optimize_join_domain_prefilters( root, stats, - _config(), + ConfigOptions.from_polars_engine(engine), ) semis = _joins(optimized, "Semi") assert any(semi.children[0] is sliced for semi in semis) - assert not any(semi.children[0] is target for semi in semis) + assert not any(semi.children[0] is target_ir for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) -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, +def test_target_replacement_does_not_rewrite_shared_domain_side( + engine: SPMDEngine, +) -> None: + shared = pl.LazyFrame( + { + "target_key": range(20), + "other": [i % 2 for i in range(20)], + } + ) + domain_source = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "other2": [0, 1], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key", "other2") + ) + domain = shared.join( domain_source, - ("other",), - ("other2",), - maintain_order="left", + left_on=pl.col("other").cast(pl.Int32), + right_on=pl.col("other2").cast(pl.Int32), ) - root = _join(shared, domain, ("target_key",), ("domain_key",)) + query = shared.join(domain, left_on="target_key", right_on="domain_key") + root = translate_query(query, engine) + + assert isinstance(root, Join) + shared_ir, domain_ir = root.children + assert isinstance(domain_ir, Join) + assert domain_ir.children[0] is shared_ir optimized = optimize_join_domain_prefilters( root, - _stats(shared=(shared, 1_000), domain_source=(domain_source, 5)), - _config(), + StatsCollector(), + ConfigOptions.from_polars_engine(engine), ) 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 + filtered, unfiltered_domain = optimized.children + assert unfiltered_domain is domain_ir + assert domain_ir.children[0] is shared_ir + semis = _joins(filtered, "Semi") + assert len(semis) == 1 + assert semis[0].children[0] is dataframe_scan(root, "target_key") + assert not _joins(unfiltered_domain, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_target_prefilter_rewrites_only_selected_self_join_edge( @@ -715,19 +844,14 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( assert len(filtered_semis) == 1 assert not _joins(unfiltered, "Semi") assert any(filtered_semis[0].children[0] is node for node in traversal([source_ir])) - - expected = query.collect() - assert sorted(expected.select("value", "value_right").rows()) == [ - (10, 10), - (10, 20), - ] assert_gpu_result_equal(query, engine=engine, check_row_order=False) -def test_no_domain_prefilter_for_outer_join() -> None: +@pytest.mark.parametrize("how", ["Left", "Right", "Cross", "Full"]) +def test_no_domain_prefilter_for_outer_join(how) -> None: part = _scan("part", ("p_partkey",), predicate=True) lineitem = _scan("lineitem", ("l_partkey",)) - root = _join(part, lineitem, ("p_partkey",), ("l_partkey",), how="Left") + root = _join(part, lineitem, ("p_partkey",), ("l_partkey",), how=how) optimized = optimize_join_domain_prefilters( root, From 1ce2e59355e260ff8c06f17e76009464790039c5 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 10:49:00 +0100 Subject: [PATCH 59/65] Remove inappropriate intimacy in tests Construct remaining prefilter tests using Polars queries --- .../streaming/test_join_domain_prefilter.py | 303 ++++++++---------- 1 file changed, 137 insertions(+), 166 deletions(-) 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 3d67286480ac..7cf8e7f13fb7 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -3,16 +3,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING 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 Cache, DataFrameScan, Distinct, Join, Scan, Select, Slice +from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.column_domain import ColumnRef from cudf_polars.engine.options import StreamingOptions @@ -33,17 +31,14 @@ from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes 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 +from cudf_polars.utils.config import ConfigOptions if TYPE_CHECKING: import concurrent.futures + from typing import Any from cudf_polars.dsl.ir import IR from cudf_polars.engine.spmd import SPMDEngine - from cudf_polars.streaming.base import SerializedDataSourceInfo - -I64 = DataType(pl.Int64()) -BOOL = DataType(pl.Boolean()) @pytest.fixture @@ -57,96 +52,10 @@ def engine(spmd_engine_factory) -> SPMDEngine: ) -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 _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, - 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( +def make_config( *, dynamic_planning: bool = True, join_domain_prefilter: bool = True ) -> ConfigOptions: - executor_options: dict[str, object] = { + executor_options: dict[str, Any] = { "join_domain_prefilter": {"trace": False} if join_domain_prefilter else None } if not dynamic_planning: @@ -159,7 +68,7 @@ def _config( ) -def _joins(ir: IR, how: str | None = None) -> list[Join]: +def find_joins(ir: IR, how: str | None = None) -> list[Join]: return [ node for node in traversal([ir]) @@ -169,7 +78,10 @@ def _joins(ir: IR, how: str | None = None) -> list[Join]: def translate_query(query: pl.LazyFrame, engine: SPMDEngine) -> IR: """Translate a public Polars query and remove logical Cache nodes.""" - return remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + t = Translator(query._ldf.visit(), engine) + root = t.translate_ir() + assert not t.errors + return remove_cache_nodes(root) def dataframe_scan(ir: IR, column: str) -> DataFrameScan: @@ -182,7 +94,9 @@ def dataframe_scan(ir: IR, column: str) -> DataFrameScan: return match -def test_simple_domain_prefilter_filters_large_side(engine: SPMDEngine) -> None: +@pytest.fixture +def simple_query() -> pl.LazyFrame: + """Return a query with a small selective join domain.""" part = ( pl.LazyFrame( { @@ -199,8 +113,13 @@ def test_simple_domain_prefilter_filters_large_side(engine: SPMDEngine) -> None: "l_suppkey": range(20), } ) - query = part.join(lineitem, left_on="p_partkey", right_on="l_partkey") - root = translate_query(query, engine) + return part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + + +def test_simple_domain_prefilter_filters_large_side( + simple_query: pl.LazyFrame, engine: SPMDEngine +) -> None: + root = translate_query(simple_query, engine) assert isinstance(root, Join) part_ir, _ = root.children @@ -214,36 +133,37 @@ def test_simple_domain_prefilter_filters_large_side(engine: SPMDEngine) -> None: assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") assert len(semis) == 1 assert semis[0].children[0] is lineitem_ir - assert not _joins(part_ir, "Semi") - assert_gpu_result_equal(query, engine=engine, check_row_order=False) + assert not find_joins(part_ir, "Semi") + assert_gpu_result_equal(simple_query, engine=engine, check_row_order=False) -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",)) +def test_domain_prefilter_is_independent_of_dynamic_planning( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + root = translate_query(simple_query, engine) optimized = optimize_join_domain_prefilters( root, - _stats(part=(part, 6), lineitem=(lineitem, 1_800)), - _config(dynamic_planning=False), + StatsCollector(), + make_config(dynamic_planning=False), ) - assert _joins(optimized, "Semi") + assert find_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",)) +def test_domain_prefilter_can_be_disabled( + simple_query: pl.LazyFrame, engine: SPMDEngine +) -> None: + root = translate_query(simple_query, engine) optimized = optimize_join_domain_prefilters( root, - _stats(part=(part, 6), lineitem=(lineitem, 1_800)), - _config(join_domain_prefilter=False), + StatsCollector(), + make_config(join_domain_prefilter=False), ) assert optimized is root @@ -279,7 +199,7 @@ def test_nullable_join_keys_preserve_results( config, ) - semi_joins = _joins(optimized, "Semi") + semi_joins = find_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) @@ -311,28 +231,36 @@ def test_prefilter_does_not_move_below_distinct_on_non_subset_column( config, ) - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") assert any(isinstance(semi.children[0], Distinct) for semi in semis) 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",)) - root = _join(supplier, lineitem, ("s_suppkey",), ("l_suppkey",)) - stats = _stats(supplier=(supplier, 30), lineitem=(lineitem, 1_800)) - +def test_no_simple_domain_prefilter_when_domain_is_not_selective( + engine: SPMDEngine, +) -> None: + supplier = pl.LazyFrame({"s_suppkey": range(3)}) + lineitem = pl.LazyFrame({"l_suppkey": [i % 3 for i in range(20)]}) + query = supplier.join( + lineitem, + left_on="s_suppkey", + right_on="l_suppkey", + ) + root = translate_query(query, engine) + stats = StatsCollector() + assert isinstance(root, Join) decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) optimized = optimize_join_domain_prefilters( root, stats, - _config(), + ConfigOptions.from_polars_engine(engine), ) assert decision == Decision(reason="no_selective_domain") assert optimized is root - assert not _joins(optimized, "Semi") + assert not find_joins(optimized, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) def test_composite_domain_prefilter_constrains_domain_first( @@ -393,7 +321,7 @@ def test_composite_domain_prefilter_constrains_domain_first( assert decision.reason == "applied" assert isinstance(decision.candidate, CompositeCandidate) - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" assert optimized.children[1] is supplier_ir @@ -446,7 +374,7 @@ def test_derived_selectivity_propagates_through_rewritten_children( ConfigOptions.from_polars_engine(engine), ) - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") expected_targets = { dataframe_scan(root, "n_nationkey"), dataframe_scan(root, "c_custkey"), @@ -502,7 +430,7 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( lineitem_ir = dataframe_scan(root, "l_orderkey") orders_ir = dataframe_scan(root, "o_orderkey") - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") assert sum(semi.children[0] is lineitem_ir for semi in semis) == 1 assert any(semi.children[0] is orders_ir for semi in semis) assert not any( @@ -562,7 +490,7 @@ def test_target_source_follows_join_key_through_rename( ConfigOptions.from_polars_engine(engine), ) - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") assert any(semi.children[0] is small_ir for semi in semis) assert not any(semi.children[0] is big_ir for semi in semis) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -612,7 +540,7 @@ def test_domain_source_follows_join_key_through_rename( ) semi = next( - semi for semi in _joins(optimized, "Semi") if semi.children[0] is target_ir + semi for semi in find_joins(optimized, "Semi") if semi.children[0] is target_ir ) selected_domain = semi.children[1] assert isinstance(selected_domain, Select) @@ -624,15 +552,20 @@ def test_domain_source_follows_join_key_through_rename( assert_gpu_result_equal(query, engine=engine, check_row_order=False) -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", +def test_composite_domain_columns_follow_renames(engine: SPMDEngine) -> None: + query = pl.LazyFrame( + { + "raw_key": range(10), + "raw_constraint": range(10), + } + ).select( + pl.col("raw_key").alias("domain_key"), + pl.col("raw_constraint").alias("domain_constraint"), ) + renamed = translate_query(query, engine) + source = dataframe_scan(renamed, "raw_key") - analyzed = analyze_plan(renamed, _stats(source=(source, 10))) + analyzed = analyze_plan(renamed, StatsCollector()) facts = PlanFacts( row_estimates={renamed: 20, source: 10}, selective_nodes=analyzed.selective_nodes, @@ -674,31 +607,47 @@ def test_composite_domain_columns_do_not_reconverge_after_join( assert_gpu_result_equal(query, engine=engine, check_row_order=False) -def test_contains_node_uses_dag_equality() -> None: - source = _scan("source", ("key",)) - equal_source = _scan("source", ("key",)) - root = _select(source, key="key") +def test_contains_node_uses_dag_equality(engine: SPMDEngine) -> None: + query = pl.LazyFrame({"key": range(3)}).filter(pl.col("key") >= 0).slice(0, 2) + root = translate_query(query, engine) + + assert isinstance(root, Slice) + source = root.children[0] + equal_source = source.reconstruct(source.children) assert source is not equal_source assert source == equal_source assert contains_node(root, equal_source) -def test_plan_facts_share_lineage_suffixes_across_shared_dag() -> None: - source = _scan("source", ("raw_key",)) - left = _select(source, left_key="raw_key") - right = _select(source, right_key="raw_key") - root = _join(left, right, ("left_key",), ("right_key",)) +def test_plan_facts_share_lineage_suffixes_across_shared_dag( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame({"raw_key": range(10)}) + query = source.select(pl.col("raw_key").alias("left_key")).join( + source.select(pl.col("raw_key").alias("right_key")), + left_on="left_key", + right_on="right_key", + ) + root = translate_query(query, engine) - facts = analyze_plan(root, _stats(source=(source, 10))) + assert isinstance(root, Join) + left, right = root.children + source_ir = dataframe_scan(root, "raw_key") + facts = analyze_plan(root, StatsCollector()) left_lineage = facts.column_lineages[ColumnRef(left, "left_key")] right_lineage = facts.column_lineages[ColumnRef(right, "right_key")] - source_lineage = facts.column_lineages[ColumnRef(source, "raw_key")] + source_lineage = facts.column_lineages[ColumnRef(source_ir, "raw_key")] assert left_lineage.column == ColumnRef(left, "left_key") assert right_lineage.column == ColumnRef(right, "right_key") - assert left_lineage.source is source_lineage - assert right_lineage.source is source_lineage + while left_lineage.source is not source_lineage: + assert left_lineage.source is not None + left_lineage = left_lineage.source + while right_lineage.source is not source_lineage: + assert right_lineage.source is not None + right_lineage = right_lineage.source + assert left_lineage.source is right_lineage.source assert source_lineage.source is None @@ -739,7 +688,7 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: ConfigOptions.from_polars_engine(engine), ) - semis = _joins(optimized, "Semi") + semis = find_joins(optimized, "Semi") assert any(semi.children[0] is sliced for semi in semis) assert not any(semi.children[0] is target_ir for semi in semis) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -788,10 +737,10 @@ def test_target_replacement_does_not_rewrite_shared_domain_side( filtered, unfiltered_domain = optimized.children assert unfiltered_domain is domain_ir assert domain_ir.children[0] is shared_ir - semis = _joins(filtered, "Semi") + semis = find_joins(filtered, "Semi") assert len(semis) == 1 assert semis[0].children[0] is dataframe_scan(root, "target_key") - assert not _joins(unfiltered_domain, "Semi") + assert not find_joins(unfiltered_domain, "Semi") assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -840,24 +789,46 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( assert isinstance(rewritten_self_join, Join) filtered, unfiltered = rewritten_self_join.children assert unfiltered is source_ir - filtered_semis = _joins(filtered, "Semi") + filtered_semis = find_joins(filtered, "Semi") assert len(filtered_semis) == 1 - assert not _joins(unfiltered, "Semi") + assert not find_joins(unfiltered, "Semi") assert any(filtered_semis[0].children[0] is node for node in traversal([source_ir])) assert_gpu_result_equal(query, engine=engine, check_row_order=False) -@pytest.mark.parametrize("how", ["Left", "Right", "Cross", "Full"]) -def test_no_domain_prefilter_for_outer_join(how) -> None: - part = _scan("part", ("p_partkey",), predicate=True) - lineitem = _scan("lineitem", ("l_partkey",)) - root = _join(part, lineitem, ("p_partkey",), ("l_partkey",), how=how) +@pytest.mark.parametrize("how", ["left", "right", "cross", "full"]) +def test_no_domain_prefilter_for_outer_join( + how: Any, + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame({"l_partkey": [i % 10 for i in range(20)]}) + if how == "cross": + query = part.join(lineitem, how=how) + else: + query = part.join( + lineitem, + left_on="p_partkey", + right_on="l_partkey", + how=how, + ) + root = translate_query(query, engine) optimized = optimize_join_domain_prefilters( root, - _stats(part=(part, 6), lineitem=(lineitem, 1_800)), - _config(), + StatsCollector(), + ConfigOptions.from_polars_engine(engine), ) assert optimized is root - assert not _joins(optimized, "Semi") + assert not find_joins(optimized, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) From f69f2f2d02c4bd5fcea14fa62d428cf53552df08 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 11:24:32 +0100 Subject: [PATCH 60/65] Update docstring with better explanation --- .../streaming/join_domain_prefilter.py | 28 ++++++++----------- 1 file changed, 11 insertions(+), 17 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 1a546b73675f..3525a687a2f8 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -4,8 +4,8 @@ Insert derived key-domain prefilters for streaming joins. For a supported inner equijoin, this optimization tries to use the join-key -values produced by one input to reduce the other input before the original -join. In relational notation, a simple rewrite is:: +values produced by one input to reduce the size of the other input before +the original join. In relational notation, a simple rewrite is:: left join[left.key = right.key] right @@ -15,9 +15,7 @@ join[left.key = right.key] right In this example, the right hand table is selected to pre-filter the left -table before performing the inner join. The inserted semi-join is therefore -an exact filter: it only removes target rows that could not match the -domain input. +table before performing the inner join. The implementation uses the following terms: @@ -25,25 +23,21 @@ A chain from a named output column towards columns in its input subplan. Each step guarantees that every value in the output column also appears in the referenced child column, although row order and multiplicity are not - preserved and the child may contain additional values. Candidate traversal - does not cross an operator with which a semi-join cannot safely commute. + preserved and the child may contain additional values. ``child edge`` One particular parent-to-child position in the IR DAG. The same child node may occur on more than one edge, so a lineage records child indices and a rewrite follows the resulting edge path to change only the chosen occurrence. -``producer`` - A node on a column lineage, together with the column name at that node and - its edge path from the join input. Producers are possible locations for - inserting a target semi-join or projecting a domain key. ``target`` - The join input to reduce. A semi-join is inserted at a producer in this - input's column lineage and replaces only the selected child-edge - occurrence. + The side of the join to filter. ``domain`` - The other join input, whose join-key values provide the semi-join domain. - The domain producer may be below projections, renames, or other operators - through which the semi-join can safely be pushed. + The side of the join used to provide key values for the filtering of + ``target``. +``producer`` + A node on a column lineage, together with the column name at that node and + its edge path from the join input. So termed because it "produces" the + key values participating in the join. ``constraint domain`` Selective values of another join key from the target input, used to reduce the domain before deriving the values that will filter the target. From 68d01bdfd1c8bc9179e9e4babd1b946c31c06351 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 14:50:19 +0100 Subject: [PATCH 61/65] Remove double checking of source being None --- .../cudf_polars/streaming/join_domain_prefilter.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 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 3525a687a2f8..106a8a8f5bac 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -736,11 +736,12 @@ def _smallest_node_containing_all( lineage.source_child_index != source_child_index for lineage in lineages[1:] ): break - sources = [lineage.source for lineage in lineages] - if any(source is None for source in sources): + sources = [lineage.source for lineage in lineages if lineage.source is not None] + if len(sources) != len(lineages): + # Some sources are None break path = (*path, source_child_index) - lineages = [source for source in sources if source is not None] + lineages = sources if not candidates: return None return min(candidates, key=lambda item: (item[0], item[1]))[2] From 7df367bba064ef264d46e7b1e294b6ccebaa5970 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 16:23:50 +0100 Subject: [PATCH 62/65] Rename to join filter pushdown, provide more explanation --- .../cudf_polars/cudf_polars/engine/options.py | 16 ++--- .../streaming/join_domain_prefilter.py | 4 +- .../cudf_polars/cudf_polars/utils/config.py | 70 +++++++++++-------- .../streaming/test_join_domain_prefilter.py | 8 +-- .../tests/streaming/test_options.py | 6 +- .../tests/streaming/test_parallel.py | 2 +- python/cudf_polars/tests/test_config.py | 62 ++++++++-------- 7 files changed, 88 insertions(+), 80 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 3866bee41f4c..295a081307fa 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -26,7 +26,7 @@ from cudf_polars.quent import QuentContext from cudf_polars.utils.config import ( DynamicPlanningOptions, - JoinDomainPrefilterOptions, + JoinFilterPushdownOptions, ParquetOptions, ) @@ -249,12 +249,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`. ``None`` + join_filter_pushdown + Config for join filter pushdown optimizations, dict or + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions`. ``None`` disables the rewrite. - Env: ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER`` and - ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__*``. + Env: ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN`` and + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__*``. Default: enabled. Category: executor. sink_to_directory @@ -355,8 +355,8 @@ class StreamingOptions: dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = ( _opt("executor") ) - join_domain_prefilter: ( - dict[str, Any] | JoinDomainPrefilterOptions | None | Unspecified + join_filter_pushdown: ( + dict[str, Any] | JoinFilterPushdownOptions | 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 106a8a8f5bac..93ef2f865002 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ -Insert derived key-domain prefilters for streaming joins. +Rewrite a plan, inserting prefilters in join DAGs. For a supported inner equijoin, this optimization tries to use the join-key values produced by one input to reduce the size of the other input before @@ -345,7 +345,7 @@ def optimize_join_domain_prefilters( The rewrite is intentionally conservative: only inner joins with simple column equality keys are considered. """ - options = config_options.executor.join_domain_prefilter + options = config_options.executor.join_filter_pushdown if options is None: return ir threshold = options.threshold diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 912187ca7e60..9f4f2c2cecc0 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -54,7 +54,7 @@ "DaskContext", "DynamicPlanningOptions", "InMemoryExecutor", - "JoinDomainPrefilterOptions", + "JoinFilterPushdownOptions", "ParquetOptions", "RayContext", "SPMDContext", @@ -363,8 +363,9 @@ class DynamicPlanningOptions: The maximum number of chunks to sample before making dynamic-planning decisions. Default is 2. join_prefilter_threshold - Row-count ratio (small / large) below which a join key prefilter is - applied. Set to 0 to disable join prefiltering. Default is 0.5. + Row-count ratio (small / large) below which one side of a join is + filtered by a bloom filter built from the other side before + performing the join. Set to 0 to disable. Default is 0.5. join_prefilter_max_key_columns 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. @@ -430,27 +431,34 @@ def __post_init__(self) -> None: # noqa: D105 @dataclasses.dataclass(frozen=True) -class JoinDomainPrefilterOptions: +class JoinFilterPushdownOptions: """ - Configuration for the logical join-domain prefilter rewrite. + Configuration options for join filter pushdown in the logical plan. - Pass ``None`` to ``StreamingExecutor(join_domain_prefilter=...)`` to + When performing a join between two tables, it is often favourable + to pre-filter one side of the join with the keys (full or partial) of + the other side. This can reduce the size of tables that actually + participate in the join. + + cudf-polars supports a form of this where we can rewrite inner joins by + selecting a side to be filtered by the keys of the other side. + + Pass ``None`` to ``StreamingExecutor(join_filter_pushdown=...)`` to disable the rewrite. These options can be configured via environment variables with the prefix - ``CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__``. + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__``. Parameters ---------- threshold - Row-count ratio (domain / target) below which a derived key-domain - semi-join filter is inserted. Default is 0.5. + Row-count ratio (key-provider-rows / to-be-filtered-table-rows) below which a + filter on is inserted on the to-be-filtered table. Default is 0.5. trace - Whether to emit plan-time trace decisions for derived key-domain - prefilters. Default is False. + Whether to emit plan-time trace decisions for filter decisions. Default is False. """ - _env_prefix = "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER" + _env_prefix = "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN" threshold: float = dataclasses.field( default_factory=_make_default_factory( @@ -740,9 +748,9 @@ class StreamingExecutor: dynamic_planning Options controlling dynamic shuffle planning. See :class:`~cudf_polars.utils.config.DynamicPlanningOptions` for more. - join_domain_prefilter + join_filter_pushdown Options controlling the logical join-domain prefilter rewrite. See - :class:`~cudf_polars.utils.config.JoinDomainPrefilterOptions` for more. + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for more. ``None`` disables the rewrite. max_io_threads Maximum number of IO threads. Default is 4. @@ -813,8 +821,8 @@ class StreamingExecutor: dynamic_planning: DynamicPlanningOptions | None = dataclasses.field( default_factory=DynamicPlanningOptions ) - join_domain_prefilter: JoinDomainPrefilterOptions | None = dataclasses.field( - default_factory=JoinDomainPrefilterOptions + join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field( + default_factory=JoinFilterPushdownOptions ) max_io_threads: int = dataclasses.field( default_factory=_make_default_factory( @@ -877,17 +885,17 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) - if isinstance(self.join_domain_prefilter, dict): + if isinstance(self.join_filter_pushdown, dict): object.__setattr__( self, - "join_domain_prefilter", - JoinDomainPrefilterOptions(**self.join_domain_prefilter), + "join_filter_pushdown", + JoinFilterPushdownOptions(**self.join_filter_pushdown), ) - if self.join_domain_prefilter is not None and not isinstance( - self.join_domain_prefilter, JoinDomainPrefilterOptions + if self.join_filter_pushdown is not None and not isinstance( + self.join_filter_pushdown, JoinFilterPushdownOptions ): raise TypeError( - "join_domain_prefilter must be a JoinDomainPrefilterOptions " + "join_filter_pushdown must be a JoinFilterPushdownOptions " "instance, dict, or None" ) @@ -923,7 +931,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"]) + d["join_filter_pushdown"] = json.dumps(d["join_filter_pushdown"]) # Hash the quent context UUIDs as ints quent_context = d["quent_context"] @@ -1059,16 +1067,16 @@ 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 + # Handle join_filter_pushdown: check user config, then env var + user_join_filter_pushdown = user_executor_options.get( + "join_filter_pushdown", None ) - if user_join_domain_prefilter is None: - env_join_domain_prefilter = os.environ.get( - "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER", "1" + if user_join_filter_pushdown is None: + env_join_filter_pushdown = os.environ.get( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1" ) - if not _bool_converter(env_join_domain_prefilter): - user_executor_options["join_domain_prefilter"] = None + if not _bool_converter(env_join_filter_pushdown): + user_executor_options["join_filter_pushdown"] = None executor = StreamingExecutor(**user_executor_options) case _: # pragma: no cover; Unreachable 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 7cf8e7f13fb7..9aa569cdd074 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py @@ -46,17 +46,17 @@ def engine(spmd_engine_factory) -> SPMDEngine: """Return an SPMD engine configured for join-domain prefilter tests.""" return spmd_engine_factory( StreamingOptions( - join_domain_prefilter={"threshold": 0.5}, + join_filter_pushdown={"threshold": 0.5}, raise_on_fail=True, ) ) def make_config( - *, dynamic_planning: bool = True, join_domain_prefilter: bool = True + *, dynamic_planning: bool = True, join_filter_pushdown: bool = True ) -> ConfigOptions: executor_options: dict[str, Any] = { - "join_domain_prefilter": {"trace": False} if join_domain_prefilter else None + "join_filter_pushdown": {"trace": False} if join_filter_pushdown else None } if not dynamic_planning: executor_options["dynamic_planning"] = None @@ -163,7 +163,7 @@ def test_domain_prefilter_can_be_disabled( optimized = optimize_join_domain_prefilters( root, StatsCollector(), - make_config(join_domain_prefilter=False), + make_config(join_filter_pushdown=False), ) assert optimized is root diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index c9d07f687ea4..d1a5a54aafa2 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -83,9 +83,9 @@ 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 +def test_executor_options_join_filter_pushdown_disabled() -> None: + result = StreamingOptions(join_filter_pushdown=None).to_executor_options() + assert result["join_filter_pushdown"] is None # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/streaming/test_parallel.py b/python/cudf_polars/tests/streaming/test_parallel.py index 742daadd6cdb..e32fab2ce3a3 100644 --- a/python/cudf_polars/tests/streaming/test_parallel.py +++ b/python/cudf_polars/tests/streaming/test_parallel.py @@ -97,7 +97,7 @@ def test_optimize_removes_cache_nodes() -> None: query = source.join(source, on="key", suffix="_right") engine = GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": None}, + executor_options={"join_filter_pushdown": None}, ) ir = Translator(query._ldf.visit(), engine).translate_ir() diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 9227da82bbfe..9d2101891d1b 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -30,7 +30,7 @@ Cluster, ConfigOptions, DynamicPlanningOptions, - JoinDomainPrefilterOptions, + JoinFilterPushdownOptions, MemoryResourceConfig, StreamingExecutor, ) @@ -614,9 +614,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.join_domain_prefilter is not None - assert config.executor.join_domain_prefilter.threshold == 0.5 - assert not config.executor.join_domain_prefilter.trace + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.5 + assert not config.executor.join_filter_pushdown.trace def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -656,31 +656,31 @@ 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 + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.5 + assert not config.executor.join_filter_pushdown.trace -def test_join_domain_prefilter_options_from_env( +def test_join_filter_pushdown_options_from_env( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__THRESHOLD", "0.125" + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" ) - monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_DOMAIN_PREFILTER__TRACE", "1") + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - 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 + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.125 + assert config.executor.join_filter_pushdown.trace -def test_join_domain_prefilter_disabled_from_env( +def test_join_filter_pushdown_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") + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "0") + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - assert config.executor.join_domain_prefilter is None + assert config.executor.join_filter_pushdown is None @pytest.mark.parametrize("value, expected", [("none", None), ("null", None), ("2", 2)]) @@ -775,62 +775,62 @@ def test_validate_join_prefilter_trace() -> None: ) -def test_validate_join_domain_prefilter_options() -> None: +def test_validate_join_filter_pushdown_options() -> None: with pytest.raises(TypeError, match="threshold must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": {"threshold": "bad"}}, + executor_options={"join_filter_pushdown": {"threshold": "bad"}}, ) ) with pytest.raises(ValueError, match="threshold must be between"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": {"threshold": 1.5}}, + executor_options={"join_filter_pushdown": {"threshold": 1.5}}, ) ) with pytest.raises(TypeError, match="trace must be"): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": {"trace": "bad"}}, + executor_options={"join_filter_pushdown": {"trace": "bad"}}, ) ) -def test_validate_join_domain_prefilter_type() -> None: +def test_validate_join_filter_pushdown_type() -> None: with pytest.raises( TypeError, - match="join_domain_prefilter must be a JoinDomainPrefilterOptions instance", + match="join_filter_pushdown must be a JoinFilterPushdownOptions instance", ): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": object()}, + executor_options={"join_filter_pushdown": object()}, ) ) -def test_join_domain_prefilter_from_instance() -> None: - options = JoinDomainPrefilterOptions(threshold=0.25, trace=True) +def test_join_filter_pushdown_from_instance() -> None: + options = JoinFilterPushdownOptions(threshold=0.25, trace=True) config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": options}, + executor_options={"join_filter_pushdown": options}, ) ) - assert config.executor.join_domain_prefilter is options + assert config.executor.join_filter_pushdown is options -def test_join_domain_prefilter_disabled_from_options() -> None: +def test_join_filter_pushdown_disabled_from_options() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"join_domain_prefilter": None}, + executor_options={"join_filter_pushdown": None}, ) ) - assert config.executor.join_domain_prefilter is None + assert config.executor.join_filter_pushdown is None assert hash(config) == hash(config) From 5c24361360b9f29107171c8d1030bbc0d3d75787 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 16:32:30 +0100 Subject: [PATCH 63/65] Rename module --- .../{join_domain_prefilter.py => join_filter_pushdown.py} | 0 python/cudf_polars/cudf_polars/streaming/parallel.py | 2 +- ...st_join_domain_prefilter.py => test_join_filter_pushdown.py} | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) rename python/cudf_polars/cudf_polars/streaming/{join_domain_prefilter.py => join_filter_pushdown.py} (100%) rename python/cudf_polars/tests/streaming/{test_join_domain_prefilter.py => test_join_filter_pushdown.py} (99%) diff --git a/python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py similarity index 100% rename from python/cudf_polars/cudf_polars/streaming/join_domain_prefilter.py rename to python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 2dcd87bdccc8..7af2ab189a65 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -117,7 +117,7 @@ def optimize_with_stats( IR The optimized IR graph. """ - from cudf_polars.streaming.join_domain_prefilter import ( + from cudf_polars.streaming.join_filter_pushdown import ( optimize_join_domain_prefilters, ) diff --git a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py similarity index 99% rename from python/cudf_polars/tests/streaming/test_join_domain_prefilter.py rename to python/cudf_polars/tests/streaming/test_join_filter_pushdown.py index 9aa569cdd074..2117da931aec 100644 --- a/python/cudf_polars/tests/streaming/test_join_domain_prefilter.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -15,7 +15,7 @@ from cudf_polars.dsl.utils.column_domain import ColumnRef from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import StatsCollector -from cudf_polars.streaming.join_domain_prefilter import ( +from cudf_polars.streaming.join_filter_pushdown import ( CompositeCandidate, Decision, PlanFacts, From 9638a582c59fc4a21b328727da4e5e190108650e Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Wed, 15 Jul 2026 16:38:39 +0100 Subject: [PATCH 64/65] Finish renaming, add more explanatory comments --- .../streaming/join_filter_pushdown.py | 34 +++++++++++----- .../cudf_polars/streaming/parallel.py | 4 +- .../streaming/test_join_filter_pushdown.py | 40 +++++++++---------- 3 files changed, 47 insertions(+), 31 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py index 93ef2f865002..5ba74515cf1b 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -334,16 +334,32 @@ def semijoin_pushdown_candidates( lineage = source -def optimize_join_domain_prefilters( +def optimize_join_filter_pushdown( ir: IR, stats: StatsCollector, config_options: ConfigOptions[StreamingExecutor], ) -> IR: """ - Insert generic semi-join key-domain prefilters before streaming lowering. + Rewrite an IR DAG to apply filter pushdown of keys. - The rewrite is intentionally conservative: only inner joins with simple - column equality keys are considered. + This optimization pass inspects joins in the DAG and attempts to push a + prefilter obtained from the keys of one side of the join onto the + inputs of the other side. This can be highly beneficial at large scale + since if we have a selective join we can avoid data movement by + prefiltering before performing the actual join. + + Parameters + ---------- + ir + DAG to rewrite. + stats + Pre-populated statistics. + config_options + Configuration options controlling the rewrite. + + Returns + ------- + Rewritten DAG. """ options = config_options.executor.join_filter_pushdown if options is None: @@ -794,19 +810,19 @@ def contains_node(root: IR, needle: IR) -> bool: def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: - join_domain_prefilter: dict[str, Any] = { + join_filter_pushdown: dict[str, Any] = { "considered": True, "threshold": threshold, "reason": decision.reason, } record = { "scope": Scope.PLAN.value, - "join_domain_prefilter": join_domain_prefilter, + "join_filter_pushdown": join_filter_pushdown, "actor_ir_id": ir.get_stable_id(), "actor_ir_type": type(ir).__name__, } if (candidate := decision.candidate) is not None: - join_domain_prefilter.update( + join_filter_pushdown.update( { "mode": candidate.mode, "target_side": candidate.target_side, @@ -819,10 +835,10 @@ def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: } ) if isinstance(candidate, CompositeCandidate): - join_domain_prefilter.update( + join_filter_pushdown.update( { "constraint_key": candidate.target_constraint_key.name, "estimated_constraint_rows": candidate.constraint_domain.rows, } ) - log("Join Domain Prefilter", **record) + log("Join Filter Pushdown", **record) diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 7af2ab189a65..2ac5c8c2eef7 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -118,11 +118,11 @@ def optimize_with_stats( The optimized IR graph. """ from cudf_polars.streaming.join_filter_pushdown import ( - optimize_join_domain_prefilters, + optimize_join_filter_pushdown, ) ir = remove_cache_nodes(ir) - return optimize_join_domain_prefilters(ir, stats, config_options) + return optimize_join_filter_pushdown(ir, stats, config_options) def _lower_ir_graph_impl( diff --git a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py index 2117da931aec..3dfe42df7fd3 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -25,7 +25,7 @@ analyze_plan, apply_candidate, contains_node, - optimize_join_domain_prefilters, + optimize_join_filter_pushdown, semijoin_pushdown_candidates, ) from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes @@ -116,7 +116,7 @@ def simple_query() -> pl.LazyFrame: return part.join(lineitem, left_on="p_partkey", right_on="l_partkey") -def test_simple_domain_prefilter_filters_large_side( +def test_simple_prefilter_filters_large_side( simple_query: pl.LazyFrame, engine: SPMDEngine ) -> None: root = translate_query(simple_query, engine) @@ -140,13 +140,13 @@ def test_simple_domain_prefilter_filters_large_side( assert_gpu_result_equal(simple_query, engine=engine, check_row_order=False) -def test_domain_prefilter_is_independent_of_dynamic_planning( +def test_filter_pushdown_is_independent_of_dynamic_planning( simple_query: pl.LazyFrame, engine: SPMDEngine, ) -> None: root = translate_query(simple_query, engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), make_config(dynamic_planning=False), @@ -155,12 +155,12 @@ def test_domain_prefilter_is_independent_of_dynamic_planning( assert find_joins(optimized, "Semi") -def test_domain_prefilter_can_be_disabled( +def test_filter_pushdown_can_be_disabled( simple_query: pl.LazyFrame, engine: SPMDEngine ) -> None: root = translate_query(simple_query, engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), make_config(join_filter_pushdown=False), @@ -193,7 +193,7 @@ def test_nullable_join_keys_preserve_results( ir = Translator(query._ldf.visit(), engine).translate_ir() config = ConfigOptions.from_polars_engine(engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( ir, collect_statistics(ir, config, parquet_stats_executor), config, @@ -225,7 +225,7 @@ def test_prefilter_does_not_move_below_distinct_on_non_subset_column( ir = Translator(query._ldf.visit(), engine).translate_ir() config = ConfigOptions.from_polars_engine(engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( ir, collect_statistics(ir, config, parquet_stats_executor), config, @@ -236,7 +236,7 @@ def test_prefilter_does_not_move_below_distinct_on_non_subset_column( assert_gpu_result_equal(query, engine=engine, check_row_order=False) -def test_no_simple_domain_prefilter_when_domain_is_not_selective( +def test_no_simple_filter_pushdown_when_domain_is_not_selective( engine: SPMDEngine, ) -> None: supplier = pl.LazyFrame({"s_suppkey": range(3)}) @@ -251,7 +251,7 @@ def test_no_simple_domain_prefilter_when_domain_is_not_selective( assert isinstance(root, Join) decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, stats, ConfigOptions.from_polars_engine(engine), @@ -263,7 +263,7 @@ def test_no_simple_domain_prefilter_when_domain_is_not_selective( assert_gpu_result_equal(query, engine=engine, check_row_order=False) -def test_composite_domain_prefilter_constrains_domain_first( +def test_composite_filter_pushdown_constrains_domain_first( engine: SPMDEngine, ) -> None: nation = ( @@ -317,7 +317,7 @@ def test_composite_domain_prefilter_constrains_domain_first( assert isinstance(order_lineitem, Join) lineitem_ir = order_lineitem.children[1] decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) - optimized = optimize_join_domain_prefilters(root, stats, config) + optimized = optimize_join_filter_pushdown(root, stats, config) assert decision.reason == "applied" assert isinstance(decision.candidate, CompositeCandidate) @@ -368,7 +368,7 @@ def test_derived_selectivity_propagates_through_rewritten_children( ) root = translate_query(query, engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), ConfigOptions.from_polars_engine(engine), @@ -422,7 +422,7 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( ) root = translate_query(query, engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), ConfigOptions.from_polars_engine(engine), @@ -484,7 +484,7 @@ def test_target_source_follows_join_key_through_rename( renamed_big_ir, small_ir = joined.children assert isinstance(renamed_big_ir, Select) big_ir = renamed_big_ir.children[0] - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), ConfigOptions.from_polars_engine(engine), @@ -533,7 +533,7 @@ def test_domain_source_follows_join_key_through_rename( target_ir, domain_ir = root.children assert isinstance(domain_ir, Join) renamed_unrelated_ir, domain_source_ir = domain_ir.children - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), ConfigOptions.from_polars_engine(engine), @@ -682,7 +682,7 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: (ColumnRef(sliced, "target_key"), ()), ) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, stats, ConfigOptions.from_polars_engine(engine), @@ -727,7 +727,7 @@ def test_target_replacement_does_not_rewrite_shared_domain_side( assert isinstance(domain_ir, Join) assert domain_ir.children[0] is shared_ir - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), ConfigOptions.from_polars_engine(engine), @@ -797,7 +797,7 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( @pytest.mark.parametrize("how", ["left", "right", "cross", "full"]) -def test_no_domain_prefilter_for_outer_join( +def test_no_filter_pushdown_for_unsupported_joins( how: Any, engine: SPMDEngine, ) -> None: @@ -823,7 +823,7 @@ def test_no_domain_prefilter_for_outer_join( ) root = translate_query(query, engine) - optimized = optimize_join_domain_prefilters( + optimized = optimize_join_filter_pushdown( root, StatsCollector(), ConfigOptions.from_polars_engine(engine), From 3d9bfe4ec1757674c63e628660adc5bcc73cb7e9 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Thu, 16 Jul 2026 11:24:45 +0100 Subject: [PATCH 65/65] Final renaming --- docs/cudf/source/cudf_polars/api.md | 2 +- docs/cudf/source/cudf_polars/options.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index aa64dc463cec..fcb8cf3c6278 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -67,7 +67,7 @@ Most users interact with them through `StreamingOptions` fields rather than dire .. automodule:: cudf_polars.utils.config :members: DynamicPlanningOptions, - JoinDomainPrefilterOptions, + JoinFilterPushdownOptions, MemoryResourceConfig, ParquetOptions, StreamingExecutor, diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index 94ca7d52e1bf..6744ad959a93 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/docs/cudf/source/cudf_polars/options.md @@ -108,7 +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 | +| `join_filter_pushdown` | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `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`