acruise opened a new pull request, #7418:
URL: https://github.com/apache/datafusion-comet/pull/7418
## Which issue does this PR close?
Closes #7070.
## Rationale for this change
`NativeBatchDecoderIterator` read each compressed shuffle block into a
thread-local direct buffer,
and `close()` shrank that buffer back to 128 KiB whenever it had grown.
`CometBlockStoreShuffleReader.read`
creates one iterator per fetched map output and closes the previous one as
it moves on, so for every
block over 128 KiB, which at the default batch size is most of them (#5678
measured 700 KiB to 1.4 MiB),
a reducer paid two `allocateDirect` calls: one to shrink on close and one to
grow again on the next
block. Each is zero-filled and goes through the JDK's direct-memory
reservation, which can force a
`System.gc()` near `MaxDirectMemorySize`. The coalescing read reset the
buffer at every block boundary
through a flag on `concatenate`, for the same effect. And the reset ran on
whichever thread called
`close()`, so a close from the task-completion thread replaced that thread's
buffer, not the reader's.
The direct-read path, `CometShuffleBlockIterator`, already keeps one buffer
per task that grows and
stays. This brings the JVM-consumer path, used by stages that end in a Spark
operator and by reads
without AQE, in line with it.
## What changes are included in this PR?
- `ShuffleReadBuffer` (new): a growable direct buffer owned by a task.
`prepare(bytes)` returns it
cleared with its limit at `bytes`, allocating 128 KiB for the first block
or twice the block when
one outgrows it; `release()` drops the reference. Releasing never frees
the memory itself: the task
thread fills the buffer outside the iterator's monitor, so that completion
can close the input it
reads from, and a read still in flight keeps its own reference until the
cleaner runs.
- `NativeBatchDecoderIterator` takes a `readBuffer` (default: its own),
reads each block through
`prepare`, and no longer resets anything in `close()`. The thread-local
and `resetDataBuf` are gone.
- `CometBlockStoreShuffleReader.read` creates one `ShuffleReadBuffer`,
passes it to every iterator,
and releases it in the completion listener after closing the last iterator
and the `NativeUtil`.
`readWithCoalescing` does the same for its single iterator.
`readAsRawStream` and `concatenate` lose
the reset flag.
- `CometCelebornShuffleReader` is unchanged: it creates one iterator per
task, whose own buffer is
therefore already per task.
- `fetchIterator` becomes `protected[shuffle]` so a test can hand the reader
its blocks.
## How are these changes tested?
- `CometBlockStoreShuffleReaderSuite`: `ShuffleReadBuffer` grows to twice
the block, keeps its size and
position/limit contract across smaller blocks, and lets go on release.
- `NativeBatchDecoderIteratorSuite`:
- `coalescing keeps the task's read buffer across fetched blocks` (was
`... releases oversized read
buffers between fetched blocks`): three frames across two fetched
blocks, the first over 128 KiB,
arrive in the same buffer object.
- `a reduce task reads every fetched block into one buffer`:
`CometBlockStoreShuffleReader.read` over
three fetched blocks, two of them over 128 KiB, allocates exactly one
direct buffer. It reports
three when the reader does not pass its buffer to the iterators.
- `CometCelebornShuffleReaderSuite`, which hosts the decoder lifecycle and
concurrency checks, passes
unchanged (78 tests across the three suites).
Throwaway harness, not included: `CometBlockStoreShuffleReader.read` over
400 fetched blocks of 169 KiB
per task (uncompressed frames of 40K ints, real native decode), best of 15
rounds after 5 warmup, three
JVMs each, Apple M3 Max, JDK 21:
| | per block (best) | per task, median |
|---|---|---|
| main | 53 to 58 us | 28 to 31 ms |
| this PR | 23 to 25 us | 10 to 11 ms |
The difference is the two zero-filled 340 KiB allocations per block. Not
measured on a cluster, where
the direct-memory reservation's `System.gc()` fallback would add to it.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
