多執行緒和記憶體映射：使用 Arroy 優化 ANN 效能

這是系列部落格文章的第二部分 最初發布在我的個人部落格上。從 第一部分開始旅程。

如果可以向你展示單執行緒、記憶體映射的鍵值儲存如何比手寫的記憶體映射解決方案更有效率，那不是很好嗎？在將 Spotify 基於磁碟的近似最近鄰居函式庫移植到 Rust，更具體地說是移植到 LMDB 時，我遇到了問題。這些問題主要歸咎於 LMDB 和記憶體安全。以下是這個故事。

提醒您，Arroy 是一個將嵌入（浮點數向量）儲存在磁碟上的函式庫。在這些向量之上會產生一些資料結構，看起來像是以常態控制的樹狀結構，用於遞迴地將向量資料集分割成子網域。但是你可以在本系列的 [第一部分](/blog/spotify-inspired-hybrid-search-and-rust/) 中了解更多關於這方面的資訊。

Annoy 如何產生樹狀節點？

Annoy，Spotify 函式庫會將節點儲存在磁碟上，項目節點（包含嵌入的節點）以及我們將在本文章中稱為樹狀節點的其他節點。這樣做的好處是雙重的

• 程式的記憶體使用量很低，並且由作業系統優化，因為所有節點都存在於記憶體映射的檔案中。
• 概念很簡單：使用其 ID 存取節點。函式庫將使用簡單的乘法找到其在磁碟上的偏移量。

但是，這樣做也有缺點。所有節點：包含嵌入的項目以及樹狀節點的大小必須相同，並且如果使用者想要使用 ID 231 來識別其嵌入，則函式庫會將檔案大小增加到 ID 乘以節點大小。例如，對於一個 768 維的向量，儲存一個 ID 為 231 的單一項目將產生一個超過 6.5 TiB 的檔案。

在產生最終的資料結構時，樹狀節點會與包含嵌入的使用者項目一起寫入同一個檔案中。這些樹狀建構過程是並行運行的，每個樹一個運行過程，因此在定義新建立的節點 ID 和為其保留的記憶體區域時，需要大量的同步，最重要的是，當記憶體映射檔案太小而無法接受更多節點時，只有一個執行緒必須擴大檔案，因此在可變記憶體映射之上並圍繞著 node_id 變數使用互斥鎖。樹狀產生的有趣特性之一是，產生過程只需要包含嵌入的原始使用者項目。

將其移植到 LMDB 時遇到的挑戰

有趣的事實是，這是我很長一段時間以來第一次寫 C++，也是我第一次要求 ChatGPT 為我編寫程式碼，因為我對 C++ 沒有信心，並且擔心會遇到一些記憶體區段錯誤。我需要一個小程式從 stdin 反序列化嵌入並將其提供給 Annoy。聊天機器人的程式碼大多有效，但它省略了一個關鍵的 空向量檢查，這導致了記憶體區段錯誤...

將其移植到 LMDB 的主要障礙是，寫入此鍵值儲存是單執行緒的。它不支援並行的寫入操作，以維持正確平衡的 BTree。好戲即將上演！

從不同執行緒讀取使用者項目節點

自一開始，我們就在 Meilisearch 中使用 LMDB。它是一個維護良好的鍵值儲存，在 Firefox 中使用，並由 OpenLDAP 開發人員維護。它是記憶體映射的，並且不維護記憶體中條目的使用者端快取，而是為您提供指向磁碟上記憶體映射區域的指標。主要優勢在於，只要您的讀取交易存在，您就可以保留指向此區域的指標。是的！因為它是一個支援原子提交的交易型鍵值儲存！

但是樹狀產生不需要參考產生的節點，而只需要使用者項目。我們之前看到 LMDB 提供直接指向記憶體映射檔案的指標，而無需維護任何中間快取（可能會失效）。LMDB 還有另一個怪癖：您無法在多個執行緒之間共享唯讀交易，即 RoTxn: !Sync，並且您無法在執行緒之間移動寫入交易，即 RwTxn: !Send + !Sync。哈！而且沒有辦法在未提交的變更上建立讀取交易。這是一個問題，因為我們想要在儲存項目的同一個交易中產生資料結構樹。

