nidx-binding 6.9.3.post626__tar.gz → 6.9.3.post636__tar.gz

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.9.3.post626
+Version: 6.9.3.post636
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/nidx_protos/pyproject.toml

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

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/nidx_vector/src/hnsw/ram_hnsw.rs

@@ -113,10 +113,31 @@ impl RAMHnsw {
     /// A bug in a previous version of this program could cause a node in layer N
     /// to link to a node in layer N-1. This breaks navigation accross layer N.
     /// This function will delete any such link from the graph.
-    pub fn fix_broken_links(&self) {
+    /// Also delete empty layers and entrypoints pointing to unexisting layers
+    pub fn fix_broken_graph(&mut self) {
+        // Fix links to a node not in this lauer
         for l in &self.layers[1..] {
             l.fix_broken_links();
         }
+
+        // Delete empty layers
+        while let Some(layer) = self.layers.last() {
+            if layer.out.is_empty() {
+                self.layers.pop();
+            } else {
+                break;
+            }
+        }
+
+        // If entrypoint point to non-existing layer, point it to the top-most layer
+        if self.entry_point.layer >= self.layers.len() {
+            self.entry_point.layer = self.layers.len() - 1;
+            let last_layer = self.layers.last().unwrap();
+            if !last_layer.contains(&self.entry_point.node) {
+                // If the current entrypoint node is not in the last layer, point to another node that is here
+                self.entry_point.node = *self.layers.last().unwrap().out.keys().next().unwrap();
+            }
+        }
     }
 }
 
@@ -172,7 +193,7 @@ mod tests {
         let mut graph = RAMHnsw::new();
         graph.layers.push(layer0);
         graph.layers.push(layer1);
-        graph.fix_broken_links();
+        graph.fix_broken_graph();
         assert!(graph.layers[1].out[&VectorAddr(0)].read().unwrap().is_empty());
     }
 }

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/nidx_vector/src/hnsw/search.rs

@@ -21,7 +21,7 @@
 use bit_set::BitSet;
 use rustc_hash::FxHashSet;
 use std::cmp::{Ordering, Reverse};
-use std::collections::{BinaryHeap, HashMap, VecDeque};
+use std::collections::{BinaryHeap, HashMap};
 use std::time::Instant;
 use tracing::trace;
 
@@ -190,33 +190,29 @@ impl<'a, DR: DataRetriever> HnswSearcher<'a, DR> {
         }
     }
 
