The Case for Learned Index Structures
and
ALEX: An Updatable Adaptive Learned Index

可行性

索引本身就是模型，B树可以看作是一个映射一个key到一个有序数组中它对应record的位置的模型；哈希索引可以看作一个映射key到一个无序数组对应位置的模型；位图索引可以看作是验证一个数据记录是否存在的模型。

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

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

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

用机器学习的角度来看B树

B树本身就是一个模型，给定一个key，然后预测它在有序的集合中的位置(如果你管这个叫预测的话)。因为效率上的原因，索引一般都不会给到这个key具体的位置，而是会查找叶子节点，然后再去二分查找或者扫描来找到具体的位置。在B树中，一个节点的大小就是一页。

从机器学习的角度来看，B树就是一个回归树，把key和它对应的位置在最小误差和最大误差内进行映射，最小误差是0，最大误差是pagesize，并且保证如果key存在的话，在最大误差范围内必须能够找到。既然B树本身就是一个回归树，那么能不能考虑用其他机器学习模型来取代回归树。

用什么ML模型

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

用ML模型取代B树也有一些问题，B树插入和查找的代价是有限的，能充分利用cache；另外B树的页也不需要是连续的，这些都是ML模型需要考虑解决的问题。好处也是显而易见的，用ML模型取代B树有可能把时间复杂度从O(logn)减少到常数级。ML尤其是神经网络，能够学习不同的数据分布。问题在于平衡模型的复杂度和精确度。

假设B树的一个节点能存100条record，可以看作每经过一个节点，所确定的要查找的key所在的范围就缩小100倍。查一个key一共要遍历log₁₀₀N个节点。用二分查找遍历一个B树节点，大概需要50个时钟周期，很难并行。这50个时钟周期是假设都在cache中的，一次cache miss又会额外需要50-100个时钟周期。每个周期能执行8-16条SIMD指令。所以一个模型用少于50*8个操作将范围缩小100倍时会比B树快。考虑到cache miss，和GPU和TPU性能和性能增长远超CPU，实际上可以使用更复杂的模型。

范围索引模型就是CDF模型

一个模型，它的工作是给定一个key预测它在有序数组中的位置，近似于累积分布函数CDF。
预测一个有序数组中的key的位置近似于CDF，预测可以写成
p = F(Key) * N

几点观察:

1. B树通过建立回归树来学习数据分布，如果用线性模型的话，它是通过线性函数的最小平方误差来学习

2. 估计数据集分布是一个广泛的问题，已经有了很多解决方法

3. 学习CDF也可以对其他索引进行优化，前面提到的数据分布对索引有优化效果

4. 经验CDF近似于理论CDF

一个简单的实验

用200M的服务器日志用Tensorflow实现了一个简单的在时间戳上的二级索引。用ReLU训练了一个两层的全连接的神经网络，每层32个神经元。

实际上效果很差，主要有以下原因：

1. Tensorflow适合训练大模型，调用开销很大，尤其是用python

2. B树或者说决策树，很容易用简单的if语句过拟合，而其他的模型虽然可以更有效的近似CDF的整体形状，但在单个数据上并不准确。也就是它可能比较容易将范围缩小到比如1000，但在这1000个数据中预测就非常不准确了。过拟合对于机器学习都是极力避免的问题，但是在这里越是过拟合，最后查询的精度越高，因为我们查询的本身就是我们存储的数据。

3. B树是cache优化和操作优化的，保持上层节点在cache中，需要其他page再取。但是神经网络需要所有的权值计算出一个预测值，可能这些大量的乘法cost比较高。

RMI

前面简单的学习型索引，效果并不好，主要在于最后一段的预测并不精确，last-mile。比如从100M将预测误差缩小到100是非常困难的。但是将100M缩小到10K的范围并不困难，也就是范围缩小了10000，所以可能取代B树前两层，用更简单的模型实现。同理，再把10K范围缩小到100。

下面是两个方程(损失函数的解释)，已经比较清楚了。

