[SPARK-58511][SQL][FOLLOWUP] Reset the adaptive partial aggregation pass-through buffer with copyFrom

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 {

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}")
+    }
+  }
 }