This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c62a58d77ab4 [SPARK-57980][SQL] Extract the HashAggregateExec hash-map
spill machinery into a shared helper
c62a58d77ab4 is described below
commit c62a58d77ab45ffc5276845d050b5b1f0d7fee8f
Author: Gengliang Wang <[email protected]>
AuthorDate: Fri Aug 21 12:44:28 2026 -0700
[SPARK-57980][SQL] Extract the HashAggregateExec hash-map spill machinery
into a shared helper
### What changes were proposed in this pull request?
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);
...
```
This follows the same "extract type-independent generated machinery into a
compiled helper" pattern
as SPARK-57909 (`ColumnarToRowExec.advanceBatch`), and mirrors the existing
`HashAggregateExec.finishAggregate` helper that already backs the
aggregate's output path. 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.
Measured with `WholeStageCodegenSizeBenchmark` (added under SPARK-57915;
plans all 135 TPC-DS
queries, 562 of the 2247 whole-stage-codegen stages carry this block),
current `master` vs. this
change:
| metric | master | this PR | delta |
| --- | --- | --- | --- |
| constant pool, summed over stages | 432,881 | 431,233 | -0.4% |
| source code size (chars) | 21,846,001 | 21,794,526 | -0.2% |
| max method bytecode, summed over stages | 841,185 | 841,048 | ~unchanged |
| inner classes | 258 | 258 | unchanged |
| codegen fallbacks | 0 | 0 | unchanged |
The block lives in the aggregate's `doAggregateWithKeys` helper method
(rarely a stage's single
largest method), so `max method bytecode (sum)` barely moves; the reduction
shows up in the
constant pool, since the `destructAndCreateExternalSorter` / `merge` method
references are now
emitted once instead of per stage.
### 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.
### 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, run with whole-stage codegen both on and off:
- `HashAggregationQueryWithControlledFallbackSuite`, which forces the hash
map to spill at controlled
points (`fallbackStartsAt` = 1, 2, 3) and therefore exercises both the
first-spill (destruct) and
later-spill (merge) branches of the extracted helper, and
- `DataFrameAggregateSuite` / `AggregationQuerySuite` for general aggregate
correctness.
No behavior changes, so no new tests are added; the extracted branch is
directly covered by the
controlled-fallback suite.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #57059 from gengliangwang/SPARK-57980-hashagg-spill.
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 59b40ecdfbd0..f3561fb91394 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
@@ -979,11 +979,8 @@ case class HashAggregateExec(
val spillMap =
s"""
- |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 been spilled, so it should have enough memory
now,
|// try to allocate buffer again.
@@ -1354,3 +1351,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]
