从威士忌到波尔本,LSM-Tree 上如何用机器学习加速索引

Posted: December 08, 2020

本文是发表在 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)log(n) 二分查找,现在就只要引入 k 和 b 两个变量 O(1)O(1) 在线上求值即可。在数据不改变的情况下,可以让机器学习算法过拟合数据,从而达到准确的预测,或者给出一个具有可控误差范围的预测(如预测位置为 pospos,可以保证 key 如果存在,一定在 [posδmin,pos+δmax][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

读数据流程示意图,来自 Bourbon 论文截图

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

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

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

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

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

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

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

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

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

感想

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

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

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

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