From 0f5b43ab3b873f06fef571e137301fe3cd4c692f Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 28 Jul 2026 14:40:24 +0100 Subject: [PATCH 01/17] Introduce PushdownFilterHint to represent semi joins in filter pushdown The joins introduced by the filter pushdown optimisation pass on plans are not required for correctness of the query. Instead they are only hints that might be useful for providing a faster implementation. Currently we cannot take advantage of this optional component of the filter because there is no way to distinguish between a semi join that is part of the user query, and one that the optimisation pass introduces. To fix this, introduce a PushdownFilterHint node. Initially, this is removed during lowering to actors. --- .../cudf_polars/dsl/utils/column_domain.py | 9 + .../cudf_polars/streaming/explain.py | 30 +++- .../cudf_polars/streaming/filter_hint.py | 73 ++++++++ .../cudf_polars/cudf_polars/streaming/join.py | 10 ++ .../streaming/join_filter_pushdown.py | 55 +++--- .../cudf_polars/streaming/parallel.py | 1 + .../tests/streaming/test_explain.py | 36 ++++ .../streaming/test_join_filter_pushdown.py | 158 ++++++++++-------- 8 files changed, 276 insertions(+), 96 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/filter_hint.py diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py index 382d0027f649..f15f5f01a0b9 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -21,6 +21,7 @@ Slice, Sort, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: from collections.abc import Mapping @@ -141,3 +142,11 @@ def _( return { name: ColumnBinding(0, name) for name in node.schema if name in child.schema } + + +@column_domain_bindings.register(PushdownFilterHint) +def _(node: PushdownFilterHint) -> Mapping[str, ColumnBinding]: + target = node.children[0] + return { + name: ColumnBinding(0, name) for name in node.schema if name in target.schema + } diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 011e3be232a7..b93a96d36ec3 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -37,8 +37,9 @@ from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import IOPartitionFlavor +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.io import StreamingScan, scan_partition_plan -from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.parallel import lower_ir_graph, optimize_with_stats from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.statistics import ( collect_statistics, @@ -134,6 +135,7 @@ def explain_query( # Include row-count statistics for the logical plan with cm: stats = collect_statistics(ir, config, executor) + ir = optimize_with_stats(ir, config, stats) return _repr_ir_tree(ir, stats=stats) else: return _repr_ir_tree(ir) @@ -469,6 +471,17 @@ def _(ir: Join, *, offset: str = "") -> str: return _repr_header(offset, f"JOIN {ir.options[0]} {left_on} {right_on}", ir.schema) +@_repr_ir.register +def _(ir: PushdownFilterHint, *, offset: str = "") -> str: + target_on = tuple(ne.name for ne in ir.target_on) + domain_on = tuple(ne.name for ne in ir.domain_on) + return _repr_header( + offset, + f"PUSHDOWN FILTER HINT {target_on} {domain_on}", + ir.schema, + ) + + _BinaryOperator = plc.binaryop.BinaryOperator _BINOP_SYMBOLS: dict[_BinaryOperator, str] = { _BinaryOperator.EQUAL: "==", @@ -580,6 +593,15 @@ def _(ir: Join) -> dict[str, Serializable]: } +@_serialize_properties.register +def _(ir: PushdownFilterHint) -> dict[str, Serializable]: + return { + "target_on": [ne.name for ne in ir.target_on], + "domain_on": [ne.name for ne in ir.domain_on], + "nulls_equal": ir.nulls_equal, + } + + @_serialize_properties.register def _(ir: GroupBy) -> dict[str, Serializable]: return { @@ -815,4 +837,10 @@ def from_query( """ config_options = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() + if not lowered and config_options.executor.name == "streaming": + with concurrent.futures.ThreadPoolExecutor( + thread_name_prefix="cudf-polars-explain" + ) as executor: + stats = collect_statistics(ir, config_options, executor) + ir = optimize_with_stats(ir, config_options, stats) return cls.from_ir(ir, config_options=config_options, lowered=lowered) diff --git a/python/cudf_polars/cudf_polars/streaming/filter_hint.py b/python/cudf_polars/cudf_polars/streaming/filter_hint.py new file mode 100644 index 000000000000..4893f5e7cab3 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/filter_hint.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Logical filter hints for the streaming runtime.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from cudf_polars.dsl.ir import IR + +if TYPE_CHECKING: + from collections.abc import Sequence + + from cudf_polars.containers import DataFrame + from cudf_polars.dsl.expr import NamedExpr + from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.typing import Schema + + +class PushdownFilterHint(IR): + """ + Optional join-key filter placed in a logical plan. + + The first child is the target to filter and the second child, the + domain, provides the keys to filter against. Applying the filter is + optional. + """ + + __slots__ = ("domain_on", "nulls_equal", "target_on") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "target_on", + "domain_on", + "nulls_equal", + ) + _n_non_child_args: ClassVar[int] = 3 + + target_on: tuple[NamedExpr, ...] + """Expressions selecting filter keys from the target.""" + domain_on: tuple[NamedExpr, ...] + """Expressions selecting filter keys from the domain.""" + nulls_equal: bool + """Whether null key values compare equal.""" + + def __init__( + self, + schema: Schema, + target_on: Sequence[NamedExpr], + domain_on: Sequence[NamedExpr], + nulls_equal: bool, # noqa: FBT001 + target: IR, + domain: IR, + ): + self.schema = schema + self.target_on = tuple(target_on) + self.domain_on = tuple(domain_on) + self.nulls_equal = nulls_equal + self._non_child_args = (self.target_on, self.domain_on, self.nulls_equal) + self.children = (target, domain) + + @classmethod + def do_evaluate( + cls, + target_on: tuple[NamedExpr, ...], + domain_on: tuple[NamedExpr, ...], + nulls_equal: bool, # noqa: FBT001 + target: DataFrame, + domain: DataFrame, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Ignore the optional filter and return the target.""" + return target diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 63d76d6c328f..c13d977c27c2 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -12,6 +12,7 @@ from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.utils import ( @@ -149,6 +150,15 @@ def _has_non_pointwise_keys(ir: Join) -> bool: return not all(expr.is_pointwise for expr in traversal(keys)) +@lower_ir_node.register(PushdownFilterHint) +def _( + ir: PushdownFilterHint, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + """Discard the optional filter without lowering its domain.""" + target, _domain = ir.children + return rec(target) + + @lower_ir_node.register(ConditionalJoin) def _( ir: ConditionalJoin, rec: LowerIRTransformer diff --git a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py index b099e1584ea6..59763e83edf8 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -5,13 +5,14 @@ For a supported inner equijoin, this optimization tries to use the join-key values produced by one input to reduce the size of the other input before -the original join. In relational notation, a simple rewrite is:: +the original join. It records that opportunity with a logical +``PushdownFilterHint``:: left join[left.key = right.key] right -> - (left semijoin[left.key = right.key] project(right.key)) + PushdownFilterHint(left, left.key, project(right.key), right.key) join[left.key = right.key] right In this example, the right hand table is selected to pre-filter the left @@ -49,14 +50,14 @@ A rewrite that projects one domain join key and uses it to filter the corresponding target key directly. ``composite candidate`` - For a multi-key join, a rewrite that first semi-joins the domain using the - constraint domain, then projects the reduced domain's key used to filter - the target. + For a multi-key join, a rewrite that first hints that the domain should be + filtered using the constraint domain, then projects the reduced domain's + key used to filter the target. Plan rewrite has three stages. ``analyze_plan`` gathers row estimates, source scan facts, selective nodes, and column value-domain lineages. Candidate selection consumes those facts and returns a decision. -``apply_candidate`` then constructs the selected semi-join rewrite. +``apply_candidate`` then constructs the selected filter-hint rewrite. Row estimates, selectivity propagation, thresholds, and candidate scores are only heuristics for deciding whether a safe rewrite is likely to improve @@ -100,6 +101,7 @@ ColumnRef, column_domain_bindings, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping, Sequence @@ -249,6 +251,12 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: rows = node.df.shape()[0] elif isinstance(node, (Select, Projection, HStack, Filter, Distinct, GroupBy)): rows = row_estimates[node.children[0]] + elif isinstance(node, PushdownFilterHint): + rows = _estimate_join_rows( + "Semi", + row_estimates[node.children[0]], + row_estimates[node.children[1]], + ) elif isinstance(node, Join): rows = _estimate_join_rows( node.options[0], @@ -329,7 +337,7 @@ def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: Returns ------- bool - True if a semijoin cannot be pushed past this node, otherwise False. + True if a filter hint cannot be pushed past this node, otherwise False. """ # TODO: Need better cost model to handle nodes that are shared. Pushing # a filter into a shared node will typically mean that it is no longer @@ -350,11 +358,11 @@ def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: ) -def semijoin_pushdown_candidates( +def filter_hint_pushdown_candidates( facts: PlanFacts, root: IR, column: str ) -> Iterator[tuple[ColumnRef, tuple[int, ...]]]: """ - Yield column domain lineage providing valid locations for semijoin pushdown. + Yield column domain lineage providing valid locations for a filter hint. Parameters ---------- @@ -452,7 +460,7 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: if node is original: facts = rec.state["facts"] else: - # Child rewrites introduce new semi joins and reconstructed ancestors. + # Child rewrites introduce new filter hints and reconstructed ancestors. # Re-analyze that current subtree so parent joins can use the derived # selectivity and cardinality when ranking their own candidates. facts = analyze_plan(node, rec.state["stats"]) @@ -473,13 +481,12 @@ def apply_candidate(ir: Join, candidate: Candidate) -> IR: left, right = ir.children domain = _make_domain(candidate, ir) target = candidate.target - target_filter = _make_semi_join( + target_filter = _make_filter_hint( target.node, expr.Col(target.node.schema[target.column], target.column), domain, expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), nulls_equal=ir.options[1], - suffix=ir.options[3], ) if candidate.target_side == "left": left = replace_at_path(left, target.path, target_filter) @@ -611,7 +618,7 @@ def _simple_candidates( continue if contains_node(target.node, domain.node): continue - if domain.is_single_source and has_filtering_semi_ancestor( + if domain.is_single_source and has_filtering_hint_ancestor( target_child, target.path ): continue @@ -704,7 +711,7 @@ def _make_domain(candidate: Candidate, ir: Join) -> IR: candidate.constraint_domain.column, candidate.target_constraint_key, ) - constrained = _make_semi_join( + constrained = _make_filter_hint( candidate.domain.node, expr.Col( candidate.domain.node.schema[candidate.domain.columns[1]], @@ -716,7 +723,6 @@ def _make_domain(candidate: Candidate, ir: Join) -> IR: candidate.target_constraint_key.name, ), nulls_equal=ir.options[1], - suffix=ir.options[3], ) return _project_bound_key( constrained, candidate.domain.column, candidate.domain_key @@ -735,20 +741,19 @@ def _project_bound_key(source: IR, bound_column: str, output_key: expr.Col) -> S ) -def _make_semi_join( +def _make_filter_hint( target: IR, target_key: expr.Col, domain: IR, domain_key: expr.Col, *, nulls_equal: bool, - suffix: str, -) -> Join: - return Join( +) -> PushdownFilterHint: + return PushdownFilterHint( target.schema, (expr.NamedExpr(target_key.name, target_key),), (expr.NamedExpr(domain_key.name, domain_key),), - ("Semi", nulls_equal, None, suffix, False, "none"), + nulls_equal, target, domain, ) @@ -784,7 +789,7 @@ def _smallest_key_producer( exclude: IR | None = None, ) -> _Producer | None: producers = [] - for reference, path in semijoin_pushdown_candidates(facts, root, column): + for reference, path in filter_hint_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name if node is exclude: continue @@ -840,7 +845,7 @@ def _smallest_node_containing_all( def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: source_candidates = [] fallback_candidates = [] - for reference, path in semijoin_pushdown_candidates(facts, root, column): + for reference, path in filter_hint_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name producer = make_producer(node, (bound_column,), path, facts) if producer is None: @@ -890,11 +895,11 @@ def domain_cost_is_small( return domain.cost / target.rows <= threshold -def has_filtering_semi_ancestor(root: IR, path: Sequence[int]) -> bool: - """Return whether a selected child edge is below a filtering semi join.""" +def has_filtering_hint_ancestor(root: IR, path: Sequence[int]) -> bool: + """Return whether a selected child edge is below a pushdown-filter hint.""" node = root for child_index in path: - if isinstance(node, Join) and node.options[0] == "Semi" and child_index == 0: + if isinstance(node, PushdownFilterHint) and child_index == 0: return True node = node.children[child_index] return False diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 2ac5c8c2eef7..55704ce37a60 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -15,6 +15,7 @@ # handlers at import time so the dispatch table is populated before any query # is lowered. import cudf_polars.streaming.distinct +import cudf_polars.streaming.filter_hint import cudf_polars.streaming.groupby import cudf_polars.streaming.io import cudf_polars.streaming.join diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5ea6be578ae4..610c35256d60 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -137,6 +137,42 @@ def test_explain_logical_plan_with_join(tmp_path, df): assert "JOIN Inner ('x',) ('x',)" in plan +def test_explain_pushdown_filter_hint_is_logical_only(): + domain = ( + pl.LazyFrame({"key": [1, 99], "active": [True, False]}) + .filter("active") + .select("key") + ) + target = pl.LazyFrame({"key": [i % 10 for i in range(20)]}) + query = domain.join(target, on="key") + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={"join_filter_pushdown": {"threshold": 0.5}}, + ) + + logical = explain_query(query, engine, physical=False) + physical = explain_query(query, engine, physical=True) + logical_serialized = serialize_query(query, engine, physical=False) + physical_serialized = serialize_query(query, engine, physical=True) + + assert "PUSHDOWN FILTER HINT ('key',) ('key',)" in logical + assert "PUSHDOWN FILTER HINT" not in physical + assert any( + node.type == "PushdownFilterHint" + and node.properties + == { + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + } + for node in logical_serialized.nodes.values() + ) + assert not any( + node.type == "PushdownFilterHint" for node in physical_serialized.nodes.values() + ) + + def test_explain_logical_plan_with_sort(tmp_path, df): make_partitioned_source(df, tmp_path, fmt="parquet", n_files=2) diff --git a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py index 7601788400db..6e7e7c475dd1 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -16,6 +16,7 @@ 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.filter_hint import PushdownFilterHint from cudf_polars.streaming.join_filter_pushdown import ( CompositeCandidate, Decision, @@ -26,10 +27,14 @@ analyze_plan, apply_candidate, contains_node, + filter_hint_pushdown_candidates, optimize_join_filter_pushdown, - semijoin_pushdown_candidates, ) -from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes +from cudf_polars.streaming.parallel import ( + lower_ir_graph, + optimize_with_stats, + remove_cache_nodes, +) from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import ConfigOptions @@ -77,6 +82,10 @@ def find_joins(ir: IR, how: str | None = None) -> list[Join]: ] +def find_hints(ir: IR) -> list[PushdownFilterHint]: + return [node for node in traversal([ir]) if isinstance(node, PushdownFilterHint)] + + def translate_query(query: pl.LazyFrame, engine: SPMDEngine) -> IR: """Translate a public Polars query and remove logical Cache nodes.""" t = Translator(query._ldf.visit(), engine) @@ -95,10 +104,12 @@ def dataframe_scan(ir: IR, column: str) -> DataFrameScan: return match -def join_key_names(join: Join) -> tuple[str, ...]: - """Return the column names used on the left of a simple-column join.""" - names = tuple(key.value.name for key in join.left_on if isinstance(key.value, Col)) - assert len(names) == len(join.left_on) +def hint_key_names(hint: PushdownFilterHint) -> tuple[str, ...]: + """Return the target column names used by a filter hint.""" + names = tuple( + key.value.name for key in hint.target_on if isinstance(key.value, Col) + ) + assert len(names) == len(hint.target_on) return names @@ -141,10 +152,11 @@ def test_simple_prefilter_filters_large_side( assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" - semis = find_joins(optimized, "Semi") - assert len(semis) == 1 - assert semis[0].children[0] is lineitem_ir - assert not find_joins(part_ir, "Semi") + assert not find_joins(optimized, "Semi") + hints = find_hints(optimized) + assert len(hints) == 1 + assert hints[0].children[0] is lineitem_ir + assert not find_hints(part_ir) assert_gpu_result_equal(simple_query, engine=engine, check_row_order=False) @@ -160,7 +172,21 @@ def test_filter_pushdown_is_independent_of_dynamic_planning( make_config(dynamic_planning=False), ) - assert find_joins(optimized, "Semi") + assert find_hints(optimized) + + +def test_filter_hints_are_discarded_during_lowering( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + root = translate_query(simple_query, engine) + config = ConfigOptions.from_polars_engine(engine) + + lowering = lower_ir_graph(root, config, StatsCollector()) + + assert find_hints(lowering.optimized) + assert not find_hints(lowering.lowered) + assert not find_joins(lowering.lowered, "Semi") def test_filter_pushdown_can_be_disabled( @@ -207,9 +233,9 @@ def test_nullable_join_keys_preserve_results( config, ) - semi_joins = find_joins(optimized, "Semi") - assert semi_joins - assert all(join.options[1] is nulls_equal for join in semi_joins) + hints = find_hints(optimized) + assert hints + assert all(hint.nulls_equal is nulls_equal for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -239,8 +265,8 @@ def test_prefilter_does_not_move_below_distinct_on_non_subset_column( config, ) - semis = find_joins(optimized, "Semi") - assert any(isinstance(semi.children[0], Distinct) for semi in semis) + hints = find_hints(optimized) + assert any(isinstance(hint.children[0], Distinct) for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -267,7 +293,7 @@ def test_no_simple_filter_pushdown_when_domain_is_not_selective( assert decision == Decision(reason="no_profitable_domain") assert optimized is root - assert not find_joins(optimized, "Semi") + assert not find_hints(optimized) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -329,12 +355,12 @@ def test_composite_filter_pushdown_constrains_domain_first( assert decision.reason == "applied" assert isinstance(decision.candidate, CompositeCandidate) - semis = find_joins(optimized, "Semi") + hints = find_hints(optimized) assert isinstance(optimized, Join) assert optimized.options[0] == "Inner" assert optimized.children[1] is supplier_ir - assert any(semi.children[0] is supplier_ir for semi in semis) - assert any(semi.children[0] is lineitem_ir for semi in semis) + assert any(hint.children[0] is supplier_ir for hint in hints) + assert any(hint.children[0] is lineitem_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -388,16 +414,16 @@ def test_prefilter_uses_cheaper_source_domain_and_skips_expensive_domain( supplier_ir = dataframe_scan(root, "s_suppkey") lineitem_ir = dataframe_scan(root, "l_orderkey") orders_ir = dataframe_scan(root, "o_orderkey") - semis = find_joins(optimized, "Semi") - partkey_semis = [ - semi - for semi in semis - if semi.children[0] is lineitem_ir and join_key_names(semi) == ("l_partkey",) + hints = find_hints(optimized) + partkey_hints = [ + hint + for hint in hints + if hint.children[0] is lineitem_ir and hint_key_names(hint) == ("l_partkey",) ] - assert partkey_semis - assert not any(semi.children[0] is orders_ir for semi in semis) - assert contains_node(partkey_semis[0].children[1], part_ir) - assert not contains_node(partkey_semis[0].children[1], supplier_ir) + assert partkey_hints + assert not any(hint.children[0] is orders_ir for hint in hints) + assert contains_node(partkey_hints[0].children[1], part_ir) + assert not contains_node(partkey_hints[0].children[1], supplier_ir) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -442,13 +468,11 @@ def test_source_only_domain_does_not_stack_on_prefiltered_source( ) lineitem_ir = dataframe_scan(root, "l_partkey") - lineitem_semis = [ - semi - for semi in find_joins(optimized, "Semi") - if semi.children[0] is lineitem_ir + lineitem_hints = [ + hint for hint in find_hints(optimized) if hint.children[0] is lineitem_ir ] - assert any(join_key_names(semi) == ("l_partkey",) for semi in lineitem_semis) - assert not any(join_key_names(semi) == ("l_orderkey",) for semi in lineitem_semis) + assert any(hint_key_names(hint) == ("l_partkey",) for hint in lineitem_hints) + assert not any(hint_key_names(hint) == ("l_orderkey",) for hint in lineitem_hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -496,13 +520,13 @@ def test_derived_selectivity_propagates_through_rewritten_children( ConfigOptions.from_polars_engine(engine), ) - semis = find_joins(optimized, "Semi") + hints = find_hints(optimized) expected_targets = { dataframe_scan(root, "n_nationkey"), dataframe_scan(root, "c_custkey"), dataframe_scan(root, "o_orderkey"), } - assert expected_targets <= {semi.children[0] for semi in semis} + assert expected_targets <= {hint.children[0] for hint in hints} assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -552,13 +576,10 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( lineitem_ir = dataframe_scan(root, "l_orderkey") orders_ir = dataframe_scan(root, "o_orderkey") - semis = find_joins(optimized, "Semi") - assert sum(semi.children[0] is lineitem_ir for semi in semis) == 1 - assert not any(semi.children[0] is orders_ir for semi in semis) - assert not any( - isinstance(semi.children[0], Join) and semi.children[0].options[0] == "Semi" - for semi in semis - ) + hints = find_hints(optimized) + assert sum(hint.children[0] is lineitem_ir for hint in hints) == 1 + assert not any(hint.children[0] is orders_ir for hint in hints) + assert not any(isinstance(hint.children[0], PushdownFilterHint) for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -612,9 +633,9 @@ def test_target_source_follows_join_key_through_rename( ConfigOptions.from_polars_engine(engine), ) - semis = find_joins(optimized, "Semi") - assert any(semi.children[0] is small_ir for semi in semis) - assert not any(semi.children[0] is big_ir for semi in semis) + hints = find_hints(optimized) + assert any(hint.children[0] is small_ir for hint in hints) + assert not any(hint.children[0] is big_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -661,14 +682,11 @@ def test_domain_source_follows_join_key_through_rename( ConfigOptions.from_polars_engine(engine), ) - semi = next( - semi for semi in find_joins(optimized, "Semi") if semi.children[0] is target_ir - ) - selected_domain = semi.children[1] + hint = next(hint for hint in find_hints(optimized) if hint.children[0] is target_ir) + selected_domain = hint.children[1] assert isinstance(selected_domain, Select) rewritten_domain_source = selected_domain.children[0] - assert isinstance(rewritten_domain_source, Join) - assert rewritten_domain_source.options[0] == "Semi" + assert isinstance(rewritten_domain_source, PushdownFilterHint) assert rewritten_domain_source.children[0] is domain_source_ir assert rewritten_domain_source.children[0] is not renamed_unrelated_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -721,7 +739,7 @@ def test_composite_domain_columns_do_not_reconverge_after_join( facts = analyze_plan(joined, StatsCollector()) producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) - candidates = tuple(semijoin_pushdown_candidates(facts, joined, "value")) + candidates = tuple(filter_hint_pushdown_candidates(facts, joined, "value")) assert candidates[0] == (ColumnRef(joined, "value"), ()) assert len(candidates) >= 2 assert all(path == (0,) * len(path) for _, path in candidates[1:]) @@ -802,7 +820,7 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: facts = analyze_plan(root, stats) lineage = facts.column_lineages[ColumnRef(sliced, "target_key")] assert lineage.column == ColumnRef(sliced, "target_key") - assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( + assert tuple(filter_hint_pushdown_candidates(facts, sliced, "target_key")) == ( (ColumnRef(sliced, "target_key"), ()), ) @@ -812,9 +830,9 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: ConfigOptions.from_polars_engine(engine), ) - semis = find_joins(optimized, "Semi") - assert any(semi.children[0] is sliced for semi in semis) - assert not any(semi.children[0] is target_ir for semi in semis) + hints = find_hints(optimized) + assert any(hint.children[0] is sliced for hint in hints) + assert not any(hint.children[0] is target_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -861,10 +879,10 @@ def test_target_replacement_does_not_rewrite_shared_domain_side( filtered, unfiltered_domain = optimized.children assert unfiltered_domain is domain_ir assert domain_ir.children[0] is shared_ir - semis = find_joins(filtered, "Semi") - assert len(semis) == 1 - assert semis[0].children[0] is shared_ir - assert not find_joins(unfiltered_domain, "Semi") + hints = find_hints(filtered) + assert len(hints) == 1 + assert hints[0].children[0] is shared_ir + assert not find_hints(unfiltered_domain) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -913,12 +931,12 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( assert isinstance(rewritten_self_join, Join) filtered, unfiltered = rewritten_self_join.children assert unfiltered is source_ir - filtered_semis = find_joins(filtered, "Semi") - assert len(filtered_semis) == 1 - assert not find_joins(unfiltered, "Semi") + filtered_hints = find_hints(filtered) + assert len(filtered_hints) == 1 + assert not find_hints(unfiltered) # The shared node is a valid insertion point, but its children are not: - # Only this consumer should be wrapped by the semi-join. - assert filtered_semis[0].children[0] is source_ir + # Only this consumer should be wrapped by the filter hint. + assert filtered_hints[0].children[0] is source_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -967,8 +985,8 @@ def test_internal_prefilter_rewrites_shared_subplan_once( rewritten_left, rewritten_right = optimized.children assert rewritten_left is rewritten_right assert rewritten_left is not original_shared - (internal_semi,) = find_joins(rewritten_left, "Semi") - assert internal_semi.children[0] is target_ir + (internal_hint,) = find_hints(rewritten_left) + assert internal_hint.children[0] is target_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -1006,5 +1024,5 @@ def test_no_filter_pushdown_for_unsupported_joins( ) assert optimized is root - assert not find_joins(optimized, "Semi") + assert not find_hints(optimized) assert_gpu_result_equal(query, engine=engine, check_row_order=False) From b99df370a477657dfd6d3f845c6a7ec089f8269d Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 4 Aug 2026 12:30:35 +0100 Subject: [PATCH 02/17] Apply adaptive prefilters to direct join inputs --- python/cudf_polars/cudf_polars/dsl/ir.py | 7 +- .../cudf_polars/streaming/actor_graph/join.py | 629 ++++++++++++++---- .../streaming/actor_graph/join_bindings.py | 88 +++ .../streaming/actor_graph/prefilter.py | 360 ++++++++++ .../streaming/actor_graph/utils.py | 29 +- .../cudf_polars/streaming/explain.py | 41 +- .../cudf_polars/streaming/filter_hint.py | 80 ++- .../cudf_polars/cudf_polars/streaming/join.py | 121 +++- .../cudf_polars/cudf_polars/utils/config.py | 22 +- .../tests/streaming/test_explain.py | 33 +- .../streaming/test_join_filter_pushdown.py | 29 +- .../tests/streaming/test_tracing.py | 111 ++++ python/cudf_polars/tests/test_config.py | 27 +- 13 files changed, 1410 insertions(+), 167 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py create mode 100644 python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 559777f43c7d..aec9fc84c55e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2543,7 +2543,12 @@ class Join(IR): """A join of two dataframes.""" __slots__ = ("left_on", "options", "right_on") - _non_child = ("schema", "left_on", "right_on", "options") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "left_on", + "right_on", + "options", + ) _n_non_child_args = 3 left_on: tuple[expr.NamedExpr, ...] """List of expressions used as keys in the left frame.""" diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index d944847e621b..cca40f3d5fec 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal +from cudf_streaming import BloomFilter, CardinalityEstimator from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -16,6 +17,7 @@ TableChunk, make_table_chunks_available_or_wait, ) +from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED from rapidsmpf.memory.memory_reservation import opaque_memory_usage from rapidsmpf.streaming.core.actor import define_actor from rapidsmpf.streaming.core.memory_reserve_or_wait import ( @@ -23,7 +25,7 @@ ) from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IR, Join +from cudf_polars.dsl.ir import IR, Join, Projection from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import ( AllGatherManager, @@ -35,8 +37,14 @@ from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) +from cudf_polars.streaming.actor_graph.join_bindings import bind_join_inputs from cudf_polars.streaming.actor_graph.nodes import default_node_multi -from cudf_polars.streaming.actor_graph.tracing import send_chunk +from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterExecution, + choose_prefilter, + count_rows_passthrough, +) +from cudf_polars.streaming.actor_graph.tracing import LOG_TRACES, send_chunk from cudf_polars.streaming.actor_graph.utils import ( CUDF_ROW_LIMIT, MAX_ROWS_PER_PARTITION, @@ -56,11 +64,12 @@ send_metadata, shutdown_on_error, ) +from cudf_polars.streaming.filter_hint import JoinWithPrefilter from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: - from collections.abc import MutableMapping + from collections.abc import Iterable, MutableMapping from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel @@ -69,8 +78,14 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.join_bindings import ( + BoundPrefilter, + JoinBindings, + JoinInputBinding, + ) from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo + from cudf_polars.streaming.filter_hint import JoinSide from cudf_polars.utils.config import StreamingExecutor @@ -98,6 +113,53 @@ class JoinStrategy: """The key expressions for the right side. Only used for shuffle joins.""" +@dataclass(frozen=True, slots=True) +class JoinCollectiveIds: + """Named collective-ID slots reserved for a dynamic join.""" + + size_estimate: int + left_redistribution: int + right_redistribution: int + + @classmethod + def from_reserved(cls, collective_ids: list[int]) -> JoinCollectiveIds: + """Construct the named slots from IDs reserved for a dynamic join.""" + if len(collective_ids) < 3: + raise ValueError( + "Dynamic join requires 3 reserved collective IDs " + "(allgather + left shuffle + right shuffle); got " + f"{len(collective_ids)} for this Join. " + "Ensure ReserveOpIDs is run with dynamic_planning enabled." + ) + return cls(*collective_ids[:3]) + + @property + def cardinality_tags(self) -> tuple[int, int]: + """Tags available for concurrent prefilter cardinality estimates.""" + return (self.size_estimate, self.left_redistribution) + + @property + def broadcast(self) -> int: + """ID used by a broadcast join after size estimation completes.""" + return self.left_redistribution + + def shuffle(self, side: JoinSide) -> int: + """Return the collective ID for one shuffle input.""" + if side == "left": + return self.left_redistribution + return self.right_redistribution + + def prefilter(self, strategy: JoinStrategy, target_side: JoinSide) -> int: + """Return the subsequent join collective reused by a prefilter.""" + if strategy.broadcast_side is not None: + if target_side != strategy.broadcast_side: + raise ValueError( + "Only the broadcast input can have an active prefilter" + ) + return self.broadcast + return self.shuffle(target_side) + + @define_actor() async def broadcast_join_actor( context: Context, @@ -154,7 +216,7 @@ async def broadcast_join_actor( ch_left, ch_right, JoinStrategy(broadcast_side=broadcast_side), - [collective_id], + collective_id, target_partition_size, tracer=tracer, ) @@ -249,7 +311,7 @@ async def _broadcast_join_large_chunk( broadcast_side: Literal["left", "right"], *, tracer: ActorTracer | None, -) -> None: +) -> int: """Join one large-side chunk with the small DataFrame(s) and send the result.""" large_df = chunk_to_frame(large_chunk, large_child) large_chunk_size = large_chunk.data_alloc_size() @@ -280,8 +342,10 @@ async def _broadcast_join_large_chunk( output_chunk = TableChunk.from_pylibcudf_table( df.table, df.stream, exclusive_view=True, br=context.br() ) + output_rows = output_chunk.shape[0] await send_chunk(context, ch_out, output_chunk, seq_num, tracer=tracer) del df, large_df + return output_rows async def _broadcast_join( @@ -293,27 +357,27 @@ async def _broadcast_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: JoinStrategy, - collective_ids: list[int], - target_partition_size: int, + collective_id: int, + target_partition_size: int | None, *, tracer: ActorTracer | None, + trace_stats: dict[str, Any] | None = None, ) -> None: """ Execute a broadcast join after initial sampling. The small side is gathered (if not already duplicated) and concatenated into a single DataFrame, then joined with each chunk from the large side. - Pops one collective ID from collective_ids for allgather when needed. + Uses ``collective_id`` for the allgather when needed. """ left_metadata, right_metadata = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), ) - collective_id = collective_ids.pop(0) if collective_ids else 0 broadcast_side = strategy.broadcast_side assert broadcast_side is not None - left, right = ir.children + left, right = ir.children[:2] if tracer is not None: tracer.decision = f"broadcast_{broadcast_side}" @@ -355,8 +419,6 @@ async def _broadcast_join( partitioning=partitioning, duplicated=output_duplicated, ) - await send_metadata(ch_out, context, metadata_out) - small_dfs, small_size = await _collect_small_side_for_broadcast( context, comm, @@ -368,17 +430,26 @@ async def _broadcast_join( concat_size_limit=(target_partition_size if ir.options[0] == "Inner" else None), ) + # Publish output metadata only once the broadcast-side collective has + # completed. Besides making the data channel ready when advertised, this + # permits a consumer to reuse the collective ID after receiving metadata. + await send_metadata(ch_out, context, metadata_out) + + input_rows = 0 + output_rows = 0 while (msg := await large_ch.recv(context)) is not None: - await _broadcast_join_large_chunk( + large_chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + input_rows += large_chunk.shape[0] + output_rows += await _broadcast_join_large_chunk( context, ir, ir_context, ch_out, small_dfs, small_child, - TableChunk.from_message(msg, br=context.br()).make_available_and_spill( - context.br(), allow_overbooking=True - ), + large_chunk, large_child, msg.sequence_number, small_size, @@ -386,9 +457,170 @@ async def _broadcast_join( tracer=tracer, ) + if trace_stats is not None: + trace_stats["input_rows"] = input_rows + trace_stats["output_rows"] = output_rows await ch_out.drain(context) +def add_bloom_prefilter( + context: Context, + comm: Communicator, + bloom_bytes: int, + execution: PrefilterExecution, + target_indices: Iterable[int], + ch_domain_keys: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the channels and actors for an approximate Bloom prefilter.""" + bloom = BloomFilter( + context, + comm, + LIBCUDF_DEFAULT_HASH_SEED, + bloom_bytes, + ) + ch_filter = context.create_channel() + execution.add_channel(ch_filter) + execution.add_task( + bloom.build( + context, + ch_domain_keys, + ch_filter, + collective_id, + ) + ) + ch_apply_input = ch_target + ch_apply_output = ch_filtered + if trace_stats is not None: + ch_counted_input: Channel[TableChunk] = context.create_channel() + ch_raw_output: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_counted_input) + execution.add_channel(ch_raw_output) + execution.add_task( + count_rows_passthrough( + context, + ch_target, + ch_counted_input, + trace_stats, + "input_rows", + ) + ) + execution.add_task( + count_rows_passthrough( + context, + ch_raw_output, + ch_filtered, + trace_stats, + "output_rows", + ) + ) + ch_apply_input = ch_counted_input + ch_apply_output = ch_raw_output + execution.add_task( + bloom.apply( + context, + ch_filter, + ch_apply_input, + ch_apply_output, + target_indices, + ) + ) + + +def make_prefilter_execution( + context: Context, + comm: Communicator, + ir_context: IRExecutionContext, + strategy: JoinStrategy, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + bindings: JoinBindings, + collective_ids: JoinCollectiveIds, +) -> PrefilterExecution: + """Create the actors and channels that realize selected prefilters.""" + execution = PrefilterExecution(context, ch_left, ch_right) + + # Prepare every required domain before connecting target-side filters. This + # is important for opposing direct filters: each filter must consume the + # replay produced while the same input's keys are copied for the other one. + for bound in bindings.prefilters: + decision = bound.decision + if decision is None: + raise ValueError("Join prefilter has no runtime decision") + prefilter = bound.prefilter + if decision.method == "skip": + continue + + indices = names_to_indices(prefilter.domain_on, bound.domain.node.schema) + bound.key_channel = execution.buffer_domain(prefilter.domain_side, indices) + + for bound in bindings.prefilters: + decision = bound.decision + assert decision is not None + if decision.method == "skip": + continue + prefilter = bound.prefilter + ch_domain_keys = bound.key_channel + assert ch_domain_keys is not None + target_side = prefilter.target_side + target = bound.target.node + ch_target = execution.join_inputs[target_side] + ch_filtered: Channel[TableChunk] = context.create_channel() + trace_stats = bound.trace + + collective_id = collective_ids.prefilter(strategy, target_side) + if decision.method == "bloom": + assert decision.bloom_bytes is not None + add_bloom_prefilter( + context, + comm, + decision.bloom_bytes, + execution, + names_to_indices(prefilter.target_on, target.schema), + ch_domain_keys, + ch_target, + ch_filtered, + collective_id, + trace_stats, + ) + else: + assert decision.method == "broadcast_semi_join" + domain_schema = {key.name: key.value.dtype for key in prefilter.domain_on} + if len(domain_schema) != len(prefilter.domain_on): + raise ValueError("Broadcast semi-join keys must have unique names") + projected_domain = Projection(domain_schema, bound.domain.node) + semi_join = Join( + target.schema, + prefilter.target_on, + prefilter.domain_on, + ("Semi", prefilter.nulls_equal, None, "", False, "none"), + target, + projected_domain, + ) + execution.add_task( + _broadcast_join( + context, + comm, + semi_join, + ir_context, + ch_filtered, + ch_target, + ch_domain_keys, + JoinStrategy(broadcast_side="right"), + collective_id, + target_partition_size=None, + tracer=None, + trace_stats=trace_stats, + ) + ) + execution.replace_join_input(target_side, ch_filtered) + + return execution + + def _get_key_indices( ir: Join, n_partitioned_keys: int | None, @@ -399,7 +631,7 @@ def _get_key_indices( tuple[NamedExpr, ...], tuple[NamedExpr, ...], ]: - left, right = ir.children + left, right = ir.children[:2] n_keys = n_partitioned_keys if n_partitioned_keys is not None else len(ir.left_on) left_keys = ir.left_on[:n_keys] right_keys = ir.right_on[:n_keys] @@ -438,7 +670,7 @@ async def _join_chunks( recv_metadata(ch_right, context), ) - left, right = ir.children + left, right = ir.children[:2] while True: left_msg, right_msg = await gather_in_task_group( ch_left.recv(context), ch_right.recv(context) @@ -534,7 +766,8 @@ async def _shuffle_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: JoinStrategy, - collective_ids: list[int], + left_collective_id: int, + right_collective_id: int, *, tracer: ActorTracer | None, ) -> None: @@ -577,7 +810,7 @@ async def _shuffle_join( strategy.left_keys, ir.children[0].schema, strategy.shuffle_modulus, - collective_ids.pop(0), + left_collective_id, ), _global_shuffle( context, @@ -588,7 +821,7 @@ async def _shuffle_join( strategy.right_keys, ir.children[1].schema, strategy.shuffle_modulus, - collective_ids.pop(0), + right_collective_id, ), _join_chunks( context, @@ -650,56 +883,41 @@ def _num_indices(partitioning: NormalizedPartitioning) -> int: ) -async def _aggregate_estimates( +async def aggregate_estimates( context: Context, comm: Communicator, - left_sample: TableSizeStats, - right_sample: TableSizeStats, - collective_ids: list[int], -) -> tuple[TableSizeStats, TableSizeStats]: + samples: tuple[TableSizeStats, ...], + collective_id: int, +) -> tuple[TableSizeStats, ...]: """Aggregate table-size and row estimates across ranks.""" - # AllGather size, row, and chunk count estimates across ranks + # AllGather size, row, chunk count, and completeness estimates across ranks. totals = await allgather_reduce( context, comm, - collective_ids.pop(0), - left_sample.total_size, - right_sample.total_size, - left_sample.total_rows, - right_sample.total_rows, - left_sample.total_chunks, - right_sample.total_chunks, - int(left_sample.is_complete), - int(right_sample.is_complete), - ) - ( - left_total, - right_total, - left_total_rows, - right_total_rows, - left_total_chunks, - right_total_chunks, - left_complete_count, - right_complete_count, - ) = totals - - new_left_sample = TableSizeStats( - chunks=left_sample.chunks, - total_size=left_total, - total_rows=left_total_rows, - total_chunks=left_total_chunks, - is_complete=left_complete_count == comm.nranks, - cardinality=left_sample.cardinality, + collective_id, + *( + value + for sample in samples + for value in ( + sample.total_size, + sample.total_rows, + sample.total_chunks, + int(sample.is_complete), + ) + ), ) - new_right_sample = TableSizeStats( - chunks=right_sample.chunks, - total_size=right_total, - total_rows=right_total_rows, - total_chunks=right_total_chunks, - is_complete=right_complete_count == comm.nranks, - cardinality=right_sample.cardinality, + totals_iter = iter(totals) + return tuple( + TableSizeStats( + chunks=sample.chunks, + total_size=next(totals_iter), + total_rows=next(totals_iter), + total_chunks=next(totals_iter), + is_complete=next(totals_iter) == comm.nranks, + cardinality=sample.cardinality, + ) + for sample in samples ) - return new_left_sample, new_right_sample def _choose_strategy_from_samples( @@ -849,78 +1067,197 @@ def _modulus(partitioning: NormalizedPartitioning) -> int | None: return max(large, min_shuffle_modulus) -async def _choose_strategy( +def join_input_requires_redistribution( + strategy: JoinStrategy, + side: Literal["left", "right"], + partitioning: NormalizedPartitioning, + metadata: ChannelMetadata, +) -> bool: + """Return whether the join strategy redistributes an input side.""" + if strategy.broadcast_side is not None: + return side == strategy.broadcast_side and not metadata.duplicated + + indices = strategy.left_indices if side == "left" else strategy.right_indices + if not indices: + return True + desired = HashScheme(indices, strategy.shuffle_modulus) + return not ( + partitioning.inter_rank_scheme == desired + and partitioning.local_scheme == "inherit" + ) + + +def choose_prefilters( + bindings: JoinBindings, + strategy: JoinStrategy, + left_partitioning: NormalizedPartitioning, + right_partitioning: NormalizedPartitioning, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> None: + """Choose strategies for optional join prefilters.""" + partitionings = { + "left": left_partitioning, + "right": right_partitioning, + } + for bound in bindings.prefilters: + target = bound.target.sample + if target is None: + raise ValueError("Join target has not been sampled") + target_side = bound.prefilter.target_side + bound.decision = choose_prefilter( + bound.prefilter, + target, + bound.domain.sample, + target_requires_redistribution=join_input_requires_redistribution( + strategy, + target_side, + partitionings[target_side], + bound.target.metadata, + ), + broadcast_limit=broadcast_limit, + bloom_filter_max_size=bloom_filter_max_size, + ) + + +async def sample_input( + context: Context, + comm: Communicator, + input_: JoinInputBinding, + prefilter: BoundPrefilter | None, + sample_chunk_count: int, + target_partition_size: int, +) -> TableSizeStats: + """Sample one join-planning input and optionally estimate cardinality.""" + if prefilter is None: + cardinality_estimator = None + cardinality_columns: tuple[int, ...] = () + else: + cardinality_estimator = CardinalityEstimator( + context, + comm, + tag=prefilter.cardinality_tag, + ) + cardinality_columns = names_to_indices( + prefilter.prefilter.domain_on, + input_.node.schema, + ) + assert len(cardinality_columns) == len(prefilter.prefilter.domain_on), ( + "Prefilter domain keys must be columns" + ) + + return await _sample_chunks( + context, + input_.channel, + sample_chunk_count, + target_partition_size, + input_.metadata.local_count, + cardinality_estimator=cardinality_estimator, + cardinality_columns=cardinality_columns, + ) + + +async def collect_join_samples( + context: Context, + comm: Communicator, + bindings: JoinBindings, + sample_chunk_count: int, + target_partition_size: int, + collective_id: int, +) -> None: + """Sample join inputs and attach estimates to their runtime bindings.""" + inputs = [bindings.left, bindings.right] + + sampling_inputs = [] + for input_ in inputs: + domain_prefilters = [ + bound for bound in bindings.prefilters if bound.domain is input_ + ] + if len(domain_prefilters) > 1: + raise ValueError("One join input cannot provide multiple prefilter domains") + sampling_inputs.append( + (input_, domain_prefilters[0] if domain_prefilters else None) + ) + + local_samples = await gather_in_task_group( + *( + sample_input( + context, + comm, + input_, + prefilter, + sample_chunk_count, + target_partition_size, + ) + for input_, prefilter in sampling_inputs + ) + ) + samples = await aggregate_estimates( + context, + comm, + tuple(local_samples), + collective_id, + ) + for input_, sample in zip(inputs, samples, strict=True): + input_.sample = sample + + +async def choose_strategy( context: Context, comm: Communicator, ir: Join, - ch_left: Channel[TableChunk], - ch_right: Channel[TableChunk], - left_metadata: ChannelMetadata, - right_metadata: ChannelMetadata, + bindings: JoinBindings, executor: StreamingExecutor, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, *, tracer: ActorTracer | None, -) -> tuple[TableSizeStats, TableSizeStats, JoinStrategy]: - """Sample both sides, aggregate estimates, and choose broadcast vs shuffle.""" +) -> JoinStrategy: + """Collect any required samples and choose broadcast vs shuffle.""" + left, right = ir.children[:2] + left_metadata = bindings.left.metadata + right_metadata = bindings.right.metadata nranks = comm.nranks left_partitioning = NormalizedPartitioning.from_keys( left_metadata.partitioning, nranks, - keys=names_to_indices(ir.left_on, ir.children[0].schema, concrete_prefix=True), + keys=names_to_indices(ir.left_on, left.schema, concrete_prefix=True), ) right_partitioning = NormalizedPartitioning.from_keys( right_metadata.partitioning, nranks, - keys=names_to_indices(ir.right_on, ir.children[1].schema, concrete_prefix=True), + keys=names_to_indices(ir.right_on, right.schema, concrete_prefix=True), ) - hash_chunkwise = isinstance( left_partitioning.inter_rank_scheme, HashScheme ) and isinstance(right_partitioning.inter_rank_scheme, HashScheme) - if hash_chunkwise and left_partitioning.is_aligned_with( + chunkwise = hash_chunkwise and left_partitioning.is_aligned_with( right_partitioning, context.br() - ): - # We can use a chunkwise join - chunkwise = True - left_sample = TableSizeStats( + ) + + if chunkwise: + bindings.left.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=left_metadata.local_count, ) - right_sample = TableSizeStats( + bindings.right.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=right_metadata.local_count, ) else: - # Need to shuffle or broadcast - Use sampled data to choose a strategy - chunkwise = False assert executor.dynamic_planning is not None - sample_chunk_count = executor.dynamic_planning.sample_chunk_count - target_partition_size = executor.target_partition_size - left_sample, right_sample = await gather_in_task_group( - _sample_chunks( - context, - ch_left, - sample_chunk_count, - target_partition_size, - left_metadata.local_count, - ), - _sample_chunks( - context, - ch_right, - sample_chunk_count, - target_partition_size, - right_metadata.local_count, - ), - ) - left_sample, right_sample = await _aggregate_estimates( + await collect_join_samples( context, comm, - left_sample, - right_sample, - collective_ids, + bindings, + executor.dynamic_planning.sample_chunk_count, + executor.target_partition_size, + collective_ids.size_estimate, ) + left_sample = bindings.left.sample + right_sample = bindings.right.sample + if left_sample is None or right_sample is None: + raise ValueError("Join inputs have not been sampled") strategy = _choose_strategy_from_samples( comm, ir, @@ -934,8 +1271,16 @@ async def _choose_strategy( chunkwise=chunkwise, tracer=tracer, ) - - return left_sample, right_sample, strategy + if executor.join_filter_pushdown is not None: + choose_prefilters( + bindings, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + executor.join_filter_pushdown.bloom_filter_max_size, + ) + return strategy @define_actor() @@ -948,7 +1293,7 @@ async def join_actor( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], executor: StreamingExecutor, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, ) -> None: """ Dynamic Join actor that selects the best strategy at runtime. @@ -991,24 +1336,55 @@ async def join_actor( recv_metadata(ch_right, context), ) - left_sample, right_sample, strategy = await _choose_strategy( - context, - comm, + bindings = bind_join_inputs( ir, ch_left, ch_right, left_metadata, right_metadata, + collective_ids.cardinality_tags, + ) + + strategy = await choose_strategy( + context, + comm, + ir, + bindings, executor, collective_ids, tracer=tracer, ) + prefilter_traces = [] + for bound in bindings.prefilters: + if bound.decision is None: + raise ValueError("Join prefilter has no runtime decision") + trace = bound.decision.trace(bound.prefilter) + prefilter_traces.append(trace) + if LOG_TRACES: + bound.trace = trace + if tracer is not None and prefilter_traces: + tracer.set_extra("join_prefilters", prefilter_traces) + left_sample = bindings.left.sample + right_sample = bindings.right.sample + if left_sample is None or right_sample is None: + raise ValueError("Join inputs have not been sampled") ch_left_replay = context.create_channel() ch_right_replay = context.create_channel() + prefilter_execution = make_prefilter_execution( + context, + comm, + ir_context, + strategy, + ch_left_replay, + ch_right_replay, + bindings, + collective_ids, + ) async with shutdown_on_error( context, ch_left_replay, ch_right_replay, + *prefilter_execution.channels, trace_ir=ir, ir_context=ir_context, ): @@ -1029,9 +1405,10 @@ async def join_actor( right_metadata, trace_ir=ir, ), + *prefilter_execution.tasks, ] - ch_left = ch_left_replay - ch_right = ch_right_replay + ch_left = prefilter_execution.left + ch_right = prefilter_execution.right if strategy.broadcast_side is not None: actor_tasks.append( @@ -1044,7 +1421,7 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.broadcast, executor.target_partition_size, tracer=tracer, ) @@ -1060,7 +1437,8 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.shuffle("left"), + collective_ids.shuffle("right"), tracer=tracer, ) ) @@ -1073,7 +1451,7 @@ def _use_pwise_join( ir: Join, ) -> bool: """Whether to use a static-planning partition-wise join.""" - left, right = ir.children + left, right = ir.children[:2] output_count = partition_info[ir].count if ( output_count == 1 @@ -1099,8 +1477,9 @@ def _use_pwise_join( @generate_ir_sub_network.register(Join) +@generate_ir_sub_network.register(JoinWithPrefilter) def _( - ir: Join, rec: SubNetGenerator + ir: Join | JoinWithPrefilter, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: # Join operation. left, right = ir.children @@ -1110,7 +1489,6 @@ def _( executor = rec.state["config_options"].executor pwise_join = _use_pwise_join(executor, partition_info, ir) - # Process children actors, channels = process_children(ir, rec) # Create output ChannelManager @@ -1140,16 +1518,13 @@ def _( and ir.options[0] in ("Inner", "Left", "Right", "Full", "Semi", "Anti") ): # Dynamic join - decide strategy at runtime - collective_ids = list(rec.state["collective_id_map"].get(ir, [])) - # Join uses up to 3 collective IDs: allgather, left shuffle, and - # right shuffle. - if len(collective_ids) < 3: - raise ValueError( - "Dynamic join requires 3 reserved collective IDs " - "(allgather + left shuffle + right shuffle); got " - f"{len(collective_ids)} for this Join. " - "Ensure ReserveOpIDs is run with dynamic_planning enabled." - ) + collective_ids = JoinCollectiveIds.from_reserved( + rec.state["collective_id_map"].get(ir, []) + ) + # Join uses up to 3 collective IDs. Cardinality allreduces complete + # before the size allgather and join collectives. Runtime prefilters + # reuse the collective ID of the target-side join redistribution, with + # their filtered output channel providing the ordering barrier. actors[ir] = [ join_actor( rec.state["context"], diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py new file mode 100644 index 000000000000..5b52fa874830 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime bindings for optional join prefilters.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cudf_polars.streaming.filter_hint import JoinWithPrefilter + +if TYPE_CHECKING: + from typing import Any + + from cudf_streaming.channel_metadata import ChannelMetadata + from cudf_streaming.table_chunk import TableChunk + from rapidsmpf.streaming.core.channel import Channel + + from cudf_polars.dsl.ir import IR, Join + from cudf_polars.streaming.actor_graph.prefilter import PrefilterDecision + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.streaming.filter_hint import Prefilter + + +@dataclass(slots=True) +class JoinInputBinding: + """Concrete runtime resources for one input to a dynamic join.""" + + node: IR + channel: Channel[TableChunk] + metadata: ChannelMetadata + sample: TableSizeStats | None = None + + +@dataclass(slots=True) +class BoundPrefilter: + """A logical prefilter bound to its concrete runtime inputs.""" + + prefilter: Prefilter + target: JoinInputBinding + domain: JoinInputBinding + cardinality_tag: int + decision: PrefilterDecision | None = None + key_channel: Channel[TableChunk] | None = None + trace: dict[str, Any] | None = None + + +@dataclass(frozen=True, slots=True) +class JoinBindings: + """Concrete runtime inputs and optional prefilters for a dynamic join.""" + + left: JoinInputBinding + right: JoinInputBinding + prefilters: tuple[BoundPrefilter, ...] = () + + +def bind_join_inputs( + ir: Join, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + left_metadata: ChannelMetadata, + right_metadata: ChannelMetadata, + cardinality_tags: tuple[int, ...], +) -> JoinBindings: + """Bind logical join inputs and prefilters to their runtime resources.""" + left = JoinInputBinding(ir.children[0], ch_left, left_metadata) + right = JoinInputBinding(ir.children[1], ch_right, right_metadata) + if not isinstance(ir, JoinWithPrefilter): + return JoinBindings(left, right) + + if len(cardinality_tags) < len(ir.prefilters): + raise ValueError("Each join prefilter requires a cardinality collective ID") + + sides = {"left": left, "right": right} + cardinality_tags_iter = iter(cardinality_tags) + prefilters = [] + for prefilter in ir.prefilters: + target = sides[prefilter.target_side] + domain = sides[prefilter.domain_side] + prefilters.append( + BoundPrefilter( + prefilter, + target, + domain, + cardinality_tag=next(cardinality_tags_iter), + ) + ) + return JoinBindings(left, right, tuple(prefilters)) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py new file mode 100644 index 000000000000..683203bcf35b --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime planning helpers for optional prefilters.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal + +import pylibcudf as plc +from cudf_streaming import BloomFilter +from cudf_streaming.channel_metadata import ChannelMetadata +from cudf_streaming.table_chunk import TableChunk +from rapidsmpf.streaming.core.message import Message + +from cudf_polars.streaming.actor_graph.utils import ( + ChunkStore, + recv_metadata, + send_metadata, + shutdown_channels_on_error, +) +from cudf_polars.streaming.filter_hint import JoinInputPrefilter + +if TYPE_CHECKING: + from collections.abc import Coroutine, Iterable, Sequence + + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.containers import DataType + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.streaming.filter_hint import JoinSide, Prefilter + + +def estimate_bytes(dtypes: Sequence[DataType], row_count: int) -> int | None: + """ + Estimate the byte count of a table containing the given datatypes. + + Parameters + ---------- + dtypes + Types of columns in the table. + row_count + Estimated total number of rows. + + Returns + ------- + Estimated table size in bytes, or ``None`` if any dtype is not fixed width. + """ + if not all(plc.traits.is_fixed_width(dtype.plc_type) for dtype in dtypes): + return None + + return int( + # Just assume everything has a validity mask + row_count * sum(plc.types.size_of(dtype.plc_type) + 1 / 8 for dtype in dtypes) + ) + + +@dataclass(frozen=True, slots=True) +class PrefilterDecision: + """Runtime decision for one optional prefilter.""" + + method: Literal["skip", "bloom", "broadcast_semi_join"] + reason: str + target_bytes: int + domain_rows: int | None + estimated_cardinality: int | None = None + bloom_bytes: int | None = None + exact_bytes: int | None = None + + def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: + """Return serializable actor-trace information.""" + result: dict[str, str | int | None] = { + "target_side": prefilter.target_side, + "method": self.method, + "reason": self.reason, + "target_bytes": self.target_bytes, + "domain_rows": self.domain_rows, + "estimated_cardinality": self.estimated_cardinality, + "bloom_bytes": self.bloom_bytes, + "exact_bytes": self.exact_bytes, + } + result["domain_side"] = prefilter.domain_side + return result + + +def project_key_chunk( + context: Context, chunk: TableChunk, indices: Iterable[int] +) -> TableChunk: + """Copy selected columns into an owning key chunk.""" + columns = chunk.table_view().columns() + key_table = plc.Table([columns[index] for index in indices]).copy( + stream=chunk.stream, mr=context.br().device_mr + ) + return TableChunk.from_pylibcudf_table( + key_table, + chunk.stream, + exclusive_view=True, + br=context.br(), + ) + + +async def buffer_and_project_keys( + context: Context, + ch_in: Channel[TableChunk], + ch_keys: Channel[TableChunk], + ch_replay: Channel[TableChunk], + indices: Iterable[int], +) -> None: + """ + Project owning key chunks while spill-buffering an input for replay. + + The key channel is produced in full before replay begins. Its consumer must + therefore run concurrently with this coroutine. + """ + chunks = ChunkStore(context) + try: + async with shutdown_channels_on_error(context, ch_in, ch_keys, ch_replay): + metadata = await recv_metadata(ch_in, context) + key_metadata = ChannelMetadata( + local_count=metadata.local_count, + partitioning=None, + duplicated=metadata.duplicated, + ) + await send_metadata(ch_replay, context, metadata) + await send_metadata(ch_keys, context, key_metadata) + indices = tuple(indices) + while (msg := await ch_in.recv(context)) is not None: + sequence_number = msg.sequence_number + chunk = await TableChunk.from_message( + msg, br=context.br() + ).make_available_or_wait(context, net_memory_delta=0) + key_chunk = project_key_chunk(context, chunk, indices) + chunks.insert(Message(sequence_number, chunk)) + await ch_keys.send(context, Message(sequence_number, key_chunk)) + + await ch_keys.drain(context) + for msg in chunks: + await ch_replay.send(context, msg) + await ch_replay.drain(context) + finally: + chunks.clear() + + +async def count_rows_passthrough( + context: Context, + ch_in: Channel[TableChunk], + ch_out: Channel[TableChunk], + trace_stats: dict[str, Any], + row_count_key: str, +) -> None: + """Forward a table-chunk channel while recording its row count.""" + async with shutdown_channels_on_error(context, ch_in, ch_out): + metadata = await recv_metadata(ch_in, context) + await send_metadata(ch_out, context, metadata) + row_count = 0 + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message(msg, br=context.br()) + row_count += chunk.shape[0] + await ch_out.send(context, Message(msg.sequence_number, chunk)) + trace_stats[row_count_key] = row_count + await ch_out.drain(context) + + +class PrefilterExecution: + """Channels and actors used to apply prefilters before a join.""" + + def __init__( + self, + context: Context, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + ) -> None: + self.context = context + self.source_inputs = {"left": ch_left, "right": ch_right} + self.join_inputs = dict(self.source_inputs) + self.tasks: list[Coroutine[Any, Any, None]] = [] + self.channels: list[Channel[Any]] = [] + self.buffered_domains: set[JoinSide] = set() + + def buffer_domain( + self, + side: JoinSide, + indices: Iterable[int], + ) -> Channel[TableChunk]: + """Buffer one original input and return its owning key channel.""" + if side in self.buffered_domains: + raise ValueError(f"Join input {side!r} is already a prefilter domain") + + ch_keys: Channel[TableChunk] = self.context.create_channel() + ch_replay: Channel[TableChunk] = self.context.create_channel() + self.tasks.append( + buffer_and_project_keys( + self.context, + self.source_inputs[side], + ch_keys, + ch_replay, + indices, + ) + ) + self.channels.extend((ch_keys, ch_replay)) + self.join_inputs[side] = ch_replay + self.buffered_domains.add(side) + return ch_keys + + def replace_join_input( + self, + side: JoinSide, + channel: Channel[TableChunk], + ) -> None: + """Replace one join-facing input with a prefilter output channel.""" + self.join_inputs[side] = channel + self.channels.append(channel) + + def add_task(self, task: Coroutine[Any, Any, None]) -> None: + """Add an actor task to the prefilter execution.""" + self.tasks.append(task) + + def add_channel(self, channel: Channel[Any]) -> None: + """Register an auxiliary channel for shutdown on failure.""" + self.channels.append(channel) + + @property + def left(self) -> Channel[TableChunk]: + """Current left join input.""" + return self.join_inputs["left"] + + @property + def right(self) -> Channel[TableChunk]: + """Current right join input.""" + return self.join_inputs["right"] + + +def estimate_cardinality(stats: TableSizeStats) -> int | None: + """Extrapolate sampled distinct count to the estimated full row count.""" + if stats.total_rows == 0: + return 0 + if stats.cardinality is None or stats.cardinality.row_count == 0: + return None + return min( + stats.total_rows, + math.ceil( + stats.cardinality.distinct_count + * stats.total_rows + / stats.cardinality.row_count + ), + ) + + +def estimate_bloom_filter_bytes( + cardinality: int, + desired_false_positive_rate: float = 0.1, +) -> int: + """Estimate Bloom-filter bytes for the block-split policy.""" + if cardinality < 0: + raise ValueError("cardinality must be non-negative") + if not 0 < desired_false_positive_rate < 1: + raise ValueError("false_positive_rate must be between zero and one") + if cardinality == 0: + return 0 + # TODO: cuco could offer this as a static utility on the policy + # Then we wouldn't have to hardcode these magic numbers. + bits = ( + -8 # number of fingerprint bits + * cardinality + / math.log(1 - desired_false_positive_rate ** (1 / 8)) + ) + return math.ceil(bits / 8) + + +def choose_prefilter( + prefilter: Prefilter, + target: TableSizeStats, + domain: TableSizeStats | None, + *, + target_requires_redistribution: bool, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> PrefilterDecision: + """Choose whether and how to apply one prefilter.""" + domain_rows = None if domain is None else domain.total_rows + if not target_requires_redistribution: + return PrefilterDecision( + "skip", + "target_not_redistributed", + target.total_size, + domain_rows, + ) + if domain is None: + raise ValueError("A redistributed target requires domain statistics") + if ( + isinstance(prefilter, JoinInputPrefilter) + and prefilter.target_side == prefilter.domain_side + ): + return PrefilterDecision( + "skip", + "same_input", + target.total_size, + domain_rows, + ) + + cardinality = estimate_cardinality(domain) + if cardinality is None: + return PrefilterDecision( + "skip", + "missing_cardinality", + target.total_size, + domain.total_rows, + ) + if cardinality == 0: + return PrefilterDecision( + "skip", + "zero_cardinality", + target.total_size, + domain.total_rows, + estimated_cardinality=0, + bloom_bytes=0, + exact_bytes=0, + ) + + bloom_bytes = max( + 32, + BloomFilter.aligned_size(estimate_bloom_filter_bytes(cardinality)), + ) + exact_bytes = estimate_bytes( + tuple(key.value.dtype for key in prefilter.domain_on), + domain.total_rows, + ) + if bloom_bytes <= min(bloom_filter_max_size, target.total_size): + return PrefilterDecision( + "bloom", + "bloom_fits", + target.total_size, + domain.total_rows, + estimated_cardinality=cardinality, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) + if exact_bytes is not None and exact_bytes <= min( + broadcast_limit, target.total_size + ): + return PrefilterDecision( + "broadcast_semi_join", + "exact_domain_fits", + target.total_size, + domain.total_rows, + estimated_cardinality=cardinality, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) + return PrefilterDecision( + "skip", + "no_viable_filter", + target.total_size, + domain.total_rows, + estimated_cardinality=cardinality, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 1b7153954f2f..b8d6a47c2c8e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -168,7 +168,7 @@ def _keys_match( class ChunkStore: - """Ordered spillable buffer for TableChunk messages.""" + """Ordered spillable buffer for Messages.""" def __init__(self, ctx: Context) -> None: self._mids: deque[int] = deque() @@ -178,6 +178,12 @@ def __len__(self) -> int: """Return the number of messages in the store.""" return len(self._mids) + def clear(self) -> None: + """Discard all messages in the store.""" + for mid in self._mids: + self._store.extract(mid=mid) + self._mids.clear() + def insert(self, msg: Message) -> None: """Insert a message into the store.""" self._mids.append(self._store.insert(msg)) @@ -317,6 +323,7 @@ async def shutdown_on_error( record["row_count"] = tracer.row_count if tracer.decision is not None: record["decision"] = tracer.decision + record.update(tracer.extra) cudf_polars.dsl.tracing.log( "Streaming Actor", start=start, stop=stop, **record ) @@ -1229,19 +1236,23 @@ async def replay_buffered_channel( ch_in The buffered input channel. buffered_chunks - The buffered chunks to yield first. + The buffered chunks to yield first. The store is empty when this + coroutine exits, including on cancellation or error. metadata The metadata to send to the output channel. trace_ir The IR node to trace. Passed through to shutdown_on_error. """ - async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir): - await send_metadata(ch_out, context, metadata) - for msg in buffered_chunks: - await ch_out.send(context, msg) - while (msg := await ch_in.recv(context)) is not None: - await ch_out.send(context, msg) - await ch_out.drain(context) + try: + async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir): + await send_metadata(ch_out, context, metadata) + for msg in buffered_chunks: + await ch_out.send(context, msg) + while (msg := await ch_in.recv(context)) is not None: + await ch_out.send(context, msg) + await ch_out.drain(context) + finally: + buffered_chunks.clear() @dataclass(frozen=True) diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index b93a96d36ec3..053b79b1edb8 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -37,7 +37,10 @@ from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import IOPartitionFlavor -from cudf_polars.streaming.filter_hint import PushdownFilterHint +from cudf_polars.streaming.filter_hint import ( + JoinWithPrefilter, + PushdownFilterHint, +) from cudf_polars.streaming.io import StreamingScan, scan_partition_plan from cudf_polars.streaming.parallel import lower_ir_graph, optimize_with_stats from cudf_polars.streaming.shuffle import Shuffle @@ -54,6 +57,7 @@ from cudf_polars.dsl.expressions.base import Expr from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import PartitionInfo, StatsCollector + from cudf_polars.streaming.filter_hint import Prefilter @dataclasses.dataclass @@ -471,6 +475,18 @@ def _(ir: Join, *, offset: str = "") -> str: return _repr_header(offset, f"JOIN {ir.options[0]} {left_on} {right_on}", ir.schema) +@_repr_ir.register +def _(ir: JoinWithPrefilter, *, offset: str = "") -> str: + left_on = tuple(ne.name for ne in ir.left_on) + right_on = tuple(ne.name for ne in ir.right_on) + prefilters = tuple(type(prefilter).__name__ for prefilter in ir.prefilters) + return _repr_header( + offset, + f"JOIN {ir.options[0]} {left_on} {right_on} {prefilters=}", + ir.schema, + ) + + @_repr_ir.register def _(ir: PushdownFilterHint, *, offset: str = "") -> str: target_on = tuple(ne.name for ne in ir.target_on) @@ -593,6 +609,29 @@ def _(ir: Join) -> dict[str, Serializable]: } +def _serialize_prefilter(prefilter: Prefilter) -> dict[str, Serializable]: + """Serialize a normalized join prefilter descriptor.""" + properties: dict[str, Serializable] = { + "type": type(prefilter).__name__, + "target_side": prefilter.target_side, + "target_on": [ne.name for ne in prefilter.target_on], + "domain_on": [ne.name for ne in prefilter.domain_on], + "nulls_equal": prefilter.nulls_equal, + } + properties["domain_side"] = prefilter.domain_side + return properties + + +@_serialize_properties.register +def _(ir: JoinWithPrefilter) -> dict[str, Serializable]: + return { + "how": ir.options[0], + "left_on": [ne.name for ne in ir.left_on], + "right_on": [ne.name for ne in ir.right_on], + "prefilters": [_serialize_prefilter(prefilter) for prefilter in ir.prefilters], + } + + @_serialize_properties.register def _(ir: PushdownFilterHint) -> dict[str, Serializable]: return { diff --git a/python/cudf_polars/cudf_polars/streaming/filter_hint.py b/python/cudf_polars/cudf_polars/streaming/filter_hint.py index 4893f5e7cab3..736d909054eb 100644 --- a/python/cudf_polars/cudf_polars/streaming/filter_hint.py +++ b/python/cudf_polars/cudf_polars/streaming/filter_hint.py @@ -4,9 +4,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING, ClassVar +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias -from cudf_polars.dsl.ir import IR +from cudf_polars.dsl.ir import IR, Join if TYPE_CHECKING: from collections.abc import Sequence @@ -17,6 +18,81 @@ from cudf_polars.typing import Schema +JoinSide: TypeAlias = Literal["left", "right"] + + +@dataclass(frozen=True, slots=True) +class JoinInputPrefilter: + """A prefilter whose domain is already an input of its owning join.""" + + target_side: JoinSide + target_on: tuple[NamedExpr, ...] + domain_side: JoinSide + domain_on: tuple[NamedExpr, ...] + nulls_equal: bool + + +Prefilter: TypeAlias = JoinInputPrefilter + + +class JoinWithPrefilter(Join): + """Lowered join with normalized prefilter descriptors.""" + + __slots__ = ("prefilters",) + _non_child = ("schema", "left_on", "right_on", "options", "prefilters") + _n_non_child_args = 4 + + prefilters: tuple[Prefilter, ...] + + def __init__( + self, + schema: Schema, + left_on: Sequence[NamedExpr], + right_on: Sequence[NamedExpr], + options: Any, + prefilters: Sequence[Prefilter], + left: IR, + right: IR, + ): + self.schema = schema + self.left_on = tuple(left_on) + self.right_on = tuple(right_on) + self.options = options + self.prefilters = tuple(prefilters) + self.children = (left, right) + self._non_child_args = ( + self.left_on, + self.right_on, + self.options, + self.prefilters, + ) + + if not self.prefilters: + raise ValueError("JoinWithPrefilter requires at least one prefilter") + + @classmethod + def do_evaluate( + cls, + left_on: tuple[NamedExpr, ...], + right_on: tuple[NamedExpr, ...], + options: Any, + prefilters: tuple[Prefilter, ...], + left: DataFrame, + right: DataFrame, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate the join while ignoring its optional prefilters.""" + del prefilters + return Join.do_evaluate( + left_on, + right_on, + options, + left, + right, + context=context, + ) + + class PushdownFilterHint(IR): """ Optional join-key filter placed in a logical plan. diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index c13d977c27c2..67be77757268 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -8,11 +8,15 @@ from functools import reduce from typing import TYPE_CHECKING -from cudf_polars.dsl.ir import ConditionalJoin, Join, Slice +from cudf_polars.dsl.ir import ConditionalJoin, Join, Projection, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node -from cudf_polars.streaming.filter_hint import PushdownFilterHint +from cudf_polars.streaming.filter_hint import ( + JoinInputPrefilter, + JoinWithPrefilter, + PushdownFilterHint, +) from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.utils import ( @@ -26,6 +30,7 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.filter_hint import JoinSide, Prefilter from cudf_polars.streaming.parallel import LowerIRTransformer @@ -150,6 +155,88 @@ def _has_non_pointwise_keys(ir: Join) -> bool: return not all(expr.is_pointwise for expr in traversal(keys)) +def _lower_join_with_prefilters( + ir: Join, + rec: LowerIRTransformer, +) -> tuple[Join, MutableMapping[IR, PartitionInfo]]: + """Lower a join and normalize its adjacent filter hints.""" + targets = tuple( + child.children[0] if isinstance(child, PushdownFilterHint) else child + for child in ir.children + ) + lowered_targets, target_partition_info = zip( + *(rec(target) for target in targets), + strict=True, + ) + partition_info: MutableMapping[IR, PartitionInfo] = reduce( + operator.or_, target_partition_info + ) + + prefilters: list[Prefilter] = [] + claimed_sides: set[JoinSide] = set() + for target_index, child in enumerate(ir.children): + if not isinstance(child, PushdownFilterHint): + continue + + _target, domain = child.children + domain, _domain_partition_info = rec(domain) + + # A key-only Projection retains an explicit edge to its source. If that + # source is a join input and contains every requested key, the join can + # project those keys itself rather than execute a separate domain input. + left, right = lowered_targets + direct_domain = domain + while True: + if direct_domain == left and direct_domain == right: + domain_side: JoinSide | None = "right" if target_index == 0 else "left" + break + if direct_domain == left: + domain_side = "left" + break + if direct_domain == right: + domain_side = "right" + break + if isinstance(direct_domain, Projection) and all( + key.name in direct_domain.children[0].schema for key in child.domain_on + ): + (direct_domain,) = direct_domain.children + continue + domain_side = None + break + + target_side: JoinSide = "left" if target_index == 0 else "right" + if domain_side in claimed_sides: + domain_side = None + elif domain_side is not None: + claimed_sides.add(domain_side) + + if domain_side is None: + continue + prefilters.append( + JoinInputPrefilter( + target_side, + child.target_on, + domain_side, + child.domain_on, + child.nulls_equal, + ) + ) + + lowered_join: Join + if prefilters: + lowered_join = JoinWithPrefilter( + ir.schema, + ir.left_on, + ir.right_on, + ir.options, + prefilters, + *lowered_targets, + ) + else: + lowered_join = ir.reconstruct(lowered_targets) + return lowered_join, partition_info + + @lower_ir_node.register(PushdownFilterHint) def _( ir: PushdownFilterHint, rec: LowerIRTransformer @@ -224,17 +311,33 @@ def _( ) return rec(Slice(ir.schema, offset, length, new_join)) - # Lower children - children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True) - partition_info = reduce(operator.or_, _partition_info) - - # Check for dynamic planning - may have more partitions at runtime config_options = rec.state["config_options"] dynamic_planning = _dynamic_planning_on(config_options) + has_non_pointwise_keys = _has_non_pointwise_keys(ir) + if ( + dynamic_planning + and ir.options[0] != "Cross" + and ir.options[5] == "none" + and not has_non_pointwise_keys + and any(isinstance(child, PushdownFilterHint) for child in ir.children) + ): + preserve_prefilters = True + else: + preserve_prefilters = False + + if preserve_prefilters: + ir, partition_info = _lower_join_with_prefilters(ir, rec) + children = ir.children + else: + # Hints not owned by an adaptive join use the generic identity lowering. + children, _partition_info = zip( + *(rec(child) for child in ir.children), + strict=True, + ) + partition_info = reduce(operator.or_, _partition_info) - left, right = children + left, right = children[:2] output_count = max(partition_info[left].count, partition_info[right].count) - has_non_pointwise_keys = _has_non_pointwise_keys(ir) if output_count == 1 and not dynamic_planning: new_node = ir.reconstruct(children) partition_info[new_node] = PartitionInfo(count=1) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 63fc734a0ca7..26a4a32ab990 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -372,7 +372,7 @@ def __post_init__(self) -> None: # noqa: D105 @dataclasses.dataclass(frozen=True) class JoinFilterPushdownOptions: """ - Configuration options for join filter pushdown in the logical plan. + Configuration options for join filter pushdown. When performing a join between two tables, it is often favourable to pre-filter one side of the join with the keys (full or partial) of @@ -380,7 +380,8 @@ class JoinFilterPushdownOptions: participate in the join. cudf-polars supports a form of this where we can rewrite inner joins by - selecting a side to be filtered by the keys of the other side. + selecting a side to be filtered by the keys of the other side. At execution + time, these options also control how optional filters are applied. Pass ``None`` to ``StreamingExecutor(join_filter_pushdown=...)`` to disable the rewrite. @@ -393,6 +394,10 @@ class JoinFilterPushdownOptions: threshold Row-count ratio (key-provider-rows / to-be-filtered-table-rows) below which a filter on is inserted on the to-be-filtered table. Default is 0.5. + bloom_filter_max_size + Maximum Bloom-filter size in bytes. If the estimated Bloom filter exceeds + this size, an exact semi-join is preferred when its projected keys fit the + broadcast limit. Set to 0 to disable Bloom filters. Default is 32 MiB. trace Whether to emit plan-time trace decisions for filter decisions. Default is False. """ @@ -404,6 +409,13 @@ class JoinFilterPushdownOptions: f"{_env_prefix}__THRESHOLD", float, default=0.5 ) ) + bloom_filter_max_size: int = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__BLOOM_FILTER_MAX_SIZE", + int, + default=32 * 1024 * 1024, + ) + ) trace: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__TRACE", _bool_converter, default=False @@ -418,6 +430,12 @@ def __post_init__(self) -> None: # noqa: D105 object.__setattr__(self, "threshold", threshold) if not 0.0 <= threshold <= 1.0: raise ValueError("threshold must be between 0 and 1") + if isinstance(self.bloom_filter_max_size, bool) or not isinstance( + self.bloom_filter_max_size, int + ): + raise TypeError("bloom_filter_max_size must be an int") + if self.bloom_filter_max_size < 0: + raise ValueError("bloom_filter_max_size must be non-negative") if not isinstance(self.trace, bool): raise TypeError("trace must be a bool") diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 610c35256d60..85c4946785e6 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -137,7 +137,7 @@ def test_explain_logical_plan_with_join(tmp_path, df): assert "JOIN Inner ('x',) ('x',)" in plan -def test_explain_pushdown_filter_hint_is_logical_only(): +def test_explain_pushdown_filter_hint_in_dynamic_physical_plan(): domain = ( pl.LazyFrame({"key": [1, 99], "active": [True, False]}) .filter("active") @@ -157,19 +157,30 @@ def test_explain_pushdown_filter_hint_is_logical_only(): physical_serialized = serialize_query(query, engine, physical=True) assert "PUSHDOWN FILTER HINT ('key',) ('key',)" in logical - assert "PUSHDOWN FILTER HINT" not in physical + assert "prefilters=('JoinInputPrefilter',)" in physical + expected_properties = { + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + } assert any( - node.type == "PushdownFilterHint" - and node.properties - == { - "target_on": ["key"], - "domain_on": ["key"], - "nulls_equal": False, - } + node.type == "PushdownFilterHint" and node.properties == expected_properties for node in logical_serialized.nodes.values() ) - assert not any( - node.type == "PushdownFilterHint" for node in physical_serialized.nodes.values() + assert any( + node.type == "JoinWithPrefilter" + and node.properties["prefilters"] + == [ + { + "type": "JoinInputPrefilter", + "target_side": "right", + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + "domain_side": "left", + } + ] + for node in physical_serialized.nodes.values() ) 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 6e7e7c475dd1..2734a0630f87 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -16,7 +16,11 @@ 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.filter_hint import PushdownFilterHint +from cudf_polars.streaming.filter_hint import ( + JoinInputPrefilter, + JoinWithPrefilter, + PushdownFilterHint, +) from cudf_polars.streaming.join_filter_pushdown import ( CompositeCandidate, Decision, @@ -165,17 +169,22 @@ def test_filter_pushdown_is_independent_of_dynamic_planning( engine: SPMDEngine, ) -> None: root = translate_query(simple_query, engine) + config = make_config(dynamic_planning=False) optimized = optimize_join_filter_pushdown( root, StatsCollector(), - make_config(dynamic_planning=False), + config, ) assert find_hints(optimized) + lowering = lower_ir_graph(root, config, StatsCollector()) + assert not any( + isinstance(node, JoinWithPrefilter) for node in traversal([lowering.lowered]) + ) -def test_filter_hints_are_discarded_during_lowering( +def test_adjacent_filter_hint_is_recorded_on_lowered_join( simple_query: pl.LazyFrame, engine: SPMDEngine, ) -> None: @@ -185,7 +194,19 @@ def test_filter_hints_are_discarded_during_lowering( lowering = lower_ir_graph(root, config, StatsCollector()) assert find_hints(lowering.optimized) - assert not find_hints(lowering.lowered) + assert isinstance(lowering.lowered, JoinWithPrefilter) + left, right = lowering.lowered.children + assert not isinstance(left, PushdownFilterHint) + assert not isinstance(right, PushdownFilterHint) + (prefilter,) = lowering.lowered.prefilters + assert isinstance(prefilter, JoinInputPrefilter) + assert prefilter.target_side == "right" + assert prefilter.domain_side == "left" + assert tuple(right.schema) == ("l_partkey", "l_suppkey") + assert tuple(ne.name for ne in prefilter.target_on) == ("l_partkey",) + assert tuple(ne.name for ne in prefilter.domain_on) == ("p_partkey",) + assert not prefilter.nulls_equal + assert tuple(left.schema) == ("p_partkey",) assert not find_joins(lowering.lowered, "Semi") diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index f584ceff6528..13a131a94fc4 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import json import os import subprocess import sys @@ -131,6 +132,116 @@ def test_structlog_contains_expected_ir_types(timeout_seconds: int): assert b"ir_type=GroupBy" in result +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,output_rows", + [ + (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 10), + (64, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 10), + ( + 1_000_000, + 32 * 1024 * 1024, + "broadcast_left", + "skip", + "target_not_redistributed", + None, + ), + ], + ids=["bloom", "exact", "skip"], +) +def test_local_join_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + join_strategy: str, + method: str, + reason: str, + output_rows: int | None, +) -> None: + """Trace a direct-input join prefilter selected through the public engine.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + import os + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + domain = ( + pl.LazyFrame({{"key": [1, 99], "active": [True, False]}}) + .filter("active") + .select("key") + ) + target = pl.LazyFrame( + {{"key": [i % 100 for i in range(1_000)], "value": range(1_000)}} + ) + query = domain.join(target, on="key") + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 100, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" and "join_prefilters" in log + ) + record = {{ + "result_rows": result.height, + "join_strategy": event["decision"], + "prefilter": event["join_prefilters"][0], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 10 + assert record["join_strategy"] == join_strategy + assert ( + record["prefilter"].items() + >= { + "target_side": "right", + "domain_side": "left", + "method": method, + "reason": reason, + "domain_rows": 1, + }.items() + ) + if output_rows is None: + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 1 + assert record["prefilter"]["input_rows"] == 1_000 + assert record["prefilter"]["output_rows"] == output_rows + + def test_structlog_disabled_by_default(timeout_seconds: int): """Test that structlog does NOT emit events when CUDF_POLARS_LOG_TRACES is not set.""" pytest.importorskip("structlog") diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index aa01207e7e51..f9dbb1604039 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -669,10 +669,15 @@ def test_join_filter_pushdown_options_from_env( monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__BLOOM_FILTER_MAX_SIZE", + "1024", + ) monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") config = ConfigOptions.from_polars_engine(pl.GPUEngine()) assert config.executor.join_filter_pushdown is not None assert config.executor.join_filter_pushdown.threshold == 0.125 + assert config.executor.join_filter_pushdown.bloom_filter_max_size == 1024 assert config.executor.join_filter_pushdown.trace @@ -707,6 +712,24 @@ def test_validate_join_filter_pushdown_options() -> None: executor_options={"join_filter_pushdown": {"trace": "bad"}}, ) ) + with pytest.raises(TypeError, match="bloom_filter_max_size must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "join_filter_pushdown": {"bloom_filter_max_size": "bad"} + }, + ) + ) + with pytest.raises(ValueError, match="bloom_filter_max_size must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "join_filter_pushdown": {"bloom_filter_max_size": -1} + }, + ) + ) def test_validate_join_filter_pushdown_type() -> None: @@ -723,7 +746,9 @@ def test_validate_join_filter_pushdown_type() -> None: def test_join_filter_pushdown_from_instance() -> None: - options = JoinFilterPushdownOptions(threshold=0.25, trace=True) + options = JoinFilterPushdownOptions( + threshold=0.25, bloom_filter_max_size=1024, trace=True + ) config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", From 6bfe63a75040cce842c470765d19bf3d73b853e5 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 4 Aug 2026 12:31:34 +0100 Subject: [PATCH 03/17] Apply adaptive prefilters from external domains --- .../cudf_polars/streaming/actor_graph/join.py | 186 ++++++++++++++---- .../streaming/actor_graph/join_bindings.py | 30 ++- .../streaming/actor_graph/prefilter.py | 15 +- .../cudf_polars/streaming/explain.py | 12 +- .../cudf_polars/streaming/filter_hint.py | 39 +++- .../cudf_polars/cudf_polars/streaming/join.py | 55 ++++-- .../tests/streaming/test_explain.py | 6 +- .../streaming/test_join_filter_pushdown.py | 8 +- .../tests/streaming/test_tracing.py | 151 ++++++++++++++ 9 files changed, 424 insertions(+), 78 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index cca40f3d5fec..c79db7034747 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -64,7 +64,11 @@ send_metadata, shutdown_on_error, ) -from cudf_polars.streaming.filter_hint import JoinWithPrefilter +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, +) from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import _concat @@ -533,6 +537,7 @@ def add_bloom_prefilter( def make_prefilter_execution( context: Context, comm: Communicator, + ir: Join, ir_context: IRExecutionContext, strategy: JoinStrategy, ch_left: Channel[TableChunk], @@ -554,8 +559,28 @@ def make_prefilter_execution( if decision.method == "skip": continue - indices = names_to_indices(prefilter.domain_on, bound.domain.node.schema) - bound.key_channel = execution.buffer_domain(prefilter.domain_side, indices) + if isinstance(prefilter.domain, JoinInputDomain): + indices = names_to_indices(prefilter.domain_on, bound.domain.node.schema) + bound.key_channel = execution.buffer_domain(prefilter.domain.side, indices) + else: + sample = bound.domain.sample + if sample is None: + raise ValueError("Active external prefilter has no domain sample") + indices = names_to_indices(prefilter.domain_on, bound.domain.node.schema) + if indices != tuple(range(len(bound.domain.node.schema))): + raise ValueError("External prefilter domains must contain only keys") + bound.key_channel = context.create_channel() + execution.add_channel(bound.key_channel) + execution.add_task( + replay_buffered_channel( + context, + bound.key_channel, + bound.domain.channel, + sample.chunks, + bound.domain.metadata, + trace_ir=ir, + ) + ) for bound in bindings.prefilters: decision = bound.decision @@ -1095,26 +1120,35 @@ def choose_prefilters( broadcast_limit: int, bloom_filter_max_size: int, ) -> None: - """Choose strategies for optional join prefilters.""" + """Choose strategies for prefilters with sufficient available statistics.""" partitionings = { "left": left_partitioning, "right": right_partitioning, } for bound in bindings.prefilters: + if bound.decision is not None: + continue target = bound.target.sample if target is None: raise ValueError("Join target has not been sampled") target_side = bound.prefilter.target_side + target_requires_redistribution = join_input_requires_redistribution( + strategy, + target_side, + partitionings[target_side], + bound.target.metadata, + ) + if ( + isinstance(bound.prefilter.domain, ExternalDomain) + and bound.domain.sample is None + and target_requires_redistribution + ): + continue bound.decision = choose_prefilter( bound.prefilter, target, bound.domain.sample, - target_requires_redistribution=join_input_requires_redistribution( - strategy, - target_side, - partitionings[target_side], - bound.target.metadata, - ), + target_requires_redistribution=target_requires_redistribution, broadcast_limit=broadcast_limit, bloom_filter_max_size=bloom_filter_max_size, ) @@ -1157,28 +1191,24 @@ async def sample_input( ) -async def collect_join_samples( +async def collect_samples( context: Context, comm: Communicator, bindings: JoinBindings, + inputs: tuple[JoinInputBinding, ...], sample_chunk_count: int, target_partition_size: int, collective_id: int, ) -> None: - """Sample join inputs and attach estimates to their runtime bindings.""" - inputs = [bindings.left, bindings.right] - + """Sample inputs and attach aggregate estimates to their bindings.""" + if not inputs: + return sampling_inputs = [] for input_ in inputs: - domain_prefilters = [ - bound for bound in bindings.prefilters if bound.domain is input_ - ] - if len(domain_prefilters) > 1: + prefilters = [bound for bound in bindings.prefilters if bound.domain is input_] + if len(prefilters) > 1: raise ValueError("One join input cannot provide multiple prefilter domains") - sampling_inputs.append( - (input_, domain_prefilters[0] if domain_prefilters else None) - ) - + sampling_inputs.append((input_, prefilters[0] if prefilters else None)) local_samples = await gather_in_task_group( *( sample_input( @@ -1198,10 +1228,78 @@ async def collect_join_samples( tuple(local_samples), collective_id, ) - for input_, sample in zip(inputs, samples, strict=True): + for (input_, _), sample in zip(sampling_inputs, samples, strict=True): input_.sample = sample +async def release_skipped_external_domains( + context: Context, bindings: JoinBindings +) -> None: + """Release buffered data and stop external domains rejected by planning.""" + channels = [] + for bound in bindings.prefilters: + if not isinstance(bound.prefilter.domain, ExternalDomain): + continue + if bound.decision is None: + raise ValueError("Join prefilter has no runtime decision") + if bound.decision.method != "skip": + continue + if bound.domain.sample is not None: + bound.domain.sample.chunks.clear() + channels.append(bound.domain.channel) + if channels: + await gather_in_task_group(*(channel.shutdown(context) for channel in channels)) + + +async def resolve_prefilters( + context: Context, + comm: Communicator, + bindings: JoinBindings, + strategy: JoinStrategy, + left_partitioning: NormalizedPartitioning, + right_partitioning: NormalizedPartitioning, + executor: StreamingExecutor, + collective_id: int, +) -> None: + """Resolve optional prefilters after selecting the join strategy.""" + config = executor.join_filter_pushdown + if config is None or not bindings.prefilters: + return + + choose_prefilters( + bindings, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + config.bloom_filter_max_size, + ) + assert executor.dynamic_planning is not None + await collect_samples( + context, + comm, + bindings, + tuple( + bound.domain + for bound in bindings.prefilters + if isinstance(bound.prefilter.domain, ExternalDomain) + and bound.decision is None + ), + executor.dynamic_planning.sample_chunk_count, + executor.target_partition_size, + collective_id, + ) + choose_prefilters( + bindings, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + config.bloom_filter_max_size, + ) + await release_skipped_external_domains(context, bindings) + + async def choose_strategy( context: Context, comm: Communicator, @@ -1245,10 +1343,11 @@ async def choose_strategy( ) else: assert executor.dynamic_planning is not None - await collect_join_samples( + await collect_samples( context, comm, bindings, + (bindings.left, bindings.right), executor.dynamic_planning.sample_chunk_count, executor.target_partition_size, collective_ids.size_estimate, @@ -1271,15 +1370,16 @@ async def choose_strategy( chunkwise=chunkwise, tracer=tracer, ) - if executor.join_filter_pushdown is not None: - choose_prefilters( - bindings, - strategy, - left_partitioning, - right_partitioning, - executor.broadcast_limit, - executor.join_filter_pushdown.bloom_filter_max_size, - ) + await resolve_prefilters( + context, + comm, + bindings, + strategy, + left_partitioning, + right_partitioning, + executor, + collective_ids.size_estimate, + ) return strategy @@ -1292,6 +1392,7 @@ async def join_actor( ch_out: Channel[TableChunk], ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], executor: StreamingExecutor, collective_ids: JoinCollectiveIds, ) -> None: @@ -1318,6 +1419,8 @@ async def join_actor( Input channel for the left side. ch_right Input channel for the right side. + ch_prefilter_domains + Input channels providing the prefilter key domains. executor Streaming executor configuration. collective_ids @@ -1328,20 +1431,28 @@ async def join_actor( ch_out, ch_left, ch_right, + *ch_prefilter_domains, trace_ir=ir, ir_context=ir_context, ) as tracer: - left_metadata, right_metadata = await gather_in_task_group( + ( + left_metadata, + right_metadata, + *prefilter_domain_metadata, + ) = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), + *(recv_metadata(ch, context) for ch in ch_prefilter_domains), ) bindings = bind_join_inputs( ir, ch_left, ch_right, + ch_prefilter_domains, left_metadata, right_metadata, + tuple(prefilter_domain_metadata), collective_ids.cardinality_tags, ) @@ -1373,6 +1484,7 @@ async def join_actor( prefilter_execution = make_prefilter_execution( context, comm, + ir, ir_context, strategy, ch_left_replay, @@ -1482,7 +1594,7 @@ def _( ir: Join | JoinWithPrefilter, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: # Join operation. - left, right = ir.children + left, right, *prefilter_domains = ir.children partition_info = rec.state["partition_info"] left_count = partition_info[left].count right_count = partition_info[right].count @@ -1534,6 +1646,10 @@ def _( channels[ir].reserve_input_slot(), channels[left].reserve_output_slot(), channels[right].reserve_output_slot(), + tuple( + channels[domain].reserve_output_slot() + for domain in prefilter_domains + ), executor, collective_ids, ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py index 5b52fa874830..72b57363e592 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py @@ -7,7 +7,11 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from cudf_polars.streaming.filter_hint import JoinWithPrefilter +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, +) if TYPE_CHECKING: from typing import Any @@ -58,25 +62,47 @@ def bind_join_inputs( ir: Join, ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], left_metadata: ChannelMetadata, right_metadata: ChannelMetadata, + prefilter_domain_metadata: tuple[ChannelMetadata, ...], cardinality_tags: tuple[int, ...], ) -> JoinBindings: """Bind logical join inputs and prefilters to their runtime resources.""" left = JoinInputBinding(ir.children[0], ch_left, left_metadata) right = JoinInputBinding(ir.children[1], ch_right, right_metadata) if not isinstance(ir, JoinWithPrefilter): + if ch_prefilter_domains or prefilter_domain_metadata: + raise ValueError("A plain Join cannot have prefilter domain inputs") return JoinBindings(left, right) + external_inputs = tuple( + JoinInputBinding(node, channel, metadata) + for node, channel, metadata in zip( + ir.children[2:], + ch_prefilter_domains, + prefilter_domain_metadata, + strict=True, + ) + ) + external_prefilter_count = sum( + isinstance(prefilter.domain, ExternalDomain) for prefilter in ir.prefilters + ) + if external_prefilter_count != len(external_inputs): + raise ValueError("Join prefilters and external domain inputs must align") if len(cardinality_tags) < len(ir.prefilters): raise ValueError("Each join prefilter requires a cardinality collective ID") sides = {"left": left, "right": right} + external_inputs_iter = iter(external_inputs) cardinality_tags_iter = iter(cardinality_tags) prefilters = [] for prefilter in ir.prefilters: target = sides[prefilter.target_side] - domain = sides[prefilter.domain_side] + if isinstance(prefilter.domain, JoinInputDomain): + domain = sides[prefilter.domain.side] + else: + domain = next(external_inputs_iter) prefilters.append( BoundPrefilter( prefilter, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py index 683203bcf35b..390991e0c74d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -20,7 +20,10 @@ send_metadata, shutdown_channels_on_error, ) -from cudf_polars.streaming.filter_hint import JoinInputPrefilter +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, +) if TYPE_CHECKING: from collections.abc import Coroutine, Iterable, Sequence @@ -81,7 +84,11 @@ def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: "bloom_bytes": self.bloom_bytes, "exact_bytes": self.exact_bytes, } - result["domain_side"] = prefilter.domain_side + if isinstance(prefilter.domain, JoinInputDomain): + result["domain_side"] = prefilter.domain.side + else: + assert isinstance(prefilter.domain, ExternalDomain) + result["domain"] = "external" return result @@ -290,8 +297,8 @@ def choose_prefilter( if domain is None: raise ValueError("A redistributed target requires domain statistics") if ( - isinstance(prefilter, JoinInputPrefilter) - and prefilter.target_side == prefilter.domain_side + isinstance(prefilter.domain, JoinInputDomain) + and prefilter.target_side == prefilter.domain.side ): return PrefilterDecision( "skip", diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 053b79b1edb8..b6102dd33329 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -38,6 +38,8 @@ from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import IOPartitionFlavor from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, JoinWithPrefilter, PushdownFilterHint, ) @@ -479,7 +481,7 @@ def _(ir: Join, *, offset: str = "") -> str: def _(ir: JoinWithPrefilter, *, offset: str = "") -> str: left_on = tuple(ne.name for ne in ir.left_on) right_on = tuple(ne.name for ne in ir.right_on) - prefilters = tuple(type(prefilter).__name__ for prefilter in ir.prefilters) + prefilters = tuple(type(prefilter.domain).__name__ for prefilter in ir.prefilters) return _repr_header( offset, f"JOIN {ir.options[0]} {left_on} {right_on} {prefilters=}", @@ -618,7 +620,13 @@ def _serialize_prefilter(prefilter: Prefilter) -> dict[str, Serializable]: "domain_on": [ne.name for ne in prefilter.domain_on], "nulls_equal": prefilter.nulls_equal, } - properties["domain_side"] = prefilter.domain_side + if isinstance(prefilter.domain, JoinInputDomain): + properties["domain"] = { + "type": type(prefilter.domain).__name__, + "side": prefilter.domain.side, + } + elif isinstance(prefilter.domain, ExternalDomain): + properties["domain"] = {"type": type(prefilter.domain).__name__} return properties diff --git a/python/cudf_polars/cudf_polars/streaming/filter_hint.py b/python/cudf_polars/cudf_polars/streaming/filter_hint.py index 736d909054eb..de6e1765ff39 100644 --- a/python/cudf_polars/cudf_polars/streaming/filter_hint.py +++ b/python/cudf_polars/cudf_polars/streaming/filter_hint.py @@ -22,19 +22,31 @@ @dataclass(frozen=True, slots=True) -class JoinInputPrefilter: - """A prefilter whose domain is already an input of its owning join.""" +class JoinInputDomain: + """A prefilter domain provided by an input of its owning join.""" + + side: JoinSide + + +@dataclass(frozen=True, slots=True) +class ExternalDomain: + """A prefilter domain provided by an additional join input.""" + + +PrefilterDomain: TypeAlias = JoinInputDomain | ExternalDomain + + +@dataclass(frozen=True, slots=True) +class Prefilter: + """Description of an optional join prefilter.""" target_side: JoinSide target_on: tuple[NamedExpr, ...] - domain_side: JoinSide + domain: PrefilterDomain domain_on: tuple[NamedExpr, ...] nulls_equal: bool -Prefilter: TypeAlias = JoinInputPrefilter - - class JoinWithPrefilter(Join): """Lowered join with normalized prefilter descriptors.""" @@ -53,13 +65,14 @@ def __init__( prefilters: Sequence[Prefilter], left: IR, right: IR, + *external_domains: IR, ): self.schema = schema self.left_on = tuple(left_on) self.right_on = tuple(right_on) self.options = options self.prefilters = tuple(prefilters) - self.children = (left, right) + self.children = (left, right, *external_domains) self._non_child_args = ( self.left_on, self.right_on, @@ -69,6 +82,15 @@ def __init__( if not self.prefilters: raise ValueError("JoinWithPrefilter requires at least one prefilter") + external_domain_count = sum( + isinstance(prefilter.domain, ExternalDomain) + for prefilter in self.prefilters + ) + if external_domain_count != len(external_domains): + raise ValueError( + "External prefilters and additional JoinWithPrefilter children " + "must align" + ) @classmethod def do_evaluate( @@ -79,10 +101,11 @@ def do_evaluate( prefilters: tuple[Prefilter, ...], left: DataFrame, right: DataFrame, + *external_domains: DataFrame, context: IRExecutionContext, ) -> DataFrame: """Evaluate the join while ignoring its optional prefilters.""" - del prefilters + del prefilters, external_domains return Join.do_evaluate( left_on, right_on, diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 67be77757268..4218ac6bfaf7 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -13,8 +13,10 @@ from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node from cudf_polars.streaming.filter_hint import ( - JoinInputPrefilter, + ExternalDomain, + JoinInputDomain, JoinWithPrefilter, + Prefilter, PushdownFilterHint, ) from cudf_polars.streaming.repartition import Repartition @@ -30,7 +32,7 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR - from cudf_polars.streaming.filter_hint import JoinSide, Prefilter + from cudf_polars.streaming.filter_hint import JoinSide from cudf_polars.streaming.parallel import LowerIRTransformer @@ -158,7 +160,7 @@ def _has_non_pointwise_keys(ir: Join) -> bool: def _lower_join_with_prefilters( ir: Join, rec: LowerIRTransformer, -) -> tuple[Join, MutableMapping[IR, PartitionInfo]]: +) -> tuple[JoinWithPrefilter, MutableMapping[IR, PartitionInfo]]: """Lower a join and normalize its adjacent filter hints.""" targets = tuple( child.children[0] if isinstance(child, PushdownFilterHint) else child @@ -173,13 +175,15 @@ def _lower_join_with_prefilters( ) prefilters: list[Prefilter] = [] + external_domains: list[IR] = [] claimed_sides: set[JoinSide] = set() for target_index, child in enumerate(ir.children): if not isinstance(child, PushdownFilterHint): continue _target, domain = child.children - domain, _domain_partition_info = rec(domain) + domain, domain_partition_info = rec(domain) + partition_info.update(domain_partition_info) # A key-only Projection retains an explicit edge to its source. If that # source is a join input and contains every requested key, the join can @@ -210,31 +214,40 @@ def _lower_join_with_prefilters( elif domain_side is not None: claimed_sides.add(domain_side) - if domain_side is None: - continue - prefilters.append( - JoinInputPrefilter( - target_side, - child.target_on, - domain_side, - child.domain_on, - child.nulls_equal, + if domain_side is not None: + prefilters.append( + Prefilter( + target_side, + child.target_on, + JoinInputDomain(domain_side), + child.domain_on, + child.nulls_equal, + ) + ) + else: + external_domains.append(domain) + prefilters.append( + Prefilter( + target_side, + child.target_on, + ExternalDomain(), + child.domain_on, + child.nulls_equal, + ) ) - ) - lowered_join: Join - if prefilters: - lowered_join = JoinWithPrefilter( + return ( + JoinWithPrefilter( ir.schema, ir.left_on, ir.right_on, ir.options, prefilters, *lowered_targets, - ) - else: - lowered_join = ir.reconstruct(lowered_targets) - return lowered_join, partition_info + *external_domains, + ), + partition_info, + ) @lower_ir_node.register(PushdownFilterHint) diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 85c4946785e6..85ff48da64e0 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -157,7 +157,7 @@ def test_explain_pushdown_filter_hint_in_dynamic_physical_plan(): physical_serialized = serialize_query(query, engine, physical=True) assert "PUSHDOWN FILTER HINT ('key',) ('key',)" in logical - assert "prefilters=('JoinInputPrefilter',)" in physical + assert "prefilters=('JoinInputDomain',)" in physical expected_properties = { "target_on": ["key"], "domain_on": ["key"], @@ -172,12 +172,12 @@ def test_explain_pushdown_filter_hint_in_dynamic_physical_plan(): and node.properties["prefilters"] == [ { - "type": "JoinInputPrefilter", + "type": "Prefilter", "target_side": "right", "target_on": ["key"], "domain_on": ["key"], "nulls_equal": False, - "domain_side": "left", + "domain": {"type": "JoinInputDomain", "side": "left"}, } ] for node in physical_serialized.nodes.values() 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 2734a0630f87..c560e417e018 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -17,8 +17,9 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.filter_hint import ( - JoinInputPrefilter, + JoinInputDomain, JoinWithPrefilter, + Prefilter, PushdownFilterHint, ) from cudf_polars.streaming.join_filter_pushdown import ( @@ -199,9 +200,10 @@ def test_adjacent_filter_hint_is_recorded_on_lowered_join( assert not isinstance(left, PushdownFilterHint) assert not isinstance(right, PushdownFilterHint) (prefilter,) = lowering.lowered.prefilters - assert isinstance(prefilter, JoinInputPrefilter) + assert isinstance(prefilter, Prefilter) + assert isinstance(prefilter.domain, JoinInputDomain) assert prefilter.target_side == "right" - assert prefilter.domain_side == "left" + assert prefilter.domain.side == "left" assert tuple(right.schema) == ("l_partkey", "l_suppkey") assert tuple(ne.name for ne in prefilter.target_on) == ("l_partkey",) assert tuple(ne.name for ne in prefilter.domain_on) == ("p_partkey",) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index 13a131a94fc4..aeb6c16c6a7c 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -242,6 +242,157 @@ def test_local_join_prefilter_trace_records_decision_and_effect( assert record["prefilter"]["output_rows"] == output_rows +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,domain_rows", + [ + (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 30), + (512, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 30), + ( + 1_000_000, + 32 * 1024 * 1024, + "broadcast_left", + "skip", + "target_not_redistributed", + None, + ), + ], + ids=["bloom", "exact", "skip"], +) +def test_external_join_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + join_strategy: str, + method: str, + reason: str, + domain_rows: int | None, +) -> None: + """Trace an external-domain prefilter selected through the public engine.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + nation = ( + pl.LazyFrame( + {{"n_nationkey": range(10), "active": [True] * 5 + [False] * 5}} + ) + .filter("active") + .select("n_nationkey") + ) + orders = pl.LazyFrame( + {{ + "o_orderkey": range(90), + "n_nationkey": [i % 10 for i in range(90)], + }} + ) + lineitem = pl.LazyFrame( + {{ + "l_orderkey": [i % 90 for i in range(180)], + "l_suppkey": [i % 60 for i in range(180)], + }} + ) + supplier = pl.LazyFrame( + {{ + "s_suppkey": range(30), + "s_nationkey": [i % 10 for i in range(30)], + }} + ) + query = ( + nation.join(orders, on="n_nationkey") + .join( + lineitem, + left_on="o_orderkey", + right_on="l_orderkey", + maintain_order="left", + ) + .join( + supplier, + left_on=("l_suppkey", "n_nationkey"), + right_on=("s_suppkey", "s_nationkey"), + ) + ) + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 100, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" + and any( + prefilter.get("domain") == "external" + for prefilter in log.get("join_prefilters", ()) + ) + ) + (prefilter,) = ( + prefilter + for prefilter in event["join_prefilters"] + if prefilter.get("domain") == "external" + ) + record = {{ + "result_rows": result.height, + "join_strategy": event["decision"], + "prefilter": prefilter, + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 45 + assert record["join_strategy"] == join_strategy + assert ( + record["prefilter"].items() + >= { + "target_side": "right", + "domain": "external", + "method": method, + "reason": reason, + "domain_rows": domain_rows, + }.items() + ) + if method == "skip": + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 30 + assert record["prefilter"]["input_rows"] == 180 + if method == "broadcast_semi_join": + assert record["prefilter"]["output_rows"] == 90 + else: + assert 90 <= record["prefilter"]["output_rows"] < 180 + + def test_structlog_disabled_by_default(timeout_seconds: int): """Test that structlog does NOT emit events when CUDF_POLARS_LOG_TRACES is not set.""" pytest.importorskip("structlog") From 4ea40b9a76841f2d4eb88c03b75480cd662734e6 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:28:30 +0100 Subject: [PATCH 04/17] Better names --- .../cudf_polars/streaming/actor_graph/join.py | 196 +++++++++--------- .../{join_bindings.py => join_planning.py} | 56 ++--- 2 files changed, 128 insertions(+), 124 deletions(-) rename python/cudf_polars/cudf_polars/streaming/actor_graph/{join_bindings.py => join_planning.py} (70%) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index c79db7034747..3c1987189d83 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -37,7 +37,7 @@ from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.streaming.actor_graph.join_bindings import bind_join_inputs +from cudf_polars.streaming.actor_graph.join_planning import make_join_planning_state from cudf_polars.streaming.actor_graph.nodes import default_node_multi from cudf_polars.streaming.actor_graph.prefilter import ( PrefilterExecution, @@ -82,10 +82,10 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator - from cudf_polars.streaming.actor_graph.join_bindings import ( - BoundPrefilter, - JoinBindings, - JoinInputBinding, + from cudf_polars.streaming.actor_graph.join_planning import ( + JoinInput, + JoinPlanningState, + PrefilterCandidate, ) from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo @@ -542,7 +542,7 @@ def make_prefilter_execution( strategy: JoinStrategy, ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], - bindings: JoinBindings, + join_state: JoinPlanningState, collective_ids: JoinCollectiveIds, ) -> PrefilterExecution: """Create the actors and channels that realize selected prefilters.""" @@ -551,50 +551,50 @@ def make_prefilter_execution( # Prepare every required domain before connecting target-side filters. This # is important for opposing direct filters: each filter must consume the # replay produced while the same input's keys are copied for the other one. - for bound in bindings.prefilters: - decision = bound.decision + for candidate in join_state.candidates: + decision = candidate.decision if decision is None: raise ValueError("Join prefilter has no runtime decision") - prefilter = bound.prefilter + spec = candidate.spec if decision.method == "skip": continue - if isinstance(prefilter.domain, JoinInputDomain): - indices = names_to_indices(prefilter.domain_on, bound.domain.node.schema) - bound.key_channel = execution.buffer_domain(prefilter.domain.side, indices) + if isinstance(spec.domain, JoinInputDomain): + indices = names_to_indices(spec.domain_on, candidate.domain.node.schema) + candidate.key_channel = execution.buffer_domain(spec.domain.side, indices) else: - sample = bound.domain.sample + sample = candidate.domain.sample if sample is None: raise ValueError("Active external prefilter has no domain sample") - indices = names_to_indices(prefilter.domain_on, bound.domain.node.schema) - if indices != tuple(range(len(bound.domain.node.schema))): + indices = names_to_indices(spec.domain_on, candidate.domain.node.schema) + if indices != tuple(range(len(candidate.domain.node.schema))): raise ValueError("External prefilter domains must contain only keys") - bound.key_channel = context.create_channel() - execution.add_channel(bound.key_channel) + candidate.key_channel = context.create_channel() + execution.add_channel(candidate.key_channel) execution.add_task( replay_buffered_channel( context, - bound.key_channel, - bound.domain.channel, + candidate.key_channel, + candidate.domain.channel, sample.chunks, - bound.domain.metadata, + candidate.domain.metadata, trace_ir=ir, ) ) - for bound in bindings.prefilters: - decision = bound.decision + for candidate in join_state.candidates: + decision = candidate.decision assert decision is not None if decision.method == "skip": continue - prefilter = bound.prefilter - ch_domain_keys = bound.key_channel + spec = candidate.spec + ch_domain_keys = candidate.key_channel assert ch_domain_keys is not None - target_side = prefilter.target_side - target = bound.target.node + target_side = spec.target_side + target = candidate.target.node ch_target = execution.join_inputs[target_side] ch_filtered: Channel[TableChunk] = context.create_channel() - trace_stats = bound.trace + trace_stats = candidate.trace collective_id = collective_ids.prefilter(strategy, target_side) if decision.method == "bloom": @@ -604,7 +604,7 @@ def make_prefilter_execution( comm, decision.bloom_bytes, execution, - names_to_indices(prefilter.target_on, target.schema), + names_to_indices(spec.target_on, target.schema), ch_domain_keys, ch_target, ch_filtered, @@ -613,15 +613,15 @@ def make_prefilter_execution( ) else: assert decision.method == "broadcast_semi_join" - domain_schema = {key.name: key.value.dtype for key in prefilter.domain_on} - if len(domain_schema) != len(prefilter.domain_on): + domain_schema = {key.name: key.value.dtype for key in spec.domain_on} + if len(domain_schema) != len(spec.domain_on): raise ValueError("Broadcast semi-join keys must have unique names") - projected_domain = Projection(domain_schema, bound.domain.node) + projected_domain = Projection(domain_schema, candidate.domain.node) semi_join = Join( target.schema, - prefilter.target_on, - prefilter.domain_on, - ("Semi", prefilter.nulls_equal, None, "", False, "none"), + spec.target_on, + spec.domain_on, + ("Semi", spec.nulls_equal, None, "", False, "none"), target, projected_domain, ) @@ -1113,7 +1113,7 @@ def join_input_requires_redistribution( def choose_prefilters( - bindings: JoinBindings, + join_state: JoinPlanningState, strategy: JoinStrategy, left_partitioning: NormalizedPartitioning, right_partitioning: NormalizedPartitioning, @@ -1125,29 +1125,29 @@ def choose_prefilters( "left": left_partitioning, "right": right_partitioning, } - for bound in bindings.prefilters: - if bound.decision is not None: + for candidate in join_state.candidates: + if candidate.decision is not None: continue - target = bound.target.sample + target = candidate.target.sample if target is None: raise ValueError("Join target has not been sampled") - target_side = bound.prefilter.target_side + target_side = candidate.spec.target_side target_requires_redistribution = join_input_requires_redistribution( strategy, target_side, partitionings[target_side], - bound.target.metadata, + candidate.target.metadata, ) if ( - isinstance(bound.prefilter.domain, ExternalDomain) - and bound.domain.sample is None + isinstance(candidate.spec.domain, ExternalDomain) + and candidate.domain.sample is None and target_requires_redistribution ): continue - bound.decision = choose_prefilter( - bound.prefilter, + candidate.decision = choose_prefilter( + candidate.spec, target, - bound.domain.sample, + candidate.domain.sample, target_requires_redistribution=target_requires_redistribution, broadcast_limit=broadcast_limit, bloom_filter_max_size=bloom_filter_max_size, @@ -1157,26 +1157,26 @@ def choose_prefilters( async def sample_input( context: Context, comm: Communicator, - input_: JoinInputBinding, - prefilter: BoundPrefilter | None, + input_: JoinInput, + candidate: PrefilterCandidate | None, sample_chunk_count: int, target_partition_size: int, ) -> TableSizeStats: """Sample one join-planning input and optionally estimate cardinality.""" - if prefilter is None: + if candidate is None: cardinality_estimator = None cardinality_columns: tuple[int, ...] = () else: cardinality_estimator = CardinalityEstimator( context, comm, - tag=prefilter.cardinality_tag, + tag=candidate.cardinality_tag, ) cardinality_columns = names_to_indices( - prefilter.prefilter.domain_on, + candidate.spec.domain_on, input_.node.schema, ) - assert len(cardinality_columns) == len(prefilter.prefilter.domain_on), ( + assert len(cardinality_columns) == len(candidate.spec.domain_on), ( "Prefilter domain keys must be columns" ) @@ -1194,32 +1194,36 @@ async def sample_input( async def collect_samples( context: Context, comm: Communicator, - bindings: JoinBindings, - inputs: tuple[JoinInputBinding, ...], + join_state: JoinPlanningState, + inputs: tuple[JoinInput, ...], sample_chunk_count: int, target_partition_size: int, collective_id: int, ) -> None: - """Sample inputs and attach aggregate estimates to their bindings.""" + """Sample inputs and attach aggregate estimates to their planning state.""" if not inputs: return sampling_inputs = [] for input_ in inputs: - prefilters = [bound for bound in bindings.prefilters if bound.domain is input_] - if len(prefilters) > 1: + candidates = [ + candidate + for candidate in join_state.candidates + if candidate.domain is input_ + ] + if len(candidates) > 1: raise ValueError("One join input cannot provide multiple prefilter domains") - sampling_inputs.append((input_, prefilters[0] if prefilters else None)) + sampling_inputs.append((input_, candidates[0] if candidates else None)) local_samples = await gather_in_task_group( *( sample_input( context, comm, input_, - prefilter, + candidate, sample_chunk_count, target_partition_size, ) - for input_, prefilter in sampling_inputs + for input_, candidate in sampling_inputs ) ) samples = await aggregate_estimates( @@ -1233,20 +1237,20 @@ async def collect_samples( async def release_skipped_external_domains( - context: Context, bindings: JoinBindings + context: Context, join_state: JoinPlanningState ) -> None: """Release buffered data and stop external domains rejected by planning.""" channels = [] - for bound in bindings.prefilters: - if not isinstance(bound.prefilter.domain, ExternalDomain): + for candidate in join_state.candidates: + if not isinstance(candidate.spec.domain, ExternalDomain): continue - if bound.decision is None: + if candidate.decision is None: raise ValueError("Join prefilter has no runtime decision") - if bound.decision.method != "skip": + if candidate.decision.method != "skip": continue - if bound.domain.sample is not None: - bound.domain.sample.chunks.clear() - channels.append(bound.domain.channel) + if candidate.domain.sample is not None: + candidate.domain.sample.chunks.clear() + channels.append(candidate.domain.channel) if channels: await gather_in_task_group(*(channel.shutdown(context) for channel in channels)) @@ -1254,7 +1258,7 @@ async def release_skipped_external_domains( async def resolve_prefilters( context: Context, comm: Communicator, - bindings: JoinBindings, + join_state: JoinPlanningState, strategy: JoinStrategy, left_partitioning: NormalizedPartitioning, right_partitioning: NormalizedPartitioning, @@ -1263,11 +1267,11 @@ async def resolve_prefilters( ) -> None: """Resolve optional prefilters after selecting the join strategy.""" config = executor.join_filter_pushdown - if config is None or not bindings.prefilters: + if config is None or not join_state.candidates: return choose_prefilters( - bindings, + join_state, strategy, left_partitioning, right_partitioning, @@ -1278,33 +1282,33 @@ async def resolve_prefilters( await collect_samples( context, comm, - bindings, + join_state, tuple( - bound.domain - for bound in bindings.prefilters - if isinstance(bound.prefilter.domain, ExternalDomain) - and bound.decision is None + candidate.domain + for candidate in join_state.candidates + if isinstance(candidate.spec.domain, ExternalDomain) + and candidate.decision is None ), executor.dynamic_planning.sample_chunk_count, executor.target_partition_size, collective_id, ) choose_prefilters( - bindings, + join_state, strategy, left_partitioning, right_partitioning, executor.broadcast_limit, config.bloom_filter_max_size, ) - await release_skipped_external_domains(context, bindings) + await release_skipped_external_domains(context, join_state) async def choose_strategy( context: Context, comm: Communicator, ir: Join, - bindings: JoinBindings, + join_state: JoinPlanningState, executor: StreamingExecutor, collective_ids: JoinCollectiveIds, *, @@ -1312,8 +1316,8 @@ async def choose_strategy( ) -> JoinStrategy: """Collect any required samples and choose broadcast vs shuffle.""" left, right = ir.children[:2] - left_metadata = bindings.left.metadata - right_metadata = bindings.right.metadata + left_metadata = join_state.left.metadata + right_metadata = join_state.right.metadata nranks = comm.nranks left_partitioning = NormalizedPartitioning.from_keys( left_metadata.partitioning, @@ -1333,11 +1337,11 @@ async def choose_strategy( ) if chunkwise: - bindings.left.sample = TableSizeStats( + join_state.left.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=left_metadata.local_count, ) - bindings.right.sample = TableSizeStats( + join_state.right.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=right_metadata.local_count, ) @@ -1346,15 +1350,15 @@ async def choose_strategy( await collect_samples( context, comm, - bindings, - (bindings.left, bindings.right), + join_state, + (join_state.left, join_state.right), executor.dynamic_planning.sample_chunk_count, executor.target_partition_size, collective_ids.size_estimate, ) - left_sample = bindings.left.sample - right_sample = bindings.right.sample + left_sample = join_state.left.sample + right_sample = join_state.right.sample if left_sample is None or right_sample is None: raise ValueError("Join inputs have not been sampled") strategy = _choose_strategy_from_samples( @@ -1373,7 +1377,7 @@ async def choose_strategy( await resolve_prefilters( context, comm, - bindings, + join_state, strategy, left_partitioning, right_partitioning, @@ -1445,7 +1449,7 @@ async def join_actor( *(recv_metadata(ch, context) for ch in ch_prefilter_domains), ) - bindings = bind_join_inputs( + join_state = make_join_planning_state( ir, ch_left, ch_right, @@ -1460,23 +1464,23 @@ async def join_actor( context, comm, ir, - bindings, + join_state, executor, collective_ids, tracer=tracer, ) prefilter_traces = [] - for bound in bindings.prefilters: - if bound.decision is None: + for candidate in join_state.candidates: + if candidate.decision is None: raise ValueError("Join prefilter has no runtime decision") - trace = bound.decision.trace(bound.prefilter) + trace = candidate.decision.trace(candidate.spec) prefilter_traces.append(trace) if LOG_TRACES: - bound.trace = trace + candidate.trace = trace if tracer is not None and prefilter_traces: tracer.set_extra("join_prefilters", prefilter_traces) - left_sample = bindings.left.sample - right_sample = bindings.right.sample + left_sample = join_state.left.sample + right_sample = join_state.right.sample if left_sample is None or right_sample is None: raise ValueError("Join inputs have not been sampled") ch_left_replay = context.create_channel() @@ -1489,7 +1493,7 @@ async def join_actor( strategy, ch_left_replay, ch_right_replay, - bindings, + join_state, collective_ids, ) async with shutdown_on_error( diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py similarity index 70% rename from python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py index 72b57363e592..463496afb893 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_bindings.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Runtime bindings for optional join prefilters.""" +"""Actor-local planning state for dynamic joins and optional prefilters.""" from __future__ import annotations @@ -27,7 +27,7 @@ @dataclass(slots=True) -class JoinInputBinding: +class JoinInput: """Concrete runtime resources for one input to a dynamic join.""" node: IR @@ -37,12 +37,12 @@ class JoinInputBinding: @dataclass(slots=True) -class BoundPrefilter: - """A logical prefilter bound to its concrete runtime inputs.""" +class PrefilterCandidate: + """An optional prefilter and the runtime inputs needed to evaluate it.""" - prefilter: Prefilter - target: JoinInputBinding - domain: JoinInputBinding + spec: Prefilter + target: JoinInput + domain: JoinInput cardinality_tag: int decision: PrefilterDecision | None = None key_channel: Channel[TableChunk] | None = None @@ -50,15 +50,15 @@ class BoundPrefilter: @dataclass(frozen=True, slots=True) -class JoinBindings: - """Concrete runtime inputs and optional prefilters for a dynamic join.""" +class JoinPlanningState: + """Actor-local input and prefilter state for planning a dynamic join.""" - left: JoinInputBinding - right: JoinInputBinding - prefilters: tuple[BoundPrefilter, ...] = () + left: JoinInput + right: JoinInput + candidates: tuple[PrefilterCandidate, ...] = () -def bind_join_inputs( +def make_join_planning_state( ir: Join, ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], @@ -67,17 +67,17 @@ def bind_join_inputs( right_metadata: ChannelMetadata, prefilter_domain_metadata: tuple[ChannelMetadata, ...], cardinality_tags: tuple[int, ...], -) -> JoinBindings: - """Bind logical join inputs and prefilters to their runtime resources.""" - left = JoinInputBinding(ir.children[0], ch_left, left_metadata) - right = JoinInputBinding(ir.children[1], ch_right, right_metadata) +) -> JoinPlanningState: + """Create actor-local planning state from a join and its runtime inputs.""" + left = JoinInput(ir.children[0], ch_left, left_metadata) + right = JoinInput(ir.children[1], ch_right, right_metadata) if not isinstance(ir, JoinWithPrefilter): if ch_prefilter_domains or prefilter_domain_metadata: raise ValueError("A plain Join cannot have prefilter domain inputs") - return JoinBindings(left, right) + return JoinPlanningState(left, right) external_inputs = tuple( - JoinInputBinding(node, channel, metadata) + JoinInput(node, channel, metadata) for node, channel, metadata in zip( ir.children[2:], ch_prefilter_domains, @@ -96,19 +96,19 @@ def bind_join_inputs( sides = {"left": left, "right": right} external_inputs_iter = iter(external_inputs) cardinality_tags_iter = iter(cardinality_tags) - prefilters = [] - for prefilter in ir.prefilters: - target = sides[prefilter.target_side] - if isinstance(prefilter.domain, JoinInputDomain): - domain = sides[prefilter.domain.side] + candidates = [] + for spec in ir.prefilters: + target = sides[spec.target_side] + if isinstance(spec.domain, JoinInputDomain): + domain = sides[spec.domain.side] else: domain = next(external_inputs_iter) - prefilters.append( - BoundPrefilter( - prefilter, + candidates.append( + PrefilterCandidate( + spec, target, domain, cardinality_tag=next(cardinality_tags_iter), ) ) - return JoinBindings(left, right, tuple(prefilters)) + return JoinPlanningState(left, right, tuple(candidates)) From fbd9d92e1684d33e126abb6178c39b7e8cd31f31 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:52:03 +0100 Subject: [PATCH 05/17] Separate prefilter eligibility from method selection --- .../streaming/actor_graph/prefilter.py | 23 +++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py index 390991e0c74d..cdcc37ef0857 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -32,6 +32,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.containers import DataType + from cudf_polars.dsl.expr import NamedExpr from cudf_polars.streaming.actor_graph.utils import TableSizeStats from cudf_polars.streaming.filter_hint import JoinSide, Prefilter @@ -285,7 +286,7 @@ def choose_prefilter( broadcast_limit: int, bloom_filter_max_size: int, ) -> PrefilterDecision: - """Choose whether and how to apply one prefilter.""" + """Choose whether one join prefilter is eligible to be applied.""" domain_rows = None if domain is None else domain.total_rows if not target_requires_redistribution: return PrefilterDecision( @@ -307,6 +308,24 @@ def choose_prefilter( domain_rows, ) + return choose_prefilter_method( + prefilter.domain_on, + target, + domain, + broadcast_limit=broadcast_limit, + bloom_filter_max_size=bloom_filter_max_size, + ) + + +def choose_prefilter_method( + domain_on: Sequence[NamedExpr], + target: TableSizeStats, + domain: TableSizeStats, + *, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> PrefilterDecision: + """Choose the implementation for an eligible prefilter.""" cardinality = estimate_cardinality(domain) if cardinality is None: return PrefilterDecision( @@ -331,7 +350,7 @@ def choose_prefilter( BloomFilter.aligned_size(estimate_bloom_filter_bytes(cardinality)), ) exact_bytes = estimate_bytes( - tuple(key.value.dtype for key in prefilter.domain_on), + tuple(key.value.dtype for key in domain_on), domain.total_rows, ) if bloom_bytes <= min(bloom_filter_max_size, target.total_size): From 6e99daa412e013ff1bde4070b197f288da615cd5 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:55:40 +0100 Subject: [PATCH 06/17] Extract reusable prefilter execution helpers --- .../cudf_polars/streaming/actor_graph/join.py | 80 ++------------ .../streaming/actor_graph/prefilter.py | 100 +++++++++++++++--- 2 files changed, 94 insertions(+), 86 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 3c1987189d83..cdec95bdafdf 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal -from cudf_streaming import BloomFilter, CardinalityEstimator +from cudf_streaming import CardinalityEstimator from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -17,7 +17,6 @@ TableChunk, make_table_chunks_available_or_wait, ) -from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED from rapidsmpf.memory.memory_reservation import opaque_memory_usage from rapidsmpf.streaming.core.actor import define_actor from rapidsmpf.streaming.core.memory_reserve_or_wait import ( @@ -40,9 +39,9 @@ from cudf_polars.streaming.actor_graph.join_planning import make_join_planning_state from cudf_polars.streaming.actor_graph.nodes import default_node_multi from cudf_polars.streaming.actor_graph.prefilter import ( - PrefilterExecution, + JoinPrefilterExecution, + add_bloom_prefilter, choose_prefilter, - count_rows_passthrough, ) from cudf_polars.streaming.actor_graph.tracing import LOG_TRACES, send_chunk from cudf_polars.streaming.actor_graph.utils import ( @@ -73,7 +72,7 @@ from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: - from collections.abc import Iterable, MutableMapping + from collections.abc import MutableMapping from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel @@ -467,73 +466,6 @@ async def _broadcast_join( await ch_out.drain(context) -def add_bloom_prefilter( - context: Context, - comm: Communicator, - bloom_bytes: int, - execution: PrefilterExecution, - target_indices: Iterable[int], - ch_domain_keys: Channel[TableChunk], - ch_target: Channel[TableChunk], - ch_filtered: Channel[TableChunk], - collective_id: int, - trace_stats: dict[str, Any] | None, -) -> None: - """Add the channels and actors for an approximate Bloom prefilter.""" - bloom = BloomFilter( - context, - comm, - LIBCUDF_DEFAULT_HASH_SEED, - bloom_bytes, - ) - ch_filter = context.create_channel() - execution.add_channel(ch_filter) - execution.add_task( - bloom.build( - context, - ch_domain_keys, - ch_filter, - collective_id, - ) - ) - ch_apply_input = ch_target - ch_apply_output = ch_filtered - if trace_stats is not None: - ch_counted_input: Channel[TableChunk] = context.create_channel() - ch_raw_output: Channel[TableChunk] = context.create_channel() - execution.add_channel(ch_counted_input) - execution.add_channel(ch_raw_output) - execution.add_task( - count_rows_passthrough( - context, - ch_target, - ch_counted_input, - trace_stats, - "input_rows", - ) - ) - execution.add_task( - count_rows_passthrough( - context, - ch_raw_output, - ch_filtered, - trace_stats, - "output_rows", - ) - ) - ch_apply_input = ch_counted_input - ch_apply_output = ch_raw_output - execution.add_task( - bloom.apply( - context, - ch_filter, - ch_apply_input, - ch_apply_output, - target_indices, - ) - ) - - def make_prefilter_execution( context: Context, comm: Communicator, @@ -544,9 +476,9 @@ def make_prefilter_execution( ch_right: Channel[TableChunk], join_state: JoinPlanningState, collective_ids: JoinCollectiveIds, -) -> PrefilterExecution: +) -> JoinPrefilterExecution: """Create the actors and channels that realize selected prefilters.""" - execution = PrefilterExecution(context, ch_left, ch_right) + execution = JoinPrefilterExecution(context, ch_left, ch_right) # Prepare every required domain before connecting target-side filters. This # is important for opposing direct filters: each filter must consume the diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py index cdcc37ef0857..fc0e8e3ab253 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -12,6 +12,7 @@ from cudf_streaming import BloomFilter from cudf_streaming.channel_metadata import ChannelMetadata from cudf_streaming.table_chunk import TableChunk +from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED from rapidsmpf.streaming.core.message import Message from cudf_polars.streaming.actor_graph.utils import ( @@ -28,6 +29,7 @@ if TYPE_CHECKING: from collections.abc import Coroutine, Iterable, Sequence + from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context @@ -172,7 +174,24 @@ async def count_rows_passthrough( class PrefilterExecution: - """Channels and actors used to apply prefilters before a join.""" + """Channels and actor tasks used to apply one or more prefilters.""" + + def __init__(self, context: Context) -> None: + self.context = context + self.tasks: list[Coroutine[Any, Any, None]] = [] + self.channels: list[Channel[Any]] = [] + + def add_task(self, task: Coroutine[Any, Any, None]) -> None: + """Add an actor task to the prefilter execution.""" + self.tasks.append(task) + + def add_channel(self, channel: Channel[Any]) -> None: + """Register an auxiliary channel for shutdown on failure.""" + self.channels.append(channel) + + +class JoinPrefilterExecution(PrefilterExecution): + """Channels and actor tasks used to apply prefilters before a join.""" def __init__( self, @@ -180,11 +199,9 @@ def __init__( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], ) -> None: - self.context = context + super().__init__(context) self.source_inputs = {"left": ch_left, "right": ch_right} self.join_inputs = dict(self.source_inputs) - self.tasks: list[Coroutine[Any, Any, None]] = [] - self.channels: list[Channel[Any]] = [] self.buffered_domains: set[JoinSide] = set() def buffer_domain( @@ -221,14 +238,6 @@ def replace_join_input( self.join_inputs[side] = channel self.channels.append(channel) - def add_task(self, task: Coroutine[Any, Any, None]) -> None: - """Add an actor task to the prefilter execution.""" - self.tasks.append(task) - - def add_channel(self, channel: Channel[Any]) -> None: - """Register an auxiliary channel for shutdown on failure.""" - self.channels.append(channel) - @property def left(self) -> Channel[TableChunk]: """Current left join input.""" @@ -240,6 +249,73 @@ def right(self) -> Channel[TableChunk]: return self.join_inputs["right"] +def add_bloom_prefilter( + context: Context, + comm: Communicator, + bloom_bytes: int, + execution: PrefilterExecution, + target_indices: Iterable[int], + ch_domain_keys: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the channels and actors for an approximate Bloom prefilter.""" + bloom = BloomFilter( + context, + comm, + LIBCUDF_DEFAULT_HASH_SEED, + bloom_bytes, + ) + ch_filter = context.create_channel() + execution.add_channel(ch_filter) + execution.add_task( + bloom.build( + context, + ch_domain_keys, + ch_filter, + collective_id, + ) + ) + ch_apply_input = ch_target + ch_apply_output = ch_filtered + if trace_stats is not None: + ch_counted_input: Channel[TableChunk] = context.create_channel() + ch_raw_output: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_counted_input) + execution.add_channel(ch_raw_output) + execution.add_task( + count_rows_passthrough( + context, + ch_target, + ch_counted_input, + trace_stats, + "input_rows", + ) + ) + execution.add_task( + count_rows_passthrough( + context, + ch_raw_output, + ch_filtered, + trace_stats, + "output_rows", + ) + ) + ch_apply_input = ch_counted_input + ch_apply_output = ch_raw_output + execution.add_task( + bloom.apply( + context, + ch_filter, + ch_apply_input, + ch_apply_output, + target_indices, + ) + ) + + def estimate_cardinality(stats: TableSizeStats) -> int | None: """Extrapolate sampled distinct count to the estimated full row count.""" if stats.total_rows == 0: From a45112f4937b7a51a094acd135bc74f1fab10ef3 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:56:33 +0100 Subject: [PATCH 07/17] Expose reusable adaptive filtering utilities --- .../cudf_polars/streaming/actor_graph/join.py | 49 +++---------------- .../streaming/actor_graph/utils.py | 36 ++++++++++++++ 2 files changed, 42 insertions(+), 43 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index cdec95bdafdf..2d67225f5ce6 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -52,7 +52,7 @@ NormalizedPartitioning, TableSizeStats, _sample_chunks, - allgather_reduce, + aggregate_table_size_stats, chunk_to_frame, empty_table_chunk, gather_in_task_group, @@ -210,7 +210,7 @@ async def broadcast_join_actor( trace_ir=ir, ir_context=ir_context, ) as tracer: - await _broadcast_join( + await broadcast_join( context, comm, ir, @@ -351,7 +351,7 @@ async def _broadcast_join_large_chunk( return output_rows -async def _broadcast_join( +async def broadcast_join( context: Context, comm: Communicator, ir: Join, @@ -558,7 +558,7 @@ def make_prefilter_execution( projected_domain, ) execution.add_task( - _broadcast_join( + broadcast_join( context, comm, semi_join, @@ -840,43 +840,6 @@ def _num_indices(partitioning: NormalizedPartitioning) -> int: ) -async def aggregate_estimates( - context: Context, - comm: Communicator, - samples: tuple[TableSizeStats, ...], - collective_id: int, -) -> tuple[TableSizeStats, ...]: - """Aggregate table-size and row estimates across ranks.""" - # AllGather size, row, chunk count, and completeness estimates across ranks. - totals = await allgather_reduce( - context, - comm, - collective_id, - *( - value - for sample in samples - for value in ( - sample.total_size, - sample.total_rows, - sample.total_chunks, - int(sample.is_complete), - ) - ), - ) - totals_iter = iter(totals) - return tuple( - TableSizeStats( - chunks=sample.chunks, - total_size=next(totals_iter), - total_rows=next(totals_iter), - total_chunks=next(totals_iter), - is_complete=next(totals_iter) == comm.nranks, - cardinality=sample.cardinality, - ) - for sample in samples - ) - - def _choose_strategy_from_samples( comm: Communicator, ir: Join, @@ -1158,7 +1121,7 @@ async def collect_samples( for input_, candidate in sampling_inputs ) ) - samples = await aggregate_estimates( + samples = await aggregate_table_size_stats( context, comm, tuple(local_samples), @@ -1460,7 +1423,7 @@ async def join_actor( if strategy.broadcast_side is not None: actor_tasks.append( - _broadcast_join( + broadcast_join( context, comm, ir, 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 b8d6a47c2c8e..1f652b52417a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -1040,6 +1040,42 @@ class TableSizeStats: """Global cardinality statistics for the sampled rows, when requested.""" +async def aggregate_table_size_stats( + context: Context, + comm: Communicator, + samples: tuple[TableSizeStats, ...], + collective_id: int, +) -> tuple[TableSizeStats, ...]: + """Aggregate table-size and row estimates across ranks.""" + totals = await allgather_reduce( + context, + comm, + collective_id, + *( + value + for sample in samples + for value in ( + sample.total_size, + sample.total_rows, + sample.total_chunks, + int(sample.is_complete), + ) + ), + ) + totals_iter = iter(totals) + return tuple( + TableSizeStats( + chunks=sample.chunks, + total_size=next(totals_iter), + total_rows=next(totals_iter), + total_chunks=next(totals_iter), + is_complete=next(totals_iter) == comm.nranks, + cardinality=sample.cardinality, + ) + for sample in samples + ) + + @dataclass(frozen=True) class ChunkSampler: """Object for obtaining statistics from a channel of TableChunks.""" From e0d48bf72506fabd8e6cebcf5bb3ba45cf2f1ca8 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:04:59 +0100 Subject: [PATCH 08/17] Execute non-adjacent pushdown filter hints --- .../streaming/actor_graph/__init__.py | 3 +- .../actor_graph/collectives/common.py | 2 + .../cudf_polars/streaming/actor_graph/core.py | 3 +- .../streaming/actor_graph/prefilter.py | 12 +- .../streaming/actor_graph/prefilter_actor.py | 339 ++++++++++++++++++ .../cudf_polars/cudf_polars/streaming/join.py | 14 +- .../tests/streaming/test_tracing.py | 116 +++++- 7 files changed, 475 insertions(+), 14 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py index a4dd049a0eae..35589eab867e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """RapidsMPF streaming-engine support.""" @@ -15,6 +15,7 @@ import cudf_polars.streaming.actor_graph.io import cudf_polars.streaming.actor_graph.join import cudf_polars.streaming.actor_graph.over +import cudf_polars.streaming.actor_graph.prefilter_actor import cudf_polars.streaming.actor_graph.repartition import cudf_polars.streaming.actor_graph.union # noqa: F401 diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py index bae39d58cc08..7cb34f02d107 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py @@ -12,6 +12,7 @@ from cudf_polars.dsl.ir import Distinct, GroupBy, Sort from cudf_polars.dsl.traversal import traversal +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.io import StreamingSink from cudf_polars.streaming.join import Join from cudf_polars.streaming.over import Over @@ -107,6 +108,7 @@ def __init__( GroupBy, Distinct, Over, + PushdownFilterHint, ) self.collective_nodes: list[IR] = [ diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index a1e78f43d42c..dc631f8f0dae 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -23,6 +23,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.io import StreamingScan from cudf_polars.streaming.over import Over from cudf_polars.utils.config import SPMDContext @@ -178,7 +179,7 @@ def _mark_children_unbounded(node: IR) -> None: for node in traversal([ir]): if node in unbounded: _mark_children_unbounded(node) - elif isinstance(node, (Union, Join, Over)): + elif isinstance(node, (Union, Join, Over, PushdownFilterHint)): # Union processes children sequentially; Join may broadcast one # side; Over buffers (or samples-then-replays) its input before # producing output. In every case the input source needs diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py index fc0e8e3ab253..ed0ff4d3edc0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -75,10 +75,9 @@ class PrefilterDecision: bloom_bytes: int | None = None exact_bytes: int | None = None - def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: - """Return serializable actor-trace information.""" - result: dict[str, str | int | None] = { - "target_side": prefilter.target_side, + def trace_details(self) -> dict[str, str | int | None]: + """Return trace information common to every prefilter placement.""" + return { "method": self.method, "reason": self.reason, "target_bytes": self.target_bytes, @@ -87,6 +86,11 @@ def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: "bloom_bytes": self.bloom_bytes, "exact_bytes": self.exact_bytes, } + + def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: + """Return serializable actor-trace information.""" + result = self.trace_details() + result["target_side"] = prefilter.target_side if isinstance(prefilter.domain, JoinInputDomain): result["domain_side"] = prefilter.domain.side else: diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py new file mode 100644 index 000000000000..e501c9427717 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -0,0 +1,339 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Standalone execution of optional pushdown-filter hints.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from cudf_streaming import CardinalityEstimator +from rapidsmpf.streaming.core.actor import define_actor + +from cudf_polars.dsl.ir import Join, Projection +from cudf_polars.dsl.utils.naming import names_to_indices +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.join import JoinStrategy, broadcast_join +from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterExecution, + add_bloom_prefilter, + choose_prefilter_method, +) +from cudf_polars.streaming.actor_graph.utils import ( + ChannelManager, + _sample_chunks, + aggregate_table_size_stats, + gather_in_task_group, + process_children, + recv_metadata, + replay_buffered_channel, + shutdown_on_error, +) +from cudf_polars.streaming.filter_hint import PushdownFilterHint + +if TYPE_CHECKING: + from cudf_streaming.channel_metadata import ChannelMetadata + from cudf_streaming.table_chunk import TableChunk + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.prefilter import PrefilterDecision + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.utils.config import StreamingExecutor + + +def make_broadcast_semi_join(ir: PushdownFilterHint) -> Join: + """Build the synthetic semi-join used for exact filtering.""" + target, domain = ir.children + domain_schema = {key.name: key.value.dtype for key in ir.domain_on} + if len(domain_schema) != len(ir.domain_on): + raise ValueError("Broadcast semi-join keys must have unique names") + projected_domain = Projection(domain_schema, domain) + return Join( + target.schema, + ir.target_on, + ir.domain_on, + ("Semi", ir.nulls_equal, None, "", False, "none"), + target, + projected_domain, + ) + + +async def sample_prefilter_inputs( + context: Context, + comm: Communicator, + ir: PushdownFilterHint, + ch_target: Channel[TableChunk], + ch_domain: Channel[TableChunk], + target_metadata: ChannelMetadata, + domain_metadata: ChannelMetadata, + executor: StreamingExecutor, + collective_id: int, +) -> tuple[TableSizeStats, TableSizeStats]: + """Sample and aggregate a standalone hint's target and domain inputs.""" + dynamic_planning = executor.dynamic_planning + if dynamic_planning is None: + raise ValueError("Standalone prefilters require dynamic planning") + local_samples = await gather_in_task_group( + _sample_chunks( + context, + ch_target, + dynamic_planning.sample_chunk_count, + executor.target_partition_size, + target_metadata.local_count, + ), + _sample_chunks( + context, + ch_domain, + dynamic_planning.sample_chunk_count, + executor.target_partition_size, + domain_metadata.local_count, + cardinality_estimator=CardinalityEstimator( + context, + comm, + tag=collective_id, + ), + cardinality_columns=names_to_indices(ir.domain_on, ir.children[1].schema), + ), + ) + target_sample, domain_sample = await aggregate_table_size_stats( + context, + comm, + tuple(local_samples), + collective_id, + ) + return target_sample, domain_sample + + +async def replay_skipped_prefilter( + context: Context, + ir: PushdownFilterHint, + ch_out: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_domain: Channel[TableChunk], + target_metadata: ChannelMetadata, + target_sample: TableSizeStats, + domain_sample: TableSizeStats, +) -> None: + """Replay an unfiltered target and stop its unused domain input.""" + domain_sample.chunks.clear() + await gather_in_task_group( + replay_buffered_channel( + context, + ch_out, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ), + ch_domain.shutdown(context), + ) + + +async def apply_prefilter( + context: Context, + comm: Communicator, + ir: PushdownFilterHint, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_domain: Channel[TableChunk], + target_metadata: ChannelMetadata, + domain_metadata: ChannelMetadata, + target_sample: TableSizeStats, + domain_sample: TableSizeStats, + decision: PrefilterDecision, + collective_id: int, + trace_stats: dict[str, Any] | None, +) -> None: + """Apply the selected standalone prefilter implementation.""" + target, domain = ir.children + domain_indices = names_to_indices(ir.domain_on, domain.schema) + if domain_indices != tuple(range(len(domain.schema))): + raise ValueError("Pushdown filter domains must contain only keys") + + execution = PrefilterExecution(context) + ch_target_replay: Channel[TableChunk] = context.create_channel() + ch_domain_replay: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_target_replay) + execution.add_channel(ch_domain_replay) + execution.add_task( + replay_buffered_channel( + context, + ch_target_replay, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ) + ) + execution.add_task( + replay_buffered_channel( + context, + ch_domain_replay, + ch_domain, + domain_sample.chunks, + domain_metadata, + trace_ir=ir, + ) + ) + + if decision.method == "bloom": + if decision.bloom_bytes is None: + raise ValueError("Bloom prefilter decision has no filter size") + add_bloom_prefilter( + context, + comm, + decision.bloom_bytes, + execution, + names_to_indices(ir.target_on, target.schema), + ch_domain_replay, + ch_target_replay, + ch_out, + collective_id, + trace_stats, + ) + elif decision.method == "broadcast_semi_join": + execution.add_task( + broadcast_join( + context, + comm, + make_broadcast_semi_join(ir), + ir_context, + ch_out, + ch_target_replay, + ch_domain_replay, + JoinStrategy(broadcast_side="right"), + collective_id, + target_partition_size=None, + tracer=None, + trace_stats=trace_stats, + ) + ) + else: + raise ValueError(f"Cannot apply prefilter method {decision.method!r}") + + async with shutdown_on_error( + context, + *execution.channels, + trace_ir=ir, + ir_context=ir_context, + ): + await gather_in_task_group(*execution.tasks) + + +@define_actor() +async def pushdown_filter_actor( + context: Context, + comm: Communicator, + ir: PushdownFilterHint, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_domain: Channel[TableChunk], + executor: StreamingExecutor, + collective_id: int, +) -> None: + """Choose and optionally execute one standalone pushdown-filter hint.""" + samples: tuple[TableSizeStats, TableSizeStats] | None = None + async with shutdown_on_error( + context, + ch_out, + ch_target, + ch_domain, + trace_ir=ir, + ir_context=ir_context, + ) as tracer: + try: + target_metadata, domain_metadata = await gather_in_task_group( + recv_metadata(ch_target, context), + recv_metadata(ch_domain, context), + ) + samples = await sample_prefilter_inputs( + context, + comm, + ir, + ch_target, + ch_domain, + target_metadata, + domain_metadata, + executor, + collective_id, + ) + target_sample, domain_sample = samples + config = executor.join_filter_pushdown + if config is None: + raise ValueError("Standalone prefilter has no runtime configuration") + decision = choose_prefilter_method( + ir.domain_on, + target_sample, + domain_sample, + broadcast_limit=executor.broadcast_limit, + bloom_filter_max_size=config.bloom_filter_max_size, + ) + trace = decision.trace_details() + trace["placement"] = "standalone" + trace_stats = trace if tracer is not None else None + if tracer is not None: + tracer.decision = decision.method + tracer.set_extra("prefilter", trace) + + if decision.method == "skip": + await replay_skipped_prefilter( + context, + ir, + ch_out, + ch_target, + ch_domain, + target_metadata, + target_sample, + domain_sample, + ) + else: + await apply_prefilter( + context, + comm, + ir, + ir_context, + ch_out, + ch_target, + ch_domain, + target_metadata, + domain_metadata, + target_sample, + domain_sample, + decision, + collective_id, + trace_stats, + ) + finally: + if samples is not None: + for sample in samples: + sample.chunks.clear() + + +@generate_ir_sub_network.register(PushdownFilterHint) +def generate_pushdown_filter_subnetwork( + ir: PushdownFilterHint, rec: SubNetGenerator +) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: + """Generate the actor subnetwork for a standalone filter hint.""" + target, domain = ir.children + actors, channels = process_children(ir, rec) + channels[ir] = ChannelManager(rec.state["context"]) + (collective_id,) = rec.state["collective_id_map"][ir] + actors[ir] = [ + pushdown_filter_actor( + rec.state["context"], + rec.state["comm"], + ir, + rec.state["ir_context"], + channels[ir].reserve_input_slot(), + channels[target].reserve_output_slot(), + channels[domain].reserve_output_slot(), + rec.state["config_options"].executor, + collective_id, + ) + ] + return actors, channels diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 4218ac6bfaf7..02d0a2893e0a 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -254,9 +254,17 @@ def _lower_join_with_prefilters( def _( ir: PushdownFilterHint, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - """Discard the optional filter without lowering its domain.""" - target, _domain = ir.children - return rec(target) + """Preserve optional filters for dynamic execution, otherwise discard them.""" + target, domain = ir.children + target, partition_info = rec(target) + if not _dynamic_planning_on(rec.state["config_options"]): + return target, partition_info + + domain, domain_partition_info = rec(domain) + partition_info.update(domain_partition_info) + lowered = ir.reconstruct((target, domain)) + partition_info[lowered] = partition_info[target] + return lowered, partition_info @lower_ir_node.register(ConditionalJoin) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index aeb6c16c6a7c..cac488b2937b 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -242,11 +242,117 @@ def test_local_join_prefilter_trace_records_decision_and_effect( assert record["prefilter"]["output_rows"] == output_rows +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,method,reason,output_rows", + [ + (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 20), + (64, 0, "broadcast_semi_join", "exact_domain_fits", 20), + (1, 0, "skip", "no_viable_filter", None), + ], + ids=["bloom", "exact", "skip"], +) +def test_standalone_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + method: str, + reason: str, + output_rows: int | None, +) -> None: + """Trace a non-adjacent prefilter selected through the public engine.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + domain = ( + pl.LazyFrame( + {{"p_partkey": range(10), "active": [True] * 2 + [False] * 8}} + ) + .filter("active") + .select("p_partkey") + ) + target = pl.LazyFrame( + {{ + "l_partkey": [i % 10 for i in range(100)], + "value": range(100), + }} + ).with_columns((pl.col("value") + 1).alias("derived")) + query = domain.join(target, left_on="p_partkey", right_on="l_partkey") + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 10, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" + and log.get("prefilter", {{}}).get("placement") == "standalone" + ) + record = {{ + "result_rows": result.height, + "decision": event["decision"], + "prefilter": event["prefilter"], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 20 + assert record["decision"] == method + assert ( + record["prefilter"].items() + >= { + "placement": "standalone", + "method": method, + "reason": reason, + "domain_rows": 2, + }.items() + ) + if output_rows is None: + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 2 + assert record["prefilter"]["input_rows"] == 100 + assert record["prefilter"]["output_rows"] == output_rows + + @pytest.mark.parametrize( "broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,domain_rows", [ - (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 30), - (512, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 30), + (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 15), + (512, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 15), ( 1_000_000, 32 * 1024 * 1024, @@ -385,12 +491,12 @@ def test_external_join_prefilter_trace_records_decision_and_effect( assert "input_rows" not in record["prefilter"] assert "output_rows" not in record["prefilter"] else: - assert record["prefilter"]["estimated_cardinality"] == 30 + assert record["prefilter"]["estimated_cardinality"] == domain_rows assert record["prefilter"]["input_rows"] == 180 if method == "broadcast_semi_join": - assert record["prefilter"]["output_rows"] == 90 + assert record["prefilter"]["output_rows"] == 45 else: - assert 90 <= record["prefilter"]["output_rows"] < 180 + assert 45 <= record["prefilter"]["output_rows"] < 180 def test_structlog_disabled_by_default(timeout_seconds: int): From c1b9a2f2c9a3da04db0badd6abb88954264e56d7 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:29:13 +0100 Subject: [PATCH 09/17] Share adaptive input sampling --- .../cudf_polars/streaming/actor_graph/join.py | 80 ++++++----------- .../streaming/actor_graph/prefilter_actor.py | 88 +++++++------------ .../streaming/actor_graph/utils.py | 20 +++++ 3 files changed, 81 insertions(+), 107 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 2d67225f5ce6..0bd27b90e5ed 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -48,11 +48,10 @@ CUDF_ROW_LIMIT, MAX_ROWS_PER_PARTITION, ChannelManager, + ChunkSampler, ChunkStore, NormalizedPartitioning, TableSizeStats, - _sample_chunks, - aggregate_table_size_stats, chunk_to_frame, empty_table_chunk, gather_in_task_group, @@ -60,6 +59,7 @@ process_children, recv_metadata, replay_buffered_channel, + sample_inputs, send_metadata, shutdown_on_error, ) @@ -84,7 +84,6 @@ from cudf_polars.streaming.actor_graph.join_planning import ( JoinInput, JoinPlanningState, - PrefilterCandidate, ) from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo @@ -1049,43 +1048,6 @@ def choose_prefilters( ) -async def sample_input( - context: Context, - comm: Communicator, - input_: JoinInput, - candidate: PrefilterCandidate | None, - sample_chunk_count: int, - target_partition_size: int, -) -> TableSizeStats: - """Sample one join-planning input and optionally estimate cardinality.""" - if candidate is None: - cardinality_estimator = None - cardinality_columns: tuple[int, ...] = () - else: - cardinality_estimator = CardinalityEstimator( - context, - comm, - tag=candidate.cardinality_tag, - ) - cardinality_columns = names_to_indices( - candidate.spec.domain_on, - input_.node.schema, - ) - assert len(cardinality_columns) == len(candidate.spec.domain_on), ( - "Prefilter domain keys must be columns" - ) - - return await _sample_chunks( - context, - input_.channel, - sample_chunk_count, - target_partition_size, - input_.metadata.local_count, - cardinality_estimator=cardinality_estimator, - cardinality_columns=cardinality_columns, - ) - - async def collect_samples( context: Context, comm: Communicator, @@ -1108,23 +1070,39 @@ async def collect_samples( if len(candidates) > 1: raise ValueError("One join input cannot provide multiple prefilter domains") sampling_inputs.append((input_, candidates[0] if candidates else None)) - local_samples = await gather_in_task_group( - *( - sample_input( + samplers = [] + for input_, candidate in sampling_inputs: + if candidate is None: + cardinality_estimator = None + cardinality_columns: tuple[int, ...] = () + else: + cardinality_estimator = CardinalityEstimator( context, comm, - input_, - candidate, - sample_chunk_count, - target_partition_size, + tag=candidate.cardinality_tag, + ) + cardinality_columns = names_to_indices( + candidate.spec.domain_on, + input_.node.schema, + ) + assert len(cardinality_columns) == len(candidate.spec.domain_on), ( + "Prefilter domain keys must be columns" + ) + samplers.append( + ChunkSampler( + context=context, + ch_in=input_.channel, + max_chunks=sample_chunk_count, + max_bytes=target_partition_size, + ch_in_chunk_count=input_.metadata.local_count, + cardinality_estimator=cardinality_estimator, + cardinality_columns=cardinality_columns, ) - for input_, candidate in sampling_inputs ) - ) - samples = await aggregate_table_size_stats( + samples = await sample_inputs( context, comm, - tuple(local_samples), + samplers, collective_id, ) for (input_, _), sample in zip(sampling_inputs, samples, strict=True): diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py index e501c9427717..c1939aed0d46 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -20,12 +20,12 @@ ) from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, - _sample_chunks, - aggregate_table_size_stats, + ChunkSampler, gather_in_task_group, process_children, recv_metadata, replay_buffered_channel, + sample_inputs, shutdown_on_error, ) from cudf_polars.streaming.filter_hint import PushdownFilterHint @@ -61,52 +61,6 @@ def make_broadcast_semi_join(ir: PushdownFilterHint) -> Join: ) -async def sample_prefilter_inputs( - context: Context, - comm: Communicator, - ir: PushdownFilterHint, - ch_target: Channel[TableChunk], - ch_domain: Channel[TableChunk], - target_metadata: ChannelMetadata, - domain_metadata: ChannelMetadata, - executor: StreamingExecutor, - collective_id: int, -) -> tuple[TableSizeStats, TableSizeStats]: - """Sample and aggregate a standalone hint's target and domain inputs.""" - dynamic_planning = executor.dynamic_planning - if dynamic_planning is None: - raise ValueError("Standalone prefilters require dynamic planning") - local_samples = await gather_in_task_group( - _sample_chunks( - context, - ch_target, - dynamic_planning.sample_chunk_count, - executor.target_partition_size, - target_metadata.local_count, - ), - _sample_chunks( - context, - ch_domain, - dynamic_planning.sample_chunk_count, - executor.target_partition_size, - domain_metadata.local_count, - cardinality_estimator=CardinalityEstimator( - context, - comm, - tag=collective_id, - ), - cardinality_columns=names_to_indices(ir.domain_on, ir.children[1].schema), - ), - ) - target_sample, domain_sample = await aggregate_table_size_stats( - context, - comm, - tuple(local_samples), - collective_id, - ) - return target_sample, domain_sample - - async def replay_skipped_prefilter( context: Context, ir: PushdownFilterHint, @@ -251,18 +205,40 @@ async def pushdown_filter_actor( recv_metadata(ch_target, context), recv_metadata(ch_domain, context), ) - samples = await sample_prefilter_inputs( + dynamic_planning = executor.dynamic_planning + if dynamic_planning is None: + raise ValueError("Standalone prefilters require dynamic planning") + collected_samples = await sample_inputs( context, comm, - ir, - ch_target, - ch_domain, - target_metadata, - domain_metadata, - executor, + ( + ChunkSampler( + context=context, + ch_in=ch_target, + max_chunks=dynamic_planning.sample_chunk_count, + max_bytes=executor.target_partition_size, + ch_in_chunk_count=target_metadata.local_count, + ), + ChunkSampler( + context=context, + ch_in=ch_domain, + max_chunks=dynamic_planning.sample_chunk_count, + max_bytes=executor.target_partition_size, + ch_in_chunk_count=domain_metadata.local_count, + cardinality_estimator=CardinalityEstimator( + context, comm, tag=collective_id + ), + cardinality_columns=names_to_indices( + ir.domain_on, ir.children[1].schema + ), + ), + ), collective_id, ) - target_sample, domain_sample = samples + if len(collected_samples) != 2: + raise ValueError("Standalone prefilters require two input samples") + target_sample, domain_sample = collected_samples + samples = (target_sample, domain_sample) config = executor.join_filter_pushdown if config is None: raise ValueError("Standalone prefilter has no runtime configuration") 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 1f652b52417a..55192ef4f755 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -1204,6 +1204,26 @@ async def sample(self) -> TableSizeStats: ) +async def sample_inputs( + context: Context, + comm: Communicator, + samplers: Sequence[ChunkSampler], + collective_id: int, +) -> tuple[TableSizeStats, ...]: + """Sample input channels concurrently and aggregate their statistics.""" + if not samplers: + return () + local_samples = await gather_in_task_group( + *(sampler.sample() for sampler in samplers) + ) + return await aggregate_table_size_stats( + context, + comm, + tuple(local_samples), + collective_id, + ) + + async def _sample_chunks( context: Context, ch: Channel[TableChunk], From 283d7dc44b15330e7df2b4d0ccb493c34ca3936a Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:42:41 +0100 Subject: [PATCH 10/17] Share adaptive prefilter execution --- .../cudf_polars/streaming/actor_graph/join.py | 133 ++++++++----- .../streaming/actor_graph/prefilter_actor.py | 177 +++++------------- 2 files changed, 137 insertions(+), 173 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 0bd27b90e5ed..962df2ca19c0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -85,9 +85,17 @@ JoinInput, JoinPlanningState, ) + from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterDecision, + PrefilterExecution, + ) from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo - from cudf_polars.streaming.filter_hint import JoinSide + from cudf_polars.streaming.filter_hint import ( + JoinSide, + Prefilter, + PushdownFilterHint, + ) from cudf_polars.utils.config import StreamingExecutor @@ -465,6 +473,70 @@ async def broadcast_join( await ch_out.drain(context) +def add_prefilter( + execution: PrefilterExecution, + comm: Communicator, + *, + spec: Prefilter | PushdownFilterHint, + decision: PrefilterDecision, + target: IR, + domain: IR, + ch_target: Channel[TableChunk], + ch_domain_keys: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + ir_context: IRExecutionContext, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the actors and channels that apply one selected prefilter.""" + context = execution.context + if decision.method == "bloom": + if decision.bloom_bytes is None: + raise ValueError("Bloom prefilter decision has no filter size") + add_bloom_prefilter( + context, + comm, + decision.bloom_bytes, + execution, + names_to_indices(spec.target_on, target.schema), + ch_domain_keys, + ch_target, + ch_filtered, + collective_id, + trace_stats, + ) + elif decision.method == "broadcast_semi_join": + domain_schema = {key.name: key.value.dtype for key in spec.domain_on} + if len(domain_schema) != len(spec.domain_on): + raise ValueError("Broadcast semi-join keys must have unique names") + semi_join = Join( + target.schema, + spec.target_on, + spec.domain_on, + ("Semi", spec.nulls_equal, None, "", False, "none"), + target, + Projection(domain_schema, domain), + ) + execution.add_task( + broadcast_join( + context, + comm, + semi_join, + ir_context, + ch_filtered, + ch_target, + ch_domain_keys, + JoinStrategy(broadcast_side="right"), + collective_id, + target_partition_size=None, + tracer=None, + trace_stats=trace_stats, + ) + ) + else: + raise ValueError(f"Cannot apply prefilter method {decision.method!r}") + + def make_prefilter_execution( context: Context, comm: Communicator, @@ -527,51 +599,20 @@ def make_prefilter_execution( ch_filtered: Channel[TableChunk] = context.create_channel() trace_stats = candidate.trace - collective_id = collective_ids.prefilter(strategy, target_side) - if decision.method == "bloom": - assert decision.bloom_bytes is not None - add_bloom_prefilter( - context, - comm, - decision.bloom_bytes, - execution, - names_to_indices(spec.target_on, target.schema), - ch_domain_keys, - ch_target, - ch_filtered, - collective_id, - trace_stats, - ) - else: - assert decision.method == "broadcast_semi_join" - domain_schema = {key.name: key.value.dtype for key in spec.domain_on} - if len(domain_schema) != len(spec.domain_on): - raise ValueError("Broadcast semi-join keys must have unique names") - projected_domain = Projection(domain_schema, candidate.domain.node) - semi_join = Join( - target.schema, - spec.target_on, - spec.domain_on, - ("Semi", spec.nulls_equal, None, "", False, "none"), - target, - projected_domain, - ) - execution.add_task( - broadcast_join( - context, - comm, - semi_join, - ir_context, - ch_filtered, - ch_target, - ch_domain_keys, - JoinStrategy(broadcast_side="right"), - collective_id, - target_partition_size=None, - tracer=None, - trace_stats=trace_stats, - ) - ) + add_prefilter( + execution, + comm, + spec=spec, + decision=decision, + target=target, + domain=candidate.domain.node, + ch_target=ch_target, + ch_domain_keys=ch_domain_keys, + ch_filtered=ch_filtered, + collective_id=collective_ids.prefilter(strategy, target_side), + ir_context=ir_context, + trace_stats=trace_stats, + ) execution.replace_join_input(target_side, ch_filtered) return execution diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py index c1939aed0d46..4381e8982184 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -9,13 +9,11 @@ from cudf_streaming import CardinalityEstimator from rapidsmpf.streaming.core.actor import define_actor -from cudf_polars.dsl.ir import Join, Projection from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network -from cudf_polars.streaming.actor_graph.join import JoinStrategy, broadcast_join +from cudf_polars.streaming.actor_graph.join import add_prefilter from cudf_polars.streaming.actor_graph.prefilter import ( PrefilterExecution, - add_bloom_prefilter, choose_prefilter_method, ) from cudf_polars.streaming.actor_graph.utils import ( @@ -39,28 +37,10 @@ from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator - from cudf_polars.streaming.actor_graph.prefilter import PrefilterDecision from cudf_polars.streaming.actor_graph.utils import TableSizeStats from cudf_polars.utils.config import StreamingExecutor -def make_broadcast_semi_join(ir: PushdownFilterHint) -> Join: - """Build the synthetic semi-join used for exact filtering.""" - target, domain = ir.children - domain_schema = {key.name: key.value.dtype for key in ir.domain_on} - if len(domain_schema) != len(ir.domain_on): - raise ValueError("Broadcast semi-join keys must have unique names") - projected_domain = Projection(domain_schema, domain) - return Join( - target.schema, - ir.target_on, - ir.domain_on, - ("Semi", ir.nulls_equal, None, "", False, "none"), - target, - projected_domain, - ) - - async def replay_skipped_prefilter( context: Context, ir: PushdownFilterHint, @@ -86,98 +66,6 @@ async def replay_skipped_prefilter( ) -async def apply_prefilter( - context: Context, - comm: Communicator, - ir: PushdownFilterHint, - ir_context: IRExecutionContext, - ch_out: Channel[TableChunk], - ch_target: Channel[TableChunk], - ch_domain: Channel[TableChunk], - target_metadata: ChannelMetadata, - domain_metadata: ChannelMetadata, - target_sample: TableSizeStats, - domain_sample: TableSizeStats, - decision: PrefilterDecision, - collective_id: int, - trace_stats: dict[str, Any] | None, -) -> None: - """Apply the selected standalone prefilter implementation.""" - target, domain = ir.children - domain_indices = names_to_indices(ir.domain_on, domain.schema) - if domain_indices != tuple(range(len(domain.schema))): - raise ValueError("Pushdown filter domains must contain only keys") - - execution = PrefilterExecution(context) - ch_target_replay: Channel[TableChunk] = context.create_channel() - ch_domain_replay: Channel[TableChunk] = context.create_channel() - execution.add_channel(ch_target_replay) - execution.add_channel(ch_domain_replay) - execution.add_task( - replay_buffered_channel( - context, - ch_target_replay, - ch_target, - target_sample.chunks, - target_metadata, - trace_ir=ir, - ) - ) - execution.add_task( - replay_buffered_channel( - context, - ch_domain_replay, - ch_domain, - domain_sample.chunks, - domain_metadata, - trace_ir=ir, - ) - ) - - if decision.method == "bloom": - if decision.bloom_bytes is None: - raise ValueError("Bloom prefilter decision has no filter size") - add_bloom_prefilter( - context, - comm, - decision.bloom_bytes, - execution, - names_to_indices(ir.target_on, target.schema), - ch_domain_replay, - ch_target_replay, - ch_out, - collective_id, - trace_stats, - ) - elif decision.method == "broadcast_semi_join": - execution.add_task( - broadcast_join( - context, - comm, - make_broadcast_semi_join(ir), - ir_context, - ch_out, - ch_target_replay, - ch_domain_replay, - JoinStrategy(broadcast_side="right"), - collective_id, - target_partition_size=None, - tracer=None, - trace_stats=trace_stats, - ) - ) - else: - raise ValueError(f"Cannot apply prefilter method {decision.method!r}") - - async with shutdown_on_error( - context, - *execution.channels, - trace_ir=ir, - ir_context=ir_context, - ): - await gather_in_task_group(*execution.tasks) - - @define_actor() async def pushdown_filter_actor( context: Context, @@ -268,22 +156,57 @@ async def pushdown_filter_actor( domain_sample, ) else: - await apply_prefilter( - context, + target, domain = ir.children + domain_indices = names_to_indices(ir.domain_on, domain.schema) + if domain_indices != tuple(range(len(domain.schema))): + raise ValueError("Pushdown filter domains must contain only keys") + + execution = PrefilterExecution(context) + ch_target_replay: Channel[TableChunk] = context.create_channel() + ch_domain_replay: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_target_replay) + execution.add_channel(ch_domain_replay) + execution.add_task( + replay_buffered_channel( + context, + ch_target_replay, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ) + ) + execution.add_task( + replay_buffered_channel( + context, + ch_domain_replay, + ch_domain, + domain_sample.chunks, + domain_metadata, + trace_ir=ir, + ) + ) + add_prefilter( + execution, comm, - ir, - ir_context, - ch_out, - ch_target, - ch_domain, - target_metadata, - domain_metadata, - target_sample, - domain_sample, - decision, - collective_id, - trace_stats, + spec=ir, + decision=decision, + target=target, + domain=domain, + ch_target=ch_target_replay, + ch_domain_keys=ch_domain_replay, + ch_filtered=ch_out, + collective_id=collective_id, + ir_context=ir_context, + trace_stats=trace_stats, ) + async with shutdown_on_error( + context, + *execution.channels, + trace_ir=ir, + ir_context=ir_context, + ): + await gather_in_task_group(*execution.tasks) finally: if samples is not None: for sample in samples: From 735340f19a4c0b71cb06d72305c6f53e2b1d3275 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:52:59 +0100 Subject: [PATCH 11/17] Tidy --- .../streaming/actor_graph/prefilter_actor.py | 56 ++++++------------- 1 file changed, 16 insertions(+), 40 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py index 4381e8982184..b4a9a4a0ca57 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -29,7 +29,8 @@ from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: - from cudf_streaming.channel_metadata import ChannelMetadata + from collections.abc import Sequence + from cudf_streaming.table_chunk import TableChunk from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel @@ -41,31 +42,6 @@ from cudf_polars.utils.config import StreamingExecutor -async def replay_skipped_prefilter( - context: Context, - ir: PushdownFilterHint, - ch_out: Channel[TableChunk], - ch_target: Channel[TableChunk], - ch_domain: Channel[TableChunk], - target_metadata: ChannelMetadata, - target_sample: TableSizeStats, - domain_sample: TableSizeStats, -) -> None: - """Replay an unfiltered target and stop its unused domain input.""" - domain_sample.chunks.clear() - await gather_in_task_group( - replay_buffered_channel( - context, - ch_out, - ch_target, - target_sample.chunks, - target_metadata, - trace_ir=ir, - ), - ch_domain.shutdown(context), - ) - - @define_actor() async def pushdown_filter_actor( context: Context, @@ -79,7 +55,7 @@ async def pushdown_filter_actor( collective_id: int, ) -> None: """Choose and optionally execute one standalone pushdown-filter hint.""" - samples: tuple[TableSizeStats, TableSizeStats] | None = None + collected_samples: Sequence[TableSizeStats] = [] async with shutdown_on_error( context, ch_out, @@ -126,7 +102,6 @@ async def pushdown_filter_actor( if len(collected_samples) != 2: raise ValueError("Standalone prefilters require two input samples") target_sample, domain_sample = collected_samples - samples = (target_sample, domain_sample) config = executor.join_filter_pushdown if config is None: raise ValueError("Standalone prefilter has no runtime configuration") @@ -145,15 +120,17 @@ async def pushdown_filter_actor( tracer.set_extra("prefilter", trace) if decision.method == "skip": - await replay_skipped_prefilter( - context, - ir, - ch_out, - ch_target, - ch_domain, - target_metadata, - target_sample, - domain_sample, + domain_sample.chunks.clear() + await gather_in_task_group( + ch_domain.shutdown(context), + replay_buffered_channel( + context, + ch_out, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ), ) else: target, domain = ir.children @@ -208,9 +185,8 @@ async def pushdown_filter_actor( ): await gather_in_task_group(*execution.tasks) finally: - if samples is not None: - for sample in samples: - sample.chunks.clear() + for sample in collected_samples: + sample.chunks.clear() @generate_ir_sub_network.register(PushdownFilterHint) From 30984cd754d3661d89e1e5744b355225092c5de0 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 17:03:33 +0100 Subject: [PATCH 12/17] Filters not connected to their direct joins are free-standing --- .../streaming/actor_graph/prefilter_actor.py | 4 +- .../cudf_polars/streaming/explain.py | 3 +- .../cudf_polars/streaming/filter_hint.py | 19 +++++- .../cudf_polars/cudf_polars/streaming/join.py | 12 +++- .../streaming/join_filter_pushdown.py | 4 ++ .../tests/streaming/test_explain.py | 1 + .../streaming/test_join_filter_pushdown.py | 1 + .../tests/streaming/test_tracing.py | 58 +++++++++---------- 8 files changed, 62 insertions(+), 40 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py index b4a9a4a0ca57..da5ac1954fb5 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -112,8 +112,10 @@ async def pushdown_filter_actor( broadcast_limit=executor.broadcast_limit, bloom_filter_max_size=config.bloom_filter_max_size, ) - trace = decision.trace_details() + trace: dict[str, Any] = decision.trace_details() trace["placement"] = "standalone" + trace["target_on"] = [key.name for key in ir.target_on] + trace["domain_on"] = [key.name for key in ir.domain_on] trace_stats = trace if tracer is not None else None if tracer is not None: tracer.decision = decision.method diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index b6102dd33329..803c057e436b 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -495,7 +495,7 @@ def _(ir: PushdownFilterHint, *, offset: str = "") -> str: domain_on = tuple(ne.name for ne in ir.domain_on) return _repr_header( offset, - f"PUSHDOWN FILTER HINT {target_on} {domain_on}", + f"PUSHDOWN FILTER HINT {target_on} {domain_on} {ir.placement}", ir.schema, ) @@ -646,6 +646,7 @@ def _(ir: PushdownFilterHint) -> dict[str, Serializable]: "target_on": [ne.name for ne in ir.target_on], "domain_on": [ne.name for ne in ir.domain_on], "nulls_equal": ir.nulls_equal, + "placement": ir.placement, } diff --git a/python/cudf_polars/cudf_polars/streaming/filter_hint.py b/python/cudf_polars/cudf_polars/streaming/filter_hint.py index de6e1765ff39..de5f59ad3cdd 100644 --- a/python/cudf_polars/cudf_polars/streaming/filter_hint.py +++ b/python/cudf_polars/cudf_polars/streaming/filter_hint.py @@ -19,6 +19,7 @@ JoinSide: TypeAlias = Literal["left", "right"] +HintPlacement: TypeAlias = Literal["join_input", "pushed_down"] @dataclass(frozen=True, slots=True) @@ -125,14 +126,15 @@ class PushdownFilterHint(IR): optional. """ - __slots__ = ("domain_on", "nulls_equal", "target_on") + __slots__ = ("domain_on", "nulls_equal", "placement", "target_on") _non_child: ClassVar[tuple[str, ...]] = ( "schema", "target_on", "domain_on", "nulls_equal", + "placement", ) - _n_non_child_args: ClassVar[int] = 3 + _n_non_child_args: ClassVar[int] = 4 target_on: tuple[NamedExpr, ...] """Expressions selecting filter keys from the target.""" @@ -140,6 +142,8 @@ class PushdownFilterHint(IR): """Expressions selecting filter keys from the domain.""" nulls_equal: bool """Whether null key values compare equal.""" + placement: HintPlacement + """Whether the hint remains at the motivating join input.""" def __init__( self, @@ -147,6 +151,7 @@ def __init__( target_on: Sequence[NamedExpr], domain_on: Sequence[NamedExpr], nulls_equal: bool, # noqa: FBT001 + placement: HintPlacement, target: IR, domain: IR, ): @@ -154,7 +159,13 @@ def __init__( self.target_on = tuple(target_on) self.domain_on = tuple(domain_on) self.nulls_equal = nulls_equal - self._non_child_args = (self.target_on, self.domain_on, self.nulls_equal) + self.placement = placement + self._non_child_args = ( + self.target_on, + self.domain_on, + self.nulls_equal, + self.placement, + ) self.children = (target, domain) @classmethod @@ -163,10 +174,12 @@ def do_evaluate( target_on: tuple[NamedExpr, ...], domain_on: tuple[NamedExpr, ...], nulls_equal: bool, # noqa: FBT001 + placement: HintPlacement, target: DataFrame, domain: DataFrame, *, context: IRExecutionContext, ) -> DataFrame: """Ignore the optional filter and return the target.""" + del placement return target diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 02d0a2893e0a..ebdc878e7df6 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -157,13 +157,18 @@ def _has_non_pointwise_keys(ir: Join) -> bool: return not all(expr.is_pointwise for expr in traversal(keys)) +def is_direct_join_prefilter(ir: IR) -> bool: + """Return whether a hint belongs to its immediately enclosing join.""" + return isinstance(ir, PushdownFilterHint) and ir.placement == "join_input" + + def _lower_join_with_prefilters( ir: Join, rec: LowerIRTransformer, ) -> tuple[JoinWithPrefilter, MutableMapping[IR, PartitionInfo]]: """Lower a join and normalize its adjacent filter hints.""" targets = tuple( - child.children[0] if isinstance(child, PushdownFilterHint) else child + child.children[0] if is_direct_join_prefilter(child) else child for child in ir.children ) lowered_targets, target_partition_info = zip( @@ -178,8 +183,9 @@ def _lower_join_with_prefilters( external_domains: list[IR] = [] claimed_sides: set[JoinSide] = set() for target_index, child in enumerate(ir.children): - if not isinstance(child, PushdownFilterHint): + if not is_direct_join_prefilter(child): continue + assert isinstance(child, PushdownFilterHint) _target, domain = child.children domain, domain_partition_info = rec(domain) @@ -340,7 +346,7 @@ def _( and ir.options[0] != "Cross" and ir.options[5] == "none" and not has_non_pointwise_keys - and any(isinstance(child, PushdownFilterHint) for child in ir.children) + and any(is_direct_join_prefilter(child) for child in ir.children) ): preserve_prefilters = True else: 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 59763e83edf8..a82c49ef9d82 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -107,6 +107,7 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence from cudf_polars.streaming.base import StatsCollector + from cudf_polars.streaming.filter_hint import HintPlacement from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -487,6 +488,7 @@ def apply_candidate(ir: Join, candidate: Candidate) -> IR: domain, expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), nulls_equal=ir.options[1], + placement="join_input" if not target.path else "pushed_down", ) if candidate.target_side == "left": left = replace_at_path(left, target.path, target_filter) @@ -748,12 +750,14 @@ def _make_filter_hint( domain_key: expr.Col, *, nulls_equal: bool, + placement: HintPlacement = "pushed_down", ) -> PushdownFilterHint: return PushdownFilterHint( target.schema, (expr.NamedExpr(target_key.name, target_key),), (expr.NamedExpr(domain_key.name, domain_key),), nulls_equal, + placement, target, domain, ) diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 85ff48da64e0..9458c507e623 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -162,6 +162,7 @@ def test_explain_pushdown_filter_hint_in_dynamic_physical_plan(): "target_on": ["key"], "domain_on": ["key"], "nulls_equal": False, + "placement": "join_input", } assert any( node.type == "PushdownFilterHint" and node.properties == expected_properties 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 c560e417e018..a4c2ad4b7338 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -202,6 +202,7 @@ def test_adjacent_filter_hint_is_recorded_on_lowered_join( (prefilter,) = lowering.lowered.prefilters assert isinstance(prefilter, Prefilter) assert isinstance(prefilter.domain, JoinInputDomain) + assert find_hints(lowering.optimized)[0].placement == "join_input" assert prefilter.target_side == "right" assert prefilter.domain.side == "left" assert tuple(right.schema) == ("l_partkey", "l_suppkey") diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index cac488b2937b..d8675e9192b8 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -259,7 +259,7 @@ def test_standalone_prefilter_trace_records_decision_and_effect( reason: str, output_rows: int | None, ) -> None: - """Trace a non-adjacent prefilter selected through the public engine.""" + """Trace a prefilter pushed below an intervening join.""" pytest.importorskip("structlog") code = textwrap.dedent(f"""\ import json @@ -279,12 +279,17 @@ def test_standalone_prefilter_trace_records_decision_and_effect( .filter("active") .select("p_partkey") ) - target = pl.LazyFrame( - {{ - "l_partkey": [i % 10 for i in range(100)], - "value": range(100), - }} - ).with_columns((pl.col("value") + 1).alias("derived")) + target = ( + pl.LazyFrame( + {{ + "l_partkey": [i % 10 for i in range(100)], + "bridge_key": range(100), + "value": range(100), + }} + ) + .join(pl.LazyFrame({{"bridge_key": range(100)}}), on="bridge_key") + .with_columns((pl.col("value") + 1).alias("derived")) + ) query = domain.join(target, left_on="p_partkey", right_on="l_partkey") options = {{ "join_filter_pushdown": {{ @@ -349,31 +354,29 @@ def test_standalone_prefilter_trace_records_decision_and_effect( @pytest.mark.parametrize( - "broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,domain_rows", + "broadcast_limit,bloom_filter_max_size,method,reason,domain_rows", [ - (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 15), - (512, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 15), + (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 15), + (512, 0, "broadcast_semi_join", "exact_domain_fits", 15), ( 1_000_000, 32 * 1024 * 1024, - "broadcast_left", - "skip", - "target_not_redistributed", - None, + "bloom", + "bloom_fits", + 15, ), ], - ids=["bloom", "exact", "skip"], + ids=["bloom", "exact", "bloom_despite_intervening_broadcast"], ) -def test_external_join_prefilter_trace_records_decision_and_effect( +def test_indirect_prefilter_trace_records_decision_and_effect( timeout_seconds: int, broadcast_limit: int, bloom_filter_max_size: int, - join_strategy: str, method: str, reason: str, domain_rows: int | None, ) -> None: - """Trace an external-domain prefilter selected through the public engine.""" + """Trace a composite prefilter pushed below an intervening join.""" pytest.importorskip("structlog") code = textwrap.dedent(f"""\ import json @@ -442,20 +445,12 @@ def test_external_join_prefilter_trace_records_decision_and_effect( log for log in logs if log.get("scope") == "actor" - and any( - prefilter.get("domain") == "external" - for prefilter in log.get("join_prefilters", ()) - ) - ) - (prefilter,) = ( - prefilter - for prefilter in event["join_prefilters"] - if prefilter.get("domain") == "external" + and log.get("prefilter", {{}}).get("placement") == "standalone" + and log.get("prefilter", {{}}).get("target_on") == ["l_suppkey"] ) record = {{ "result_rows": result.height, - "join_strategy": event["decision"], - "prefilter": prefilter, + "prefilter": event["prefilter"], }} print("PREFILTER_TRACE=" + json.dumps(record)) """) @@ -476,12 +471,11 @@ def test_external_join_prefilter_trace_records_decision_and_effect( record = json.loads(payload) assert record["result_rows"] == 45 - assert record["join_strategy"] == join_strategy + assert record["prefilter"]["target_on"] == ["l_suppkey"] assert ( record["prefilter"].items() >= { - "target_side": "right", - "domain": "external", + "placement": "standalone", "method": method, "reason": reason, "domain_rows": domain_rows, From 5065c2b10fcf9137c4be164b5e0ad699a01b17ce Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 18:05:39 +0100 Subject: [PATCH 13/17] Ensure piecewise joins do not include prefilters --- .../cudf_polars/streaming/actor_graph/join.py | 6 ++ .../cudf_polars/cudf_polars/streaming/join.py | 28 +++++++- .../streaming/test_join_filter_pushdown.py | 65 +++++++++++++++++-- 3 files changed, 92 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 962df2ca19c0..5d0baefc6df1 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -1519,6 +1519,12 @@ def _( executor = rec.state["config_options"].executor pwise_join = _use_pwise_join(executor, partition_info, ir) + if pwise_join and isinstance(ir, JoinWithPrefilter): + raise AssertionError( + "Partition-wise JoinWithPrefilter should have been simplified " + "during IR lowering" + ) + actors, channels = process_children(ir, rec) # Create output ChannelManager diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index ebdc878e7df6..a9a266a5c166 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -162,10 +162,10 @@ def is_direct_join_prefilter(ir: IR) -> bool: return isinstance(ir, PushdownFilterHint) and ir.placement == "join_input" -def _lower_join_with_prefilters( +def lower_join_with_prefilters( ir: Join, rec: LowerIRTransformer, -) -> tuple[JoinWithPrefilter, MutableMapping[IR, PartitionInfo]]: +) -> tuple[Join, MutableMapping[IR, PartitionInfo]]: """Lower a join and normalize its adjacent filter hints.""" targets = tuple( child.children[0] if is_direct_join_prefilter(child) else child @@ -179,6 +179,28 @@ def _lower_join_with_prefilters( operator.or_, target_partition_info ) + if all( + isinstance(target, Repartition) and partition_info[target].count == 1 + for target in lowered_targets + ): + # This join will execute partition-wise, so its optional prefilters + # are unnecessary. Moreover, the piecewise join special case + # execution at runtime never has a chance to shut down prefilter + # channels that would be produced, which would leave an actor graph + # in a deadlocked state. Since they are unnecessary, drop them + # before lowering their domains and before the actor graph derives + # fanout from the lowered DAG. + return ( + Join( + ir.schema, + ir.left_on, + ir.right_on, + ir.options, + *lowered_targets, + ), + partition_info, + ) + prefilters: list[Prefilter] = [] external_domains: list[IR] = [] claimed_sides: set[JoinSide] = set() @@ -353,7 +375,7 @@ def _( preserve_prefilters = False if preserve_prefilters: - ir, partition_info = _lower_join_with_prefilters(ir, rec) + ir, partition_info = lower_join_with_prefilters(ir, rec) children = ir.children else: # Hints not owned by an adaptive join use the generic identity lowering. 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 a4c2ad4b7338..90e5d52fa92e 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import pytest @@ -11,17 +11,30 @@ from cudf_polars import Translator from cudf_polars.dsl.expr import Col -from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice -from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.ir import ( + IR, + Cache, + DataFrameScan, + Distinct, + Join, + Projection, + Select, + Slice, +) +from cudf_polars.dsl.traversal import CachingVisitor, traversal from cudf_polars.dsl.utils.column_domain import ColumnRef from cudf_polars.engine.options import StreamingOptions -from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.base import PartitionInfo, StatsCollector from cudf_polars.streaming.filter_hint import ( JoinInputDomain, JoinWithPrefilter, Prefilter, PushdownFilterHint, ) +from cudf_polars.streaming.join import ( + is_direct_join_prefilter, + lower_join_with_prefilters, +) from cudf_polars.streaming.join_filter_pushdown import ( CompositeCandidate, Decision, @@ -40,6 +53,7 @@ optimize_with_stats, remove_cache_nodes, ) +from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import ConfigOptions @@ -213,6 +227,49 @@ def test_adjacent_filter_hint_is_recorded_on_lowered_join( assert not find_joins(lowering.lowered, "Semi") +def test_partition_wise_join_discards_prefilters_before_lowering_domains( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + """Partition-wise joins must not retain optional prefilter inputs.""" + root = translate_query(simple_query, engine) + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_with_stats(root, config, StatsCollector()) + assert isinstance(optimized, Join) + + children = list(optimized.children) + (hint_index,) = ( + index for index, child in enumerate(children) if is_direct_join_prefilter(child) + ) + hint = children[hint_index] + assert isinstance(hint, PushdownFilterHint) + domain = Projection(hint.children[1].schema, hint.children[1]) + children[hint_index] = hint.reconstruct((hint.children[0], domain)) + optimized = optimized.reconstruct(children) + + targets = tuple( + child.children[0] if is_direct_join_prefilter(child) else child + for child in optimized.children + ) + repartitions = tuple(Repartition(target.schema, target) for target in targets) + lowered_targets = dict(zip(targets, repartitions, strict=True)) + + def lower_target(child: IR, rec: Any) -> tuple[IR, dict[IR, PartitionInfo]]: + assert child in lowered_targets, "prefilter domain was lowered" + lowered = lowered_targets[child] + return lowered, {lowered: PartitionInfo(count=1)} + + rec: Any = CachingVisitor( + lower_target, + state={"config_options": config}, + ) + lowered, partition_info = lower_join_with_prefilters(optimized, rec) + + assert type(lowered) is Join + assert lowered.children == repartitions + assert domain not in partition_info + + def test_filter_pushdown_can_be_disabled( simple_query: pl.LazyFrame, engine: SPMDEngine ) -> None: From 61fd0e5c705a8f153483b3d2660b68bcdad9c911 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:24:41 +0100 Subject: [PATCH 14/17] Use asdict --- .../streaming/actor_graph/prefilter.py | 16 ++-------------- .../streaming/actor_graph/prefilter_actor.py | 3 ++- 2 files changed, 4 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py index ed0ff4d3edc0..fee3bfbbbfdf 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -5,7 +5,7 @@ from __future__ import annotations import math -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Literal import pylibcudf as plc @@ -75,21 +75,9 @@ class PrefilterDecision: bloom_bytes: int | None = None exact_bytes: int | None = None - def trace_details(self) -> dict[str, str | int | None]: - """Return trace information common to every prefilter placement.""" - return { - "method": self.method, - "reason": self.reason, - "target_bytes": self.target_bytes, - "domain_rows": self.domain_rows, - "estimated_cardinality": self.estimated_cardinality, - "bloom_bytes": self.bloom_bytes, - "exact_bytes": self.exact_bytes, - } - def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: """Return serializable actor-trace information.""" - result = self.trace_details() + result = asdict(self) result["target_side"] = prefilter.target_side if isinstance(prefilter.domain, JoinInputDomain): result["domain_side"] = prefilter.domain.side diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py index da5ac1954fb5..76ecee442821 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -4,6 +4,7 @@ from __future__ import annotations +from dataclasses import asdict from typing import TYPE_CHECKING, Any from cudf_streaming import CardinalityEstimator @@ -112,7 +113,7 @@ async def pushdown_filter_actor( broadcast_limit=executor.broadcast_limit, bloom_filter_max_size=config.bloom_filter_max_size, ) - trace: dict[str, Any] = decision.trace_details() + trace = asdict(decision) trace["placement"] = "standalone" trace["target_on"] = [key.name for key in ir.target_on] trace["domain_on"] = [key.name for key in ir.domain_on] From 1e815642e0d6362865ef4e98d25fbc0db8a816cf Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:35:35 +0100 Subject: [PATCH 15/17] TableSizeStats offers distinct_count method --- .../streaming/actor_graph/prefilter.py | 30 +++++-------------- .../streaming/actor_graph/utils.py | 16 ++++++++++ 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py index fee3bfbbbfdf..a14b660aae09 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -308,22 +308,6 @@ def add_bloom_prefilter( ) -def estimate_cardinality(stats: TableSizeStats) -> int | None: - """Extrapolate sampled distinct count to the estimated full row count.""" - if stats.total_rows == 0: - return 0 - if stats.cardinality is None or stats.cardinality.row_count == 0: - return None - return min( - stats.total_rows, - math.ceil( - stats.cardinality.distinct_count - * stats.total_rows - / stats.cardinality.row_count - ), - ) - - def estimate_bloom_filter_bytes( cardinality: int, desired_false_positive_rate: float = 0.1, @@ -394,15 +378,15 @@ def choose_prefilter_method( bloom_filter_max_size: int, ) -> PrefilterDecision: """Choose the implementation for an eligible prefilter.""" - cardinality = estimate_cardinality(domain) - if cardinality is None: + distinct_count = domain.distinct_count() + if distinct_count is None: return PrefilterDecision( "skip", "missing_cardinality", target.total_size, domain.total_rows, ) - if cardinality == 0: + if distinct_count == 0: return PrefilterDecision( "skip", "zero_cardinality", @@ -415,7 +399,7 @@ def choose_prefilter_method( bloom_bytes = max( 32, - BloomFilter.aligned_size(estimate_bloom_filter_bytes(cardinality)), + BloomFilter.aligned_size(estimate_bloom_filter_bytes(distinct_count)), ) exact_bytes = estimate_bytes( tuple(key.value.dtype for key in domain_on), @@ -427,7 +411,7 @@ def choose_prefilter_method( "bloom_fits", target.total_size, domain.total_rows, - estimated_cardinality=cardinality, + estimated_cardinality=distinct_count, bloom_bytes=bloom_bytes, exact_bytes=exact_bytes, ) @@ -439,7 +423,7 @@ def choose_prefilter_method( "exact_domain_fits", target.total_size, domain.total_rows, - estimated_cardinality=cardinality, + estimated_cardinality=distinct_count, bloom_bytes=bloom_bytes, exact_bytes=exact_bytes, ) @@ -448,7 +432,7 @@ def choose_prefilter_method( "no_viable_filter", target.total_size, domain.total_rows, - estimated_cardinality=cardinality, + estimated_cardinality=distinct_count, bloom_bytes=bloom_bytes, exact_bytes=exact_bytes, ) 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 55192ef4f755..3360379566f5 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -7,6 +7,7 @@ import asyncio import contextlib import itertools +import math import operator import struct import time @@ -1039,6 +1040,21 @@ class TableSizeStats: cardinality: CardinalityEstimate | None = None """Global cardinality statistics for the sampled rows, when requested.""" + def distinct_count(self) -> int | None: + """Extrapolate sampled distinct count to the estimated full row count.""" + if self.total_rows == 0: + return 0 + if self.cardinality is None or self.cardinality.row_count == 0: + return None + return min( + self.total_rows, + math.ceil( + self.cardinality.distinct_count + * self.total_rows + / self.cardinality.row_count + ), + ) + async def aggregate_table_size_stats( context: Context, From e50c9d6318afcffc28790fe683ef8c5287fe312d Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:37:56 +0100 Subject: [PATCH 16/17] Refactor to introduce JoinPlanningState.create --- .../cudf_polars/streaming/actor_graph/join.py | 9 +- .../streaming/actor_graph/join_planning.py | 109 +++++++++--------- 2 files changed, 58 insertions(+), 60 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 5d0baefc6df1..0c39f9271326 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -36,7 +36,7 @@ from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.streaming.actor_graph.join_planning import make_join_planning_state +from cudf_polars.streaming.actor_graph.join_planning import JoinPlanningState from cudf_polars.streaming.actor_graph.nodes import default_node_multi from cudf_polars.streaming.actor_graph.prefilter import ( JoinPrefilterExecution, @@ -81,10 +81,7 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator - from cudf_polars.streaming.actor_graph.join_planning import ( - JoinInput, - JoinPlanningState, - ) + from cudf_polars.streaming.actor_graph.join_planning import JoinInput from cudf_polars.streaming.actor_graph.prefilter import ( PrefilterDecision, PrefilterExecution, @@ -1363,7 +1360,7 @@ async def join_actor( *(recv_metadata(ch, context) for ch in ch_prefilter_domains), ) - join_state = make_join_planning_state( + join_state = JoinPlanningState.create( ir, ch_left, ch_right, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py index 463496afb893..1b0972bb762e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py @@ -14,7 +14,7 @@ ) if TYPE_CHECKING: - from typing import Any + from typing import Any, Self from cudf_streaming.channel_metadata import ChannelMetadata from cudf_streaming.table_chunk import TableChunk @@ -57,58 +57,59 @@ class JoinPlanningState: right: JoinInput candidates: tuple[PrefilterCandidate, ...] = () - -def make_join_planning_state( - ir: Join, - ch_left: Channel[TableChunk], - ch_right: Channel[TableChunk], - ch_prefilter_domains: tuple[Channel[TableChunk], ...], - left_metadata: ChannelMetadata, - right_metadata: ChannelMetadata, - prefilter_domain_metadata: tuple[ChannelMetadata, ...], - cardinality_tags: tuple[int, ...], -) -> JoinPlanningState: - """Create actor-local planning state from a join and its runtime inputs.""" - left = JoinInput(ir.children[0], ch_left, left_metadata) - right = JoinInput(ir.children[1], ch_right, right_metadata) - if not isinstance(ir, JoinWithPrefilter): - if ch_prefilter_domains or prefilter_domain_metadata: - raise ValueError("A plain Join cannot have prefilter domain inputs") - return JoinPlanningState(left, right) - - external_inputs = tuple( - JoinInput(node, channel, metadata) - for node, channel, metadata in zip( - ir.children[2:], - ch_prefilter_domains, - prefilter_domain_metadata, - strict=True, - ) - ) - external_prefilter_count = sum( - isinstance(prefilter.domain, ExternalDomain) for prefilter in ir.prefilters - ) - if external_prefilter_count != len(external_inputs): - raise ValueError("Join prefilters and external domain inputs must align") - if len(cardinality_tags) < len(ir.prefilters): - raise ValueError("Each join prefilter requires a cardinality collective ID") - - sides = {"left": left, "right": right} - external_inputs_iter = iter(external_inputs) - cardinality_tags_iter = iter(cardinality_tags) - candidates = [] - for spec in ir.prefilters: - target = sides[spec.target_side] - if isinstance(spec.domain, JoinInputDomain): - domain = sides[spec.domain.side] - else: - domain = next(external_inputs_iter) - candidates.append( - PrefilterCandidate( - spec, - target, - domain, - cardinality_tag=next(cardinality_tags_iter), + @classmethod + def create( + cls, + ir: Join, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], + left_metadata: ChannelMetadata, + right_metadata: ChannelMetadata, + prefilter_domain_metadata: tuple[ChannelMetadata, ...], + cardinality_tags: tuple[int, ...], + ) -> Self: + """Create actor-local planning state from a join and its runtime inputs.""" + left = JoinInput(ir.children[0], ch_left, left_metadata) + right = JoinInput(ir.children[1], ch_right, right_metadata) + if not isinstance(ir, JoinWithPrefilter): + if ch_prefilter_domains or prefilter_domain_metadata: + raise ValueError("A plain Join cannot have prefilter domain inputs") + return cls(left, right) + + external_inputs = tuple( + JoinInput(node, channel, metadata) + for node, channel, metadata in zip( + ir.children[2:], + ch_prefilter_domains, + prefilter_domain_metadata, + strict=True, ) ) - return JoinPlanningState(left, right, tuple(candidates)) + external_prefilter_count = sum( + isinstance(prefilter.domain, ExternalDomain) for prefilter in ir.prefilters + ) + if external_prefilter_count != len(external_inputs): + raise ValueError("Join prefilters and external domain inputs must align") + if len(cardinality_tags) < len(ir.prefilters): + raise ValueError("Each join prefilter requires a cardinality collective ID") + + sides = {"left": left, "right": right} + external_inputs_iter = iter(external_inputs) + cardinality_tags_iter = iter(cardinality_tags) + candidates = [] + for spec in ir.prefilters: + target = sides[spec.target_side] + if isinstance(spec.domain, JoinInputDomain): + domain = sides[spec.domain.side] + else: + domain = next(external_inputs_iter) + candidates.append( + PrefilterCandidate( + spec, + target, + domain, + cardinality_tag=next(cardinality_tags_iter), + ) + ) + return cls(left, right, tuple(candidates)) From 09dfe8f1d66732af451654d72f7ea50000a905dc Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:42:30 +0100 Subject: [PATCH 17/17] Update comment --- python/cudf_polars/cudf_polars/streaming/actor_graph/core.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index dc631f8f0dae..a93694fe2ded 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -182,8 +182,9 @@ def _mark_children_unbounded(node: IR) -> None: elif isinstance(node, (Union, Join, Over, PushdownFilterHint)): # Union processes children sequentially; Join may broadcast one # side; Over buffers (or samples-then-replays) its input before - # producing output. In every case the input source needs - # unbounded fanout so other consumers don't block it. + # producing output; PushdownFilterHint similarly might buffer + # then replay. In every case the input source needs unbounded + # fanout so other consumers don't block it. _mark_children_unbounded(node) elif len(node.children) > 1: # Check if this node is doing any broadcasting.