基于 egg 用不到千行代码实现一个小巧精悍的 SQL 语言优化器

作者：王润基 RisingWave Labs 内核开发工程师

今天我们来介绍一个有趣的程序优化器框架 egg。基于 egg 我们可以用不到一千行 Rust 代码，实现一个小巧精悍的 SQL 语言优化器。

我们的优化器同时支持基于规则的优化（RBO）和基于代价的优化（CBO），并且实现了各种经典优化规则，例如：表达式化简，常量折叠，谓词下推，投影下推（列裁剪），连接重排序，代价估计等。我们将这款优化器装备到了教学数据库 RisingLight 项目中，作为其查询引擎的核心组成部分。最终用它运行 TPC-H 测试，取得了比较不错的效果。据我所知，我们应该也是首个 all-in egg 的开源数据库项目。

由于内容较多，本文计划拆成上下两篇。上篇首先介绍 egg 框架的基本原理和使用方式，下篇具体介绍如何用 egg 实现各种 SQL 优化规则。阅读本文需要读者有基础的 Rust 和 SQL 语言背景知识，同时欢迎熟悉编译优化和数据库的读者前来围观指导。

Egg 与 E-graph

根据官方介绍，Egg 这个名字的含义是 E-graphs good。它是一个基于 E-graph 的程序优化合成框架。其中 E-graph 的 e 表示 equivalence，是一种包含等价类的图结构。说着抽象，让我们直接来看官方提供的图示吧：（egraphs-good.github.io）

这四个图表示的是同一个表达式结构，从左到右依次展示了将表达式 a * 2 / 2 优化为 a 的过程。

图中每个节点称为 e-node，可以表示符号（a）、操作（*）或者常量数字（2）。每个节点可以拥有若干子节点，例如操作符（*）的子节点是它的操作数（a 和 2）。

多个逻辑上相等的节点可以形成一个等价类 e-class，例如 a * 2 和 a << 1 的效果是等价的，因此对应的两个节点 * 和 <<被画在同一个虚线框中。此时我们再观察那些指向子节点的箭头，就会发现它们实际上都指向 e-class 而不是 e-node。原因也不难理解，因为根据定义，一个 e-class 中的所有节点都是互相等价的，因此对于父节点来说，它的孩子选其中的哪个节点都一样。

所有的 e-node，e-class，和它们的父子关系，共同组成了一个图结构，也就是 e-graph。e-graph 的威力在于它是动态可变的。我们可以向其中插入新的节点，或者将其中两个节点 union 起来形成新的等价类或者合并等价类。

上图就展示了向图中插入新的表达式 a << 1 ，并与 a * 2 合并等价类的过程。

这样的变换就是一种重写规则。在 egg 中，我们可以使用 Lisp 风格的 S 表达式结合模式匹配来描述重写规则。例如上面的规则可以简洁地表示为 (* ?x 2) => (<< ?x 1)，其中 ?x 表示一个可以匹配任意节点的变量。当定义好全部规则后，egg 的执行器就可以在图中进行搜索、匹配和重写，找出同一个表达式的所有可能的等价表示，最后从中选出一种最优方案。这一过程在学术上被称为 Equality Saturation （没有找到正式的中文翻译，我根据自己的理解翻译为 等价类饱和扩展）。

相比于简单的节点替换，Equality Saturation 保留了变换前的方案，并在最后通过动态规划的方式选出全局最优表示，通常能够得到比贪心法更优的结果。这种设计非常适合实现数据库中基于代价的优化（CBO），同时也能兼容基于规则的优化（RBO）。不过相比于更加专业的 Volcano 和 Cascade 优化器框架，egg 目前还不能在代价函数的指导下进行启发式搜索，因此容易出现状态爆炸、找不到最优解的问题。

以上就是 egg 进行程序优化的理论基础和基本原理。更多细节可以阅读 egg 提供的教程文档。下面我们就来实际操作一下吧！

定义语言

使用 egg 的第一步是定义语言，在本文中就是 SQL 语句和其中的表达式。为了简单入手，我们先定义表达式，并复现官方样例中的表达式优化规则。待到熟悉用法后再考虑 SQL 语句。

egg 充分利用了 Rust 语言中支持代数数据类型的枚举（enum）来定义语言。我们可以用它提供的 define_language 宏定义一个 enum 来描述一个节点的各种类型：

define_language! {
    pub enum Expr {
        // values
        Constant(DataValue),         // null, true, 1, 1.0, "hello", ...
        Column(ColumnRefId),         // $1.2, $2.1, ...
        // operations
        "+" = Add([Id; 2]),          // (+ a b)
        "and" = And([Id; 2]),        // (and a b)
        // list
        "list" = List(Box<[Id]>)     // (list ...)
    }
}

