[diskann-garnet] Optimize reranking (#1308)

This addresses several performance issues with reranking.

1. The existence check turned out to be expensive as it consults the
FSM. This reads 1 extra key per Lsearch which is unnecessary as we skip
missing vectors anyway.
2. The full vector reads were serial. Using multiread allows Garnet to
parallelize these.
3. The reranked candidates buffer was allocated in the post processor;
pooling this allocation removes it from the query path.

Since we don't have good recall based tests yet, I hand verified the
recall was still ok with this change.

Author: metajack

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "diskann-garnet"
-version = "4.0.3"
+version = "4.0.4"
 edition = "2024"
 authors.workspace = true
 license.workspace = true

diskann-garnet/Cargo.toml

@@ -2,7 +2,7 @@
 <package>
   <metadata>
     <id>diskann-garnet</id>
-    <version>4.0.3</version>
+    <version>4.0.4</version>
     <readme>docs/README.md</readme>
     <authors>Microsoft</authors>
     <projectUrl>https://github.com/microsoft/DiskANN</projectUrl>

diskann-garnet/diskann-garnet.nuspec

@@ -62,6 +62,9 @@ use crate::{
 /// bytes are the serialized quant table.
 const QUANT_STATE_KEY: u32 = u32::from_be_bytes(*b"_qnt");
 
+/// Starting capacity of the pre-allocated rerank buffers.
+const RERANK_BUFFER_LENGTH: usize = 1024;
+
 #[derive(Clone)]
 struct AdjList(AdjacencyList<u32>);
 
@@ -138,6 +141,8 @@ pub(crate) struct GarnetProvider<T: VectorRepr> {
     /// Pool of pre-allocated buffers to use for filter decisions during
     /// filtered search beam expansion
     filtered_decisions_pool: ObjectPool<Vec<bool>>,
+    /// Pool of pre-allocated buffers to use for reranking
+    rerank_pool: ObjectPool<Vec<Neighbor<u32>>>,
     /// Pool of pre-allocated buffers to use for quantizing vectors
     quant_buffer_pool: ObjectPool<Vec<u8>>,
     /// Small cache for the start points' neighbors
@@ -173,6 +178,11 @@ impl<T: VectorRepr> GarnetProvider<T> {
             parallelism,
             Some(parallelism),
         );
+        let rerank_pool = ObjectPool::new(
+            Undef::new(RERANK_BUFFER_LENGTH),
+            parallelism,
+            Some(parallelism),
+        );
 
         let start_point_cache =
             DashMap::with_capacity_and_hasher(1, foldhash::fast::RandomState::default());
@@ -293,6 +303,7 @@ impl<T: VectorRepr> GarnetProvider<T> {
             id_buffer_pool,
             filtered_ids_pool,
             filtered_decisions_pool,
+            rerank_pool,
             quant_buffer_pool,
             start_point_cache,
             start_point_quant_cache,
@@ -1302,34 +1313,40 @@ impl<'a, 'b, T: VectorRepr> SearchPostProcessStep<DynamicAccessor<'a, T>, &'b [T
                 .map_err(|e| GarnetProviderError::PostProcessing(Box::new(e)));
         }
 
-        let provider = &accessor.provider;
+        let provider = accessor.provider;
         let f = T::distance(provider.metric_type, Some(provider.dim));
-        let mut v = Poly::broadcast(0u8, provider.dim * mem::size_of::<T>(), AlignToEight)?;
-
-        // Filter before computing the full precision distances.
-        let mut reranked: Vec<_> = candidates
-            .filter_map(|n| {
-                if !provider.vector_iid_exists(accessor.context, *n.id()) {
-                    None
-                } else if provider.callbacks.read_single_iid(
-                    &accessor.context.term(Term::Vector),
-                    *n.id(),
-                    &mut v,
-                ) {
-                    Some(Neighbor::new(
-                        *n.id(),
-                        f.evaluate_similarity(query, bytemuck::cast_slice::<u8, T>(&v)),
-                    ))
-                } else {
-                    None
-                }
-            })
-            .collect();
+
+        let mut reranked = provider
+            .rerank_pool
+            .get_ref(Undef::new(RERANK_BUFFER_LENGTH));
+        reranked.clear();
+
+        // Use the accessor.filtered_ids pre-allocated buffer to do a multi read from Garnet, placing the results in
+        // the rerank buffer.
+        accessor.filtered_ids.clear();
+        for nbor in candidates {
+            accessor.filtered_ids.push(4);
+            accessor.filtered_ids.push(*nbor.id());
+        }
+
+        if !accessor.filtered_ids.is_empty() {
+            provider.callbacks.read_multi_lpiid(
+                &accessor.context.term(Term::Vector),
+                &accessor.filtered_ids,
+                |i, v| {
+                    let dist = f.evaluate_similarity(query, bytemuck::cast_slice::<u8, T>(v));
+                    reranked.push(Neighbor::new(
+                        accessor.filtered_ids[i as usize * 2 + 1],
+                        dist,
+                    ));
+                },
+            );
+        }
 
         // Sort the full precision distances.
         reranked.sort_unstable_by(diskann::neighbor::ord::fast_distance);
 
-        next.post_process(accessor, query, reranked.into_iter(), output)
+        next.post_process(accessor, query, reranked.iter().copied(), output)
             .await
             .map_err(|e| GarnetProviderError::PostProcessing(Box::new(e)))
     }