This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 02daddede [CELEBORN-2421] Avoid redundant ByteBuffer allocation in
LocalPartitionDataReader
02daddede is described below
commit 02daddede2be4f0cd66a5392aa2de66d2475262e
Author: r7raul1984 <[email protected]>
AuthorDate: Thu Aug 20 22:56:03 2026 +0800
[CELEBORN-2421] Avoid redundant ByteBuffer allocation in
LocalPartitionDataReader
### What changes were proposed in this pull request?
Replace the temporary heap `ByteBuffer` allocation in
`readBufferIntoReadBuffer` with a direct write loop using
`ByteBuf.writeBytes(FileChannel, pos, len)`, which writes from the channel
straight into the output `ByteBuf` with no intermediate allocation. The channel
position is advanced manually since the absolute-read API does not move the
cursor.
Additionally, fix two correctness issues in the read loop:
- Throw `EOFException` when `writeBytes` returns `-1` (premature EOF due to
file truncation) instead of silently breaking and returning a partially filled
`ByteBuf` to the caller.
- Throw `IOException` when `writeBytes` returns `0` to prevent the loop
from spinning indefinitely.
### Why are the changes needed?
Previously, `readBufferIntoReadBuffer` allocated a temporary heap
`ByteBuffer` on every call, performed a `FileChannel` read into it, then copied
the data into the output `ByteBuf` — two copies and one allocation per read.
The new approach eliminates the intermediate allocation and one copy.
The original loop also had two correctness issues: `dataFileSize` is cached
at construction time, so a subsequent file truncation could cause `writeBytes`
to return `-1` after the integrity check passes, resulting in the caller
receiving partial data without any error. A return value of `0` would cause an
infinite spin.
### Does this PR resolve a correctness bug?
- [x] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
Benchmarks at 1 MB buffer size over 500 iterations show ~2x throughput
improvement and ~40% fewer GC collections compared to the previous
implementation.
Closes #3800 from r7raul1984/jj_fix.
Authored-by: r7raul1984 <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../worker/storage/LocalPartitionDataReader.java | 25 +++++++++++++++++-----
1 file changed, 20 insertions(+), 5 deletions(-)
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/LocalPartitionDataReader.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/LocalPartitionDataReader.java
index 6777c525f..e15e2c8f1 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/LocalPartitionDataReader.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/LocalPartitionDataReader.java
@@ -17,6 +17,7 @@
package org.apache.celeborn.service.deploy.worker.storage;
+import java.io.EOFException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
@@ -71,12 +72,26 @@ public class LocalPartitionDataReader extends
PartitionDataReader {
public void readBufferIntoReadBuffer(ByteBuf buf, long fileSize, int length,
String filePath)
throws IOException {
Utils.checkFileIntegrity(fileSize - dataFileChanel.position(), length,
filePath);
- ByteBuffer tmpBuffer = ByteBuffer.allocate(length);
- while (tmpBuffer.hasRemaining()) {
- dataFileChanel.read(tmpBuffer);
+ long position = dataFileChanel.position();
+ int totalRead = 0;
+ while (totalRead < length) {
+ int read = buf.writeBytes(dataFileChanel, position + totalRead, length -
totalRead);
+ if (read < 0) {
+ throw new EOFException(
+ "Unexpected EOF in "
+ + filePath
+ + ": expected "
+ + length
+ + " bytes but only read "
+ + totalRead);
+ }
+ if (read == 0) {
+ throw new IOException(
+ "Zero bytes read from " + filePath + ", file may be corrupted or
truncated");
+ }
+ totalRead += read;
}
- tmpBuffer.flip();
- buf.writeBytes(tmpBuffer);
+ dataFileChanel.position(position + totalRead);
}
@Override
