nidx-binding 6.9.1.post600__tar.gz → 6.9.1.post616__tar.gz

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.9.1.post600
+Version: 6.9.1.post616
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_protos/nidx.proto

@@ -7,10 +7,10 @@ import "nidx_protos/nodewriter.proto";
 
 service NidxApi {
   rpc GetShard(nodereader.GetShardRequest) returns (noderesources.Shard) {}
-
   rpc NewShard(nodewriter.NewShardRequest) returns (noderesources.ShardCreated) {}
   rpc DeleteShard(noderesources.ShardId) returns (noderesources.ShardId) {}
   rpc ListShards(noderesources.EmptyQuery) returns (noderesources.ShardIds) {}
+  rpc ConfigureShards(ShardsConfig) returns (noderesources.EmptyQuery) {}
 
   rpc AddVectorSet(nodewriter.NewVectorSetRequest) returns (nodewriter.OpStatus) {}
   rpc RemoveVectorSet(noderesources.VectorSetID) returns (nodewriter.OpStatus) {}
@@ -43,3 +43,17 @@ message Notification {
     }
     Action action = 6;
 }
+
+
+message ShardsConfig {
+    repeated ShardConfig configs = 1;
+}
+
+message ShardConfig {
+    string shard_id = 1;
+
+    // Pre-warm is a feature to minimize cold starts on searches. It'll keep
+    // disk data in memory to avoid IO latency, at expenses of RAM usage. Use
+    // with care.
+    bool prewarm_enabled = 2;
+}

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_protos/nodewriter.proto

@@ -64,6 +64,7 @@ message NewShardRequest {
     bool normalize_vectors = 4 [deprecated = true];
     VectorIndexConfig config = 5 [deprecated = true];
     map<string, VectorIndexConfig> vectorsets_configs = 6;
+    bool prewarm_enabled = 7;
 
 }
 

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_protos/pyproject.toml

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

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_protos/src/lib.rs

@@ -42,6 +42,7 @@ pub mod kb_usage {
     tonic::include_proto!("kb_usage");
 }
 
+pub use nidx::*;
 pub use nodereader::*;
 pub use noderesources::*;
 pub use nodewriter::*;

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/data_store/v1.rs

