ulysses-you commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3774039458
##########
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:
@peter-toth thanks for the measurements, I think the third option is better
rather than the revert. The queue-behind-the-map design is what this head
implements (e9313e65eb3), including your two implementation notes: `poll()`
carries the explicit `(UnsafeRow[])` cast (Janino does not infer the generic),
and both the key and the pass-through buffer are copied at queue time since
they are reused. The drain and the queue flush are one generated function
(`outputMapAndFlush`), and the emit site is now `handlePassThroughRow`.
On why this option rather than the revert you said was legitimate: the two
alternatives each give up one of the two properties this design keeps.
Reverting to `if` keeps the trigger-first order but sacrifices path-agreement —
in the split fan-out shape the trigger batch overtakes the frozen map,
`first`/`last` diverge between the generated and interpreted paths, and
`spark.sql.codegen.wholeStage` becomes observable, although it is not a block,
it is still good to solve it without more cost. The queue holds both properties
at once: the buffer is bounded by the fan-out batch width rather than
`minRows`, so the peak does not grow with the input, and the merge order is
preserved by construction in every shape, so no claim on either axis is given
up. It also collapses the `drainFused` / `isWholeStageRoot` / `needCopyResult`
special-casing the drain-loop needed into one generated function that serves
the fused, 1:1-split and fan-out-split shapes.
--
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]
