peter-toth commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3772649542
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -845,29 +1229,86 @@ case class HashAggregateExec(
}
}
- val declareRowBuffer: String = if (isFastHashMapEnabled) {
- val fastRowType = if (isVectorizedHashMapEnabled) {
- classOf[MutableColumnarRow].getName
+ val declareRowBuffer: String = {
+ val declareBuffers = if (isFastHashMapEnabled) {
+ val fastRowType = if (isVectorizedHashMapEnabled) {
+ classOf[MutableColumnarRow].getName
+ } else {
+ "UnsafeRow"
+ }
+ s"""
+ |UnsafeRow $unsafeRowBuffer = null;
+ |$fastRowType $fastRowBuffer = null;
+ """.stripMargin
+ } else {
+ s"UnsafeRow $unsafeRowBuffer = null;"
+ }
+ val declareBypassed = if (adaptivePartialAggEnabled) {
+ s"boolean $adaptiveRowBypassedTerm = false;"
} else {
- "UnsafeRow"
+ ""
}
s"""
- |UnsafeRow $unsafeRowBuffer = null;
- |$fastRowType $fastRowBuffer = null;
+ |$declareBuffers
+ |$declareBypassed
""".stripMargin
- } else {
- s"UnsafeRow $unsafeRowBuffer = null;"
}
// We try to do hash map based in-memory aggregation first. If there is
not enough memory (the
// hash map will return null for new key), we spill the hash map to disk
to free memory, then
// continue to do in-memory aggregation and spilling until all the rows
had been processed.
// Finally, sort the spilled aggregate buffers by key, and merge them
together for same key.
+ //
+ // With adaptive partial aggregation, once pass-through is active
`updateRowInHashMap` fills the
+ // single-row buffer built above; we then emit `key ++ buffer` straight to
the parent so the row
+ // skips both the fast map and the regular map.
+ //
+ // The maps were frozen when pass-through fired, so they are drained to
preserve the merge order
+ // (trigger row, then the maps, then the rest of the input). Where the
drain runs depends on
+ // the plan shape. Split by an exchange, the partial is the whole-stage
root and its output
+ // function appends a reference to the reusable output row; draining in
the same call would
+ // overwrite that row while still buffered and silently drop the trigger,
so the drain is left
+ // to `doProduceWithKeys` on the next `processNext` (the buffered trigger
is pulled first).
+ // A one-to-many child that cannot yield mid-fan-out changes the trade:
without a stop check
+ // its whole batch is appended before the deferred drain runs, flipping
the merge order for a
+ // group that straddles the freeze point. Such children report
`needCopyResult`, so every row
+ // is a copy, the aliasing hazard is gone, and the drain runs here right
after the trigger,
+ // draining the whole frozen map before the fan-out batch continues. Fused
with the Final (no
+ // exchange), the output is consumed directly by the Final's `doConsume`
and never buffered, so
+ // there is no aliasing hazard either; the drain runs here too, because
nothing else would
+ // yield the build loop (with no append there is no `shouldStop()`), and
the maps would
+ // otherwise come out only after all the remaining streamed rows, flipping
the merge order.
+ val emitPassThroughRow = if (adaptivePartialAggEnabled) {
+ val numBypassingRows = metricTerm(ctx, "numBypassingRows")
+ val drainFused = if (isWholeStageRoot && !needCopyResult) {
+ ""
+ } else {
+ // `outputMap` returns on `shouldStop()`, already true here because
the bypassed row was
+ // just appended, so drain in a loop to emit the whole frozen map
before the fan-out batch
+ // continues. Fused, nothing is buffered and `shouldStop()` stays
false, so the loop runs
+ // once.
+ s"""
+ |while (!$adaptiveMapOutputDoneTerm) {
Review Comment:
Fair, and your memory concern is right -- I measured it rather than leaving
it as prose. Instrumenting `BufferedRowIterator.append` to track the peak
`currentRows` size, `Generate` fan-out 2 over `range(0, 400000, 1, 2)` with an
`Exchange` above the partial:
| | peak rows in `currentRows` |
| --- | --- |
| feature off | 1 |
| on, `minRows = 100000` (default) | 100,002 |
| on, `minRows = 200000` | 200,002 |
So the buffer is exactly the frozen map, and with `minRows = 0` it is sized
by whatever the map held when it hit the memory limit rather than by a knob.
That is worse than what I wrote above.
**Reverting to `if` is a legitimate answer to that and I am not against
it.** What I would ask is that we write down what it gives up, because a few
sentences currently say the opposite:
- the PR description, the suite class doc and the tuning doc all state that
the generated and interpreted paths always emit the same order. That stops
being true in the split fan-out shape (an `Exchange` above the partial, a
`Generate` / `Expand` / join below it), so `spark.sql.codegen.wholeStage`
becomes observable there. That is the part I would most want said plainly
rather than quietly dropped.
- the two tests added in `77a8aad0997` go with it.
- and the docs should say that enabling the feature can reorder a group's
buffers within a single task, not only across a shuffle.
On that last point, two corrections that matter for how the sentence gets
worded:
1. **The shuffle does not make the partial's order moot.** It randomizes
order *across* map tasks, but the rows one task writes for one key land in one
block in write order and the `Final` merges them in read order. So for any key
whose rows sit in a single input partition -- pre-partitioned or bucketed
input, small data, one split per key -- the partial's emit order decides the
answer. My `dupKey = 6` case is exactly that: feature off gives `(306, 406)` on
every run and the pre-`77a8aad0997` code gave `(406, 306)` on every run, at
`splits = 2` with the `Exchange` in the plan. Reproducible, not random.
2. **The fused case is a tie, not a reason.** Nothing is appended there, so
`shouldStop()` stays false and a single `outputMap()` call drains the whole map
with either `if` or `while`. The loop only ever mattered in the split shape.
3. Agreed on `first`/`last` -- `First`'s own scaladoc says partial+final
makes it non-deterministic unless the input is sorted in a single partition. I
would only note that "this function is order-dependent" and "the same query
returns a different answer when you flip `spark.sql.codegen.wholeStage`" are
different things from a user's side.
One more thing to fold into whatever we write: the current order already
does not match a feature-off run when the **first** bypassed row collides with
a key in the frozen map -- on both paths, in every shape, including fused.
Measured on `b039db81c79`:
| repro | reference | generated | interpreted |
| --- | --- | --- | --- |
| 1:1 child, first bypassed row collides | (0,8) | (8,0) | (8,0) |
| `explode` fan-out 2, first bypassed row collides | (102,402) | (402,102) |
(402,102) |
You already document this for `dupAt = 8` in the suite comment, so it is not
news -- but it means the honest sentence is "the merge order of a bypassed
group is not preserved", not "only under fan-out".
---
**If you would rather keep the order and lose the buffer, there is a third
option.** Hold the bypassed row *behind* the map instead of racing the map
ahead of it: queue the row, and let queuing it advance the map output by one
row. That one appended map row is what makes `shouldStop()` true so the child
still yields at the end of its batch, and it is also what limits the queue,
since a child that checks `shouldStop()` between rows never gets to add a
second entry. The map then drains one row per `processNext` (what the 1:1 shape
already does) and the queue is flushed as soon as the map is done. Fused,
nothing is buffered, so the single call empties the map and the held row
follows immediately.
```diff
@@ doProduceWithKeys: mutable state
adaptiveFastKeysAtSpillTerm =
ctx.addMutableState(CodeGenerator.JAVA_INT,
"adaptiveFastKeysAtSpill")
+ adaptivePendingRowsTerm = ctx.addMutableState(
+ "java.util.LinkedList<UnsafeRow[]>", "adaptivePendingRows",
+ v => s"$v = new java.util.LinkedList<UnsafeRow[]>();", forceInline
= true)
}
@@ doProduceWithKeys: after adaptiveOutputMapFuncName
+ adaptiveFlushPendingFuncName = if (adaptivePartialAggEnabled) {
+ val name = ctx.freshName("flushPendingRows")
+ val pair = ctx.freshName("pendingRow")
+ ctx.addNewFunction(name,
+ s"""
+ |private void $name() throws java.io.IOException {
+ | while (!$adaptivePendingRowsTerm.isEmpty()) {
+ | UnsafeRow[] $pair = (UnsafeRow[])
$adaptivePendingRowsTerm.poll();
+ | $outputFunc($pair[0], $pair[1]);
+ | }
+ |}
+ """.stripMargin)
+ } else {
+ ""
+ }
@@ adaptiveOutputMap
|if (!$adaptiveMapOutputDoneTerm) {
| $adaptiveOutputMapFuncName();
+ | if ($adaptiveMapOutputDoneTerm) {
$adaptiveFlushPendingFuncName(); }
| if (shouldStop()) return;
|}
@@ emitPassThroughRow
s"""
|if ($adaptiveRowBypassedTerm) {
| $numBypassingRows.add(1);
- | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer);
- | $drainFused
+ | if ($adaptiveMapOutputDoneTerm) {
+ | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer);
+ | } else {
+ | $adaptivePendingRowsTerm.add(new UnsafeRow[] {
+ | ${unsafeRowKeyCode.value}.copy(), $unsafeRowBuffer.copy()
});
+ | $adaptiveOutputMapFuncName();
+ | if ($adaptiveMapOutputDoneTerm) {
$adaptiveFlushPendingFuncName(); }
+ | }
|}
""".stripMargin
```
On the interpreted side it is a deletion: drop the
`passThroughTriggerPending` branch from `next()` and the remaining branch order
is already map-then-pass-through, so the flag and its two assignments go too.
What it buys, measured (arms = reference / generated / interpreted, over
`splits` in {1, 2} x `twoLevelMap` in {on, off}):
- all 16 order cells match the feature-off reference, including the two rows
above that are wrong today;
- peak `currentRows` drops from 100,002 / 200,002 to **2**, and stops
depending on `minRows`;
- `numBypassingRows` unchanged (600,000 / 400,000), so the decision path is
untouched;
- your full suite: **51 succeeded, 0 failed**, all three order tests
included;
- `drainFused` and the `isWholeStageRoot` / `needCopyResult` split disappear
-- the same code serves fused, 1:1-split and fan-out-split.
What it costs, honestly:
- the bound is the batch width, not a constant. When one input row's fan-out
is wider than the frozen map it is a **tie** with today -- measured peaks of 50
and 200 for widths 50 and 200 on both builds -- and the queue holds up to
`min(batch, map)` rows on top of the buffer;
- the map comes out one row per `processNext` in the fan-out shape instead
of one burst. Same pattern the 1:1 shape already uses, but I have not measured
the cost;
- it is still a plain `LinkedList` the memory manager cannot see or spill --
smaller, not gone;
- it does not touch the one-way flip or the freeze-versus-probe trade;
- results change relative to this head (they become reference-matching), so
the `dupAt = 8` comment needs rewriting and that case wants a test against the
reference;
- one new invariant with nothing enforcing it: held rows are flushed only
when `outputMap()` sets the done flag, so a future path that ends iteration
without reaching `adaptiveFinalOutput` would drop them.
Two implementation notes if you take it: Janino does not infer the generic,
so `poll()` needs the explicit `(UnsafeRow[])` cast, and the copies are needed
because both the key row and the pass-through buffer are reused.
Either way works for me.
--
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]