@@ -98,7 +98,7 @@ impl DataStoreV1 {
         {
             let advice = match reason {
                 OpenReason::Create => memmap2::Advice::Sequential,
-                OpenReason::Search => memmap2::Advice::Random,
+                OpenReason::Search { .. } => memmap2::Advice::Random,
             };
             nodes.advise(advice)?;
         }

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/data_store/v2/paragraph_store.rs

@@ -90,7 +90,7 @@ impl ParagraphStore {
             pos.advise(memmap2::Advice::WillNeed)?;
             let advice = match reason {
                 OpenReason::Create => memmap2::Advice::Sequential,
-                OpenReason::Search => memmap2::Advice::Random,
+                OpenReason::Search { .. } => memmap2::Advice::Random,
             };
             data.advise(advice)?;
         }

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/data_store/v2/quant_vector_store.rs

@@ -19,7 +19,7 @@
 //
 
 use lazy_static::lazy_static;
-use memmap2::Mmap;
+use memmap2::{Mmap, MmapOptions};
 use std::{
     fs::File,
     io::{BufWriter, Write as _},
@@ -41,13 +41,17 @@ pub struct QuantVectorStore {
 
 impl QuantVectorStore {
     pub fn open(path: &Path, vector_len_bytes: usize, reason: &OpenReason) -> std::io::Result<Self> {
-        let data = unsafe { Mmap::map(&File::open(path.join(FILENAME))?)? };
+        let mut options = MmapOptions::new();
+        if matches!(reason, OpenReason::Search { prewarm: true }) {
+            options.populate();
+        }
+        let data = unsafe { options.map(&File::open(path.join(FILENAME))?)? };
 
         #[cfg(not(target_os = "windows"))]
         {
             let advice = match reason {
                 OpenReason::Create => memmap2::Advice::Sequential,
-                OpenReason::Search => memmap2::Advice::Random,
+                OpenReason::Search { .. } => memmap2::Advice::Random,
             };
             data.advise(advice)?;
         }

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/data_store/v2/vector_store.rs

@@ -60,7 +60,7 @@ impl VectorStore {
         {
             let advice = match reason {
                 OpenReason::Create => memmap2::Advice::Sequential,
-                OpenReason::Search => memmap2::Advice::Random,
+                OpenReason::Search { .. } => memmap2::Advice::Random,
             };
             data.advise(advice)?;
         }

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/data_store.rs

@@ -31,7 +31,7 @@ use v2::StoredParagraph;
 use crate::{ParagraphAddr, VectorAddr, vector_types::rabitq};
 
 pub enum OpenReason {
-    Search,
+    Search { prewarm: bool },
     Create,
 }
 

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/hnsw/ops_hnsw.rs

@@ -29,6 +29,8 @@ use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque};
 use std::time::Instant;
 use tracing::trace;
 
+use crate::config::{VectorCardinality, VectorConfig};
+use crate::hnsw::params::EF_SEARCH;
 use crate::inverted_index::FilterBitSet;
 use crate::vector_types::rabitq;
 use crate::{ParagraphAddr, VectorAddr};
@@ -136,31 +138,48 @@ impl PartialOrd for CnxWithBound {
 pub type Neighbours = Vec<(VectorAddr, f32)>;
 
 /// Guides an algorithm to the valid nodes.
-struct NodeFilter<'a, DR> {
-    retriever: &'a DR,
+pub struct NodeFilter<'a> {
     filter: &'a FilterBitSet,
-    paragraphs: FxHashSet<ParagraphAddr>,
+    paragraphs: Option<FxHashSet<ParagraphAddr>>,
     vec_counter: RepCounter<'a>,
 }
 
-impl<DR: DataRetriever> NodeFilter<'_, DR> {
-    pub fn passes_formula(&self, n: ParagraphAddr) -> bool {
-        self.filter.contains(n)
+impl<'a> NodeFilter<'a> {
+    pub fn new(filter: &'a FilterBitSet, with_duplicates: bool, config: &VectorConfig) -> NodeFilter<'a> {
+        let paragraphs = matches!(config.vector_cardinality, VectorCardinality::Multi).then_some(Default::default());
+        let vec_counter = RepCounter::new(!with_duplicates);
+        NodeFilter {
+            filter,
+            paragraphs,
+            vec_counter,
+        }
     }
 
-    pub fn is_valid(&self, n: VectorAddr, score: f32) -> bool {
-        !score.is_nan()
-        // Reject the candidate if we already have a result for the same paragraph
-        && !self.paragraphs.contains(&self.retriever.paragraph(n))
-        // Reject the candidate if we already have a result with an identical vector
-        && self.vec_counter.get(self.retriever.get_vector(n)) == 0
-    }
+    pub fn passes(&mut self, retriever: &'a impl DataRetriever, v: VectorAddr) -> bool {
+        let vector = retriever.get_vector(v);
+        let paragraph = retriever.paragraph(v);
 
-    /// Adds a result so that further candidates with the same vector
-    /// or paragraph will get rejected.
-    pub fn add_result(&mut self, n: VectorAddr) {
-        self.paragraphs.insert(self.retriever.paragraph(n));
-        self.vec_counter.add(self.retriever.get_vector(n));
+        // Matches query filters?
+        if !self.filter.contains(paragraph) {
+            return false;
+        }
+
+        // Is a duplicated vector?
+        if self.vec_counter.get(vector) > 0 {
+            return false;
+        };
+
+        // Is a vector for an existing paragraph? (only relevant in multi-vector search)
+        if let Some(paragraphs) = &mut self.paragraphs
+            && !paragraphs.insert(paragraph)
+        {
+            return false;
+        }
+
+        // Track the vector for duplicates (the paragraphs was tracked in the above check)
+        self.vec_counter.add(vector);
+
+        true
     }
 }
 
@@ -228,7 +247,7 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         query: &SearchVector,
         layer: L,
         number_of_results: usize,
-        mut filter: NodeFilter<'a, DR>,
+        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
@@ -253,9 +272,7 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
                 break;
             }
 
-            let paragraph_addr = self.retriever.paragraph(candidate);
-            if filter.is_valid(candidate, candidate_similarity) && filter.passes_formula(paragraph_addr) {
-                filter.add_result(candidate);
+            if !candidate_similarity.is_nan() && filter.passes(self.retriever, candidate) {
                 results.push((candidate, candidate_similarity));
             }
 
@@ -338,7 +355,10 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
                 _ => (),
             }
         }
-        ms_neighbours.into_iter().map(|Reverse(CnxWithBound(n, d))| (n, d))
+        ms_neighbours
+            .into_sorted_vec()
+            .into_iter()
+            .map(|Reverse(CnxWithBound(n, d))| (n, d))
     }
 
     fn layer_insert(
@@ -408,8 +428,7 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         query: &SearchVector,
         hnsw: H,
         k_neighbours: usize,
-        with_filter: &FilterBitSet,
-        with_duplicates: bool,
+        filter: NodeFilter,
     ) -> Neighbours {
         if k_neighbours == 0 {
             return Neighbours::with_capacity(0);
@@ -442,7 +461,10 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
         let last_layer_k = if original_query.is_some() {
             std::cmp::min(k_neighbours * rabitq::RERANKING_FACTOR, rabitq::RERANKING_LIMIT)
         } else {
-            k_neighbours
+            // We search for at least EF_SEARCH neighbours to guarantee good results
+            // In the case k > EF we could use EF search and expand the search with
+            // closest_up_nodes()
+            std::cmp::max(k_neighbours, EF_SEARCH)
         };
 
         // Find the best k nodes in the last layer
@@ -466,20 +488,12 @@ impl<'a, DR: DataRetriever> HnswOps<'a, DR> {
             neighbours.map(|(addr, _)| addr).collect()
         };
 
-        let filter = NodeFilter {
-            filter: with_filter,
-            retriever: self.retriever,
-            paragraphs: Default::default(),
-            vec_counter: RepCounter::new(!with_duplicates),
-        };
-        let layer_zero = hnsw.get_layer(0);
-
         // Find k nodes that match the filter in the last layer
         let t = Instant::now();
         let mut filtered_result = self.closest_up_nodes(
             entry_points,
             original_query.unwrap_or(query),
-            layer_zero,
+            hnsw.get_layer(0),
             k_neighbours,
             filter,
         );

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/hnsw/params.rs

@@ -44,5 +44,8 @@ pub const M_MAX: usize = 30;
 /// Number of bi-directional links created for every new element.
 pub const M: usize = 30;
 
-/// Number of neighbours that are searched for before adding a new embedding.
+/// Number of neighbours that are explored when searching for the insertion place of a new node
 pub const EF_CONSTRUCTION: usize = 100;
+
+/// Number of neighbours that are explored when searching for a query
+pub const EF_SEARCH: usize = 30;

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/hnsw.rs

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

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/inverted_index/fst_index.rs

@@ -20,7 +20,7 @@
 use std::{collections::HashSet, fs::File, io::BufWriter, path::Path, sync::Arc};
 
 use fst::{Automaton, IntoStreamer, Map, Streamer};
-use memmap2::Mmap;
+use memmap2::{Mmap, MmapOptions};
 
 use crate::VectorR;
 
@@ -59,9 +59,14 @@ pub struct FstIndexReader {
 }
 
 impl FstIndexReader {
-    pub fn open(path: &Path, map_reader: Arc<InvertedMapReader>) -> VectorR<Self> {
+    pub fn open(path: &Path, map_reader: Arc<InvertedMapReader>, prewarm: bool) -> VectorR<Self> {
+        let mut options = MmapOptions::new();
+        if prewarm {
+            options.populate();
+        }
+
         Ok(Self {
-            fst: Map::new(unsafe { Mmap::map(&File::open(path)?)? })?,
+            fst: Map::new(unsafe { options.map(&File::open(path)?)? })?,
             map_reader,
         })
     }
@@ -132,8 +137,8 @@ mod tests {
         map_writer.finish()?;
 
         // Check the map has the same contents we initialized
-        let map_reader = Arc::new(InvertedMapReader::open(tmp_map.path())?);
-        let reader = FstIndexReader::open(tmp.path(), map_reader)?;
+        let map_reader = Arc::new(InvertedMapReader::open(tmp_map.path(), false)?);
+        let reader = FstIndexReader::open(tmp.path(), map_reader, false)?;
 
         for (key, value) in entries {
             let indexed = reader.get(&key);

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/inverted_index/map.rs

@@ -23,7 +23,7 @@ use std::{
     path::Path,
 };
 
-use memmap2::Mmap;
+use memmap2::{Mmap, MmapOptions};
 use stream_vbyte::{decode::decode, encode::encode, scalar::Scalar};
 
 use crate::VectorR;
@@ -67,9 +67,14 @@ pub struct InvertedMapReader {
 }
 
 impl InvertedMapReader {
-    pub fn open(path: &Path) -> VectorR<Self> {
+    pub fn open(path: &Path, prewarm: bool) -> VectorR<Self> {
+        let mut options = MmapOptions::new();
+        if prewarm {
+            options.populate();
+        }
+
         Ok(Self {
-            data: unsafe { Mmap::map(&File::open(path)?)? },
+            data: unsafe { options.map(&File::open(path)?)? },
         })
     }
 
@@ -120,7 +125,7 @@ mod tests {
         writer.finish()?;
 
         // Check the map has the same contents we initialized
-        let reader = InvertedMapReader::open(tmp.path())?;
+        let reader = InvertedMapReader::open(tmp.path(), false)?;
         for i in 0..20 {
             let indexed = reader.get(indexes[i]);
             assert_eq!(indexed, entries[i]);

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/inverted_index.rs

@@ -151,6 +151,10 @@ pub fn build_indexes(work_path: &Path, data_store: &impl DataStore) -> VectorR<(
     Ok(())
 }
 
+pub struct OpenOptions {
+    pub prewarm: bool,
+}
+
 pub struct InvertedIndexes {
     field_index: FstIndexReader,
     label_index: FstIndexReader,
@@ -166,10 +170,13 @@ impl InvertedIndexes {
         path.join(file::INDEX_MAP).exists()
     }
 
-    pub fn open(work_path: &Path, records: usize) -> VectorR<Self> {
-        let map = Arc::new(InvertedMapReader::open(&work_path.join(file::INDEX_MAP))?);
-        let field_index = FstIndexReader::open(&work_path.join(file::FIELD_INDEX), map.clone())?;
-        let label_index = FstIndexReader::open(&work_path.join(file::LABEL_INDEX), map)?;
+    pub fn open(work_path: &Path, records: usize, options: OpenOptions) -> VectorR<Self> {
+        let map = Arc::new(InvertedMapReader::open(
+            &work_path.join(file::INDEX_MAP),
+            options.prewarm,
+        )?);
+        let field_index = FstIndexReader::open(&work_path.join(file::FIELD_INDEX), map.clone(), options.prewarm)?;
+        let label_index = FstIndexReader::open(&work_path.join(file::LABEL_INDEX), map, options.prewarm)?;
 
         Ok(Self {
             field_index,

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/nidx_vector/src/segment.rs

@@ -24,13 +24,13 @@ mod tests;
 use crate::config::{VectorConfig, flags};
 use crate::data_store::{DataStore, DataStoreV1, DataStoreV2, OpenReason, ParagraphRef, VectorRef};
 use crate::formula::Formula;
-use crate::hnsw::*;
 use crate::inverted_index::{FilterBitSet, InvertedIndexes, build_indexes};
 use crate::vector_types::rabitq;
 use crate::{ParagraphAddr, VectorAddr, VectorErr, VectorR, VectorSegmentMeta, VectorSegmentMetadata};
+use crate::{hnsw::*, inverted_index};
 use core::f32;
 use io::{BufWriter, Write};
-use memmap2::Mmap;
+use memmap2::{Mmap, MmapOptions};
 use std::cmp::Reverse;
 use std::collections::HashSet;
 use std::fs::File;
@@ -48,9 +48,12 @@ mod file_names {
 }
 
 pub fn open(metadata: VectorSegmentMetadata, config: &VectorConfig) -> VectorR<OpenSegment> {
+    // TODO: we should get this flag from the VectorConfig or some other place
+    let prewarm = false;
+
     let path = &metadata.path;
     let data_store: Box<dyn DataStore> = if DataStoreV1::exists(path)? {
-        let data_store = DataStoreV1::open(path, &config.vector_type, OpenReason::Search)?;
+        let data_store = DataStoreV1::open(path, &config.vector_type, OpenReason::Search { prewarm })?;
         // Build the index at runtime if they do not exist. This can
         // be removed once we have migrated all existing indexes
         if !InvertedIndexes::exists(path) {
@@ -58,7 +61,7 @@ pub fn open(metadata: VectorSegmentMetadata, config: &VectorConfig) -> VectorR<O
         }
         Box::new(data_store)
     } else {
-        let data_store = DataStoreV2::open(path, &config.vector_type, OpenReason::Search)?;
+        let data_store = DataStoreV2::open(path, &config.vector_type, OpenReason::Search { prewarm })?;
         // Build the index at runtime if they do not exist. This can
         // be removed once we have migrated all existing indexes
         if !InvertedIndexes::exists(path) {
@@ -66,9 +69,13 @@ pub fn open(metadata: VectorSegmentMetadata, config: &VectorConfig) -> VectorR<O
         }
         Box::new(data_store)
     };
-    let hnsw_file = File::open(path.join(file_names::HNSW))?;
 
-    let index = unsafe { Mmap::map(&hnsw_file)? };
+    let hnsw_file = File::open(path.join(file_names::HNSW))?;
+    let mut index_options = MmapOptions::new();
+    if prewarm {
+        index_options.populate();
+    }
+    let index = unsafe { index_options.map(&hnsw_file)? };
 
     // Telling the OS our expected access pattern
     #[cfg(not(target_os = "windows"))]
@@ -77,7 +84,7 @@ pub fn open(metadata: VectorSegmentMetadata, config: &VectorConfig) -> VectorR<O
         index.advise(memmap2::Advice::WillNeed)?;
     }
 
-    let inverted_indexes = InvertedIndexes::open(path, metadata.records)?;
+    let inverted_indexes = InvertedIndexes::open(path, metadata.records, inverted_index::OpenOptions { prewarm })?;
     let alive_bitset = FilterBitSet::new(metadata.records, true);
 
     Ok(OpenSegment {
@@ -189,7 +196,11 @@ fn merge_indexes<DS: DataStore + 'static>(
     };
 
     build_indexes(segment_path, &data_store)?;
-    let inverted_indexes = InvertedIndexes::open(segment_path, merged_vectors_count as usize)?;
+    let inverted_indexes = InvertedIndexes::open(
+        segment_path,
+        merged_vectors_count as usize,
+        inverted_index::OpenOptions { prewarm: false },
+    )?;
     let alive_bitset = FilterBitSet::new(metadata.records, true);
 
     Ok(OpenSegment {
@@ -289,7 +300,11 @@ fn create_indexes<DS: DataStore + 'static>(
     };
 
     build_indexes(path, &data_store)?;
-    let inverted_indexes = InvertedIndexes::open(path, vector_count as usize)?;
+    let inverted_indexes = InvertedIndexes::open(
+        path,
+        vector_count as usize,
+        inverted_index::OpenOptions { prewarm: false },
+    )?;
     let alive_bitset = FilterBitSet::new(metadata.records, true);
 
     Ok(OpenSegment {
@@ -533,7 +548,8 @@ impl OpenSegment {
         } else {
             method = "hnsw";
             let ops = HnswOps::new(&retriever, true);
-            let neighbours = ops.search(encoded_query, self.index.as_ref(), top_k, bitset, with_duplicates);
+            let filter = NodeFilter::new(bitset, with_duplicates, config);
+            let neighbours = ops.search(encoded_query, self.index.as_ref(), top_k, filter);
             Box::new(
                 neighbours
                     .into_iter()

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/pyproject.toml

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

{nidx_binding-6.9.1.post600 → nidx_binding-6.9.1.post616}/src/api/grpc.rs

@@ -23,8 +23,6 @@ use std::io::Write;
 use std::str::FromStr;
 use std::sync::Arc;
 
-use crate::errors::NidxError;
-use crate::metadata::{Index, IndexId, IndexKind, Shard};
 use axum::body::{Body, Bytes};
 use axum::extract::{Path, State};
 use axum::http::StatusCode;
@@ -41,6 +39,8 @@ use tracing::error;
 use uuid::Uuid;
 
 use crate::api::shards;
+use crate::errors::NidxError;
+use crate::metadata::{Index, IndexId, IndexKind, Shard};
 use crate::{NidxMetadata, Settings, import_export};
 
 #[derive(Clone)]
@@ -139,6 +139,10 @@ impl NidxApi for ApiServer {
         }))
     }
 
+    async fn configure_shards(&self, _request: Request<ShardsConfig>) -> Result<Response<EmptyQuery>> {
+        Ok(Response::new(EmptyQuery {}))
+    }
+
     async fn add_vector_set(&self, request: Request<NewVectorSetRequest>) -> Result<Response<OpStatus>> {
         let request = request.into_inner();
         let Some(VectorSetId {