深入理解SPARK SQL 中HashAggregateExec和ObjectHashAggregateExec以及UnsafeRow

Posted 2022-03-28 鸿乃江边鸟

感悟和理解

翻译这篇文章是为了更好的理解ObjectHashAggregateExec和UnsafeRow，关于UnsafeRow的文章，可以参考Spark源码修改系列 - UnsafeRow内存布局和代码优化。
其实明显，UnsafeRow是模拟了offheap的分配方式而在堆上进行的操作，它并没有脱离了jvm，而是使用了Java unsafe API（Java提供的底层非安全操作接口），允许用户像操作C++指针一样来申请和访问内存地址，JVM也不会对这部分非安全数据进行额外的类型检查和垃圾回收。
事实上用户可以选择是在jvm堆内还是堆外进行Unsafe Api内存分配操作，默认情况下是堆内。这在MemoryManager代码中也有体现：

/**
   * Tracks whether Tungsten memory will be allocated on the JVM heap or off-heap using
   * sun.misc.Unsafe.
   */
  final val tungstenMemoryMode: MemoryMode = 
    if (conf.get(MEMORY_OFFHEAP_ENABLED)) 
      require(conf.get(MEMORY_OFFHEAP_SIZE) > 0,
        "spark.memory.offHeap.size must be > 0 when spark.memory.offHeap.enabled == true")
      require(Platform.unaligned(),
        "No support for unaligned Unsafe. Set spark.memory.offHeap.enabled to false.")
      MemoryMode.OFF_HEAP
     else 
      MemoryMode.ON_HEAP
    
  

这个博客继续去探索Aggerate 策略，我们主要关注SPARK SQL提供的两种hash join的聚合操作，HashAggerateExec和ObjectHashAggerateExec。

ObjectHashAggregateExec

基于hash的聚合是优于基于排序的聚合，基于排序的聚合我们上篇文章中已经说到。相对于排序聚合，基于hash的聚合不需要在聚合前进行额外的排序步骤。对于HashAggerateExec操作，使用offheap的的内存方式存储聚合buffer能够通过减少GC来进一步的提高性能。

HashAggerateExec

当aggregateExpressions（从聚合逻辑计划中获取的）中所有的aggBufferAttributes属性都是可变的数据类型，Spark就会选择HashAggerateExec。从更高层次上来说，HashAggregateExec使用了名为UnsafeFixedWidthAggregationMap的offheap hashmap来存储对应的组和对应的聚合buffer。当hashmap变得很大并且没有内存可以从memory Manager分配的时候，这个hashMap将会溢写到磁盘，并且会新建一个hashMap来处理剩下的数据。当所有的输入数据都被处理完了以后，所有的溢出文件都会被合并，并且基于排序的聚合将会用来计算最终的结果。
当HashAggregateExec操作执行的时候，它会为每个分区创建一个TungstenAggregationIterator实例。这个TungstenAggregationIterator内部封装了基于hash聚合的核心操作，如：buffer溢写，回退到基于排序的聚合。

TungstenAggregationIterator 内部有个UnsafeFixedWidthAggregationMap实例，它用来存储所有的组和对应的中间聚合结果buffer。UnsafeFixedWidthAggregationMap内部会创建BytesToBytesMap实例，它是用来保存hashmap key-value对的数据结构，key和value的值被存储在内存中，如下所示：
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-svLNm9eR-1648284846269)(https://malinxiao.files.wordpress.com/2022/01/image-21.png)]
UnsafeFixedWithAggregationMap 会以UnsafeRow格式来对key组和聚合buffer编码。BytesToBytesMap一个实体保存一个对应的key组和聚合buffer对。UnsafeFixedWidthAggregationMap的getAggregationBuffer方法被key组调用后会返回对应的聚合buffer（如果key组存在的话）。如果key组不存在则会首先添加对应的key组和聚合buffer，然后再返回刚创建的聚合buffer。
当每个分区的TungstenAggregationIterator实例创建出来以后，包含输入row的迭代器就会传到TungstenAggregationIterator实例中。TungstenAggregationIterator就会调用processInputs方法开始处理输入的rows，同时回退阈值（默认Int.MaxValue，也就是2,147,483,647）也会传到processInputs方法中,用来甄别是不是应该回退到基于排序的聚合中.

 /**
   * Start processing input rows.
   */
  processInputs(testFallbackStartsAt.getOrElse((Int.MaxValue, Int.MaxValue)))

不像是基于排序的聚合，基于hash的聚合和输入rows不是有序的。processInputs方法一行一行的读取并且处理数据直到最后一行数据。当处理每一行输入数据的时候，key组首先被编码为UnsafeRow格式（这个是用来查询对应的聚合buffer），假如key组不在hashmap中，这个key组和一个空聚合buffer将会添加到hashmap中。processRow方法用来更新更新对应的buffer值。

如果对应的key组在hashmap中，聚合buffer将会被更新。

每处理完一行，行计数就会和回退阈值进行比对，如果已经到达了阈值或者没有可用的内存了，destructAndCreateExternalSorter方法将会被调用，它将会排序已有的key组，并且溢写到磁盘，并且返回一个UnsafeKVExternalSorter，这个UnsafeKVExternalSorter保存了在磁盘上的spill的信息。一个新的hashmap将会处理剩下的数据，如果又有溢出，新的UnsafeKVExternalSorter将会合并到已有的UnsafeKVExternalSorter中。

当最后一行数据处理完后，如果还是有溢出，还是会按照之前的方式处理。排序了的合并完的UnsafeKVExternalSorter将会用作基于排序的聚合的输入。

ObjectHashAggregateExec

虽然HashAggregateExec（基于Tungsten执行引擎）对于聚合操作比较适合，但是它只支持固定长度的可变的原生数据类型。对于用户自定义的聚合函数（UDAFs）和一些collect函数（如collect_list和collect_set）,HashAggregateExec是不支持的，在spark 2.2.0之前，我们必须回退到性能差一点的SortAggregateExec，但spark 2.2.0以后ObjectHashAggregateExec弥补了这个空白，它可以对 HashAggregateExec 不支持的数据类型进行基于哈希的高性能聚合。
不像HashAggregateExec存储聚合buffer在堆外内存的UnsafeRow中，ObjectHashAggregateExec存储在SpecificInternalRow中，它在 Java堆内存中内部保存聚合缓冲区字段的Java Array集合。ObjectHashAggregateExec使用ObjectAggregationMap实例来替代UnsafeFixedWidthAggregationMap实例，ObjectAggregationMap实例支持存储二进制java对象作为聚合buffer。

ObjectHashAggregateExec执行流程和HashAggregateExec流程一样。将使用基于散列的聚合从头到尾一一读取和处理输入行。当哈希映射变得太大时，按组key对哈希映射进行排序并将其溢出到磁盘。处理完所有输入行后，如果发生任何溢出，则回退到基于排序的聚合。唯一的区别是 ObjectHashAggregateExec 的回退阈值以不同的方式定义，它是根据hashmap中的key数，而不是处理的输入行数。 ObjectHashAggregateExec 的阈值可以通过spark.sql.objectHashAggregate.sortBased.fallbackThreshold 属性进行配置，默认设置为 128。

本文翻译自Spark SQL Query Engine Deep Dive (10) – HashAggregateExec & ObjectHashAggregateExec