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 b099e1584ea..6c7e83cf2ad 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -69,9 +69,13 @@ from functools import singledispatch from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict +import polars as pl + +from cudf_polars.containers import DataType from cudf_polars.dsl import expr from cudf_polars.dsl.ir import ( IR, + Cache, ConditionalJoin, DataFrameScan, Distinct, @@ -180,6 +184,19 @@ def score(self) -> tuple[int, DomainScore, DomainScore]: return (0, self.constraint_domain.domain_score, self.domain.domain_score) +@dataclass(frozen=True) +class _AggregateReuseCandidate: + """A join detail side that can be replaced by an existing aggregate domain.""" + + join: Join + detail_side: Literal["left", "right"] + value_column: expr.Col + replacement: IR + domain_node_type: str + replacement_rows: int + replacement_cost: int + + Candidate: TypeAlias = SimpleCandidate | CompositeCandidate DecisionReason: TypeAlias = Literal[ "applied", @@ -468,6 +485,22 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: return apply_candidate(node, decision.candidate) +@_rewrite.register(GroupBy) +def _(node: GroupBy, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + original = node + rewritten = reuse_if_unchanged(node, rec) + assert isinstance(rewritten, GroupBy) + node = rewritten + if node is original: + facts = rec.state["facts"] + else: + facts = analyze_plan(node, rec.state["stats"]) + rewritten, aggregate_candidate = _rewrite_aggregate_domain_reuse(node, facts) + if rec.state["trace"]: + _trace_aggregate_reuse(node, aggregate_candidate) + return rewritten + + def apply_candidate(ir: Join, candidate: Candidate) -> IR: """Apply a selected join-domain prefilter candidate to a join.""" left, right = ir.children @@ -582,6 +615,689 @@ def _select_candidate( return Decision(reason="applied", candidate=min(candidates, key=lambda c: c.score)) +def _rewrite_aggregate_domain_reuse( + ir: GroupBy, facts: PlanFacts +) -> tuple[GroupBy, _AggregateReuseCandidate | None]: + """Replace one detail join with a compatible existing aggregate domain.""" + if ir.maintain_order or ir.zlice is not None: + return ir, None + + aggregate_column = _single_summed_column(ir) + if aggregate_column is None: + return ir, None + + candidates = list( + _aggregate_reuse_candidates(ir.children[0], aggregate_column, facts) + ) + if not candidates: + return ir, None + + for candidate in sorted( + candidates, + key=lambda item: (item.replacement_rows, item.replacement_cost), + ): + child = _replace_on_aggregate_path( + ir.children[0], + aggregate_column.name, + candidate.join, + candidate.replacement, + ) + if child is not None: + rewritten = ir.reconstruct((child,)) + assert isinstance(rewritten, GroupBy) + return rewritten, candidate + return ir, None + + +def _single_summed_column(ir: GroupBy) -> expr.Col | None: + """Return the sole directly summed, non-key column.""" + if len(ir.agg_requests) != 1: + return None + value = ir.agg_requests[0].value + if not isinstance(value, expr.Agg) or value.name != "sum": + return None + if len(value.children) != 1: + return None + (child,) = value.children + if not isinstance(child, expr.Col): + return None + + if any( + isinstance(node, expr.Col) and node.name == child.name + for node in traversal([key.value for key in ir.keys]) + ): + return None + return child + + +def _aggregate_reuse_candidates( + root: IR, + summed_column: expr.Col, + facts: PlanFacts, +) -> Iterable[_AggregateReuseCandidate]: + """Yield candidates along the exact input lineage of the final sum.""" + bindings = tuple(_exact_column_bindings(root, summed_column.name)) + for index, (node, bound_value_column) in enumerate(bindings): + if ( + isinstance(node, Join) + and node.options[0] == "Inner" + and node.options[2] is None + and node.options[5] == "none" + ): + yield from _aggregate_reuse_candidates_for_join( + node, + bound_value_column, + facts, + ) + + if index + 1 < len(bindings): + child, input_column = bindings[index + 1] + if not _aggregate_reuse_edge_is_safe( + node, bound_value_column, child, input_column + ): + return + + +def _aggregate_reuse_candidates_for_join( + node: Join, + summed_column: str, + facts: PlanFacts, +) -> Iterable[_AggregateReuseCandidate]: + """Yield aggregate replacements for one join on the sum lineage.""" + left_keys = _simple_keys(node.left_on) + right_keys = _simple_keys(node.right_on) + if len(left_keys) != len(node.left_on) or len(right_keys) != len(node.right_on): + return + assert len(left_keys) == len(right_keys) + + value_binding = _join_input_binding(node, summed_column) + if value_binding is None: + return + detail_child, detail_value_column = value_binding + detail_indices = [ + index for index, child in enumerate(node.children) if child is detail_child + ] + if len(detail_indices) != 1: + return + (detail_index,) = detail_indices + detail_side: Literal["left", "right"] = "left" if detail_index == 0 else "right" + detail_keys = left_keys if detail_index == 0 else right_keys + domain_child = node.children[1 - detail_index] + domain_keys = right_keys if detail_index == 0 else left_keys + allowed_columns = {detail_value_column} + allowed_columns.update(key.name for key in detail_keys) + bound_detail_columns = set() + for output_column in node.schema: + binding = _join_input_binding(node, output_column) + if binding is not None and binding[0] is detail_child: + bound_detail_columns.add(binding[1]) + if not bound_detail_columns.issubset(allowed_columns): + return + + value_column = expr.Col( + detail_child.schema[detail_value_column], detail_value_column + ) + replacement_detail = _aggregate_reuse_detail_replacement( + detail_child, + domain_child, + domain_keys, + detail_keys, + value_column, + node.options[1], + node.options[3], + facts, + ) + if replacement_detail is None: + return + replacement, domain_node_type, replacement_rows, replacement_cost = ( + replacement_detail + ) + children = list(node.children) + children[detail_index] = replacement + yield _AggregateReuseCandidate( + join=node, + detail_side=detail_side, + value_column=value_column, + replacement=node.reconstruct(children), + domain_node_type=domain_node_type, + replacement_rows=replacement_rows, + replacement_cost=replacement_cost, + ) + + +def _aggregate_reuse_detail_replacement( + detail_child: IR, + domain_child: IR, + domain_keys: tuple[expr.Col, ...], + detail_keys: tuple[expr.Col, ...], + value_column: expr.Col, + nulls_equal: bool, # noqa: FBT001 + suffix: str, + facts: PlanFacts, +) -> tuple[IR, str, int, int] | None: + """Build a null-correct aggregate replacement for a semi-filtered detail side.""" + if len(detail_keys) != 1: + return None + (detail_key,) = detail_keys + if isinstance(detail_child, Join) and detail_child.options[0] == "Semi": + return _aggregate_reuse_detail_replacement_from_semi_detail( + detail_child, + detail_key, + value_column, + facts, + ) + if len(domain_keys) != 1: + return None + domain = _aggregate_reuse_domain_from_sibling_semi( + domain_child, + domain_keys[0], + detail_child, + detail_key, + value_column, + nulls_equal, + facts, + ) + if domain is None: + return None + return _aggregate_reuse_detail_replacement_from_domain( + domain_root=domain.root, + aggregate_domain=domain.aggregate, + detail_key=detail_key, + value_column=value_column, + nulls_equal=nulls_equal, + suffix=suffix, + ) + + +def _aggregate_reuse_detail_replacement_from_semi_detail( + detail_child: Join, + detail_key: expr.Col, + value_column: expr.Col, + facts: PlanFacts, +) -> tuple[IR, str, int, int] | None: + """Build an aggregate replacement when the detail side is already semi-filtered.""" + if detail_child.options[2] is not None or detail_child.options[5] != "none": + return None + + semi_left_keys = _simple_keys(detail_child.left_on) + semi_right_keys = _simple_keys(detail_child.right_on) + if len(semi_left_keys) != len(detail_child.left_on) or len(semi_right_keys) != len( + detail_child.right_on + ): + return None + if semi_left_keys != (detail_key,) or len(semi_right_keys) != 1: + return None + (domain_key,) = semi_right_keys + + aggregate_domain = _selective_aggregate_domain( + detail_child.children[1], + domain_key, + detail_child.children[0], + detail_key, + value_column, + facts, + ) + if aggregate_domain is None: + return None + return _aggregate_reuse_detail_replacement_from_domain( + domain_root=detail_child.children[1], + aggregate_domain=aggregate_domain, + detail_key=detail_key, + value_column=value_column, + nulls_equal=detail_child.options[1], + suffix=detail_child.options[3], + ) + + +@dataclass(frozen=True) +class _AggregateSiblingDomain: + """An aggregate domain found through a sibling semi join.""" + + root: IR + aggregate: _Producer + + +def _aggregate_reuse_domain_from_sibling_semi( + domain_child: IR, + domain_key: expr.Col, + detail_child: IR, + detail_key: expr.Col, + value_column: expr.Col, + nulls_equal: bool, # noqa: FBT001 + facts: PlanFacts, +) -> _AggregateSiblingDomain | None: + """Find an aggregate domain used to semi-filter the other join side.""" + if not isinstance(domain_child, Join) or domain_child.options[0] != "Semi": + return None + if domain_child.options[2] is not None or domain_child.options[5] != "none": + return None + if domain_child.options[1] != nulls_equal: + return None + + semi_left_keys = _simple_keys(domain_child.left_on) + semi_right_keys = _simple_keys(domain_child.right_on) + if len(semi_left_keys) != len(domain_child.left_on) or len(semi_right_keys) != len( + domain_child.right_on + ): + return None + if semi_left_keys != (domain_key,) or len(semi_right_keys) != 1: + return None + (aggregate_key,) = semi_right_keys + + aggregate_domain = _selective_aggregate_domain( + domain_child.children[1], + aggregate_key, + detail_child, + detail_key, + value_column, + facts, + ) + if aggregate_domain is None: + return None + return _AggregateSiblingDomain(domain_child.children[1], aggregate_domain) + + +def _aggregate_reuse_detail_replacement_from_domain( + *, + domain_root: IR, + aggregate_domain: _Producer, + detail_key: expr.Col, + value_column: expr.Col, + nulls_equal: bool, + suffix: str, +) -> tuple[IR, str, int, int]: + """Project aggregate key/value columns into a detail-side replacement.""" + domain = aggregate_domain + aggregate_values = _project_bound_key_and_value( + domain.node, + domain.columns[0], + domain.columns[1], + detail_key, + value_column, + ) + replacement: IR + if _domain_key_set_is_preserved(domain_root, domain.node, domain.columns[0]): + replacement = ( + aggregate_values + if nulls_equal + else _drop_null_key(aggregate_values, detail_key) + ) + else: + replacement = _make_semi_join( + aggregate_values, + expr.Col(aggregate_values.schema[detail_key.name], detail_key.name), + domain_root, + expr.Col(domain_root.schema[domain.columns[0]], domain.columns[0]), + nulls_equal=nulls_equal, + suffix=suffix, + ) + return replacement, type(domain.node).__name__, domain.rows, domain.cost + + +def _selective_aggregate_domain( + root: IR, + domain_key: expr.Col, + detail_source: IR, + detail_key: expr.Col, + value_column: expr.Col, + facts: PlanFacts, +) -> _Producer | None: + """Find a selective aggregate derived from the identical detail source.""" + candidates = [] + for groupby, bound_domain_key in _exact_column_bindings(root, domain_key.name): + if not isinstance(groupby, GroupBy): + continue + if ( + not _same_detail_source(groupby.children[0], detail_source) + or len(groupby.keys) != 1 + ): + continue + (groupby_key,) = groupby.keys + if ( + groupby_key.name != bound_domain_key + or not isinstance(groupby_key.value, expr.Col) + or groupby_key.value != detail_key + ): + continue + aggregate_column = _sum_aggregate_output_column(groupby, value_column) + if aggregate_column is None: + continue + domain = _smallest_selective_node_containing_all( + root, + bound_domain_key, + aggregate_column.name, + groupby, + facts, + ) + if domain is not None: + candidates.append((domain.rows, domain.cost, domain)) + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +def _same_detail_source(left: IR, right: IR) -> bool: + """Return whether two nodes refer to the same materialized detail source.""" + if left is right: + return True + if isinstance(left, Scan) and isinstance(right, Scan): + return ( + left.typ == right.typ + and left.reader_options == right.reader_options + and left.cloud_options == right.cloud_options + and left.paths == right.paths + and left.skip_rows == right.skip_rows + and left.n_rows == right.n_rows + and left.row_index == right.row_index + and left.include_file_paths == right.include_file_paths + and left.predicate == right.predicate + and left.parquet_options == right.parquet_options + ) + if isinstance(left, DataFrameScan) and isinstance(right, DataFrameScan): + return left == right + return ( + isinstance(left, Cache) + and isinstance(right, Cache) + and left.key == right.key + and left.refcount == right.refcount + and left.schema == right.schema + ) + + +def _sum_aggregate_output_column( + ir: GroupBy, value_column: expr.Col +) -> expr.Col | None: + """Return the output of a direct sum over the requested value column.""" + for request in ir.agg_requests: + value = request.value + if not isinstance(value, expr.Agg) or value.name != "sum": + continue + if value.dtype != value_column.dtype or len(value.children) != 1: + continue + (child,) = value.children + if isinstance(child, expr.Col) and child.name == value_column.name: + if request.name not in ir.schema: + return None + return expr.Col(ir.schema[request.name], request.name) + return None + + +def _smallest_selective_node_containing_all( + root: IR, + key_column: str, + value_column: str, + anchor: IR, + facts: PlanFacts, +) -> _Producer | None: + """Find the smallest selective node whose columns bind to an anchor.""" + candidates = [] + for node in traversal([root]): + bound_columns = _bound_aggregate_output_columns( + node, anchor, key_column, value_column + ) + if bound_columns is None or node not in facts.selective_nodes: + continue + producer = make_producer(node, bound_columns, (), facts) + if producer is not None: + candidates.append( + (producer.rows, producer.cost, len(node.schema), producer) + ) + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1], item[2]))[3] + + +def _bound_aggregate_output_columns( + node: IR, anchor: IR, key_column: str, value_column: str +) -> tuple[str, ...] | None: + """Map an aggregate key and normalized sum output through a domain node.""" + outputs = [] + for anchor_column, binding_fn in ( + (key_column, _exact_column_bindings), + (value_column, _aggregate_value_bindings), + ): + matches = [ + output_column + for output_column in node.schema + if any( + candidate is anchor and bound_column == anchor_column + for candidate, bound_column in binding_fn(node, output_column) + ) + ] + if len(matches) != 1: + return None + outputs.append(matches[0]) + return tuple(outputs) + + +def _aggregate_value_bindings(root: IR, column: str) -> Iterable[tuple[IR, str]]: + """Yield sum-value lineage through direct bindings and zero-fill normalization.""" + node = root + while column in node.schema: + yield node, column + binding = _input_binding(node, column) + if binding is None and isinstance(node, Select): + selected = next((item for item in node.exprs if item.name == column), None) + binding = _zero_fill_binding(node.children[0], selected) + if binding is None: + return + node, column = binding + + +def _zero_fill_binding( + child: IR, expression: expr.NamedExpr | None +) -> tuple[IR, str] | None: + """Bind Polars' ``sum`` null normalization to its aggregate output.""" + if expression is None or not isinstance(expression.value, expr.UnaryFunction): + return None + value = expression.value + if value.name != "fill_null" or value.options or len(value.children) != 2: + return None + source, fill = value.children + if ( + not isinstance(source, expr.Col) + or not isinstance(fill, expr.Literal) + or fill.value != 0 + or source.name not in child.schema + or source.dtype != child.schema[source.name] + or fill.dtype != source.dtype + or value.dtype != source.dtype + ): + return None + return child, source.name + + +def _project_bound_key_and_value( + source: IR, + source_key: str, + source_value: str, + output_key: expr.Col, + output_value: expr.Col, +) -> Select: + """Project aggregate key and value columns under detail-side names.""" + assert source.schema[source_key] == output_key.dtype + assert source.schema[source_value] == output_value.dtype + return Select( + {output_key.name: output_key.dtype, output_value.name: output_value.dtype}, + ( + expr.NamedExpr(output_key.name, expr.Col(output_key.dtype, source_key)), + expr.NamedExpr( + output_value.name, expr.Col(output_value.dtype, source_value) + ), + ), + True, # noqa: FBT003 + source, + ) + + +def _domain_key_set_is_preserved(root: IR, domain: IR, column: str) -> bool: + """Return whether a domain reaches the semi join through row-preserving nodes.""" + for node, _ in _exact_column_bindings(root, column): + if node is domain: + return True + if not isinstance(node, (Cache, HStack, Projection, Select)): + return False + return False + + +def _drop_null_key(source: IR, key: expr.Col) -> Filter: + """Apply the null-key behavior of a semi join with ``nulls_equal=False``.""" + bool_dtype = DataType(pl.Boolean()) + predicate = expr.NamedExpr( + "__join_filter_pushdown_key_is_not_null", + expr.BooleanFunction( + bool_dtype, + expr.BooleanFunction.Name.IsNotNull, + (), + expr.Col(key.dtype, key.name), + ), + ) + return Filter(source.schema, predicate, source) + + +def _aggregate_reuse_edge_is_safe( + node: IR, + output_column: str, + child: IR, + input_column: str, +) -> bool: + """Return whether aggregate values can move across one lineage edge.""" + child_indices = [ + index for index, candidate in enumerate(node.children) if candidate is child + ] + if len(child_indices) != 1: + return False + (child_index,) = child_indices + if _input_binding(node, output_column) != (child, input_column): + return False + if isinstance(node, (Cache, HStack, Projection, Select)): + return True + if not isinstance(node, Join): + return False + if ( + node.options[0] != "Inner" + or node.options[2] is not None + or node.options[5] != "none" + ): + return False + keys = node.left_on if child_index == 0 else node.right_on + return not any( + isinstance(item, expr.Col) and item.name == input_column + for item in traversal([key.value for key in keys]) + ) + + +def _replace_on_aggregate_path( + root: IR, + column: str, + target: Join, + replacement: IR, +) -> IR | None: + """Replace a join only on the direct lineage of an aggregate column.""" + bindings = tuple(_exact_column_bindings(root, column)) + path: list[tuple[IR, int]] = [] + for index, (node, _) in enumerate(bindings): + if node is target: + rewritten = replacement + for parent, child_index in reversed(path): + children = list(parent.children) + children[child_index] = rewritten + rewritten = parent.reconstruct(children) + return rewritten + if index + 1 == len(bindings): + break + child = bindings[index + 1][0] + child_indices = [ + child_index + for child_index, candidate in enumerate(node.children) + if candidate is child + ] + if len(child_indices) != 1: + return None + path.append((node, child_indices[0])) + return None + + +def _exact_column_bindings(root: IR, column: str) -> Iterable[tuple[IR, str]]: + """Yield direct output-to-input bindings for a column through a subplan.""" + node = root + while column in node.schema: + yield node, column + binding = _input_binding(node, column) + if binding is None: + return + node, column = binding + + +def _input_binding(node: IR, column: str) -> tuple[IR, str] | None: + """Return a proven direct input binding, stopping at ambiguous operations.""" + child = node.children[0] if len(node.children) == 1 else None + if isinstance(node, Select): + selected = next((item for item in node.exprs if item.name == column), None) + return _column_expression_binding(child, selected) + if isinstance(node, HStack): + stacked = next((item for item in node.columns if item.name == column), None) + if stacked is not None: + return _column_expression_binding(child, stacked) + return _passthrough_binding(child, column) + if isinstance(node, GroupBy): + if node.zlice is not None: + return None + key = next((item for item in node.keys if item.name == column), None) + return _column_expression_binding(child, key) + if isinstance(node, Join): + return _join_input_binding(node, column) + if isinstance(node, Distinct): + return None if node.zlice is not None else _passthrough_binding(child, column) + if isinstance(node, Sort): + return None if node.zlice is not None else _passthrough_binding(child, column) + if isinstance(node, (Cache, Filter, Projection)): + return _passthrough_binding(child, column) + return None + + +def _column_expression_binding( + child: IR | None, expression: expr.NamedExpr | None +) -> tuple[IR, str] | None: + if ( + child is not None + and expression is not None + and isinstance(expression.value, expr.Col) + and expression.value.name in child.schema + ): + return child, expression.value.name + return None + + +def _passthrough_binding(child: IR | None, column: str) -> tuple[IR, str] | None: + if child is not None and column in child.schema: + return child, column + return None + + +def _join_input_binding(node: Join, column: str) -> tuple[IR, str] | None: + if node.options[2] is not None: + return None + left, right = node.children + if node.options[0] in ("Semi", "Anti"): + return _passthrough_binding(left, column) + if node.options[0] != "Inner": + return None + bindings = [] + if column in left.schema: + bindings.append((left, column)) + suffix = node.options[3] + for right_column in right.schema: + output_column = ( + f"{right_column}{suffix}" if right_column in left.schema else right_column + ) + if output_column == column and output_column in node.schema: + bindings.append((right, right_column)) + if len(bindings) == 1: + return bindings[0] + return None + + def _simple_keys(keys: Sequence[expr.NamedExpr]) -> tuple[expr.Col, ...]: return tuple(key.value for key in keys if isinstance(key.value, expr.Col)) @@ -936,3 +1652,28 @@ def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: } ) log("Join Filter Pushdown", **record) + + +def _trace_aggregate_reuse( + ir: GroupBy, candidate: _AggregateReuseCandidate | None +) -> None: + aggregate_domain_reuse: dict[str, Any] = { + "reason": "applied" if candidate is not None else "no_reusable_domain", + } + record = { + "scope": Scope.PLAN.value, + "join_aggregate_domain_reuse": aggregate_domain_reuse, + "actor_ir_id": ir.get_stable_id(), + "actor_ir_type": type(ir).__name__, + } + if candidate is not None: + aggregate_domain_reuse.update( + { + "detail_side": candidate.detail_side, + "value_column": candidate.value_column.name, + "domain_node_type": candidate.domain_node_type, + "estimated_replacement_rows": candidate.replacement_rows, + "estimated_replacement_cost": candidate.replacement_cost, + } + ) + log("Join Aggregate Domain Reuse", **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 7601788400d..b330f8267c1 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -11,7 +11,7 @@ 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.ir import Cache, DataFrameScan, Distinct, Join, Scan, 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 @@ -36,6 +36,7 @@ if TYPE_CHECKING: import concurrent.futures + import pathlib from typing import Any from cudf_polars.dsl.ir import IR @@ -562,6 +563,77 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( assert_gpu_result_equal(query, engine=engine, check_row_order=False) +@pytest.mark.parametrize( + "nulls_equal", [False, True], ids=["nulls_not_equal", "nulls_equal"] +) +def test_reuses_existing_aggregate_domain_to_replace_detail_join( + nulls_equal: bool, # noqa: FBT001 + tmp_path: pathlib.Path, + engine: SPMDEngine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + lineitem_path = tmp_path / "lineitem.parquet" + orders_path = tmp_path / "orders.parquet" + customer_path = tmp_path / "customer.parquet" + pl.DataFrame( + { + "l_orderkey": [1, 1, 2, 3], + "l_quantity": [2.0, 3.0, 7.0, 1.0], + } + ).write_parquet(lineitem_path) + pl.DataFrame({"o_orderkey": [1, 2, 3], "o_custkey": [10, 20, 30]}).write_parquet( + orders_path + ) + pl.DataFrame({"c_custkey": [10, 20, 30], "c_name": ["a", "b", "c"]}).write_parquet( + customer_path + ) + + lineitem = pl.scan_parquet(lineitem_path) + orders = pl.scan_parquet(orders_path) + customer = pl.scan_parquet(customer_path) + selected_keys = ( + lineitem.group_by("l_orderkey") + .agg(pl.col("l_quantity").sum().alias("sum_quantity")) + .filter(pl.col("sum_quantity") > 4) + .select("l_orderkey") + ) + query = ( + orders.join( + selected_keys, + how="semi", + left_on="o_orderkey", + right_on="l_orderkey", + nulls_equal=nulls_equal, + ) + .join( + lineitem, + left_on="o_orderkey", + right_on="l_orderkey", + nulls_equal=nulls_equal, + ) + .join(customer, left_on="o_custkey", right_on="c_custkey") + .group_by("c_name", "o_custkey", "o_orderkey") + .agg(pl.col("l_quantity").sum()) + ) + ir = remove_cache_nodes(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, + ) + + detail_semis = [ + join + for join in find_joins(optimized, "Semi") + if isinstance(join.children[0], Scan) + and "l_quantity" in join.children[0].schema + ] + assert not detail_semis + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + def test_target_source_follows_join_key_through_rename( engine: SPMDEngine, ) -> None: