Change loading of nearest neighbor index to use direct I/O instead of mmapping.

This should reduce memory spike during loading.

12 files changed, 102 insertions, 82 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index 397347b7651..f694abaae8a 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -94,17 +94,15 @@ public:
 class MockIndexLoader : public NearestNeighborIndexLoader {
 private:
     int& _index_value;
-    std::unique_ptr<search::fileutil::LoadedBuffer> _buf;
+    search::FileReader<int> _reader;
 
 public:
-    MockIndexLoader(int& index_value,
-                    std::unique_ptr<search::fileutil::LoadedBuffer> buf)
+    MockIndexLoader(int& index_value, FastOS_FileInterface& file)
         : _index_value(index_value),
-          _buf(std::move(buf))
+          _reader(file)
     {}
     bool load_next() override {
-        ASSERT_EQUAL(sizeof(int), _buf->size());
-        _index_value = (reinterpret_cast<const int*>(_buf->buffer()))[0];
+        _index_value = _reader.readHostOrder();
         return false;
     }
 };
@@ -240,8 +238,8 @@ public:
         }
         return std::unique_ptr<NearestNeighborIndexSaver>();
     }
-    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::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file) override {
+        return std::make_unique<MockIndexLoader>(_index_value, file);
     }
     std::vector<Neighbor> find_top_k(uint32_t k, vespalib::eval::TypedCells vector, uint32_t explore_k,
                                      double distance_threshold) const override
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
index 74b82649c98..27504986a72 100644
--- a/searchlib/src/tests/tensor/hnsw_saver/hnsw_save_load_test.cpp
+++ b/searchlib/src/tests/tensor/hnsw_saver/hnsw_save_load_test.cpp
@@ -2,7 +2,7 @@
 
 #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/searchlib/tensor/hnsw_index_loader.hpp>
 #include <vespa/searchlib/util/bufferwriter.h>
 #include <vespa/searchlib/util/fileutil.h>
 #include <vespa/vespalib/gtest/gtest.h>
@@ -32,6 +32,22 @@ public:
     }
 };
 
+class VectorBufferReader {
+private:
+    const std::vector<char>& _data;
+    size_t _pos;
+
+public:
+    VectorBufferReader(const std::vector<char>& data) : _data(data), _pos(0) {}
+    uint32_t readHostOrder() {
+        uint32_t result = 0;
+        assert(_pos + sizeof(uint32_t) <= _data.size());
+        std::memcpy(&result, _data.data() + _pos, sizeof(uint32_t));
+        _pos += sizeof(uint32_t);
+        return result;
+    }
+};
+
 using V = std::vector<uint32_t>;
 
 void populate(HnswGraph &graph) {
@@ -103,7 +119,7 @@ public:
         return vector_writer.output;
     }
     void load_copy(std::vector<char> data) {
-        HnswIndexLoader loader(copy, std::make_unique<LoadedBuffer>(&data[0], data.size()));
+        HnswIndexLoader<VectorBufferReader> loader(copy, std::make_unique<VectorBufferReader>(data));
         while (loader.load_next()) {}
     }
 
diff --git a/searchlib/src/vespa/searchlib/attribute/load_utils.cpp b/searchlib/src/vespa/searchlib/attribute/load_utils.cpp
index 5e9bc80f46a..8f87e398035 100644
--- a/searchlib/src/vespa/searchlib/attribute/load_utils.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/load_utils.cpp
@@ -24,8 +24,6 @@ LoadUtils::openFile(const AttributeVector& attr, const vespalib::string& suffix)
     return FileUtil::openFile(attr.getBaseFileName() + "." + suffix);
 }
 
