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..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,15 +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. - val emptyAggBufferCode = if (adaptivePartialAggEnabled) { - GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) - } else { - null - } // Per-row local flag marking that the current row is being streamed through (held by no map). val adaptiveRowBypassedTerm = ctx.freshName("adaptiveRowBypassed") @@ -1106,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) { @@ -1121,8 +1131,8 @@ case class HashAggregateExec( | } |} else if ($adaptivePassThroughTerm) { | $adaptiveRowBypassedTerm = true; - | ${emptyAggBufferCode.code} - | $unsafeRowBuffer = ${emptyAggBufferCode.value}; + | $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}") + } + } }