andygrove opened a new pull request, #2782:
URL: https://github.com/apache/datafusion-ballista/pull/2782
# Which issue does this PR close?
Closes #2756.
# Rationale for this change
With `ballista.planner.coalesce.enabled` on, `CoalescePartitionsRule` skips
the rewrite whenever the bin-pack produces one group. A stage whose whole input
fits in `ballista.planner.coalesce.target_partition_bytes` then still runs one
partition per upstream partition, most of them reading next to nothing. Those
are exactly the stages coalescing exists for: final aggregates, small joins and
the sorts below an `ORDER BY` merge. Spark's `CoalesceShufflePartitions` goes
all the way down to one partition.
The issue asked four questions. Here is what I found.
**Why was `K <= 1` excluded?** It was a conservative default rather than a
guard for an operator requirement. #1684 calls the one-group case degenerate
and keeps such a stage at full parallelism. None of the operators the rule
allows (`tolerates_partition_count_change`) need more than one partition.
The guard was hiding two bugs, though. Both already fire on main when a
shuffle packs into more than one but fewer than all of its partitions (`1 < K <
M`):
- **Sorted shuffles.** The adapter reads an exchange whose input is sorted
with `RangeShuffleReaderExec`, one k-way merge per partition. Once coalesced it
builds a plain `ShuffleReaderExec` instead, which concatenates a group's files.
The `SortPreservingMergeExec` above then merges unsorted streams, and `ORDER
BY` results come back out of order. With one group allowed and nothing else
changed, every small `ORDER BY` over more than one partition would come back
unsorted.
- **Passthrough stage widths.** A stage below a `CoalescePartitionsExec` or
`SortPreservingMergeExec` writes one partition per partition it reads. Since
#2714 its output is resolved to the width its writer produced, so after
coalescing it declares `M` partitions but resolves `K` location lists. If the
stage reading it doesn't coalesce too, its reader still exposes `M` partitions
and indexes past the end. On main with coalescing on, TPC-H q17 at SF100 fails
this way (`index out of bounds: the len is 8 but the index is 10`).
**Can `K = 1` be expressed with the existing `CoalescePlan`?** Yes. One
group spanning all `M` upstream partitions gives a one-partition reader. The
stage needs no new shape: `CoalescePartitionsExec` and
`SortPreservingMergeExec` pass a single partition straight through, and
Ballista turns off round-robin repartitioning, so there is no `RepartitionExec`
to drop.
**How much does it matter?** I ran TPC-H and TPC-DS with coalescing on, on
main and on this PR, with the benchmark runners' `--verify` checking every
answer against single-process DataFusion. The setup was one machine with 32
cores, a scheduler and two executors with 8 vcores and a 16 GiB pool each, 16
target partitions, and Parquet data. Each stage's tasks and partitions come
from the scheduler's REST API.
| | TPC-H SF10 | TPC-H SF100 | TPC-DS SF10 |
|---|---:|---:|---:|
| Queries | 22 | 21 | 98 |
| Stages | 124 | 130 | 1,102 |
| Stages coalesced, main / this PR | 11 / 51 | 9 / 27 | 26 / 441 |
| ... to one partition, main / this PR | 0 / 40 | 0 / 18 | 0 / 417 |
| Partitions run, main / this PR | 1,535 / 1,175 (-23%) | 1,682 / 1,457
(-13%) | 16,072 / 12,949 (-19%) |
| Tasks launched, main / this PR | 217 / 192 (-12%) | 233 / 216 (-7%) |
2,164 / 2,040 (-6%) |
TPC-H SF100 leaves out q17, which fails on main as described above. TPC-DS
leaves out q23, which hangs on main with coalescing on or off. Every other
query returned the right answer on both builds, and so did q17 at SF100 with
this PR.
In the stages this PR reads as one partition, partitions fall from 569 to
220 at TPC-H SF10, from 243 to 18 at SF100 and from 6,331 to 3,265 at TPC-DS
SF10.
Ballista already packs a stage's partitions into tasks by free vcores
(`ballista.scheduler.max_partitions_per_task`), so the task count falls less
than the partition count. The partitions are what a single group removes: each
one holds a vcore, runs the stage's operators and writes its own shuffle
output, which the next stage then fetches. When a stage reads a shuffle through
a gather, such as a `CollectLeft` build side or the final merge, its partition
count comes from elsewhere, and coalescing only cuts the number of streams the
gather reads.
Wall time is flat. Over three alternating runs per build at TPC-H SF100, the
per-query medians add up to 177.6 s on main and 176.1 s here. Most queries are
within 3%. q11 is the exception, 0.53 s to 0.71 s: its 3.1M-row join stage now
runs as one task taking 219 ms, where 16 partitions took 36 ms. That is the
parallelism a 64 MiB target trades away, the same trade every coalesced
partition already makes. Spark keeps more of it by default with
`parallelismFirst`, which caps the target at the shuffle's size divided by the
cluster's parallelism. That could be a follow-up if latency matters more than
the saved partitions. The SF10 runs showed the same totals within noise.
**Are there alternatives?** Packing partitions into tasks, described above,
is the closest, and it doesn't remove the partitions. Running small stages on
the scheduler would put the data path (shuffle reads, memory, object stores) in
the scheduler, where one slow stage could hold up planning for every job.
Fusing a small stage into its consumer removes a shuffle, but that has to be
decided before the upstream stage runs, from estimates, which is the
static-planning problem AQE is there to avoid. Broadcast joins already cover
the common small-side case. Coalescing to one partition is a small change to a
rule that already exists.
# What changes are included in this PR?
Two commits, so the fixes can be reviewed, or split out, on their own:
1. **`fix(aqe): keep sorted shuffles and passthrough widths intact when
coalescing`**
- `CoalescePartitionsRule` leaves a stage alone when any leaf's input
declares an ordering, like the broadcast and range-repartition cases. Such a
stage is usually the final merge of an `ORDER BY`, which runs as one task
anyway. The same check keeps the prefix-window rewrite's final stage
uncoalesced: `PrefixMergeExec` keys its state by partition, and the shuffle it
reads is sorted.
- `AdaptivePlanner::take_stage_output_partitions` resolves every
partition a stage's boundary declares, and more if the writer produced more (a
skew-split final stage), padding with empty location lists.
2. **`feat(aqe): coalesce shuffle partitions down to a single partition`**
- The rule keeps `K = 1` and skips only `K >= M`.
- The tuning guide says coalescing can go down to a single task.
Tests:
- `state::aqe::test::coalesce_rule` gains a one-partition aggregate, both
legs of a join coalesced to one partition together, a sorted leaf at `K = 1`
and `K = 2`, and a coalesced passthrough stage that still resolves all 8 of its
declared partitions. All four fail on main.
- `ballista/client/tests/coalesce.rs` runs four query shapes (a sorted
`GROUP BY`, a sorted join, a global aggregate over a join and a sorted scan) on
a standalone cluster with coalescing on, at five targets and with broadcast
joins on and off, and compares each result with single-process DataFusion, row
by row. On main 16 of the 40 cases fail: the global aggregate panics and the
sorted results come back out of order. All pass here.
- I also ran `ballista/client/tests/prefix_window.rs` with coalescing on at
targets from 128 bytes to 64 MiB. On main it hangs at 2048 bytes and returns
wrong running sums at 1024. Here it passes at every target.
# Are there any user-facing changes?
Only with `ballista.planner.coalesce.enabled=true`, which is off by default.
A shuffle that fits in the coalesce target is now read as a single partition,
so its stage runs as one task. The out-of-order `ORDER BY` results and the
passthrough panic described above are fixed. There are no API or configuration
changes.
--
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]
