feat(index): add write function for Serialize

Signed-off-by: LHT129 <tianlan.lht@antgroup.com>

Author: LHT129

include/vsag/index.h

@@ -51,6 +51,10 @@ enum class IndexType { HNSW, DISKANN, HGRAPH, IVF, PYRAMID, BRUTEFORCE, SPARSE,
 #define DATA_FLAG_ATTRIBUTE 0x20
 #define DATA_FLAG_ID 0x40
 
+using OffsetType = uint64_t;
+using SizeType = uint64_t;
+using WriteFuncType = std::function<void(OffsetType, SizeType, const void*)>;
+
 class Index {
 public:
     // [basic methods]
@@ -637,6 +641,16 @@ class Index {
     [[nodiscard]] virtual tl::expected<BinarySet, Error>
     Serialize() const = 0;
 
+    /**
+      * @brief Serialize index by write function
+      *
+      * @param write_func is a function to write serialized index
+      */
+    [[nodiscard]] virtual tl::expected<void, Error>
+    Serialize(WriteFuncType write_func) const {
+        throw std::runtime_error("Index doesn't support Serialize with write function");
+    }
+
     /**
       * @brief Deserialize index from a set of byte array. Causing exception if this index is not empty
       *

include/vsag/index_features.h

@@ -35,6 +35,7 @@ enum IndexFeature {
     SUPPORT_METRIC_TYPE_COSINE,          /**< Supports cosine metric type */
     SUPPORT_SERIALIZE_FILE,              /**< Supports serialization to a file */
     SUPPORT_SERIALIZE_BINARY_SET,        /**< Supports serialization to a binary set */
+    SUPPORT_SERIALIZE_WRITE_FUNC,        /**< Supports serialization to a write function */
     SUPPORT_DESERIALIZE_FILE,            /**< Supports deserialization from a file */
     SUPPORT_DESERIALIZE_BINARY_SET,      /**< Supports deserialization from a binary set */
     SUPPORT_DESERIALIZE_READER_SET,      /**< Supports deserialization from a reader set */

src/algorithm/brute_force.cpp

@@ -422,6 +422,7 @@ BruteForce::InitFeatures() {
         IndexFeature::SUPPORT_DESERIALIZE_READER_SET,
         IndexFeature::SUPPORT_SERIALIZE_BINARY_SET,
         IndexFeature::SUPPORT_SERIALIZE_FILE,
+        IndexFeature::SUPPORT_SERIALIZE_WRITE_FUNC,
     });
 
     // others

src/algorithm/hgraph.cpp

@@ -1116,6 +1116,7 @@ HGraph::InitFeatures() {
         IndexFeature::SUPPORT_DESERIALIZE_READER_SET,
         IndexFeature::SUPPORT_SERIALIZE_BINARY_SET,
         IndexFeature::SUPPORT_SERIALIZE_FILE,
+        IndexFeature::SUPPORT_SERIALIZE_WRITE_FUNC,
     });
     // other
     this->index_feature_list_->SetFeatures({

src/algorithm/inner_index_interface.cpp

@@ -143,6 +143,16 @@ InnerIndexInterface::Serialize() const {
     return bs;
 }
 
+void
+InnerIndexInterface::Serialize(const WriteFuncType& write_func) const {
+    std::string time_record_name = this->GetName() + " Serialize";
+    SlowTaskTimer t(time_record_name);
+
+    uint64_t num_bytes = this->CalSerializeSize();
+    WriteFuncStreamWriter writer(write_func, 0);
+    this->Serialize(writer);
+}
+
 void
 InnerIndexInterface::Deserialize(const BinarySet& binary_set) {
     std::string time_record_name = this->GetName() + " Deserialize";

src/algorithm/inner_index_interface.h

@@ -329,6 +329,9 @@ class InnerIndexInterface {
     virtual void
     Serialize(std::ostream& out_stream) const;
 
+    virtual void
+    Serialize(const WriteFuncType& write_func) const;
+
     virtual void
     Serialize(StreamWriter& writer) const = 0;
 

src/algorithm/ivf.cpp

@@ -296,6 +296,7 @@ IVF::InitFeatures() {
         IndexFeature::SUPPORT_DESERIALIZE_READER_SET,
         IndexFeature::SUPPORT_SERIALIZE_BINARY_SET,
         IndexFeature::SUPPORT_SERIALIZE_FILE,
+        IndexFeature::SUPPORT_SERIALIZE_WRITE_FUNC,
     });
 
     auto name = this->bucket_->GetQuantizerName();

src/index/index_impl.h

@@ -376,6 +376,11 @@ class IndexImpl : public Index {
         SAFE_CALL(return this->inner_index_->Serialize());
     }
 
+    [[nodiscard]] tl::expected<void, Error>
+    Serialize(WriteFuncType write_func) const override {
+        SAFE_CALL(this->inner_index_->Serialize(write_func));
+    }
+
     tl::expected<void, Error>
     Serialize(std::ostream& out_stream) override {
         SAFE_CALL(this->inner_index_->Serialize(out_stream));

tests/test_hgraph.cpp

@@ -678,6 +678,8 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::HGraphTestIndex,
     TestIndex::TestSerializeBinarySet(index, index2, dataset, search_param, true);
     index2 = TestIndex::TestFactory(name, param, true);
     TestIndex::TestSerializeReaderSet(index, index2, dataset, search_param, name, true);
+    index2 = TestIndex::TestFactory(name, param, true);
+    TestIndex::TestSerializeWriteFunc(index, index2, dataset, search_param, true);
     vsag::Options::Instance().set_block_size_limit(origin_size);
 }
 static void
@@ -1413,6 +1415,8 @@ TestHGraphSerialize(const fixtures::HGraphTestIndexPtr& test_index,
                 index2 = TestIndex::TestFactory(test_index->name, param, true);
                 TestIndex::TestSerializeReaderSet(
                     index, index2, dataset, search_param, test_index->name, true);
+                index2 = TestIndex::TestFactory(test_index->name, param, true);
+                TestIndex::TestSerializeWriteFunc(index, index2, dataset, search_param, true);
                 vsag::Options::Instance().set_block_size_limit(origin_size);
             }
         }

tests/test_index.cpp

@@ -1086,6 +1086,58 @@ TestIndex::TestSerializeReaderSet(const IndexPtr& index_from,
     }
 }
 
+void
+TestIndex::TestSerializeWriteFunc(const IndexPtr& index_from,
+                                  const IndexPtr& index_to,
+                                  const TestDatasetPtr& dataset,
+                                  const std::string& search_param,
+                                  bool expected_success) {
+    if (not index_from->CheckFeature(vsag::SUPPORT_SERIALIZE_WRITE_FUNC)) {
+        return;
+    }
+    auto dir = fixtures::TempDir("serialize");
+    auto path = dir.GenerateRandomFile();
+    std::ofstream outfile(path, std::ios::out | std::ios::binary);
+    vsag::WriteFuncType write_func =
+        [&outfile](vsag::OffsetType offset, vsag::SizeType size, const void* data) -> void {
+        outfile.seekp(offset);
+        outfile.write(reinterpret_cast<const char*>(data), size);
+    };
+    auto serialize_index = index_from->Serialize(write_func);
+    REQUIRE(serialize_index.has_value() == expected_success);
+    outfile.close();
+
+    std::ifstream infile(path, std::ios::in | std::ios::binary);
+    auto deserialize_index = index_to->Deserialize(infile);
+    REQUIRE(deserialize_index.has_value() == expected_success);
+    infile.close();
+    if (index_to->GetNumElements() == 0) {
+        return;
+    }
+
+    const auto& queries = dataset->query_;
+    auto query_count = queries->GetNumElements();
+    auto dim = queries->GetDim();
+    auto topk = 10;
+    for (auto i = 0; i < query_count; ++i) {
+        auto query = vsag::Dataset::Make();
+        query->NumElements(1)
+            ->Dim(dim)
+            ->Paths(queries->GetPaths() + i)
+            ->SparseVectors(queries->GetSparseVectors() + i)
+            ->Float32Vectors(queries->GetFloat32Vectors() + i * dim)
+            ->Owner(false);
+        auto res_from = index_from->KnnSearch(query, topk, search_param);
+        auto res_to = index_to->KnnSearch(query, topk, search_param);
+        REQUIRE(res_from.has_value());
+        REQUIRE(res_to.has_value());
+        REQUIRE(res_from.value()->GetDim() == res_to.value()->GetDim());
+        for (auto j = 0; j < topk; ++j) {
+            REQUIRE(res_to.value()->GetIds()[j] == res_from.value()->GetIds()[j]);
+        }
+    }
+}
+
 void
 TestIndex::TestConcurrentAddSearch(const TestIndex::IndexPtr& index,
                                    const TestDatasetPtr& dataset,