Merge pull request #18753 from vespa-engine/balder/more-efficient-save-of-hnsw-index

Instead of having one large array of individually allocated vectors use

3 files changed, 47 insertions, 16 deletions

diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp
index 6593a60d6b5..8ac613c0976 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.cpp
@@ -6,7 +6,31 @@
 
 namespace search::tensor {
 
-HnswIndexSaver::~HnswIndexSaver() {}
+namespace {
+
+size_t
+count_valid_link_arrays(const HnswGraph & graph) {
+    size_t count(0);
+    size_t num_nodes = graph.node_refs.size();
+    for (size_t i = 0; i < num_nodes; ++i) {
+        auto node_ref = graph.node_refs[i].load_acquire();
+        if (node_ref.valid()) {
+            count += graph.nodes.get(node_ref).size();
+        }
+    }
+    return count;
+}
+
+}
+
+HnswIndexSaver::MetaData::MetaData()
+    : entry_docid(0),
+      entry_level(-1),
+      refs(),
+      nodes()
+{}
+HnswIndexSaver::MetaData::~MetaData() = default;
+HnswIndexSaver::~HnswIndexSaver() = default;
 
 HnswIndexSaver::HnswIndexSaver(const HnswGraph &graph)
     : _graph_links(graph.links), _meta_data()
@@ -15,19 +39,22 @@ HnswIndexSaver::HnswIndexSaver(const HnswGraph &graph)
     _meta_data.entry_docid = entry.docid;
     _meta_data.entry_level = entry.level;
     size_t num_nodes = graph.node_refs.size();
-    _meta_data.nodes.reserve(num_nodes);
+    assert (num_nodes <= (std::numeric_limits<uint32_t>::max() - 1));
+    size_t link_array_count = count_valid_link_arrays(graph);
+    assert (link_array_count <= std::numeric_limits<uint32_t>::max());
+    _meta_data.refs.reserve(link_array_count);
+    _meta_data.nodes.reserve(num_nodes+1);
     for (size_t i = 0; i < num_nodes; ++i) {
-        LevelVector node;
+        _meta_data.nodes.push_back(_meta_data.refs.size());
         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.refs.push_back(links_ref.load_acquire());
             }
         }
-        _meta_data.nodes.emplace_back(std::move(node));
     }
+    _meta_data.nodes.push_back(_meta_data.refs.size());
 }
 
 void
@@ -35,12 +62,15 @@ 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();
+    uint32_t num_nodes = _meta_data.nodes.size() - 1;
     writer.write(&num_nodes, sizeof(uint32_t));
-    for (const auto &node : _meta_data.nodes) {
-        uint32_t num_levels = node.size();
+    for (uint32_t i(0); i < num_nodes; i++) {
+        uint32_t offset = _meta_data.nodes[i];
+        uint32_t next_offset = _meta_data.nodes[i+1];
+        uint32_t num_levels = next_offset - offset;
         writer.write(&num_levels, sizeof(uint32_t));
-        for (auto links_ref : node) {
+        for (; offset < next_offset; offset++) {
+            auto links_ref = _meta_data.refs[offset];
             if (links_ref.valid()) {
                 vespalib::ConstArrayRef<uint32_t> link_array = _graph_links.get(links_ref);
                 uint32_t num_links = link_array.size();
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h
index 387e35efd8b..9052c3f1909 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_saver.h
@@ -5,6 +5,7 @@
 #include "nearest_neighbor_index_saver.h"
 #include "hnsw_graph.h"
 #include <vespa/vespalib/datastore/entryref.h>
+#include <vespa/vespalib/stllike/allocator.h>
 #include <vector>
 
 namespace search::tensor {
@@ -17,18 +18,19 @@ namespace search::tensor {
  **/
 class HnswIndexSaver : public NearestNeighborIndexSaver {
 public:
-    using LevelVector = std::vector<vespalib::datastore::EntryRef>;
-
     HnswIndexSaver(const HnswGraph &graph);
-    ~HnswIndexSaver();
+    ~HnswIndexSaver() override;
     void save(BufferWriter& writer) const override;
 
 private:
     struct MetaData {
+        using EntryRef = vespalib::datastore::EntryRef;
         uint32_t entry_docid;
         int32_t  entry_level;
-        std::vector<LevelVector> nodes;
-        MetaData() : entry_docid(0), entry_level(-1), nodes() {}
+        std::vector<EntryRef, vespalib::allocator_large<EntryRef>> refs;
+        std::vector<uint32_t, vespalib::allocator_large<uint32_t>> nodes;
+        MetaData();
+        ~MetaData();
     };
     const HnswGraph::LinkStore &_graph_links;
     MetaData _meta_data;
diff --git a/vespalib/src/vespa/vespalib/stllike/allocator.h b/vespalib/src/vespa/vespalib/stllike/allocator.h
index 0a890973e09..70d8f435286 100644
--- a/vespalib/src/vespa/vespalib/stllike/allocator.h
+++ b/vespalib/src/vespa/vespalib/stllike/allocator.h
@@ -13,7 +13,6 @@ namespace vespalib {
  */
 template <typename T>
 class allocator_large {
-    using PtrAndSize = alloc::MemoryAllocator::PtrAndSize;
 public:
     allocator_large() noexcept : _allocator(alloc::MemoryAllocator::select_allocator()) {}
     using value_type = T;