NoahKusaba opened a new pull request, #25354: URL: https://github.com/apache/datafusion/pull/25354
## Which issue does this PR close? - [Related to apache/datafusion-ballista#2454.](https://github.com/apache/datafusion-ballista/issues/2454) ## Rationale for this change `JoinSelection` can swap a `CollectLeft` hash join so the former build side becomes the probe side. That input may have been collapsed to a single partition because it was the build side, for example a broadcast exchange in a distributed engine. After the swap it still reports one partition, so the join runs as a single task. In Ballista this ran TPC-H q8 at SF1000 (stage 5) through one task with 79.9M input rows. `JoinSelection` has no way to let such an input undo the collapse, and a later optimizer rule can only patch the plan after the swap has happened. ## What changes are included in this PR? - `ExecutionPlan::as_probe_side`, a new method defaulting to `Ok(None)`. `JoinSelection` calls it on the input it moves to the probe side during a `CollectLeft` swap, and the input may return a replacement plan. - `swap_to_collect_left`, used by every swap that produces a `CollectLeft` join: both branches of `try_collect_left` and the null-aware swap in the `Partitioned` branch. Plans that return `None`, which is the default, behave exactly as before. ## What is the testing strategy for this PR? `test_collect_left_swap_uses_as_probe_side` in `core/tests/physical_optimizer/join_selection.rs` asserts that the swapped join's probe side is the plan the hook returned. It runs three cases, one per call site: - `PartitionMode::CollectLeft`, reaching the `(true, true)` branch of `try_collect_left` - `PartitionMode::Auto`, reaching the `(false, true)` branch - `PartitionMode::Partitioned` with `null_aware`, reaching the null-aware swap Each case was checked by reverting its own call site to `swap_inputs`: only that case failed. The existing swap tests cover the default `None` path, since `as_probe_side` returns `None` for every other plan. ## Are there any user-facing changes? A new `ExecutionPlan` method with a default implementation, so this is not a breaking change. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
