You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ge...@apache.org on 2026/08/22 02:42:59 UTC
(spark) branch branch-4.3 updated: [SPARK-57980][4.3][SQL] Extract the HashAggregateExec hash-map spill machinery into a shared helper
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 <ge...@apache.org>
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 <ge...@apache.org>
Signed-off-by: Gengliang Wang <ge...@apache.org>
---
.../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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org