Spark SQL实现原理-逻辑计划优化规则:ColumnPruning(列裁剪)规则

该逻辑计划优化规则,尝试从逻辑计划中去掉不需要的列,从而减少读取数据的量。

列裁剪效果

列裁剪规则会在多种情况下生效,下面通过几个例子来理解该优化规则的行为:

排序并进行列裁剪

当有groupBy等聚合操作时,会把不需要的列在读取数据时去掉,以减少数据的读取量。

case class Person(id: Long, name: String, city: String)
var ds2 = Seq(Person(0, "Ker", "Smith")).toDS
ds2.orderBy('id.asc).select('name).explain(true)

但我们在scala终端中,通过sc.setLogLevel(“TRACE”)打开跟踪日志后就可以看到逻辑计划的优化过程。但执行以上代码时,最开始的逻辑计划是:

Project [id#134L]                                                             
!+- Aggregate [id#134L, name#135], [id#134L, name#135, count(1) AS count#187L] 
!   +- LocalRelation [id#134L, name#135, city#136]                             
!

注意:逻辑计划中的LocalRelation节点表示从本地数据源读取数据。

通过ColumnPruning优化规则优化后的逻辑计划是:

=== Applying Rule org.apache.spark.sql.catalyst.optimizer.ColumnPruning ===
 Project [id#134L]
 +- Aggregate [id#134L, name#135], [id#134L]
    +- Project [id#134L, name#135]
       +- LocalRelation [id#134L, name#135, city#136]

可以看到在读取数据时,添加了Project计划,该计划中只有会使用到的两个字段。也就是说:在读取数据时,只需要读取id和name这两个字段就可以了。再进一步通过Project折叠的优化规则后,得到最终的优化结果:

== Optimized Logical Plan ==
Aggregate [id#134L, name#135], [id#134L]
+- LocalRelation [id#134L, name#135]

可以看到,通过逻辑计划的优化,在从数据源读取数据时把没有用到的列去掉了。这样就减少了需要传输的数据量。

groupBy时进行列裁剪
case class Person(id: Long, age: Long, city: String)
var ds2 = Seq(Person(0, 10, "Smith")).toDS
ds2.groupBy('id).avg("age").orderBy('id.asc).explain(true)

在优化之前的规则为:

Sort [id#84L ASC NULLS FIRST], true                          
 +- Aggregate [id#84L], [id#84L, avg(age#85L) AS avg(age)#92] 
!   +- LocalRelation [id#84L, age#85L, city#86]

优化之后的规则

=== Applying Rule org.apache.spark.sql.catalyst.optimizer.ColumnPruning ===
Sort [id#84L ASC NULLS FIRST], true
+- Aggregate [id#84L], [id#84L, avg(age#85L) AS avg(age)#92]
   +- Project [id#84L, age#85L]
      +- LocalRelation [id#84L, age#85L, city#86]

由于在进行聚合操作时,其结果只需要保留其中的几个字段,所以优化的逻辑计划中添加了Project [id#84L, age#85L]计划,该计划的意思是指获取需要的几个字段。

其他情况下的列裁剪

其他情况还包括,在Union+select的组合去掉不需要的属性;在进行Join+select时去掉不需要的列;在进行window+select操作时,去掉没有对其进行操作的列,等情况。

这些情况可以自己编写具体的代码,并设置日志为TRACE,执行代码来查看跟踪日志。

另外,该规则有可能会和其他规则产生冲突,所以,但遇到冲突时,可能会进行优化。

列裁剪的实现

列裁剪功能在ColumnPruning类中实现。具体的实现代码的细节可以在该类中查看。

object ColumnPruning extends Rule[LogicalPlan] {
	  def apply(plan: LogicalPlan): LogicalPlan = 
	  			removeProjectBeforeFilter(plan transform {
	  			case p @ Project(_, Project(...))
	  				...
          // 处理聚合操作+select操作的情况;  
          case p @ Project(_, a: Aggregate) if (a.outputSet -- p.references).nonEmpty
          	...
          // 处理t1.union(t2).select(...)操作的组合
          case p @ Project(_, u: Union) => 
            ...
          // 处理window操作+select操作的情况
          case p @ Project(_, w: Window) if (w.windowOutputSet -- p.references).nonEmpty
            ...
	  )}
}

在该类中会对Aggregate,Expand,Join等操作中不需要的列进行裁剪。主要的实现逻辑是:把这些操作使用的列的集合和进行操作后需要保留列的集合取交集(这样去掉不需要的列),并把需要保留的列尽可能下推到数据源读取时。

小结

本文对逻辑计划优化的列裁剪优化规则进行了分析。该规则会根据各种操作的目标列,对列进行裁剪的逻辑计划优化,并尽可能把这种优化推到数据源处。