一种将 Tree-LSTM 的强化学习用于连接顺序选择的方法【SQL查询】

【导读】

本篇博客讲解的是 2020 年由清华大学李国良教授团队发表在 ICDE 上的论文，介绍它所提出的算法与实验结果，并结合实际情况给出一些思考。

原文链接：

http://dbgroup.cs.tsinghua.edu.cn/ligl/papers/icde2020-learnedjoinorder.pdf

Part 1 - 背景

连接顺序选择 (Join Order Selection, JOS) 是一个为 SQL 查询寻找最佳连接顺序的问题。作为数据库查询优化器的主要关注点已经被广泛研究了几十年。这个问题很难，因为它的解空间很大，彻底遍历解空间来寻找最佳连接顺序的代价是非常昂贵的。传统的方法通常基于基数估计和代价模型，与启发式剪枝相结合，通过一些剪枝技术来缩小搜索可能的连接顺序的解空间。

尽管经过几十年的努力，传统的优化器在处理复杂的 SQL 查询时仍然存在低可扩展性或低准确性的问题。基于动态规划 (DP) 的算法通常会选择最优方案，但代价非常昂贵。启发式方法，如 GEQO、QuickPick-1000 和 GOO，可以更快地计算计划，但往往产生糟糕的计划。

近年来，基于机器学习 (ML) 和深度学习 (DL) 的基于学习的优化器 (Learned Optimizer) 方法在数据库界越来越受欢迎。特别是基于深度强化学习 (Deep Reinforcement Learning, DRL) 的方法，如 ReJOIN 和 DQ，已经显示出了令人满意的结果 —— 它们可以生成与原生查询优化器一样好的计划，但在学习后执行的速度相比原生查询优化器更快。但是，现有的基于 DRL 的学习优化器将连接树编码为固定长度的向量，其中向量的长度由数据库中的表和列决定。这将导致两个问题：

1. 这些向量不能捕获连接树的结构信息，可能会导致连接顺序选择到一个糟糕的计划 ;
2. 当数据库模式发生变化，如添加列 / 表或多别名表名时，这个基于学习的优化器将会失效，这需要一个新的不同长度的输入向量，然后再重新训练神经网络。

在本文中，作者提出了一种新的基于学习的优化器 RTOS，它使用强化学习和树状结构长短期记忆 (tree-LSTM) 来进行连接顺序选择。RTOS 对现有的基于 DRL 的方法进行了两方面的改进:

(1) 采用图神经网络来捕获连接树的结构 ;

(2) 支持对数据库模式和多别名表名的修改。

在 Join Order Benchmark (JOB) 和 TPC-H 上的大量实验表明，RTOS 优于传统优化器和现有的基于 drl 的学习优化器。

Part 2 - 相关工作

DRL，深度强化学习相关概念：

关于强化学习的概念性解释在早先的博客《Constraint-aware SQL Generation Using Reinforcement Learning》中也有提到，这里不再重复赘述 Agent、Environment、State、Action、Reward 等概念，感兴趣的读者也可查找相关资料加深了解。

在 RTOS 中 Agent 对应优化器，通过与 Environment(对应 DBMS) 的试错交互从反馈中学习。对于图中所示的一条 SQL 查询语句，随着强化学习过程的推进会进入不同的 State(每个 State 对应着连接计划)。RTOS 读取当前计划并利用 Tree-LSTM 来计算估计的长期 Reward(例如，哪个表应该连接起来以使当前状态更完整)。然后它会选择具有最大 Reward(对应最小 Cost) 的预期最佳 Action。当所有的表都被连接 (即，导出一个完整的连接计划) 时，连接计划将被发送到 DBMS 执行。

DQ，《Learning to Optimize Join Queries With Deep Reinforcement Learning》中提出的基于 DRL 学习优化器方法：

DQ 将连接顺序选择问题抽象为一个马尔可夫决策过程。DQ 使用了 DQN(Deep Q Network) 作为强化学习模型来指导连接顺序的选择，由于 DQN 中的代价估计函数 Q-function 计算代价过高，因此使用一个两层的 MLP(Multilayer Perceptron) 来学习 Q-function。神经网络的输入是当前的 State，包括了 SQL 语句中的查询信息以及通过连接左侧右侧状态表示的连接操作信息。