上面代码是一个简化版本，其中包含了三种节点类型：值类型，操作符，和一个特殊的列表节点。

1. 值类型，表达式中的叶子节点。它可以是一个常量 Constant，或者一个变量 Column 引用数据库表中的一列。egg 要求它们关联的 Rust 类型（如 DataValue）实现 Debug, Display, Eq, Ord 和 FromStr trait，以实现显示、字符串解析、比较去重等操作。
2. 操作符，通常是表达式中间节点。与 Rust 标准语法相比前面多了个字符串常量，用来在显示和解析字符串时描述这个节点。它关联的 Rust 类型可以是 Id，[Id; N]或者 Box<[Id]>，分别表示 1 个，N 个，或不定长个子节点。其中 Id 是一个节点的索引，内部是个整数。我们必须拿着 Id 去 EGraph 或 RecExpr 容器中查询才能得到具体的节点内容。因此这类中间节点不能脱离容器而独立存在。
3. 列表，这是我们引入的一种特殊节点，表示若干节点组成的有序列表。它可以有任意多个子节点。

定义好节点后，我们可以向图中依次插入节点来构造一个表达式：

/// Construct an expression for `[1, 2 + 3]`.
fn construct_expr() -> RecExpr {
    // Internally, `RecExpr` is a simple array of `Expr`,
    // whose last element is the root node.
    let mut e = RecExpr::default();
    let a: Id = e.add(Expr::Constant(1));
    let b: Id = e.add(Expr::Constant(2));
    let c: Id = e.add(Expr::Constant(3));
    let d: Id = e.add(Expr::Add([b, c]));
    e.add(Expr::List([a, d].into()));
    e
}

同时，egg 支持使用 Lisp 风格的 S-expression 来描述表达式：

/// Parse an expression from the string.
fn parse_expr() -> RecExpr {
    "(list 1 (+ 2 3))".parse().unwrap()
}
// It displays in the same format.
assert_eq!(parse_expr().to_string(), "(list 1 (+ 2 3))");

相比之下用这种方式更加简洁直观，在后文中我们统一使用这种表示方式。

表达式重写

构造好表达式后，我们可以进一步定义规则来对它进行重写。

比如对一开始提到的例子，可以用 rewrite 宏描述如下：

rewrite!("mul2-to-shl1"; "(* ?a 2)" => "(<< ?a 1)")

=> 两侧的字符串分别表示变换前后的表达式模式（Pattern），其中 ?a 表示可以替换成任何表达式的变量。当 egg 执行器搜索到符合这个模式的表达式后，会按照右边的模式生成一个新的表达式，并将其和原来的节点 union 起来形成等价类。另外，这里的 => 也可以替换成 <=> 表示双向变换。

对于另一些规则来说，它们还需要附加一定的执行条件。例如，除法规则仅在除数非 0 时才有效：

rewrite!("div-to-mul"; "(/ ?a ?b)" => "(* ?a (/ 1 ?b))" if is_not_zero("?b"))

这里我们使用了 is_not_zero 函数来判断变量 ?b 是否非 0。不过我们如何才能知道一个表达式是否非 0 呢？这时就需要引入表达式分析。

表达式分析

egg 的表达式分析功能允许我们为每个等价类关联任意值来描述它的某些特性。例如：它是不是常数，它的数据类型，其中引用了哪些列......

常量分析

比如为了解决上面的问题，我们需要引入常量分析。顾名思义，常量分析会判断每个节点对应的表达式是不是常数，以及具体的常数值。例如对于表达式 (+ 1 1) ，它的关联值是 Some(2)，因为我们知道它等价于常数 2；而对于表达式 (+ x 1)，它的关联值就是 None，因为我们无法判断它是不是常数。不过，我们暂时也无法断言它不是常数。

注意观察，这里的分析值是关联到 e-class 而不是 e-node 的。因为根据定义，同一个等价类中的所有节点在逻辑上都是等价的，因此他们共享相同的属性。然而，同一个等价类中的不同节点分析出的值可能是不同的！随着新节点的加入，等价类的分析值会逐渐变得精确。例如，对表达式 (- x x) 进行常量分析，它的结果会是 None 。然而当我们应用规则 (- ?x ?x) => 0 以后，它会和 0 产生合并，合并后的分析结果则是 Some(0)。这一更新会进一步传播给上游节点（如 (- (- x x))），从而使得整个图的分析结果都更加精确。

表达式分析会随着节点的加入和合并同步进行。因此当一个节点被加入图中以后，我们立刻就能拿到它的分析数据。每一种分析都实现在 Analysis trait 下面，具体细节可以参考 egg 的 API 文档，其中给的例子就是常量分析。有了常量分析的结果以后，我们就可以实现 is_not_zero函数：

