优化 Vitess 查询规划器中的聚合

发表于： 2025年9月14日 2025年9月14日
分类： Vitess
标签： vitess

引言

最近，我遇到了一个有趣的 bug。一位用户报告说，他们的查询会导致 VTGate 获取大量数据，有时甚至会触发内存不足（OOM）错误。为了更深入地了解 Vitess 中的分组和聚合处理方式，我推荐你阅读之前的一篇博客文章。

问题查询

出现问题的查询如下：

SELECT SUM(user.type)
FROM user
    JOIN user_extra ON user.team_id = user_extra.id
GROUP BY user_extra.id
ORDER BY user_extra.id;

查询规划器无法将聚合委托给 MySQL，因此需要从多个分片中获取大量数据进行聚合处理，导致性能问题。

规划和树重写

在查询规划阶段，我们会对查询树进行广泛的重写处理，以尽可能多地将工作下推到分片级别（Routes）。这种重写过程会反复进行，直到整个查询树在一次完整扫描后不再发生变化，这种状态被称为**固定点（fixed-point）**。重写的目标是通过将操作尽可能靠近数据源，从而优化查询的执行计划。

初始计划

在首次进行“视线扩展”（horizon expansion）后，生成的初始计划如下：

Ordering (user_extra.id)
└── Aggregator (ORG sum(`user`.type), user_extra.id group by user_extra.id)
    └── ApplyJoin on [`user`.team_id | :user_team_id = user_extra.id | `user`.team_id = user_extra.id]
        ├── Route (Scatter on user)
        │   └── Table (user.user)
        └── Route (Scatter on user)
            └── Filter (:user_team_id = user_extra.id)
                └── Table (user.user_extra)

在这一阶段，我们无法直接通过 JOIN 将聚合下推到 MySQL 中，因此需要获取大量数据来处理聚合。不过，可以尝试将排序下推到聚合操作下方。

下推排序操作

通过将排序操作下推到聚合操作下方，计划变更如下：

Aggregator (ORG sum(`user`.type), user_extra.id group by user_extra.id)
└── Ordering (user_extra.id)
    └── ApplyJoin on `user`.team_id = user_extra.id

然而，排序操作无法继续下推，因为查询是基于 JOIN 的右侧（user_extra.id）进行排序的。而排序仅能下推到 JOIN 的左侧。这种情形非常不利——排序阻碍了聚合的下推处理，这意味着需要获取所有数据并对其进行排序后才能完成聚合。

解决方案

在这种情况下，我通常会利用规划器中的不同阶段（phases）来解决问题。

阶段划分

规划器中包含多个阶段，这些阶段按顺序运行。在完成一个阶段后，会运行一些“下推重写器”（push-down rewriters），然后进入下一阶段。
重写器的功能：

特定任务：通过重写器运行查询计划以完成特定任务。例如，“拉取 UNION 的 DISTINCT”重写器会从 UNION 中提取 DISTINCT 部分，并单独使用一个操作符处理它。
控制下推操作的时机：一些重写器只有在达到特定阶段后才会启动。

通过将“排序下推至聚合下方”的重写器延迟到“分裂聚合阶段”（split aggregation phase），我们可以将聚合操作下推到 JOIN 内部。这并不会阻止排序重写器发挥作用，而只是让它稍后再执行。

最终计划

最终查询树如下：

Aggregator (SUM(`user`.type) group by user_extra.col)
└── Projection (SUM(`user`.type) * count(*), user_extra.col)
    └── Ordering (user_extra.col)
        └── ApplyJoin (on [`user`.team_id = user_extra.id])
            ├── Route (Scatter on user)
            │   └── Aggregator (SUM(type) group by team_id)
            │       └── Table (user)
            └── Route (Scatter on user_extra)
                └── Aggregator (count(*) group by user_extra.col)
                    └── Filter (:user_team_id = user_extra.id)
                        └── Table (user_extra)

现在，大部分聚合操作都已被下推到 MySQL 中处理。在 VTGate 层面，我们只需对各分片返回的 SUM 值进行最终 SUM（汇总）操作。这显著减少了 VTGate 需要处理的数据量。