dongjoon-hyun opened a new pull request, #60196:
URL: https://github.com/apache/spark/pull/60196
### What changes were proposed in this pull request?
This PR lets `BytesToBytesMap` start with a small data page and grow its
pages by doubling, and uses it in `UnsafeHashedRelation`, the hash relation for
join keys other than a single `LONG` (for example a `STRING` key or multiple
key columns).
- `BytesToBytesMap`: a new constructor takes `initialPageSizeBytes`. The
first data page has that size (at most `pageSizeBytes`), and each later page
doubles until it reaches `pageSizeBytes`. A record larger than the next page
still gets a page of its own size. The existing constructors pass
`pageSizeBytes` as the initial page size, so the other users, such as hash
aggregation (`UnsafeFixedWidthAggregationMap`), behave as before.
- `MemoryConsumer`: a new `allocatePage(size, required)` requests `size`
bytes and, like `allocatePage(required)`, accepts a smaller page of at least
`required` bytes. `allocatePage(required)` now calls it with `max(pageSize,
required)`. `BytesToBytesMap` needs both values so that a partial memory grant
still works when the requested page size grows.
- `UnsafeHashedRelation`: both building (`apply`) and deserializing (`read`)
use a 1 MiB initial page, `UnsafeHashedRelation.INITIAL_PAGE_SIZE_BYTES`, the
same size as `LongToUnsafeRowMap.INITIAL_PAGE_SIZE_BYTES` from SPARK-60998.
For a large relation, the pages are 1, 2, 4, 8, 16 and 32 MiB, then 64 MiB
each, so it needs 6 more pages than before. That is small next to the limit of
8192 pages per task.
The growth is opt-in because changing every `BytesToBytesMap` would also
change when hash aggregation falls back to sort-based aggregation. A smaller
fixed page size in `UnsafeHashedRelation` alone is not enough either: each
relation would still reserve a whole page, and a large relation would need many
more pages.
### Why are the changes needed?
`BytesToBytesMap` allocates its data pages through
`MemoryConsumer.allocatePage`, which rounds every request up to `max(pageSize,
required)`. `UnsafeHashedRelation` passes
`SparkEnv.get.memoryManager.pageSizeBytes`, which is `spark.buffer.pageSize` or
a value of up to 64 MiB derived from the executor memory and cores. So every
relation reserves a full page, even when it holds only a few rows.
- A shuffled hash join keeps its relation until the task ends, and the
relation cannot spill. A task with a chain of shuffled hash joins on such keys
runs out of execution memory long before the relations hold that much data:
```
[UNABLE_TO_ACQUIRE_MEMORY] Unable to acquire 2048 bytes of memory, got 0.
Memory used by task 0 grouped by consumer:
org.apache.spark.unsafe.map.BytesToBytesMap@5aa9a000: 64.0 MiB
org.apache.spark.unsafe.map.BytesToBytesMap@6bc9c3f9: 64.0 MiB
org.apache.spark.unsafe.map.BytesToBytesMap@5240e37f: 64.0 MiB
org.apache.spark.unsafe.map.BytesToBytesMap@31cf8e3c: 64.0 MiB
```
This is the output of the new test `SPARK-61082: many
UnsafeHashedRelations fit in one task when each holds little data` without this
fix: 4 relations of 10 rows each take the whole 256 MiB, and the 5th cannot
allocate its 2 KiB pointer array.
- A broadcast hash join builds the relation on the driver and deserializes
it on each executor with an on-heap memory manager that is not bounded by
execution memory. So every broadcast relation allocates a full page as a
`long[]` on the heap, 64 MiB even for a single row.
SPARK-60998 fixed the same problem for `LongToUnsafeRowMap`, which is used
for a single `LONG` key. For `UnsafeHashedRelation`, this is not a regression:
it has rounded its pages up to the page size since Spark 2.0, when it started
to build a `BytesToBytesMap` directly (SPARK-14052).
### Does this PR introduce _any_ user-facing change?
Yes. Queries that keep many `UnsafeHashedRelation`s in one task, such as a
chain of shuffled hash joins on non-`LONG` keys, no longer fail with
`UNABLE_TO_ACQUIRE_MEMORY` when the relations are small. Each relation reserves
memory close to what it stores, and a small broadcast relation uses much less
driver and executor heap. For small relations, the `data size of build side`
metric of `ShuffledHashJoin` and the `data size` metric of `BroadcastExchange`
are smaller accordingly. Hash aggregation and the other users of
`BytesToBytesMap` do not change.
### How was this patch tested?
New tests:
- `HashedRelationSuite`, which runs in both `HashedRelationOnHeapSuite` and
`HashedRelationOffHeapSuite`:
- `SPARK-61082: many UnsafeHashedRelations fit in one task when each holds
little data`: 16 relations of 10 rows each on a `STRING` key, in 256 MiB of
execution memory.
- `SPARK-61082: UnsafeHashedRelation does not round its page up to the
page size`: a relation with one row and its deserialized copy each reserve less
than 2 MiB, instead of about 64 MiB.
- `AbstractBytesToBytesMapSuite`, which runs in both
`BytesToBytesMapOnHeapSuite` and `BytesToBytesMapOffHeapSuite`:
- `dataPagesGrowFromInitialPageSize`: with a 1 KiB initial page and a 4
KiB page size, the pages are 1, 2, 4, 4 and 4 KiB.
- `initialCapacityBoundsChecking` also checks that the initial page size
must be positive.
Both `HashedRelationSuite` tests fail without the fix in both memory modes.
The existing suites for `BytesToBytesMap`, `TaskMemoryManager`, hash
aggregation and joins also pass.
```
build/sbt "core/testOnly
org.apache.spark.unsafe.map.BytesToBytesMapOnHeapSuite
org.apache.spark.unsafe.map.BytesToBytesMapOffHeapSuite
org.apache.spark.memory.TaskMemoryManagerSuite"
build/sbt "sql/testOnly
org.apache.spark.sql.execution.joins.HashedRelationOnHeapSuite
org.apache.spark.sql.execution.joins.HashedRelationOffHeapSuite
org.apache.spark.sql.execution.UnsafeFixedWidthAggregationMapSuite
org.apache.spark.sql.execution.UnsafeKVExternalSorterSuite"
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
--
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]