+    /// Breadth-first search to find the closest nodes to the entry-points that fulfill the filtering conditions
     fn closest_up_nodes<L: SearchableLayer>(
         &'a self,
-        entry_points: Vec<VectorAddr>,
+        entry_points: Vec<Cnx>,
         query: &SearchVector,
         layer: L,
         number_of_results: usize,
         mut filter: NodeFilter<'a>,
     ) -> Vec<(VectorAddr, f32)> {
-        // We just need to perform BFS, the replacement is the closest node to the actual
-        // best solution. This algorithm takes a lazy approach to computing the similarity of
-        // candidates.
-
         const MAX_VECTORS_TO_PRELOAD: u32 = 20_000;
         let mut results = Vec::new();
-        let inner_entry_points_iter = entry_points.iter().map(|VectorAddr(inner)| *inner as usize);
+        let inner_entry_points_iter = entry_points.iter().map(|Cnx(VectorAddr(inner), _)| *inner as usize);
         let mut visited_nodes: BitSet = BitSet::from_iter(inner_entry_points_iter);
-        let mut candidates = VecDeque::from(entry_points);
+        let mut candidates = entry_points;
+        candidates.sort_unstable();
 
         let mut preloaded = 0;
 
         loop {
-            let Some(candidate) = candidates.pop_front() else {
+            let Some(Cnx(candidate, candidate_similarity)) = candidates.pop() else {
                 break;
             };
 
-            let candidate_similarity = self.retriever.similarity(candidate, query);
-
             if candidate_similarity < self.retriever.min_score() {
                 break;
             }
@@ -226,22 +222,28 @@ impl<'a, DR: DataRetriever> HnswSearcher<'a, DR> {
             }
 
             if results.len() == number_of_results {
-                return results;
+                break;
             }
 
-            let mut sorted_out: Vec<_> = layer.get_out_edges(candidate).collect();
-            sorted_out.sort_by(|a, b| b.1.total_cmp(&a.1));
-            sorted_out.into_iter().for_each(|(new_candidate, _)| {
-                if !visited_nodes.contains(new_candidate.0 as usize) {
-                    visited_nodes.insert(new_candidate.0 as usize);
-                    candidates.push_back(new_candidate);
-
-                    if self.preload_nodes && preloaded < MAX_VECTORS_TO_PRELOAD {
+            if self.preload_nodes && preloaded < MAX_VECTORS_TO_PRELOAD {
+                for (new_candidate, _) in layer.get_out_edges(candidate) {
+                    if !visited_nodes.contains(new_candidate.0 as usize) {
                         self.retriever.will_need(new_candidate);
                         preloaded += 1;
                     }
                 }
-            });
+            }
+
+            for (new_candidate, _) in layer.get_out_edges(candidate) {
+                if visited_nodes.insert(new_candidate.0 as usize) {
+                    let new_similarity = self.retriever.similarity(new_candidate, query);
+
+                    if new_similarity > self.retriever.min_score() {
+                        candidates.push(Cnx(new_candidate, new_similarity));
+                    }
+                }
+            }
+            candidates.sort_unstable();
         }
 
         results
@@ -364,13 +366,13 @@ impl<'a, DR: DataRetriever> HnswSearcher<'a, DR> {
             let t = Instant::now();
             let reranked = rabitq::rerank_top(neighbours.collect(), k_neighbours, self.retriever, query)
                 .into_iter()
-                .map(|Reverse(Cnx(addr, _))| addr)
+                .map(|Reverse(c)| c)
                 .collect();
             let time = t.elapsed();
             trace!(?time, "HNSW search: reranking");
             reranked
         } else {
-            neighbours.map(|(addr, _)| addr).collect()
+            neighbours.map(|(addr, score)| Cnx(addr, score.score)).collect()
         };
 
         // Find k nodes that match the filter in the last layer

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/nidx_vector/src/hnsw.rs

@@ -26,5 +26,6 @@ mod search;
 
 pub use build::HnswBuilder;
 pub use disk_hnsw::DiskHnsw;
+pub use params::M;
 pub use ram_hnsw::RAMHnsw;
 pub use search::{Cnx, DataRetriever, EstimatedScore, HnswSearcher, NodeFilter, SearchVector};

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/nidx_vector/src/segment.rs

@@ -26,7 +26,7 @@ use crate::data_store::{DataStore, DataStoreV1, DataStoreV2, OpenReason, Paragra
 use crate::formula::Formula;
 use crate::inverted_index::{FilterBitSet, InvertedIndexes, build_indexes};
 use crate::vector_types::rabitq;
-use crate::{ParagraphAddr, VectorAddr, VectorErr, VectorR, VectorSegmentMeta, VectorSegmentMetadata};
+use crate::{ParagraphAddr, VectorAddr, VectorErr, VectorR, VectorSegmentMeta, VectorSegmentMetadata, hnsw};
 use crate::{hnsw::*, inverted_index};
 use core::f32;
 use io::{BufWriter, Write};
@@ -42,9 +42,6 @@ use std::path::Path;
 use std::time::Instant;
 use tracing::{debug, trace};
 
-/// How much expensive is to find a node via HNSW compared to a simple brute force scan
-const HNSW_COST_FACTOR: usize = 100;
-
 mod file_names {
     pub const HNSW: &str = "index.hnsw";
 }
@@ -158,7 +155,7 @@ fn merge_indexes<DS: DataStore + 'static>(
         // If there are no deletions, we can reuse the first segment
         // HNSW since its indexes will match the the ones in data_store
         index = DiskHnsw::deserialize(&operants[0].index);
-        index.fix_broken_links();
+        index.fix_broken_graph();
         start_vector_index = operants[0].data_store.stored_vector_count();
     }
     let merged_vectors_count = data_store.stored_vector_count();
@@ -543,14 +540,10 @@ impl OpenSegment {
         if matching == 0 {
             return Box::new(empty());
         }
-        let expected_traversal_scan = top_k * self.metadata.records / matching;
 
         let t = Instant::now();
         let method;
-        let results = if matching < expected_traversal_scan * HNSW_COST_FACTOR {
-            method = "brute force";
-            self.brute_force_search(bitset, top_k, retriever, encoded_query, &raw_query)
-        } else {
+        let results = if use_hnsw(self.metadata.records, matching, top_k, rabitq) {
             method = "hnsw";
             let ops = HnswSearcher::new(&retriever, true);
             let filter = NodeFilter::new(bitset, with_duplicates, config);
@@ -561,6 +554,9 @@ impl OpenSegment {
                     .map(|(address, dist)| ScoredVector::new(address, self.data_store.as_ref(), dist))
                     .take(top_k),
             )
+        } else {
+            method = "brute force";
+            self.brute_force_search(bitset, top_k, retriever, encoded_query, &raw_query)
         };
 
         let time = t.elapsed();
@@ -627,6 +623,42 @@ impl OpenSegment {
     }
 }
 
+fn use_hnsw(total_nodes: usize, matching_nodes: usize, top_k: usize, has_rabitq: bool) -> bool {
+    // Cost of a full vector comparison compared to a quantized vector comparison
+    let full_cost: usize;
+    // How many more vectors are evaluated in layer_search for quantized vectors
+    let search_mult: usize;
+    // How many vectors are reranked
+    let rerank_mult: usize;
+    if has_rabitq {
+        full_cost = 16;
+        search_mult = rabitq::RERANKING_FACTOR * 3 / 4;
+        rerank_mult = rabitq::RERANKING_FACTOR / 2;
+    } else {
+        full_cost = 1;
+        search_mult = 1;
+        rerank_mult = 0;
+    }
+
+    // Estimated vectors visited during hnsw search (quantized vectors)
+    let hnsw_rq = ((total_nodes as f32).ln() - 2.0).powi(2) * (top_k as f32).ln() * search_mult as f32;
+    // Estimated vectors visited in closest_up_nodes (full vectors)
+    let hnsw_full = (top_k * rerank_mult) + (top_k * hnsw::M * total_nodes / matching_nodes);
+
+    // Quantized vectors evaluated with brute-force
+    let bf_rq = matching_nodes;
+    // Estimated full vectors to be evaluated during rerank in brute-force
+    let bf_full = top_k * rerank_mult;
+
+    let hnsw_cost = hnsw_rq as usize + hnsw_full * full_cost;
+    let bf_cost = bf_rq + bf_full * full_cost;
+
+    let use_hnsw = hnsw_cost < bf_cost;
+    debug!(hnsw_cost, bf_cost, use_hnsw, "Estimated search costs");
+
+    use_hnsw
+}
+
 #[cfg(test)]
 mod test {
     use std::collections::{BTreeMap, HashSet};

{nidx_binding-6.9.3.post626 → nidx_binding-6.9.3.post636}/pyproject.toml

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