andygrove opened a new pull request, #5913: URL: https://github.com/apache/datafusion-comet/pull/5913
## Which issue does this PR close? Part of #5905 (finding R5). Does not close it. ## Rationale for this change `CometBlockStoreShuffleReader` creates one `NativeBatchDecoderIterator` per fetched map output and closes the previous one as it moves on. The iterator read each compressed block into a thread-local direct `ByteBuffer`, and `close()` reset that buffer to its 128 KB initial size whenever it had grown. With 8192-row batches almost every block is larger than 128 KB compressed, so for every map output a reducer paid two `ByteBuffer.allocateDirect` calls: one to shrink on close, one to regrow on the next block. Each is zero-filled and goes through the JDK's direct-memory reservation, which can trigger `System.gc()` when close to `MaxDirectMemorySize`. A reducer over thousands of map outputs did this thousands of times per task. The reset was also running on whichever thread called `close()`. When task completion closed the iterator from another thread, it replaced that thread's thread-local buffer rather than the reader's. The direct-read path (`CometShuffleBlockIterator`) already keeps one buffer per task that grows and stays; this brings the JVM-consumer path in line with it. ## What changes are included in this PR? - New `ShuffleBlockBuffer`, a task-scoped growable direct buffer: `acquire(n)` returns the buffer positioned at zero with the limit set, allocating only when the block does not fit (doubling, starting at 128 KB). It never shrinks; it is released with the task. - `NativeBatchDecoderIterator` takes a `dataBuffer: ShuffleBlockBuffer` (defaulting to a fresh one) instead of using a thread-local, and `close()` no longer touches the buffer. The thread-local and its reset are removed. - `CometBlockStoreShuffleReader.read()` creates one `ShuffleBlockBuffer` and passes it to every iterator it creates for the task. `CometCelebornShuffleReader` creates one iterator per task already, so its behaviour is unchanged. Memory: at most one direct buffer of twice the largest compressed block per running task, held for the task's duration, versus the old steady state of the same buffer plus a 128 KB one being churned per map output. ## How are these changes tested? - New lifecycle check `reusesTaskScopedBufferAcrossIterators` (run from `CometCelebornShuffleReaderSuite` like the other decoder lifecycle checks): the buffer allocates once at the initial size, hands the same instance back while blocks fit, grows to twice the block when one does not, and is shared by three successive iterators whose `close()` calls do not reallocate it. - `CometCelebornShuffleReaderSuite`, `CometNativeShuffleSuite` and `CometShuffleSuite` pass (171 tests), covering the block-store reader path end to end. -- 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]
