diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 559777f43c7d..aec9fc84c55e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2543,7 +2543,12 @@ class Join(IR): """A join of two dataframes.""" __slots__ = ("left_on", "options", "right_on") - _non_child = ("schema", "left_on", "right_on", "options") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "left_on", + "right_on", + "options", + ) _n_non_child_args = 3 left_on: tuple[expr.NamedExpr, ...] """List of expressions used as keys in the left frame.""" 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 382d0027f649..f15f5f01a0b9 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -21,6 +21,7 @@ Slice, Sort, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: from collections.abc import Mapping @@ -141,3 +142,11 @@ def _( return { name: ColumnBinding(0, name) for name in node.schema if name in child.schema } + + +@column_domain_bindings.register(PushdownFilterHint) +def _(node: PushdownFilterHint) -> Mapping[str, ColumnBinding]: + target = node.children[0] + return { + name: ColumnBinding(0, name) for name in node.schema if name in target.schema + } diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py index a4dd049a0eae..35589eab867e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.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 """RapidsMPF streaming-engine support.""" @@ -15,6 +15,7 @@ import cudf_polars.streaming.actor_graph.io import cudf_polars.streaming.actor_graph.join import cudf_polars.streaming.actor_graph.over +import cudf_polars.streaming.actor_graph.prefilter_actor import cudf_polars.streaming.actor_graph.repartition import cudf_polars.streaming.actor_graph.union # noqa: F401 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 bae39d58cc08..7cb34f02d107 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 @@ -12,6 +12,7 @@ from cudf_polars.dsl.ir import Distinct, GroupBy, Sort from cudf_polars.dsl.traversal import traversal +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.io import StreamingSink from cudf_polars.streaming.join import Join from cudf_polars.streaming.over import Over @@ -107,6 +108,7 @@ def __init__( GroupBy, Distinct, Over, + PushdownFilterHint, ) self.collective_nodes: list[IR] = [ diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index a1e78f43d42c..a93694fe2ded 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -23,6 +23,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.io import StreamingScan from cudf_polars.streaming.over import Over from cudf_polars.utils.config import SPMDContext @@ -178,11 +179,12 @@ def _mark_children_unbounded(node: IR) -> None: for node in traversal([ir]): if node in unbounded: _mark_children_unbounded(node) - elif isinstance(node, (Union, Join, Over)): + elif isinstance(node, (Union, Join, Over, PushdownFilterHint)): # Union processes children sequentially; Join may broadcast one # side; Over buffers (or samples-then-replays) its input before - # producing output. In every case the input source needs - # unbounded fanout so other consumers don't block it. + # producing output; PushdownFilterHint similarly might buffer + # then replay. In every case the input source needs unbounded + # fanout so other consumers don't block it. _mark_children_unbounded(node) elif len(node.children) > 1: # Check if this node is doing any broadcasting. 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 d944847e621b..0c39f9271326 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal +from cudf_streaming import CardinalityEstimator from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -23,7 +24,7 @@ ) from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IR, Join +from cudf_polars.dsl.ir import IR, Join, Projection from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import ( AllGatherManager, @@ -35,17 +36,22 @@ from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) +from cudf_polars.streaming.actor_graph.join_planning import JoinPlanningState from cudf_polars.streaming.actor_graph.nodes import default_node_multi -from cudf_polars.streaming.actor_graph.tracing import send_chunk +from cudf_polars.streaming.actor_graph.prefilter import ( + JoinPrefilterExecution, + add_bloom_prefilter, + choose_prefilter, +) +from cudf_polars.streaming.actor_graph.tracing import LOG_TRACES, send_chunk from cudf_polars.streaming.actor_graph.utils import ( CUDF_ROW_LIMIT, MAX_ROWS_PER_PARTITION, ChannelManager, + ChunkSampler, ChunkStore, NormalizedPartitioning, TableSizeStats, - _sample_chunks, - allgather_reduce, chunk_to_frame, empty_table_chunk, gather_in_task_group, @@ -53,9 +59,15 @@ process_children, recv_metadata, replay_buffered_channel, + sample_inputs, send_metadata, shutdown_on_error, ) +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, +) from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import _concat @@ -69,8 +81,18 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.join_planning import JoinInput + from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterDecision, + PrefilterExecution, + ) from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo + from cudf_polars.streaming.filter_hint import ( + JoinSide, + Prefilter, + PushdownFilterHint, + ) from cudf_polars.utils.config import StreamingExecutor @@ -98,6 +120,53 @@ class JoinStrategy: """The key expressions for the right side. Only used for shuffle joins.""" +@dataclass(frozen=True, slots=True) +class JoinCollectiveIds: + """Named collective-ID slots reserved for a dynamic join.""" + + size_estimate: int + left_redistribution: int + right_redistribution: int + + @classmethod + def from_reserved(cls, collective_ids: list[int]) -> JoinCollectiveIds: + """Construct the named slots from IDs reserved for a dynamic join.""" + if len(collective_ids) < 3: + raise ValueError( + "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." + ) + return cls(*collective_ids[:3]) + + @property + def cardinality_tags(self) -> tuple[int, int]: + """Tags available for concurrent prefilter cardinality estimates.""" + return (self.size_estimate, self.left_redistribution) + + @property + def broadcast(self) -> int: + """ID used by a broadcast join after size estimation completes.""" + return self.left_redistribution + + def shuffle(self, side: JoinSide) -> int: + """Return the collective ID for one shuffle input.""" + if side == "left": + return self.left_redistribution + return self.right_redistribution + + def prefilter(self, strategy: JoinStrategy, target_side: JoinSide) -> int: + """Return the subsequent join collective reused by a prefilter.""" + if strategy.broadcast_side is not None: + if target_side != strategy.broadcast_side: + raise ValueError( + "Only the broadcast input can have an active prefilter" + ) + return self.broadcast + return self.shuffle(target_side) + + @define_actor() async def broadcast_join_actor( context: Context, @@ -145,7 +214,7 @@ async def broadcast_join_actor( trace_ir=ir, ir_context=ir_context, ) as tracer: - await _broadcast_join( + await broadcast_join( context, comm, ir, @@ -154,7 +223,7 @@ async def broadcast_join_actor( ch_left, ch_right, JoinStrategy(broadcast_side=broadcast_side), - [collective_id], + collective_id, target_partition_size, tracer=tracer, ) @@ -249,7 +318,7 @@ async def _broadcast_join_large_chunk( broadcast_side: Literal["left", "right"], *, tracer: ActorTracer | None, -) -> None: +) -> int: """Join one large-side chunk with the small DataFrame(s) and send the result.""" large_df = chunk_to_frame(large_chunk, large_child) large_chunk_size = large_chunk.data_alloc_size() @@ -280,11 +349,13 @@ async def _broadcast_join_large_chunk( output_chunk = TableChunk.from_pylibcudf_table( df.table, df.stream, exclusive_view=True, br=context.br() ) + output_rows = output_chunk.shape[0] await send_chunk(context, ch_out, output_chunk, seq_num, tracer=tracer) del df, large_df + return output_rows -async def _broadcast_join( +async def broadcast_join( context: Context, comm: Communicator, ir: Join, @@ -293,27 +364,27 @@ async def _broadcast_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: JoinStrategy, - collective_ids: list[int], - target_partition_size: int, + collective_id: int, + target_partition_size: int | None, *, tracer: ActorTracer | None, + trace_stats: dict[str, Any] | None = None, ) -> None: """ Execute a broadcast join after initial sampling. The small side is gathered (if not already duplicated) and concatenated into a single DataFrame, then joined with each chunk from the large side. - Pops one collective ID from collective_ids for allgather when needed. + Uses ``collective_id`` for the allgather when needed. """ left_metadata, right_metadata = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), ) - collective_id = collective_ids.pop(0) if collective_ids else 0 broadcast_side = strategy.broadcast_side assert broadcast_side is not None - left, right = ir.children + left, right = ir.children[:2] if tracer is not None: tracer.decision = f"broadcast_{broadcast_side}" @@ -355,8 +426,6 @@ async def _broadcast_join( partitioning=partitioning, duplicated=output_duplicated, ) - await send_metadata(ch_out, context, metadata_out) - small_dfs, small_size = await _collect_small_side_for_broadcast( context, comm, @@ -368,17 +437,26 @@ async def _broadcast_join( concat_size_limit=(target_partition_size if ir.options[0] == "Inner" else None), ) + # Publish output metadata only once the broadcast-side collective has + # completed. Besides making the data channel ready when advertised, this + # permits a consumer to reuse the collective ID after receiving metadata. + await send_metadata(ch_out, context, metadata_out) + + input_rows = 0 + output_rows = 0 while (msg := await large_ch.recv(context)) is not None: - await _broadcast_join_large_chunk( + large_chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + input_rows += large_chunk.shape[0] + output_rows += await _broadcast_join_large_chunk( context, ir, ir_context, ch_out, small_dfs, small_child, - TableChunk.from_message(msg, br=context.br()).make_available_and_spill( - context.br(), allow_overbooking=True - ), + large_chunk, large_child, msg.sequence_number, small_size, @@ -386,9 +464,157 @@ async def _broadcast_join( tracer=tracer, ) + if trace_stats is not None: + trace_stats["input_rows"] = input_rows + trace_stats["output_rows"] = output_rows await ch_out.drain(context) +def add_prefilter( + execution: PrefilterExecution, + comm: Communicator, + *, + spec: Prefilter | PushdownFilterHint, + decision: PrefilterDecision, + target: IR, + domain: IR, + ch_target: Channel[TableChunk], + ch_domain_keys: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + ir_context: IRExecutionContext, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the actors and channels that apply one selected prefilter.""" + context = execution.context + if decision.method == "bloom": + if decision.bloom_bytes is None: + raise ValueError("Bloom prefilter decision has no filter size") + add_bloom_prefilter( + context, + comm, + decision.bloom_bytes, + execution, + names_to_indices(spec.target_on, target.schema), + ch_domain_keys, + ch_target, + ch_filtered, + collective_id, + trace_stats, + ) + elif decision.method == "broadcast_semi_join": + domain_schema = {key.name: key.value.dtype for key in spec.domain_on} + if len(domain_schema) != len(spec.domain_on): + raise ValueError("Broadcast semi-join keys must have unique names") + semi_join = Join( + target.schema, + spec.target_on, + spec.domain_on, + ("Semi", spec.nulls_equal, None, "", False, "none"), + target, + Projection(domain_schema, domain), + ) + execution.add_task( + broadcast_join( + context, + comm, + semi_join, + ir_context, + ch_filtered, + ch_target, + ch_domain_keys, + JoinStrategy(broadcast_side="right"), + collective_id, + target_partition_size=None, + tracer=None, + trace_stats=trace_stats, + ) + ) + else: + raise ValueError(f"Cannot apply prefilter method {decision.method!r}") + + +def make_prefilter_execution( + context: Context, + comm: Communicator, + ir: Join, + ir_context: IRExecutionContext, + strategy: JoinStrategy, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + join_state: JoinPlanningState, + collective_ids: JoinCollectiveIds, +) -> JoinPrefilterExecution: + """Create the actors and channels that realize selected prefilters.""" + execution = JoinPrefilterExecution(context, ch_left, ch_right) + + # Prepare every required domain before connecting target-side filters. This + # is important for opposing direct filters: each filter must consume the + # replay produced while the same input's keys are copied for the other one. + for candidate in join_state.candidates: + decision = candidate.decision + if decision is None: + raise ValueError("Join prefilter has no runtime decision") + spec = candidate.spec + if decision.method == "skip": + continue + + if isinstance(spec.domain, JoinInputDomain): + indices = names_to_indices(spec.domain_on, candidate.domain.node.schema) + candidate.key_channel = execution.buffer_domain(spec.domain.side, indices) + else: + sample = candidate.domain.sample + if sample is None: + raise ValueError("Active external prefilter has no domain sample") + indices = names_to_indices(spec.domain_on, candidate.domain.node.schema) + if indices != tuple(range(len(candidate.domain.node.schema))): + raise ValueError("External prefilter domains must contain only keys") + candidate.key_channel = context.create_channel() + execution.add_channel(candidate.key_channel) + execution.add_task( + replay_buffered_channel( + context, + candidate.key_channel, + candidate.domain.channel, + sample.chunks, + candidate.domain.metadata, + trace_ir=ir, + ) + ) + + for candidate in join_state.candidates: + decision = candidate.decision + assert decision is not None + if decision.method == "skip": + continue + spec = candidate.spec + ch_domain_keys = candidate.key_channel + assert ch_domain_keys is not None + target_side = spec.target_side + target = candidate.target.node + ch_target = execution.join_inputs[target_side] + ch_filtered: Channel[TableChunk] = context.create_channel() + trace_stats = candidate.trace + + add_prefilter( + execution, + comm, + spec=spec, + decision=decision, + target=target, + domain=candidate.domain.node, + ch_target=ch_target, + ch_domain_keys=ch_domain_keys, + ch_filtered=ch_filtered, + collective_id=collective_ids.prefilter(strategy, target_side), + ir_context=ir_context, + trace_stats=trace_stats, + ) + execution.replace_join_input(target_side, ch_filtered) + + return execution + + def _get_key_indices( ir: Join, n_partitioned_keys: int | None, @@ -399,7 +625,7 @@ def _get_key_indices( tuple[NamedExpr, ...], tuple[NamedExpr, ...], ]: - left, right = ir.children + left, right = ir.children[:2] n_keys = n_partitioned_keys if n_partitioned_keys is not None else len(ir.left_on) left_keys = ir.left_on[:n_keys] right_keys = ir.right_on[:n_keys] @@ -438,7 +664,7 @@ async def _join_chunks( recv_metadata(ch_right, context), ) - left, right = ir.children + left, right = ir.children[:2] while True: left_msg, right_msg = await gather_in_task_group( ch_left.recv(context), ch_right.recv(context) @@ -534,7 +760,8 @@ async def _shuffle_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: JoinStrategy, - collective_ids: list[int], + left_collective_id: int, + right_collective_id: int, *, tracer: ActorTracer | None, ) -> None: @@ -577,7 +804,7 @@ async def _shuffle_join( strategy.left_keys, ir.children[0].schema, strategy.shuffle_modulus, - collective_ids.pop(0), + left_collective_id, ), _global_shuffle( context, @@ -588,7 +815,7 @@ async def _shuffle_join( strategy.right_keys, ir.children[1].schema, strategy.shuffle_modulus, - collective_ids.pop(0), + right_collective_id, ), _join_chunks( context, @@ -650,58 +877,6 @@ def _num_indices(partitioning: NormalizedPartitioning) -> int: ) -async def _aggregate_estimates( - context: Context, - comm: Communicator, - left_sample: TableSizeStats, - right_sample: TableSizeStats, - collective_ids: list[int], -) -> tuple[TableSizeStats, TableSizeStats]: - """Aggregate table-size and row estimates across ranks.""" - # AllGather size, row, and chunk count estimates across ranks - totals = await allgather_reduce( - context, - comm, - collective_ids.pop(0), - left_sample.total_size, - right_sample.total_size, - left_sample.total_rows, - right_sample.total_rows, - left_sample.total_chunks, - right_sample.total_chunks, - int(left_sample.is_complete), - int(right_sample.is_complete), - ) - ( - left_total, - right_total, - left_total_rows, - right_total_rows, - left_total_chunks, - right_total_chunks, - left_complete_count, - right_complete_count, - ) = totals - - new_left_sample = TableSizeStats( - chunks=left_sample.chunks, - total_size=left_total, - total_rows=left_total_rows, - total_chunks=left_total_chunks, - is_complete=left_complete_count == comm.nranks, - cardinality=left_sample.cardinality, - ) - new_right_sample = TableSizeStats( - chunks=right_sample.chunks, - total_size=right_total, - total_rows=right_total_rows, - total_chunks=right_total_chunks, - is_complete=right_complete_count == comm.nranks, - cardinality=right_sample.cardinality, - ) - return new_left_sample, new_right_sample - - def _choose_strategy_from_samples( comm: Communicator, ir: Join, @@ -849,78 +1024,254 @@ def _modulus(partitioning: NormalizedPartitioning) -> int | None: return max(large, min_shuffle_modulus) -async def _choose_strategy( +def join_input_requires_redistribution( + strategy: JoinStrategy, + side: Literal["left", "right"], + partitioning: NormalizedPartitioning, + metadata: ChannelMetadata, +) -> bool: + """Return whether the join strategy redistributes an input side.""" + if strategy.broadcast_side is not None: + return side == strategy.broadcast_side and not metadata.duplicated + + indices = strategy.left_indices if side == "left" else strategy.right_indices + if not indices: + return True + desired = HashScheme(indices, strategy.shuffle_modulus) + return not ( + partitioning.inter_rank_scheme == desired + and partitioning.local_scheme == "inherit" + ) + + +def choose_prefilters( + join_state: JoinPlanningState, + strategy: JoinStrategy, + left_partitioning: NormalizedPartitioning, + right_partitioning: NormalizedPartitioning, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> None: + """Choose strategies for prefilters with sufficient available statistics.""" + partitionings = { + "left": left_partitioning, + "right": right_partitioning, + } + for candidate in join_state.candidates: + if candidate.decision is not None: + continue + target = candidate.target.sample + if target is None: + raise ValueError("Join target has not been sampled") + target_side = candidate.spec.target_side + target_requires_redistribution = join_input_requires_redistribution( + strategy, + target_side, + partitionings[target_side], + candidate.target.metadata, + ) + if ( + isinstance(candidate.spec.domain, ExternalDomain) + and candidate.domain.sample is None + and target_requires_redistribution + ): + continue + candidate.decision = choose_prefilter( + candidate.spec, + target, + candidate.domain.sample, + target_requires_redistribution=target_requires_redistribution, + broadcast_limit=broadcast_limit, + bloom_filter_max_size=bloom_filter_max_size, + ) + + +async def collect_samples( + context: Context, + comm: Communicator, + join_state: JoinPlanningState, + inputs: tuple[JoinInput, ...], + sample_chunk_count: int, + target_partition_size: int, + collective_id: int, +) -> None: + """Sample inputs and attach aggregate estimates to their planning state.""" + if not inputs: + return + sampling_inputs = [] + for input_ in inputs: + candidates = [ + candidate + for candidate in join_state.candidates + if candidate.domain is input_ + ] + if len(candidates) > 1: + raise ValueError("One join input cannot provide multiple prefilter domains") + sampling_inputs.append((input_, candidates[0] if candidates else None)) + samplers = [] + for input_, candidate in sampling_inputs: + if candidate is None: + cardinality_estimator = None + cardinality_columns: tuple[int, ...] = () + else: + cardinality_estimator = CardinalityEstimator( + context, + comm, + tag=candidate.cardinality_tag, + ) + cardinality_columns = names_to_indices( + candidate.spec.domain_on, + input_.node.schema, + ) + assert len(cardinality_columns) == len(candidate.spec.domain_on), ( + "Prefilter domain keys must be columns" + ) + samplers.append( + ChunkSampler( + context=context, + ch_in=input_.channel, + max_chunks=sample_chunk_count, + max_bytes=target_partition_size, + ch_in_chunk_count=input_.metadata.local_count, + cardinality_estimator=cardinality_estimator, + cardinality_columns=cardinality_columns, + ) + ) + samples = await sample_inputs( + context, + comm, + samplers, + collective_id, + ) + for (input_, _), sample in zip(sampling_inputs, samples, strict=True): + input_.sample = sample + + +async def release_skipped_external_domains( + context: Context, join_state: JoinPlanningState +) -> None: + """Release buffered data and stop external domains rejected by planning.""" + channels = [] + for candidate in join_state.candidates: + if not isinstance(candidate.spec.domain, ExternalDomain): + continue + if candidate.decision is None: + raise ValueError("Join prefilter has no runtime decision") + if candidate.decision.method != "skip": + continue + if candidate.domain.sample is not None: + candidate.domain.sample.chunks.clear() + channels.append(candidate.domain.channel) + if channels: + await gather_in_task_group(*(channel.shutdown(context) for channel in channels)) + + +async def resolve_prefilters( + context: Context, + comm: Communicator, + join_state: JoinPlanningState, + strategy: JoinStrategy, + left_partitioning: NormalizedPartitioning, + right_partitioning: NormalizedPartitioning, + executor: StreamingExecutor, + collective_id: int, +) -> None: + """Resolve optional prefilters after selecting the join strategy.""" + config = executor.join_filter_pushdown + if config is None or not join_state.candidates: + return + + choose_prefilters( + join_state, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + config.bloom_filter_max_size, + ) + assert executor.dynamic_planning is not None + await collect_samples( + context, + comm, + join_state, + tuple( + candidate.domain + for candidate in join_state.candidates + if isinstance(candidate.spec.domain, ExternalDomain) + and candidate.decision is None + ), + executor.dynamic_planning.sample_chunk_count, + executor.target_partition_size, + collective_id, + ) + choose_prefilters( + join_state, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + config.bloom_filter_max_size, + ) + await release_skipped_external_domains(context, join_state) + + +async def choose_strategy( context: Context, comm: Communicator, ir: Join, - ch_left: Channel[TableChunk], - ch_right: Channel[TableChunk], - left_metadata: ChannelMetadata, - right_metadata: ChannelMetadata, + join_state: JoinPlanningState, executor: StreamingExecutor, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, *, tracer: ActorTracer | None, -) -> tuple[TableSizeStats, TableSizeStats, JoinStrategy]: - """Sample both sides, aggregate estimates, and choose broadcast vs shuffle.""" +) -> JoinStrategy: + """Collect any required samples and choose broadcast vs shuffle.""" + left, right = ir.children[:2] + left_metadata = join_state.left.metadata + right_metadata = join_state.right.metadata nranks = comm.nranks left_partitioning = NormalizedPartitioning.from_keys( left_metadata.partitioning, nranks, - keys=names_to_indices(ir.left_on, ir.children[0].schema, concrete_prefix=True), + keys=names_to_indices(ir.left_on, left.schema, concrete_prefix=True), ) right_partitioning = NormalizedPartitioning.from_keys( right_metadata.partitioning, nranks, - keys=names_to_indices(ir.right_on, ir.children[1].schema, concrete_prefix=True), + keys=names_to_indices(ir.right_on, right.schema, concrete_prefix=True), ) - hash_chunkwise = isinstance( left_partitioning.inter_rank_scheme, HashScheme ) and isinstance(right_partitioning.inter_rank_scheme, HashScheme) - if hash_chunkwise and left_partitioning.is_aligned_with( + chunkwise = hash_chunkwise and left_partitioning.is_aligned_with( right_partitioning, context.br() - ): - # We can use a chunkwise join - chunkwise = True - left_sample = TableSizeStats( + ) + + if chunkwise: + join_state.left.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=left_metadata.local_count, ) - right_sample = TableSizeStats( + join_state.right.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=right_metadata.local_count, ) else: - # Need to shuffle or broadcast - Use sampled data to choose a strategy - chunkwise = False assert executor.dynamic_planning is not None - sample_chunk_count = executor.dynamic_planning.sample_chunk_count - target_partition_size = executor.target_partition_size - left_sample, right_sample = await gather_in_task_group( - _sample_chunks( - context, - ch_left, - sample_chunk_count, - target_partition_size, - left_metadata.local_count, - ), - _sample_chunks( - context, - ch_right, - sample_chunk_count, - target_partition_size, - right_metadata.local_count, - ), - ) - left_sample, right_sample = await _aggregate_estimates( + await collect_samples( context, comm, - left_sample, - right_sample, - collective_ids, + join_state, + (join_state.left, join_state.right), + executor.dynamic_planning.sample_chunk_count, + executor.target_partition_size, + collective_ids.size_estimate, ) + left_sample = join_state.left.sample + right_sample = join_state.right.sample + if left_sample is None or right_sample is None: + raise ValueError("Join inputs have not been sampled") strategy = _choose_strategy_from_samples( comm, ir, @@ -934,8 +1285,17 @@ async def _choose_strategy( chunkwise=chunkwise, tracer=tracer, ) - - return left_sample, right_sample, strategy + await resolve_prefilters( + context, + comm, + join_state, + strategy, + left_partitioning, + right_partitioning, + executor, + collective_ids.size_estimate, + ) + return strategy @define_actor() @@ -947,8 +1307,9 @@ async def join_actor( ch_out: Channel[TableChunk], ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], executor: StreamingExecutor, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, ) -> None: """ Dynamic Join actor that selects the best strategy at runtime. @@ -973,6 +1334,8 @@ async def join_actor( Input channel for the left side. ch_right Input channel for the right side. + ch_prefilter_domains + Input channels providing the prefilter key domains. executor Streaming executor configuration. collective_ids @@ -983,32 +1346,72 @@ async def join_actor( ch_out, ch_left, ch_right, + *ch_prefilter_domains, trace_ir=ir, ir_context=ir_context, ) as tracer: - left_metadata, right_metadata = await gather_in_task_group( + ( + left_metadata, + right_metadata, + *prefilter_domain_metadata, + ) = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), + *(recv_metadata(ch, context) for ch in ch_prefilter_domains), ) - left_sample, right_sample, strategy = await _choose_strategy( - context, - comm, + join_state = JoinPlanningState.create( ir, ch_left, ch_right, + ch_prefilter_domains, left_metadata, right_metadata, + tuple(prefilter_domain_metadata), + collective_ids.cardinality_tags, + ) + + strategy = await choose_strategy( + context, + comm, + ir, + join_state, executor, collective_ids, tracer=tracer, ) + prefilter_traces = [] + for candidate in join_state.candidates: + if candidate.decision is None: + raise ValueError("Join prefilter has no runtime decision") + trace = candidate.decision.trace(candidate.spec) + prefilter_traces.append(trace) + if LOG_TRACES: + candidate.trace = trace + if tracer is not None and prefilter_traces: + tracer.set_extra("join_prefilters", prefilter_traces) + left_sample = join_state.left.sample + right_sample = join_state.right.sample + if left_sample is None or right_sample is None: + raise ValueError("Join inputs have not been sampled") ch_left_replay = context.create_channel() ch_right_replay = context.create_channel() + prefilter_execution = make_prefilter_execution( + context, + comm, + ir, + ir_context, + strategy, + ch_left_replay, + ch_right_replay, + join_state, + collective_ids, + ) async with shutdown_on_error( context, ch_left_replay, ch_right_replay, + *prefilter_execution.channels, trace_ir=ir, ir_context=ir_context, ): @@ -1029,13 +1432,14 @@ async def join_actor( right_metadata, trace_ir=ir, ), + *prefilter_execution.tasks, ] - ch_left = ch_left_replay - ch_right = ch_right_replay + ch_left = prefilter_execution.left + ch_right = prefilter_execution.right if strategy.broadcast_side is not None: actor_tasks.append( - _broadcast_join( + broadcast_join( context, comm, ir, @@ -1044,7 +1448,7 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.broadcast, executor.target_partition_size, tracer=tracer, ) @@ -1060,7 +1464,8 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.shuffle("left"), + collective_ids.shuffle("right"), tracer=tracer, ) ) @@ -1073,7 +1478,7 @@ def _use_pwise_join( ir: Join, ) -> bool: """Whether to use a static-planning partition-wise join.""" - left, right = ir.children + left, right = ir.children[:2] output_count = partition_info[ir].count if ( output_count == 1 @@ -1099,18 +1504,24 @@ def _use_pwise_join( @generate_ir_sub_network.register(Join) +@generate_ir_sub_network.register(JoinWithPrefilter) def _( - ir: Join, rec: SubNetGenerator + ir: Join | JoinWithPrefilter, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: # Join operation. - left, right = ir.children + left, right, *prefilter_domains = ir.children partition_info = rec.state["partition_info"] left_count = partition_info[left].count right_count = partition_info[right].count executor = rec.state["config_options"].executor pwise_join = _use_pwise_join(executor, partition_info, ir) - # Process children + if pwise_join and isinstance(ir, JoinWithPrefilter): + raise AssertionError( + "Partition-wise JoinWithPrefilter should have been simplified " + "during IR lowering" + ) + actors, channels = process_children(ir, rec) # Create output ChannelManager @@ -1140,16 +1551,13 @@ def _( and ir.options[0] in ("Inner", "Left", "Right", "Full", "Semi", "Anti") ): # Dynamic join - decide strategy at runtime - collective_ids = list(rec.state["collective_id_map"].get(ir, [])) - # Join uses up to 3 collective IDs: allgather, left shuffle, and - # right shuffle. - if len(collective_ids) < 3: - raise ValueError( - "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." - ) + collective_ids = JoinCollectiveIds.from_reserved( + rec.state["collective_id_map"].get(ir, []) + ) + # Join uses up to 3 collective IDs. Cardinality allreduces complete + # before the size allgather and join collectives. Runtime prefilters + # reuse the collective ID of the target-side join redistribution, with + # their filtered output channel providing the ordering barrier. actors[ir] = [ join_actor( rec.state["context"], @@ -1159,6 +1567,10 @@ def _( channels[ir].reserve_input_slot(), channels[left].reserve_output_slot(), channels[right].reserve_output_slot(), + tuple( + channels[domain].reserve_output_slot() + for domain in prefilter_domains + ), executor, collective_ids, ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py new file mode 100644 index 000000000000..1b0972bb762e --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Actor-local planning state for dynamic joins and optional prefilters.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, +) + +if TYPE_CHECKING: + from typing import Any, Self + + from cudf_streaming.channel_metadata import ChannelMetadata + from cudf_streaming.table_chunk import TableChunk + from rapidsmpf.streaming.core.channel import Channel + + from cudf_polars.dsl.ir import IR, Join + from cudf_polars.streaming.actor_graph.prefilter import PrefilterDecision + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.streaming.filter_hint import Prefilter + + +@dataclass(slots=True) +class JoinInput: + """Concrete runtime resources for one input to a dynamic join.""" + + node: IR + channel: Channel[TableChunk] + metadata: ChannelMetadata + sample: TableSizeStats | None = None + + +@dataclass(slots=True) +class PrefilterCandidate: + """An optional prefilter and the runtime inputs needed to evaluate it.""" + + spec: Prefilter + target: JoinInput + domain: JoinInput + cardinality_tag: int + decision: PrefilterDecision | None = None + key_channel: Channel[TableChunk] | None = None + trace: dict[str, Any] | None = None + + +@dataclass(frozen=True, slots=True) +class JoinPlanningState: + """Actor-local input and prefilter state for planning a dynamic join.""" + + left: JoinInput + right: JoinInput + candidates: tuple[PrefilterCandidate, ...] = () + + @classmethod + def create( + cls, + ir: Join, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], + left_metadata: ChannelMetadata, + right_metadata: ChannelMetadata, + prefilter_domain_metadata: tuple[ChannelMetadata, ...], + cardinality_tags: tuple[int, ...], + ) -> Self: + """Create actor-local planning state from a join and its runtime inputs.""" + left = JoinInput(ir.children[0], ch_left, left_metadata) + right = JoinInput(ir.children[1], ch_right, right_metadata) + if not isinstance(ir, JoinWithPrefilter): + if ch_prefilter_domains or prefilter_domain_metadata: + raise ValueError("A plain Join cannot have prefilter domain inputs") + return cls(left, right) + + external_inputs = tuple( + JoinInput(node, channel, metadata) + for node, channel, metadata in zip( + ir.children[2:], + ch_prefilter_domains, + prefilter_domain_metadata, + strict=True, + ) + ) + external_prefilter_count = sum( + isinstance(prefilter.domain, ExternalDomain) for prefilter in ir.prefilters + ) + if external_prefilter_count != len(external_inputs): + raise ValueError("Join prefilters and external domain inputs must align") + if len(cardinality_tags) < len(ir.prefilters): + raise ValueError("Each join prefilter requires a cardinality collective ID") + + sides = {"left": left, "right": right} + external_inputs_iter = iter(external_inputs) + cardinality_tags_iter = iter(cardinality_tags) + candidates = [] + for spec in ir.prefilters: + target = sides[spec.target_side] + if isinstance(spec.domain, JoinInputDomain): + domain = sides[spec.domain.side] + else: + domain = next(external_inputs_iter) + candidates.append( + PrefilterCandidate( + spec, + target, + domain, + cardinality_tag=next(cardinality_tags_iter), + ) + ) + return cls(left, right, tuple(candidates)) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py new file mode 100644 index 000000000000..a14b660aae09 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -0,0 +1,438 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime planning helpers for optional prefilters.""" + +from __future__ import annotations + +import math +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Literal + +import pylibcudf as plc +from cudf_streaming import BloomFilter +from cudf_streaming.channel_metadata import ChannelMetadata +from cudf_streaming.table_chunk import TableChunk +from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED +from rapidsmpf.streaming.core.message import Message + +from cudf_polars.streaming.actor_graph.utils import ( + ChunkStore, + recv_metadata, + send_metadata, + shutdown_channels_on_error, +) +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, +) + +if TYPE_CHECKING: + from collections.abc import Coroutine, Iterable, Sequence + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.containers import DataType + from cudf_polars.dsl.expr import NamedExpr + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.streaming.filter_hint import JoinSide, Prefilter + + +def estimate_bytes(dtypes: Sequence[DataType], row_count: int) -> int | None: + """ + Estimate the byte count of a table containing the given datatypes. + + Parameters + ---------- + dtypes + Types of columns in the table. + row_count + Estimated total number of rows. + + Returns + ------- + Estimated table size in bytes, or ``None`` if any dtype is not fixed width. + """ + if not all(plc.traits.is_fixed_width(dtype.plc_type) for dtype in dtypes): + return None + + return int( + # Just assume everything has a validity mask + row_count * sum(plc.types.size_of(dtype.plc_type) + 1 / 8 for dtype in dtypes) + ) + + +@dataclass(frozen=True, slots=True) +class PrefilterDecision: + """Runtime decision for one optional prefilter.""" + + method: Literal["skip", "bloom", "broadcast_semi_join"] + reason: str + target_bytes: int + domain_rows: int | None + estimated_cardinality: int | None = None + bloom_bytes: int | None = None + exact_bytes: int | None = None + + def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: + """Return serializable actor-trace information.""" + result = asdict(self) + result["target_side"] = prefilter.target_side + if isinstance(prefilter.domain, JoinInputDomain): + result["domain_side"] = prefilter.domain.side + else: + assert isinstance(prefilter.domain, ExternalDomain) + result["domain"] = "external" + return result + + +def project_key_chunk( + context: Context, chunk: TableChunk, indices: Iterable[int] +) -> TableChunk: + """Copy selected columns into an owning key chunk.""" + columns = chunk.table_view().columns() + key_table = plc.Table([columns[index] for index in indices]).copy( + stream=chunk.stream, mr=context.br().device_mr + ) + return TableChunk.from_pylibcudf_table( + key_table, + chunk.stream, + exclusive_view=True, + br=context.br(), + ) + + +async def buffer_and_project_keys( + context: Context, + ch_in: Channel[TableChunk], + ch_keys: Channel[TableChunk], + ch_replay: Channel[TableChunk], + indices: Iterable[int], +) -> None: + """ + Project owning key chunks while spill-buffering an input for replay. + + The key channel is produced in full before replay begins. Its consumer must + therefore run concurrently with this coroutine. + """ + chunks = ChunkStore(context) + try: + async with shutdown_channels_on_error(context, ch_in, ch_keys, ch_replay): + metadata = await recv_metadata(ch_in, context) + key_metadata = ChannelMetadata( + local_count=metadata.local_count, + partitioning=None, + duplicated=metadata.duplicated, + ) + await send_metadata(ch_replay, context, metadata) + await send_metadata(ch_keys, context, key_metadata) + indices = tuple(indices) + while (msg := await ch_in.recv(context)) is not None: + sequence_number = msg.sequence_number + chunk = await TableChunk.from_message( + msg, br=context.br() + ).make_available_or_wait(context, net_memory_delta=0) + key_chunk = project_key_chunk(context, chunk, indices) + chunks.insert(Message(sequence_number, chunk)) + await ch_keys.send(context, Message(sequence_number, key_chunk)) + + await ch_keys.drain(context) + for msg in chunks: + await ch_replay.send(context, msg) + await ch_replay.drain(context) + finally: + chunks.clear() + + +async def count_rows_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 recording its row count.""" + async with shutdown_channels_on_error(context, ch_in, ch_out): + 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) + + +class PrefilterExecution: + """Channels and actor tasks used to apply one or more prefilters.""" + + def __init__(self, context: Context) -> None: + self.context = context + self.tasks: list[Coroutine[Any, Any, None]] = [] + self.channels: list[Channel[Any]] = [] + + def add_task(self, task: Coroutine[Any, Any, None]) -> None: + """Add an actor task to the prefilter execution.""" + self.tasks.append(task) + + def add_channel(self, channel: Channel[Any]) -> None: + """Register an auxiliary channel for shutdown on failure.""" + self.channels.append(channel) + + +class JoinPrefilterExecution(PrefilterExecution): + """Channels and actor tasks used to apply prefilters before a join.""" + + def __init__( + self, + context: Context, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + ) -> None: + super().__init__(context) + self.source_inputs = {"left": ch_left, "right": ch_right} + self.join_inputs = dict(self.source_inputs) + self.buffered_domains: set[JoinSide] = set() + + def buffer_domain( + self, + side: JoinSide, + indices: Iterable[int], + ) -> Channel[TableChunk]: + """Buffer one original input and return its owning key channel.""" + if side in self.buffered_domains: + raise ValueError(f"Join input {side!r} is already a prefilter domain") + + ch_keys: Channel[TableChunk] = self.context.create_channel() + ch_replay: Channel[TableChunk] = self.context.create_channel() + self.tasks.append( + buffer_and_project_keys( + self.context, + self.source_inputs[side], + ch_keys, + ch_replay, + indices, + ) + ) + self.channels.extend((ch_keys, ch_replay)) + self.join_inputs[side] = ch_replay + self.buffered_domains.add(side) + return ch_keys + + def replace_join_input( + self, + side: JoinSide, + channel: Channel[TableChunk], + ) -> None: + """Replace one join-facing input with a prefilter output channel.""" + self.join_inputs[side] = channel + self.channels.append(channel) + + @property + def left(self) -> Channel[TableChunk]: + """Current left join input.""" + return self.join_inputs["left"] + + @property + def right(self) -> Channel[TableChunk]: + """Current right join input.""" + return self.join_inputs["right"] + + +def add_bloom_prefilter( + context: Context, + comm: Communicator, + bloom_bytes: int, + execution: PrefilterExecution, + target_indices: Iterable[int], + ch_domain_keys: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the channels and actors for an approximate Bloom prefilter.""" + bloom = BloomFilter( + context, + comm, + LIBCUDF_DEFAULT_HASH_SEED, + bloom_bytes, + ) + ch_filter = context.create_channel() + execution.add_channel(ch_filter) + execution.add_task( + bloom.build( + context, + ch_domain_keys, + ch_filter, + collective_id, + ) + ) + ch_apply_input = ch_target + ch_apply_output = ch_filtered + if trace_stats is not None: + ch_counted_input: Channel[TableChunk] = context.create_channel() + ch_raw_output: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_counted_input) + execution.add_channel(ch_raw_output) + execution.add_task( + count_rows_passthrough( + context, + ch_target, + ch_counted_input, + trace_stats, + "input_rows", + ) + ) + execution.add_task( + count_rows_passthrough( + context, + ch_raw_output, + ch_filtered, + trace_stats, + "output_rows", + ) + ) + ch_apply_input = ch_counted_input + ch_apply_output = ch_raw_output + execution.add_task( + bloom.apply( + context, + ch_filter, + ch_apply_input, + ch_apply_output, + target_indices, + ) + ) + + +def estimate_bloom_filter_bytes( + cardinality: int, + desired_false_positive_rate: float = 0.1, +) -> int: + """Estimate Bloom-filter bytes for the block-split policy.""" + if cardinality < 0: + raise ValueError("cardinality must be non-negative") + if not 0 < desired_false_positive_rate < 1: + raise ValueError("false_positive_rate must be between zero and one") + if cardinality == 0: + return 0 + # TODO: cuco could offer this as a static utility on the policy + # Then we wouldn't have to hardcode these magic numbers. + bits = ( + -8 # number of fingerprint bits + * cardinality + / math.log(1 - desired_false_positive_rate ** (1 / 8)) + ) + return math.ceil(bits / 8) + + +def choose_prefilter( + prefilter: Prefilter, + target: TableSizeStats, + domain: TableSizeStats | None, + *, + target_requires_redistribution: bool, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> PrefilterDecision: + """Choose whether one join prefilter is eligible to be applied.""" + domain_rows = None if domain is None else domain.total_rows + if not target_requires_redistribution: + return PrefilterDecision( + "skip", + "target_not_redistributed", + target.total_size, + domain_rows, + ) + if domain is None: + raise ValueError("A redistributed target requires domain statistics") + if ( + isinstance(prefilter.domain, JoinInputDomain) + and prefilter.target_side == prefilter.domain.side + ): + return PrefilterDecision( + "skip", + "same_input", + target.total_size, + domain_rows, + ) + + return choose_prefilter_method( + prefilter.domain_on, + target, + domain, + broadcast_limit=broadcast_limit, + bloom_filter_max_size=bloom_filter_max_size, + ) + + +def choose_prefilter_method( + domain_on: Sequence[NamedExpr], + target: TableSizeStats, + domain: TableSizeStats, + *, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> PrefilterDecision: + """Choose the implementation for an eligible prefilter.""" + distinct_count = domain.distinct_count() + if distinct_count is None: + return PrefilterDecision( + "skip", + "missing_cardinality", + target.total_size, + domain.total_rows, + ) + if distinct_count == 0: + return PrefilterDecision( + "skip", + "zero_cardinality", + target.total_size, + domain.total_rows, + estimated_cardinality=0, + bloom_bytes=0, + exact_bytes=0, + ) + + bloom_bytes = max( + 32, + BloomFilter.aligned_size(estimate_bloom_filter_bytes(distinct_count)), + ) + exact_bytes = estimate_bytes( + tuple(key.value.dtype for key in domain_on), + domain.total_rows, + ) + if bloom_bytes <= min(bloom_filter_max_size, target.total_size): + return PrefilterDecision( + "bloom", + "bloom_fits", + target.total_size, + domain.total_rows, + estimated_cardinality=distinct_count, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) + if exact_bytes is not None and exact_bytes <= min( + broadcast_limit, target.total_size + ): + return PrefilterDecision( + "broadcast_semi_join", + "exact_domain_fits", + target.total_size, + domain.total_rows, + estimated_cardinality=distinct_count, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) + return PrefilterDecision( + "skip", + "no_viable_filter", + target.total_size, + domain.total_rows, + estimated_cardinality=distinct_count, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py new file mode 100644 index 000000000000..76ecee442821 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Standalone execution of optional pushdown-filter hints.""" + +from __future__ import annotations + +from dataclasses import asdict +from typing import TYPE_CHECKING, Any + +from cudf_streaming import CardinalityEstimator +from rapidsmpf.streaming.core.actor import define_actor + +from cudf_polars.dsl.utils.naming import names_to_indices +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.join import add_prefilter +from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterExecution, + choose_prefilter_method, +) +from cudf_polars.streaming.actor_graph.utils import ( + ChannelManager, + ChunkSampler, + gather_in_task_group, + process_children, + recv_metadata, + replay_buffered_channel, + sample_inputs, + shutdown_on_error, +) +from cudf_polars.streaming.filter_hint import PushdownFilterHint + +if TYPE_CHECKING: + from collections.abc import Sequence + + from cudf_streaming.table_chunk import TableChunk + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.utils.config import StreamingExecutor + + +@define_actor() +async def pushdown_filter_actor( + context: Context, + comm: Communicator, + ir: PushdownFilterHint, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_domain: Channel[TableChunk], + executor: StreamingExecutor, + collective_id: int, +) -> None: + """Choose and optionally execute one standalone pushdown-filter hint.""" + collected_samples: Sequence[TableSizeStats] = [] + async with shutdown_on_error( + context, + ch_out, + ch_target, + ch_domain, + trace_ir=ir, + ir_context=ir_context, + ) as tracer: + try: + target_metadata, domain_metadata = await gather_in_task_group( + recv_metadata(ch_target, context), + recv_metadata(ch_domain, context), + ) + dynamic_planning = executor.dynamic_planning + if dynamic_planning is None: + raise ValueError("Standalone prefilters require dynamic planning") + collected_samples = await sample_inputs( + context, + comm, + ( + ChunkSampler( + context=context, + ch_in=ch_target, + max_chunks=dynamic_planning.sample_chunk_count, + max_bytes=executor.target_partition_size, + ch_in_chunk_count=target_metadata.local_count, + ), + ChunkSampler( + context=context, + ch_in=ch_domain, + max_chunks=dynamic_planning.sample_chunk_count, + max_bytes=executor.target_partition_size, + ch_in_chunk_count=domain_metadata.local_count, + cardinality_estimator=CardinalityEstimator( + context, comm, tag=collective_id + ), + cardinality_columns=names_to_indices( + ir.domain_on, ir.children[1].schema + ), + ), + ), + collective_id, + ) + if len(collected_samples) != 2: + raise ValueError("Standalone prefilters require two input samples") + target_sample, domain_sample = collected_samples + config = executor.join_filter_pushdown + if config is None: + raise ValueError("Standalone prefilter has no runtime configuration") + decision = choose_prefilter_method( + ir.domain_on, + target_sample, + domain_sample, + broadcast_limit=executor.broadcast_limit, + bloom_filter_max_size=config.bloom_filter_max_size, + ) + trace = asdict(decision) + trace["placement"] = "standalone" + trace["target_on"] = [key.name for key in ir.target_on] + trace["domain_on"] = [key.name for key in ir.domain_on] + trace_stats = trace if tracer is not None else None + if tracer is not None: + tracer.decision = decision.method + tracer.set_extra("prefilter", trace) + + if decision.method == "skip": + domain_sample.chunks.clear() + await gather_in_task_group( + ch_domain.shutdown(context), + replay_buffered_channel( + context, + ch_out, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ), + ) + else: + target, domain = ir.children + domain_indices = names_to_indices(ir.domain_on, domain.schema) + if domain_indices != tuple(range(len(domain.schema))): + raise ValueError("Pushdown filter domains must contain only keys") + + execution = PrefilterExecution(context) + ch_target_replay: Channel[TableChunk] = context.create_channel() + ch_domain_replay: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_target_replay) + execution.add_channel(ch_domain_replay) + execution.add_task( + replay_buffered_channel( + context, + ch_target_replay, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ) + ) + execution.add_task( + replay_buffered_channel( + context, + ch_domain_replay, + ch_domain, + domain_sample.chunks, + domain_metadata, + trace_ir=ir, + ) + ) + add_prefilter( + execution, + comm, + spec=ir, + decision=decision, + target=target, + domain=domain, + ch_target=ch_target_replay, + ch_domain_keys=ch_domain_replay, + ch_filtered=ch_out, + collective_id=collective_id, + ir_context=ir_context, + trace_stats=trace_stats, + ) + async with shutdown_on_error( + context, + *execution.channels, + trace_ir=ir, + ir_context=ir_context, + ): + await gather_in_task_group(*execution.tasks) + finally: + for sample in collected_samples: + sample.chunks.clear() + + +@generate_ir_sub_network.register(PushdownFilterHint) +def generate_pushdown_filter_subnetwork( + ir: PushdownFilterHint, rec: SubNetGenerator +) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: + """Generate the actor subnetwork for a standalone filter hint.""" + target, domain = ir.children + actors, channels = process_children(ir, rec) + channels[ir] = ChannelManager(rec.state["context"]) + (collective_id,) = rec.state["collective_id_map"][ir] + actors[ir] = [ + pushdown_filter_actor( + rec.state["context"], + rec.state["comm"], + ir, + rec.state["ir_context"], + channels[ir].reserve_input_slot(), + channels[target].reserve_output_slot(), + channels[domain].reserve_output_slot(), + rec.state["config_options"].executor, + collective_id, + ) + ] + return actors, channels 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 1b7153954f2f..3360379566f5 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -7,6 +7,7 @@ import asyncio import contextlib import itertools +import math import operator import struct import time @@ -168,7 +169,7 @@ def _keys_match( class ChunkStore: - """Ordered spillable buffer for TableChunk messages.""" + """Ordered spillable buffer for Messages.""" def __init__(self, ctx: Context) -> None: self._mids: deque[int] = deque() @@ -178,6 +179,12 @@ def __len__(self) -> int: """Return the number of messages in the store.""" return len(self._mids) + def clear(self) -> None: + """Discard all messages in the store.""" + for mid in self._mids: + self._store.extract(mid=mid) + self._mids.clear() + def insert(self, msg: Message) -> None: """Insert a message into the store.""" self._mids.append(self._store.insert(msg)) @@ -317,6 +324,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 ) @@ -1032,6 +1040,57 @@ class TableSizeStats: cardinality: CardinalityEstimate | None = None """Global cardinality statistics for the sampled rows, when requested.""" + def distinct_count(self) -> int | None: + """Extrapolate sampled distinct count to the estimated full row count.""" + if self.total_rows == 0: + return 0 + if self.cardinality is None or self.cardinality.row_count == 0: + return None + return min( + self.total_rows, + math.ceil( + self.cardinality.distinct_count + * self.total_rows + / self.cardinality.row_count + ), + ) + + +async def aggregate_table_size_stats( + context: Context, + comm: Communicator, + samples: tuple[TableSizeStats, ...], + collective_id: int, +) -> tuple[TableSizeStats, ...]: + """Aggregate table-size and row estimates across ranks.""" + totals = await allgather_reduce( + context, + comm, + collective_id, + *( + value + for sample in samples + for value in ( + sample.total_size, + sample.total_rows, + sample.total_chunks, + int(sample.is_complete), + ) + ), + ) + totals_iter = iter(totals) + return tuple( + TableSizeStats( + chunks=sample.chunks, + total_size=next(totals_iter), + total_rows=next(totals_iter), + total_chunks=next(totals_iter), + is_complete=next(totals_iter) == comm.nranks, + cardinality=sample.cardinality, + ) + for sample in samples + ) + @dataclass(frozen=True) class ChunkSampler: @@ -1161,6 +1220,26 @@ async def sample(self) -> TableSizeStats: ) +async def sample_inputs( + context: Context, + comm: Communicator, + samplers: Sequence[ChunkSampler], + collective_id: int, +) -> tuple[TableSizeStats, ...]: + """Sample input channels concurrently and aggregate their statistics.""" + if not samplers: + return () + local_samples = await gather_in_task_group( + *(sampler.sample() for sampler in samplers) + ) + return await aggregate_table_size_stats( + context, + comm, + tuple(local_samples), + collective_id, + ) + + async def _sample_chunks( context: Context, ch: Channel[TableChunk], @@ -1229,19 +1308,23 @@ async def replay_buffered_channel( ch_in The buffered input channel. buffered_chunks - The buffered chunks to yield first. + The buffered chunks to yield first. The store is empty when this + coroutine exits, including on cancellation or error. metadata The metadata to send to the output channel. trace_ir The IR node to trace. Passed through to shutdown_on_error. """ - async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir): - await send_metadata(ch_out, context, metadata) - for msg in buffered_chunks: - await ch_out.send(context, msg) - while (msg := await ch_in.recv(context)) is not None: - await ch_out.send(context, msg) - await ch_out.drain(context) + try: + async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir): + await send_metadata(ch_out, context, metadata) + for msg in buffered_chunks: + await ch_out.send(context, msg) + while (msg := await ch_in.recv(context)) is not None: + await ch_out.send(context, msg) + await ch_out.drain(context) + finally: + buffered_chunks.clear() @dataclass(frozen=True) diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 011e3be232a7..803c057e436b 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -37,8 +37,14 @@ from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import IOPartitionFlavor +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, + PushdownFilterHint, +) from cudf_polars.streaming.io import StreamingScan, scan_partition_plan -from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.parallel import lower_ir_graph, optimize_with_stats from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.statistics import ( collect_statistics, @@ -53,6 +59,7 @@ from cudf_polars.dsl.expressions.base import Expr from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import PartitionInfo, StatsCollector + from cudf_polars.streaming.filter_hint import Prefilter @dataclasses.dataclass @@ -134,6 +141,7 @@ def explain_query( # Include row-count statistics for the logical plan with cm: stats = collect_statistics(ir, config, executor) + ir = optimize_with_stats(ir, config, stats) return _repr_ir_tree(ir, stats=stats) else: return _repr_ir_tree(ir) @@ -469,6 +477,29 @@ def _(ir: Join, *, offset: str = "") -> str: return _repr_header(offset, f"JOIN {ir.options[0]} {left_on} {right_on}", ir.schema) +@_repr_ir.register +def _(ir: JoinWithPrefilter, *, offset: str = "") -> str: + left_on = tuple(ne.name for ne in ir.left_on) + right_on = tuple(ne.name for ne in ir.right_on) + prefilters = tuple(type(prefilter.domain).__name__ for prefilter in ir.prefilters) + return _repr_header( + offset, + f"JOIN {ir.options[0]} {left_on} {right_on} {prefilters=}", + ir.schema, + ) + + +@_repr_ir.register +def _(ir: PushdownFilterHint, *, offset: str = "") -> str: + target_on = tuple(ne.name for ne in ir.target_on) + domain_on = tuple(ne.name for ne in ir.domain_on) + return _repr_header( + offset, + f"PUSHDOWN FILTER HINT {target_on} {domain_on} {ir.placement}", + ir.schema, + ) + + _BinaryOperator = plc.binaryop.BinaryOperator _BINOP_SYMBOLS: dict[_BinaryOperator, str] = { _BinaryOperator.EQUAL: "==", @@ -580,6 +611,45 @@ def _(ir: Join) -> dict[str, Serializable]: } +def _serialize_prefilter(prefilter: Prefilter) -> dict[str, Serializable]: + """Serialize a normalized join prefilter descriptor.""" + properties: dict[str, Serializable] = { + "type": type(prefilter).__name__, + "target_side": prefilter.target_side, + "target_on": [ne.name for ne in prefilter.target_on], + "domain_on": [ne.name for ne in prefilter.domain_on], + "nulls_equal": prefilter.nulls_equal, + } + if isinstance(prefilter.domain, JoinInputDomain): + properties["domain"] = { + "type": type(prefilter.domain).__name__, + "side": prefilter.domain.side, + } + elif isinstance(prefilter.domain, ExternalDomain): + properties["domain"] = {"type": type(prefilter.domain).__name__} + return properties + + +@_serialize_properties.register +def _(ir: JoinWithPrefilter) -> dict[str, Serializable]: + return { + "how": ir.options[0], + "left_on": [ne.name for ne in ir.left_on], + "right_on": [ne.name for ne in ir.right_on], + "prefilters": [_serialize_prefilter(prefilter) for prefilter in ir.prefilters], + } + + +@_serialize_properties.register +def _(ir: PushdownFilterHint) -> dict[str, Serializable]: + return { + "target_on": [ne.name for ne in ir.target_on], + "domain_on": [ne.name for ne in ir.domain_on], + "nulls_equal": ir.nulls_equal, + "placement": ir.placement, + } + + @_serialize_properties.register def _(ir: GroupBy) -> dict[str, Serializable]: return { @@ -815,4 +885,10 @@ def from_query( """ config_options = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() + if not lowered and config_options.executor.name == "streaming": + with concurrent.futures.ThreadPoolExecutor( + thread_name_prefix="cudf-polars-explain" + ) as executor: + stats = collect_statistics(ir, config_options, executor) + ir = optimize_with_stats(ir, config_options, stats) return cls.from_ir(ir, config_options=config_options, lowered=lowered) diff --git a/python/cudf_polars/cudf_polars/streaming/filter_hint.py b/python/cudf_polars/cudf_polars/streaming/filter_hint.py new file mode 100644 index 000000000000..de5f59ad3cdd --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/filter_hint.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Logical filter hints for the streaming runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias + +from cudf_polars.dsl.ir import IR, Join + +if TYPE_CHECKING: + from collections.abc import Sequence + + from cudf_polars.containers import DataFrame + from cudf_polars.dsl.expr import NamedExpr + from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.typing import Schema + + +JoinSide: TypeAlias = Literal["left", "right"] +HintPlacement: TypeAlias = Literal["join_input", "pushed_down"] + + +@dataclass(frozen=True, slots=True) +class JoinInputDomain: + """A prefilter domain provided by an input of its owning join.""" + + side: JoinSide + + +@dataclass(frozen=True, slots=True) +class ExternalDomain: + """A prefilter domain provided by an additional join input.""" + + +PrefilterDomain: TypeAlias = JoinInputDomain | ExternalDomain + + +@dataclass(frozen=True, slots=True) +class Prefilter: + """Description of an optional join prefilter.""" + + target_side: JoinSide + target_on: tuple[NamedExpr, ...] + domain: PrefilterDomain + domain_on: tuple[NamedExpr, ...] + nulls_equal: bool + + +class JoinWithPrefilter(Join): + """Lowered join with normalized prefilter descriptors.""" + + __slots__ = ("prefilters",) + _non_child = ("schema", "left_on", "right_on", "options", "prefilters") + _n_non_child_args = 4 + + prefilters: tuple[Prefilter, ...] + + def __init__( + self, + schema: Schema, + left_on: Sequence[NamedExpr], + right_on: Sequence[NamedExpr], + options: Any, + prefilters: Sequence[Prefilter], + left: IR, + right: IR, + *external_domains: IR, + ): + self.schema = schema + self.left_on = tuple(left_on) + self.right_on = tuple(right_on) + self.options = options + self.prefilters = tuple(prefilters) + self.children = (left, right, *external_domains) + self._non_child_args = ( + self.left_on, + self.right_on, + self.options, + self.prefilters, + ) + + if not self.prefilters: + raise ValueError("JoinWithPrefilter requires at least one prefilter") + external_domain_count = sum( + isinstance(prefilter.domain, ExternalDomain) + for prefilter in self.prefilters + ) + if external_domain_count != len(external_domains): + raise ValueError( + "External prefilters and additional JoinWithPrefilter children " + "must align" + ) + + @classmethod + def do_evaluate( + cls, + left_on: tuple[NamedExpr, ...], + right_on: tuple[NamedExpr, ...], + options: Any, + prefilters: tuple[Prefilter, ...], + left: DataFrame, + right: DataFrame, + *external_domains: DataFrame, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate the join while ignoring its optional prefilters.""" + del prefilters, external_domains + return Join.do_evaluate( + left_on, + right_on, + options, + left, + right, + context=context, + ) + + +class PushdownFilterHint(IR): + """ + Optional join-key filter placed in a logical plan. + + The first child is the target to filter and the second child, the + domain, provides the keys to filter against. Applying the filter is + optional. + """ + + __slots__ = ("domain_on", "nulls_equal", "placement", "target_on") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "target_on", + "domain_on", + "nulls_equal", + "placement", + ) + _n_non_child_args: ClassVar[int] = 4 + + target_on: tuple[NamedExpr, ...] + """Expressions selecting filter keys from the target.""" + domain_on: tuple[NamedExpr, ...] + """Expressions selecting filter keys from the domain.""" + nulls_equal: bool + """Whether null key values compare equal.""" + placement: HintPlacement + """Whether the hint remains at the motivating join input.""" + + def __init__( + self, + schema: Schema, + target_on: Sequence[NamedExpr], + domain_on: Sequence[NamedExpr], + nulls_equal: bool, # noqa: FBT001 + placement: HintPlacement, + target: IR, + domain: IR, + ): + self.schema = schema + self.target_on = tuple(target_on) + self.domain_on = tuple(domain_on) + self.nulls_equal = nulls_equal + self.placement = placement + self._non_child_args = ( + self.target_on, + self.domain_on, + self.nulls_equal, + self.placement, + ) + self.children = (target, domain) + + @classmethod + def do_evaluate( + cls, + target_on: tuple[NamedExpr, ...], + domain_on: tuple[NamedExpr, ...], + nulls_equal: bool, # noqa: FBT001 + placement: HintPlacement, + target: DataFrame, + domain: DataFrame, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Ignore the optional filter and return the target.""" + del placement + return target diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 63d76d6c328f..a9a266a5c166 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -8,10 +8,17 @@ from functools import reduce from typing import TYPE_CHECKING -from cudf_polars.dsl.ir import ConditionalJoin, Join, Slice +from cudf_polars.dsl.ir import ConditionalJoin, Join, Projection, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, + Prefilter, + PushdownFilterHint, +) from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.utils import ( @@ -25,6 +32,7 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.filter_hint import JoinSide from cudf_polars.streaming.parallel import LowerIRTransformer @@ -149,6 +157,144 @@ def _has_non_pointwise_keys(ir: Join) -> bool: return not all(expr.is_pointwise for expr in traversal(keys)) +def is_direct_join_prefilter(ir: IR) -> bool: + """Return whether a hint belongs to its immediately enclosing join.""" + return isinstance(ir, PushdownFilterHint) and ir.placement == "join_input" + + +def lower_join_with_prefilters( + ir: Join, + rec: LowerIRTransformer, +) -> tuple[Join, MutableMapping[IR, PartitionInfo]]: + """Lower a join and normalize its adjacent filter hints.""" + targets = tuple( + child.children[0] if is_direct_join_prefilter(child) else child + for child in ir.children + ) + lowered_targets, target_partition_info = zip( + *(rec(target) for target in targets), + strict=True, + ) + partition_info: MutableMapping[IR, PartitionInfo] = reduce( + operator.or_, target_partition_info + ) + + if all( + isinstance(target, Repartition) and partition_info[target].count == 1 + for target in lowered_targets + ): + # This join will execute partition-wise, so its optional prefilters + # are unnecessary. Moreover, the piecewise join special case + # execution at runtime never has a chance to shut down prefilter + # channels that would be produced, which would leave an actor graph + # in a deadlocked state. Since they are unnecessary, drop them + # before lowering their domains and before the actor graph derives + # fanout from the lowered DAG. + return ( + Join( + ir.schema, + ir.left_on, + ir.right_on, + ir.options, + *lowered_targets, + ), + partition_info, + ) + + prefilters: list[Prefilter] = [] + external_domains: list[IR] = [] + claimed_sides: set[JoinSide] = set() + for target_index, child in enumerate(ir.children): + if not is_direct_join_prefilter(child): + continue + assert isinstance(child, PushdownFilterHint) + + _target, domain = child.children + domain, domain_partition_info = rec(domain) + partition_info.update(domain_partition_info) + + # A key-only Projection retains an explicit edge to its source. If that + # source is a join input and contains every requested key, the join can + # project those keys itself rather than execute a separate domain input. + left, right = lowered_targets + direct_domain = domain + while True: + if direct_domain == left and direct_domain == right: + domain_side: JoinSide | None = "right" if target_index == 0 else "left" + break + if direct_domain == left: + domain_side = "left" + break + if direct_domain == right: + domain_side = "right" + break + if isinstance(direct_domain, Projection) and all( + key.name in direct_domain.children[0].schema for key in child.domain_on + ): + (direct_domain,) = direct_domain.children + continue + domain_side = None + break + + target_side: JoinSide = "left" if target_index == 0 else "right" + if domain_side in claimed_sides: + domain_side = None + elif domain_side is not None: + claimed_sides.add(domain_side) + + if domain_side is not None: + prefilters.append( + Prefilter( + target_side, + child.target_on, + JoinInputDomain(domain_side), + child.domain_on, + child.nulls_equal, + ) + ) + else: + external_domains.append(domain) + prefilters.append( + Prefilter( + target_side, + child.target_on, + ExternalDomain(), + child.domain_on, + child.nulls_equal, + ) + ) + + return ( + JoinWithPrefilter( + ir.schema, + ir.left_on, + ir.right_on, + ir.options, + prefilters, + *lowered_targets, + *external_domains, + ), + partition_info, + ) + + +@lower_ir_node.register(PushdownFilterHint) +def _( + ir: PushdownFilterHint, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + """Preserve optional filters for dynamic execution, otherwise discard them.""" + target, domain = ir.children + target, partition_info = rec(target) + if not _dynamic_planning_on(rec.state["config_options"]): + return target, partition_info + + domain, domain_partition_info = rec(domain) + partition_info.update(domain_partition_info) + lowered = ir.reconstruct((target, domain)) + partition_info[lowered] = partition_info[target] + return lowered, partition_info + + @lower_ir_node.register(ConditionalJoin) def _( ir: ConditionalJoin, rec: LowerIRTransformer @@ -214,17 +360,33 @@ def _( ) return rec(Slice(ir.schema, offset, length, new_join)) - # Lower children - children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True) - partition_info = reduce(operator.or_, _partition_info) - - # Check for dynamic planning - may have more partitions at runtime config_options = rec.state["config_options"] dynamic_planning = _dynamic_planning_on(config_options) + has_non_pointwise_keys = _has_non_pointwise_keys(ir) + if ( + dynamic_planning + and ir.options[0] != "Cross" + and ir.options[5] == "none" + and not has_non_pointwise_keys + and any(is_direct_join_prefilter(child) for child in ir.children) + ): + preserve_prefilters = True + else: + preserve_prefilters = False - left, right = children + if preserve_prefilters: + ir, partition_info = lower_join_with_prefilters(ir, rec) + children = ir.children + else: + # Hints not owned by an adaptive join use the generic identity lowering. + children, _partition_info = zip( + *(rec(child) for child in ir.children), + strict=True, + ) + partition_info = reduce(operator.or_, _partition_info) + + left, right = children[:2] output_count = max(partition_info[left].count, partition_info[right].count) - has_non_pointwise_keys = _has_non_pointwise_keys(ir) if output_count == 1 and not dynamic_planning: new_node = ir.reconstruct(children) partition_info[new_node] = PartitionInfo(count=1) 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 b099e1584ea6..a82c49ef9d82 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -5,13 +5,14 @@ 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 -the original join. In relational notation, a simple rewrite is:: +the original join. It records that opportunity with a logical +``PushdownFilterHint``:: left join[left.key = right.key] right -> - (left semijoin[left.key = right.key] project(right.key)) + PushdownFilterHint(left, left.key, project(right.key), right.key) join[left.key = right.key] right In this example, the right hand table is selected to pre-filter the left @@ -49,14 +50,14 @@ 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. + For a multi-key join, a rewrite that first hints that the domain should be + filtered 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, source scan facts, 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. +``apply_candidate`` then constructs the selected filter-hint rewrite. Row estimates, selectivity propagation, thresholds, and candidate scores are only heuristics for deciding whether a safe rewrite is likely to improve @@ -100,11 +101,13 @@ ColumnRef, column_domain_bindings, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping, Sequence from cudf_polars.streaming.base import StatsCollector + from cudf_polars.streaming.filter_hint import HintPlacement from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -249,6 +252,12 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: rows = node.df.shape()[0] elif isinstance(node, (Select, Projection, HStack, Filter, Distinct, GroupBy)): rows = row_estimates[node.children[0]] + elif isinstance(node, PushdownFilterHint): + rows = _estimate_join_rows( + "Semi", + row_estimates[node.children[0]], + row_estimates[node.children[1]], + ) elif isinstance(node, Join): rows = _estimate_join_rows( node.options[0], @@ -329,7 +338,7 @@ def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: Returns ------- bool - True if a semijoin cannot be pushed past this node, otherwise False. + True if a filter hint cannot be pushed past this node, otherwise False. """ # TODO: Need better cost model to handle nodes that are shared. Pushing # a filter into a shared node will typically mean that it is no longer @@ -350,11 +359,11 @@ def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: ) -def semijoin_pushdown_candidates( +def filter_hint_pushdown_candidates( facts: PlanFacts, root: IR, column: str ) -> Iterator[tuple[ColumnRef, tuple[int, ...]]]: """ - Yield column domain lineage providing valid locations for semijoin pushdown. + Yield column domain lineage providing valid locations for a filter hint. Parameters ---------- @@ -452,7 +461,7 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: if node is original: facts = rec.state["facts"] else: - # Child rewrites introduce new semi joins and reconstructed ancestors. + # Child rewrites introduce new filter hints and reconstructed ancestors. # 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"]) @@ -473,13 +482,13 @@ def apply_candidate(ir: Join, candidate: Candidate) -> IR: left, right = ir.children domain = _make_domain(candidate, ir) target = candidate.target - target_filter = _make_semi_join( + target_filter = _make_filter_hint( 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], + placement="join_input" if not target.path else "pushed_down", ) if candidate.target_side == "left": left = replace_at_path(left, target.path, target_filter) @@ -611,7 +620,7 @@ def _simple_candidates( continue if contains_node(target.node, domain.node): continue - if domain.is_single_source and has_filtering_semi_ancestor( + if domain.is_single_source and has_filtering_hint_ancestor( target_child, target.path ): continue @@ -704,7 +713,7 @@ def _make_domain(candidate: Candidate, ir: Join) -> IR: candidate.constraint_domain.column, candidate.target_constraint_key, ) - constrained = _make_semi_join( + constrained = _make_filter_hint( candidate.domain.node, expr.Col( candidate.domain.node.schema[candidate.domain.columns[1]], @@ -716,7 +725,6 @@ def _make_domain(candidate: Candidate, ir: Join) -> IR: candidate.target_constraint_key.name, ), nulls_equal=ir.options[1], - suffix=ir.options[3], ) return _project_bound_key( constrained, candidate.domain.column, candidate.domain_key @@ -735,20 +743,21 @@ def _project_bound_key(source: IR, bound_column: str, output_key: expr.Col) -> S ) -def _make_semi_join( +def _make_filter_hint( target: IR, target_key: expr.Col, domain: IR, domain_key: expr.Col, *, nulls_equal: bool, - suffix: str, -) -> Join: - return Join( + placement: HintPlacement = "pushed_down", +) -> PushdownFilterHint: + return PushdownFilterHint( target.schema, (expr.NamedExpr(target_key.name, target_key),), (expr.NamedExpr(domain_key.name, domain_key),), - ("Semi", nulls_equal, None, suffix, False, "none"), + nulls_equal, + placement, target, domain, ) @@ -784,7 +793,7 @@ def _smallest_key_producer( exclude: IR | None = None, ) -> _Producer | None: producers = [] - for reference, path in semijoin_pushdown_candidates(facts, root, column): + for reference, path in filter_hint_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name if node is exclude: continue @@ -840,7 +849,7 @@ def _smallest_node_containing_all( def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: source_candidates = [] fallback_candidates = [] - for reference, path in semijoin_pushdown_candidates(facts, root, column): + for reference, path in filter_hint_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name producer = make_producer(node, (bound_column,), path, facts) if producer is None: @@ -890,11 +899,11 @@ def domain_cost_is_small( return domain.cost / target.rows <= threshold -def has_filtering_semi_ancestor(root: IR, path: Sequence[int]) -> bool: - """Return whether a selected child edge is below a filtering semi join.""" +def has_filtering_hint_ancestor(root: IR, path: Sequence[int]) -> bool: + """Return whether a selected child edge is below a pushdown-filter hint.""" node = root for child_index in path: - if isinstance(node, Join) and node.options[0] == "Semi" and child_index == 0: + if isinstance(node, PushdownFilterHint) and child_index == 0: return True node = node.children[child_index] return False diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 2ac5c8c2eef7..55704ce37a60 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -15,6 +15,7 @@ # handlers at import time so the dispatch table is populated before any query # is lowered. import cudf_polars.streaming.distinct +import cudf_polars.streaming.filter_hint import cudf_polars.streaming.groupby import cudf_polars.streaming.io import cudf_polars.streaming.join diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 63fc734a0ca7..26a4a32ab990 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -372,7 +372,7 @@ def __post_init__(self) -> None: # noqa: D105 @dataclasses.dataclass(frozen=True) class JoinFilterPushdownOptions: """ - Configuration options for join filter pushdown in the logical plan. + Configuration options for join filter pushdown. 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 @@ -380,7 +380,8 @@ class JoinFilterPushdownOptions: 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. + selecting a side to be filtered by the keys of the other side. At execution + time, these options also control how optional filters are applied. Pass ``None`` to ``StreamingExecutor(join_filter_pushdown=...)`` to disable the rewrite. @@ -393,6 +394,10 @@ class JoinFilterPushdownOptions: threshold 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. + bloom_filter_max_size + Maximum Bloom-filter size in bytes. If the estimated Bloom filter exceeds + this size, an exact semi-join is preferred when its projected keys fit the + broadcast limit. Set to 0 to disable Bloom filters. Default is 32 MiB. trace Whether to emit plan-time trace decisions for filter decisions. Default is False. """ @@ -404,6 +409,13 @@ class JoinFilterPushdownOptions: f"{_env_prefix}__THRESHOLD", float, default=0.5 ) ) + bloom_filter_max_size: int = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__BLOOM_FILTER_MAX_SIZE", + int, + default=32 * 1024 * 1024, + ) + ) trace: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__TRACE", _bool_converter, default=False @@ -418,6 +430,12 @@ def __post_init__(self) -> None: # noqa: D105 object.__setattr__(self, "threshold", threshold) if not 0.0 <= threshold <= 1.0: raise ValueError("threshold must be between 0 and 1") + if isinstance(self.bloom_filter_max_size, bool) or not isinstance( + self.bloom_filter_max_size, int + ): + raise TypeError("bloom_filter_max_size must be an int") + if self.bloom_filter_max_size < 0: + raise ValueError("bloom_filter_max_size must be non-negative") if not isinstance(self.trace, bool): raise TypeError("trace must be a bool") diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5ea6be578ae4..9458c507e623 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -137,6 +137,54 @@ def test_explain_logical_plan_with_join(tmp_path, df): assert "JOIN Inner ('x',) ('x',)" in plan +def test_explain_pushdown_filter_hint_in_dynamic_physical_plan(): + domain = ( + pl.LazyFrame({"key": [1, 99], "active": [True, False]}) + .filter("active") + .select("key") + ) + target = pl.LazyFrame({"key": [i % 10 for i in range(20)]}) + query = domain.join(target, on="key") + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={"join_filter_pushdown": {"threshold": 0.5}}, + ) + + logical = explain_query(query, engine, physical=False) + physical = explain_query(query, engine, physical=True) + logical_serialized = serialize_query(query, engine, physical=False) + physical_serialized = serialize_query(query, engine, physical=True) + + assert "PUSHDOWN FILTER HINT ('key',) ('key',)" in logical + assert "prefilters=('JoinInputDomain',)" in physical + expected_properties = { + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + "placement": "join_input", + } + assert any( + node.type == "PushdownFilterHint" and node.properties == expected_properties + for node in logical_serialized.nodes.values() + ) + assert any( + node.type == "JoinWithPrefilter" + and node.properties["prefilters"] + == [ + { + "type": "Prefilter", + "target_side": "right", + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + "domain": {"type": "JoinInputDomain", "side": "left"}, + } + ] + for node in physical_serialized.nodes.values() + ) + + def test_explain_logical_plan_with_sort(tmp_path, df): make_partitioned_source(df, tmp_path, fmt="parquet", n_files=2) 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 7601788400db..90e5d52fa92e 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import pytest @@ -11,11 +11,30 @@ from cudf_polars import Translator from cudf_polars.dsl.expr import Col -from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice -from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.ir import ( + IR, + Cache, + DataFrameScan, + Distinct, + Join, + Projection, + Select, + Slice, +) +from cudf_polars.dsl.traversal import CachingVisitor, traversal 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.base import PartitionInfo, StatsCollector +from cudf_polars.streaming.filter_hint import ( + JoinInputDomain, + JoinWithPrefilter, + Prefilter, + PushdownFilterHint, +) +from cudf_polars.streaming.join import ( + is_direct_join_prefilter, + lower_join_with_prefilters, +) from cudf_polars.streaming.join_filter_pushdown import ( CompositeCandidate, Decision, @@ -26,10 +45,15 @@ analyze_plan, apply_candidate, contains_node, + filter_hint_pushdown_candidates, optimize_join_filter_pushdown, - semijoin_pushdown_candidates, ) -from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes +from cudf_polars.streaming.parallel import ( + lower_ir_graph, + optimize_with_stats, + remove_cache_nodes, +) +from cudf_polars.streaming.repartition import Repartition 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 @@ -77,6 +101,10 @@ def find_joins(ir: IR, how: str | None = None) -> list[Join]: ] +def find_hints(ir: IR) -> list[PushdownFilterHint]: + return [node for node in traversal([ir]) if isinstance(node, PushdownFilterHint)] + + def translate_query(query: pl.LazyFrame, engine: SPMDEngine) -> IR: """Translate a public Polars query and remove logical Cache nodes.""" t = Translator(query._ldf.visit(), engine) @@ -95,10 +123,12 @@ def dataframe_scan(ir: IR, column: str) -> DataFrameScan: return match -def join_key_names(join: Join) -> tuple[str, ...]: - """Return the column names used on the left of a simple-column join.""" - names = tuple(key.value.name for key in join.left_on if isinstance(key.value, Col)) - assert len(names) == len(join.left_on) +def hint_key_names(hint: PushdownFilterHint) -> tuple[str, ...]: + """Return the target column names used by a filter hint.""" + names = tuple( + key.value.name for key in hint.target_on if isinstance(key.value, Col) + ) + assert len(names) == len(hint.target_on) return names @@ -141,10 +171,11 @@ def test_simple_prefilter_filters_large_side( assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" - semis = find_joins(optimized, "Semi") - assert len(semis) == 1 - assert semis[0].children[0] is lineitem_ir - assert not find_joins(part_ir, "Semi") + assert not find_joins(optimized, "Semi") + hints = find_hints(optimized) + assert len(hints) == 1 + assert hints[0].children[0] is lineitem_ir + assert not find_hints(part_ir) assert_gpu_result_equal(simple_query, engine=engine, check_row_order=False) @@ -153,14 +184,90 @@ def test_filter_pushdown_is_independent_of_dynamic_planning( engine: SPMDEngine, ) -> None: root = translate_query(simple_query, engine) + config = make_config(dynamic_planning=False) optimized = optimize_join_filter_pushdown( root, StatsCollector(), - make_config(dynamic_planning=False), + config, + ) + + assert find_hints(optimized) + lowering = lower_ir_graph(root, config, StatsCollector()) + assert not any( + isinstance(node, JoinWithPrefilter) for node in traversal([lowering.lowered]) + ) + + +def test_adjacent_filter_hint_is_recorded_on_lowered_join( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + root = translate_query(simple_query, engine) + config = ConfigOptions.from_polars_engine(engine) + + lowering = lower_ir_graph(root, config, StatsCollector()) + + assert find_hints(lowering.optimized) + assert isinstance(lowering.lowered, JoinWithPrefilter) + left, right = lowering.lowered.children + assert not isinstance(left, PushdownFilterHint) + assert not isinstance(right, PushdownFilterHint) + (prefilter,) = lowering.lowered.prefilters + assert isinstance(prefilter, Prefilter) + assert isinstance(prefilter.domain, JoinInputDomain) + assert find_hints(lowering.optimized)[0].placement == "join_input" + assert prefilter.target_side == "right" + assert prefilter.domain.side == "left" + assert tuple(right.schema) == ("l_partkey", "l_suppkey") + assert tuple(ne.name for ne in prefilter.target_on) == ("l_partkey",) + assert tuple(ne.name for ne in prefilter.domain_on) == ("p_partkey",) + assert not prefilter.nulls_equal + assert tuple(left.schema) == ("p_partkey",) + assert not find_joins(lowering.lowered, "Semi") + + +def test_partition_wise_join_discards_prefilters_before_lowering_domains( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + """Partition-wise joins must not retain optional prefilter inputs.""" + root = translate_query(simple_query, engine) + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_with_stats(root, config, StatsCollector()) + assert isinstance(optimized, Join) + + children = list(optimized.children) + (hint_index,) = ( + index for index, child in enumerate(children) if is_direct_join_prefilter(child) ) + hint = children[hint_index] + assert isinstance(hint, PushdownFilterHint) + domain = Projection(hint.children[1].schema, hint.children[1]) + children[hint_index] = hint.reconstruct((hint.children[0], domain)) + optimized = optimized.reconstruct(children) - assert find_joins(optimized, "Semi") + targets = tuple( + child.children[0] if is_direct_join_prefilter(child) else child + for child in optimized.children + ) + repartitions = tuple(Repartition(target.schema, target) for target in targets) + lowered_targets = dict(zip(targets, repartitions, strict=True)) + + def lower_target(child: IR, rec: Any) -> tuple[IR, dict[IR, PartitionInfo]]: + assert child in lowered_targets, "prefilter domain was lowered" + lowered = lowered_targets[child] + return lowered, {lowered: PartitionInfo(count=1)} + + rec: Any = CachingVisitor( + lower_target, + state={"config_options": config}, + ) + lowered, partition_info = lower_join_with_prefilters(optimized, rec) + + assert type(lowered) is Join + assert lowered.children == repartitions + assert domain not in partition_info def test_filter_pushdown_can_be_disabled( @@ -207,9 +314,9 @@ def test_nullable_join_keys_preserve_results( config, ) - semi_joins = find_joins(optimized, "Semi") - assert semi_joins - assert all(join.options[1] is nulls_equal for join in semi_joins) + hints = find_hints(optimized) + assert hints + assert all(hint.nulls_equal is nulls_equal for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -239,8 +346,8 @@ def test_prefilter_does_not_move_below_distinct_on_non_subset_column( config, ) - semis = find_joins(optimized, "Semi") - assert any(isinstance(semi.children[0], Distinct) for semi in semis) + hints = find_hints(optimized) + assert any(isinstance(hint.children[0], Distinct) for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -267,7 +374,7 @@ def test_no_simple_filter_pushdown_when_domain_is_not_selective( assert decision == Decision(reason="no_profitable_domain") assert optimized is root - assert not find_joins(optimized, "Semi") + assert not find_hints(optimized) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -329,12 +436,12 @@ def test_composite_filter_pushdown_constrains_domain_first( assert decision.reason == "applied" assert isinstance(decision.candidate, CompositeCandidate) - semis = find_joins(optimized, "Semi") + hints = find_hints(optimized) assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" 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 any(hint.children[0] is supplier_ir for hint in hints) + assert any(hint.children[0] is lineitem_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -388,16 +495,16 @@ def test_prefilter_uses_cheaper_source_domain_and_skips_expensive_domain( supplier_ir = dataframe_scan(root, "s_suppkey") lineitem_ir = dataframe_scan(root, "l_orderkey") orders_ir = dataframe_scan(root, "o_orderkey") - semis = find_joins(optimized, "Semi") - partkey_semis = [ - semi - for semi in semis - if semi.children[0] is lineitem_ir and join_key_names(semi) == ("l_partkey",) + hints = find_hints(optimized) + partkey_hints = [ + hint + for hint in hints + if hint.children[0] is lineitem_ir and hint_key_names(hint) == ("l_partkey",) ] - assert partkey_semis - assert not any(semi.children[0] is orders_ir for semi in semis) - assert contains_node(partkey_semis[0].children[1], part_ir) - assert not contains_node(partkey_semis[0].children[1], supplier_ir) + assert partkey_hints + assert not any(hint.children[0] is orders_ir for hint in hints) + assert contains_node(partkey_hints[0].children[1], part_ir) + assert not contains_node(partkey_hints[0].children[1], supplier_ir) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -442,13 +549,11 @@ def test_source_only_domain_does_not_stack_on_prefiltered_source( ) lineitem_ir = dataframe_scan(root, "l_partkey") - lineitem_semis = [ - semi - for semi in find_joins(optimized, "Semi") - if semi.children[0] is lineitem_ir + lineitem_hints = [ + hint for hint in find_hints(optimized) if hint.children[0] is lineitem_ir ] - assert any(join_key_names(semi) == ("l_partkey",) for semi in lineitem_semis) - assert not any(join_key_names(semi) == ("l_orderkey",) for semi in lineitem_semis) + assert any(hint_key_names(hint) == ("l_partkey",) for hint in lineitem_hints) + assert not any(hint_key_names(hint) == ("l_orderkey",) for hint in lineitem_hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -496,13 +601,13 @@ def test_derived_selectivity_propagates_through_rewritten_children( ConfigOptions.from_polars_engine(engine), ) - semis = find_joins(optimized, "Semi") + hints = find_hints(optimized) 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 expected_targets <= {hint.children[0] for hint in hints} assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -552,13 +657,10 @@ 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 = find_joins(optimized, "Semi") - assert sum(semi.children[0] is lineitem_ir for semi in semis) == 1 - assert not 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 - ) + hints = find_hints(optimized) + assert sum(hint.children[0] is lineitem_ir for hint in hints) == 1 + assert not any(hint.children[0] is orders_ir for hint in hints) + assert not any(isinstance(hint.children[0], PushdownFilterHint) for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -612,9 +714,9 @@ def test_target_source_follows_join_key_through_rename( ConfigOptions.from_polars_engine(engine), ) - 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) + hints = find_hints(optimized) + assert any(hint.children[0] is small_ir for hint in hints) + assert not any(hint.children[0] is big_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -661,14 +763,11 @@ def test_domain_source_follows_join_key_through_rename( ConfigOptions.from_polars_engine(engine), ) - semi = next( - semi for semi in find_joins(optimized, "Semi") if semi.children[0] is target_ir - ) - selected_domain = semi.children[1] + hint = next(hint for hint in find_hints(optimized) if hint.children[0] is target_ir) + selected_domain = hint.children[1] assert isinstance(selected_domain, Select) rewritten_domain_source = selected_domain.children[0] - assert isinstance(rewritten_domain_source, Join) - assert rewritten_domain_source.options[0] == "Semi" + assert isinstance(rewritten_domain_source, PushdownFilterHint) 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) @@ -721,7 +820,7 @@ 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) - candidates = tuple(semijoin_pushdown_candidates(facts, joined, "value")) + candidates = tuple(filter_hint_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:]) @@ -802,7 +901,7 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: facts = analyze_plan(root, stats) lineage = facts.column_lineages[ColumnRef(sliced, "target_key")] assert lineage.column == ColumnRef(sliced, "target_key") - assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( + assert tuple(filter_hint_pushdown_candidates(facts, sliced, "target_key")) == ( (ColumnRef(sliced, "target_key"), ()), ) @@ -812,9 +911,9 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: ConfigOptions.from_polars_engine(engine), ) - 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) + hints = find_hints(optimized) + assert any(hint.children[0] is sliced for hint in hints) + assert not any(hint.children[0] is target_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -861,10 +960,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 = find_joins(filtered, "Semi") - assert len(semis) == 1 - assert semis[0].children[0] is shared_ir - assert not find_joins(unfiltered_domain, "Semi") + hints = find_hints(filtered) + assert len(hints) == 1 + assert hints[0].children[0] is shared_ir + assert not find_hints(unfiltered_domain) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -913,12 +1012,12 @@ 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 = find_joins(filtered, "Semi") - assert len(filtered_semis) == 1 - assert not find_joins(unfiltered, "Semi") + filtered_hints = find_hints(filtered) + assert len(filtered_hints) == 1 + assert not find_hints(unfiltered) # The shared node is a valid insertion point, but its children are not: - # Only this consumer should be wrapped by the semi-join. - assert filtered_semis[0].children[0] is source_ir + # Only this consumer should be wrapped by the filter hint. + assert filtered_hints[0].children[0] is source_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -967,8 +1066,8 @@ def test_internal_prefilter_rewrites_shared_subplan_once( rewritten_left, rewritten_right = optimized.children assert rewritten_left is rewritten_right assert rewritten_left is not original_shared - (internal_semi,) = find_joins(rewritten_left, "Semi") - assert internal_semi.children[0] is target_ir + (internal_hint,) = find_hints(rewritten_left) + assert internal_hint.children[0] is target_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -1006,5 +1105,5 @@ def test_no_filter_pushdown_for_unsupported_joins( ) assert optimized is root - assert not find_joins(optimized, "Semi") + assert not find_hints(optimized) assert_gpu_result_equal(query, engine=engine, check_row_order=False) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index f584ceff6528..d8675e9192b8 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import json import os import subprocess import sys @@ -131,6 +132,367 @@ def test_structlog_contains_expected_ir_types(timeout_seconds: int): assert b"ir_type=GroupBy" in result +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,output_rows", + [ + (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 10), + (64, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 10), + ( + 1_000_000, + 32 * 1024 * 1024, + "broadcast_left", + "skip", + "target_not_redistributed", + None, + ), + ], + ids=["bloom", "exact", "skip"], +) +def test_local_join_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + join_strategy: str, + method: str, + reason: str, + output_rows: int | None, +) -> None: + """Trace a direct-input join prefilter selected through the public engine.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + import os + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + domain = ( + pl.LazyFrame({{"key": [1, 99], "active": [True, False]}}) + .filter("active") + .select("key") + ) + target = pl.LazyFrame( + {{"key": [i % 100 for i in range(1_000)], "value": range(1_000)}} + ) + query = domain.join(target, on="key") + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 100, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" and "join_prefilters" in log + ) + record = {{ + "result_rows": result.height, + "join_strategy": event["decision"], + "prefilter": event["join_prefilters"][0], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 10 + assert record["join_strategy"] == join_strategy + assert ( + record["prefilter"].items() + >= { + "target_side": "right", + "domain_side": "left", + "method": method, + "reason": reason, + "domain_rows": 1, + }.items() + ) + if output_rows is None: + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 1 + assert record["prefilter"]["input_rows"] == 1_000 + assert record["prefilter"]["output_rows"] == output_rows + + +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,method,reason,output_rows", + [ + (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 20), + (64, 0, "broadcast_semi_join", "exact_domain_fits", 20), + (1, 0, "skip", "no_viable_filter", None), + ], + ids=["bloom", "exact", "skip"], +) +def test_standalone_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + method: str, + reason: str, + output_rows: int | None, +) -> None: + """Trace a prefilter pushed below an intervening join.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + domain = ( + pl.LazyFrame( + {{"p_partkey": range(10), "active": [True] * 2 + [False] * 8}} + ) + .filter("active") + .select("p_partkey") + ) + target = ( + pl.LazyFrame( + {{ + "l_partkey": [i % 10 for i in range(100)], + "bridge_key": range(100), + "value": range(100), + }} + ) + .join(pl.LazyFrame({{"bridge_key": range(100)}}), on="bridge_key") + .with_columns((pl.col("value") + 1).alias("derived")) + ) + query = domain.join(target, left_on="p_partkey", right_on="l_partkey") + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 10, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" + and log.get("prefilter", {{}}).get("placement") == "standalone" + ) + record = {{ + "result_rows": result.height, + "decision": event["decision"], + "prefilter": event["prefilter"], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 20 + assert record["decision"] == method + assert ( + record["prefilter"].items() + >= { + "placement": "standalone", + "method": method, + "reason": reason, + "domain_rows": 2, + }.items() + ) + if output_rows is None: + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 2 + assert record["prefilter"]["input_rows"] == 100 + assert record["prefilter"]["output_rows"] == output_rows + + +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,method,reason,domain_rows", + [ + (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 15), + (512, 0, "broadcast_semi_join", "exact_domain_fits", 15), + ( + 1_000_000, + 32 * 1024 * 1024, + "bloom", + "bloom_fits", + 15, + ), + ], + ids=["bloom", "exact", "bloom_despite_intervening_broadcast"], +) +def test_indirect_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + method: str, + reason: str, + domain_rows: int | None, +) -> None: + """Trace a composite prefilter pushed below an intervening join.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + nation = ( + pl.LazyFrame( + {{"n_nationkey": range(10), "active": [True] * 5 + [False] * 5}} + ) + .filter("active") + .select("n_nationkey") + ) + orders = pl.LazyFrame( + {{ + "o_orderkey": range(90), + "n_nationkey": [i % 10 for i in range(90)], + }} + ) + lineitem = pl.LazyFrame( + {{ + "l_orderkey": [i % 90 for i in range(180)], + "l_suppkey": [i % 60 for i in range(180)], + }} + ) + 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"), + ) + ) + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 100, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" + and log.get("prefilter", {{}}).get("placement") == "standalone" + and log.get("prefilter", {{}}).get("target_on") == ["l_suppkey"] + ) + record = {{ + "result_rows": result.height, + "prefilter": event["prefilter"], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 45 + assert record["prefilter"]["target_on"] == ["l_suppkey"] + assert ( + record["prefilter"].items() + >= { + "placement": "standalone", + "method": method, + "reason": reason, + "domain_rows": domain_rows, + }.items() + ) + if method == "skip": + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == domain_rows + assert record["prefilter"]["input_rows"] == 180 + if method == "broadcast_semi_join": + assert record["prefilter"]["output_rows"] == 45 + else: + assert 45 <= record["prefilter"]["output_rows"] < 180 + + def test_structlog_disabled_by_default(timeout_seconds: int): """Test that structlog does NOT emit events when CUDF_POLARS_LOG_TRACES is not set.""" pytest.importorskip("structlog") diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index aa01207e7e51..f9dbb1604039 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -669,10 +669,15 @@ def test_join_filter_pushdown_options_from_env( monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__BLOOM_FILTER_MAX_SIZE", + "1024", + ) monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.executor.join_filter_pushdown is not None assert config.executor.join_filter_pushdown.threshold == 0.125 + assert config.executor.join_filter_pushdown.bloom_filter_max_size == 1024 assert config.executor.join_filter_pushdown.trace @@ -707,6 +712,24 @@ def test_validate_join_filter_pushdown_options() -> None: executor_options={"join_filter_pushdown": {"trace": "bad"}}, ) ) + with pytest.raises(TypeError, match="bloom_filter_max_size must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "join_filter_pushdown": {"bloom_filter_max_size": "bad"} + }, + ) + ) + with pytest.raises(ValueError, match="bloom_filter_max_size must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "join_filter_pushdown": {"bloom_filter_max_size": -1} + }, + ) + ) def test_validate_join_filter_pushdown_type() -> None: @@ -723,7 +746,9 @@ def test_validate_join_filter_pushdown_type() -> None: def test_join_filter_pushdown_from_instance() -> None: - options = JoinFilterPushdownOptions(threshold=0.25, trace=True) + options = JoinFilterPushdownOptions( + threshold=0.25, bloom_filter_max_size=1024, trace=True + ) config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming",