PreQR: 如何对数据库中的查询进行表征？

本文是对 SIGMOD 22 上的工作 ### PreQR: Pre-training representation for SQL understanding 的浅析。

简介

Query Representation 是 AI4DB 很多任务的必经之路之一。比如基于机器学习的基数估计在预测前需要把查询表征为一个向量；自动调参时需要抽取 Workload 中的关键特征以根据特征推荐参数。

已有的 Query Representation 一般基于 One-hot encoding 的方法简单根据 SQL 关键字进行编码，这种方法存在如下问题：

忽略了 SQL 的结构信息
对表、列的编码忽略了 SQL 的上下文语义，包括数据库 schema、主键外键信息等
SQL 中的数值都被 normalize 到了 [0,1] 区间上，忽略了数据分布的特征

PreQR 基于 NLP 技术，利用 BERT 设计了一个预训练模型用于表征 SQL，主要贡献点如下：

第一个面向 SQL 表征的预训练模型
PreQR 的 Embedding 方法可以表征查询结构、列上的数据分布
设计了一个 Graph Based 的模型用于表征数据库 schema 定义
基于 Attention 机制识别查询的结构化、schema 信息

PreQR 的基本架构如图所示，可以分为 Input Embedding --> Query-Aware Schema --> SQLBert 三大块：

Input Embedding: 生成一个初始 encoding 结果，相似的 SQL 会共享相同的 state representation，另外还会生成 token/position embedding 结果。
Query-Aware Schema: 生成 schema representation。由于一条查询往往只涉及数据库中极少的表和列，因此设计了一个自适应模型去识别并生成 schema 的 sub-graph 的 embedding
SQLBert: 将上述 Embedding 信息作为输入，合并并生成最终的 embedding 结果

问题定义

给定数据库 D，schema S，训练一个模型 F 使：Q x D x S -> Y. Q 代表数据库 D 上查询最频繁的 SQL 集合，Y 代表查询的最终编码表示。

本文还简化了一些问题设定：

认为数据库中的数据分布不会随着时间而改变（这点似乎不那么 make sense）
schema 不会被修改
最频繁的查询可以通过有限数量的查询模板来表示

Input Embedding

为了表示 SQL 多结构，PreQR 首先把 SQL 转换成了有限自动机（其实就是做了个 Parse）。具体来说，T = {t1, t2...} 表示查询 q 的 token，有限状态机有对应 token 的起始状态 a(start), a(end )。a_end 可能有多个。

对每个 group（如 JOB、JOB-Light...），PreQR 会对 SQL 做聚类然后生成一个或多个查询模板。查询模板数量不多，JOB-Light 上只有 1 个，JOB 上有 3 个。

对每个模板构建一个子自动机(sub-automaton)，最后生成的自动机是把所有子自动机合并的结果，合并采用了 maximal prefix strategy（就是把自动机的相同前缀合并）

由于 SQL 中的词汇和自然语言差距极大，PreQR 采用了两个字典来做 token embedding。一个是 WorldPiece 字典（词汇量30000），在 training 过程中，我们随机移除了一些 input markers，然后去预测 mask 掉的是什么（常用的 NLP 训练技巧），然后 mask layer 里用的是 database-specific vocabularies。

对于 SQL 语句中的 values，PreQR 将其转换为离散的 token，如年份可以离散化到 [1900,2000], [2000,2010]... 等区间里。最后 SQL Embedding 被看作一个 token embedding 后的序列，顺序由自动机决定。

这里似乎也带来一个问题：如果一个 Dataset 中突然来了一个和已有模板均不匹配的查询，似乎就没办法** Embedding** 了？这大概就是要增加限制条件 3 的原因

Query-Aware Schema

High Level 的角度来说，Query-Aware Schema 环节主要工作是：

把 schema 转换为一个图
采用一种图编码方法生成当前 schema 的图表示
对 graph encoding 的结果降维
把 query-aware schema 的编码结果输入到 SQLBert 模型中

图本质上是 G={V,E,R} 的集合，包含顶点、边、边的标签。顶点有 Column、Table 两种类型，边的标签，即 R 表明了节点和节点的关系，如下表所示。

将节点 token 化之后用双向 LSTM 对节点进行 Encoding：

然后采用经典的 GNN 前向传播模型根据邻居节点中的信息逐步更新 Graph 中每个节点的状态：

接着用平均池化（Average Pooling) 技术降低 embedding 维度后得到 Graph 的全局表示：

L 是神经网络最后一层的输出

PreQR 还自行在每一层和其下一层间创建了边，以保证 L 层能学到 L+1 层中的表示。

由于 Schema 有可能非常复杂，包含成百上千个表和列，其中可能只有少部分表和列对查询是重要的。因此在把 Graph Embedding 的结果作为 SQLBert 的输入前，PreQR 还做了一次 Self Attention。

SQLBert

使用了最经典的 Trm_g 模型，和原文比几乎没有任何改进，只是说明了一下它 layer 数量什么的，连图不到半页。

SQLBert 的输入就是 Input Embedding 和 Schema Embedding 结果。稍有不同的地方在于：

Schema Emebding 作为 K 输入，将 Input Embedding 做完 Self-Attention 之后的结果作为 Q 输入，V 本质上是根据当前查询特征再次对 Schema 中特征的权重做了一次调整。

最终结果前，将 Input Embedding 这一层在全连接神经网络输出后（e_q）。直接和 Schema Embedding 做全连接神经网络后的结果（e_g）通过 ResNet 拼接，再做 Linear 和 Softmax 操作。

另外本文稍微介绍了一下训练时的设置。三个模块训练时都使用了 NLP 中经典的 "Masked LM"方法，随机 mask 掉 15% 的 token，然后让模型去预测，以此训练网络。

模型更新

以上讨论都基于数据库 schema 和数据分布以及查询模板不变的假设，PreQR 还说了一下假如它们变了怎么办：

数据分布发生巨大变化。则只影响 Input Embedding 中的 token embedding，只需要对 value 分布做一个增量学习即可，大概需要几分钟的训练时间。
schema 被更新了。需要对 schema graph model 做增量训练，需要耗费若干小时。
查询模板发生了变化，即出现了当前模板无法处理的新查询。则需要在 Input Embedding 阶段重新训练，生成新的查询模板，耗费 5-10 小时。
PreQR 整个模型的重新训练时间大约 < 20 小时

但是这里并没有说明数据量和硬件环境。

实验结果

在四类需要做 Query Encoding 的任务上评估 Encoding 效果的好坏：

查询聚类、基数估计、代价估计、SQL2Text