Rework join filter pushdown rewrites to leave hints and apply adaptively - #23584
Rework join filter pushdown rewrites to leave hints and apply adaptively#23584wence- wants to merge 17 commits into
Conversation
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.
github-actions
Bot
added
Python
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds optional runtime filtering for streaming joins. It introduces filter-hint IR nodes, runtime sampling and method selection, Bloom and broadcast semi-join execution, configurable Bloom-filter limits, explain serialization, tracing, and expanded tests. ChangesJoin prefilter planning and execution
Estimated code review effort: 5 (Critical) | ~90 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (9)
python/cudf_polars/tests/streaming/test_explain.py (1)
140-185: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the remaining prefilter serialization branches.
The test covers only a
JoinInputDomainprefilter withplacement == "join_input". Two new branches stay untested:
- The
ExternalDomainbranch of_serialize_prefilterinpython/cudf_polars/cudf_polars/streaming/explain.py(line 628), which emits{"type": "ExternalDomain"}with nosidekey.- The
placement == "pushed_down"value rendered by thePushdownFilterHintrepresentation and serializer.A composite-candidate query, such as the one in
test_composite_filter_pushdown_constrains_domain_first, produces both. Add a case that asserts those outputs.As per coding guidelines: "Ensure test files provide comprehensive edge case coverage (empty, all-null, single-element, mixed types) and do not depend on external datasets".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/tests/streaming/test_explain.py` around lines 140 - 185, Extend test_explain_pushdown_filter_hint_in_dynamic_physical_plan or add a focused composite-candidate test using the pattern from test_composite_filter_pushdown_constrains_domain_first. Assert the logical PushdownFilterHint serialization uses placement "pushed_down", and the physical prefilter serialization includes an ExternalDomain object with type "ExternalDomain" and no side key, while retaining the existing JoinInputDomain assertions.Source: Coding guidelines
python/cudf_polars/cudf_polars/streaming/actor_graph/core.py (1)
182-187: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtend the explanatory comment to cover
PushdownFilterHint.The comment lists
Union,Join, andOverbut notPushdownFilterHint. Add the reason for the new type so the branch stays self-explanatory.♻️ Proposed comment update
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 consumes its domain in + # full before it emits filtered target chunks. In every case the + # input source needs unbounded fanout so other consumers don't + # block it. _mark_children_unbounded(node)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/actor_graph/core.py` around lines 182 - 187, Extend the comment above the isinstance branch in the actor-graph logic to explicitly explain why PushdownFilterHint also requires unbounded fanout, while preserving the existing rationale for Union, Join, and Over.python/cudf_polars/cudf_polars/streaming/filter_hint.py (1)
84-94: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the positional alignment contract for external domains.
The constructor validates only the count of external domains. Consumers such as
make_join_planning_stateinpython/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.pypairchildren[2:]with the external prefilters in iteration order. Record this order requirement in the class docstring so future callers keep the pairing correct.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/filter_hint.py` around lines 84 - 94, Update the JoinWithPrefilter class docstring to document that external domains in children[2:] must correspond positionally, in iteration order, to the external prefilters. Keep the existing validation unchanged and explicitly describe this ordering contract for callers such as make_join_planning_state.python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py (1)
400-466: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConfirm the interaction between
bloom_filter_max_sizeandtarget.total_size.Line 436 compares
bloom_bytesagainstmin(bloom_filter_max_size, target.total_size). A small target therefore disables the Bloom filter even when the configured limit is large. The same applies to the exact path at line 446. The behavior looks intentional, because a filter should not cost more than the data it filters. Add a short comment that states this rule, so the two limits are not confused later.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py` around lines 400 - 466, The size checks in choose_prefilter_method should document that a prefilter must not consume more memory than the target data it filters. Add a concise comment near the bloom and exact size comparisons explaining that min(bloom_filter_max_size or broadcast_limit, target.total_size) enforces both the configured limit and the target-size cap.python/cudf_polars/tests/streaming/test_tracing.py (1)
209-222: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the subprocess trace harness into a helper.
The three new tests repeat the same steps: copy the environment, set
CUDF_POLARS_LOG_TRACES, runsubprocess.check_output, and parse thePREFILTER_TRACE=line. Move these steps into one module-level helper that accepts the generated code and returns the decoded record. This removes about forty duplicated lines and keeps the parsing rule in one place.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/tests/streaming/test_tracing.py` around lines 209 - 222, Extract the repeated subprocess trace setup and PREFILTER_TRACE parsing from the tracing tests into a module-level helper that accepts the generated code and returns the decoded JSON record. Update all three tests to call this helper, preserving the copied environment, CUDF_POLARS_LOG_TRACES setting, subprocess options, and single-trace-line parsing behavior.python/cudf_polars/cudf_polars/streaming/join.py (1)
344-353: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAssign
preserve_prefiltersdirectly from the condition.The
if/elseblock only assignsTrueorFalse. A direct assignment is shorter and keeps the condition in one place.♻️ Proposed simplification
- if ( + preserve_prefilters = ( dynamic_planning and ir.options[0] != "Cross" and ir.options[5] == "none" and not has_non_pointwise_keys and any(is_direct_join_prefilter(child) for child in ir.children) - ): - preserve_prefilters = True - else: - preserve_prefilters = False + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/join.py` around lines 344 - 353, In the join planning logic, replace the if/else assignment to preserve_prefilters with a direct boolean assignment using the existing dynamic_planning, join-type, options, has_non_pointwise_keys, and is_direct_join_prefilter conditions. Preserve the condition and resulting True/False behavior unchanged.python/cudf_polars/tests/test_config.py (1)
715-732: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd edge cases for
bloom_filter_max_size.The validation code rejects
boolexplicitly at line 433 ofpython/cudf_polars/cudf_polars/utils/config.py, and the documentation defines0as the disabling value. Neither case is covered here. Add one case that passesTrueand expectsTypeError, and one case that passes0and expects success.Based on path instructions: "Ensure test files provide comprehensive edge case coverage (empty, all-null, single-element, mixed types)".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/tests/test_config.py` around lines 715 - 732, Extend the bloom_filter_max_size tests in ConfigOptions.from_polars_engine coverage with a True input that expects TypeError and a 0 input that succeeds, preserving the existing invalid-string and negative-value cases.Source: Path instructions
python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py (1)
168-181: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueDefine a shared Protocol for
add_prefilter.spec
add_prefilterreads onlytarget_on,domain_on, andnulls_equal, which bothPrefilterandPushdownFilterHintprovide. Use a Protocol for these fields to make the structural contract explicit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py` around lines 168 - 181, Define a shared structural Protocol for the add_prefilter spec argument containing target_on, domain_on, and nulls_equal, and annotate add_prefilter to accept that Protocol. Ensure both Prefilter and PushdownFilterHint satisfy the contract without duplicating or narrowing their existing field definitions.python/cudf_polars/cudf_polars/streaming/parallel.py (1)
18-18: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
filter_hintimport.filter_hint.pyregisters no handlers, andjoin.pyimports its classes before registeringPushdownFilterHint.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/parallel.py` at line 18, Remove the unused cudf_polars.streaming.filter_hint import from the module imports in parallel.py; leave the existing join.py registration and class imports unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py`:
- Around line 138-141: Update the domain validation in the prefilter actor
around names_to_indices so malformed optional filters do not abort query
execution: use an assert only if lowering guarantees key-only domains; otherwise
route non-key domains through the same replay/fallback behavior as the skip
branch instead of raising ValueError.
In `@python/cudf_polars/cudf_polars/streaming/explain.py`:
- Line 144: Update the docstrings for explain_query and
SerializablePlan.from_query (including serialize_query) to state that
non-physical output is the optimized pre-lowering logical plan, after
optimize_with_stats, and may contain PushdownFilterHint nodes; replace the
inaccurate “logical (pre-lowering)” wording while preserving the physical-plan
documentation.
In `@python/cudf_polars/tests/streaming/test_tracing.py`:
- Around line 356-377: Update
test_indirect_prefilter_trace_records_decision_and_effect to remove the
unreachable method == "skip" branch, since all parameterized cases select bloom
or broadcast_semi_join, and narrow the domain_rows annotation from int | None to
int to match the supplied values.
---
Nitpick comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/core.py`:
- Around line 182-187: Extend the comment above the isinstance branch in the
actor-graph logic to explicitly explain why PushdownFilterHint also requires
unbounded fanout, while preserving the existing rationale for Union, Join, and
Over.
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py`:
- Around line 168-181: Define a shared structural Protocol for the add_prefilter
spec argument containing target_on, domain_on, and nulls_equal, and annotate
add_prefilter to accept that Protocol. Ensure both Prefilter and
PushdownFilterHint satisfy the contract without duplicating or narrowing their
existing field definitions.
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py`:
- Around line 400-466: The size checks in choose_prefilter_method should
document that a prefilter must not consume more memory than the target data it
filters. Add a concise comment near the bloom and exact size comparisons
explaining that min(bloom_filter_max_size or broadcast_limit, target.total_size)
enforces both the configured limit and the target-size cap.
In `@python/cudf_polars/cudf_polars/streaming/filter_hint.py`:
- Around line 84-94: Update the JoinWithPrefilter class docstring to document
that external domains in children[2:] must correspond positionally, in iteration
order, to the external prefilters. Keep the existing validation unchanged and
explicitly describe this ordering contract for callers such as
make_join_planning_state.
In `@python/cudf_polars/cudf_polars/streaming/join.py`:
- Around line 344-353: In the join planning logic, replace the if/else
assignment to preserve_prefilters with a direct boolean assignment using the
existing dynamic_planning, join-type, options, has_non_pointwise_keys, and
is_direct_join_prefilter conditions. Preserve the condition and resulting
True/False behavior unchanged.
In `@python/cudf_polars/cudf_polars/streaming/parallel.py`:
- Line 18: Remove the unused cudf_polars.streaming.filter_hint import from the
module imports in parallel.py; leave the existing join.py registration and class
imports unchanged.
In `@python/cudf_polars/tests/streaming/test_explain.py`:
- Around line 140-185: Extend
test_explain_pushdown_filter_hint_in_dynamic_physical_plan or add a focused
composite-candidate test using the pattern from
test_composite_filter_pushdown_constrains_domain_first. Assert the logical
PushdownFilterHint serialization uses placement "pushed_down", and the physical
prefilter serialization includes an ExternalDomain object with type
"ExternalDomain" and no side key, while retaining the existing JoinInputDomain
assertions.
In `@python/cudf_polars/tests/streaming/test_tracing.py`:
- Around line 209-222: Extract the repeated subprocess trace setup and
PREFILTER_TRACE parsing from the tracing tests into a module-level helper that
accepts the generated code and returns the decoded JSON record. Update all three
tests to call this helper, preserving the copied environment,
CUDF_POLARS_LOG_TRACES setting, subprocess options, and single-trace-line
parsing behavior.
In `@python/cudf_polars/tests/test_config.py`:
- Around line 715-732: Extend the bloom_filter_max_size tests in
ConfigOptions.from_polars_engine coverage with a True input that expects
TypeError and a 0 input that succeeds, preserving the existing invalid-string
and negative-value cases.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5743318a-7d20-4592-ac8a-7495dbfa3acc
📒 Files selected for processing (20)
python/cudf_polars/cudf_polars/dsl/ir.pypython/cudf_polars/cudf_polars/dsl/utils/column_domain.pypython/cudf_polars/cudf_polars/streaming/actor_graph/__init__.pypython/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.pypython/cudf_polars/cudf_polars/streaming/actor_graph/core.pypython/cudf_polars/cudf_polars/streaming/actor_graph/join.pypython/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.pypython/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.pypython/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.pypython/cudf_polars/cudf_polars/streaming/actor_graph/utils.pypython/cudf_polars/cudf_polars/streaming/explain.pypython/cudf_polars/cudf_polars/streaming/filter_hint.pypython/cudf_polars/cudf_polars/streaming/join.pypython/cudf_polars/cudf_polars/streaming/join_filter_pushdown.pypython/cudf_polars/cudf_polars/streaming/parallel.pypython/cudf_polars/cudf_polars/utils/config.pypython/cudf_polars/tests/streaming/test_explain.pypython/cudf_polars/tests/streaming/test_join_filter_pushdown.pypython/cudf_polars/tests/streaming/test_tracing.pypython/cudf_polars/tests/test_config.py
| # Include row-count statistics for the logical plan | ||
| with cm: | ||
| stats = collect_statistics(ir, config, executor) | ||
| ir = optimize_with_stats(ir, config, stats) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the docstrings for the new logical-plan semantics.
explain_query and SerializablePlan.from_query now apply optimize_with_stats before rendering the non-physical plan. The physical parameter is still documented as "If False, show the logical (pre-lowering) plan." in explain_query (line 108) and in serialize_query. The output is now the optimized pre-lowering plan, which contains PushdownFilterHint nodes that the translator never produces. Update both docstrings so callers know the rewrite is applied.
Also applies to: 888-893
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf_polars/cudf_polars/streaming/explain.py` at line 144, Update the
docstrings for explain_query and SerializablePlan.from_query (including
serialize_query) to state that non-physical output is the optimized pre-lowering
logical plan, after optimize_with_stats, and may contain PushdownFilterHint
nodes; replace the inaccurate “logical (pre-lowering)” wording while preserving
the physical-plan documentation.
| @pytest.mark.parametrize( | ||
| "broadcast_limit,bloom_filter_max_size,method,reason,domain_rows", | ||
| [ | ||
| (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 15), | ||
| (512, 0, "broadcast_semi_join", "exact_domain_fits", 15), | ||
| ( | ||
| 1_000_000, | ||
| 32 * 1024 * 1024, | ||
| "bloom", | ||
| "bloom_fits", | ||
| 15, | ||
| ), | ||
| ], | ||
| ids=["bloom", "exact", "bloom_despite_intervening_broadcast"], | ||
| ) | ||
| def test_indirect_prefilter_trace_records_decision_and_effect( | ||
| timeout_seconds: int, | ||
| broadcast_limit: int, | ||
| bloom_filter_max_size: int, | ||
| method: str, | ||
| reason: str, | ||
| domain_rows: int | None, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove the unreachable skip branch.
No parameter set of test_indirect_prefilter_trace_records_decision_and_effect produces method == "skip". The branch at lines 484-486 never runs. The annotation domain_rows: int | None is also wider than the supplied values, which are all 15. Either add a skip case or delete the branch and narrow the annotation to int.
Also applies to: 484-486
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf_polars/tests/streaming/test_tracing.py` around lines 356 - 377,
Update test_indirect_prefilter_trace_records_decision_and_effect to remove the
unreachable method == "skip" branch, since all parameterized cases select bloom
or broadcast_semi_join, and narrow the domain_rows annotation from int | None to
int to match the supplied values.
TomAugspurger
left a comment
TomAugspurger
left a comment
There was a problem hiding this comment.
Partial review, but overall the strategy and implementation make sense so far.
| if node in unbounded: | ||
| _mark_children_unbounded(node) | ||
| elif isinstance(node, (Union, Join, Over)): | ||
| elif isinstance(node, (Union, Join, Over, PushdownFilterHint)): |
There was a problem hiding this comment.
Nitpick: the comment below describes why Union, Join, and Over need unbounded fanout.
Maybe it's just a person preference thing, but I'd prefer an IR.needs_unbounaded_fanout property / class variable, and then we can override that for specific IR nodes, and document the reasoning where we set the property.
But this is fine too.
There was a problem hiding this comment.
Whether one needs unbounded fanout is a property of the implementation choice, not the IR node per-se. So I prefer it here.
| chunks.clear() | ||
|
|
||
|
|
||
| async def count_rows_passthrough( |
There was a problem hiding this comment.
Maybe this could live in actor_graph/tracing.py or actor_graph/utils.py. Only the join prefiltering is using it today, but I could see that changing in the near future.
3 participants
Description
In #22996 and #22997 we added join filter pushdown optimisations for cudf-polars plans. These are represented as semi joins using the existing IR structure. As such, we are on the hook to execute them, even if that would not be beneficial for runtime execution. Examples are cases where the join that is being filtered will be performed via broadcast, or is already compatibly shuffled. In such cases carrying out the semi join merely adds extra work.
To fix this, introduce a special
PushdownFilterHintnode that is, optionally, applied at runtime. To decide whether, and how, to apply these hints we now sample join inputs and estimate cardinality in addition to size. We use the cardinality estimate to decide whether or not a bloom filter would be effective: calculating the estimated false positive rate.During lowering, we build a join planning state that is updated at runtime with cardinality estimates and other relevant information such as partitioning. When we then come to execute a filter hint we can inspect this state and take an appropriate action.
Filter hints are classified into two types:
Join(Hint(target, domain), domain)Join1(Join2(Hint(target, j1_domain), j2_domain), j1_domain)The former check at runtime whether or not the domain is already compatibly distributed (or will be broadcast) and then elide the hint. The latter cannot do so without inducing a cycle in the execution DAG, which is undesirable. We therefore always apply these "indirect" hints, but still only if the hint does not induce a shuffle. Better cost models for placement might elide some of these, or give us enough information to decide whether to reject them based on sampling.
Checklist