add set_node for unit testing

3 files changed, 54 insertions, 4 deletions

diff --git a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
index 1204ae1e9bc..9a7589e1962 100644
--- a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
+++ b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
@@ -266,5 +266,33 @@ TEST_F(HnswIndexTest, 2d_vectors_inserted_in_hierarchic_graph_with_heuristic_sel
     expect_level_0(7, {3, 6});
 }
 
+TEST_F(HnswIndexTest, manual_insert)
+{
+    init(false);
+
+    std::vector<uint32_t> nbl;
+    HnswNode empty{nbl};
+    index->set_node(1, empty);
+    index->set_node(2, empty);
+
+    HnswNode three{{1,2}};
+    index->set_node(3, three);
+    expect_level_0(1, {3});
+    expect_level_0(2, {3});
+    expect_level_0(3, {1,2});
+
+    expect_entry_point(1, 0);
+
+    HnswNode twolevels{{{1},nbl}};
+    index->set_node(4, twolevels);
+
+    expect_entry_point(4, 1);
+    expect_level_0(1, {3,4});
+
+    HnswNode five{{{1,2}, {4}}};
+    index->set_node(5, five);
+    expect_levels(4, {{1}, {5}});
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
 
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 0d308206761..1c62e257d58 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -136,7 +136,7 @@ HnswIndex::select_neighbors(const HnswCandidateVector& neighbors, uint32_t max_l
 }
 
 void
-HnswIndex::connect_new_node(uint32_t docid, const LinkArray& neighbors, uint32_t level)
+HnswIndex::connect_new_node(uint32_t docid, const LinkArrayRef &neighbors, uint32_t level)
 {
     set_link_array(docid, level, neighbors);
     for (uint32_t neighbor_docid : neighbors) {
@@ -371,5 +371,28 @@ HnswIndex::get_node(uint32_t docid) const
     return HnswNode(result);
 }
 
+void
+HnswIndex::set_node(uint32_t docid, HnswNode node)
+{
+    _node_refs.ensure_size(docid + 1, AtomicEntryRef());
+    // A document cannot be added twice.
+    assert(!_node_refs[docid].load_acquire().valid());
+
+    // make new node
+    size_t num_levels = node.size();
+    assert(num_levels > 0);
+    LevelArray levels(num_levels, AtomicEntryRef());
+    auto node_ref = _nodes.add(levels);
+    _node_refs[docid].store_release(node_ref);
+
+    for (size_t level = 0; level < num_levels; ++level) {
+        connect_new_node(docid, node.level(level), level);
+    }
+    int max_level = num_levels - 1;
+    if (_entry_level < max_level) {
+        _entry_docid = docid;
+        _entry_level = max_level;
+    }
 }
 
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
index 800b88923b5..a66fdfafbc4 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
@@ -110,7 +110,7 @@ protected:
     LinkArray select_neighbors_heuristic(const HnswCandidateVector& neighbors, uint32_t max_links) const;
     LinkArray select_neighbors_simple(const HnswCandidateVector& neighbors, uint32_t max_links) const;
     LinkArray select_neighbors(const HnswCandidateVector& neighbors, uint32_t max_links) const;
-    void connect_new_node(uint32_t docid, const LinkArray& neighbors, uint32_t level);
+    void connect_new_node(uint32_t docid, const LinkArrayRef &neighbors, uint32_t level);
     void remove_link_to(uint32_t remove_from, uint32_t remove_id, uint32_t level);
 
     inline TypedCells get_vector(uint32_t docid) const {
@@ -145,8 +145,7 @@ public:
 
     // Should only be used by unit tests.
     HnswNode get_node(uint32_t docid) const;
-
-    // TODO: Implement set_node() as well for use in unit tests.
+    void set_node(uint32_t docid, HnswNode node);
 };
 
 }