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

Query optimizers look tidy from the outside. Feed in a SQL query, get back a plan, move on. But under the hood, the search space blows up fast, and even a database as battle-tested as Postgres is still making educated guesses. That’s the core of this experiment: can a small open-weights model learn to steer Postgres toward better plans than the default planner picks on its own?

The setup is a good example of why planning is such a mess. Join ordering is NP-hard, so Postgres leans on statistics and heuristics instead of perfect knowledge. It estimates cardinalities from pg_statistic, assumes values are spread uniformly across joins, and then prunes the possible plans with dynamic programming — or a genetic algorithm once queries get large enough. That works until it doesn’t. A bad early estimate can poison the rest of the tree.

The author makes that concrete with an IMDb-style query about Japanese companies in the 2000s. Swap the order of the joins and you can end up pushing 400k rows into the second join instead of 100k. Same tables. Same result. Different amount of work. And once you add join methods, scan types, commutativity, and other planner choices, even a modest query can explode into thousands of possible execution paths.

Instead of trying to out-guess Postgres’s cost model directly, the project uses pg_hint_plan, a Postgres extension that reads hints from SQL comments and nudges the planner toward specific joins and scans. That gives the model a concrete target: produce hints that lead to faster execution. The author settled on a small 4B model for local training and inference, because the real goal wasn’t one-off cleverness. It was repeated analytic queries, where a model can take some upfront pain and then save time every time the same query comes back.

The headline result is blunt: a 4B model initially unable to produce a query plan for 99 of 113 join-heavy queries was eventually used to cut latency by 44.7% across that set. The project also involved a custom Postgres measurement rig, a GRPO variant for noisy rollouts, training spread across two machines, and off-policy distillation from hundreds of GPT-6 Astra trajectories. But the interesting part is simpler than the machinery. Postgres still guesses. A trained model can learn a better guess.