SPA: A SQL-Plan-Aware Reinforcement Learning Framework for Query Rewriting with LLMs
For database administrators and query optimizers, SPA addresses the problem of generating effective query rewrites that improve runtime performance, moving beyond rule-based and naive LLM approaches.
SPA is a reinforcement learning framework that trains LLMs to rewrite SQL queries using physical execution feedback, outperforming rule-based and LLM baselines in end-to-end runtime and reducing harmful slowdown rewrites on both IID and OOD workloads.
SQL query rewriting is a well-established technique for improving database performance without schema or index changes, yet finding effective rewrites for modern analytical workloads remains difficult: rule-based methods are limited to predefined transformations, while LLM-based approaches often produce rewrites that are semantically valid but compile to equivalent physical plans or degrade runtime performance. We present SPA, a SQL-Plan-Aware reinforcement learning framework that trains LLMs to rewrite queries using physical execution feedback. SPA formulates rewriting as a policy optimization problem and extends GRPO with rewards spanning semantic equivalence, textual rewrite distance, physical-plan divergence, and runtime speedup. To handle reward sparsity across query difficulty, SPA introduces Probability-Gated Adaptive Reward Shaping, a query-level curriculum that unlocks higher-level rewards only once a rollout group achieves sufficient mastery of lower-level objectives, and further improves sample efficiency through on-policy self-improvement by recycling slowdown rewrites from the current policy as targeted training signals. On both IID and OOD workloads, SPA outperforms rule-based and strong LLM baselines in end-to-end runtime, substantially reduces harmful slowdown rewrites, and yields strong tail-latency gains.