Training a 4B model to produce 81% faster query plans than Postgres
Training a 4B model to produce 81% faster query plans than Postgres
训练一个 4B 模型,生成比 Postgres 快 81% 的查询计划
How good are query optimizers, really? Leis et al. asked this exact question in 2015. Then, they asked it again 10 years later. Despite an enormous body of research spanning a decade since their original exploration, they found that query optimizers continue to leave much to be desired. 查询优化器到底有多好?Leis 等人在 2015 年提出了这个问题,并在 10 年后再次提出。尽管在他们最初的研究之后,过去十年间涌现了大量的研究成果,但他们发现查询优化器仍有很大的改进空间。
I was surprised when I first learned about this. A Postgres database should know everything about the stuff that lives in its tables, no? How hard can it be? As it turns out: enormously hard. In fact, one particular task a query optimizer needs to do, join ordering, is known to be NP-hard. 当我第一次了解到这一点时感到很惊讶。Postgres 数据库难道不应该了解其表中存储的所有内容吗?这能有多难?事实证明:非常难。事实上,查询优化器需要执行的一项特定任务——连接排序(join ordering)——被公认为是 NP-hard 问题。
So query optimizers are hard. What’s not as hard is verifying whether a query plan an optimizer picks is good or not. Put simply, a good query optimizer produces plans that run fast, and a bad one produces slow plans. Language models are particularly good at learning how to do tasks with easily verifiable outputs. 所以查询优化器很难。但验证优化器选择的查询计划好坏并不难。简单来说,好的查询优化器生成的计划运行速度快,而差的优化器生成的计划运行速度慢。语言模型特别擅长学习如何执行那些输出易于验证的任务。
Because there’s a single axis to optimize for—execution time of a query—the problem beautifully reduces to reinforcing the behaviors that guide a model to produce faster query plans. What follows is a breakdown of an experiment I ran to explore the question: can a small, open-weights model be post-trained via supervised fine-tuning (SFT) and agentic reinforcement learning (RL) to produce Postgres query plans that beat Postgres’s default plans? 由于只有一个优化维度——查询执行时间,这个问题可以完美地简化为:强化那些引导模型生成更快查询计划的行为。接下来我将详细介绍我进行的一项实验,旨在探讨:是否可以通过监督微调(SFT)和智能体强化学习(RL)对小型开源权重模型进行后训练,使其生成的 Postgres 查询计划优于 Postgres 的默认计划?
The answer to our question is a resounding yes. Highlights include: 答案是肯定的。实验亮点包括:
- Attaining a 44.7% latency reduction across 113 join-heavy queries from a 4B model initially unable to produce a query plan for 99 of them
- 在 113 个重连接查询中实现了 44.7% 的延迟降低,而该 4B 模型最初甚至无法为其中 99 个查询生成查询计划
- Constructing a Postgres measurement rig that minimizes Linux page cache contention noise across concurrent containers
- 构建了一个 Postgres 测量装置,最大限度地减少了并发容器中 Linux 页面缓存竞争带来的噪声
- Designing a custom GRPO variant for scoring RL rollouts in an inherently noisy environment
- 设计了一种自定义的 GRPO 变体,用于在本质上充满噪声的环境中对 RL 展开(rollouts)进行评分
- Splitting RL across two machines: vLLM and the trainer on a rented 2x H100 node and four Postgres containers running on my desk
- 将 RL 分布在两台机器上:vLLM 和训练器运行在租用的 2x H100 节点上,而四个 Postgres 容器运行在我桌面的机器上
- Running off-policy distillation across half a thousand GPT-6 Astra agent trajectories
- 在 500 条 GPT-6 Astra 智能体轨迹上运行离策略蒸馏(off-policy distillation)
Let’s start from the beginning. Inside a query optimizer 让我们从头开始。深入查询优化器内部
Consider the following slice of the IMDb dataset: 考虑 IMDb 数据集中的以下片段:
-- An IMDb title (movie, series, episode, etc.) [~1M rows]
title ( id integer PRIMARY KEY, title text, production_year integer, kind_id integer -- FK -> kind_type )
-- Movie <> company junction table [~2M rows]
movie_companies ( id integer PRIMARY KEY, movie_id integer, -- FK -> title.id company_id integer, -- FK -> company_name.id company_type_id integer, -- FK -> company_type.id note text )
-- A company's name, origin, etc. [~100k rows]
company_name ( id integer PRIMARY KEY, name text, country_code text -- '[us]', '[jp]', ... )
-- Lookup table of company roles for a title [4 rows]
company_type ( id integer PRIMARY KEY, kind text -- 'production companies', 'distributors', ... )
-- Lookup table for what a title _is_ [7 rows]
kind_type ( id integer PRIMARY KEY, kind text -- 'movie', 'tv series', 'episode', ... )
Let’s say I’m trying to answer the question: “Which Japanese companies put out the most titles in the 2000s?” We might write the following query: 假设我想回答这个问题:“哪些日本公司在 2000 年代发行的影片最多?”我们可能会写出以下查询:
SELECT cn.name, COUNT(*) AS titles
FROM title AS t, movie_companies AS mc, company_name AS cn
WHERE t.id = mc.movie_id AND mc.company_id = cn.id
AND cn.country_code = '[jp]' AND t.production_year BETWEEN 2000 AND 2009
GROUP BY cn.name ORDER BY titles DESC LIMIT 10;
Running this query outputs 10 Japanese companies with the number of titles they were associated with between 2000 and 2009, sorted from highest to lowest. But how did Postgres get these results? The path Postgres took to get this data for us is not a foregone conclusion, and it has everything to do with what we call selective predicates (i.e. the filtering conditions in a WHERE clause). 运行此查询将输出 10 家日本公司及其在 2000 年至 2009 年间关联的影片数量,按从高到低排序。但 Postgres 是如何得到这些结果的呢?Postgres 获取这些数据的路径并非预先注定的,这与我们所说的选择性谓词(即 WHERE 子句中的过滤条件)密切相关。
To illustrate this, let’s imagine our same query without the Japanese company filter or the date range filter: 为了说明这一点,让我们想象一下去掉日本公司过滤器或日期范围过滤器后的同一个查询:
SELECT cn.name, COUNT(*) AS titles
FROM title AS t, movie_companies AS mc, company_name AS cn
WHERE t.id = mc.movie_id AND mc.company_id = cn.id
GROUP BY cn.name ORDER BY titles DESC LIMIT 10;
mc can only join with cn via mc.company_id = cn.id, and t can only join with mc via t.id = mc.movie_id. These constraints produce two valid join trees: mc 只能通过 mc.company_id = cn.id 与 cn 连接,而 t 只能通过 t.id = mc.movie_id 与 mc 连接。这些约束产生了两个有效的连接树:
(Note: There are technically eight join trees if we take commutativity into account. In this case, we don’t because it doesn’t affect the size of the relations resulting from the joins.) (注:如果考虑交换律,技术上共有八种连接树。但在本例中,我们不考虑交换律,因为它不会影响连接结果的关系大小。)
The cardinality of a table or query result is the number of rows it contains. Assume the relevant tables have the following cardinalities: 表或查询结果的基数(cardinality)是指它包含的行数。假设相关表具有以下基数:
- cn = 100k
- mc = 2m
- t = 1m
Taking into account our joins, we get the following cardinalities: 考虑到我们的连接,我们得到以下基数:
- (cn ⋈ mc) = 2m, then ⋈ t = 2m
- (t ⋈ mc) = 2m, then ⋈ cn = 2m
Regardless of the order in which these three tables are joined, the same 2m rows are always passed into the second join. Now let’s add back our selective predicates: 无论这三个表以什么顺序连接,相同的 200 万行数据总是会被传递到第二次连接中。现在让我们加回选择性谓词:
-
cn’ = 5k (assuming 5% of our 100k companies are Japanese)
-
mc = 2m (does not change)
-
t’ = 200k (assuming 20% of our 1m titles were made in the 2000s)
-
(cn’ ⋈ mc) ≈ 100k, then ⋈ t’ ≈ 20k
-
(t’ ⋈ mc) ≈ 400k, then ⋈ cn’ ≈ 20k
The first join ordering filters the 2m movie_companies entries down to the 5% slice of companies that are Japanese. Assuming uniform distribution (we’ll discuss later why we assume this), this join results in approximately 100k rows. Joining the result with the filtered title table keeps only the 20% of those rows from the 2000s. 第一种连接顺序将 200 万条 movie_companies 条目过滤为 5% 的日本公司切片。假设分布均匀(我们稍后会讨论为什么要这样假设),这次连接的结果大约是 10 万行。将该结果与过滤后的 title 表连接,仅保留其中 20% 来自 2000 年代的行。