但是魔力無處不在，從以下小型且有趣的資料結構開始。原理是將指向內部使用者項目的指標及其嵌入保留在 Vec<*const u8> 中。感謝 Rust，我們可以確保在編譯時透過將生命週期保留在結構中，指標將存在足夠長的時間。透過使用 Deref 使用 &'t RwTxn 取得 &'t RoTxn 也確保我們無法在使用 &'t mut RwTxn 讀取時修改資料庫。 根據 LMDB 的主要開發人員表示，在執行緒之間共享這些指標是安全的，這也是我為此結構實作 Sync 的原因。

pub struct ImmutableItems<'t, D> {
    item_ids: RoaringBitmap,
    constant_length: Option<usize>,
    offsets: Vec<*const u8>,
    _marker: marker::PhantomData<(&'t (), D)>,
}

impl<'t, D: Distance> ImmutableItems<'t, D> {
    pub fn new(rtxn: &'t RoTxn, database: Database<D>, index: u16) -> heed::Result<Self> {
        let mut item_ids = RoaringBitmap::new();
        let mut constant_length = None;
        let mut offsets = Vec::new();

        for result in database.items()? {
            let (item_id, bytes) = result?;
            assert_eq!(*constant_length.get_or_insert(bytes.len()), bytes.len());
            item_ids.push(item_id);
            offsets.push(bytes.as_ptr());
        }

        Ok(ImmutableLeafs { item_ids, constant_length, offsets, _marker: marker::PhantomData })
    }

    pub fn get(&self, item_id: ItemId) -> heed::Result<Option<Leaf<'t, D>>> {
        let len = match self.constant_length {
            Some(len) => len,
            None => return Ok(None),
        };
        let ptr = match self.item_ids.rank(item_id).checked_sub(1).and_then(|offset| self.offsets.get(offset)) {
            Some(ptr) => *ptr,
            None => return Ok(None),
        };
        let bytes = unsafe { slice::from_raw_parts(ptr, len) };
        ItemCodec::bytes_decode(bytes).map(Some)
    }
}

unsafe impl<D> Sync for ImmutableItems<'_, D> {}

您可能還注意到了此簡化版本資料結構中的其他一些有趣的優化。我們也知道使用者項目節點具有恆定的長度，因此我決定只儲存一次，將偏移向量的大小縮小了兩倍。考慮到我們的目標是儲存 1 億個向量，並且這個向量位於記憶體中，我們將其大小從 1526 MiB 縮小到 763 MiB，這不是很多，但總比沒有好。

並行寫入樹狀節點

好的！現在我們知道如何在不進行任何使用者端同步的情況下儲存指向項目的指標並在執行緒之間共享它們，我們需要使用它來產生樹狀節點。我們已經知道如何在 Meilisearch 中處理 LMDB，並決定實作相同的方法。為了並行寫入 LMDB，請寫入不同的獨立檔案，然後將所有內容合併。這利用了 無共享原則 並隔離了演算法。這大大減少了同步點的數量 與原始 C++ 程式碼相比（請查看 .lock/unlock 呼叫），在我們的程式碼中只有一行：原子性增加的全域樹狀節點 ID。

pub fn next_node_id(&self) -> u64 {
    self.0.fetch_add(1, Ordering::Relaxed)
}

我們基於使用者項目節點產生常態分割節點的函數現在只是將節點寫入獨立的檔案中。節點會附加到檔案中，而磁碟上的偏移量和界限會儲存到向量中，每個節點一個 usize。使用 Rayon，我們可以並行執行所有樹狀產生函數，並且在完成後，檢索檔案和界限，以將唯一識別的節點循序寫入 LMDB 中。

效能比較

我們在 Meilisearch 的目標是支援約 768 維的 1 億個嵌入。如果我們將它們儲存為沒有任何 降維 的 f32 向量，則相當於 100M * 768 * 4 = 307B，換句話說，儲存原始向量需要 286 GiB，而無需任何內部樹狀節點，也就是說，沒有辦法有效地在其中搜尋。

如果您不指定要產生的樹狀結構數量，則當樹狀節點的數量小於項目向量的數量時，該演算法將繼續建立樹狀結構。最後，樹狀節點和項目的數量必須大致相同。

探索我們可以索引的向量限制

