This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 5644a6ae3774 [SPARK-57980][4.3][SQL] Extract the HashAggregateExec
hash-map spill machinery into a shared helper
5644a6ae3774 is described below
commit 5644a6ae377434d41b32e01394d7c6ab080045cf
Author: Gengliang Wang <[email protected]>
AuthorDate: Fri Aug 21 19:42:48 2026 -0700
[SPARK-57980][4.3][SQL] Extract the HashAggregateExec hash-map spill
machinery into a shared helper
### What changes were proposed in this pull request?
This backports #57059 ([SPARK-57980][SQL]) to `branch-4.3`, cherry-picked
from master commit
c62a58d77ab45ffc5276845d050b5b1f0d7fee8f. The PR is already on `master` and
`branch-4.x`.
When `HashAggregateExec` cannot allocate an aggregation buffer from the
in-memory hash map, its
generated code spills the map to disk: the first spill destructs the map
into a new
`UnsafeKVExternalSorter`, and later spills merge the map into the existing
sorter. This
type-independent branch is re-emitted into every HashAggregateExec
whole-stage-codegen stage.
This PR extracts it into a shared helper
`HashAggregateExec.spillHashMapToSorter(hashMap, sorter):
UnsafeKVExternalSorter` and calls it from
the generated code. Returning the sorter lets the generated code keep it in
a single mutable field.
Before (generated code, per HashAggregateExec stage):
```java
if (hashAgg_unsafeRowAggBuffer_0 == null) {
if (hashAgg_sorter_0 == null) {
hashAgg_sorter_0 = hashAgg_hashMap_0.destructAndCreateExternalSorter();
} else {
hashAgg_sorter_0.merge(hashAgg_hashMap_0.destructAndCreateExternalSorter());
}
...
```
After:
```java
if (hashAgg_unsafeRowAggBuffer_0 == null) {
hashAgg_sorter_0 =
org.apache.spark.sql.execution.aggregate.HashAggregateExec
.spillHashMapToSorter(hashAgg_hashMap_0, hashAgg_sorter_0);
...
```
The helper is a static method (it needs nothing from the plan instance), so
-- unlike routing
through the `references[]` plan object -- it adds no per-stage
constant-pool entry. Only the fully
type-independent spill-and-create block is moved. The surrounding logic
that is interleaved with the
test-only controlled-fallback counter (`resetCounter`) and the retry lookup
(which uses the
schema-specific key/hash variables) is left inline unchanged.
**Conflict resolution.** The cherry-pick conflicted in
`HashAggregateExec.scala`. On `master`,
`findOrInsertRegularHashMap` has since been restructured by the
adaptive-partial-aggregation work,
which split the spill block out into its own `spillMap` fragment reused by
two branches. `branch-4.3`
predates that and still has the single flat interpolated string with the
spill block inline. The
resolution keeps `branch-4.3`'s structure and replaces only the inline
`destructAndCreateExternalSorter`/`merge` block with the helper call; no
part of the
adaptive-partial-aggregation refactor is pulled in. The new helper and the
new test suite are
verbatim from the master commit.
### Why are the changes needed?
This is part of the umbrella SPARK-56908 (reduce the size of code generated
by whole-stage codegen).
The spill branch is real bytecode and constant-pool method-references (the
`destructAndCreateExternalSorter` and `merge` calls) that Janino cannot
fold away; collapsing it to a
single helper call compiles it once per JVM instead of re-emitting it into
every HashAggregateExec
stage. Planning all 135 TPC-DS queries produces 562 HashAggregateExec
whole-stage-codegen stages that
carry this block.
On `master`, `WholeStageCodegenSizeBenchmark` measured this as -0.4% summed
constant pool
(432,881 -> 431,233) and -0.2% source code size (21,846,001 -> 21,794,526
chars), with max method
bytecode, inner classes, and codegen fallbacks unchanged. Those numbers are
not re-measured here: the
benchmark harness (SPARK-57915) is not on `branch-4.3`.
Keeping `branch-4.3` in sync with `master` and `branch-4.x` on this file
also keeps subsequent
SPARK-56908 backports conflict-free.
### Does this PR introduce _any_ user-facing change?
No. The helper performs exactly the same spill/merge operations in the same
order as the previous
inline code; this is a pure code-organization change with no behavioral
difference.
### How was this patch tested?
Existing tests, plus the test suite that came with the original PR, run on
`branch-4.3`:
- `HashAggregateExecSuite` (added by the original PR) -- directly covers
both branches of the
extracted helper, the first-spill destruct and the later-spill merge. 2
tests, passed.
- `HashAggregationQueryWithControlledFallbackSuite`, which forces the hash
map to spill at
controlled points (`fallbackStartsAt` = 1, 2, 3) and therefore exercises
the rewritten generated
code with whole-stage codegen both on and off. 26 tests, passed.
- `sql/compile` + `sql/Test/compile` clean.
No behavior changes, so no new tests are added beyond the ones
cherry-picked with the commit.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
Closes #58220 from gengliangwang/SPARK-57980-branch-4.3.
Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
---
.../execution/aggregate/HashAggregateExec.scala | 30 ++++-
.../aggregate/HashAggregateExecSuite.scala | 137 +++++++++++++++++++++
2 files changed, 162 insertions(+), 5 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 62c4f896f2ee..93d8becccaae 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
@@ -676,11 +676,8 @@ case class HashAggregateExec(
|// 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());
- | }
+ | $sorterTerm =
org.apache.spark.sql.execution.aggregate.HashAggregateExec
+ | .spillHashMapToSorter($hashMapTerm, $sorterTerm);
| $resetCounter
| // the hash map had be spilled, it should have enough memory now,
| // try to allocate buffer again.
@@ -897,3 +894,26 @@ case class HashAggregateExec(
override protected def withNewChildInternal(newChild: SparkPlan):
HashAggregateExec =
copy(child = newChild)
}
+
+object HashAggregateExec {
+ /**
+ * Spills the in-memory hash map to disk and returns the sorter holding the
spilled data. Called
+ * by the generated code of [[HashAggregateExec]] (through the static
forwarder) when a buffer
+ * cannot be allocated from the hash map: the first spill destructs the map
into a new sorter, and
+ * later spills merge into the existing one. Returning the sorter lets the
generated code keep it
+ * in a single mutable field.
+ *
+ * Extracting this type-independent spill machinery into a shared helper
keeps it compiled once
+ * per JVM instead of being re-emitted into every HashAggregateExec stage's
generated code.
+ */
+ def spillHashMapToSorter(
+ hashMap: UnsafeFixedWidthAggregationMap,
+ sorter: UnsafeKVExternalSorter): UnsafeKVExternalSorter = {
+ if (sorter == null) {
+ hashMap.destructAndCreateExternalSorter()
+ } else {
+ sorter.merge(hashMap.destructAndCreateExternalSorter())
+ sorter
+ }
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExecSuite.scala
new file mode 100644
index 000000000000..2da4b5e0d66f
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExecSuite.scala
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.aggregate
+
+import java.util.Properties
+
+import scala.util.Random
+
+import org.mockito.Mockito._
+
+import org.apache.spark.{SparkConf, SparkFunSuite, TaskContext,
TaskContextImpl}
+import org.apache.spark.internal.config.MEMORY_OFFHEAP_ENABLED
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.{UnsafeFixedWidthAggregationMap,
UnsafeKVExternalSorter}
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField,
StructType}
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Test suite for the static helpers of [[HashAggregateExec]] that its
generated code calls into.
+ */
+class HashAggregateExecSuite extends SparkFunSuite with SharedSparkSession {
+
+ private val groupKeySchema = StructType(StructField("product", StringType)
:: Nil)
+ private val aggBufferSchema = StructType(StructField("salePrice",
IntegerType) :: Nil)
+ private val PAGE_SIZE_BYTES: Long = 1L << 26 // 64 megabytes
+
+ /**
+ * Runs `f` against a fresh aggregation map, with a task context in place
because both the map and
+ * the sorter it is destructed into allocate through the task memory
manager. Asserts on the way
+ * out that the spilled sorters left nothing behind.
+ */
+ private def withAggregationMap(f: UnsafeFixedWidthAggregationMap => Unit):
Unit = {
+ val conf = new SparkConf().set(MEMORY_OFFHEAP_ENABLED.key, "false")
+ val taskMemoryManager = new TaskMemoryManager(new TestMemoryManager(conf),
0)
+ // The map registers a completion listener, which a mock swallows: this
test drives the spills
+ // directly rather than running a task to completion. The sorter, in turn,
reads the task memory
+ // manager off the thread-local context, so a real one has to be in place
as well.
+ val taskContext = mock(classOf[TaskContext])
+ when(taskContext.taskMemoryManager()).thenReturn(taskMemoryManager)
+ TaskContext.setTaskContext(new TaskContextImpl(
+ stageId = 0,
+ stageAttemptNumber = 0,
+ partitionId = 0,
+ numPartitions = 1,
+ taskAttemptId = Random.nextInt(10000),
+ attemptNumber = 0,
+ taskMemoryManager = taskMemoryManager,
+ localProperties = new Properties,
+ metricsSystem = null))
+
+ val map = new UnsafeFixedWidthAggregationMap(
+ InternalRow(0), // empty aggregation buffer
+ aggBufferSchema,
+ groupKeySchema,
+ taskContext,
+ 128, // initial capacity
+ PAGE_SIZE_BYTES)
+
+ try {
+ f(map)
+ } finally {
+ map.free()
+ TaskContext.unset()
+ }
+ assert(taskMemoryManager.cleanUpAllAllocatedMemory() === 0)
+ }
+
+ /** Inserts `keys` into the map, using each key's length as its aggregation
buffer value. */
+ private def insert(map: UnsafeFixedWidthAggregationMap, keys: Seq[String]):
Unit = {
+ keys.foreach { key =>
+ val buffer =
map.getAggregationBuffer(InternalRow(UTF8String.fromString(key)))
+ assert(buffer != null)
+ buffer.setInt(0, key.length)
+ }
+ }
+
+ /**
+ * Drains the sorter and returns the keys it held, checking along the way
that each key still
+ * carries its own aggregation buffer -- draining to exhaustion also
releases the sorter's memory.
+ */
+ private def drainKeys(sorter: UnsafeKVExternalSorter): Seq[String] = {
+ val keys = Seq.newBuilder[String]
+ val iter = sorter.sortedIterator()
+ while (iter.next()) {
+ assert(iter.getKey.getString(0).length === iter.getValue.getInt(0))
+ keys += iter.getKey.getString(0)
+ }
+ keys.result()
+ }
+
+ test("spillHashMapToSorter destructs the map into a new sorter on the first
spill") {
+ withAggregationMap { map =>
+ val keys = Seq("apple", "banana", "cherry")
+ insert(map, keys)
+
+ val sorter = HashAggregateExec.spillHashMapToSorter(map, null)
+
+ assert(sorter != null)
+ assert(drainKeys(sorter) === keys.sorted)
+ }
+ }
+
+ test("spillHashMapToSorter merges into the existing sorter and returns it") {
+ withAggregationMap { map =>
+ val firstKeys = Seq("apple", "banana", "cherry")
+ val secondKeys = Seq("damson", "elderberry", "fig")
+ insert(map, firstKeys)
+ val sorter = HashAggregateExec.spillHashMapToSorter(map, null)
+
+ // The map keeps accepting keys after being destructed, so the second
spill hits the merge
+ // branch with a sorter already in hand.
+ insert(map, secondKeys)
+ val merged = HashAggregateExec.spillHashMapToSorter(map, sorter)
+
+ // The same sorter comes back, now holding the keys of both spills.
+ assert(merged eq sorter)
+ assert(drainKeys(merged) === (firstKeys ++ secondKeys).sorted)
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
