← Back to feed

A researcher has published the results of post-training a 4-billion-parameter open-weights model to produce Postgres query plans that beat the database’s native optimizer on join-heavy workloads. The model was trained using supervised fine-tuning followed by agentic reinforcement learning, with query execution latency as the reward signal.

The headline result: a 1.81x geometric mean speedup per query across 113 join-heavy test queries sourced from the IMDb dataset — 81% faster on average. Summed across all 113 queries, total workload latency dropped 44.7%. The starting point was stark: the base model could not produce a valid plan for 99 of the 113 queries before training.

Why query optimization is hard (and measurable)

Database query optimization, specifically join ordering, is NP-hard. Postgres’s query planner makes heuristic choices about which tables to join in what order, using statistics about row counts and selectivity. These estimates are imperfect, and on complex joins with many tables they compound into plans that are significantly slower than optimal.

What makes this a tractable RL target is that the reward signal is unambiguous: execution time. Unlike tasks where verifying a correct output is itself difficult, a faster query plan produces a measurably shorter wall-clock time. This is the same property that makes coding benchmarks amenable to RL — a test suite either passes or it does not.

Training setup

The work combined supervised fine-tuning on labeled plan examples with a custom GRPO (Group Relative Policy Optimization) variant adapted for noisy environments. Query execution time is inherently noisy: Linux page cache state, concurrent load, and I/O patterns all affect individual measurements. The researcher built a measurement rig across four Postgres containers to minimize cache contention across concurrent runs.

Rollout generation was split across hardware: a rented 2xH100 node ran vLLM and the training loop, while local machines hosted the Postgres evaluation containers. Off-policy distillation was applied across roughly 500 GPT-6 Astra agent trajectories.

Scope and limitations

The benchmark is join-heavy queries from the IMDb dataset — a standard query optimization research corpus. The 113 queries were specifically selected to stress join ordering decisions. The result does not generalize automatically to all query types; queries without complex joins are less likely to benefit. The Postgres planner performs well on simpler workloads where cardinality estimates are accurate.

The model also operates at inference time inside the query planning loop, adding latency to plan selection itself. For workloads where queries execute once and plan overhead is amortized across millions of executions, this matters less. For short-lived ad-hoc queries, the planning overhead could offset gains.

The broader signal

The work is a clean demonstration that small specialized models can match or exceed handcrafted heuristics on problems with a verifiable reward signal. Database query optimization has decades of expert engineering behind it. A 4B model with a few weeks of RL closed the gap and opened it in favor of the learned approach — at least on the tested workload. The full write-up and code are at rohanbansal.com/qorl.