Timeout Enabling - Brute Force (#161)

* enable timeout function
  implement using it with brute force algorithm

* freeing result list if not null

* added tests

* review changes

* moved compute block scores to the BF algorithm

* added `vec_sim_index.cpp` file.

* more review fixes

* comment fix

Author: GuyAv46

.gitignore

@@ -10,6 +10,10 @@
 /1/
 **/build/
 
+# Ignore benchmark fetched data but not the source file
+/tests/benchmark/data/*
+!/tests/benchmark/data/hnsw_indices.txt
+
 # Prerequisites
 *.d
 

src/CMakeLists.txt

@@ -55,6 +55,7 @@ add_library(VectorSimilarity ${VECSIM_LIBTYPE}
     VecSim/algorithms/brute_force/vector_block.cpp
 	VecSim/algorithms/hnsw/visited_nodes_handler.cpp
 	VecSim/algorithms/hnsw/hnsw_wrapper.cpp
+    VecSim/vec_sim_index.cpp
     VecSim/vec_sim.cpp
     VecSim/query_results.cpp
 	VecSim/info_iterator.cpp

src/VecSim/algorithms/brute_force/bf_batch_iterator.cpp

@@ -12,21 +12,21 @@ unsigned char BF_BatchIterator::next_id = 0;
 
 // heuristics: decide if using heap or select search, based on the ratio between the
 // number of remaining results and the index size.
-VecSimQueryResult *BF_BatchIterator::searchByHeuristics(size_t n_res,
-                                                        VecSimQueryResult_Order order) {
+VecSimQueryResult_List BF_BatchIterator::searchByHeuristics(size_t n_res,
+                                                            VecSimQueryResult_Order order) {
     if ((this->index->indexSize() - this->getResultsCount()) / 1000 > n_res) {
         // Heap based search always returns the results ordered by score
         return this->heapBasedSearch(n_res);
     }
-    VecSimQueryResult *res = this->selectBasedSearch(n_res);
+    VecSimQueryResult_List rl = this->selectBasedSearch(n_res);
     if (order == BY_SCORE) {
-        sort_results_by_score(res);
+        sort_results_by_score(rl);
     }
-    return res;
+    return rl;
 }
 
-void BF_BatchIterator::swapScores(const unordered_map<size_t, size_t> &TopCandidatesIndices,
-                                  size_t res_num) {
+void BF_BatchIterator::swapScores(
+    const vecsim_stl::unordered_map<size_t, size_t> &TopCandidatesIndices, size_t res_num) {
     // Create a set of the indices in the scores array for every results that we return.
     set<size_t> indices;
     for (auto pos : TopCandidatesIndices) {
@@ -55,11 +55,12 @@ void BF_BatchIterator::swapScores(const unordered_map<size_t, size_t> &TopCandid
     this->scores_valid_start_pos = next_scores_valid_start_pos;
 }
 
-VecSimQueryResult *BF_BatchIterator::heapBasedSearch(size_t n_res) {
+VecSimQueryResult_List BF_BatchIterator::heapBasedSearch(size_t n_res) {
+    VecSimQueryResult_List rl = {0};
     float upperBound = std::numeric_limits<float>::lowest();
     vecsim_stl::max_priority_queue<pair<float, labelType>> TopCandidates(this->allocator);
     // map vector's label to its index in the scores vector.
-    unordered_map<size_t, size_t> TopCandidatesIndices(n_res);
+    vecsim_stl::unordered_map<size_t, size_t> TopCandidatesIndices(n_res, this->allocator);
     for (size_t i = this->scores_valid_start_pos; i < this->scores.size(); i++) {
         if (TopCandidates.size() < n_res) {
             TopCandidates.emplace(this->scores[i].first, this->scores[i].second);
@@ -69,28 +70,30 @@ VecSimQueryResult *BF_BatchIterator::heapBasedSearch(size_t n_res) {
             if (this->scores[i].first >= upperBound) {
                 continue;
             } else {
-                TopCandidates.emplace(this->scores[i].first, this->scores[i].second);
-                TopCandidatesIndices[this->scores[i].second] = i;
                 // remove the furthest vector from the candidates and from the label->index mappings
+                // we first remove the worst candidate so we wont exceed the allocated size
                 TopCandidatesIndices.erase(TopCandidates.top().second);
                 TopCandidates.pop();
+                TopCandidatesIndices[this->scores[i].second] = i;
+                TopCandidates.emplace(this->scores[i].first, this->scores[i].second);
                 upperBound = TopCandidates.top().first;
             }
         }
     }
 
     // Save the top results to return.
-    auto *results = array_new_len<VecSimQueryResult>(TopCandidates.size(), TopCandidates.size());
+    rl.results = array_new_len<VecSimQueryResult>(TopCandidates.size(), TopCandidates.size());
     for (int i = (int)TopCandidates.size() - 1; i >= 0; --i) {
-        VecSimQueryResult_SetId(results[i], TopCandidates.top().second);
-        VecSimQueryResult_SetScore(results[i], TopCandidates.top().first);
+        VecSimQueryResult_SetId(rl.results[i], TopCandidates.top().second);
+        VecSimQueryResult_SetScore(rl.results[i], TopCandidates.top().first);
         TopCandidates.pop();
     }
-    swapScores(TopCandidatesIndices, array_len(results));
-    return results;
+    swapScores(TopCandidatesIndices, array_len(rl.results));
+    return rl;
 }
 
-VecSimQueryResult *BF_BatchIterator::selectBasedSearch(size_t n_res) {
+VecSimQueryResult_List BF_BatchIterator::selectBasedSearch(size_t n_res) {
+    VecSimQueryResult_List rl = {0};
     size_t remaining_vectors_count = this->scores.size() - this->scores_valid_start_pos;
     // Get an iterator to the effective first element in the scores array, which is the first
     // element that hasn't been returned in previous iterations.
@@ -105,20 +108,22 @@ VecSimQueryResult *BF_BatchIterator::selectBasedSearch(size_t n_res) {
     // will be placed before it, and all the rest will be placed after.
     std::nth_element(valid_begin_it, n_th_element_pos, this->scores.end());
 
-    auto *results = array_new<VecSimQueryResult>(n_res);
+    rl.results = array_new<VecSimQueryResult>(n_res);
     for (size_t i = this->scores_valid_start_pos; i < this->scores_valid_start_pos + n_res; i++) {
-        results = array_append(results, VecSimQueryResult{});
-        VecSimQueryResult_SetId(results[array_len(results) - 1], this->scores[i].second);
-        VecSimQueryResult_SetScore(results[array_len(results) - 1], this->scores[i].first);
+        rl.results = array_append(rl.results, VecSimQueryResult{});
+        VecSimQueryResult_SetId(rl.results[array_len(rl.results) - 1], this->scores[i].second);
+        VecSimQueryResult_SetScore(rl.results[array_len(rl.results) - 1], this->scores[i].first);
     }
     // Update the valid results start position after returning the results.
-    this->scores_valid_start_pos += array_len(results);
-    return results;
+    this->scores_valid_start_pos += array_len(rl.results);
+    return rl;
 }
 
 BF_BatchIterator::BF_BatchIterator(void *query_vector, const BruteForceIndex *bf_index,
+                                   VecSimQueryParams *queryParams,
                                    std::shared_ptr<VecSimAllocator> allocator)
-    : VecSimBatchIterator(query_vector, allocator), index(bf_index), scores_valid_start_pos(0) {
+    : VecSimBatchIterator(query_vector, queryParams ? queryParams->timeoutCtx : nullptr, allocator),
+      index(bf_index), scores_valid_start_pos(0) {
     BF_BatchIterator::next_id++;
 }
 
@@ -131,20 +136,29 @@ VecSimQueryResult_List BF_BatchIterator::getNextResults(size_t n_res,
         assert(getResultsCount() == 0);
         this->scores.reserve(this->index->indexSize());
         vecsim_stl::vector<VectorBlock *> blocks = this->index->getVectorBlocks();
+        VecSimQueryResult_Code rc;
         for (auto &block : blocks) {
             // compute the scores for the vectors in every block and extend the scores array.
-            vecsim_stl::vector<std::pair<float, labelType>> block_scores =
-                block->computeBlockScores(getIndex()->distFunc(), getQueryBlob());
-            this->scores.insert(this->scores.end(), block_scores.begin(), block_scores.end());
+            auto block_scores = this->index->computeBlockScores(block, this->getQueryBlob(),
+                                                                this->getTimeoutCtx(), &rc);
+            if (VecSim_OK != rc) {
+                return {NULL, rc};
+            }
+            for (size_t i = 0; i < block_scores.size(); i++) {
+                this->scores.emplace_back(block_scores[i], block->getMember(i)->label);
+            }
         }
     }
-    VecSimQueryResult *results = searchByHeuristics(n_res, order);
+    if (__builtin_expect(VecSimIndex::timeoutCallback(this->getTimeoutCtx()), 0)) {
+        return {NULL, VecSim_QueryResult_TimedOut};
+    }
+    VecSimQueryResult_List rl = searchByHeuristics(n_res, order);
 
-    this->updateResultsCount(array_len(results));
+    this->updateResultsCount(array_len(rl.results));
     if (order == BY_ID) {
-        sort_results_by_id(results);
+        sort_results_by_id(rl);
     }
-    return results;
+    return rl;
 }
 
 bool BF_BatchIterator::isDepleted() {

src/VecSim/algorithms/brute_force/bf_batch_iterator.h

@@ -14,14 +14,15 @@ class BF_BatchIterator : public VecSimBatchIterator {
                                    // that hasn't been returned already.
     static unsigned char next_id;  // this holds the next available id to be used by a new instance.
 
-    VecSimQueryResult *searchByHeuristics(size_t n_res, VecSimQueryResult_Order order);
-    VecSimQueryResult *selectBasedSearch(size_t n_res);
-    VecSimQueryResult *heapBasedSearch(size_t n_res);
-    void swapScores(const unordered_map<size_t, size_t> &TopCandidatesIndices, size_t res_num);
+    VecSimQueryResult_List searchByHeuristics(size_t n_res, VecSimQueryResult_Order order);
+    VecSimQueryResult_List selectBasedSearch(size_t n_res);
+    VecSimQueryResult_List heapBasedSearch(size_t n_res);
+    void swapScores(const vecsim_stl::unordered_map<size_t, size_t> &TopCandidatesIndices,
+                    size_t res_num);
 
 public:
     BF_BatchIterator(void *query_vector, const BruteForceIndex *index,
-                     std::shared_ptr<VecSimAllocator> allocator);
+                     VecSimQueryParams *queryParams, std::shared_ptr<VecSimAllocator> allocator);
 
     inline const BruteForceIndex *getIndex() const { return index; };
 

src/VecSim/algorithms/brute_force/brute_force.cpp

@@ -180,9 +180,30 @@ double BruteForceIndex::getDistanceFrom(size_t label, const void *vector_data) {
 
 size_t BruteForceIndex::indexSize() const { return this->count; }
 
+// Compute the score for every vector in the block by using the given distance function.
+vecsim_stl::vector<float> BruteForceIndex::computeBlockScores(VectorBlock *block,
+                                                              const void *queryBlob,
+                                                              void *timeoutCtx,
+                                                              VecSimQueryResult_Code *rc) const {
+    size_t len = block->getLength();
+    vecsim_stl::vector<float> scores(len, this->allocator);
+    for (size_t i = 0; i < len; i++) {
+        if (__builtin_expect(VecSimIndex::timeoutCallback(timeoutCtx), 0)) {
+            *rc = VecSim_QueryResult_TimedOut;
+            return scores;
+        }
+        scores[i] = this->dist_func(block->getVector(i), queryBlob, &this->dim);
+    }
+    *rc = VecSim_QueryResult_OK;
+    return scores;
+}
+
 VecSimQueryResult_List BruteForceIndex::topKQuery(const void *queryBlob, size_t k,
                                                   VecSimQueryParams *queryParams) {
 
+    VecSimQueryResult_List rl = {0};
+    void *timeoutCtx = queryParams ? queryParams->timeoutCtx : NULL;
+
     this->last_mode = STANDARD_KNN;
     float normalized_blob[this->dim]; // This will be use only if metric == VecSimMetric_Cosine
     if (this->metric == VecSimMetric_Cosine) {
@@ -196,10 +217,9 @@ VecSimQueryResult_List BruteForceIndex::topKQuery(const void *queryBlob, size_t
     vecsim_stl::max_priority_queue<pair<float, labelType>> TopCandidates(this->allocator);
     // For every block, compute its vectors scores and update the Top candidates max heap
     for (auto vectorBlock : this->vectorBlocks) {
-        size_t block_size = vectorBlock->getLength();
-        vecsim_stl::vector<float> scores(block_size, this->allocator);
-        for (size_t i = 0; i < block_size; i++) {
-            scores[i] = this->dist_func(vectorBlock->getVector(i), queryBlob, &this->dim);
+        auto scores = computeBlockScores(vectorBlock, queryBlob, timeoutCtx, &rl.code);
+        if (VecSim_OK != rl.code) {
+            return rl;
         }
         for (size_t i = 0; i < scores.size(); i++) {
             // Always choose the current candidate if we have less than k.
@@ -219,13 +239,14 @@ VecSimQueryResult_List BruteForceIndex::topKQuery(const void *queryBlob, size_t
             }
         }
     }
-    auto *results = array_new_len<VecSimQueryResult>(TopCandidates.size(), TopCandidates.size());
+    rl.results = array_new_len<VecSimQueryResult>(TopCandidates.size(), TopCandidates.size());
     for (int i = (int)TopCandidates.size() - 1; i >= 0; --i) {
-        VecSimQueryResult_SetId(results[i], TopCandidates.top().second);
-        VecSimQueryResult_SetScore(results[i], TopCandidates.top().first);
+        VecSimQueryResult_SetId(rl.results[i], TopCandidates.top().second);
+        VecSimQueryResult_SetScore(rl.results[i], TopCandidates.top().first);
         TopCandidates.pop();
     }
-    return results;
+    rl.code = VecSim_QueryResult_OK;
+    return rl;
 }
 
 VecSimIndexInfo BruteForceIndex::info() const {
@@ -290,7 +311,8 @@ VecSimBatchIterator *BruteForceIndex::newBatchIterator(const void *queryBlob,
         float_vector_normalize((float *)queryBlobCopy, dim);
     }
     // Ownership of queryBlobCopy moves to BF_BatchIterator that will free it at the end.
-    return new (this->allocator) BF_BatchIterator(queryBlobCopy, this, this->allocator);
+    return new (this->allocator)
+        BF_BatchIterator(queryBlobCopy, this, queryParams, this->allocator);
 }
 
 bool BruteForceIndex::preferAdHocSearch(size_t subsetSize, size_t k, bool initial_check) {

src/VecSim/algorithms/brute_force/brute_force.h

@@ -21,6 +21,9 @@ class BruteForceIndex : public VecSimIndex {
     virtual int deleteVector(size_t id) override;
     virtual double getDistanceFrom(size_t label, const void *vector_data) override;
     virtual size_t indexSize() const override;
+    vecsim_stl::vector<float> computeBlockScores(VectorBlock *block, const void *queryBlob,
+                                                 void *timeoutCtx,
+                                                 VecSimQueryResult_Code *rc) const;
     virtual VecSimQueryResult_List topKQuery(const void *queryBlob, size_t k,
                                              VecSimQueryParams *queryParams) override;
     virtual VecSimIndexInfo info() const override;

src/VecSim/algorithms/brute_force/vector_block.cpp

@@ -31,13 +31,3 @@ void VectorBlock::addVector(VectorBlockMember *vectorBlockMember, const void *ve
     memcpy(this->vectors + (this->length * this->dim), vectorData, this->dim * sizeof(float));
     this->length++;
 }
-
-vecsim_stl::vector<std::pair<float, labelType>>
-VectorBlock::computeBlockScores(DISTFUNC<float> DistFunc, const void *queryBlob) {
-    vecsim_stl::vector<std::pair<float, labelType>> scores(this->length, this->allocator);
-    for (size_t i = 0; i < this->length; i++) {
-        scores[i] = {DistFunc(this->getVector(i), queryBlob, &this->dim),
-                     this->getMember(i)->label};
-    }
-    return scores;
-}

src/VecSim/algorithms/brute_force/vector_block.h

@@ -42,11 +42,6 @@ struct VectorBlock : public VecsimBaseObject {
         member->block = this;
     }
 
-    // Compute the score for every vector in the block by using the given distance function.
-    // Return a collection of (score, label) pairs for every vector in the block.
-    vecsim_stl::vector<std::pair<float, labelType>> computeBlockScores(DISTFUNC<float> DistFunc,
-                                                                       const void *queryBlob);
-
     virtual ~VectorBlock();
 
 private:

src/VecSim/algorithms/hnsw/hnsw_batch_iterator.cpp

@@ -16,24 +16,25 @@ inline bool HNSW_BatchIterator::hasVisitedNode(idType node_id) const {
 
 VecSimQueryResult_List HNSW_BatchIterator::prepareResults(candidatesMaxHeap top_candidates,
                                                           size_t n_res) {
+    VecSimQueryResult_List rl = {0};
     // size_t initial_results_num = array_len(batch_results);
     // Put the "spare" results (if exist) in the results heap.
     while (top_candidates.size() > n_res) {
         this->top_candidates_extras.emplace(top_candidates.top().first,
                                             top_candidates.top().second); // (distance, id)
         top_candidates.pop();
     }
-    auto *batch_results =
-        array_new_len<VecSimQueryResult>(top_candidates.size(), top_candidates.size());
+    rl.results = array_new_len<VecSimQueryResult>(top_candidates.size(), top_candidates.size());
     // Return results from the top candidates heap, put them in reverse order in the batch results
     // array.
     for (int i = (int)(top_candidates.size() - 1); i >= 0; i--) {
         labelType label = this->hnsw_index->getExternalLabel(top_candidates.top().second);
-        VecSimQueryResult_SetId(batch_results[i], label);
-        VecSimQueryResult_SetScore(batch_results[i], top_candidates.top().first);
+        VecSimQueryResult_SetId(rl.results[i], label);
+        VecSimQueryResult_SetScore(rl.results[i], top_candidates.top().first);
         top_candidates.pop();
     }
-    return batch_results;
+    rl.code = VecSim_QueryResult_OK;
+    return rl;
 }
 
 candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,
@@ -131,8 +132,10 @@ candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,
 HNSW_BatchIterator::HNSW_BatchIterator(void *query_vector, HNSWIndex *index_wrapper,
                                        VecSimQueryParams *queryParams,
                                        std::shared_ptr<VecSimAllocator> allocator)
-    : VecSimBatchIterator(query_vector, std::move(allocator)), index_wrapper(index_wrapper),
-      depleted(false), top_candidates_extras(this->allocator), candidates(this->allocator) {
+    : VecSimBatchIterator(query_vector, queryParams ? queryParams->timeoutCtx : nullptr,
+                          std::move(allocator)),
+      index_wrapper(index_wrapper), depleted(false), top_candidates_extras(this->allocator),
+      candidates(this->allocator) {
     this->space = index_wrapper->getSpace();
 
     this->hnsw_index = index_wrapper->getHNSWIndex();
@@ -169,7 +172,7 @@ VecSimQueryResult_List HNSW_BatchIterator::getNextResults(size_t n_res,
     // 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);
 
-    this->updateResultsCount(array_len(batch_results));
+    this->updateResultsCount(VecSimQueryResult_Len(batch_results));
     if (this->getResultsCount() == this->index_wrapper->indexSize()) {
         this->depleted = true;
     }

src/VecSim/algorithms/hnsw/hnsw_wrapper.cpp

@@ -94,6 +94,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};
     try {
         this->last_mode = STANDARD_KNN;
         float normalized_data[this->dim]; // This will be use only if metric == VecSimMetric_Cosine
@@ -112,20 +113,21 @@ VecSimQueryResult_List HNSWIndex::topKQuery(const void *query_data, size_t k,
             }
         }
         auto knn_res = hnsw->searchKnn(query_data, k);
-        auto *results = array_new_len<VecSimQueryResult>(knn_res.size(), knn_res.size());
+        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(results[i], knn_res.top().second);
-            VecSimQueryResult_SetScore(results[i], knn_res.top().first);
+            VecSimQueryResult_SetId(rl.results[i], knn_res.top().second);
+            VecSimQueryResult_SetScore(rl.results[i], knn_res.top().first);
             knn_res.pop();
         }
         // Restore efRuntime
         hnsw->setEf(originalEF);
         assert(hnsw->getEf() == originalEF);
 
-        return results;
+        rl.code = VecSim_QueryResult_OK;
     } catch (...) {
-        return NULL;
+        rl.code = VecSim_QueryResult_Err;
     }
+    return rl;
 }
 
 VecSimIndexInfo HNSWIndex::info() const {