viirya opened a new pull request, #60167:
URL: https://github.com/apache/spark/pull/60167
### What changes were proposed in this pull request?
This PR passes a 1 MiB page size floor to `MemoryConsumer` in
`LongToUnsafeRowMap`, instead of the default page size of the
`TaskMemoryManager`.
`MemoryConsumer.allocatePage` rounds every request up to `max(pageSize,
required)`, and `pageSize` is only used there. With the 1 MiB floor:
- the initial page is 1 MiB again, as before SPARK-54116,
- the page grows by doubling (2 MiB, 4 MiB, ...) instead of jumping to the
page size,
- a deserialized map allocates what it reads, with at least 1 MiB.
`LongToUnsafeRowMap.MIN_PAGE_SIZE_BYTES` is added for the floor, and
`init()` uses it for the initial page. The on-heap and off-heap support added
by SPARK-54116 does not change.
### Why are the changes needed?
Since SPARK-54116, `LongToUnsafeRowMap` allocates its page through
`MemoryConsumer.allocatePage`. Before it, the map used a 1 MiB `long[]` and
acquired execution memory for exactly that size. Now every map reserves at
least a full page, which is `spark.buffer.pageSize`, or a value derived from
the executor memory and cores of up to 64 MiB. That happens even when the map
holds only a few rows.
When one task keeps many maps alive at once, for example a stage with a long
chain of shuffled hash joins, the task runs out of execution memory long before
the maps hold that much data:
```
[UNABLE_TO_ACQUIRE_MEMORY] Unable to acquire 512 bytes of memory, got 0.
Memory used by task 0 grouped by consumer:
org.apache.spark.sql.execution.joins.LongToUnsafeRowMap@47d52df1: 64.0 MiB
org.apache.spark.sql.execution.joins.LongToUnsafeRowMap@454d5d4a: 64.0 MiB
org.apache.spark.sql.execution.joins.LongToUnsafeRowMap@c17cf36: 64.0 MiB
org.apache.spark.sql.execution.joins.LongToUnsafeRowMap@239e2389: 64.0 MiB
```
This is the output of the new test `SPARK-60998: many LongToUnsafeRowMaps
fit in one task when each holds little data` without this fix: 4 maps of 10
rows each take the whole 256 MiB.
### Does this PR introduce _any_ user-facing change?
Yes. It fixes a regression in Spark 4.2.0: queries that keep many
`LongToUnsafeRowMap`s in one task, such as a chain of shuffled hash joins,
could fail with `UNABLE_TO_ACQUIRE_MEMORY`. Each map now reserves memory close
to what it stores, as in Spark 4.1 and earlier.
### How was this patch tested?
New tests in `HashedRelationSuite`, which run in both
`HashedRelationOnHeapSuite` and `HashedRelationOffHeapSuite` with a 64 MiB page
size:
- `SPARK-60998: LongToUnsafeRowMap does not round its page up to the page
size`: the map reserves the array plus 1 MiB, and 2 MiB after a 1.5 MiB value
is appended.
- `SPARK-60998: many LongToUnsafeRowMaps fit in one task when each holds
little data`: 16 small maps in 256 MiB of execution memory.
- `SPARK-60998: deserialized LongToUnsafeRowMap does not round its page up
to the page size`: for a non-empty map and an empty map.
All three tests fail without the fix, in both memory modes.
```
build/sbt "sql/testOnly
org.apache.spark.sql.execution.joins.HashedRelationOnHeapSuite
org.apache.spark.sql.execution.joins.HashedRelationOffHeapSuite"
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]