aboutsummaryrefslogtreecommitdiffstats
path: root/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp
blob: 56af2ed3b354a3e95837ed7783d8d03a57a9ed43 (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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "hnsw_graph.h"
#include "hnsw_index.h"
#include <vespa/vespalib/datastore/array_store.hpp>

namespace search::tensor {

HnswGraph::HnswGraph()
  : node_refs(),
    node_refs_size(1u),
    nodes(HnswIndex::make_default_node_store_config(), {}),
    links(HnswIndex::make_default_link_store_config(), {}),
    entry_docid_and_level()
{
    node_refs.ensure_size(1, AtomicEntryRef());
    EntryNode entry;
    set_entry_node(entry);
}

HnswGraph::~HnswGraph() = default;

HnswGraph::NodeRef
HnswGraph::make_node_for_document(uint32_t docid, uint32_t num_levels)
{
    node_refs.ensure_size(docid + 1, AtomicEntryRef());
    // A document cannot be added twice.
    assert(!get_node_ref(docid).valid());
    // Note: The level array instance lives as long as the document is present in the index.
    std::vector<AtomicEntryRef> levels(num_levels, AtomicEntryRef());
    auto node_ref = nodes.add(levels);
    node_refs[docid].store_release(node_ref);
    if (docid >= node_refs_size.load(std::memory_order_relaxed)) {
        node_refs_size.store(docid + 1, std::memory_order_release);
    }
    return node_ref;
}

void
HnswGraph::remove_node_for_document(uint32_t docid)
{
    auto node_ref = get_node_ref(docid);
    assert(node_ref.valid());
    auto levels = nodes.get(node_ref);
    vespalib::datastore::EntryRef invalid;
    node_refs[docid].store_release(invalid);
    // Ensure data referenced through the old ref can be recycled:
    nodes.remove(node_ref);
    for (size_t i = 0; i < levels.size(); ++i) {
        auto old_links_ref = levels[i].load_relaxed();
        links.remove(old_links_ref);
    }
    if (docid + 1 == node_refs_size.load(std::memory_order_relaxed)) {
        trim_node_refs_size();
    }
}

void
HnswGraph::trim_node_refs_size()
{
    uint32_t check_doc_id = node_refs_size.load(std::memory_order_relaxed) - 1;
    while (check_doc_id > 0u && !get_node_ref(check_doc_id).valid()) {
        --check_doc_id;
    }
    node_refs_size.store(check_doc_id + 1, std::memory_order_release);
}

void     
HnswGraph::set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& new_links)
{
    auto new_links_ref = links.add(new_links);
    auto node_ref = get_node_ref(docid);
    assert(node_ref.valid());
    auto levels = nodes.get_writable(node_ref);
    assert(level < levels.size());
    auto old_links_ref = levels[level].load_relaxed();
    levels[level].store_release(new_links_ref);
    links.remove(old_links_ref);
}

HnswGraph::Histograms
HnswGraph::histograms() const
{
    Histograms result;
    size_t num_nodes = node_refs_size.load(std::memory_order_acquire);
    for (size_t i = 0; i < num_nodes; ++i) {
        auto node_ref = acquire_node_ref(i);
        if (node_ref.valid()) {
            uint32_t levels = 0;
            uint32_t l0links = 0;
            auto level_array = nodes.get(node_ref);
            levels = level_array.size();
            if (levels > 0) {
                auto links_ref = level_array[0].load_acquire();
                auto link_array = links.get(links_ref);
                l0links = link_array.size();
            }
            while (result.level_histogram.size() <= levels) {
                result.level_histogram.push_back(0);
            }
            ++result.level_histogram[levels];
            while (result.links_histogram.size() <= l0links) {
                result.links_histogram.push_back(0);
            }
            ++result.links_histogram[l0links];
        }
    }
    return result;
}

void
HnswGraph::set_entry_node(EntryNode node) {
    uint64_t value = node.level;
    value <<= 32;
    value |= node.docid;
    if (node.node_ref.valid()) {
        assert(node.level >= 0);
        assert(node.docid > 0);
    } else {
        assert(node.level == -1);
        assert(node.docid == 0);
    }
    entry_docid_and_level.store(value, std::memory_order_release);
}

} // namespace

namespace vespalib {

template class RcuVectorBase<vespalib::datastore::AtomicEntryRef>;

}