A new interface taking a template comparator is added, so it can not be virtual.

orrorcol · orrorcol · commit 9a3c03d93d66 · 2019-09-26T13:19:18.000+08:00
I modified the sift_1b(not commited) to test the new interface, the result is ok.

Test result of sift_1b on 1 million data.
Loading GT:
Loading queries:
Loading index from sift1b_1m_ef_40_M_16.bin:
Actual memory usage: 417 Mb
Parsing gt:
10000
Loaded gt
1	0.2371	13.319 us
2	0.3712	15.691 us
3	0.4615	18.5166 us
4	0.5273	20.6371 us
5	0.5758	22.1235 us
6	0.6179	24.4141 us
7	0.6502	25.9906 us
8	0.6796	28.2004 us
9	0.7042	29.8559 us
10	0.7243	31.3286 us
11	0.7432	36.0276 us
12	0.7605	34.9448 us
13	0.7754	36.4176 us
14	0.7874	37.7606 us
15	0.8013	44.6698 us
16	0.8116	47.4424 us
17	0.8239	46.9154 us
18	0.8312	45.9322 us
19	0.8379	49.3406 us
20	0.8442	49.124 us
21	0.8507	52.1223 us
22	0.8566	52.4161 us
23	0.8622	56.9665 us
24	0.8675	71.5782 us
25	0.8731	72.4451 us
26	0.8768	57.0935 us
27	0.8812	58.3525 us
28	0.8845	59.5751 us
29	0.889	61.7516 us
30	0.8935	62.6091 us
40	0.9224	76.8735 us
50	0.9412	92.5431 us
60	0.9541	107.141 us
70	0.9632	121.24 us
80	0.9708	135.862 us
90	0.9756	163.516 us
100	0.9792	180.539 us
140	0.9883	228.747 us
180	0.9921	281.199 us
220	0.9942	338.32 us
260	0.9956	388.501 us
300	0.9962	445.776 us
340	0.9968	477.474 us
380	0.9975	534.054 us
420	0.9982	582.327 us
460	0.9983	625.824 us
Actual memory usage: 419 Mb
diff --git a/hnswlib/bruteforce.h b/hnswlib/bruteforce.h
@@ -2,6 +2,7 @@
 #include <unordered_map>
 #include <fstream>
 #include <mutex>
