cloud-fan commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3713151770
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -154,6 +200,31 @@ case class HashAggregateExec(
private var hashMapTerm: String = _
private var sorterTerm: String = _
+ // Codegen state for adaptive partial aggregation. When the pre-shuffle
reduction ratio of the
+ // regular (second-level) hash map is too low, the operator stops populating
the map and instead
+ // streams each remaining row through as a single-row partial buffer for the
Final aggregate to
+ // merge. Only the regular map is governed: the append-only fast hash map
keeps absorbing hot
Review Comment:
Please measure all processed rows against the total number of in-memory
keys, including both maps. DBR uses the operator-level invariant—equivalent to
`fastMap.rowCount + regularMap.getNumKeys` over all processed rows—so limiting
both sides to regular-map traffic makes two-level-map routing change the
decision and creates a second policy to maintain.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -626,13 +792,27 @@ case class HashAggregateExec(
| long $beforeAgg = System.nanoTime();
| $doAggFuncName(partitionIndex);
| $aggTime.add((System.nanoTime() - $beforeAgg) / $NANOS_PER_MILLIS);
+ | $adaptiveStopCheck
|}
- |// output the result
- |$outputFromFastHashMap
- |$outputFromRegularHashMap
+ |$adaptiveResumeBuild
+ |$adaptiveFinalOutput
""".stripMargin
}
+ // Blocking operators normally suppress the child's `shouldStop()` check
because they buffer all
+ // output. With adaptive partial aggregation, pass-through rows are appended
to the output buffer
+ // while consuming child input, so the stop check must be re-enabled to keep
the buffer bounded.
+ override def needStopCheck: Boolean =
+ adaptivePartialAggConfig.isDefined || super.needStopCheck
+
+ // Blocking operators normally do not copy their result because every output
row is drained (via
+ // `shouldStop()`) before the next one is produced. Adaptive pass-through
breaks that assumption:
+ // when an `Expand` sits below, one input row fans out into several
pass-through rows that are all
+ // appended in the same child loop iteration before any drain, and they all
alias the single
+ // result `UnsafeRow`. Copy the result so the buffered rows do not collapse
into the last one.
+ override def needCopyResult: Boolean =
+ adaptivePartialAggConfig.isDefined || super.needCopyResult
Review Comment:
Please propagate the child's copy requirement instead of enabling copies for
every adaptive aggregate. DBR gates this on the child, and `ExpandExec` already
reports `needCopyResult = true`, so doing the same preserves the aliasing fix
without adding `row.copy()` to ordinary output.
##########
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:
Please keep adaptive pass-through available after earlier spills. DBR
intentionally resets the processed-row count after each spill and evaluates the
new in-memory map epoch; if it activates, the existing sorter and current map
drain before the remaining rows pass through. Permanently gating on `sorter ==
null` materially diverges for inputs whose cardinality becomes unfavorable
later.
--
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]
