/** * hnswlib-node provides Node.js bindings for Hnswlib. * * Copyright (c) 2022-2024 Atsushi Tatsuma * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ #include #include #include #include #include #include #include #include #include "hnswlib/hnswlib.h" namespace { void normalizePoint(std::vector& vec) { const size_t dim = vec.size(); const float norm = std::sqrt(std::fabs(std::inner_product(vec.begin(), vec.end(), vec.begin(), 0.0f))); if (norm > 0.0f) { for (size_t i = 0; i < dim; i++) vec[i] /= norm; } } } // namespace class L2Space : public Napi::ObjectWrap { public: uint32_t dim_; std::unique_ptr l2space_; L2Space(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return; } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return; } dim_ = info[0].As().Uint32Value(); l2space_ = std::unique_ptr(new hnswlib::L2Space(static_cast(dim_))); } static Napi::Object Init(Napi::Env env, Napi::Object exports) { // clang-format off Napi::Function func = DefineClass(env, "L2Space", { InstanceMethod("distance", &L2Space::distance), InstanceMethod("getNumDimensions", &L2Space::getNumDimensions) }); // clang-format on Napi::FunctionReference* constructor = new Napi::FunctionReference(); *constructor = Napi::Persistent(func); env.SetInstanceData(constructor); exports.Set("L2Space", func); return exports; } private: Napi::Value distance(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 2) { Napi::Error::New(env, "Expected 2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsArray()) { Napi::TypeError::New(env, "Invalid the first argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsArray()) { Napi::TypeError::New(env, "Invalid the second argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } Napi::Array arr_a = info[0].As(); Napi::Array arr_b = info[1].As(); if (arr_a.Length() != dim_) { Napi::Error::New(env, "Invalid the first array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr_a.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } if (arr_b.Length() != dim_) { Napi::Error::New(env, "Invalid the second array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr_b.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } std::vector vec_a(dim_, 0.0); std::vector vec_b(dim_, 0.0); for (uint32_t i = 0; i < dim_; i++) { Napi::Value val_a = arr_a[i]; Napi::Value val_b = arr_b[i]; vec_a[i] = val_a.As().FloatValue(); vec_b[i] = val_b.As().FloatValue(); } hnswlib::DISTFUNC df = l2space_->get_dist_func(); const float d = df(vec_a.data(), vec_b.data(), l2space_->get_dist_func_param()); return Napi::Number::New(info.Env(), d); } Napi::Value getNumDimensions(const Napi::CallbackInfo& info) { return Napi::Number::New(info.Env(), dim_); } }; class InnerProductSpace : public Napi::ObjectWrap { public: uint32_t dim_; std::unique_ptr ipspace_; InnerProductSpace(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return; } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return; } dim_ = info[0].As().Uint32Value(); ipspace_ = std::unique_ptr(new hnswlib::InnerProductSpace(static_cast(dim_))); } static Napi::Object Init(Napi::Env env, Napi::Object exports) { // clang-format off Napi::Function func = DefineClass(env, "InnerProductSpace", { InstanceMethod("distance", &InnerProductSpace::distance), InstanceMethod("getNumDimensions", &InnerProductSpace::getNumDimensions) }); // clang-format on Napi::FunctionReference* constructor = new Napi::FunctionReference(); *constructor = Napi::Persistent(func); env.SetInstanceData(constructor); exports.Set("InnerProductSpace", func); return exports; } private: Napi::Value distance(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 2) { Napi::Error::New(env, "Expected 2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsArray()) { Napi::TypeError::New(env, "Invalid the first argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsArray()) { Napi::TypeError::New(env, "Invalid the second argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } Napi::Array arr_a = info[0].As(); Napi::Array arr_b = info[1].As(); if (arr_a.Length() != dim_) { Napi::Error::New(env, "Invalid the first array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr_a.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } if (arr_b.Length() != dim_) { Napi::Error::New(env, "Invalid the second array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr_b.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } std::vector vec_a(dim_, 0.0); std::vector vec_b(dim_, 0.0); for (uint32_t i = 0; i < dim_; i++) { Napi::Value val_a = arr_a[i]; Napi::Value val_b = arr_b[i]; vec_a[i] = val_a.As().FloatValue(); vec_b[i] = val_b.As().FloatValue(); } hnswlib::DISTFUNC df = ipspace_->get_dist_func(); const float d = df(vec_a.data(), vec_b.data(), ipspace_->get_dist_func_param()); return Napi::Number::New(info.Env(), d); } Napi::Value getNumDimensions(const Napi::CallbackInfo& info) { return Napi::Number::New(info.Env(), dim_); } }; class CustomFilterFunctor : public hnswlib::BaseFilterFunctor { public: CustomFilterFunctor(const Napi::Env& env, const Napi::Function& callback) : env_(env), callback_(callback) {} bool operator()(hnswlib::labeltype id) { return callback_.Call({Napi::Number::New(env_, id)}).ToBoolean().Value(); } private: Napi::Env env_; Napi::Function callback_; }; class LoadBruteforceSearchIndexWorker : public Napi::AsyncWorker { public: LoadBruteforceSearchIndexWorker(const std::string& filename, hnswlib::BruteforceSearch** index, hnswlib::SpaceInterface** space, Napi::Promise::Deferred deferred, Napi::Function& callback) : Napi::AsyncWorker(callback), deferred(deferred), filename_(filename), result_(false), index_(index), space_(space) {} ~LoadBruteforceSearchIndexWorker() {} void Execute() { try { std::ifstream ifs(filename_); if (ifs.is_open()) { ifs.close(); } else { throw std::runtime_error("failed to open file: " + filename_); } if (*index_) delete *index_; *index_ = new hnswlib::BruteforceSearch(*space_, filename_); result_ = true; } catch (const std::exception& e) { result_ = false; SetError("Hnswlib Error: " + std::string(e.what())); } } void OnOK() { Napi::HandleScope scope(Env()); Napi::Boolean result = Napi::Boolean::New(Env(), result_); deferred.Resolve(result); } void OnError(const Napi::Error& e) { deferred.Reject(e.Value()); } Napi::Promise::Deferred deferred; private: std::string filename_; bool result_; hnswlib::BruteforceSearch** index_; hnswlib::SpaceInterface** space_; }; class SaveBruteforceSearchIndexWorker : public Napi::AsyncWorker { public: SaveBruteforceSearchIndexWorker(const std::string& filename, hnswlib::BruteforceSearch** index, Napi::Promise::Deferred deferred, Napi::Function& callback) : Napi::AsyncWorker(callback), deferred(deferred), filename_(filename), result_(false), index_(index) {} ~SaveBruteforceSearchIndexWorker() {} void Execute() { try { if (*index_ == nullptr) throw std::runtime_error("search index is not constructed."); (*index_)->saveIndex(filename_); result_ = true; } catch (const std::exception& e) { result_ = false; SetError("Hnswlib Error: " + std::string(e.what())); } } void OnOK() { Napi::HandleScope scope(Env()); Napi::Boolean result = Napi::Boolean::New(Env(), result_); deferred.Resolve(result); } void OnError(const Napi::Error& e) { deferred.Reject(e.Value()); } Napi::Promise::Deferred deferred; private: std::string filename_; bool result_; hnswlib::BruteforceSearch** index_; }; class BruteforceSearch : public Napi::ObjectWrap { public: uint32_t dim_; hnswlib::BruteforceSearch* index_; hnswlib::SpaceInterface* space_; bool normalize_; BruteforceSearch(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info), index_(nullptr), space_(nullptr), normalize_(false) { Napi::Env env = info.Env(); if (info.Length() != 2) { Napi::Error::New(env, "Expected 2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return; } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return; } if (!info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a number.").ThrowAsJavaScriptException(); return; } const std::string space_name = info[0].As().ToString(); if (space_name != "l2" && space_name != "ip" && space_name != "cosine") { Napi::Error::New(env, "Wrong space name, expected \"l2\", \"ip\", or \"cosine\".").ThrowAsJavaScriptException(); return; } dim_ = info[1].As().Uint32Value(); try { if (space_name == "l2") { space_ = new hnswlib::L2Space(static_cast(dim_)); } else if (space_name == "ip") { space_ = new hnswlib::InnerProductSpace(static_cast(dim_)); } else { space_ = new hnswlib::InnerProductSpace(static_cast(dim_)); normalize_ = true; } } catch (const std::bad_alloc& err) { Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return; } } ~BruteforceSearch() { if (space_) delete space_; if (index_) delete index_; } static Napi::Object Init(Napi::Env env, Napi::Object exports) { // clang-format off Napi::Function func = DefineClass(env, "BruteforceSearch", { InstanceMethod("initIndex", &BruteforceSearch::initIndex), InstanceMethod("readIndex", &BruteforceSearch::readIndex), InstanceMethod("readIndexSync", &BruteforceSearch::readIndexSync), InstanceMethod("writeIndex", &BruteforceSearch::writeIndex), InstanceMethod("writeIndexSync", &BruteforceSearch::writeIndexSync), InstanceMethod("addPoint", &BruteforceSearch::addPoint), InstanceMethod("removePoint", &BruteforceSearch::removePoint), InstanceMethod("searchKnn", &BruteforceSearch::searchKnn), InstanceMethod("getMaxElements", &BruteforceSearch::getMaxElements), InstanceMethod("getCurrentCount", &BruteforceSearch::getCurrentCount), InstanceMethod("getNumDimensions", &BruteforceSearch::getNumDimensions) }); // clang-format on Napi::FunctionReference* constructor = new Napi::FunctionReference(); *constructor = Napi::Persistent(func); env.SetInstanceData(constructor); exports.Set("BruteforceSearch", func); return exports; } private: Napi::Value initIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } const uint32_t max_elements = info[0].As().Uint32Value(); if (index_) delete index_; try { index_ = new hnswlib::BruteforceSearch(space_, static_cast(max_elements)); } catch (const std::bad_alloc& err) { index_ = nullptr; Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value readIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); Napi::Function callback = Napi::Function::New(env, [](const Napi::CallbackInfo& info) { return info.Env().Undefined(); }); Napi::Promise::Deferred deferred = Napi::Promise::Deferred::New(info.Env()); LoadBruteforceSearchIndexWorker* worker = new LoadBruteforceSearchIndexWorker(filename, &index_, &space_, deferred, callback); worker->Queue(); return worker->deferred.Promise(); } Napi::Value readIndexSync(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); try { std::ifstream ifs(filename); if (ifs.is_open()) { ifs.close(); } else { throw std::runtime_error("failed to open file: " + filename); } if (index_) delete index_; index_ = new hnswlib::BruteforceSearch(space_, filename); } catch (const std::exception& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value writeIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); Napi::Function callback = Napi::Function::New(env, [](const Napi::CallbackInfo& info) { return info.Env().Undefined(); }); Napi::Promise::Deferred deferred = Napi::Promise::Deferred::New(info.Env()); SaveBruteforceSearchIndexWorker* worker = new SaveBruteforceSearchIndexWorker(filename, &index_, deferred, callback); worker->Queue(); return worker->deferred.Promise(); } Napi::Value writeIndexSync(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); if (index_ == nullptr) { Napi::Error::New(env, "Search index is not constructed.").ThrowAsJavaScriptException(); return env.Null(); } index_->saveIndex(filename); return env.Null(); } Napi::Value addPoint(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 2) { Napi::Error::New(env, "Expected 2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsArray()) { Napi::TypeError::New(env, "Invalid the first argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } Napi::Array arr = info[0].As(); if (arr.Length() != dim_) { Napi::Error::New(env, "Invalid the given array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } std::vector vec(dim_, 0.0); for (uint32_t i = 0; i < dim_; i++) { Napi::Value val = arr[i]; vec[i] = val.As().FloatValue(); } if (normalize_) normalizePoint(vec); const uint32_t idx = info[1].As().Uint32Value(); try { index_->addPoint(reinterpret_cast(vec.data()), static_cast(idx)); } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value removePoint(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } const uint32_t idx = info[0].As().Uint32Value(); index_->removePoint(static_cast(idx)); return env.Null(); } Napi::Value searchKnn(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() < 2) { Napi::Error::New(env, "Expected 2-3 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsArray()) { Napi::TypeError::New(env, "Invalid the first argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[2].IsUndefined() && !info[2].IsFunction()) { Napi::TypeError::New(env, "Invalid the third argument type, must be a function.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } CustomFilterFunctor* filterFn = nullptr; if (info[2].IsFunction()) { try { filterFn = new CustomFilterFunctor(env, info[2].As()); } catch (const std::bad_alloc& err) { Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return env.Null(); } } Napi::Array arr = info[0].As(); if (arr.Length() != dim_) { Napi::Error::New(env, "Invalid the given array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } const uint32_t k = info[1].As().Uint32Value(); if (k > index_->maxelements_) { Napi::Error::New(env, "Invalid the number of k-nearest neighbors (cannot be given a value greater than `maxElements`: " + std::to_string(index_->maxelements_) + ").") .ThrowAsJavaScriptException(); return env.Null(); } if (k <= 0) { Napi::Error::New(env, "Invalid the number of k-nearest neighbors (must be a positive number).") .ThrowAsJavaScriptException(); return env.Null(); } std::vector vec(dim_, 0.0); for (uint32_t i = 0; i < dim_; i++) { Napi::Value val = arr[i]; vec[i] = val.As().FloatValue(); } if (normalize_) normalizePoint(vec); std::priority_queue> knn = index_->searchKnn(reinterpret_cast(vec.data()), static_cast(k), filterFn); const size_t n_results = knn.size(); Napi::Array arr_distances = Napi::Array::New(env, n_results); Napi::Array arr_neighbors = Napi::Array::New(env, n_results); for (int32_t i = static_cast(n_results) - 1; i >= 0; i--) { const std::pair& nn = knn.top(); arr_distances[i] = Napi::Number::New(env, nn.first); arr_neighbors[i] = Napi::Number::New(env, nn.second); knn.pop(); } if (filterFn) delete filterFn; Napi::Object results = Napi::Object::New(env); results.Set("distances", arr_distances); results.Set("neighbors", arr_neighbors); return results; } Napi::Value getMaxElements(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (index_ == nullptr) return Napi::Number::New(env, 0); return Napi::Number::New(env, index_->maxelements_); } Napi::Value getCurrentCount(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (index_ == nullptr) return Napi::Number::New(env, 0); return Napi::Number::New(env, index_->cur_element_count); } Napi::Value getNumDimensions(const Napi::CallbackInfo& info) { return Napi::Number::New(info.Env(), dim_); } }; class LoadHierarchicalNSWIndexWorker : public Napi::AsyncWorker { public: LoadHierarchicalNSWIndexWorker(const std::string& filename, const bool allow_replace_deleted, hnswlib::HierarchicalNSW** index, hnswlib::SpaceInterface** space, Napi::Promise::Deferred deferred, Napi::Function& callback) : Napi::AsyncWorker(callback), deferred(deferred), filename_(filename), allow_replace_deleted_(allow_replace_deleted), result_(false), index_(index), space_(space) {} ~LoadHierarchicalNSWIndexWorker() {} void Execute() { try { std::ifstream ifs(filename_); if (ifs.is_open()) { ifs.close(); } else { throw std::runtime_error("failed to open file: " + filename_); } if (*index_) delete *index_; *index_ = new hnswlib::HierarchicalNSW(*space_, filename_, false, 0, allow_replace_deleted_); result_ = true; } catch (const std::exception& e) { result_ = false; SetError("Hnswlib Error: " + std::string(e.what())); } } void OnOK() { Napi::HandleScope scope(Env()); Napi::Boolean result = Napi::Boolean::New(Env(), result_); deferred.Resolve(result); } void OnError(const Napi::Error& e) { deferred.Reject(e.Value()); } Napi::Promise::Deferred deferred; private: std::string filename_; bool allow_replace_deleted_; bool result_; hnswlib::HierarchicalNSW** index_; hnswlib::SpaceInterface** space_; }; class SaveHierarchicalNSWIndexWorker : public Napi::AsyncWorker { public: SaveHierarchicalNSWIndexWorker(const std::string& filename, hnswlib::HierarchicalNSW** index, Napi::Promise::Deferred deferred, Napi::Function& callback) : Napi::AsyncWorker(callback), deferred(deferred), filename_(filename), result_(false), index_(index) {} ~SaveHierarchicalNSWIndexWorker() {} void Execute() { try { if (*index_ == nullptr) throw std::runtime_error("search index is not constructed."); (*index_)->saveIndex(filename_); result_ = true; } catch (const std::exception& e) { result_ = false; SetError("Hnswlib Error: " + std::string(e.what())); } } void OnOK() { Napi::HandleScope scope(Env()); Napi::Boolean result = Napi::Boolean::New(Env(), result_); deferred.Resolve(result); } void OnError(const Napi::Error& e) { deferred.Reject(e.Value()); } Napi::Promise::Deferred deferred; private: std::string filename_; bool result_; hnswlib::HierarchicalNSW** index_; }; class HierarchicalNSW : public Napi::ObjectWrap { public: uint32_t dim_; hnswlib::HierarchicalNSW* index_; hnswlib::SpaceInterface* space_; bool normalize_; HierarchicalNSW(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info), index_(nullptr), space_(nullptr), normalize_(false) { Napi::Env env = info.Env(); if (info.Length() != 2) { Napi::Error::New(env, "Expected 2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return; } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return; } if (!info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a number.").ThrowAsJavaScriptException(); return; } const std::string space_name = info[0].As().ToString(); if (space_name != "l2" && space_name != "ip" && space_name != "cosine") { Napi::Error::New(env, "Wrong space name, expected \"l2\", \"ip\", or \"cosine\".").ThrowAsJavaScriptException(); return; } dim_ = info[1].As().Uint32Value(); try { if (space_name == "l2") { space_ = new hnswlib::L2Space(static_cast(dim_)); } else if (space_name == "ip") { space_ = new hnswlib::InnerProductSpace(static_cast(dim_)); } else { space_ = new hnswlib::InnerProductSpace(static_cast(dim_)); normalize_ = true; } } catch (const std::bad_alloc& err) { Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return; } } ~HierarchicalNSW() { if (space_) delete space_; if (index_) delete index_; } static Napi::Object Init(Napi::Env env, Napi::Object exports) { // clang-format off Napi::Function func = DefineClass(env, "HierarchicalNSW", { InstanceMethod("initIndex", &HierarchicalNSW::initIndex), InstanceMethod("readIndex", &HierarchicalNSW::readIndex), InstanceMethod("readIndexSync", &HierarchicalNSW::readIndexSync), InstanceMethod("writeIndex", &HierarchicalNSW::writeIndex), InstanceMethod("writeIndexSync", &HierarchicalNSW::writeIndexSync), InstanceMethod("resizeIndex", &HierarchicalNSW::resizeIndex), InstanceMethod("addPoint", &HierarchicalNSW::addPoint), InstanceMethod("markDelete", &HierarchicalNSW::markDelete), InstanceMethod("unmarkDelete", &HierarchicalNSW::unmarkDelete), InstanceMethod("searchKnn", &HierarchicalNSW::searchKnn), InstanceMethod("getIdsList", &HierarchicalNSW::getIdsList), InstanceMethod("getPoint", &HierarchicalNSW::getPoint), InstanceMethod("getMaxElements", &HierarchicalNSW::getMaxElements), InstanceMethod("getCurrentCount", &HierarchicalNSW::getCurrentCount), InstanceMethod("getNumDimensions", &HierarchicalNSW::getNumDimensions), InstanceMethod("getEf", &HierarchicalNSW::getEf), InstanceMethod("setEf", &HierarchicalNSW::setEf) }); // clang-format on Napi::FunctionReference* constructor = new Napi::FunctionReference(); *constructor = Napi::Persistent(func); env.SetInstanceData(constructor); exports.Set("HierarchicalNSW", func); return exports; } private: Napi::Value initIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() < 1) { Napi::Error::New(env, "Expected 1-5 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } const bool given_named_args = info.Length() == 1 && info[0].IsObject(); const Napi::Object named_args = given_named_args ? info[0].As() : Napi::Object::New(env); if (given_named_args) { if (!named_args.Has("maxElements")) { Napi::TypeError::New(env, "Missing named argument `maxElements`, must be specified.").ThrowAsJavaScriptException(); return env.Null(); } if (!named_args.Get("maxElements").IsNumber()) { Napi::TypeError::New(env, "Invalid the named argument type `maxElements`, must be a number.") .ThrowAsJavaScriptException(); return env.Null(); } if (named_args.Has("m") && !named_args.Get("m").IsNumber()) { Napi::TypeError::New(env, "Invalid the named argument type `m`, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (named_args.Has("efConstruction") && !named_args.Get("efConstruction").IsNumber()) { Napi::TypeError::New(env, "Invalid the named argument type `efConstruction`, must be a number.") .ThrowAsJavaScriptException(); return env.Null(); } if (named_args.Has("randomSeed") && !named_args.Get("randomSeed").IsNumber()) { Napi::TypeError::New(env, "Invalid the named argument type `randomSeed`, must be a number.") .ThrowAsJavaScriptException(); return env.Null(); } if (named_args.Has("allowReplaceDeleted") && !named_args.Get("allowReplaceDeleted").IsBoolean()) { Napi::TypeError::New(env, "Invalid the named argument type `allowRepalceDeleted`, must be a boolean.") .ThrowAsJavaScriptException(); return env.Null(); } } else { if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsUndefined() && !info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the sencond argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[2].IsUndefined() && !info[2].IsNumber()) { Napi::TypeError::New(env, "Invalid the third argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[3].IsUndefined() && !info[3].IsNumber()) { Napi::TypeError::New(env, "Invalid the fourth argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[4].IsUndefined() && !info[4].IsBoolean()) { Napi::TypeError::New(env, "Invalid the fifth argument type, must be a boolean.").ThrowAsJavaScriptException(); return env.Null(); } } const uint32_t max_elements_ = given_named_args ? named_args.Get("maxElements").As().Uint32Value() : info[0].As().Uint32Value(); const uint32_t m_ = given_named_args ? (!named_args.Has("m") ? 16 : named_args.Get("m").As().Uint32Value()) : (info[1].IsUndefined() ? 16 : info[1].As().Uint32Value()); const uint32_t ef_construction_ = given_named_args ? (!named_args.Has("efConstruction") ? 200 : named_args.Get("efConstruction").As().Uint32Value()) : (info[2].IsUndefined() ? 200 : info[2].As().Uint32Value()); const uint32_t random_seed_ = given_named_args ? (!named_args.Has("randomSeed") ? 100 : named_args.Get("randomSeed").As().Uint32Value()) : (info[3].IsUndefined() ? 100 : info[3].As().Uint32Value()); const bool allow_replace_deleted_ = given_named_args ? (!named_args.Has("allowReplaceDeleted") ? false : named_args.Get("allowReplaceDeleted").As().Value()) : (info[4].IsUndefined() ? false : info[4].As().Value()); if (index_) delete index_; try { index_ = new hnswlib::HierarchicalNSW(space_, static_cast(max_elements_), static_cast(m_), static_cast(ef_construction_), static_cast(random_seed_), allow_replace_deleted_); } catch (const std::bad_alloc& err) { index_ = nullptr; Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value readIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() < 1) { Napi::Error::New(env, "Expected 1-2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsUndefined() && !info[1].IsBoolean()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a boolean.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); const bool allow_replace_deleted = info[1].IsUndefined() ? false : info[1].As().Value(); Napi::Function callback = Napi::Function::New(env, [](const Napi::CallbackInfo& info) { return info.Env().Undefined(); }); Napi::Promise::Deferred deferred = Napi::Promise::Deferred::New(info.Env()); LoadHierarchicalNSWIndexWorker* worker = new LoadHierarchicalNSWIndexWorker(filename, allow_replace_deleted, &index_, &space_, deferred, callback); worker->Queue(); return worker->deferred.Promise(); } Napi::Value readIndexSync(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() < 1) { Napi::Error::New(env, "Expected 1-2 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsUndefined() && !info[1].IsBoolean()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a boolean.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); const bool replace_deleted = info[1].IsUndefined() ? false : info[1].As().Value(); if (index_) delete index_; try { std::ifstream ifs(filename); if (ifs.is_open()) { ifs.close(); } else { throw std::runtime_error("failed to open file: " + filename); } index_ = new hnswlib::HierarchicalNSW(space_, filename, false, 0, replace_deleted); } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } catch (const std::bad_alloc& err) { Napi::Error::New(env, "Hnswlib Error: " + std::string(err.what())).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value writeIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); Napi::Function callback = Napi::Function::New(env, [](const Napi::CallbackInfo& info) { return info.Env().Undefined(); }); Napi::Promise::Deferred deferred = Napi::Promise::Deferred::New(info.Env()); SaveHierarchicalNSWIndexWorker* worker = new SaveHierarchicalNSWIndexWorker(filename, &index_, deferred, callback); worker->Queue(); return worker->deferred.Promise(); } Napi::Value writeIndexSync(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsString()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a string.").ThrowAsJavaScriptException(); return env.Null(); } const std::string filename = info[0].As().ToString(); if (index_ == nullptr) { Napi::Error::New(env, "Search index is not constructed.").ThrowAsJavaScriptException(); return env.Null(); } index_->saveIndex(filename); return env.Null(); } Napi::Value resizeIndex(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } const uint32_t new_max_elements = info[0].As().Uint32Value(); try { index_->resizeIndex(static_cast(new_max_elements)); } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } catch (const std::bad_alloc& err) { Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value addPoint(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() < 2) { Napi::Error::New(env, "Expected 2-3 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsArray()) { Napi::TypeError::New(env, "Invalid the first argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[2].IsUndefined() && !info[2].IsBoolean()) { Napi::TypeError::New(env, "Invalid the third argument type, must be a boolean.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } Napi::Array arr = info[0].As(); if (arr.Length() != dim_) { Napi::Error::New(env, "Invalid the given array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } std::vector vec(dim_, 0.0); for (uint32_t i = 0; i < dim_; i++) { Napi::Value val = arr[i]; vec[i] = val.As().FloatValue(); } if (normalize_) normalizePoint(vec); const uint32_t idx = info[1].As().Uint32Value(); const bool replace_deleted = info[2].IsUndefined() ? false : info[2].As().Value(); try { index_->addPoint(reinterpret_cast(vec.data()), static_cast(idx), replace_deleted); } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value markDelete(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } const uint32_t idx = info[0].As().Uint32Value(); try { index_->markDelete(static_cast(idx)); } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value unmarkDelete(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } const uint32_t idx = info[0].As().Uint32Value(); try { index_->unmarkDelete(static_cast(idx)); } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return env.Null(); } return env.Null(); } Napi::Value searchKnn(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() < 2) { Napi::Error::New(env, "Expected 2-3 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsArray()) { Napi::TypeError::New(env, "Invalid the first argument type, must be an Array.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[1].IsNumber()) { Napi::TypeError::New(env, "Invalid the second argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (!info[2].IsUndefined() && !info[2].IsFunction()) { Napi::TypeError::New(env, "Invalid the third argument type, must be a function.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } Napi::Array arr = info[0].As(); if (arr.Length() != dim_) { Napi::Error::New(env, "Invalid the given array length (expected " + std::to_string(dim_) + ", but got " + std::to_string(arr.Length()) + ").") .ThrowAsJavaScriptException(); return env.Null(); } CustomFilterFunctor* filterFn = nullptr; if (info[2].IsFunction()) { try { filterFn = new CustomFilterFunctor(env, info[2].As()); } catch (const std::bad_alloc& err) { Napi::Error::New(env, err.what()).ThrowAsJavaScriptException(); return env.Null(); } } const uint32_t k = info[1].As().Uint32Value(); if (k > index_->max_elements_) { Napi::Error::New(env, "Invalid the number of k-nearest neighbors (cannot be given a value greater than `maxElements`: " + std::to_string(index_->max_elements_) + ").") .ThrowAsJavaScriptException(); return env.Null(); } std::vector vec(dim_, 0.0); for (uint32_t i = 0; i < dim_; i++) { Napi::Value val = arr[i]; vec[i] = val.As().FloatValue(); } if (normalize_) normalizePoint(vec); std::priority_queue> knn = index_->searchKnn(reinterpret_cast(vec.data()), static_cast(k), filterFn); const size_t n_results = knn.size(); Napi::Array arr_distances = Napi::Array::New(env, n_results); Napi::Array arr_neighbors = Napi::Array::New(env, n_results); for (int32_t i = static_cast(n_results) - 1; i >= 0; i--) { const std::pair& nn = knn.top(); arr_distances[i] = Napi::Number::New(env, nn.first); arr_neighbors[i] = Napi::Number::New(env, nn.second); knn.pop(); } if (filterFn) delete filterFn; Napi::Object results = Napi::Object::New(env); results.Set("distances", arr_distances); results.Set("neighbors", arr_neighbors); return results; } Napi::Value getIdsList(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (index_ == nullptr) return Napi::Array::New(env, 0); uint32_t n_elements = index_->label_lookup_.size(); Napi::Array ids = Napi::Array::New(env, n_elements); uint32_t counter = 0; for (auto kv : index_->label_lookup_) { ids[counter] = Napi::Number::New(env, kv.first); counter++; } return ids; } Napi::Value getPoint(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) return Napi::Array::New(env, 0); try { const uint32_t label = info[0].As().Uint32Value(); std::vector vec = index_->getDataByLabel(static_cast(label)); Napi::Array point = Napi::Array::New(env, vec.size()); for (size_t i = 0; i < vec.size(); i++) point[i] = Napi::Number::New(env, vec[i]); return point; } catch (const std::runtime_error& e) { Napi::Error::New(env, "Hnswlib Error: " + std::string(e.what())).ThrowAsJavaScriptException(); return Napi::Array::New(env, 0); } } Napi::Value getMaxElements(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (index_ == nullptr) return Napi::Number::New(env, 0); return Napi::Number::New(env, index_->max_elements_); } Napi::Value getCurrentCount(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (index_ == nullptr) return Napi::Number::New(env, 0); return Napi::Number::New(env, index_->cur_element_count); } Napi::Value getNumDimensions(const Napi::CallbackInfo& info) { return Napi::Number::New(info.Env(), dim_); } Napi::Value getEf(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } return Napi::Number::New(env, index_->ef_); } Napi::Value setEf(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (info.Length() != 1) { Napi::Error::New(env, "Expected 1 arguments, but got " + std::to_string(info.Length()) + ".") .ThrowAsJavaScriptException(); return env.Null(); } if (!info[0].IsNumber()) { Napi::TypeError::New(env, "Invalid the first argument type, must be a number.").ThrowAsJavaScriptException(); return env.Null(); } if (index_ == nullptr) { Napi::Error::New(env, "Search index has not been initialized, call `initIndex` in advance.").ThrowAsJavaScriptException(); return env.Null(); } const uint32_t ef = info[0].As().Uint32Value(); index_->setEf(static_cast(ef)); return env.Null(); } }; Napi::Object Init(Napi::Env env, Napi::Object exports) { L2Space::Init(env, exports); InnerProductSpace::Init(env, exports); BruteforceSearch::Init(env, exports); HierarchicalNSW::Init(env, exports); return exports; } NODE_API_MODULE(addon, Init)