从 SortExec 的排序来谈 Spark Tungsten 计划中的缓存友好特性

背景

本文基于Spark 3.5.3
Spark Tungsten（钨丝计划）是 Databricks 引入的 Spark SQL 优化引擎。主要致力于提高Spark的执行效率：
它主要从以下三个方面：

1. 内存管理优化
2. 全阶段代码生成
3. 缓存敏感计算

前两个好理解，如 内存管理优化 主要是通过手动管理内存，使用紧凑二进制的UnsafeRow的数据结构来替换jvm的数据结构
全阶段代码生成 通过消除虚函数的调用(会有二次寻址)，来提高JVM运行效率，
那第三个是什么呢？缓存敏感计算到底指的是什么？这一篇文章就来说清楚。

结论

先说结论：在Sort中 Spark通过将排序键和指向数据的指针存储在一起，主要是先利用排序键做前缀索引，来进行Record的数据比较，如果排序键相同的话，再从内存中获取对应的数据进行比较，这种从内存中获取数据是随机访问，导致CPU的局部缓存命中率低，不像把数据排序的键存在一起，可以高效的利用CPU的局部缓存。
支持前缀比较器的类型有: BooleanType | ByteType | ShortType | IntegerType | LongType | DateType |TimestampType | TimestampNTZType | FloatType | DoubleType | _: AnsiIntervalType｜ DecimalType 且 精度位数小于 18
注意这里的StringType不支持前缀比较器,因为String类型的长度是不定的，不能用long类型来表示，从而无法进行比较。

分析

先看SortExec涉及到排序的数据流：

SortExec.doExecute
   ||
   \/
SortExec.createSorter
   ||
   \/
insertRow  ->   UnsafeExternalRowSorter.sort
                 ||
                 \/
                UnsafeExternalSorter.getSortedIterator
                 ||
                 \/
                UnsafeInMemorySorter.getSortedIterator

这里的UnsafeInMemorySorter包括了 插入数据的recordPointer + keyPrefix ；
注意：这里的数据保存在LongArray中（从taskMemoryManager分配过来的），这里的LongArray以8字节作为一个步长，用Unsafe.putLong和setLong进行操作，也就是存储8字节的数据
这里的 recordPointer 来源于 memory page以及page内的offset

 private MemoryBlock currentPage;
 ...
 recordPointer = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor)

keyPrefix 来源于SortPrefixUtils.getPrefixComparator的computePrefix的方法，这里会根据不同的类型返回不同的前缀比较器。
SortPrefixUtils.getPrefixComparator的代码如下：

def getPrefixComparator(sortOrder: SortOrder): PrefixComparator = {
    sortOrder.dataType match {
      case StringType => stringPrefixComparator(sortOrder)
      case BinaryType => binaryPrefixComparator(sortOrder)
      case BooleanType | ByteType | ShortType | IntegerType | LongType | DateType | TimestampType |
          TimestampNTZType | _: AnsiIntervalType =>
        longPrefixComparator(sortOrder)
      case dt: DecimalType if dt.precision - dt.scale <= Decimal.MAX_LONG_DIGITS =>
        longPrefixComparator(sortOrder)
      case FloatType | DoubleType => doublePrefixComparator(sortOrder)
      case dt: DecimalType => doublePrefixComparator(sortOrder)
      case _ => NoOpPrefixComparator
    }
  }

UnsafeInMemorySorter.getSortedIterator的流程如下:

                ┌─────────────────────────────────────┐
                │  UnsafeInMemorySorter.getSortedIterator()  │
                └─────────────────────────────────────┘
                                    │
                                    ▼
                ┌─────────────────────────────────────┐
                │    prefixComparator 支持 Radix Sort？│
                └─────────────────────────────────────┘
                        │                      │
                       Yes                     No
                        │                      │
                        ▼                      ▼
            ┌───────────────────────┐  ┌───────────────────────┐
            │    RadixSort          │  │    TimSort (Sorter)   │
            ├───────────────────────┤  ├───────────────────────┤
            │ • 直接排序 key prefix  │  │ • 需要比较器           │
            │ • O(n * 8) 复杂度      │  │ • O(n log n) 复杂度   │
            │                       │  │ • 需要完整记录比较      │
            │ • 缓存友好             │  │ • 可能缓存未命中        │
            └───────────────────────┘  └───────────────────────┘
                          │                      │
                          └──────────┬───────────┘
                                     │
                                     ▼
                    ┌─────────────────────────────────────┐
                    │      返回 UnsafeSorterIterator      │
                    └─────────────────────────────────────┘ 

这里会根据 prefixComparator 是否支持Radix sort（基数排序）来判断，是直接走 Radix sort还是 TimSort
如果是Radix sort,则直接根据前缀比较器进行比较；
如果是TimSort的话，则先直接根据前缀比较器进行比较，如果相等，则在根据pointer获取数据进行比较。

注意无论是 Radix sort 还是 TimSort ，如果只是用到前缀比较器的话（没有用到pointer获取数据）,因为排序键的前缀都存在一个LongArray类型中，这样CPU需要数据的时候，会以数据块为单位从内存加载到缓存，根据局部性的原理（若一数据被访问，邻近数据也大概率被访问），相邻的排序键的前缀也都会加载到缓存中，这样CPU的缓存命中率就会大大提升：

 offset = RadixSort.sortKeyPrefixArray(
          array, nullBoundaryPos, (pos - nullBoundaryPos) / 2L, 0, 7,
          radixSortSupport.sortDescending(), radixSortSupport.sortSigned());
 ...
 MemoryBlock unused = new MemoryBlock(
          array.getBaseObject(),
          array.getBaseOffset() + pos * 8L,
          (array.size() - pos) * 8L);
 LongArray buffer = new LongArray(unused);
 Sorter<RecordPointerAndKeyPrefix, LongArray> sorter =
   new Sorter<>(new UnsafeSortDataFormat(buffer));
 sorter.sort(array, 0, pos / 2, sortComparator);