Reduce memory spike during loading of hnsw index by committing at regular intervals.

This ensures that memory buffers on hold lists are trimmed while loading,
keeping the excess memory usage at a minimum.

1 files changed, 22 insertions, 4 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index c17bdf06854..397347b7651 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -17,6 +17,7 @@
 #include <vespa/searchlib/tensor/hnsw_index.h>
 #include <vespa/searchlib/tensor/nearest_neighbor_index.h>
 #include <vespa/searchlib/tensor/nearest_neighbor_index_factory.h>
+#include <vespa/searchlib/tensor/nearest_neighbor_index_loader.h>
 #include <vespa/searchlib/tensor/nearest_neighbor_index_saver.h>
 #include <vespa/searchlib/tensor/serialized_fast_value_attribute.h>
 #include <vespa/searchlib/tensor/tensor_attribute.h>
@@ -51,6 +52,7 @@ using search::tensor::HnswIndex;
 using search::tensor::HnswNode;
 using search::tensor::NearestNeighborIndex;
 using search::tensor::NearestNeighborIndexFactory;
+using search::tensor::NearestNeighborIndexLoader;
 using search::tensor::NearestNeighborIndexSaver;
 using search::tensor::PrepareResult;
 using search::tensor::TensorAttribute;
@@ -89,6 +91,24 @@ public:
     }
 };
 
+class MockIndexLoader : public NearestNeighborIndexLoader {
+private:
+    int& _index_value;
+    std::unique_ptr<search::fileutil::LoadedBuffer> _buf;
+
+public:
+    MockIndexLoader(int& index_value,
+                    std::unique_ptr<search::fileutil::LoadedBuffer> buf)
+        : _index_value(index_value),
+          _buf(std::move(buf))
+    {}
+    bool load_next() override {
+        ASSERT_EQUAL(sizeof(int), _buf->size());
+        _index_value = (reinterpret_cast<const int*>(_buf->buffer()))[0];
+        return false;
+    }
+};
+
 class MockPrepareResult : public PrepareResult {
 public:
     uint32_t docid;
@@ -220,10 +240,8 @@ public:
         }
         return std::unique_ptr<NearestNeighborIndexSaver>();
     }
-    bool load(const search::fileutil::LoadedBuffer& buf) override {
-        ASSERT_EQUAL(sizeof(int), buf.size());
-        _index_value = (reinterpret_cast<const int*>(buf.buffer()))[0];
-        return true;
+    std::unique_ptr<NearestNeighborIndexLoader> make_loader(std::unique_ptr<search::fileutil::LoadedBuffer> buf) override {
+        return std::make_unique<MockIndexLoader>(_index_value, std::move(buf));
     }
     std::vector<Neighbor> find_top_k(uint32_t k, vespalib::eval::TypedCells vector, uint32_t explore_k,
                                      double distance_threshold) const override