Dandandan opened a new pull request, #2342:
URL: https://github.com/apache/datafusion-ballista/pull/2342
## Which issue does this PR close?
None filed. Found while checking whether AQE can hold several broadcast
joins in one stage — it can (q5 stage 3 and q8 stage 4 each hold four), but q7
splits two apart and runs the second one single-threaded.
## Rationale for this change
The `CollectLeft` arm of `SelectJoinRule` puts a multi-partition build side
behind a broadcast exchange, which becomes a stage of its own. That pays off
only when the probe is partitioned as well:
- A **single-partition build** needs no stage — every task can read it in
place.
- A **single-partition probe** pins the join to one task whatever the build
does, so the stage buys nothing.
And the stage is not free. The build/probe choice comes from a size
comparison, and AQE replans after every stage completion, so a later round can
see statistics the first round did not have and answer that comparison the
other way — swapping the inputs. The exchange created by the first answer is
already a stage by then, and it is now on the **probe** side, where a broadcast
read exposes a single partition. Everything above it collapses into one task.
TPC-H q7 hit exactly this. Its stage 5 read 1.46 M rows through one task,
running the `nation` join, the projection and the partial aggregate
single-threaded while 15 partitions of input sat ready:
```
-- before: stage 5, 1 task, task_ms 39
AggregateExec: mode=Partial, gby=[supp_nation, cust_nation, l_year]
HashJoinExec: mode=CollectLeft, on=[(n_nationkey@0, c_nationkey@3)],
filter=...
FilterExec: n_name = GERMANY OR n_name = FRANCE
DataSourceExec: file_groups={1}
ShuffleReaderExec: upstream_stage: 4, broadcast: true,
upstream_partition_count: 16
```
I confirmed the sequence by instrumenting the rule: the broadcast exchange
is created once, from the branch that wraps a multi-partition build side, and
the same join's swap decision is logged twice with opposite answers in the same
query — `swap=Ok(false)` then `swap=Ok(true)` — as the estimates change between
rounds.
## What changes are included in this PR?
One guard, extracted as `broadcast_build_side_pays_off`: wrap the build side
only when both sides have more than one partition. With it, q7's swap puts
`nation` on the build side, the probe keeps its 16 partitions, and no exchange
is planted:
```
-- after: stage 5, 4 tasks, task_ms 10/13/14
ShuffleReaderExec: upstream_stage: 4, partitioning: Hash([l_orderkey@3],
16)
```
## Are these changes tested?
A unit test covers the three cases of the predicate (both sides partitioned,
single-partition probe, single-partition build). Full `ballista-scheduler`
suite passes (363 + 25); clippy clean.
End to end on TPC-H SF10, one scheduler and two executors x 4 vcores, AQE
with `collect_statistics=true` and packing on:
| | before | after |
| --- | --- | --- |
| q7 stage 5 | 1 task, `broadcast: true`, 1.46 M rows | 4 tasks,
`Hash([l_orderkey@3], 16)` |
| q7 job time | 1777 ms | 1249 ms (-30%) |
| q7 broadcast readers | 1 | 0 |
All 22 queries return identical row counts, and the other 21 plans are
unchanged — same stage counts, same broadcast-reader counts — so the guard is
confined to the case it describes.
**Draft** for two reasons. The machine was running another workload during
these measurements, so treat the -30% as indicative rather than precise; a
clean run (ideally at SF1000, where the collapsed stage would carry ~150 M rows
through one core) is worth having before merge. And q8 keeps three broadcast
readers with this guard: the same collapse shape appears there via chained
dimension joins, so there is likely a second case to handle.
## Are there any user-facing changes?
No API or configuration change. One fewer stage boundary and more
parallelism in plans that join a partitioned input against a single-partition
side.
--
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]
