SPA: A SQL-Plan-Aware Reinforcement Learning Framework for Query Rewriting with LLMs
Abstract
SQL query rewriting improves database performance without schema changes, but LLM rewriters often produce trivial rewrites that preserve the physical plan or plan-changing rewrites that slow execution. Although database execution reveals these failures, runtime alone provides sparse and noisy supervision: invalid or inequivalent rewrites cannot receive meaningful runtime credit, while timing differences under an unchanged plan largely reflect measurement noise. We present SPA, a SQL-Plan-Aware reinforcement learning framework that makes execution feedback more reliable and discriminative. SPA combines a physical-plan-divergence reward that filters plan-preserving rewrites, PGARS, a per-query curriculum for rewrite difficulty that cannot be predicted in advance, and on-policy self-improvement that trains the policy to recover from slowdown rewrites produced by its own rollouts. Across in-distribution and out-of-distribution workloads, SPA outperforms rule-based baselines and achieves runtime competitive with strong proprietary LLMs while maintaining high semantic equivalence and substantially lower inference cost under batched local serving.