REJOIN，《Deep Reinforcement Learning for Join Order Enumeration》中提出的基于 DRL 学习优化器方法：

ReJOIN 主要使用了邻近策略优化算法 (Proximal Policy Optimization) 来指导连接顺序的选择。其中的关键组成是一个用于策略选取的神经网络。这个神经网络通过输入向量化的状态信息进行训练，这些信息包含：由深度信息组成的树结构向量、连接谓词向量和选择性谓词向量。对于不同的 SQL 语句而言，存在不同的神经网络参数和 Reward，ReJOIN 会根据这条 SQL 语句之前的参数和 Reward，来估算以后的 Reward，从而达到对测试集中的 SQL 语句进行连接顺序选择的目的。

现有基于 DRL 方法的不足：

假设一个数据库有 4 个表 T1、T2、T3、T4。DQ 使用独热编码 (1 表示表在树中，否则为 0) 来编码连接树：图中可以发现 ((T1,T2), T3) 和 ((T1,T3),T2) 具有相同的特征向量 [1,1,1,0]。ReJOIN 使用连接表的深度构造特征向量：图 (b) 中可以发现，每个表的深度是 d = 2，((T1,T2),(T3 T4)) 和 ((T1, T3), (T2,T4)) 具有相同的特征向量 [1/4,1/4,1/4,1/4]。

Part 3 - RTOS 框架

对于上述提到的基于 DRL 的学习优化器方法， 由于它们无法捕获连接树的结构特征，进而导致无法总是获得代价低、执行时间短的方案。此外这些方法无法够适应一些改变数据库模式的操作，如添加列 / 表或多别名表名后，需要重新对神经网络进行训练。因此本文作者提出了 RTOS，致力于捕捉连接树的结构特征，对计划进行更好的估计得到执行时间低的计划，同时能够更快的适应一些会改动数据库有的操作。

整个 RTOS 框架由两部分组成：DRL Optimizer 和 DBMS。

1. DRL Optimizer 中包含 State、Optimizer 和 Memory Pool。State 维护连接进程的当前状态信息，具体内容在 Part4 部分解释；Optimizer 对应于 RL 中的 Agent，它是整个系统的核心部分。对于给定的状态，可以将每个候选连接条件视为一个 Action；Memory Pool 记录 RTOS 生成的计划的状态和来自 DBMS 的反馈。
2. DBMS 用于为 RTOS 生成的连接计划进行开销估算。RTOS 为给定的查询生成一个连接计划，然后将其传递给 DBMS，例如 PostgreSQL。我们使用 DBMS 中的两个组件，Estimator 和 Executor。Estimator 可以给出计划的 Cost，使用统计数据来估算开销，而不需要执行。而 Executor 用户获得计划实际执行的 Latency。

另外，在模型训练方面，RTOS 使用了两阶段训练的方法。具体来说，先用 Cost 作为反馈来训练得到开销较低的模型；再用真实执行的 Latency 作为反馈来调整模型参数。

Part 4 - States 表示

State 的表示由三部分组成：

1. 输入 SQL 查询的表示：用于描述查询的全剧信息，通过神经网络进行表示，包含了查询需要连接的表以及查询中的谓词信息等。
2. 表和列的表示：通过神经网络表示查询中的表和列信息，依据谓词计算特征，对参与计算的列进行表示，然后通过池化得到表的表示。
3. 连接数和连接的表示：通过多种 Tree-LSTM 的组合表示的连接树和连接状态信息。

前两种信息包含了之前方法：DQ，ReJOIN 中提到的特征向量信息，最后一种表示保留了连接顺序以及完整的连接树的结构信息。通过神经网络来表示的状态信息可以实现在需要修改数据库模式，例如新增加一列或者一个表时，可以通过申请新的参数来实现，而不需要重新训练模型。

RTOS 强化学习策略：

