Normalize query vector in batch search cosine (#144)

alonre24 · web-flow · commit 0fdca34c7886 · 2022-03-10T14:35:01.000+02:00
* Allocate and copy query blob in batch iterator creation + normalize it in case of cosine
diff --git a/src/VecSim/algorithms/brute_force/bf_batch_iterator.cpp b/src/VecSim/algorithms/brute_force/bf_batch_iterator.cpp
@@ -116,7 +116,7 @@ VecSimQueryResult *BF_BatchIterator::selectBasedSearch(size_t n_res) {
     return results;
 }
 
-BF_BatchIterator::BF_BatchIterator(const void *query_vector, const BruteForceIndex *bf_index,
+BF_BatchIterator::BF_BatchIterator(void *query_vector, const BruteForceIndex *bf_index,
                                    std::shared_ptr<VecSimAllocator> allocator)
     : VecSimBatchIterator(query_vector, allocator), index(bf_index), scores_valid_start_pos(0) {
     BF_BatchIterator::next_id++;
diff --git a/src/VecSim/algorithms/brute_force/bf_batch_iterator.h b/src/VecSim/algorithms/brute_force/bf_batch_iterator.h
@@ -20,7 +20,7 @@ class BF_BatchIterator : public VecSimBatchIterator {
     void swapScores(const unordered_map<size_t, size_t> &TopCandidatesIndices, size_t res_num);
 
 public:
-    BF_BatchIterator(const void *query_vector, const BruteForceIndex *index,
+    BF_BatchIterator(void *query_vector, const BruteForceIndex *index,
                      std::shared_ptr<VecSimAllocator> allocator);
 
     inline const BruteForceIndex *getIndex() const { return index; };
diff --git a/src/VecSim/algorithms/brute_force/brute_force.cpp b/src/VecSim/algorithms/brute_force/brute_force.cpp
@@ -9,6 +9,7 @@
 #include <memory>
 #include <cstring>
 #include <queue>
+#include <cassert>
 
 using namespace std;
 
@@ -229,7 +230,7 @@ VecSimQueryResult_List BruteForceIndex::topKQuery(const void *queryBlob, size_t
     return results;
 }
 
-VecSimIndexInfo BruteForceIndex::info() {
+VecSimIndexInfo BruteForceIndex::info() const {
 
     VecSimIndexInfo info;
     info.algo = VecSimAlgo_BF;
@@ -281,7 +282,16 @@ VecSimInfoIterator *BruteForceIndex::infoIterator() {
 }
 
 VecSimBatchIterator *BruteForceIndex::newBatchIterator(const void *queryBlob) {
-    return new (this->allocator) BF_BatchIterator(queryBlob, this, this->allocator);
+    // As this is the only supported type, we always allocate 4 bytes for every element in the
+    // vector.
+    assert(this->vecType == VecSimType_FLOAT32);
+    auto *queryBlobCopy = this->allocator->allocate(sizeof(float) * this->dim);
+    memcpy(queryBlobCopy, queryBlob, dim * sizeof(float));
+    if (metric == VecSimMetric_Cosine) {
+        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);
 }
 
 bool BruteForceIndex::preferAdHocSearch(size_t subsetSize, size_t k) {
diff --git a/src/VecSim/algorithms/brute_force/brute_force.h b/src/VecSim/algorithms/brute_force/brute_force.h
@@ -23,7 +23,7 @@ class BruteForceIndex : public VecSimIndex {
                                             VecSimQueryParams *qparams) override;
     virtual VecSimQueryResult_List topKQuery(const void *queryBlob, size_t k,
                                              VecSimQueryParams *queryParams) override;
-    virtual VecSimIndexInfo info() override;
+    virtual VecSimIndexInfo info() const override;
     virtual VecSimInfoIterator *infoIterator() override;
     virtual VecSimBatchIterator *newBatchIterator(const void *queryBlob) override;
     bool preferAdHocSearch(size_t subsetSize, size_t k) override;
diff --git a/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.cpp b/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.cpp
@@ -127,7 +127,7 @@ candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,
     return top_candidates;
 }
 
-HNSW_BatchIterator::HNSW_BatchIterator(const void *query_vector, HNSWIndex *index_wrapper,
+HNSW_BatchIterator::HNSW_BatchIterator(void *query_vector, HNSWIndex *index_wrapper,
                                        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) {
diff --git a/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h b/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h
@@ -36,7 +36,7 @@ class HNSW_BatchIterator : public VecSimBatchIterator {
     inline bool hasVisitedNode(idType node_id) const;
 
 public:
-    HNSW_BatchIterator(const void *query_vector, HNSWIndex *index,
+    HNSW_BatchIterator(void *query_vector, HNSWIndex *index,
                        std::shared_ptr<VecSimAllocator> allocator);
 
     VecSimQueryResult_List getNextResults(size_t n_res, VecSimQueryResult_Order order) override;
diff --git a/src/VecSim/algorithms/hnsw/hnsw_wrapper.cpp b/src/VecSim/algorithms/hnsw/hnsw_wrapper.cpp
@@ -133,7 +133,7 @@ VecSimQueryResult_List HNSWIndex::topKQuery(const void *query_data, size_t k,
     }
 }
 
-VecSimIndexInfo HNSWIndex::info() {
+VecSimIndexInfo HNSWIndex::info() const {
 
     VecSimIndexInfo info;
     info.algo = VecSimAlgo_HNSWLIB;
@@ -152,7 +152,16 @@ VecSimIndexInfo HNSWIndex::info() {
 }
 
 VecSimBatchIterator *HNSWIndex::newBatchIterator(const void *queryBlob) {
-    return new (this->allocator) HNSW_BatchIterator(queryBlob, this, this->allocator);
+    // As this is the only supported type, we always allocate 4 bytes for every element in the
+    // vector.
+    assert(this->vecType == VecSimType_FLOAT32);
+    auto *queryBlobCopy = this->allocator->allocate(sizeof(float) * this->dim);
+    memcpy(queryBlobCopy, queryBlob, dim * sizeof(float));
+    if (metric == VecSimMetric_Cosine) {
+        float_vector_normalize((float *)queryBlobCopy, dim);
+    }
+    // Ownership of queryBlobCopy moves to HNSW_BatchIterator that will free it at the end.
+    return new (this->allocator) HNSW_BatchIterator(queryBlobCopy, this, this->allocator);
 }
 
 VecSimInfoIterator *HNSWIndex::infoIterator() {
diff --git a/src/VecSim/algorithms/hnsw/hnsw_wrapper.h b/src/VecSim/algorithms/hnsw/hnsw_wrapper.h
@@ -20,7 +20,7 @@ class HNSWIndex : public VecSimIndex {
                                             VecSimQueryParams *qparams) override;
     virtual VecSimQueryResult_List topKQuery(const void *queryBlob, size_t k,
                                              VecSimQueryParams *queryParams) override;
-    virtual VecSimIndexInfo info() override;
+    virtual VecSimIndexInfo info() const override;
     virtual VecSimInfoIterator *infoIterator() override;
     virtual VecSimBatchIterator *newBatchIterator(const void *queryBlob) override;
     bool preferAdHocSearch(size_t subsetSize, size_t k) override;
diff --git a/src/VecSim/batch_iterator.h b/src/VecSim/batch_iterator.h
@@ -9,12 +9,11 @@
  */
 struct VecSimBatchIterator : public VecsimBaseObject {
 private:
-    const void *query_vector;
+    void *query_vector;
     size_t returned_results_count;
 
 public:
-    explicit VecSimBatchIterator(const void *query_vector,
-                                 std::shared_ptr<VecSimAllocator> allocator)
+    explicit VecSimBatchIterator(void *query_vector, std::shared_ptr<VecSimAllocator> allocator)
         : VecsimBaseObject(allocator), query_vector(query_vector), returned_results_count(0){};
 
     inline const void *getQueryBlob() const { return query_vector; }
@@ -35,5 +34,5 @@ struct VecSimBatchIterator : public VecsimBaseObject {
     // Reset the iterator to the initial state, before any results has been returned.
     virtual void reset() = 0;
 
-    virtual ~VecSimBatchIterator() = default;
+    virtual ~VecSimBatchIterator() { allocator->free_allocation(this->query_vector); };
 };
diff --git a/src/VecSim/vec_sim_index.h b/src/VecSim/vec_sim_index.h
@@ -90,7 +90,7 @@ struct VecSimIndex : public VecsimBaseObject {
      *
      * @return Index general and specific meta-data.
      */
-    virtual VecSimIndexInfo info() = 0;
+    virtual VecSimIndexInfo info() const = 0;
 
     /**
      * @brief Returns an index information in an iterable structure.
diff --git a/tests/unit/test_bruteforce.cpp b/tests/unit/test_bruteforce.cpp
@@ -1124,3 +1124,70 @@ TEST_F(BruteForceTest, batchIteratorSwapIndices) {
     VecSimBatchIterator_Free(batchIterator);
     VecSimIndex_Free(index);
 }
+
+TEST_F(BruteForceTest, testCosine) {
+    size_t dim = 128;
+    size_t n = 100;
+
+    VecSimParams params{.algo = VecSimAlgo_BF,
+                        .bfParams = BFParams{.type = VecSimType_FLOAT32,
+                                             .dim = dim,
+                                             .metric = VecSimMetric_Cosine,
+                                             .initialCapacity = n}};
+    VecSimIndex *index = VecSimIndex_New(&params);
+
+    for (size_t i = 1; i <= n; i++) {
+        float f[dim];
+        f[0] = (float)i / n;
+        for (size_t j = 1; j < dim; j++) {
+            f[j] = 1.0f;
+        }
+        VecSimIndex_AddVector(index, (const void *)f, i);
+    }
+    ASSERT_EQ(VecSimIndex_IndexSize(index), n);
+    float query[dim];
+    for (size_t i = 0; i < dim; i++) {
+        query[i] = 1.0f;
+    }
+    auto verify_res = [&](size_t id, float score, size_t index) {
+        ASSERT_EQ(id, (n - index));
+        float first_coordinate = (float)id / n;
+        // By cosine definition: 1 - ((A \dot B) / (norm(A)*norm(B))), where A is the query vector
+        // and B is the current result vector.
+        float expected_score =
+            1.0f -
+            ((first_coordinate + (float)dim - 1.0f) /
+             (sqrtf((float)dim) * sqrtf((float)(dim - 1) + first_coordinate * first_coordinate)));
+        // Verify that abs difference between the actual and expected score is at most 1/10^6.
+        ASSERT_NEAR(score, expected_score, 1e-6);
+    };
+    runTopKSearchTest(index, query, 10, verify_res);
+
+    // Test with batch iterator.
+    VecSimBatchIterator *batchIterator = VecSimBatchIterator_New(index, query);
+    size_t iteration_num = 0;
+
+    // get the 10 vectors whose ids are the maximal among those that hasn't been returned yet,
+    // in every iteration. The order should be from the largest to the lowest id.
+    size_t n_res = 10;
+    while (VecSimBatchIterator_HasNext(batchIterator)) {
+        std::vector<size_t> expected_ids(n_res);
+        auto verify_res_batch = [&](size_t id, float score, size_t index) {
+            ASSERT_EQ(id, (n - n_res * iteration_num - index));
+            float first_coordinate = (float)id / n;
+            // By cosine definition: 1 - ((A \dot B) / (norm(A)*norm(B))), where A is the query
+            // vector and B is the current result vector.
+            float expected_score =
+                1.0f - ((first_coordinate + (float)dim - 1.0f) /
+                        (sqrtf((float)dim) *
+                         sqrtf((float)(dim - 1) + first_coordinate * first_coordinate)));
+            // Verify that abs difference between the actual and expected score is at most 1/10^6.
+            ASSERT_NEAR(score, expected_score, 1e-6);
+        };
+        runBatchIteratorSearchTest(batchIterator, n_res, verify_res_batch);
+        iteration_num++;
+    }
+    ASSERT_EQ(iteration_num, n / n_res);
+    VecSimBatchIterator_Free(batchIterator);
+    VecSimIndex_Free(index);
+}
diff --git a/tests/unit/test_hnswlib.cpp b/tests/unit/test_hnswlib.cpp
@@ -1242,4 +1242,71 @@ TEST_F(HNSWLibTest, preferAdHocOptimization) {
     }
     VecSimIndex_Free(index);
 }
+
+TEST_F(HNSWLibTest, testCosine) {
+    size_t dim = 128;
+    size_t n = 100;
+
+    VecSimParams params{.algo = VecSimAlgo_HNSWLIB,
+                        .hnswParams = HNSWParams{.type = VecSimType_FLOAT32,
+                                                 .dim = dim,
+                                                 .metric = VecSimMetric_Cosine,
+                                                 .initialCapacity = n}};
+    VecSimIndex *index = VecSimIndex_New(&params);
+
+    for (size_t i = 1; i <= n; i++) {
+        float f[dim];
+        f[0] = (float)i / n;
+        for (size_t j = 1; j < dim; j++) {
+            f[j] = 1.0f;
+        }
+        VecSimIndex_AddVector(index, (const void *)f, i);
+    }
+    ASSERT_EQ(VecSimIndex_IndexSize(index), n);
+    float query[dim];
+    for (size_t i = 0; i < dim; i++) {
+        query[i] = 1.0f;
+    }
+    auto verify_res = [&](size_t id, float score, size_t index) {
+        ASSERT_EQ(id, (n - index));
+        float first_coordinate = (float)id / n;
+        // By cosine definition: 1 - ((A \dot B) / (norm(A)*norm(B))), where A is the query vector
+        // and B is the current result vector.
+        float expected_score =
+            1.0f -
+            ((first_coordinate + (float)dim - 1.0f) /
+             (sqrtf((float)dim) * sqrtf((float)(dim - 1) + first_coordinate * first_coordinate)));
+        // Verify that abs difference between the actual and expected score is at most 1/10^6.
+        ASSERT_NEAR(score, expected_score, 1e-6);
+    };
+    runTopKSearchTest(index, query, 10, verify_res);
+
+    // Test with batch iterator.
+    VecSimBatchIterator *batchIterator = VecSimBatchIterator_New(index, query);
+    size_t iteration_num = 0;
+
+    // get the 10 vectors whose ids are the maximal among those that hasn't been returned yet,
+    // in every iteration. The order should be from the largest to the lowest id.
+    size_t n_res = 10;
+    while (VecSimBatchIterator_HasNext(batchIterator)) {
+        std::vector<size_t> expected_ids(n_res);
+        auto verify_res_batch = [&](size_t id, float score, size_t index) {
+            ASSERT_EQ(id, (n - n_res * iteration_num - index));
+            float first_coordinate = (float)id / n;
+            // By cosine definition: 1 - ((A \dot B) / (norm(A)*norm(B))), where A is the query
+            // vector and B is the current result vector.
+            float expected_score =
+                1.0f - ((first_coordinate + (float)dim - 1.0f) /
+                        (sqrtf((float)dim) *
+                         sqrtf((float)(dim - 1) + first_coordinate * first_coordinate)));
+            // Verify that abs difference between the actual and expected score is at most 1/10^6.
+            ASSERT_NEAR(score, expected_score, 1e-6);
+        };
+        runBatchIteratorSearchTest(batchIterator, n_res, verify_res_batch);
+        iteration_num++;
+    }
+    ASSERT_EQ(iteration_num, n / n_res);
+    VecSimBatchIterator_Free(batchIterator);
+    VecSimIndex_Free(index);
+}
 } // namespace hnswlib

Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,7 @@ VecSimQueryResult *BF_BatchIterator::selectBasedSearch(size_t n_res) {`
`116`	`116`	`return results;`
`117`	`117`	`}`
`118`	`118`
`119`		`-BF_BatchIterator::BF_BatchIterator(const void query_vector, const BruteForceIndex bf_index,`
	`119`	`+BF_BatchIterator::BF_BatchIterator(void query_vector, const BruteForceIndex bf_index,`
`120`	`120`	`std::shared_ptr<VecSimAllocator> allocator)`
`121`	`121`	`: VecSimBatchIterator(query_vector, allocator), index(bf_index), scores_valid_start_pos(0) {`
`122`	`122`	`BF_BatchIterator::next_id++;`
Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,7 @@ candidatesMaxHeap HNSW_BatchIterator::scanGraph(candidatesMinHeap &candidates,`
`127`	`127`	`return top_candidates;`
`128`	`128`	`}`
`129`	`129`
`130`		`-HNSW_BatchIterator::HNSW_BatchIterator(const void query_vector, HNSWIndex index_wrapper,`
	`130`	`+HNSW_BatchIterator::HNSW_BatchIterator(void query_vector, HNSWIndex index_wrapper,`
`131`	`131`	`std::shared_ptr<VecSimAllocator> allocator)`
`132`	`132`	`: VecSimBatchIterator(query_vector, std::move(allocator)), index_wrapper(index_wrapper),`
`133`	`133`	`depleted(false), top_candidates_extras(this->allocator), candidates(this->allocator) {`
Original file line number	Diff line number	Diff line change
`@@ -133,7 +133,7 @@ VecSimQueryResult_List HNSWIndex::topKQuery(const void *query_data, size_t k,`
`133`	`133`	`}`
`134`	`134`	`}`
`135`	`135`
`136`		`-VecSimIndexInfo HNSWIndex::info() {`
	`136`	`+VecSimIndexInfo HNSWIndex::info() const {`
`137`	`137`
`138`	`138`	`VecSimIndexInfo info;`
`139`	`139`	`info.algo = VecSimAlgo_HNSWLIB;`
`@@ -152,7 +152,16 @@ VecSimIndexInfo HNSWIndex::info() {`
`152`	`152`	`}`
`153`	`153`
`154`	`154`	`VecSimBatchIterator HNSWIndex::newBatchIterator(const void queryBlob) {`
`155`		`- return new (this->allocator) HNSW_BatchIterator(queryBlob, this, this->allocator);`
	`155`	`+ // As this is the only supported type, we always allocate 4 bytes for every element in the`
	`156`	`+ // vector.`
	`157`	`+ assert(this->vecType == VecSimType_FLOAT32);`
	`158`	`+ auto queryBlobCopy = this->allocator->allocate(sizeof(float) this->dim);`
	`159`	`+ memcpy(queryBlobCopy, queryBlob, dim * sizeof(float));`
	`160`	`+ if (metric == VecSimMetric_Cosine) {`
	`161`	`+ float_vector_normalize((float *)queryBlobCopy, dim);`
	`162`	`+ }`
	`163`	`+ // Ownership of queryBlobCopy moves to HNSW_BatchIterator that will free it at the end.`
	`164`	`+ return new (this->allocator) HNSW_BatchIterator(queryBlobCopy, this, this->allocator);`
`156`	`165`	`}`
`157`	`166`
`158`	`167`	`VecSimInfoIterator *HNSWIndex::infoIterator() {`
Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ struct VecSimIndex : public VecsimBaseObject {`
`90`	`90`	`*`
`91`	`91`	`* @return Index general and specific meta-data.`
`92`	`92`	`*/`
`93`		`- virtual VecSimIndexInfo info() = 0;`
	`93`	`+ virtual VecSimIndexInfo info() const = 0;`
`94`	`94`
`95`	`95`	`/**`
`96`	`96`	`* @brief Returns an index information in an iterable structure.`