gengliangwang opened a new pull request, #57059:
URL: https://github.com/apache/spark/pull/57059
### 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)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
