200 行 Rust 代码编写一个向量搜索库代码已开源！

　　由于 AI 与机器学习的快速发展，如今向量数据库已无处不在。虽然向量搜索能够支持复杂的 AI 或机器学习应用程序，但这个概念本身并不难。在文本中，我们来一起看看向量数据库的工作原理，并使用不到 200 行 Rust 代码构建一个简单的向量搜索库。

　　本文使用的方法基于流行库 annoy 中使用的一种名叫Locality Sensitive Hashing（局部敏感哈希，LSH）的系列算法。本文的目的不是介绍一种新奇的算法或库，而是分享如何编写一个向量搜索。

　　我们很难使用传统的数据库表示和查询文档、图像、视频等复杂的非结构化数据，尤其是无法查找相似数据。那么，YouTube 又是如何选择下一个播放视频的呢？Spotify 又是如何根据当前歌曲自定义播放列表的呢？

　　2010 年代 AI 的进步（从 Word2Vec 和 GloVe 开始）使我们能够构建这些对象的语义表示，即使用笛卡尔空间中的点表示这些对象。假设一个视频被映射到点 [0.1, -5.1, 7.55]，另一个被映射到点 [5.3, -0.1, 2.7]。这些机器学习算法的神奇之处在于，这些点的选择保留了语义信息：两个视频的相似度越高，它们的向量之间的距离越小。

　　请注意，这些向量（或更专业的叫法是嵌入）不一定是三维的，它们可以、而且通常确实位于更高维的空间中（比如 128 维或 750 维）。而且距离也不一定是欧几里得距离，其他形式的距离也可以，比如点积。无论采用哪种形式，重要的是它们之间的距离对应于相似性。

　　下面，假设我们可以访问所有 YouTube 视频的此类向量。那么我们如何找到与某个视频相似度最高的其他视频呢？很简单。遍历所有视频，计算它们之间的距离，并选择距离最小的视频，这个过程又称之为查找视频的最近邻居。这种方法确实有效，除了一个问题，线性O(N)搜索的成本过高。所以，我们需要一种更快的次线性方法来查找视频的最近邻居。这通常是不可能的，必须付出一些代价。

　　然而，在实际情况下，我们并不需要找到最近的视频，只要足够接近就可以了。这就是近似最近邻（Approximate Nearest Neighbor，ANN）算法，又称为向量搜索。我们的目标是通过次线性方法找到空间中任何足够近的最近邻。那么，实际如何实现呢？

　　向量搜索算法背后的基本思想都是相同的：通过一些预处理来识别彼此足够接近的点（有点类似于建立索引）。查询时，使用这个索引来排除大部分点。只留下少量点，然后再进行线性扫描。

　　然而，这个简单的想法有很多可以实现的方法。目前有几种最先进的向量搜索算法，例如 HNSW（Hierarchical Navigable Small World，分层的可导航小世界）是一种图，它连接了互相接近的顶点，而且还保存了距离固定入口点的长距离边。此外还有一些开源项目，例如 Facebook 的 FAISS，以及一些高可用向量数据库的 PaaS 产品，例如 Pinecone 和 Weaviate。

　　在本文中，假设有N个指定的点，我们要在这些点上构建一个简化的向量搜索索引，步骤如下：

　　构建一个通过 C 并垂直于连接 A 和 B 的线段的超平面（即更高维度的线）。

　　分别对两组向量做以下处理：如果组的大小大于可配置参数最大节点数，则在该组之上递归调用此过程，构建子树。否则，使用所有向量（或它们的唯一 ID）构建一个叶节点。

　　我们将使用这个随机过程来构建一棵树，其中的每个节点都是一个超平面定义，左边的子树是超平面下方的所有向量，右边的子树是超平面上方的所有向量。向量集不断递归拆分，直到叶节点包含的向量不超过最大节点数。下图是包含5个点的示例：

