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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.9.1.post600
+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.post600 → 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.post600 → 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.post600 → 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.post600"
+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.post600 → 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.post600 → 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.post600 → nidx_binding-6.9.1.post603}/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.post603}/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.post603}/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.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.post600 → 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.post600 → 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.post600 → 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.post600 → 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 {
@@ -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 {

{nidx_binding-6.9.1.post600 → 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.post600"
+version = "6.9.1.post603"
 
 [project.urls]
 Homepage = "https://nuclia.com"

{nidx_binding-6.9.1.post600 → 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 {