The Case for Learned Index Structures

1. 介绍

无论何时需要有效的数据访问，索引结构都是答案，并且存在各种各样的选择来满足各种访问模式的不同需求。 例如，B 树是范围查找的最佳选择（例如，在特定时间范围内检索一段值记录 (Value)）; HashMap 在单 Key 查找这个领域是无敌的 ; 而 Bloom-filter 通常用于检查值记录 (Value) 是否存在。 由于数据库和许多其他应用的索引非常重要，因此在过去的几十年里，它们已经得到了广泛的优化，以获得更高的内存、缓存和 CPU 效率 [ 28,48,22,11]。

然而，所有这些索引仍然是通用数据结构，假设数据的最坏情况分布，并没有利用现实世界数据中存在的更常见模式。 例如，如果目标是建立高度特定的系统，用来存储和查询具有连续整数键 (Key) 的固定长度值记录 (Value)（例如，键 (Key) 从 1 到 100M），那么设计者就不会使用常规的 B 树索引，因为键 (Key) 本身可以用作偏移量来作查找或者范围查询，达到 O(1) 而不是 O(long n) 的时间复杂度。 而且，索引内存大小将从 O(n) 减小到 O(1) 。也许令人惊讶的是，对于其他数据结构，相同的优化仍然是可能的。 换句话说，了解确切的数据分布可以高度优化数据库系统使用的几乎所有索引。

当然，在大多数现实世界的用例中，数据并不完全遵循已知的模式，为每个用例构建专门解决方案的代价都太高了。 然而，我们认为机器学习为挖掘数据里面的模式和相关性提供了一个机会，从而能够以低工程成本，自动合成我们称为学习索引的索引结构。

在本文中，我们探讨了学习模型（包括神经网络）在多大程度上可以用来代替传统的 B 树到 Bloom-filter 的索引结构。 这似乎与直觉相反，因为机器学习并不提供传统的索引数据结构的输入输出，并且因为最强大的机器学习模型，神经网络的计算一般认为是非常昂贵的。 然而，我们认为，这些明显的障碍都不像它们看起来那么坑爹。 相反，我们使用学习模型的方式可能会带来巨大的好处，特别是在下一代硬件上。

就输入输出的语义来说，索引在很大程度上已经是学习模型，使得用神经网络等其他类型的模型取代它们变得非常简单。 例如，B 树可以被看作是一个模型，它将一个键 (Key) 字作为输入并预测数据值记录 (Value) 的位置。 Bloom-Filter 是一个二元分类器，它基于一个键 (Key) 来预测键 (Key) 是否存在于一个集合中。 显然，这就存在微妙但重要的差异。 例如，Bloom-filter 可能有假阳性 (false positives)，但没有假阴性 (false negatives)。然而，正如我们将在本文中展示的那样，可以通过新颖的学习技术和 / 或简单的辅助数据结构解决这些差异。

在性能方面，我们观察到每个 CPU 都具有强大的 SIMD 功能，并且我们推测许多笔记本电脑和手机很快将拥有图形处理单元（GPU）或张量处理单元（TPU）。 推测 CPU-SIMD / GPU / TPU 的功能将越来越强大，这是合理的，因为比通用指令集更容易扩展神经网络使用的有限的（并行）数学运算。 这样，今后执行神经网络的高成本在未来可能实际上可以忽略不计。例如，Nvidia 和 Google 的 TPU 已经能够在单个指令周期中执行数千次（如果不是数万次）神经网络操作 [ 3 ] 。 此外，有人表示，到 2025 年，GPU 的性能将提高 1000 倍 ，而 CPU 发展已经停滞，不按摩尔定律发展 [ 5 ] 。通过用神经网络取代重分支的索引结构，数据库可以从这些硬件趋势中受益。

重要的是要指出，我们并不主张用学习索引结构来完全取代传统的索引结构。 相反，我们概述了一种建立索引的新方法，它补充了现有的工作，并且可以说为一个有数十年历史的领域开辟了一个全新的研究方向。 虽然我们专注于分析只读工作负载，但我们还概述了如何将这种想法扩展到对写入频繁工作负载的索引做加速。 此外，我们简要概述如何使用相同的原则来替换数据库及其他组件的操作，包括排序和联表 (join)。 如果成功，这可能导致未来数据库的开发方式和现在彻底不同。

