nidx-binding 6.9.3.post630__tar.gz → 6.9.4.post638__tar.gz

{nidx_binding-6.9.3.post630 → nidx_binding-6.9.4.post638}/PKG-INFO

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

{nidx_binding-6.9.3.post630 → nidx_binding-6.9.4.post638}/nidx_protos/pyproject.toml

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

{nidx_binding-6.9.3.post630 → nidx_binding-6.9.4.post638}/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.post630 → nidx_binding-6.9.4.post638}/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";
 }
@@ -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.post630 → nidx_binding-6.9.4.post638}/pyproject.toml

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