Skip to content

[ISSUE-301][Subtask][Improvement][AQE] Merge continuous ShuffleDataSegment into single one #303

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
roryqi merged 1 commit into from
Nov 7, 2022
Merged

[ISSUE-301][Subtask][Improvement][AQE] Merge continuous ShuffleDataSegment into single one #303

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions common/src/main/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) {
long totalLen = 0;

long lastTaskAttemptId = -1;
long lastExpectedBlockIndex = -1;

/**
* One ShuffleDataSegment should meet following requirements:
Expand All @@ -79,6 +80,7 @@ public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) {
* 3. ShuffleDataSegment's blocks should be continuous
*
*/
int index = 0;
while (byteBuffer.hasRemaining()) {
try {
long offset = byteBuffer.getLong();
Expand All @@ -98,7 +100,8 @@ public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) {
break;
}

if ((taskAttemptId < lastTaskAttemptId && bufferSegments.size() > 0) || bufferOffset >= readBufferSize) {
if ((taskAttemptId < lastTaskAttemptId && bufferSegments.size() > 0 && index - lastExpectedBlockIndex != 1)
|| bufferOffset >= readBufferSize) {
ShuffleDataSegment sds = new ShuffleDataSegment(fileOffset, bufferOffset, bufferSegments);
dataFileSegments.add(sds);
bufferSegments = Lists.newArrayList();
Expand All @@ -107,14 +110,20 @@ public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) {
}

if (expectTaskIds.contains(taskAttemptId)) {
if (bufferOffset != 0 && index - lastExpectedBlockIndex > 1) {
throw new RssException("There are discontinuous blocks which should not happen when using LOCAL_ORDER.");
}

if (fileOffset == -1) {
fileOffset = offset;
}
bufferSegments.add(new BufferSegment(blockId, bufferOffset, length, uncompressLength, crc, taskAttemptId));
bufferOffset += length;
lastExpectedBlockIndex = index;
}

lastTaskAttemptId = taskAttemptId;
index++;
} catch (BufferUnderflowException ue) {
throw new RssException("Read index data under flow", ue);
}
Expand All @@ -124,7 +133,6 @@ public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) {
ShuffleDataSegment sds = new ShuffleDataSegment(fileOffset, bufferOffset, bufferSegments);
dataFileSegments.add(sds);
}

return dataFileSegments;
}
}
88 changes: 88 additions & 0 deletions common/src/test/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,102 @@
import org.junit.jupiter.api.Test;
import org.roaringbitmap.longlong.Roaring64NavigableMap;

import org.apache.uniffle.common.BufferSegment;
import org.apache.uniffle.common.ShuffleDataSegment;
import org.apache.uniffle.common.ShuffleIndexResult;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

public class LocalOrderSegmentSplitterTest {

@Test
public void testSplitWithDiscontinuousBlocksShouldThrowException() {
Roaring64NavigableMap taskIds = Roaring64NavigableMap.bitmapOf(1, 2, 4);
LocalOrderSegmentSplitter splitter = new LocalOrderSegmentSplitter(taskIds, 32);
byte[] data = generateData(
Pair.of(1, 1),
Pair.of(1, 2),
Pair.of(1, 3),
Pair.of(1, 4)
);
try {
splitter.split(new ShuffleIndexResult(data, -1));
fail();
} catch (Exception e) {
// ignore
}
}

@Test
public void testSplitForMergeContinuousSegments() {
/**
* case1: (32, 5) (16, 1) (10, 1) (16, 2) (6, 1) (8, 1) (10, 3) (9, 1)
*
* It will skip the (32, 5) and merge others into one dataSegment when no exceeding the
* read buffer size.
*/
Roaring64NavigableMap taskIds = Roaring64NavigableMap.bitmapOf(1, 2);
LocalOrderSegmentSplitter splitter = new LocalOrderSegmentSplitter(taskIds, 1000);
byte[] data = generateData(
Pair.of(32, 5),
Pair.of(16, 1),
Pair.of(10, 1),
Pair.of(16, 2),
Pair.of(6, 1),
Pair.of(8, 1),
Pair.of(10, 3),
Pair.of(9, 1)
);
List<ShuffleDataSegment> dataSegments = splitter.split(new ShuffleIndexResult(data, -1));
assertEquals(2, dataSegments.size());
assertEquals(32, dataSegments.get(0).getOffset());
assertEquals(56, dataSegments.get(0).getLength());

List<BufferSegment> bufferSegments = dataSegments.get(0).getBufferSegments();
assertEquals(0, bufferSegments.get(0).getOffset());
assertEquals(16, bufferSegments.get(0).getLength());

assertEquals(16, bufferSegments.get(1).getOffset());
assertEquals(10, bufferSegments.get(1).getLength());

assertEquals(26, bufferSegments.get(2).getOffset());
assertEquals(16, bufferSegments.get(2).getLength());

assertEquals(42, bufferSegments.get(3).getOffset());
assertEquals(6, bufferSegments.get(3).getLength());

assertEquals(48, bufferSegments.get(4).getOffset());
assertEquals(8, bufferSegments.get(4).getLength());

assertEquals(98, dataSegments.get(1).getOffset());
assertEquals(9, dataSegments.get(1).getLength());
bufferSegments = dataSegments.get(1).getBufferSegments();
assertEquals(1, bufferSegments.size());
assertEquals(0, bufferSegments.get(0).getOffset());
assertEquals(9, bufferSegments.get(0).getLength());

/**
* case2: (16, 1) (16, 2) (6, 1)
*
* It will skip merging into one dataSegment when exceeding the
* read buffer size.
*/
data = generateData(
Pair.of(16, 1),
Pair.of(15, 2),
Pair.of(1, 1),
Pair.of(6, 1)
);
dataSegments = new LocalOrderSegmentSplitter(taskIds, 32).split(new ShuffleIndexResult(data, -1));
assertEquals(2, dataSegments.size());
assertEquals(0, dataSegments.get(0).getOffset());
assertEquals(32, dataSegments.get(0).getLength());
assertEquals(32, dataSegments.get(1).getOffset());
assertEquals(6, dataSegments.get(1).getLength());
}

@Test
public void testSplit() {
Roaring64NavigableMap taskIds = Roaring64NavigableMap.bitmapOf(1);
Expand Down