Training a 4B model to produce 81% faster query plans than Postgres
Four RL rollouts commence for a single query. Qwen produces a candidate strategy per rollout and sends it off to Postgres for measurement against its own default plan. Scalar rewards are assigned to each rollout, which flow backwards to update Qwen’s weights.
Weights nudged toward the faster plans.
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.
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.
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?
The answer to our question is a resounding yes. Highlights include:
Let’s start from the beginning.
Consider the following slice of the IMDb dataset:
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:
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).
To illustrate this, let’s imagine our same query without the Japanese company filter or the date range filter:
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 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. valid join trees:
The cardinality of a table or query result is the number of rows it contains. Assume the relevant tables have the following cardinalities:
Taking into account our joins, we get the following cardinalities:
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:
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.
The second join ordering filters the 2m movie_companies entries down to the 20% slice of titles that were made in the 2000s. The same uniformity assumption holds, so the first join results in 400k rows, meaning we’re passing 400k rows into the second join.
We do 4x the work if we picked the second join ordering.
Unfortunately, it doesn’t stop there.
Each join can use any of:
Factoring commutativity back in now While commutativity doesn’t change the number of rows produced, it must be considered now because it does affect performance regarding the join algorithm used. , there are 4 different outer/inner join orientations, resulting in 8 possible combinations:
(cn⋈mc)⋈t(cn bowtie mc) bowtie t(cn⋈mc)⋈t t⋈(cn⋈mc)t bowtie (cn bowtie mc)t⋈(cn⋈mc)
(mc⋈cn)⋈t(mc bowtie cn) bowtie t(mc⋈cn)⋈t t⋈(mc⋈cn)t bowtie (mc bowtie cn)t⋈(mc⋈cn)
(t⋈mc)⋈cn(t bowtie mc) bowtie cn(t⋈mc)⋈cn cn⋈(t⋈mc)cn bowtie (t bowtie mc)cn⋈(t⋈mc)
(mc⋈t)⋈cn(mc bowtie t) bowtie cn(mc⋈t)⋈cn cn⋈(mc⋈t)cn bowtie (mc bowtie t)cn⋈(mc⋈t)
Lastly, each table can be scanned in different ways. Considering just four types of scans:
2 Join trees: which pair of tables joins first. × 22 Orientations: each of the 2 joins can swap which input is outer and which is inner. × 32 Algorithms: each of the 2 joins picks hash, merge, or nested loop. × 43 Scans: each of the 3 tables is either read sequentially or via index, index-only or bitmap scans. = 4,608
There are 4,608 different ways to run this query This is actually an undercount. Plans can run in parallel, aggregates can be hashed or sorted, etc.It’s also worth noting that Postgres doesn’t evaluate all of these plans. It uses dynamic programming (and a genetic algorithm for queries involving 12+ joins) to prune the search space. !
To make matters worse, every join combinatorially explodes the search space:
1 Join trees: with two tables there is only one way to join them. × 21 Orientation: 1 join means there are only 2 orientations. × 31 Algorithm: the join algorithm can be a hash join, merge join or nested loop. × 42 Scans: each of the 2 tables is either read sequentially or via index, index-only or bitmap scans. = 96
2 Join trees: the ways 3 tables can be joined up, before any swapping of inputs. × 22 Orientations: each of the 2 joins can swap which input is outer and which is inner. × 32 Algorithms: each of the 2 joins picks hash, merge, or nested loop. × 43 Scans: each of the 3 tables is either read sequentially or via index, index-only or bitmap scans. = 4,608
8 Join trees: the ways 4 tables can be joined up, before any swapping of inputs. × 23 Orientations: each of the 3 joins can swap which input is outer and which is inner. × 33 Algorithms: each of the 3 joins picks hash, merge, or nested loop. × 44 Scans: each of the 4 tables is either read sequentially or via index, index-only or bitmap scans. = 442,368
25 Join trees: the ways 5 tables can be joined up, before any swapping of inputs. × 24 Orientations: each of the 4 joins can swap which input is outer and which is inner. × 34 Algorithms: each of the 4 joins picks hash, merge, or nested loop. × 45 Scans: each of the 5 tables is either read sequentially or via index, index-only or bitmap scans. = 33,177,600
56 Join trees: the ways 6 tables can be joined up, before any swapping of inputs. × 25 Orientations: each of the 5 joins can swap which input is outer and which is inner. × 35 Algorithms: each of the 5 joins picks hash, merge, or nested loop. × 46 Scans: each of the 6 tables is either read sequentially or via index, index-only or bitmap scans. = 1,783,627,776
696 Join trees: the ways 7 tables can be joined up, before any swapping of inputs. × 26 Orientations: each of the 6 joins can swap which input is outer and which is inner. × 36 Algorithms: each of the 6 joins picks hash, merge, or nested loop. × 47 Scans: each of the 7 tables is either read sequentially or via index, index-only or bitmap scans. = 532,030,685,184
4,698 Join trees: the ways 8 tables can be joined up, before any swapping of inputs. × 27 Orientations: each of the 7 joins can swap which input is outer and which is inner. × 37 Algorithms: each of the 7 joins picks hash, merge, or nested loop. × 48 Scans: each of the 8 tables is either read sequentially or via index, index-only or bitmap scans. = 86,188,970,999,808
20,340 Join trees: the ways 9 tables can be joined up, before any swapping of inputs. × 28 Orientations: each of the 8 joins can swap which input is outer and which is inner. × 38 Algorithms: each of the 8 joins picks hash, merge, or nested loop. × 49 Scans: each of the 9 tables is either read sequentially or via index, index-only or bitmap scans. = 8,955,727,561,359,360
242,160 Join trees: the ways 10 tables can be joined up, before any swapping of inputs. × 29 Orientations: each of the 9 joins can swap which input is outer and which is inner. × 39 Algorithms: each of the 9 joins picks hash, merge, or nested loop. × 410 Scans: each of the 10 tables is either read sequentially or via index, index-only or bitmap scans. = 2,558,960,455,762,575,360
1,490,850 Join trees: the ways 11 tables can be joined up, before any swapping of inputs. × 210 Orientations: each of the 10 joins can swap which input is outer and which is inner. × 310 Algorithms: each of the 10 joins picks hash, merge, or nested loop. × 411 Scans: each of the 11 tables is either read sequentially or via index, index-only or bitmap scans. = 378,099,722,048,923,238,400
11,932,560 Join trees: the ways 12 tables can be joined up, before any swapping of inputs. × 211 Orientations: each of the 11 joins can swap which input is outer and which is inner. × 311 Algorithms: each of the 11 joins picks hash, merge, or nested loop. × 412 Scans: each of the 12 tables is either read sequentially or via index, index-only or bitmap scans. = 72,630,206,166,931,876,085,760 lots!
Postgres is in a tough spot here. It would be reasonable to think it could simply count cardinalities and pick the plan that minimizes the number of rows passed through to successive joins.
But this would imply Postgres can count cardinalities during query planning. It can’t. In order to know this, it would need to actually run each join and count the resulting rows. This defeats the whole point of a fast query optimizer. A query optimizer does not aim to be exact in its cost minimization… it aims to be good enough across many types of queries.
Instead, Postgres uses statistics to estimate cardinalities. The planner queries the pg_statistic table, getting back common values for each column and their frequencies, and a histogram for the rest. Things get a bit more complicated when you tack on joins. Postgres doesn’t know how the rows in one table are distributed over the other. To get around this, it assumes that the frequency of a given value in the first table can simply be applied over the second table. This is the uniform distribution assumption I mentioned earlier.
Assuming a uniform distribution is fine as a heuristic, but when it fails, it fails hard. Looking back at an earlier join ordering (cn′⋈mc)≈100k, then ⋈ t′≈20k(cn’ bowtie mc) approx 100text{k}, text{ then } bowtie t’ approx 20text{k}(cn′⋈mc)≈100k, then ⋈ t′≈20k, we filtered 2m movie_companies entries on the assumption that 5% of them were from Japanese companies. But what if the 5% of companies that are Japanese were actually responsible for 50% of the movies? The first join would produce 1m rows! The cost model says pick the first join ordering; in reality, the second one is actually better since it only sends 400k rows through to the second join.
One bad estimate in an early join can cascade through the rest of the join tree, corrupting all other estimates.
Postgres always picks the plan with the lowest cost, and we can’t change its cost model without modifying its source code, so how can we actually steer it to pick different plans that have higher costs?
pg_hint_plan is a beautifully simple third-party extension: just by adding structured “hints” as comments above SQL statements, you can nudge Postgres towards plans that use the instructions provided in the hint. For example:
Example from pg_hint_plan’s documentation.
The hint mandates usage of a HashJoin for joining pgbench_accounts and pgbench_branches, and doing a sequential scan of the pgbench_accounts table; the actual query plan follows suit nicely.
Given that we can influence Postgres to pick different—and potentially better—query plans using pg_hint_plan hints, the question we’re starting with is:
Can a language model learn to produce hints that result in better query plans?
What might make this a worthwhile problem to solve?
My first idea was to give the model the query and the exact same set of information Postgres’s planner has. This amounts to seeing if we could build a better cardinality estimator. I came to the conclusion this is not a worthwhile avenue to explore; we would be fighting decades of cardinality estimation research. Furthermore, the inference latency alone would far outweigh any learned usefulness compared to Postgres’s ultra-fast query optimizer.
The second idea—and what I believe is the correct formulation—lies in a specific database usage pattern: heavy analytic w