diff options
author | Tor Egge <Tor.Egge@online.no> | 2023-05-23 11:32:21 +0200 |
---|---|---|
committer | Tor Egge <Tor.Egge@online.no> | 2023-05-23 11:32:21 +0200 |
commit | a3f5a86d596c10b57f7a404e9a17fad56c50f20c (patch) | |
tree | 251b6b8e8a00e723d66bb8af5d72fe3666c913a1 | |
parent | 14583ad0719a537ef8a400262052ddf1966d6329 (diff) |
Store max squared norm in file header during hnsw index save when using
dotproduct distance metric. Adjust hnsw index load.
9 files changed, 103 insertions, 21 deletions
diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp index e3c9e05073e..841d7f92b62 100644 --- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp +++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp @@ -9,6 +9,7 @@ #include <vespa/searchlib/tensor/doc_vector_access.h> #include <vespa/searchlib/tensor/distance_functions.h> #include <vespa/searchlib/tensor/hnsw_index.h> +#include <vespa/searchlib/tensor/mips_distance_transform.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> @@ -54,6 +55,7 @@ using search::tensor::DocVectorAccess; using search::tensor::HnswIndex; using search::tensor::HnswIndexType; using search::tensor::HnswTestNode; +using search::tensor::MipsDistanceFunctionFactoryBase; using search::tensor::NearestNeighborIndex; using search::tensor::NearestNeighborIndexFactory; using search::tensor::NearestNeighborIndexLoader; @@ -285,13 +287,15 @@ public: void populate_address_space_usage(AddressSpaceUsage&) const override {} void get_state(const vespalib::slime::Inserter&) const override {} void shrink_lid_space(uint32_t) override { } - std::unique_ptr<NearestNeighborIndexSaver> make_saver() const override { + std::unique_ptr<NearestNeighborIndexSaver> make_saver(vespalib::GenericHeader& header) const override { + (void) header; if (_index_value != 0) { return std::make_unique<MockIndexSaver>(_index_value); } return std::unique_ptr<NearestNeighborIndexSaver>(); } - std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file) override { + std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file, const vespalib::GenericHeader& header) override { + (void) header; return std::make_unique<MockIndexLoader>(_index_value, file); } std::vector<Neighbor> find_top_k(uint32_t k, @@ -342,12 +346,15 @@ class MockNearestNeighborIndexFactory : public NearestNeighborIndexFactory { const vespalib::string test_dir = "test_data/"; const vespalib::string attr_name = test_dir + "my_attr"; +const vespalib::string hnsw_max_squared_norm = "hnsw.max_squared_norm"; + struct FixtureTraits { bool use_dense_tensor_attribute = false; bool use_direct_tensor_attribute = false; bool enable_hnsw_index = false; bool use_mock_index = false; bool use_mmap_file_allocator = false; + bool use_mips_distance = false; FixtureTraits dense() && { use_dense_tensor_attribute = true; @@ -381,6 +388,14 @@ struct FixtureTraits { return *this; } + FixtureTraits mips_hnsw() && { + use_dense_tensor_attribute = true; + enable_hnsw_index = true; + use_mock_index = false; + use_mips_distance = true; + return *this; + } + FixtureTraits direct() && { use_dense_tensor_attribute = false; use_direct_tensor_attribute = true; @@ -606,8 +621,9 @@ Fixture::Fixture(const vespalib::string &typeSpec, FixtureTraits traits) _mmap_allocator_base_dir("mmap-file-allocator-factory-dir") { if (traits.enable_hnsw_index) { - _cfg.set_distance_metric(DistanceMetric::Euclidean); - _cfg.set_hnsw_index_params(HnswIndexParams(4, 20, DistanceMetric::Euclidean)); + auto dm = traits.use_mips_distance ? DistanceMetric::Dotproduct : DistanceMetric::Euclidean; + _cfg.set_distance_metric(dm); + _cfg.set_hnsw_index_params(HnswIndexParams(4, 20, dm)); } vespalib::alloc::MmapFileAllocatorFactory::instance().setup(_mmap_allocator_base_dir); setup(); @@ -1254,6 +1270,23 @@ TEST_F("Nearest neighbor index type is added to attribute file header", DenseTen EXPECT_EQUAL("hnsw", header.getTag("nearest_neighbor_index").asString()); } +class DenseTensorAttributeMipsIndex : public Fixture { +public: + DenseTensorAttributeMipsIndex() : Fixture(vec_2d_spec, FixtureTraits().mips_hnsw()) {} +}; + +TEST_F("Nearest neighbor index with mips distance metrics stores square of max distance", DenseTensorAttributeMipsIndex) +{ + f.set_example_tensors(); + f.save(); + auto header = f.get_file_header(); + EXPECT_TRUE(header.hasTag(hnsw_max_squared_norm)); + EXPECT_EQUAL(130.0, header.getTag(hnsw_max_squared_norm).asFloat()); + f.load(); + auto& norm_store = dynamic_cast<MipsDistanceFunctionFactoryBase&>(f.hnsw_index().distance_function_factory()).get_max_squared_norm_store(); + EXPECT_EQUAL(130.0, norm_store.get_max()); +} + template <typename ParentT> class NearestNeighborBlueprintFixtureBase : public ParentT { private: diff --git a/searchlib/src/vespa/searchlib/attribute/attribute_header.cpp b/searchlib/src/vespa/searchlib/attribute/attribute_header.cpp index 122c2c0c55e..cde1686828f 100644 --- a/searchlib/src/vespa/searchlib/attribute/attribute_header.cpp +++ b/searchlib/src/vespa/searchlib/attribute/attribute_header.cpp @@ -5,6 +5,8 @@ #include <vespa/vespalib/data/databuffer.h> #include <vespa/vespalib/util/exceptions.h> +using vespalib::GenericHeader; + namespace search::attribute { namespace { @@ -57,7 +59,8 @@ AttributeHeader::AttributeHeader(const vespalib::string &fileName) _uniqueValueCount(0), _totalValueCount(0), _createSerialNum(0u), - _version(0) + _version(0), + _extra_tags() { } @@ -244,6 +247,10 @@ AttributeHeader::addTags(vespalib::GenericHeader &header) const header.putTag(Tag(predicateLowerBoundTag, params.lower_bound())); header.putTag(Tag(predicateUpperBoundTag, params.upper_bound())); } + for (uint32_t i = 0; i < _extra_tags.getNumTags(); ++i) { + auto& tag = _extra_tags.getTag(i); + header.putTag(tag); + } } bool diff --git a/searchlib/src/vespa/searchlib/attribute/attribute_header.h b/searchlib/src/vespa/searchlib/attribute/attribute_header.h index 7c0b8f3084b..8c5a0edc6a6 100644 --- a/searchlib/src/vespa/searchlib/attribute/attribute_header.h +++ b/searchlib/src/vespa/searchlib/attribute/attribute_header.h @@ -2,16 +2,15 @@ #pragma once -#include <vespa/vespalib/stllike/string.h> +#include <vespa/eval/eval/value_type.h> #include <vespa/searchcommon/attribute/basictype.h> #include <vespa/searchcommon/attribute/collectiontype.h> #include <vespa/searchcommon/attribute/hnsw_index_params.h> #include <vespa/searchcommon/attribute/predicate_params.h> -#include <vespa/eval/eval/value_type.h> +#include <vespa/vespalib/data/fileheader.h> +#include <vespa/vespalib/stllike/string.h> #include <optional> -namespace vespalib { class GenericHeader; } - namespace search::attribute { /** @@ -34,6 +33,7 @@ private: uint64_t _totalValueCount; uint64_t _createSerialNum; uint32_t _version; + vespalib::GenericHeader _extra_tags; void internalExtractTags(const vespalib::GenericHeader &header); public: @@ -71,6 +71,7 @@ public: const std::optional<HnswIndexParams>& get_hnsw_index_params() const { return _hnsw_index_params; } static AttributeHeader extractTags(const vespalib::GenericHeader &header, const vespalib::string &file_name); void addTags(vespalib::GenericHeader &header) const; + vespalib::GenericHeader& get_extra_tags() noexcept { return _extra_tags; } }; } diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp index 3fdad3d507b..3e0cd71be8b 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp @@ -6,6 +6,7 @@ #include "hash_set_visited_tracker.h" #include "hnsw_index_loader.hpp" #include "hnsw_index_saver.h" +#include "mips_distance_transform.h" #include "random_level_generator.h" #include "vector_bundle.h" #include <vespa/searchlib/attribute/address_space_components.h> @@ -31,6 +32,7 @@ using search::StateExplorerUtils; using search::queryeval::GlobalFilter; using vespalib::datastore::CompactionStrategy; using vespalib::datastore::EntryRef; +using vespalib::GenericHeader; namespace { @@ -41,6 +43,29 @@ constexpr size_t max_level_array_size = 16; constexpr size_t max_link_array_size = 193; constexpr vespalib::duration MAX_COUNT_DURATION(100ms); +const vespalib::string hnsw_max_squared_norm = "hnsw.max_squared_norm"; + +void save_mips_max_distance(GenericHeader& header, DistanceFunctionFactory& dff) { + auto* mips_dff = dynamic_cast<MipsDistanceFunctionFactoryBase*>(&dff); + if (mips_dff != nullptr) { + auto& norm_store = mips_dff->get_max_squared_norm_store(); + header.putTag(GenericHeader::Tag(hnsw_max_squared_norm, norm_store.get_max())); + } +} + +void load_mips_max_distance(const GenericHeader& header, DistanceFunctionFactory& dff) { + auto* mips_dff = dynamic_cast<MipsDistanceFunctionFactoryBase*>(&dff); + if (mips_dff != nullptr) { + auto& norm_store = mips_dff->get_max_squared_norm_store(); + if (header.hasTag(hnsw_max_squared_norm)) { + auto& tag = header.getTag(hnsw_max_squared_norm); + if (tag.getType() == GenericHeader::Tag::Type::TYPE_FLOAT) { + (void) norm_store.get_max(tag.asFloat()); + } + } + } +} + bool has_link_to(vespalib::ConstArrayRef<uint32_t> links, uint32_t id) { for (uint32_t link : links) { if (link == id) return true; @@ -836,16 +861,18 @@ HnswIndex<type>::shrink_lid_space(uint32_t doc_id_limit) template <HnswIndexType type> std::unique_ptr<NearestNeighborIndexSaver> -HnswIndex<type>::make_saver() const +HnswIndex<type>::make_saver(GenericHeader& header) const { + save_mips_max_distance(header, distance_function_factory()); return std::make_unique<HnswIndexSaver<type>>(_graph); } template <HnswIndexType type> std::unique_ptr<NearestNeighborIndexLoader> -HnswIndex<type>::make_loader(FastOS_FileInterface& file) +HnswIndex<type>::make_loader(FastOS_FileInterface& file, const vespalib::GenericHeader& header) { assert(get_entry_nodeid() == 0); // cannot load after index has data + load_mips_max_distance(header, distance_function_factory()); using ReaderType = FileReader<uint32_t>; using LoaderType = HnswIndexLoader<ReaderType, type>; return std::make_unique<LoaderType>(_graph, _id_mapping, std::make_unique<ReaderType>(&file)); diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h index 0809dcf4fe3..7858cb65bf9 100644 --- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h +++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h @@ -211,8 +211,8 @@ public: void get_state(const vespalib::slime::Inserter& inserter) const override; void shrink_lid_space(uint32_t doc_id_limit) override; - std::unique_ptr<NearestNeighborIndexSaver> make_saver() const override; - std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file) override; + std::unique_ptr<NearestNeighborIndexSaver> make_saver(vespalib::GenericHeader& header) const override; + std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file, const vespalib::GenericHeader& header) override; std::vector<Neighbor> find_top_k( uint32_t k, diff --git a/searchlib/src/vespa/searchlib/tensor/mips_distance_transform.h b/searchlib/src/vespa/searchlib/tensor/mips_distance_transform.h index fabd6bfcc57..833fa3e689b 100644 --- a/searchlib/src/vespa/searchlib/tensor/mips_distance_transform.h +++ b/searchlib/src/vespa/searchlib/tensor/mips_distance_transform.h @@ -37,6 +37,18 @@ public: } }; +class MipsDistanceFunctionFactoryBase : public DistanceFunctionFactory { +protected: + std::shared_ptr<MaximumSquaredNormStore> _sq_norm_store; +public: + MipsDistanceFunctionFactoryBase() + : _sq_norm_store(std::make_shared<MaximumSquaredNormStore>()) + { + } + ~MipsDistanceFunctionFactoryBase() = default; + MaximumSquaredNormStore& get_max_squared_norm_store() noexcept { return *_sq_norm_store; } +}; + /** * Factory for distance functions which can apply a transformation * mapping Maximum Inner Product Search to a nearest neighbor @@ -45,10 +57,10 @@ public: * to the longest vector inserted so far, or at least length 1. */ template<typename FloatType> -class MipsDistanceFunctionFactory : public DistanceFunctionFactory { - std::shared_ptr<MaximumSquaredNormStore> _sq_norm_store; +class MipsDistanceFunctionFactory : public MipsDistanceFunctionFactoryBase { public: - MipsDistanceFunctionFactory() : _sq_norm_store(std::make_shared<MaximumSquaredNormStore>()) {} + MipsDistanceFunctionFactory() : MipsDistanceFunctionFactoryBase() { } + ~MipsDistanceFunctionFactory() = default; BoundDistanceFunction::UP for_query_vector(const vespalib::eval::TypedCells& lhs) override; diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h index 4b7b934fee0..9cd1065a356 100644 --- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h +++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h @@ -14,6 +14,7 @@ class FastOS_FileInterface; +namespace vespalib { class GenericHeader; } namespace vespalib::datastore { class CompactionSpec; class CompactionStrategy; @@ -88,14 +89,14 @@ public: * This function is always called by the attribute write thread, * 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 std::unique_ptr<NearestNeighborIndexSaver> make_saver(vespalib::GenericHeader& header) const = 0; /** * Creates a loader that is used to load the index from the given file. * * This might throw std::runtime_error. */ - virtual std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file) = 0; + virtual std::unique_ptr<NearestNeighborIndexLoader> make_loader(FastOS_FileInterface& file, const vespalib::GenericHeader& header) = 0; virtual std::vector<Neighbor> find_top_k(uint32_t k, const BoundDistanceFunction &df, diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp index 5e554f76779..f499695a584 100644 --- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp +++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp @@ -357,10 +357,11 @@ TensorAttribute::onInitSave(vespalib::stringref fileName) { vespalib::GenerationHandler::Guard guard(getGenerationHandler(). takeGuard()); - auto index_saver = (_index ? _index->make_saver() : std::unique_ptr<NearestNeighborIndexSaver>()); + auto header = this->createAttributeHeader(fileName); + auto index_saver = (_index ? _index->make_saver(header.get_extra_tags()) : std::unique_ptr<NearestNeighborIndexSaver>()); return std::make_unique<TensorAttributeSaver> (std::move(guard), - this->createAttributeHeader(fileName), + std::move(header), attribute::make_entry_ref_vector_snapshot(_refVector, getCommittedDocIdLimit()), _tensorStore, std::move(index_saver)); diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute_loader.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute_loader.cpp index aada583627b..2ea28fd822d 100644 --- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute_loader.cpp +++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute_loader.cpp @@ -273,7 +273,7 @@ TensorAttributeLoader::load_index() { FileWithHeader index_file(LoadUtils::openFile(_attr, TensorAttributeSaver::index_file_suffix())); try { - auto index_loader = _index->make_loader(index_file.file()); + auto index_loader = _index->make_loader(index_file.file(), index_file.header()); size_t cnt = 0; while (index_loader->load_next()) { if ((++cnt % LOAD_COMMIT_INTERVAL) == 0) { |