Add dynamic join key prefilter planning - #22995
Conversation
Add generic dynamic-planning support for join key prefilters in the streaming actor graph. The planner evaluates join type, key compatibility, size estimates, and configured selectivity thresholds to decide when a small side can build a bloom/key prefilter for the larger side before shuffle. The implementation supports prefix key selection for multi-key joins, records structured trace metadata and skip reasons, preserves the original full join after the row-reduction stage, and exposes conservative dynamic-planning options for enabling, sizing, and tracing the prefilter path.
pentschev
added
feature request
pentschev
added
non-breaking
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughDynamic planning now exposes join-prefilter threshold, key-column cap, and tracing settings. Join execution selects a prefilter decision, records tracer metadata, and routes the shuffle/filter path from that decision. Tests cover configuration parsing, tracer extras, and join prefilter outcomes. ChangesJoin prefilter flow
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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: 1
🤖 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/utils/config.py`:
- Around line 390-396: The validation in the join_prefilter_max_key_columns
setter currently accepts bool values because isinstance(True, int) passes, so
update the type check to reject booleans explicitly while still allowing only
real ints or None. Adjust the validation logic in the config class around
join_prefilter_max_key_columns so that a boolean raises the same TypeError as
other invalid types, and keep the existing lower-bound check for valid integers.
🪄 Autofix (Beta)
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: d0830731-e945-42bc-b766-a27a2e2247fd
📒 Files selected for processing (8)
python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.pypython/cudf_polars/cudf_polars/streaming/actor_graph/join.pypython/cudf_polars/cudf_polars/streaming/actor_graph/tracing.pypython/cudf_polars/cudf_polars/streaming/actor_graph/utils.pypython/cudf_polars/cudf_polars/utils/config.pypython/cudf_polars/tests/streaming/test_join.pypython/cudf_polars/tests/streaming/test_tracing.pypython/cudf_polars/tests/test_config.py
wence-
left a comment
wence-
left a comment
There was a problem hiding this comment.
So IIUC, this change adds better tracing to the bloom filter selection: good.
The rest of the changes don't appear to actually make any changes to the runtime decision making, but I may have misread things.
I think possibly the only thing is that we allow configuring a key_prefix for the filter, rather than always considering all keys.
| "Dynamic join requires 3 reserved collective IDs " | ||
| "(allgather + left shuffle + right shuffle + bloom filter); got " | ||
| "Dynamic join requires 4 reserved collective IDs " | ||
| "(size allgather + strategy allgather/bloom prefilter " |
There was a problem hiding this comment.
If the two allgathers are not concurrent we can reuse the tag.
There was a problem hiding this comment.
OK, so the substantive change to this comment is really updating the (incorrect) 3 to 4.
There was a problem hiding this comment.
Yes, the size allgather completes before the later strategy-specific collective, so the tag could be reused if we want to minimize the reserved ID count. The current version keeps a distinct ID per possible collective call, which is simpler to reason about but more conservative than strictly necessary.
The main correction here is that the old comment/check was stale: dynamic joins can consume size allgather + bloom/allgather + left shuffle + right shuffle.
Are you ok keeping the explicit 4-ID reservation with corrected wording here? As you pointed to, the alternative is reworking this to reuse the size-estimate tag after strategy selection, but I'd prefer to punt "good-to-have" changes to a separate PR if possible. Let me know what you think.
There was a problem hiding this comment.
The changed comment is wrong because you never allgather a strategy. So can you just flip the "3" to "4"?
| right_rows=right_rows, | ||
| threshold=threshold, | ||
| ) | ||
| if key_column_count != len(right_key_indices): |
There was a problem hiding this comment.
When can the left and right key length not be the same?
There was a problem hiding this comment.
For a valid Join IR, they should not differ. This check is defensive rather than a case I expect from normal planning. A mismatch would indicate malformed join metadata or a future caller passing inconsistent key index tuples.
If you prefer I can make this an assert instead of a skipped-planning reason, since it is not really a runtime profitability/legality decision.
There was a problem hiding this comment.
Please assert. If we reach this state we are in an impossible place and so we should actually raise an error.
| join_prefilter_max_key_columns | ||
| Maximum number of join-key columns to use for the prefilter. Set to | ||
| ``None`` to use all join keys. Default is 1. |
There was a problem hiding this comment.
What is the idea here? Under what circumstances do we not want to consider all keys in the prefilter?
I suppose if you have a composite key where total composite key cardinality is high but the prefix of the key is low cardinality then this might give you better selectivity.
But do we specifically need it?
There was a problem hiding this comment.
The idea is to decouple the prefilter key set from the full join key set. The prefilter is only a conservative row-reduction step; the original join still runs afterward with the complete join keys, so using a prefix can keep extra rows but should not remove rows that could satisfy the full join.
For Q9 this is material. The relevant join is the composite-key join:
(p_partkey, ps_suppkey) = (l_partkey, l_suppkey)
With the current default join_prefilter_max_key_columns=1, the local prefilter uses only the p_partkey/l_partkey prefix. The trace from Q9/SF30K/8 nodes shows that this reduces the large lineitem side substantially before the shuffle/full join: one actor filters about 3.30B rows to 1.70B rows, and another filters about 5.70B rows to 2.94B rows. The query completes and validates in that configuration.
Forcing all join keys with join_prefilter_max_key_columns=None, Q9 OOMs. So this is not just a tuning convenience, for this workload using all keys changes the memory/runtime behavior enough to lose the Q9 improvement that this PR is trying to provide.
So yes, I think we specifically need the ability to use a key prefix rather than always using all keys. Whether this remains a user-facing option or becomes an internal conservative default is a separate question (let me know what your thoughts are), but it seems clear we cannot use all keys and still avoid OOMing.
There was a problem hiding this comment.
OK, but what if the selective key was the match on the second part of the key, not the prefix?
There was a problem hiding this comment.
Please also document that this selects the size of the prefix of join keys considered for filtering.
There was a problem hiding this comment.
OK, but what if the selective key was the match on the second part of the key, not the prefix?
For this PR specifically, we do not handle that case automatically. Here, the local bloom prefilter uses a key prefix, so join_prefilter_max_key_columns=1 means the first join key only.
There is a later IR-level optimizer branch in #22996 , that starts to address this more generally. That branch enumerates join-key positions when looking for derived domain prefilter candidates, so a selective second key can be used if the selective domain is visible in the IR. The follow-up to that #22997 , keeps that capability and adds profitability/source-cost/stacking guards so those derived filters do not introduce regressions.
So I agree this PR is not a general "best key subset" selector. It preserves the validated prefix behavior needed for Q9 in the local dynamic bloom path. More general key-position/subset selection belongs in the follow-up domain-prefilter optimizer branches, or in a future costed selector for this local bloom path.
There was a problem hiding this comment.
Please also document that this selects the size of the prefix of join keys considered for filtering.
Done in fbe886b. Please let me know if that still isn't clear.
| def _optional_float_converter(v: str) -> float | None: | ||
| if v.lower() in {"none", "null"}: | ||
| return None | ||
| return float(v) | ||
|
|
||
|
|
||
| def _optional_int_converter(v: str) -> int | None: | ||
| if v.lower() in {"none", "null"}: | ||
| return None | ||
| return int(v) |
There was a problem hiding this comment.
T = TypeVar("T")
def _optional_converter(v: str, parse: Callable[[str], T]) -> T | None:
if v.lower() in {"none", "null"}:
return None
return parse(v)
| if join_prefilter_threshold is None: | ||
| join_prefilter_threshold = self.bloom_filter_threshold | ||
| object.__setattr__( | ||
| self, "join_prefilter_threshold", join_prefilter_threshold | ||
| ) | ||
| elif not isinstance(join_prefilter_threshold, float): |
There was a problem hiding this comment.
This introduces a new way of configuring thresholds for bloom filter application. I think we are just safe to cull the old codepath
There was a problem hiding this comment.
Agreed. The fallback was intended to preserve the old bloom-filter threshold behavior, but after this change there is really one join-prefilter decision path. Keeping both knobs makes the config harder to explain.
As part of 519cd1d, I’ve simplified this so the new join prefilter threshold is the single threshold used by this path, and removed the old bloom-filter-specific code.
Reject boolean values for join_prefilter_max_key_columns instead of accepting them through Python's bool-is-int relationship. Add explicit config validation coverage so only None and positive integer limits remain valid.
Use a single typed helper for optional environment-value parsing and define the optional float and int converters in terms of it. This keeps the None/null handling in one place without changing accepted values.
Make join_prefilter_threshold the single dynamic-planning threshold for join prefilters, preserving the existing 0.5 default. Remove the obsolete bloom_filter_threshold option, its fallback/override tests, and the unused use_bloom_filter helper now replaced by _select_join_prefilter.
| right_rows=right_rows, | ||
| threshold=threshold, | ||
| ) | ||
| if key_column_count != len(right_key_indices): |
There was a problem hiding this comment.
For a valid Join IR, they should not differ. This check is defensive rather than a case I expect from normal planning. A mismatch would indicate malformed join metadata or a future caller passing inconsistent key index tuples.
If you prefer I can make this an assert instead of a skipped-planning reason, since it is not really a runtime profitability/legality decision.
| "Dynamic join requires 3 reserved collective IDs " | ||
| "(allgather + left shuffle + right shuffle + bloom filter); got " | ||
| "Dynamic join requires 4 reserved collective IDs " | ||
| "(size allgather + strategy allgather/bloom prefilter " |
There was a problem hiding this comment.
Yes, the size allgather completes before the later strategy-specific collective, so the tag could be reused if we want to minimize the reserved ID count. The current version keeps a distinct ID per possible collective call, which is simpler to reason about but more conservative than strictly necessary.
The main correction here is that the old comment/check was stale: dynamic joins can consume size allgather + bloom/allgather + left shuffle + right shuffle.
Are you ok keeping the explicit 4-ID reservation with corrected wording here? As you pointed to, the alternative is reworking this to reuse the size-estimate tag after strategy selection, but I'd prefer to punt "good-to-have" changes to a separate PR if possible. Let me know what you think.
| if join_prefilter_threshold is None: | ||
| join_prefilter_threshold = self.bloom_filter_threshold | ||
| object.__setattr__( | ||
| self, "join_prefilter_threshold", join_prefilter_threshold | ||
| ) | ||
| elif not isinstance(join_prefilter_threshold, float): |
There was a problem hiding this comment.
Agreed. The fallback was intended to preserve the old bloom-filter threshold behavior, but after this change there is really one join-prefilter decision path. Keeping both knobs makes the config harder to explain.
As part of 519cd1d, I’ve simplified this so the new join prefilter threshold is the single threshold used by this path, and removed the old bloom-filter-specific code.
| join_prefilter_max_key_columns | ||
| Maximum number of join-key columns to use for the prefilter. Set to | ||
| ``None`` to use all join keys. Default is 1. |
There was a problem hiding this comment.
The idea is to decouple the prefilter key set from the full join key set. The prefilter is only a conservative row-reduction step; the original join still runs afterward with the complete join keys, so using a prefix can keep extra rows but should not remove rows that could satisfy the full join.
For Q9 this is material. The relevant join is the composite-key join:
(p_partkey, ps_suppkey) = (l_partkey, l_suppkey)
With the current default join_prefilter_max_key_columns=1, the local prefilter uses only the p_partkey/l_partkey prefix. The trace from Q9/SF30K/8 nodes shows that this reduces the large lineitem side substantially before the shuffle/full join: one actor filters about 3.30B rows to 1.70B rows, and another filters about 5.70B rows to 2.94B rows. The query completes and validates in that configuration.
Forcing all join keys with join_prefilter_max_key_columns=None, Q9 OOMs. So this is not just a tuning convenience, for this workload using all keys changes the memory/runtime behavior enough to lose the Q9 improvement that this PR is trying to provide.
So yes, I think we specifically need the ability to use a key prefix rather than always using all keys. Whether this remains a user-facing option or becomes an internal conservative default is a separate question (let me know what your thoughts are), but it seems clear we cannot use all keys and still avoid OOMing.
| def _optional_float_converter(v: str) -> float | None: | ||
| if v.lower() in {"none", "null"}: | ||
| return None | ||
| return float(v) | ||
|
|
||
|
|
||
| def _optional_int_converter(v: str) -> int | None: | ||
| if v.lower() in {"none", "null"}: | ||
| return None | ||
| return int(v) |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/utils/config.py`:
- Around line 318-320: `DynamicPlanningOptions.join_prefilter_threshold` is
documented as accepting 0 to disable prefiltering, but `__post_init__` currently
rejects that value by enforcing only `float`. Update the validation in
`DynamicPlanningOptions.__post_init__` (and any related type checks/default
handling) so numeric zero is accepted as a valid disable sentinel while
preserving existing float-range validation for other values. Make sure the
docstring and runtime behavior stay aligned for `join_prefilter_threshold`.
🪄 Autofix (Beta)
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: d8c25a79-c130-41a7-baac-3bb261f04775
📒 Files selected for processing (3)
python/cudf_polars/cudf_polars/streaming/actor_graph/join.pypython/cudf_polars/cudf_polars/utils/config.pypython/cudf_polars/tests/test_config.py
💤 Files with no reviewable changes (1)
- python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
Allow numeric integer values such as 0 for join_prefilter_threshold and normalize the stored value to float during DynamicPlanningOptions validation. Reject booleans explicitly and add config coverage for the documented disable value.
| key_column_count: int = 0, | ||
| ) -> JoinPrefilterDecision: | ||
| return JoinPrefilterDecision( | ||
| considered=True, |
There was a problem hiding this comment.
We never set considered=False so it seems the existence of a JoinPrefilterDecision means we considered it.
| if key_column_count == 0: | ||
| return _skipped_prefilter( | ||
| "no_join_keys", | ||
| left_rows=left_rows, | ||
| right_rows=right_rows, | ||
| threshold=threshold, | ||
| ) |
There was a problem hiding this comment.
When are we joining with no join keys?
There was a problem hiding this comment.
This was in fact not needed for the supported prefilter path. In 6bee038 I removed the no_join_keys branch and added coverage for the actual keyless case (Cross), which now skips via the existing unsupported_join_type path.
| build_side = "left" | ||
| filter_side = "right" |
There was a problem hiding this comment.
We don't need to track build side and filter side because they are each others inverse.
There was a problem hiding this comment.
You're right, simplified that in 97466a3.
| "Dynamic join requires 3 reserved collective IDs " | ||
| "(allgather + left shuffle + right shuffle + bloom filter); got " | ||
| "Dynamic join requires 4 reserved collective IDs " | ||
| "(size allgather + strategy allgather/bloom prefilter " |
There was a problem hiding this comment.
The changed comment is wrong because you never allgather a strategy. So can you just flip the "3" to "4"?
| def _skipped_prefilter( | ||
| reason: str, | ||
| *, | ||
| left_rows: int, | ||
| right_rows: int, | ||
| threshold: float, | ||
| ratio: float | None = None, | ||
| key_column_count: int = 0, | ||
| ) -> JoinPrefilterDecision: | ||
| return JoinPrefilterDecision( |
There was a problem hiding this comment.
This function seems unnecessary, we can just inline the construction of the JoinPrefilterDecision at the call sites and not lose any abstraction.
| join_prefilter_max_key_columns | ||
| Maximum number of join-key columns to use for the prefilter. Set to | ||
| ``None`` to use all join keys. Default is 1. |
There was a problem hiding this comment.
Please also document that this selects the size of the prefix of join keys considered for filtering.
| prefilter_trace_stats["apply_side_prepartitioned"] = _is_already_partitioned( | ||
| apply_meta, | ||
| prefilter_decision.apply_indices, | ||
| strategy.shuffle_modulus, | ||
| comm.nranks, | ||
| ) |
There was a problem hiding this comment.
Why is this relevant information for the filter? Should it not be tracked by the join? Or I guess, the prefilter doesn't make any adaptive decisions based on this information so tracking it is not that useful?
There was a problem hiding this comment.
Good point. This was added in the spirit of making the dynamic path observable, but looking at it again, apply_side_prepartitioned is not really prefilter decision metadata, the prefilter does not use it to decide whether to run, which side to filter, or which keys to use.
The prefilter trace still records the decision inputs/outcome, so I removed this field from join_prefilter metadata in 543a73d.
| """Return structured trace metadata for this decision.""" | ||
| metadata: dict[str, Any] = { | ||
| "considered": self.considered, | ||
| "estimated_left_rows": self.left_rows, | ||
| "estimated_right_rows": self.right_rows, | ||
| "threshold": self.threshold, | ||
| "filtered_side": self.filter_side, | ||
| "build_side": self.build_side, | ||
| "key_column_count": self.key_column_count, | ||
| } | ||
| if self.small_large_ratio is not None: | ||
| metadata["small_large_ratio"] = self.small_large_ratio | ||
| if self.reason_skipped is not None: | ||
| metadata["reason_skipped"] = self.reason_skipped | ||
| return metadata |
There was a problem hiding this comment.
Seems like, given we have a dataclass, we could just to to_dict() and not implement this.
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
Clarify that join_prefilter_max_key_columns controls the size of the join-key prefix used by the prefilter, rather than selecting an arbitrary key subset.
Replace the hand-written JoinPrefilterDecision trace metadata mapping with dataclasses.asdict, so the trace output follows the dataclass fields without duplicating the field list.
Drop the always-true JoinPrefilterDecision.considered field since the presence of a decision already means the prefilter was considered. Inline skipped JoinPrefilterDecision construction at the return sites so the selector does not carry a helper that only forwards dataclass arguments.
Drop the no_join_keys branch from the join prefilter selector. Keyless cross joins are already unsupported by the selector, so cover that behavior directly instead of carrying a separate skip reason.
Replace the defensive mismatched-key skip reason with an assertion. A valid Join IR must provide the same number of left and right join keys, so reaching this state indicates malformed join metadata rather than a prefilter planning decision.
Remove the redundant build_side field from JoinPrefilterDecision. The bloom-filter build side is the inverse of filter_side, so task construction now derives that relationship from the side being filtered.
Drop apply_side_prepartitioned from join prefilter trace metadata. The prefilter does not make adaptive decisions from this value and the information is not consumed elsewhere, so keeping it in the filter trace adds noise without affecting planning.
Keep the dynamic-join reserved collective count at four, but remove the inaccurate strategy-allgather wording from the comment and error message. The four reserved IDs are the allgather, left shuffle, right shuffle, and bloom filter.
Clarify that ActorTracer extra metadata is for nested runtime decisions that do not have their own IR node but should still be logged with the parent actor trace.
| join_prefilter_max_key_columns | ||
| Maximum number of join-key columns to use for the prefilter. Set to | ||
| ``None`` to use all join keys. Default is 1. |
There was a problem hiding this comment.
OK, but what if the selective key was the match on the second part of the key, not the prefix?
For this PR specifically, we do not handle that case automatically. Here, the local bloom prefilter uses a key prefix, so join_prefilter_max_key_columns=1 means the first join key only.
There is a later IR-level optimizer branch in #22996 , that starts to address this more generally. That branch enumerates join-key positions when looking for derived domain prefilter candidates, so a selective second key can be used if the selective domain is visible in the IR. The follow-up to that #22997 , keeps that capability and adds profitability/source-cost/stacking guards so those derived filters do not introduce regressions.
So I agree this PR is not a general "best key subset" selector. It preserves the validated prefix behavior needed for Q9 in the local dynamic bloom path. More general key-position/subset selection belongs in the follow-up domain-prefilter optimizer branches, or in a future costed selector for this local bloom path.
| join_prefilter_max_key_columns | ||
| Maximum number of join-key columns to use for the prefilter. Set to | ||
| ``None`` to use all join keys. Default is 1. |
There was a problem hiding this comment.
Please also document that this selects the size of the prefix of join keys considered for filtering.
Done in fbe886b. Please let me know if that still isn't clear.
| """Return structured trace metadata for this decision.""" | ||
| metadata: dict[str, Any] = { | ||
| "considered": self.considered, | ||
| "estimated_left_rows": self.left_rows, | ||
| "estimated_right_rows": self.right_rows, | ||
| "threshold": self.threshold, | ||
| "filtered_side": self.filter_side, | ||
| "build_side": self.build_side, | ||
| "key_column_count": self.key_column_count, | ||
| } | ||
| if self.small_large_ratio is not None: | ||
| metadata["small_large_ratio"] = self.small_large_ratio | ||
| if self.reason_skipped is not None: | ||
| metadata["reason_skipped"] = self.reason_skipped | ||
| return metadata |
| key_column_count: int = 0, | ||
| ) -> JoinPrefilterDecision: | ||
| return JoinPrefilterDecision( | ||
| considered=True, |
| if key_column_count == 0: | ||
| return _skipped_prefilter( | ||
| "no_join_keys", | ||
| left_rows=left_rows, | ||
| right_rows=right_rows, | ||
| threshold=threshold, | ||
| ) |
There was a problem hiding this comment.
This was in fact not needed for the supported prefilter path. In 6bee038 I removed the no_join_keys branch and added coverage for the actual keyless case (Cross), which now skips via the existing unsupported_join_type path.
| build_side = "left" | ||
| filter_side = "right" |
There was a problem hiding this comment.
You're right, simplified that in 97466a3.
| prefilter_trace_stats["apply_side_prepartitioned"] = _is_already_partitioned( | ||
| apply_meta, | ||
| prefilter_decision.apply_indices, | ||
| strategy.shuffle_modulus, | ||
| comm.nranks, | ||
| ) |
There was a problem hiding this comment.
Good point. This was added in the spirit of making the dynamic path observable, but looking at it again, apply_side_prepartitioned is not really prefilter decision metadata, the prefilter does not use it to decide whether to run, which side to filter, or which keys to use.
The prefilter trace still records the decision inputs/outcome, so I removed this field from join_prefilter metadata in 543a73d.
| "Dynamic join requires 3 reserved collective IDs " | ||
| "(allgather + left shuffle + right shuffle + bloom filter); got " | ||
| "Dynamic join requires 4 reserved collective IDs " | ||
| "(size allgather + strategy allgather/bloom prefilter " |
| def _skipped_prefilter( | ||
| reason: str, | ||
| *, | ||
| left_rows: int, | ||
| right_rows: int, | ||
| threshold: float, | ||
| ratio: float | None = None, | ||
| key_column_count: int = 0, | ||
| ) -> JoinPrefilterDecision: | ||
| return JoinPrefilterDecision( |
Cast the join_prefilter_max_key_columns default to int | None so _make_default_factory infers the same optional type as _optional_int_converter. This avoids mypy narrowing the default to int and rejecting the converter.
wence-
left a comment
wence-
left a comment
There was a problem hiding this comment.
I have some small suggestions for (maybe) tidying up the logic, overall looking good
| f"{_env_prefix}__BLOOM_FILTER_THRESHOLD", float, default=0.5 | ||
| f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", | ||
| _optional_int_converter, | ||
| default=cast("int | None", 1), |
There was a problem hiding this comment.
Why do you need this cast?, 1 is definitely an int | None.
There was a problem hiding this comment.
Agreed, the cast should not be necessary. It was only compensating for _make_default_factory using the same type variable for both the converter result and default, which caused mypy to infer int from default=1 and reject the Callable[[str], int | None] converter.
In 702c6ad I now removed the cast and updated _make_default_factory to model the converter result and fallback default with separate type variables. Its return type is now their union, so this case type-checks naturally without changing runtime behavior.
| Select a safe join-key prefilter. | ||
|
|
||
| The prefilter only removes rows that cannot participate in the original | ||
| join. The full join still runs afterward with the complete key set. | ||
| """ |
There was a problem hiding this comment.
This docstring is not very informative. Also the bit about the full join running afterwards is kind of not relevant. Since that is not a decision this function makes.
Perhaps something like:
Determine if a pre-filter should apply to a join.
Parameters
----------
join_type
Type of join
left_rows
Estimate of the number of rows in the left table
right_rows
Estimate of the number of rows in the right table
left_key_indices
Column indices of the keys in the left table
right_key_indices
Column indices of the keys in the right table
threshold
Size ratio above which filtering is turned off
max_key_columns
Number of columns to use in the key prefix for the filter
Returns
-------
JoinPrefilterDecision
The determination of whether a pre-filter should apply.
There was a problem hiding this comment.
I replaced the docstring in a37608d with one describing the selector’s parameters, returned type, and removed the execution details outside this function’s responsibility as you suggested. I also clarified the exact threshold boundary and the behavior of max_key_columns=None.
| small_large_ratio=ratio, | ||
| key_column_count=key_column_count, | ||
| reason_skipped="no_legal_large_side", | ||
| ) |
There was a problem hiding this comment.
This whole function is "straightforward" but very hard to follow because it's a combination of early exit, and fallthrough.
I think that it would be easier to follow if there were a single exit at the end and all of these branches just select the decision parameters.
We can early exit for unsupported join types and/or prefiltering disabled by parameters, so something like:
small_rows, large_rows = sorted([left_rows, right_rows])
ratio = small_rows / large_rows if large_rows > 0 else None
if threshold == 0:
return JPD(reason_skipped="disabled")
elif join_type not in ("Inner", "Semi", "Left", "Anti", "Right"):
return JPD(reason_skipped="unsupported_join_type")
elif ratio is None:
return JPD(reason_skipped="no_large_side")
elif ratio >= threshold:
return JPD(reason_skipped="ratio_above_threshold")
reason_skipped = None
filter_side = None
if join_type in ("Inner", "Semi"):
filter_side = "right" if left_rows <= right_rows else "left"
elif join_type in ("Left", "Anti")
if left_rows >= right_rows:
reason_skipped="no_legal_large_side"
else:
filter_side="right"
elif join_type == "Right":
if right_rows >= left_rows:
reason_skipped="no_legal_large_side"
else:
filter_side = "left"
else:
assert_never(join_type)
if filter_side == "right":
build_indices = ...
apply_indices = ...
else:
...
return JPD(..., reason_skipped=reason_skipped)
WDYT?
It's not 100% clear to me that this will come out clearer though, so let's try.
There was a problem hiding this comment.
Thanks for the suggestion. I've implemented that in 6bc7559, and I think that is indeed much clearer.
Model the environment converter result and fallback default as separate type variables. This lets optional converters use concrete non-optional defaults without casts while preserving the factory's complete return type.
Describe the inputs that drive join prefilter planning and the decision returned by _select_join_prefilter. Remove execution details that are outside the selector's responsibility.
Separate unsupported and disabled cases from supported join planning, then collect filter-side, ratio, and skip-reason state into one final JoinPrefilterDecision. Preserve the existing selection behavior and trace metadata while making the control flow easier to follow.
| f"{_env_prefix}__BLOOM_FILTER_THRESHOLD", float, default=0.5 | ||
| f"{_env_prefix}__JOIN_PREFILTER_MAX_KEY_COLUMNS", | ||
| _optional_int_converter, | ||
| default=cast("int | None", 1), |
There was a problem hiding this comment.
Agreed, the cast should not be necessary. It was only compensating for _make_default_factory using the same type variable for both the converter result and default, which caused mypy to infer int from default=1 and reject the Callable[[str], int | None] converter.
In 702c6ad I now removed the cast and updated _make_default_factory to model the converter result and fallback default with separate type variables. Its return type is now their union, so this case type-checks naturally without changing runtime behavior.
| Select a safe join-key prefilter. | ||
|
|
||
| The prefilter only removes rows that cannot participate in the original | ||
| join. The full join still runs afterward with the complete key set. | ||
| """ |
There was a problem hiding this comment.
I replaced the docstring in a37608d with one describing the selector’s parameters, returned type, and removed the execution details outside this function’s responsibility as you suggested. I also clarified the exact threshold boundary and the behavior of max_key_columns=None.
| small_large_ratio=ratio, | ||
| key_column_count=key_column_count, | ||
| reason_skipped="no_legal_large_side", | ||
| ) |
There was a problem hiding this comment.
Thanks for the suggestion. I've implemented that in 6bc7559, and I think that is indeed much clearer.
wence-
left a comment
wence-
left a comment
There was a problem hiding this comment.
Still missing one line of coverage in the config. Looks good!
|
Thanks @wence- ! |
|
/merge |
…truct pre-filters for inner joins (#22996) Add a streaming optimizer pass that attempts to pre-filter one side of an input to inner joins before actor-graph lowering. The pass uses existing dynamic-planning scan statistics and join metadata to determine where it is beneficial to push a semi-join against a join key onto the other side of a join. The simplest example of such a rewrite is that we turn ```python left.join(right, on="key", how="inner") ``` into, assuming we somehow determine that `right` is selective, ```python ( left.join(right.select("key"), on="key", how="semi") .join(right, on="key", how="inner") ) ``` The optimization pass handles the case where a "domain" key, used to provide the right-hand side of the semi join, is "simple" and derived directly from some input node, as well as the more complex case where a domain key is already constrained by some other semi-join filter. Only inner joins are rewritten, and only if all the keys are simple column keys. If heuristics determine that simple keys are not selective, we also don't perform the rewrite. Material results of this change running NDSH SF30K on 8xNVL4 nodes are (previous results come from the change in #22995): * **Q5 doesn't OOM on 8 nodes anymore and improved runtime performance: 9.35s lukewarm, 5.14s hot** (previously 40.25s lukewarm, OOM on hot) * Q9 unchanged performance or slight regression: 47.12s lukewarm, 32.68s hot (previously 43.63s lukewarm, 30.56s hot) Authors: - Peter Andreas Entschev (https://github.com/pentschev) - Lawrence Mitchell (https://github.com/wence-) Approvers: - Lawrence Mitchell (https://github.com/wence-) - Mads R. B. Kristensen (https://github.com/madsbk) - Tom Augspurger (https://github.com/TomAugspurger) URL: #22996
4 participants
Add generic dynamic-planning support for join key prefilters in the streaming actor graph. The planner evaluates join type, key compatibility, size estimates, and configured selectivity thresholds to decide when a small side can build a bloom/key prefilter for the larger side before shuffle.
The implementation supports prefix key selection for multi-key joins, records structured trace metadata and skip reasons, preserves the original full join after the row-reduction stage, and exposes conservative dynamic-planning options for enabling, sizing, and tracing the prefilter path.
Material results of this change running NDSH SF30K on 8xNVL4 nodes are: