You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ju...@apache.org on 2025/02/06 03:23:27 UTC
(flink) branch release-2.0 updated: [FLINK-37218][network] Avoid repeatedly reading the broadcast buffer from the sort-merge shuffle data file while accumulating small buffers from multiple sub-partitions.
This is an automated email from the ASF dual-hosted git repository.
junrui pushed a commit to branch release-2.0
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/release-2.0 by this push:
new 4e2bc4d05f4 [FLINK-37218][network] Avoid repeatedly reading the broadcast buffer from the sort-merge shuffle data file while accumulating small buffers from multiple sub-partitions.
4e2bc4d05f4 is described below
commit 4e2bc4d05f4444fc2911c0347cc0c7ec30699b62
Author: JunRuiLee <jr...@gmail.com>
AuthorDate: Fri Jan 24 13:59:27 2025 +0800
[FLINK-37218][network] Avoid repeatedly reading the broadcast buffer from the sort-merge shuffle data file while accumulating small buffers from multiple sub-partitions.
---
.../network/partition/PartitionedFileReader.java | 106 ++++++++++++++++-----
.../partition/SortMergeSubpartitionReader.java | 12 ++-
.../partition/PartitionedFileWriteReadTest.java | 38 +++++---
3 files changed, 114 insertions(+), 42 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionedFileReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionedFileReader.java
index 582e1637c44..74acc9cba88 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionedFileReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionedFileReader.java
@@ -19,7 +19,6 @@
package org.apache.flink.runtime.io.network.partition;
import org.apache.flink.annotation.VisibleForTesting;
-import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferHeader;
@@ -32,6 +31,7 @@ import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.util.ArrayDeque;
import java.util.Queue;
+import java.util.function.BiConsumer;
import java.util.function.Consumer;
import static org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil.HEADER_LENGTH;
@@ -74,8 +74,10 @@ class PartitionedFileReader {
/** Number of remaining bytes in the current data region read. */
private long currentRegionRemainingBytes;
- /** A queue storing pairs of file offsets and sizes to be read. */
- private final Queue<Tuple2<Long, Long>> offsetAndSizesToRead = new ArrayDeque<>();
+ /** A queue storing {@link BufferPositionDescriptor} to be read. */
+ private final Queue<BufferPositionDescriptor> readBufferPositions = new ArrayDeque<>();
+
+ private BufferPositionDescriptor currentBufferPositionDescriptor;
PartitionedFileReader(
PartitionedFile partitionedFile,
@@ -100,14 +102,15 @@ class PartitionedFileReader {
private void moveToNextReadablePosition(ByteBuffer indexEntryBuf) throws IOException {
while (currentRegionRemainingBytes <= 0 && hasNextPositionToRead()) {
- if (!offsetAndSizesToRead.isEmpty()) {
- Tuple2<Long, Long> offsetAndSize = offsetAndSizesToRead.poll();
- nextOffsetToRead = offsetAndSize.f0;
- currentRegionRemainingBytes = offsetAndSize.f1;
+ if (!readBufferPositions.isEmpty()) {
+ BufferPositionDescriptor descriptor = readBufferPositions.poll();
+ nextOffsetToRead = descriptor.offset;
+ currentRegionRemainingBytes = descriptor.size;
+ currentBufferPositionDescriptor = descriptor;
} else {
// move to next region which has buffers
if (nextRegionToRead < partitionedFile.getNumRegions()) {
- updateReadableOffsetAndSize(indexEntryBuf, offsetAndSizesToRead);
+ updateReadableOffsetAndSize(indexEntryBuf, readBufferPositions);
++nextRegionToRead;
}
}
@@ -115,8 +118,7 @@ class PartitionedFileReader {
}
private boolean hasNextPositionToRead() {
- return !offsetAndSizesToRead.isEmpty()
- || nextRegionToRead < partitionedFile.getNumRegions();
+ return !readBufferPositions.isEmpty() || nextRegionToRead < partitionedFile.getNumRegions();
}
/**
@@ -144,12 +146,12 @@ class PartitionedFileReader {
*
* @param indexEntryBuf A ByteBuffer containing index entries which provide offset and size
* information.
- * @param offsetAndSizesToRead A queue to store the updated offsets and sizes.
+ * @param readBufferPositions A queue to store the buffer position descriptors.
* @throws IOException If an I/O error occurs when accessing the index file channel.
*/
@VisibleForTesting
void updateReadableOffsetAndSize(
- ByteBuffer indexEntryBuf, Queue<Tuple2<Long, Long>> offsetAndSizesToRead)
+ ByteBuffer indexEntryBuf, Queue<BufferPositionDescriptor> readBufferPositions)
throws IOException {
int startSubpartition = subpartitionIndexSet.getStartIndex();
int endSubpartition = subpartitionIndexSet.getEndIndex();
@@ -157,18 +159,18 @@ class PartitionedFileReader {
if (startSubpartition >= subpartitionOrderRotationIndex
|| endSubpartition < subpartitionOrderRotationIndex) {
updateReadableOffsetAndSize(
- startSubpartition, endSubpartition, indexEntryBuf, offsetAndSizesToRead);
+ startSubpartition, endSubpartition, indexEntryBuf, readBufferPositions);
} else {
updateReadableOffsetAndSize(
subpartitionOrderRotationIndex,
endSubpartition,
indexEntryBuf,
- offsetAndSizesToRead);
+ readBufferPositions);
updateReadableOffsetAndSize(
startSubpartition,
subpartitionOrderRotationIndex - 1,
indexEntryBuf,
- offsetAndSizesToRead);
+ readBufferPositions);
}
}
@@ -181,7 +183,7 @@ class PartitionedFileReader {
* @param startSubpartition The starting index of the subpartition range to be processed.
* @param endSubpartition The ending index of the subpartition range to be processed.
* @param indexEntryBuf A ByteBuffer containing the index entries to read offsets and sizes.
- * @param offsetAndSizesToRead A queue to store the updated offsets and sizes.
+ * @param readBufferPositions A queue to store the buffer position descriptors.
* @throws IOException If an I/O error occurs during reading of index entries.
* @throws IllegalStateException If offsets are not contiguous and not from a single buffer.
*/
@@ -189,7 +191,7 @@ class PartitionedFileReader {
int startSubpartition,
int endSubpartition,
ByteBuffer indexEntryBuf,
- Queue<Tuple2<Long, Long>> offsetAndSizesToRead)
+ Queue<BufferPositionDescriptor> readBufferPositions)
throws IOException {
partitionedFile.getIndexEntry(
indexFileChannel, indexEntryBuf, nextRegionToRead, startSubpartition);
@@ -202,18 +204,28 @@ class PartitionedFileReader {
long endPartitionSize = indexEntryBuf.getLong();
if (startPartitionOffset != endPartitionOffset || startPartitionSize != endPartitionSize) {
- offsetAndSizesToRead.add(
- Tuple2.of(
+ readBufferPositions.add(
+ new BufferPositionDescriptor(
startPartitionOffset,
- endPartitionOffset + endPartitionSize - startPartitionOffset));
+ endPartitionOffset + endPartitionSize - startPartitionOffset,
+ 1));
} else if (startPartitionSize != 0) {
// this branch is for broadcast subpartitions
- for (int i = startSubpartition; i <= endSubpartition; i++) {
- offsetAndSizesToRead.add(Tuple2.of(startPartitionOffset, startPartitionSize));
- }
+ readBufferPositions.add(
+ new BufferPositionDescriptor(
+ startPartitionOffset,
+ startPartitionSize,
+ endSubpartition - startSubpartition + 1));
}
}
+ @VisibleForTesting
+ void readCurrentRegion(
+ Queue<MemorySegment> freeSegments, BufferRecycler recycler, Consumer<Buffer> consumer)
+ throws IOException {
+ readCurrentRegion(freeSegments, recycler, (buffer, repeatCount) -> consumer.accept(buffer));
+ }
+
/**
* Reads a buffer from the current region of the target {@link PartitionedFile} and moves the
* read position forward.
@@ -226,7 +238,9 @@ class PartitionedFileReader {
* @return Whether the file reader has remaining data to read.
*/
boolean readCurrentRegion(
- Queue<MemorySegment> freeSegments, BufferRecycler recycler, Consumer<Buffer> consumer)
+ Queue<MemorySegment> freeSegments,
+ BufferRecycler recycler,
+ BiConsumer<Buffer, Integer> consumer)
throws IOException {
if (currentRegionRemainingBytes == 0) {
return false;
@@ -300,7 +314,7 @@ class PartitionedFileReader {
ByteBuffer byteBuffer,
Buffer buffer,
BufferAndHeader partialBuffer,
- Consumer<Buffer> consumer) {
+ BiConsumer<Buffer, Integer> consumer) {
BufferHeader header = partialBuffer.header;
CompositeBuffer targetBuffer = partialBuffer.buffer;
while (byteBuffer.hasRemaining()) {
@@ -331,7 +345,7 @@ class PartitionedFileReader {
}
header = null;
- consumer.accept(targetBuffer);
+ consumer.accept(targetBuffer, currentBufferPositionDescriptor.repeatCount);
targetBuffer = null;
}
return new BufferAndHeader(targetBuffer, header);
@@ -366,4 +380,44 @@ class PartitionedFileReader {
this.header = header;
}
}
+
+ /**
+ * Represents the position and size of a buffer along with the repeat count. For a regular
+ * buffer, the repeat count is typically one. For a broadcast buffer, the repeat count
+ * corresponds to the number of subpartitions.
+ */
+ @VisibleForTesting
+ static class BufferPositionDescriptor {
+ private final long offset;
+ private final long size;
+ private final int repeatCount;
+
+ /**
+ * Constructs a BufferPositionDescriptor with specified offset, size, and repeat count.
+ *
+ * @param offset the offset of the buffer
+ * @param size the size of the buffer
+ * @param repeatCount the repeat count for the buffer
+ */
+ BufferPositionDescriptor(long offset, long size, int repeatCount) {
+ this.offset = offset;
+ this.size = size;
+ this.repeatCount = repeatCount;
+ }
+
+ @VisibleForTesting
+ long getOffset() {
+ return offset;
+ }
+
+ @VisibleForTesting
+ long getSize() {
+ return size;
+ }
+
+ @VisibleForTesting
+ int getRepeatCount() {
+ return repeatCount;
+ }
+ }
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeSubpartitionReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeSubpartitionReader.java
index af0e7cec906..9f68ad942ba 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeSubpartitionReader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeSubpartitionReader.java
@@ -111,14 +111,14 @@ class SortMergeSubpartitionReader
}
}
- private void addBuffer(Buffer buffer) {
+ private void addBuffer(Buffer buffer, int repeatCount) {
boolean needRecycleBuffer = false;
synchronized (lock) {
if (isReleased) {
needRecycleBuffer = true;
} else {
- addBufferToFullyFilledBuffer(buffer);
+ addBufferToFullyFilledBuffer(buffer, repeatCount);
}
}
@@ -128,6 +128,14 @@ class SortMergeSubpartitionReader
}
}
+ private void addBufferToFullyFilledBuffer(Buffer buffer, int repeatCount) {
+ for (int i = 0; i < repeatCount; i++) {
+ addBufferToFullyFilledBuffer(buffer);
+ buffer.retainBuffer();
+ }
+ buffer.recycleBuffer();
+ }
+
private void addBufferToFullyFilledBuffer(Buffer buffer) {
if (toFilledBuffer == null) {
toFilledBuffer =
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionedFileWriteReadTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionedFileWriteReadTest.java
index 587b27727d6..1a61a8217f1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionedFileWriteReadTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionedFileWriteReadTest.java
@@ -212,15 +212,23 @@ class PartitionedFileWriteReadTest {
createAndConfigIndexEntryBuffer(),
subpartitionOrderRotationIndex);
- Queue<Tuple2<Long, Long>> offsetAndSizesToRead = new ArrayDeque<>();
+ Queue<PartitionedFileReader.BufferPositionDescriptor> offsetAndSizesToRead =
+ new ArrayDeque<>();
fileReader.updateReadableOffsetAndSize(
createAndConfigIndexEntryBuffer(), offsetAndSizesToRead);
if (isBroadcastRegion) {
- assertThat(offsetAndSizesToRead).hasSize(end - start + 1);
- for (Tuple2<Long, Long> tuple2 : offsetAndSizesToRead) {
- assertThat(tuple2.f0).isEqualTo(regionStat[start].get(0).f0);
- assertThat(tuple2.f1).isEqualTo(regionStat[start].get(0).f1);
+ assertThat(
+ offsetAndSizesToRead.stream()
+ .map(
+ PartitionedFileReader.BufferPositionDescriptor
+ ::getRepeatCount)
+ .reduce(Integer::sum)
+ .get())
+ .isEqualTo(end - start + 1);
+ for (PartitionedFileReader.BufferPositionDescriptor descriptor : offsetAndSizesToRead) {
+ assertThat(descriptor.getOffset()).isEqualTo(regionStat[start].get(0).f0);
+ assertThat(descriptor.getSize()).isEqualTo(regionStat[start].get(0).f1);
}
return;
}
@@ -228,34 +236,36 @@ class PartitionedFileWriteReadTest {
if (start >= subpartitionOrderRotationIndex || end <= subpartitionOrderRotationIndex - 1) {
assertThat(offsetAndSizesToRead).hasSize(1);
- Tuple2<Long, Long> offsetAndSize = offsetAndSizesToRead.poll();
- assertThat(offsetAndSize.f0).isEqualTo(regionStat[start].get(0).f0);
+ PartitionedFileReader.BufferPositionDescriptor descriptor = offsetAndSizesToRead.poll();
+ assertThat(descriptor.getOffset()).isEqualTo(regionStat[start].get(0).f0);
long expectedSize = 0L;
for (int i = start; i <= end; i++) {
expectedSize += regionStat[i].get(0).f1;
}
- assertThat(offsetAndSize.f1).isEqualTo(expectedSize);
+ assertThat(descriptor.getSize()).isEqualTo(expectedSize);
} else {
assertThat(offsetAndSizesToRead).hasSize(2);
- Tuple2<Long, Long> offsetAndSize1 = offsetAndSizesToRead.poll();
- Tuple2<Long, Long> offsetAndSize2 = offsetAndSizesToRead.poll();
- assertThat(offsetAndSize1.f0)
+ PartitionedFileReader.BufferPositionDescriptor descriptor1 =
+ offsetAndSizesToRead.poll();
+ PartitionedFileReader.BufferPositionDescriptor descriptor2 =
+ offsetAndSizesToRead.poll();
+ assertThat(descriptor1.getOffset())
.isEqualTo(regionStat[subpartitionOrderRotationIndex].get(0).f0);
- assertThat(offsetAndSize2.f0).isEqualTo(regionStat[start].get(0).f0);
+ assertThat(descriptor2.getOffset()).isEqualTo(regionStat[start].get(0).f0);
long expectedSize = 0L;
for (int i = subpartitionOrderRotationIndex; i <= end; i++) {
expectedSize += regionStat[i].get(0).f1;
}
- assertThat(offsetAndSize1.f1).isEqualTo(expectedSize);
+ assertThat(descriptor1.getSize()).isEqualTo(expectedSize);
expectedSize = 0L;
for (int i = start; i < subpartitionOrderRotationIndex; i++) {
expectedSize += regionStat[i].get(0).f1;
}
- assertThat(offsetAndSize2.f1).isEqualTo(expectedSize);
+ assertThat(descriptor2.getSize()).isEqualTo(expectedSize);
}
}