feat(warp): add warp algorithm for multi-vector similarity search

- Add warp algorithm for multi-vector similarity search
- Add warp_parameter for configuration
- Add serialize/deserialize support
- Add test cases for warp index
- Optimize compute_maxsin_similarity by moving vec_indices and dists allocation outside the query loop
- Implement binary tree merge for parallel heap merging in SearchWithRequest
- Add parallel support for RangeSearch using atomic index distribution
- Use BatchInsertVector in add_one_doc for batch insertion
- Simplify serialization/deserialization using WriteVector/ReadVector for doc_offsets
- Fix Train to use total vector count instead of document count
- Pre-calculate and reserve vector capacity in Add to avoid frequent resizes
- Remove unused variables (num_elements, query_num)
- Fix nullptr pointer arithmetic for extra_info

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: jinjiabao.jjb <jinjiabao.jjb@antgroup.com>

Author: inabao

examples/cpp/110_index_warp.cpp

@@ -0,0 +1,117 @@
+
+// Copyright 2024-present the vsag project
+//
+// 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.
+
+#include <vsag/vsag.h>
+
+#include <iostream>
+#include <vector>
+
+int
+main(int argc, char** argv) {
+    vsag::init();
+
+    /******************* Prepare Multi-Vector Base Dataset *****************/
+    // In ColBERT-style retrieval, each document has multiple vectors (one per token)
+    // We'll create 100 documents, each with variable number of vectors (5-10)
+    int64_t num_docs = 100;
+    int64_t dim = 128;
+
+    std::vector<int64_t> ids(num_docs);
+    std::vector<uint32_t> vector_counts(num_docs);
+    std::vector<float> datas;
+    std::mt19937 rng(47);
+    std::uniform_real_distribution<float> distrib_real;
+    std::uniform_int_distribution<int> vec_count_dist(5, 10);
+
+    // Generate document IDs and variable vector counts
+    uint64_t total_vectors = 0;
+    for (int64_t i = 0; i < num_docs; ++i) {
+        ids[i] = i;
+        vector_counts[i] = vec_count_dist(rng);
+        total_vectors += vector_counts[i];
+    }
+
+    // Generate all vectors
+    datas.reserve(total_vectors * dim);
+    for (uint64_t i = 0; i < total_vectors * dim; ++i) {
+        datas.push_back(distrib_real(rng));
+    }
+
+    // Create dataset with VectorCounts for multi-vector support
+    auto base = vsag::Dataset::Make();
+    base->NumElements(num_docs)
+        ->Dim(dim)
+        ->Ids(ids.data())
+        ->Float32Vectors(datas.data())
+        ->VectorCounts(vector_counts.data())  // Specify number of vectors per document
+        ->Owner(false);
+
+    std::cout << "Created multi-vector dataset with " << num_docs << " documents" << std::endl;
+    std::cout << "Total vectors: " << total_vectors << " (avg " << (double)total_vectors / num_docs
+              << " vectors per doc)" << std::endl;
+
+    /******************* Create WARP Index *****************/
+    // WARP index for ColBERT-style maxsin similarity
+    std::string warp_build_parameters = R"(
+    {
+        "dtype": "float32",
+        "metric_type": "ip",
+        "dim": 128
+    }
+    )";
+    auto index = vsag::Factory::CreateIndex("warp", warp_build_parameters).value();
+
+    /******************* Build WARP Index *****************/
+    if (auto build_result = index->Build(base); build_result.has_value()) {
+        std::cout << "After Build(), Index WARP contains: " << index->GetNumElements()
+                  << " documents" << std::endl;
+    } else {
+        std::cerr << "Failed to build index: " << build_result.error().message << std::endl;
+        exit(-1);
+    }
+
+    /******************* Prepare Multi-Vector Query *****************/
+    // Query also has multiple vectors (representing query tokens)
+    uint32_t query_vec_count = 3;
+    std::vector<float> query_vectors(query_vec_count * dim);
+    for (uint32_t i = 0; i < query_vec_count * dim; ++i) {
+        query_vectors[i] = distrib_real(rng);
+    }
+
+    auto query = vsag::Dataset::Make();
+    query->NumElements(1)
+        ->Dim(dim)
+        ->Float32Vectors(query_vectors.data())
+        ->VectorCounts(&query_vec_count)  // Specify query has multiple vectors
+        ->Owner(false);
+
+    std::cout << "Query has " << query_vec_count << " vectors" << std::endl;
+
+    /******************* KnnSearch For WARP Index *****************/
+    // WARP performs maxsin similarity: for each query vector, find max similarity
+    // with any document vector, sum across query vectors
+    auto warp_search_parameters = R"({})";
+    int64_t topk = 5;
+    auto result = index->KnnSearch(query, topk, warp_search_parameters).value();
+
+    /******************* Print Search Result *****************/
+    std::cout << "Top-" << topk << " results (maxsin similarity): " << std::endl;
+    for (int64_t i = 0; i < result->GetDim(); ++i) {
+        std::cout << "  Document " << result->GetIds()[i]
+                  << ": score = " << result->GetDistances()[i] << std::endl;
+    }
+
+    return 0;
+}

examples/cpp/CMakeLists.txt

@@ -41,6 +41,9 @@ target_link_libraries (108_index_gno_imi vsag)
 add_executable (109_index_sindi 109_index_sindi.cpp)
 target_link_libraries (109_index_sindi vsag)
 
+add_executable (110_index_warp 110_index_warp.cpp)
+target_link_libraries (110_index_warp vsag)
+
 add_executable (201_custom_allocator 201_custom_allocator.cpp)
 target_link_libraries (201_custom_allocator vsag)
 

include/vsag/constants.h

@@ -25,6 +25,7 @@ extern const char* const INDEX_SPARSE;
 extern const char* const INDEX_SINDI;
 extern const char* const INDEX_BRUTE_FORCE;
 extern const char* const INDEX_IVF;
+extern const char* const INDEX_WARP;
 extern const char* const DIM;
 extern const char* const NUM_ELEMENTS;
 extern const char* const IDS;
@@ -36,6 +37,7 @@ extern const char* const ATTRIBUTE_SETS;
 extern const char* const DATASET_PATHS;
 extern const char* const EXTRA_INFOS;
 extern const char* const EXTRA_INFO_SIZE;
+extern const char* const VECTOR_COUNTS;
 
 extern const char* const HNSW_DATA;
 extern const char* const CONJUGATE_GRAPH_DATA;

include/vsag/dataset.h

@@ -354,6 +354,23 @@ class Dataset : public std::enable_shared_from_this<Dataset> {
      */
     virtual std::vector<std::string>
     GetStatistics(const std::vector<std::string>& stat_keys) const = 0;
+
+    /**
+     * @brief Sets the vector counts array for the dataset (for multi-vector documents).
+     *
+     * @param counts Pointer to the array of vector counts per document.
+     * @return DatasetPtr A shared pointer to the dataset with updated vector counts.
+     */
+    virtual DatasetPtr
+    VectorCounts(const uint32_t* counts) = 0;
+
+    /**
+     * @brief Retrieves the vector counts array of the dataset.
+     *
+     * @return const uint32_t* Pointer to the array of vector counts per document.
+     */
+    virtual const uint32_t*
+    GetVectorCounts() const = 0;
 };
 
 };  // namespace vsag

include/vsag/index.h

@@ -43,7 +43,7 @@ struct MergeUnit {
     IdMapFunction id_map_func = nullptr;
 };
 
-enum class IndexType { HNSW, DISKANN, HGRAPH, IVF, PYRAMID, BRUTEFORCE, SPARSE, SINDI };
+enum class IndexType { HNSW, DISKANN, HGRAPH, IVF, PYRAMID, BRUTEFORCE, SPARSE, SINDI, WARP };
 
 #define DATA_FLAG_FLOAT32_VECTOR 0x01
 #define DATA_FLAG_INT8_VECTOR 0x02