Meilisearch 如何在不到一分钟内更新包含数百万向量嵌入的数据库
我们如何在向量存储中实现增量索引。
这是博客文章系列的第 4 部分 最初发布在 Clément Renault 的博客上。从 第 1 部分、 第 2 部分和 第 3 部分开始旅程。这篇博文假设您已经阅读了第 1 部分和第 2 部分。
在这篇博文中,我们将讨论如何在 Arroy 中实现增量索引。增量是指在树中插入或删除项目时,我们可以仅更新所需的部分,而不是重建所有内容。
在 Meilisearch,这一点至关重要,因为我们的用户通常会定期发送更新,并且他们的内容不断变化。从头开始重建所有内容成本非常高,而且不会带来良好的体验;这就是为什么我们已经在最新版本中花费了大量时间来使所有 数据结构支持增量索引,而 arroy 就是其中之一,我们也必须实现它!
理论
下图显示了在现有树中添加和删除文档时应发生的情况的高级视图。
提醒一下,arroy 是一个基于 LMDB 的向量存储,LMDB 是一个嵌入式事务性内存映射键值存储。它将您的向量存储在由三种节点组成的树中
- 拆分节点 - 存储有关其左右之间分散的哪些项目 ID 的信息。
- 后代节点 - 存储多个项目 ID。
- 项目 - 单个向量。
对于以下示例,我们将考虑后代节点不能包含多于两个项目,并且 ID 越接近,项目在空间中越接近。例如,项目 1
接近项目 2
,但远离项目 10
。
在这里,我们尝试插入项目 2
和 7
,并删除项目 3
和 11
。
按照我们在搜索中所做的工作,我们将要插入的项目在左右拆分节点之间拆分。但是,已删除的项目不再存在于数据库中!我们没有它们关联的向量,因此我们需要迭代所有叶子来查找和删除它们。
幸运的是,由于我们使用咆哮位图,此列表不会消耗太多内存,并且由于我们从不更新它,因此我们可以在不复制它的情况下与每个线程共享它 🎉
请注意,要插入的项目是如何在左右拆分节点之间进行平衡的。在下一步中,我们将遵循相同的过程,并在拆分节点上拆分两个要插入的项目列表。
正是在这一步中,所有好的事情都发生了。从左到右
- 第一个后代节点删除了其项目
3
,并添加了项目2
。 - 第二个后代节点变得太大,必须替换为一个新的拆分节点。
- 项目
8
成为一个包含项目7
和8
的后代。 - 在最后一个后代节点中删除项目
10
后,我们将其替换为指向项目的直接链接,以减少树中的节点数量(并通过减少查找次数来缩短搜索时间)。
现在您已经看到了我们增量索引过程的所有步骤,我们仍然要记下发生的一些事情
- 在前三个步骤中,我们必须读取所有树节点,默认情况下,这等于数据库中的项目总数。
- 在最后一步中,我们必须写入四个新的树节点。以前,我们会重写整个树。写入次数(这是数据库中最慢的操作)减少到最低限度。
- 该过程在单个树上运行,这意味着我们仍然可以多线程计算每棵树。
- ID 不再是顺序的,这与 并行写入树 中描述的 ID 生成策略不太兼容。
我们是如何修复 ID 生成的?
在索引过程开始时,我们需要收集所有现有的树 ID。
我们需要生成新节点的信息是
- 用于知道何时停止构建树的树节点总数。
- 我们想要重用的已用 ID 中的“空洞”。当我们编辑树时,会产生这种碎片。
- 树中存在的最高 ID,用于生成下一个新 ID。
“简单”的想法是为树节点的总数设置一个计数器。我们可以增加一个计数器来生成新的 ID。最具挑战性的部分是在线程之间共享一个可用的 ID,并且没有大的互斥锁。
我花了一些时间才找到解决最后一点的方案;我实现了一个完整的复杂解决方案 (150loc),最后我在一个小时内 从头开始重写了整个方案,以获得一个简单且更高效的解决方案。我们将看到在找到最终解决方案之前我们经历的所有想法。
在可用 ID 上共享迭代器
首先想到的最直接的想法是创建一个迭代器并在所有可用 ID 上共享它。
let mut available_ids = available_ids.into_iter(); roots.into_par_iter().for_each(|root| { // stuff let next_id = available_ids.next(); // stuff });
如果您熟悉 Rust,您会立即注意到 for_each
无法捕获 available_ids
上的可变引用,因为我们使用了来自 Rayon 的 into_par_iter
方法,该方法以多线程方式在迭代器上执行下一个方法。这意味着我们无法调用 .next()
。
通过使用互斥锁(互斥的缩写),我们可以使其编译
let mut available_ids = Mutex::new(available_ids.into_iter()); roots.into_par_iter().for_each(|root| { // stuff let next_id = available_ids.lock().unwrap().next(); // stuff });
这是安全的,并且会产生正确的结果,但不会很好地扩展,因为所有线程都必须在锁上相互等待。
让每个线程拥有自己的 ID 列表
当您需要删除锁时,一个常见的解决方案是事先拆分工作,以便每个线程都可以全力以赴地工作,而无需知道其他线程上正在发生的事情。这是我 最初实现的 解决方案。这个想法是将咆哮位图分解成与要更新的树一样多的较小的咆哮位图。
以下函数可以将咆哮位图拆分为大小相等的 n
个子位图
pub fn split_ids_in(available: RoaringBitmap, n: usize) -> Vec<RoaringBitmap> { // Define the number of elements per sub-bitmap let chunk_size = available.len() / n as u64; let mut ret = Vec::new(); let mut iter = available.into_iter(); for _ in 0..(n - 1) { // Create a roaring bitmap of `chunk_size` elements from the iterator without consuming it let available: RoaringBitmap = iter.by_ref().take(chunk_size as usize).collect(); ret.push(available); } // the last element is going to contain everything remaining ret.extend(iter); ret }
有了这个函数,就可以很方便地使用每个树一个位图,并通过 Rayon 并行迭代器的 zip
方法进行更新。
let available_ids = split_ids_in(available_ids, roots.len()); roots.into_par_iter().zip(available_ids).for_each(|(root, available_ids)| { let mut available_ids = available_ids.into_iter(); // stuff let next_id = available_ids.next(); // Here, we can use the available_ids without any locks or inter-thread synchronization // stuff });
当我们更新已存在的根树时,这种方法效果很好。但之后,我们可能需要从头创建新的树,而我们无法预先知道需要创建多少新树。这是一个问题,因为没有这些信息,我们就不知道需要创建多少个子位图。
在这一点上,我没有看到解决这个问题的简单方法,但我假设我们很少创建大量的新树,并且所有的新树都会使用大量的 ID。这意味着,将所有 ID 都给**第一棵**新树可能*会奏效*(在没有在生产环境中监控的情况下,很难 100% 确定)。
最终解决方案
之前的解决方案很复杂,甚至不能完美工作。当我在浏览 Roaring Bitmap 的文档时,我看到了 select
方法,并立即明白了如何使最初的方法可行。
此方法允许您以高效的方式通过索引获取位图中的值。
例如,如果位图中有以下值:13, 14, 15, 98, 99, 100
select(0)
返回13
。select(3)
返回98
。select(5)
返回100
。select(6)
返回None
。
考虑到这一点,如果我们以**只读**方式与所有线程共享位图,并将位图中的索引作为原子数字共享,我们可以让多个线程同时获取可用的 ID,而无需锁的同步。
一旦 select
方法返回 None
,这意味着我们可以停止从位图中选择 ID,而是使用旧方法从头开始生成新的 ID。
即使速度非常快,调用 select
方法仍然需要一些时间,因此一旦线程从该方法收到 None
返回值,我们将更新一个原子布尔值,告诉我们位图中是否还有值需要查找。如果没有,我们可以跳过调用 select
,直接生成新的 ID。
考虑到这一点,我们在第二部分中向您展示的酷炫的 ConcurrentNodeIds
结构变得稍微复杂了一些,但它仍然允许我们在没有任何锁的情况下公平地生成 ID!
/// A concurrent ID generator that will never return the same ID twice. pub struct ConcurrentNodeIds { /// The current tree node ID we should use if no other IDs are available. current: AtomicU32, /// The total number of tree node IDs used. used: AtomicU64, /// A list of IDs to exhaust before picking IDs from `current`. available: RoaringBitmap, /// The current Nth ID to select in the bitmap. select_in_bitmap: AtomicU32, /// Tells if you should look in the roaring bitmap or if all the IDs are already exhausted. look_into_bitmap: AtomicBool, } impl ConcurrentNodeIds { /// Creates an ID generator returning unique IDs, avoiding the specified used IDs. pub fn new(used: RoaringBitmap) -> ConcurrentNodeIds { let last_id = used.max().map_or(0, |id| id + 1); let used_ids = used.len(); let available = RoaringBitmap::from_sorted_iter(0..last_id).unwrap() - used; ConcurrentNodeIds { current: AtomicU32::new(last_id), used: AtomicU64::new(used_ids), select_in_bitmap: AtomicU32::new(0), look_into_bitmap: AtomicBool::new(!available.is_empty()), available, } } /// Returns a new unique ID and increases the count of IDs used. pub fn next(&self) -> Result<u32> { if self.look_into_bitmap.load(Ordering::Relaxed) { let current = self.select_in_bitmap.fetch_add(1, Ordering::Relaxed); match self.available.select(current) { Some(id) => Ok(id), None => { self.look_into_bitmap.store(false, Ordering::Relaxed); Ok(self.current.fetch_add(1, Ordering::Relaxed)) } } } else { Ok(self.current.fetch_add(1, Ordering::Relaxed)) } } }
实践中的性能
我还没有进行大量的基准测试(我很想这样做,并且可能会在以后更新这篇文章),但在我对几十万个项目进行的基准测试中,我得到了以下结果:
- 平均而言,在索引了三批项目后,我们获得了超过 10 倍的加速。
- 我们之后观察到的小幅上升,使我们从约 700 毫秒变为约 3 秒,这是由于随着项目数量的增加而创建新树所致。
- 我预计此基准测试代表最坏的情况
- 我所有的项目 ID 都是随机生成的,这意味着很少有重复/更新,但始终有新的插入和更多的写入。
- 我的向量也是以均匀分布随机生成的,这意味着不应该像在真实数据集中观察到的那样出现聚类。
此功能现在已在包含数百万个文档的生产数据中使用,我们看到在不到一分钟的时间内,在 200 万个项目的数据库中添加了 9000 个项目的数量级。
Meilisearch 是一个开源搜索引擎,使开发人员能够在享受简单直观的 DX 的同时构建最先进的体验。
有关 Meilisearch 的更多信息,您可以加入 Discord 上的社区,或订阅 新闻简报。您可以通过查看 路线图 并参与 产品讨论 来了解有关该产品的更多信息。