From b938131466c2bcc28738bc677d6b04d0f17a578c Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 28 Jul 2026 14:40:24 +0100 Subject: [PATCH 01/22] 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 d2098b20cb7b..1f2dd6bcae33 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 d26273af6829321965c109907837a5b6ec350b50 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 4 Aug 2026 12:30:35 +0100 Subject: [PATCH 02/22] Apply adaptive prefilters to direct join inputs --- python/cudf_polars/cudf_polars/dsl/ir.py | 7 +- .../cudf_polars/streaming/actor_graph/join.py | 642 ++++++++++++++---- .../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 | 110 +++ python/cudf_polars/tests/test_config.py | 27 +- 13 files changed, 1427 insertions(+), 162 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 98c33c57fc0a..468ca2f54e6d 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2731,7 +2731,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 7a7b905384ab..a807249f8ddd 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, TypeAlias, assert_never +from cudf_streaming import BloomFilter, CardinalityEstimator from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -18,6 +19,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 ( @@ -26,7 +28,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, @@ -42,8 +44,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, @@ -64,11 +72,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 cudf_streaming.channel_metadata import Ordering from rapidsmpf.communicator.communicator import Communicator @@ -78,8 +87,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 @@ -139,6 +154,53 @@ class OrderedJoinStrategy: ) +@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 isinstance(strategy, BroadcastJoinStrategy): + if target_side != strategy.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, @@ -195,7 +257,7 @@ async def broadcast_join_actor( ch_left, ch_right, BroadcastJoinStrategy(side=broadcast_side), - [collective_id], + collective_id, target_partition_size, tracer=tracer, ) @@ -290,7 +352,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() @@ -321,8 +383,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( @@ -334,26 +398,26 @@ async def _broadcast_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: BroadcastJoinStrategy, - 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.side - left, right = ir.children + left, right = ir.children[:2] if tracer is not None: tracer.decision = f"broadcast_{broadcast_side}" @@ -395,8 +459,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, @@ -408,6 +470,13 @@ 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: # Unknown: the large chunk is freed but the join output replaces # it, and its size depends on selectivity we cannot estimate here. @@ -417,7 +486,8 @@ async def _broadcast_join( reserve_extra=0, net_memory_delta=missing_net_memory_delta, ) - await _broadcast_join_large_chunk( + input_rows += large_chunk.shape[0] + output_rows += await _broadcast_join_large_chunk( context, ir, ir_context, @@ -432,9 +502,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, + BroadcastJoinStrategy(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, @@ -445,7 +676,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] @@ -561,7 +792,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) @@ -662,7 +893,8 @@ async def _shuffle_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: ShuffleJoinStrategy, - collective_ids: list[int], + left_collective_id: int, + right_collective_id: int, *, tracer: ActorTracer | None, ) -> None: @@ -705,7 +937,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, @@ -716,7 +948,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, @@ -785,7 +1017,7 @@ async def _ordered_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: OrderedJoinStrategy, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, *, tracer: ActorTracer | None, ) -> None: @@ -823,7 +1055,7 @@ async def _ordered_join( ch_left, strategy.left_input_ordering, strategy.left_output_ordering, - collective_id=collective_ids.pop(0), + collective_id=collective_ids.shuffle("left"), ), _adjust_ordered_join_side( context, @@ -834,7 +1066,7 @@ async def _ordered_join( ch_right, strategy.right_input_ordering, strategy.right_output_ordering, - collective_id=collective_ids.pop(0), + collective_id=collective_ids.shuffle("right"), ), _join_chunks( context, @@ -891,47 +1123,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 - ( - left_total, - right_total, - left_total_rows, - right_total_rows, - left_total_chunks, - right_total_chunks, - ) = await allgather_reduce( + # 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, - ) - - new_left_sample = TableSizeStats( - chunks=left_sample.chunks, - total_size=left_total, - total_rows=left_total_rows, - total_chunks=left_total_chunks, + 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, + 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( @@ -1077,45 +1303,185 @@ 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 isinstance(strategy, BroadcastJoinStrategy): + return side == strategy.side and not metadata.duplicated + if isinstance(strategy, OrderedJoinStrategy): + # Ordered inputs already have a viable join strategy without adaptive + # sampling. Keep that strategy and avoid introducing a sampled + # prefilter pipeline solely to optimize boundary alignment. + return False + + assert isinstance(strategy, ShuffleJoinStrategy) + 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, ) @@ -1133,45 +1499,39 @@ async def _choose_strategy( ): if tracer is not None: tracer.decision = "ordered" - left_sample = TableSizeStats( + 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, ) - return left_sample, right_sample, ordered_strategy + if executor.join_filter_pushdown is not None: + choose_prefilters( + bindings, + ordered_strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + executor.join_filter_pushdown.bloom_filter_max_size, + ) + return ordered_strategy 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, @@ -1185,8 +1545,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() @@ -1199,7 +1567,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. @@ -1242,24 +1610,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, ): @@ -1280,9 +1679,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 isinstance(strategy, BroadcastJoinStrategy): actor_tasks.append( @@ -1295,7 +1695,7 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.broadcast, executor.target_partition_size, tracer=tracer, ) @@ -1326,7 +1726,8 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.shuffle("left"), + collective_ids.shuffle("right"), tracer=tracer, ) ) @@ -1341,7 +1742,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 @@ -1367,8 +1768,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 @@ -1378,7 +1780,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 @@ -1408,16 +1809,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 5a77f8ead634..b63e02fc7841 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -175,7 +175,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() @@ -185,6 +185,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)) @@ -324,6 +330,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 ) @@ -1264,19 +1271,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 47c8a9f6dfdb..06ef89836ab4 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -509,7 +509,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 @@ -517,7 +517,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. @@ -530,6 +531,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. """ @@ -541,6 +546,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 @@ -555,6 +567,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 a14ff016e985..004960c502d3 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -205,6 +205,116 @@ def test_io_tasks_wait_for_memory_admission( assert second["admitted"] >= first["stop"] +@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 4f0a953dd6cc..dce6bd635ca1 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -856,10 +856,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 @@ -894,6 +899,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: @@ -910,7 +933,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 a6ef807180baab6a27eb457d502d9a291005073c Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 4 Aug 2026 12:31:34 +0100 Subject: [PATCH 03/22] 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 a807249f8ddd..dbabff9eb504 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -72,7 +72,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 @@ -578,6 +582,7 @@ def add_bloom_prefilter( def make_prefilter_execution( context: Context, comm: Communicator, + ir: Join, ir_context: IRExecutionContext, strategy: JoinStrategy, ch_left: Channel[TableChunk], @@ -599,8 +604,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 @@ -1337,26 +1362,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, ) @@ -1399,28 +1433,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( @@ -1440,10 +1470,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, @@ -1519,10 +1617,11 @@ async def choose_strategy( return ordered_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, @@ -1545,15 +1644,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 @@ -1566,6 +1666,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: @@ -1592,6 +1693,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 @@ -1602,20 +1705,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, ) @@ -1647,6 +1758,7 @@ async def join_actor( prefilter_execution = make_prefilter_execution( context, comm, + ir, ir_context, strategy, ch_left_replay, @@ -1773,7 +1885,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 @@ -1825,6 +1937,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 004960c502d3..ae5b8822749e 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -315,6 +315,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 4460009550747770e12f5b9b967f127c5b970b4f Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:28:30 +0100 Subject: [PATCH 04/22] 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 dbabff9eb504..3d05ef618929 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -44,7 +44,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, @@ -91,10 +91,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 @@ -587,7 +587,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.""" @@ -596,50 +596,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": @@ -649,7 +649,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, @@ -658,15 +658,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, ) @@ -1355,7 +1355,7 @@ def join_input_requires_redistribution( def choose_prefilters( - bindings: JoinBindings, + join_state: JoinPlanningState, strategy: JoinStrategy, left_partitioning: NormalizedPartitioning, right_partitioning: NormalizedPartitioning, @@ -1367,29 +1367,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, @@ -1399,26 +1399,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" ) @@ -1436,32 +1436,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( @@ -1475,20 +1479,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)) @@ -1496,7 +1500,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, @@ -1505,11 +1509,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, @@ -1520,33 +1524,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, *, @@ -1554,8 +1558,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, @@ -1575,11 +1579,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, ) @@ -1620,15 +1624,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( @@ -1647,7 +1651,7 @@ async def choose_strategy( await resolve_prefilters( context, comm, - bindings, + join_state, strategy, left_partitioning, right_partitioning, @@ -1719,7 +1723,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, @@ -1734,23 +1738,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() @@ -1763,7 +1767,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 3f83bcaeb6dd0f08aa280fb55088567f621826dc Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:52:03 +0100 Subject: [PATCH 05/22] 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 25939064b2d63b3caf3cbb863c8a76b95a6919ba Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:55:40 +0100 Subject: [PATCH 06/22] 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 3d05ef618929..542240ad59cb 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, TypeAlias, assert_never -from cudf_streaming import BloomFilter, CardinalityEstimator +from cudf_streaming import CardinalityEstimator from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -19,7 +19,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 ( @@ -47,9 +46,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 ( @@ -81,7 +80,7 @@ from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: - from collections.abc import Iterable, MutableMapping + from collections.abc import MutableMapping from cudf_streaming.channel_metadata import Ordering from rapidsmpf.communicator.communicator import Communicator @@ -512,73 +511,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, @@ -589,9 +521,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 5974338705dc7803a2759c6e08f6fd8e209fe0f6 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 14:56:33 +0100 Subject: [PATCH 07/22] 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 542240ad59cb..0a1d2b752ec8 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -60,7 +60,7 @@ TableSizeStats, _sample_chunks, _update_ordering_indices, - allgather_reduce, + aggregate_table_size_stats, chunk_to_frame, empty_table_chunk, gather_in_task_group, @@ -251,7 +251,7 @@ async def broadcast_join_actor( trace_ir=ir, ir_context=ir_context, ) as tracer: - await _broadcast_join( + await broadcast_join( context, comm, ir, @@ -392,7 +392,7 @@ async def _broadcast_join_large_chunk( return output_rows -async def _broadcast_join( +async def broadcast_join( context: Context, comm: Communicator, ir: Join, @@ -603,7 +603,7 @@ def make_prefilter_execution( projected_domain, ) execution.add_task( - _broadcast_join( + broadcast_join( context, comm, semi_join, @@ -1080,43 +1080,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, @@ -1400,7 +1363,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), @@ -1734,7 +1697,7 @@ async def join_actor( if isinstance(strategy, BroadcastJoinStrategy): 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 b63e02fc7841..b1ee799284cf 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -1075,6 +1075,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 668c49a0c37ee494239835381696102c21694e71 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:04:59 +0100 Subject: [PATCH 08/22] Execute non-adjacent pushdown filter hints --- .../streaming/actor_graph/__init__.py | 1 + .../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, 474 insertions(+), 13 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 c78c8084ae54..c3733c08eec2 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py @@ -16,6 +16,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 fbec424d3104..5d98190e3f80 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -22,6 +22,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.over import Over from cudf_polars.utils.config import SPMDContext @@ -176,7 +177,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 ae5b8822749e..f66161445d4b 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -315,11 +315,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, @@ -458,12 +564,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 aec4f561c0d6dda7dda9a868987d350d72bb0420 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:29:13 +0100 Subject: [PATCH 09/22] 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 0a1d2b752ec8..59d7c04b9111 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -55,12 +55,11 @@ CUDF_ROW_LIMIT, MAX_ROWS_PER_PARTITION, ChannelManager, + ChunkSampler, ChunkStore, NormalizedPartitioning, TableSizeStats, - _sample_chunks, _update_ordering_indices, - aggregate_table_size_stats, chunk_to_frame, empty_table_chunk, gather_in_task_group, @@ -68,6 +67,7 @@ process_children, recv_metadata, replay_buffered_channel, + sample_inputs, send_metadata, shutdown_on_error, ) @@ -93,7 +93,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 @@ -1291,43 +1290,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, @@ -1350,23 +1312,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 b1ee799284cf..fe2f57bfc4f0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -1239,6 +1239,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 89b0dd2d462539e4cc04065a10a7bbe0cd35e3d3 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:42:41 +0100 Subject: [PATCH 10/22] Share adaptive prefilter execution --- .../cudf_polars/streaming/actor_graph/join.py | 156 +++++++++------ .../streaming/actor_graph/prefilter_actor.py | 177 +++++------------- 2 files changed, 149 insertions(+), 184 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 59d7c04b9111..166b5a7bd9c7 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -94,9 +94,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 @@ -510,6 +518,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, + BroadcastJoinStrategy(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, @@ -572,51 +644,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, - BroadcastJoinStrategy(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 @@ -1474,23 +1515,24 @@ async def choose_strategy( ): if tracer is not None: tracer.decision = "ordered" - 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, ) - if executor.join_filter_pushdown is not None: - choose_prefilters( - bindings, - ordered_strategy, - left_partitioning, - right_partitioning, - executor.broadcast_limit, - executor.join_filter_pushdown.bloom_filter_max_size, - ) + await resolve_prefilters( + context, + comm, + join_state, + ordered_strategy, + left_partitioning, + right_partitioning, + executor, + collective_ids.size_estimate, + ) return ordered_strategy else: assert executor.dynamic_planning is not None 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 555dfc0fe1b40f5132d3c61d28ed8c3efaa542dc Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 15:52:59 +0100 Subject: [PATCH 11/22] 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 f44c9f40c522e113dd760fccd7efa64f351c39f6 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 17:03:33 +0100 Subject: [PATCH 12/22] 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 f66161445d4b..bd3f8dee4ec1 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -332,7 +332,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 @@ -352,12 +352,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": {{ @@ -422,31 +427,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 @@ -515,20 +518,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)) """) @@ -549,12 +544,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 7f699869b6818c19bfe0f2532e5b0c1410c3bb88 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 7 Aug 2026 18:05:39 +0100 Subject: [PATCH 13/22] 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 166b5a7bd9c7..70cc6355f6f4 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -1811,6 +1811,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 dbd1c99831be1ca4009aea945f4e589da1d8752c Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:24:41 +0100 Subject: [PATCH 14/22] 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 fe6ab880f52c0873886eafea5c7455bc03b0a0eb Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:35:35 +0100 Subject: [PATCH 15/22] 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 fe2f57bfc4f0..aebf1fb71d6a 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 @@ -1074,6 +1075,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 18b7d57f0e8d36024d430e0c553202e4887e415e Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:37:56 +0100 Subject: [PATCH 16/22] 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 70cc6355f6f4..90c88c7dcc51 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -43,7 +43,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, @@ -90,10 +90,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, @@ -1638,7 +1635,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 628cea65c4ecb48fbe231e5c4a56ac19fc884d4a Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 10 Aug 2026 11:42:30 +0100 Subject: [PATCH 17/22] 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 5d98190e3f80..82688495f859 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -180,8 +180,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. From 87d39a73ba8d5935e298f7ccf33ab603409b7247 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 1 Sep 2026 16:56:25 +0100 Subject: [PATCH 18/22] Add test that prefilters are skipped if ordered join is picked --- .../tests/streaming/test_tracing.py | 86 +++++++++++++------ 1 file changed, 62 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index bd3f8dee4ec1..b1d7e4f68993 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -206,32 +206,61 @@ def test_io_tasks_wait_for_memory_admission( @pytest.mark.parametrize( - "broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,output_rows", + "ordered,broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,domain_rows,output_rows", [ - (1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 10), - (64, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 10), + (False, 1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 1, 10), + (False, 64, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 1, 10), ( + False, 1_000_000, 32 * 1024 * 1024, "broadcast_left", "skip", "target_not_redistributed", + 1, + None, + ), + ( + True, + 1, + 32 * 1024 * 1024, + "ordered", + "skip", + "target_not_redistributed", + None, None, ), ], - ids=["bloom", "exact", "skip"], + ids=["bloom", "exact", "broadcast-skip", "ordered-skip"], ) def test_local_join_prefilter_trace_records_decision_and_effect( + tmp_path: pathlib.Path, timeout_seconds: int, + ordered: bool, # noqa: FBT001 broadcast_limit: int, bloom_filter_max_size: int, join_strategy: str, method: str, reason: str, + domain_rows: int | None, output_rows: int | None, ) -> None: """Trace a direct-input join prefilter selected through the public engine.""" pytest.importorskip("structlog") + domain_path = tmp_path / "domain.parquet" + target_path = tmp_path / "target.parquet" + pl.DataFrame( + { + "key": range(100), + "active": [i % 10 == 0 for i in range(100)], + } + ).write_parquet(domain_path) + pl.DataFrame( + { + "key": range(1_000), + "value": range(1_000), + } + ).write_parquet(target_path) code = textwrap.dedent(f"""\ import json import os @@ -244,14 +273,24 @@ def test_local_join_prefilter_trace_records_decision_and_effect( 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)}} - ) + ordered = {ordered!r} + if ordered: + domain = ( + pl.scan_parquet({str(domain_path)!r}) + .filter("active") + .select("key") + .set_sorted("key") + ) + target = pl.scan_parquet({str(target_path)!r}).set_sorted("key") + else: + 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": {{ @@ -259,8 +298,8 @@ def test_local_join_prefilter_trace_records_decision_and_effect( "bloom_filter_max_size": {bloom_filter_max_size}, }}, "broadcast_limit": {broadcast_limit}, - "target_partition_size": 64, - "max_rows_per_partition": 100, + "target_partition_size": 1 << 30 if ordered else 64, + "max_rows_per_partition": 1_000_000 if ordered else 100, }} with SPMDEngine(executor_options=options) as engine: with structlog.testing.capture_logs() as logs: @@ -296,16 +335,15 @@ def test_local_join_prefilter_trace_records_decision_and_effect( 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() - ) + expected_prefilter: dict[str, str | int] = { + "target_side": "right", + "domain_side": "left", + "method": method, + "reason": reason, + } + if domain_rows is not None: + expected_prefilter["domain_rows"] = domain_rows + assert record["prefilter"].items() >= expected_prefilter.items() if output_rows is None: assert "input_rows" not in record["prefilter"] assert "output_rows" not in record["prefilter"] From 58f53ed5124d7437631ba1727f0cea39a16918cc Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 1 Sep 2026 17:05:25 +0100 Subject: [PATCH 19/22] Fix docs --- docs/cudf/source/cudf_polars/options.md | 2 +- python/cudf_polars/cudf_polars/utils/config.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index d4af20497ff9..c28d50abd0d7 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/docs/cudf/source/cudf_polars/options.md @@ -109,7 +109,7 @@ Environment variables follow these patterns: | `target_partition_size` | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto. | auto | | `max_concurrent_io_tasks` | Number of concurrent IO producer tasks for each scan node. Tune with an integer or a `{"local": ..., "remote": ...}` dict. | auto | | `dynamic_planning` | Dynamic planning configuration, dict or {class}`~cudf_polars.utils.config.DynamicPlanningOptions`. `None` disables. | enabled | -| `join_filter_pushdown` | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `None` disables. | enabled | +| `join_filter_pushdown` | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `None` disables. | disabled | | `sink_to_directory` | Whether `.sink_*()` writes its output as a directory. The `spmd`, `ray`, and `dask` engines always use `True`; passing `False` raises `ValueError`. | `True` | ### Category: `engine` diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 06ef89836ab4..9eb52f3e24e4 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -844,8 +844,8 @@ class StreamingExecutor: :class:`~cudf_polars.utils.config.DynamicPlanningOptions` for more. join_filter_pushdown Options controlling the logical join-domain prefilter rewrite. See - :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for more. - ``None`` disables the rewrite. + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for + more. Disabled by default (or by explicitly providing ``None``). Enable through environment variables with ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN=1``. From 34a2b9bf96bf8ca2de26dc47fb74aafd0535565c Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 1 Sep 2026 17:33:15 +0100 Subject: [PATCH 20/22] Awaitable reservation for copy of key columns --- .../streaming/actor_graph/prefilter.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 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 a14b660aae09..0a518e12f3ed 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -13,6 +13,8 @@ from cudf_streaming.channel_metadata import ChannelMetadata from cudf_streaming.table_chunk import TableChunk from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED +from rapidsmpf.memory.memory_reservation import opaque_memory_usage +from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory from rapidsmpf.streaming.core.message import Message from cudf_polars.streaming.actor_graph.utils import ( @@ -87,16 +89,18 @@ def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: return result -def project_key_chunk( +async 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 - ) + columns = tuple(chunk.table_view().columns()[index] for index in indices) + bytes = sum(column.device_buffer_size() for column in columns) + with opaque_memory_usage( + await reserve_memory(context, size=bytes, net_memory_delta=0) + ): + table = plc.Table(columns).copy(stream=chunk.stream, mr=context.br().device_mr) return TableChunk.from_pylibcudf_table( - key_table, + table, chunk.stream, exclusive_view=True, br=context.br(), @@ -133,7 +137,7 @@ async def buffer_and_project_keys( 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) + key_chunk = await project_key_chunk(context, chunk, indices) chunks.insert(Message(sequence_number, chunk)) await ch_keys.send(context, Message(sequence_number, key_chunk)) From abe559b8cb2a1616f069668fd84564c8e5b45dfa Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 1 Sep 2026 17:49:52 +0100 Subject: [PATCH 21/22] Minor fixes --- .../cudf_polars/streaming/benchmarks/utils.py | 7 ++++++- .../cudf_polars/cudf_polars/streaming/explain.py | 7 ++++++- .../cudf_polars/tests/streaming/test_tracing.py | 16 ++++++---------- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index ac8c5453ec72..4a03eeebc6e6 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -814,7 +814,12 @@ def print_query_plan( elif CUDF_POLARS_AVAILABLE: assert isinstance(engine, pl.GPUEngine) if args.explain_logical: - logical_plan = explain_query(q, engine, physical=False) + logical_plan = explain_query( + q, + engine, + optimized=run_config.frontend in _STREAMING_FRONTENDS, + physical=False, + ) if args.explain and run_config.frontend in _STREAMING_FRONTENDS: plan = explain_query(q, engine) else: diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 803c057e436b..fda95175eda6 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -91,6 +91,7 @@ def explain_query( q: pl.LazyFrame, engine: pl.GPUEngine, *, + optimized: bool = True, physical: bool = True, executor: concurrent.futures.Executor | None = None, ) -> str: @@ -103,6 +104,9 @@ def explain_query( The LazyFrame to explain. engine : pl.GPUEngine The configured GPU engine to use. + optimized + If True and showing the logical plan, run cudf-polars specific + query optimization. physical : bool, default True If True, show the physical (lowered) plan. If False, show the logical (pre-lowering) plan. @@ -141,7 +145,8 @@ 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) + if optimized: + ir = optimize_with_stats(ir, config, stats) return _repr_ir_tree(ir, stats=stats) else: return _repr_ir_tree(ir) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index b1d7e4f68993..5d59ff249220 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -485,7 +485,7 @@ def test_indirect_prefilter_trace_records_decision_and_effect( bloom_filter_max_size: int, method: str, reason: str, - domain_rows: int | None, + domain_rows: int, ) -> None: """Trace a composite prefilter pushed below an intervening join.""" pytest.importorskip("structlog") @@ -592,16 +592,12 @@ def test_indirect_prefilter_trace_records_decision_and_effect( "domain_rows": domain_rows, }.items() ) - if method == "skip": - assert "input_rows" not in record["prefilter"] - assert "output_rows" not in record["prefilter"] + assert record["prefilter"]["estimated_cardinality"] == domain_rows + assert record["prefilter"]["input_rows"] == 180 + if method == "broadcast_semi_join": + assert record["prefilter"]["output_rows"] == 45 else: - assert record["prefilter"]["estimated_cardinality"] == domain_rows - assert record["prefilter"]["input_rows"] == 180 - if method == "broadcast_semi_join": - assert record["prefilter"]["output_rows"] == 45 - else: - assert 45 <= record["prefilter"]["output_rows"] < 180 + assert 45 <= record["prefilter"]["output_rows"] < 180 def test_structlog_disabled_by_default(timeout_seconds: int): From 2b2963c0d99760d02de71340aeb6d11d73f0e22f Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 1 Sep 2026 17:51:47 +0100 Subject: [PATCH 22/22] Correct defaults --- python/cudf_polars/cudf_polars/utils/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 9eb52f3e24e4..41ed4dbf317c 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -944,7 +944,7 @@ class StreamingExecutor: default_factory=DynamicPlanningOptions ) join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field( - default_factory=JoinFilterPushdownOptions + default=None ) max_concurrent_io_tasks: MaxConcurrentIOTasks = dataclasses.field( default_factory=_make_default_factory(