本文的其余部分概述如下：在下一节中，我们以 B 树为例介绍学习索引的总体思路。 在第 4 节中，我们将这个想法扩展到 Hash 索引，并在第 5 节中扩展到 Bloom-Filters。 所有部分都包含单独的评估和列出未解决的挑战。最后在第 6 部分我们讨论相关的工作，并在第 7 部分结束。

正在上传…重新上传取消

图1：为什么B 树是模型

2. 范围索引

索引结构已经是模型，因为它们可以“预测”给定键 (Key) 的值的位置。 要看到这一点，请在主键 (Key) 已排序的分析内存数据库（即只读）中考虑一个 B 树索引，如图 [1]（a）所示。 在这种情况下，B-Tree 提供从查找键 (Key) 到排序的值记录 (Value) 阵列内的位置的映射，并保证值记录 (Value) 位置大于等于查找到的位置。 请注意，必须对数据进行排序以允许范围请求。 还要注意，这个相同的一般概念适用于次级索引，其中底层将是 <key，pointer> 对的列表，其中键 (Key) 是索引属性的值，指针是对值记录 (Value) 的引用。

出于效率的原因，通常不会对已排序值记录 (Value) 的每个关键 (Key) 字进行索引，而只是每个 n 个值记录 (Value) 的一个键 (Key)，即每页面的第一个键 (Key)。 [2] 这有助于显着减少索引必须存储的键 (Key) 数量，而不会有任何显着的性能损失。 因此，B 树是一个模型，在 ML 术语中是回归树：它将键 (Key) 映射到具有最小和最大误差的位置之间（最小误差 0，最大误差页面大小）并保证可以在该地区找到该键 (Key) 锁对应的值记录 (Value)（如果存在）。 因此，我们可以用其他类型的机器学习模型（包括深度学习模型）取代 B 树索引，只要它们也能够提供类似的有关最小误差和最大误差的有力保证。

乍一看，可能很难为其他类型的 ML 模型提供相同的错误保证，但它实际上非常简单。 B-Tree 仅为存储的数据提供这种保证，而不是针对所有可能的数据。 对于新数据，B 树需要重新平衡，或者在机器学习的术语里面叫重新训练，通过训练来提供相同的误差保证。 这就极大地简化了问题：最小误差和最大误差是经过训练的（即存储的）数据的最大误差。 也就是说，我们唯一需要做的就是对每个键 (Key) 执行训练，并记住一个位置的最好和最差的位置预测。 给定一个键 (Key)，该模型预测哪里能找到相应的值记录 (Value); 如果键 (Key) 存在，则保证处于由最小和最大误差定义的预测范围内。 因此，我们能够用任何其他类型的回归模型（包括线性回归或神经网络）代替 B 树（见图 [1]（b））。

现在，我们需要解决其他技术挑战，然后才能使用学习好的索引替代 B 树。 例如，B 树具有插入和查找的有限成本，并且在利用缓存方面特别好。 此外，B 树可以将键 (Key) 映射到未连续映射到内存或磁盘的页面。 此外，如果查找关键 (Key) 字不存在于集合中，某些模型可能会返回最小 / 最大错误范围之外的位置，如果它们不是单调递增的模型。所有这些都是有趣的挑战 / 研究问题，会在本节中与潜在解决方案一起详细解释。

同时，使用其他类型的模型，特别是深度学习模型作为索引可以提供巨大的好处。 最重要的是，它有可能把 B 树 log n 查找成本为一个常数。 例如，假定数据集具有 1M 个唯一键 (Key)，大小在 1M 和 2M 之间（因此 1,000,009 存储在第 10 个位置上）。 在这种情况下，一个简单的线性模型，由一个单一的乘法和加法组成，可以完美地预测任何键 (Key) 的位置，而 B 树会需要做一个 log n 操作。 机器学习，尤其是神经网络的优点在于，他们能够学习各种各样的数据分布 / 混合和其他数据特征和模式。 显然，挑战在于平衡模型的复杂性与准确性。

2.1 我们可以承受模型有多复杂？ 来做个估算吧

为了更好地理解模型的复杂性，需要知道同样的一段时间内，遍历 B 树可以执行多少操作，以及学习索引需要达到什么样的精度来超过 B 树的精度。

