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 b3ea004d00d7 [SPARK-57907][SQL] Extract schema-independent
fast-hash-map machinery into a shared base class
b3ea004d00d7 is described below
commit b3ea004d00d7f1f0840a737b254d98fbe5469723
Author: Gengliang Wang <[email protected]>
AuthorDate: Mon Jul 6 10:44:07 2026 -0700
[SPARK-57907][SQL] Extract schema-independent fast-hash-map machinery into
a shared base class
### What changes were proposed in this pull request?
Most of every generated `hashAgg_FastHashMap_N` class is identical,
schema-independent boilerplate:
the field declarations, the constructor body (batch allocation,
empty-buffer projection, bucket
array init), `rowIterator`, and `close`. Only `findOrInsert`, `equals`, and
`hash` depend on the
key/value schema. That boilerplate is re-emitted into (and re-JIT'd for)
every aggregation stage.
Before -- the generated class carries all of it inline:
```java
public class hashAgg_FastHashMap_0 {
private org.apache.spark.sql.catalyst.expressions.RowBasedKeyValueBatch
batch;
private int[] buckets;
private int capacity = 1 << 16;
private double loadFactor = 0.5;
private int numBuckets = (int) (capacity / loadFactor);
private int maxSteps = 2;
private int numRows = 0;
private Object emptyVBase;
private long emptyVOff;
private int emptyVLen;
private boolean isBatchFull = false;
private ... UnsafeRowWriter agg_rowWriter;
public hashAgg_FastHashMap_0(TaskMemoryManager taskMemoryManager,
InternalRow emptyAggBuffer) {
batch = RowBasedKeyValueBatch.allocate(...);
final UnsafeProjection valueProjection = UnsafeProjection.create(...);
final byte[] emptyBuffer =
valueProjection.apply(emptyAggBuffer).getBytes();
emptyVBase = emptyBuffer; emptyVOff = ...; emptyVLen = ...;
agg_rowWriter = new UnsafeRowWriter(...);
buckets = new int[numBuckets];
java.util.Arrays.fill(buckets, -1);
}
public UnsafeRow findOrInsert(...) {
...
UnsafeRow agg_result = agg_rowWriter.getRow();
Object kbase = agg_result.getBaseObject(); long koff = ...; int klen =
...;
UnsafeRow vRow = batch.appendRow(kbase, koff, klen, emptyVBase,
emptyVOff, emptyVLen);
if (vRow == null) { isBatchFull = true; } else { buckets[idx] =
numRows++; }
return vRow;
...
}
private boolean equals(int idx, ...) { ... }
private long hash(...) { ... }
public ... KVIterator<UnsafeRow, UnsafeRow> rowIterator() { return
batch.rowIterator(); }
public void close() { batch.close(); }
}
```
After -- all schema-independent members move to a new hand-written base
class
`RowBasedAggregateHashMap` (compiled by javac and JIT'd once per JVM), and
the generated class
shrinks to a forwarding constructor plus the three typed methods; the
insert slow path delegates to
the base's `final appendCurrentKey(idx)`:
```java
public class hashAgg_FastHashMap_0
extends
org.apache.spark.sql.execution.aggregate.RowBasedAggregateHashMap {
public hashAgg_FastHashMap_0(TaskMemoryManager taskMemoryManager,
InternalRow emptyAggBuffer) {
super(keySchema, valueSchema, taskMemoryManager, emptyAggBuffer, 1 <<
16, numKeyFields, ...);
}
public UnsafeRow findOrInsert(...) {
...
return appendCurrentKey(idx); // shared slow path in the base class
...
}
private boolean equals(int idx, ...) { ... }
private long hash(...) { ... }
}
```
The bucket-scan/hash/equals hot path is emitted byte-identically to before.
Note this is a
relocation, not deletion of logic: the boilerplate still exists as
bytecode, but in one precompiled
base class instead of being re-emitted into every stage's generated class
-- which is what shrinks
Janino input, generated-class count, and per-stage JIT/profiling work.
Class hierarchy:
```
RowBasedAggregateHashMap (new, hand-written, AutoCloseable)
<- generated subclass (findOrInsert / equals / hash + forwarding
constructor)
```
### 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),
this removes about **-4%** of total generated source. The extracted
machinery is JIT-compiled and
profiled once per JVM instead of re-emitted and re-profiled per aggregation
stage.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests: `SingleLevelAggregateHashMapSuite`,
`TwoLevelAggregateHashMapSuite`,
`TwoLevelAggregateHashMapWithVectorizedMapSuite`, `WholeStageCodegenSuite`,
`DataFrameAggregateSuite`. The generated hot path and the extracted insert
path perform the same
operations in the same order as before.
### Related
Part of the
[SPARK-56908](https://issues.apache.org/jira/browse/SPARK-56908) series. The
generated
row-based map class continues to expose `close()` (the base class
implements `AutoCloseable`), so it
composes with the sibling close-hook change under the same umbrella.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #56975 from gengliangwang/SPARK-57907-fastmap-base-class.
Authored-by: Gengliang Wang <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
---
.../aggregate/RowBasedAggregateHashMap.java | 117 +++++++++++++++++++++
.../aggregate/RowBasedHashMapGenerator.scala | 97 ++++++++---------
2 files changed, 159 insertions(+), 55 deletions(-)
diff --git
a/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java
b/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java
new file mode 100644
index 000000000000..3f67d050a8c8
--- /dev/null
+++
b/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java
@@ -0,0 +1,117 @@
+/*
+ * 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 org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.catalyst.expressions.RowBasedKeyValueBatch;
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection;
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.unsafe.KVIterator;
+import org.apache.spark.unsafe.Platform;
+
+/**
+ * Shared, type-independent machinery for the row-based fast hash map that
{@code HashAggregateExec}
+ * generates as the level-1 cache during hash aggregation (see {@code
RowBasedHashMapGenerator}).
+ *
+ * The generated subclass supplies only the schema-specific, strongly-typed
methods that must be
+ * codegen'd for speed -- {@code findOrInsert}, {@code equals}, and {@code
hash} -- while the
+ * boilerplate that does not depend on the key/value schema (the batch and
bucket bookkeeping, the
+ * empty-buffer setup, the row iterator, and {@code close}) lives here so it
is written and JIT'd
+ * once per JVM rather than re-emitted into every aggregation stage's
generated class.
+ *
+ * The hot path is unchanged: the bucket scan, hash, equals, and typed key
write still run in the
+ * generated subclass; the one helper the subclass calls into here ({@link
#appendCurrentKey(int)})
+ * is {@code final} so HotSpot can inline it.
+ *
+ * NOTE: like the generated map it replaces, this does not support nullable
keys; the caller falls
+ * back to the {@code BytesToBytesMap} for those.
+ */
+public abstract class RowBasedAggregateHashMap implements AutoCloseable {
+ // NOTE: these field names are part of the contract with the generated
subclass --
+ // RowBasedHashMapGenerator emits code that references them by name (e.g.
`rowWriter`,
+ // `buckets`, `numBuckets`), so renaming one only fails when Janino compiles
the generated
+ // code at runtime, not at build time.
+ protected final RowBasedKeyValueBatch batch;
+ protected final int[] buckets;
+ protected final int capacity;
+ protected final int numBuckets;
+ protected final int maxSteps = 2;
+ protected int numRows = 0;
+ protected final Object emptyVBase;
+ protected final long emptyVOff;
+ protected final int emptyVLen;
+ protected boolean isBatchFull = false;
+ protected final UnsafeRowWriter rowWriter;
+
+ protected RowBasedAggregateHashMap(
+ StructType keySchema,
+ StructType valueSchema,
+ TaskMemoryManager taskMemoryManager,
+ InternalRow emptyAggregationBuffer,
+ int capacity,
+ int numKeyFields,
+ int keyVarLenBufferSize) {
+ this.capacity = capacity;
+ double loadFactor = 0.5;
+ this.numBuckets = (int) (capacity / loadFactor);
+ this.batch =
+ RowBasedKeyValueBatch.allocate(keySchema, valueSchema,
taskMemoryManager, capacity);
+
+ UnsafeProjection valueProjection = UnsafeProjection.create(valueSchema);
+ byte[] emptyBuffer =
valueProjection.apply(emptyAggregationBuffer).getBytes();
+ this.emptyVBase = emptyBuffer;
+ this.emptyVOff = Platform.BYTE_ARRAY_OFFSET;
+ this.emptyVLen = emptyBuffer.length;
+
+ this.rowWriter = new UnsafeRowWriter(numKeyFields, keyVarLenBufferSize);
+
+ this.buckets = new int[numBuckets];
+ java.util.Arrays.fill(this.buckets, -1);
+ }
+
+ /**
+ * Appends the key currently staged in {@link #rowWriter} (the generated
subclass has already
+ * called {@code reset()} and written the typed key columns) together with
the empty aggregation
+ * buffer, and records the new row's bucket. Returns the value row to
aggregate into, or
+ * {@code null} when the batch is full.
+ */
+ protected final UnsafeRow appendCurrentKey(int idx) {
+ UnsafeRow aggResult = rowWriter.getRow();
+ UnsafeRow vRow = batch.appendRow(
+ aggResult.getBaseObject(), aggResult.getBaseOffset(),
aggResult.getSizeInBytes(),
+ emptyVBase, emptyVOff, emptyVLen);
+ if (vRow == null) {
+ isBatchFull = true;
+ } else {
+ buckets[idx] = numRows++;
+ }
+ return vRow;
+ }
+
+ public final KVIterator<UnsafeRow, UnsafeRow> rowIterator() {
+ return batch.rowIterator();
+ }
+
+ @Override
+ public final void close() {
+ batch.close();
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/RowBasedHashMapGenerator.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/RowBasedHashMapGenerator.scala
index 286aa1acd3cb..1086a8d5fb80 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/RowBasedHashMapGenerator.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/RowBasedHashMapGenerator.scala
@@ -28,6 +28,13 @@ import org.apache.spark.sql.types._
* `BytesToBytesMap` if a given key isn't found). This is 'codegened' in
HashAggregate to speed
* up aggregates w/ key.
*
+ * The schema-independent machinery (fields, batch/bucket bookkeeping, the row
iterator, and
+ * `close`) lives in the hand-written base class
+ * [[org.apache.spark.sql.execution.aggregate.RowBasedAggregateHashMap]]; this
generator emits a
+ * thin subclass containing a forwarding constructor and the strongly-typed,
schema-specific
+ * methods (`findOrInsert`, `equals`, `hash`), so the boilerplate is written
and JIT'd once per
+ * JVM instead of being re-emitted into every aggregation stage's generated
code.
+ *
* We also have VectorizedHashMapGenerator, which generates a append-only
vectorized hash map.
* We choose one of the two as the 1st level, fast hash map during aggregation.
*
@@ -44,6 +51,30 @@ class RowBasedHashMapGenerator(
extends HashMapGenerator (ctx, aggregateExpressions, generatedClassName,
groupingKeySchema, bufferSchema) {
+ /**
+ * Emits a thin subclass of [[RowBasedAggregateHashMap]] holding a
forwarding constructor and
+ * the typed, schema-specific methods. The shared machinery (fields,
batch/bucket setup,
+ * `rowIterator`, `close`) is inherited from the base class.
+ */
+ override def generate(): String = {
+ s"""
+ |public class $generatedClassName
+ | extends
org.apache.spark.sql.execution.aggregate.RowBasedAggregateHashMap {
+ |${initializeAggregateHashMap()}
+ |
+ |${generateFindOrInsert()}
+ |
+ |${generateEquals()}
+ |
+ |${generateHashFunction()}
+ |}
+ """.stripMargin
+ }
+
+ /**
+ * Generates the constructor, which forwards the key/value schema, the empty
aggregation buffer,
+ * and the row-writer sizing to the base class that owns the shared state.
+ */
override protected def initializeAggregateHashMap(): String = {
val keySchema = ctx.addReferenceObj("keySchemaTerm", groupingKeySchema)
val valueSchema = ctx.addReferenceObj("valueSchemaTerm", bufferSchema)
@@ -55,38 +86,11 @@ class RowBasedHashMapGenerator(
}
s"""
- | private
org.apache.spark.sql.catalyst.expressions.RowBasedKeyValueBatch batch;
- | private int[] buckets;
- | private int capacity = 1 << $bitMaxCapacity;
- | private double loadFactor = 0.5;
- | private int numBuckets = (int) (capacity / loadFactor);
- | private int maxSteps = 2;
- | private int numRows = 0;
- | private Object emptyVBase;
- | private long emptyVOff;
- | private int emptyVLen;
- | private boolean isBatchFull = false;
- | private
org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter agg_rowWriter;
- |
- |
| public $generatedClassName(
| org.apache.spark.memory.TaskMemoryManager taskMemoryManager,
| InternalRow emptyAggregationBuffer) {
- | batch =
org.apache.spark.sql.catalyst.expressions.RowBasedKeyValueBatch
- | .allocate($keySchema, $valueSchema, taskMemoryManager, capacity);
- |
- | final UnsafeProjection valueProjection =
UnsafeProjection.create($valueSchema);
- | final byte[] emptyBuffer =
valueProjection.apply(emptyAggregationBuffer).getBytes();
- |
- | emptyVBase = emptyBuffer;
- | emptyVOff = Platform.BYTE_ARRAY_OFFSET;
- | emptyVLen = emptyBuffer.length;
- |
- | agg_rowWriter = new
org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(
- | ${groupingKeySchema.length}, ${numVarLenFields * 32});
- |
- | buckets = new int[numBuckets];
- | java.util.Arrays.fill(buckets, -1);
+ | super($keySchema, $valueSchema, taskMemoryManager,
emptyAggregationBuffer,
+ | 1 << $bitMaxCapacity, ${groupingKeySchema.length},
${numVarLenFields * 32});
| }
""".stripMargin
}
@@ -118,27 +122,28 @@ class RowBasedHashMapGenerator(
* [[org.apache.spark.sql.catalyst.expressions.UnsafeRow]] which keeps track
of the
* aggregate value(s) for a given set of keys. If the corresponding row
doesn't exist, the
* generated method adds the corresponding row in the associated
- * [[org.apache.spark.sql.catalyst.expressions.RowBasedKeyValueBatch]].
+ * [[org.apache.spark.sql.catalyst.expressions.RowBasedKeyValueBatch]] via
the inherited
+ * `appendCurrentKey` helper.
*
*/
protected def generateFindOrInsert(): String = {
val createUnsafeRowForKey = groupingKeys.zipWithIndex.map { case (key:
Buffer, ordinal: Int) =>
key.dataType match {
case t: DecimalType =>
- s"agg_rowWriter.write(${ordinal}, ${key.name}, ${t.precision},
${t.scale})"
+ s"rowWriter.write(${ordinal}, ${key.name}, ${t.precision},
${t.scale})"
case t: DataType =>
if (!t.isInstanceOf[StringType] &&
!t.isInstanceOf[CalendarIntervalType] &&
!CodeGenerator.isPrimitiveType(t)) {
throw new IllegalArgumentException(s"cannot generate code for
unsupported type: $t")
}
- s"agg_rowWriter.write(${ordinal}, ${key.name})"
+ s"rowWriter.write(${ordinal}, ${key.name})"
}
}.mkString(";\n")
val resetNullBits = if (groupingKeySchema.map(_.nullable).forall(_ ==
false)) {
""
} else {
- "agg_rowWriter.zeroOutNullBytes();"
+ "rowWriter.zeroOutNullBytes();"
}
s"""
@@ -151,23 +156,10 @@ class RowBasedHashMapGenerator(
| // Return bucket index if it's either an empty slot or already
contains the key
| if (buckets[idx] == -1) {
| if (numRows < capacity && !isBatchFull) {
- | agg_rowWriter.reset();
+ | rowWriter.reset();
| $resetNullBits
| ${createUnsafeRowForKey};
- | org.apache.spark.sql.catalyst.expressions.UnsafeRow agg_result
- | = agg_rowWriter.getRow();
- | Object kbase = agg_result.getBaseObject();
- | long koff = agg_result.getBaseOffset();
- | int klen = agg_result.getSizeInBytes();
- |
- | UnsafeRow vRow
- | = batch.appendRow(kbase, koff, klen, emptyVBase,
emptyVOff, emptyVLen);
- | if (vRow == null) {
- | isBatchFull = true;
- | } else {
- | buckets[idx] = numRows++;
- | }
- | return vRow;
+ | return appendCurrentKey(idx);
| } else {
| // No more space
| return null;
@@ -184,11 +176,6 @@ class RowBasedHashMapGenerator(
""".stripMargin
}
- protected def generateRowIterator(): String = {
- """
- |public org.apache.spark.unsafe.KVIterator<UnsafeRow, UnsafeRow>
rowIterator() {
- | return batch.rowIterator();
- |}
- """.stripMargin
- }
+ // The row iterator and close are inherited from RowBasedAggregateHashMap.
+ protected def generateRowIterator(): String = ""
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
