受 Spotify 启发：使用混合搜索和 Rust 提升 Meilisearch

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

统一关键字和语义搜索

在 Meilisearch，我们正在努力进行混合搜索，将广泛使用的关键字搜索算法与不太常见的语义搜索相结合。前者非常擅长精确匹配，而后者则能找到你更能描述的东西。

我将解释我们为什么要谈论Meilisearch/arroy，以及它对我们为何如此重要。语义搜索对我们来说是新的东西。原理很简单：文档与向量（浮点数列表）相关联，并且它们彼此越接近，在语义上就越接近。当用户查询引擎时，你会计算查询的嵌入（向量），并进行近似最近邻搜索，以获得与用户查询第 n 个最接近的项。有趣的事实：经过这么长时间，Annoy仍然是顶尖的竞争者之一，所以想象一下我们最终发布 Arroy 时的情景。

看起来很简单，对吧？你真的认为这一系列博客文章会谈论一个非常简单的主题吗？加入我们的旅程，@irevoire 和我 (@Kerollmops)，在 @dureuill 的协助下，将这个前沿库移植到 Rust。

优化高维嵌入的存储和搜索

嵌入的维度介于 768 和 1536 之间。Meilisearch 的客户希望存储超过 1 亿个嵌入。任何降维算法都没有修改过的嵌入使用 32 位浮点数。这意味着存储这些数据将占用 286 GiB 到 572 GiB 之间的空间，具体取决于维度。

是的，它可以放入 RAM 中，但代价是什么？存储在磁盘上便宜得多。嘿！这只是原始向量。我向你保证，在 O(n) 中对所有向量进行近似邻居搜索会非常慢。因此，我们决定将它们存储在磁盘上。在选择 Spotify/Annoy 之前，我们研究了许多不同的解决方案。最初，我们使用了 instant-distance Rust crate 中的 HNSW，这是另一个数据结构，用于进行 ANNs（近似最近邻）搜索。但是，它不是基于磁盘的；一切都在内存中，这不实用。

Spotify 的超平面树用于高效 ANNs

Spotify 开发了一个很棒的 C++ 库，用于在大量的向量集中进行搜索。该算法会生成多个引用向量的树。树节点是随机超平面，将向量子集一分为二。根节点将整个向量集一分为二，左分支和右分支再次递归地分割向量的子集，直到子集足够小。当我们执行 ANNs 搜索时，我们会遍历所有根树，并根据提供的干草堆决定是继续超平面的左侧还是右侧。优点是每个节点和向量都直接存储在内存映射文件中。

Annoy 和 Arroy 生成超平面以分割向量空间的方式

在树的末尾，我们可以看到后代的节点。这些定义了适合上述递归拆分节点这一侧的叶向量的最终列表。这是用户提供的无符号整数列表。Annoy 将它们表示为 u32 的切片，但我们决定使用 RoaringBitmaps 来减小它们的大小。Annoy 无法压缩它们，因为它们有一个有趣的约束：所有节点、用户叶向量、拆分节点和后代的大小必须相同，以便可以使用磁盘上的偏移量访问。

上面的图像显示了拆分平面可以表示的方式。一个超平面在这里以二维方式分割节点子集，但想象一下有 768 到 1536 个维度的情况。每个超平面都会创建两个点子集，这些点子集通过另一个超平面递归拆分。一旦要拆分的点数足够小，我们就创建一个后代节点，其中包含与这些点对应的项 ID。此外，点的嵌入永远不会在内存中重复；我们通过它们的 ID 引用它们。

将 Annoy 适配到 LMDB

那么，如果它运行得这么好，我们为什么要将它移植到 Rust 呢？首先，因为我遵循用 Rust 重写的教条 😛，其次是因为这是一个带有许多可怕的黑客和未定义行为的 C++ 库，第三是因为 Meilisearch 基于 LMDB，一个原子化的、基于事务的、内存映射的键值存储。此外，自 2015 年以来，他们就想使用 LMDB，但他们尚未实现；这需要大量时间来相应地更改数据结构。

LMDB 使用 BTree 对条目进行排序，并且它为中间数据结构占用的空间比 Annoy 多，Annoy 直接使用文件中的偏移量来识别向量。使用此键值存储的另一个优点是管理增量更新。插入和删除向量更容易。假设你在 Annoy 中插入一个由高 32 位整数标识的向量。在这种情况下，它将分配大量内存来将其存储为专用偏移量，并在必要时增加文件大小。

在 Meilisearch 和 Arroy 中，我们使用 heed，一个类型化的 LMDB Rust 包装器，以避免未定义的行为、错误以及与 C/C++ 库相关的东西。因此，我们广泛使用 &mut 和 & 来确保我们不会修改键值存储条目，同时保留其中的引用，并确保读写事务不能在线程之间引用或发送。但这将是另一个故事。

我们最初认为使用此键值存储会使 Arroy 比 Annoy 慢。但是，LMDB 在将页面写入磁盘之前先将其写入内存，这似乎比直接写入可变内存映射区域快得多。另一方面，LMDB 不保证 Annoy 的文件格式允许的值的内存对齐；我们稍后会讨论这个问题。

使用 SIMD 优化向量处理

Arroy 应该做的最占用 CPU 的任务是尝试将向量云一分为二。这需要在热循环中计算大量的距离。但是我们从内存映射的磁盘读取嵌入。会出什么问题呢？

fn align_floats(unaligned_bytes: &[u8]) -> Vec<f32> {
    let count = unaligned_bytes.len() / mem::size_of::<f32>();
    let mut floats = vec![f32::NAN; count];
    cast_slice_mut(&mut floats).copy_from_slice(unaligned_bytes);
    floats
}

我们在 macOS 上使用 Instrument 来分析我们的程序，发现大量时间都花在了移动东西上，即 _platform_memmove。原因是，从磁盘读取未对齐的浮点数是未定义的行为，因此我们将浮点数复制到对齐的内存区域，如上所示。成本：通过读取进行分配，再加上调用 memmove。

分析显示，为了对齐字节，需要进行大量的内存复制

在将距离函数从 C++ 移植到 Rust 时，我们直接使用了 Qdrant 经过 SIMD 优化的距离函数，没有事先进行修改。尽管我们意识到读取未对齐内存的性能成本，但我们决定在未对齐的浮点数上执行距离函数，并确保不将其表示为 &[f32]，因为这会是未定义行为。这些函数接受原始字节切片，并使用指针和 SIMD 函数来处理它们。我们解锁了性能。距离函数的速度变慢了，但这弥补了 memmove 和分配的开销！

直接处理未对齐的内存消除了复制操作

与 memmove 调用类似，我们可以看到 _platform_memcmp 函数在这里占用了大量空间。原因是 LMDB 使用此标准函数来比较键字节，并确定树中一个键的字典顺序位于另一个键之上还是之下。每当我们读取或写入 LMDB 时都会用到它。@irevoire 大大减小了键的大小，并看到了性能的显著提升。我们进一步尝试使用 MDB_INTEGERKEY，这使得 LMDB 使用计算机的字节序来比较内存，但这使用起来很复杂，并且没有显示出显著的性能提升。

即将到来的挑战

在将这个很酷的算法移植到 Rust 和 LMDB 时，我们缺少最重要的功能之一：多线程树构建。这个功能缺失的主要原因是 LMDB 不支持并发写入磁盘。这是这个库的优点之一；写入是确定性的。我们已经非常了解 LMDB，我们将在本系列的第二部分中解释我们如何利用 LMDB 的强大功能以及我们如何击败 Spotify 库。

除了要达到当前 Annoy 的功能之外，我们还需要更多 Meilisearch 的功能。我们在 Arroy 中实现了微软的 Filtered-DiskANN 功能。通过提供我们想要检索的项目 ID 子集，我们可以避免搜索整个树来获取近似最近邻。我们将在即将发布的文章中讨论这一点。

在 Meilisearch v1.6 中，我们优化了仅更新文档部分（如投票数或浏览量）的性能。所描述的 Arroy 单线程版本会为每个向量调整重新构建树节点。@irevoire 将在他的下一篇文章中探讨 Arroy 的增量索引，这是 Annoy 不具备的功能。

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

本系列的[第二部分](/blog/refining-ann-performance-with-rust/)和[第三部分](/blog/arroy-filtered-disk-ann/)现已发布，探讨了使用 Rust 实现 ANN 算法的进一步进展。

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

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