aboutsummaryrefslogtreecommitdiffstats
path: root/searchlib/src/vespa/searchlib/tensor/hnsw_single_best_neighbors.h
blob: acb14f79b7a26d337a5d31dcbc626205da824fe5 (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
29
30
31
32
33
34
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#pragma once

#include "hnsw_index_utils.h"
#include "nearest_neighbor_index.h"

namespace search::tensor {

/*
 * A priority queue of best neighbors for hnsw index. Used for search
 * when hnsw index has a single node per document.
 */
class HnswSingleBestNeighbors {
    using EntryRef = vespalib::datastore::EntryRef;
    FurthestPriQ _candidates;
public:
    HnswSingleBestNeighbors()
        : _candidates()
    {
    }
    ~HnswSingleBestNeighbors() = default;

    std::vector<NearestNeighborIndex::Neighbor> get_neighbors(uint32_t k, double distance_threshold);
    void push(const HnswCandidate& candidate) { _candidates.push(candidate); }
    void pop() { _candidates.pop(); }
    const HnswCandidateVector& peek() const { return _candidates.peek(); }
    bool empty() const { return _candidates.empty(); }
    const HnswCandidate& top() const { return _candidates.top(); }
    size_t size() const { return _candidates.size(); }
    void emplace(uint32_t nodeid, uint32_t docid, EntryRef ref, double distance) { _candidates.emplace(nodeid, docid, ref, distance); }
};

}