受 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++ 函式庫相關的東西。因此，我們大量使用 &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 中實作了 微軟的 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 上的社群，或訂閱電子報。您可以查看 路線圖 並參與產品討論，以了解更多關於產品的資訊。