本文介绍了一项由Ziniu Wu、Ryan Marcus、Zhengchun Liu、Parimarjan Negi、Vikram Nathan、Pascal Pfeil、Gaurav Saxena、Mohammad Rahman、Balakrishnan (Murali) Narayanaswamy和Tim Kraska等研究人员共同完成的研究，题为《STAGE: Query Execution Time Prediction in Amazon Redshift》。该研究于2024年6月9日至15日在ACM SIGMOD/PODS国际数据管理会议（SIGMOD’24）上发表，主要研究领域为数据库管理系统（DBMS）中的查询执行时间预测。

研究背景与动机

在现代数据库管理系统中，查询执行时间（execution time, exec-time）的预测是一个关键组件，尤其是在Amazon Redshift这样的云数据仓库中。准确的执行时间预测对于许多下游任务至关重要，包括查询优化、工作负载调度、资源管理等。然而，现有的执行时间预测技术（包括Redshift中使用的技术）存在冷启动问题、估计不准确以及对工作负载和数据变化的鲁棒性不足等问题。为了解决这些问题，研究人员提出了一种新颖的分层执行时间预测器——STAGE预测器。

研究方法与流程

STAGE预测器由三个模型状态组成： 1. 执行时间缓存（Execution Time Cache）：记录最近执行的查询及其执行时间，当相同的查询再次提交时，直接返回缓存中的执行时间。 2. 轻量级本地模型（Local Model）：针对每个Redshift实例进行优化，使用贝叶斯集成（Bayesian Ensemble）的XGBoost模型来预测查询的执行时间，并提供不确定性测量。 3. 复杂全局模型（Global Model）：基于图神经网络（Graph Neural Network, GCN）的全局模型，能够在所有Redshift实例之间进行知识迁移，适用于冷启动或本地模型不确定性较高的情况。

STAGE预测器的工作流程如下： 1. 当新查询到达时，首先检查执行时间缓存。如果缓存中有匹配的查询，则直接返回缓存中的执行时间。 2. 如果缓存未命中，则使用本地模型进行预测。如果本地模型对预测结果有较高的置信度，则返回该预测。 3. 如果本地模型对预测结果不确定，则使用全局模型进行预测。全局模型虽然推理延迟较高，但由于仅在必要时使用，整体推理开销被分摊。

实验结果与结论

研究团队在Amazon Redshift的实际生产环境中对STAGE预测器进行了端到端评估。实验结果表明，与Redshift中现有的执行时间预测器相比，STAGE预测器能够将查询执行延迟平均降低20%。此外，研究还通过消融实验验证了STAGE预测器中每个组件的性能和鲁棒性。

STAGE预测器的主要贡献包括： 1. 提出了一个分层模型框架，结合了缓存、本地模型和全局模型的优势，能够在保证低推理延迟的同时提供高精度的执行时间预测。 2. 通过实验验证了STAGE预测器在Redshift中的实际效果，显著提升了查询执行效率。 3. 提出了未来研究方向，包括在其他数据库任务中应用分层模型、改进执行时间预测中的环境因素建模等。

研究亮点

1. 分层模型设计：STAGE预测器首次在数据库管理系统中引入了分层模型的概念，结合了缓存、本地模型和全局模型的优势，有效解决了冷启动、推理延迟和预测准确性的问题。
2. 实际生产环境验证：研究在Redshift的实际生产环境中进行了大规模实验，验证了STAGE预测器的实际效果，证明了其在提升查询执行效率方面的显著优势。
3. 不确定性测量：本地模型通过贝叶斯集成XGBoost模型提供了可靠的不确定性测量，能够在预测不确定性较高时自动切换到全局模型，提升了预测的鲁棒性。

未来研究方向

研究团队指出，STAGE预测器的设计理念可以扩展到其他数据库任务中，例如查询优化和假设推理。此外，未来的研究可以进一步探索如何将环境因素（如内存和CPU利用率）纳入执行时间预测模型中，以进一步提升预测的准确性。

总结

STAGE预测器通过分层模型的设计，有效解决了Amazon Redshift中查询执行时间预测的冷启动、推理延迟和预测准确性问题。实验结果表明，STAGE预测器能够显著提升查询执行效率，并为未来的数据库研究提供了新的方向。