// 无需过分关注函数签名，照着文档仿写就好
/// Returns true if the expression is a non-zero constant.
fn is_not_zero(var: &str) -> impl Fn(&mut EGraph, Id, &Subst) -> bool {
    let v = var.parse::<Var>().unwrap();
    move |egraph, _, subst| {
        matches!(egraph[subst[v]].data.constant, Some(n) if !n.is_zero())
    }
}

这样我们就依赖表达式分析实现了一个条件重写规则。

常量折叠

在表达式分析的过程中，egg 还允许我们修改 e-graph，进行节点添加或合并。利用这一机制我们就可以实现 **常量折叠（Constant Folding）**优化：将确定是常量的表达式替换为单个值。

如左图所示，在常量分析中我们发现 (/ 2 2) 节点等价于常数 1。此时我们创建一个新节点 1 （会发现已经存在于图中）并将其与 (/ 2 2) 合并。这样就完成了常量折叠优化。这也是唯一一个利用表达式分析的副作用，而不是通过重写规则完成的优化。（具体代码可参考文档）

类型分析

类似地，我们还可以对表达式进行类型分析。也就是推断每个表达式返回值的类型，并在过程中进行类型检查。由于类型检查的结果可能是错误，因此它关联的数据是一个 Result<Type, TypeError>。例如 (+ INT INT) 的类型是 Ok(INT)，而 (+ INT VARCHAR) 的类型则是 Err(..)。当发现子节点出现类型错误时，父节点则应该继续抛出这个错误，以便正确定位错误源头。

另外，类型分析还需要依赖外部信息。对于常量节点我们可以直接推断它的类型，但对于引用某一列的变量节点，我们就需要去数据库的 Catalog 中查询对应列的类型。在 egg 中这一需求可以通过在 Analyze 结构中添加 catalog 字段来实现。（具体代码可参考 RisingLight）

总的来说，类型分析在数据库中主要有两个用途：一是在解析 SQL 语句后的 binding 阶段执行类型检查，二是在生成 executor 时确定每个算子每一列的输出类型。在采用新优化器前，RisingLight 的类型检查分散在 binder 的各个函数中，并且每一个表达式节点都要额外维护一个类型字段；采用新优化器后，类型分析被实现为一个 egg Analysis。相比之下代码更加集中，并且节点本身不再需要维护类型信息，需要的时候随时去访问关联数据即可。

定义 SQL 算子

到这里我们用 egg 实现了基础的表达式分析和化简，下面开始进入正题，看看如何使用 egg 描述和处理 SQL 查询语句。

首先我们需要引入 SQL 算子的概念，它背后的理论基础是关系代数。和常见的表达式算子（如加减乘除）类似，SQL 算子作用于一个或多个关系（也就是数据库中的表），并生成一个新的关系。一条 SQL 查询语句，会首先转换成由若干 SQL 算子组成的执行计划树，经过优化器的化简后，再交给执行器执行。

比如说我们想查询：2022 年给 RisingLight 项目提交超过 10 次的用户，并且按提交数排序取前 10 名。这一需求可以表达成以下 SQL 语句：

SELECT users.name, count(commits.id)
FROM users, repos, commits
WHERE users.id = commits.user_id
    AND repos.id = commits.repo_id
    AND repos.name = 'RisingLight'
    AND commits.time BETWEEN date '2022-01-01' AND date '2022-12-31'
GROUP BY users.name
HAVING count(commits.id) >= 10
ORDER BY count(commits.id) DESC
LIMIT 10

数据库查询引擎会为它生成这样的执行计划树：

Projection: [users.name, count(commits.id)]
└─Limit: 10, offset=0
  └─Order: count(commits.id) desc
    └─Filter: count(commits.id) >= 10
      └─Aggregate: count(commits.id), groupby=users.name
        └─Filter: users.id = commits.user_id
          │       AND repos.id = commits.repo_id
          │       AND repos.name = 'RisingLight'
          │       AND commits.time BETWEEN date '2022-01-01' AND date '2022-12-31'
          └─Join: inner, on=true
            ├─Join: inner, on=true
            │ ├─Scan: users [id, name, email]
            │ └─Scan: repos [id, name, url]
            └─Scan: commits [id, user_id, repo_id, time, message]

这里面涉及了常用的各种 SQL 算子。从下往上来看分别是：

• Scan：从指定表中扫描指定列的全部数据
• Join：以指定条件连接两个表
• Filter：以指定条件过滤表中的每一行
• Aggregate：以指定 key 分组，并将每组中的数据按指定函数聚合成一行
• Order：以指定 key 对表按行排序
• Limit：跳过 offset 行后保留 limit 行
• Projection：以指定表达式对每一行进行投影变换

这里边除了 Scan 节点直接从表中读取数据以外，其它节点的输入都来自子节点的输出。

