[#26937] DocDB: YbHnsw persistence

Summary:
The YbHnsw library currently generates a block-based representation for the HNSW (Hierarchical Navigable Small World) graph structure. However, this representation is only stored in memory and lacks persistence, meaning it is lost when the application terminates or restarts.

To make YbHnsw a viable solution for implementing a vector index, we need to add functionality for saving the in-memory HNSW graph data to the file system and reloading it when needed. This persistence layer will ensure that the graph remains available across application sessions, improving usability and efficiency.

Also added basic wiring (not final interface) for block cache. To provide ability to load and unload blocks from the diff in followup diffs.
Jira: DB-16382

Test Plan: YbHnswTest.Persistence

Reviewers: arybochkin

Reviewed By: arybochkin

Subscribers: ybase

Tags: #jenkins-ready

Differential Revision: https://phorge.dev.yugabyte.com/D43473

Author: spolitov

CMakeLists.txt

@@ -897,6 +897,7 @@ set(YB_SUBDIR_NAMES
     fs
     gen_yrpc
     gutil
+    hnsw
     integration-tests
     master
     qlexpr

src/yb/hnsw/CMakeLists.txt

@@ -0,0 +1,34 @@
+# Copyright (c) YugabyteDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License.  You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied.  See the License for the specific language governing permissions and limitations
+# under the License.
+#
+
+set(YB_PCH_PREFIX hnsw)
+
+set(NHSW_SRCS
+  hnsw_block_cache.cc
+  hnsw.cc
+  )
+
+set(NHSW_LIBS
+  vector_index
+  yb_common
+  yb_util
+  )
+
+ADD_YB_LIBRARY(hnsw
+    SRCS ${NHSW_SRCS}
+    DEPS ${NHSW_LIBS})
+
+set(YB_TEST_LINK_LIBS
+    hnsw yb_common_test_util ${YB_MIN_TEST_LIBS})
+
+ADD_YB_TEST(hnsw-test)

src/yb/hnsw/block_writer.h

@@ -0,0 +1,91 @@
+// Copyright (c) YugabyteDB, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied.  See the License for the specific language governing permissions and limitations
+// under the License.
+//
+
+#pragma once
+
+#include "yb/hnsw/types.h"
+
+#include "yb/util/logging.h"
+
+namespace yb::hnsw {
+
+class BlockWriter {
+ public:
+  BlockWriter() = default;
+  BlockWriter(std::byte* out, std::byte* end) : out_(out), end_(end) {
+    DCHECK_LE(out_, end_);
+  }
+
+  explicit BlockWriter(DataBlock& block)
+      : BlockWriter(block.data(), block.data() + block.size()) {}
+
+  BlockWriter(BlockWriter&& rhs)
+      : out_(std::exchange(rhs.out_, nullptr)), end_(std::exchange(rhs.end_, nullptr)) {
+  }
+
+  void operator=(BlockWriter&& rhs) {
+    DCHECK_EQ(out_, end_);
+    out_ = rhs.out_;
+    end_ = rhs.end_;
+    rhs.out_ = rhs.end_;
+  }
+
+  ~BlockWriter() {
+    DCHECK_EQ(out_, end_);
+  }
+
+  std::byte* out() const {
+    return out_;
+  }
+
+  std::byte* end() const {
+    return end_;
+  }
+
+  size_t SpaceLeft() const {
+    return end_ - out_;
+  }
+
+  BlockWriter Split(size_t size) {
+    auto old_end = end_;
+    end_ = out_ + size;
+    DCHECK_LE(end_, old_end);
+    return BlockWriter(end_, old_end);
+  }
+
+  std::byte* Prepare(size_t size) {
+    auto result = out_;
+    out_ += size;
+    DCHECK_LE(out_, end_);
+    return result;
+  }
+
+  template <class Value>
+  void Append(Value value) {
+    Store<Value, HnswEndian>(out_, value);
+    out_ += sizeof(Value);
+  }
+
+  template <class... Args>
+  void AppendBytes(Args&&... args) {
+    Slice slice(std::forward<Args>(args)...);
+    memcpy(out_, slice.data(), slice.size());
+    out_ += slice.size();
+  }
+
+ private:
+  std::byte* out_ = nullptr;
+  std::byte* end_ = nullptr;
+};
+
+} // namespace yb::hnsw

src/yb/vector_index/yb_hnsw-test.cc renamed to src/yb/hnsw/hnsw-test.cc

@@ -11,19 +11,20 @@
 // under the License.
 //
 
-#include "yb/util/test_util.h"
+#include "yb/hnsw/hnsw.h"
+#include "yb/hnsw/hnsw_block_cache.h"
 
 #include "yb/util/random_util.h"
+#include "yb/util/test_util.h"
 #include "yb/util/tsan_util.h"
 
 #include "yb/vector_index/vector_index_fwd.h"
 #include "yb/vector_index/distance.h"
 #include "yb/vector_index/usearch_include_wrapper_internal.h"
-#include "yb/vector_index/yb_hnsw.h"
 
-namespace yb::vector_index {
+namespace yb::hnsw {
 
-using IndexImpl = unum::usearch::index_dense_gt<VectorId>;
+using IndexImpl = unum::usearch::index_dense_gt<vector_index::VectorId>;
 using Vector = std::vector<float>;
 
 unum::usearch::index_dense_config_t CreateIndexDenseConfig() {
@@ -36,7 +37,7 @@ unum::usearch::index_dense_config_t CreateIndexDenseConfig() {
 }
 
 struct AcceptAllVectors {
-  bool operator()(const VectorId& id) const {
+  bool operator()(const vector_index::VectorId& id) const {
     return true;
   }
 };
@@ -62,7 +63,7 @@ class YbHnswTest : public YBTest {
 
   void InsertRandomVector(Vector& holder) {
     RandomVector(holder);
-    ASSERT_TRUE(index_.add(VectorId::GenerateRandom(), holder.data()));
+    ASSERT_TRUE(index_.add(vector_index::VectorId::GenerateRandom(), holder.data()));
   }
 
   void InsertRandomVectors(size_t count) {
@@ -79,7 +80,7 @@ class YbHnswTest : public YBTest {
   }
 
   void VerifySearch(const Vector& query_vector, size_t max_results) {
-    VectorFilter filter = AcceptAllVectors();
+    vector_index::VectorFilter filter = AcceptAllVectors();
     auto usearch_results = index_.filtered_search(query_vector.data(), max_results, filter);
     auto yb_hnsw_results = yb_hnsw_.Search(query_vector.data(), max_results, filter, context_);
     ASSERT_EQ(usearch_results.count, yb_hnsw_results.size());
@@ -91,24 +92,42 @@ class YbHnswTest : public YBTest {
   }
 
   std::vector<Vector> PrepareRandom(size_t num_vectors, size_t num_searches);
+  Status InitYbHnsw(bool load);
+
   void TestPerf();
+  void TestSimple(bool load);
 
   size_t dimensions_ = 8;
   size_t max_vectors_ = 65536;
   std::mt19937_64 rng_{42};
   unum::usearch::metric_punned_t metric_;
   IndexImpl index_;
+  BlockCachePtr block_cache_ = std::make_shared<BlockCache>(*Env::Default());
   YbHnsw yb_hnsw_;
   YbHnswSearchContext context_;
 };
 
-TEST_F(YbHnswTest, Simple) {
+Status YbHnswTest::InitYbHnsw(bool load) {
+  auto path = GetTestPath("0.yb_hnsw");
+  if (load) {
+    {
+      YbHnsw temp(metric_);
+      RETURN_NOT_OK(temp.Import(index_, path, block_cache_));
+    }
+    RETURN_NOT_OK(yb_hnsw_.Init(path, block_cache_));
+  } else {
+    RETURN_NOT_OK(yb_hnsw_.Import(index_, path, block_cache_));
+  }
+  return Status::OK();
+}
+
+void YbHnswTest::TestSimple(bool load) {
   constexpr size_t kNumVectors = 100;
   constexpr size_t kNumSearches = 10;
   constexpr size_t kMaxResults = 10;
 
   InsertRandomVectors(kNumVectors);
-  yb_hnsw_.Import(index_);
+  ASSERT_OK(InitYbHnsw(load));
 
   Vector query_vector;
   for (size_t i = 0; i != kNumSearches; ++i) {
@@ -117,10 +136,18 @@ TEST_F(YbHnswTest, Simple) {
   }
 }
 
+TEST_F(YbHnswTest, Simple) {
+  TestSimple(/* load= */ false);
+}
+
+TEST_F(YbHnswTest, Persistence) {
+  TestSimple(/* load= */ true);
+}
+
 std::vector<Vector> YbHnswTest::PrepareRandom(size_t num_vectors, size_t num_searches) {
   EXPECT_LE(num_vectors, max_vectors_);
   InsertRandomVectors(num_vectors);
-  yb_hnsw_.Import(index_);
+  EXPECT_OK(InitYbHnsw(false));
 
   std::vector<Vector> query_vectors(num_searches);
   for (auto& vector : query_vectors) {
@@ -151,7 +178,7 @@ void YbHnswTest::TestPerf() {
 
   auto query_vectors = PrepareRandom(num_vectors, num_searches);
   YbHnswSearchContext context;
-  VectorFilter filter = AcceptAllVectors();
+  vector_index::VectorFilter filter = AcceptAllVectors();
   MonoTime start = MonoTime::Now();
   for (const auto& query_vector : query_vectors) {
     index_.filtered_search(query_vector.data(), kMaxResults, filter);
@@ -183,4 +210,4 @@ TEST_F(YbHnswTest, Perf2048Dims) {
   TestPerf();
 }
 
-}  // namespace yb::vector_index
+}  // namespace yb::hnsw