This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 7a83a4d2a8b6 [SPARK-57908][SQL] Replace generated fast-hash-map
TaskCompletionListener with a close-hook helper
7a83a4d2a8b6 is described below
commit 7a83a4d2a8b6bee41733be30629e501a31fcd17f
Author: Gengliang Wang <[email protected]>
AuthorDate: Mon Jul 6 09:26:57 2026 -0700
[SPARK-57908][SQL] Replace generated fast-hash-map TaskCompletionListener
with a close-hook helper
### What changes were proposed in this pull request?
Janino cannot compile lambdas, so the per-map close hook that
`HashAggregateExec` whole-stage
codegen emits for the fast hash map was written as an anonymous
`TaskCompletionListener` class --
which Janino compiles into a separate generated inner class, one per fast
hash map.
Before -- the generated code emits an anonymous class implementing
`TaskCompletionListener`:
```java
$thisPlan.getTaskContext().addTaskCompletionListener(
new org.apache.spark.util.TaskCompletionListener() {
Override
public void onTaskCompletion(org.apache.spark.TaskContext context) {
fastHashMap.close();
}
});
```
After -- a one-line call to a new compiled helper
`HashAggregateExec.addFastHashMapCloseHook(AutoCloseable)` (the listener is
a lambda in compiled
Scala, which Janino never sees), and the generated map class (row-based and
vectorized) implements
`java.lang.AutoCloseable`:
```java
$thisPlan.addFastHashMapCloseHook(fastHashMap);
```
The hook registers at the same point with the same effect, once per task;
only the anonymous
generated inner class is removed.
### Why are the changes needed?
Part of [SPARK-56908](https://issues.apache.org/jira/browse/SPARK-56908)
(reduce generated Java size
in whole-stage codegen). On a TPC-DS codegen dump (150 queries, 1,572
whole-stage-codegen subtrees),
generated inner classes drop from 542 to 271 (**-50%**) -- the anonymous
listener accounts for
exactly half of all generated inner classes. Fewer generated classes means
less Janino work, less
class loading/verification, and less metaspace per query per executor.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests (`SingleLevelAggregateHashMapSuite`,
`TwoLevelAggregateHashMapSuite`,
`TwoLevelAggregateHashMapWithVectorizedMapSuite`, `WholeStageCodegenSuite`,
`DataFrameAggregateSuite`) exercise the generated fast hash map through
both the row-based and
vectorized paths; the close hook still fires once per task.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #56976 from gengliangwang/SPARK-57908-fastmap-close-hook.
Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
(cherry picked from commit 78d458bcff45acf8a2c98eaa2f03b1fe2e132f4c)
Signed-off-by: Gengliang Wang <[email protected]>
---
.../sql/execution/aggregate/HashAggregateExec.scala | 18 +++++++++++-------
.../sql/execution/aggregate/HashMapGenerator.scala | 2 +-
2 files changed, 12 insertions(+), 8 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 e1bfb50634d0..62c4f896f2ee 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
@@ -177,6 +177,16 @@ case class HashAggregateExec(
TaskContext.get()
}
+ /**
+ * Registers a task-completion hook to close the generated fast hash map, so
that the close hook
+ * is a plain method call on this plan, with the listener being a lambda in
compiled Scala code,
+ * rather than an anonymous `TaskCompletionListener` emitted per fast hash
map (one fewer
+ * generated inner class per map). This is called by the generated Java
class, should be public.
+ */
+ def addFastHashMapCloseHook(fastHashMap: AutoCloseable): Unit = {
+ TaskContext.get().addTaskCompletionListener[Unit](_ => fastHashMap.close())
+ }
+
def getEmptyAggregationBuffer(): InternalRow = {
val initExpr = declFunctions.flatMap(f => f.initialValues)
val initialBuffer = UnsafeProjection.create(initExpr)(EmptyRow)
@@ -488,13 +498,7 @@ case class HashAggregateExec(
// output (e.g. aggregate followed by limit).
val addHookToCloseFastHashMap = if (isFastHashMapEnabled) {
s"""
- |$thisPlan.getTaskContext().addTaskCompletionListener(
- | new org.apache.spark.util.TaskCompletionListener() {
- | @Override
- | public void onTaskCompletion(org.apache.spark.TaskContext
context) {
- | $fastHashMapTerm.close();
- | }
- |});
+ |$thisPlan.addFastHashMapCloseHook($fastHashMapTerm);
""".stripMargin
} else ""
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashMapGenerator.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashMapGenerator.scala
index af9dcb44cf1f..74e87106f498 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashMapGenerator.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashMapGenerator.scala
@@ -64,7 +64,7 @@ abstract class HashMapGenerator(
def generate(): String = {
s"""
- |public class $generatedClassName {
+ |public class $generatedClassName implements java.lang.AutoCloseable {
|${initializeAggregateHashMap()}
|
|${generateFindOrInsert()}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
