ulysses-you commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3718360865
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -663,46 +855,124 @@ case class HashAggregateExec(
case _ => ("true", "", "")
}
- val findOrInsertRegularHashMap: String =
- s"""
- |// generate grouping key
- |${unsafeRowKeyCode.code}
- |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode();
- |if ($checkFallbackForBytesToBytesMap) {
- | // try to get the buffer from hash map
- | $unsafeRowBuffer =
- | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys,
$unsafeRowKeyHash);
- |}
- |// Can't allocate buffer from the hash map. Spill the map and
fallback to sort-based
- |// aggregation after processing all input rows.
- |if ($unsafeRowBuffer == null) {
- | if ($sorterTerm == null) {
- | $sorterTerm = $hashMapTerm.destructAndCreateExternalSorter();
- | } else {
- |
$sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter());
- | }
- | $resetCounter
- | // the hash map had be spilled, it should have enough memory now,
- | // try to allocate buffer again.
- | $unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow(
- | $unsafeRowKeys, $unsafeRowKeyHash);
- | if ($unsafeRowBuffer == null) {
- | // failed to allocate the first page
- | throw QueryExecutionErrors.aggregateOutOfMemoryError();
- | }
- |}
- """.stripMargin
+ val findOrInsertRegularHashMap: String = {
+ // Assumes the grouping key projection (`unsafeRowKeyCode.code`) has
already run for this row,
+ // so `unsafeRowKeyCode.value` holds the current key. The projection is
emitted exactly once
+ // per regular-map row (see below); emitting it in more than one runtime
branch is unsafe
+ // because the projection's subexpression/writer state assigned in one
branch would be read
+ // stale from another (e.g. the adaptive pass-through path would reuse
the last probed key).
+ val probeRegularMap =
+ s"""
+ |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode();
+ |if ($checkFallbackForBytesToBytesMap) {
+ | // try to get the buffer from hash map
+ | $unsafeRowBuffer =
+ | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys,
$unsafeRowKeyHash);
+ |}
+ """.stripMargin
+
+ val spillMap =
+ s"""
+ |if ($sorterTerm == null) {
+ | $sorterTerm = $hashMapTerm.destructAndCreateExternalSorter();
+ |} else {
+ |
$sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter());
+ |}
+ |$resetCounter
+ |// the hash map had been spilled, so it should have enough memory
now,
+ |// try to allocate buffer again.
+ |$unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow(
+ | $unsafeRowKeys, $unsafeRowKeyHash);
+ |if ($unsafeRowBuffer == null) {
+ | // failed to allocate the first page
+ | throw QueryExecutionErrors.aggregateOutOfMemoryError();
+ |}
+ """.stripMargin
+
+ if (adaptivePartialAggConfig.isDefined) {
+ val cfg = adaptivePartialAggConfig.get
+ // Adaptive partial aggregation governs only this regular
(second-level) map. Count the
+ // rows that enter it (a fast-map miss, or every row when the fast map
is off) and use
+ // `regularMap.getNumKeys() / regularRows` as the pre-shuffle
reduction ratio.
+ // - Tier 2 (on-spill): when the map cannot allocate for a new key
(it would otherwise
+ // spill), bypass instead if the ratio is at least
`spillReductionRatioThreshold`.
+ // - Tier 1 (no-spill): from `sampleRows` regular rows on, bypass if
the ratio is at
+ // least `noSpillReductionRatioThreshold`. The sampling window
doubles after each
+ // sub-threshold check, so low-cardinality input is re-evaluated
only rarely while a
+ // late high-cardinality tail can still trigger the bypass.
+ // Both tiers fire only before any spill (`sorter == null`): once the
map has spilled, the
+ // reduction-ratio estimate no longer covers the spilled rows, and
pass-through must never
+ // coexist with sort-based aggregation. When the map is full after a
spill, the map spills
+ // again as usual.
+ // The key projection runs once here so `unsafeRowKeyCode.value` is
valid for both the
+ // probe below and the pass-through buffer built by the caller.
+ s"""
+ |// generate grouping key
+ |${unsafeRowKeyCode.code}
+ |if (!$adaptivePassThroughTerm) {
+ | $probeRegularMap
+ | if ($unsafeRowBuffer == null) {
+ | if ($sorterTerm == null && $regularMapRowCountTerm > 0 &&
+ | (double) $hashMapTerm.getNumKeys() >=
+ | $regularMapRowCountTerm *
${cfg.spillReductionRatioThreshold}D) {
+ | $adaptivePassThroughTerm = true;
+ | } else {
+ | $spillMap
+ | }
+ | }
+ | if ($unsafeRowBuffer != null) {
+ | $regularMapRowCountTerm += 1;
+ | if ($sorterTerm == null &&
Review Comment:
Addressed. The permanent `sorter == null` gate is gone. The processed-row
count and the next check point are now reset after each spill, so the new
in-memory map epoch is judged on its own rows and pass-through stays available
after earlier spills. When it activates, the existing sorter and the current
map drain before the remaining rows pass through.
Added a test that asserts both happen in the same task --
`numTasksFallBacked = 1` and `numBypassingRows > 0` together, which was
impossible under the old gate.
--
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]
