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.
翻译:SQL查询重写是一种在不更改模式或索引的情况下提升数据库性能的成熟技术,然而为现代分析型负载找到有效的重写方案仍面临挑战:基于规则的方法局限于预定义变换,而基于大语言模型的方法生成的重写结果虽语义正确,却可能编译为等效的物理计划或导致运行时性能退化。本文提出SPA——一种基于SQL-计划感知的强化学习框架,通过物理执行反馈训练大语言模型进行查询重写。SPA将重写形式化为策略优化问题,并扩展GRPO算法,引入涵盖语义等价性、文本重写距离、物理计划差异及运行时加速的奖励机制。为应对不同查询难度下的奖励稀疏性,SPA提出概率门控自适应奖励塑形——一种查询级课程学习策略:只有当生成组在低层级目标上达到足够熟练度时,才解锁高层级奖励。此外,SPA通过基于当前策略的"减速重写"回收作为针对性训练信号,实现策略内自我改进以提升样本效率。在独立同分布与非独立同分布工作负载上,SPA在端到端运行时性能上超越基于规则的基线方法与强基准大语言模型,显著减少有害的减速重写,并在尾部延迟方面实现显著改进。