likes
comments
collection
share

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

作者站长头像
站长
· 阅读数 19

本文是对 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 结果

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

问题定义

给定数据库 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(就是把自动机的相同前缀合并)

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

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

由于 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 环节主要工作是:

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

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

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

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

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

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

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

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

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

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

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

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

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

SQLBert

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

使用了最经典的 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 还说了一下假如它们变了怎么办:

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

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

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

实验结果

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

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

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

转载自:https://juejin.cn/post/7282666367237259264
评论
请登录