nidx-binding 6.9.1.post596__tar.gz → 6.9.1.post603__tar.gz

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/PKG-INFO

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

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/nidx_protos/pyproject.toml

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

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/nidx_vector/src/data_store/v2/paragraph_store.rs

@@ -19,7 +19,11 @@
 //
 
 use memmap2::Mmap;
-use std::{fs::File, io::Write as _, path::Path};
+use std::{
+    fs::File,
+    io::{BufWriter, Write as _},
+    path::Path,
+};
 
 use crate::{
     VectorR,
@@ -86,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)?;
         }
@@ -112,8 +116,8 @@ impl ParagraphStore {
 }
 
 pub struct ParagraphStoreWriter {
-    data: File,
-    pos: File,
+    data: BufWriter<File>,
+    pos: BufWriter<File>,
     data_pos: u32,
     addr: u32,
 }
@@ -121,8 +125,8 @@ pub struct ParagraphStoreWriter {
 impl ParagraphStoreWriter {
     pub fn new(path: &Path) -> std::io::Result<Self> {
         Ok(Self {
-            pos: File::create(path.join(FILENAME_POS))?,
-            data: File::create(path.join(FILENAME_DATA))?,
+            pos: BufWriter::new(File::create(path.join(FILENAME_POS))?),
+            data: BufWriter::new(File::create(path.join(FILENAME_DATA))?),
             data_pos: 0,
             addr: 0,
         })
@@ -159,9 +163,9 @@ impl ParagraphStoreWriter {
         Ok(self.addr)
     }
 
-    pub fn close(self) -> std::io::Result<()> {
-        self.data.sync_all()?;
-        self.pos.sync_all()?;
+    pub fn close(&mut self) -> std::io::Result<()> {
+        self.data.flush()?;
+        self.pos.flush()?;
         Ok(())
     }
 }

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/nidx_vector/src/data_store/v2/quant_vector_store.rs

@@ -19,8 +19,12 @@
 //
 
 use lazy_static::lazy_static;
-use memmap2::Mmap;
-use std::{fs::File, io::Write as _, path::Path};
+use memmap2::{Mmap, MmapOptions};
+use std::{
+    fs::File,
+    io::{BufWriter, Write as _},
+    path::Path,
+};
 
 use crate::{
     data_store::{OpenReason, VectorAddr},
@@ -37,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)?;
         }
@@ -86,13 +94,13 @@ impl QuantVectorStore {
 }
 
 pub struct QuantVectorStoreWriter {
-    output: File,
+    output: BufWriter<File>,
 }
 
 impl QuantVectorStoreWriter {
     pub fn new(path: &Path) -> std::io::Result<Self> {
         Ok(Self {
-            output: File::create(path.join(FILENAME))?,
+            output: BufWriter::new(File::create(path.join(FILENAME))?),
         })
     }
 
@@ -100,7 +108,7 @@ impl QuantVectorStoreWriter {
         self.output.write_all(vector)
     }
 
-    pub fn close(self) -> std::io::Result<()> {
-        self.output.sync_all()
+    pub fn close(&mut self) -> std::io::Result<()> {
+        self.output.flush()
     }
 }

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/nidx_vector/src/data_store/v2/vector_store.rs

@@ -22,7 +22,7 @@ use lazy_static::lazy_static;
 use memmap2::Mmap;
 use std::{
     fs::File,
-    io::{Seek, SeekFrom, Write as _},
+    io::{BufWriter, Seek, SeekFrom, Write as _},
     path::Path,
 };
 
@@ -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)?;
         }
@@ -119,7 +119,7 @@ impl VectorStore {
 }
 
 pub struct VectorStoreWriter {
-    output: File,
+    output: BufWriter<File>,
     addr: u32,
     padding_bytes: usize,
 }
@@ -127,7 +127,7 @@ pub struct VectorStoreWriter {
 impl VectorStoreWriter {
     pub fn new(path: &Path, vector_type: &VectorType) -> std::io::Result<Self> {
         Ok(Self {
-            output: File::create(path.join(FILENAME))?,
+            output: BufWriter::new(File::create(path.join(FILENAME))?),
             addr: 0,
             padding_bytes: padding_bytes(vector_type),
         })
@@ -151,7 +151,7 @@ impl VectorStoreWriter {
         Ok((first_addr, last_addr))
     }
 
-    pub fn close(self) -> std::io::Result<()> {
-        self.output.sync_all()
+    pub fn close(mut self) -> std::io::Result<()> {
+        self.output.flush()
     }
 }

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/nidx_vector/src/data_store/v2.rs

@@ -77,7 +77,7 @@ impl DataStoreV2 {
 
         paragraphs.close()?;
         vectors.close()?;
-        if let Some(quantized) = quantized {
+        if let Some(mut quantized) = quantized {
             quantized.close()?;
         }
 

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/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.post596 → nidx_binding-6.9.1.post603}/nidx_vector/src/searcher.rs

@@ -156,10 +156,11 @@ impl<'a> Fssc<'a> {
             buff: HashSet::with_capacity(size),
         }
     }
-    fn add(&mut self, candidate: ScoredParagraph<'a>, vector: Vec<u8>) {
-        if !self.with_duplicates && self.seen.contains(&vector) {
+    fn add(&mut self, candidate: ScoredParagraph<'a>, vector: &[u8]) {
+        if !self.with_duplicates && self.seen.contains(vector) {
             return;
         } else if !self.with_duplicates {
+            let vector = vector.to_vec();
             self.seen.insert(vector);
         }
 
@@ -260,37 +261,23 @@ impl Searcher {
         let min_score = request.min_score();
         let mut ffsv = Fssc::new(request.no_results(), with_duplicates);
 
-        let pool = rayon::ThreadPoolBuilder::new().build().unwrap();
-        let results: Vec<_> = pool.install(|| {
-            self.open_segments
-                .par_iter()
-                .flat_map(|open_segment| {
-                    // Skip this segment if it doesn't match the segment filter
-                    if !segment_filter
-                        .as_ref()
-                        .is_none_or(|f| segment_matches(f, open_segment.tags()))
-                    {
-                        return vec![];
-                    }
-                    let partial_solution =
-                        open_segment.search(query, filter, with_duplicates, no_results, &self.config, min_score);
-
-                    partial_solution
-                        .map(|candidate| {
-                            let addr = candidate.paragraph();
-                            let paragraph = open_segment.get_paragraph(addr);
-                            (
-                                ScoredParagraph::new(open_segment, addr, paragraph, candidate.score()),
-                                candidate.vector().to_vec(),
-                            )
-                        })
-                        .collect()
-                })
-                .collect()
-        });
-
-        for (paragraph, vector) in results {
-            ffsv.add(paragraph, vector);
+        for open_segment in &self.open_segments {
+            // Skip this segment if it doesn't match the segment filter
+            if !segment_filter
+                .as_ref()
+                .is_none_or(|f| segment_matches(f, open_segment.tags()))
+            {
+                continue;
+            }
+            let partial_solution =
+                open_segment.search(query, filter, with_duplicates, no_results, &self.config, min_score);
+
+            for candidate in partial_solution {
+                let addr = candidate.paragraph();
+                let paragraph = open_segment.get_paragraph(addr);
+                let scored_paragraph = ScoredParagraph::new(open_segment, addr, paragraph, candidate.score());
+                ffsv.add(scored_paragraph, candidate.vector());
+            }
         }
 
         Ok(ffsv.into())

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/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 {
@@ -180,8 +187,6 @@ fn merge_indexes<DS: DataStore + 'static>(
         index.advise(memmap2::Advice::Random)?;
     }
 
-    build_indexes(segment_path, &data_store)?;
-
     let metadata = VectorSegmentMetadata {
         path: segment_path.to_path_buf(),
         records: data_store.stored_paragraph_count() as usize,
@@ -191,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 {
@@ -291,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 {

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/pyproject.toml

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

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/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 {

{nidx_binding-6.9.1.post596 → nidx_binding-6.9.1.post603}/src/segment_store.rs

@@ -107,6 +107,7 @@ pub async fn download_segment(
 
     let temp_dir2 = temp_dir.clone();
     let mut tar = tar::Archive::new(reader);
+    tar.set_preserve_mtime(false); // We don't care about metadata and we can save a few syscalls
     let result = tokio::task::spawn_blocking(move || tar.unpack(temp_dir2)).await?;
 
     if let Err(e) = result {