| 
 | 1 | +#include "../../hnswlib/hnswlib.h"  | 
 | 2 | + | 
 | 3 | +#include <assert.h>  | 
 | 4 | + | 
 | 5 | +#include <vector>  | 
 | 6 | +#include <iostream>  | 
 | 7 | +#include <cstdio>  | 
 | 8 | +#include <thread>  | 
 | 9 | +#include <chrono>  | 
 | 10 | + | 
 | 11 | +namespace {  | 
 | 12 | + | 
 | 13 | +const size_t M = 32;  | 
 | 14 | +const size_t ef_construction = 500;  | 
 | 15 | +const size_t random_seed = 100;  | 
 | 16 | +const bool allow_replace_deleted = false;  | 
 | 17 | + | 
 | 18 | +const size_t dimension = 1024;  | 
 | 19 | +const size_t total_items = 100 * 10000;  | 
 | 20 | +const size_t num_query = 500 * 10000;  | 
 | 21 | +size_t topk = 10;  | 
 | 22 | +const size_t max_thread_num = 48;  | 
 | 23 | +const std::string index_path = "./hnsw.index";  | 
 | 24 | + | 
 | 25 | +std::vector<float> data(total_items * dimension);  | 
 | 26 | +std::vector<float> query(num_query * dimension);  | 
 | 27 | + | 
 | 28 | + | 
 | 29 | +void check_knn_closer(hnswlib::AlgorithmInterface<float>* alg_hnsw) {  | 
 | 30 | +    for (size_t j = 0; j < num_query; ++j) {  | 
 | 31 | +        const void* p = query.data() + j * dimension;  | 
 | 32 | +        auto gd = alg_hnsw->searchKnn(p, topk);  | 
 | 33 | +        auto res = alg_hnsw->searchKnnCloserFirst(p, topk);  | 
 | 34 | +        assert(gd.size() == res.size());  | 
 | 35 | +        size_t t = gd.size();  | 
 | 36 | +        while (!gd.empty()) {  | 
 | 37 | +            assert(gd.top() == res[--t]);  | 
 | 38 | +            gd.pop();  | 
 | 39 | +        }  | 
 | 40 | +    }  | 
 | 41 | +    std::cout << "test hnsw search knn closer first success..." << std::endl;  | 
 | 42 | +}  | 
 | 43 | + | 
 | 44 | +void test_compatibility(bool hnsw_first_use_blocks_memory,  | 
 | 45 | +                        bool hnsw_second_use_blocks_memory) {  | 
 | 46 | + | 
 | 47 | +    std::cout << "================== test compatibility ==================" << std::endl;  | 
 | 48 | +    hnswlib::L2Space space(dimension);  | 
 | 49 | +    hnswlib::AlgorithmInterface<float>* alg_hnsw_first = new hnswlib::HierarchicalNSW<float>(&space, 2 * total_items,  | 
 | 50 | +            M, ef_construction, random_seed, allow_replace_deleted, hnsw_first_use_blocks_memory);  | 
 | 51 | + | 
 | 52 | +    for (size_t i = 0; i < total_items; ++i) {  | 
 | 53 | +        alg_hnsw_first->addPoint(data.data() + dimension * i, i);  | 
 | 54 | +    }  | 
 | 55 | +    check_knn_closer(alg_hnsw_first);  | 
 | 56 | + | 
 | 57 | +    // save hnsw index  | 
 | 58 | +    std::remove(index_path.data());  | 
 | 59 | +    alg_hnsw_first->saveIndex(index_path);  | 
 | 60 | +    std::cout << "save hnsw(use_small_blocks_memory = " << hnsw_first_use_blocks_memory << ") index success" << std::endl;  | 
 | 61 | +    delete alg_hnsw_first;  | 
 | 62 | + | 
 | 63 | +    // load hnsw index  | 
 | 64 | +    hnswlib::AlgorithmInterface<float>* alg_hnsw_second = new hnswlib::HierarchicalNSW<float>(&space, false,  | 
 | 65 | +            0, allow_replace_deleted, hnsw_second_use_blocks_memory);  | 
 | 66 | +    std::cout << "load hnsw(use_small_blocks_memory = " << hnsw_second_use_blocks_memory << ") index success" << std::endl;  | 
 | 67 | +    std::remove(index_path.data());  | 
 | 68 | +    check_knn_closer(alg_hnsw_second);  | 
 | 69 | + | 
 | 70 | +    delete alg_hnsw_second;  | 
 | 71 | +}  | 
 | 72 | + | 
 | 73 | +void test_performace(bool use_small_blocks_memory) {  | 
 | 74 | +    if (total_items == 0) {  | 
 | 75 | +      return;  | 
 | 76 | +    }  | 
 | 77 | + | 
 | 78 | +    std::cout << "================== test preformace("  | 
 | 79 | +              << "dimension: " << dimension  | 
 | 80 | +              << ", M: " << M  | 
 | 81 | +              << ", ef_construction: " << ef_construction  | 
 | 82 | +              << ", topk: " << topk  | 
 | 83 | +              << ", use_small_blocks_memory: " << (use_small_blocks_memory ? "ture" : "false" )  | 
 | 84 | +              << ") ==================" << std::endl;  | 
 | 85 | +    hnswlib::L2Space space(dimension);  | 
 | 86 | +    hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * total_items,  | 
 | 87 | +            M, ef_construction, random_seed, allow_replace_deleted, use_small_blocks_memory);  | 
 | 88 | + | 
 | 89 | +    std::vector<std::thread> threads;  | 
 | 90 | +    size_t num_threads = (total_items >= max_thread_num ? max_thread_num : total_items);  | 
 | 91 | +    size_t batch_num = (total_items / (num_threads <= 1 ? 1 : (num_threads - 1))) + 1;  | 
 | 92 | +    auto start_time = std::chrono::system_clock::now();  | 
 | 93 | +    for (size_t idx = 0; idx < total_items; idx += batch_num) {  | 
 | 94 | +        size_t start = idx;  | 
 | 95 | +        size_t end = std::min(idx + batch_num, total_items);  | 
 | 96 | +        threads.push_back(  | 
 | 97 | +            std::thread(  | 
 | 98 | +                [alg_hnsw, start, end] {  | 
 | 99 | +                    for (size_t i = start; i < end; i++) {  | 
 | 100 | +                       alg_hnsw->addPoint(data.data() + i * dimension, i);  | 
 | 101 | +                    }  | 
 | 102 | +                }  | 
 | 103 | +            )  | 
 | 104 | +        );  | 
 | 105 | +    }  | 
 | 106 | +    for (auto &thread : threads) {  | 
 | 107 | +       thread.join();  | 
 | 108 | +    }  | 
 | 109 | +    threads.clear();  | 
 | 110 | +    auto end_time = std::chrono::system_clock::now();  | 
 | 111 | +    double duration_in_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();  | 
 | 112 | +    double duration_in_seconds = static_cast<double>((std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time)).count()) / 1000.0;  | 
 | 113 | +    size_t qps = (duration_in_seconds == 0 ? total_items : total_items / duration_in_seconds);  | 
 | 114 | +    double latency = (total_items == 0 ? 0 : duration_in_ms / total_items);  | 
 | 115 | +    std::cout << "Start " << num_threads << " thread to add " << total_items << " items to hnsw index, cost "  | 
 | 116 | +              << duration_in_seconds << " seconds, qps: " << qps  << ", latency: " << latency << "ms" << std::endl;  | 
 | 117 | + | 
 | 118 | + | 
 | 119 | +    num_threads = (num_query >= max_thread_num ? max_thread_num : num_query);  | 
 | 120 | +    batch_num = (num_query / (num_threads <= 1 ? 1 : (num_threads - 1))) + 1;  | 
 | 121 | +    start_time = std::chrono::system_clock::now();  | 
 | 122 | +    for (size_t idx = 0; idx < num_query; idx += batch_num) {  | 
 | 123 | +        size_t start = idx;  | 
 | 124 | +        size_t end = std::min(idx + batch_num, num_query);  | 
 | 125 | +        threads.push_back(  | 
 | 126 | +            std::thread(  | 
 | 127 | +                [alg_hnsw, start, end] {  | 
 | 128 | +                    for (size_t i = start; i < end; i++) {  | 
 | 129 | +                        const void* p = query.data() + i * dimension;  | 
 | 130 | +                        auto gd = alg_hnsw->searchKnn(p, topk);  | 
 | 131 | +                    }  | 
 | 132 | +                }  | 
 | 133 | +            )  | 
 | 134 | +        );  | 
 | 135 | +    }  | 
 | 136 | +    for (auto &thread : threads) {  | 
 | 137 | +       thread.join();  | 
 | 138 | +    }  | 
 | 139 | +    threads.clear();  | 
 | 140 | +    end_time = std::chrono::system_clock::now();  | 
 | 141 | +    duration_in_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();  | 
 | 142 | +    duration_in_seconds = static_cast<double>((std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time)).count()) / 1000.0;  | 
 | 143 | +    qps = (duration_in_seconds == 0 ? num_query : num_query / duration_in_seconds);  | 
 | 144 | +    latency = (num_query == 0 ? 0 : duration_in_ms / num_query);  | 
 | 145 | +    std::cout << "Start " << num_threads << " thread to exec " << num_query << " searchKnn, cost "  | 
 | 146 | +              << duration_in_seconds << " seconds, qps: " << qps  << ", latency: " << latency << "ms" << std::endl;  | 
 | 147 | + | 
 | 148 | +    delete alg_hnsw;  | 
 | 149 | +}  | 
 | 150 | + | 
 | 151 | +}  // namespace  | 
 | 152 | + | 
 | 153 | +int main() {  | 
 | 154 | + | 
 | 155 | +    std::mt19937 rng;  | 
 | 156 | +    rng.seed(47);  | 
 | 157 | +    std::uniform_real_distribution<> distrib;  | 
 | 158 | + | 
 | 159 | +    for (size_t i = 0; i < total_items * dimension; ++i) {  | 
 | 160 | +        data[i] = distrib(rng);  | 
 | 161 | +    }  | 
 | 162 | +    for (size_t i = 0; i < num_query * dimension; ++i) {  | 
 | 163 | +        query[i] = distrib(rng);  | 
 | 164 | +    }  | 
 | 165 | + | 
 | 166 | +    test_compatibility(true, true);  | 
 | 167 | +    test_compatibility(false, false);  | 
 | 168 | +    test_compatibility(true, false);  | 
 | 169 | +    test_compatibility(false, true);  | 
 | 170 | + | 
 | 171 | +    test_performace(true);  | 
 | 172 | +    test_performace(false);  | 
 | 173 | + | 
 | 174 | +    return 0;  | 
 | 175 | +}  | 
0 commit comments