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++ 库相关的未定义行为、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，我们将在本系列的第 2 部分中解释我们如何利用 LMDB 的强大功能以及我们如何击败 Spotify 库。

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

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

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

本系列的[第2部分](/blog/refining-ann-performance-with-rust/ 和 [第3部分](/blog/arroy-filtered-disk-ann/ 现已推出，探索使用 Rust 进行 ANN 算法的进一步改进。

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

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