如何在SPARKSQL中通过Expand操作实现GroupingSets（逻辑计划层面详解）

背景

本文基于spark 3.1.2
之前在做bug调试的时候遇到了expand的问题&＃xff0c;在此记录一下

分析

运行该sql&＃xff1a;

create table test_a_pt(col1 int, col2 int,pt string) USING parquet PARTITIONED BY (pt);
insert into table test_a_pt values(1,2,&＃39;20220101&＃39;),(3,4,&＃39;20220101&＃39;),(1,2,&＃39;20220101&＃39;),(3,4,&＃39;20220101&＃39;),(1,2,&＃39;20220101&＃39;),(3,4,&＃39;20220101&＃39;);
select count(*),col1 as alias
from test_a_pt
group by col1,col2
grouping sets (col1,col2)
order by col1,col2 ;


可以看到如下逻辑计划的变化&＃xff08;只截取grouping sets相关的&＃xff09;&＃xff1a;

&＃61;&＃61;&＃61; Applying Rule org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveRelations &＃61;&＃61;&＃61;&＃39;Sort [&＃39;col1 ASC NULLS FIRST], true &＃39;Sort [&＃39;col1 ASC NULLS FIRST], true&＃43;- &＃39;GroupingSets [ArrayBuffer(&＃39;col1), ArrayBuffer(&＃39;col2)], [&＃39;col1, &＃39;col2], [&＃39;col1, &＃39;count(1) AS alias#221] &＃43;- &＃39;GroupingSets [ArrayBuffer(&＃39;col1), ArrayBuffer(&＃39;col2)], [&＃39;col1, &＃39;col2], [&＃39;col1, &＃39;count(1) AS alias#221]
! &＃43;- &＃39;UnresolvedRelation [test_table], [], false &＃43;- &＃39;SubqueryAlias spark_catalog.default.test_table
! &＃43;- &＃39;UnresolvedCatalogRelation &＃96;default&＃96;.&＃96;test_table&＃96;, [], false

对于GroupingSets里面的信息做一下解释&＃xff1a;

&＃39;GroupingSets [ArrayBuffer(&＃39;col1), ArrayBuffer(&＃39;col2)], [&＃39;col1, &＃39;col2], [&＃39;col1, &＃39;count(1) AS alias#221]

• *&＃96;*表示还未解析的计划&＃xff0c;

• [ArrayBuffer(&＃39;col1), ArrayBuffer(&＃39;col2)] 是grouping sets里面的两个值col1和col2

• [&＃39;col1, &＃39;col2] 是group by后面的值col1和col2

• [&＃39;col1, &＃39;count(1) AS alias#221] 是聚合表达式的值&＃xff0c;也就是select后面的值 count(*),col1 as alias

接下来就是&＃xff1a;
ResolveGroupingAnalytics计划&＃xff1a;

06:49:07.323 WARN org.apache.spark.sql.catalyst.rules.PlanChangeLogger:
&＃61;&＃61;&＃61; Applying Rule org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveGroupingAnalytics &＃61;&＃61;&＃61;&＃39;Sort [&＃39;col1 ASC NULLS FIRST], true &＃39;Sort [&＃39;col1 ASC NULLS FIRST], true
!&＃43;- &＃39;GroupingSets [ArrayBuffer(col1#223), ArrayBuffer(col2#224)], [col1#223, col2#224], [col1#223, count(1) AS alias#221L] &＃43;- Aggregate [col1#229, col2#230, spark_grouping_id#228L], [col1#229, count(1) AS alias#221L]
! &＃43;- SubqueryAlias spark_catalog.default.test_table &＃43;- Expand [List(col1#223, col2#224, pt#225, col1#226, null, 1), List(col1#223, col2#224, pt#225, null, col2#227, 2)], [col1#223, col2#224, pt#225, col1#229, col2#230, spark_grouping_id#228L]
! &＃43;- Relation[col1#223,col2#224,pt#225] parquet &＃43;- Project [col1#223, col2#224, pt#225, col1#223 AS col1#226, col2#224 AS col2#227]
! &＃43;- SubqueryAlias spark_catalog.default.test_table
! &＃43;- Relation[col1#223,col2#224,pt#225] parquet


代码自己可以去看&＃xff0c;我们从逻辑来上分析&＃xff1a;

&＃39;GroupingSets [ArrayBuffer(col1#223), ArrayBuffer(col2#224)], [col1#223, col2#224], [col1#223, count(1) AS alias#221L] ||\/&＃43;- Aggregate [col1#229, col2#230, spark_grouping_id#228L], [col1#229, count(1) AS alias#221L] &＃43;- Expand [List(col1#223, col2#224, pt#225, col1#226, null, 1), List(col1#223, col2#224, pt#225, null, col2#227, 2)], [col1#223, col2#224, pt#225, col1#229, col2#230, spark_grouping_id#228L]&＃43;- Project [col1#223, col2#224, pt#225, col1#223 AS col1#226, col2#224 AS col2#227]

把最重要的转换提取出来做解释&＃xff1a;

&＃43;- Project [col1#223, col2#224, pt#225, col1#223 AS col1#226, col2#224 AS col2#227]


• 前三个expression col1#223, col2#224, pt#225 是根据 Relation(也就是从表test_a_pt直接获取到的&＃xff0c;和表的字段保持一致)

• 后面的expression col1#223 AS col1#226, col2#224 AS col2#227 是根据grouping sets和group by的值整合过来的(并且会加上别名&＃xff0c;取别名是为了Expand用的)&＃xff0c;如果没有group by 这个表达式才会取grouping sets的值&＃xff0c;否则就取group by后面的值(目前spark 3.1.2的做法是group by的属性肯定包含了grouping sets里面的属性,SPARK-33229可以支持)&＃xff1a;

如&＃xff1a;group by col1,col2 grouping sets (col1,col2)
则取 col1,col2 如&＃xff1a;grouping sets (col1,col2)
则取 col1,col2


对于Expand:

Expand [List(col1#223, col2#224, pt#225, col1#226, null, 1), List(col1#223, col2#224, pt#225, null, col2#227, 2)], [col1#223, col2#224, pt#225, col1#229, col2#230, spark_grouping_id#228L]


List(col1#223, col2#224, pt#225, col1#226, null, 1), List(col1#223, col2#224, pt#225, null, col2#227, 2) 这些是expand的输入expression&＃xff0c;其中

• List(col1#223, col2#224, pt#225, col1#226, null, 1) 中的
  col1#223, col2#224, pt#225 也是从表test_a_pt直接获取到的字段&＃xff0c;和表的字段保持一致
  col1#226 是从 Project的col1#223 AS col1#226取到的&＃xff08;作为Expand的输入表达式&＃xff09;&＃xff0c;
  null 根据grouping sets的特性而增加的一行值&＃xff08;作为Expand的输入表达式&＃xff09;
  1 也是增加的一行值&＃xff08;作为Expand的输入表达式&＃xff09;
• List(col1#223, col2#224, pt#225, null, col2#227, 2) 解释也和上面一样&＃xff0c;只不过null的位置发生了变化&＃xff0c;而1变成了2&＃xff0c;这是为了做聚合的时候进行区分

[col1#223, col2#224, pt#225, col1#229, col2#230, spark_grouping_id#228L] 这些是expand的输出expression&＃xff0c;其中

• col1#223, col2#224, pt#225 和表test_a_pt的字段值一样
• col1#229, col2#230, spark_grouping_id#228L 是expand做的的扩展字段&＃xff0c;
  因为col1和col2的值可能为null&＃xff0c;所以exprId和表test_a_pt不一致&＃xff0c;
  spark_grouping_id#228L 纯属于虚拟字段

而且expand的输入字段是一个Seq(Seq)&＃xff0c;这在ExpandExec的时候&＃xff0c;会进行row的倍数扩大&＃xff0c;Seq里的元素有几个&＃xff0c;就会扩展多少倍。

对于Aggregate

Aggregate [col1#229, col2#230, spark_grouping_id#228L], [col1#229, count(1) AS alias#221L]


其中&＃xff0c;

• [col1#229, col2#230, spark_grouping_id#228L]就是把Expand的输出字段&＃xff0c;按照这三个表达式进行group by 聚合
• [col1#229, count(1) AS alias#221L] 是聚合表达式&＃xff0c;包括聚合的部分字段和部分聚合函数&＃xff0c;也就是select语句count(*),col1 as alias

至此Grouping sets 转Expand就分析完了。