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

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

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

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

Annoy 如何生成树节点？

Spotify 的库 Annoy 将节点存储在磁盘上，包括项目节点（包含嵌入的节点）和其他节点，在本文中我们将这些节点称为树节点。这样做的好处是双重的

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

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

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

我们将其移植到 LMDB 时遇到的挑战

一个有趣的事实是，这是很长一段时间以来我第一次编写 C++ 代码，也是第一次请求 ChatGPT 为我编写代码，因为我对编写 C++ 没有信心，并且害怕陷入段错误。我需要一个小程序来从 stdin 反序列化嵌入，并将它们提供给 Annoy。聊天机器人的代码大部分是有效的，但它忽略了一个关键的 空向量检查，这导致了段错误...

将其移植到 LMDB 的主要障碍是，写入此键值存储是单线程的。它不支持并发写入操作，以保持 BTree 的正确平衡。有趣的事情来了！

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

自 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> {}

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

并行写入树节点

好的！现在我们知道如何存储指向项目的指针并在线程之间共享它们而无需任何用户端同步，我们需要使用它来生成树节点。我们已经知道如何在 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 向量，而没有任何 降维，则相当于 100M * 768 * 4 = 307B，换句话说，286 GiB 用于存储原始向量，没有任何内部树节点，即无法有效地搜索它们。

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

发现我们可以索引的向量的限制

Arroy 和 Annoy 都广泛使用内存映射，但方式略有不同。在 @dureuill 之前的文章中，我们看到操作系统没有无限的内存映射空间。因此，让我们深入研究性能结果。

当我使用 2500 万个 768 维向量运行 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 则相反，在内存映射文件中保留空间并直接写入其中。通过这样做，操作系统需要在内存映射区域中保留更多空间，并且被迫使缓存中的项目节点失效，以使刚刚写入的树节点在缓存中保持热状态，从而无缘无故地减慢系统速度，因为算法不需要树节点。

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

无共享原则获胜

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

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

您可以在 Lobste.rs、Hacker News、Rust Subreddit 或 X (以前的 Twitter) 上评论这篇文章。

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

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

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