让Postgres分析性能提升300倍：批处理、算子融合与SIMD
- 如今许多数据集都能装进内存，大部分磁盘 I/O 不再存在；
- 对于装不进内存的数据集，工作负载也有所不同：数据分析通常是大批量扫描，瓶颈往往不再是磁盘吞吐，而是 CPU 吞吐或内存吞吐；
- 近些年磁盘本身也快得多了，NVMe 比机械硬盘快几百倍。
为了让你直观感受 Postgres 查询引擎有多慢，我们来看一个简单的查询——对前 5 亿个数求和：
CREATE TABLE my_table AS select col::float8 from generate_series(1.0, 500000000.0) g(col); SELECT SUM(col) FROM my_table;我在 Postgres 上跑这条查询大约需要 20 秒。测试环境是 c8g.4xl，且关闭了并行查询。 作为对比，下面是 Rust 中等价代码的耗时：
let table: Vec<f64> = (1..=500_000_000usize).map(|i| i as f64).collect();
let mut sum = 0.0;
for &value in &table {
sum += value;
}
这个查询耗时 358 毫秒，比之前快了大约 55 倍。信不信由你，我们还能比 358 毫秒更快。不过这个例子并不完全公平——Postgres 内部做的事情远比这复杂得多。但话说回来，数据库优化的核心就是想方设法消除这些开销。（如果你好奇的话，Postgres 中两个最大的开销来源是：1. 加锁；2. 解析 Postgres 存储格式并提取与查询相关的元组。）
为了把焦点收窄到查询引擎本身的影响，我们来构建一个 Postgres 查询引擎的迷你版。先简单介绍一下什么是查询引擎。Postgres 处理 SQL 查询时，首先会把查询转换成一种内部表示，也就是"查询计划"（Query Plan），用来描述 Postgres 将如何执行查询。对于上面那个例子，Postgres 生成的查询计划大致如下：
这个计划表达的意思就是"从 my_table 中取出所有行，对这些行里的值求和"。由于查询本身很简单，这个查询计划也很朴素；但一旦涉及 join、排序、子查询等操作，查询计划就会变得复杂得多。Postgres 一共有超过 40 种不同类型的计划节点。
生成查询计划之后，Postgres 会把它交给查询引擎。查询引擎是 Postgres 中真正负责按照计划取出数据行并执行聚合运算的部分。Postgres 采用的是一种称为"火山模型"（Volcano model）的执行器。为了让你直观地理解它的工作方式，下面给出一个 Postgres 查询引擎的迷你实现：
use std::hint::black_box;
trait Node {
fn next(&mut self) -> Option;
}
struct SeqScan<'a> {
table: &'a [f64],
pos: usize,
}
impl Node for SeqScan<'_> {
fn next(&mut self) -> Option {
if self.pos >= self.table.len() {
return None; // end of table
}
let value = self.table[self.pos];
self.pos += 1;
Some(value)
}
}
struct SumAggregate<'a> {
child: Box,
total: f64,
done: bool,
}
impl Node for SumAggregate<'_> {
fn next(&mut self) -> Option {
if self.done {
return None;
}
while let Some(value) = self.child.next() {
self.total += value;
}
self.done = true;
Some(self.total)
}
}
let table: Vec = (1..=500_000_000usize).map(|i| i as f64).collect();
let mut plan = SumAggregate {
child: black_box(Box::new(SeqScan { table: &table, pos: 0 })),
total: 0.0,
done: false,
};
let sum = plan.next().unwrap();
（需要用 black_box 来防止编译器优化干扰基准测试结果）
Volcano 模型的核心特性是 `next()` 方法，查询计划中的所有节点都支持该方法。`next()` 的职责是返回单行数据：顺序扫描中的 `next()` 返回下一行数据；聚合操作中的 `next()` 则会完成整个聚合计算，然后返回单行结果。执行查询计划只需在根节点上反复调用 `next()`，直到它不再返回任何行。Volcano 模型的优点是非常简单：每个计划节点只需实现一个方法就够了。虽然做了简化，但上面的代码与 Postgres 内部实现非常接近。
虽然 Volcano 模型简单，但同时也带来了不少开销。我运行这个示例用了 1.3 秒。虽然去掉了大量非查询引擎的部分后比 Postgres 原版快很多，但由于 Volcano 模型的开销，仍然比裸的 for 循环慢。
上面这段代码最大的性能瓶颈在于 `next()` 一次只处理一行，没有批处理机制。`SeqScan.next()` 每行都会被调用一次，这会带来很大的额外开销——尤其因为很多 CPU 优化（比如流水线技术）在调用一个要到运行时才能确定的函数时很难发挥作用。我们可以做的第一个优化就是引入批处理：
const BATCH: usize = 1024;
trait BatchNode {
fn next_batch(&mut self, out: &mut [f64; BATCH]) -> usize;
}
struct BatchSeqScan<'a> {
table: &'a [f64],
pos: usize,
}
impl BatchNode for BatchSeqScan<'_> {
fn next_batch(&mut self, out: &mut [f64; BATCH]) -> usize {
let n = (self.table.len() - self.pos).min(BATCH);
out[..n].copy_from_slice(&self.table[self.pos..self.pos + n]);
self.pos += n;
n
}
}
struct BatchSumAggregate<'a> {
child: Box,
total: f64,
}
impl BatchSumAggregate<'_> {
fn run(&mut self) -> f64 {
let mut buf = [0.0f64; BATCH];
loop {
let n = self.child.next_batch(&mut buf);
if n == 0 {
break;
}
for &value in &buf[..n] {
self.total += value;
}
}
self.total
}
}
let mut plan = BatchSumAggregate {
child: black_box(Box::new(BatchSeqScan { table: &table, pos: 0 })),
total: 0.0,
};
let sum = plan.run();
光是批处理就消除了大部分开销，查询执行时间从 1.3 秒降到了大约 480ms。虽然还是比 for 循环慢，但已经接近了很多。一个非常关键的细节是：批处理缓冲区分配在栈上，因此聚合节点在运行过程中完全不需要分配内存。内存分配通常是较慢的操作之一，所以当你编写极致性能的代码时，应当尽量减少内存分配的次数。
现在，如果你对批处理版本做性能分析，会发现瓶颈转移到了 `copy_from_slice` 上。即便我们已经做了批处理，把数据拷贝进缓冲区这一步仍然有开销。这可以用所谓的「算子融合」来消除：如果我们事先知道某些操作总是会一起执行，就可以合并成一个节点来替代原来的两个节点。在这个例子里，我们可以做一个 `SumAggregateSequentialScan` 节点，把顺序扫描和求和的逻辑合在一起：
struct SumAggregateSequentialScan<'a> {
table: &'a [f64],
done: bool,
}
impl Node for SumAggregateSequentialScan<'_> {
fn next(&mut self) -> Option<f64> {
if self.done {
return None;
}
self.done = true;
let mut total = 0.0;
for &value in self.table {
total += value;
}
Some(total)
}
}
这段代码的性能和直接的 for 循环完全一样，因为它本质上就是同一个循环。这看起来有点「作弊」的味道，事实上也确实是作弊——我们针对已知的特定查询做了硬编码优化。用算子融合的方式硬编码几个最常见的场景是合理的，但你很快就会遇到没预先准备过的情形。
这可以用 JIT 编译来解决。有了 JIT 编译，你就能为查询生成理想的代码，并在每一条查询上都「作弊」。无论查询长什么样，JIT 编译都能生成最优代码并始终执行算子融合。可惜这篇文章已经够长了，关于 pgrust 如何利用 JIT 编译就只能下次再聊了。
最后再看一个优化方向：SIMD。SIMD 指的是一组 CPU 指令，能在多条数据上同时执行同一个操作。用 SIMD 对多行数据并行处理，通常比逐行处理快很多。如果把上面的代码改成 SIMD 版本：
#[cfg(target_arch = "aarch64")]
struct SumAggregateSequentialScanSimd<'a> {
table: &'a [f64],
done: bool,
}
#[cfg(target_arch = "aarch64")]
impl Node for SumAggregateSequentialScanSimd<'_> {
fn next(&mut self) -> Option<f64> {
if self.done {
return None;
}
self.done = true;
use std::arch::aarch64::*;
let mut acc = unsafe { [vdupq_n_f64(0.0); 4] };
let (chunks, rest) = self.table.as_chunks::<8>();
for chunk in chunks {
for lane in 0..4 {
unsafe {
let v = vld1q_f64(chunk.as_ptr().add(2 * lane));
acc[lane] = vaddq_f64(acc[lane], v);
}
}
}
let mut tail = 0.0;
for &value in rest {
tail += value;
}
Some(unsafe {
let s01 = vaddq_f64(acc[0], acc[1]);
let s23 = vaddq_f64(acc[2], acc[3]);
vaddvq_f64(vaddq_f64(s01, s23)) + tail
})
}
}
我们的代码现在耗时 135 毫秒，比 for 循环快了近 3 倍，比最初的 Volcano 代码快了 10 倍。虽然编译器通常会用 SIMD 替代 for 循环，但我特意选了这样一个例子来避免这种情况。编译器在处理浮点数时一般不会引入 SIMD，因为结果会略有不同。这是因为浮点运算不满足结合律，改变求和顺序可能导致结果出现细微差异。
—
总的来说，通过三个简单的优化，我们让查询速度提升了 10 倍。最终结果如下：
| 实现方式 | 耗时 | 加速比 |
|---|---|---|
| Postgres | 约 20 秒 | — |
| Volcano 模型 | 1.3 秒 | 1× |
| + 批处理 | 480 毫秒 | 2.7× |
| + 算子融合 | 358 毫秒 | 3.6× |
| + SIMD | 135 毫秒 | 9.6× |
这些优化手段，再加上其他更多技巧，让 pgrust 在处理分析型查询时比 Postgres 快出数百倍。
基准测试环境：AWS c8g.4xlarge（Graviton4，16 vCPU），PostgreSQL 18.4，`max_parallel_workers_per_gather = 0`，数据预热于共享缓冲区，取5次运行的中位数。Rust 通过 `cargo build –release` 构建，每个实现运行4次，所有测量均在同一台机器的同一进程内完成。
感谢阅读，若想支持该项目，支持 pgrust 的最佳方式是在 GitHub 上为我们点星。如需持续关注：
1. GitHub
- Discord
- 邮件列表 — 每周更新 pgrust 动态，包括 JIT 编译的后续进展
- pgrust.com