nidx-binding 6.9.0.post567__tar.gz → 6.9.0.post576__tar.gz

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.9.0.post567
+Version: 6.9.0.post576
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/nidx_protos/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "nidx_protos"
-version = "6.9.0.post567"
+version = "6.9.0.post576"
 license = "AGPL-3.0-or-later"
 description = "Protobuf definitions for nucliadb/nidx"
 authors = [{ name = "Nuclia", email = "nucliadb@nuclia.com" }]

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/nidx_vector/src/hnsw/ops_hnsw.rs

@@ -37,7 +37,7 @@ use super::*;
 /// Implementors of this trait can guide the hnsw search
 pub trait DataRetriever: std::marker::Sync {
     fn similarity(&self, x: VectorAddr, y: &SearchVector) -> f32;
-    fn similarity_upper_bound(&self, x: VectorAddr, y: &SearchVector) -> f32;
+    fn similarity_upper_bound(&self, x: VectorAddr, y: &SearchVector) -> EstimatedScore;
     fn paragraph(&self, x: VectorAddr) -> ParagraphAddr;
     fn get_vector(&self, x: VectorAddr) -> &[u8];
     /// Embeddings with smaller similarity should not be considered.
@@ -67,6 +67,28 @@ pub trait Hnsw {
     fn get_layer(&self, i: usize) -> Self::L;
 }
 
+#[derive(Clone, Copy)]
+pub struct EstimatedScore {
+    pub score: f32,
+    pub upper_bound: f32,
+}
+
+impl EstimatedScore {
+    pub fn new_with_error(score: f32, error: f32) -> Self {
+        Self {
+            score,
+            upper_bound: score + error,
+        }
+    }
+
+    pub fn new_exact(score: f32) -> Self {
+        Self {
+            score,
+            upper_bound: score,
+        }
+    }
+}
+
 ///  Tuples ([`VectorAddr`], [`f32`]) can not be stored in a [`BinaryHeap`] because [`f32`] does not
 /// implement [`Ord`]. [`Cnx`] is an application of the new-type pattern that lets us bypass the
 /// orphan rules and store such tuples in a [`BinaryHeap`].
@@ -89,6 +111,25 @@ impl PartialOrd for Cnx {
     }
 }
 
