# Bao: Learning to Steer Query Optimizers

arXiv: https://arxiv.org/abs/2004.03814 · HTML: https://mlsanity.com/paper/2004.03814/

## Key facts

- Cited by 62 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 8 Apr 2020.

## Built on by (most cited first, 30 of 62)

- [Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Training](https://mlsanity.com/paper/2208.06102/) (2022)
- [Unifying the Perspectives of NLP and Software Engineering: A Survey on Language Models for Code](https://mlsanity.com/paper/2311.07989/) (2023)
- [Lero: A Learning-to-Rank Query Optimizer](https://mlsanity.com/paper/2302.06873/) (2023)
- [Balsa: Learning a Query Optimizer Without Expert Demonstrations](https://mlsanity.com/paper/2201.01441/) (2022)
- [Zero-Shot Cost Models for Out-of-the-box Learned Cost Prediction](https://mlsanity.com/paper/2201.00561/) (2022)
- [Glia: A Human-Inspired AI for Automated Systems Design and Optimization](https://mlsanity.com/paper/2510.27176/) (2025)
- [Cortex AISQL: A Production SQL Engine for Unstructured Data](https://mlsanity.com/paper/2511.07663/) (2025)
- [A Survey on Deep Reinforcement Learning for Data Processing and Analytics](https://mlsanity.com/paper/2108.04526/) (2021)
- [E2ETune: End-to-End Knob Tuning via Fine-tuned Generative Language Model](https://mlsanity.com/paper/2404.11581/) (2024)
- [QUITE: A Query Rewrite System Beyond Rules with LLM Agents](https://mlsanity.com/paper/2506.07675/) (2025)
- [Detect, Distill and Update: Learned DB Systems Facing Out of Distribution Data](https://mlsanity.com/paper/2210.05508/) (2022)
- [Graph Neural Networks for Databases: A Survey](https://mlsanity.com/paper/2502.12908/) (2025)
- [NeurDB: An AI-powered Autonomous Data System](https://mlsanity.com/paper/2405.03924/) (2024)
- [No DBA? No regret! Multi-armed bandits for index tuning of analytical and HTAP workloads with provable guarantees](https://mlsanity.com/paper/2108.10130/) (2021)
- [Real-time Workload Pattern Analysis for Large-scale Cloud Databases](https://mlsanity.com/paper/2307.02626/) (2023)
- [Sibyl: Forecasting Time-Evolving Query Workloads](https://mlsanity.com/paper/2401.03723/) (2024)
- [The Unreasonable Effectiveness of LLMs for Query Optimization](https://mlsanity.com/paper/2411.02862/) (2024)
- [Kepler: Robust Learning for Faster Parametric Query Optimization](https://mlsanity.com/paper/2306.06798/) (2023)
- [AnnotatedTables: A Large Tabular Dataset with Language Model Annotations](https://mlsanity.com/paper/2406.16349/) (2024)
- [HERO: Hint-Based Efficient and Reliable Query Optimizer](https://mlsanity.com/paper/2412.02372/) (2024)
- [In-Context Adaptation to Concept Drift for Learned Database Operations](https://mlsanity.com/paper/2505.04404/) (2025)
- [Conditionally Risk-Averse Contextual Bandits](https://mlsanity.com/paper/2210.13573/) (2022)
- [Deploying a Steered Query Optimizer in Production at Microsoft](https://mlsanity.com/paper/2210.13625/) (2022)
- [Learned-Database Systems Security](https://mlsanity.com/paper/2212.10318/) (2022)
- [ML-Powered Index Tuning: An Overview of Recent Progress and Open Challenges](https://mlsanity.com/paper/2308.13641/) (2023)
- [Supporting Dynamic Agentic Workloads: How Data and Agents Interact](https://mlsanity.com/paper/2512.09548/) (2025)
- [AI-Driven Research for Databases](https://mlsanity.com/paper/2604.06566/) (2026)
- [TENSILE: A Tensor granularity dynamic GPU memory scheduling method toward multiple dynamic workloads system](https://mlsanity.com/paper/2105.13336/) (2021)
- [JoinGym: An Efficient Query Optimization Environment for Reinforcement Learning](https://mlsanity.com/paper/2307.11704/) (2023)
- [FOSS: A Self-Learned Doctor for Query Optimizer](https://mlsanity.com/paper/2312.06357/) (2023)

## Abstract (arXiv)

Query optimization remains one of the most challenging problems in data management systems. Recent efforts to apply machine learning techniques to query optimization challenges have been promising, but have shown few practical gains due to substantive training overhead, inability to adapt to changes, and poor tail performance. Motivated by these difficulties and drawing upon a long history of research in multi-armed bandits, we introduce Bao (the BAndit Optimizer). Bao takes advantage of the wisdom built into existing query optimizers by providing per-query optimization hints. Bao combines modern tree convolutional neural networks with Thompson sampling, a decades-old and well-studied reinforcement learning algorithm. As a result, Bao automatically learns from its mistakes and adapts to changes in query workloads, data, and schema. Experimentally, we demonstrate that Bao can quickly (an order of magnitude faster than previous approaches) learn strategies that improve end-to-end query execution performance, including tail latency. In cloud environments, we show that Bao can offer both reduced costs and better performance compared with a sophisticated commercial system.