Timeout enabling for HNSW (#163)

* enabling timeout on hnsw

* added timeout support for hnsw batch iterator

* added tests

* checks at least once every graph scan call

* improve test

* added comment

* moved candidate iteration to a function

* NULLing VecSimQueryResult_List after freeing

* improved tests

* review fixes

* added timeout check to searchBottomLayerEP

* test improvement

Author: GuyAv46

src/VecSim/algorithms/hnsw/hnsw_batch_iterator.cpp

@@ -33,13 +33,13 @@ VecSimQueryResult_List HNSW_BatchIterator::prepareResults(candidatesMaxHeap top_
         VecSimQueryResult_SetScore(rl.results[i], top_candidates.top().first);
         top_candidates.pop();
     }
-    rl.code = VecSim_QueryResult_OK;
     return rl;
 }
 
 candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,
                                                 candidatesMinHeap &top_candidates_extras,
-                                                float &lower_bound, idType entry_point) {
+                                                float &lower_bound, idType entry_point,
+                                                VecSimQueryResult_Code *rc) {
 
     candidatesMaxHeap top_candidates(this->allocator);
     if (entry_point == HNSW_INVALID_ID) {
@@ -58,6 +58,11 @@ candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,
         this->visitNode(entry_point);
         candidates.emplace(dist, entry_point);
     }
+    // Checks that we didn't got timeout between iterations.
+    if (__builtin_expect(VecSimIndex::timeoutCallback(this->getTimeoutCtx()), 0)) {
+        *rc = VecSim_QueryResult_TimedOut;
+        return top_candidates;
+    }
 
     // Move extras from previous iteration to the top candidates.
     while (top_candidates.size() < this->hnsw_index->getEf() && !top_candidates_extras.empty()) {
@@ -77,6 +82,10 @@ candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,
         if (curr_node_dist > lower_bound && top_candidates.size() >= this->hnsw_index->getEf()) {
             break;
         }
+        if (__builtin_expect(VecSimIndex::timeoutCallback(this->getTimeoutCtx()), 0)) {
+            *rc = VecSim_QueryResult_TimedOut;
+            return top_candidates;
+        }
         if (top_candidates.size() < this->hnsw_index->getEf() || lower_bound > curr_node_dist) {
             top_candidates.emplace(curr_node_dist, curr_node_id);
             if (top_candidates.size() > this->hnsw_index->getEf()) {
@@ -152,6 +161,7 @@ HNSW_BatchIterator::HNSW_BatchIterator(void *query_vector, HNSWIndex *index_wrap
 VecSimQueryResult_List HNSW_BatchIterator::getNextResults(size_t n_res,
                                                           VecSimQueryResult_Order order) {
 
+    VecSimQueryResult_List batch = {0};
     // If ef_runtime lower than the number of results to return, increase it. Therefore, we assume
     // that the number of results that return from the graph scan is at least n_res (if exist).
     size_t orig_ef = this->hnsw_index->getEf();
@@ -162,26 +172,33 @@ VecSimQueryResult_List HNSW_BatchIterator::getNextResults(size_t n_res,
     // In the first iteration, we search the graph from top bottom to find the initial entry point,
     // and then we scan the graph to get results (layer 0).
     if (this->getResultsCount() == 0) {
-        idType bottom_layer_ep = this->hnsw_index->searchBottomLayerEP(this->getQueryBlob());
+        idType bottom_layer_ep = this->hnsw_index->searchBottomLayerEP(
+            this->getQueryBlob(), this->getTimeoutCtx(), &batch.code);
+        if (VecSim_OK != batch.code) {
+            return batch;
+        }
         this->entry_point = bottom_layer_ep;
     }
     // We ask for at least n_res candidate from the scan. In fact, at most ef results will return,
     // and it could be that ef > n_res.
     auto top_candidates = this->scanGraph(this->candidates, this->top_candidates_extras,
-                                          this->lower_bound, this->entry_point);
+                                          this->lower_bound, this->entry_point, &batch.code);
+    if (VecSim_OK != batch.code) {
+        return batch;
+    }
     // Move the spare results to the "extras" queue if needed, and create the batch results array.
-    auto batch_results = this->prepareResults(top_candidates, n_res);
+    batch = this->prepareResults(top_candidates, n_res);
 
-    this->updateResultsCount(VecSimQueryResult_Len(batch_results));
+    this->updateResultsCount(VecSimQueryResult_Len(batch));
     if (this->getResultsCount() == this->index_wrapper->indexSize()) {
         this->depleted = true;
     }
     // By default, results are ordered by score.
     if (order == BY_ID) {
-        sort_results_by_id(batch_results);
+        sort_results_by_id(batch);
     }
     this->hnsw_index->setEf(orig_ef);
-    return batch_results;
+    return batch;
 }
 
 bool HNSW_BatchIterator::isDepleted() {

src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h

@@ -29,7 +29,7 @@ class HNSW_BatchIterator : public VecSimBatchIterator {
 
     candidatesMaxHeap scanGraph(candidatesMinHeap &candidates,
                                 candidatesMinHeap &spare_top_candidates, float &lower_bound,
-                                idType entry_point);
+                                idType entry_point, VecSimQueryResult_Code *rc);
     VecSimQueryResult_List
     prepareResults(vecsim_stl::max_priority_queue<pair<float, idType>> top_candidates,
                    size_t n_res);

src/VecSim/algorithms/hnsw/hnsw_wrapper.cpp

@@ -95,6 +95,7 @@ void HNSWIndex::setEf(size_t ef) { this->hnsw->setEf(ef); }
 VecSimQueryResult_List HNSWIndex::topKQuery(const void *query_data, size_t k,
                                             VecSimQueryParams *queryParams) {
     VecSimQueryResult_List rl = {0};
+    void *timeoutCtx = nullptr;
     try {
         this->last_mode = STANDARD_KNN;
         float normalized_data[this->dim]; // This will be use only if metric == VecSimMetric_Cosine
@@ -108,11 +109,15 @@ VecSimQueryResult_List HNSWIndex::topKQuery(const void *query_data, size_t k,
         size_t originalEF = this->hnsw->getEf();
 
         if (queryParams) {
+            timeoutCtx = queryParams->timeoutCtx;
             if (queryParams->hnswRuntimeParams.efRuntime != 0) {
                 hnsw->setEf(queryParams->hnswRuntimeParams.efRuntime);
             }
         }
-        auto knn_res = hnsw->searchKnn(query_data, k);
+        auto knn_res = hnsw->searchKnn(query_data, k, timeoutCtx, &rl.code);
+        if (VecSim_OK != rl.code) {
+            return rl;
+        }
         rl.results = array_new_len<VecSimQueryResult>(knn_res.size(), knn_res.size());
         for (int i = (int)knn_res.size() - 1; i >= 0; --i) {
             VecSimQueryResult_SetId(rl.results[i], knn_res.top().second);
@@ -123,7 +128,6 @@ VecSimQueryResult_List HNSWIndex::topKQuery(const void *query_data, size_t k,
         hnsw->setEf(originalEF);
         assert(hnsw->getEf() == originalEF);
 
-        rl.code = VecSim_QueryResult_OK;
     } catch (...) {
         rl.code = VecSim_QueryResult_Err;
     }

src/VecSim/algorithms/hnsw/hnswlib.h

@@ -35,6 +35,8 @@ typedef unsigned int linklistsizeint;
 
 template <typename dist_t>
 using candidatesMaxHeap = vecsim_stl::max_priority_queue<pair<dist_t, tableint>>;
+template <typename dist_t>
+using candidatesLabelsMaxHeap = vecsim_stl::max_priority_queue<pair<dist_t, labeltype>>;
 
 template <typename dist_t>
 class HierarchicalNSW : public VecsimBaseObject {
@@ -110,8 +112,17 @@ class HierarchicalNSW : public VecsimBaseObject {
                           size_t Mcurmax, tableint *node_neighbors,
                           const vecsim_stl::set<tableint> &orig_neighbors, tableint *removed_links,
                           size_t *removed_links_num);
+    inline dist_t processCandidate(tableint curNodeId, const void *data_point, size_t layer,
+                                   size_t ef, tag_t visited_tag,
+                                   candidatesMaxHeap<dist_t> &top_candidates,
+                                   candidatesMaxHeap<dist_t> &candidates_set,
+                                   dist_t lowerBound) const;
     candidatesMaxHeap<dist_t> searchLayer(tableint ep_id, const void *data_point, size_t layer,
                                           size_t ef) const;
+    candidatesLabelsMaxHeap<dist_t> searchBottomLayer_WithTimeout(tableint ep_id,
+                                                                  const void *data_point, size_t ef,
+                                                                  size_t k, void *timeoutCtx,
+                                                                  VecSimQueryResult_Code *rc) const;
     void getNeighborsByHeuristic2(candidatesMaxHeap<dist_t> &top_candidates, size_t M);
     tableint mutuallyConnectNewElement(tableint cur_c, candidatesMaxHeap<dist_t> &top_candidates,
                                        size_t level);
@@ -144,9 +155,10 @@ class HierarchicalNSW : public VecsimBaseObject {
     bool removePoint(labeltype label);
     void addPoint(const void *data_point, labeltype label);
     dist_t getDistanceByLabelFromPoint(labeltype label, const void *data_point);
-    tableint searchBottomLayerEP(const void *query_data) const;
-    vecsim_stl::max_priority_queue<pair<dist_t, labeltype>> searchKnn(const void *query_data,
-                                                                      size_t k) const;
+    tableint searchBottomLayerEP(const void *query_data, void *timeoutCtx,
+                                 VecSimQueryResult_Code *rc) const;
+    vecsim_stl::max_priority_queue<pair<dist_t, labeltype>>
+    searchKnn(const void *query_data, size_t k, void *timeoutCtx, VecSimQueryResult_Code *rc) const;
 };
 
 /**
@@ -337,6 +349,57 @@ void HierarchicalNSW<dist_t>::removeExtraLinks(linklistsizeint *node_ll,
     *removed_links_num = removed_idx;
 }
 
+template <typename dist_t>
+dist_t HierarchicalNSW<dist_t>::processCandidate(tableint curNodeId, const void *data_point,
+                                                 size_t layer, size_t ef, tag_t visited_tag,
+                                                 candidatesMaxHeap<dist_t> &top_candidates,
+                                                 candidatesMaxHeap<dist_t> &candidate_set,
+                                                 dist_t lowerBound) const {
+
+#ifdef ENABLE_PARALLELIZATION
+    std::unique_lock<std::mutex> lock(link_list_locks_[curNodeId]);
+#endif
+    linklistsizeint *node_ll = get_linklist_at_level(curNodeId, layer);
+    size_t links_num = getListCount(node_ll);
+    auto *node_links = (tableint *)(node_ll + 1);
+#ifdef USE_SSE
+    _mm_prefetch((char *)(visited_nodes_handler->getElementsTags() + *(node_ll + 1)), _MM_HINT_T0);
+    _mm_prefetch((char *)(visited_nodes_handler->getElementsTags() + *(node_ll + 1) + 64),
+                 _MM_HINT_T0);
+    _mm_prefetch(getDataByInternalId(*node_links), _MM_HINT_T0);
+    _mm_prefetch(getDataByInternalId(*(node_links + 1)), _MM_HINT_T0);
+#endif
+
+    for (size_t j = 0; j < links_num; j++) {
+        tableint candidate_id = *(node_links + j);
+#ifdef USE_SSE
+        _mm_prefetch((char *)(visited_nodes_handler->getElementsTags() + *(node_links + j + 1)),
+                     _MM_HINT_T0);
+        _mm_prefetch(getDataByInternalId(*(node_links + j + 1)), _MM_HINT_T0);
+#endif
+        if (this->visited_nodes_handler->getNodeTag(candidate_id) == visited_tag)
+            continue;
+        this->visited_nodes_handler->tagNode(candidate_id, visited_tag);
+        char *currObj1 = (getDataByInternalId(candidate_id));
+
+        dist_t dist1 = fstdistfunc_(data_point, currObj1, dist_func_param_);
+        if (top_candidates.size() < ef || lowerBound > dist1) {
+            candidate_set.emplace(-dist1, candidate_id);
+#ifdef USE_SSE
+            _mm_prefetch(getDataByInternalId(candidate_set.top().second), _MM_HINT_T0);
+#endif
+            top_candidates.emplace(dist1, candidate_id);
+
+            if (top_candidates.size() > ef)
+                top_candidates.pop();
+
+            if (!top_candidates.empty())
+                lowerBound = top_candidates.top().first;
+        }
+    }
+    return lowerBound;
+}
+
 template <typename dist_t>
 candidatesMaxHeap<dist_t> HierarchicalNSW<dist_t>::searchLayer(tableint ep_id,
                                                                const void *data_point, size_t layer,
@@ -366,50 +429,10 @@ candidatesMaxHeap<dist_t> HierarchicalNSW<dist_t>::searchLayer(tableint ep_id,
         }
         candidate_set.pop();
 
-        tableint curNodeNum = curr_el_pair.second;
-#ifdef ENABLE_PARALLELIZATION
-        std::unique_lock<std::mutex> lock(link_list_locks_[curNodeNum]);
-#endif
-        linklistsizeint *node_ll = get_linklist_at_level(curNodeNum, layer);
-        size_t links_num = getListCount(node_ll);
-        auto *node_links = (tableint *)(node_ll + 1);
-#ifdef USE_SSE
-        _mm_prefetch((char *)(visited_nodes_handler->getElementsTags() + *(node_ll + 1)),
-                     _MM_HINT_T0);
-        _mm_prefetch((char *)(visited_nodes_handler->getElementsTags() + *(node_ll + 1) + 64),
-                     _MM_HINT_T0);
-        _mm_prefetch(getDataByInternalId(*node_links), _MM_HINT_T0);
-        _mm_prefetch(getDataByInternalId(*(node_links + 1)), _MM_HINT_T0);
-#endif
-
-        for (size_t j = 0; j < links_num; j++) {
-            tableint candidate_id = *(node_links + j);
-#ifdef USE_SSE
-            _mm_prefetch((char *)(visited_nodes_handler->getElementsTags() + *(node_links + j + 1)),
-                         _MM_HINT_T0);
-            _mm_prefetch(getDataByInternalId(*(node_links + j + 1)), _MM_HINT_T0);
-#endif
-            if (this->visited_nodes_handler->getNodeTag(candidate_id) == visited_tag)
-                continue;
-            this->visited_nodes_handler->tagNode(candidate_id, visited_tag);
-            char *currObj1 = (getDataByInternalId(candidate_id));
-
-            dist_t dist1 = fstdistfunc_(data_point, currObj1, dist_func_param_);
-            if (top_candidates.size() < ef || lowerBound > dist1) {
-                candidate_set.emplace(-dist1, candidate_id);
-#ifdef USE_SSE
-                _mm_prefetch(getDataByInternalId(candidate_set.top().second), _MM_HINT_T0);
-#endif
-                top_candidates.emplace(dist1, candidate_id);
-
-                if (top_candidates.size() > ef)
-                    top_candidates.pop();
-
-                if (!top_candidates.empty())
-                    lowerBound = top_candidates.top().first;
-            }
-        }
+        lowerBound = processCandidate(curr_el_pair.second, data_point, layer, ef, visited_tag,
+                                      top_candidates, candidate_set, lowerBound);
     }
+
 #ifdef ENABLE_PARALLELIZATION
     visited_nodes_handler_pool->returnVisitedNodesHandlerToPool(this->visited_nodes_handler);
 #endif
@@ -1034,7 +1057,8 @@ void HierarchicalNSW<dist_t>::addPoint(const void *data_point, const labeltype l
 }
 
 template <typename dist_t>
-tableint HierarchicalNSW<dist_t>::searchBottomLayerEP(const void *query_data) const {
+tableint HierarchicalNSW<dist_t>::searchBottomLayerEP(const void *query_data, void *timeoutCtx,
+                                                      VecSimQueryResult_Code *rc) const {
 
     if (cur_element_count == 0) {
         return entrypoint_node_;
@@ -1045,6 +1069,10 @@ tableint HierarchicalNSW<dist_t>::searchBottomLayerEP(const void *query_data) co
     for (size_t level = maxlevel_; level > 0; level--) {
         bool changed = true;
         while (changed) {
+            if (__builtin_expect(VecSimIndex::timeoutCallback(timeoutCtx), 0)) {
+                *rc = VecSim_QueryResult_TimedOut;
+                return HNSW_INVALID_ID;
+            }
             changed = false;
             linklistsizeint *node_ll = get_linklist(currObj, level);
             unsigned short links_count = getListCount(node_ll);
@@ -1064,30 +1092,81 @@ tableint HierarchicalNSW<dist_t>::searchBottomLayerEP(const void *query_data) co
             }
         }
     }
+    *rc = VecSim_QueryResult_OK;
     return currObj;
 }
 
 template <typename dist_t>
-vecsim_stl::max_priority_queue<pair<dist_t, labeltype>>
-HierarchicalNSW<dist_t>::searchKnn(const void *query_data, size_t k) const {
+candidatesLabelsMaxHeap<dist_t>
+HierarchicalNSW<dist_t>::searchBottomLayer_WithTimeout(tableint ep_id, const void *data_point,
+                                                       size_t ef, size_t k, void *timeoutCtx,
+                                                       VecSimQueryResult_Code *rc) const {
+    candidatesLabelsMaxHeap<dist_t> results(this->allocator);
+
+#ifdef ENABLE_PARALLELIZATION
+    this->visited_nodes_handler =
+        this->visited_nodes_handler_pool->getAvailableVisitedNodesHandler();
+#endif
 
-    vecsim_stl::max_priority_queue<std::pair<dist_t, labeltype>> result(this->allocator);
-    if (cur_element_count == 0)
-        return result;
+    tag_t visited_tag = this->visited_nodes_handler->getFreshTag();
 
-    tableint bottom_layer_ep = searchBottomLayerEP(query_data);
-    vecsim_stl::max_priority_queue<pair<dist_t, tableint>> top_candidates =
-        searchLayer(bottom_layer_ep, query_data, 0, std::max(ef_, k));
+    candidatesMaxHeap<dist_t> top_candidates(this->allocator);
+    candidatesMaxHeap<dist_t> candidate_set(this->allocator);
+
+    dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_);
+    dist_t lowerBound = dist;
+    top_candidates.emplace(dist, ep_id);
+    candidate_set.emplace(-dist, ep_id);
+
+    this->visited_nodes_handler->tagNode(ep_id, visited_tag);
+
+    while (!candidate_set.empty()) {
+        std::pair<dist_t, tableint> curr_el_pair = candidate_set.top();
+        if ((-curr_el_pair.first) > lowerBound) {
+            break;
+        }
+        if (__builtin_expect(VecSimIndex::timeoutCallback(timeoutCtx), 0)) {
+            *rc = VecSim_QueryResult_TimedOut;
+            return results;
+        }
+        candidate_set.pop();
 
+        lowerBound = processCandidate(curr_el_pair.second, data_point, 0, ef, visited_tag,
+                                      top_candidates, candidate_set, lowerBound);
+    }
+#ifdef ENABLE_PARALLELIZATION
+    visited_nodes_handler_pool->returnVisitedNodesHandlerToPool(this->visited_nodes_handler);
+#endif
     while (top_candidates.size() > k) {
         top_candidates.pop();
     }
     while (top_candidates.size() > 0) {
-        std::pair<dist_t, tableint> rez = top_candidates.top();
-        result.push(std::pair<dist_t, labeltype>(rez.first, getExternalLabel(rez.second)));
+        auto &res = top_candidates.top();
+        results.emplace(res.first, getExternalLabel(res.second));
         top_candidates.pop();
     }
-    return result;
+    *rc = VecSim_QueryResult_OK;
+    return results;
+}
+
+template <typename dist_t>
+candidatesLabelsMaxHeap<dist_t>
+HierarchicalNSW<dist_t>::searchKnn(const void *query_data, size_t k, void *timeoutCtx,
+                                   VecSimQueryResult_Code *rc) const {
+
+    if (cur_element_count == 0) {
+        *rc = VecSim_QueryResult_OK;
+        return candidatesLabelsMaxHeap<dist_t>(this->allocator);
+    }
+
+    tableint bottom_layer_ep = searchBottomLayerEP(query_data, timeoutCtx, rc);
+    if (VecSim_OK != *rc) {
+        return candidatesLabelsMaxHeap<dist_t>(this->allocator);
+    }
+    candidatesLabelsMaxHeap<dist_t> results = searchBottomLayer_WithTimeout(
+        bottom_layer_ep, query_data, std::max(ef_, k), k, timeoutCtx, rc);
+
+    return results;
 }
 
 } // namespace hnswlib