+#[derive(Clone, Copy)]
+pub struct CnxWithBound(VectorAddr, EstimatedScore);
+impl Eq for CnxWithBound {}
+impl Ord for CnxWithBound {
+    fn cmp(&self, other: &Self) -> Ordering {
+        f32::total_cmp(&self.1.score, &other.1.score)
+    }
+}
+impl PartialEq for CnxWithBound {
+    fn eq(&self, other: &Self) -> bool {
+        f32::eq(&self.1.score, &other.1.score)
+    }
+}
+impl PartialOrd for CnxWithBound {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
 /// A list of neighbours containing a pointer to the embedding and their similarity.
 pub type Neighbours = Vec<(VectorAddr, f32)>;
 
@@ -243,7 +284,7 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         layer: L,
         k_neighbours: usize,
         entry_points: &[VectorAddr],
-    ) -> Neighbours {
+    ) -> impl Iterator<Item = (VectorAddr, EstimatedScore)> {
         // Nodes already visited
         let mut visited = FxHashSet::default();
         // Nodes to visit
@@ -254,9 +295,9 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         // The initial candidates are the entry points
         for ep in entry_points.iter().copied() {
             visited.insert(ep);
-            let similarity = self.retriever.similarity(ep, query);
-            candidates.push(Cnx(ep, similarity));
-            ms_neighbours.push(Reverse(Cnx(ep, similarity)));
+            let similarity = self.retriever.similarity_upper_bound(ep, query);
+            candidates.push(CnxWithBound(ep, similarity));
+            ms_neighbours.push(Reverse(CnxWithBound(ep, similarity)));
         }
 
         loop {
@@ -264,9 +305,14 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
                 // No more candidates, done
                 (None, _) => break,
                 // Candidate is worse than worse result, done
-                (Some(Cnx(_, cs)), Some(Reverse(Cnx(_, ws)))) if cs < ws => break,
+                (
+                    Some(CnxWithBound(_, EstimatedScore { score: cs, .. })),
+                    Some(Reverse(CnxWithBound(_, EstimatedScore { score: ws, .. }))),
+                ) if cs < ws => {
+                    break;
+                }
                 // Candidate is better than worse result
-                (Some(Cnx(cn, _)), Some(Reverse(Cnx(_, mut ws)))) => {
+                (Some(CnxWithBound(cn, _)), Some(Reverse(CnxWithBound(_, mut ws)))) => {
                     for (y, _) in layer.get_out_edges(cn) {
                         if self.preload_nodes && !visited.contains(&y) {
                             self.retriever.will_need_vector(y);
@@ -275,10 +321,10 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
                     for (y, _) in layer.get_out_edges(cn) {
                         if !visited.contains(&y) {
                             visited.insert(y);
-                            let similarity = self.retriever.similarity(y, query);
-                            if similarity > ws || ms_neighbours.len() < k_neighbours {
-                                candidates.push(Cnx(y, similarity));
-                                ms_neighbours.push(Reverse(Cnx(y, similarity)));
+                            let similarity = self.retriever.similarity_upper_bound(y, query);
+                            if similarity.score > ws.score || ms_neighbours.len() < k_neighbours {
+                                candidates.push(CnxWithBound(y, similarity));
+                                ms_neighbours.push(Reverse(CnxWithBound(y, similarity)));
                                 if ms_neighbours.len() > k_neighbours {
                                     ms_neighbours.pop();
                                 }
@@ -290,11 +336,7 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
                 _ => (),
             }
         }
-        ms_neighbours
-            .into_sorted_vec()
-            .into_iter()
-            .map(|Reverse(Cnx(n, d))| (n, d))
-            .collect()
+        ms_neighbours.into_iter().map(|Reverse(CnxWithBound(n, d))| (n, d))
     }
 
     fn layer_insert(
@@ -305,9 +347,11 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         mmax: usize,
     ) -> Vec<VectorAddr> {
         use params::*;
-        let neighbours =
-            self.layer_search::<&RAMLayer>(&SearchVector::Stored(x), layer, ef_construction(), entry_points);
-        let neighbours = self.select_neighbours_heuristic(m(), neighbours, layer);
+        let search_neighbours = self
+            .layer_search::<&RAMLayer>(&SearchVector::Stored(x), layer, EF_CONSTRUCTION, entry_points)
+            .map(|(addr, score)| (addr, score.score))
+            .collect();
+        let neighbours = self.select_neighbours_heuristic(M, search_neighbours, layer);
         let mut needs_repair = HashSet::new();
         let mut result = Vec::with_capacity(neighbours.len());
         layer.add_node(x);
@@ -343,7 +387,11 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
                 for l in (0..=top_layer).rev() {
                     if l > level {
                         // Above insertion point, just search
-                        eps[0] = self.layer_search(&SearchVector::Stored(x), &hnsw.layers[l], 1, &eps)[0].0;
+                        let new_ep = self
+                            .layer_search(&SearchVector::Stored(x), &hnsw.layers[l], 1, &eps)
+                            .next()
+                            .unwrap();
+                        eps[0] = new_ep.0;
                     } else {
                         eps = self.layer_insert(x, &mut hnsw.layers[l], &eps, params::m_max_for_layer(l));
                     }
@@ -370,13 +418,13 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         };
 
         let mut crnt_layer = entry_point.layer;
-        let mut neighbours = vec![(entry_point.node, 0.)];
+        let mut entry_points = vec![entry_point.node];
 
+        // Traverse upper layer finding the best match on each one
         while crnt_layer != 0 {
             let layer = hnsw.get_layer(crnt_layer);
-            let entry_points: Vec<_> = neighbours.into_iter().map(|(node, _)| node).collect();
             let layer_res = self.layer_search(query, layer, 1, &entry_points);
-            neighbours = layer_res;
+            entry_points = layer_res.map(|(addr, _)| addr).collect();
             crnt_layer -= 1;
         }
 
@@ -387,23 +435,24 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
             None
         };
         let last_layer_k = if original_query.is_some() {
-            std::cmp::min(k_neighbours * 100, 2000)
+            std::cmp::min(k_neighbours * rabitq::RERANKING_FACTOR, rabitq::RERANKING_LIMIT)
         } else {
             k_neighbours
         };
 
-        let entry_points: Vec<_> = neighbours.into_iter().map(|(node, _)| node).collect();
+        // Find the best k nodes in the last layer
         let layer = hnsw.get_layer(crnt_layer);
-        let mut neighbors = self.layer_search(query, layer, last_layer_k, &entry_points);
+        let neighbours = self.layer_search(query, layer, last_layer_k, &entry_points);
 
-        // If using RabitQ, rerank using the original vectors
-        if let Some(query) = original_query {
-            neighbors = neighbors
+        let entry_points = if let Some(query) = original_query {
+            // If using RabitQ, rerank using the original vectors
+            rabitq::rerank_top(neighbours.collect(), k_neighbours, self.retriever, query)
                 .into_iter()
-                .map(|(addr, _)| (addr, self.retriever.similarity(addr, query)))
-                .collect();
-            neighbors.sort_unstable_by(|x, y| y.1.partial_cmp(&x.1).unwrap_or(Ordering::Equal));
-        }
+                .map(|Reverse(Cnx(addr, _))| addr)
+                .collect()
+        } else {
+            neighbours.map(|(addr, _)| addr).collect()
+        };
 
         let filter = NodeFilter {
             filter: with_filter,
@@ -412,8 +461,8 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
             vec_counter: RepCounter::new(!with_duplicates),
         };
         let layer_zero = hnsw.get_layer(0);
-        let entry_points: Vec<_> = neighbors.into_iter().take(k_neighbours).map(|(node, _)| node).collect();
 
+        // Find k nodes that match the filter in the last layer
         let mut filtered_result = self.closest_up_nodes(
             entry_points,
             original_query.unwrap_or(query),

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/nidx_vector/src/hnsw/params.rs

@@ -23,11 +23,11 @@
 
 /// Factor by which the layer distribution should deviate.
 pub fn level_factor() -> f64 {
-    1.0 / (m() as f64).ln()
+    1.0 / (M as f64).ln()
 }
 
 pub const fn m_max_for_layer(layer: usize) -> usize {
-    if layer == 0 { m_max0() } else { m_max() }
+    if layer == 0 { M_MAX_0 } else { M_MAX }
 }
 
 /// M to use when pruning neighbours
@@ -36,21 +36,13 @@ pub const fn prune_m(m: usize) -> usize {
 }
 
 /// Upper limit to the number of out-edges a embedding can have.
-pub const fn m_max0() -> usize {
-    60
-}
+pub const M_MAX_0: usize = 60;
 
 /// Upper limit to the number of out-edges a embedding can have.
-pub const fn m_max() -> usize {
-    30
-}
+pub const M_MAX: usize = 30;
 
 /// Number of bi-directional links created for every new element.
-pub const fn m() -> usize {
-    30
-}
+pub const M: usize = 30;
 
 /// Number of neighbours that are searched for before adding a new embedding.
-pub const fn ef_construction() -> usize {
-    100
-}
+pub const EF_CONSTRUCTION: usize = 100;

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/nidx_vector/src/hnsw.rs

@@ -24,8 +24,5 @@ mod params;
 mod ram_hnsw;
 
 pub use disk_hnsw::DiskHnsw;
-pub use ops_hnsw::Cnx;
-pub use ops_hnsw::DataRetriever;
-pub use ops_hnsw::HnswOps;
-pub use ops_hnsw::SearchVector;
+pub use ops_hnsw::{Cnx, DataRetriever, EstimatedScore, HnswOps, SearchVector};
 pub use ram_hnsw::RAMHnsw;

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/nidx_vector/src/segment.rs

@@ -352,14 +352,14 @@ impl<DS: DataStore> DataRetriever for Retriever<'_, DS> {
         }
     }
 
-    fn similarity_upper_bound(&self, x: VectorAddr, y: &SearchVector) -> f32 {
+    fn similarity_upper_bound(&self, x: VectorAddr, y: &SearchVector) -> EstimatedScore {
         match y {
             SearchVector::RabitQ(query) => {
                 let x = self.data_store.get_quantized_vector(x);
                 let (est, err) = query.similarity(x);
-                est + err
+                EstimatedScore::new_with_error(est, err)
             }
-            _ => self.similarity(x, y),
+            _ => EstimatedScore::new_exact(self.similarity(x, y)),
         }
     }
 
@@ -492,16 +492,16 @@ impl OpenSegment {
     }
 
     #[allow(clippy::too_many_arguments)]
-    fn _search<DS: DataStore>(
-        &self,
-        data_store: &DS,
+    fn _search<'a, DS: DataStore>(
+        &'a self,
+        data_store: &'a DS,
         query: &[f32],
         filter: &Formula,
         with_duplicates: bool,
         results: usize,
         config: &VectorConfig,
         min_score: f32,
-    ) -> Box<dyn Iterator<Item = ScoredVector<'_>> + '_> {
+    ) -> Box<dyn Iterator<Item = ScoredVector<'a>> + '_> {
         let raw_query = SearchVector::Query(config.vector_type.encode(query));
         let rabitq = data_store.has_quantized() && !config.flags.contains(&flags::DISABLE_RABITQ_SEARCH.to_string());
         let encoded_query = if rabitq {
@@ -539,49 +539,50 @@ impl OpenSegment {
         }
     }
 
-    fn brute_force_search<DS: DataStore>(
-        &self,
+    fn brute_force_search<'a, DS: DataStore>(
+        &'a self,
         bitset: &FilterBitSet,
         results: usize,
-        retriever: Retriever<DS>,
+        retriever: Retriever<'a, DS>,
         encoded_query: &SearchVector,
         raw_query: &SearchVector,
-    ) -> Box<dyn Iterator<Item = ScoredVector<'_>> + '_> {
+    ) -> Box<dyn Iterator<Item = ScoredVector<'a>> + '_> {
         let mut scored_results = Vec::new();
         for paragraph_addr in bitset.iter() {
             let paragraph = retriever.data_store.get_paragraph(paragraph_addr);
 
-            // Only return the best vector match per paragraph
+            // Only return the best vector match per paragraph (only relevant with multi-vectors)
             let best_vector_score = paragraph
                 .vectors(&paragraph_addr)
                 .map(|va| {
                     let score = retriever.similarity_upper_bound(va, encoded_query);
-                    Cnx(va, score)
+                    (va, score)
                 })
-                .max_by(|v, w| v.1.total_cmp(&w.1))
+                .max_by(|v, w| v.1.score.total_cmp(&w.1.score))
                 .unwrap();
 
-            if best_vector_score.1 >= retriever.min_score {
-                scored_results.push(Reverse(best_vector_score));
+            if best_vector_score.1.upper_bound >= retriever.min_score {
+                scored_results.push(best_vector_score);
             }
         }
-        scored_results.sort_unstable();
-        // If using RabitQ, rerank top results using the raw vectors
+
         if matches!(encoded_query, SearchVector::RabitQ(_)) {
-            scored_results = scored_results
-                .into_iter()
-                .take(std::cmp::min(results * 100, 2000))
-                .map(|Reverse(Cnx(addr, _))| Reverse(Cnx(addr, retriever.similarity(addr, raw_query))))
-                .collect();
-            scored_results.sort_unstable();
+            // If using RabitQ, rerank top results using the raw vectors
+            Box::new(
+                rabitq::rerank_top(scored_results, results, &retriever, raw_query)
+                    .into_iter()
+                    .map(|Reverse(Cnx(addr, score))| ScoredVector::new(addr, self.data_store.as_ref(), score)),
+            )
+        } else {
+            // If using raw vectors, sort by score and take top_k
+            scored_results.sort_unstable_by(|a, b| b.1.score.total_cmp(&a.1.score));
+            Box::new(
+                scored_results
+                    .into_iter()
+                    .map(|a| ScoredVector::new(a.0, self.data_store.as_ref(), a.1.score))
+                    .take(results),
+            )
         }
-
-        Box::new(
-            scored_results
-                .into_iter()
-                .map(|Reverse(a)| ScoredVector::new(a.0, self.data_store.as_ref(), a.1))
-                .take(results),
-        )
     }
 }
 

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/nidx_vector/src/vector_types/rabitq.rs

@@ -18,12 +18,28 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 //
 
+use std::{cmp::Reverse, collections::BinaryHeap};
+
 use simsimd::SpatialSimilarity;
 
-use crate::config::VectorType;
+use crate::{
+    VectorAddr,
+    config::VectorType,
+    hnsw::{Cnx, DataRetriever, EstimatedScore, SearchVector},
+};
 
+/// Constant to adjust the error bound of the estimated similarity.
+// The paper recommends 1.9, reasonable values are from 0.0 to 4.0
+// Higher numbers give more error bound so we need to load more raw
+// vectors but provide slightly better recall.
 const EPSILON: f32 = 1.9;
 
+/// How many vectors to evaluate per each expected result.
+/// This is multiplied by top_k to get the number of vectors to search for.
+pub const RERANKING_FACTOR: usize = 100;
+/// The maximum number of vectors to evaluate during reranking.
+pub const RERANKING_LIMIT: usize = 2000;
+
 pub struct EncodedVector<'a> {
     data: &'a [u8],
 }
@@ -207,6 +223,31 @@ impl QueryVector {
     }
 }
 
+/// Rerank results from RabitQ search using the raw vectors. Uses the error bound to know when to stop reranking
+pub fn rerank_top(
+    candidates: Vec<(VectorAddr, EstimatedScore)>,
+    top_k: usize,
+    retriever: &impl DataRetriever,
+    query: &SearchVector,
+) -> BinaryHeap<Reverse<Cnx>> {
+    let mut best = BinaryHeap::new();
+    let mut best_k = 0.0;
+    for (addr, EstimatedScore { upper_bound, .. }) in candidates {
+        if best.len() < top_k || best_k < upper_bound {
+            // If the candidate score could be better than what we have so far, calculate the accurate similarity
+            let real_score = retriever.similarity(addr, query);
+            if best.len() < top_k || best_k < real_score {
+                best.push(Reverse(Cnx(addr, real_score)));
+                if best.len() > top_k {
+                    best.pop();
+                }
+                best_k = best.peek().unwrap().0.1;
+            }
+        }
+    }
+    best
+}
+
 #[cfg(test)]
 mod tests {
     use rand::{Rng, SeedableRng, rngs::SmallRng};

{nidx_binding-6.9.0.post567 → nidx_binding-6.9.0.post576}/pyproject.toml

@@ -13,7 +13,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
-version = "6.9.0.post567"
+version = "6.9.0.post576"
 
 [project.urls]
 Homepage = "https://nuclia.com"