NoahKusaba opened a new issue, #2454:
URL: https://github.com/apache/datafusion-ballista/issues/2454
**Is your feature request related to a problem or challenge? Please describe
what you are trying to do.**
Found while analysing TPC-H q8 at SF1000 in #2419.
AQE plans a broadcast from *estimated* statistics. If the broadcast side
then *measures* larger, the replan's `join_selection` swaps the `CollectLeft`
join. The broadcast `ExchangeExec` lands on the **probe** side but still
reports one partition, so the join stage runs as a single task. In q8, ~80M
probe rows went through one task.
Repro shape: `guessed_small`, estimated at 10 MB and broadcast, joined with
`mid`, a known 500 MB. `guessed_small`'s stage then reports 8 GB, and the
replan swaps it onto the probe side. Expected join stage:
```
ShuffleWriterExec: partitioning: UnknownPartitioning(4)
HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)],
projection=[val@2]
CoalescePartitionsExec
ShuffleReaderExec: upstream_stage: 1, partitioning:
UnknownPartitioning(4)
ShuffleReaderExec: upstream_stage: 0, partitioning:
UnknownPartitioning(4)
```
**Describe the solution you'd like**
Add a hook to DataFusion's `ExecutionPlan`, which `join_selection` calls on
the input it moves to the probe side during a `CollectLeft` swap:
```rust
fn as_probe_side(&self) -> Result<Option<Arc<dyn ExecutionPlan>>> { Ok(None)
}
```
Ballista implements it on `ExchangeExec`, so no extra optimizer rule is
needed:
```rust
fn as_probe_side(&self) -> Result<Option<Arc<dyn ExecutionPlan>>> {
if !self.broadcast ||
self.input.properties().output_ordering().is_some() {
return Ok(None);
}
Ok(Some(Arc::new(self.to_partitioned())))
}
```
The fix happens at the swap itself, with no rule ordering to get right. It
has been prototyped against DataFusion `main` and passes the probe-side tests
in Ballista. It needs an upstream DataFusion PR and a Ballista upgrade to a
release that includes it.
**Describe alternatives you've considered**
- **Ballista-only rule.** Run a rule after `join_selection` that replaces a
broadcast `ExchangeExec` on a `CollectLeft` probe side with `to_partitioned()`.
It works on the current DataFusion, so it could bridge until the hook lands.
However, it repairs the plan after the swap and depends on where it sits in the
rule order.
- **Remove `ExchangeExec`'s `broadcast` flag.** Derive broadcast from the
exchange's side of the join, so a swap can't misplace it. Cleanest, but a
larger refactor.
**Additional context**
- Null-aware anti joins (#2198) and broadcasts over ordered inputs must keep
a single partition.
- Related: #2434, #2024, #2121.
--
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]