　　我们随机选择向量 A1=(4,2)和 B1=(5,7)，二者的中点是 (4.5,4.5)，我们构建一条线，要求通过中点且垂直于线)。这条线（图中蓝色的线）给了我们两组向量：一组包含 2 个向量，另一组包含 4 个。假设最大节点数设置为 2。那么，第一组向量不需要进一步拆分，而第二组向量需要进一步递归——如图所示，我们又构建了一个红色超平面。对于大型数据集，如此重复拆分可以将超空间拆分为几个不同的区域，如下所示。

　　此处的每个区域代表一个叶节点，而且感觉足够接近的点很可能最终会出现在同一个叶节点中。接下来，给定一个查询点，我们可以在对数时间内自上而下遍历树，找到它所属的叶子，然后对叶节点中的所有（数量不多）点进行线性扫描。很明显，这个算法并不是万无一失的，即便是距离非常近的两个点也完全有可能被一个超平面分开，最终导致二者彼此相距很远。但是，这个问题可以通过构建多棵独立的树来解决，这样，如果两个点之间的距离足够近，它们出现在某些树同一个叶节点中的概率就很高。在查询时，我们自上而下遍历所有树，定位相关的叶节点，然后求所有叶节点的所有候选节点的并集，并对所有节点进行线性扫描。

　　首先，我们为 Rust 的 Vector 类型定义一些工具函数，包括求点积、均值、哈希值以及平方 L2 距离。感谢 Rust 的优雅的类型系统，我们可以传递泛型类型参数 N，以在编译时强制索引中的所有向量具有相同的维度。

　　接下来，我们来生成随机超平面，构建最近邻树的森林。我们应该如何表示树中的点呢？

　　我们可以直接将 D 维向量存储在叶节点中。但是，如果这个维度（D）非常大，那么会导致内存碎片化（严重影响到性能），而且当多棵树引用同一个向量时，还会造成森林重复保存。因此，我们将向量存储在全局可访问的连续位置中，并在叶节点中保存类型为 usize 的索引（在 64 位系统上仅占用 8 字节，相反如果直接保存四维向量，每个维度的 f32 类型就会占用 4 字节）。以下是用于表示树内部节点和叶节点的数据类型。

　　我们从索引中随机采样两个，分别对应于向量 A 和 B，计算 n = A - B，并找到 A 和 B 的中点（point_on_plane）。存储超平面使用的结构包含系数（向量 n）和常量（n 和 point_on_plane 的点积），形式为：n(x-x0) = nx - nx0。我们可以求向量和 n 之间的点积，然后减去常数，就可以将向量放在超平面的上方或下方。由于树中的内部节点包含超平面定义，叶节点包含向量 ID，因此我们可以使用 ADT 对树进行类型检查：

　　请注意，由于该算法不允许重复，构建两点之间的超平面要求这两个点是唯一的，因此在建立索引之前，我们必须去除向量集中的重复项。

　　索引建立之后，下一步我们如何使用它搜索某棵树中输入向量的 K 个近似最近邻？我们的超平面存储在非叶节点处，因此我们可以从树的根开始搜索：这个向量是在这个超平面的上方还是下方？我们可以用点积计算，复杂度为 O(D)。接下来，我们可以根据响应，递归搜索左子树或右子树，直到找到叶节点。请记住，叶节点中存储的向量数最大为最大节点数，这些向量位于输入向量的近似邻域中（它们将落入所有超平面下超空间的同一分区中）。如果这个叶节点的向量索引数超过 K，我们就需要根据到输入向量的 L2 距离，对所有这些向量进行排序，并返回最接近的 K 个向量。

　　假设我们的索引最后得到了一棵平衡树，对于维度 D、向量数量 N 和最大节点数 M N，搜索耗费的时间为 O(Dlog(N) + DM + Mlog(M))，最坏的情况下我们需要比较 log(N) 个超平面（即树的高度）才能找到叶节点，每次比较耗费的时间为 O(D) 个点积，计算一个叶节点的所有候选向量的 L2 距离需要 O(DM)，最后在 O(Mlog(M)) 时间内对它们进行排序，并返回 前 K 个。

