受 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 搜索时，我们遍历所有根树，并根据提供的 haystack 决定是继续在超平面的左侧还是右侧进行。优点是每个节点和向量都直接存储在内存映射文件中。

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

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

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

将 Annoy 适配到 LMDB

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

LMDB 使用 BTree 来排序条目，并且它比 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，我们将在本系列的第 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 上的社区或订阅 新闻通讯。您可以通过查看 路线图 并参与 产品讨论 来了解更多关于该产品的信息。