优点： 该模型体系结构具有以下优点：（1）将模型规模和复杂度与执行成本分离开来。（2） 它利用了这样一个事实，即很容易了解数据分布的总体形状。（3） 它有效地将空间划分为更小的子范围，如B-树，以便更容易地以更少的操作实现所需的“最后一英里”精度。（4） 两个阶段之间不需要搜索过程。例如，Model1.1的输出直接用于在下一阶段选择模型。这不仅减少了管理结构的指令数，还允许将整个索引表示为TPU/GPU的稀疏矩阵乘法。

搜索策略与单调性

负载较小时，在有序数组上二分搜索或扫描通常是最快的。

两种搜索策略：

• 模型偏向搜索：把模型预测值作为二分搜索的中间值

• 偏向四分搜索：四分搜索将原来的二分搜索用一个点把数据段分成两份换成用三个点把数据段分成四份，这样可以在三个点上开始prefetch，在数据不在cache中有更好的性能。最初的三个点定为pos -a, pos, pos + a

留下来一些问题，主要就是插入与更新。越是过拟合，新插入的数据越是不同于之前训练好的模型需要重新训练；还有就是对于数据分布的改变，能检测出来，能保证在比如B树的O(logn)的查找和插入代价吗？
对于插入有一个简单的方法，就是用delta index，把所有插入的放在buffer中，周期性的与模型合并并重新训练。

ALEX

学习型索引的核心思想是把索引看作是一个预测位置的模型。它是静态的，只读的。

本文在前面的基础上提出了一个新的，支持点查询，小范围查询，插入更新删除操作的学习型索引ALEX。

两方面的修改：1.提出了一个空间与时间的平衡，不只是提出了可更新的学习型索引，在查找上也更快了。为了支持可更新，这里在叶子节点上用了Gapped Array。2.学习型索引只支持静态的RMI(SRMI)，RMI的层数和每一层的模型数量在初始化时就是固定的。SRMI在插入时表现不好如果数据分布与之前的model不同。ALEX可以在运行时动态的高效的根据工作负载来更新RMI的结构。

通过以上主要是两点，希望能够：插入可以与B+树相比；查找比B+树和learned index更快；索引的空间比B+树和learned index更小；数据存储空间(叶子节点)可以与动态B+树相比。

和learned index的几点不同：

1. 存数据的数据结构是叶子节点，像B+树。允许单个节点灵活的扩展或分裂，在插入的时候也限制了数据移动的数量。在B+树中，每个节点就是一个key和payload的数组，数组最后是空出来的空间，来吸收新插入的数据。类似的，ALEX用的是类似的设计，但在自由空间的使用上会更谨慎，它会把自动空间会放在数组中元素中间，来达到更快的插入。
2. 第二个不同是ALEX用的是指数查找。learned index使用的是二分查找，在错误边界内进行二分查找。实验表明指数查找更快，因为如果模型很好，预测距离正确的位置很近。这也不需要模型存储错误边界了。
3. 基于模型的插入。ALEX会把数据插入到模型预测的位置，减少了模型的错误预测。
4. ALEX会根据工作负载动态的调整RMI的形状和高度。
5. ALEX不需要针对每个数据集和工作负载进行调整参数(前面说过，RMI是静态的，高度和每一层的模型的数量都是在初始化是固定的)。ALEX可以自动的批量加载和用代价模型调整RMI的结构以达到更高的性能。

节点的布局

Data node
像B+树，分为中间节点和数据节点(叶子节点)。数据节点内存的是一个线性回归模型(两个double，斜率和截距)和两个Gapped Array，一个是key，一个是payload。每个gap都会用离它最近的数据来填。为了快速的跳过gap，还维持一个bitmap记录哪些是空的，哪些地方存了数据。

中间节点
把RMI中的模型视为中间节点。中间节点内存了一个现行回归模型和一个包含指向孩子节点的指针数组。和B+树不同的是，定位下一层的孩子节点是计算出来，而不是比较出来的。

