From 3c2f98129e223c0dbf713ad046d7e913228873a0 Mon Sep 17 00:00:00 2001 From: Arne Juul Date: Thu, 26 Mar 2020 07:54:18 +0000 Subject: Add save and load of HNSW index * split graph data out into a simpler struct * implement HnswIndexSaver * implement HnswIndexLoader * add unit test checking that saving and loading works as expected --- searchlib/CMakeLists.txt | 1 + .../src/tests/tensor/hnsw_saver/CMakeLists.txt | 9 + .../tensor/hnsw_saver/hnsw_save_load_test.cpp | 150 +++++++++++++++++ .../src/vespa/searchlib/tensor/CMakeLists.txt | 4 + .../src/vespa/searchlib/tensor/hnsw_graph.cpp | 51 ++++++ searchlib/src/vespa/searchlib/tensor/hnsw_graph.h | 74 +++++++++ .../src/vespa/searchlib/tensor/hnsw_index.cpp | 181 ++++++++------------- searchlib/src/vespa/searchlib/tensor/hnsw_index.h | 48 ++---- .../vespa/searchlib/tensor/hnsw_index_loader.cpp | 48 ++++++ .../src/vespa/searchlib/tensor/hnsw_index_loader.h | 20 +++ .../vespa/searchlib/tensor/hnsw_index_saver.cpp | 57 +++++++ .../src/vespa/searchlib/tensor/hnsw_index_saver.h | 34 ++++ .../tensor/nearest_neighbor_index_saver.cpp | 3 + .../tensor/nearest_neighbor_index_saver.h | 2 +- 14 files changed, 532 insertions(+), 150 deletions(-) create mode 100644 searchlib/src/tests/tensor/hnsw_saver/CMakeLists.txt create mode 100644 searchlib/src/tests/tensor/hnsw_saver/hnsw_save_load_test.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/hnsw_graph.h create mode 100644 searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h create mode 100644 searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h create mode 100644 searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.cpp (limited to 'searchlib') diff --git a/searchlib/CMakeLists.txt b/searchlib/CMakeLists.txt index aaf8f91387e..389bf918b5c 100644 --- a/searchlib/CMakeLists.txt +++ b/searchlib/CMakeLists.txt @@ -214,6 +214,7 @@ vespa_define_module( src/tests/tensor/dense_tensor_store src/tests/tensor/distance_functions src/tests/tensor/hnsw_index + src/tests/tensor/hnsw_saver src/tests/transactionlog src/tests/transactionlogstress src/tests/true diff --git a/searchlib/src/tests/tensor/hnsw_saver/CMakeLists.txt b/searchlib/src/tests/tensor/hnsw_saver/CMakeLists.txt new file mode 100644 index 00000000000..90202e222a7 --- /dev/null +++ b/searchlib/src/tests/tensor/hnsw_saver/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +vespa_add_executable(searchlib_hnsw_save_load_test_app TEST + SOURCES + hnsw_save_load_test.cpp + DEPENDS + searchlib + gtest +) +vespa_add_test(NAME searchlib_hnsw_save_load_test_app COMMAND searchlib_hnsw_save_load_test_app) diff --git a/searchlib/src/tests/tensor/hnsw_saver/hnsw_save_load_test.cpp b/searchlib/src/tests/tensor/hnsw_saver/hnsw_save_load_test.cpp new file mode 100644 index 00000000000..b9e27d413f3 --- /dev/null +++ b/searchlib/src/tests/tensor/hnsw_saver/hnsw_save_load_test.cpp @@ -0,0 +1,150 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include +#include +#include +#include +#include +#include +#include + +#include +LOG_SETUP("hnsw_save_load_test"); + +using namespace search::tensor; +using search::BufferWriter; +using search::fileutil::LoadedBuffer; + +class VectorBufferWriter : public BufferWriter { +private: + char tmp[1024]; +public: + std::vector output; + VectorBufferWriter() { + setup(tmp, 1024); + } + ~VectorBufferWriter() {} + void flush() override { + for (size_t i = 0; i < usedLen(); ++i) { + output.push_back(tmp[i]); + } + rewind(); + } +}; + +using V = std::vector; + +void populate(HnswGraph &graph) { + // no 0 + graph.make_node_for_document(1, 1); + graph.make_node_for_document(2, 2); + // no 3 + graph.make_node_for_document(4, 2); + graph.make_node_for_document(5, 0); + graph.make_node_for_document(6, 1); + + graph.set_link_array(1, 0, V{2, 4, 6}); + graph.set_link_array(2, 0, V{1, 4, 6}); + graph.set_link_array(4, 0, V{1, 2, 6}); + graph.set_link_array(6, 0, V{1, 2, 4}); + graph.set_link_array(2, 1, V{4}); + graph.set_link_array(4, 1, V{2}); + graph.set_entry_node(2, 1); +} + +void modify(HnswGraph &graph) { + graph.remove_node_for_document(2); + graph.remove_node_for_document(6); + graph.make_node_for_document(7, 2); + + graph.set_link_array(1, 0, V{7, 4}); + graph.set_link_array(4, 0, V{7, 2}); + graph.set_link_array(7, 0, V{4, 2}); + graph.set_link_array(4, 1, V{7}); + graph.set_link_array(7, 1, V{4}); + + graph.set_entry_node(4, 1); +} + + +class CopyGraphTest : public ::testing::Test { +public: + HnswGraph original; + HnswGraph copy; + + void expect_empty_d(uint32_t docid) const { + EXPECT_FALSE(copy.node_refs[docid].load_acquire().valid()); + } + + void expect_level_0(uint32_t docid, const V& exp_links) const { + auto levels = copy.get_level_array(docid); + EXPECT_GE(levels.size(), 1); + auto links = copy.get_link_array(docid, 0); + EXPECT_EQ(exp_links.size(), links.size()); + for (size_t i = 0; i < exp_links.size() && i < links.size(); ++i) { + EXPECT_EQ(exp_links[i], links[i]); + } + } + + void expect_level_1(uint32_t docid, const V& exp_links) const { + auto levels = copy.get_level_array(docid); + EXPECT_EQ(2, levels.size()); + auto links = copy.get_link_array(docid, 1); + EXPECT_EQ(exp_links.size(), links.size()); + for (size_t i = 0; i < exp_links.size() && i < links.size(); ++i) { + EXPECT_EQ(exp_links[i], links[i]); + } + } + + std::vector save_original() const { + HnswIndexSaver saver(original); + VectorBufferWriter vector_writer; + saver.save(vector_writer); + return vector_writer.output; + } + void load_copy(std::vector data) { + HnswIndexLoader loader(copy); + LoadedBuffer buffer(&data[0], data.size()); + loader.load(buffer); + } + + void expect_copy_as_populated() const { + EXPECT_EQ(copy.size(), 7); + EXPECT_EQ(copy.entry_docid, 2); + EXPECT_EQ(copy.entry_level, 1); + + expect_empty_d(0); + expect_empty_d(3); + expect_empty_d(5); + + expect_level_0(1, {2, 4, 6}); + expect_level_0(2, {1, 4, 6}); + expect_level_0(4, {1, 2, 6}); + expect_level_0(6, {1, 2, 4}); + + expect_level_1(2, {4}); + expect_level_1(4, {2}); + } +}; + +TEST_F(CopyGraphTest, reconstructs_graph) +{ + populate(original); + auto data = save_original(); + load_copy(data); + expect_copy_as_populated(); +} + +TEST_F(CopyGraphTest, later_changes_ignored) +{ + populate(original); + HnswIndexSaver saver(original); + modify(original); + VectorBufferWriter vector_writer; + saver.save(vector_writer); + auto data = vector_writer.output; + load_copy(data); + expect_copy_as_populated(); +} + +GTEST_MAIN_RUN_ALL_TESTS() diff --git a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt index 7090158c773..0f106f693f8 100644 --- a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt +++ b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt @@ -9,11 +9,15 @@ vespa_add_library(searchlib_tensor OBJECT generic_tensor_attribute.cpp generic_tensor_attribute_saver.cpp generic_tensor_store.cpp + hnsw_graph.cpp hnsw_index.cpp + hnsw_index_loader.cpp + hnsw_index_saver.cpp imported_tensor_attribute_vector.cpp imported_tensor_attribute_vector_read_guard.cpp inv_log_level_generator.cpp nearest_neighbor_index.cpp + nearest_neighbor_index_saver.cpp tensor_attribute.cpp tensor_store.cpp DEPENDS diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp new file mode 100644 index 00000000000..a77cc5b8f63 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp @@ -0,0 +1,51 @@ +// Copyright Verizon Media. 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 +#include + +namespace search::tensor { + +HnswGraph::HnswGraph() + : node_refs(), + nodes(HnswIndex::make_default_node_store_config()), + links(HnswIndex::make_default_link_store_config()), + entry_docid(0), // Note that docid 0 is reserved and never used + entry_level(-1) +{} + +HnswGraph::~HnswGraph() {} + +void +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(!node_refs[docid].load_acquire().valid()); + // Note: The level array instance lives as long as the document is present in the index. + vespalib::Array levels(num_levels, AtomicEntryRef()); + auto node_ref = nodes.add(levels); + node_refs[docid].store_release(node_ref); +} + +void +HnswGraph::remove_node_for_document(uint32_t docid) { + auto node_ref = node_refs[docid].load_acquire(); + nodes.remove(node_ref); + search::datastore::EntryRef invalid; + node_refs[docid].store_release(invalid); +} + +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 = node_refs[docid].load_acquire(); + assert(node_ref.valid()); + auto levels = nodes.get_writable(node_ref); + auto old_links_ref = levels[level].load_acquire(); + levels[level].store_release(new_links_ref); + links.remove(old_links_ref); +} + +} // namespace diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_graph.h b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.h new file mode 100644 index 00000000000..64892d06f09 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.h @@ -0,0 +1,74 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +#include +#include +#include +#include + +namespace search::tensor { + +/** + * Stroage of a hierarchical navigable small world graph (HNSW) + * that is used for approximate K-nearest neighbor search. + */ +struct HnswGraph { + using AtomicEntryRef = search::datastore::AtomicEntryRef; + + // This uses 10 bits for buffer id -> 1024 buffers. + // As we have very short arrays we get less fragmentation with fewer and larger buffers. + using EntryRefType = search::datastore::EntryRefT<22>; + + // Provides mapping from document id -> node reference. + // The reference is used to lookup the node data in NodeStore. + using NodeRefVector = vespalib::RcuVector; + + // This stores the level arrays for all nodes. + // Each node consists of an array of levels (from level 0 to n) where each entry is a reference to the link array at that level. + using NodeStore = search::datastore::ArrayStore; + using StoreConfig = search::datastore::ArrayStoreConfig; + using LevelArrayRef = NodeStore::ConstArrayRef; + + // This stores all link arrays. + // A link array consists of the document ids of the nodes a particular node is linked to. + using LinkStore = search::datastore::ArrayStore; + using LinkArrayRef = LinkStore::ConstArrayRef; + + NodeRefVector node_refs; + NodeStore nodes; + LinkStore links; + uint32_t entry_docid; + int32_t entry_level; + + HnswGraph(); + + ~HnswGraph(); + + void make_node_for_document(uint32_t docid, uint32_t num_levels); + + void remove_node_for_document(uint32_t docid); + + LevelArrayRef get_level_array(uint32_t docid) const { + auto node_ref = node_refs[docid].load_acquire(); + assert(node_ref.valid()); + return nodes.get(node_ref); + } + + LinkArrayRef get_link_array(uint32_t docid, uint32_t level) const { + auto levels = get_level_array(docid); + assert(level < levels.size()); + return links.get(levels[level].load_acquire()); + } + + void set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& new_links); + + void set_entry_node(uint32_t docid, int32_t level) { + entry_docid = docid; + entry_level = level; + } + + size_t size() const { return node_refs.size(); } +}; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp index 19b02d18893..7ed2a160107 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp @@ -2,7 +2,8 @@ #include "distance_function.h" #include "hnsw_index.h" -#include "nearest_neighbor_index_saver.h" +#include "hnsw_index_loader.h" +#include "hnsw_index_saver.h" #include "random_level_generator.h" #include #include @@ -67,54 +68,6 @@ HnswIndex::max_links_for_level(uint32_t level) const return (level == 0) ? _cfg.max_links_at_level_0() : _cfg.max_links_on_inserts(); } -void -HnswIndex::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(!_node_refs[docid].load_acquire().valid()); - - // Note: The level array instance lives as long as the document is present in the index. - LevelArray levels(num_levels, AtomicEntryRef()); - auto node_ref = _nodes.add(levels); - _node_refs[docid].store_release(node_ref); -} - -void -HnswIndex::remove_node_for_document(uint32_t docid) -{ - auto node_ref = _node_refs[docid].load_acquire(); - _nodes.remove(node_ref); - EntryRef invalid; - _node_refs[docid].store_release(invalid); -} - -HnswIndex::LevelArrayRef -HnswIndex::get_level_array(uint32_t docid) const -{ - auto node_ref = _node_refs[docid].load_acquire(); - return _nodes.get(node_ref); -} - -HnswIndex::LinkArrayRef -HnswIndex::get_link_array(uint32_t docid, uint32_t level) const -{ - auto levels = get_level_array(docid); - assert(level < levels.size()); - return _links.get(levels[level].load_acquire()); -} - -void -HnswIndex::set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& links) -{ - auto new_links_ref = _links.add(links); - auto node_ref = _node_refs[docid].load_acquire(); - auto levels = _nodes.get_writable(node_ref); - auto old_links_ref = levels[level].load_acquire(); - levels[level].store_release(new_links_ref); - _links.remove(old_links_ref); -} - bool HnswIndex::have_closer_distance(HnswCandidate candidate, const LinkArrayRef& result) const { @@ -183,7 +136,7 @@ HnswIndex::select_neighbors(const HnswCandidateVector& neighbors, uint32_t max_l void HnswIndex::shrink_if_needed(uint32_t docid, uint32_t level) { - auto old_links = get_link_array(docid, level); + auto old_links = _graph.get_link_array(docid, level); uint32_t max_links = max_links_for_level(level); if (old_links.size() > max_links) { HnswCandidateVector neighbors; @@ -192,7 +145,7 @@ HnswIndex::shrink_if_needed(uint32_t docid, uint32_t level) neighbors.emplace_back(neighbor_docid, dist); } auto split = select_neighbors(neighbors, max_links); - set_link_array(docid, level, split.used); + _graph.set_link_array(docid, level, split.used); for (uint32_t removed_docid : split.unused) { remove_link_to(removed_docid, docid, level); } @@ -202,9 +155,9 @@ HnswIndex::shrink_if_needed(uint32_t docid, uint32_t level) void HnswIndex::connect_new_node(uint32_t docid, const LinkArrayRef &neighbors, uint32_t level) { - set_link_array(docid, level, neighbors); + _graph.set_link_array(docid, level, neighbors); for (uint32_t neighbor_docid : neighbors) { - auto old_links = get_link_array(neighbor_docid, level); + auto old_links = _graph.get_link_array(neighbor_docid, level); add_link_to(neighbor_docid, level, old_links, docid); } for (uint32_t neighbor_docid : neighbors) { @@ -216,11 +169,11 @@ void HnswIndex::remove_link_to(uint32_t remove_from, uint32_t remove_id, uint32_t level) { LinkArray new_links; - auto old_links = get_link_array(remove_from, level); + auto old_links = _graph.get_link_array(remove_from, level); for (uint32_t id : old_links) { if (id != remove_id) new_links.push_back(id); } - set_link_array(remove_from, level, new_links); + _graph.set_link_array(remove_from, level, new_links); } @@ -245,7 +198,7 @@ HnswIndex::find_nearest_in_layer(const TypedCells& input, const HnswCandidate& e bool keep_searching = true; while (keep_searching) { keep_searching = false; - for (uint32_t neighbor_docid : get_link_array(nearest.docid, level)) { + for (uint32_t neighbor_docid : _graph.get_link_array(nearest.docid, level)) { double dist = calc_distance(input, neighbor_docid); if (dist < nearest.distance) { nearest = HnswCandidate(neighbor_docid, dist); @@ -260,7 +213,7 @@ void HnswIndex::search_layer(const TypedCells& input, uint32_t neighbors_to_find, FurthestPriQ& best_neighbors, uint32_t level) const { NearestPriQ candidates; - uint32_t doc_id_limit = _node_refs.size(); + uint32_t doc_id_limit = _graph.node_refs.size(); auto visited = _visited_set_pool.get(doc_id_limit); for (const auto &entry : best_neighbors.peek()) { assert(entry.docid < doc_id_limit); @@ -275,7 +228,7 @@ HnswIndex::search_layer(const TypedCells& input, uint32_t neighbors_to_find, Fur break; } candidates.pop(); - for (uint32_t neighbor_docid : get_link_array(cand.docid, level)) { + for (uint32_t neighbor_docid : _graph.get_link_array(cand.docid, level)) { if ((neighbor_docid >= doc_id_limit) || visited.is_marked(neighbor_docid)) { continue; } @@ -295,15 +248,12 @@ HnswIndex::search_layer(const TypedCells& input, uint32_t neighbors_to_find, Fur HnswIndex::HnswIndex(const DocVectorAccess& vectors, DistanceFunction::UP distance_func, RandomLevelGenerator::UP level_generator, const Config& cfg) - : _vectors(vectors), + : + _graph(), + _vectors(vectors), _distance_func(std::move(distance_func)), _level_generator(std::move(level_generator)), - _cfg(cfg), - _node_refs(), - _nodes(make_default_node_store_config()), - _links(make_default_link_store_config()), - _entry_docid(0), // Note that docid 0 is reserved and never used - _entry_level(-1) + _cfg(cfg) { } @@ -315,16 +265,16 @@ HnswIndex::add_document(uint32_t docid) auto input = get_vector(docid); // TODO: Add capping on num_levels int level = _level_generator->max_level(); - make_node_for_document(docid, level + 1); - if (_entry_docid == 0) { - _entry_docid = docid; - _entry_level = level; + _graph.make_node_for_document(docid, level + 1); + uint32_t entry_docid = get_entry_docid(); + if (entry_docid == 0) { + _graph.set_entry_node(docid, level); return; } - int search_level = _entry_level; - double entry_dist = calc_distance(input, _entry_docid); - HnswCandidate entry_point(_entry_docid, entry_dist); + int search_level = get_entry_level(); + double entry_dist = calc_distance(input, entry_docid); + HnswCandidate entry_point(entry_docid, entry_dist); while (search_level > level) { entry_point = find_nearest_in_layer(input, entry_point, search_level); --search_level; @@ -332,7 +282,7 @@ HnswIndex::add_document(uint32_t docid) FurthestPriQ best_neighbors; best_neighbors.push(entry_point); - search_level = std::min(level, _entry_level); + search_level = std::min(level, search_level); // Insert the added document in each level it should exist in. while (search_level >= 0) { @@ -342,9 +292,8 @@ HnswIndex::add_document(uint32_t docid) connect_new_node(docid, neighbors.used, search_level); --search_level; } - if (level > _entry_level) { - _entry_docid = docid; - _entry_level = level; + if (level > get_entry_level()) { + _graph.set_entry_node(docid, level); } } @@ -354,7 +303,7 @@ HnswIndex::mutual_reconnect(const LinkArrayRef &cluster, uint32_t level) std::vector pairs; for (uint32_t i = 0; i + 1 < cluster.size(); ++i) { uint32_t n_id_1 = cluster[i]; - LinkArrayRef n_list_1 = get_link_array(n_id_1, level); + LinkArrayRef n_list_1 = _graph.get_link_array(n_id_1, level); for (uint32_t j = i + 1; j < cluster.size(); ++j) { uint32_t n_id_2 = cluster[j]; if (has_link_to(n_list_1, n_id_2)) continue; @@ -363,10 +312,10 @@ HnswIndex::mutual_reconnect(const LinkArrayRef &cluster, uint32_t level) } std::sort(pairs.begin(), pairs.end()); for (const PairDist & pair : pairs) { - LinkArrayRef old_links_1 = get_link_array(pair.id_first, level); + LinkArrayRef old_links_1 = _graph.get_link_array(pair.id_first, level); if (old_links_1.size() >= _cfg.max_links_on_inserts()) continue; - LinkArrayRef old_links_2 = get_link_array(pair.id_second, level); + LinkArrayRef old_links_2 = _graph.get_link_array(pair.id_second, level); if (old_links_2.size() >= _cfg.max_links_on_inserts()) continue; add_link_to(pair.id_first, level, old_links_1, pair.id_second); @@ -377,27 +326,25 @@ HnswIndex::mutual_reconnect(const LinkArrayRef &cluster, uint32_t level) void HnswIndex::remove_document(uint32_t docid) { - bool need_new_entrypoint = (docid == _entry_docid); + bool need_new_entrypoint = (docid == get_entry_docid()); LinkArray empty; - LevelArrayRef node_levels = get_level_array(docid); + LevelArrayRef node_levels = _graph.get_level_array(docid); for (int level = node_levels.size(); level-- > 0; ) { - LinkArrayRef my_links = get_link_array(docid, level); + LinkArrayRef my_links = _graph.get_link_array(docid, level); for (uint32_t neighbor_id : my_links) { if (need_new_entrypoint) { - _entry_docid = neighbor_id; - _entry_level = level; + _graph.set_entry_node(neighbor_id, level); need_new_entrypoint = false; } remove_link_to(neighbor_id, docid, level); } mutual_reconnect(my_links, level); - set_link_array(docid, level, empty); + _graph.set_link_array(docid, level, empty); } if (need_new_entrypoint) { - _entry_docid = 0; - _entry_level = -1; + _graph.set_entry_node(0, -1); } - remove_node_for_document(docid); + _graph.remove_node_for_document(docid); } void @@ -405,26 +352,26 @@ HnswIndex::transfer_hold_lists(generation_t current_gen) { // Note: RcuVector transfers hold lists as part of reallocation based on current generation. // We need to set the next generation here, as it is incremented on a higher level right after this call. - _node_refs.setGeneration(current_gen + 1); - _nodes.transferHoldLists(current_gen); - _links.transferHoldLists(current_gen); + _graph.node_refs.setGeneration(current_gen + 1); + _graph.nodes.transferHoldLists(current_gen); + _graph.links.transferHoldLists(current_gen); } void HnswIndex::trim_hold_lists(generation_t first_used_gen) { - _node_refs.removeOldGenerations(first_used_gen); - _nodes.trimHoldLists(first_used_gen); - _links.trimHoldLists(first_used_gen); + _graph.node_refs.removeOldGenerations(first_used_gen); + _graph.nodes.trimHoldLists(first_used_gen); + _graph.links.trimHoldLists(first_used_gen); } vespalib::MemoryUsage HnswIndex::memory_usage() const { vespalib::MemoryUsage result; - result.merge(_node_refs.getMemoryUsage()); - result.merge(_nodes.getMemoryUsage()); - result.merge(_links.getMemoryUsage()); + result.merge(_graph.node_refs.getMemoryUsage()); + result.merge(_graph.nodes.getMemoryUsage()); + result.merge(_graph.links.getMemoryUsage()); result.merge(_visited_set_pool.memory_usage()); return result; } @@ -439,13 +386,17 @@ HnswIndex::get_state(const vespalib::slime::Inserter& inserter) const std::unique_ptr HnswIndex::make_saver() const { - return std::unique_ptr(); + return std::make_unique(_graph); } void HnswIndex::load(const fileutil::LoadedBuffer& buf) { - (void) buf; + assert(get_entry_docid() == 0); // cannot load after index has data + HnswIndexLoader loader(_graph); + if (! loader.load(buf)) { + fprintf(stderr, "ERROR loading HNSW graph\n"); + } } struct NeighborsByDocId { @@ -476,12 +427,13 @@ FurthestPriQ HnswIndex::top_k_candidates(const TypedCells &vector, uint32_t k) const { FurthestPriQ best_neighbors; - if (_entry_level < 0) { + if (get_entry_level() < 0) { return best_neighbors; } - double entry_dist = calc_distance(vector, _entry_docid); - HnswCandidate entry_point(_entry_docid, entry_dist); - int search_level = _entry_level; + uint32_t entry_docid = get_entry_docid(); + int search_level = get_entry_level(); + double entry_dist = calc_distance(vector, entry_docid); + HnswCandidate entry_point(entry_docid, entry_dist); while (search_level > 0) { entry_point = find_nearest_in_layer(vector, entry_point, search_level); --search_level; @@ -494,14 +446,14 @@ HnswIndex::top_k_candidates(const TypedCells &vector, uint32_t k) const HnswNode HnswIndex::get_node(uint32_t docid) const { - auto node_ref = _node_refs[docid].load_acquire(); + auto node_ref = _graph.node_refs[docid].load_acquire(); if (!node_ref.valid()) { return HnswNode(); } - auto levels = _nodes.get(node_ref); + auto levels = _graph.nodes.get(node_ref); HnswNode::LevelArray result; for (const auto& links_ref : levels) { - auto links = _links.get(links_ref.load_acquire()); + auto links = _graph.links.get(links_ref.load_acquire()); HnswNode::LinkArray result_links(links.begin(), links.end()); std::sort(result_links.begin(), result_links.end()); result.push_back(result_links); @@ -514,14 +466,13 @@ HnswIndex::set_node(uint32_t docid, const HnswNode &node) { size_t num_levels = node.size(); assert(num_levels > 0); - make_node_for_document(docid, num_levels); + _graph.make_node_for_document(docid, num_levels); for (size_t level = 0; level < num_levels; ++level) { connect_new_node(docid, node.level(level), level); } int max_level = num_levels - 1; - if (_entry_level < max_level) { - _entry_docid = docid; - _entry_level = max_level; + if (get_entry_level() < max_level) { + _graph.set_entry_node(docid, max_level); } } @@ -529,15 +480,15 @@ bool HnswIndex::check_link_symmetry() const { bool all_sym = true; - for (size_t docid = 0; docid < _node_refs.size(); ++docid) { - auto node_ref = _node_refs[docid].load_acquire(); + for (size_t docid = 0; docid < _graph.node_refs.size(); ++docid) { + auto node_ref = _graph.node_refs[docid].load_acquire(); if (node_ref.valid()) { - auto levels = _nodes.get(node_ref); + auto levels = _graph.nodes.get(node_ref); uint32_t level = 0; for (const auto& links_ref : levels) { - auto links = _links.get(links_ref.load_acquire()); + auto links = _graph.links.get(links_ref.load_acquire()); for (auto neighbor_docid : links) { - auto neighbor_links = get_link_array(neighbor_docid, level); + auto neighbor_links = _graph.get_link_array(neighbor_docid, level); if (! has_link_to(neighbor_links, docid)) { all_sym = false; } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h index 1185acd9624..12eecb7e44c 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h @@ -8,6 +8,7 @@ #include "hnsw_node.h" #include "nearest_neighbor_index.h" #include "random_level_generator.h" +#include "hnsw_graph.h" #include #include #include @@ -57,54 +58,30 @@ public: }; protected: - using AtomicEntryRef = search::datastore::AtomicEntryRef; + using AtomicEntryRef = HnswGraph::AtomicEntryRef; + using NodeStore = HnswGraph::NodeStore; - // This uses 10 bits for buffer id -> 1024 buffers. - // As we have very short arrays we get less fragmentation with fewer and larger buffers. - using EntryRefType = search::datastore::EntryRefT<22>; - - // Provides mapping from document id -> node reference. - // The reference is used to lookup the node data in NodeStore. - using NodeRefVector = vespalib::RcuVector; + using LinkStore = HnswGraph::LinkStore; + using LinkArrayRef = HnswGraph::LinkArrayRef; + using LinkArray = vespalib::Array; - // This stores the level arrays for all nodes. - // Each node consists of an array of levels (from level 0 to n) where each entry is a reference to the link array at that level. - using NodeStore = search::datastore::ArrayStore; - using LevelArrayRef = NodeStore::ConstArrayRef; + using LevelArrayRef = HnswGraph::LevelArrayRef; using LevelArray = vespalib::Array; - // This stores all link arrays. - // A link array consists of the document ids of the nodes a particular node is linked to. - using LinkStore = search::datastore::ArrayStore; - using LinkArrayRef = LinkStore::ConstArrayRef; - using LinkArray = vespalib::Array; - using TypedCells = vespalib::tensor::TypedCells; + HnswGraph _graph; const DocVectorAccess& _vectors; DistanceFunction::UP _distance_func; RandomLevelGenerator::UP _level_generator; Config _cfg; - NodeRefVector _node_refs; - NodeStore _nodes; - LinkStore _links; mutable vespalib::ReusableSetPool _visited_set_pool; - uint32_t _entry_docid; - int _entry_level; - - static search::datastore::ArrayStoreConfig make_default_node_store_config(); - static search::datastore::ArrayStoreConfig make_default_link_store_config(); uint32_t max_links_for_level(uint32_t level) const; - void make_node_for_document(uint32_t docid, uint32_t num_levels); - void remove_node_for_document(uint32_t docid); - LevelArrayRef get_level_array(uint32_t docid) const; - LinkArrayRef get_link_array(uint32_t docid, uint32_t level) const; - void set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& links); void add_link_to(uint32_t docid, uint32_t level, const LinkArrayRef& old_links, uint32_t new_link) { LinkArray new_links(old_links.begin(), old_links.end()); new_links.push_back(new_link); - set_link_array(docid, level, new_links); + _graph.set_link_array(docid, level, new_links); } /** @@ -163,13 +140,16 @@ public: FurthestPriQ top_k_candidates(const TypedCells &vector, uint32_t k) const; - uint32_t get_entry_docid() const { return _entry_docid; } - uint32_t get_entry_level() const { return _entry_level; } + uint32_t get_entry_docid() const { return _graph.entry_docid; } + int32_t get_entry_level() const { return _graph.entry_level; } // Should only be used by unit tests. HnswNode get_node(uint32_t docid) const; void set_node(uint32_t docid, const HnswNode &node); bool check_link_symmetry() const; + + static search::datastore::ArrayStoreConfig make_default_node_store_config(); + static search::datastore::ArrayStoreConfig make_default_link_store_config(); }; } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp new file mode 100644 index 00000000000..6626515afdb --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp @@ -0,0 +1,48 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "hnsw_index_loader.h" +#include "hnsw_graph.h" +#include + +namespace search::tensor { + +HnswIndexLoader::~HnswIndexLoader() {} + +HnswIndexLoader::HnswIndexLoader(HnswGraph &graph) + : _graph(graph) +{ +} + +bool +HnswIndexLoader::load(const fileutil::LoadedBuffer& buf) +{ + size_t num_readable = buf.size(sizeof(uint32_t)); + if (num_readable < 3) return false; + const uint32_t *ptr = static_cast(buf.buffer()); + const uint32_t *end = ptr + num_readable; + uint32_t entry_docid = *ptr++; + int32_t entry_level = *ptr++; + uint32_t num_nodes = *ptr++; + std::vector link_array; + for (uint32_t docid = 0; docid < num_nodes; ++docid) { + if (ptr == end) return false; + uint32_t num_levels = *ptr++; + if (num_levels > 0) { + _graph.make_node_for_document(docid, num_levels); + for (uint32_t level = 0; level < num_levels; ++level) { + uint32_t num_links = *ptr++; + link_array.clear(); + while (num_links-- > 0) { + if (ptr == end) return false; + link_array.push_back(*ptr++); + } + _graph.set_link_array(docid, level, link_array); + } + } + } + _graph.set_entry_node(entry_docid, entry_level); + return true; +} + + +} diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h new file mode 100644 index 00000000000..e7653b37946 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h @@ -0,0 +1,20 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +namespace search::fileutil { class LoadedBuffer; } + +namespace search::tensor { + +class HnswGraph; + +class HnswIndexLoader { +public: + HnswIndexLoader(HnswGraph &graph); + ~HnswIndexLoader(); + bool load(const fileutil::LoadedBuffer& buf); +private: + HnswGraph &_graph; +}; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp new file mode 100644 index 00000000000..e9d70a17794 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp @@ -0,0 +1,57 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "hnsw_index_saver.h" +#include "hnsw_graph.h" +#include + +namespace search::tensor { + +HnswIndexSaver::~HnswIndexSaver() {} + +HnswIndexSaver::HnswIndexSaver(const HnswGraph &graph) + : _graph(graph), _meta_data() +{ + _meta_data.entry_docid = _graph.entry_docid; + _meta_data.entry_level = _graph.entry_level; + size_t num_nodes = _graph.node_refs.size(); + _meta_data.nodes.reserve(num_nodes); + for (size_t i = 0; i < num_nodes; ++i) { + LevelVector node; + auto node_ref = _graph.node_refs[i].load_acquire(); + if (node_ref.valid()) { + auto levels = _graph.nodes.get(node_ref); + for (const auto& links_ref : levels) { + auto level = links_ref.load_acquire(); + node.push_back(level); + } + } + _meta_data.nodes.emplace_back(std::move(node)); + } +} + +void +HnswIndexSaver::save(BufferWriter& writer) const +{ + writer.write(&_meta_data.entry_docid, sizeof(uint32_t)); + writer.write(&_meta_data.entry_level, sizeof(int32_t)); + uint32_t num_nodes = _meta_data.nodes.size(); + writer.write(&num_nodes, sizeof(uint32_t)); + for (const auto &node : _meta_data.nodes) { + uint32_t num_levels = node.size(); + writer.write(&num_levels, sizeof(uint32_t)); + for (auto links_ref : node) { + if (links_ref.valid()) { + vespalib::ConstArrayRef link_array = _graph.links.get(links_ref); + uint32_t num_links = link_array.size(); + writer.write(&num_links, sizeof(uint32_t)); + writer.write(link_array.cbegin(), sizeof(uint32_t)*num_links); + } else { + uint32_t num_links = 0; + writer.write(&num_links, sizeof(uint32_t)); + } + } + } + writer.flush(); +} + +} diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h new file mode 100644 index 00000000000..70dc37e38c5 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h @@ -0,0 +1,34 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +#include "nearest_neighbor_index_saver.h" +#include +#include + +namespace search::tensor { + +class HnswGraph; + +class HnswIndexSaver : public NearestNeighborIndexSaver { +public: + using LevelVector = std::vector; + + struct MetaData { + uint32_t entry_docid; + int32_t entry_level; + std::vector nodes; + MetaData() : entry_docid(0), entry_level(-1), nodes() {} + }; + + HnswIndexSaver(const HnswGraph &graph); + ~HnswIndexSaver(); + void save(BufferWriter& writer) const override; + +private: + const HnswGraph &_graph; + MetaData _meta_data; + +}; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.cpp b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.cpp new file mode 100644 index 00000000000..4b293488737 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.cpp @@ -0,0 +1,3 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "nearest_neighbor_index_saver.h" diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.h index 7d599ac31c8..cee48d63359 100644 --- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.h +++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_saver.h @@ -21,7 +21,7 @@ namespace search::tensor { class NearestNeighborIndexSaver { public: virtual ~NearestNeighborIndexSaver() {} - virtual void save(BufferWriter& writer) const; + virtual void save(BufferWriter& writer) const = 0; }; } -- cgit v1.2.3 From d217c94fcbf129e4a557264b64b76f8563c1d892 Mon Sep 17 00:00:00 2001 From: Arne Juul Date: Fri, 27 Mar 2020 12:59:14 +0000 Subject: review follow-up: * add documentation comments * style fixes * return bool from load() for error handling --- .../tensorattribute/tensorattribute_test.cpp | 2 +- .../src/vespa/searchlib/tensor/hnsw_graph.cpp | 6 ++++-- .../src/vespa/searchlib/tensor/hnsw_index.cpp | 6 ++---- searchlib/src/vespa/searchlib/tensor/hnsw_index.h | 2 +- .../vespa/searchlib/tensor/hnsw_index_loader.cpp | 23 +++++++++++----------- .../src/vespa/searchlib/tensor/hnsw_index_loader.h | 15 ++++++++++++++ .../vespa/searchlib/tensor/hnsw_index_saver.cpp | 14 ++++++------- .../src/vespa/searchlib/tensor/hnsw_index_saver.h | 23 ++++++++++++---------- .../searchlib/tensor/nearest_neighbor_index.h | 2 +- 9 files changed, 55 insertions(+), 38 deletions(-) (limited to 'searchlib') diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp index b6bdee4f94d..00450eab21a 100644 --- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp +++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp @@ -148,7 +148,7 @@ public: std::unique_ptr make_saver() const override { return std::unique_ptr(); } - void load(const search::fileutil::LoadedBuffer&) override {} + bool load(const search::fileutil::LoadedBuffer&) override { return false; } std::vector find_top_k(uint32_t k, vespalib::tensor::TypedCells vector, uint32_t explore_k) const override { (void) k; (void) vector; diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp index a77cc5b8f63..db8d1067980 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp @@ -18,7 +18,8 @@ HnswGraph::HnswGraph() HnswGraph::~HnswGraph() {} void -HnswGraph::make_node_for_document(uint32_t docid, uint32_t num_levels) { +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(!node_refs[docid].load_acquire().valid()); @@ -29,7 +30,8 @@ HnswGraph::make_node_for_document(uint32_t docid, uint32_t num_levels) { } void -HnswGraph::remove_node_for_document(uint32_t docid) { +HnswGraph::remove_node_for_document(uint32_t docid) +{ auto node_ref = node_refs[docid].load_acquire(); nodes.remove(node_ref); search::datastore::EntryRef invalid; diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp index 7ed2a160107..de6daba650c 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp @@ -389,14 +389,12 @@ HnswIndex::make_saver() const return std::make_unique(_graph); } -void +bool HnswIndex::load(const fileutil::LoadedBuffer& buf) { assert(get_entry_docid() == 0); // cannot load after index has data HnswIndexLoader loader(_graph); - if (! loader.load(buf)) { - fprintf(stderr, "ERROR loading HNSW graph\n"); - } + return loader.load(buf); } struct NeighborsByDocId { diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h index 12eecb7e44c..b5d57c2ebfd 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h @@ -133,7 +133,7 @@ public: void get_state(const vespalib::slime::Inserter& inserter) const override; std::unique_ptr make_saver() const override; - void load(const fileutil::LoadedBuffer& buf) override; + bool load(const fileutil::LoadedBuffer& buf) override; std::vector find_top_k(uint32_t k, TypedCells vector, uint32_t explore_k) const override; const DistanceFunction *distance_function() const override { return _distance_func.get(); } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp index 6626515afdb..75c62e5202f 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp @@ -9,7 +9,7 @@ namespace search::tensor { HnswIndexLoader::~HnswIndexLoader() {} HnswIndexLoader::HnswIndexLoader(HnswGraph &graph) - : _graph(graph) + : _graph(graph), _ptr(nullptr), _end(nullptr), _failed(false) { } @@ -17,29 +17,28 @@ bool HnswIndexLoader::load(const fileutil::LoadedBuffer& buf) { size_t num_readable = buf.size(sizeof(uint32_t)); - if (num_readable < 3) return false; - const uint32_t *ptr = static_cast(buf.buffer()); - const uint32_t *end = ptr + num_readable; - uint32_t entry_docid = *ptr++; - int32_t entry_level = *ptr++; - uint32_t num_nodes = *ptr++; + _ptr = static_cast(buf.buffer()); + _end = _ptr + num_readable; + uint32_t entry_docid = nextVal(); + int32_t entry_level = nextVal(); + uint32_t num_nodes = nextVal(); std::vector link_array; for (uint32_t docid = 0; docid < num_nodes; ++docid) { - if (ptr == end) return false; - uint32_t num_levels = *ptr++; + uint32_t num_levels = nextVal(); if (num_levels > 0) { _graph.make_node_for_document(docid, num_levels); for (uint32_t level = 0; level < num_levels; ++level) { - uint32_t num_links = *ptr++; + uint32_t num_links = nextVal(); link_array.clear(); while (num_links-- > 0) { - if (ptr == end) return false; - link_array.push_back(*ptr++); + link_array.push_back(nextVal()); } _graph.set_link_array(docid, level, link_array); } } } + if (_failed) return false; + _graph.node_refs.ensure_size(num_nodes); _graph.set_entry_node(entry_docid, entry_level); return true; } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h index e7653b37946..abc68889a1b 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h @@ -2,12 +2,17 @@ #pragma once +#include + namespace search::fileutil { class LoadedBuffer; } namespace search::tensor { class HnswGraph; +/** + * Implements loading of HNSW graph structure from binary format. + **/ class HnswIndexLoader { public: HnswIndexLoader(HnswGraph &graph); @@ -15,6 +20,16 @@ public: bool load(const fileutil::LoadedBuffer& buf); private: HnswGraph &_graph; + const uint32_t *_ptr; + const uint32_t *_end; + bool _failed; + uint32_t nextVal() { + if (__builtin_expect((_ptr == _end), false)) { + _failed = true; + return 0; + } + return *_ptr++; + } }; } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp index e9d70a17794..acff30f8cbf 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp @@ -9,17 +9,17 @@ namespace search::tensor { HnswIndexSaver::~HnswIndexSaver() {} HnswIndexSaver::HnswIndexSaver(const HnswGraph &graph) - : _graph(graph), _meta_data() + : _graph_links(graph.links), _meta_data() { - _meta_data.entry_docid = _graph.entry_docid; - _meta_data.entry_level = _graph.entry_level; - size_t num_nodes = _graph.node_refs.size(); + _meta_data.entry_docid = graph.entry_docid; + _meta_data.entry_level = graph.entry_level; + size_t num_nodes = graph.node_refs.size(); _meta_data.nodes.reserve(num_nodes); for (size_t i = 0; i < num_nodes; ++i) { LevelVector node; - auto node_ref = _graph.node_refs[i].load_acquire(); + auto node_ref = graph.node_refs[i].load_acquire(); if (node_ref.valid()) { - auto levels = _graph.nodes.get(node_ref); + auto levels = graph.nodes.get(node_ref); for (const auto& links_ref : levels) { auto level = links_ref.load_acquire(); node.push_back(level); @@ -41,7 +41,7 @@ HnswIndexSaver::save(BufferWriter& writer) const writer.write(&num_levels, sizeof(uint32_t)); for (auto links_ref : node) { if (links_ref.valid()) { - vespalib::ConstArrayRef link_array = _graph.links.get(links_ref); + vespalib::ConstArrayRef link_array = _graph_links.get(links_ref); uint32_t num_links = link_array.size(); writer.write(&num_links, sizeof(uint32_t)); writer.write(link_array.cbegin(), sizeof(uint32_t)*num_links); diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h index 70dc37e38c5..d1d8e0db19d 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h @@ -3,32 +3,35 @@ #pragma once #include "nearest_neighbor_index_saver.h" +#include "hnsw_graph.h" #include #include namespace search::tensor { -class HnswGraph; - +/** + * Implements saving of HNSW graph structure in binary format. + * The constructor takes a snapshot of all meta-data, but + * the links will be fetched from the graph in the save() + * method. + **/ class HnswIndexSaver : public NearestNeighborIndexSaver { public: using LevelVector = std::vector; + HnswIndexSaver(const HnswGraph &graph); + ~HnswIndexSaver(); + void save(BufferWriter& writer) const override; + +private: struct MetaData { uint32_t entry_docid; int32_t entry_level; std::vector nodes; MetaData() : entry_docid(0), entry_level(-1), nodes() {} }; - - HnswIndexSaver(const HnswGraph &graph); - ~HnswIndexSaver(); - void save(BufferWriter& writer) const override; - -private: - const HnswGraph &_graph; + const HnswGraph::LinkStore &_graph_links; MetaData _meta_data; - }; } diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h index bb6ef012a56..aca2ce2af66 100644 --- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h +++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h @@ -47,7 +47,7 @@ public: * and the caller ensures that an attribute read guard is held during the lifetime of the saver. */ virtual std::unique_ptr make_saver() const = 0; - virtual void load(const fileutil::LoadedBuffer& buf) = 0; + virtual bool load(const fileutil::LoadedBuffer& buf) = 0; virtual std::vector find_top_k(uint32_t k, vespalib::tensor::TypedCells vector, -- cgit v1.2.3