-
-
 FileInterfaceUP
 LoadUtils::openDAT(const AttributeVector& attr)
 {
diff --git a/searchlib/src/vespa/searchlib/attribute/load_utils.h b/searchlib/src/vespa/searchlib/attribute/load_utils.h
index 6833ab6b0b7..82be1844a19 100644
--- a/searchlib/src/vespa/searchlib/attribute/load_utils.h
+++ b/searchlib/src/vespa/searchlib/attribute/load_utils.h
@@ -16,10 +16,7 @@ public:
     using FileInterfaceUP = std::unique_ptr<FastOS_FileInterface>;
     using LoadedBufferUP = std::unique_ptr<fileutil::LoadedBuffer>;
 
-private:
     static FileInterfaceUP openFile(const AttributeVector& attr, const vespalib::string& suffix);
-
-public:
     static FileInterfaceUP openDAT(const AttributeVector& attr);
     static FileInterfaceUP openIDX(const AttributeVector& attr);
     static FileInterfaceUP openWeight(const AttributeVector& attr);
diff --git a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt
index e10db433a58..b5f7492a3b8 100644
--- a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt
+++ b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt
@@ -17,7 +17,6 @@ vespa_add_library(searchlib_tensor OBJECT
     hash_set_visited_tracker.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
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index 00aede95ca4..21f03e6003f 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -11,7 +11,6 @@
 #include <vespa/searchlib/attribute/load_utils.h>
 #include <vespa/searchlib/attribute/readerbase.h>
 #include <vespa/vespalib/data/slime/inserter.h>
-#include <vespa/vespalib/util/exceptions.h>
 #include <vespa/vespalib/util/lambdatask.h>
 #include <vespa/vespalib/util/memory_allocator.h>
 #include <vespa/vespalib/util/mmap_file_allocator_factory.h>
@@ -39,17 +38,50 @@ class BlobSequenceReader : public ReaderBase
 private:
     static constexpr uint8_t tensorIsNotPresent = 0;
     static constexpr uint8_t tensorIsPresent = 1;
+    bool _use_index_file;
+    FileWithHeader _index_file;
+
 public:
-    BlobSequenceReader(AttributeVector &attr);
+    BlobSequenceReader(AttributeVector& attr, bool has_index);
     ~BlobSequenceReader();
     bool is_present();
     void readTensor(void *buf, size_t len) { _datFile.file().ReadBuf(buf, len); }
+    bool use_index_file() const { return _use_index_file; }
+    FastOS_FileInterface& index_file() { return _index_file.file(); }
 };
 
-BlobSequenceReader::BlobSequenceReader(AttributeVector &attr)
-    : ReaderBase(attr)
+bool
+can_use_index_save_file(const search::attribute::Config &config, const search::attribute::AttributeHeader &header)
 {
+    if (!config.hnsw_index_params().has_value() || !header.get_hnsw_index_params().has_value()) {
+        return false;
+    }
+    const auto &config_params = config.hnsw_index_params().value();
+    const auto &header_params = header.get_hnsw_index_params().value();
+    if ((config_params.max_links_per_node() != header_params.max_links_per_node()) ||
+        (config_params.distance_metric() != header_params.distance_metric())) {
+        return false;
+    }
+    return true;
 }
+
+bool
+has_index_file(AttributeVector& attr)
+{
+    return LoadUtils::file_exists(attr, DenseTensorAttributeSaver::index_file_suffix());
+}
+
+BlobSequenceReader::BlobSequenceReader(AttributeVector& attr, bool has_index)
+    : ReaderBase(attr),
+      _use_index_file(has_index && has_index_file(attr) &&
+                      can_use_index_save_file(attr.getConfig(),
+                                              search::attribute::AttributeHeader::extractTags(getDatHeader()))),
+      _index_file(_use_index_file ?
+                  attribute::LoadUtils::openFile(attr, DenseTensorAttributeSaver::index_file_suffix()) :
+                  std::unique_ptr<Fast_BufferedFile>())
+{
+}
+
 BlobSequenceReader::~BlobSequenceReader() = default;
 
 bool
@@ -65,20 +97,7 @@ BlobSequenceReader::is_present() {
     return true;
 }
 
-bool
-can_use_index_save_file(const search::attribute::Config &config, const search::attribute::AttributeHeader &header)
-{
-    if (!config.hnsw_index_params().has_value() || !header.get_hnsw_index_params().has_value()) {
-        return false;
-    }
-    const auto &config_params = config.hnsw_index_params().value();
-    const auto &header_params = header.get_hnsw_index_params().value();
-    if ((config_params.max_links_per_node() != header_params.max_links_per_node()) ||
-        (config_params.distance_metric() != header_params.distance_metric())) {
-        return false;
-    }
-    return true;
-}
+
 
 std::unique_ptr<vespalib::alloc::MemoryAllocator>
 make_memory_allocator(const vespalib::string& name, bool swappable)
@@ -336,22 +355,19 @@ private:
 bool
 DenseTensorAttribute::onLoad(vespalib::Executor *executor)
 {
-    BlobSequenceReader tensorReader(*this);
-    if (!tensorReader.hasData()) {
+    BlobSequenceReader reader(*this, _index.get() != nullptr);
+    if (!reader.hasData()) {
         return false;
     }
-    bool has_index_file = LoadUtils::file_exists(*this, DenseTensorAttributeSaver::index_file_suffix());
-    bool use_index_file = has_index_file && _index && can_use_index_save_file(getConfig(), search::attribute::AttributeHeader::extractTags(tensorReader.getDatHeader()));
-
-    setCreateSerialNum(tensorReader.getCreateSerialNum());
-    assert(tensorReader.getVersion() == DENSE_TENSOR_ATTRIBUTE_VERSION);
+    setCreateSerialNum(reader.getCreateSerialNum());
+    assert(reader.getVersion() == DENSE_TENSOR_ATTRIBUTE_VERSION);
     assert(getConfig().tensorType().to_spec() ==
-           tensorReader.getDatHeader().getTag(tensorTypeTag).asString());
-    uint32_t numDocs(tensorReader.getDocIdLimit());
+           reader.getDatHeader().getTag(tensorTypeTag).asString());
+    uint32_t numDocs(reader.getDocIdLimit());
     _refVector.reset();
     _refVector.unsafe_reserve(numDocs);
     std::unique_ptr<Loader> loader;
-    if (_index && !use_index_file) {
+    if (_index && !reader.use_index_file()) {
         if (executor != nullptr) {
             loader = std::make_unique<ThreadedLoader>(*this, *executor);
         } else {
@@ -359,9 +375,9 @@ DenseTensorAttribute::onLoad(vespalib::Executor *executor)
         }
     }
     for (uint32_t lid = 0; lid < numDocs; ++lid) {
-        if (tensorReader.is_present()) {
+        if (reader.is_present()) {
             auto raw = _denseTensorStore.allocRawBuffer();
-            tensorReader.readTensor(raw.data, _denseTensorStore.getBufSize());
+            reader.readTensor(raw.data, _denseTensorStore.getBufSize());
             _refVector.push_back(raw.ref);
             if (loader) {
                 loader->load(lid, raw.ref);
@@ -376,18 +392,17 @@ DenseTensorAttribute::onLoad(vespalib::Executor *executor)
     commit();
     setNumDocs(numDocs);
     setCommittedDocIdLimit(numDocs);
-    if (_index && use_index_file) {
-        auto buffer = LoadUtils::loadFile(*this, DenseTensorAttributeSaver::index_file_suffix());
+    if (_index && reader.use_index_file()) {
         try {
-            auto index_loader = _index->make_loader(std::move(buffer));
+            auto index_loader = _index->make_loader(reader.index_file());
             size_t cnt = 0;
             while (index_loader->load_next()) {
                 if ((++cnt % LOAD_COMMIT_INTERVAL) == 0) {
                     commit();
                 }
             }
-        } catch (const vespalib::IoException& ex) {
-            LOG(error, "IoException while loading nearest neighbor index for tensor attribute '%s': %s",
+        } catch (const std::runtime_error& ex) {
+            LOG(error, "Exception while loading nearest neighbor index for tensor attribute '%s': %s",
                 getName().c_str(), ex.what());
             return false;
         }
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 8da8c4ba01f..7e0efb78ce6 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -4,7 +4,7 @@
 #include "distance_function.h"
 #include "hash_set_visited_tracker.h"
 #include "hnsw_index.h"
-#include "hnsw_index_loader.h"
+#include "hnsw_index_loader.hpp"
 #include "hnsw_index_saver.h"
 #include "random_level_generator.h"
 #include "reusable_set_visited_tracker.h"
@@ -696,10 +696,12 @@ HnswIndex::make_saver() const
 }
 
 std::unique_ptr<NearestNeighborIndexLoader>
-HnswIndex::make_loader(std::unique_ptr<fileutil::LoadedBuffer> buf)
+HnswIndex::make_loader(FastOS_FileInterface& file)
 {
     assert(get_entry_docid() == 0); // cannot load after index has data
-    return std::make_unique<HnswIndexLoader>(_graph, std::move(buf));
+    using ReaderType = FileReader<uint32_t>;
+    using LoaderType = HnswIndexLoader<ReaderType>;
+    return std::make_unique<LoaderType>(_graph, std::make_unique<ReaderType>(file));
 }
 
 struct NeighborsByDocId {
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
index 4cb7afd1a24..dde4d377c4c 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
@@ -183,7 +183,7 @@ public:
     void shrink_lid_space(uint32_t doc_id_limit) override;
 
     std::unique_ptr<NearestNeighborIndexSaver> make_saver() const override;
-    std::unique_ptr<NearestNeighborIndexLoader> make_loader(std::unique_ptr<fileutil::LoadedBuffer> buf) override;
+    std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file) override;
 
     std::vector<Neighbor> find_top_k(uint32_t k, TypedCells vector, uint32_t explore_k,
                                      double distance_threshold) const override;
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h
index 0b6658e42ec..80910fb2afe 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h
@@ -8,7 +8,7 @@
 #include <memory>
 #include <vector>
 
-namespace search::fileutil { class LoadedBuffer; }
+class FastOS_FileInterface;
 
 namespace search::tensor {
 
@@ -17,12 +17,11 @@ struct HnswGraph;
 /**
  * Implements loading of HNSW graph structure from binary format.
  **/
+template <typename ReaderType>
 class HnswIndexLoader : public NearestNeighborIndexLoader {
 private:
     HnswGraph& _graph;
-    std::unique_ptr<fileutil::LoadedBuffer> _buf;
-    const uint32_t* _ptr;
-    const uint32_t* _end;
+    std::unique_ptr<ReaderType> _reader;
     uint32_t _entry_docid;
     int32_t _entry_level;
     uint32_t _num_nodes;
@@ -32,17 +31,11 @@ private:
 
     void init();
     uint32_t next_int() {
-        if (__builtin_expect((_ptr == _end), false)) {
-            throw vespalib::IoException
-                    (vespalib::IoException::createMessage("Already at the end of buffer when trying to get next int",
-                                                          vespalib::IoException::CORRUPT_DATA),
-                     vespalib::IoException::CORRUPT_DATA, "");
-        }
-        return *_ptr++;
+        return _reader->readHostOrder();
     }
 
 public:
-    HnswIndexLoader(HnswGraph& graph, std::unique_ptr<fileutil::LoadedBuffer> buf);
+    HnswIndexLoader(HnswGraph& graph, std::unique_ptr<ReaderType> reader);
     virtual ~HnswIndexLoader();
     bool load_next() override;
 };
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.hpp
index 53b702a4d79..a1ec5d591f4 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.hpp
@@ -1,30 +1,29 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
+#pragma once
+
 #include "hnsw_index_loader.h"
 #include "hnsw_graph.h"
 #include <vespa/searchlib/util/fileutil.h>
 
 namespace search::tensor {
 
+template <typename ReaderType>
 void
-HnswIndexLoader::init()
+HnswIndexLoader<ReaderType>::init()
 {
-    size_t num_readable = _buf->size(sizeof(uint32_t));
-    _ptr = static_cast<const uint32_t *>(_buf->buffer());
-    _end = _ptr + num_readable;
     _entry_docid = next_int();
     _entry_level = next_int();
     _num_nodes = next_int();
 }
 
-HnswIndexLoader::~HnswIndexLoader() {}
-
+template <typename ReaderType>
+HnswIndexLoader<ReaderType>::~HnswIndexLoader() {}
 
-HnswIndexLoader::HnswIndexLoader(HnswGraph& graph, std::unique_ptr<fileutil::LoadedBuffer> buf)
+template <typename ReaderType>
+HnswIndexLoader<ReaderType>::HnswIndexLoader(HnswGraph& graph, std::unique_ptr<ReaderType> reader)
     : _graph(graph),
-      _buf(std::move(buf)),
-      _ptr(nullptr),
-      _end(nullptr),
+      _reader(std::move(reader)),
       _entry_docid(0),
       _entry_level(0),
       _num_nodes(0),
@@ -35,8 +34,9 @@ HnswIndexLoader::HnswIndexLoader(HnswGraph& graph, std::unique_ptr<fileutil::Loa
     init();
 }
 
+template <typename ReaderType>
 bool
-HnswIndexLoader::load_next()
+HnswIndexLoader<ReaderType>::load_next()
 {
     assert(!_complete);
     if (_docid < _num_nodes) {
diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
index f75cdae8a92..dbb439913f5 100644
--- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
@@ -10,6 +10,8 @@
 #include <memory>
 #include <vector>
 
+class FastOS_FileInterface;
+
 namespace vespalib::slime { struct Inserter; }
 
 namespace search::fileutil { class LoadedBuffer; }
@@ -80,11 +82,11 @@ public:
     virtual std::unique_ptr<NearestNeighborIndexSaver> make_saver() const = 0;
 
     /**
-     * Creates a loader that is used to load the index from the given buffer.
+     * Creates a loader that is used to load the index from the given file.
      *
-     * This might throw vespalib::IoException.
+     * This might throw std::runtime_error.
      */
-    virtual std::unique_ptr<NearestNeighborIndexLoader> make_loader(std::unique_ptr<fileutil::LoadedBuffer> buf) = 0;
+    virtual std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file) = 0;
 
     virtual std::vector<Neighbor> find_top_k(uint32_t k,
                                              vespalib::eval::TypedCells vector,
diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_loader.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_loader.h
index 703f8f863d1..41864429360 100644
--- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_loader.h
+++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_loader.h
@@ -15,7 +15,7 @@ public:
      * Loads the next part of the index (e.g. the node corresponding to a given document)
      * and returns whether there is more data to load.
      *
-     * This might throw vespalib::IoException.
+     * This might throw std::runtime_error.
      */
     virtual bool load_next() = 0;
 };