【Calcite 系列】深入理解 Calcite 的 AggregateUnionTransposeRule

AggregateUnionTransposeRule 会把 Aggregate 下推到 UNION ALL 的每个输入之上，再在 union 顶部加一层汇总聚合把子结果合并回来。这和 AggregateJoinTransposeRule 的思路类似，都是先在分支内部做局部汇总，再做顶层合并。本文结合源码实现，分析它为什么只支持 UNION ALL、只支持一组明确列出的聚合函数，以及它为什么要特别关心各个 union 输入的列空值性差异。

1. 规则要解决什么问题

典型输入：

sql 复制代码

SELECT deptno, SUM(sal)
FROM (
  SELECT deptno, sal FROM t1
  UNION ALL
  SELECT deptno, sal FROM t2
) u
GROUP BY deptno;

规则可以把它改写成：

sql 复制代码

SELECT deptno, SUM(sum_sal)
FROM (
  SELECT deptno, SUM(sal) AS sum_sal FROM t1 GROUP BY deptno
  UNION ALL
  SELECT deptno, SUM(sal) AS sum_sal FROM t2 GROUP BY deptno
) x
GROUP BY deptno;

这样每个分支先本地聚合，顶部只合并局部结果。

2. 为什么只能用于 UNION ALL

源码注释给了一个非常直观的反例：如果是 UNION DISTINCT，先对子查询分别聚合再 union，和先 union 去重后再聚合，结果可能完全不同。

因此规则一开始就要求：

java 复制代码

if (!union.all) {
  return;
}

3. 为什么只支持有限几类聚合函数

源码里维护了一个 SUPPORTED_AGGREGATES：

MIN/MAX
COUNT
SUM
SUM0
ANY_VALUE
位运算聚合

只要出现：

DISTINCT
不在白名单中的聚合函数

规则就直接退出。

这说明该规则只处理那些具有明确"局部聚合 + 顶层合并"性质的函数。

4. 为什么还要先检查所有输入是否已经唯一

源码会遍历 union 每个输入，检查其在当前 group set 上是否已经唯一。

如果所有输入都已经唯一，那么把 aggregate 下推到各输入上没有任何收益，还可能导致规则来回触发形成循环。因此会直接退出。

5. 子聚合是如何构造的

对于 union 每个输入，规则都会构造一个 child aggregate：

java 复制代码

relBuilder.push(input);
relBuilder.aggregate(relBuilder.groupKey(aggRel.getGroupSet()), childAggCalls);

也就是说，分支级聚合沿用原 group key，但聚合调用可能根据输入列的 nullability 重新推断类型。

6. 为什么要关心不同输入的空值性

源码有一个很重要的细节：

如果 union 某一列在整体 row type 中是 nullable，但某个具体输入里是 not null，或者反过来，那么把 aggregate 下推到子输入时，聚合返回类型可能和顶层不同。

因此它会在构造 child aggregate call 时，必要时重新创建 AggregateCall，让返回类型和当前输入环境匹配。这是为了避免 row type mismatch。

相关测试在 8500 行附近也专门覆盖了这一点。

7. 顶层为什么 COUNT 会变成 SUM0

在 transformAggCalls(...) 中，如果原始聚合是 COUNT，顶层不会继续用 COUNT，而是改成 SUM0。

原因和 join transpose、merge 等规则类似：

子聚合已经算出了每个分支的局部计数；
顶层要做的是把这些局部计数加起来；
因此语义上应是 SUM0(local_count)。

8. 顶层 group key 为什么要重映射

子聚合之后，union 输出中的字段布局已经变成：

前面是 group key；
后面是每个子聚合产出的 measure 列。

因此顶部 aggregate 的 group key 必须被重新映射到新的列坐标系上。源码通过 INVERSE_SURJECTION mapping 完成这个重定位。

9. 代表性测试

RelOptRulesTest 中和这条规则直接相关的测试主要在：

6000 行附近的基础用例；
8500 行附近的类型、断言和唯一性边界用例。

这些测试覆盖了：

正常 UNION ALL 下推；
row type mismatch 修复；
子 aggregate 构造断言；
"如果所有输入已经唯一则不应下推"。

总结

AggregateUnionTransposeRule 的核心思路是：

仅对 UNION ALL 做局部聚合下推；
子分支先按原 group key 汇总；
顶层再对这些局部结果做二次聚合；
对 COUNT、类型和空值性做额外修正。

它是一条典型的"先分支内压缩，再全局合并"的聚合下推规则。