pgrust v0.4-beta 发布：全新查询引擎 Gravity 性能超越 ClickHouse
Programming at Light Speed: Building the Gravity Query Engine
pgrust 发布 v0.4-beta 版本，推出全新查询引擎 Gravity。Gravity 采用多阶段编译器架构，将 Postgres 查询计划经由 Lisp 中间表示层（IR）转为优化机器码，在支持完整 Postgres SQL 方言与类型系统的同时，在 ClickBench 上比 ClickHouse 快多达 3 倍，较原生 PostgreSQL 17 有显著性能提升。
文章详细公开了如何通过自研编译器与特定列存优化实现超越 ClickHouse 的性能，为数据库内核开发提供了可借鉴的从底层极限倒推的工程方法。
随着 pgrust v0.4-beta 的发布，我们的新查询引擎 Gravity 现在在 ClickBench 上的速度最高可比 ClickHouse 快 3 倍，而 ClickBench 是 ClickHouse 官方专门用来展示其运行速度的基准测试。更棒的是，Gravity 在实现这一速度的同时，还支持完整的 Postgres SQL 方言和类型系统。
在之前的几篇博客中，我曾介绍过我们如何通过批处理、算子融合和 SIMD 技术，使 pgrust 在处理分析型查询时比 Postgres 快数百倍，以及我们如何构建了一个能在 5 微秒内编译模板的复制修补（copy-and-patch）JIT 编译器。经过多个版本的积累，所有这些工作最终汇聚成了 Gravity。
复制修补是一个很好的起点，但把预编译的模板拼接在一起，其能力终究是有限的。由于我们没有真正的中间表示或整个循环范围内的寄存器分配，与 rustc 或 Clang 这样的优化编译器相比，我们最终放弃了大量的性能空间。对于 pgrust v0.4，我们希望一竿子插到底，构建一个能为任何 SQL query 生成最优机器码的查询编译器。这就是我们构建 Gravity 的原因。
通常情况下，当工程师试图让软件运行得更快时，他们会采取自底向上的方式。他们运行基准测试、查看火焰图、寻找热点，然后研究如何将代码提速 15% 或 20%。自底向上方法的局限在于，你不知道上限在哪里。你可能距离硬件能运行的最快代码只差 20%，也可能差了 100 倍。
当我们开始构建 Gravity 时，我们颠覆了这个问题。我们问自己：数据库的极限速度究竟能有多快？我们不想循序渐进地加快现有查询引擎的速度，而是希望从硬件和文件格式的物理速度极限出发，倒推回来。我喜欢把这称为“以光速编程”。
为了解“光速”究竟是什么样，我们让 Claude 针对一组分析型查询，直接针对我们的列式存储库编写了独立且高度优化的 Rust 程序。这些程序绕过了整个 SQL 执行器，直接在我们的磁盘格式上运行了经过手工调优、支持多线程并带有 SIMD 的 Rust 循环。当我们运行这些手写的 Rust programs 时，它们的查询执行速度全面比 ClickHouse 等现有分析型数据库快大约 4 倍。随后，我们从这些手调的 Rust 循环中挖掘出可以整合到 Gravity 中的优化方案。
我们对 Gravity 的目标变得很简单：对于任何 SQL query，都在运行时进行 JIT 编译生成机器码，使其性能能够达到 Claude 直接针对我们的文件格式手写的 Rust 代码的 80% 以内（即差距在 20% 以内）。
Gravity 的总体架构
我们在设计 Gravity 架构时的主要目标，是为生成的输出赋予极大的灵活性。我们知道，如果能轻松修改 Gravity 以输出查询所需的任何底层优化，我们就能更容易地缩小差距并逼近光速。
Gravity 没有采用硬编码的执行模式或拼接固定模板的做法，而是被构造成一个多阶段的编译器管道，能够将高级的 Postgres plan 逐步降级为原始汇编代码：
SQL Query
-> Postgres Plan
-> Pipeline Graph
-> Lisp IR
-> Assembly为了看清一个查询在编译过程各个阶段的样子，让我们来看下面这个查询：
SELECT UserID, extract(minute FROM EventTime) AS m, MIN(URL), COUNT(*)
FROM hits
WHERE URL LIKE '%google%'
GROUP BY UserID, m
ORDER BY COUNT(*) DESC
LIMIT 10;阶段 1：Postgres 计划
Gravity 并没有取代 Postgres 的解析器或查询规划器。当查询到达时，pgrust 规划器会分析表统计信息并生成标准的查询计划。
Limit (cost=... rows=10)
-> Sort (Sort Key: count(*) DESC)
-> HashAggregate (Group Key: userid, extract(minute FROM eventtime))
-> Seq Scan on hits (Filter: (url ~~ '%google%'))阶段 2：流水线图
接下来，Gravity 将查询计划转换为“流水线图”（pipeline graph）。这是一种中间表示，用于表达数据在整个查询执行过程中的流动方式。每个流水线代表在整个查询过程中行是如何被转换和操作的。
以下是我们的查询对应的流水线图的草图：
graph pipelines=3 breakers=2 root=p2
p0 eligible end=breaker(slot=0 HashAgg)
source range rel=hits am=pgrc
lanes Attr(UserID):byval8 Attr(EventTime):byval8 Attr(URL):dict4
filter like(URL, param0)
project (key0=UserID, key1=extract_minute(EventTime), agg0=min(URL), agg1=count_star())
p1 serial end=breaker(slot=1 TopK)
source breaker slot=0
topk k=10 order=agg1:DESC
p2 serial end=sink
source breaker slot=1
project (UserID, m, min_url, count) -> wire在这种情况下，查询被拆分为三个流水线：
- p0 从 hits 表中取出数据行，过滤出 url 包含“google”的行，然后将这些行插入到哈希表中，以计算按 user_id 和事件时间分钟分组的 COUNT(*) 和 MIN(url)
- p1 获取生成的哈希表，并将其放入大小为 10 的堆中，以计算 COUNT(*) 最高的 10 个分组
- p2 获取生成的堆并提取出行，然后将其通过网络发送出去
Stage 3: Lisp IR
接下来，Gravity 将每个流水线编译为我们基于 Lisp 的、类似汇编的中间表示中的单个函数。由于 IR 是通用的，我们可以将 p0 内部的扫描、过滤、投影和聚合融合到我们想要的任何循环结构中。以下是 p0 的部分代码示例：
;; ── loop head: one trip per row ─────────────────────────────── (Seq Scan)
(block b210 ((v2574 i64) ...) ; v2574 = row index; loop state rides along as block params
(v2585 (icmp slt v2574 v103)) ; more rows in this batch?
(condbr v2585 (b211) (b35 ...))) ; no -> fetch the next batch
(block b213 ((v2587 i64)) ; ── read the row's columns ── (Seq Scan)
(v2591 (ptr_add v105 v2587 8 0)) ; &EventTime[row]
(v2592 (ptr_add v109 v2587 8 0)) ; &UserID[row]
(v2593 (ptr_add v113 v2587 4 0)) ; &URL[row] (dictionary code)
(v2612 (loadz i32 v2593 0)) ; ── WHERE URL LIKE '%google%' ── (Filter)
(v2616 (load i64 (ptr_add v2605 (ushr v2612 6) 8 0)))
(v2621 (iand (ushr v2616 (iand v2612 63)) 1)) ; one bit per distinct URL
;; … ; row fails -> straight to the next row
)
(block b227 () ; ── extract(minute FROM EventTime) ── (Projection)
(v2695 (load i64 v2591 0))
(v2711 (iadd v2695 211813660800000000))
(v2715 (umulh (ushr v2711 10) 10746007285463371))
(v2720 (isub v2711 (imul (ushr v2715 11) 3600000000)))
(v2726 (ushr (imul (ushr v2720 8) 18764999) 42)) ; minute
(v2732 (load i64 v2592 0)) ; UserID
;; …
)
(block b229 () ; ── GROUP BY UserID, m ── (HashAggregate)
(v2897 (icmp eq v2576 v2732)) ; same UserID as the group in registers?
(v2898 (icmp eq v2577 v2739)) ; same minute?
(v2899 (iand v2897 v2898))
(condbr v2901 (b249) (b250))) ; no -> write it to the hash table (b250), start a new group
(block b249 () ; ── MIN(URL), COUNT(*) ── (HashAggregate)
(v2903 (loadz i32 v2593 0)) ; this row's URL code
(v2911 (icmp ule v2903 (iand v2581 4294967295)))
(v2912 (select v2911 v2908 v2581)) ; MIN: keep the smaller code
(v2921 (iadd v2584 1)) ; COUNT(*) += 1
(br b228 v2576 v2577 v2578 v2579 v12 v2912 v2913 v2915 v2921))
(block b228 (...) ; ── next row ──
(br b214 ...)) ; -> b214: row index + 1 -> back to b210Stage 4: Assembly
最后，我们为每个目标平台提供了后端，将 IR 编译为汇编代码：
;; ── b210 loop head ─────────────────────────────────────────── (Seq Scan)
2b3c cmp x10, x24 ; v2585 = row < rows_in_batch
2b40 b.lt 0x2b70 ; condbr v2585 (b211) (b35)
;; ── b213 read the row's columns ───────────────────────────── (Seq Scan)
; v2591/v2592/v2593 (ptr_add) emit no instructions:
; they fold into the indexed loads below ([base, row, lsl #n])
2b8c ldr w6, [x27, x5, lsl #2] ; v2612 = URL[row] (loadz i32)
;; ── WHERE URL LIKE '%google%' ──────────────────────────────── (Filter)
2b90 lsr x7, x6, #6 ; code >> 6
2b94 ldr x8, [x0, x7, lsl #3] ; v2616 = words[code >> 6]
2b98 and x0, x6, #0x3f ; code & 63
2b9c lsr x7, x8, x0 ; word >> (code & 63)
2ba0 and x7, x7, #0x1 ; v2621 = matched bit
;; ── b227 extract(minute FROM EventTime) ───────────────────── (Projection)
2e04 ldr x1, [x25, x5, lsl #3] ; v2695 = EventTime[row]
2e38 mov x7, #0x2000 ; ┐
2e3c movk x7, #0xfa8f, lsl #16 ; │ constant 211813660800000000
2e40 movk x7, #0x8366, lsl #32 ; │
2e44 movk x7, #0x2f0, lsl #48 ; ┘
2e48 add x1, x1, x7 ; v2711
2e4c lsr x7, x1, #10 ; v2711 >> 10
2e50 mov x8, #0xf94b ; ┐
2e54 movk x8, #0xcb00, lsl #16 ; │ constant 10746007285463371
2e58 movk x8, #0x2d6f, lsl #32 ; │
2e5c movk x8, #0x26, lsl #48 ; ┘
2e60 umulh x7, x7, x8 ; v2715
2e64 lsr x7, x7, #11 ; v2715 >> 11 (hours)
2e68 mov x8, #0xa400 ; ┐ constant 3600000000
2e6c movk x8, #0xd693, lsl #16 ; ┘
2e70 mul x7, x7, x8 ; hours * 3600000000
2e74 sub x1, x1, x7 ; v2720 = µs into the hour
2e78 lsr x1, x1, #8 ; v2720 >> 8
2e7c mov x7, #0x54c7 ; ┐ constant 18764999
2e80 movk x7, #0x11e, lsl #16 ; ┘
2e84 mul x1, x1, x7
2e88 lsr x1, x1, #42 ; v2726 = minute
2e9c ldr x0, [x26, x5, lsl #3] ; v2732 = UserID[row]
;; ── b229 same group as the run in registers? ──────────────── (HashAggregate)
2eb0 cmp x12, x0 ; v2897 = run.UserID == UserID
2eb4 ccmp x13, x1, #0x0, eq ; v2898, v2899 = && run.minute == minute
2eb8 ccmp x14, x6, #0x0, eq ; v2901 = && run.flags == flags (the third key compare, cut from the IR sample)
2ec8 b.ne 0x2efc ; condbr v2901 (b249) (b250)
;; ── b249 MIN(URL), COUNT(*) ───────────────────────────────── (HashAggregate)
2ecc ldr w7, [x27, x5, lsl #2] ; v2903 = URL[row] (loaded again)
2ed0 ldr w8, [x19, #264] ; ┐ v2908 = (part << 32) | code
2ed4 orr x0, x7, x8, lsl #32 ; ┘ (the select's operand; its definition is cut from the IR sample)
2ed8 and x8, x30, #0xffffffff ; v2581 & 0xFFFFFFFF (run.min lives in x30)
2edc cmp x7, x8 ; ┐ v2911 = code <= run.min's code
2ee0 cset x7, ls ; ┘
2ee4 cmp x7, #0x0 ; ┐ v2912 = select(v2911, v2908, v2581)
2ee8 csel x30, x0, x30, ne ; ┘
2ef0 add x4, x4, #0x1 ; v2921 = COUNT(*) + 1 (count lives in x4)
2ef8 b 0x337c ; br b228 …
;; ── b228 → b214 next row ──────────────────────────────────────
337c mov x18, #0x1 ; block argument v12 ("a run is open")
3380 add x10, x10, #0x1 ; row + 1
3384 b 0x2b3c ; back to b210从底层开始：首先设计 Lisp IR 并进行基准测试
我们构建 Gravity 的第一部分是后端。在构建 SQL 编译器之前，我们希望证明我们的 IR 和代码生成器确实能够媲美 rustc -O3。一旦 Lisp IR 和后端开始工作，我们就通过让 Claude 编写优化后的 IR 程序（该程序镜像了之前手写的 Rust 代码）来对其进行测试。
对这些手写的 IR 程序进行基准测试带给我们一个重大的好消息：我们的 Lisp IR 在不到一毫秒的时间内完成编译，其运行速度达到了由 rustc 编译的手写 Rust 代码的 80% 到 90%（即性能差距在 10% 到 20% 以内）。更重要的是，在我们测试的几乎所有查询中，我们都不需要发明任何花哨的新 IR 指令。该 IR 具有足够的表达能力，我们的简易 SQL 引擎与“光速”之间的整个差距，归根结底就在于将 SQL 查询编译为 IR。
为规划器引入新算法
当我们把 Gravity 与 Claude 手写的 Rust 代码在我们的查询语料库中进行对比时，我们发现有些查询中手写的代码速度仍然快了几个数量级，这并不是因为循环的编译方式，而是因为它运行的是根本上不同的算法。尽管我们保留了 Postgres 的规划器，但标准的 Postgres 规划器节点无法表达这些光速算法。我们必须确保我们的规划器能够对它们进行编码。
在我们优化的查询中，我们发现有三个算法需要添加到规划器中。这三个算法都是针对列式格式（columnar format）特有的，因此 Postgres 规划器无法支持它们也是合情合理的。请记住，对于列式表，通常会指定一个排序列，因此数据通常按相似值彼此靠近的顺序排列。这三种算法是：
- ZoneWalk - 对于 ORDER BY … LIMIT n 这样的查询，Postgres 通常会通过使用索引或扫描整张表来执行。pgrcolumnar 启用了一种执行该查询的额外方式。pgrcolumnar 将表分成每批 8192 行的块，称为 granule。每个 granule 都有元数据，用于跟踪该批数据的最小值和最大值。我们不扫描整表并排序来查找前 n 行，而是使用最小/最大边界按顺序读取 granule，并在解析出 n 行后立即停止。
- ZoneTopK - 对于形如 GROUP BY ... LIMIT n 的查询，我们可以使用一种与 ZoneWalk 非常相似的算法。我们不需要扫描整张表并对每一行进行哈希运算只为了返回 n 个分组，而是可以利用元数据按 GROUP BY 键的顺序遍历 granule。一旦我们看到了 n 个完整的分组，并且它们的最大可能键小于所有剩余批次的最小键，我们就可以立即停止扫描。ClickHouse 26.10 具有类似的的优化
enable_group_by_top_k_dynamic_filtering，它同样要求完整的分组。 - TwoPass - For a query of the shape GROUP BY ... ORDER BY count(*) DESC LIMIT 10, if one of the columns in the group by has high cardinality, it can be faster to execute the query in two passes:
- 通过计算每个键的行数来计算每个分组大小的上界：SELECT unique_key, COUNT(*) ... GROUP BY unique_key
- 然后返回并仅获取那些上界足够大、有可能进入前 10 名的分组
之所以可行，是因为查询中最昂贵的部分是第一趟扫描。通过在第一趟扫描中只获取一列，然后在第二趟中获取其余列，你能够节省大量的内存带宽
这三个计划都作为成本估算计划被添加到了 pgrust 规划器中。在规划阶段，规划器会估计遍历在停止前实际需要读取多少个 granule，并且只有在重写方案比标准扫描更便宜时才会选择它。
既然我们的查询使用了与光速实现相同的算法，并且我们知道我们的 IR 编译查询计划的速度足够快，那么最后一步就是将 SQL 编译为我们的 IR。
优化 SQL→IR 编译
对于 SQL→IR 编译，最重要的是我们能把 IR 优化到什么程度。因为我们使用了与光速 Rust 程序相同的算法，所以我们可以将我们的代码与 Rust 代码进行 diff 对比，看看差异从何而来。这带来了 100 多项优化。以下是其中三项，让你对 Gravity 为何如此之快有个概念。
1. 内联函数调用和聚合
在标准的 Postgres 中，对每一行评估像 extract(minute FROM EventTime) 这样的 SQL 函数或更新像 AVG 或 MIN 这样的聚合，都会通过 Postgres 的函数管理器（fmgr）进行函数调用。每行进行一次函数调用（甚至每小批次一次）会彻底摧毁性能。
为了解决这个问题，我们将 Postgres 的所有内置函数按层级分类，并针对 238 个核心函数和聚合，将它们直接内联到我们的 Lisp IR 中。带有 COUNT(*) 和 SUM 的查询将会变成类似下面这样：
(block b32 ()
(v313 (load i64 v307 24)) ; load count from aggregate state
(v314 (iadd v313 v194)) ; count + 1
((store v307 24 v314))
(v315 (load i64 v307 32)) ; load sum from aggregate state
(v316 (iadd v315 v309)) ; sum + val
((store v307 32 v316))
(br b22 v197))对 AVG 的性能进行基准测试，内联使循环从每行 392 条指令降至每行 120 条指令，并使查询速度提升了 2.56 倍。
2. 延迟物化
在我们的列式格式中，诸如 URL 之类的文本列采用字典编码：表的每个部分存储一个唯一字符串的字典，而行本身存储该字典中的整数代码。
早期版本的 Gravity 仅根据表达式的 SQL 类型（TEXT、INT4、NUMERIC）来决定如何为表达式生成 IR。因为 url 具有 SQL 类型 TEXT，所以 Gravity 将字典代码解码为字符串并将其复制到每一行中。当我们查看 Claude 为这些查询手写的 Rust 代码时，它在行循环中从不触及字符串。它使用了延迟物化。它尽可能长时间地对字典代码进行操作，直到最后才将实际的字符串物化。
我们重新设计了 Gravity，使每个值都尽可能晚地进行实例化。在计算每一行的过滤器（例如 url LIKE '%google%'）时，编译器会生成一个步骤，针对字典中的每个唯一条目运行一次高度优化的子串匹配器，并创建一个位图，记录哪些字典代码符合该过滤条件。默认情况下，v0.4-beta 会将该位图保存在服务器内存中，以包含常数的过滤器和表分区作为键，因此随后在未更改的数据上使用相同 LIKE 的查询会直接重用它，而无需再次运行匹配器。设置 pgrust.dict_filter_cache = off，或将 pgrust.condition_cache 设置为 granule 或 off，则会使每个查询构建自己的位图。
为了计算像 MIN(URL) 这样的聚合，Gravity 不会逐行比较字符串。压缩字典是经过排序的，因此对于从同一字典编码的两个字符串，代码较小的字符串对应的字符串也较小。正因为如此，当在用同一字典压缩的一系列行中计算最小字符串时，它主要变成了整数比较，而不是字符串比较。
3. 在寄存器中折叠
即使我们的过滤器和表达式被内联，在数百万行数据上探测多键哈希表代价依然高昂。
当我们把 Gravity 的输出与 Claude 手写的 Rust 在 GROUP BY ... 查询上的表现进行对比时，我们注意到了一些有趣的事情。在真实世界的表中，像 UserID 或事件时间戳这样的列的值通常在磁盘上部分聚簇，这意味着我们经常会看到属于完全相同组的连续行。Claude 的 Rust 代码利用了这一点：将当前组的键和运行计数保存在局部变量中（rustc 将其保留在 CPU 寄存器中）。只要下一行的键与上一行相同，它就只在寄存器中递增计数器，并且仅在连续行结束时才探测哈希表。在一个拥有 1亿 行的表中，这把 1亿 次哈希表探测压缩到了 2200 万次。
我们修改了 Gravity，使其能够追踪先前的组键，并且仅在遇到新的组键时才探测哈希表。
确保其具备通用性
我们在构建 Gravity 时面临的一个挑战是，Claude 倾向于将优化过度拟合到我们的语料库查询中。通常，Claude 编写的优化仅针对我们正在进行基准测试的精确查询。如果用户编写 strpos(URL, 'google') > 0 而不是 URL LIKE '%google%'，或者对包裹在 Postgres DOMAIN 类型中的列进行分组，那么脆弱的模式匹配器就会错过快速路径，并静默回退到慢速代码。
为了解决这个问题，我们建立了一个包含所有优化的目录。然后，我们编写了一个模糊测试器（fuzzer），它获取单个优化并在许多不同的查询形态上进行测试。它会尝试各种类型（int2、int4、int8、numeric、text、domain）、group by 等等，以确认该优化在许多查询中都能生效，而不仅仅是一个查询。Dan Luu wrote about 介绍了 pgrust 需要进行多少测试，这也暴露了 Postgres 本身真实的 Bug。
Gravity 的现状
首先针对手写的 Rust 设计 Lisp IR，然后教导 Gravity 生成内联、延迟实例化且寄存器优化的 IR，这让我们非常接近“光速”目标。在广泛的分析查询中，Gravity 的运行时 JIT 在一毫秒内编译查询，其执行性能达到了 Claude 直接针对我们的文件格式手写的 Rust 代码的 10% 到 20% 以内。
当我们隔离查询引擎，在开启 Gravity 的情况下对完全相同的列式表运行完全相同的 pgrust 二进制文件，并与我们的 Volcano 行执行器进行对比时，Gravity 在冷启动时快 102 倍，热启动时快 499 倍，其运行速度比 ClickHouse 快多达 3 倍，同时支持 Postgres 完整的 SQL 方言和类型系统。对于任何来自 Postgres 的用户：在 ClickBench 的 c6a.4xlarge 机器上，pgrust v0.4-beta 的综合得分比 PostgreSQL 17 的得分高出 1,000 倍以上，在冷启动运行中则高出 1,365 倍。
如果你想深入了解 Lisp IR、TPDE 后端或降级器，可以查看代码并在 github.com/malisper/pgrust 上关注我们的定期提交。
关于我们早期的 v0.4-preview ClickBench 条目的一点说明：我撰文记录了 Hacker News 上发生的事情。
来源：Hacker News-首页 · pgrust.com