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

背景

本文基于Spark 3.5.3

Spark Tungsten（钨丝计划）是 Databricks 引入的 Spark SQL 优化引擎。主要致力于提高Spark的执行效率：

它主要从以下三个方面：

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

前两个好理解，如 内存管理优化 主要是通过手动管理内存，使用紧凑二进制的UnsafeRow的数据结构来替换jvm的数据结构
全阶段代码生成 通过消除虚函数的调用(会有二次寻址)，来提高JVM运行效率，

那第三个是什么呢？缓存敏感计算到底指的是什么？这一篇文章就来说清楚。

结论

先说结论：在Sort中 Spark通过将排序键和指向数据的指针存储在一起，主要是先利用排序键做前缀索引，来进行Record的数据比较，如果排序键相同的话，再从内存中获取对应的数据进行比较，这种从内存中获取数据是随机访问，导致CPU的局部缓存命中率低，不像把数据排序的键存在一起，可以高效的利用CPU的局部缓存。

注意这里的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);