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

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

向您展示一个单线程、内存映射的键值存储如何比手写的内存映射解决方案更高效，这难道不是一件很棒的事吗？在将 Spotify 基于磁盘的近似最近邻库移植到 Rust，更具体地说是移植到 LMDB 时，我遇到了问题。这些问题主要归因于 LMDB 和内存安全性。以下是这个故事。

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

Annoy 如何生成树节点？

Spotify 的 Annoy 库将节点存储在磁盘上，包括项节点（包含嵌入的节点）和本文中我们将称之为树节点。这样做有双重优势：

• 程序的内存使用率低，并由操作系统优化，因为所有节点都存在于内存映射文件中。
• 概念很简单：通过使用其 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 通过 Deref 获取 &'t RoTxn 也确保了我们无法在使用 &'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> {}

您可能也注意到此数据结构的简化版本中还有其他有趣的优化。我们还知道用户项节点的长度是常量，所以我决定只存储一次，从而将偏移量向量的大小减少了一半。考虑到我们的目标是存储 100M 个向量，并且此向量在内存中，我们将其大小从 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 向量，则相当于 100M * 768 * 4 = 307B，换句话说，存储原始向量需要 286 GiB，不包含任何内部树节点，即无法高效搜索它们。

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

探索可索引向量的限制

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

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

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

所以，经过一番二分法，我找到了 arroy 成功利用所有内存而没有受限于磁盘速度的点。它可以在 63 GiB 的机器上索引 15.625M。您可以在这个 htop 截图中看到磁盘读取速度为零，因为所有向量都适合 RAM，arroy 正在将树节点写入磁盘，并且 CPU 正在尽力工作。处理时间不到七小时。

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

对于那些查看过 C++ 代码库并可能被前面描述的内存映射问题暗示的人，您可能已经注意到这个细微的差异：Arroy 在 Annoy 相反的情况下将生成的树节点写入不同的原始文件，而 Annoy 则在内存映射文件中预留空间并直接写入。通过这种技巧，操作系统需要在内存映射区域中保留更多的空间，并且被迫从缓存中使项节点失效，以使刚写入的树节点在缓存中保持热度，这无缘无故地降低了系统速度，因为树节点对于算法来说不是必需的。

因此，在对 63 GiB 机器上的 Annoy 限制进行更多二分法探索后，我发现它可以在五小时内大致索引 1000 万个向量。

“无共享”原则的胜利

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

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

您可以在 Lobste.rs、Hacker News、Rust Subreddit 或 X（前 Twitter）上评论本文。

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

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

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