diff options
9 files changed, 147 insertions, 53 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 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 2db6437664e..74b82649c98 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 @@ -103,9 +103,8 @@ public: return vector_writer.output; } void load_copy(std::vector<char> data) { - HnswIndexLoader loader(copy); - LoadedBuffer buffer(&data[0], data.size()); - loader.load(buffer); + HnswIndexLoader loader(copy, std::make_unique<LoadedBuffer>(&data[0], data.size())); + while (loader.load_next()) {} } void expect_copy_as_populated() const { diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp index 7c05699b8e1..fd86fbf1c73 100644 --- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp +++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp @@ -3,6 +3,7 @@ #include "dense_tensor_attribute.h" #include "dense_tensor_attribute_saver.h" #include "nearest_neighbor_index.h" +#include "nearest_neighbor_index_loader.h" #include "nearest_neighbor_index_saver.h" #include "tensor_attribute.hpp" #include <vespa/eval/eval/value.h> @@ -10,10 +11,11 @@ #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> #include <vespa/vespalib/util/threadstackexecutor.h> -#include <vespa/vespalib/util/lambdatask.h> #include <thread> #include <vespa/log/log.h> @@ -29,6 +31,7 @@ namespace search::tensor { namespace { constexpr uint32_t DENSE_TENSOR_ATTRIBUTE_VERSION = 1; +constexpr uint32_t LOAD_COMMIT_INTERVAL = 256; const vespalib::string tensorTypeTag("tensortype"); class BlobSequenceReader : public ReaderBase @@ -266,7 +269,7 @@ private: _attr.setCommittedDocIdLimit(std::max(_attr.getCommittedDocIdLimit(), lid + 1)); _attr._index->complete_add_document(lid, std::move(prepared)); --_pending; - if ((lid % 256) == 0) { + if ((lid % LOAD_COMMIT_INTERVAL) == 0) { _attr.commit(); }; } @@ -319,7 +322,7 @@ public: // This ensures that get_vector() (via getTensor()) is able to find the newly added tensor. _attr.setCommittedDocIdLimit(lid + 1); _attr._index->add_document(lid); - if ((lid % 256) == 0) { + if ((lid % LOAD_COMMIT_INTERVAL) == 0) { _attr.commit(); } } @@ -375,7 +378,17 @@ DenseTensorAttribute::onLoad(vespalib::Executor *executor) setCommittedDocIdLimit(numDocs); if (_index && use_index_file) { auto buffer = LoadUtils::loadFile(*this, DenseTensorAttributeSaver::index_file_suffix()); - if (!_index->load(*buffer)) { + try { + auto index_loader = _index->make_loader(std::move(buffer)); + 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", + 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 49aa64212ae..8da8c4ba01f 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp @@ -1,16 +1,17 @@ // Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +#include "bitvector_visited_tracker.h" #include "distance_function.h" +#include "hash_set_visited_tracker.h" #include "hnsw_index.h" #include "hnsw_index_loader.h" #include "hnsw_index_saver.h" #include "random_level_generator.h" -#include "bitvector_visited_tracker.h" -#include "hash_set_visited_tracker.h" #include "reusable_set_visited_tracker.h" #include <vespa/searchcommon/common/compaction_strategy.h> #include <vespa/searchlib/attribute/address_space_components.h> #include <vespa/searchlib/attribute/address_space_usage.h> +#include <vespa/searchlib/util/fileutil.h> #include <vespa/searchlib/util/state_explorer_utils.h> #include <vespa/vespalib/data/slime/cursor.h> #include <vespa/vespalib/data/slime/inserter.h> @@ -694,12 +695,11 @@ HnswIndex::make_saver() const return std::make_unique<HnswIndexSaver>(_graph); } -bool -HnswIndex::load(const fileutil::LoadedBuffer& buf) +std::unique_ptr<NearestNeighborIndexLoader> +HnswIndex::make_loader(std::unique_ptr<fileutil::LoadedBuffer> buf) { assert(get_entry_docid() == 0); // cannot load after index has data - HnswIndexLoader loader(_graph); - return loader.load(buf); + return std::make_unique<HnswIndexLoader>(_graph, std::move(buf)); } struct NeighborsByDocId { diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h index 4503459a88a..4cb7afd1a24 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; - bool load(const fileutil::LoadedBuffer& buf) override; + std::unique_ptr<NearestNeighborIndexLoader> make_loader(std::unique_ptr<fileutil::LoadedBuffer> buf) 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.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp index c0aec9ff91a..53b702a4d79 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.cpp @@ -6,45 +6,64 @@ namespace search::tensor { +void +HnswIndexLoader::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() {} -HnswIndexLoader::HnswIndexLoader(HnswGraph &graph) - : _graph(graph), _ptr(nullptr), _end(nullptr), _failed(false) + +HnswIndexLoader::HnswIndexLoader(HnswGraph& graph, std::unique_ptr<fileutil::LoadedBuffer> buf) + : _graph(graph), + _buf(std::move(buf)), + _ptr(nullptr), + _end(nullptr), + _entry_docid(0), + _entry_level(0), + _num_nodes(0), + _docid(0), + _link_array(), + _complete(false) { + init(); } bool -HnswIndexLoader::load(const fileutil::LoadedBuffer& buf) +HnswIndexLoader::load_next() { - size_t num_readable = buf.size(sizeof(uint32_t)); - _ptr = static_cast<const uint32_t *>(buf.buffer()); - _end = _ptr + num_readable; - uint32_t entry_docid = next_int(); - int32_t entry_level = next_int(); - uint32_t num_nodes = next_int(); - std::vector<uint32_t> link_array; - for (uint32_t docid = 0; docid < num_nodes; ++docid) { + assert(!_complete); + if (_docid < _num_nodes) { uint32_t num_levels = next_int(); if (num_levels > 0) { - _graph.make_node_for_document(docid, num_levels); + _graph.make_node_for_document(_docid, num_levels); for (uint32_t level = 0; level < num_levels; ++level) { uint32_t num_links = next_int(); - link_array.clear(); + _link_array.clear(); while (num_links-- > 0) { - link_array.push_back(next_int()); + _link_array.push_back(next_int()); } - _graph.set_link_array(docid, level, link_array); + _graph.set_link_array(_docid, level, _link_array); } } } - if (_failed) return false; - _graph.node_refs.ensure_size(std::max(num_nodes, 1u)); - _graph.node_refs_size.store(std::max(num_nodes, 1u), std::memory_order_release); - _graph.trim_node_refs_size(); - auto entry_node_ref = _graph.get_node_ref(entry_docid); - _graph.set_entry_node({entry_docid, entry_node_ref, entry_level}); - return true; + if (++_docid < _num_nodes) { + return true; + } else { + _graph.node_refs.ensure_size(std::max(_num_nodes, 1u)); + _graph.node_refs_size.store(std::max(_num_nodes, 1u), std::memory_order_release); + _graph.trim_node_refs_size(); + auto entry_node_ref = _graph.get_node_ref(_entry_docid); + _graph.set_entry_node({_entry_docid, entry_node_ref, _entry_level}); + _complete = true; + return false; + } } - } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h index 9f5ae66011f..0b6658e42ec 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_loader.h @@ -2,7 +2,11 @@ #pragma once +#include "nearest_neighbor_index_loader.h" +#include <vespa/vespalib/util/exceptions.h> #include <cstdint> +#include <memory> +#include <vector> namespace search::fileutil { class LoadedBuffer; } @@ -13,23 +17,34 @@ struct HnswGraph; /** * Implements loading of HNSW graph structure from binary format. **/ -class HnswIndexLoader { -public: - HnswIndexLoader(HnswGraph &graph); - ~HnswIndexLoader(); - bool load(const fileutil::LoadedBuffer& buf); +class HnswIndexLoader : public NearestNeighborIndexLoader { private: - HnswGraph &_graph; - const uint32_t *_ptr; - const uint32_t *_end; - bool _failed; + HnswGraph& _graph; + std::unique_ptr<fileutil::LoadedBuffer> _buf; + const uint32_t* _ptr; + const uint32_t* _end; + uint32_t _entry_docid; + int32_t _entry_level; + uint32_t _num_nodes; + uint32_t _docid; + std::vector<uint32_t> _link_array; + bool _complete; + + void init(); uint32_t next_int() { if (__builtin_expect((_ptr == _end), false)) { - _failed = true; - return 0; + 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++; } + +public: + HnswIndexLoader(HnswGraph& graph, std::unique_ptr<fileutil::LoadedBuffer> buf); + virtual ~HnswIndexLoader(); + bool load_next() override; }; } diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h index b8f30a53ddf..f75cdae8a92 100644 --- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h +++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h @@ -22,6 +22,7 @@ class CompactionStrategy; namespace search::tensor { +class NearestNeighborIndexLoader; class NearestNeighborIndexSaver; /** @@ -77,7 +78,13 @@ public: * and the caller ensures that an attribute read guard is held during the lifetime of the saver. */ virtual std::unique_ptr<NearestNeighborIndexSaver> make_saver() const = 0; - virtual bool load(const fileutil::LoadedBuffer& buf) = 0; + + /** + * Creates a loader that is used to load the index from the given buffer. + * + * This might throw vespalib::IoException. + */ + virtual std::unique_ptr<NearestNeighborIndexLoader> make_loader(std::unique_ptr<fileutil::LoadedBuffer> buf) = 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 new file mode 100644 index 00000000000..703f8f863d1 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index_loader.h @@ -0,0 +1,23 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +namespace search::tensor { + +/** + * Interface that is used to load a nearest neighbor index from binary form. + */ +class NearestNeighborIndexLoader { +public: + virtual ~NearestNeighborIndexLoader() {} + + /** + * 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. + */ + virtual bool load_next() = 0; +}; + +} |