Arroy 和 Annoy 都廣泛使用記憶體映射，但方式略有不同。在來自 @dureuill 的上一篇文章中，我們看到 [作業系統沒有無限的記憶體映射空間](/blog/dynamic-virtual-address-management/)。所以，讓我們深入了解效能結果。

當使用 768 維的 2500 萬個向量運行 Arroy 時，我注意到了一些奇怪的事情。CPU 沒有被使用，並且程式似乎大量讀取 SSD/NVMe，太多了 🤔 樹狀產生演算法正在將向量空間分割成具有較少向量的子網域。但是，它必須首先找到合適的常態來分割整個網域，為此，它會隨機選擇許多項目。不幸的是，我的機器只有 63 GiB 的記憶體，而索引這 2500 萬個項目需要超過 72 Gib。Annoy 也以相同的方式掙扎。

經過調查，我了解了為什麼會達到整個交換空間和記憶體映射限制。不僅僅是因為我們將範數和其他資料與向量一起儲存，使得項目節點大小不只是768 * 4 位元組，而且在 Arroy 的情況下，LMDB 需要維護關於條目的 BTree 資料結構，這些也會佔用記憶體映射空間。兩個程式都會請求記憶體中不可用的隨機項目節點，因此作業系統會從磁碟中提取它們，這需要時間。而且每個執行緒都是平行進行的。CPU 只是在等待磁碟。

因此，經過一些二分法搜尋後，我找到了 Arroy 成功使用所有記憶體而不受限於磁碟速度的臨界點。它可以在一台 63 GiB 的機器上為 1562.5 萬個項目建立索引。從這個 htop 截圖中可以看到，由於所有向量都已載入記憶體中，磁碟讀取速度為零，Arroy 正在將樹節點寫入磁碟，並且 CPU 正盡其所能地工作。處理過程花費不到七個小時。

但是... Annoy 無法為相同數量的文件建立索引。它遇到與之前相同的問題：磁碟讀取速度高而 CPU 使用率低。但為什麼呢？我需要釐清原因，因為節點格式相同、要建立索引的項目向量數量相同，並且演算法也已移植。那麼，兩種解決方案之間的差異是什麼？

對於那些研究過 C++ 程式碼庫並且可能受到之前描述的記憶體映射問題啟發的人來說，您可能注意到了這個細微的差異：Arroy 是將產生的樹節點寫入不同的原始檔案中，而 Annoy 則相反，它是在記憶體映射的檔案中保留空間並直接寫入。通過這種技巧，作業系統需要在記憶體映射區域中保留更多的空間，並且被迫使快取中的項目節點失效，以保持剛寫入的樹節點在快取中，這會無緣無故地拖慢系統速度，因為該演算法並不需要樹節點。

因此，在經過更多二分法搜尋以找出 Annoy 在 63 GiB 機器上的限制後，我發現它大約可以在五個小時內為 1000 萬個向量建立索引。

「共享無」（share nothing）原則獲勝。

所以 Arroy 可以在同一台機器上為多 36% 的向量建立索引，但這需要多長時間？與其以單執行緒方式複製所有這些樹節點，以平行方式寫入同一個檔案是否會更快？這一段會簡短得多，因為我只做了一些小測試，但 Arroy 總是更快！

現在，您可能會告訴我，我需要大約 400GiB 的記憶體才能為 1 億個向量建立索引，您可能是對的，但@irevoire將在未來的文章中談論增量索引。我為了好玩做了一個線性迴歸。在 12 個執行緒和所需記憶體的情況下，我預測這將需要一天 22 小時 43 分鐘😬。喔！由於我們也在 Meilisearch 中使用這個向量儲存，我們需要在搜尋時提供過濾這個資料結構的方法。這是本系列的下一篇文章。

您可以在 Lobste.rs、Hacker News、Rust Subreddit 或 X (前身為 Twitter)上評論這篇文章。

本系列的第 3 部分現已推出，探討了我們使用 Arroy 的過濾式磁碟 ANN 進行搜尋的進一步進展。

Meilisearch 是一個開源搜尋引擎，不僅為終端使用者提供最先進的體驗，還提供簡單直觀的開發人員體驗。

如需更多關於 Meilisearch 的資訊，您可以加入 Discord 上的社群，或訂閱電子報。您可以查看路線圖並參與產品討論，以了解更多關於該產品的資訊。