本文是发表在 OSDI2020 上的 From WiscKey to Bourbon: A Learned Index for Log-Structured Merge Trees 一文的阅读笔记。这篇文章出自 University of Wisconsin-Madison 做存储闻名的 Remzi 组（他也是经典操作系统教材 OSTEP 的作者），传承了该组优秀的文笔和扎实的实验功底。

基于机器学习的习得索引（Learned Index）算法自 17 年 The Case for Learned Index Structures 介绍到数据库领域以来，吸引了大量的研究目光，但苦于在更新面前维护成本过高，一直没有在生产中应用的可能性。而此文巧妙地利用了 LSM-Tree 中 SST 数据不变（immutable）的特性，使得习得索引有了用武之地，在 WiscKey 系统上可以获得 1.23x-1.78x 的性能提升，令人瞩目。

Learned Index

简单地说，习得索引是用来预测数据所在位置的。

何谓索引？无论是 B 树还是 LSM-Tree 上基于二分查找的索引，都是给定一个 key，然后查找对应的值所在存储介质中的位置，可以认为传统的索引是一个准确的映射。而习得索引则可以通过机器学习算法来给出一个预测。

一句话说明习得索引的优势：在数据有序的前提下，用值为 x 轴，位置为 y 轴，用简单的线性模型去拟合一条线，那么本来需要 $log(n)$ 二分查找，现在就只要引入 k 和 b 两个变量 $O(1)$ 在线上求值即可。在数据不改变的情况下，可以让机器学习算法过拟合数据，从而达到准确的预测，或者给出一个具有可控误差范围的预测（如预测位置为 $pos$，可以保证 key 如果存在，一定在 $[pos-\delta_{min},pos+\delta_{max}]$），在读取的时候做小范围的局部搜索。不过如何增量更新习得索引还是个未解之谜，之前提出习得索引的那篇论文把增量更新的算法留给了未来的研究。

LSM-Tree, LevelDB and WiscKey

LSM-Tree 是非常著名的键值存储结构（KV Storage）设计，而 LevelDB 是 Google 开源的一个实现。如果对 LSM-Tree 不太了解，建议阅读 WiscKey: Separating Keys from Values in SSD-conscious Storage 这篇论文的 2.1 和 2.2 小节——这也顺便引入了 WiscKey。

简单地说，LSM-Tree 的思想是，因为在硬盘上顺序读写要远快于随机读写，因此对于数据的更新就只 append，然后通过周期性的 compaction 来形成有序的结构（通常称一个这样的结构为 SST，Sorting String Table），通过遍历所有 SST，然后在 SST 内部二分来进行查找；而 LevelDB 则是通过对 compaction 结果进行分层，平衡了 compaction 和查询的代价。值得一提的是，工业界目前用的比较多的是 RocksDB，代码虽然基于 LevelDB，但 Facebook 的工程师们在上面加了非常多的优化，也在大量业界产品中受到了严格的验证和打磨。

LSM-Tree 的结构在硬盘存储上有着非常好的性能，但却面临着较为严重的写放大问题（Write Amplification），因为每一次 compaction 都需要把键值重新写一次。除了时间和空间上的开销外，对于 SSD 这种存在写疲劳（wearing）的存储介质，写放大势必会大幅缩短 SSD 的生命。

在 LevelDB 类似的实现上，WiscKey 观察到，我们实际上只需要维护键的索引即可，每次 compaction 只需要修改键，而值则可以单独存储到值日志（Value Log）中。从而 LSM-Tree 维护的就是从键到对应的值在值日志中的位置。对于值较大的情况，可以大大减少 compaction 的读写量，而读取和写入却不会有任何额外开销，甚至可以跳过将值写入 WAL，直接写入值日志，进一步减少写入量。

读取

为了引入习得索引的优化，这里再介绍一下 LevelDB/WiscKey 的读取流程。

首先，如果 key 在 memtable 中，那么直接读取返回；否则自上而下地读 SST

1. FindFiles，找 SST 文件，如果在 L0，那么每个文件都得读，因为 L0 不保证 SST 的键区间不相交；如果在更深的层，那么键区间保证不相交，每层只需要读一个 SST 文件即可。L1 开始，每层可以在内存中维护一个 SST 的有序区间索引，在这个索引上二分即可。
2. LoadIB + FB。IB 和 FB 分别是 index block 和 filter block 的缩写，index block 是 SST 内部切分的 block 的 index，用于找对应的 block；filter block 则是一个布隆过滤器（bloom filter），可以快速排除 key 不在的情况，因此我们首先加载这两个结构。
3. SearchIB。二分查 index block 找到对应的 block
4. SearchFB。用布隆过滤器过滤，如果没有滤掉
5. LoadDB。则把这个 block 加载到内存
6. SearchDB。在这个 block 中继续二分查找
7. ReadValue。找到 key 后读数据，如果考虑 WiscKey 键值分离的情况，需要去 value log 中读。