本文使用了 DQN 的深度强化学习方法，通过设计一个 Q-network 来估计每个状态的 Q 值。在对训练的反馈进行收集时，由于 Latency 需要实际执行 SQL 查询来获取，收集的代价较高。而 Cost 作为 Latency 的近似的估计，不需要实际执行 SQL 查询，收集的代价较低，并且可以实现和 Latency 一样为指导连接顺序选择。因在 RTOS 的强化学习的反馈和训练的过程分为两个阶段：

第一阶段，使用 SQL 语句的 Cost 作为反馈，生成大量的训练数据，使得 RTOS 的基于 DRL 的学习优化器能够达到传统方法，例如动态规划相同的效果，选择出 Cost 较低的连接顺序的计划。

第二阶段，使用实际执行 SQL 语句得到的 Latency 作为反馈微调 RTOS 的模型。通过实际的 Latency 微调可以矫正 RTOS 中对 Cost 估计错误的连接顺序的计划。即表现为估算的 Cost 值很高而实际执行的 Latency 较低的计划，反之亦然。

Part 5 - 实验结果

论文基于 PostgreSQL 构建了 RTOS，实验采用的 benchmark 是 JOB(Join Order Benchmark) 和 TPC-H。JOB 是一个基于 IMDB 的真实数据集，提供真实的工作负载，它有来自 33 个模板的 113 个查询，包含了 3.6GB 的数据 (计算索引时为 11GB) 和 21 张表。TPC-H 是一个标准的行业数据库基准测试，有 8 张表。论文使用了 22 个模板生成了 4GB 数据和 110 个查询。除了上述提到的 baseline(例如 DQ，ReJOIN，DP) 以外，论文还比较了 SkinnerDB(一种基于强化学习策略的方法) 和 QuickPick(一种贪心算法)。实验结果如下：

COST 结果：

LATENCY 结果：

JOB 数据集上结果：

第一个指标：论文以 DP 作为基线，并按照之前的工作报告其他方法的 MRC(mean relative cost, 平均相对成本 )，其中 MRC = 1 表示与 DP 性能相同。观察 Cost 和 Latency 的结果可以发现，RTOS 在 Cost 上要优于其他四种方法，且几乎与 DP 相等，在 Latency 更佳，是所有方法中延迟最低的。

第二个指标：论文把数据集中的查询按照模板进行分类，计算 GMRL( Geometric Mean Relevant Latency, 几何平均相关延迟)，从图 9 中可以发现，大多数情况下 RTOS 的性能都不弱于 DP，且优于其他方法。对于 T20,T26,T28,T29 等 RTOS 的性能要远远好于 DP，是因为实验数据中的测试模板多数是比较长或者是带有多别名的查询，传统的开销模型往往会出现估算不准确的问题，导致基于 Cost 选择出的计划可能不是最优，而 RTOS 可以解决这一问题，通过从查询执行的 latency 中进行学习，可以修正选择出的计划，实现更高性能。

TPC-H 数据集上结果：

图 10 显示了 TPC-H 不同模板上的，图中没有展示 DRL 方法与 DP 方法执行相同的模板，以获得更好的视图。可以发现对于所有模板，RTOS 生成的规划并不比 DP 差 (在 GMRL= 1 处的水平线)，也比其他方法好。

Part 6 - 总结与思考

本文介绍了将机器学习应用在连接顺序选择问题上的方法：使用 Tree-LSTM 学习连接计划的树形结构的 RTOS。首先，RTOS 使用了深度强化学习技术解决 JOS 问题，在此基础之上，采用了 Tree-LSTM 模型为连接状态进行编码。基于前两个工作设计了基于 DRL 的学习优化器，利用 SQL 解析的知识、DRL 的连接顺序选择、DNN 的计划代价估算，可以在成本和延迟两个基准上生成良好的计划；并通过证明了该方法能够很好地学习连接树的结构，获取连接操作的信息。论文还证明，该代价可以预先训练神经网络，减少延迟调优时间。然而，论文中有提到，在实际使用场景下还存在一些问题，例如：训练时间过长、不能根据不同数据库状态适应性的选择连接顺序等问题。