From 0938b6faa072bf4e0a95f1c88647337323919ccb Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Fri, 23 Oct 2026 16:07:44 +0000 Subject: [PATCH 1/2] [SPARK-58511][SQL][FOLLOWUP] Generate the adaptive partial aggregation initial buffer in its own method When adaptive partial aggregation bypasses the partial aggregation, each passed-through row starts from a fresh initial aggregation buffer, which the generated code built inline in the per-row consume method. The projection grows with the number of buffer slots, so generate it in its own method to keep a wide aggregation's consume method under the JIT's huge-method limit. Co-authored-by: Isaac --- .../aggregate/HashAggregateExec.scala | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala index 1dec77de3bbf..4fca03d4e1ff 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala @@ -918,11 +918,22 @@ case class HashAggregateExec( // For adaptive partial aggregation pass-through, each bypassed row is emitted as a single-row // partial buffer: start from the initial aggregation buffer, apply the update expressions once, // and output `key ++ buffer` for the Final aggregate to merge. This projects the initial - // buffer. - val emptyAggBufferCode = if (adaptivePartialAggEnabled) { - GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) + // buffer, in its own method: the projection's code grows with the number of buffer slots, and + // inlining it into the per-row consume method could push a wide aggregation's consume method + // past the JIT's huge-method limit. + val makeEmptyAggBufferFunc = if (adaptivePartialAggEnabled) { + val emptyAggBufferCode = + GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) + val makeEmptyAggBuffer = ctx.freshName("makeEmptyAggBuffer") + ctx.addNewFunction(makeEmptyAggBuffer, + s""" + |private UnsafeRow $makeEmptyAggBuffer() { + | ${emptyAggBufferCode.code} + | return ${emptyAggBufferCode.value}; + |} + """.stripMargin) } else { - null + "" } // Per-row local flag marking that the current row is being streamed through (held by no map). val adaptiveRowBypassedTerm = ctx.freshName("adaptiveRowBypassed") @@ -1121,8 +1132,7 @@ case class HashAggregateExec( | } |} else if ($adaptivePassThroughTerm) { | $adaptiveRowBypassedTerm = true; - | ${emptyAggBufferCode.code} - | $unsafeRowBuffer = ${emptyAggBufferCode.value}; + | $unsafeRowBuffer = $makeEmptyAggBufferFunc(); |} """.stripMargin } else { From bbf464a404bbe12763bd7095dce0a61d7f7e4fa2 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Sat, 24 Oct 2026 06:57:07 +0000 Subject: [PATCH 2/2] Reset the pass-through buffer with copyFrom from buffers registered once Build the initial aggregation buffer once per task and reset a reusable buffer from it with `UnsafeRow.copyFrom` for each passed-through row, as the interpreted path does, instead of projecting the initial values for every row. Register both buffers in `doProduceWithKeys` so the children of a fused `UnionExec` share them, and add a generated-code test. Co-authored-by: Isaac --- .../aggregate/HashAggregateExec.scala | 44 +++++++++---------- .../AdaptivePartialAggregationSuite.scala | 38 +++++++++++++++- 2 files changed, 59 insertions(+), 23 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala index 4fca03d4e1ff..f64cec4bf244 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala @@ -265,6 +265,15 @@ case class HashAggregateExec( // passed-through rows a child emits before it honours `shouldStop()` (see // `handlePassThroughRow`). private var adaptivePendingRowsTerm: String = _ + // The initial aggregation buffer, and the reusable buffer each passed-through row starts from. + // The reusable buffer is reset with `copyFrom`, as on the interpreted path, so the per-row reset + // is a single call however many buffer slots there are. Both are created once per task and + // registered in `doProduceWithKeys`, which runs once, rather than in `doConsumeWithKeys`, which + // runs once per child of a fused `UnionExec`. Reusing one buffer is safe because a + // passed-through row is copied before it is queued, and `outputFunc` reads the buffer before it + // returns. + private var adaptiveEmptyBufferTerm: String = _ + private var adaptivePassThroughBufferTerm: String = _ /** * This is called by generated Java class, should be public. @@ -589,6 +598,14 @@ case class HashAggregateExec( val thisPlan = ctx.addReferenceObj("plan", this) + if (adaptivePartialAggEnabled) { + adaptiveEmptyBufferTerm = ctx.addMutableState("UnsafeRow", "adaptiveEmptyBuffer", + v => s"$v = (UnsafeRow) $thisPlan.getEmptyAggregationBuffer();", forceInline = true) + adaptivePassThroughBufferTerm = ctx.addMutableState("UnsafeRow", + "adaptivePassThroughBuffer", v => s"$v = $adaptiveEmptyBufferTerm.copy();", + forceInline = true) + } + // Create a name for the iterator from the fast hash map, and the code to create fast hash map. val (iterTermForFastHashMap, createFastHashMap) = if (isFastHashMapEnabled) { // Generates the fast hash map class and creates the fast hash map term. @@ -915,26 +932,6 @@ case class HashAggregateExec( val unsafeRowBuffer = ctx.freshName("unsafeRowAggBuffer") val fastRowBuffer = ctx.freshName("fastAggBuffer") - // For adaptive partial aggregation pass-through, each bypassed row is emitted as a single-row - // partial buffer: start from the initial aggregation buffer, apply the update expressions once, - // and output `key ++ buffer` for the Final aggregate to merge. This projects the initial - // buffer, in its own method: the projection's code grows with the number of buffer slots, and - // inlining it into the per-row consume method could push a wide aggregation's consume method - // past the JIT's huge-method limit. - val makeEmptyAggBufferFunc = if (adaptivePartialAggEnabled) { - val emptyAggBufferCode = - GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) - val makeEmptyAggBuffer = ctx.freshName("makeEmptyAggBuffer") - ctx.addNewFunction(makeEmptyAggBuffer, - s""" - |private UnsafeRow $makeEmptyAggBuffer() { - | ${emptyAggBufferCode.code} - | return ${emptyAggBufferCode.value}; - |} - """.stripMargin) - } else { - "" - } // Per-row local flag marking that the current row is being streamed through (held by no map). val adaptiveRowBypassedTerm = ctx.freshName("adaptiveRowBypassed") @@ -1117,7 +1114,9 @@ case class HashAggregateExec( } // The grouping key was already projected in `findOrInsertRegularHashMap` // (`unsafeRowKeyCode.code`), so `unsafeRowKeyCode.value` holds this row's key. Only build - // the single-row partial buffer here. + // the single-row partial buffer here: reset the reusable buffer to the initial aggregation + // buffer, so the update expressions below apply this row alone and `key ++ buffer` is + // emitted for the Final aggregate to merge. s""" |if ($heldByAMap) { | if (!$adaptivePassThroughTerm) { @@ -1132,7 +1131,8 @@ case class HashAggregateExec( | } |} else if ($adaptivePassThroughTerm) { | $adaptiveRowBypassedTerm = true; - | $unsafeRowBuffer = $makeEmptyAggBufferFunc(); + | $adaptivePassThroughBufferTerm.copyFrom($adaptiveEmptyBufferTerm); + | $unsafeRowBuffer = $adaptivePassThroughBufferTerm; |} """.stripMargin } else { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala index d20848fd5345..83665ec51c52 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.aggregate import org.apache.spark.sql.{DataFrame, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.aggregate.{Partial, PartialMerge} +import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -34,7 +35,7 @@ import org.apache.spark.tags.ExtendedSQLTest * behind the map and flushed only after it drains, so every group merges its buffers in the same * order as a run that never bypasses, including order-sensitive aggregates such as `first`/`last`. * - * The suite has two halves: + * The suite has three parts: * 1. Correctness: aggregate results are identical to the reference (feature-off) run across the * full matrix of codegen on/off, two-level map on/off, and spill/no-spill, over a range of * aggregate shapes, key types, and `Expand`-bearing plans (ROLLUP / CUBE / GROUPING SETS / @@ -43,6 +44,7 @@ import org.apache.spark.tags.ExtendedSQLTest * 2. Triggering: the `numBypassingRows` metric proves the bypass actually fires when (and only * when) it should -- high-cardinality input bypasses, low-cardinality input keeps aggregating, * the feature switch and eligibility rules are honored, and both check points work. + * 3. Generated code: the pass-through path keeps the code layout that bounds its size. */ @ExtendedSQLTest class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession @@ -1420,4 +1422,38 @@ class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession } } } + + ///////////////////////////////////////////////////////////////////////////// + // Part 3: Generated code -- the layout that bounds the pass-through code size. + ///////////////////////////////////////////////////////////////////////////// + + test("pass-through resets its buffer with copyFrom from buffers registered once") { + // Each passed-through row resets a reusable buffer with one `copyFrom` of the initial + // aggregation buffer rather than projecting the initial values again, so the per-row code does + // not grow with the number of buffer slots. Under a fused Union the consume code is generated + // once per child, so each child resets the buffer, but all of them must share the same pair + // of buffers rather than registering a pair each. + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") ++ fixedPlanConfs): _*) { + val df = spark.range(0, 100, 1, 2) + .union(spark.range(100, 200, 1, 2)) + .groupBy("id") + .agg(sum($"id") as "s", avg($"id") as "a", max($"id") as "m") + df.collect() + val code = collect(df.queryExecution.executedPlan) { + case w: WholeStageCodegenExec => w.doCodeGen()._2.body + }.mkString("\n") + val resets = """(\w+)\.copyFrom\((\w+)\);""".r.findAllMatchIn(code) + .map(m => (m.group(1), m.group(2))) + .filter(_._2.contains("adaptiveEmptyBuffer")) + .toSeq + assert(resets.length == 2, + s"expected one reset per child of the fused Union, got ${resets.length}:\n$code") + assert(resets.distinct.length == 1, + s"expected every child to reset the same pair of buffers, got ${resets.distinct}") + } + } }