Spark的UnsafeShuffleWriter是Tungsten-Project（内存管理）引入的新的Shuffle Writer。
该writer在写数据到磁盘时，会将数据有序写入（仅仅分区间有序，分区内无序）。 本文主要介绍其写数据实现，并讨论其缓存友好设计及实现。以下介绍其相关实现：
UnsafeShuffleWriter的实现
写数据相关类
各函数类的职责：
UnsafeShuffleWriter ： 数据写的入口，并将最后内存数据以及spill至磁盘的数据读取并以分区merge后再次写入磁盘
ShuffleExternalSorter ：内存申请，管理，写入数据至内存页，代理ShuffleInMemorySorter的写入
ShuffleInMemorySorter ： 将分区号及数据指针写入LongArray
其核心流程主要包括如下：
- 将输入数据（record）遍历
- 将
分区号和数据指针写入ShuffleInMemorySorter的array:LongArray,序列化的数据写入page页 - 如果其中内存不足时，会将序列化数据根据分区号spill至磁盘
- 数据操作完毕后，将当前内存中的数据以及spill至磁盘的数据重新merge，按分区号排序写如磁盘。
UnsafeShuffleWriter 遍历数据
核心代码逻辑如下：
public void write(scala.collection.Iterator> records) throws IOException {
boolean success = false;
try {
while (records.hasNext()) {
// 插入数据
insertRecordIntoSorter(records.next());
}
// 将内存数据及spill之后的数据重新merge并写入磁盘，同时也生成index文件
closeAndWriteOutput();
success = true;
} finally {
...
}
}
其中insertRecordIntoSorter实现如下：
void insertRecordIntoSorter(Product2 record) throws IOException {
assert(sorter != null);
// 将key,value写入serBuffer
final K key = record._1();
final int partitionId = partitioner.getPartition(key);
serBuffer.reset();
serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
serOutputStream.flush();
final int serializedRecordSize = serBuffer.size();
// 调用ShuffleExternalSorter的insertRecord方法插入数据
sorter.insertRecord(
serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);
}
调用ShuffleExternalSorter的insertRecord的方法存入数据及分区号，数据指针，主要包含以下步骤：
- 如果空间不足，将触发spill落盘
- 申请LongArray用于存储partitionId和指针
- 当前page也不足以缓存数据时，申请新page页
- 将数据存入缓存页
- 缓存指针和partitionId
public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
throws IOException {
if (inMemSorter.numRecords() >= numElementsForSpillThreshold) {
logger.info("Spilling data because number of spilledRecords crossed the threshold " +
numElementsForSpillThreshold);
// 1. 如果存储空间不足，则触发spill操作
spill();
}
// 2. 申请LongArray用于存储partitionId和指针
growPointerArrayIfNecessary();
// Need 4 bytes to store the record length.
final int required = length + 4;
// 3. 为缓存数据 获取申请新page页
acquireNewPageIfNecessary(required);
final Object base = currentPage.getBaseObject();
final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
Platform.putInt(base, pageCursor, length);
pageCursor += 4;
// 4. 将数据存入page中
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
pageCursor += length;
// 5. 将数据指针和partitionId写入inMemSorter中的Array中
inMemSorter.insertRecord(recordAddress, partitionId);
}
spill数据
核心调用关系：ShuffleExternalSorter.spill -> ShuffleExternalSorter.writeSortedFile(false).
其核心实现逻辑是：
- 根据分区id对inMemSort中的数据LongArray的进行排序
- 根据排序好的LongArrary中数据找出对应的真实数据对应的page页及在该页的位置，也就是找出真实数据
- 将数据落盘
writeSortedFile 方法实现如下：
private void writeSortedFile(boolean isLastFile) {
//...
// This call performs the actual sort.
// 将数据跟分区号排序
// This call performs the actual sort.
final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords =
inMemSorter.getSortedIterator();
// 用于存放spill信息
final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
// 创建writer
final DiskBlockObjectWriter writer =
blockManager.getDiskWriter(blockId, file, ser, fileBufferSizeBytes, writeMetricsToUse);
int currentPartition = -1;
final int uaoSize = UnsafeAlignedOffset.getUaoSize();
while (sortedRecords.hasNext()) {
sortedRecords.loadNext();
// 计算数据分区号
final int partition = sortedRecords.packedRecordPointer.getPartitionId();
assert (partition >= currentPartition);
if (partition != currentPartition) {
// Switch to the new partition
if (currentPartition != -1) {
final FileSegment fileSegment = writer.commitAndGet();
spillInfo.partitionLengths[currentPartition] = fileSegment.length();
}
currentPartition = partition;
}
//通过指针（page页及offset）计算数据地址
final long recordPointer = sortedRecords.packedRecordPointer.getRecordPointer();
final Object recordPage = taskMemoryManager.getPage(recordPointer);
final long recordOffsetInPage = taskMemoryManager.getOffsetInPage(recordPointer);
int dataRemaining = UnsafeAlignedOffset.getSize(recordPage, recordOffsetInPage);
long recordReadPosition = recordOffsetInPage + uaoSize; // skip over record length
//// 将数据通过writer写出, 由于可能使用zero copy技术，每次数据写入字节数有限制，因此此处使用循环调用
while (dataRemaining > 0) {
final int toTransfer = Math.min(diskWriteBufferSize, dataRemaining);
Platform.copyMemory(
recordPage, recordReadPosition, writeBuffer, Platform.BYTE_ARRAY_OFFSET, toTransfer);
writer.write(writeBuffer, 0, toTransfer);
recordReadPosition += toTransfer;
dataRemaining -= toTransfer;
}
writer.recordWritten();
}
final FileSegment committedSegment = writer.commitAndGet();
writer.close();
if (currentPartition != -1) {
spillInfo.partitionLengths[currentPartition] = committedSegment.length();
spills.add(spillInfo);
}
...
}
为LongArray申请内存空间
当LongArray剩余空间不足以满足新增数据时，会对LongArary扩容，扩容会会将原有LongArray中的数据copy至新的Array中，并释放源有LongArray
private void growPointerArrayIfNecessary() throws IOException {
assert(inMemSorter != null);
if (!inMemSorter.hasSpaceForAnotherRecord()) {
long used = inMemSorter.getMemoryUsage();
LongArray array;
try {
// could trigger spilling // 申请内存空间
array = allocateArray(used / 8 * 2);
} catch (OutOfMemoryError e) { //...
}
return;
}
// check if spilling is triggered or not
if (inMemSorter.hasSpaceForAnotherRecord()) {
freeArray(array);
} else {
// 完成扩容后的老数据迁移
inMemSorter.expandPointerArray(array);
}
}
}
申请内存页用户存放真实数据
如果currentPage为空或者currentPage不足以存放数据时，则重新申请缓存page，并将其设置为currentPage. 内存页的申请，我们后续单独讨论
private void acquireNewPageIfNecessary(int required) {
if (currentPage == null ||
pageCursor + required > currentPage.getBaseOffset() + currentPage.size() ) {
// TODO: try to find space in previous pages
currentPage = allocatePage(required);
pageCursor = currentPage.getBaseOffset();
allocatedPages.add(currentPage);
}
}
数据写入page页
- 先将数据长度写入page页
- 将真实数据写入page页
final Object base = currentPage.getBaseObject();
// 先放入长度，占用4个字节
Platform.putInt(base, pageCursor, length);
pageCursor += 4;
//放入真实的数据, 占用length个字段
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
pageCursor += length;
将分区id,page页数,page页的offset写入LongArray
调用inMemSorter.insertRecord(recordAddress, partitionId)方法完成，本质是放入LongArray数组中
public void insertRecord(long recordPointer, int partitionId) {
if (!hasSpaceForAnotherRecord()) {
throw new IllegalStateException("There is no space for new record");
}
array.set(pos, PackedRecordPointer.packPointer(recordPointer, partitionId));
pos++;
}
缓存友好/感知及其实现
CPU从高速缓存中读取数据的效率与从内存中读取数据的效率相差较大。从高速缓存读取数据的性能是内存中读取数据的几十倍上百倍。缓存友好/感知的含义是指CPU访问数据时，更多/更高概率的从高速缓存中（L1，L2,L3级缓存效率依次递减）命中数据，此时CPU的数据读取性能更好。在大批量数据处理时，如果能够提升高速缓存命中的概率，则可以显著提升应用运行性能。
在Spark的Tungsten-Project的内存管理中就对提升数据的缓存做了优化，极大提升了在高速缓存中被命中的概率。通常情况下，计算机的缓存空间较小，如果将整条数据都进行缓存，则有缓存数据的记录数就很少。因此Spark采用只缓存记录的部分关键字段/信息的方式。
缓存友好的实现
Spark在内存中inMemSort中使用一个LongArray来保存每条每条记录的partitionId和该记录的指针，其中partitionId用于对记录进行分区排序，记录的指针用于指向数据的地址，通过指针可以访问到该数据。因此只需要在LongArray中放入一个long型数据，即可同时存放partitionId和数据地址
LongArray设计
LongArray中存储的数据的结构如下：
从上图可以看出，有24个bit位用于存放分区号，这也是为什么使用UnsafeShuffleWriter需要保证分区数小于2^24
从上图可以看出page页使用13个bit位表示因此可以推算出 page页不能多于2^13
从上图可以offset使用的是27个bit位，也就是单个page也不能存储超过其所能存储的数据
在将数据spill磁盘时，我们只需要对LongArray中的记录进行排序（每个记录中都包含该一个真实数据的partitionId和指向该数据的地址指针），有序（基于分区号有序）的LongArray的顺序访问即可保证读取的page也中的数据是有序的。
基于分区的排序
有了上图中的设计之后，我们就可以做到：
无序访问真实的数据，仅仅通过访问LongArray中的数据即可完成基于分区的排序
在高速缓存中不用缓存真实数据（一般远远大于LongArray中存储的一个long型数据），仅仅缓存LongArray（包括分区号，真实数据的地址指针），极大提升了高速缓存命中率，加速了排序过程
从以上分析可以看出，此处的实现结合了计算机硬件的设计，高效利用硬件，完成了程序的性能提升。如果没有该设计，那么在对数据spill的过程中，在数据量较大的情况下，由于无法对大量数据进行高速缓存，与以上方案相比，性能应当有明显的降低。