+#include <algorithm>
 
 namespace hnswlib {
     template<typename dist_t>
@@ -84,11 +85,10 @@ namespace hnswlib {
         }
 
 
-        retType<dist_t> searchKnn(const void *query_data, size_t k) const {
-            retType<dist_t> result;
-            if (cur_element_count == 0) return result;
-
+        std::priority_queue<std::pair<dist_t, labeltype >>
+        searchKnn(const void *query_data, size_t k) const {
             std::priority_queue<std::pair<dist_t, labeltype >> topResults;
+            if (cur_element_count == 0) return topResults;
             for (int i = 0; i < k; i++) {
                 dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_);
                 topResults.push(std::pair<dist_t, labeltype>(dist, *((labeltype *) (data_ + size_per_element_ * i +
@@ -106,13 +106,30 @@ namespace hnswlib {
                 }
 
             }
-            while (!topResults.empty()) {
-                auto each = topResults.top();
-                result.push(each);
-                topResults.pop();
+            return topResults;
+        };
+
+        template <typename Comp>
+        std::vector<std::pair<dist_t, labeltype>>
+        searchKnn(const void* query_data, size_t k, Comp comp) {
+            std::vector<std::pair<dist_t, labeltype>> result;
+            if (cur_element_count == 0) return result;
+
+            auto ret = searchKnn(query_data, k);
+
+            while (!ret.empty()) {
+                result.push_back(ret.top());
+                ret.pop();
             }
+
+            if (result.size() > 1) {
+                if (!comp(result.front(), result.back())) {
+                    std::reverse(result.begin(), result.end());
+                }
+            }
+
             return result;
-        };
+        }
 
         void saveIndex(const std::string &location) {
             std::ofstream output(location, std::ios::binary);
diff --git a/hnswlib/hnswalg.h b/hnswlib/hnswalg.h
@@ -484,6 +484,8 @@ namespace hnswlib {
 
 
         std::priority_queue<std::pair<dist_t, tableint>> searchKnnInternal(void *query_data, int k) {
+            std::priority_queue<std::pair<dist_t, tableint  >> top_candidates;
+            if (cur_element_count == 0) return top_candidates;
             tableint currObj = enterpoint_node_;
             dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_);
 
@@ -510,8 +512,6 @@ namespace hnswlib {
                 }
             }
 
-
-            std::priority_queue<std::pair<dist_t, tableint  >> top_candidates;
             if (has_deletions_) {
                 std::priority_queue<std::pair<dist_t, tableint  >> top_candidates1=searchBaseLayerST<true>(currObj, query_data,
                                                                                                              ef_);
@@ -895,8 +895,9 @@ namespace hnswlib {
             return cur_c;
         };
 
-        retType<dist_t> searchKnn(const void *query_data, size_t k) const {
-            retType<dist_t> result;
+        std::priority_queue<std::pair<dist_t, labeltype >>
+        searchKnn(const void *query_data, size_t k) const {
+            std::priority_queue<std::pair<dist_t, labeltype >> result;
             if (cur_element_count == 0) return result;
 
             tableint currObj = enterpoint_node_;
@@ -948,6 +949,27 @@ namespace hnswlib {
             return result;
         };
 
+        template <typename Comp>
+        std::vector<std::pair<dist_t, labeltype>>
+        searchKnn(const void* query_data, size_t k, Comp comp) {
+            std::vector<std::pair<dist_t, labeltype>> result;
+            if (cur_element_count == 0) return result;
+
+            auto ret = searchKnn(query_data, k);
+
+            while (!ret.empty()) {
+                result.push_back(ret.top());
+                ret.pop();
+            }
+
+            if (result.size() > 1) {
+                if (!comp(result.front(), result.back())) {
+                    std::reverse(result.begin(), result.end());
+                }
+            }
+
+            return result;
+        }
 
     };
 
diff --git a/hnswlib/hnswlib.h b/hnswlib/hnswlib.h
@@ -39,14 +39,6 @@ namespace hnswlib {
         }
     };
 
-    template <typename T>
-    using retType = std::priority_queue<
-                std::pair<T, labeltype>,
-                std::vector<std::pair<T, labeltype>>,
-                pairGreater<std::pair<T, labeltype>>
-            >;
-
-
     template<typename T>
     static void writeBinaryPOD(std::ostream &out, const T &podRef) {
         out.write((char *) &podRef, sizeof(T));
@@ -78,7 +70,10 @@ namespace hnswlib {
     class AlgorithmInterface {
     public:
         virtual void addPoint(const void *datapoint, labeltype label)=0;
-        virtual retType<dist_t> searchKnn(const void *, size_t) const = 0;
+        virtual std::priority_queue<std::pair<dist_t, labeltype >> searchKnn(const void *, size_t) const = 0;
+        template <typename Comp>
+        std::vector<std::pair<dist_t, labeltype>> searchKnn(const void*, size_t, Comp) {
+        }
         virtual void saveIndex(const std::string &location)=0;
         virtual ~AlgorithmInterface(){
         }
diff --git a/sift_1b.cpp b/sift_1b.cpp
@@ -1,5 +1,6 @@
 #include <iostream>
 #include <fstream>
+#include <queue>
 #include <chrono>
 #include "hnswlib/hnswlib.h"
 
@@ -146,10 +147,10 @@ static size_t getCurrentRSS() {
 
 static void
 get_gt(unsigned int *massQA, unsigned char *massQ, unsigned char *mass, size_t vecsize, size_t qsize, L2SpaceI &l2space,
-       size_t vecdim, vector<retType<int>> &answers, size_t k) {
+       size_t vecdim, vector<std::priority_queue<std::pair<int, labeltype >>> &answers, size_t k) {
 
 
-    (vector<retType<int>>(qsize)).swap(answers);
+    (vector<std::priority_queue<std::pair<int, labeltype >>>(qsize)).swap(answers);
     DISTFUNC<int> fstdistfunc_ = l2space.get_dist_func();
     cout << qsize << "\n";
     for (int i = 0; i < qsize; i++) {
@@ -161,15 +162,15 @@ get_gt(unsigned int *massQA, unsigned char *massQ, unsigned char *mass, size_t v
 
 static float
 test_approx(unsigned char *massQ, size_t vecsize, size_t qsize, HierarchicalNSW<int> &appr_alg, size_t vecdim,
-            vector<retType<int>> &answers, size_t k) {
+            vector<std::priority_queue<std::pair<int, labeltype >>> &answers, size_t k) {
     size_t correct = 0;
     size_t total = 0;
     //uncomment to test in parallel mode:
     //#pragma omp parallel for
     for (int i = 0; i < qsize; i++) {
 
-        retType<int> result = appr_alg.searchKnn(massQ + vecdim * i, k);
-        retType<int> gt(answers[i]);
+        std::priority_queue<std::pair<int, labeltype >> result = appr_alg.searchKnn(massQ + vecdim * i, k);
+        std::priority_queue<std::pair<int, labeltype >> gt(answers[i]);
         unordered_set<labeltype> g;
         total += gt.size();
 
@@ -195,7 +196,7 @@ test_approx(unsigned char *massQ, size_t vecsize, size_t qsize, HierarchicalNSW<
 
 static void
 test_vs_recall(unsigned char *massQ, size_t vecsize, size_t qsize, HierarchicalNSW<int> &appr_alg, size_t vecdim,
-               vector<retType<int>> &answers, size_t k) {
+               vector<std::priority_queue<std::pair<int, labeltype >>> &answers, size_t k) {
     vector<size_t> efs;// = { 10,10,10,10,10 };
     for (int i = k; i < 30; i++) {
         efs.push_back(i);
@@ -230,7 +231,7 @@ inline bool exists_test(const std::string &name) {
 void sift_test1B() {
 	
 	
-	int subset_size_milllions = 1;
+	int subset_size_milllions = 200;
 	int efConstruction = 40;
 	int M = 16;
 	
@@ -350,7 +351,7 @@ void sift_test1B() {
     }
 
 
-    vector<retType<int>> answers;
+    vector<std::priority_queue<std::pair<int, labeltype >>> answers;
     size_t k = 1;
     cout << "Parsing gt:\n";
     get_gt(massQA, massQ, mass, vecsize, qsize, l2space, vecdim, answers, k);