ALEX的内部节点和Learned index是为了不同的目的。在learned index中，需要选模型来符合数据。一个完美的内部节点会把数据平均的分给它的孩子，一个完美的RMI有完美的内部节点，会产生大小相等的数据节点。 但是RMI的目标不是产生大小相等的数据节点，而是尽可能使数据节点内的key的分布基本上是线性的，是的一个线性模型能够比较精确的fit数据。

因此ALEX内部节点的角色是提供更灵活的方式来划分key space。[0,1)这个key space中……

多个指针指向同一个节点

ALEX算法

查找和范围查询

为了查找一个key，需要从RMI的根节点，不断用模型计算指针数组中的定位找到孩子节点，直到数据节点。中间节点是完美精度的，不需要在内部节点中遍历。用数据节点中的模型预测要查找的key的位置，然后进行指数搜索。用bitmap跳过中间的gap，必要的话用存下来指向下一个数据节点的指针。

插入

两种情况
插入不满的数据节点，前面提到了，基于模型的插入。如果模型预测的位置不对，先用指数查询找到对的位置(因为要保证有序)，如果是gap，直接插入；否则向最近的gap移动一位。

插入已满的数据节点，两种机制创建更大的空间，扩展和分裂。ALEX用一个简单的代价模型来选择这两种机制。

节点满的标准：ALEX不是等到节点100%满的时候才再去扩展。因为随着gap的减少，插入性能会越来越低(gap越少，要移动的数据可能就越多越远)。所以引入了一个gapped array的最小最大密度的概念。一个节点超过du时认为这个节点已经满了。

节点扩展机制：扩展一个包含n条key的数据节点，分配一个新的更大的gapped array包含n/dl个槽。然后调整或者重新训练这个线性回归模型，再基于这个回归模型进行插入。

节点分裂机制：把一个数据节点分成两个，需要把原节点的数据分给这两个新节点，每个新节点负责一半。两种方式分裂：

1. 像B+树一样旁路分裂，两种情况，如果父内部节点未满，用两个指向新节点的指针取代父节点中指向这个分裂节点的指针，父内部节点有冗余指针指向分裂数据节点时，直接用一般的冗余指针指向两个新内存节点；否则，父节点size double，给每个指针再进行一个冗余的拷贝，直到节点达到最大值。如果父节点达到最大，分裂这个父内部节点，和B+树一样，分裂向上传递。由于限制了所有的内部节点大小是2的幂，可以以保持边界的方式分裂节点，不需要重新训练模型。
2. 把数据节点变成一个内部节点，包含两个孩子，指向这两个新的数据节点。

代价模型：为了决定选择哪种机制，用了简单的线性代价模型来预测平均查找时间和插入时间基于每个数据节点的样本统计：平均指数查找迭代次数和每次插入平均移动的次数。

统计量在创建数据节点时并不知道，为了得到新节点的代价，会假设这个节点内的数据是均匀访问的，插入是根据现有的key的分布插入的。具体就是每个key错误距离的平均log2，现有的key距离最近的gap的平均距离

遍历到叶子节点代价模型预测从根节点到数据节点的时间，两个统计量，要遍历的叶子节点的深度，所有内部节点和数据节点元数据的大小

删除，更新和其他操作

删除不需要移动key，如果删除后小于dl，需要进行节点的收缩或合并。也要用内部节点代价模型来决定

边界上的插入

插入到最左边或最右边
假设插入到右边

批量加载

ALEX支持批量加载，用来初始化索引大量数据或重建索引。目标是找到一个代价最小，操作时间最少的RMI。任何RMI操作都可以分成遍历到叶子节点和在叶子节点操作两部分。所以RMI代价也结合了TraverToLeaf和intra-node cost两部分。

批量加载算法： 从根节点开始，每个节点独立的决定这个节点是中间节点还是数据节点。如果是中间节点，fanout是多少。fanout必须是2的幂，孩子节点平分当前节点的key space。因为用的时候线性模型，所以每个节点造成的影响对于全局是加法。
如果一个节点决定是内部节点，它的孩子节点递归进行这个操作。
The fanout tree