Learned index 综述

十二月 10, 2023 3141

Learned index 综述

学习索引 (learned index) 技术

摘要

索引技术在数据库中扮演着重要角色, 在数据库中有 b tree, hash 索引等. 但是这些数据结构没有利用到被索引数据分布的信息.

本文总结了谷歌发表于 NIPS 2017 和 SIGMOD 2018 的两篇文章, 探讨如何利用机器学习技术得到更好的范围索引, 使用机器学习模型替代传统的 B-trees 索引, 以及在这之上的改进, 使得学习索引技术可扩展, 解决学习索引更新问题.

简介

数据库依靠索引结构来高效地执行众多核心操作. B-trees 用于范围检索, Hash-Maps 用于按照键值搜索, 布隆过滤器用于判断 record 是否存在.

  • 但是这些索引是广义目的的数据结构, 没有利用到被索引数据的信息, 在一些极端情况下会表现得较差,
  • 比如当数据键值为从 1 递增到 n, 如果使用 b-tree 索引, 查询的时间复杂度为 o(logn), 实际上只需要 o(1) 的复杂度. 同样索引的空间复杂度也只需要 o(1), 而不是 o(n).

如果能够学习到一个模型反映数据的模式, 就有可能自动构建索引结构即学习索引 (learned index) , 进而得到显著性能提升.

  • Kraska[1], [2] 等人提出使用机器学习模型代替传统的 B 树索引，并在真实数据集上取得了不错的效果，但其提出的模型假设工作负载是静态的、只读的，对于索引更新问题没有提出很好的解决办法. [3] 提出了基于中间层的可扩展的学习索引模型 Dabble，用来解决索引更新引发的模型重新训练问题.

内容

B-Trees 可以看成模型

范围索引, 比如 B-Trees, 可以视为一个机器学习模型, 给定输入 Key 返回 position.

一个极端的例子与数据分布

一个极端的例子，一亿条数据，主键是整型从1到1亿，如果用B树的话，从根节点到一个个中间节点，非常麻烦，如果知道数据是这样分布的话，那肯定不会选择建索引，直接通过偏移量访问了。索引，一般索引都是不知道数据分布的，这样更具有通用性，但是通常来说，知道数据分布对于索引的优化(空间和查询效率)是通常都是非常有效的(比如对于稠密数据，用ART肯定没有用普通的span为8的前缀树表现好，尤其是索引的构建上)。但是对于这些传统索引来说，无论是了解数据分布，还是让索引为每种数据分布做专门的优化都过于复杂了。

于是有了用机器学习来学习出一个能够反映数据分布，并能自动生成出特定的索引结构的想法。

前面提到了B树需要保证要要查询的key必须在一个区间内，第一感觉用机器学习模型不简单，但事实并非如此。因为这个保证是给已经存储的key，而不是任意key，对于新来的key，B树需要重新平衡，对于学习型索引需要重新训练。再有就是这个错误边界并不是一定要有的，因为数据是有序的，误差也很容易纠正通过在预测值附近进行local search，比如用指数查找。因此，可以用其他回归模型，比如线性回归或神经网络来取代B树。

学习索引 (learned index)

学习索引要做的事是学习到一个映射函数, f(key)->pos, 将 key 写成 x, pos 写成 y, 希望学习到一个模型 f(x) 约等于 y.

  • 因为 x 是排序过的, 所以 f 实际上是在给数据的 CDF 建模 [4]. 下图是学习索引的一个示例.

模型架构

上面的问题实质上是一个回归问题（给$x_i$求y。通过输入特征来预测连续数值输出）, 可以使用均方差作为损失函数, 训练模型.

同时为了提高模型的准确率, 使用层级模型. 损失函数如下所示

模型架构如下所示:

层级模型和集成模型的思路相似, 都是通过集成弱分类器得到性能更好的分类器.

学习索引 (learned index) 如何预测

不同于一般的机器学习任务, 在使用学习索引 (learned index) 的时候, 需要考虑模型输出偏差 delta

  • 需要找到最大的误差 delta_max, 在 [y_pred-delta_max, y_pred+delta_max] 范围内搜寻起始点, 但是这个搜寻起始点可以使用二分查找, 时间复杂度不是很高.

Dabble 模型

之前其提出的模型假设工作负载是静态的、只读的，对于索引更新问题没有提出很好的解决办法. 因此提出了基于中间层的可扩展的学习索引模型 Dabble

  • 用来解决索引更新引发的模型重训练问题

实验：

Dabble 模型中设计了 3 个机制:

  • (1) 利用聚类算法将数据集按照数据的聚集方式划分为 K 个部分, 得到 K 个数据区域, 使得在每一个数据区域内部, 数据分布尽可能相同;
  • (2) 对于每一个数据区域, 分别利用神经网络对其进行训练, 使网络能够对该数据区域的分布得到一个比较好的拟合;
  • (3) 在神经网络训练阶段, 重点关注访问频率高的热点数据, 从而使神经网络对这些数据的预测精度更高. 另外, 为了更快地处理动态插入以及更新操作, Dabble 模型提出了一种基于中间层的数据插入机制, 从而使得不同的神经网络模型之间解耦, 模型之间保持独立性. 当数据插入时, 只需要重新训练有数据插入的那个数据区域对应的模型即可, 不需要对整个 Dabble 重新训练, 从而提高了模型的可扩展性.

Dabble 模型架构图如下所示

实验

论文 [1][2] 在 2 亿网站日志数据 weblogs 数据集上跑了实验. 将经过 cache 优化过的 B-Tree 页大小为 128 作为 baseline, 对于学习索引, 使用两阶段模型, 在没有 GPU 和 TPU 的机器上运行结果如下图所示. 可以看到学习索引比 B-Tree 更快的同时, 内存占用更少.

论文 [3] 在 weblogs 和 lognormal 数据集上进行评测, 结果如下图.

总结

本文总结了三篇文章, 主要探讨了结合机器学习技术的范围索引架构. 相较于传统的索引结构, 可学习型的索引展示了更好的可适应性.

参考文献

[1] Alex Beutel, Tim Kraska∗, Ed H. Chi, Jeffrey Dean, Neoklis Polyzotis “A Machine Learning Approach to Databases Indexes”NIPS 2017

[2] Alex Beutel, Tim Kraska∗, Ed H. Chi, Jeffrey Dean, Neoklis Polyzotis “The Case for Learned Index Structures”SIGMOD 2018

[3] 高远宁, 叶金标, 杨念祖, 高晓沨, 陈贵海. 基于中间层的可扩展学习索引技术 [J]. 软件学报, 2020, 31(3): 620-633. http://www.jos.org.cn/1000-9825/5910.htm

[4] M. Magdon-Ismail and A. F. Atiya. Neural networks for density estimation. In M. J. Kearns, S. A. Solla, and D. A. Cohn, editors, Advances in Neural Information Processing Systems 11, pages 522–528. MIT Press, 1999.

查看评论