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

14 files changed, 532 insertions, 150 deletions

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 <vespa/searchlib/tensor/hnsw_graph.h>
+#include <vespa/searchlib/tensor/hnsw_index_saver.h>
+#include <vespa/searchlib/tensor/hnsw_index_loader.h>
+#include <vespa/vespalib/util/bufferwriter.h>
+#include <vespa/searchlib/util/fileutil.h>
+#include <vespa/vespalib/gtest/gtest.h>
+#include <vector>
+
+#include <vespa/log/log.h>
+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<char> 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<uint32_t>;
+
+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<char> save_original() const {
+        HnswIndexSaver saver(original);
+        VectorBufferWriter vector_writer;
+        saver.save(vector_writer);
+        return vector_writer.output;
+    }
+    void load_copy(std::vector<char> 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 <vespa/vespalib/datastore/array_store.hpp>
+#include <vespa/vespalib/util/rcuvector.hpp>
+
+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<AtomicEntryRef> 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 <vespa/vespalib/datastore/array_store.h>
+#include <vespa/vespalib/datastore/atomic_entry_ref.h>
+#include <vespa/vespalib/datastore/entryref.h>
+#include <vespa/vespalib/util/rcuvector.h>
+
+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<AtomicEntryRef>;
+
+    // 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<AtomicEntryRef, EntryRefType>;
+    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<uint32_t, EntryRefType>;
+    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 <vespa/searchlib/util/state_explorer_utils.h>
 #include <vespa/eval/tensor/dense/typed_cells.h>
@@ -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<PairDist> 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<NearestNeighborIndexSaver>
 HnswIndex::make_saver() const
 {
-    return std::unique_ptr<NearestNeighborIndexSaver>();
+    return std::make_unique<HnswIndexSaver>(_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 <vespa/eval/tensor/dense/typed_cells.h>
 #include <vespa/searchlib/common/bitvector.h>
 #include <vespa/vespalib/datastore/array_store.h>
@@ -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<AtomicEntryRef>;
+    using LinkStore = HnswGraph::LinkStore;
+    using LinkArrayRef = HnswGraph::LinkArrayRef;
+    using LinkArray = vespalib::Array<uint32_t>;
 
-    // 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<AtomicEntryRef, EntryRefType>;
-    using LevelArrayRef = NodeStore::ConstArrayRef;
+    using LevelArrayRef = HnswGraph::LevelArrayRef;
     using LevelArray = vespalib::Array<AtomicEntryRef>;
 
-    // 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<uint32_t, EntryRefType>;
-    using LinkArrayRef = LinkStore::ConstArrayRef;
-    using LinkArray = vespalib::Array<uint32_t>;
-
     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 <vespa/searchlib/util/fileutil.h>
+
+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<const uint32_t *>(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<uint32_t> 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 <vespa/vespalib/util/bufferwriter.h>
+
+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<uint32_t> 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 <vespa/vespalib/datastore/entryref.h>
+#include <vector>
+
+namespace search::tensor {
+
+class HnswGraph;
+
+class HnswIndexSaver : public NearestNeighborIndexSaver {
+public:
+    using LevelVector = std::vector<search::datastore::EntryRef>;
+
+    struct MetaData {
+        uint32_t entry_docid;
+        int32_t  entry_level;
+        std::vector<LevelVector> 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;
 };
 
 }