　　但是，如果我们找到的叶节点少于 K 个向量，会发生什么情况？如果最大节点数太小，或超平面分割不均匀，则会导致某个子树中的向量非常少。为了解决这个问题，我们可以在树搜索中添加一些简单的回溯功能。例如，如果返回的候选向量数量不足，我们可以在内部节点再递归调用一次，访问另一个分枝。代码如下：

　　请注意，为了进一步优化递归调用，我们还可以保存子树中的向量总数，并在每个内部节点中保存指向所有叶节点的指针列表，但为了简单起见，此处略过。

　　将此搜索扩展到一片森林非常简单，只需从所有树木中单独收集前 K 个候选者，按距离对它们进行排序，然后返回总体的前 K 个匹配项。请注意，随着树的数量增加，内存的使用和搜索时间都会呈线性增长，但我们可以获得更好的更近邻居，因为我们收集了来自不同树的候选者。

　　其实，这个实现非常简单，以至于我们一度怀疑相较于最先进的算法，这个算法非常拙略。因此，我们做了一些基准测试来证实我们的怀疑。

　　算法可以通过延迟和质量来评估。质量通常通过召回率来衡量：求近似最近邻搜索返回结果与线性搜索返回结果的百分比。严格来说，有时返回的结果并不在前 K，但非常接近 K，因此也没关系，为了量化这一点，我们可以计算邻居的平均欧几里德距离，并将其与平均距离进行比较。

　　所有的基准测试都是在搭载了 2.3 GHz 四核英特尔酷睿 i5 处理器的机器上运行的，使用了 999,994 条维基百科数据 FastText 嵌入，300 个维度。我们设置返回前 K 个查询结果。

　　两种近似最近邻方法快了三个数量级。天选看起来在这个微型基准测试中，我们的 Rust 实现比 HNSW 快 3 倍。

　　我们使用的这个语料库包含 999,994 个单词，我们可视化了在 HNSW 和我们的自定义 Rust 索引下，每个单词与其前 K=20 个近似邻居的平均欧几里得距离的分布：

　　通过上面的数字可以看出，虽然我们的算法在质量上无法与尖端技术相媲美，但它的速度非常快。为什么会这样？

　　老实说，在构建这个算法时，我们兴奋得昏了头脑，性能优化也只是为了好玩。下面是一些重要的优化：

　　将去除文档的重复数据的过程转移到索引冷路径上。通过索引（而不是浮点数组）引用向量可以大幅提高搜索速度，因为跨树查找唯一候选者只需要对 8 字节索引进行散列（而不是 300 维 f32 数据）。

　　在将点添加到全局候选列表之前，散列并查找唯一向量，通过递归搜索调用中的可变引用传递数据，以避免栈帧之间和栈帧内的复制。

　　将 N 作为通用类型参数传递进去，这样类型检查就会将 300 维的数据当作长度为 300 的 f32 数组，这不仅可以改善缓存局部性，还可以减少内存占用。

　　我们怀疑内部操作（例如点积）被 Rust 编译器向量化了，但我们没有检查。

　　当搜索涉及多个树时应当使用并行。不要顺序收集候选者，而是应该并行进行，因为每个树都会访问完全不同的内存，因此每个树都可以在各自的线程上单独运行，候选者通过通道发送到主线程。线程可以在创建索引时建立，并使用模拟搜索进行预热（将树加载到缓存），从而减小搜索的额外开销。这样搜索就不会根据树的数量线性增长了。

　　大型树可能无法完整地加载到内存中，可能需要从磁盘中读取，所以可以将特定的子图放到磁盘上，并精心设计算法，保证搜索正常工作的同时，尽可能减小文件I/O。

　　继续深入，如果树无法保存在一个实例的磁盘上，可以将子树分布到多个实例中，并让递归搜索在本地数据不存在时发出RPC请求。

　　树包含许多内存重定向（基于指针的树对L1缓存不友好）。平衡树可以写成数组，但我们的树只是接衡，其超平面是随机的。是否可能采用新的数据结构？

　　上述问题的解决方案，对于新数据即时编制索引也是成立的（可能需要对大型树进行分片）。如果特定索引序列会导致非常不平衡的树，那么是否应该重建树？天选团队