hnswlib 0.8.1 → 0.9.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 33dd2f9cd8656dfe469b4b9baf319ae4e0b78826dec58da8ae78516077a6fbc2
-  data.tar.gz: ad32e54a325136587ae12716243b88ed1269c7bc9cd4e92e9aa8896b2517d4d2
+  metadata.gz: 39458b40736b4a330a0c1769a3adcf37d9d40614b1270d33885932244e943b74
+  data.tar.gz: d49bf8e158c55235fdeca0656cd3ddab6a07e5abc1468eb8917b21f9baf5d89c
 SHA512:
-  metadata.gz: 361822ebf216d8f8ba3abf8314e52b1a72887fb2a03aea29572940376326b128442fc007fa7de6a6e925facb0b9295076fb10a98ff053aad7e0fbe41e4973c2a
-  data.tar.gz: 38f9f4b10665d9e87940b23fae51a15640229367f6bb74a7dba176bab6e798c8fad8abce106c07e38e0acc9046949f80208f9f9b30e3341aceb5a6193dac3aca
+  metadata.gz: ec8560f2b7aa9bb5df098b9c863e4602465a42e8ea27a19be55b9d62e9ba0d32fa70ccbcb13323b771370df5ded9370f9f91d17c8d63aecc18d2b45e29df0149
+  data.tar.gz: c9b434fb313f41e2f0570070944eb6eb3b4432a65fce2c63ca6592ab1de30f9b04834d87e1bd40db465dceea0eaac2279378eda537f6b7d51c2bdb0aa06e1ab6

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## [[0.9.0](https://github.com/yoshoku/hnswlib.rb/compare/v0.8.1...v0.9.0)] - 2023-12-16
+
+- Update bundled hnswlib version to 0.8.0.
+  - Multi-vector document search and epsilon search, which are added only to the C++ version, are not supported. These features will be supported in future release.
+
 ## [0.8.1] - 2023-03-18
 
 - Update the type declarations of HierarchicalNSW and BruteforceSearch along with recent changes.

data/ext/hnswlib/hnswlibext.hpp

@@ -528,10 +528,6 @@ private:
       free(index->linkLists_);
       index->linkLists_ = nullptr;
     }
