Test hnsw index prepare_add_document and find_top_k with missing tensor.

1 files changed, 35 insertions, 1 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 b697effeab4..fcce2f5eb17 100644
--- a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
+++ b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
@@ -62,8 +62,17 @@ public:
         _vectors[docid] = vec;
         return *this;
     }
+    void clear(uint32_t docid) {
+        if (docid < _vectors.size()) {
+            _vectors[docid].clear();
+        }
+    }
     vespalib::eval::TypedCells get_vector(uint32_t docid, uint32_t subspace) const noexcept override {
-        return get_vectors(docid).cells(subspace);
+        auto bundle = get_vectors(docid);
+        if (subspace < bundle.subspaces()) {
+            return bundle.cells(subspace);
+        }
+        return { nullptr, _subspace_type.cell_type(), 0 };
     }
     VectorBundle get_vectors(uint32_t docid) const noexcept override {
         ArrayRef ref(_vectors[docid]);
@@ -277,6 +286,12 @@ public:
         return index->get_active_nodes();
     }
 
+    /*
+     * Simulate race where writer has cleared a tensor while read thread still
+     * use old graph.
+     */
+    void writer_clears_tensor(uint32_t docid) { vectors.clear(docid); }
+
     static constexpr bool is_single = std::is_same_v<IndexType, HnswIndex<HnswIndexType::SINGLE>>;
 };
 
@@ -827,6 +842,14 @@ TYPED_TEST(HnswIndexTest, hnsw_graph_can_be_saved_and_loaded)
     this->check_savetest_index("after load");
 }
 
+TYPED_TEST(HnswIndexTest, search_during_remove)
+{
+    this->init(false);
+    this->make_savetest_index();
+    this->writer_clears_tensor(4);
+    this->expect_top_3_by_docid("{0, 0}", {0, 0}, {7});
+}
+
 using HnswMultiIndexTest = HnswIndexTest<HnswIndex<HnswIndexType::MULTI>>;
 
 namespace {
@@ -1022,4 +1045,15 @@ TYPED_TEST(TwoPhaseTest, two_phase_add)
     this->expect_levels(nodeids[0], {{2}, {4}});
 }
 
+TYPED_TEST(TwoPhaseTest, prepare_insert_during_remove)
+{
+    this->init(false);
+    this->make_savetest_index();
+    this->writer_clears_tensor(4);
+    auto prepared = this->prepare_add(2, 1);
+    this->remove_document(4);
+    this->complete_add(2, std::move(prepared));
+    EXPECT_EQ(2, this->get_active_nodes());
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()