Source-linked AI summary
Bao: Learning to Steer Query Optimizers
Ryan Marcus, Parimarjan Negi, Hongzi Mao, Nesime Tatbul, Mohammad Alizadeh, Tim Kraska
TL;DR
Query optimization is difficult, while learned approaches have faced training overhead, brittleness under changes, and poor tail performance. Bao steers existing optimizers with query-specific hints using tree convolution and Thompson sampling, and it reports improved performance, adaptation, and cost outcomes.
Problem
Learned query optimization has faced substantial training overhead, limited adaptation to workload, data, and schema changes, and poor tail performance.
Method
Bao uses tree convolutional neural networks and Thompson sampling to select query-specific optimizer hints while leveraging an existing query optimizer.
Results
Bao outperforms open-source and commercial optimizers with approximately one hour of training, while improving median and tail latencies amid dynamic workloads, data, and schema.
Takeaways & Limitations
Bao provides a learned query-optimization system that adapts to changing workloads, data, and schemas while rarely incurring catastrophic executions.
Takeaways & Limitations
Bao’s prototype can require 210ms for optimization, which may not suit applications requiring faster planning time.
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
Bao addresses practical weaknesses in learned query optimization by steering an existing optimizer with query-specific hints. It combines contextual bandits and tree convolution to adapt across workload, data, and schema changes while improving performance and avoiding catastrophic tail regressions.
- Query optimization remains difficult because cardinality estimation and cost modeling have resisted decades of research.
- Prior learned approaches face three practical drawbacks: inefficient training, brittleness under changes, and catastrophic tail regressions.
- Bao outperforms open-source and commercial optimizers with approximately one hour of training while adapting to workload, data, and schema changes.
- Bao steers an existing query optimizer by selecting coarse-grained hint sets for each incoming query rather than replacing the optimizer.
- Different hints can help or harm different queries: disabling loop joins improves query 16b by 3x but causes an almost 50x regression for query 24b.
- Bao combines tree convolutional neural networks with Thompson sampling to predict effective plans and balance exploration with exploitation.
2. SYSTEM MODEL
Bao selects query hints for each incoming query, uses the underlying optimizer to generate candidate plans, and learns from observed execution performance to improve future selections.
- The underlying optimizer produces one query plan for each hint set, and Bao transforms each plan into a vector tree.Each vector-tree node is a feature vector used as input to Bao’s predictive model.
- A tree convolutional neural network predicts the execution outcome of each candidate plan, such as its wall-clock time.
- Bao selects a hint set for each incoming query to steer the underlying optimizer toward a suitable plan.Hints alter optimizer behavior; Bao does not stitch plans from different hint sets.
- Thompson sampling balances exploring new plans with exploiting plans already predicted or observed to be fast.
- After execution, Bao records the selected plan and observed performance, periodically retraining its predictive model in a feedback loop.The process can be executed in parallel to maintain reasonable optimization times.
3. SELECTING QUERY HINTS
Bao selects per-query hint sets by framing query optimization as contextual bandit learning over plans produced by an underlying optimizer. It represents plans as vector trees for a tree convolutional model, while practical deviations address training cost and dependence between queries.
- System model: Bao assumes a finite family of hint sets, known query-plan operator types, and a user-defined metric such as execution time or disk operations.Each hint set is passed with the query to the underlying optimizer, which produces a plan tree.
- Action space: Bao selects one complete plan from a single hint set rather than stitching plans together, keeping the per-query action space at O(k) for k hint sets.The authors experimentally tested plan stitching but were unable to get the model to convergence.
- Bandit formulation: Bao treats each hint set as an arm and the optimizer-produced plans under those hints as the contextual information for selecting a plan.The selected plan is executed, its performance becomes a reward, and the system seeks to minimize regret over time.
- Learning loop: Bao updates a predictive model from observed plan-performance pairs and uses Thompson sampling to balance exploration with exploitation.The model maps plan trees to estimated performance, while posterior sampling avoids always choosing randomly or always exploiting the current best estimate.
- Practical considerations: Hint selection is not exactly an independent bandit process because one query plan can alter cache state for subsequent queries.The authors argue this effect is typically short-lived in OLAP workloads and report experimental evidence supporting Thompson sampling in such settings.
- Prediction model: The prediction model applies stacked tree convolutions, dynamic pooling, and fully connected layers to map vectorized plan trees to performance predictions.Tree convolutions support pattern recognition over structures such as join pipelines and potentially problematic operators.
- Plan representation: Bao converts query-plan trees into vector trees by binarizing operators and encoding each node with operator type, cardinality, and cost information.The representation can optionally include cache information and omits direct schema identifiers, helping it remain applicable across schema changes.
- Practical considerations: Classical Thompson sampling is modified because resampling neural-network parameters after every query is time-consuming and experience can grow without bound.Training time increases with the number of stored training examples, motivating practical deviations from the classical procedure.
4. RELATED WORK
Related work applies machine learning to cardinality estimation, query-plan construction, adaptive processing, and broader database management tasks. Bao builds on this landscape while using Thompson sampling and machine-programming ideas in its approach.
- Learned cardinality estimators use supervised deep learning or linear mixture models, but require extensive training data and may not adapt to data or schema changes.
- Reinforcement-learning approaches have constructed query optimizers that can find lower-cost plans after sufficient training, including methods targeting query latency directly.
- Adaptive query-processing research has reported promising reinforcement-learning results, but those techniques do not apply to non-adaptive systems.
- Thompson sampling is a well-established method for updating beliefs from experience and has been used in reinforcement learning and cloud workload management.
- Reinforcement learning has also been applied to elastic-cluster management, scheduling, physical design, and visions of entire database systems built from learned components.
- Bao belongs to a broader machine-programming trend that seeks easy-to-use, adaptive, and inventive systems, including applications beyond data management.
5. EXPERIMENTS
Bao adapts to dynamic workloads, data, and schema changes while improving cost, latency, and tail performance against PostgreSQL and a commercial system. Its online learning converges quickly, learns diverse query-specific hint strategies, and remains resilient when hint-set behavior changes.
- Changing schema, workload, and data: Bao achieves lower cost and workload latency than PostgreSQL across changing workloads, including 50% and 40% reductions on Stack and Corp, respectively.The Corp workload includes a significant schema change, while Stack represents changing data.
- Tail latency analysis: Bao drastically reduces tail latency versus PostgreSQL and always reduces it versus the commercial system, with the strongest commercial-system gains on smaller VM types.On N1-8, 99% latency falls from 130 seconds with PostgreSQL to under 20 seconds with Bao.
- Training time and convergence: From a cold start, Bao matches PostgreSQL within one hour and exceeds it within two hours, while matching the commercial optimizer within 90 minutes and exceeding it within three hours.The same convergence pattern is reported across VM types, with similar behavior on Stack and Corp.
- Training time and convergence: Bao maintains performance as the dynamic IMDb workload shifts, with its performance curve remaining straight after a short initial period.This indicates that workload changes did not produce a significant change in query performance.
- Predictive model and strategy: Bao learns a diverse, query-specific hint strategy: 35 hint sets were selected over 100 times for IMDb, while no single hint set outperformed PostgreSQL.The single best hint set, disabling loop joins, performed significantly worse than PostgreSQL, whereas Bao’s strategy performed better than that baseline.
- Resiliency to changing hint behavior: When the Temp hint set changed from optimal plans to cross joins after one hour, Bao selected it over 95% of the time beforehand and quickly stopped using it afterward.This experiment tests adaptation to a hint set whose performance changes over time.
6. CONCLUSION AND FUTURE WORK
Bao steers query optimizers with reinforcement learning and query-specific hints, combining Thompson sampling with tree convolutional neural networks. It reduces median and tail latencies under dynamic workloads, data, and schema, while future work targets cloud integration and predictive cost modeling.
- Bao steers a query optimizer using reinforcement learning and selects query-specific optimizer hints with Thompson sampling and tree convolutional neural networks.
- Bao can match open-source and commercial optimizer performance with as little as one hour of training time.
- Bao reduces median and tail latencies even when workloads, data, and schema change.
- Future work will investigate Bao for resource utilization in multi-tenant cloud environments with scarce disk, RAM, and CPU resources.
- Future work will test whether Bao’s predictive model can serve as a cost model in traditional database optimizers.