SparkSQL如何实现聚合下推

简介

在之前性能分析的文章中，我们用火焰图看到了程序的一个瓶颈点，Spark的聚合操作执行，

其中GeneratedIterator#agg_doAggregateWithKeys是使用Code Generation技术生成的代码，生成的代码可参考这里，或者这样来看，

scala> val pairsDF = Seq((1,1), (2,2), (3,3)).toDF("a", "b")
pairsDF: org.apache.spark.sql.DataFrame = [a: int, b: int]

scala> pairsDF.createOrReplaceTempView("pairs")

scala> val groupedDF = spark.sql("SELECT count(*) FROM pairs GROUP BY a")
groupedDF: org.apache.spark.sql.DataFrame = [count(1): bigint]

scala> groupedDF.queryExecution.debug.codegen
Found 2 WholeStageCodegen subtrees.
== Subtree 1 / 2 ==
*HashAggregate(keys=[a#5], functions=[partial_count(1)], output=[a#5, count#16L])
+- LocalTableScan [a#5]

Generated code:
//// CodeGen 生成的代码

== Subtree 2 / 2 ==
*HashAggregate(keys=[a#5], functions=[count(1)], output=[count(1)#12L])
+- Exchange hashpartitioning(a#5, 200)
   +- *HashAggregate(keys=[a#5], functions=[partial_count(1)], output=[a#5, count#16L])
      +- LocalTableScan [a#5]

Generated code:
//// CodeGen 生成的代码

之前我们也是留下了TODO准备来实现聚合操作下推到数据源去执行。现在这个优化已经完成了，今天就来分享下是如何实现的。

为什么要实现聚合下推

这个问题似乎问得比较蠢：）很明显，如果数据源能够支持聚合操作，那么将聚合下推就不必传输大量数据给到SparkSQL再进行聚合，而是直接返回聚合结果就行了。而且数据源本身可能就对聚合有很多优化（缓存什么的），所以聚合下推才是一个较优的选择。

聚合下推的需求其实在社区也已经提了很久了，当前SparkSQL只支持了Filter跟Project下推，下面几个issue都是希望SparkSQL能够支持更多的operator下推，

• [SPARK-10899] Support JDBC pushdown for additional commands
• [SPARK-12449] Pushing down arbitrary logical plans to data sources
• [SPARK-12686] Support group-by push down into data sources

而且在SparkSummit上也有人分享过他们所做的实现，The Pushdown of Everything。在评论他们的实现之前，我们先来看下，原生的SparkSQL是怎么实现聚合的，

    spark.read.format("org.apache.spark.examples.sql.DefaultSource")
      .option("from", "1").option("to", "10")
      .load().createOrReplaceTempView("tt")
    // scalastyle:off println
    var df = spark.sql("SELECT count(*) FROM tt WHERE a>1 GROUP BY c")
    println(df.queryExecution)

DefaultSource的代码在这里，输出结果如下，

INFO SparkSqlParser: Parsing command: SELECT count(*) FROM tt WHERE a>1 GROUP BY c

== Parsed Logical Plan ==
'Aggregate ['c], [unresolvedalias('count(1), None)]
+- 'Filter ('a > 1)
   +- 'UnresolvedRelation `tt`

== Analyzed Logical Plan ==
count(1): bigint
Aggregate [c#2], [count(1) AS count(1)#21L]
+- Filter (a#0 > 1)
   +- SubqueryAlias tt
      +- Relation[a#0,b#1L,c#2,d#3,e#4,g#5,f#6,i#7,j#8] ComplicatedScan(1,10)

== Optimized Logical Plan ==
Aggregate [c#2], [count(1) AS count(1)#21L]
+- Project [c#2]
   +- Filter (isnotnull(a#0) && (a#0 > 1))
      +- Relation[a#0,b#1L,c#2,d#3,e#4,g#5,f#6,i#7,j#8] ComplicatedScan(1,10)

== Physical Plan ==
*HashAggregate(keys=[c#2], functions=[count(1)], output=[count(1)#21L])
+- Exchange hashpartitioning(c#2, 200)
   +- *HashAggregate(keys=[c#2], functions=[partial_count(1)], output=[c#2, count#25L])
      +- *Project [c#2]
         +- *Scan ComplicatedScan(1,10) [c#2] PushedFilters: [*IsNotNull(a), *GreaterThan(a,1)], ReadSchema: struct

通过Physical Plan可以看到数据源通过PrunedFilteredScan#buildScan接口返回数据给到SparkSQL，下层的HashAggregate执行部分聚合，Exchange进行shuffle，最后由上层的HashAggregate进行最终聚合。

回到那个SparkSummit上的分享，他们实现了什么呢？看下他们的slide，对于聚合操作是这样的，

可以理解为将上面Physical Plan中的这一部分，

   +- *HashAggregate(keys=[c#2], functions=[partial_count(1)], output=[c#2, count#25L])
      +- *Project [c#2]
         +- *Scan ComplicatedScan(1,10) [c#2] PushedFilters: [*IsNotNull(a), *GreaterThan(a,1)], ReadSchema: struct

替换成了CatalystSource，由数据源来实现这个接口，

也就是说数据源需要去解析LogicalPlan，然后实现部分聚合。

这个方案在SPARK-12449里面有一番讨论，最主要的问题是，对于数据源来说，实现CatalystSource，解析LogicalPlan的成本太高，而且LogicalPlan是SparkSQL内部的数据结构，如果暴露给数据源，API compatibility会是一个大问题。

如何实现聚合下推

那么我们又是怎么实现的呢？实际上跟上面的方案也是类似的，来看下，

    spark.conf.set(SQLConf.AGGREGATION_PUSHDOWN_ENABLED.key, true)
    println("==========AGGREGATION_PUSHDOWN_ENABLED==========")
    df = spark.sql("SELECT count(*) FROM tt WHERE a>1 GROUP BY c")
    println(df.queryExecution)
    // scalastyle:on println

AGGREGATION_PUSHDOWN_ENABLED是增加的一个配置项。输出如下，

==========AGGREGATION_PUSHDOWN_ENABLED==========
INFO SparkSqlParser: Parsing command: SELECT count(*) FROM tt WHERE a>1 GROUP BY c
== Parsed Logical Plan ==
'Aggregate ['c], [unresolvedalias('count(1), None)]
+- 'Filter ('a > 1)
   +- 'UnresolvedRelation `tt`

== Analyzed Logical Plan ==
count(1): bigint
Aggregate [c#2], [count(1) AS count(1)#27L]
+- Filter (a#0 > 1)
   +- SubqueryAlias tt
      +- Relation[a#0,b#1L,c#2,d#3,e#4,g#5,f#6,i#7,j#8] ComplicatedScan(1,10)

== Optimized Logical Plan ==
Aggregate [c#2], [count(1) AS count(1)#27L]
+- Project [c#2]
   +- Filter (isnotnull(a#0) && (a#0 > 1))
      +- Relation[a#0,b#1L,c#2,d#3,e#4,g#5,f#6,i#7,j#8] ComplicatedScan(1,10)

== Physical Plan ==
*HashAggregate(keys=[c#2], functions=[count(1)], output=[count(1)#27L])
+- Exchange hashpartitioning(c#2, 200)
   +- *Scan ComplicatedScan(1,10) [c#2,count#31L] AggregateFunctions: [CountStar()], GroupingColumns: [c], PushedFilters: [*IsNotNull(a), *GreaterThan(a,1)]

通过Physical Plan可以看到，数据源通过AggregatedFilteredScan#buildScan直接返回了部分聚合的结果。这个AggregatedFilteredScan是我新增的一个接口，定义如下，

/**
 * A BaseRelation that can perform aggregation and filter using selected predicates.
 *
 * Row fields MUST be as below:
 * ([GroupingColumn1, GroupingColumn2 ... ,]
 * AggregateFunction1Result[, AggregateFunction2Result ...])
 */
@InterfaceStability.Unstable
trait AggregatedFilteredScan {
  def buildScan(groupingColumns: Array[String],
      aggregateFunctions: Array[AggregateFunc],
      filters: Array[Filter]): RDD[Row]
}

相比于CatalystSource，实现AggregatedFilteredScan非常简单，SparkSQL会将Filter，GroupBy字段以及聚合函数直接下推给到数据源，数据源根据这些信息执行聚合操作并返回聚合结果就可以了。

OK，上面只是直接展示了实现的结果，还没有说到是如何实现的。不过实际上也已经看到了，要实现下推主要就是需要修改SparkSQL的Physical Plan的生成逻辑，也就是SparkPlanner。这里有必要先介绍下SparkSQL大致的运行流程，如下图，

SparkSQL基于ANTLRv4的SQL Parser（语法文件戳这里）将SQL查询转换成Unresolved Logical Plan，此时的表，字段都是unresolved的；然后Analyzer使用元信息将其转换成Resolved Logical Plan；接着SparkOptimizer进行一系列优化，包括常量折叠，谓词下推，Join重排等等；然后就是上面我们提到的SparkPlanner将Optimized Logical Plan转换成Physical Plan，例如Logical Plan中有Join操作，那么这一步就是要决定是使用HashJoin还是BroadcastJoin等最终的物理操作；图中在Physical Plan转换成RDD之前还有一步基于代价来选择Physical Plan，这实际上就是Cost-Based Optimization（CBO），然而目前的SparkSQL是还没有实现的，计划是在2.3.0版本实现，可参考[SPARK-16026] Cost-based Optimizer framework，貌似主要是华为的同学贡献的代码。OK，最后就是将Physical Plan转换成RDD对应的API，运行RDD就可以了。啰嗦一句，RDD的树状结构真是天然可以match到SQL语法的树状结构，从这个层面来讲，Spark真是太适合作为一个分布式的SQL引擎了。

回到聚合下推的实现上来，Logical Plan通过SparkPlanner转换成Physical Plan，SparkPlanner内部基于一系列策略来完成转换操作，

  def strategies: Seq[Strategy] =
      extraStrategies ++ (
      FileSourceStrategy ::
      DataSourceStrategy ::
      DDLStrategy ::
      SpecialLimits ::
      Aggregation ::
      JoinSelection ::
      InMemoryScans ::
      BasicOperators :: Nil)

而我们做的其实就是通过修改Aggregation这个策略来将原本Physical Plan的这部分，

   +- *HashAggregate
      +- *Project
         +- *Scan PrunedFilteredScan

替换成Scan AggregatedFilteredScan就可以了。

完整的实现可以看看我提的这个PR。这个PR没有得到反馈，我猜大概是因为这个issue吧：[SPARK-15689] Data source API v2，也就是2.3.0版本准备实现一套新的DataSource API，为什么需要一套新的API？主要是两个原因：
1. 老的API是面向行存的（RowDataSourceScanExec），并且需要进行数据源与SparkSQL之间的数据类型转换；
2. 老的API太过于依赖SparkSQL内部的实现，这样一来如果SparkSQL内部要做一些大的改动，还需要考虑API的兼容问题。这里其实跟上面提到的CatalystSource暴露出LogicalPlan给数据源是同样的问题；

然后前几天在微博上看到Spark PMC的一位大大说了，在2.3.0版本，SparkSQL将会原生支持聚合下推。这对Spark用户来说是个好消息，通过下面给出的性能对比可以看到聚合下推所带来的性能提升。

简单的性能对比

这是在测试环境上所做的一个简单的性能对比，测试了多次，耗时都差不多。

聚合不下推的情况下，

下推的情况下，

1. 第一条SQL，对2亿条数据的count操作，性能提升了十倍以上，主要还是Hxxx本身的性能，妥妥的；
2. 第二条SQL，数据量变小，2千万左右，聚合操作增多的情况下，性能提升没有第一条SQL那么夸张，但是依然能有3倍多的提升；

暂时只有这么简单的性能对比，后面再找时间做一次完整的测试。

Limit及OrderBy下推

聚合下推实现了之后，我又如法炮制，修改SpecialLimits策略，实现了Limit及OrderBy的下推^_^

不测不知道，一测吓一跳哈，性能提升能有几十倍。

Limit及OrderBy不下推的情况下，

Limit及OrderBy下推的情况下，

在不下推的情况下，SparkSQL需要请求数据源返回所有的数据，也就是2亿条，然后进行排序。可想而知这是相当耗时的；而下推给到数据源，则数据源本地直接排序返回LIMIT条数即可。

可以猜想，数据量越大的情况下，Limit及OrderBy下推的性能提升就越大。

这块代码还没有去提交PR，等到2.3.0新版的DataSource API出来之后再看看哈。也欢迎感兴趣的同学一起交流。alright，今天就先到这了，have fun ^_^