在这个读取流程中，我们可以有一个观察（论文 3.1 小节）。如下图所示，最左边是作为参照的全内存读取的理想情况，而从硬盘、NVMe（SSD）再到 Optane（NVM），数据读取速度变快，但查找索引（即第一二三步的 FindeFiles、SearchIB+SearchDB 和 Search FB）的时间不变，因此在整个读取流程中的占比会越来越大。随着存储介质访问速度的提升，优化索引读取是个收益增加的事情。

Bourbon：利用习得索引进行优化

LSM-Tree 是为了优化写入而产生的存储结构，留给优化的空间已经很局限了，但习得索引可以加速读，而且 compaction 产生的 SST 是只读的，因此可以不考虑更新地在这上面直接构建习得索引，这两者实在是天造地设。

Bourbon 采用了 piece-wise linear regression 作为学习算法，在 WiscKey 结构上优化键的查找。一个 WiscKey 上的 SST 的结构长这样，$k_1,k_2,k_3,\cdots,k_n$，其中 $k_i$ 由键 $k$ 和值日志中的位置 $p$ 拼合而成，满足 $k_1<k_2<\cdots<k_n$。查找则是给出一个键，返回这个键在 SST 中的位置下标——然后直接读取这个位置，就可以知道值在值日志中的位置，从而完成读取。

因为 $k$ 和 $p$ 的长度都可控，因此随着 $k_i$ 的递增，其位置也规律性地增加，这对拟合线性函数非常友好。简单地说，我们首先确定一个误差范围，然后维护一个连续的键集合，保证在这个集合去做线性回归来拟合键的位置不会超过误差范围。我们顺序考虑每一个键，如果这个键地加入会破坏误差，则放到一个新的集合中并结束对当前集合的维护，否则加入当前集合。这样我们就可以把整个 SST 分成多个线性函数组成的小段，每个段内可以 $O(1)$ 地给定键，返回一个位置区间，在这个区间中做局部搜索即可读到想要的键，而查找小段则可以直接二分，或者继续构建线性回归函数来查找。这就是为什么该学习算法被称为 piece-wise 的原因——因为拟合出来的结果是分段的线性函数。

相比于直接二分查找，习得索引只需要维护非常小的额外空间（即线性函数的两个参数），就可以把 $O(\log n)$ 优化成 $O(1)$。

通常情况下，小段数量并不会特别的多，如果真的特别多——我们可以回归二分查找，放弃在这个 SST 上使用习得索引。这也就引出 Bourbon 中最重要的话题——如何在是否使用习得索引上做权衡（trade off）。

对于在层次较高的 SST 来说，变动可能比较频繁，因此如果要采用习得索引，那么每次 compaction 过后都需要重新训练，这个代价会比较高；而对于底层的 SST，随着层数增加，变动的频率会降低，那么采用习得索引的收益就会超过训练代价，但在此基础上也得考虑到读取频率的降低。

为了达到权衡的目的，Bourbon 设计了实验观察在不同工作负载下、不同层上 SST 的生命周期和访问频率，设计了一个判断是否运用习得索引的模型。鉴于本文已经足够长，也完成了算法主体的介绍，感兴趣的读者可以阅读原论文了解细节。

感想

我以前对在数据库领域使用机器学习的想法是不太感冒的，但这一篇文章真正让我看到了其可能性和价值。原文第一作者曾到我们公司进行过分享，而现在他正和我们 CEO 一起组队参加 TiDB Hackathon，队伍名称就叫 Bourbon，同样作为参赛选手的我表示瑟瑟发抖。

我其实很早就注意到这篇文章了，但写这篇阅读笔记却拖延了很久，主要是对文章后面实验设计与评测那段（trade off 的设计），不知道该如何写笔记。对于科研工作来说，其实最重要的部分就是这里了，我虽有能力理解（并表示赞同和佩服）但实在无法狗尾续貂，因此只能留空给原文——我想有能力和兴趣深挖的读者必然会去做这件事的。

最后，Remzi 这个组是真的有趣。WiscKey 一文也是他们提出来的，这个名字是 Wisconsin Key 的缩写也是威士忌的谐音；而基于 WiscKey 的 Bourbon 一词则是一种美国著名威士忌品类

这可能就是顶尖计算机科学家兼老酒鬼的恶趣味吧！