Report address space usage for hnsw nodeid mapping.

5 files changed, 30 insertions, 8 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index d00342d0e51..2f51459ebfa 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -887,6 +887,7 @@ public:
     }
     void test_setup();
     void test_save_load(bool multi_node);
+    void test_address_space_usage();
 };
 
 template <HnswIndexType type>
@@ -935,6 +936,23 @@ TensorAttributeHnswIndex<type>::test_save_load(bool multi_node)
     expect_level_0(1, index_b.get_node(2));
 }
 
+template <HnswIndexType type>
+void
+TensorAttributeHnswIndex<type>::test_address_space_usage()
+{
+    bool dense = type == HnswIndexType::SINGLE;
+    search::AddressSpaceUsage usage = _attr->getAddressSpaceUsage();
+    const auto& all = usage.get_all();
+    EXPECT_EQUAL(dense ? 3u : 5u, all.size());
+    EXPECT_EQUAL(1u, all.count("tensor-store"));
+    EXPECT_EQUAL(1u, all.count("hnsw-levels-store"));
+    EXPECT_EQUAL(1u, all.count("hnsw-links-store"));
+    if (!dense) {
+        EXPECT_EQUAL(1u, all.count("hnsw-nodeid-mapping"));
+        EXPECT_EQUAL(1u, all.count("shared-string-repo"));
+    }
+}
+
 class DenseTensorAttributeHnswIndex : public TensorAttributeHnswIndex<HnswIndexType::SINGLE> {
 public:
     DenseTensorAttributeHnswIndex() : TensorAttributeHnswIndex<HnswIndexType::SINGLE>(vec_2d_spec, FixtureTraits().hnsw()) {}
@@ -970,16 +988,15 @@ TEST_F("Hnsw index is integrated in mixed tensor attribute and can be saved and
     f.test_save_load(true);
 }
 
-TEST_F("Populates address space usage", DenseTensorAttributeHnswIndex)
+TEST_F("Populates address space usage in dense tensor attribute with hnsw index", DenseTensorAttributeHnswIndex)
 {
-    search::AddressSpaceUsage usage = f._attr->getAddressSpaceUsage();
-    const auto& all = usage.get_all();
-    EXPECT_EQUAL(3u, all.size());
-    EXPECT_EQUAL(1u, all.count("tensor-store"));
-    EXPECT_EQUAL(1u, all.count("hnsw-levels-store"));
-    EXPECT_EQUAL(1u, all.count("hnsw-links-store"));
+    f.test_address_space_usage();
 }
 
+TEST_F("Populates address space usage in mixed tensor attribute with hnsw index", MixedTensorAttributeHnswIndex)
+{
+    f.test_address_space_usage();
+}
 
 class DenseTensorAttributeMockIndex : public Fixture {
 public:
diff --git a/searchlib/src/vespa/searchlib/attribute/address_space_components.cpp b/searchlib/src/vespa/searchlib/attribute/address_space_components.cpp
index 244e01a3874..75c198f71e2 100644
--- a/searchlib/src/vespa/searchlib/attribute/address_space_components.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/address_space_components.cpp
@@ -21,5 +21,6 @@ const vespalib::string AddressSpaceComponents::tensor_store = "tensor-store";
 const vespalib::string AddressSpaceComponents::shared_string_repo = "shared-string-repo";
 const vespalib::string AddressSpaceComponents::hnsw_levels_store = "hnsw-levels-store";
 const vespalib::string AddressSpaceComponents::hnsw_links_store = "hnsw-links-store";
+const vespalib::string AddressSpaceComponents::hnsw_nodeid_mapping = "hnsw-nodeid-mapping";
 
 }
diff --git a/searchlib/src/vespa/searchlib/attribute/address_space_components.h b/searchlib/src/vespa/searchlib/attribute/address_space_components.h
index a8a41a024f8..20302a51f0d 100644
--- a/searchlib/src/vespa/searchlib/attribute/address_space_components.h
+++ b/searchlib/src/vespa/searchlib/attribute/address_space_components.h
@@ -20,6 +20,7 @@ public:
     static const vespalib::string shared_string_repo;
     static const vespalib::string hnsw_levels_store;
     static const vespalib::string hnsw_links_store;
+    static const vespalib::string hnsw_nodeid_mapping;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 5483a0d2de4..1529611753b 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -706,6 +706,9 @@ HnswIndex<type>::populate_address_space_usage(search::AddressSpaceUsage& usage)
 {
     usage.set(AddressSpaceComponents::hnsw_levels_store, _graph.levels_store.addressSpaceUsage());
     usage.set(AddressSpaceComponents::hnsw_links_store, _graph.links_store.addressSpaceUsage());
+    if constexpr (type == HnswIndexType::MULTI) {
+        usage.set(AddressSpaceComponents::hnsw_nodeid_mapping, _id_mapping.address_space_usage());
+    }
 }
 
 template <HnswIndexType type>
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_nodeid_mapping.h b/searchlib/src/vespa/searchlib/tensor/hnsw_nodeid_mapping.h
index 8adf621b9b3..b544eb6f09f 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_nodeid_mapping.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_nodeid_mapping.h
@@ -55,7 +55,7 @@ public:
     void assign_generation(generation_t current_gen);
     void reclaim_memory(generation_t oldest_used_gen);
     void on_load(vespalib::ConstArrayRef<HnswNode> nodes);
-    // TODO: Add support for compaction
+    vespalib::AddressSpace address_space_usage() const { return _nodeids.addressSpaceUsage(); }
     vespalib::MemoryUsage memory_usage() const;
     vespalib::MemoryUsage update_stat(const vespalib::datastore::CompactionStrategy& compaction_strategy);
     bool consider_compact() const noexcept { return _nodeids.consider_compact(); }