sunchao commented on code in PR #57742:
URL: https://github.com/apache/spark/pull/57742#discussion_r3722889760
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -845,29 +1158,56 @@ 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.
+ val emitPassThroughRow = if (adaptivePartialAggEnabled) {
+ val numBypassingRows = metricTerm(ctx, "numBypassingRows")
+ s"""
+ |if ($adaptiveRowBypassedTerm) {
+ | $numBypassingRows.add(1);
+ | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer);
Review Comment:
**[P1] Drain frozen partial buffers before emitting bypassed rows**
This immediately sends a suffix row to the downstream `Final` aggregate,
while the earlier rows for the same group remain in the frozen hash maps until
a later output phase. That reverses the merge order even for a deterministic
single-partition plan with no shuffle.
A focused regression with `minRows=2` changes `first(v), last(v)` for one
group from **`(0, 20)` to `(20, 0)`**. With ANSI enabled, another valid group
containing `[-1, Long.MaxValue, 1]` newly fails with `ARITHMETIC_OVERFLOW`
because the final aggregate receives `Long.MaxValue` and `1` before the frozen
`-1`; the merge base returns `Long.MaxValue`. Please output the frozen prefix
before the first bypassed row and cover both an order-sensitive aggregate and
ANSI overflow.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -535,19 +613,30 @@ case class HashAggregateExec(
// `addNewFunction` spills this helper into a nested class (as can happen
// once the outer class passes the code-size threshold), the bare field
// reference fails with `IllegalAccessError`.
+
+ // Generate code for output. This must happen before the `doAgg` helper
below, because with
+ // adaptive partial aggregation enabled, `doConsumeWithKeys` (invoked from
the child's produce
+ // inside `doAgg`) emits pass-through rows by calling this output function
directly.
+ val keyTerm = ctx.freshName("aggKey")
+ val bufferTerm = ctx.freshName("aggBuffer")
+ outputFunc = generateResultFunction(ctx)
+
+ // After the child input is consumed, finish the build: with adaptive
partial aggregation mark
+ // that the child is fully consumed (to support re-entry; the map
iterators are set up inside
+ // the map-output function), otherwise set up the map iterators for the
output below.
+ val postChildProduce = if (adaptivePartialAggEnabled) {
+ s"$adaptiveChildrenConsumedTerm = true;"
Review Comment:
**[P1] Do not mark a yielded fused-union child as consumed**
`UnionExec` emits each child inside its own `unionChildProcess` helper. Once
adaptive pass-through appends a row, the leaf's `shouldStop()` returns from
that nested helper, but execution then reaches this assignment even though the
owning partition's input iterator still contains rows. The next `processNext()`
skips `adaptiveResumeBuild`, silently discarding that input.
With the normal two-level-map and union-partitioning settings and
`minRows=2`, `spark.range(0, 100, 1, 2).union(spark.range(100, 200, 1,
2)).groupBy("id").count()` returns only **12 of 200 groups** on this head; the
exact merge-base query returns all 200. At the default threshold the same
failure becomes reachable after the first 100,000 high-cardinality rows of a
fused-union partition. Please propagate a yield out of `doAgg` before setting
`adaptiveChildrenConsumed`, and add a genuinely fused `UnionExec` regression.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -626,13 +788,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 = adaptivePartialAggEnabled
Review Comment:
**[P2] Bound bypass output for non-resumable generator children**
Enabling `needStopCheck` does not bound a `GenerateExec` child: its
collection/iterator fan-out loop never checks `shouldStop()`, so it cannot
yield until the whole input collection has been expanded. Once adaptive
pass-through activates, each generated row is copied into the unbounded
`BufferedRowIterator` linked list before the upstream stop check can run.
For `explode(sequence(1, 1000)) -> partial HashAggregate` with `minRows=2`,
the exact merge base buffers **one** output row per task, but this head buffers
**998** before returning the first. With the default 100,000-row threshold, a
multi-million-element array can therefore retain millions of copied rows and
exhaust executor memory. The generator's standalone behavior predates this PR;
the newly regressed plan is specifically the previously blocking/spill-managed
generator-to-hash-aggregate pipeline. Please make the fan-out resumable or
disable pass-through beneath non-resumable one-to-many children.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -4156,6 +4156,48 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val ADAPTIVE_PARTIAL_AGGREGATION_ENABLED =
+
buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.enabled")
+ .doc("When true, hash aggregation adaptively bypasses the pre-shuffle
partial aggregation " +
+ "at runtime when it observes that the partial aggregation is not
reducing the number of " +
+ "rows enough to be worthwhile. Once bypassed, the remaining input rows
are passed " +
+ "through as single-row partial aggregation buffers for the final
aggregation to merge, " +
+ "which avoids the cost of maintaining and spilling a large aggregation
map with little " +
+ "reduction benefit. This applies only to hash aggregation with
grouping keys.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(true)
+
+ val ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS =
+
buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.minRows")
+ .doc("The number of rows to process before adaptive partial aggregation
(see " +
+ s"'${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}') evaluates the
compaction ratio. The " +
+ "ratio is evaluated once this many rows have been processed since the
previous " +
+ "evaluation, so a decision is never made on too few rows. A value of 0
disables the " +
+ "periodic evaluation entirely, leaving only the check made when the
aggregation map is " +
+ "about to spill.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .longConf
+ .checkValue(_ >= 0, "The minimum row count must not be negative.")
+ .createWithDefault(100000)
+
+ val ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION =
+
buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction")
+ .doc("The minimum compaction ratio required to keep the pre-shuffle
partial aggregation " +
+ s"(see '${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). The compaction
ratio is the " +
+ "number of processed rows divided by the number of keys held in the
aggregation maps, " +
+ "so a ratio of 10 means the partial aggregation collapses ten rows
into one. When the " +
+ s"ratio is below this value after
'${ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key}' rows, " +
+ "or when the aggregation map is about to spill, the partial
aggregation is bypassed for " +
+ "the rest of the input. A larger value bypasses more aggressively.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .doubleConf
+ .checkValue(_ >= 1.0, "The minimum compaction ratio must be at least
1.0.")
Review Comment:
**[P2] Reject non-finite compaction thresholds**
`Infinity` and overflowing literals such as `1e309` pass this `>= 1.0`
validation, but `HashAggregateExec` interpolates the resulting double directly
into generated Java as `InfinityD`, which is not a valid identifier or numeric
literal. A simple grouped query with
`spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction=Infinity`
and `spark.sql.codegen.fallback=false` reproducibly fails with `Unknown
variable or type "InfinityD"`; when fallback is enabled the whole stage
silently deoptimizes instead. This configuration did not exist at the merge
base. Please require a finite value or explicitly render
`Double.POSITIVE_INFINITY`, and add a configuration regression.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -663,46 +851,123 @@ case class HashAggregateExec(
case _ => ("true", "", "")
}
- val findOrInsertRegularHashMap: String =
+ // The compaction ratio is measured at the operator level: all processed
rows against the keys
+ // held by both maps, so two-level-map routing does not change the
decision. The same predicate
+ // decides both check points -- periodically every `minRows` rows, and
right before the map
+ // would spill (in which case the spill is skipped entirely). `minRows =
0` disables the
+ // periodic check: the row count is only ever compared after being
incremented past 0, so it
+ // never matches and only the spill check remains.
+ val adaptiveIneffective = if (adaptivePartialAggEnabled) {
+ val totalKeys = if (isFastHashMapEnabled) {
+ s"($fastHashMapTerm.getNumKeys() + $hashMapTerm.getNumKeys())"
+ } else {
+ s"$hashMapTerm.getNumKeys()"
+ }
+ s"$processedRowsTerm < (double) $totalKeys * ${adaptiveMinCompaction}D"
+ } else {
+ ""
+ }
+
+ // After a spill the map starts a new in-memory epoch, so the counters
restart and the ratio of
+ // that epoch alone decides whether the remaining rows are passed through.
+ val adaptiveResetEpoch = if (adaptivePartialAggEnabled) {
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();
- | }
- |}
+ |$processedRowsTerm = 0L;
Review Comment:
**[P2] Exclude persistent fast-map keys when starting a new spill epoch**
A regular-map spill resets `processedRows` here, but the compaction
predicate still includes every key in the fast map, which was not spilled and
survives across epochs. The next check therefore compares epoch-local rows
against a denominator containing unrelated pre-epoch keys.
With four old fast-map keys and eight new-epoch rows across four new
regular-map keys, the real compaction ratio is **2.0**, but the code computes
`8 / (4 + 4) = 1.0` and incorrectly bypasses. A controlled spill regression on
this head bypasses **92** effectively aggregatable rows; the merge base
continues aggregating them. At default fast-map capacity, 65,536 retained keys
can make a healthy **3.94x** epoch appear to be **1.0997x**, multiplying
shuffle volume. Snapshot/subtract pre-spill fast-map keys when resetting the
epoch, and test a post-spill input that should keep aggregating.
--
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]
