usearch 2.17.1 → 2.19.4

package/javascript/dist/esm/usearch.js

@@ -229,10 +229,12 @@ export class Index {
      *        If a single key is provided, it is associated with all provided vectors.
      * @param {Float32Array|Float64Array|Int8Array} vectors - Input matrix representing vectors,
      *        matrix of size n * d, where n is the number of vectors, and d is their dimensionality.
+     * @param {number} [threads=0] - Optional, default is 0. Number of threads to use for indexing.
+     *        If set to 0, the number of threads is determined automatically.
      * @throws Will throw an error if the length of keys doesn't match the number of vectors
      *         or if it's not a single key.
      */
-    add(keys, vectors) {
+    add(keys, vectors, threads = 0) {
         let normalizedKeys = normalizeKeys(keys);
         let normalizedVectors = normalizeVectors(vectors, this.#compiledIndex.dimensions());
         let countVectors = normalizedVectors.length / this.#compiledIndex.dimensions();
@@ -244,8 +246,11 @@ export class Index {
         else if (normalizedKeys.length !== countVectors) {
             throw new Error(`The length of keys (${normalizedKeys.length}) must match the number of vectors (${countVectors}) or be a single key.`);
         }
+        if ((!Number.isNaN(threads) && typeof threads !== "number") || threads < 0) {
+            throw new Error("`threads` must be a non-negative integer representing the number of threads to use for searching.");
+        }
         // Call the compiled method
-        this.#compiledIndex.add(normalizedKeys, normalizedVectors);
+        this.#compiledIndex.add(normalizedKeys, normalizedVectors, threads);
     }
     /**
      * Perform a k-nearest neighbor search on the index.
@@ -264,17 +269,21 @@ export class Index {
      *
      * @param {Float32Array|Float64Array|Int8Array|Array<Array<number>>} vectors - Input matrix representing query vectors, can be a TypedArray or an array of TypedArray.
      * @param {number} k - The number of nearest neighbors to search for each query vector.
+     * @param {number} [threads=0] - Optional, default is 0. Number of threads to use for searching. If set to 0, the number of threads is determined automatically.
      * @return {Matches|BatchMatches} - Search results for one or more queries, containing keys, distances, and counts of the matches found.
      * @throws Will throw an error if `k` is not a positive integer or if the size of the vectors is not a multiple of dimensions.
      * @throws Will throw an error if `vectors` is not a valid input type (TypedArray or an array of TypedArray) or if its flattened size is not a multiple of dimensions.
      */
-    search(vectors, k) {
+    search(vectors, k, threads = 0) {
         if ((!Number.isNaN(k) && typeof k !== "number") || k <= 0) {
             throw new Error("`k` must be a positive integer representing the number of nearest neighbors to search for.");
         }
+        if ((!Number.isNaN(threads) && typeof threads !== "number") || threads < 0) {
+            throw new Error("`threads` must be a non-negative integer representing the number of threads to use for searching.");
+        }
         const normalizedVectors = normalizeVectors(vectors, this.#compiledIndex.dimensions());
         // Call the compiled method and create Matches or BatchMatches object with the result
-        const result = this.#compiledIndex.search(normalizedVectors, k);
+        const result = this.#compiledIndex.search(normalizedVectors, k, threads);
         const countInQueries = normalizedVectors.length / Number(this.#compiledIndex.dimensions());
         const batchMatches = new BatchMatches(...result, k);
         if (countInQueries === 1) {
@@ -402,6 +411,7 @@ export class Index {
  * @param {number} dimensions - The dimensionality of the vectors in both the dataset and the queries. It defines the number of elements in each vector.
  * @param {number} count - The number of nearest neighbors to return for each query. If the dataset contains fewer vectors than the specified count, the result will contain only the available vectors.
  * @param {MetricKind} metric - The distance metric to be used for the search.
+ * @param {number} [threads=0] - Optional, default is 0. The number of threads to use for the search. If set to 0, the number of threads is determined automatically.
  * @return {Matches|BatchMatches} - Returns a `Matches` or `BatchMatches` object containing the results of the search.
  * @throws Will throw an error if `dimensions` and `count` are not positive integers.
  * @throws Will throw an error if `metric` is not a valid MetricKind.
@@ -415,21 +425,25 @@ export class Index {
  * const dimensions = 2; // The number of elements in each vector.
  * const count = 1; // The number of nearest neighbors to return for each query.
  * const metric = MetricKind.IP; // Using the Inner Product distance metric.
+ * const threads = 0; // How many threads to use to perform the search.
  *
- * const result = exactSearch(dataset, queries, dimensions, count, metric);
+ * const result = exactSearch(dataset, queries, dimensions, count, metric, threads);
  * // result might be:
  * // {
  * //    keys: BigUint64Array [ 1n ],
  * //    distances: Float32Array [ some_value ],
  * // }
  */
-function exactSearch(dataset, queries, dimensions, count, metric) {
+export function exactSearch(dataset, queries, dimensions, count, metric, threads = 0) {
     // Validate and normalize the dimensions and count
     dimensions = Number(dimensions);
     count = Number(count);
     if (count <= 0 || dimensions <= 0) {
         throw new Error("Dimensions and count must be positive integers.");
     }
+    if ((!Number.isNaN(threads) && typeof threads !== "number") || threads < 0) {
+        throw new Error("`threads` must be a non-negative integer representing the number of threads to use for searching.");
+    }
     // Validate metric
     if (!Object.values(MetricKind).includes(metric)) {
         throw new Error(`Invalid metric: ${metric}. It must be one of: ${Object.values(MetricKind).join(", ")}`);
@@ -450,7 +464,7 @@ function exactSearch(dataset, queries, dimensions, count, metric) {
         throw new Error("Count must be equal or smaller than the number of vectors in the dataset.");
     }
     // Call the compiled function with the normalized input
-    const result = compiled.exactSearch(dataset, queries, dimensions, count, metric);
+    const result = compiled.exactSearch(dataset, queries, dimensions, count, metric, threads);
     // Create and return a Matches or BatchMatches object with the result
     if (countInQueries == 1) {
         return new Matches(result[0], result[1]);

package/javascript/lib.cpp

@@ -9,7 +9,8 @@
  *  @see NodeJS docs: https://nodejs.org/api/addons.html#hello-world
  *
  */
-#include <new> // `std::bad_alloc`
+#include <new>    // `std::bad_alloc`
+#include <thread> // `std::thread::hardware_concurrency()`
 
 #define NAPI_CPP_EXCEPTIONS
 #include <napi.h>
@@ -20,7 +21,10 @@
 using namespace unum::usearch;
 using namespace unum;
 
+using index_error_t = usearch::error_t;
 using add_result_t = typename index_dense_t::add_result_t;
+using search_result_t = typename index_dense_t::search_result_t;
+using state_result_t = typename index_dense_t::state_result_t;
 
 class CompiledIndex : public Napi::ObjectWrap<CompiledIndex> {
   public:
@@ -44,6 +48,7 @@ class CompiledIndex : public Napi::ObjectWrap<CompiledIndex> {
     Napi::Value Count(Napi::CallbackInfo const& ctx);
 
     std::unique_ptr<index_dense_t> native_;
+    std::mutex mtx;
 };
 
 Napi::Object CompiledIndex::Init(Napi::Env env, Napi::Object exports) {
@@ -77,7 +82,6 @@ std::size_t napi_argument_to_size(Napi::Value v) {
 }
 
 CompiledIndex::CompiledIndex(Napi::CallbackInfo const& ctx) : Napi::ObjectWrap<CompiledIndex>(ctx) {
-
     // Directly assign the parameters without checks
     std::size_t dimensions = napi_argument_to_size(ctx[0]);
     metric_kind_t metric_kind = metric_from_name(ctx[1].As<Napi::String>().Utf8Value().c_str());
@@ -95,7 +99,12 @@ CompiledIndex::CompiledIndex(Napi::CallbackInfo const& ctx) : Napi::ObjectWrap<C
 
     index_dense_config_t config(connectivity, expansion_add, expansion_search);
     config.multi = multi;
-    native_.reset(new index_dense_t(index_dense_t::make(metric, config)));
+    state_result_t result = index_dense_t::make(metric, config);
+    if (!result) {
+        Napi::TypeError::New(ctx.Env(), result.error.release()).ThrowAsJavaScriptException();
+        return;
+    }
+    native_.reset(new index_dense_t(std::move(result.index)));
     if (!native_)
         Napi::Error::New(ctx.Env(), "Out of memory!").ThrowAsJavaScriptException();
 }
@@ -155,45 +164,92 @@ void CompiledIndex::View(Napi::CallbackInfo const& ctx) {
 }
 
 void CompiledIndex::Add(Napi::CallbackInfo const& ctx) {
+    Napi::Env env = ctx.Env();
+
+    // Check the number of arguments
+    if (ctx.Length() != 3) {
+        Napi::TypeError::New(env, "`Add` expects 3 arguments: keys, vectors[, threads]").ThrowAsJavaScriptException();
+        return;
+    }
+
     // Extract keys and vectors from arguments
     Napi::BigUint64Array keys = ctx[0].As<Napi::BigUint64Array>();
-    std::size_t tasks = keys.ElementLength();
+    Napi::TypedArray vectors = ctx[1].As<Napi::TypedArray>();
+
+    // Optional arguments
+    std::size_t threads = napi_argument_to_size(ctx[2]);
+    if (threads == 0)
+        threads = std::thread::hardware_concurrency();
 
-    // Ensure there is enough capacity
+    // Ensure there is enough capacity and memory
+    std::size_t tasks = keys.ElementLength();
     if (native_->size() + tasks >= native_->capacity())
-        native_->reserve(ceil2(native_->size() + tasks));
+        if (!native_->try_reserve({ceil2(native_->size() + tasks), threads})) {
+            Napi::TypeError::New(env, "Failed to reserve memory").ThrowAsJavaScriptException();
+            return;
+        }
 
-    // Create an instance of the executor with the default number of threads
+    // Run insertions concurrently
     auto run_parallel = [&](auto vectors) {
-        executor_stl_t executor;
+        // Errors can be set only from the main thread, so before spawning workers
+        // we need temporary space to keep the message
+        index_error_t first_error{};
+        std::atomic<bool> failed{false};
+        executor_default_t executor{threads};
         executor.fixed(tasks, [&](std::size_t /*thread_idx*/, std::size_t task_idx) {
-            add_result_t result = native_->add(static_cast<default_key_t>(keys[task_idx]),
-                                               vectors + task_idx * native_->dimensions());
-            if (!result)
-                Napi::Error::New(ctx.Env(), result.error.release()).ThrowAsJavaScriptException();
+            if (failed.load())
+                return;
+            auto key = static_cast<default_key_t>(keys[task_idx]);
+            auto vector = vectors + task_idx * native_->dimensions();
+            add_result_t result = native_->add(key, vector);
+            if (!result) {
+                if (!failed.exchange(true)) {
+                    first_error = std::move(result.error);
+                } else {
+                    result.error.release();
+                }
+            }
         });
+        if (failed)
+            Napi::TypeError::New(env, first_error.release()).ThrowAsJavaScriptException();
     };
 
-    Napi::TypedArray vectors = ctx[1].As<Napi::TypedArray>();
-    if (vectors.TypedArrayType() == napi_float32_array) {
-        run_parallel(vectors.As<Napi::Float32Array>().Data());
-    } else if (vectors.TypedArrayType() == napi_float64_array) {
-        run_parallel(vectors.As<Napi::Float64Array>().Data());
-    } else if (vectors.TypedArrayType() == napi_int8_array) {
-        run_parallel(vectors.As<Napi::Int8Array>().Data());
-    } else {
-        Napi::TypeError::New(ctx.Env(),
-                             "Unsupported TypedArray. Supported types are Float32Array, Float64Array, and Int8Array.")
-            .ThrowAsJavaScriptException();
+    // Dispatch the parallel tasks based on the `TypedArray` type
+    try {
+        if (vectors.TypedArrayType() == napi_float32_array) {
+            run_parallel(vectors.As<Napi::Float32Array>().Data());
+        } else if (vectors.TypedArrayType() == napi_float64_array) {
+            run_parallel(vectors.As<Napi::Float64Array>().Data());
+        } else if (vectors.TypedArrayType() == napi_int8_array) {
+            run_parallel(vectors.As<Napi::Int8Array>().Data());
+        } else {
+            Napi::TypeError::New(
+                env, "Unsupported TypedArray. Supported types are Float32Array, Float64Array, and Int8Array.")
+                .ThrowAsJavaScriptException();
+        }
+    } catch (...) {
+        Napi::TypeError::New(env, "Insertion failed").ThrowAsJavaScriptException();
     }
 }
 
 Napi::Value CompiledIndex::Search(Napi::CallbackInfo const& ctx) {
     Napi::Env env = ctx.Env();
+
+    // Check the number of arguments
+    if (ctx.Length() != 3) {
+        Napi::TypeError::New(env, "`Search` expects 3 arguments: queries, k[, threads]").ThrowAsJavaScriptException();
+        return env.Null();
+    }
+
+    // Extract mandatory arguments
     Napi::TypedArray queries = ctx[0].As<Napi::TypedArray>();
-    std::size_t tasks = queries.ElementLength() / native_->dimensions();
     std::size_t wanted = napi_argument_to_size(ctx[1]);
+    std::size_t threads = napi_argument_to_size(ctx[2]);
+    if (threads == 0)
+        threads = std::thread::hardware_concurrency();
 
+    // Run queries concurrently
+    std::size_t tasks = queries.ElementLength() / native_->dimensions();
     auto run_parallel = [&](auto vectors) -> Napi::Value {
         Napi::Array result_js = Napi::Array::New(env, 3);
         Napi::BigUint64Array matches_js = Napi::BigUint64Array::New(env, tasks * wanted);
@@ -204,29 +260,30 @@ Napi::Value CompiledIndex::Search(Napi::CallbackInfo const& ctx) {
         auto distances_data = distances_js.Data();
         auto counts_data = counts_js.Data();
 
-        try {
-            bool failed = false;
-            executor_stl_t executor;
-            executor.fixed(tasks, [&](std::size_t /*thread_idx*/, std::size_t task_idx) {
-                auto result = native_->search(vectors + task_idx * native_->dimensions(), wanted);
-                if (!result) {
-                    failed = true;
-                    Napi::TypeError::New(env, result.error.release()).ThrowAsJavaScriptException();
+        // Errors can be set only from the main thread, so before spawning workers
+        // we need temporary space to keep the message
+        index_error_t first_error{};
+        std::atomic<bool> failed{false};
+        executor_default_t executor{threads};
+        executor.fixed(tasks, [&](std::size_t /*thread_idx*/, std::size_t task_idx) {
+            if (failed.load())
+                return;
+            auto vector = vectors + task_idx * native_->dimensions();
+            search_result_t result = native_->search(vector, wanted);
+            if (!result) {
+                if (!failed.exchange(true)) {
+                    first_error = std::move(result.error);
                 } else {
-                    counts_data[task_idx] = result.dump_to(matches_data + task_idx * native_->dimensions(),
-                                                           distances_data + task_idx * native_->dimensions());
+                    result.error.release();
                 }
-            });
-
-            if (failed)
-                return env.Null();
-
-        } catch (std::bad_alloc const&) {
-            Napi::TypeError::New(env, "Out of memory").ThrowAsJavaScriptException();
-            return env.Null();
-
-        } catch (...) {
-            Napi::TypeError::New(env, "Search failed").ThrowAsJavaScriptException();
+            } else {
+                auto matches = matches_data + task_idx * wanted;
+                auto distances = distances_data + task_idx * wanted;
+                counts_data[task_idx] = result.dump_to(matches, distances);
+            }
+        });
+        if (failed) {
+            Napi::TypeError::New(env, first_error.release()).ThrowAsJavaScriptException();
             return env.Null();
         }
 
@@ -236,16 +293,22 @@ Napi::Value CompiledIndex::Search(Napi::CallbackInfo const& ctx) {
         return result_js;
     };
 
-    if (queries.TypedArrayType() == napi_float32_array) {
-        return run_parallel(queries.As<Napi::Float32Array>().Data());
-    } else if (queries.TypedArrayType() == napi_float64_array) {
-        return run_parallel(queries.As<Napi::Float64Array>().Data());
-    } else if (queries.TypedArrayType() == napi_int8_array) {
-        return run_parallel(queries.As<Napi::Int8Array>().Data());
-    } else {
-        Napi::TypeError::New(env,
-                             "Unsupported TypedArray. Supported types are Float32Array, Float64Array, and Int8Array.")
-            .ThrowAsJavaScriptException();
+    // Dispatch the parallel tasks based on the `TypedArray` type
+    try {
+        if (queries.TypedArrayType() == napi_float32_array) {
+            return run_parallel(queries.As<Napi::Float32Array>().Data());
+        } else if (queries.TypedArrayType() == napi_float64_array) {
+            return run_parallel(queries.As<Napi::Float64Array>().Data());
+        } else if (queries.TypedArrayType() == napi_int8_array) {
+            return run_parallel(queries.As<Napi::Int8Array>().Data());
+        } else {
+            Napi::TypeError::New(
+                env, "Unsupported TypedArray. Supported types are Float32Array, Float64Array, and Int8Array.")
+                .ThrowAsJavaScriptException();
+            return env.Null();
+        }
+    } catch (...) {
+        Napi::TypeError::New(env, "Search failed").ThrowAsJavaScriptException();
         return env.Null();
     }
 }
@@ -287,21 +350,36 @@ Napi::Value CompiledIndex::Count(Napi::CallbackInfo const& ctx) {
 Napi::Value exactSearch(Napi::CallbackInfo const& ctx) {
     Napi::Env env = ctx.Env();
 
+    // Check the number of arguments
+    if (ctx.Length() != 6) {
+        Napi::TypeError::New(env,
+                             "`exactSearch` expects 6 arguments: dataset, queries, dimensions, k, metric[, threads].")
+            .ThrowAsJavaScriptException();
+        return env.Null();
+    }
+
     // Extracting parameters directly without additional type checks.
     Napi::TypedArray dataset = ctx[0].As<Napi::TypedArray>();
     Napi::ArrayBuffer datasetBuffer = dataset.ArrayBuffer();
     Napi::TypedArray queries = ctx[1].As<Napi::TypedArray>();
     Napi::ArrayBuffer queriesBuffer = queries.ArrayBuffer();
-    std::uint64_t dimensions = napi_argument_to_size(ctx[2]);
-    std::uint64_t wanted = napi_argument_to_size(ctx[3]);
+    std::size_t dimensions = napi_argument_to_size(ctx[2]);
+    std::size_t wanted = napi_argument_to_size(ctx[3]);
     metric_kind_t metric_kind = metric_from_name(ctx[4].As<Napi::String>().Utf8Value().c_str());
+    std::size_t threads = napi_argument_to_size(ctx[5]);
+    if (threads == 0)
+        threads = std::thread::hardware_concurrency();
 
+    // Check the types used
     scalar_kind_t quantization;
     std::size_t bytes_per_scalar;
     switch (queries.TypedArrayType()) {
     case napi_float64_array: quantization = scalar_kind_t::f64_k, bytes_per_scalar = 8; break;
+    case napi_float32_array: quantization = scalar_kind_t::f32_k, bytes_per_scalar = 4; break;
     case napi_int8_array: quantization = scalar_kind_t::i8_k, bytes_per_scalar = 1; break;
-    default: quantization = scalar_kind_t::f32_k, bytes_per_scalar = 4; break;
+    default:
+        Napi::TypeError::New(env, "Unsupported TypedArray for queries.").ThrowAsJavaScriptException();
+        return env.Null();
     }
 
     metric_punned_t metric(dimensions, metric_kind, quantization);
@@ -310,7 +388,7 @@ Napi::Value exactSearch(Napi::CallbackInfo const& ctx) {
         return env.Null();
     }
 
-    executor_default_t executor;
+    executor_default_t executor{threads};
     exact_search_t search;
 
     // Performing the exact search.