Source:
public void testWriteAndReadSortBuffer() throws Exception { int numSubpartitions = 10; int bufferSize = 1024; int bufferPoolSize = 1000; Random random = new Random(); Queue<DataAndType>[] dataWritten = new Queue[numSubpartitions]; Queue<Buffer>[] buffersRead = new Queue[numSubpartitions]; for (int i = 0; i < numSubpartitions; ++i) { dataWritten[i] = new ArrayDeque<>(); buffersRead[i] = new ArrayDeque<>(); } int[] numBytesWritten = new int[numSubpartitions]; int[] numBytesRead = new int[numSubpartitions]; Arrays.fill(numBytesWritten, 0); Arrays.fill(numBytesRead, 0); int totalBytesWritten = 0; SortBuffer sortBuffer = createSortBuffer(bufferPoolSize, bufferSize, numSubpartitions); while (true) { int recordSize = random.nextInt(bufferSize * 4 - 1) + 1; byte[] bytes = new byte[recordSize]; random.nextBytes(bytes); ByteBuffer record = ByteBuffer.wrap(bytes); int subpartition = random.nextInt(numSubpartitions); boolean isBuffer = random.nextBoolean(); Buffer.DataType dataType = isBuffer ? Buffer.DataType.DATA_BUFFER : Buffer.DataType.EVENT_BUFFER; if (!sortBuffer.append(record, subpartition, dataType)) { sortBuffer.finish(); break; } record.rewind(); dataWritten[subpartition].add(new DataAndType(record, dataType)); numBytesWritten[subpartition] += recordSize; totalBytesWritten += recordSize; } while (sortBuffer.hasRemaining()) { MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledSegment(bufferSize); SortBuffer.BufferWithChannel bufferAndChannelIndex = sortBuffer.copyData(readBuffer); int subpartition = bufferAndChannelIndex.getChannelIndex(); buffersRead[subpartition].add(bufferAndChannelIndex.getBuffer()); numBytesRead[subpartition] += bufferAndChannelIndex.getBuffer().readableBytes(); } assertEquals(totalBytesWritten, sortBuffer.numBytes()); checkWriteAndReadResult(numSubpartitions, numBytesWritten, numBytesRead, dataWritten, buffersRead); }

Target:
public void testWriteAndReadSortBuffer() throws Exception { int numSubpartitions = 10; int bufferSize = 1024; int bufferPoolSize = 1000; Random random = new Random(1111); Queue<DataAndType>[] dataWritten = new Queue[numSubpartitions]; Queue<Buffer>[] buffersRead = new Queue[numSubpartitions]; for (int i = 0; i < numSubpartitions; ++i) { dataWritten[i] = new ArrayDeque<>(); buffersRead[i] = new ArrayDeque<>(); } int[] numBytesWritten = new int[numSubpartitions]; int[] numBytesRead = new int[numSubpartitions]; Arrays.fill(numBytesWritten, 0); Arrays.fill(numBytesRead, 0); int totalBytesWritten = 0; SortBuffer sortBuffer = createSortBuffer( bufferPoolSize, bufferSize, numSubpartitions, getRandomSubpartitionOrder(numSubpartitions)); while (true) { int recordSize = random.nextInt(bufferSize * 4 - 1) + 1; byte[] bytes = new byte[recordSize]; random.nextBytes(bytes); ByteBuffer record = ByteBuffer.wrap(bytes); int subpartition = random.nextInt(numSubpartitions); boolean isBuffer = random.nextBoolean(); DataType dataType = isBuffer ? DataType.DATA_BUFFER : DataType.EVENT_BUFFER; if (!sortBuffer.append(record, subpartition, dataType)) { sortBuffer.finish(); break; } record.rewind(); dataWritten[subpartition].add(new DataAndType(record, dataType)); numBytesWritten[subpartition] += recordSize; totalBytesWritten += recordSize; } while (sortBuffer.hasRemaining()) { MemorySegment readBuffer = MemorySegmentFactory.allocateUnpooledSegment(bufferSize); SortBuffer.BufferWithChannel bufferAndChannel = sortBuffer.copyIntoSegment(readBuffer); int subpartition = bufferAndChannel.getChannelIndex(); buffersRead[subpartition].add(bufferAndChannel.getBuffer()); numBytesRead[subpartition] += bufferAndChannel.getBuffer().readableBytes(); } assertEquals(totalBytesWritten, sortBuffer.numBytes()); checkWriteReadResult(numSubpartitions, numBytesWritten, numBytesRead, dataWritten, buffersRead); }