-    if (index->visited_list_pool_) {
-      delete index->visited_list_pool_;
-      index->visited_list_pool_ = nullptr;
-    }
 
     try {
       index->loadIndex(filename, space);

data/ext/hnswlib/src/bruteforce.h

@@ -85,10 +85,16 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
 
 
     void removePoint(labeltype cur_external) {
-        size_t cur_c = dict_external_to_internal[cur_external];
+        std::unique_lock<std::mutex> lock(index_lock);
 
-        dict_external_to_internal.erase(cur_external);
+        auto found = dict_external_to_internal.find(cur_external);
+        if (found == dict_external_to_internal.end()) {
+            return;
+        }
+
+        dict_external_to_internal.erase(found);
 
+        size_t cur_c = found->second;
         labeltype label = *((labeltype*)(data_ + size_per_element_ * (cur_element_count-1) + data_size_));
         dict_external_to_internal[label] = cur_c;
         memcpy(data_ + size_per_element_ * cur_c,
@@ -107,7 +113,7 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
             dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_);
             labeltype label = *((labeltype*) (data_ + size_per_element_ * i + data_size_));
             if ((!isIdAllowed) || (*isIdAllowed)(label)) {
-                topResults.push(std::pair<dist_t, labeltype>(dist, label));
+                topResults.emplace(dist, label);
             }
         }
         dist_t lastdist = topResults.empty() ? std::numeric_limits<dist_t>::max() : topResults.top().first;
@@ -116,7 +122,7 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
             if (dist <= lastdist) {
                 labeltype label = *((labeltype *) (data_ + size_per_element_ * i + data_size_));
                 if ((!isIdAllowed) || (*isIdAllowed)(label)) {
-                    topResults.push(std::pair<dist_t, labeltype>(dist, label));
+                    topResults.emplace(dist, label);
                 }
                 if (topResults.size() > k)
                     topResults.pop();

data/ext/hnswlib/src/hnswalg.h

@@ -8,6 +8,7 @@
 #include <assert.h>
 #include <unordered_set>
 #include <list>
+#include <memory>
 
 namespace hnswlib {
 typedef unsigned int tableint;
@@ -33,7 +34,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
     double mult_{0.0}, revSize_{0.0};
     int maxlevel_{0};
 
-    VisitedListPool *visited_list_pool_{nullptr};
+    std::unique_ptr<VisitedListPool> visited_list_pool_{nullptr};
 
     // Locks operations with element by label value
     mutable std::vector<std::mutex> label_op_locks_;
@@ -93,8 +94,8 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         size_t ef_construction = 200,
         size_t random_seed = 100,
         bool allow_replace_deleted = false)
-        : link_list_locks_(max_elements),
-            label_op_locks_(MAX_LABEL_OPERATION_LOCKS),
+        : label_op_locks_(MAX_LABEL_OPERATION_LOCKS),
+            link_list_locks_(max_elements),
             element_levels_(max_elements),
             allow_replace_deleted_(allow_replace_deleted) {
         max_elements_ = max_elements;
@@ -102,7 +103,13 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         data_size_ = s->get_data_size();
         fstdistfunc_ = s->get_dist_func();
         dist_func_param_ = s->get_dist_func_param();
-        M_ = M;
+        if ( M <= 10000 ) {
+            M_ = M;
+        } else {
+            HNSWERR << "warning: M parameter exceeds 10000 which may lead to adverse effects." << std::endl;
+            HNSWERR << "         Cap to 10000 will be applied for the rest of the processing." << std::endl;
+            M_ = 10000;
+        }
         maxM_ = M_;
         maxM0_ = M_ * 2;
         ef_construction_ = std::max(ef_construction, M_);
@@ -123,7 +130,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
         cur_element_count = 0;
 
-        visited_list_pool_ = new VisitedListPool(1, max_elements);
+        visited_list_pool_ = std::unique_ptr<VisitedListPool>(new VisitedListPool(1, max_elements));
 
         // initializations for special treatment of the first node
         enterpoint_node_ = -1;
@@ -139,13 +146,20 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
 
     ~HierarchicalNSW() {
+        clear();
+    }
+
+    void clear() {
         free(data_level0_memory_);
+        data_level0_memory_ = nullptr;
         for (tableint i = 0; i < cur_element_count; i++) {
             if (element_levels_[i] > 0)
                 free(linkLists_[i]);
         }
         free(linkLists_);
-        delete visited_list_pool_;
+        linkLists_ = nullptr;
+        cur_element_count = 0;
+        visited_list_pool_.reset(nullptr);
     }
 
 
@@ -292,9 +306,15 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
     }
 
 
-    template <bool has_deletions, bool collect_metrics = false>
+    // bare_bone_search means there is no check for deletions and stop condition is ignored in return of extra performance
+    template <bool bare_bone_search = true, bool collect_metrics = false>
     std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst>
-    searchBaseLayerST(tableint ep_id, const void *data_point, size_t ef, BaseFilterFunctor* isIdAllowed = nullptr) const {
+    searchBaseLayerST(
+        tableint ep_id,
+        const void *data_point,
+        size_t ef,
+        BaseFilterFunctor* isIdAllowed = nullptr,
+        BaseSearchStopCondition<dist_t>* stop_condition = nullptr) const {
         VisitedList *vl = visited_list_pool_->getFreeVisitedList();
         vl_type *visited_array = vl->mass;
         vl_type visited_array_tag = vl->curV;
@@ -303,10 +323,15 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> candidate_set;
 
         dist_t lowerBound;
-        if ((!has_deletions || !isMarkedDeleted(ep_id)) && ((!isIdAllowed) || (*isIdAllowed)(getExternalLabel(ep_id)))) {
-            dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_);
+        if (bare_bone_search ||
+            (!isMarkedDeleted(ep_id) && ((!isIdAllowed) || (*isIdAllowed)(getExternalLabel(ep_id))))) {
+            char* ep_data = getDataByInternalId(ep_id);
+            dist_t dist = fstdistfunc_(data_point, ep_data, dist_func_param_);
             lowerBound = dist;
             top_candidates.emplace(dist, ep_id);
+            if (!bare_bone_search && stop_condition) {
+                stop_condition->add_point_to_result(getExternalLabel(ep_id), ep_data, dist);
+            }
             candidate_set.emplace(-dist, ep_id);
         } else {
             lowerBound = std::numeric_limits<dist_t>::max();
@@ -317,9 +342,19 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
         while (!candidate_set.empty()) {
             std::pair<dist_t, tableint> current_node_pair = candidate_set.top();
+            dist_t candidate_dist = -current_node_pair.first;
 
-            if ((-current_node_pair.first) > lowerBound &&
-                (top_candidates.size() == ef || (!isIdAllowed && !has_deletions))) {
+            bool flag_stop_search;
+            if (bare_bone_search) {
+                flag_stop_search = candidate_dist > lowerBound;
+            } else {
+                if (stop_condition) {
+                    flag_stop_search = stop_condition->should_stop_search(candidate_dist, lowerBound);
+                } else {
+                    flag_stop_search = candidate_dist > lowerBound && top_candidates.size() == ef;
+                }
+            }
+            if (flag_stop_search) {
                 break;
             }
             candidate_set.pop();
@@ -354,7 +389,14 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
                     char *currObj1 = (getDataByInternalId(candidate_id));
                     dist_t dist = fstdistfunc_(data_point, currObj1, dist_func_param_);
 
-                    if (top_candidates.size() < ef || lowerBound > dist) {
+                    bool flag_consider_candidate;
+                    if (!bare_bone_search && stop_condition) {
+                        flag_consider_candidate = stop_condition->should_consider_candidate(dist, lowerBound);
+                    } else {
+                        flag_consider_candidate = top_candidates.size() < ef || lowerBound > dist;
+                    }
+
+                    if (flag_consider_candidate) {
                         candidate_set.emplace(-dist, candidate_id);
 #ifdef USE_SSE
                         _mm_prefetch(data_level0_memory_ + candidate_set.top().second * size_data_per_element_ +
@@ -362,11 +404,30 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
                                         _MM_HINT_T0);  ////////////////////////
 #endif
 
-                        if ((!has_deletions || !isMarkedDeleted(candidate_id)) && ((!isIdAllowed) || (*isIdAllowed)(getExternalLabel(candidate_id))))
+                        if (bare_bone_search ||
+                            (!isMarkedDeleted(candidate_id) && ((!isIdAllowed) || (*isIdAllowed)(getExternalLabel(candidate_id))))) {
                             top_candidates.emplace(dist, candidate_id);
+                            if (!bare_bone_search && stop_condition) {
+                                stop_condition->add_point_to_result(getExternalLabel(candidate_id), currObj1, dist);
+                            }
+                        }
 
-                        if (top_candidates.size() > ef)
+                        bool flag_remove_extra = false;
+                        if (!bare_bone_search && stop_condition) {
+                            flag_remove_extra = stop_condition->should_remove_extra();
+                        } else {
+                            flag_remove_extra = top_candidates.size() > ef;
+                        }
+                        while (flag_remove_extra) {
+                            tableint id = top_candidates.top().second;
                             top_candidates.pop();
+                            if (!bare_bone_search && stop_condition) {
+                                stop_condition->remove_point_from_result(getExternalLabel(id), getDataByInternalId(id), dist);
+                                flag_remove_extra = stop_condition->should_remove_extra();
+                            } else {
+                                flag_remove_extra = top_candidates.size() > ef;
+                            }
+                        }
 
                         if (!top_candidates.empty())
                             lowerBound = top_candidates.top().first;
@@ -381,8 +442,8 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
 
     void getNeighborsByHeuristic2(
-            std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> &top_candidates,
-    const size_t M) {
+        std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> &top_candidates,
+        const size_t M) {
         if (top_candidates.size() < M) {
             return;
         }
@@ -574,8 +635,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         if (new_max_elements < cur_element_count)
             throw std::runtime_error("Cannot resize, max element is less than the current number of elements");
 
-        delete visited_list_pool_;
-        visited_list_pool_ = new VisitedListPool(1, new_max_elements);
+        visited_list_pool_.reset(new VisitedListPool(1, new_max_elements));
 
         element_levels_.resize(new_max_elements);
 
@@ -596,6 +656,32 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         max_elements_ = new_max_elements;
     }
 
+    size_t indexFileSize() const {
+        size_t size = 0;
+        size += sizeof(offsetLevel0_);
+        size += sizeof(max_elements_);
+        size += sizeof(cur_element_count);
+        size += sizeof(size_data_per_element_);
+        size += sizeof(label_offset_);
+        size += sizeof(offsetData_);
+        size += sizeof(maxlevel_);
+        size += sizeof(enterpoint_node_);
+        size += sizeof(maxM_);
+
+        size += sizeof(maxM0_);
+        size += sizeof(M_);
+        size += sizeof(mult_);
+        size += sizeof(ef_construction_);
+
+        size += cur_element_count * size_data_per_element_;
+
+        for (size_t i = 0; i < cur_element_count; i++) {
+            unsigned int linkListSize = element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0;
+            size += sizeof(linkListSize);
+            size += linkListSize;
+        }
+        return size;
+    }
 
     void saveIndex(const std::string &location) {
         std::ofstream output(location, std::ios::binary);
@@ -634,6 +720,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         if (!input.is_open())
             throw std::runtime_error("Cannot open file");
 
+        clear();
         // get file size:
         input.seekg(0, input.end);
         std::streampos total_filesize = input.tellg();
@@ -699,7 +786,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         std::vector<std::mutex>(max_elements).swap(link_list_locks_);
         std::vector<std::mutex>(MAX_LABEL_OPERATION_LOCKS).swap(label_op_locks_);
 
-        visited_list_pool_ = new VisitedListPool(1, max_elements);
+        visited_list_pool_.reset(new VisitedListPool(1, max_elements));
 
         linkLists_ = (char **) malloc(sizeof(void *) * max_elements);
         if (linkLists_ == nullptr)
@@ -753,7 +840,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         size_t dim = *((size_t *) dist_func_param_);
         std::vector<data_t> data;
         data_t* data_ptr = (data_t*) data_ptrv;
-        for (int i = 0; i < dim; i++) {
+        for (size_t i = 0; i < dim; i++) {
             data.push_back(*data_ptr);
             data_ptr += 1;
         }
@@ -1217,11 +1304,12 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         }
 
         std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> top_candidates;
-        if (num_deleted_) {
-            top_candidates = searchBaseLayerST<true, true>(
+        bool bare_bone_search = !num_deleted_ && !isIdAllowed;
+        if (bare_bone_search) {
+            top_candidates = searchBaseLayerST<true>(
                     currObj, query_data, std::max(ef_, k), isIdAllowed);
         } else {
-            top_candidates = searchBaseLayerST<false, true>(
+            top_candidates = searchBaseLayerST<false>(
                     currObj, query_data, std::max(ef_, k), isIdAllowed);
         }
 
@@ -1237,6 +1325,60 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
     }
 
 
+    std::vector<std::pair<dist_t, labeltype >>
+    searchStopConditionClosest(
+        const void *query_data,
+        BaseSearchStopCondition<dist_t>& stop_condition,
+        BaseFilterFunctor* isIdAllowed = nullptr) const {
+        std::vector<std::pair<dist_t, labeltype >> result;
+        if (cur_element_count == 0) return result;
+
+        tableint currObj = enterpoint_node_;
+        dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_);
+
+        for (int level = maxlevel_; level > 0; level--) {
+            bool changed = true;
+            while (changed) {
+                changed = false;
+                unsigned int *data;
+
+                data = (unsigned int *) get_linklist(currObj, level);
+                int size = getListCount(data);
+                metric_hops++;
+                metric_distance_computations+=size;
+
+                tableint *datal = (tableint *) (data + 1);
+                for (int i = 0; i < size; i++) {
+                    tableint cand = datal[i];
+                    if (cand < 0 || cand > max_elements_)
+                        throw std::runtime_error("cand error");
+                    dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_);
+
+                    if (d < curdist) {
+                        curdist = d;
+                        currObj = cand;
+                        changed = true;
+                    }
+                }
+            }
+        }
+
+        std::priority_queue<std::pair<dist_t, tableint>, std::vector<std::pair<dist_t, tableint>>, CompareByFirst> top_candidates;
+        top_candidates = searchBaseLayerST<false>(currObj, query_data, 0, isIdAllowed, &stop_condition);
+
+        size_t sz = top_candidates.size();
+        result.resize(sz);
+        while (!top_candidates.empty()) {
+            result[--sz] = top_candidates.top();
+            top_candidates.pop();
+        }
+
+        stop_condition.filter_results(result);
+
+        return result;
+    }
+
+
     void checkIntegrity() {
         int connections_checked = 0;
         std::vector <int > inbound_connections_num(cur_element_count, 0);
@@ -1247,7 +1389,6 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
                 tableint *data = (tableint *) (ll_cur + 1);
                 std::unordered_set<tableint> s;
                 for (int j = 0; j < size; j++) {
-                    assert(data[j] > 0);
                     assert(data[j] < cur_element_count);
                     assert(data[j] != i);
                     inbound_connections_num[data[j]]++;

data/ext/hnswlib/src/hnswlib.h

@@ -1,4 +1,13 @@
 #pragma once
+
+// https://github.com/nmslib/hnswlib/pull/508
+// This allows others to provide their own error stream (e.g. RcppHNSW)
+#ifndef HNSWLIB_ERR_OVERRIDE
+  #define HNSWERR std::cerr
+#else
+  #define HNSWERR HNSWLIB_ERR_OVERRIDE
+#endif
+
 #ifndef NO_MANUAL_VECTORIZATION
 #if (defined(__SSE__) || _M_IX86_FP > 0 || defined(_M_AMD64) || defined(_M_X64))
 #define USE_SSE
@@ -15,7 +24,7 @@
 #ifdef _MSC_VER
 #include <intrin.h>
 #include <stdexcept>
-void cpuid(int32_t out[4], int32_t eax, int32_t ecx) {
+static void cpuid(int32_t out[4], int32_t eax, int32_t ecx) {
     __cpuidex(out, eax, ecx);
 }
 static __int64 xgetbv(unsigned int x) {
@@ -122,6 +131,24 @@ class BaseFilterFunctor {
     virtual ~BaseFilterFunctor() {};
 };
 
+template<typename dist_t>
+class BaseSearchStopCondition {
+ public:
+    virtual void add_point_to_result(labeltype label, const void *datapoint, dist_t dist) = 0;
+
+    virtual void remove_point_from_result(labeltype label, const void *datapoint, dist_t dist) = 0;
+
+    virtual bool should_stop_search(dist_t candidate_dist, dist_t lowerBound) = 0;
+
+    virtual bool should_consider_candidate(dist_t candidate_dist, dist_t lowerBound) = 0;
+
+    virtual bool should_remove_extra() = 0;
+
+    virtual void filter_results(std::vector<std::pair<dist_t, labeltype >> &candidates) = 0;
+
+    virtual ~BaseSearchStopCondition() {}
+};
+
 template <typename T>
 class pairGreater {
  public:
@@ -196,5 +223,6 @@ AlgorithmInterface<dist_t>::searchKnnCloserFirst(const void* query_data, size_t
 
 #include "space_l2.h"
 #include "space_ip.h"
+#include "stop_condition.h"
 #include "bruteforce.h"
 #include "hnswalg.h"

data/ext/hnswlib/src/space_ip.h

@@ -157,19 +157,44 @@ InnerProductSIMD16ExtAVX512(const void *pVect1v, const void *pVect2v, const void
 
     __m512 sum512 = _mm512_set1_ps(0);
 
-    while (pVect1 < pEnd1) {
-        //_mm_prefetch((char*)(pVect2 + 16), _MM_HINT_T0);
+    size_t loop = qty16 / 4;
 
+    while (loop--) {
         __m512 v1 = _mm512_loadu_ps(pVect1);
-        pVect1 += 16;
         __m512 v2 = _mm512_loadu_ps(pVect2);
+        pVect1 += 16;
+        pVect2 += 16;
+
+        __m512 v3 = _mm512_loadu_ps(pVect1);
+        __m512 v4 = _mm512_loadu_ps(pVect2);
+        pVect1 += 16;
+        pVect2 += 16;
+
+        __m512 v5 = _mm512_loadu_ps(pVect1);
+        __m512 v6 = _mm512_loadu_ps(pVect2);
+        pVect1 += 16;
         pVect2 += 16;
-        sum512 = _mm512_add_ps(sum512, _mm512_mul_ps(v1, v2));
+
+        __m512 v7 = _mm512_loadu_ps(pVect1);
+        __m512 v8 = _mm512_loadu_ps(pVect2);
+        pVect1 += 16;
+        pVect2 += 16;
+
+        sum512 = _mm512_fmadd_ps(v1, v2, sum512);
+        sum512 = _mm512_fmadd_ps(v3, v4, sum512);
+        sum512 = _mm512_fmadd_ps(v5, v6, sum512);
+        sum512 = _mm512_fmadd_ps(v7, v8, sum512);
     }
 
-    _mm512_store_ps(TmpRes, sum512);
-    float sum = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7] + TmpRes[8] + TmpRes[9] + TmpRes[10] + TmpRes[11] + TmpRes[12] + TmpRes[13] + TmpRes[14] + TmpRes[15];
+    while (pVect1 < pEnd1) {
+        __m512 v1 = _mm512_loadu_ps(pVect1);
+        __m512 v2 = _mm512_loadu_ps(pVect2);
+        pVect1 += 16;
+        pVect2 += 16;
+        sum512 = _mm512_fmadd_ps(v1, v2, sum512);
+    }
 
+    float sum = _mm512_reduce_add_ps(sum512);
     return sum;
 }
 

data/ext/hnswlib/src/stop_condition.h

@@ -0,0 +1,276 @@
+#pragma once
+#include "space_l2.h"
+#include "space_ip.h"
+#include <assert.h>
+#include <unordered_map>
+
+namespace hnswlib {
+
+template<typename DOCIDTYPE>
+class BaseMultiVectorSpace : public SpaceInterface<float> {
+ public:
+    virtual DOCIDTYPE get_doc_id(const void *datapoint) = 0;
+
+    virtual void set_doc_id(void *datapoint, DOCIDTYPE doc_id) = 0;
+};
+
+
+template<typename DOCIDTYPE>
+class MultiVectorL2Space : public BaseMultiVectorSpace<DOCIDTYPE> {
+    DISTFUNC<float> fstdistfunc_;
+    size_t data_size_;
+    size_t vector_size_;
+    size_t dim_;
+
+ public:
+    MultiVectorL2Space(size_t dim) {
+        fstdistfunc_ = L2Sqr;
+#if defined(USE_SSE) || defined(USE_AVX) || defined(USE_AVX512)
+    #if defined(USE_AVX512)
+        if (AVX512Capable())
+            L2SqrSIMD16Ext = L2SqrSIMD16ExtAVX512;
+        else if (AVXCapable())
+            L2SqrSIMD16Ext = L2SqrSIMD16ExtAVX;
+    #elif defined(USE_AVX)
+        if (AVXCapable())
+            L2SqrSIMD16Ext = L2SqrSIMD16ExtAVX;
+    #endif
+
+        if (dim % 16 == 0)
+            fstdistfunc_ = L2SqrSIMD16Ext;
+        else if (dim % 4 == 0)
+            fstdistfunc_ = L2SqrSIMD4Ext;
+        else if (dim > 16)
+            fstdistfunc_ = L2SqrSIMD16ExtResiduals;
+        else if (dim > 4)
+            fstdistfunc_ = L2SqrSIMD4ExtResiduals;
+#endif
+        dim_ = dim;
+        vector_size_ = dim * sizeof(float);
+        data_size_ = vector_size_ + sizeof(DOCIDTYPE);
+    }
+
+    size_t get_data_size() override {
+        return data_size_;
+    }
+
+    DISTFUNC<float> get_dist_func() override {
+        return fstdistfunc_;
+    }
+
+    void *get_dist_func_param() override {
+        return &dim_;
+    }
+
+    DOCIDTYPE get_doc_id(const void *datapoint) override {
+        return *(DOCIDTYPE *)((char *)datapoint + vector_size_);
+    }
+
+    void set_doc_id(void *datapoint, DOCIDTYPE doc_id) override {
+        *(DOCIDTYPE*)((char *)datapoint + vector_size_) = doc_id;
+    }
+
+    ~MultiVectorL2Space() {}
+};
+
+
+template<typename DOCIDTYPE>
+class MultiVectorInnerProductSpace : public BaseMultiVectorSpace<DOCIDTYPE> {
+    DISTFUNC<float> fstdistfunc_;
+    size_t data_size_;
+    size_t vector_size_;
+    size_t dim_;
+
+ public:
+    MultiVectorInnerProductSpace(size_t dim) {
+        fstdistfunc_ = InnerProductDistance;
+#if defined(USE_AVX) || defined(USE_SSE) || defined(USE_AVX512)
+    #if defined(USE_AVX512)
+        if (AVX512Capable()) {
+            InnerProductSIMD16Ext = InnerProductSIMD16ExtAVX512;
+            InnerProductDistanceSIMD16Ext = InnerProductDistanceSIMD16ExtAVX512;
+        } else if (AVXCapable()) {
+            InnerProductSIMD16Ext = InnerProductSIMD16ExtAVX;
+            InnerProductDistanceSIMD16Ext = InnerProductDistanceSIMD16ExtAVX;
+        }
+    #elif defined(USE_AVX)
+        if (AVXCapable()) {
+            InnerProductSIMD16Ext = InnerProductSIMD16ExtAVX;
+            InnerProductDistanceSIMD16Ext = InnerProductDistanceSIMD16ExtAVX;
+        }
+    #endif
+    #if defined(USE_AVX)
+        if (AVXCapable()) {
+            InnerProductSIMD4Ext = InnerProductSIMD4ExtAVX;
+            InnerProductDistanceSIMD4Ext = InnerProductDistanceSIMD4ExtAVX;
+        }
+    #endif
+
+        if (dim % 16 == 0)
+            fstdistfunc_ = InnerProductDistanceSIMD16Ext;
+        else if (dim % 4 == 0)
+            fstdistfunc_ = InnerProductDistanceSIMD4Ext;
+        else if (dim > 16)
+            fstdistfunc_ = InnerProductDistanceSIMD16ExtResiduals;
+        else if (dim > 4)
+            fstdistfunc_ = InnerProductDistanceSIMD4ExtResiduals;
+#endif
+        vector_size_ = dim * sizeof(float);
+        data_size_ = vector_size_ + sizeof(DOCIDTYPE);
+    }
+
+    size_t get_data_size() override {
+        return data_size_;
+    }
+
+    DISTFUNC<float> get_dist_func() override {
+        return fstdistfunc_;
+    }
+
+    void *get_dist_func_param() override {
+        return &dim_;
+    }
+
+    DOCIDTYPE get_doc_id(const void *datapoint) override {
+        return *(DOCIDTYPE *)((char *)datapoint + vector_size_);
+    }
+
+    void set_doc_id(void *datapoint, DOCIDTYPE doc_id) override {
+        *(DOCIDTYPE*)((char *)datapoint + vector_size_) = doc_id;
+    }
+
+    ~MultiVectorInnerProductSpace() {}
+};
+
+
+template<typename DOCIDTYPE, typename dist_t>
+class MultiVectorSearchStopCondition : public BaseSearchStopCondition<dist_t> {
+    size_t curr_num_docs_;
+    size_t num_docs_to_search_;
+    size_t ef_collection_;
+    std::unordered_map<DOCIDTYPE, size_t> doc_counter_;
+    std::priority_queue<std::pair<dist_t, DOCIDTYPE>> search_results_;
+    BaseMultiVectorSpace<DOCIDTYPE>& space_;
+
+ public:
+    MultiVectorSearchStopCondition(
+        BaseMultiVectorSpace<DOCIDTYPE>& space,
+        size_t num_docs_to_search,
+        size_t ef_collection = 10)
+        : space_(space) {
+            curr_num_docs_ = 0;
+            num_docs_to_search_ = num_docs_to_search;
+            ef_collection_ = std::max(ef_collection, num_docs_to_search);
+        }
+
+    void add_point_to_result(labeltype label, const void *datapoint, dist_t dist) override {
+        DOCIDTYPE doc_id = space_.get_doc_id(datapoint);
+        if (doc_counter_[doc_id] == 0) {
+            curr_num_docs_ += 1;
+        }
+        search_results_.emplace(dist, doc_id);
+        doc_counter_[doc_id] += 1;
+    }
+
+    void remove_point_from_result(labeltype label, const void *datapoint, dist_t dist) override {
+        DOCIDTYPE doc_id = space_.get_doc_id(datapoint);
+        doc_counter_[doc_id] -= 1;
+        if (doc_counter_[doc_id] == 0) {
+            curr_num_docs_ -= 1;
+        }
+        search_results_.pop();
+    }
+
+    bool should_stop_search(dist_t candidate_dist, dist_t lowerBound) override {
+        bool stop_search = candidate_dist > lowerBound && curr_num_docs_ == ef_collection_;
+        return stop_search;
+    }
+
+    bool should_consider_candidate(dist_t candidate_dist, dist_t lowerBound) override {
+        bool flag_consider_candidate = curr_num_docs_ < ef_collection_ || lowerBound > candidate_dist;
+        return flag_consider_candidate;
+    }
+
+    bool should_remove_extra() override {
+        bool flag_remove_extra = curr_num_docs_ > ef_collection_;
+        return flag_remove_extra;
+    }
+
+    void filter_results(std::vector<std::pair<dist_t, labeltype >> &candidates) override {
+        while (curr_num_docs_ > num_docs_to_search_) {
+            dist_t dist_cand = candidates.back().first;
+            dist_t dist_res = search_results_.top().first;
+            assert(dist_cand == dist_res);
+            DOCIDTYPE doc_id = search_results_.top().second;
+            doc_counter_[doc_id] -= 1;
+            if (doc_counter_[doc_id] == 0) {
+                curr_num_docs_ -= 1;
+            }
+            search_results_.pop();
+            candidates.pop_back();
+        }
+    }
+
+    ~MultiVectorSearchStopCondition() {}
+};
+
+
+template<typename dist_t>
+class EpsilonSearchStopCondition : public BaseSearchStopCondition<dist_t> {
+    float epsilon_;
+    size_t min_num_candidates_;
+    size_t max_num_candidates_;
+    size_t curr_num_items_;
+
+ public:
+    EpsilonSearchStopCondition(float epsilon, size_t min_num_candidates, size_t max_num_candidates) {
+        assert(min_num_candidates <= max_num_candidates);
+        epsilon_ = epsilon;
+        min_num_candidates_ = min_num_candidates;
+        max_num_candidates_ = max_num_candidates;
+        curr_num_items_ = 0;
+    }
+
+    void add_point_to_result(labeltype label, const void *datapoint, dist_t dist) override {
+        curr_num_items_ += 1;
+    }
+
+    void remove_point_from_result(labeltype label, const void *datapoint, dist_t dist) override {
+        curr_num_items_ -= 1;
+    }
+
+    bool should_stop_search(dist_t candidate_dist, dist_t lowerBound) override {
+        if (candidate_dist > lowerBound && curr_num_items_ == max_num_candidates_) {
+            // new candidate can't improve found results
+            return true;
+        }
+        if (candidate_dist > epsilon_ && curr_num_items_ >= min_num_candidates_) {
+            // new candidate is out of epsilon region and
+            // minimum number of candidates is checked
+            return true;
+        }
+        return false;
+    }
+
+    bool should_consider_candidate(dist_t candidate_dist, dist_t lowerBound) override {
+        bool flag_consider_candidate = curr_num_items_ < max_num_candidates_ || lowerBound > candidate_dist;
+        return flag_consider_candidate;
+    }
+
+    bool should_remove_extra() {
+        bool flag_remove_extra = curr_num_items_ > max_num_candidates_;
+        return flag_remove_extra;
+    }
+
+    void filter_results(std::vector<std::pair<dist_t, labeltype >> &candidates) override {
+        while (!candidates.empty() && candidates.back().first > epsilon_) {
+            candidates.pop_back();
+        }
+        while (candidates.size() > max_num_candidates_) {
+            candidates.pop_back();
+        }
+    }
+
+    ~EpsilonSearchStopCondition() {}
+};
+}  // namespace hnswlib

data/lib/hnswlib/version.rb

@@ -3,8 +3,8 @@
 # Hnswlib.rb provides Ruby bindings for the Hnswlib.
 module Hnswlib
   # The version of Hnswlib.rb you install.
-  VERSION = '0.8.1'
+  VERSION = '0.9.0'
 
   # The version of Hnswlib included with gem.
-  HSWLIB_VERSION = '0.7.0'
+  HSWLIB_VERSION = '0.8.0'
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: hnswlib
 version: !ruby/object:Gem::Version
-  version: 0.8.1
+  version: 0.9.0
 platform: ruby
 authors:
 - yoshoku
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-03-18 00:00:00.000000000 Z
+date: 2023-12-16 00:00:00.000000000 Z
 dependencies: []
 description: Hnswlib.rb provides Ruby bindings for the Hnswlib.
 email:
@@ -30,6 +30,7 @@ files:
 - ext/hnswlib/src/hnswlib.h
 - ext/hnswlib/src/space_ip.h
 - ext/hnswlib/src/space_l2.h
+- ext/hnswlib/src/stop_condition.h
 - ext/hnswlib/src/visited_list_pool.h
 - lib/hnswlib.rb
 - lib/hnswlib/version.rb
@@ -41,6 +42,7 @@ metadata:
   homepage_uri: https://github.com/yoshoku/hnswlib.rb
   source_code_uri: https://github.com/yoshoku/hnswlib.rb
   changelog_uri: https://github.com/yoshoku/hnswlib.rb/blob/main/CHANGELOG.md
+  documentation_uri: https://yoshoku.github.io/hnswlib.rb/doc/
   rubygems_mfa_required: 'true'
 post_install_message:
 rdoc_options: []
@@ -57,7 +59,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.26
+rubygems_version: 3.4.22
 signing_key:
 specification_version: 4
 summary: Ruby bindings for the Hnswlib.