jayzhan211 opened a new pull request, #24795:
URL: https://github.com/apache/datafusion/pull/24795
## Which issue does this PR close?
- Part of #24704.
- Part of #19906.
## Rationale for this change
The blocked / chunked state management work in #24704 needs a way to tell
whether it worked. The suites we normally judge aggregation PRs on —
ClickBench
and h2o groupby — measure whole-query time, and the phase this work changes
is a
small fraction of that, so they can neither show the win nor size it. They
stay
as the regression gate; this benchmark is what shows the improvement.
It measures only the drain phase: what `AggregateExec` does after its input
is
exhausted and it starts producing output. Today both implementations drain by
calling `emit(EmitTo::All)` — materializing every group into one giant
`RecordBatch` — and then handing that batch downstream as `batch.slice(..)`
chunks (`aggregate_hash_table/common.rs:265`, labelled "temporary solution
until
blocked state management is implemented", and `grouped_hash_stream.rs:1337`).
Two consequences, and the benchmark quantifies both:
1. **The whole drain is one poll.** Time-to-first-batch equals total drain
time,
so the runtime is blocked for the entire span (#19906).
2. **Nothing is released until the drain ends.** Every slice shares the giant
batch's buffers, so memory held at the halfway point is 100% of what was
held
at the start (#24704, symptom 1).
Baseline on `main` (61bf6b96c), 10M groups, five aggregates:
```
groups key agg build_ms drain_ms ttfb_ms max_gap_ms
peak_pool peak_live pool_% live_%
10000000 int64 wide 284.0 16.6 16.2 16.2
1158.2MB 1168.2MB 100.0% 100.0%
10000000 utf8 wide 512.6 16.1 15.8 15.8
1574.3MB 1616.2MB 100.0% 100.0%
10000000 dict wide 1035.4 33.2 32.8 32.8
1702.4MB 1798.4MB 100.0% 100.0%
10000000 liststr wide 1091.4 1687.1 1686.7 1686.7
2966.3MB 5226.8MB 100.0% 100.0%
```
Two things to read off it.
**`max_gap_ms ≈ ttfb_ms ≈ drain_ms` in every row.** The entire drain happens
inside one `poll_next`, exactly as #19906 describes.
**Key layout decides the size of that poll, and the spread is 100x.** This
was
the surprise. For `Int64` and `Utf8` keys, `emit(EmitTo::All)` is close to a
buffer move — `take_needed` is a `mem::take` — so 10M groups drain in ~16 ms
and
there is very little to win. Keys that go through arrow's row format have to
decode every group on the way out: `GroupValuesRows::emit` calls
`convert_rows`
over the whole table (`group_values/row.rs:215`), and nested keys use a
row-backed `GroupColumn` inside the vectorized path. A `List(Utf8)` key at
10M
groups blocks the runtime for **1687 ms in a single poll** — the ">1s stall
at
~10M groups" reported in #19906, reproduced.
Worth stating explicitly for anyone measuring this work: on flat keys the
latency win is ~16 ms, so a benchmark run on `Int64` keys alone will show
nothing. The four default shapes are chosen to span that range.
`pool_%` and `live_%` are 100% everywhere: nothing is released until the
drain
ends, on either code path, in either accounting. Shapes not in the default
set
behave the same way — `Struct(Int64, Utf8)` drains in 174 ms at 10M groups
and
`List(Int64)` in 664 ms, both at 100% — and a ten-column flat key is cheap
(~2 ms at 1M groups), so column count is not what drives the stall.
## What changes are included in this PR?
One new benchmark, `datafusion/physical-plan/benches/aggregate_drain.rs`, and
its `[[bench]]` entry. No changes to any non-test code.
| metric | meaning |
|---|---|
| `drain_ms` | input exhausted → last output batch |
| `ttfb_ms` | input exhausted → first output batch |
| `max_gap_ms` | longest interval between consecutive output batches — the
long-poll proxy |
| `peak_*` | peak memory over the run |
| `pool_%` / `live_%` | memory at the 50%-drained mark over memory after the
first batch: ~100% means nothing is released until the drain ends, ~50% means
memory is released as output is produced |
Three design decisions worth flagging for review:
**Criterion is not used.** The quantities of interest are within-run timings
and
memory samples, not a throughput distribution. The bench is `harness = false`
with a plain `main` that prints a table.
**Memory is reported two ways, because they disagree.** `pool` is what the
`MemoryPool` has reserved; `live` is bytes actually live on the heap, from a
counting global allocator. The pool does not track the materialized output
batch, so on the legacy path with `utf8` keys it reports **0.1 MB** reserved
after the first output batch while the process is holding **960 MB** —
symptom 3
in #24704. A pool-only measurement would be blind exactly where the problem
is.
**The ratio is anchored to the first output batch, not to the peak.** Peak is
reached while the hash table is still being built, and the current code
takes a
one-time step down from build state to materialized output. `at50 / peak`
therefore reads ~50% today and looks like incremental release already works;
`at50 / first` correctly reads 100%.
Input is one row per group, so the build phase is as short as possible and
the
drain is what is being measured. A wrapper `ExecutionPlan` records when the
input is exhausted, which is what separates build from drain; a sampler
thread
reads memory every 250 µs so peaks reached *inside* a single long poll are
not
missed.
```sh
cargo bench -p datafusion-physical-plan --features test_utils --bench
aggregate_drain
# just the shape that stalls
cargo bench -p datafusion-physical-plan --features test_utils --bench
aggregate_drain -- \
--keys liststr --groups 10000000 --aggs wide
# the pre-migration path
cargo bench -p datafusion-physical-plan --features test_utils --bench
aggregate_drain -- --legacy
```
Defaults are `int64, utf8, dict, liststr` keys × `sum, wide` aggregates ×
`10k, 10M` groups — a cheap floor, the common case, the `GroupValuesRows`
fallback, and the shape that stalls. Note that `liststr` at 10M groups holds
several GB.
Grouped aggregation is mid-migration (#22710), and both implementations drain
the same way, so both are covered: with
`execution.enable_migration_aggregate`
on (the default) a single grouping set runs on `SingleHashAggregateStream`,
and
`--legacy` turns the flag off to measure `GroupedHashAggregateStream`.
### How to read the output
One line per shape. The drain emits `groups / batch_size` batches — 1221 of
them
at 10M groups with the default 8192 — so a healthy drain spreads its work
across
1221 polls, and today's does not:
- **`max_gap_ms` is the headline.** It is the longest single stretch the
tokio
worker was blocked. Compare it against the mean gap, `drain_ms / batches`.
For
`liststr` the mean gap is 1.4 ms but `max_gap_ms` is 1687 ms, i.e. one poll
does 100% of the work. That ratio *is* the long poll.
- **`ttfb_ms ≈ drain_ms` is the same fact from the consumer's side.** The
first
row downstream costs as much as all of them.
- **`pool_%` / `live_%` at 100%** mean the operator is still holding
everything
it held at the start of the drain when it is half-finished.
- **`peak_live` vs `peak_pool`** shows whether output is materialized on top
of
state rather than moved out of it. For `liststr`, `peak_live` (5227 MB) far
exceeds `peak_pool` (2966 MB) because `convert_rows` builds the whole
output
while the row buffer is still alive; for flat keys the two nearly match
because emit is a move.
- **`build_ms`** is not part of the gate, but watch it: blocked storage puts
a
`(block, offset)` indirection on every group lookup, and that cost lands
here.
ClickBench and h2o are the real guard for it.
### What the follow-up work should show
Targets for #24704 at 10M groups, derived from the baseline above. Hard gates
are marked; the rest are informational but should move in the stated
direction.
| metric | today (`liststr`) | expected after | why |
|---|---|---|---|
| `max_gap_ms` | 1687 | **< 20** (gate) | one block of work per poll, not
the whole table. Mean gap is 1.4 ms; anything under ~10x that is healthy |
| `ttfb_ms` | 1687 | **< 20** (gate) | first batch costs one block |
| `live_%` | 100% | **≤ 60%** (gate) | half the groups emitted ⇒ about half
the memory released. 50% is ideal; the hash table itself may not shrink in step
|
| `pool_%` | 100% | **≤ 60%** (gate) | must move *with* `live_%` — see
failure modes below |
| `drain_ms` | 1687 | 1400–2000 (gate: ≤ 110% of baseline) | same total
work, spread out. This is where #19562 died |
| `peak_live` | 5227 MB | ~3000–3500 MB | the full-output spike disappears;
only state plus one block is live. Flat keys will barely move, since they have
no spike to remove |
And on the flat shapes, which have almost nothing to win and everything to
lose:
| shape | metric | today | requirement |
|---|---|---|---|
| `utf8` @ 10M | `drain_ms` | 16.1 | ≤ 18 (gate) |
| `int64` @ 10M | `drain_ms` | 16.6 | ≤ 18 (gate) |
| any @ 10k | `drain_ms` | ~0 | no measurable regression (gate) |
The strongest single check is not a threshold at all: **`max_gap_ms` should
stop
depending on group count.** Run `--groups 10000000,20000000` — today
`max_gap_ms`
roughly doubles, because the poll materializes everything. After blocked
emission it should be flat, because a poll materializes one block regardless
of
how many groups exist.
Failure modes this benchmark is designed to catch:
- `drain_ms` up 2x or more — `EmitTo::First(n)` shifting the remaining
elements
on contiguous storage, the O(remaining) trap from #19562.
- `live_%` still 100% — blocks are being emitted but not dropped, or the
output
batches still share one allocation.
- `live_%` drops but `pool_%` stays at 100% (or the reverse) — the
reservation
no longer describes reality. Both must move together, or downstream spill
decisions get worse rather than better.
- `max_gap_ms` down but `drain_ms` up — work was spread out by making more of
it. A win on latency paid for with throughput.
## Are these changes tested?
The benchmark is the test artifact; it adds no product code. `cargo test
--benches` runs the binary with `--test`, which is handled as a fast smoke
run
at 1k groups that asserts each configuration emits exactly one row per group.
`cargo fmt --all` and `cargo clippy --all-targets --all-features -- -D
warnings`
are clean.
## Are there any user-facing changes?
No. New benchmark only; no public API or behavior 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]