在 egg 中，我们可以将它们更加具体地定义如下。右边的注释说明了每个子节点的具体含义，其中[some..] 表示某种类型的 list 节点，child left right 为 SQL 算子，其它为表达式节点：

define_language! {
    pub enum Expr {
        // ...
        // plans
        "scan" = Scan([Id; 2),          // (scan table [column..])
        "proj" = Proj([Id; 2]),         // (proj [expr..] child)
        "filter" = Filter([Id; 2]),     // (filter condition child)
        "join" = Join([Id; 4]),         // (join type condition left right)
            "inner" = Inner,
            "left_outer" = LeftOuter,
            "right_outer" = RightOuter,
            "full_outer" = FullOuter,
        "agg" = Agg([Id; 3]),           // (agg aggs=[expr..] group_keys=[expr..] child)
                                            // expressions must be agg
                                            // output = aggs || group_keys
        "order" = Order([Id; 2]),       // (order [order_key..] child)
            "asc" = Asc(Id),                // (asc key)
            "desc" = Desc(Id),              // (desc key)
        "limit" = Limit([Id; 3]),       // (limit limit offset child)
    }
}

注意到我们将这些 SQL 算子和之前的表达式节点定义在了同一个 enum 中，在 Rust 类型系统层面并没有对它们作区分，并且所有的子节点都是同样的 Id 类型。这也就意味着我们无法在编译时对子节点类型做出约束，而只能以注释的形式描述它们，并在写代码时自觉遵守。诸如 (scan 1 2) 这样的表达式在 egg 看来是完全合法的，但却不符合我们定义的规范。而一旦意外出现这样的非法节点，就会导致重写规则失效，程序 panic 或产生未定义行为的后果。换句话说，这是一种动态类型语言。我们可以享受动态类型带来的灵活性，同时也需要小心缺乏编译期类型检查带来的风险。

基于以上定义，我们的执行计划在程序中将表示为这样的 egg 表达式：

(proj (list $1.2 (count $3.1))
  (limit 10 0
    (order (list (desc (count $3.1)))
      (filter (>= (count $3.1) 10)
        (agg (list (count $3.1)) (list $1.2)
          (filter (and (and (and (= $1.1 $3.2) (= $2.1 $3.3)) (= $2.2 'RisingLight'))
                    (and (>= $3.4 '2022-01-01') (<= $3.4 '2022-12-31')))
            (join inner true
              (join inner true
                (scan $1 (list $1.1 $1.2 $1.3))
                (scan $2 (list $2.1 $2.2 $2.3)))
              (scan $3 (list $3.1 $3.2 $3.3 $3.4 $3.5))
)))))))

这里所有的列名都已经转成了 ID 形式（如 $1.2 表示第 1 表的第 2 列）。各位是不是已经看括号看得眼花缭乱了，这就是 Lisp 啊 [doge]

不过显然，这个初始的执行计划并不是最优的。假设我们的 users 表有 1000 行，repos 表有 50 行，commits 表有 100,000 行数据。那么在这个计划中，底层的两个 join 就需要消耗 1000 x 50 x 100,000 次计算，并生成同样规模的输出。而实际上，由于各个表的 ID 都是主键（唯一且非空），因此可以使用 HashJoin 算法并调整连接顺序，使得两次连接在 100,000 的常数倍时间内完成。

而为了找到这样更优的计划，我们需要引入一系列的优化策略来对计划树进行重写。这将是下篇文章的主要内容。

总结与预告

这篇文章我们介绍了程序优化器框架 egg。它基于 E-graph 数据结构实现了 Equality Saturation 算法，根据用户给定的重写规则对表达式进行等价类扩展，从而找到更优表示，实现优化的效果。我们基于 egg 提供的分析功能实现了对表达式的常量分析，并顺便完成了第一个常量折叠优化。最后我们介绍了常用的 SQL 算子并在 egg 中对它们进行定义。

下篇文章我们将重点探讨如何在 egg 中实现各种 SQL 算子优化策略，其中包括：

• 算子消除
• 谓词下推
• 投影下推与列裁剪
• 连接重排序与 HashJoin
• 代价函数（CBO）

以及在数据库工程实现中比较烦人的几个话题：

• 如何从 SELECT 语句中提取聚合函数
• 如何将对列的引用转换为对输入数据的物理下标
• 如何处理子查询

欢迎继续关注！

关于 RisingWave

RisingWave 是一款分布式 SQL 流处理数据库，旨在帮助用户降低实时应用的的开发成本。作为专为云上分布式流处理而设计的系统，RisingWave 为用户提供了与 PostgreSQL 类似的使用体验，并且具备比 Flink 高出 10 倍的性能以及更低的成本。了解更多：

GitHub: risingwave.com/github

官网: risingwave.com

公众号： RisingWave 中文开源社区