diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index a1c6fcb76ff0..5c099b8b7177 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -76,6 +76,7 @@ Most users interact with them through `StreamingOptions` fields rather than dire .. automodule:: cudf_polars.utils.config :members: DynamicPlanningOptions, + JoinFilterPushdownOptions, MemoryResourceConfig, ParquetOptions, StreamingExecutor, diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index 5d813e74bbe1..6744ad959a93 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/docs/cudf/source/cudf_polars/options.md @@ -108,6 +108,7 @@ Environment variables follow these patterns: | `broadcast_limit` | Maximum number of bytes for broadcast joins. | auto | | `target_partition_size` | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto. | auto | | `dynamic_planning` | Dynamic planning configuration, dict or {class}`~cudf_polars.utils.config.DynamicPlanningOptions`. `None` disables. | enabled | +| `join_filter_pushdown` | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `None` disables. | enabled | | `sink_to_directory` | Whether `.sink_*()` writes its output as a directory. The `spmd`, `ray`, and `dask` engines always use `True`; passing `False` raises `ValueError`. | `True` | ### Category: `engine` diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py new file mode 100644 index 000000000000..382d0027f649 --- /dev/null +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for tracking column value domains between IR nodes.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import singledispatch +from typing import TYPE_CHECKING + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, + Slice, + Sort, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + + from cudf_polars.dsl.ir import IR + +__all__ = [ + "ColumnBinding", + "ColumnLineage", + "ColumnRef", + "column_domain_bindings", +] + + +@dataclass(frozen=True) +class ColumnBinding: + """A direct binding to a named column on a specific child edge.""" + + child_index: int + name: str + + +@dataclass(frozen=True) +class ColumnRef: + """A named column produced by an IR node.""" + + node: IR + name: str + + +@dataclass(frozen=True) +class ColumnLineage: + """Persistent value-domain lineage, sharing suffixes across DAG branches.""" + + column: ColumnRef + source: ColumnLineage | None = None + source_child_index: int | None = None + """Child edge leading to ``source``, or None if there is no source.""" + + +@singledispatch +def column_domain_bindings(node: IR) -> Mapping[str, ColumnBinding]: + """ + Map output columns to child columns containing their value domains. + + For every ``output_name -> ColumnBinding(child_index, input_name)`` binding, + every value appearing in ``node[output_name]`` is guaranteed to appear in + ``node.children[child_index][input_name]``. Row order, multiplicity, and + cardinality are not preserved. + + If a name in ``node.schema`` does not appear in the mapping it means + that it was not possible to derive a relationship between the domain of + the output and input values for that column. + """ + return {} + + +@column_domain_bindings.register(Select) +def _(node: Select) -> Mapping[str, ColumnBinding]: + return { + item.name: ColumnBinding(0, item.value.name) + for item in node.exprs + if isinstance(item.value, expr.Col) + } + + +@column_domain_bindings.register(HStack) +def _(node: HStack) -> Mapping[str, ColumnBinding]: + child = node.children[0] + replaced = {item.name for item in node.columns} + return { + name: ColumnBinding(0, name) for name in child.schema if name not in replaced + } | { + item.name: ColumnBinding(0, item.value.name) + for item in node.columns + if isinstance(item.value, expr.Col) + } + + +@column_domain_bindings.register(GroupBy) +def _(node: GroupBy) -> Mapping[str, ColumnBinding]: + return { + key.name: ColumnBinding(0, key.value.name) + for key in node.keys + if isinstance(key.value, expr.Col) + } + + +@column_domain_bindings.register(Join) +def _(node: Join) -> Mapping[str, ColumnBinding]: + left, right = node.children + how = node.options[0] + if how in ("Semi", "Anti"): + return { + name: ColumnBinding(0, name) for name in node.schema if name in left.schema + } + if how != "Inner": + return {} + + bindings = {name: ColumnBinding(0, name) for name in left.schema} + suffix = node.options[3] + for name in right.schema: + output_name = f"{name}{suffix}" if name in left.schema else name + if output_name in node.schema: + bindings[output_name] = ColumnBinding(1, name) + return bindings + + +@column_domain_bindings.register(Distinct) +@column_domain_bindings.register(Filter) +@column_domain_bindings.register(Projection) +@column_domain_bindings.register(Slice) +@column_domain_bindings.register(Sort) +def _( + node: Distinct | Filter | Projection | Slice | Sort, +) -> Mapping[str, ColumnBinding]: + child = node.children[0] + return { + name: ColumnBinding(0, name) for name in node.schema if name in child.schema + } diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 80359eab925d..b382f12099a4 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -734,11 +734,17 @@ def evaluate_on_rank( """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) + lowering, node_map = lower_ir_graph_with_node_map( + ir, config_options, stats, rank=comm.rank, nranks=comm.nranks + ) + optimized = lowering.optimized + ir = lowering.lowered + partition_info = lowering.partition_info if config_options.executor.quent_context is not None: assert local_quent_context is not None - logical_plan_id = ir.get_stable_plan_id() + logical_plan_id = optimized.get_stable_plan_id() plan, ops, ports, logical_op_by_id = build_plan( - ir, + optimized, config_options, query=local_quent_context.context.query, plan_id=logical_plan_id, @@ -752,10 +758,6 @@ def evaluate_on_rank( local_quent_context.logger, plan, ops, ports ) - ir, partition_info, node_map = lower_ir_graph_with_node_map( - ir, config_options, stats, rank=comm.rank, nranks=comm.nranks - ) - if comm.rank == 0: log_query_plan(ir, config_options) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index c959f24f98d9..295a081307fa 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -26,6 +26,7 @@ from cudf_polars.quent import QuentContext from cudf_polars.utils.config import ( DynamicPlanningOptions, + JoinFilterPushdownOptions, ParquetOptions, ) @@ -248,6 +249,14 @@ class StreamingOptions: Env: ``CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING``. Default: enabled. Category: executor. + join_filter_pushdown + Config for join filter pushdown optimizations, dict or + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions`. ``None`` + disables the rewrite. + Env: ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN`` and + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__*``. + Default: enabled. + Category: executor. sink_to_directory Whether multi-partition sink operations should write to a directory rather than a single file. The ``spmd``/``ray``/``dask`` engines @@ -346,6 +355,9 @@ class StreamingOptions: dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = ( _opt("executor") ) + join_filter_pushdown: ( + dict[str, Any] | JoinFilterPushdownOptions | None | Unspecified + ) = _opt("executor") sink_to_directory: bool | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__SINK_TO_DIRECTORY", parse_boolean ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index d02077206739..256e45440d8d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -38,8 +38,9 @@ import cudf_polars.dsl.tracing from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction -from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select +from cudf_polars.dsl.ir import Filter, GroupBy, HStack, Join, Projection, Select from cudf_polars.dsl.tracing import Scope +from cudf_polars.dsl.utils.column_domain import column_domain_bindings from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk @@ -446,9 +447,8 @@ def _derived_ordering( def _select_column_targets(select: Select) -> dict[str, dict[str, None]]: old_to_new_names: defaultdict[str, dict[str, None]] = defaultdict(dict) - for ne in select.exprs: - if isinstance(ne.value, Col): - old_to_new_names[ne.value.name][ne.name] = None + for output_name, source in column_domain_bindings(select).items(): + old_to_new_names[source.name][output_name] = None return dict(old_to_new_names) @@ -603,7 +603,7 @@ def maybe_remap_partitioning( ), local=_remap_scheme_simple(ir, partitioning.local, ir.children[0]), ) - if isinstance(ir, (Cache, Join, Projection, Filter)): + if isinstance(ir, (Join, Projection, Filter)): child = child_ir if child_ir is not None else ir.children[0] return Partitioning( inter_rank=_remap_scheme_simple(ir, partitioning.inter_rank, child), diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 8e62e161cd7c..b585138d583a 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -123,8 +123,10 @@ def explain_query( if physical: with cm: stats = collect_statistics(ir, config, executor) - lowered_ir, partition_info = lower_ir_graph(ir, config, stats) - return _repr_ir_tree(lowered_ir, partition_info, stats=stats, config=config) + lowered = lower_ir_graph(ir, config, stats) + return _repr_ir_tree( + lowered.lowered, lowered.partition_info, stats=stats, config=config + ) else: if config.executor.name == "streaming": # Include row-count statistics for the logical plan @@ -150,7 +152,9 @@ def collect_partition_plan( with concurrent.futures.ThreadPoolExecutor() as executor: stats = collect_statistics(ir, config, executor) - lowered_ir, partition_info = lower_ir_graph(ir, config, stats) + lowered = lower_ir_graph(ir, config, stats) + lowered_ir = lowered.lowered + partition_info = lowered.partition_info seen: set[tuple] = set() rows: list[PartitionPlanRow] = [] @@ -755,7 +759,9 @@ def from_ir( if lowered: with cm: stats = collect_statistics(ir, config_options, executor) - ir, partition_info_d = lower_ir_graph(ir, config_options, stats) + lowering = lower_ir_graph(ir, config_options, stats) + ir = lowering.lowered + partition_info_d = lowering.partition_info partition_info_dict = {} nodes: dict[str, SerializableIRNode] = {} diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 70b1585ba348..63d76d6c328f 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -75,7 +75,7 @@ def _make_hash_join( partition_info, output_count, ) - # Always reconstruct in case children contain Cache nodes + # Reconstruct with the lowered and possibly shuffled children. ir = ir.reconstruct([left, right]) # Record new partitioning info diff --git a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py new file mode 100644 index 000000000000..5ba74515cf1b --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -0,0 +1,844 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Rewrite a plan, inserting prefilters in join DAGs. + +For a supported inner equijoin, this optimization tries to use the join-key +values produced by one input to reduce the size of the other input before +the original join. In relational notation, a simple rewrite is:: + + left join[left.key = right.key] right + + -> + + (left semijoin[left.key = right.key] project(right.key)) + join[left.key = right.key] right + +In this example, the right hand table is selected to pre-filter the left +table before performing the inner join. + +The implementation uses the following terms: + +``column lineage`` + A chain from a named output column towards columns in its input subplan. + Each step guarantees that every value in the output column also appears in + the referenced child column, although row order and multiplicity are not + preserved and the child may contain additional values. +``child edge`` + One particular parent-to-child position in the IR DAG. The same child node + may occur on more than one edge, so a lineage records child indices and a + rewrite follows the resulting edge path to change only the chosen + occurrence. +``target`` + The side of the join to filter. +``domain`` + The side of the join used to provide key values for the filtering of + ``target``. +``producer`` + A node on a column lineage, together with the column name at that node and + its edge path from the join input. So termed because it "produces" the + key values participating in the join. +``constraint domain`` + Selective values of another join key from the target input, used to reduce + the domain before deriving the values that will filter the target. +``simple candidate`` + A rewrite that projects one domain join key and uses it to filter the + corresponding target key directly. +``composite candidate`` + For a multi-key join, a rewrite that first semi-joins the domain using the + constraint domain, then projects the reduced domain's key used to filter + the target. + +Plan rewrite has three stages. ``analyze_plan`` gathers row estimates, +selective nodes, and column value-domain lineages. Candidate selection +consumes those facts and returns a decision. ``apply_candidate`` then +constructs the selected semi-join rewrite. + +Row estimates, selectivity propagation, thresholds, and candidate scores are +only heuristics for deciding whether a safe rewrite is likely to improve +execution. Poor estimates can choose an unprofitable rewrite. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import singledispatch +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + IR, + ConditionalJoin, + DataFrameScan, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Rolling, + Scan, + Select, + Slice, + Sort, + Union, +) +from cudf_polars.dsl.tracing import Scope, log +from cudf_polars.dsl.traversal import ( + CachingVisitor, + post_traversal, + reuse_if_unchanged, + traversal, +) +from cudf_polars.dsl.utils.column_domain import ( + ColumnLineage, + ColumnRef, + column_domain_bindings, +) + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator, Mapping, Sequence + + from cudf_polars.streaming.base import StatsCollector + from cudf_polars.typing import GenericTransformer + from cudf_polars.utils.config import ConfigOptions, StreamingExecutor + + +@dataclass(frozen=True) +class _Producer: + """A subtree and its bound column names at an insertion point.""" + + node: IR + columns: tuple[str, ...] + rows: int + path: tuple[int, ...] = () + """Child-edge path from the candidate root to ``node``.""" + + @property + def column(self) -> str: + """First bound column in the producer.""" + return self.columns[0] + + +@dataclass(frozen=True) +class SimpleCandidate: + """A direct key-domain prefilter candidate.""" + + mode = "simple" + target_side: Literal["left", "right"] + target: _Producer + target_key: expr.Col + domain: _Producer + domain_key: expr.Col + + @property + def score(self) -> tuple[int, int, int]: + """Rank after composite candidates, then by domain size.""" + return (1, self.domain.rows, self.domain.rows) + + +@dataclass(frozen=True) +class CompositeCandidate: + """A key-domain prefilter constrained by another join key.""" + + mode = "composite" + target_side: Literal["left", "right"] + target: _Producer + target_key: expr.Col + domain: _Producer + domain_key: expr.Col + constraint_domain: _Producer + domain_constraint_key: expr.Col + target_constraint_key: expr.Col + + @property + def score(self) -> tuple[int, int, int]: + """Prefer smaller constraint and domain inputs.""" + return (0, self.constraint_domain.rows, self.domain.rows) + + +Candidate: TypeAlias = SimpleCandidate | CompositeCandidate +DecisionReason: TypeAlias = Literal[ + "applied", + "maintain_order", + "no_selective_domain", + "non_column_join_key", + "not_inner_join", + "sliced_join", +] + + +@dataclass(frozen=True) +class Decision: + """Result of considering a join for a domain prefilter.""" + + reason: DecisionReason + candidate: Candidate | None = None + + +@dataclass(frozen=True) +class PlanFacts: + """Facts derived in one bottom-up traversal of an IR DAG.""" + + row_estimates: Mapping[IR, int | None] + selective_nodes: frozenset[IR] + column_lineages: Mapping[ColumnRef, ColumnLineage] + + +class _RewriteState(TypedDict): + """State shared by the join-domain prefilter DAG rewrite.""" + + threshold: float + trace: bool + stats: StatsCollector + facts: PlanFacts + + +def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: + """ + Derive row, selectivity, and column-domain facts for an IR DAG. + + Parameters + ---------- + ir + Root node to gather facts for. + stats + Pre-populated statistics + + Returns + ------- + Gather facts about the plan. + """ + row_estimates: dict[IR, int | None] = {} + selective_nodes: set[IR] = set() + column_lineages: dict[ColumnRef, ColumnLineage] = {} + + for node in post_traversal([ir]): + if isinstance(node, (Scan, DataFrameScan)): + source_info = stats.scan_stats.get(node) + rows = None if source_info is None else source_info.row_count + if rows is None and isinstance(node, DataFrameScan): + rows = node.df.shape()[0] + elif isinstance(node, (Select, Projection, HStack, Filter, Distinct, GroupBy)): + rows = row_estimates[node.children[0]] + elif isinstance(node, Join): + rows = _estimate_join_rows( + node.options[0], + row_estimates[node.children[0]], + row_estimates[node.children[1]], + ) + else: + child_estimates = [ + estimate + for child in node.children + if (estimate := row_estimates[child]) is not None + ] + rows = max(child_estimates, default=None) + row_estimates[node] = rows + + if ( + (isinstance(node, Scan) and node.predicate is not None) + or isinstance(node, Filter) + or any(child in selective_nodes for child in node.children) + ): + selective_nodes.add(node) + + bindings = column_domain_bindings(node) + for name in node.schema: + column = ColumnRef(node, name) + binding = bindings.get(name) + if binding is None: + source_lineage = None + source_child_index = None + else: + source_child_index = binding.child_index + source = ColumnRef( + node.children[source_child_index], + binding.name, + ) + source_lineage = column_lineages[source] + column_lineages[column] = ColumnLineage( + column, source_lineage, source_child_index + ) + + return PlanFacts( + row_estimates=row_estimates, + selective_nodes=frozenset(selective_nodes), + column_lineages=column_lineages, + ) + + +def blocks_pushdown(node: IR) -> bool: + """ + Return whether a node blocks filter pushdown. + + Parameters + ---------- + node + Node to check + + Returns + ------- + bool + True if a semijoin cannot be pushed past this node, otherwise False. + """ + return ( + # TODO: Distinct and Rolling only block pushdown in some + # circumstances, but we'd need to make the logic more complicated: + # - We can push through distinct if the filter applies to the columns + # that are being used to determine distinct rows + # - We can push through rolling if the filter applies to the + # groupby keys. + # TODO: We can push through an unsliced Union, but need to + # distribute the filter onto every child. + isinstance(node, (Distinct, Rolling, Slice, Union)) + # Can't push through anything that is sliced. + or (isinstance(node, (GroupBy, Sort)) and node.zlice is not None) + or (isinstance(node, (ConditionalJoin, Join)) and node.options[2] is not None) + ) + + +def semijoin_pushdown_candidates( + facts: PlanFacts, root: IR, column: str +) -> Iterator[tuple[ColumnRef, tuple[int, ...]]]: + """ + Yield column domain lineage providing valid locations for semijoin pushdown. + + Parameters + ---------- + facts + Gathered facts about the plan + root + Root node to search from + column + Name of column we're finding the lineage of. + + Returns + ------- + Iterator + Of valid insertion points and their child-edge paths from ``root``. + """ + try: + lineage = facts.column_lineages[ColumnRef(root, column)] + except KeyError: + return + path: tuple[int, ...] = () + while True: + yield lineage.column, path + source = lineage.source + source_child_index = lineage.source_child_index + if blocks_pushdown(lineage.column.node) or source is None: + return + assert source_child_index is not None + path = (*path, source_child_index) + lineage = source + + +def optimize_join_filter_pushdown( + ir: IR, + stats: StatsCollector, + config_options: ConfigOptions[StreamingExecutor], +) -> IR: + """ + Rewrite an IR DAG to apply filter pushdown of keys. + + This optimization pass inspects joins in the DAG and attempts to push a + prefilter obtained from the keys of one side of the join onto the + inputs of the other side. This can be highly beneficial at large scale + since if we have a selective join we can avoid data movement by + prefiltering before performing the actual join. + + Parameters + ---------- + ir + DAG to rewrite. + stats + Pre-populated statistics. + config_options + Configuration options controlling the rewrite. + + Returns + ------- + Rewritten DAG. + """ + options = config_options.executor.join_filter_pushdown + if options is None: + return ir + threshold = options.threshold + trace = options.trace + if threshold == 0: + return ir + + state = _RewriteState( + threshold=threshold, + trace=trace, + stats=stats, + facts=analyze_plan(ir, stats), + ) + mapper: GenericTransformer[IR, IR, _RewriteState] = CachingVisitor( + _rewrite, state=state + ) + return mapper(ir) + + +@singledispatch +def _rewrite(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + raise AssertionError + + +@_rewrite.register(IR) +def _(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + return reuse_if_unchanged(node, rec) + + +@_rewrite.register(Join) +def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + original = node + rewritten = reuse_if_unchanged(node, rec) + assert isinstance(rewritten, Join) + node = rewritten + if node is original: + facts = rec.state["facts"] + else: + # Child rewrites introduce new semi joins and reconstructed ancestors. + # Re-analyze that current subtree so parent joins can use the derived + # selectivity and cardinality when ranking their own candidates. + facts = analyze_plan(node, rec.state["stats"]) + decision = _select_candidate( + node, + rec.state["threshold"], + facts, + ) + if rec.state["trace"]: + _trace_decision(node, rec.state["threshold"], decision) + if decision.candidate is None: + return node + return apply_candidate(node, decision.candidate) + + +def apply_candidate(ir: Join, candidate: Candidate) -> IR: + """Apply a selected join-domain prefilter candidate to a join.""" + left, right = ir.children + domain = _make_domain(candidate, ir) + target = candidate.target + target_filter = _make_semi_join( + target.node, + expr.Col(target.node.schema[target.column], target.column), + domain, + expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), + nulls_equal=ir.options[1], + suffix=ir.options[3], + ) + if candidate.target_side == "left": + left = replace_at_path(left, target.path, target_filter) + else: + right = replace_at_path(right, target.path, target_filter) + return ir.reconstruct((left, right)) + + +def replace_at_path(root: IR, path: Sequence[int], replacement: IR) -> IR: + """ + Replace a specific child in a DAG starting at root. + + Parameters + ---------- + root + Root of DAG to carry out replacement. + path + Breadcrumb trail selecting which child at every level to recurse + into. + replacement + Replacement node to return when the path becomes empty. + + Returns + ------- + IR + New DAG with the selected child replaced with replacement. + + Notes + ----- + This specifically does not use replacement by equality so that we can + disambiguate between shared children in the DAG where we only want to + replace one. + """ + if not path: + return replacement + index, *path = path + children = list(root.children) + children[index] = replace_at_path(children[index], path, replacement) + return root.reconstruct(children) + + +def _select_candidate( + ir: Join, + threshold: float, + facts: PlanFacts, +) -> Decision: + if ir.options[0] != "Inner": + return Decision(reason="not_inner_join") + if ir.options[2] is not None: + return Decision(reason="sliced_join") + if ir.options[5] != "none": + return Decision(reason="maintain_order") + + left_keys = _simple_keys(ir.left_on) + right_keys = _simple_keys(ir.right_on) + if len(left_keys) != len(ir.left_on) or len(right_keys) != len(ir.right_on): + return Decision(reason="non_column_join_key") + + candidates: list[Candidate] = [] + left: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( + "left", + ir.children[0], + left_keys, + ) + right: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( + "right", + ir.children[1], + right_keys, + ) + for (target_side, target_child, target_keys), ( + _, + domain_child, + domain_keys, + ) in ((left, right), (right, left)): + candidates.extend( + _composite_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + threshold, + facts, + ) + ) + candidates.extend( + _simple_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + threshold, + facts, + ) + ) + + if not candidates: + return Decision(reason="no_selective_domain") + return Decision(reason="applied", candidate=min(candidates, key=lambda c: c.score)) + + +def _simple_keys(keys: Sequence[expr.NamedExpr]) -> tuple[expr.Col, ...]: + return tuple(key.value for key in keys if isinstance(key.value, expr.Col)) + + +def _simple_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], + threshold: float, + facts: PlanFacts, +) -> Iterable[SimpleCandidate]: + for target_key, domain_key in zip(target_keys, domain_keys, strict=True): + target = _largest_key_source(target_child, target_key.name, facts) + if target is None: + continue + domain = _smallest_key_producer( + domain_child, + domain_key.name, + facts, + require_selective=True, + ) + if domain is None: + continue + if domain.rows / target.rows > threshold: + continue + if contains_node(target.node, domain.node): + continue + yield SimpleCandidate( + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=domain_key, + ) + + +def _composite_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], + threshold: float, + facts: PlanFacts, +) -> Iterable[CompositeCandidate]: + if len(target_keys) < 2: + return + + for filter_index, (target_key, domain_key) in enumerate( + zip(target_keys, domain_keys, strict=True) + ): + target = _largest_key_source(target_child, target_key.name, facts) + if target is None: + continue + + for constraint_index, ( + target_constraint_key, + domain_constraint_key, + ) in enumerate(zip(target_keys, domain_keys, strict=True)): + if constraint_index == filter_index: + continue + domain = _smallest_node_containing_all( + domain_child, + (domain_key.name, domain_constraint_key.name), + facts, + ) + if domain is None: + continue + if domain.rows / target.rows > threshold: + continue + constraint_domain = _smallest_key_producer( + target_child, + target_constraint_key.name, + facts, + require_selective=True, + exclude=target.node, + ) + if constraint_domain is None: + continue + if constraint_domain.rows / domain.rows > threshold: + continue + if contains_node(target.node, domain.node) or contains_node( + target.node, constraint_domain.node + ): + continue + yield CompositeCandidate( + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=domain_key, + constraint_domain=constraint_domain, + domain_constraint_key=domain_constraint_key, + target_constraint_key=target_constraint_key, + ) + + +def _make_domain(candidate: Candidate, ir: Join) -> IR: + if isinstance(candidate, SimpleCandidate): + return _project_bound_key( + candidate.domain.node, + candidate.domain.column, + candidate.domain_key, + ) + + constraint_domain = _project_bound_key( + candidate.constraint_domain.node, + candidate.constraint_domain.column, + candidate.target_constraint_key, + ) + constrained = _make_semi_join( + candidate.domain.node, + expr.Col( + candidate.domain.node.schema[candidate.domain.columns[1]], + candidate.domain.columns[1], + ), + constraint_domain, + expr.Col( + constraint_domain.schema[candidate.target_constraint_key.name], + 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 + ) + + +def _project_bound_key(source: IR, bound_column: str, output_key: expr.Col) -> Select: + """Project a bound source column under its join-visible key name.""" + dtype = source.schema[bound_column] + assert dtype == output_key.dtype + return Select( + {output_key.name: dtype}, + (expr.NamedExpr(output_key.name, expr.Col(dtype, bound_column)),), + True, # noqa: FBT003 + source, + ) + + +def _make_semi_join( + target: IR, + target_key: expr.Col, + domain: IR, + domain_key: expr.Col, + *, + nulls_equal: bool, + suffix: str, +) -> Join: + return Join( + target.schema, + (expr.NamedExpr(target_key.name, target_key),), + (expr.NamedExpr(domain_key.name, domain_key),), + ("Semi", nulls_equal, None, suffix, False, "none"), + target, + domain, + ) + + +def _smallest_key_producer( + root: IR, + column: str, + facts: PlanFacts, + *, + require_selective: bool, + exclude: IR | None = None, +) -> _Producer | None: + candidates = [] + for reference, path in semijoin_pushdown_candidates(facts, root, column): + node, bound_column = reference.node, reference.name + if node is exclude: + continue + rows = facts.row_estimates.get(node) + if rows is None or rows <= 0: + continue + if require_selective and node not in facts.selective_nodes: + continue + candidates.append( + (rows, len(node.schema), _Producer(node, (bound_column,), rows, path)) + ) + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +def _smallest_node_containing_all( + root: IR, columns: Sequence[str], facts: PlanFacts +) -> _Producer | None: + candidates = [] + lineages: list[ColumnLineage] = [] + for column in columns: + lineage = facts.column_lineages.get(ColumnRef(root, column)) + if lineage is None: + return None + lineages.append(lineage) + if not lineages: + return None + path: tuple[int, ...] = () + while True: + node = lineages[0].column.node + if any(lineage.column.node != node for lineage in lineages[1:]): + break + bound_columns = tuple(lineage.column.name for lineage in lineages) + rows = facts.row_estimates.get(node) + if rows is not None and rows > 0: + candidates.append( + ( + rows, + len(node.schema), + _Producer(node, bound_columns, rows, path), + ) + ) + if blocks_pushdown(node): + break + source_child_index = lineages[0].source_child_index + if source_child_index is None or any( + lineage.source_child_index != source_child_index for lineage in lineages[1:] + ): + break + sources = [lineage.source for lineage in lineages if lineage.source is not None] + if len(sources) != len(lineages): + # Some sources are None + break + path = (*path, source_child_index) + lineages = sources + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +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): + node, bound_column = reference.node, reference.name + rows = facts.row_estimates.get(node) + if rows is None or rows <= 0: + continue + item = ( + rows, + len(node.schema), + _Producer(node, (bound_column,), rows, path), + ) + if isinstance(node, (Scan, DataFrameScan)): + source_candidates.append(item) + else: + fallback_candidates.append(item) + candidates = source_candidates or fallback_candidates + if not candidates: + return None + return max(candidates, key=lambda item: (item[0], -item[1]))[2] + + +def _estimate_join_rows( + how: str, left_rows: int | None, right_rows: int | None +) -> int | None: + if left_rows is None: + return right_rows + if right_rows is None: + return left_rows + if how in ("Inner", "Semi", "Anti"): + return min(left_rows, right_rows) + if how == "Left": + return left_rows + if how == "Right": + return right_rows + if how == "Full": + return max(left_rows, right_rows) + return None + + +def contains_node(root: IR, needle: IR) -> bool: + """Return whether an equal node occurs in a DAG rooted at ``root``.""" + return needle in traversal([root]) + + +def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: + join_filter_pushdown: dict[str, Any] = { + "considered": True, + "threshold": threshold, + "reason": decision.reason, + } + record = { + "scope": Scope.PLAN.value, + "join_filter_pushdown": join_filter_pushdown, + "actor_ir_id": ir.get_stable_id(), + "actor_ir_type": type(ir).__name__, + } + if (candidate := decision.candidate) is not None: + join_filter_pushdown.update( + { + "mode": candidate.mode, + "target_side": candidate.target_side, + "target_key": candidate.target_key.name, + "domain_key": candidate.domain_key.name, + "estimated_target_rows": candidate.target.rows, + "estimated_domain_rows": candidate.domain.rows, + "target_node_type": type(candidate.target.node).__name__, + "domain_node_type": type(candidate.domain.node).__name__, + } + ) + if isinstance(candidate, CompositeCandidate): + join_filter_pushdown.update( + { + "constraint_key": candidate.target_constraint_key.name, + "estimated_constraint_rows": candidate.constraint_domain.rows, + } + ) + log("Join Filter Pushdown", **record) diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 295798116d26..2ac5c8c2eef7 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -4,8 +4,9 @@ from __future__ import annotations +import dataclasses import operator -from functools import partial, reduce +from functools import reduce from typing import TYPE_CHECKING import polars as pl @@ -35,7 +36,7 @@ Slice, Union, ) -from cudf_polars.dsl.traversal import CachingVisitor, traversal +from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged, traversal from cudf_polars.dsl.utils.naming import unique_names from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node @@ -52,6 +53,7 @@ from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.dispatch import LowerIRTransformer, State + from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -65,6 +67,64 @@ def _( ) +@lower_ir_node.register(Cache) +def _( + ir: Cache, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: # pragma: no cover + raise AssertionError("Cache nodes should have been removed before lowering") + + +@dataclasses.dataclass +class LoweringInfo: + """Information produced by optimizing and lowering an IR graph.""" + + optimized: IR # IR after optimization + lowered: IR # optimized IR after lowering + partition_info: MutableMapping[ + IR, PartitionInfo + ] # Partition mapping for nodes in the lowered IR. + + +def remove_cache_nodes(ir: IR) -> IR: + """Remove logical cache nodes while preserving shared DAG structure.""" + + def rewrite(node: IR, rec: GenericTransformer[IR, IR, None]) -> IR: + if isinstance(node, Cache): + return rec(node.children[0]) + return reuse_if_unchanged(node, rec) + + mapper: GenericTransformer[IR, IR, None] = CachingVisitor(rewrite, state=None) + return mapper(ir) + + +def optimize_with_stats( + ir: IR, config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector +) -> IR: + """ + Optimize an IR graph given some statistics. + + Parameters + ---------- + ir + Root of the graph to optimize. + config_options + GPUEngine configuration options. + stats + Pre-computed statistics. + + Returns + ------- + IR + The optimized IR graph. + """ + from cudf_polars.streaming.join_filter_pushdown import ( + optimize_join_filter_pushdown, + ) + + ir = remove_cache_nodes(ir) + return optimize_join_filter_pushdown(ir, stats, config_options) + + def _lower_ir_graph_impl( ir: IR, config_options: ConfigOptions[StreamingExecutor], @@ -72,15 +132,19 @@ def _lower_ir_graph_impl( *, rank: int = 0, nranks: int = 1, -) -> tuple[tuple[IR, MutableMapping[IR, PartitionInfo]], LowerIRTransformer]: +) -> tuple[LoweringInfo, LowerIRTransformer]: state: State = { "config_options": config_options, "stats": stats, "rank": rank, "nranks": nranks, } + optimized = optimize_with_stats(ir, config_options, stats) mapper: LowerIRTransformer = CachingVisitor(lower_ir_node, state=state) - return mapper(ir), mapper + lowered, partition_info = mapper(optimized) + return LoweringInfo( + optimized=optimized, lowered=lowered, partition_info=partition_info + ), mapper def lower_ir_graph( @@ -90,7 +154,7 @@ def lower_ir_graph( *, rank: int = 0, nranks: int = 1, -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: +) -> LoweringInfo: """ Rewrite an IR graph and extract partitioning information. @@ -109,9 +173,7 @@ def lower_ir_graph( Returns ------- - new_ir, partition_info - The rewritten graph and a mapping from unique nodes - in the new graph to associated partitioning information. + LoweringInfo Notes ----- @@ -132,7 +194,7 @@ def lower_ir_graph_with_node_map( *, rank: int = 0, nranks: int = 1, -) -> tuple[IR, MutableMapping[IR, PartitionInfo], dict[str, list[str]]]: +) -> tuple[LoweringInfo, dict[str, list[str]]]: """ Lower an IR graph and return a mapping from physical to logical stable IDs. @@ -155,10 +217,8 @@ def lower_ir_graph_with_node_map( Returns ------- - new_ir - The rewritten IR graph. - partition_info - Mapping from unique nodes in the new graph to partitioning info. + LoweringInfo + Information about the lowered IR graph. node_map Mapping ``{physical_stable_id: [logical_stable_id, ...]}`` built from the internal :class:`CachingVisitor` cache. Nodes inserted @@ -173,7 +233,7 @@ def lower_ir_graph_with_node_map( old_key = str(old_node.get_stable_id()) node_map.setdefault(new_key, []).append(old_key) - return *result, node_map + return result, node_map def evaluate_streaming( @@ -282,8 +342,6 @@ def _lower_ir_pwise( return new_node, partition_info -_lower_ir_pwise_preserve = partial(_lower_ir_pwise, preserve_partitioning=True) -lower_ir_node.register(Cache, _lower_ir_pwise_preserve) lower_ir_node.register(HConcat, _lower_ir_pwise) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 3451010b464a..8a265d875b39 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -54,6 +54,7 @@ "DaskContext", "DynamicPlanningOptions", "InMemoryExecutor", + "JoinFilterPushdownOptions", "ParquetOptions", "RayContext", "SPMDContext", @@ -362,8 +363,9 @@ class DynamicPlanningOptions: The maximum number of chunks to sample before making dynamic-planning decisions. Default is 2. join_prefilter_threshold - Row-count ratio (small / large) below which a join key prefilter is - applied. Set to 0 to disable join prefiltering. Default is 0.5. + Row-count ratio (small / large) below which one side of a join is + filtered by a bloom filter built from the other side before + performing the join. Set to 0 to disable. Default is 0.5. join_prefilter_max_key_columns Maximum number of columns from the join-key prefix to use for the prefilter. Set to ``None`` to use the full join-key list. Default is 1. @@ -428,6 +430,59 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("join_prefilter_trace must be a bool") +@dataclasses.dataclass(frozen=True) +class JoinFilterPushdownOptions: + """ + Configuration options for join filter pushdown in the logical plan. + + When performing a join between two tables, it is often favourable + to pre-filter one side of the join with the keys (full or partial) of + the other side. This can reduce the size of tables that actually + participate in the join. + + cudf-polars supports a form of this where we can rewrite inner joins by + selecting a side to be filtered by the keys of the other side. + + Pass ``None`` to ``StreamingExecutor(join_filter_pushdown=...)`` to + disable the rewrite. + + These options can be configured via environment variables with the prefix + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__``. + + Parameters + ---------- + 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. + trace + Whether to emit plan-time trace decisions for filter decisions. Default is False. + """ + + _env_prefix = "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN" + + threshold: float = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__THRESHOLD", float, default=0.5 + ) + ) + trace: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__TRACE", _bool_converter, default=False + ) + ) + + def __post_init__(self) -> None: # noqa: D105 + threshold = self.threshold + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)): + raise TypeError("threshold must be a float or int") + threshold = float(threshold) + object.__setattr__(self, "threshold", threshold) + if not 0.0 <= threshold <= 1.0: + raise ValueError("threshold must be between 0 and 1") + if not isinstance(self.trace, bool): + raise TypeError("trace must be a bool") + + @dataclasses.dataclass(frozen=True, eq=True) class MemoryResourceConfig: """ @@ -693,6 +748,10 @@ class StreamingExecutor: dynamic_planning Options controlling dynamic shuffle planning. See :class:`~cudf_polars.utils.config.DynamicPlanningOptions` for more. + join_filter_pushdown + Options controlling the logical join-domain prefilter rewrite. See + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for more. + ``None`` disables the rewrite. max_io_threads Maximum number of IO threads. Default is 4. This controls the parallelism of IO operations when reading data. @@ -762,6 +821,9 @@ class StreamingExecutor: dynamic_planning: DynamicPlanningOptions | None = dataclasses.field( default_factory=DynamicPlanningOptions ) + join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field( + default_factory=JoinFilterPushdownOptions + ) max_io_threads: int = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__MAX_IO_THREADS", int, default=4 @@ -823,6 +885,20 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) + if isinstance(self.join_filter_pushdown, dict): + object.__setattr__( + self, + "join_filter_pushdown", + JoinFilterPushdownOptions(**self.join_filter_pushdown), + ) + if self.join_filter_pushdown is not None and not isinstance( + self.join_filter_pushdown, JoinFilterPushdownOptions + ): + raise TypeError( + "join_filter_pushdown must be a JoinFilterPushdownOptions " + "instance, dict, or None" + ) + if self.cluster in ("spmd", "ray", "dask"): if self.sink_to_directory is False: raise ValueError( @@ -855,6 +931,7 @@ def __hash__(self) -> int: # noqa: D105 # to json and hash that. d = dataclasses.asdict(self) d["dynamic_planning"] = json.dumps(d["dynamic_planning"]) + d["join_filter_pushdown"] = json.dumps(d["join_filter_pushdown"]) # Hash the quent context UUIDs as ints quent_context = d["quent_context"] @@ -1019,6 +1096,17 @@ def from_polars_engine( if not _bool_converter(env_dynamic_planning): user_executor_options["dynamic_planning"] = None + # Handle join_filter_pushdown: check user config, then env var + user_join_filter_pushdown = user_executor_options.get( + "join_filter_pushdown", None + ) + if user_join_filter_pushdown is None: + env_join_filter_pushdown = os.environ.get( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1" + ) + if not _bool_converter(env_join_filter_pushdown): + user_executor_options["join_filter_pushdown"] = None + executor = StreamingExecutor(**user_executor_options) case _: # pragma: no cover; Unreachable raise ValueError(f"Unsupported executor: {user_executor}") diff --git a/python/cudf_polars/tests/dsl/test_column_domain.py b/python/cudf_polars/tests/dsl/test_column_domain.py new file mode 100644 index 000000000000..999f3f2463d2 --- /dev/null +++ b/python/cudf_polars/tests/dsl/test_column_domain.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import polars as pl + +import pylibcudf as plc + +from cudf_polars.containers import DataType +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + DataFrameScan, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, + Slice, + Sort, +) +from cudf_polars.dsl.utils.column_domain import ( + ColumnBinding, + column_domain_bindings, +) + +I64 = DataType(pl.Int64()) +BOOL = DataType(pl.Boolean()) + + +def make_scan(*names: str) -> DataFrameScan: + frame = pl.DataFrame({name: [1] for name in names}) + return DataFrameScan(dict.fromkeys(names, I64), frame._df, None) + + +def col(name: str) -> expr.Col: + return expr.Col(I64, name) + + +def named_col(output: str, source: str) -> expr.NamedExpr: + return expr.NamedExpr(output, col(source)) + + +def test_source_has_no_column_domain_bindings() -> None: + assert column_domain_bindings(make_scan("a")) == {} + + +def test_select_binds_aliases_and_omits_derived_columns() -> None: + child = make_scan("a", "b") + node = Select( + {"renamed": I64, "derived": I64}, + ( + named_col("renamed", "a"), + expr.NamedExpr("derived", expr.Literal(I64, 1)), + ), + True, # noqa: FBT003 + child, + ) + + assert column_domain_bindings(node) == { + "renamed": ColumnBinding(0, "a"), + } + + +def test_hstack_binds_passthrough_alias_and_override() -> None: + child = make_scan("a", "b") + node = HStack( + {"a": I64, "b": I64, "alias": I64}, + ( + expr.NamedExpr("a", expr.Literal(I64, 1)), + named_col("alias", "b"), + ), + True, # noqa: FBT003 + child, + ) + + assert column_domain_bindings(node) == { + "b": ColumnBinding(0, "b"), + "alias": ColumnBinding(0, "b"), + } + + +def test_groupby_binds_only_direct_keys() -> None: + child = make_scan("a", "b") + node = GroupBy( + {"key": I64, "value": I64}, + (named_col("key", "a"),), + (named_col("value", "b"),), + False, # noqa: FBT003 + None, + child, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "a"), + } + + +def test_inner_join_binds_left_right_and_suffixed_columns() -> None: + left = make_scan("key", "left_value") + right = make_scan("key", "right_value") + node = Join( + { + "key": I64, + "left_value": I64, + "key_right": I64, + "right_value": I64, + }, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Inner", False, (0, 1), "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "key"), + "left_value": ColumnBinding(0, "left_value"), + "key_right": ColumnBinding(1, "key"), + "right_value": ColumnBinding(1, "right_value"), + } + + +def test_inner_join_omits_coalesced_right_key() -> None: + left = make_scan("key", "left_value") + right = make_scan("key", "right_value") + node = Join( + {"key": I64, "left_value": I64, "right_value": I64}, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Inner", False, None, "_right", True, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "key"), + "left_value": ColumnBinding(0, "left_value"), + "right_value": ColumnBinding(1, "right_value"), + } + + +def test_semi_join_binds_only_left_columns() -> None: + left = make_scan("key", "value") + right = make_scan("key") + node = Join( + left.schema, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Semi", False, None, "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "key"), + "value": ColumnBinding(0, "value"), + } + + +def test_outer_join_has_no_column_domain_bindings() -> None: + left = make_scan("key") + right = make_scan("other") + node = Join( + {**left.schema, **right.schema}, + (named_col("key", "key"),), + (named_col("other", "other"),), + ("Left", False, None, "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == {} + + +def test_passthrough_nodes_bind_same_named_columns() -> None: + child = make_scan("a", "b") + mask = expr.NamedExpr("mask", expr.Literal(BOOL, True)) # noqa: FBT003 + nodes = ( + Filter(child.schema, mask, child), + Projection({"b": I64}, child), + Slice(child.schema, 0, 1, child), + Distinct( + child.schema, + plc.stream_compaction.DuplicateKeepOption.KEEP_ANY, + None, + (0, 1), + False, # noqa: FBT003 + child, + ), + Sort( + child.schema, + (named_col("a", "a"),), + (plc.types.Order.ASCENDING,), + (plc.types.NullOrder.AFTER,), + False, # noqa: FBT003 + (0, 1), + child, + ), + ) + + for node in nodes: + assert column_domain_bindings(node) == { + name: ColumnBinding(0, name) for name in node.schema + } diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index be1aedc770b8..2b7217d0689c 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -449,10 +449,9 @@ def test_lower_ir_graph_with_node_map() -> None: ir, config_options, concurrent.futures.ThreadPoolExecutor() ) - _lowered_ir, _partition_info, node_map = lower_ir_graph_with_node_map( - ir, config_options, stats - ) + lowering, node_map = lower_ir_graph_with_node_map(ir, config_options, stats) + assert lowering.optimized is ir assert len(node_map) > 0 for physical_sid, logical_sids in node_map.items(): assert isinstance(physical_sid, str) diff --git a/python/cudf_polars/tests/streaming/test_dataframescan.py b/python/cudf_polars/tests/streaming/test_dataframescan.py index 014851f689e7..0fd27f2b57e8 100644 --- a/python/cudf_polars/tests/streaming/test_dataframescan.py +++ b/python/cudf_polars/tests/streaming/test_dataframescan.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -59,7 +59,7 @@ def test_parallel_dataframescan( ) qir = Translator(df._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, info = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -68,6 +68,8 @@ def test_parallel_dataframescan( parquet_stats_executor, ), ) + ir = lowering.lowered + info = lowering.partition_info count = info[ir].count if max_rows_per_partition < total_row_count: assert count > 1 @@ -106,7 +108,7 @@ def test_join_in_memory_lazy_stable_id_pickle( right = pl.LazyFrame({"k": [2, 3, 4], "y": [1, 2, 3]}).collect(engine=engine).lazy() qir = Translator(left.join(right, on="k")._ldf.visit(), engine).translate_ir() config_options = ConfigOptions.from_polars_engine(engine) - ir, _ = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -115,6 +117,7 @@ def test_join_in_memory_lazy_stable_id_pickle( parquet_stats_executor, ), ) + ir = lowering.lowered _assert_stable_ids_match(ir, pickle.loads(pickle.dumps(ir))) @@ -128,7 +131,7 @@ def test_dataframescan_pickle( ) qir = Translator(df._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, _ = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -137,6 +140,7 @@ def test_dataframescan_pickle( parquet_stats_executor, ), ) + ir = lowering.lowered # Pickle and unpickle the IR (which contains DataFrameScan) pickled = pickle.dumps(ir) diff --git a/python/cudf_polars/tests/streaming/test_hstack.py b/python/cudf_polars/tests/streaming/test_hstack.py index f7a3c6213340..232ad33e255e 100644 --- a/python/cudf_polars/tests/streaming/test_hstack.py +++ b/python/cudf_polars/tests/streaming/test_hstack.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Tests for CSE HStack handling in the streaming executor.""" @@ -150,14 +150,16 @@ def test_cse_agg_shared_decomposition( assert len(inner_hstacks) == (1 if comm_subexpr_elim else 0) config_options = ConfigOptions.from_polars_engine(engine) - lowered, _ = lower_ir_graph( + lowering = lower_ir_graph( ir, config_options, collect_statistics(ir, config_options, parquet_stats_executor), ) # Both paths must lower to a single Repartition computing one aggregation. - repartitions = [n for n in traversal([lowered]) if isinstance(n, Repartition)] + repartitions = [ + n for n in traversal([lowering.lowered]) if isinstance(n, Repartition) + ] assert len(repartitions) == 1 assert len(repartitions[0].children[0].exprs) == 1 # type: ignore[attr-defined] assert_gpu_result_equal(q, engine=engine, collect_kwargs={"optimizations": opts}) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index cdc544c43169..2ac2edc88903 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -491,18 +491,17 @@ def test_broadcast_limit( q = left.join(right, on="y", how="inner") ir = Translator(q._ldf.visit(), engine).translate_ir() config_options = ConfigOptions.from_polars_engine(engine) - shuffle_nodes = [ - type(node) - for node in lower_ir_graph( + lowering = lower_ir_graph( + ir, + config_options, + collect_statistics( ir, config_options, - collect_statistics( - ir, - config_options, - parquet_stats_executor, - ), - )[1] - if isinstance(node, Shuffle) + parquet_stats_executor, + ), + ) + shuffle_nodes = [ + type(node) for node in lowering.partition_info if isinstance(node, Shuffle) ] # NOTE: Expect small table to have 3 partitions (9 / 3). @@ -516,7 +515,7 @@ def test_broadcast_limit( assert len(shuffle_nodes) == 0 -def test_cache_preserves_partitioning_join( +def test_shared_join_preserves_partitioning( parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ): engine = pl.GPUEngine( @@ -542,20 +541,24 @@ def test_cache_preserves_partitioning_join( config_options = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() - lowered_ir, partition_info = lower_ir_graph( + lowering = lower_ir_graph( ir, config_options, collect_statistics(ir, config_options, parquet_stats_executor), ) + lowered_ir = lowering.lowered + partition_info = lowering.partition_info + + assert not any(isinstance(node, Cache) for node in traversal([lowered_ir])) - # Cache should preserve partitioning on 'key' - cache_partitioning = [ + # Removing Cache should preserve the shared join's partitioning on 'key'. + join_partitioning = [ [ne.name for ne in partition_info[node].partitioned_on] for node in traversal([lowered_ir]) - if isinstance(node, Cache) + if isinstance(node, Join) ] - assert cache_partitioning == [["key"]], ( - f"Cache should preserve partitioning on 'key', got {cache_partitioning}" + assert join_partitioning == [["key"]], ( + f"Shared join should be partitioned on 'key', got {join_partitioning}" ) # Only 2 shuffles needed (for join sides, not for groupby) diff --git a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py new file mode 100644 index 000000000000..3dfe42df7fd3 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -0,0 +1,834 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +import polars as pl + +from cudf_polars import Translator +from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice +from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.utils.column_domain import ColumnRef +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.join_filter_pushdown import ( + CompositeCandidate, + Decision, + PlanFacts, + SimpleCandidate, + _select_candidate, + _smallest_node_containing_all, + analyze_plan, + apply_candidate, + contains_node, + optimize_join_filter_pushdown, + semijoin_pushdown_candidates, +) +from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes +from cudf_polars.streaming.statistics import collect_statistics +from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.utils.config import ConfigOptions + +if TYPE_CHECKING: + import concurrent.futures + from typing import Any + + from cudf_polars.dsl.ir import IR + from cudf_polars.engine.spmd import SPMDEngine + + +@pytest.fixture +def engine(spmd_engine_factory) -> SPMDEngine: + """Return an SPMD engine configured for join-domain prefilter tests.""" + return spmd_engine_factory( + StreamingOptions( + join_filter_pushdown={"threshold": 0.5}, + raise_on_fail=True, + ) + ) + + +def make_config( + *, dynamic_planning: bool = True, join_filter_pushdown: bool = True +) -> ConfigOptions: + executor_options: dict[str, Any] = { + "join_filter_pushdown": {"trace": False} if join_filter_pushdown else None + } + if not dynamic_planning: + executor_options["dynamic_planning"] = None + return ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options=executor_options, + ) + ) + + +def find_joins(ir: IR, how: str | None = None) -> list[Join]: + return [ + node + for node in traversal([ir]) + if isinstance(node, Join) and (how is None or node.options[0] == how) + ] + + +def translate_query(query: pl.LazyFrame, engine: SPMDEngine) -> IR: + """Translate a public Polars query and remove logical Cache nodes.""" + t = Translator(query._ldf.visit(), engine) + root = t.translate_ir() + assert not t.errors + return remove_cache_nodes(root) + + +def dataframe_scan(ir: IR, column: str) -> DataFrameScan: + """Return the unique in-memory scan containing ``column``.""" + (match,) = ( + node + for node in traversal([ir]) + if isinstance(node, DataFrameScan) and column in node.schema + ) + return match + + +@pytest.fixture +def simple_query() -> pl.LazyFrame: + """Return a query with a small selective join domain.""" + part = ( + pl.LazyFrame( + { + "p_partkey": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 10 for i in range(20)], + "l_suppkey": range(20), + } + ) + return part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + + +def test_simple_prefilter_filters_large_side( + simple_query: pl.LazyFrame, engine: SPMDEngine +) -> None: + root = translate_query(simple_query, engine) + + assert isinstance(root, Join) + part_ir, _ = root.children + lineitem_ir = dataframe_scan(root, "l_partkey") + facts = analyze_plan(root, StatsCollector()) + decision = _select_candidate(root, 0.5, facts) + + assert decision.reason == "applied" + assert isinstance(decision.candidate, SimpleCandidate) + optimized = apply_candidate(root, decision.candidate) + + assert isinstance(optimized, Join) + assert optimized.options[0] == "Inner" + 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_gpu_result_equal(simple_query, engine=engine, check_row_order=False) + + +def test_filter_pushdown_is_independent_of_dynamic_planning( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + root = translate_query(simple_query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + make_config(dynamic_planning=False), + ) + + assert find_joins(optimized, "Semi") + + +def test_filter_pushdown_can_be_disabled( + simple_query: pl.LazyFrame, engine: SPMDEngine +) -> None: + root = translate_query(simple_query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + make_config(join_filter_pushdown=False), + ) + + assert optimized is root + + +@pytest.mark.parametrize( + "nulls_equal", [False, True], ids=["nulls_not_equal", "nulls_equal"] +) +def test_nullable_join_keys_preserve_results( + nulls_equal: bool, # noqa: FBT001 + engine: SPMDEngine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + domain = pl.LazyFrame( + { + "key": [None, 1, 2, 9], + "active": [True, True, True, False], + } + ).filter("active") + target = pl.LazyFrame( + { + "key": [None, 1, 2, 3] * 10, + "value": range(40), + } + ) + query = domain.join(target, on="key", nulls_equal=nulls_equal) + + ir = Translator(query._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_join_filter_pushdown( + ir, + collect_statistics(ir, config, parquet_stats_executor), + config, + ) + + semi_joins = find_joins(optimized, "Semi") + assert semi_joins + assert all(join.options[1] is nulls_equal for join in semi_joins) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_prefilter_does_not_move_below_distinct_on_non_subset_column( + engine: SPMDEngine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + target = pl.LazyFrame( + { + "group": [1, 1] * 100, + "key": [0, 1] * 100, + } + ).unique(subset="group", keep="first", maintain_order=True) + domain = pl.LazyFrame( + { + "key": [1, 2], + "active": [True, False], + } + ).filter("active") + query = target.join(domain, on="key") + + ir = Translator(query._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_join_filter_pushdown( + ir, + collect_statistics(ir, config, parquet_stats_executor), + config, + ) + + semis = find_joins(optimized, "Semi") + assert any(isinstance(semi.children[0], Distinct) for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_no_simple_filter_pushdown_when_domain_is_not_selective( + engine: SPMDEngine, +) -> None: + supplier = pl.LazyFrame({"s_suppkey": range(3)}) + lineitem = pl.LazyFrame({"l_suppkey": [i % 3 for i in range(20)]}) + query = supplier.join( + lineitem, + left_on="s_suppkey", + right_on="l_suppkey", + ) + root = translate_query(query, engine) + stats = StatsCollector() + assert isinstance(root, Join) + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) + + optimized = optimize_join_filter_pushdown( + root, + stats, + ConfigOptions.from_polars_engine(engine), + ) + + assert decision == Decision(reason="no_selective_domain") + assert optimized is root + assert not find_joins(optimized, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_composite_filter_pushdown_constrains_domain_first( + engine: SPMDEngine, +) -> None: + 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 % 30 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"), + ) + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + config = ConfigOptions.from_polars_engine(engine) + stats = StatsCollector() + + assert isinstance(root, Join) + order_lineitem, supplier_ir = root.children + assert isinstance(order_lineitem, Join) + lineitem_ir = order_lineitem.children[1] + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) + optimized = optimize_join_filter_pushdown(root, stats, config) + + assert decision.reason == "applied" + assert isinstance(decision.candidate, CompositeCandidate) + semis = find_joins(optimized, "Semi") + 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_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_derived_selectivity_propagates_through_rewritten_children( + engine: SPMDEngine, +) -> None: + region = ( + pl.LazyFrame( + { + "r_regionkey": [0, 1], + "active": [True, False], + } + ) + .filter("active") + .select("r_regionkey") + ) + nation = pl.LazyFrame( + { + "n_nationkey": range(10), + "n_regionkey": [i % 2 for i in range(10)], + } + ) + customer = pl.LazyFrame( + { + "c_custkey": range(40), + "c_nationkey": [i % 10 for i in range(40)], + } + ) + orders = pl.LazyFrame( + { + "o_orderkey": range(200), + "o_custkey": [i % 40 for i in range(200)], + } + ) + query = ( + region.join(nation, left_on="r_regionkey", right_on="n_regionkey") + .join(customer, left_on="n_nationkey", right_on="c_nationkey") + .join(orders, left_on="c_custkey", right_on="o_custkey") + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + semis = find_joins(optimized, "Semi") + expected_targets = { + dataframe_scan(root, "n_nationkey"), + dataframe_scan(root, "c_custkey"), + dataframe_scan(root, "o_orderkey"), + } + assert expected_targets <= {semi.children[0] for semi in semis} + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_rewritten_domain_filters_other_side_instead_of_stacking( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(6), + "part_active": [True] * 3 + [False] * 3, + } + ) + .filter("part_active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_orderkey": [i % 15 for i in range(180)], + "l_partkey": [i % 6 for i in range(180)], + "l_suppkey": [i % 3 for i in range(180)], + } + ) + supplier = pl.LazyFrame({"s_suppkey": range(3)}) + orders = ( + pl.LazyFrame( + { + "o_orderkey": range(15), + "order_active": [True] * 8 + [False] * 7, + } + ) + .filter("order_active") + .select("o_orderkey") + ) + query = ( + part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + .join(supplier, left_on="l_suppkey", right_on="s_suppkey") + .join(orders, left_on="l_orderkey", right_on="o_orderkey") + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + 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 any(semi.children[0] is orders_ir for semi in semis) + assert not any( + isinstance(semi.children[0], Join) and semi.children[0].options[0] == "Semi" + for semi in semis + ) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_target_source_follows_join_key_through_rename( + engine: SPMDEngine, +) -> None: + big = pl.LazyFrame( + { + "left_key": range(20), + "other": [i % 5 for i in range(20)], + } + ) + renamed_big = big.select(pl.col("left_key").alias("foo"), "other") + small = pl.LazyFrame( + { + "left_key": range(10), + "other2": [i % 5 for i in range(10)], + } + ) + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + query = renamed_big.join( + small, + left_on="other", + right_on="other2", + maintain_order="left", + ).join( + domain, + left_on="left_key", + right_on="domain_key", + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + + assert isinstance(root, Join) + joined = root.children[0] + assert isinstance(joined, Join) + renamed_big_ir, small_ir = joined.children + assert isinstance(renamed_big_ir, Select) + big_ir = renamed_big_ir.children[0] + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + 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) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_domain_source_follows_join_key_through_rename( + engine: SPMDEngine, +) -> None: + target = pl.LazyFrame({"target_key": range(20)}) + unrelated = pl.LazyFrame( + { + "domain_key": [100, 101], + "other": [0, 1], + "active": [True, False], + } + ) + renamed_unrelated = unrelated.filter("active").select( + pl.col("domain_key").alias("foo"), "other" + ) + domain_source = pl.LazyFrame( + { + "domain_key": range(1, 6), + "other2": range(5), + } + ) + domain = renamed_unrelated.join( + domain_source, + left_on="other", + right_on="other2", + maintain_order="left", + ) + query = target.join( + domain, + left_on="target_key", + right_on="domain_key", + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + + assert isinstance(root, Join) + target_ir, domain_ir = root.children + assert isinstance(domain_ir, Join) + renamed_unrelated_ir, domain_source_ir = domain_ir.children + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + 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] + 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 rewritten_domain_source.children[0] is domain_source_ir + assert rewritten_domain_source.children[0] is not renamed_unrelated_ir + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_composite_domain_columns_follow_renames(engine: SPMDEngine) -> None: + query = pl.LazyFrame( + { + "raw_key": range(10), + "raw_constraint": range(10), + } + ).select( + pl.col("raw_key").alias("domain_key"), + pl.col("raw_constraint").alias("domain_constraint"), + ) + renamed = translate_query(query, engine) + source = dataframe_scan(renamed, "raw_key") + + analyzed = analyze_plan(renamed, StatsCollector()) + facts = PlanFacts( + row_estimates={renamed: 20, source: 10}, + selective_nodes=analyzed.selective_nodes, + column_lineages=analyzed.column_lineages, + ) + producer = _smallest_node_containing_all( + renamed, ("domain_key", "domain_constraint"), facts + ) + + assert producer is not None + assert producer.node is source + assert producer.columns == ("raw_key", "raw_constraint") + + +def test_composite_domain_columns_do_not_reconverge_after_join( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame({"key": [1, 1, 2], "value": [10, 20, 30]}) + query = source.join(source, on="key", suffix="_right") + joined = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(joined, Join) + assert isinstance(joined.children[0], Cache) + assert joined.children[0] is joined.children[1] + joined = remove_cache_nodes(joined) + assert isinstance(joined, Join) + assert joined.children[0] is joined.children[1] + + facts = analyze_plan(joined, StatsCollector()) + producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) + + candidates = tuple(semijoin_pushdown_candidates(facts, joined, "value")) + assert candidates[0] == (ColumnRef(joined, "value"), ()) + assert len(candidates) >= 2 + assert all(path == (0,) * len(path) for _, path in candidates[1:]) + assert producer is not None + assert producer.node is joined + assert producer.columns == ("value", "value_right") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_contains_node_uses_dag_equality(engine: SPMDEngine) -> None: + query = pl.LazyFrame({"key": range(3)}).filter(pl.col("key") >= 0).slice(0, 2) + root = translate_query(query, engine) + + assert isinstance(root, Slice) + source = root.children[0] + equal_source = source.reconstruct(source.children) + + assert source is not equal_source + assert source == equal_source + assert contains_node(root, equal_source) + + +def test_plan_facts_share_lineage_suffixes_across_shared_dag( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame({"raw_key": range(10)}) + query = source.select(pl.col("raw_key").alias("left_key")).join( + source.select(pl.col("raw_key").alias("right_key")), + left_on="left_key", + right_on="right_key", + ) + root = translate_query(query, engine) + + assert isinstance(root, Join) + left, right = root.children + source_ir = dataframe_scan(root, "raw_key") + facts = analyze_plan(root, StatsCollector()) + left_lineage = facts.column_lineages[ColumnRef(left, "left_key")] + right_lineage = facts.column_lineages[ColumnRef(right, "right_key")] + source_lineage = facts.column_lineages[ColumnRef(source_ir, "raw_key")] + + assert left_lineage.column == ColumnRef(left, "left_key") + assert right_lineage.column == ColumnRef(right, "right_key") + while left_lineage.source is not source_lineage: + assert left_lineage.source is not None + left_lineage = left_lineage.source + while right_lineage.source is not source_lineage: + assert right_lineage.source is not None + right_lineage = right_lineage.source + assert left_lineage.source is right_lineage.source + assert source_lineage.source is None + + +def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: + target = ( + pl.LazyFrame({"target_key": range(20)}) + .filter(pl.col("target_key") >= 0) + .slice(0, 10) + ) + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + query = target.join(domain, left_on="target_key", right_on="domain_key") + root = translate_query(query, engine) + + assert isinstance(root, Join) + sliced = root.children[0] + assert isinstance(sliced, Slice) + target_ir = dataframe_scan(root, "target_key") + stats = StatsCollector() + facts = analyze_plan(root, stats) + lineage = facts.column_lineages[ColumnRef(sliced, "target_key")] + assert lineage.column == ColumnRef(sliced, "target_key") + assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( + (ColumnRef(sliced, "target_key"), ()), + ) + + optimized = optimize_join_filter_pushdown( + root, + stats, + 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) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_target_replacement_does_not_rewrite_shared_domain_side( + engine: SPMDEngine, +) -> None: + shared = pl.LazyFrame( + { + "target_key": range(20), + "other": [i % 2 for i in range(20)], + } + ) + domain_source = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "other2": [0, 1], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key", "other2") + ) + domain = shared.join( + domain_source, + left_on=pl.col("other").cast(pl.Int32), + right_on=pl.col("other2").cast(pl.Int32), + ) + query = shared.join(domain, left_on="target_key", right_on="domain_key") + root = translate_query(query, engine) + + assert isinstance(root, Join) + shared_ir, domain_ir = root.children + assert isinstance(domain_ir, Join) + assert domain_ir.children[0] is shared_ir + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + assert isinstance(optimized, Join) + 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 dataframe_scan(root, "target_key") + assert not find_joins(unfiltered_domain, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_target_prefilter_rewrites_only_selected_self_join_edge( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame( + { + "key": [1, 1, 2, 2], + "value": [10, 20, 30, 40], + } + ) + domain = ( + pl.LazyFrame( + { + "domain_value": [10, 999], + "active": [True, False], + } + ) + .filter("active") + .select("domain_value") + ) + query = source.join(source, on="key", suffix="_right").join( + domain, + left_on="value", + right_on="domain_value", + ) + translated = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(translated, Join) + translated_self_join = translated.children[0] + assert isinstance(translated_self_join, Join) + shared_cache = translated_self_join.children[0] + assert isinstance(shared_cache, Cache) + assert translated_self_join.children[1] is shared_cache + source_ir = shared_cache.children[0] + + optimized = optimize_with_stats( + translated, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + rewritten_self_join = optimized.children[0] + assert isinstance(rewritten_self_join, Join) + filtered, unfiltered = rewritten_self_join.children + assert unfiltered is source_ir + filtered_semis = find_joins(filtered, "Semi") + assert len(filtered_semis) == 1 + assert not find_joins(unfiltered, "Semi") + assert any(filtered_semis[0].children[0] is node for node in traversal([source_ir])) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +@pytest.mark.parametrize("how", ["left", "right", "cross", "full"]) +def test_no_filter_pushdown_for_unsupported_joins( + how: Any, + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame({"l_partkey": [i % 10 for i in range(20)]}) + if how == "cross": + query = part.join(lineitem, how=how) + else: + query = part.join( + lineitem, + left_on="p_partkey", + right_on="l_partkey", + how=how, + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + assert optimized is root + assert not find_joins(optimized, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index c5a42062c9a7..d1a5a54aafa2 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -83,6 +83,11 @@ def test_executor_options_sink_to_directory_absent_when_unspecified() -> None: assert "sink_to_directory" not in StreamingOptions().to_executor_options() +def test_executor_options_join_filter_pushdown_disabled() -> None: + result = StreamingOptions(join_filter_pushdown=None).to_executor_options() + assert result["join_filter_pushdown"] is None + + # --------------------------------------------------------------------------- # to_engine_options # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/streaming/test_parallel.py b/python/cudf_polars/tests/streaming/test_parallel.py index 44d1b4ce0757..e32fab2ce3a3 100644 --- a/python/cudf_polars/tests/streaming/test_parallel.py +++ b/python/cudf_polars/tests/streaming/test_parallel.py @@ -12,9 +12,13 @@ from polars.testing import assert_frame_equal from cudf_polars import Translator +from cudf_polars.dsl.ir import Cache, Join from cudf_polars.dsl.traversal import traversal from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.parallel import optimize_with_stats from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.utils.config import ConfigOptions @pytest.mark.parametrize("column", ["a", "b"]) @@ -88,6 +92,30 @@ def test_evaluate_streaming(streaming_engine): assert_frame_equal(expected, got_streaming) +def test_optimize_removes_cache_nodes() -> None: + source = pl.LazyFrame({"key": [1, 1, 2], "value": [10, 20, 30]}) + query = source.join(source, on="key", suffix="_right") + engine = GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": None}, + ) + ir = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(ir, Join) + assert isinstance(ir.children[0], Cache) + assert ir.children[0] is ir.children[1] + + optimized = optimize_with_stats( + ir, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + assert optimized.children[0] is optimized.children[1] + assert not any(isinstance(node, Cache) for node in traversal([optimized])) + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_parallel.py (round 3) # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 15ba081e69fb..9cb44ce6c0a7 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -204,7 +204,7 @@ def test_target_partition_size( ) qir = Translator(q._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, info = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -213,6 +213,8 @@ def test_target_partition_size( parquet_stats_executor, ), ) + ir = lowering.lowered + info = lowering.partition_info count = info[ir].count if blocksize <= 12_000: assert count > n_files diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index a0718e9be5bf..1d1395b48f0f 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -31,6 +31,7 @@ ConfigOptions, DynamicPlanningOptions, InMemoryExecutor, + JoinFilterPushdownOptions, MemoryResourceConfig, StreamingExecutor, ) @@ -614,6 +615,9 @@ def test_dynamic_planning_defaults() -> None: assert config.executor.dynamic_planning.join_prefilter_threshold == 0.5 assert config.executor.dynamic_planning.join_prefilter_max_key_columns == 1 assert not config.executor.dynamic_planning.join_prefilter_trace + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.5 + assert not config.executor.join_filter_pushdown.trace def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -653,6 +657,31 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non assert config.executor.dynamic_planning.join_prefilter_threshold == 0.25 assert config.executor.dynamic_planning.join_prefilter_max_key_columns is None assert config.executor.dynamic_planning.join_prefilter_trace + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.5 + assert not config.executor.join_filter_pushdown.trace + + +def test_join_filter_pushdown_options_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" + ) + 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.trace + + +def test_join_filter_pushdown_disabled_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "0") + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.join_filter_pushdown is None @pytest.mark.parametrize("value, expected", [("none", None), ("null", None), ("2", 2)]) @@ -747,6 +776,65 @@ def test_validate_join_prefilter_trace() -> None: ) +def test_validate_join_filter_pushdown_options() -> None: + with pytest.raises(TypeError, match="threshold must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": {"threshold": "bad"}}, + ) + ) + with pytest.raises(ValueError, match="threshold must be between"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": {"threshold": 1.5}}, + ) + ) + with pytest.raises(TypeError, match="trace must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": {"trace": "bad"}}, + ) + ) + + +def test_validate_join_filter_pushdown_type() -> None: + with pytest.raises( + TypeError, + match="join_filter_pushdown must be a JoinFilterPushdownOptions instance", + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": object()}, + ) + ) + + +def test_join_filter_pushdown_from_instance() -> None: + options = JoinFilterPushdownOptions(threshold=0.25, trace=True) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": options}, + ) + ) + assert config.executor.join_filter_pushdown is options + + +def test_join_filter_pushdown_disabled_from_options() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": None}, + ) + ) + assert config.executor.join_filter_pushdown is None + assert hash(config) == hash(config) + + def test_dynamic_planning_from_instance() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine(