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

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

结合关键词和语义搜索

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

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

听起来很简单，对吧？你**真的**认为这个系列博客文章会讨论一个死简单的主题吗？加入我们，与@irevoire和我（@Kerollmops）一起，在@dureuill的协助下，将这个尖端库移植到Rust。

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

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

是的，它能放入内存，但代价是什么？存储在磁盘上便宜得多。而且！这还只是原始向量。我向你保证，在所有向量中进行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使用B树来排序条目，它比Annoy占用更多的中间数据结构空间，Annoy直接使用文件中的偏移量来标识向量。使用这种键值存储的另一个优点是管理增量更新。插入和删除向量更容易。如果你在Annoy中插入一个由高32位整数标识的向量，它将分配大量内存来存储它作为专用偏移量，并在必要时增加文件大小。

在 Meilisearch 和 Arroy 中，我们使用 heed，一个类型化的 LMDB Rust 封装器，以避免与 C/C++ 库相关的未定义行为、bug 和其他问题。因此，我们大量使用&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）上评论此文章。

本系列的第二部分和第三部分现已发布，进一步探讨了Rust中ANN算法的进步。

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

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