考虑一个 B 树索引 100M 值记录 (Value)，页面大小为 100（译注：也就是单个节点的子节点数量，国内有翻译为阶的）。我们可以将每个 B-Tree 节点视为划分空间的一种方式，减少“误差”并缩小区域以查找数据。 因此，我们说 B-Tree 的页面大小为 100，每个节点的查找精度为 1/100 ，所以我们需要遍历 log(100, N) 个节点。 因此，第一个节点将查找空间从 100 M 缩小到 100 M / 100 = 1 M ，第二个节点从 1 M 到 1 M / 100 = 10 k 等等，直到找到值记录 (Value) 为止。 同时，遍历单个 B-Tree 页面需要大约 50 个时钟周期（我们测量了对超过 100 个缓存驻留记录的二分查找与遍历查找具有大致相同的性能），并且非常难以并行化 [3]。 相比之下，现代 CPU 可以在每个周期执行 8-16 个 SIMD 操作。 因此，只要学习索引模型的 查找精度 / 运算数 超过 (1/100) / 50 * 8 = 400 个算术运算，学习索引模型就会更快（译注：这里隐含了一个公式，查找速度 = 查找精度 / 运算数，也就是单位时间内的 查找概率）。 请注意，这个估算仍假定所有 B 数的页都在缓存中。 单个缓存未命中花费 50-100 个额外的周期，因此可以允许更复杂一些的模型。

此外，机器学习的快速发展正在彻底改变游戏。 它们允许在相同的时间内运行更复杂的模型，并从 CPU 中卸载计算（到 GPU 或 TPU 上）。 例如，NVIDIA 最新的 Tesla V100 GPU 能够实现 120 TeraFlops 的低精度的深度学习算术运算（每个时钟周期≈60,000 次运算） [ 7 ]。 假设整个学习索引都载入了 GPU 的内存（我们在 3.6 节中展示这是一个非常合理的假设），在 30 个时钟周期内，我们可以执行 100 万次神经网络操作。 当然，传输输入和从 GPU 获取结果的延迟仍然明显较高，大约为 2 微秒或数千个时钟周期，但这个问题并非不可克服，可以通过批处理方式，或者更加紧密的集成 CPU/ GPU / TPU [ 4 ] 。 最后，可以预期，GPU / TPU 每秒的浮点 / 整型操作的能力和数量将继续增加，而提高 CPU 执行 if 语句性能的进展基本上停滞不前 [ 5 ]。 尽管我们认为 GPU / TPU 是实践中采用学习索引的主要原因，但本文中我们将重点放在有限的 CPU 能力上，以便更好地研究通过机器学习取代 / 增强索引的影响，排除硬件更改的因素。

正在上传…重新上传取消

图2：作为CDF 的索引

2.2 范围索引模型是 CDF 模型

正如本节开头所述，索引是一个模型，它将一个键 (Key) 字作为输入并预测值记录 (Value) 的位置。 而对于单点查询，值记录 (Value) 的顺序并不重要，对于范围查询，必须根据查找关键 (Key) 字对数据进行排序，以便可以有效地检索范围内（例如，时间范围内）的所有数据项。 这导致了一个有趣的观察：预测给定排序数组内键 (Key) 的位置的模型有效地近似于累积分布函数（CDF）。 我们可以建模数据的 CDF 来预测位置：

正在上传…重新上传取消

其中p 是位置估计， F (Key) 是小于或等于查找键(Key) 的数据出现的概率，也就是累积分布函数(CDF, estimated cumulative distribution function, 详见这篇)，N 是键(Key) 的总数量（另见图[2]）。 这个观察开辟了一套全新的有趣方向：首先，它意味着任何一种索引字面上都需要学习数据分布。 B 树通过构建回归树来“学习”数据分布。 线性回归模型将通过最小化线性函数的(平方) 误差来学习数据分布。 其次，估计数据集的分布是一个众所周知的问题，学习索引可以从之前数十年的研究中受益。 第三，学习CDF 对优化其他类型的索引结构和潜在算法也起着关键作用，我们将在本文后面概述。

2.3 第一个，粗糙的学习索引

为了更好地理解用学习索引取代传统 B 树有哪些技术要求，我们使用了 200M Web 服务器日志值记录，目标是使用 Tensorflow [ 9 ] 在时间戳上建立二级索引。 我们使用 ReLU 激活函数训练了每层 32 个神经元（即 32 个宽度）的双层全连接神经网络 ; 时间戳是输入要素，位置是标签。

