From 4dff822253c81d6ff63bee1971e8a4e9c247542a Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 28 Jul 2026 16:06:57 +0100 Subject: [PATCH] Remove bloom prefilter application We are going to use the new filterpushdown hint nodes to handle this more coherently. --- .../actor_graph/collectives/common.py | 4 +- .../cudf_polars/streaming/actor_graph/join.py | 418 +----------------- .../streaming/actor_graph/utils.py | 1 - .../cudf_polars/cudf_polars/utils/config.py | 61 --- .../cudf_polars/tests/streaming/test_join.py | 195 +------- .../tests/streaming/test_tracing.py | 6 - python/cudf_polars/tests/test_config.py | 117 ----- 7 files changed, 9 insertions(+), 793 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 c528cf8e42e7..bae39d58cc08 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,13 +138,11 @@ 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: allgather, left shuffle, right shuffle, - # and bloom filter. + # Join needs 3 IDs: allgather, left shuffle, right shuffle. self.collective_id_map[node] = [ _get_new_collective_id_unsafe(), _get_new_collective_id_unsafe(), _get_new_collective_id_unsafe(), - _get_new_collective_id_unsafe(), ] elif isinstance(node, Sort): if self.dynamic_planning_enabled: 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 a984e69957d1..688c4f69d401 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -4,11 +4,9 @@ from __future__ import annotations -from dataclasses import asdict, dataclass +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal -import pylibcudf as plc -from cudf_streaming.bloom_filter import BloomFilter from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -18,13 +16,11 @@ TableChunk, make_table_chunks_available_or_wait, ) -from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED from rapidsmpf.memory.memory_reservation import opaque_memory_usage from rapidsmpf.streaming.core.actor import define_actor from rapidsmpf.streaming.core.memory_reserve_or_wait import ( reserve_memory, ) -from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR, Join @@ -64,9 +60,8 @@ from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: - from collections.abc import Coroutine, Iterable, MutableMapping + from collections.abc import MutableMapping - from cudf_streaming.bloom_filter import BloomFilterChunk from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context @@ -103,30 +98,6 @@ class JoinStrategy: """The key expressions 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, @@ -554,321 +525,6 @@ def _log_shuffle_strategy_decision( tracer.decision = "shuffle" -async def passthrough_split( - context: Context, - ch_in: Channel[TableChunk], - ch_split: Channel[TableChunk], - ch_out: Channel[TableChunk], - *, - indices: Iterable[int], -) -> None: - """ - Pass all messages from ch_in to ch_out, copying key columns to ch_split. - - Parameters - ---------- - context - Streaming context - ch_in - Channel to consume - ch_split - Channel to send key columns to - ch_out - Channel to forward ch_in to - indices - Column indices of the input table to send to ch_split - - Notes - ----- - This sends everything to ch_split before forwarding to ch_out, so the - consumer must consume all of ch_split before consuming ch_out. - """ - meta = await recv_metadata(ch_in, context) - await send_metadata(ch_out, context, meta) - buffer = context.spillable_messages() - mids = [] - while (msg := await ch_in.recv(context)) is not None: - chunk = await TableChunk.from_message( - msg, br=context.br() - ).make_available_or_wait(context, net_memory_delta=0) - columns = chunk.table_view().columns() - key_table = TableChunk.from_pylibcudf_table( - plc.Table( - [ - columns[i].copy(chunk.stream, mr=context.br().device_mr) - for i in indices - ] - ), - chunk.stream, - exclusive_view=True, - br=context.br(), - ) - mids.append(buffer.insert(Message(msg.sequence_number, chunk))) - await ch_split.send(context, Message(msg.sequence_number, key_table)) - await ch_split.drain(context) - for mid in mids: - await ch_out.send(context, buffer.extract(mid=mid)) - await ch_out.drain(context) - - -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: - """ - 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. - """ - if threshold == 0.0: - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - reason_skipped="disabled", - ) - - if join_type not in ("Inner", "Semi", "Left", "Anti", "Right"): - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - reason_skipped="unsupported_join_type", - ) - - if len(left_key_indices) != len(right_key_indices) or len(left_key_indices) == 0: - return JoinPrefilterDecision( - left_rows=left_rows, - right_rows=right_rows, - threshold=threshold, - reason_skipped="expression_keys", - ) - key_column_count = len(left_key_indices) - if max_key_columns is not None: - key_column_count = min(key_column_count, max_key_columns) - - 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( - context: Context, - comm: Communicator, - *, - ch_left: Channel[TableChunk], - ch_right: Channel[TableChunk], - decision: JoinPrefilterDecision, - tag: int, - trace_stats: dict[str, Any] | None, -) -> tuple[ - Channel[TableChunk], - Channel[TableChunk], - list[Coroutine[Any, Any, None]], - list[Channel], -]: - """ - Create bloom filter tasks for a pair of channels participating in a shuffle join. - - Parameters - ---------- - context - Streaming context - comm - Communicator - ch_left - Left input channel - ch_right - Right input channel - 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.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.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: Make the filter size configurable. - filter_size = 32 * 1024 * 1024 - filter = BloomFilter(context, comm, LIBCUDF_DEFAULT_HASH_SEED, filter_size) - 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, - bloom_build_input, - passthrough_output, - indices=build_indices, - ), - filter.build( - context, - bloom_build_input, - bloom_build_output, - tag, - ), - filter.apply( - context, - bloom_build_output, - apply_input, - apply_output, - apply_indices, - ), - ] - return ch_left, ch_right, filter_tasks, chs_to_shutdown - - async def _shuffle_join( context: Context, comm: Communicator, @@ -880,11 +536,7 @@ async def _shuffle_join( strategy: JoinStrategy, collective_ids: list[int], *, - row_counts: tuple[int, int], tracer: ActorTracer | None, - prefilter_threshold: float, - prefilter_max_key_columns: int | None, - prefilter_trace: bool, ) -> None: """Execute a shuffle (hash) join.""" # Send output metadata @@ -902,37 +554,6 @@ async def _shuffle_join( duplicated=False, ) await send_metadata(ch_out, context, metadata_out) - left_rows, right_rows = row_counts - bloom_tag = collective_ids.pop(0) - prefilter_decision = _select_join_prefilter( - ir.options[0], - left_rows, - right_rows, - strategy.left_indices, - strategy.right_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'}_prefiltered" - ch_left, ch_right, filter_tasks, chs_to_shutdown = make_filter_tasks( - context, - comm, - ch_left=ch_left, - ch_right=ch_right, - decision=prefilter_decision, - tag=bloom_tag, - trace_stats=prefilter_trace_stats if prefilter_trace else None, - ) - else: - filter_tasks = [] - chs_to_shutdown = [] # Construct a shuffle-shuffle-join pipeline. # The shuffle operations will pass chunks through unchanged # if the data is already partitioned correctly. @@ -941,14 +562,12 @@ async def _shuffle_join( # note: this is an actor inside of an actor. How should we log that in our traces? async with shutdown_on_error( context, - *chs_to_shutdown, ch_left_shuffle, ch_right_shuffle, trace_ir=ir, ir_context=ir_context, ): actor_tasks = [ - *filter_tasks, _global_shuffle( context, comm, @@ -1422,22 +1041,6 @@ 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, @@ -1449,14 +1052,7 @@ async def join_actor( ch_right, strategy, collective_ids, - row_counts=( - left_sample.total_rows, - right_sample.total_rows, - ), tracer=tracer, - prefilter_threshold=prefilter_threshold, - prefilter_max_key_columns=prefilter_max_key_columns, - prefilter_trace=prefilter_trace, ) ) await gather_in_task_group(*actor_tasks) @@ -1536,12 +1132,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: allgather, left shuffle, right - # shuffle, and bloom filter. - if len(collective_ids) < 4: + # Join uses up to 3 collective IDs: allgather, left shuffle, and + # right shuffle. + if len(collective_ids) < 3: raise ValueError( - "Dynamic join requires 4 reserved collective IDs " - "(allgather + left shuffle + right shuffle + bloom filter); got " + "Dynamic join requires 3 reserved collective IDs " + "(allgather + 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/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 256e45440d8d..454ca23895f9 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -311,7 +311,6 @@ 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 ab9dc846b8a5..63fc734a0ca7 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -171,16 +171,6 @@ 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) - - def _quent_context_converter(v: str) -> QuentContext | None: from cudf_polars.quent._context import QuentContext @@ -362,16 +352,6 @@ class DynamicPlanningOptions: sample_chunk_count 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 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. - join_prefilter_trace - Whether to collect input/output row counts around applied join - prefilters. Default is False. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -381,53 +361,12 @@ class DynamicPlanningOptions: f"{_env_prefix}__SAMPLE_CHUNK_COUNT", int, default=2 ) ) - 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}__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, - ) - ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.sample_chunk_count, int): 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") - 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) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index 2ac2edc88903..2733b051925a 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -15,10 +15,7 @@ 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 ( - _select_join_prefilter, - _use_pwise_join, -) +from cudf_polars.streaming.actor_graph.join import _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 @@ -235,196 +232,6 @@ def test_join_and_slice(request, zlice, streaming_engine_factory): assert_gpu_result_equal(q, engine=streaming_engine) -@pytest.mark.parametrize("how", ["inner", "semi", "left", "right"]) -def test_bloom_filter_join(how, streaming_engine_factory): - streaming_engine = streaming_engine_factory( - StreamingOptions( - max_rows_per_partition=2, - broadcast_limit=10, - target_partition_size=10, - ), - ) - dim = pl.LazyFrame({"key": range(10), "val": range(10)}) - fact = pl.LazyFrame({"key": range(200), "data": range(200)}) - left, right = (dim, fact) if how == "right" else (fact, dim) - q = left.join(right, on="key", how=how) - 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_skips_mismatched_key_count() -> None: - decision = _select_join_prefilter( - "Inner", - 10, - 1_000, - (0,), - (0, 1), - threshold=0.5, - max_key_columns=1, - ) - assert not decision.enabled - assert decision.reason_skipped == "expression_keys" - - @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 fa7c508670bd..f584ceff6528 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -42,12 +42,6 @@ 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 d77fe56ec685..f04bd5e9a67b 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -611,9 +611,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.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_filter_pushdown is None @@ -638,28 +635,6 @@ def test_dynamic_planning_sample_chunk_count_from_env( assert config.executor.dynamic_planning.sample_chunk_count == 3 -def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1") - 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 - 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_filter_pushdown_options_from_env( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -683,98 +658,6 @@ def test_join_filter_pushdown_disabled_from_env( assert config.executor.join_filter_pushdown 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 -) -> 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": {"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(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": {"join_prefilter_max_key_columns": 0} - }, - ) - ) - - -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_validate_join_filter_pushdown_options() -> None: with pytest.raises(TypeError, match="threshold must be"): ConfigOptions.from_polars_engine(