多线程与内存映射：使用Arroy优化ANN性能

这是系列博客文章的第二部分，最初发布在我的个人博客上。从第一部分开始。

如果能向您展示一个单线程、内存映射的键值存储如何比手写的内存映射解决方案更高效，那岂不是很好？在将Spotify的基于磁盘的近似最近邻库移植到Rust，特别是移植到LMDB时，我遇到了问题。这些问题主要源于LMDB和内存安全。以下是这个故事。

提醒一下，Arroy是一个将嵌入（浮点向量）存储在磁盘上的库。在这些向量之上生成了一些数据结构，它们看起来像是受法线控制的树，用于递归地将向量数据集分割成子域。您可以在本系列的[第一部分](/blog/spotify-inspired-hybrid-search-and-rust/)中阅读更多相关内容。

Annoy如何生成树节点？

Annoy，这个Spotify库将节点存储在磁盘上，包括项目节点（包含嵌入的节点）以及我们本文中称之为树节点的其他节点。这样做的好处有两点：

• 由于所有节点都位于内存映射文件中，程序的内存使用量较低，并由操作系统进行优化。
• 概念很简单：通过ID访问节点。库将通过简单的乘法找到其在磁盘上的偏移量。

然而，这样做也有缺点。所有节点：包含嵌入的项目以及树节点必须具有相同的大小，并且如果用户想使用ID 231来标识其嵌入，库会将文件大小增加到ID乘以节点大小。例如，对于一个768维的向量，使用ID 231存储单个项目将生成一个超过6.5 TiB的文件。

在生成最终数据结构时，树节点都与包含嵌入的用户项目一起写入到同一个磁盘文件中。这些树构建过程并行运行，每棵树一个运行进程，因此在定义新创建的节点ID及其保留的内存区域时，需要大量的同步。最重要的是，当内存映射文件太小无法接受更多节点时，只有一个线程可以增大文件，因此在可变内存映射之上和node_id变量周围使用了互斥锁。树生成的一个有趣特性是，生成过程只需要原始的带有嵌入的用户项目。

移植到LMDB时遇到的挑战

一个有趣的事实是，这是我很久以来第一次编写C++代码，也是我第一次请ChatGPT为我编写代码，因为我对C++缺乏信心，害怕遇到段错误。我需要一个小程序来从标准输入反序列化嵌入并将其提供给Annoy。聊天机器人的代码大部分是有效的，但它遗漏了一个关键的空向量检查，导致了段错误……

移植到LMDB的主要障碍是，写入这个键值存储是单线程的。它不支持并发写入操作来维护一个正确平衡的B树。有意思的事情来了！

从不同线程读取用户项目节点

我们Meilisearch从一开始就使用了LMDB。它是一个维护良好的键值存储，被Firefox使用并由OpenLDAP开发者维护。它是内存映射的，并且不维护内存中的用户端条目缓存，而是为您提供一个指向磁盘上内存映射区域的指针。主要优点是，只要您的读取事务存在，您就可以一直保留指向该区域的指针。太棒了！因为它也是一个支持原子提交的事务性键值存储！

但是树的生成不需要引用已生成的节点，只需要用户项目。我们之前看到LMDB提供了指向内存映射文件的直接指针，而不维护任何可能失效的中间缓存。LMDB还有一个奇怪之处：您不能在多个线程之间共享只读事务，即RoTxn: !Sync；您也不能在线程之间移动写入事务，即RwTxn: !Send + !Sync。噢！而且无法对未提交的更改创建读取事务。这是一个问题，因为我们希望在存储项目的同时生成数据结构树。

但魔法无处不在，从下面这个小而有趣的数据结构开始。其原理是将带有嵌入的内部用户项目的指针保存在Vec<*const u8>中。感谢Rust，我们可以在编译时通过将生命周期保留在结构体中来确保这些指针的生命周期足够长。通过使用&'t RwTxn来获取&'t RoTxn，通过使用Deref也确保了我们无法在使用&'t mut RwTxn进行读取时修改数据库。根据LMDB主要开发者的说法，在线程之间共享这些指针是安全的，这也是我为该结构体实现Sync的原因。

pub struct ImmutableItems<'t, D> {
    item_ids: RoaringBitmap,
    constant_length: Option<usize>,
    offsets: Vec<*const u8>,
    _marker: marker::PhantomData<(&'t (), D)>,
}

impl<'t, D: Distance> ImmutableItems<'t, D> {
    pub fn new(rtxn: &'t RoTxn, database: Database<D>, index: u16) -> heed::Result<Self> {
        let mut item_ids = RoaringBitmap::new();
        let mut constant_length = None;
        let mut offsets = Vec::new();

        for result in database.items()? {
            let (item_id, bytes) = result?;
            assert_eq!(*constant_length.get_or_insert(bytes.len()), bytes.len());
            item_ids.push(item_id);
            offsets.push(bytes.as_ptr());
        }

        Ok(ImmutableLeafs { item_ids, constant_length, offsets, _marker: marker::PhantomData })
    }

    pub fn get(&self, item_id: ItemId) -> heed::Result<Option<Leaf<'t, D>>> {
        let len = match self.constant_length {
            Some(len) => len,
            None => return Ok(None),
        };
        let ptr = match self.item_ids.rank(item_id).checked_sub(1).and_then(|offset| self.offsets.get(offset)) {
            Some(ptr) => *ptr,
            None => return Ok(None),
        };
        let bytes = unsafe { slice::from_raw_parts(ptr, len) };
        ItemCodec::bytes_decode(bytes).map(Some)
    }
}

unsafe impl<D> Sync for ImmutableItems<'_, D> {}