之后，我们使用 Tensorflow 和 Python 作为前端，测量随机选择的键的查找时间 (排除一开始跑的一些数据)。 在这种情况下，我们实现了每秒≈1250 次预测，即，使用 Tensorflow 执行模型需要≈80,000 纳秒 (ns)，甚至没算搜索时间。 预测对全量遍历几乎没有什么帮助。 作为比较，B 树遍历同样数据只需要≈300 ns， 小两个数量级，搜索键 (Key) 空间时，并且快 2-3 倍 。 其原因是多方面的：

1. Tensorflow 旨在有效地运行较大的模型，而不是小型模型，因此具有显着的调用开销，尤其是在 Python 作为前端时。
2. 一般来说，B 树，或者一般意义上的决策树，使用少量的操作就能过拟合数据，因为它们使用简单的 if 语句递归地分割空间。 相比之下，其他模型可以更有效地估计 CDF 的总体形状，但在单个数据实例级别的精确定位上有障碍。 要看到这个，请再次看看图 [2] 。 该图表明，从大的视图看，CDF 函数看起来非常平滑和规则。 但是，如果放大单个值记录 (Value)，越来越多的非规律显示出来 ; 一个众所周知的统计效应。 许多数据集恰恰具有这种行为：从大局看，数据分布显得非常平滑，而越放大越难接近 CDF，由于个体层面上的“随机性”。 因此，像神经网络，多项式回归等模型可能需要更多的 CPU 和空间从整个数据集缩小到到数千项中选择单个（译注，这里也就是在上文中指的预测误差，min/max error），单个神经网络通常需要更多的空间和 CPU 时间为“最后一公里”减少从数千到数百的误差。
3. 典型的 ML 优化目标是最小化平均误差。 但是，对于索引，我们不仅需要猜测项目的最佳位置，而且还需要实际找到它，但前面讨论的最小和最大误差率更重要。
4. B- 树的缓存效率非常高，因为它们始终将顶层节点保存在缓存中，并在需要时访问其他页面。 但是，其他模型并不像缓存和操作高效。 例如，标准神经网络需要所有权重参数来做预测，这种预测需要大量的乘法和权重参数，必须从内存中读入到缓存 (译注，这里和上面的缓存都是指 CPU 高速缓存)。

3. RM 索引 (Recursive-model，递归模型)

为了克服挑战并探索模型作为索引替代或丰富的潜力，我们开发了学习索引框架 （LIF），递归模型索引（RMI）和基于标准误差的搜索策略。 我们主要关注简单的全连接神经网络，因为它们的简单性，但其他许多类型模型也是可能的。

3.1 学习索引框架（LIF）

LIF 可以看作是一个索引合成系统或者说索引工厂 ; 给定一个索引规范， LIF 生成不同的索引配置，优化它们并自动化测试它们。 虽然 LIF 可以即时学习简单模型（例如，线性回归模型），但它依赖于更复杂模型（例如 NN）的 Tensorflow。 但是，它从不使用 Tensorflow 进行推理。 相反，给定一个经过训练的 Tensorflow 模型， LIF 会自动从模型中提取所有权重，并根据模型规范在 C ++ 中生成高效的索引结构。 虽然使用 XLA 的 Tensorflow 已经支持代码编译，但其重点主要放在大规模计算上，其中模型的执行时间量级是微秒或毫秒。 相比之下，我们的代码生成专注于小型模型，因此必须消除 Tensorflow 管理大模型的开销。 这里我们利用来自 [ 21 ] 的想法，它已经展示了如何避免 Spark 运行时不必要的开销。 因此，我们能够执行 30 纳秒级的简单模型。

3.2 递归模型索引

如第 2.3 节所述，构建替代 B 树的替代学习模型的关键挑战之一是最后一英里的准确性。 例如，使用单个模型把误差从 100M 减少到几百这个数量级是非常困难的。 同时，将误差从 100M 降低到 10K，例如 100 * 100 = 10000 的精度实现起来还简单些，这样就可以用简单模型替换 B-Tree 的前两层。 同样的，将误差从 10k 降低到 100 是一个更简单的问题，因为模型只需要关注数据的一个子集。

基于这种观察并受到专家们的共同努力 [ 51] 的启发，我们提出了递归回归模型（参见图 [3]）。 也就是说，我们构建了一个模型层次