Skip to content

Commit 6bec2b8

Browse files
dlenevdlenev
committed
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.
1 parent 65261fc commit 6bec2b8

4 files changed

Lines changed: 817 additions & 119 deletions

File tree

‎unittest/gunit/hnsw-t.cc‎

Lines changed: 162 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ TEST_F(HnswTest, BruteForceRecall) {
161161
const auto points =
162162
make_pseudo_random_points(kNumPoints, kRecallDims, &state);
163163
for (size_t i = 0; i < kNumPoints; ++i) {
164-
index.insert(i, /*base_pk=*/i, as_bytes(points[i]));
164+
index.insert(i + 1, /*base_pk=*/i, as_bytes(points[i]));
165165
}
166166

167167
const auto queries =
@@ -257,7 +257,7 @@ TEST_F(HnswTest, StreamMultipleBatchesNoDuplicates) {
257257
TestHnsw index(kDims, euclidean, kM, kEfConstruction);
258258
for (uint64_t i = 0; i < 30; ++i) {
259259
const auto v = make_vec({static_cast<float>(i), 0.0f});
260-
index.insert(i, 1000 + i, as_bytes(v));
260+
index.insert(i + 1, 1000 + i, as_bytes(v));
261261
}
262262
const auto query = make_vec({0.0f, 0.0f});
263263
// Small batch/ef forces continuation refills across batches.
@@ -275,7 +275,7 @@ TEST_F(HnswTest, StreamDistancesNonDecreasingAcrossBatches) {
275275
std::vector<std::vector<float>> points;
276276
for (uint64_t i = 0; i < 40; ++i) {
277277
points.push_back(make_vec({static_cast<float>(i), 0.0f}));
278-
index.insert(i, /*base_pk=*/i, as_bytes(points.back()));
278+
index.insert(i + 1, /*base_pk=*/i, as_bytes(points.back()));
279279
}
280280

281281
const auto query = make_vec({0.0f, 0.0f});
@@ -302,7 +302,7 @@ TEST_F(HnswTest, StreamDrainsEntireGraph) {
302302

303303
for (uint64_t i = 0; i < kNumPoints; ++i) {
304304
const auto v = make_vec({static_cast<float>(i), 0.0f});
305-
index.insert(i, /*base_pk=*/1000 + i, as_bytes(v));
305+
index.insert(i + 1, /*base_pk=*/1000 + i, as_bytes(v));
306306
}
307307

308308
const auto query = make_vec({0.0f, 0.0f});
@@ -326,7 +326,7 @@ TEST_F(HnswTest, StreamDrainsEntireGraphWithEfSmallerThanGraph) {
326326

327327
for (uint64_t i = 0; i < kNumPoints; ++i) {
328328
const auto v = make_vec({static_cast<float>(i), 0.0f});
329-
index.insert(i, /*base_pk=*/1000 + i, as_bytes(v));
329+
index.insert(i + 1, /*base_pk=*/1000 + i, as_bytes(v));
330330
}
331331

332332
const auto query = make_vec({0.0f, 0.0f});
@@ -354,7 +354,7 @@ TEST_F(HnswTest, StreamNoDuplicatesInMultiDimensionalGraph) {
354354
const auto points =
355355
make_pseudo_random_points(kNumPoints, kStreamDims, &state);
356356
for (size_t i = 0; i < kNumPoints; ++i) {
357-
index.insert(i, /*base_pk=*/i, as_bytes(points[i]));
357+
index.insert(i + 1, /*base_pk=*/i, as_bytes(points[i]));
358358
}
359359

360360
const auto queries =
@@ -382,30 +382,33 @@ TEST_F(HnswTest, StreamYieldsEachNodeAtMostOnce) {
382382

383383
for (uint64_t i = 0; i < 5; ++i) {
384384
index.insert(
385-
i, i,
385+
i + 1, i + 1,
386386
as_bytes(make_vec({100.0f + 0.5f * static_cast<float>(i), 0.0f})));
387387
}
388388
for (uint64_t i = 5; i < 10; ++i) {
389389
index.insert(
390-
i, i,
390+
i + 1, i + 1,
391391
as_bytes(make_vec({100.0f + 0.5f * static_cast<float>(i - 5), 1.0f})));
392392
}
393393
for (uint64_t i = 10; i < 15; ++i) {
394394
index.insert(
395-
i, i,
395+
i + 1, i + 1,
396396
as_bytes(make_vec({100.0f + 0.5f * static_cast<float>(i - 10), 2.0f})));
397397
}
398398
for (uint64_t i = 100; i < 105; ++i) {
399399
index.insert(
400-
i, i, as_bytes(make_vec({0.5f * static_cast<float>(i - 100), 0.0f})));
400+
i + 1, i + 1,
401+
as_bytes(make_vec({0.5f * static_cast<float>(i - 100), 0.0f})));
401402
}
402403
for (uint64_t i = 105; i < 110; ++i) {
403404
index.insert(
404-
i, i, as_bytes(make_vec({0.5f * static_cast<float>(i - 105), 1.0f})));
405+
i + 1, i + 1,
406+
as_bytes(make_vec({0.5f * static_cast<float>(i - 105), 1.0f})));
405407
}
406408
for (uint64_t i = 110; i < 115; ++i) {
407409
index.insert(
408-
i, i, as_bytes(make_vec({0.5f * static_cast<float>(i - 110), 2.0f})));
410+
i + 1, i + 1,
411+
as_bytes(make_vec({0.5f * static_cast<float>(i - 110), 2.0f})));
409412
}
410413

411414
const auto query = make_vec({0.0f, 0.0f});
@@ -466,7 +469,7 @@ TEST_F(HnswTest, StreamBruteForceRecall) {
466469
const auto points =
467470
make_pseudo_random_points(kNumPoints, kRecallDims, &state);
468471
for (size_t i = 0; i < kNumPoints; ++i) {
469-
index.insert(i, /*base_pk=*/i, as_bytes(points[i]));
472+
index.insert(i + 1, /*base_pk=*/i, as_bytes(points[i]));
470473
}
471474

472475
const auto queries =
@@ -506,14 +509,153 @@ TEST_F(HnswTest, StreamBruteForceRecall) {
506509
<< "avg stream recall@125=" << avg_recall;
507510
}
508511

512+
namespace {
513+
514+
void assert_round_trip_knn(RoundTripFixture *fixture, size_t dims, size_t M,
515+
size_t ef_construction, size_t k, size_t ef_search,
516+
bool validate_built [[maybe_unused]]) {
517+
LoadTestHnsw built(dims, euclidean, M, ef_construction);
518+
populate_round_trip_index(built, fixture);
519+
ASSERT_GT(fixture->store.entry_point, 0U);
520+
#ifndef NDEBUG
521+
if (validate_built) {
522+
EXPECT_TRUE(built.validate());
523+
}
524+
#endif
525+
526+
LoadTestHnsw reloaded(dims, euclidean, M, ef_construction);
527+
reloaded.init_from_entry_point(fixture->store.entry_point, &fixture->store);
528+
529+
const auto from_built = built.k_nn_search(as_bytes(fixture->query), k,
530+
ef_search, &fixture->store);
531+
const auto from_reloaded = reloaded.k_nn_search(as_bytes(fixture->query), k,
532+
ef_search, &fixture->store);
533+
534+
ASSERT_EQ(from_built.size(), from_reloaded.size());
535+
EXPECT_EQ(from_built, from_reloaded);
536+
}
537+
538+
void assert_round_trip_stream(RoundTripFixture *fixture, size_t dims, size_t M,
539+
size_t ef_construction, size_t batch_size,
540+
size_t ef_search, size_t max_results) {
541+
LoadTestHnsw built(dims, euclidean, M, ef_construction);
542+
populate_round_trip_index(built, fixture);
543+
ASSERT_GT(fixture->store.entry_point, 0U);
544+
545+
LoadTestHnsw reloaded(dims, euclidean, M, ef_construction);
546+
reloaded.init_from_entry_point(fixture->store.entry_point, &fixture->store);
547+
548+
const auto stream_built = drain_stream(built, as_bytes(fixture->query),
549+
batch_size, ef_search, max_results);
550+
const auto stream_reloaded =
551+
drain_stream(reloaded, as_bytes(fixture->query), batch_size, ef_search,
552+
max_results, &fixture->store);
553+
554+
ASSERT_EQ(stream_built.size(), stream_reloaded.size());
555+
EXPECT_EQ(stream_built, stream_reloaded);
556+
}
557+
558+
} // namespace
559+
560+
TEST_F(HnswTest, RoundTripSmallGraph) {
561+
RoundTripFixture fixture = make_fixed_round_trip_fixture(kDims);
562+
assert_round_trip_knn(&fixture, kDims, kM, kEfConstruction,
563+
/*k=*/3, /*ef_search=*/16, /*validate_built=*/true);
564+
}
565+
566+
TEST_F(HnswTest, RoundTripRandomGraph) {
567+
constexpr size_t kNumPoints = 1000;
568+
constexpr uint64_t kSeed = 4242;
569+
RoundTripFixture fixture =
570+
make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
571+
assert_round_trip_knn(&fixture, kDims, kM, kEfConstruction,
572+
/*k=*/10, /*ef_search=*/50, /*validate_built=*/false);
573+
}
574+
575+
TEST_F(HnswTest, RoundTripStreamSearch) {
576+
RoundTripFixture fixture = make_fixed_round_trip_fixture(kDims);
577+
assert_round_trip_stream(&fixture, kDims, kM, kEfConstruction,
578+
/*batch_size=*/2, /*ef_search=*/16,
579+
/*max_results=*/3);
580+
}
581+
582+
TEST_F(HnswTest, RoundTripStreamRandomGraph) {
583+
constexpr size_t kNumPoints = 1000;
584+
constexpr uint64_t kSeed = 5150;
585+
RoundTripFixture fixture =
586+
make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
587+
assert_round_trip_stream(&fixture, kDims, kM, kEfConstruction,
588+
/*batch_size=*/5, /*ef_search=*/50,
589+
/*max_results=*/25);
590+
}
591+
592+
TEST_F(HnswTest, InitFromEntryPointLoadsEP) {
593+
RoundTripFixture fixture = make_fixed_round_trip_fixture(kDims);
594+
LoadTestHnsw built(kDims, euclidean, kM, kEfConstruction);
595+
populate_round_trip_index(built, &fixture);
596+
ASSERT_GT(fixture.store.entry_point, 0U);
597+
598+
LoadTestHnsw reloaded(kDims, euclidean, kM, kEfConstruction);
599+
reloaded.init_from_entry_point(fixture.store.entry_point, &fixture.store);
600+
601+
EXPECT_EQ(1U, fixture.store.load_counts.size());
602+
EXPECT_EQ(1U, fixture.store.load_counts.at(fixture.store.entry_point));
603+
for (uint64_t id : fixture.graph_ids) {
604+
if (id == fixture.store.entry_point) {
605+
continue;
606+
}
607+
EXPECT_EQ(0U, fixture.store.load_counts.count(id));
608+
}
609+
}
610+
611+
TEST_F(HnswTest, SearchLoadsNeighborsOnDemand) {
612+
constexpr size_t kNumPoints = 1000;
613+
constexpr uint64_t kSeed = 7777;
614+
RoundTripFixture fixture =
615+
make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
616+
LoadTestHnsw built(kDims, euclidean, kM, kEfConstruction);
617+
populate_round_trip_index(built, &fixture);
618+
ASSERT_GT(fixture.store.entry_point, 0U);
619+
620+
LoadTestHnsw reloaded(kDims, euclidean, kM, kEfConstruction);
621+
reloaded.init_from_entry_point(fixture.store.entry_point, &fixture.store);
622+
ASSERT_EQ(1U, fixture.store.load_counts.size());
623+
624+
const auto hits =
625+
reloaded.k_nn_search(as_bytes(fixture.query), 10, 50, &fixture.store);
626+
ASSERT_GE(hits.size(), 1U);
627+
EXPECT_GT(fixture.store.load_counts.size(), 1U);
628+
EXPECT_LT(fixture.store.load_counts.size(), fixture.graph_ids.size());
629+
}
630+
631+
TEST_F(HnswTest, LoadNodeIdempotent) {
632+
constexpr size_t kNumPoints = 1000;
633+
constexpr uint64_t kSeed = 8888;
634+
RoundTripFixture fixture =
635+
make_random_round_trip_fixture(kDims, kNumPoints, kSeed);
636+
LoadTestHnsw built(kDims, euclidean, kM, kEfConstruction);
637+
populate_round_trip_index(built, &fixture);
638+
ASSERT_GT(fixture.store.entry_point, 0U);
639+
640+
LoadTestHnsw reloaded(kDims, euclidean, kM, kEfConstruction);
641+
reloaded.init_from_entry_point(fixture.store.entry_point, &fixture.store);
642+
643+
const auto query = as_bytes(fixture.query);
644+
(void)reloaded.k_nn_search(query, 10, 50, &fixture.store);
645+
const auto counts_after_first = fixture.store.load_counts;
646+
647+
(void)reloaded.k_nn_search(query, 10, 50, &fixture.store);
648+
EXPECT_EQ(counts_after_first, fixture.store.load_counts);
649+
}
650+
509651
#ifndef NDEBUG
510652
TEST_F(HnswTest, GraphInvariants) {
511653
TestHnsw index(kDims, euclidean, kM, kEfConstruction);
512654
EXPECT_TRUE(index.validate());
513655

514656
for (uint64_t i = 0; i < 50; ++i) {
515657
const auto v = make_vec({static_cast<float>(i), static_cast<float>(i % 7)});
516-
index.insert(i, 1000 + i, as_bytes(v));
658+
index.insert(i + 1, 1000 + i, as_bytes(v));
517659
EXPECT_TRUE(index.validate()) << "after insert i=" << i;
518660
}
519661
}
@@ -525,6 +667,12 @@ TEST(HnswDeathTest, SearchKZeroAsserts) {
525667
EXPECT_DEATH_IF_SUPPORTED(index.k_nn_search(as_bytes(v), /*k=*/0, 16), "");
526668
}
527669

670+
TEST(HnswDeathTest, ZeroIdAsserts) {
671+
TestHnsw index(2, euclidean, 4, 16);
672+
const auto v = make_vec({1.0f, 2.0f});
673+
EXPECT_DEATH_IF_SUPPORTED(index.insert(0, 100, as_bytes(v)), "");
674+
}
675+
528676
TEST(HnswDeathTest, DuplicateIdAsserts) {
529677
TestHnsw index(2, euclidean, 4, 16);
530678
const auto v = make_vec({1.0f, 2.0f});

‎unittest/gunit/hnsw_bench-t.cc‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ class BenchFixture {
193193
m_ef_construction);
194194
}
195195
for (size_t i = 0; i < m_num_points; ++i) {
196-
m_index->insert(i, /*base_pk=*/i, as_bytes(m_points[i]));
196+
m_index->insert(i + 1, /*base_pk=*/i, as_bytes(m_points[i]));
197197
}
198198
m_build_seconds = seconds_since(build_start);
199199

@@ -236,7 +236,7 @@ class BenchFixture {
236236
size_t arena_requested_bytes() const { return m_arena.bytes_requested; }
237237
size_t arena_requests() const { return m_arena.requests_number; }
238238

239-
const BorrowedHnsw &index() const { return *m_index; }
239+
BorrowedHnsw &index() { return *m_index; }
240240
const std::vector<float> &query(size_t q) const { return m_queries[q]; }
241241
const std::unordered_set<uint64_t> &ground_truth(size_t q) const {
242242
return m_ground_truth[q];
@@ -288,8 +288,8 @@ class BenchFixture {
288288

289289
/** Built on first use, so a --gtest_filter run only pays for what it selects.
290290
*/
291-
static const BenchFixture &fixture() {
292-
static const BenchFixture instance;
291+
static BenchFixture &fixture() {
292+
static BenchFixture instance;
293293
return instance;
294294
}
295295

@@ -299,7 +299,7 @@ struct RecallStats {
299299
double worst = 1.0;
300300
};
301301

302-
static RecallStats measure_knn_recall(const BenchFixture &f, size_t ef_search) {
302+
static RecallStats measure_knn_recall(BenchFixture &f, size_t ef_search) {
303303
RecallStats stats;
304304
for (size_t q = 0; q < f.num_queries(); ++q) {
305305
const std::vector<uint64_t> found =
@@ -312,8 +312,7 @@ static RecallStats measure_knn_recall(const BenchFixture &f, size_t ef_search) {
312312
return stats;
313313
}
314314

315-
static RecallStats measure_stream_recall(const BenchFixture &f,
316-
size_t ef_search) {
315+
static RecallStats measure_stream_recall(BenchFixture &f, size_t ef_search) {
317316
RecallStats stats;
318317
for (size_t q = 0; q < f.num_queries(); ++q) {
319318
const std::vector<uint64_t> found =
@@ -329,7 +328,7 @@ static RecallStats measure_stream_recall(const BenchFixture &f,
329328
}
330329

331330
TEST(HnswBenchmark, RecallVsEfSearch) {
332-
const BenchFixture &f = fixture();
331+
BenchFixture &f = fixture();
333332
f.print_config();
334333

335334
std::printf(" %9s | %9s %9s | %9s %9s | %8s\n", "ef_search", "knn_avg",
@@ -380,7 +379,7 @@ TEST(HnswBenchmark, RecallVsEfSearch) {
380379
}
381380

382381
TEST(HnswBenchmark, StreamFullDrainQuality) {
383-
const BenchFixture &f = fixture();
382+
BenchFixture &f = fixture();
384383

385384
// A full drain costs O(N * Mmax) distance evaluations, so sample a few
386385
// queries rather than the whole set.
@@ -480,7 +479,7 @@ static void BM_HnswKnnSearch(size_t num_iterations) {
480479
"(configure with -DWITH_DEBUG=OFF for meaningful results)";
481480
#endif
482481
StopBenchmarkTiming();
483-
const BenchFixture &f = fixture();
482+
BenchFixture &f = fixture();
484483
constexpr size_t kEfSearch = 64;
485484

486485
StartBenchmarkTiming();
@@ -500,7 +499,7 @@ static void BM_HnswStreamSearch(size_t num_iterations) {
500499
"(configure with -DWITH_DEBUG=OFF for meaningful results)";
501500
#endif
502501
StopBenchmarkTiming();
503-
const BenchFixture &f = fixture();
502+
BenchFixture &f = fixture();
504503
constexpr size_t kEfSearch = 64;
505504

506505
StartBenchmarkTiming();

0 commit comments

Comments
 (0)