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

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

如果我能向你展示，一个单线程、内存映射的键值存储如何比手写的内存映射解决方案更有效，那不是很好吗？我在将 Spotify 基于磁盘的近似最近邻库移植到 Rust，更具体地说是 LMDB 时，遇到了问题。这些问题主要是由 LMDB 和内存安全引起的。下面是具体情况。

提醒你一下，Arroy 是一个将嵌入（浮点数向量）存储在磁盘上的库。一些数据结构在这些向量之上生成，看起来像树，由法线控制，用于将向量数据集递归地分割成子域。但你可以在本系列的[第 1 部分](/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++ 没有信心，并且担心出现一些段错误。我需要一个小程序来从 stdin 反序列化嵌入并将它们提供给 Annoy。聊天机器人的代码在很大程度上是有效的，但它忽略了一个关键的 空向量检查，这导致了段错误...

将其移植到 LMDB 的主要障碍是，写入此键值存储是单线程的。它不支持并发写入操作来维护正确平衡的 B 树。接下来的事情会很有趣！

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

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

但是树生成不需要引用生成的节点，而只需要用户项。我们之前看到，LMDB 直接给出指向内存映射文件的指针，而不维护任何中间缓存（可能会失效）。LMDB 还有另一个怪癖：你不能在多个线程之间共享只读事务，即 RoTxn: !Sync，并且你不能在线程之间移动写入事务，即 RwTxn: !Send + !Sync。哎！而且没有办法在未提交的更改上创建读取事务。这是一个问题，因为我们希望在我们存储项的同一个事务中生成数据结构树。

但是魔法无处不在，从以下小型有趣的数据结构开始。原理是使用 Vec<*const u8> 来保留指向内部用户项（包含嵌入）的指针。借助 Rust，我们可以通过在结构体中保留生命周期，在编译时确保指针的生命周期足够长。通过使用 Deref 使用 &'t RwTxn 获取 &'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> {}

你可能还注意到此数据结构的简化版本中还有一些其他的有趣优化。我们还知道用户项节点具有恒定的长度，因此我决定只存储一次，从而将偏移向量的大小缩小了两倍。考虑到我们的目标是存储 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 的目标是支持大约 768 维的 1 亿个嵌入向量。如果我们以 f32 向量的形式存储它们，而不进行任何降维，则相当于 1 亿 * 768 * 4 = 3070 亿字节，换句话说，存储原始向量需要 286 GiB，这还不包括任何内部树节点，即无法高效地搜索它们。

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

探索我们可以索引的向量的极限

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

当我使用 768 维的 2500 万个向量运行 Arroy 时，我注意到了一些奇怪的事情。CPU 没有被使用，并且程序似乎在大量读取 SSD/NVMe，有点太多了🤔。树生成算法正在将向量空间划分为具有较少向量的子域。但是，它必须首先找到合适的法线来分割整个域，为此，它会随机选择许多项目。不幸的是，我的机器只有 63 GiB 的内存，而索引这 2500 万个项目需要超过 72 GiB。Annoy 也遇到了同样的问题。

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

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

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

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

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

无共享原则获胜

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

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

您可以在 Lobste.rs、Hacker News、Rust 子版块 或 X（以前的 Twitter） 上评论这篇文章。

本系列的第 3 部分现已发布，探讨了我们使用 Arroy 的过滤磁盘 ANN 进行搜索的进一步进展。

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

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