Postgres分析查询提速300倍：批处理、算子融合与SIMD

上周我们发布了 pgrust 的 0.2 版本。这个版本主要关注性能。它比之前的 pgrust 版本快 10 倍。在 OLTP 基准测试中，pgrust 比 Postgres 快 30%，而在 Clickbench（Clickhouse 的 analytical 数据库基准测试）中，pgrust 比 Postgres 快 300 倍。它甚至领先于 Clickhouse！查询引擎是我们为实现更好性能所做的最大改动之一。查询引擎本身贡献了 300 倍中的约 10 倍。我们将从 Postgres 查询引擎的微型版本开始，然后逐一添加我们为使 pgrust 查询引擎如此之快而做的优化。为了说明为什么相比 Postgres 有如此大的改进空间，Postgres 是在一个不同的时代创建的。最初的 Postgres 项目可以追溯到 80 年代。它是在数据库性能的主要瓶颈是磁盘 I/O 的时代构建的。三个趋势使这种情况不再成立：- 许多数据集现在可以放入内存，消除了大部分磁盘 I/O - 对于无法放入内存的数据集，工作负载有所不同。数据分析会批量扫描数据。瓶颈往往不再是磁盘吞吐量，而通常是 CPU 吞吐量或内存吞吐量 - 近年来磁盘变得更快。NVMe 比硬盘快数百倍。这三个趋势都使得 CPU 和内存速度比历史上更重要。我们做的许多优化都针对这一点。查询引擎是数据库中 CPU 的主要使用者。我们优化了 pgrust 查询引擎，使其在处理相同查询时比 Postgres 使用更少的 CPU 和内存带宽。为了让您了解 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 = (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 查询引擎的微型版本。首先，简要解释一下什么是查询引擎。在处理你的 SQL 查询时，Postgres 首先将你的查询转换为一种内部表示，称为“查询计划”，它描述了 Postgres 将如何执行查询。在上面的例子中，Postgres 将生成一个查询计划，可能看起来像下面这样：这实际上表示“从 my_table 获取行，并对这些行中的值求和”。鉴于查询的性质，这个查询计划非常简单，但当你开始处理连接、排序、子查询等时，它们会变得复杂得多。总的来说，Postgres 有超过 40 种不同类型的计划节点。生成查询计划后，Postgres 将其传递给查询引擎。Postgres 查询引擎是 Postgres 的一部分，它接收查询计划并实际检索行并执行聚合。Postgres 使用一种称为“Volcano 模型”的执行器风格。为了了解它是如何工作的，这里有一个 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; // 表结束 } 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秒缩短到约480毫秒。虽然仍比for循环慢，但已接近很多。一个非常重要的细节是，批处理缓冲区是在栈上分配的。这意味着聚合节点在运行期间无需分配任何内存。内存分配往往是较慢的操作之一。因此，在编写超高速代码时，应尽量减少内存分配的次数。