Training a 4B model to produce 81% faster query plans than Postgres

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. 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: 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 Constructing a Postgres measurement rig that minimizes Linux page cache contention noise across concurrent containers Designing a custom GRPO variant for scoring RL rollouts in an inherently noisy environment Splitting RL across two machines: vLLM and the trainer on a rented 2x H100 node and four Postgres containers running on my desk Running off-policy distillation across half a thousand GPT-6 Astra agent trajectories Let’s start from the beginning. Inside a query optimizer Consider the following slice of the IMDb dataset: -- 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: 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). 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 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: ⋈ ⋈ t cn mc (cn ⋈ mc) ⋈ t ⋈ ⋈ cn t mc (t ⋈ mc) ⋈ cn The two join trees for our query. The lower join runs first; the result is an input into the root join. The cardinality of a table or query result is the number of rows it contains. Assume the relevant tables have the following cardinalities: cn=100kcn = 100\text{k} mc=2mmc = 2\text{m} t=1mt = 1\text{m} Taking into account our joins, we get the following cardinalities: (cn⋈mc)=2m, then ⋈t=2m(cn \bowtie mc) = 2\text{m}, \text{ then } \bowtie t = 2\text{m} (t⋈mc)=2m, then ⋈cn=2m(t \bowtie mc) = 2\text{m}, \text{ then } \bowtie cn = 2\text{m} 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: cn′=5kcn' = 5\text{k} (assuming 5% of our 100k companies are Japanese) mc=2mmc = 2\text{m} (does not change) t′=200kt' = 200\text{k} (assuming 20% of our 1m titles were made in the 2000s) (cn′⋈mc)≈100k, then ⋈ t′≈20k(cn' \bowtie mc) \approx 100\text{k}, \text{ then } \bowtie\ t' \approx 20\text{k} (t′⋈mc)≈400k, then ⋈ cn′≈20k(t' \bowtie mc) \approx 400\text{k}, \text{ then } \bowtie\ cn' \approx 20\text{k} 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. A combinatorial explosion Each join can use any of: Hash join Merge join Nested-loop join 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 t⋈(cn⋈mc)t \bowtie (cn \bowtie mc)(mc⋈cn)⋈t(mc \bowtie cn) \bowtie t t⋈(mc⋈cn)t \bowtie (mc \bowtie cn)(t⋈mc)⋈cn(t \bowtie mc) \bowtie cn cn⋈(t⋈mc)cn \bowtie (t \bowtie mc)(mc⋈t)⋈cn(mc \bowtie t) \bowtie cn cn⋈(mc⋈t)cn \bowtie (mc \bowtie t) Lastly, each table can be scanned in different ways. Considering just four types of scans: Sequential Index Index-only Bitmap 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: 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; Click to select the number of tables being joined together. From four tables onwards, queries on the left are from JOB. On the right is a rough estimate of the size of the search space. Estimating, not counting 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 100\text{k}, \text{ then } \bowtie\ t' \approx 20\text{k}, 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. Postgres assumption - 100k rows Actual - 100k rows ⋈ ⋈ t′ cn′ mc (cn′ ⋈ mc) ⋈ t′ ⋈ ⋈ cn′ t′ mc (t′ ⋈ mc) ⋈ cn′ Share of movie_companies rows that belong to Japanese companies: (uniform) Drag the slider to make Japanese companies more productive. Notice how Postgres’s estimate is static while actual row counts get affected. One bad estimate in