Add: reserve_capacity_and_threads for Rust

Author: ashvardanian

.vscode/settings.json

@@ -135,7 +135,8 @@
     "xstring": "cpp",
     "xtr1common": "cpp",
     "xtree": "cpp",
-    "xutility": "cpp"
+    "xutility": "cpp",
+    "*.rs": "cpp"
   },
   "cSpell.words": [
     "allclose",

Cargo.lock

@@ -35,7 +35,7 @@ cxx = "1.0.160"
 cxx-build = "1.0.160"
 
 [dev-dependencies]
-fork_union = "2.1.0"  # for concurrency tests
+fork_union = "2.1.1"  # for concurrency tests
 rand_distr = "0.5.1"  # uniform floats distribution
 rand_chacha = "0.9.0" # random number generator
 rand = "0.9.1"

Cargo.toml

@@ -134,6 +134,9 @@ size_t NativeIndex::count(vector_key_t key) const { return index_->count(key); }
 bool NativeIndex::contains(vector_key_t key) const { return index_->contains(key); }
 
 void NativeIndex::reserve(size_t capacity) const { index_->reserve(capacity); }
+void NativeIndex::reserve_capacity_and_threads(size_t capacity, size_t threads) const {
+    index_->reserve({capacity, threads});
+}
 
 size_t NativeIndex::dimensions() const { return index_->dimensions(); }
 size_t NativeIndex::connectivity() const { return index_->connectivity(); }

rust/lib.cpp

@@ -25,6 +25,7 @@ class NativeIndex {
     NativeIndex(std::unique_ptr<index_dense_t> index);
 
     void reserve(size_t) const;
+    void reserve_capacity_and_threads(size_t, size_t) const;
 
     void add_b1x8(vector_key_t key, rust::Slice<uint8_t const> vector) const;
     void add_i8(vector_key_t key, rust::Slice<int8_t const> vector) const;

rust/lib.hpp

@@ -332,6 +332,7 @@ pub mod ffi {
 
         pub fn new_native_index(options: &IndexOptions) -> Result<UniquePtr<NativeIndex>>;
         pub fn reserve(self: &NativeIndex, capacity: usize) -> Result<()>;
+        pub fn reserve_capacity_and_threads(self: &NativeIndex, capacity: usize, threads: usize) -> Result<()>;
         pub fn dimensions(self: &NativeIndex) -> usize;
         pub fn connectivity(self: &NativeIndex) -> usize;
         pub fn size(self: &NativeIndex) -> usize;
@@ -1143,6 +1144,20 @@ impl Index {
         self.inner.reserve(capacity)
     }
 
+    /// Reserves memory for a specified number of incoming vectors & active threads.
+    ///
+    /// # Arguments
+    ///
+    /// * `capacity` - The desired total capacity, including the current size.
+    /// * `threads` - The number of threads to use for the operation.
+    pub fn reserve_capacity_and_threads(
+        self: &Index,
+        capacity: usize,
+        threads: usize,
+    ) -> Result<(), cxx::Exception> {
+        self.inner.reserve_capacity_and_threads(capacity, threads)
+    }
+
     /// Retrieves the number of dimensions in the vectors indexed.
     pub fn dimensions(self: &Index) -> usize {
         self.inner.dimensions()
@@ -1758,7 +1773,7 @@ mod tests {
         };
 
         let index = Arc::new(Index::new(&options).unwrap());
-        index.reserve(VECTOR_COUNT).unwrap();
+        index.reserve_capacity_and_threads(VECTOR_COUNT, THREAD_COUNT).unwrap();
 
         // Generate deterministic vectors using rand crate for reproducible testing
         let seed = 42; // Fixed seed for reproducibility