pgrust v0.4-beta 发布 Gravity 查询引擎,ClickBench 上最高比 ClickHouse 快 3 倍
Programming at Light Speed: Building the Gravity Query Engine
pgrust v0.4-beta 发布新查询引擎 Gravity,在支持完整 Postgres SQL 方言和类型系统的同时,ClickBench 性能最高比 ClickHouse 快 3 倍。
文章拆解了 Gravity 从 Postgres 查询计划到 Lisp IR 和机器码的编译链路,并说明函数内联、延迟物化与列式规划算法的实现。
随着 pgrust v0.4-beta 发布,我们的新查询引擎 Gravity 现在在 ClickBench 上比 ClickHouse 快至 3 倍。ClickBench 是 ClickHouse 团队专门设计、用来展示 ClickHouse 有多快的基准测试。更出色的是,Gravity 在支持完整 Postgres SQL 方言和类型系统的同时,仍能达到这一速度。
在此前的几篇文章中,我写过我们如何利用批处理、算子融合和 SIMD,让 pgrust 执行分析型查询的速度比 Postgres 快数百倍,以及我们如何构建了一个 copy-and-patch JIT 编译器,仅需 5 微秒即可编译模板。所有这些经过多个版本积累的工作,最终成就了 Gravity。
copy-and-patch 是一个很好的起点,但拼接预编译模板所能带来的提升终究有限。由于我们没有真正的中间表示,也没有在整个循环中进行寄存器分配,因此与 rustc 或 Clang 这类优化型编译器相比,我们浪费了大量性能空间。对于 pgrust v0.4,我们决定更进一步,构建一个能为任意 SQL 查询生成最优机器代码的查询编译器。这正是我们打造 Gravity 的原因。
通常,工程师试图提升软件性能时,会采用自底向上的方式。他们运行基准测试、查看火焰图、找出热点,再想办法将代码提速 15% 或 20%。问题在于,采用自底向上的方式时,你不知道性能上限在哪里。你可能距离硬件所能运行的最快代码只差 20%,也可能还相差 100 倍。
开始打造 Gravity 时,我们把这个问题反过来问:我们问自己,数据库究竟能有多快?我们不想只是逐步提升现有查询引擎的速度,而是要从硬件和文件格式的物理速度极限出发,反向推导实现方案。我喜欢把这称为“光速编程”。
为了弄清楚“光速”究竟是什么样子,我们让 Claude 针对一组分析型查询,直接基于我们的列式存储 crate 编写独立且经过深度优化的 Rust 程序。这些程序绕过整个 SQL 执行器,直接在我们的磁盘格式上运行经过手工调优、使用 SIMD 的多线程 Rust 循环。运行这些手写 Rust 程序时,我们发现它们执行查询的速度整体上大约比 ClickHouse 等现有分析型数据库快 4 倍。随后,我们从这些手工调优的 Rust 循环中挖掘可纳入 Gravity 的优化。
Gravity 的目标变得很简单:针对任意 SQL 查询,在运行时 JIT 编译机器代码,使其与 Claude 直接针对我们的文件格式编写的手工 Rust 代码之间的性能差距不超过 20%。
Gravity 的总体架构
设计 Gravity 架构时,我们的主要目标是让生成的代码拥有极大的灵活性。我们知道,如果能轻松修改 Gravity,使其输出查询所需的任何底层优化,那么缩小差距、接近光速就会容易得多。
我们没有将执行模式写死,也没有拼接固定的模板,而是把 Gravity 构建为一条多阶段编译器流水线,逐步将高层 Postgres 计划降低为原始汇编代码:
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 接收生成的堆,提取其中的行并通过网络发送出去
阶段 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 b210阶段 4:汇编
最后,我们为每个目标平台配备后端,用于将 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 组件是后端。我们希望先证明 IR 和代码生成器确实能够匹配 rustc -O3,然后才在其上构建 SQL 编译器。当 Lisp IR 和后端能够正常运行后,我们让 Claude 编写经过优化的 IR 程序,对应此前手写的 Rust 代码,并以此进行测试。
对这些手写 IR 程序进行基准测试后,我们得到了一条令人振奋的消息:我们的 Lisp IR 编译耗时不到一毫秒,运行性能与 rustc 编译的手写 Rust 代码仅相差 10% 到 20%。更重要的是,在几乎所有测试过的查询中,我们都不需要设计任何花哨的新 IR 指令。IR 的表达能力已经足够强,使我们朴素的 SQL 引擎与“light speed”之间的全部差距都归结为将 SQL 查询编译到 IR。
让规划器学习新算法
当我们在查询语料库中将 Gravity 与 Claude 手写的 Rust 代码进行对比时,发现有些查询中手写代码仍然快上几个数量级,原因不在于其循环的编译方式,而在于它采用了完全不同的算法。尽管我们保留了 Postgres 的规划器,但标准的 Postgres 计划节点无法表达 light-speed 算法。我们必须确保计划能够对这些算法进行编码。
在我们优化过的查询中,我们发现需要向规划器添加三种算法。这三种算法都专用于列式格式,因此 Postgres 规划器无法支持它们也就不足为奇。请记住,对于列式表,通常需要指定排序键,因此数据通常会按将相似值彼此靠近的顺序排列。这三种算法是:
- ZoneWalk - 对于类似 ORDER BY … LIMIT n 的查询,Postgres 会通过使用索引或扫描整张表来执行。pgrcolumnar 提供了执行该查询的另一种方式。pgrcolumnar 将表拆分成每批 8192 行的批次,称为 granule。每个 granule 都包含元数据,用于跟踪该批次的最小值和最大值。我们不再扫描整张表并排序来查找前 n 行,而是利用最小值/最大值边界按顺序读取各个 granule,并在确定 n 行后立即停止。
- ZoneTopK - 对于形如 GROUP BY ... LIMIT n 的查询,我们可以使用一种与 ZoneWalk 非常相似的算法。我们不必为了返回 n 个分组而扫描整张表并对每一行进行哈希处理,而是可以利用元数据,按 GROUP BY 键的顺序遍历各个数据粒。一旦我们已经看到 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 规划器。在制定计划时,规划器会估算一次遍历在停止前实际必须读取多少个数据粒,并且只有在重写方案比标准扫描更便宜时才会选择它。
现在,我们的查询已经采用与 light speed 实现相同的算法,而且我们已知 IR 足够快,能够编译查询计划,最后一步便是将 SQL 编译到我们的 IR。
优化 SQL→IR 编译
对于 SQL→IR 编译,唯一重要的就是能把 IR 优化到何种程度。由于我们采用了与 light-speed Rust 程序相同的算法,因此可以将我们的代码与 Rust 代码进行差异比较,找出差异的来源。这带来了 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. 在寄存器中运行 Collapse
即使过滤器和表达式已经内联,在数百万行数据上探测多键哈希表仍然代价高昂。
我们将 Gravity 在 GROUP BY ... 查询中的输出与 Claude 手写的 Rust 实现进行比较时,发现了一个有趣的现象。在真实世界的表中,UserID 或事件时间戳等列的值通常会在磁盘上部分聚簇,这意味着我们经常会看到连续多行属于完全相同的组。Claude 的 Rust 代码利用了这一点:它将当前组的键和运行计数保存在局部变量中(rustc 会将这些变量保留在 CPU 寄存器中)。只要下一行的键与上一行相同,它就只需递增寄存器中的计数器,仅在连续行段结束时探测哈希表。在一张拥有一亿行的表上,这把一亿次哈希表探测减少到了两千二百万次。
我们修改了 Gravity,使其跟踪上一个组的键,并且只在遇到新的组键时才探测哈希表。
确保其能够泛化
构建 Gravity 时,我们面临的一个挑战是,Claude 往往会将这些优化过拟合到我们的语料查询上。Claude 经常只为正在基准测试的那一个查询编写优化。如果用户编写的是 strpos(URL, 'google') > 0,而不是 URL LIKE '%google%',或者按使用 Postgres DOMAIN 类型包装的列进行分组,那么脆弱的模式匹配器就无法命中快速路径,并会悄无声息地回退到慢速代码。
为应对这一问题,我们为每项优化建立了一份目录。随后,我们编写了一个模糊测试器,对每项优化进行测试,并将其应用于多种不同的查询形式。它会尝试各种类型(int2、int4、int8、numeric、text、domain)的 GROUP BY,以及更多形式,以确认该优化能在多种查询中生效,而不只是对一种查询生效。Dan Luu 撰文介绍了 pgrust 的测试工作量,而 pgrust 也发现了 Postgres 本身的真实缺陷。
Gravity 目前的状况
我们先围绕手写的 Rust 实现设计 Lisp IR,随后又让 Gravity 生成经过内联、延迟物化和寄存器优化的 IR,这使我们与“速度极限”目标已经非常接近。在广泛的分析型查询中,Gravity 的运行时 JIT 会在不到一毫秒内完成查询编译,执行表现与 Claude 直接针对我们的文件格式手写的 Rust 代码相比,差距在 10% 到 20% 以内。
当我们通过在完全相同的列式表上运行完全相同的 pgrust 二进制文件,并比较开启 Gravity 与使用我们的 Volcano 行执行器来隔离查询引擎时,Gravity 在冷运行中快 102 倍,在热运行中快 499 倍;同时,它支持 Postgres 的完整 SQL 方言和类型系统,运行速度最高比 ClickHouse 快 3 倍。对于来自 Postgres 的用户:在 ClickBench 的 c6a.4xlarge 机器上,pgrust v0.4-beta 的综合得分比 PostgreSQL 17 高出 1,000 多倍,冷运行中更是高出 1,365 倍。
如果你想深入了解 Lisp IR、TPDE 后端或 lowerer,可以查看代码并关注我们在 github.com/malisper/pgrust 的日常提交。
关于我们之前发布的 v0.4-preview ClickBench 条目,有一点说明:我在 Hacker News 上写下了事情的经过。
来源:Hacker News · pgrust.com