训练一个4B模型，生成比Postgres快81%的查询计划

# 训练一个4B模型生成比Postgres快81%的查询计划 来源: https://rohanbansal.com/qorl ## 查询优化器究竟有多好？ Leis等人（https://vldb.org/pvldb/vol9/p204-leis.pdf）在2015年提出了这个确切的问题。十年后，他们再次追问（https://www.vldb.org/pvldb/vol18/p5531-viktor.pdf）。 尽管在他们最初的探索之后十年间涌现了大量研究，他们发现查询优化器仍然不尽如人意。 初次了解到这点时我感到惊讶。Postgres数据库*应该*非常了解其表中的数据，对吧？能有多难呢？ 事实证明：难如登天。事实上，查询优化器需要完成的某项特定任务——连接顺序选择——是已知的NP难问题（https://dl.acm.org/doi/10.1145/1270.1498）。 因此查询优化器很难。相比之下，*验证*优化器选择的查询计划是否良好则没那么困难。简单来说，好的查询优化器产生运行快的计划，差的则产生慢计划。语言模型特别擅长学习如何完成具有易于验证输出的任务。因为只有一个优化维度——查询的执行时间——这个问题可以完美地简化为强化引导模型产生更快查询计划的行为。 接下来，我将详细说明一个探索性实验：能否通过监督微调（SFT）和智能体强化学习（RL）对一个小型开源权重模型进行后训练，使其产生的Postgres查询计划优于Postgres默认计划？ 我们问题的答案是响亮的“可以”。主要成果包括： - 一个4B模型在113个连接密集型查询上实现了**44.7%的延迟降低**，而该模型最初无法为其中99个查询生成计划 - 构建了Postgres测量装置，最大限度地减少并发容器间Linux页缓存争用噪声 - 为内在噪声环境中的RL展开评分设计了定制的GRPO变体 - 将RL分布在两台机器上：租用的2×H100节点运行vLLM和训练器，我的书桌上运行四个Postgres容器 - 在500条GPT-6 Astra智能体轨迹上运行了离线策略蒸馏 让我们从头说起。 ## 查询优化器内部 考虑以下IMDb数据集片段（https://en.wikipedia.org/wiki/IMDb）： ``` -- IMDb作品（电影、剧集等）[约100万行] title ( id integer PRIMARY KEY, title text, production_year integer, kind_id integer -- 外键 -> kind_type ) -- 电影与公司关联表[约200万行] movie_companies ( id integer PRIMARY KEY, movie_id integer, -- 外键 -> title.id company_id integer, -- 外键 -> company_name.id company_type_id integer, -- 外键 -> company_type.id note text ) -- 公司名称、来源等信息[约10万行] company_name ( id integer PRIMARY KEY, name text, country_code text -- '[us]', '[jp]' 等 ) -- 作品公司角色查找表[4行] company_type ( id integer PRIMARY KEY, kind text -- '制作公司', '发行商' 等 ) -- 作品类型查找表[7行] kind_type ( id integer PRIMARY KEY, kind text -- '电影', '电视剧', '剧集' 等 ) ``` 假设我想回答这个问题：“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; ``` 运行此查询会输出10家日本公司及其在2000年至2009年间关联的作品数量，按数量降序排列。 但Postgres*如何*得到这些结果？ Postgres为我们获取这些数据的路径并非既定结论，这完全取决于我们所说的选择性谓词（即`WHERE`子句中的过滤条件）。 为了说明这一点，让我们想象一个去掉日本公司过滤和日期范围过滤的相同查询： ``` 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`只能通过`mc.company_id = cn.id`与`cn`连接，`t`只能通过`t.id = mc.movie_id`与`mc`连接。 这些约束产生了两个https://rohanbansal.com/qorl#annotation-a1-commutativityhttps://rohanbansal.com/qorl#annotation-reference-a1-commutativity如果考虑交换律，技术上存在八种连接树。但在此情况下我们不考虑，因为它不影响连接结果关系的大小。有效的连接树： ⋈⋈tcnmc(cn ⋈ mc) ⋈ t ⋈⋈cntmc(t ⋈ mc) ⋈ cn 我们查询的两种连接树。下方的连接先执行；其结果作为根连接的输入。表或查询结果的*基数*是其包含的行数。假设相关表具有以下基数： 1. cn=100k 2. mc=2m 3. t=1m 考虑我们的连接，得到以下基数： (cn⋈mc)=2m，然后⋈t=2m (t⋈mc)=2m，然后⋈cn=2m 无论这三个表以何种顺序连接，总是有相同的200万行传递到第二次连接。 现在让我们加回选择性谓词： 1. cn'=5k（假设10万家公司中有5%是日本公司） 2. mc=2m（不变） 3. t'=200k（假设100万部作品中有20%是2000年代制作的） (cn'⋈mc)≈100k，然后⋈t'≈20k (t'⋈mc)≈400k，然后⋈cn'≈20k 第一种连接顺序将200万`movie_companies`条目过滤为5%的日本公司部分。假设均匀分布（我们稍后会讨论*为什么*做此假设），此连接产生约10万行。将结果与过滤后的`title`表连接，仅保留2000年代作品的20%行。 第二种连接顺序将200万`movie_companies`条目过滤为20%的2000年代作品部分。同样的均匀性假设成立，因此第一次连接产生40万行，意味着我们要将40万行传递到第二次连接。 如果我们选择第二种连接顺序，工作量增加了**4倍**。 不幸的是，这还不是全部。 ### 组合爆炸 每次连接可以使用以下任一方式： 1. 哈希连接 2. 归并连接 3. 嵌套循环连接 现在重新考虑交换律https://rohanbansal.com/qorl#annotation-a2-commutativityhttps://rohanbansal.com/qorl#annotation-reference-a2-commutativity虽然交换律不影响产生的行数，但现在必须考虑，因为它*确实*影响所用连接算法的性能。，存在4种不同的外层/内层连接*方向*，导致8种可能的组合： (cn⋈mc)⋈t t⋈(cn⋈mc) (mc⋈cn)⋈t t⋈(mc⋈cn) (t⋈mc)⋈cn cn⋈(t⋈mc) (mc⋈t)⋈cn cn⋈(mc⋈t) 最后，每个表可以用不同方式扫描。仅考虑四种扫描类型： 1. 顺序扫描 2. 索引扫描 3. 仅索引扫描 4. 位图扫描 这个查询有4,608种不同的执行方式https://rohanbansal.com/qorl#annotation-a-pruninghttps://rohanbansal.com/qorl#annotation-reference-a-pruning这实际上是低估了。计划可以并行运行，聚合可以哈希或排序等等。 值得注意的是，Postgres不会评估所有这些计划。它使用动态规划（以及对涉及12+连接的查询使用遗传算法(https://www.postgresql.org/docs/current/geqo-pg-intro.html)）来修剪搜索空间。 更糟的是，每个连接都会组合爆炸式扩大搜索空间： ``` SELECT MIN(t.title) AS movie_title FROM company_name AS cn, keyword AS k, movie_companies AS mc, movie_keyword AS mk, title AS t WHERE cn.country_code ='[de]' AND k.keyword ='character-name-in-title' AND cn.id = mc.company_id AND mc.movie_id = t.id AND t.id = mk.movie_id AND mk.keyword_id = k.id AND mc.movie_id = mk.movie_id; ``` 点击查看连接的表数量。从四表开始，左侧查询来自JOB。右侧是搜索空间大小的粗略估计。 ### 估计而非计数 Postgres在这里处境艰难。可能会认为它只需计算基数，选择最小化传递到后续连接的行数的计划。 但这将意味着Postgres*能*在查询规划期间计算基数。它不能。要知道这点，它需要实际运行每个连接并计算结果行数。这违背了快速查询优化器的全部意义。查询优化器的目标不是在成本最小化方面精确...而是对多种查询类型都足够好。 相反，Postgres使用统计数据来估计基数。规划器查询`pg_statistic`表，获取每个列的常见值及其频率，以及其余部分的直方图。当添加连接时，事情变得稍微复杂些。Postgres不知道一个表中的行在另一个表上如何分布。为了解决这个问题，它假设第一个表中给定值的频率可以直接应用于第二个表。这就是我之前提到的均匀分布假设。 假设均匀分布作为启发式方法是可以的，但当它失效时，就会彻底失败。回顾之前的连接顺序(cn'⋈mc)≈100k，然后⋈t'≈20k，我们基于5%公司是日本公司的假设过滤了200万`movie_companies`条目。但如果这5%的日本公司实际上贡献了**50%**的电影呢？第一次连接将产生100万行！成本模型说选择第一种连接顺序；实际上，第二种更好，因为它只传递40万行到第二次连接。 Postgres假设：10万行 实际：100万行 ⋈⋈t'cn'mc(cn' ⋈ mc) ⋈ t' ⋈⋈cn't'mc(t' ⋈ mc) ⋈ cn' `movie_companies`行中属于日本公司的比例：（均匀分布） 拖动滑块使日本公司更具生产力。注意Postgres的估计如何保持静态而实际行数受到影响。早期连接中的一个错误估计可能会级联影响整个连接树，破坏所有其他估计。 ## 如何引导大象 Postgres总是选择成本最低的计划，而我们无法在不修改源代码的情况下改变其成本模型，那么我们如何实际引导它选择具有更高成本的不同计划？ 这就需要pg_hint_plan（https://github.com/ossc-db/pg_hint_plan）。 `pg_hint_plan`是一个极其简单的第三方扩展：只需在SQL语句上方添加结构化“提示”作为注释，就可以引导Postgres使用提示中提供的指令来制定计划。例如： ``` /*+ HashJoin(a b) SeqScan(a) */ EXPLAIN SELECT * FROM pgbench_branches b JOIN pgbench_accounts a ON b.bid = a.bid ORDER BY a.aid; ``` ``` QUERY PLAN --------------------------------------------------------------------------------- Sort (cost=31465.84..31715.84 rows=100000 width=197) Sort Key: a.aid -> Hash Join (cost=1.02..4016.02 rows=100000 width=197) Hash Cond: (a.bid = b.bid) -> Seq Scan on pgbench_accounts a (cost=0.00..2640.00 rows=100000 width=97) -> Hash (cost=1.01..1.01 rows=1 width=100) -> Seq Scan on pgbench_branches b (cost=0.00..1.01 rows=1 width=100) (7 rows) ``` 来自`pg_hint_plan`文档的示例（https://github.com/ossc-db/pg_hint_plan/blob/master/docs/description.md）。 该提示强制使用`HashJoin`连接`pgbench_accounts`和`pgbench_branches`，并对`pgbench_accounts`表进行顺序扫描；实际查询计划很好地遵循了这一点。 ## 我们问题的表述 鉴于我们可以使用`pg_hint_plan`提示影响Postgres选择不同——可能更好——的查询计划，我们开始的问题是： > 语言模型能否学会产生能带来更好查询计划的提示？ ### 有研究价值 什么可能使这成为一个值得解决的问题？ 我的第一个想法是将查询和Postgres规划器拥有的完全相同的信息提供给模型。这相当于查看我们是否能构建更好的基数估计器。我得出的结论是，这不是一个值得探索的途径；我们将与几十年的基数估计研究抗争。此外，仅推理延迟就远远超过相比Postgres超快查询优化器的任何学习收益。 第二个想法——也是我认为正确的表述——在于一种特定的数据库使用模式：繁重的分析工作负载。如果查询使用次优的默认Postgres计划被运行数千次，效率损失是显而易见的。相反，可以训练一个模型来为特定查询找到更好的执行方式。训练过程可能需要预先执行该查询数十到数百次，但所有查询运行的分摊成本将大大降低。 目标不是尝试在一次性查询的时间/效率帕累托前沿上击败Postgres，但也许我们可以在反复运行的查询上击败它。 ## 模型及其工具 我决定从一个4B小模型开始，因为它最容易在我家的2×RTX 3090设备（亲切地命名为FLOPper）上自行训练/推理。 在我开始这个项目时，Qwen 3.8系列模型发布了，不幸的是没有4B变体。然而，我遇到了德国一个小实验室Empero（https://empero.org/）的Qwen 3.8 4B蒸馏版，并感到好奇。他们使用Qwen 3.8的2.4T模型作为教师模型，将知识蒸馏到Qwen 3.5 4B中，产生了`empero-ai/Qwen3.8-4B-Distill`（https://huggingface.co/empero-ai/Qwen3.8-4B-Distill）。这个蒸馏模型并不完全优于其基础3.5模型；它在MMLU（https://en.wikipedia.org/wiki/MMLU）任务上表现更好，在GSM8K（https://huggingface.co/datasets/openai/gsm8k）任务上略差。换句话说，这个蒸馏在评估广泛知识时表现更好，在多步数学推理上略差。至于哪个对我们任务更好，我不知道；无论如何我决定坚持使用蒸馏模型。 模型确定后，我构建了一个轻量级智能体工具`qo-agent`，用于协调提示生成。它被赋予以下六个工具： 1. `inspect_relation`—列出表的列类型和可空性、索引定义以及估计的行数和字节数 2. `get_column_stats`—获取1-8列的Postgres规划器统计信息 3. `get_plan`—获取默认计划的估计或提交的候选者的存储计划 4. `evaluate_candidate`—验证提议的计划动作，然后