support index serialization self-description

Signed-off-by: Xiangyu Wang <[email protected]>

Author: wxyucs

examples/cpp/402_persistent_streaming.cpp

@@ -25,6 +25,7 @@
 #include <unordered_set>
 
 #include "vsag/binaryset.h"
+#include "vsag/options.h"
 
 int
 main(int32_t argc, char** argv) {
@@ -52,14 +53,15 @@ main(int32_t argc, char** argv) {
         "dtype": "float32",
         "metric_type": "l2",
         "dim": 128,
-        "hnsw": {
+        "index_param": {
+            "base_quantization_type": "fp32",
             "max_degree": 16,
             "ef_construction": 100
         }
     }
     )";
     vsag::IndexPtr index = nullptr;
-    if (auto create_index = engine.CreateIndex("hnsw", index_paramesters);
+    if (auto create_index = engine.CreateIndex("hgraph", index_paramesters);
         not create_index.has_value()) {
         std::cout << "create index failed: " << create_index.error().message << std::endl;
         abort();
@@ -76,6 +78,7 @@ main(int32_t argc, char** argv) {
     std::cout << "index contains vectors: " << index->GetNumElements() << std::endl;
 
     /******************* Save Index to OStream *****************/
+    vsag::Options::Instance().set_new_version(true);
     std::ofstream out_stream("/tmp/vsag-persistent-streaming.index");
     auto serialize_result = index->Serialize(out_stream);
     out_stream.close();
@@ -86,13 +89,14 @@ main(int32_t argc, char** argv) {
 
     /******************* Load Index from IStream *****************/
     index = nullptr;
-    if (auto create_index = engine.CreateIndex("hnsw", index_paramesters);
+    if (auto create_index = engine.CreateIndex("hgraph", index_paramesters);
         not create_index.has_value()) {
         std::cout << "create index failed: " << create_index.error().message << std::endl;
         abort();
     } else {
         index = *create_index;
     }
+
     std::ifstream in_stream("/tmp/vsag-persistent-streaming.index");
     if (auto deserialize = index->Deserialize(in_stream); not deserialize.has_value()) {
         std::cerr << "load index failed: " << deserialize.error().message << std::endl;
@@ -109,7 +113,7 @@ main(int32_t argc, char** argv) {
     query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector)->Owner(false);
     auto search_parameters = R"(
     {
-        "hnsw": {
+        "hgraph": {
             "ef_search": 100
         }
     }

include/vsag/constants.h

@@ -165,4 +165,8 @@ extern const char* const IVF_PRECISE_QUANTIZATION_TYPE;
 extern const char* const IVF_PRECISE_IO_TYPE;
 extern const char* const IVF_PRECISE_FILE_PATH;
 
+// serialization
+extern const char* const SERIAL_MAGIC_BEGIN;
+extern const char* const SERIAL_MAGIC_END;
+
 }  // namespace vsag

include/vsag/options.h

@@ -19,6 +19,7 @@
 #include <mutex>
 #include <string>
 
+#include "utils/function_exists_check.h"
 #include "vsag/allocator.h"
 #include "vsag/logger.h"
 
@@ -43,6 +44,20 @@ class Options {
     static Options&
     Instance();
 
+public:
+    inline bool
+    new_version() const {
+        return new_version_;
+    }
+
+    inline void
+    set_new_version(bool new_version) {
+        new_version_ = new_version;
+    }
+
+private:
+    bool new_version_ = true;
+
 public:
     /**
      * @brief Gets the number of threads for IO operations.

src/algorithm/hgraph.cpp

@@ -27,8 +27,13 @@
 #include "impl/odescent_graph_builder.h"
 #include "impl/pruning_strategy.h"
 #include "index/iterator_filter.h"
+#include "logger.h"
+#include "serialization.h"
+#include "stream_reader.h"
+#include "typing.h"
 #include "utils/standard_heap.h"
 #include "utils/util_functions.h"
+#include "vsag/options.h"
 
 namespace vsag {
 
@@ -607,7 +612,7 @@ HGraph::RangeSearch(const DatasetPtr& query,
 }
 
 void
-HGraph::serialize_basic_info(StreamWriter& writer) const {
+HGraph::serialize_basic_info_v0_14(StreamWriter& writer) const {
     StreamWriter::WriteObj(writer, this->use_reorder_);
     StreamWriter::WriteObj(writer, this->dim_);
     StreamWriter::WriteObj(writer, this->metric_);
@@ -629,12 +634,120 @@ HGraph::serialize_basic_info(StreamWriter& writer) const {
     }
 }
 
+void
+HGraph::deserialize_basic_info_v0_14(StreamReader& reader) {
+    StreamReader::ReadObj(reader, this->use_reorder_);
+    StreamReader::ReadObj(reader, this->dim_);
+    StreamReader::ReadObj(reader, this->metric_);
+    uint64_t max_level;
+    StreamReader::ReadObj(reader, max_level);
+    for (uint64_t i = 0; i < max_level; ++i) {
+        this->route_graphs_.emplace_back(this->generate_one_route_graph());
+    }
+    StreamReader::ReadObj(reader, this->entry_point_id_);
+    StreamReader::ReadObj(reader, this->ef_construct_);
+    StreamReader::ReadObj(reader, this->mult_);
+    InnerIdType capacity;
+    StreamReader::ReadObj(reader, capacity);
+    this->max_capacity_.store(capacity);
+    StreamReader::ReadVector(reader, this->label_table_->label_table_);
+
+    uint64_t size;
+    StreamReader::ReadObj(reader, size);
+    for (uint64_t i = 0; i < size; ++i) {
+        LabelType key;
+        StreamReader::ReadObj(reader, key);
+        InnerIdType value;
+        StreamReader::ReadObj(reader, value);
+        this->label_table_->label_remap_.emplace(key, value);
+    }
+}
+
+#define TO_JSON(json_obj, var) json_obj[#var] = this->var##_;
+
+#define TO_JSON_BASE64(json_obj, var) json_obj[#var] = base64_encode_obj(this->var##_);
+
+#define TO_JSON_ATOMIC(json_obj, var) json_obj[#var] = this->var##_.load();
+
+JsonType
+HGraph::serialize_basic_info() const {
+    JsonType jsonify_basic_info;
+    TO_JSON(jsonify_basic_info, use_reorder);
+    TO_JSON(jsonify_basic_info, dim);
+    TO_JSON(jsonify_basic_info, metric);
+    TO_JSON(jsonify_basic_info, entry_point_id);
+    TO_JSON(jsonify_basic_info, ef_construct);
+    // logger::debug("mult: {}", this->mult_);
+    TO_JSON_BASE64(jsonify_basic_info, mult);
+    TO_JSON_ATOMIC(jsonify_basic_info, max_capacity);
+    jsonify_basic_info["max_level"] = this->route_graphs_.size();
+
+    return jsonify_basic_info;
+}
+
+#define FROM_JSON(json_obj, var) this->var##_ = (json_obj)[#var];
+
+#define FROM_JSON_BASE64(json_obj, var) base64_decode_obj((json_obj)[#var], this->var##_);
+
+#define FROM_JSON_ATOMIC(json_obj, var) this->var##_.store((json_obj)[#var]);
+
+void
+HGraph::deserialize_basic_info(JsonType jsonify_basic_info) {
+    FROM_JSON(jsonify_basic_info, use_reorder);
+    FROM_JSON(jsonify_basic_info, dim);
+    FROM_JSON(jsonify_basic_info, metric);
+    FROM_JSON(jsonify_basic_info, entry_point_id);
+    FROM_JSON(jsonify_basic_info, ef_construct);
+    FROM_JSON_BASE64(jsonify_basic_info, mult);
+    // logger::debug("mult: {}", this->mult_);
+    FROM_JSON_ATOMIC(jsonify_basic_info, max_capacity);
+
+    uint64_t max_level = jsonify_basic_info["max_level"];
+    for (uint64_t i = 0; i < max_level; ++i) {
+        this->route_graphs_.emplace_back(this->generate_one_route_graph());
+    }
+}
+
+void
+HGraph::serialize_label_info(StreamWriter& writer) const {
+    StreamWriter::WriteVector(writer, this->label_table_->label_table_);
+    uint64_t size = this->label_table_->label_remap_.size();
+    StreamWriter::WriteObj(writer, size);
+    for (const auto& pair : this->label_table_->label_remap_) {
+        auto key = pair.first;
+        StreamWriter::WriteObj(writer, key);
+        StreamWriter::WriteObj(writer, pair.second);
+    }
+}
+
+void
+HGraph::deserialize_label_info(StreamReader& reader) const {
+    StreamReader::ReadVector(reader, this->label_table_->label_table_);
+    uint64_t size;
+    StreamReader::ReadObj(reader, size);
+    for (uint64_t i = 0; i < size; ++i) {
+        LabelType key;
+        StreamReader::ReadObj(reader, key);
+        InnerIdType value;
+        StreamReader::ReadObj(reader, value);
+        this->label_table_->label_remap_.emplace(key, value);
+    }
+}
+
 void
 HGraph::Serialize(StreamWriter& writer) const {
     if (this->ignore_reorder_) {
         this->use_reorder_ = false;
     }
-    this->serialize_basic_info(writer);
+
+    // basic info moved to metadata since version 0.15
+    // only for test
+    if (Options::Instance().new_version()) {
+        this->serialize_label_info(writer);
+    } else {
+        this->serialize_basic_info_v0_14(writer);
+    }
+
     this->basic_flatten_codes_->Serialize(writer);
     this->bottom_graph_->Serialize(writer);
     if (this->use_reorder_) {
@@ -646,11 +759,32 @@ HGraph::Serialize(StreamWriter& writer) const {
     if (this->extra_info_size_ > 0 && this->extra_infos_ != nullptr) {
         this->extra_infos_->Serialize(writer);
     }
+
+    // serialize footer (introduce since v0.15)
+    if (Options::Instance().new_version()) {
+        auto metadata = std::make_shared<Metadata>();
+        auto jsonify_basic_info = this->serialize_basic_info();
+        metadata->Set("basic_info", jsonify_basic_info);
+        logger::debug(jsonify_basic_info.dump());
+        auto footer = std::make_shared<Footer>(metadata);
+        footer->Write(writer);
+    }
 }
 
 void
 HGraph::Deserialize(StreamReader& reader) {
-    this->deserialize_basic_info(reader);
+    // try to deserialize footer (only in new version)
+    auto footer = Footer::Parse(reader);
+    if (footer != nullptr) {
+        logger::debug("parse with new version format");
+        auto metadata = footer->GetMetadata();
+        this->deserialize_basic_info(metadata->Get("basic_info"));
+        this->deserialize_label_info(reader);
+    } else {
+        logger::debug("parse with v0.14 version format");
+        this->deserialize_basic_info_v0_14(reader);
+    }
+
     this->basic_flatten_codes_->Deserialize(reader);
     this->bottom_graph_->Deserialize(reader);
     if (this->use_reorder_) {
@@ -676,35 +810,6 @@ HGraph::Deserialize(StreamReader& reader) {
     }
 }
 
-void
-HGraph::deserialize_basic_info(StreamReader& reader) {
-    StreamReader::ReadObj(reader, this->use_reorder_);
-    StreamReader::ReadObj(reader, this->dim_);
-    StreamReader::ReadObj(reader, this->metric_);
-    uint64_t max_level;
-    StreamReader::ReadObj(reader, max_level);
-    for (uint64_t i = 0; i < max_level; ++i) {
-        this->route_graphs_.emplace_back(this->generate_one_route_graph());
-    }
-    StreamReader::ReadObj(reader, this->entry_point_id_);
-    StreamReader::ReadObj(reader, this->ef_construct_);
-    StreamReader::ReadObj(reader, this->mult_);
-    InnerIdType capacity;
-    StreamReader::ReadObj(reader, capacity);
-    this->max_capacity_.store(capacity);
-    StreamReader::ReadVector(reader, this->label_table_->label_table_);
-
-    uint64_t size;
-    StreamReader::ReadObj(reader, size);
-    for (uint64_t i = 0; i < size; ++i) {
-        LabelType key;
-        StreamReader::ReadObj(reader, key);
-        InnerIdType value;
-        StreamReader::ReadObj(reader, value);
-        this->label_table_->label_remap_.emplace(key, value);
-    }
-}
-
 float
 HGraph::CalcDistanceById(const float* query, int64_t id) const {
     auto flat = this->basic_flatten_codes_;

src/algorithm/hgraph.h

@@ -198,11 +198,25 @@ class HGraph : public InnerIndexInterface {
                      InnerSearchParam& inner_search_param,
                      IteratorFilterContext* iter_ctx) const;
 
+    // since v0.12
     void
-    serialize_basic_info(StreamWriter& writer) const;
+    serialize_basic_info_v0_14(StreamWriter& writer) const;
 
     void
-    deserialize_basic_info(StreamReader& reader);
+    deserialize_basic_info_v0_14(StreamReader& reader);
+
+    // since v0.15
+    JsonType
+    serialize_basic_info() const;
+
+    void
+    deserialize_basic_info(JsonType jsonify_basic_info);
+
+    void
+    serialize_label_info(StreamWriter& writer) const;
+
+    void
+    deserialize_label_info(StreamReader& reader) const;
 
     void
     reorder(const void* query,

src/algorithm/inner_index_interface.cpp

@@ -140,7 +140,7 @@ InnerIndexInterface::Deserialize(const BinarySet& binary_set) {
 
     try {
         uint64_t cursor = 0;
-        auto reader = ReadFuncStreamReader(func, cursor);
+        auto reader = ReadFuncStreamReader(func, cursor, b.size);
         this->Deserialize(reader);
     } catch (const std::runtime_error& e) {
         throw VsagException(ErrorType::READ_ERROR, "failed to Deserialize: ", e.what());
@@ -161,7 +161,7 @@ InnerIndexInterface::Deserialize(const ReaderSet& reader_set) {
             index_reader->Read(offset, len, dest);
         };
         uint64_t cursor = 0;
-        auto reader = ReadFuncStreamReader(func, cursor);
+        auto reader = ReadFuncStreamReader(func, cursor, index_reader->Size());
         this->Deserialize(reader);
         return;
     } catch (const std::bad_alloc& e) {

src/constants.cpp

@@ -167,4 +167,7 @@ const char* const IVF_BASE_FILE_PATH = "base_file_path";
 const char* const IVF_PRECISE_QUANTIZATION_TYPE = "precise_quantization_type";
 const char* const IVF_PRECISE_IO_TYPE = "precise_io_type";
 const char* const IVF_PRECISE_FILE_PATH = "precise_file_path";
+
+const char* const SERIAL_MAGIC_BEGIN = "vsag0000";
+const char* const SERIAL_MAGIC_END = "0000gasv";
 };  // namespace vsag

src/impl/conjugate_graph.cpp

@@ -183,7 +183,7 @@ ConjugateGraph::Deserialize(const Binary& binary) {
     };
 
     int64_t cursor = 0;
-    ReadFuncStreamReader reader(func, cursor);
+    ReadFuncStreamReader reader(func, cursor, binary.size);
     BufferStreamReader buffer_reader(&reader, binary.size, allocator_);
     return this->Deserialize(buffer_reader);
 }

src/index/hnsw.cpp

@@ -563,7 +563,7 @@ HNSW::deserialize(const BinarySet& binary_set) {
     try {
         std::unique_lock lock(rw_mutex_);
         int64_t cursor = 0;
-        ReadFuncStreamReader reader(func, cursor);
+        ReadFuncStreamReader reader(func, cursor, b.size);
         BufferStreamReader buffer_reader(&reader, b.size, allocator_.get());
         alg_hnsw_->loadIndex(buffer_reader, this->space_.get());
         if (use_conjugate_graph_) {
@@ -617,7 +617,7 @@ HNSW::deserialize(const ReaderSet& reader_set) {
         std::unique_lock lock(rw_mutex_);
 
         int64_t cursor = 0;
-        ReadFuncStreamReader reader(func, cursor);
+        ReadFuncStreamReader reader(func, cursor, hnsw_data->Size());
         BufferStreamReader buffer_reader(&reader, hnsw_data->Size(), allocator_.get());
         alg_hnsw_->loadIndex(buffer_reader, this->space_.get());
     } catch (const std::runtime_error& e) {