Basic HNSW class (phase 2 - support for persistence and loading) [part 1].

https://perconadev.atlassian.net/browse/PS-11267

Added persistence and on-demand loading APIs to HNSW implementation.

Extended unit tests.

(cherry picked from commit 6bec2b8,
dlenev/vector-mvp-11267, PR percona#6131)

Carried here because the aux-storage write path has nothing to hook into
without it: the class on our base takes a single template parameter and
its insert() has no context argument. This commit adds the second
parameter and the four callbacks the design targets. It applies cleanly
because it sits directly on 65261fc, the same phase-1 commit our
base builds on.

To be dropped from this branch once PS-11267 lands in vector-mvp.

gunit hnsw-t 34/34, percona vector 17/17.

Author: satya-bodapati

unittest/gunit/hnsw-t.cc

@@ -161,7 +161,7 @@ TEST_F(HnswTest, BruteForceRecall) {
   const auto points =
       make_pseudo_random_points(kNumPoints, kRecallDims, &state);
   for (size_t i = 0; i < kNumPoints; ++i) {
-    index.insert(i, /*base_pk=*/i, as_bytes(points[i]));
+    index.insert(i + 1, /*base_pk=*/i, as_bytes(points[i]));
   }
 
   const auto queries =
@@ -257,7 +257,7 @@ TEST_F(HnswTest, StreamMultipleBatchesNoDuplicates) {
   TestHnsw index(kDims, euclidean, kM, kEfConstruction);
   for (uint64_t i = 0; i < 30; ++i) {
     const auto v = make_vec({static_cast<float>(i), 0.0f});
-    index.insert(i, 1000 + i, as_bytes(v));
+    index.insert(i + 1, 1000 + i, as_bytes(v));
   }
   const auto query = make_vec({0.0f, 0.0f});
   // Small batch/ef forces continuation refills across batches.
@@ -275,7 +275,7 @@ TEST_F(HnswTest, StreamDistancesNonDecreasingAcrossBatches) {
   std::vector<std::vector<float>> points;
   for (uint64_t i = 0; i < 40; ++i) {
     points.push_back(make_vec({static_cast<float>(i), 0.0f}));
-    index.insert(i, /*base_pk=*/i, as_bytes(points.back()));
+    index.insert(i + 1, /*base_pk=*/i, as_bytes(points.back()));
   }
 
   const auto query = make_vec({0.0f, 0.0f});
@@ -302,7 +302,7 @@ TEST_F(HnswTest, StreamDrainsEntireGraph) {
 
   for (uint64_t i = 0; i < kNumPoints; ++i) {
     const auto v = make_vec({static_cast<float>(i), 0.0f});
-    index.insert(i, /*base_pk=*/1000 + i, as_bytes(v));
+    index.insert(i + 1, /*base_pk=*/1000 + i, as_bytes(v));
   }
 
   const auto query = make_vec({0.0f, 0.0f});
@@ -326,7 +326,7 @@ TEST_F(HnswTest, StreamDrainsEntireGraphWithEfSmallerThanGraph) {
 
   for (uint64_t i = 0; i < kNumPoints; ++i) {
     const auto v = make_vec({static_cast<float>(i), 0.0f});
-    index.insert(i, /*base_pk=*/1000 + i, as_bytes(v));
+    index.insert(i + 1, /*base_pk=*/1000 + i, as_bytes(v));
   }
 
   const auto query = make_vec({0.0f, 0.0f});
@@ -354,7 +354,7 @@ TEST_F(HnswTest, StreamNoDuplicatesInMultiDimensionalGraph) {
   const auto points =
       make_pseudo_random_points(kNumPoints, kStreamDims, &state);
   for (size_t i = 0; i < kNumPoints; ++i) {
-    index.insert(i, /*base_pk=*/i, as_bytes(points[i]));
+    index.insert(i + 1, /*base_pk=*/i, as_bytes(points[i]));
   }
 
   const auto queries =
@@ -382,30 +382,33 @@ TEST_F(HnswTest, StreamYieldsEachNodeAtMostOnce) {
 
   for (uint64_t i = 0; i < 5; ++i) {
     index.insert(
-        i, i,
+        i + 1, i + 1,
         as_bytes(make_vec({100.0f + 0.5f * static_cast<float>(i), 0.0f})));
   }
   for (uint64_t i = 5; i < 10; ++i) {
     index.insert(
-        i, i,
+        i + 1, i + 1,
         as_bytes(make_vec({100.0f + 0.5f * static_cast<float>(i - 5), 1.0f})));
   }
   for (uint64_t i = 10; i < 15; ++i) {
     index.insert(
-        i, i,
+        i + 1, i + 1,
         as_bytes(make_vec({100.0f + 0.5f * static_cast<float>(i - 10), 2.0f})));
   }
   for (uint64_t i = 100; i < 105; ++i) {
     index.insert(
-        i, i, as_bytes(make_vec({0.5f * static_cast<float>(i - 100), 0.0f})));
+        i + 1, i + 1,
+        as_bytes(make_vec({0.5f * static_cast<float>(i - 100), 0.0f})));
   }
   for (uint64_t i = 105; i < 110; ++i) {
     index.insert(
-        i, i, as_bytes(make_vec({0.5f * static_cast<float>(i - 105), 1.0f})));
+        i + 1, i + 1,
+        as_bytes(make_vec({0.5f * static_cast<float>(i - 105), 1.0f})));
   }
   for (uint64_t i = 110; i < 115; ++i) {
     index.insert(
-        i, i, as_bytes(make_vec({0.5f * static_cast<float>(i - 110), 2.0f})));
+        i + 1, i + 1,
+        as_bytes(make_vec({0.5f * static_cast<float>(i - 110), 2.0f})));
   }
 
   const auto query = make_vec({0.0f, 0.0f});
@@ -466,7 +469,7 @@ TEST_F(HnswTest, StreamBruteForceRecall) {
   const auto points =
       make_pseudo_random_points(kNumPoints, kRecallDims, &state);
   for (size_t i = 0; i < kNumPoints; ++i) {
-    index.insert(i, /*base_pk=*/i, as_bytes(points[i]));
+    index.insert(i + 1, /*base_pk=*/i, as_bytes(points[i]));
   }
 
   const auto queries =
@@ -506,14 +509,153 @@ TEST_F(HnswTest, StreamBruteForceRecall) {
       << "avg stream recall@125=" << avg_recall;
 }
 
+namespace {
+
+void assert_round_trip_knn(RoundTripFixture *fixture, size_t dims, size_t M,
+                           size_t ef_construction, size_t k, size_t ef_search,
+                           bool validate_built [[maybe_unused]]) {
+  LoadTestHnsw built(dims, euclidean, M, ef_construction);
+  populate_round_trip_index(built, fixture);
+  ASSERT_GT(fixture->store.entry_point, 0U);
+#ifndef NDEBUG
+  if (validate_built) {
+    EXPECT_TRUE(built.validate());
+  }
+#endif
+
+  LoadTestHnsw reloaded(dims, euclidean, M, ef_construction);
+  reloaded.init_from_entry_point(fixture->store.entry_point, &fixture->store);
+
+  const auto from_built = built.k_nn_search(as_bytes(fixture->query), k,
+                                            ef_search, &fixture->store);
+  const auto from_reloaded = reloaded.k_nn_search(as_bytes(fixture->query), k,
+                                                  ef_search, &fixture->store);
+
+  ASSERT_EQ(from_built.size(), from_reloaded.size());
+  EXPECT_EQ(from_built, from_reloaded);
+}
+
+void assert_round_trip_stream(RoundTripFixture *fixture, size_t dims, size_t M,
+                              size_t ef_construction, size_t batch_size,
+                              size_t ef_search, size_t max_results) {
+  LoadTestHnsw built(dims, euclidean, M, ef_construction);
+  populate_round_trip_index(built, fixture);
+  ASSERT_GT(fixture->store.entry_point, 0U);
+
+  LoadTestHnsw reloaded(dims, euclidean, M, ef_construction);
+  reloaded.init_from_entry_point(fixture->store.entry_point, &fixture->store);
+
+  const auto stream_built = drain_stream(built, as_bytes(fixture->query),
+                                         batch_size, ef_search, max_results);
+  const auto stream_reloaded =
+      drain_stream(reloaded, as_bytes(fixture->query), batch_size, ef_search,
+                   max_results, &fixture->store);
+
+  ASSERT_EQ(stream_built.size(), stream_reloaded.size());
+  EXPECT_EQ(stream_built, stream_reloaded);
+}
+
+}  // namespace
+
+TEST_F(HnswTest, RoundTripSmallGraph) {
+  RoundTripFixture fixture = make_fixed_round_trip_fixture(kDims);
+  assert_round_trip_knn(&fixture, kDims, kM, kEfConstruction,
+                        /*k=*/3, /*ef_search=*/16, /*validate_built=*/true);
+}
+
+TEST_F(HnswTest, RoundTripRandomGraph) {
+  constexpr size_t kNumPoints = 1000;
+  constexpr uint64_t kSeed = 4242;
+  RoundTripFixture fixture =
+      make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
+  assert_round_trip_knn(&fixture, kDims, kM, kEfConstruction,
+                        /*k=*/10, /*ef_search=*/50, /*validate_built=*/false);
+}
+
+TEST_F(HnswTest, RoundTripStreamSearch) {
+  RoundTripFixture fixture = make_fixed_round_trip_fixture(kDims);
+  assert_round_trip_stream(&fixture, kDims, kM, kEfConstruction,
+                           /*batch_size=*/2, /*ef_search=*/16,
+                           /*max_results=*/3);
+}
+
+TEST_F(HnswTest, RoundTripStreamRandomGraph) {
+  constexpr size_t kNumPoints = 1000;
+  constexpr uint64_t kSeed = 5150;
+  RoundTripFixture fixture =
+      make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
+  assert_round_trip_stream(&fixture, kDims, kM, kEfConstruction,
+                           /*batch_size=*/5, /*ef_search=*/50,
+                           /*max_results=*/25);
+}
+
+TEST_F(HnswTest, InitFromEntryPointLoadsEP) {
+  RoundTripFixture fixture = make_fixed_round_trip_fixture(kDims);
+  LoadTestHnsw built(kDims, euclidean, kM, kEfConstruction);
+  populate_round_trip_index(built, &fixture);
+  ASSERT_GT(fixture.store.entry_point, 0U);
+
+  LoadTestHnsw reloaded(kDims, euclidean, kM, kEfConstruction);
+  reloaded.init_from_entry_point(fixture.store.entry_point, &fixture.store);
+
+  EXPECT_EQ(1U, fixture.store.load_counts.size());
+  EXPECT_EQ(1U, fixture.store.load_counts.at(fixture.store.entry_point));
+  for (uint64_t id : fixture.graph_ids) {
+    if (id == fixture.store.entry_point) {
+      continue;
+    }
+    EXPECT_EQ(0U, fixture.store.load_counts.count(id));
+  }
+}
+
+TEST_F(HnswTest, SearchLoadsNeighborsOnDemand) {
+  constexpr size_t kNumPoints = 1000;
+  constexpr uint64_t kSeed = 7777;
+  RoundTripFixture fixture =
+      make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
+  LoadTestHnsw built(kDims, euclidean, kM, kEfConstruction);
+  populate_round_trip_index(built, &fixture);
+  ASSERT_GT(fixture.store.entry_point, 0U);
+
+  LoadTestHnsw reloaded(kDims, euclidean, kM, kEfConstruction);
+  reloaded.init_from_entry_point(fixture.store.entry_point, &fixture.store);
+  ASSERT_EQ(1U, fixture.store.load_counts.size());
+
+  const auto hits =
+      reloaded.k_nn_search(as_bytes(fixture.query), 10, 50, &fixture.store);
+  ASSERT_GE(hits.size(), 1U);
+  EXPECT_GT(fixture.store.load_counts.size(), 1U);
+  EXPECT_LT(fixture.store.load_counts.size(), fixture.graph_ids.size());
+}
+
+TEST_F(HnswTest, LoadNodeIdempotent) {
+  constexpr size_t kNumPoints = 1000;
+  constexpr uint64_t kSeed = 8888;
+  RoundTripFixture fixture =
+      make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
+  LoadTestHnsw built(kDims, euclidean, kM, kEfConstruction);
+  populate_round_trip_index(built, &fixture);
+  ASSERT_GT(fixture.store.entry_point, 0U);
+
+  LoadTestHnsw reloaded(kDims, euclidean, kM, kEfConstruction);
+  reloaded.init_from_entry_point(fixture.store.entry_point, &fixture.store);
+
+  const auto query = as_bytes(fixture.query);
+  (void)reloaded.k_nn_search(query, 10, 50, &fixture.store);
+  const auto counts_after_first = fixture.store.load_counts;
+
+  (void)reloaded.k_nn_search(query, 10, 50, &fixture.store);
+  EXPECT_EQ(counts_after_first, fixture.store.load_counts);
+}
+
 #ifndef NDEBUG
 TEST_F(HnswTest, GraphInvariants) {
   TestHnsw index(kDims, euclidean, kM, kEfConstruction);
   EXPECT_TRUE(index.validate());
 
   for (uint64_t i = 0; i < 50; ++i) {
     const auto v = make_vec({static_cast<float>(i), static_cast<float>(i % 7)});
-    index.insert(i, 1000 + i, as_bytes(v));
+    index.insert(i + 1, 1000 + i, as_bytes(v));
     EXPECT_TRUE(index.validate()) << "after insert i=" << i;
   }
 }
@@ -525,6 +667,12 @@ TEST(HnswDeathTest, SearchKZeroAsserts) {
   EXPECT_DEATH_IF_SUPPORTED(index.k_nn_search(as_bytes(v), /*k=*/0, 16), "");
 }
 
+TEST(HnswDeathTest, ZeroIdAsserts) {
+  TestHnsw index(2, euclidean, 4, 16);
+  const auto v = make_vec({1.0f, 2.0f});
+  EXPECT_DEATH_IF_SUPPORTED(index.insert(0, 100, as_bytes(v)), "");
+}
+
 TEST(HnswDeathTest, DuplicateIdAsserts) {
   TestHnsw index(2, euclidean, 4, 16);
   const auto v = make_vec({1.0f, 2.0f});

unittest/gunit/hnsw_bench-t.cc

@@ -193,7 +193,7 @@ class BenchFixture {
                       m_ef_construction);
     }
     for (size_t i = 0; i < m_num_points; ++i) {
-      m_index->insert(i, /*base_pk=*/i, as_bytes(m_points[i]));
+      m_index->insert(i + 1, /*base_pk=*/i, as_bytes(m_points[i]));
     }
     m_build_seconds = seconds_since(build_start);
 
@@ -236,7 +236,7 @@ class BenchFixture {
   size_t arena_requested_bytes() const { return m_arena.bytes_requested; }
   size_t arena_requests() const { return m_arena.requests_number; }
 
-  const BorrowedHnsw &index() const { return *m_index; }
+  BorrowedHnsw &index() { return *m_index; }
   const std::vector<float> &query(size_t q) const { return m_queries[q]; }
   const std::unordered_set<uint64_t> &ground_truth(size_t q) const {
     return m_ground_truth[q];
@@ -288,8 +288,8 @@ class BenchFixture {
 
 /** Built on first use, so a --gtest_filter run only pays for what it selects.
  */
-static const BenchFixture &fixture() {
-  static const BenchFixture instance;
+static BenchFixture &fixture() {
+  static BenchFixture instance;
   return instance;
 }
 
@@ -299,7 +299,7 @@ struct RecallStats {
   double worst = 1.0;
 };
 
-static RecallStats measure_knn_recall(const BenchFixture &f, size_t ef_search) {
+static RecallStats measure_knn_recall(BenchFixture &f, size_t ef_search) {
   RecallStats stats;
   for (size_t q = 0; q < f.num_queries(); ++q) {
     const std::vector<uint64_t> found =
@@ -312,8 +312,7 @@ static RecallStats measure_knn_recall(const BenchFixture &f, size_t ef_search) {
   return stats;
 }
 
-static RecallStats measure_stream_recall(const BenchFixture &f,
-                                         size_t ef_search) {
+static RecallStats measure_stream_recall(BenchFixture &f, size_t ef_search) {
   RecallStats stats;
   for (size_t q = 0; q < f.num_queries(); ++q) {
     const std::vector<uint64_t> found =
@@ -329,7 +328,7 @@ static RecallStats measure_stream_recall(const BenchFixture &f,
 }
 
 TEST(HnswBenchmark, RecallVsEfSearch) {
-  const BenchFixture &f = fixture();
+  BenchFixture &f = fixture();
   f.print_config();
 
   std::printf("  %9s | %9s %9s | %9s %9s | %8s\n", "ef_search", "knn_avg",
@@ -380,7 +379,7 @@ TEST(HnswBenchmark, RecallVsEfSearch) {
 }
 
 TEST(HnswBenchmark, StreamFullDrainQuality) {
-  const BenchFixture &f = fixture();
+  BenchFixture &f = fixture();
 
   // A full drain costs O(N * Mmax) distance evaluations, so sample a few
   // queries rather than the whole set.
@@ -480,7 +479,7 @@ static void BM_HnswKnnSearch(size_t num_iterations) {
                   "(configure with -DWITH_DEBUG=OFF for meaningful results)";
 #endif
   StopBenchmarkTiming();
-  const BenchFixture &f = fixture();
+  BenchFixture &f = fixture();
   constexpr size_t kEfSearch = 64;
 
   StartBenchmarkTiming();
@@ -500,7 +499,7 @@ static void BM_HnswStreamSearch(size_t num_iterations) {
                   "(configure with -DWITH_DEBUG=OFF for meaningful results)";
 #endif
   StopBenchmarkTiming();
-  const BenchFixture &f = fixture();
+  BenchFixture &f = fixture();
   constexpr size_t kEfSearch = 64;
 
   StartBenchmarkTiming();