您可能还注意到这个简化版数据结构中的其他一些有趣的优化。我们也知道用户项目节点具有恒定的长度，因此我决定只存储一次，将偏移向量的大小减半。考虑到我们的目标是存储1亿个向量，并且这个向量在内存中，我们将其大小从1526 MiB缩小到763 MiB，虽然不多，但总比没有好。

并行写入树节点

好的！既然我们知道了如何在不进行任何用户端同步的情况下存储指向项目并跨线程共享它们，那么我们现在需要使用它来生成树节点。我们Meilisearch已经知道如何处理LMDB，并决定实现相同的变通方法。为了并行写入LMDB，我们写入不同的、独立的文件，然后将所有内容合并。这利用了“无共享原则”并隔离了算法。与原始C++代码（查看.lock/unlock调用）相比，这大大减少了同步点数量，在我们的代码中只剩下一行：原子递增的全局树节点ID。

pub fn next_node_id(&self) -> u64 {
    self.0.fetch_add(1, Ordering::Relaxed)
}

我们基于用户项目节点生成普通分割节点的函数现在只是简单地将节点写入独立的文件。节点被追加到文件中，并且磁盘上的偏移量和边界存储在一个向量中，每个节点一个usize。使用Rayon，我们并行运行所有树生成函数，完成后，检索文件和边界，然后将唯一标识的节点按顺序写入LMDB。

性能比较

我们Meilisearch的目标是支持1亿个约768维的嵌入。如果我们将其存储为不进行任何降维的f32向量，则相当于1亿 * 768 * 4 = 307B，换句话说，需要286 GiB来存储原始向量，不包含任何内部树节点，即无法有效地搜索它们。

如果您不指定要生成的树的数量，算法将继续创建树，直到树节点的数量小于项目向量的数量。最终，树节点和项目的数量大致相同。

发现可索引向量的限制

Arroy和Annoy都大量使用内存映射，但方式略有不同。在@dureuill之前的一篇文章中，我们看到[操作系统没有无限的内存映射空间](/blog/dynamic-virtual-address-management/)。所以，让我们深入研究性能结果。

当我使用Arroy处理2500万个768维向量时，我注意到了一些奇怪的事情。CPU没有被充分利用，程序似乎大量读取SSD/NVMe，太多了🤔 树生成算法将向量空间分割成包含更少向量的子域。然而，它首先必须找到合适的法线来分割整个域，为此，它随机选择许多项目。不幸的是，我的机器只有63 GiB内存，而索引这2500万个项目需要超过72 Gib。Annoy也以同样的方式苦苦挣扎。

经过调查，我明白了为什么会达到整个交换空间和内存映射的限制。项目节点不仅是768 * 4字节，因为我们除了向量还存储了范数和其他内容，但在Arroy的情况下，LMDB需要维护围绕条目的B树数据结构，这些也占用了内存映射空间。两个程序都请求内存中不可用的随机项目节点，因此操作系统会从磁盘中获取它们，这需要时间。而且每个线程都在并行执行此操作。CPU只是在等待磁盘。

因此，经过几次二分法查找，我找到了Arroy成功使用所有内存而不受磁盘速度限制的点。它可以在63 GiB的机器上索引1562.5万个向量。您可以从htop截图上看到，磁盘读取速度为零，因为所有向量都适合RAM，Arroy正在将树节点写入磁盘，并且CPU正在尽力工作。它处理了不到七个小时。

但是……Annoy无法索引相同数量的文档。它面临我们之前看到的问题：高磁盘读取和低CPU使用率。但为什么呢？我需要澄清，因为节点格式相同，要索引的项目向量数量相同，算法也已移植。那么，两种解决方案之间有什么区别呢？

对于那些查看过C++代码库并可能受到之前描述的内存映射问题启发的人，您可能已经注意到这个细微的差异：Arroy将生成的树节点写入不同的原始文件，而Annoy则相反，它在内存映射文件中预留空间并直接写入其中。通过这种技巧，操作系统需要在内存映射区域中保留更多空间，并且被迫从缓存中清除项目节点，以使刚写入的树节点在缓存中保持“热”状态，这会无缘无故地降低系统速度，因为算法并不需要这些树节点。

因此，经过更多的二分法查找，以找出Annoy在63 GiB机器上的限制，我发现它大约能在五小时内索引1000万个向量。

无共享原则的胜利

因此，Arroy在同一台机器上可以索引多36%的向量，但这需要多长时间呢？并行写入同一个文件会不会比以单线程方式复制所有这些树节点更快？这将是一个短得多的段落，因为我只做了一些小型测试，但Arroy总是更快！

现在，您可能会告诉我，索引1亿个向量大约需要400 GiB的内存，您说得可能没错，但@irevoire将在未来的文章中讨论增量索引。我为了好玩做了一个线性回归。在12个线程和所需内存的情况下，我预测它需要1天22小时43分钟😬。噢！由于我们也在Meilisearch中使用这个向量存储，我们需要提供在搜索时过滤此数据结构的方法。这是本系列的下一篇文章。

您可以在Lobste.rs、Hacker News、Rust Subreddit或X（原Twitter）上评论这篇文章。

本系列的第三部分现已发布，深入探讨了Arroy的过滤磁盘ANN在搜索方面的进一步进展。

Meilisearch是一款开源搜索引擎，它不仅为最终用户提供最先进的体验，还提供简单直观的开发者体验。

要了解更多关于Meilisearch的信息，您可以加入Discord社区或订阅我们的新闻通讯。您可以通过查看产品路线图和参与产品讨论来了解更多关于该产品的信息。