aboutsummaryrefslogtreecommitdiffstats
path: root/searchlib/src/vespa/searchlib/tensor/hnsw_multi_best_neighbors.cpp
blob: 86cfc74324723d5e5c7ca07f7621096acdf33b37 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "hnsw_multi_best_neighbors.h"

namespace search::tensor {

HnswMultiBestNeighbors::~HnswMultiBestNeighbors() = default;

std::vector<NearestNeighborIndex::Neighbor>
HnswMultiBestNeighbors::get_neighbors(uint32_t k, double distance_threshold)
{
    while (_docids.size() > k) {
        pop();
    }
    std::vector<NearestNeighborIndex::Neighbor> result;
    result.reserve(_docids.size());
    while (!_candidates.empty()) {
        auto& hit = _candidates.top();
        uint32_t docid = hit.docid;
        if (remove_docid(docid) && (!(hit.distance > distance_threshold))) {
            result.emplace_back(docid, hit.distance);
        }
        _candidates.pop();
    }
    return result;
}

}