Spark Aggregations execution
Spark Aggregations execution
一、Distinct aggregation 算法
包含 distinct 关键字的 aggregation 由 4 个物理执行步骤组成。我们使用以下 query 来介绍：
val dataset = Seq(
(1, "a"), (1, "a"), (1, "a"), (2, "b"), (2, "b"), (3, "c"), (3, "c")
).toDF("nr", "letter")
dataset.groupBy($"nr").agg(functions.countDistinct("letter")).explain(true)① partial aggregation 步骤
第一步是创建一个 partial aggregate，此 partial aggregate 的 grouping key 将不仅包括 query 中定义的 grouping key（nr），还包含 distinct 的列（letter），效果如 group by nr、letter，执行计划如下：
HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- LocalTableScan [nr#5, letter#6]② partial merge aggregation 步骤
这一步将通过 shuffle 将具有相同 grouping key（此处为 nr、letter）的数据划分为同一分区：
+- HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- Exchange hashpartitioning(nr#5, letter#6, 200)
+- HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- LocalTableScan [nr#5, letter#6]③ partial aggregation for distinct 步骤
第三步，Spark 最终开始执行聚合，执行的是 partial aggregate：
+- HashAggregate(keys=[nr#5], functions=[partial_count(distinct letter#6)], output=[nr#5, count#18L])
+- HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- Exchange hashpartitioning(nr#5, letter#6, 200)
+- HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- LocalTableScan [nr#5, letter#6]④ final aggregation 步骤
第四步，partial aggregate（第三步）的结果将合并到最终结果中，并进行返回。它涉及 shuffle：
HashAggregate(keys=[nr#5], functions=[count(distinct letter#6)], output=[nr#5, count(DISTINCT letter)#12L])
+- Exchange hashpartitioning(nr#5, 200)
+- HashAggregate(keys=[nr#5], functions=[partial_count(distinct letter#6)], output=[nr#5, count#18L])
+- HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- Exchange hashpartitioning(nr#5, letter#6, 200)
+- HashAggregate(keys=[nr#5, letter#6], functions=[], output=[nr#5, letter#6])
+- LocalTableScan [nr#5, letter#6]我们用下面的这张图来总结上述几个步骤：
二、无 Distinct aggregation 算法
无 Distinct aggregation 会简单一些，仅包含两个步骤，我们通过下面的例子来说明：
val dataset = Seq(
(1, "a"), (1, "a"), (1, "a"), (2, "b"), (2, "b"), (3, "c"), (3, "c")
).toDF("nr", "letter")
dataset.groupBy($"nr").count().explain(true)①、partial aggregations 步骤
第一步即进行局部聚合：
HashAggregate(keys=[nr#5], functions=[partial_count(1)], output=[nr#5, count#17L])
+- PlanLater LocalRelation [nr#5]②、final aggregation 步骤
第二步，毫无疑问，对部分结果进行了最终汇总：
HashAggregate(keys=[nr#5], functions=[count(1)], output=[nr#5, count#12L])
+- HashAggregate(keys=[nr#5], functions=[partial_count(1)], output=[nr#5, count#17L])
+- PlanLater LocalRelation [nr#5]三、Hash-based 和 Sort-based aggregation
上述两种模式都会调用到 createAggregate 方法，该方法为以下 3 种策略创建物理执行计划：
- hash-based
- object-hash-based
- sort-based
这 3 中策略有一些共性。一个 Spark Sql aggregation 主要由两部分组成：
- 一个 agg buffer（聚合缓冲区：包含 grouping keys 和 agg value）
- 一个 agg state（聚合状态：仅 agg value）
每次调用 GROUP BY key 并对其使用一些聚合时，框架都会创建一个聚合缓冲区，保留给定的聚合（GROUP BY key）。指定 key（COUNT，SUM等）所涉及的聚合都在此聚合缓冲区存储其部分（partial）或最终聚合结果，称为聚合状态。该状态的存储格式取决于聚合：
- 对于 AVG，它将是2个值，一个是出现次数，另一个是值的总和
- 对于 MIN，它将是到目前为止所看到的最小值 依此类推
hash-based 策略使用可变的、原始的、固定 size 的类型来作为 agg state，包括：
- NullType
- BooleanType
- ByteType
- ShortType
- IntegerType
- LongType
- FloatType
- DoubleType
- DateType
- TimestampType 这里的可变能力非常重要，因为 Spark 会直接修改该值（如对于 count 来说，遇到新的 row，就会把 count 的值（agg state）加上 1）。
对于 agg state 的值是其他类型的情况，使用 object-hash-based 策略，该策略自 2.2.0 版本引入，目的是为了解决 hash-based 策略的局限性（必须使用可变的、原始的、固定 size 的类型来作为 agg state）。在 2.2.0 之前，针对 HashAggregateExec 不支持的其他类型执行的聚合都会转换为 sort-based 的策略。大部分情况下，sort-based 的性能会比 hash-based 的差，因为在聚合前会进行额外的排序。通过参数 spark.sql.execution.useObjectHashAggregateExec 来控制是否使用 object-hash-based 聚合，默认为 true。我们通过下面的例子来理解 sort-based 和 object-hash-based 的区别：
查询
val dataset2 = Seq(
(1, "a"), (1, "aa"), (1, "a"), (2, "b"), (2, "b"), (3, "c"), (3, "c")
).toDF("nr", "letter")
dataset2.groupBy("nr").agg(functions.collect_list("letter").as("collected_letters")).explain(true)
如你所见，上图两个物理执行计划均只进行一次 shuffle，但 sort-based 聚合相对于 object-hash-based 额外多了两次排序，带来性能开销。
另一个值得关注的点是，hash-based 和 object-hash-based 运行过程中如果内存不够用，会切换成 sort-based 聚合。对于 object-hash-based 聚合，通过参数 spark.sql.objectHashAggregate.sortBased.fallbackThreshold 控内存中（一种 hashMap）最多持有多少个 agg buffer（一个 grouping key 的组合一个），若超过该值，则切换为 sort-based agg，该配置默认值为 128。如果切换为 sort-based agg，会打印如下日志：
ObjectAggregationIterator: Aggregation hash map reaches threshold capacity (128 entries), spilling and falling back to sort based aggregation. You may change the threshold by adjust option spark.sql.objectHashAggregate.sortBased.fallbackThreshold对于 hash-based，该值为 Integer.MaxValue