diff --git a/ci/run_cudf_polars_polars_tests.sh b/ci/run_cudf_polars_polars_tests.sh index 3073b61812d7..959f4d033921 100755 --- a/ci/run_cudf_polars_polars_tests.sh +++ b/ci/run_cudf_polars_polars_tests.sh @@ -62,7 +62,7 @@ DESELECTED_TESTS_STR=$(printf -- " --deselect %s" "${DESELECTED_TESTS[@]}") # Don't quote the `DESELECTED_...` variable because `pytest` can't handle # multiple quoted arguments inline # shellcheck disable=SC2086 -# Fail fast (-x) because failed tests pollute the state +# Fail fast (-x) rather than trying to continue because failed tests pollute the state echo "Run polars tests with injected in-memory GPU engine" python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ python -m pytest \ @@ -90,7 +90,6 @@ python "${TIMEOUT_TOOL_PATH}" --enable-python 4800 \ --import-mode=importlib \ --cache-clear \ -x \ - -v \ -m "" \ -p cudf_polars.testing.inject_gpu_engine \ -W ignore::ResourceWarning \ diff --git a/ci/test_python_other.sh b/ci/test_python_other.sh index 0efebe1afafe..ecb83d66df22 100755 --- a/ci/test_python_other.sh +++ b/ci/test_python_other.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -42,8 +42,9 @@ timeout 30m ./ci/run_custreamz_pytests.sh \ --cov-report=term rapids-logger "pytest cudf-polars" +# Fail fast (-x) rather than trying to continue because failed tests pollute the state ./ci/run_cudf_polars_pytests.sh \ - -vv \ + -x \ --junitxml="${RAPIDS_TESTS_DIR}/junit-cudf-polars.xml" \ --numprocesses=4 \ --dist=worksteal \ diff --git a/ci/test_wheel_cudf_polars.sh b/ci/test_wheel_cudf_polars.sh index d9df9c05d4eb..1869ec8800a9 100755 --- a/ci/test_wheel_cudf_polars.sh +++ b/ci/test_wheel_cudf_polars.sh @@ -91,7 +91,6 @@ for version in "${VERSIONS[@]}"; do # Fail fast (-x) rather than trying to continue because failed tests pollute the state ./ci/run_cudf_polars_pytests.sh \ - -vv \ "${COVERAGE_ARGS[@]}" \ --numprocesses=4 \ --dist=worksteal \ diff --git a/python/cudf_polars/cudf_polars/dsl/traversal.py b/python/cudf_polars/cudf_polars/dsl/traversal.py index 095a9719597b..654b6fba17f4 100644 --- a/python/cudf_polars/cudf_polars/dsl/traversal.py +++ b/python/cudf_polars/cudf_polars/dsl/traversal.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Traversal and visitor utilities for nodes.""" from __future__ import annotations -from collections import deque +from collections import Counter, deque from typing import TYPE_CHECKING, Generic from cudf_polars.typing import ( @@ -15,19 +15,40 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Generator, MutableMapping, Sequence + from collections.abc import Callable, Generator, Mapping, MutableMapping, Sequence from cudf_polars.typing import GenericTransformer, NodeT __all__: list[str] = [ "CachingVisitor", + "collect_refcount", "make_recursive", + "post_traversal", "reuse_if_unchanged", "traversal", ] +def collect_refcount(nodes: Sequence[NodeT]) -> Mapping[NodeT, int]: + """ + Determine reference counts of all nodes in a DAG. + + Parameters + ---------- + nodes + Sequence of root nodes + + Returns + ------- + Mapping from nodes to frequency of occurrence in the DAG. + """ + refcount = Counter(nodes) + for node in traversal(nodes): + refcount.update(node.children) + return refcount + + def traversal(nodes: Sequence[NodeT]) -> Generator[NodeT, None, None]: """ Pre-order traversal of nodes in an expression. 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 5ba74515cf1b..b099e1584ea6 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -38,6 +38,10 @@ 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. +``source cost`` + An estimate of the cost required to materialize a producer. This guards + against treating a small intermediate result as a cheap domain when + producing it requires scanning large inputs. ``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. @@ -49,10 +53,10 @@ 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. +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. Row estimates, selectivity propagation, thresholds, and candidate scores are only heuristics for deciding whether a safe rewrite is likely to improve @@ -86,6 +90,7 @@ from cudf_polars.dsl.tracing import Scope, log from cudf_polars.dsl.traversal import ( CachingVisitor, + collect_refcount, post_traversal, reuse_if_unchanged, traversal, @@ -104,6 +109,17 @@ from cudf_polars.utils.config import ConfigOptions, StreamingExecutor +DomainScore: TypeAlias = tuple[int, int, int] + + +@dataclass(frozen=True) +class SourceFacts: + """Source-derived facts for an IR node.""" + + cost: int | None + is_single_source: bool + + @dataclass(frozen=True) class _Producer: """A subtree and its bound column names at an insertion point.""" @@ -111,6 +127,8 @@ class _Producer: node: IR columns: tuple[str, ...] rows: int + cost: int + is_single_source: bool path: tuple[int, ...] = () """Child-edge path from the candidate root to ``node``.""" @@ -119,6 +137,11 @@ def column(self) -> str: """First bound column in the producer.""" return self.columns[0] + @property + def domain_score(self) -> DomainScore: + """Scoring function for a domain.""" + return (self.cost, self.rows, len(self.node.schema)) + @dataclass(frozen=True) class SimpleCandidate: @@ -132,9 +155,9 @@ class SimpleCandidate: 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) + def score(self) -> tuple[int, DomainScore]: + """Rank after composite candidates, then by domain cost.""" + return (1, self.domain.domain_score) @dataclass(frozen=True) @@ -152,16 +175,16 @@ class CompositeCandidate: 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) + def score(self) -> tuple[int, DomainScore, DomainScore]: + """Prefer cheaper constraint and domain inputs.""" + return (0, self.constraint_domain.domain_score, self.domain.domain_score) Candidate: TypeAlias = SimpleCandidate | CompositeCandidate DecisionReason: TypeAlias = Literal[ "applied", "maintain_order", - "no_selective_domain", + "no_profitable_domain", "non_column_join_key", "not_inner_join", "sliced_join", @@ -181,8 +204,10 @@ class PlanFacts: """Facts derived in one bottom-up traversal of an IR DAG.""" row_estimates: Mapping[IR, int | None] + source_facts: Mapping[IR, SourceFacts] selective_nodes: frozenset[IR] column_lineages: Mapping[ColumnRef, ColumnLineage] + refcounts: Mapping[IR, int] class _RewriteState(TypedDict): @@ -210,8 +235,11 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: Gather facts about the plan. """ row_estimates: dict[IR, int | None] = {} + source_facts: dict[IR, SourceFacts] = {} + source_nodes: dict[IR, frozenset[IR]] = {} selective_nodes: set[IR] = set() column_lineages: dict[ColumnRef, ColumnLineage] = {} + refcounts = collect_refcount([ir]) for node in post_traversal([ir]): if isinstance(node, (Scan, DataFrameScan)): @@ -236,6 +264,23 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: rows = max(child_estimates, default=None) row_estimates[node] = rows + if isinstance(node, (Scan, DataFrameScan)): + sources: frozenset[IR] = frozenset((node,)) + else: + sources = frozenset( + source for child in node.children for source in source_nodes[child] + ) + source_nodes[node] = sources + source_rows = [ + source_rows + for source in sources + if (source_rows := row_estimates[source]) is not None and source_rows > 0 + ] + source_facts[node] = SourceFacts( + cost=sum(source_rows) if source_rows else rows, + is_single_source=len(sources) == 1, + ) + if ( (isinstance(node, Scan) and node.predicate is not None) or isinstance(node, Filter) @@ -263,26 +308,33 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: return PlanFacts( row_estimates=row_estimates, + source_facts=source_facts, selective_nodes=frozenset(selective_nodes), column_lineages=column_lineages, + refcounts=refcounts, ) -def blocks_pushdown(node: IR) -> bool: +def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: """ Return whether a node blocks filter pushdown. Parameters ---------- node - Node to check + Node to check. + facts + Facts about the plan. Returns ------- bool True if a semijoin cannot be pushed past this node, otherwise False. """ - return ( + # 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 + # shared, since the same filter will not come from every consumer. + return facts.refcounts[node] > 1 or ( # 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 @@ -327,7 +379,7 @@ def semijoin_pushdown_candidates( yield lineage.column, path source = lineage.source source_child_index = lineage.source_child_index - if blocks_pushdown(lineage.column.node) or source is None: + if blocks_pushdown(lineage.column.node, facts) or source is None: return assert source_child_index is not None path = (*path, source_child_index) @@ -526,7 +578,7 @@ def _select_candidate( ) if not candidates: - return Decision(reason="no_selective_domain") + return Decision(reason="no_profitable_domain") return Decision(reason="applied", candidate=min(candidates, key=lambda c: c.score)) @@ -559,6 +611,12 @@ def _simple_candidates( continue if contains_node(target.node, domain.node): continue + if domain.is_single_source and has_filtering_semi_ancestor( + target_child, target.path + ): + continue + if not domain_cost_is_small(domain, target, threshold): + continue yield SimpleCandidate( target_side=target_side, target=target, @@ -617,6 +675,10 @@ def _composite_candidates( target.node, constraint_domain.node ): continue + if not domain_cost_is_small(domain, target, threshold): + continue + if not domain_cost_is_small(constraint_domain, domain, threshold): + continue yield CompositeCandidate( target_side=target_side, target=target, @@ -692,6 +754,27 @@ def _make_semi_join( ) +def make_producer( + node: IR, + columns: tuple[str, ...], + path: tuple[int, ...], + facts: PlanFacts, +) -> _Producer | None: + """Construct a producer from gathered plan facts, if fully estimated.""" + rows = facts.row_estimates.get(node) + source = facts.source_facts[node] + if rows is None or rows <= 0 or source.cost is None: + return None + return _Producer( + node=node, + columns=columns, + rows=rows, + cost=source.cost, + is_single_source=source.is_single_source, + path=path, + ) + + def _smallest_key_producer( root: IR, column: str, @@ -700,28 +783,25 @@ def _smallest_key_producer( require_selective: bool, exclude: IR | None = None, ) -> _Producer | None: - candidates = [] + producers = [] 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: + producer = make_producer(node, (bound_column,), path, facts) + if producer is not None: + producers.append(producer) + if not producers: return None - return min(candidates, key=lambda item: (item[0], item[1]))[2] + return min(producers, key=lambda p: p.domain_score) def _smallest_node_containing_all( root: IR, columns: Sequence[str], facts: PlanFacts ) -> _Producer | None: - candidates = [] + producers = [] lineages: list[ColumnLineage] = [] for column in columns: lineage = facts.column_lineages.get(ColumnRef(root, column)) @@ -736,16 +816,10 @@ def _smallest_node_containing_all( 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): + producer = make_producer(node, bound_columns, path, facts) + if producer is not None: + producers.append(producer) + if blocks_pushdown(node, facts): break source_child_index = lineages[0].source_child_index if source_child_index is None or any( @@ -758,9 +832,9 @@ def _smallest_node_containing_all( break path = (*path, source_child_index) lineages = sources - if not candidates: + if not producers: return None - return min(candidates, key=lambda item: (item[0], item[1]))[2] + return min(producers, key=lambda p: p.domain_score) def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: @@ -768,13 +842,13 @@ def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | 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: + producer = make_producer(node, (bound_column,), path, facts) + if producer is None: continue item = ( - rows, + producer.rows, len(node.schema), - _Producer(node, (bound_column,), rows, path), + producer, ) if isinstance(node, (Scan, DataFrameScan)): source_candidates.append(item) @@ -809,6 +883,23 @@ def contains_node(root: IR, needle: IR) -> bool: return needle in traversal([root]) +def domain_cost_is_small( + domain: _Producer, target: _Producer, threshold: float +) -> bool: + """Return whether building a domain is cheap enough for its target.""" + 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.""" + node = root + for child_index in path: + if isinstance(node, Join) and node.options[0] == "Semi" and child_index == 0: + return True + node = node.children[child_index] + return False + + def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: join_filter_pushdown: dict[str, Any] = { "considered": True, @@ -830,6 +921,8 @@ def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: "domain_key": candidate.domain_key.name, "estimated_target_rows": candidate.target.rows, "estimated_domain_rows": candidate.domain.rows, + "estimated_target_cost": candidate.target.cost, + "estimated_domain_cost": candidate.domain.cost, "target_node_type": type(candidate.target.node).__name__, "domain_node_type": type(candidate.domain.node).__name__, } @@ -839,6 +932,7 @@ def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: { "constraint_key": candidate.target_constraint_key.name, "estimated_constraint_rows": candidate.constraint_domain.rows, + "estimated_constraint_cost": candidate.constraint_domain.cost, } ) log("Join Filter Pushdown", **record) 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 3dfe42df7fd3..7601788400db 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -10,6 +10,7 @@ import polars as pl 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.utils.column_domain import ColumnRef @@ -94,6 +95,13 @@ 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) + return names + + @pytest.fixture def simple_query() -> pl.LazyFrame: """Return a query with a small selective join domain.""" @@ -257,7 +265,7 @@ def test_no_simple_filter_pushdown_when_domain_is_not_selective( ConfigOptions.from_polars_engine(engine), ) - assert decision == Decision(reason="no_selective_domain") + assert decision == Decision(reason="no_profitable_domain") assert optimized is root assert not find_joins(optimized, "Semi") assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -330,6 +338,120 @@ def test_composite_filter_pushdown_constrains_domain_first( assert_gpu_result_equal(query, engine=engine, check_row_order=False) +def test_prefilter_uses_cheaper_source_domain_and_skips_expensive_domain( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(60), + "part_active": [True] * 30 + [False] * 30, + } + ) + .filter("part_active") + .select("p_partkey") + ) + partsupp = pl.LazyFrame( + { + "ps_partkey": [i % 60 for i in range(120)], + "ps_suppkey": [i % 30 for i in range(120)], + } + ) + supplier = pl.LazyFrame({"s_suppkey": range(30)}) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 60 for i in range(1_800)], + "l_suppkey": [i % 30 for i in range(1_800)], + "l_orderkey": [i % 900 for i in range(1_800)], + } + ) + orders = pl.LazyFrame({"o_orderkey": range(900)}) + query = ( + part.join(partsupp, left_on="p_partkey", right_on="ps_partkey") + .join(supplier, left_on="ps_suppkey", right_on="s_suppkey") + .join( + lineitem, + left_on=("p_partkey", "ps_suppkey"), + right_on=("l_partkey", "l_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), + ) + + part_ir = dataframe_scan(root, "p_partkey") + 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",) + ] + 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_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_source_only_domain_does_not_stack_on_prefiltered_source( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(60), + "part_active": [True] * 30 + [False] * 30, + } + ) + .filter("part_active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 60 for i in range(1_800)], + "l_orderkey": [i % 150 for i in range(1_800)], + } + ) + orders = ( + pl.LazyFrame( + { + "o_orderkey": range(150), + "order_active": [True] * 75 + [False] * 75, + } + ) + .filter("order_active") + .select("o_orderkey") + ) + query = part.join(lineitem, left_on="p_partkey", right_on="l_partkey").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_partkey") + lineitem_semis = [ + semi + for semi in find_joins(optimized, "Semi") + if semi.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_gpu_result_equal(query, engine=engine, check_row_order=False) + + def test_derived_selectivity_propagates_through_rewritten_children( engine: SPMDEngine, ) -> None: @@ -432,7 +554,7 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( 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(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 @@ -568,8 +690,10 @@ def test_composite_domain_columns_follow_renames(engine: SPMDEngine) -> None: analyzed = analyze_plan(renamed, StatsCollector()) facts = PlanFacts( row_estimates={renamed: 20, source: 10}, + source_facts=analyzed.source_facts, selective_nodes=analyzed.selective_nodes, column_lineages=analyzed.column_lineages, + refcounts=analyzed.refcounts, ) producer = _smallest_node_containing_all( renamed, ("domain_key", "domain_constraint"), facts @@ -739,7 +863,7 @@ def test_target_replacement_does_not_rewrite_shared_domain_side( 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 semis[0].children[0] is shared_ir assert not find_joins(unfiltered_domain, "Semi") assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -792,7 +916,59 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( 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])) + # 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 + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_internal_prefilter_rewrites_shared_subplan_once( + engine: SPMDEngine, +) -> None: + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + target = pl.LazyFrame( + { + "target_key": [i % 10 for i in range(20)], + "value": range(20), + } + ) + shared = domain.join( + target, + left_on="domain_key", + right_on="target_key", + ) + query = shared.join(shared, on="domain_key", suffix="_right") + translated = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(translated, Join) + shared_cache = translated.children[0] + assert isinstance(shared_cache, Cache) + assert translated.children[1] is shared_cache + original_shared = shared_cache.children[0] + assert len(find_joins(original_shared, "Inner")) == 1 + target_ir = dataframe_scan(original_shared, "target_key") + + optimized = optimize_with_stats( + translated, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + 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 assert_gpu_result_equal(query, engine=engine, check_row_order=False)