The Case for Learned Index Structures

17年的Google和MIT的旧文[1]，最近因为SageDB而重读。使用机器学习方法来优化只读场景下的纯内存索引系统。注意本文非直译，要看原文可以到SIGMOD找pdf。

The Case for Learned Idx Structs

文章主要思路是通过学习key的顺序、结构等来预测record在位置、存在与否等。效果方面，据称部分场景下，相对b-tree可以优化70%的内存占用。

最大价值其实在于使用ML来优化（索引）系统这个新的方向，某种程度上打破了系统优化的沉寂。

Range Index

审视下btree查找完成的功能：输入一个key，每次选出一个可能的范围(分支节点)，直到最后命中(叶子节点)。这其实跟ML中模型类似。

换句话说，若能估算出数据的累积分布（记作F），那么查询key所在位置，也可以看成是 pos = F(key) * N 这样一个过程。

基于此，文章首先尝试了一个较为简单的朴素方案：使用tf训练并运行一个2层全连接的神经网络，每层32个单元，使用ReLU作为激发函数。

可以预料到，这个方案运行效果很差，单次查找耗时比btree高了2个数量级。原因是多方面的：

1.TF并不针对小数据集优化。在ns级响应需求下，简单使用TF不够。

2.单一的神经网络在最后的精细部分(last mile)，需要花费大量的计算与存储资源，而效果却不见得能有多好。

3.btree的设计考量了内存优化等，而朴素方案显然还较为粗糙(全连接层)，有提升空间。

The RM-Index

针对上文朴素方案的三个问题，文章进行了一系列的优化。

首先不再使用TF进行推断查找，而是开发发了一个叫LIF的框架，从TF模型中提取权重参数，直接生成专为小数据集优化的高效C++代码。

其次，使用RMI递归（而非单一）模型，在逐步缩小key的范围。由于每次问题被分解到小范围内进行，资源消耗得到改善的同时，模型的精度可以更好提升。RMI每层的输出是下一层的输入，有利于使用TPU/GPU进行优化。RMI中可以在不同的stage混合使用btree在内的不同类型模型来达到最佳效果，所以理论上不会比单纯的btree差。叠加多个模型的方案在传统设计中使用较少。:)

并且，由于模型实际上已经预测出key的位置（position），而不仅仅是范围（range），本文使用了两种新的查找算法（MBS、BQS）利用该信息来更高效地进行查找。

针对三个问题的一连串的优化之后，模型的训练和运行有了明显的改善。训练过程使用sgd只需要一次或少量的访问就可以。2亿条记录能在秒级完成。

对于整型数据集，相同内存占用时下，RMI经常能较btree有数倍性能提升。或者说性能相同时，内存占用会有数量级的优化。见下图：

在尽量公平相似的场景下，跟其他相关方案（FAST等）比也有明显优化：

在字符串场景下，优化场景并不显著，原因可能在于字符串比较太耗时，model执行时间过长等。后者用TPU/GPU可能会有优化空间。

Point Index

point idx（hash索引）的优化空间在于，典型的数据冲突可能会有33%（如生日）。然而，实际冲突减少和运行优化效果取决于两个主要方面：

1.数据本身的分布情况。比如均匀分布场景下，learned idx不会比普通的随机hash函数好多少；

2.其他payload等

从文章的数据集来说，还是有效果的：

Existence Index

存在性的索引，本文优化方面主要在于内存占用：10亿记录（Google的测试数据集）典型bloom filter需要1.76GB，如果要FPR为0.01%，则需要2.23GB。

前边在索引需要学习数据分布，而存在性索引，需要让合法的key相关，非法的key相关，而合法key与非法key间不相关。这其实就很像分类问题了。

另一方面，由于ML的特别，FPR下降时，FNR通常会上升。这跟bloom filter要求的FPR尽量小，FNR为0有冲突。解决方案是，在模型判断为false时，另行使用一个overflow bloom filter进行判断。由于bloom filter的大小于数据集相关，因为后接的bloom filter大小于FNR（即false部分相关），这要远小于传统方案中的大小。

需要注意的是，该文章只研究了查询存在一定规律的场景，并在此这上建立模型。这在实际使用时要视业务场景而定。

在符合条件的情况下，1.7M条URL，

1.FPR 0.5%，FNR 55%时，2.04MB->1.31MB，减少36%

2.FPR 0.1%，FNR 76%时，3.06MB->2.59MB，减少15%。

Conclusion and Future Work

文章研究的是单一维度索引，如果能支持多维索引，对现实系统将会有更多帮助。

[1] T. Kraska, A. Beutel, E. H. Chi, J. Dean, and

N. Polyzotis. The case for learned index structures. In SIGMOD, pages 489–504, 2018.