AI4DB试玩-Bao/Balsa适配PG18

为了准备IFClub社区10.25西安站活动分享话题，老杨最近利用空闲时间研究了几篇Learned Based Optimizer相关论文。第一感触就是好多都是Kpi工程，代码质量和风格一言难尽，感觉就是发了paper后就不管了，没有任何维护痕迹...

我找了两个自己感兴趣的方案：Bao和Balsa，并都适配到了PG18，本篇主要讲讲适配过程以及对AI4DB的一点愚见。

一、传统优化器的局限性

传统优化器，以Cost Based Optimizer为例大多依赖静态统计信息进行基数预估和代价计算。

因此会存在一些局限性，导致代价预估偏差，直接的表现就是优化器选择的最优计划实际执行效果很不理想。

1、静态统计信息不能完全反馈出数据特征

2、数据分布均匀假设，其实很多时候数据分布可能是倾斜的

3、算子的代价预估算法本身有缺陷，导致预估误差

4、越来越复杂的现代数据类型和模型，没有合适的统计信息来表达其特征

5、缺少对动态系统负载的考虑 ...

这其实是能够理解的，CBO存在误差，就好比大模型存在幻觉一样。

二、如何防止次优执行计划

1、提高CBO代价预估准确性

采用更全面的统计信息、更优秀的代价预估算法、考虑系统负载，等等这些都可以提升预估准确性，在执行计划生成阶段解决。

但是，我个人认为这可能不是最优解，因为预估总可能会有误差。

2、引入执行计划管理机制

采用计划基线机制，baseline永远是一组接受过验证的执行计划，基本确保实际执行的计划一直在baseline里，而不是永远执行CBO直接生成的计划，减少cost预估误差导致的性能劣化。

可参考经典的Oracle计划决策树：

在我看来执行计划管理本身就是机器学习赋能传统优化器的一种实现，不断地对样本计划进行分析训练，反馈实际性能表现去进行进一步的学习和决策。不难发现在各种CBO中SPM是依赖Hint等基础特性来实现的。

可惜的是直到目前PG社区并没有支持执行计划管理。不过已经有很多专家学者尝试使用强化学习模型来增强数据库选择best plan的能力，很多论文将这种方案命名为Learned Based Optimizer，我个人更愿意称为SPM++

Bao & Balsa

这两个模型也相对古老了，都是适配的PG12，不过我认为思想是比较经典的，所以打算适配下PG18玩玩。

Bao: Making Learned Query Optimization Practical

Bao是适用PostgreSQL的学习型优化器[1]，使用基于树卷积神经网络TCNN强化学习的成本模型来预测执行计划的性能，但其结合的hint集合有明显的局限性，仅支持ScanMethod(SeqScan,IndexScan,Index Only Scan,Bitmap Index Scan)及JoinMethod(Nestloop,HashJoin,MergeJoin)。

Bao有两个组件：Bao server（一个独立的Python 应用程序）和PostgreSQL扩展pg_bao（直接与PostgreSQL集成并与Bao server通信）

适配18还得费点功夫，主要是pg_bao扩展的调整。

PG18和12版本在Planner和Executor函数存在差异，以及部分buffer结构及函数变动，需要调整的函数如下：

// main.c
static PlannedStmt* bao_planner(Query *parse, const char *query_string,
                                int cursorOptions, ParamListInfo boundParams);
static void bao_ExecutorStart(QueryDesc *queryDesc, int eflags);
static void bao_ExecutorEnd(QueryDesc *queryDesc);
static void bao_ExplainOneQuery(Query* query, int cursorOptions, IntoClause* into,
                                ExplainState* es, const char* queryString,
                                ParamListInfo params, QueryEnvironment *queryEnv);
                              
// bao_bufferstate.h                     
static char* buffer_state();

// bao_planner.h
BaoPlan* plan_query(Query *parse, const char* query_string, int cursorOptions, ParamListInfo boundParams)

Bao的部署参考官方文档： https://rmarcus.info/bao_docs/introduction.html

值得注意的是训练需要的数据imdb，目前需要从 https://event.cwi.nl/da/job/imdb.tgz下载。

从简单测试结果看，使用Bao时整体性能提升一倍左右。

Balsa

Balsa: Learning a Query Optimizer Without Expert Demonstrations

Balsa 是一个学习型查询优化器[2]。它使用深度强化学习和模拟到现实学习，通过反复试验来学习优化SQL查询。

值得注意的是，Balsa是第一个端到端学习优化器，它不依赖于现有专家优化器计划的学习，同时能够超越专家计划的性能，有时甚至超越相当大的幅度。

简单来说就是使用神经网络强化学习结合pg_hint_plan来影响和预估执行计划性能。

适配不费事，直接安装最新版本PG18和pg_hint_plan-REL18_1_8_0扩展即可，支持的hint相对多一点。同样使用上边的链接下载imdb训练数据。

部署过程请参考：https://github.com/balsa-project/balsa/tree/main?tab=readme-ov-file#setup

从测试结果来看，经过一段时间训练后Balsa下workload的latency要优于PG专家优化器。

浅谈AI4DB

我认为AI4DB确实能覆盖DB的各方各面，比如刚提到的AI Based Optimizer让“最优计划"更加得稳定；还比如DBAIOPS可能经过对应模型的训练后，就能够真正实现数据库自治，真正的实现“自动驾驶”。不是case by case式的能力，而是泛化的可持续自动演进增强的能力。

so，各位DBA小伙伴，大家慌吗

Reference

[1]

https://files.zotero.net/eyJleHBpcmVzIjoxNzYwODgxNjM2LCJoYXNoIjoiOWI2ZmNlM2EyNTE0MWU3OTg3ZWU3MzdlN2JhMjA2ODgiLCJpdGVtIjoiMzY2NTI0MlwvNklIWlM0N0IiLCJmaWxlbmFtZSI6IlJ5YW4gTWFyY3VzIGV0IGFsLiAtIDIwMjEgLSBCYW8gTWFraW5nIExlYXJuZWQgUXVlcnkgT3B0aW1pemF0aW9uIFByYWN0aWNhbC5wZGYiLCJjb250ZW50VHlwZSI6ImFwcGxpY2F0aW9uXC9wZGYifQ%3D%3D/b838cd9c2746c2979b4cbe90de1ed5244569e3f5d8f19ed123c20ebb317cabb1/Ryan%20Marcus%20et%20al.%20-%202021%20-%20Bao%20Making%20Learned%20Query%20Optimization%20Practical.pdf

[2]

https://zongheng.me/pubs/balsa-sigmod2022.pdf