fix: Don't crash on reaching hnsw capacity (#2046)

* fix: Don't crash on reaching hnsw capacity Signed-off-by: Vladislav Oleshko <[email protected]> --------- Signed-off-by: Vladislav Oleshko <[email protected]>
dragonflydb · Oct 20, 2023 · 4a2dd30 · 4a2dd30
1 parent 1d02e12
commit 4a2dd30
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/src/core/search/indices.cc b/src/core/search/indices.cc
@@ -163,6 +163,8 @@ struct HnswlibAdapter {
   }
 
   void Add(float* data, DocId id) {
+    if (world_.cur_element_count + 1 >= world_.max_elements_)
+      world_.resizeIndex(world_.cur_element_count * 2);
     world_.addPoint(data, id);
   }
 

diff --git a/src/core/search/search_test.cc b/src/core/search/search_test.cc
@@ -500,6 +500,23 @@ TEST_P(KnnTest, Simple2D) {
   }
 }
 
+TEST_P(KnnTest, AutoResize) {
+  // Make sure index resizes automatically even with a small initial capacity
+  const size_t kInitialCapacity = 5;
+
+  auto schema = MakeSimpleSchema({{"pos", SchemaField::VECTOR}});
+  schema.fields["pos"].special_params =
+      SchemaField::VectorParams{GetParam(), 1, VectorSimilarity::L2, kInitialCapacity};
+  FieldIndices indices{schema, PMR_NS::get_default_resource()};
+
+  for (size_t i = 0; i < 100; i++) {
+    MockedDocument doc{Map{{"pos", ToBytes({float(i)})}}};
+    indices.Add(i, &doc);
+  }
+
+  EXPECT_EQ(indices.GetAllDocs().size(), 100);
+}
+
 INSTANTIATE_TEST_SUITE_P(KnnFlat, KnnTest, testing::Values(false));
 INSTANTIATE_TEST_SUITE_P(KnnHnsw, KnnTest, testing::Values(true));