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]

Reply via email to