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..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 @@ -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: 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 5303f9a8d991..031a7658f984 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 @@ -42,7 +42,6 @@ ChannelManager, NormalizedPartitioning, TableSizeStats, - _is_already_partitioned, _sample_chunks, allgather_reduce, chunk_to_frame, @@ -59,8 +58,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 +92,30 @@ 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.""" + + left_rows: int + right_rows: int + threshold: float + filter_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.""" + return asdict(self) + + @define_actor() async def broadcast_join_actor( context: Context, @@ -565,22 +587,134 @@ async def passthrough_split( await ch_out.drain(context) -def use_bloom_filter( +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, -) -> 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 + max_key_columns: int | None, +) -> JoinPrefilterDecision: + """ + Determine whether to apply a prefilter to a join. + + 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), ( + "left and right join key counts must match" + ) + if threshold == 0.0: + return JoinPrefilterDecision( + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + reason_skipped="disabled", + ) + + if max_key_columns is not None: + key_column_count = min(key_column_count, max_key_columns) + + if join_type not in ("Inner", "Semi", "Left", "Anti", "Right"): + return JoinPrefilterDecision( + left_rows=left_rows, + right_rows=right_rows, + threshold=threshold, + key_column_count=key_column_count, + reason_skipped="unsupported_join_type", + ) + + 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 + + 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] + 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, + right_rows=right_rows, + threshold=threshold, + filter_side=filter_side, + build_indices=build_indices, + apply_indices=apply_indices, + key_column_count=key_column_count, + small_large_ratio=ratio, + reason_skipped=reason_skipped, + ) + + +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( @@ -589,14 +723,13 @@ def make_filter_tasks( *, 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,57 +745,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. """ - if left_rows < right_rows: - passthrough_input = ch_left - build_indices = strategy.left_indices - bloom_apply_input = ch_right - apply_indices = strategy.right_indices - apply_meta = strategy.right_meta - else: - passthrough_input = ch_right - build_indices = strategy.right_indices - bloom_apply_input = ch_left - apply_indices = strategy.left_indices - 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, [], [] - + assert decision.enabled + 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 left_rows < right_rows: + if decision.filter_side == "right": + passthrough_input = ch_left ch_left = passthrough_output + build_indices = decision.build_indices + bloom_apply_input = ch_right + apply_indices = decision.apply_indices ch_right = context.create_channel() bloom_apply_output = ch_right else: + passthrough_input = ch_right ch_right = passthrough_output + build_indices = decision.build_indices + bloom_apply_input = ch_left + apply_indices = decision.apply_indices ch_left = context.create_channel() bloom_apply_output = ch_left + # TODO: configure based on GPU L2 size nblocks = BloomFilter.fitting_num_blocks(32 * 1024 * 1024) filter = BloomFilter(context, comm, LIBCUDF_DEFAULT_HASH_SEED, nblocks) + 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, @@ -679,16 +834,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 @@ -705,7 +855,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 @@ -725,18 +877,32 @@ 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 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 = [] @@ -1211,6 +1377,22 @@ async def join_actor( ) ) else: + dynamic_options = executor.dynamic_planning + prefilter_threshold = ( + dynamic_options.join_prefilter_threshold + if dynamic_options 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, @@ -1227,11 +1409,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) @@ -1311,10 +1491,11 @@ 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: allgather, left shuffle, right + # shuffle, and bloom filter. if len(collective_ids) < 4: raise ValueError( - "Dynamic join requires 3 reserved collective IDs " + "Dynamic join requires 4 reserved collective IDs " "(allgather + left shuffle + right shuffle + bloom filter); 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..c318affa4d2a 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,15 @@ 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 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 + 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..0a74cade39b8 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,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 @@ -168,6 +169,16 @@ def _bool_converter(v: str) -> bool: raise ValueError(f"Invalid boolean value: '{v}'") +def _optional_converter(v: str, parse: Callable[[str], T]) -> T | None: + if v.lower() in {"none", "null"}: + return None + return parse(v) + + +def _optional_int_converter(v: str) -> int | None: + return _optional_converter(v, int) + + @dataclasses.dataclass(frozen=True) class ParquetOptions: """ @@ -304,10 +315,15 @@ 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 the large side of an inner or semi shuffle join. - Set to 0 to disable bloom filtering. Default is 0.5. + 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. + 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. + join_prefilter_trace + Whether to collect input/output row counts around applied join + prefilters. Default is False. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -317,9 +333,25 @@ class DynamicPlanningOptions: f"{_env_prefix}__SAMPLE_CHUNK_COUNT", int, default=2 ) ) - bloom_filter_threshold: float = dataclasses.field( + join_prefilter_threshold: float = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_PREFILTER_THRESHOLD", + float, + default=0.5, + ) + ) + join_prefilter_max_key_columns: int | None = dataclasses.field( default_factory=_make_default_factory( - f"{_env_prefix}__BLOOM_FILTER_THRESHOLD", float, default=0.5 + f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", + _optional_int_converter, + default=1, + ) + ) + join_prefilter_trace: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__JOIN_PREFILTER_TRACE", + _bool_converter, + default=False, ) ) @@ -328,10 +360,26 @@ 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 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: + 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( + "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..8ac799115032 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,181 @@ 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.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.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.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" + + +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" + + +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"] ) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index f584ceff6528..fa7c508670bd 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", {"enabled": True}) + assert tracer.extra == {"join_prefilter": {"enabled": 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..ad00a5b69d59 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, @@ -672,7 +673,9 @@ 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 def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -696,40 +699,117 @@ 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"): +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 + + +@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 == expected + + +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( + executor="streaming", + executor_options={ + "dynamic_planning": {"join_prefilter_threshold": "bad"} + }, + ) + ) + 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( executor="streaming", executor_options={ - "dynamic_planning": {"bloom_filter_threshold": "bad"} + "dynamic_planning": {"join_prefilter_threshold": 1.5} }, ) ) -def test_validate_bloom_filter_threshold_range() -> None: - with pytest.raises(ValueError, match="bloom_filter_threshold must be between"): +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(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" + ): ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"dynamic_planning": {"bloom_filter_threshold": 1.5}}, + executor_options={ + "dynamic_planning": {"join_prefilter_max_key_columns": 0} + }, ) ) -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 +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: - from cudf_polars.utils.config import DynamicPlanningOptions - config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming",