feat: add HSNW index deserialization

Author: BorysTheDev

src/core/search/hnsw_index.cc

@@ -133,6 +133,19 @@ struct HnswlibAdapter {
     return metadata;
   }
 
+  void SetMetadata(const HnswIndexMetadata& metadata) {
+    MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kWriteLock);
+    if (world_.max_elements_ < metadata.max_elements) {
+      world_.resizeIndex(metadata.max_elements);
+    }
+    world_.cur_element_count.store(metadata.cur_element_count);
+    world_.maxlevel_ = metadata.maxlevel;
+    world_.enterpoint_node_ = metadata.enterpoint_node;
+    world_.ef_construction_ = metadata.ef_construction;
+    world_.mult_ = metadata.mult;
+    // Note: M, maxM, maxM0 are set at construction and shouldn't change
+  }
+
   size_t GetNodeCount() const {
     MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kReadLock);
     return world_.cur_element_count.load();
@@ -208,6 +221,99 @@ struct HnswlibAdapter {
     return out;
   }
 
+ public:
+  // Restore HNSW graph structure from serialized nodes with metadata
+  void RestoreFromNodes(const std::vector<HnswNodeData>& nodes, const HnswIndexMetadata& metadata) {
+    MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kWriteLock);
+    absl::WriterMutexLock resize_lock(&resize_mutex_);
+
+    if (nodes.empty()) {
+      return;
+    }
+
+    // Ensure we have enough capacity
+    size_t required_capacity = metadata.cur_element_count;
+    if (world_.max_elements_ < required_capacity) {
+      world_.resizeIndex(required_capacity);
+    }
+
+    // Restore each node - directly set up memory and fields
+    for (const auto& node : nodes) {
+      size_t internal_id = node.internal_id;
+
+      // Register label in lookup table
+      world_.label_lookup_[node.global_id] = internal_id;
+
+      // Set the level
+      world_.element_levels_[internal_id] = node.level;
+
+      // Clear level 0 memory and set label
+      memset(world_.data_level0_memory_ + internal_id * world_.size_data_per_element_ +
+                 world_.offsetLevel0_,
+             0, world_.size_data_per_element_);
+      world_.setExternalLabel(internal_id, node.global_id);
+
+      // Allocate upper layer links if needed
+      if (node.level > 0) {
+        world_.linkLists_[internal_id] =
+            (char*)mi_malloc(world_.size_links_per_element_ * node.level + 1);
+        memset(world_.linkLists_[internal_id], 0, world_.size_links_per_element_ * node.level + 1);
+      }
+
+      // Restore links for layer 0
+      if (!node.levels_links.empty()) {
+        auto* ll0 = world_.get_linklist0(internal_id);
+        world_.setListCount(ll0, node.levels_links[0].size());
+        auto* links0 = reinterpret_cast<uint32_t*>(ll0 + 1);
+        std::copy(node.levels_links[0].begin(), node.levels_links[0].end(), links0);
+      }
+
+      // Restore links for upper layers
+      for (int lvl = 1; lvl <= node.level && lvl < static_cast<int>(node.levels_links.size());
+           ++lvl) {
+        auto* ll = world_.get_linklist(internal_id, lvl);
+        world_.setListCount(ll, node.levels_links[lvl].size());
+        auto* links = reinterpret_cast<uint32_t*>(ll + 1);
+        std::copy(node.levels_links[lvl].begin(), node.levels_links[lvl].end(), links);
+      }
+    }
+
+    // Set the metadata for the graph
+    world_.cur_element_count.store(metadata.cur_element_count);
+    world_.maxlevel_ = metadata.maxlevel;
+    world_.enterpoint_node_ = metadata.enterpoint_node;
+
+    VLOG(1) << "Restored HNSW index with " << metadata.cur_element_count
+            << " nodes, maxlevel=" << metadata.maxlevel
+            << ", enterpoint=" << metadata.enterpoint_node;
+  }
+
+  // Update vector data for an existing node (used after RestoreFromNodes)
+  void UpdateVectorData(GlobalDocId id, const void* data) {
+    MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kWriteLock);
+
+    // Find the internal id for this label
+    auto it = world_.label_lookup_.find(id);
+    if (it == world_.label_lookup_.end()) {
+      LOG(WARNING) << "UpdateVectorData: label " << id << " not found in index";
+      return;
+    }
+
+    size_t internal_id = it->second;
+
+    // Copy/store the vector data based on copy_vector_ mode
+    if (world_.copy_vector_) {
+      // Owned mode: copy data into world's vector memory
+      char* data_ptr = world_.data_vector_memory_ + internal_id * world_.data_size_;
+      memcpy(data_ptr, data, world_.data_size_);
+    } else {
+      // Borrowed mode: store pointer to external data
+      char* ptr_location = world_.getDataPtrByInternalId(internal_id);
+      memcpy(ptr_location, &data, sizeof(void*));
+    }
+  }
+
+ private:
   HnswSpace space_;
   HierarchicalNSW<float> world_;
   absl::Mutex resize_mutex_;
@@ -264,6 +370,10 @@ HnswIndexMetadata HnswVectorIndex::GetMetadata() const {
   return adapter_->GetMetadata();
 }
 
+void HnswVectorIndex::SetMetadata(const HnswIndexMetadata& metadata) {
+  adapter_->SetMetadata(metadata);
+}
+
 size_t HnswVectorIndex::GetNodeCount() const {
   return adapter_->GetNodeCount();
 }
@@ -272,4 +382,28 @@ std::vector<HnswNodeData> HnswVectorIndex::GetNodesRange(size_t start, size_t en
   return adapter_->GetNodesRange(start, end);
 }
 
+void HnswVectorIndex::RestoreFromNodes(const std::vector<HnswNodeData>& nodes,
+                                       const HnswIndexMetadata& metadata) {
+  adapter_->RestoreFromNodes(nodes, metadata);
+  restored_ = true;
+}
+
+bool HnswVectorIndex::UpdateVectorData(GlobalDocId id, const DocumentAccessor& doc,
+                                       std::string_view field) {
+  auto vector_ptr = doc.GetVector(field, dim_);
+  if (!vector_ptr) {
+    return false;
+  }
+
+  const void* data = nullptr;
+  if (std::holds_alternative<OwnedFtVector>(*vector_ptr)) {
+    data = std::get<OwnedFtVector>(*vector_ptr).first.get();
+  } else {
+    data = std::get<BorrowedFtVector>(*vector_ptr);
+  }
+
+  adapter_->UpdateVectorData(id, data);
+  return true;
+}
+
 }  // namespace dfly::search

src/core/search/hnsw_index.h

@@ -64,18 +64,40 @@ class HnswVectorIndex {
   // Get metadata for serialization
   HnswIndexMetadata GetMetadata() const;
 
+  // Set metadata (used during restoration)
+  void SetMetadata(const HnswIndexMetadata& metadata);
+
   // Get total number of nodes in the index
   size_t GetNodeCount() const;
 
   // Get nodes in the specified range [start, end)
   // Returns vector of node data for serialization
   std::vector<HnswNodeData> GetNodesRange(size_t start, size_t end) const;
 
+  // Restore graph structure from serialized nodes with metadata
+  // This restores the HNSW graph links but NOT the vector data
+  // Vector data must be populated separately via UpdateVectorData
+  void RestoreFromNodes(const std::vector<HnswNodeData>& nodes, const HnswIndexMetadata& metadata);
+
+  // Update vector data for an existing node (used after RestoreFromNodes)
+  // This populates the vector data for a node that already has graph links
+  bool UpdateVectorData(GlobalDocId id, const DocumentAccessor& doc, std::string_view field);
+
+  // Mark index as restored from RDB (should use UpdateVectorData instead of Add)
+  void SetRestored(bool restored) {
+    restored_ = restored;
+  }
+
+  bool IsRestored() const {
+    return restored_;
+  }
+
  private:
   bool copy_vector_;
   size_t dim_;
   VectorSimilarity sim_;
   std::unique_ptr<HnswlibAdapter> adapter_;
+  bool restored_ = false;
 };
 
 }  // namespace dfly::search

src/server/rdb_load.cc

@@ -5,6 +5,7 @@
 #include "server/rdb_load.h"
 
 #include "absl/strings/escaping.h"
+#include "server/search/global_hnsw_index.h"
 #include "server/tiered_storage.h"
 
 extern "C" {
@@ -2213,8 +2214,8 @@ error_code RdbLoader::Load(io::Source* src) {
     }
 
     if (type == RDB_OPCODE_VECTOR_INDEX) {
-      // Stub: read and ignore HNSW vector index data
-      // Binary format: [index_name, elements_number,
+      // Read HNSW vector index graph data and restore directly
+      // Binary format: [index_key, elements_number,
       //   then for each node (little-endian):
       //     internal_id (4 bytes), global_id (8 bytes), level (4 bytes),
       //     for each level (0 to level): links_num (4 bytes) + links (4 bytes each)]
@@ -2224,26 +2225,93 @@ error_code RdbLoader::Load(io::Source* src) {
       uint64_t elements_number;
       SET_OR_RETURN(LoadLen(nullptr), elements_number);
 
+      // Only restore if shard count matches (GlobalDocId encodes shard_id)
+      bool should_restore =
+          shard_count_ > 0 && shard_set != nullptr && shard_count_ == shard_set->size();
+
+      // Extract index_name and field_name from index_key
+      size_t colon_pos = index_key.find(':');
+      string index_name = (colon_pos != string::npos) ? index_key.substr(0, colon_pos) : index_key;
+      string field_name = (colon_pos != string::npos) ? index_key.substr(colon_pos + 1) : "";
+
+      // Check if we can get the HNSW index (it should exist from FT.CREATE in aux)
+      auto hnsw_index = should_restore
+                            ? GlobalHnswIndexRegistry::Instance().Get(index_name, field_name)
+                            : nullptr;
+      if (should_restore && !hnsw_index) {
+        LOG(WARNING) << "HNSW index not found for restoration: " << index_key;
+        should_restore = false;
+      }
+
+      std::vector<search::HnswNodeData> nodes;
+      if (should_restore) {
+        nodes.reserve(elements_number);
+      }
+
       for (uint64_t elem = 0; elem < elements_number; ++elem) {
-        [[maybe_unused]] uint32_t internal_id;
+        uint32_t internal_id;
         SET_OR_RETURN(FetchInt<uint32_t>(), internal_id);
-        [[maybe_unused]] uint64_t global_id;
+        uint64_t global_id;
         SET_OR_RETURN(FetchInt<uint64_t>(), global_id);
         uint32_t level;
         SET_OR_RETURN(FetchInt<uint32_t>(), level);
 
+        search::HnswNodeData node;
+        if (should_restore) {
+          node.internal_id = internal_id;
+          node.global_id = global_id;
+          node.level = level;
+          node.levels_links.resize(level + 1);
+        }
+
         for (uint32_t lvl = 0; lvl <= level; ++lvl) {
           uint32_t links_num;
           SET_OR_RETURN(FetchInt<uint32_t>(), links_num);
+
+          if (should_restore) {
+            node.levels_links[lvl].reserve(links_num);
+          }
+
           for (uint32_t i = 0; i < links_num; ++i) {
-            [[maybe_unused]] uint32_t link;
+            uint32_t link;
             SET_OR_RETURN(FetchInt<uint32_t>(), link);
+            if (should_restore) {
+              node.levels_links[lvl].push_back(link);
+            }
           }
         }
+
+        if (should_restore) {
+          nodes.push_back(std::move(node));
+        }
       }
 
-      VLOG(2) << "Ignoring HNSW vector index: " << index_key
-              << " elements_number=" << elements_number;
+      if (should_restore && !nodes.empty()) {
+        // Get metadata - it was set via SetMetadataForIndex after FT.CREATE
+        search::HnswIndexMetadata metadata = hnsw_index->GetMetadata();
+
+        if (metadata.cur_element_count == 0) {
+          // Create default metadata from graph data
+          metadata.cur_element_count = nodes.size();
+          metadata.maxlevel = -1;
+          metadata.enterpoint_node = 0;
+          for (const auto& node : nodes) {
+            if (node.level > metadata.maxlevel) {
+              metadata.maxlevel = node.level;
+              metadata.enterpoint_node = node.internal_id;
+            }
+          }
+        }
+
+        // Restore the HNSW graph directly and mark as restored
+        hnsw_index->RestoreFromNodes(nodes, metadata);
+
+        LOG(INFO) << "Restored HNSW index " << index_key << " with " << nodes.size() << " nodes";
+      } else if (elements_number > 0) {
+        VLOG(2) << "Skipping HNSW vector index restore: " << index_key
+                << " elements_number=" << elements_number << " shard_count_=" << shard_count_
+                << " current_shards=" << (shard_set ? shard_set->size() : 0);
+      }
       continue;
     }
 
@@ -2975,7 +3043,6 @@ std::vector<std::string> RdbLoader::pending_synonym_cmds_;
 // Static synchronization for thread-safe search index creation
 base::SpinLock RdbLoader::search_index_mu_;
 absl::flat_hash_set<std::string> RdbLoader::created_search_indices_;
-
 std::vector<std::string> RdbLoader::TakePendingSynonymCommands() {
   std::vector<std::string> result;
   result.swap(pending_synonym_cmds_);
@@ -2985,6 +3052,8 @@ std::vector<std::string> RdbLoader::TakePendingSynonymCommands() {
 void RdbLoader::LoadSearchIndexDefFromAux(string&& def) {
   string index_name;
   string full_cmd;
+  string hnsw_field_name;
+  std::optional<search::HnswIndexMetadata> hnsw_meta;
 
   // Check if this is new JSON format (starts with '{') or old format ("index_name cmd")
   if (!def.empty() && def[0] == '{') {
@@ -2998,9 +3067,28 @@ void RdbLoader::LoadSearchIndexDefFromAux(string&& def) {
       const auto& json = *json_opt;
       index_name = json["name"].as<string>();
       string cmd = json["cmd"].as<string>();
-
-      // TODO: restore HNSW metadata from json["hnsw_metadata"] if present
-      // Currently we just restore the index definition, HNSW graph will be rebuilt
+      hnsw_field_name = json["field"].as<string>();
+
+      // Parse HNSW metadata if present
+      if (json.contains("hnsw_metadata")) {
+        const auto& meta = json["hnsw_metadata"];
+        search::HnswIndexMetadata m;
+        m.max_elements = meta["max_elements"].as<size_t>();
+        m.cur_element_count = meta["cur_element_count"].as<size_t>();
+        m.maxlevel = meta["maxlevel"].as<int>();
+        m.enterpoint_node = meta["enterpoint_node"].as<size_t>();
+        m.M = meta["M"].as<size_t>();
+        m.maxM = meta["maxM"].as<size_t>();
+        m.maxM0 = meta["maxM0"].as<size_t>();
+        m.ef_construction = meta["ef_construction"].as<size_t>();
+        m.mult = meta["mult"].as<double>();
+        hnsw_meta = m;
+
+        VLOG(1) << "Parsed HNSW metadata for index " << index_name << " field " << hnsw_field_name
+                << ": max_elements=" << m.max_elements
+                << " cur_element_count=" << m.cur_element_count << " maxlevel=" << m.maxlevel
+                << " M=" << m.M;
+      }
 
       full_cmd = absl::StrCat(index_name, " ", cmd);
     } catch (const std::exception& e) {
@@ -3032,6 +3120,13 @@ void RdbLoader::LoadSearchIndexDefFromAux(string&& def) {
   }
 
   LoadSearchCommandFromAux(service_, std::move(full_cmd), "FT.CREATE", "index definition");
+
+  // Store metadata on HNSW index after index creation (for later graph restoration)
+  if (hnsw_meta && !hnsw_field_name.empty()) {
+    if (auto index = GlobalHnswIndexRegistry::Instance().Get(index_name, hnsw_field_name); index) {
+      index->SetMetadata(*hnsw_meta);
+    }
+  }
 }
 
 void RdbLoader::LoadSearchSynonymsFromAux(string&& def) {
@@ -3055,7 +3150,7 @@ void RdbLoader::PerformPostLoad(Service* service, bool is_error) {
   if (is_error)
     return;
 
-  // Rebuild all search indices as only their definitions are extracted from the snapshot
+  // Rebuild all search indices - for restored HNSW indices, this will populate vectors
   shard_set->AwaitRunningOnShardQueue([](EngineShard* es) {
     OpArgs op_args{es, nullptr,
                    DbContext{&namespaces->GetDefaultNamespace(), 0, GetCurrentTimeMs()}};