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) {