Support gno-imi partition strategy

examples/cpp/106_index_ivf.cpp

@@ -49,7 +49,9 @@ main(int argc, char** argv) {
         "dim": 128,
         "index_param": {
             "buckets_count": 50,
-            "base_quantization_type": "fp32"
+            "base_quantization_type": "fp32",
+            "partition_strategy_type": "ivf",
+            "ivf_train_type": "kmeans"
         }
     }
     )";

examples/cpp/108_index_gno_imi.cpp

@@ -0,0 +1,107 @@
+
+// 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 <fstream>
+#include <iostream>
+
+int
+main(int argc, char** argv) {
+    vsag::init();
+
+    /******************* Prepare Base Dataset *****************/
+    int64_t num_vectors = 10000;
+    int64_t dim = 128;
+    std::vector<int64_t> ids(num_vectors);
+    std::vector<float> datas(num_vectors * dim);
+    std::mt19937 rng(47);
+    std::uniform_real_distribution<float> distrib_real;
+    for (int64_t i = 0; i < num_vectors; ++i) {
+        ids[i] = i;
+    }
+    for (int64_t i = 0; i < dim * num_vectors; ++i) {
+        datas[i] = distrib_real(rng);
+    }
+
+    auto base = vsag::Dataset::Make();
+    base->NumElements(num_vectors)
+        ->Dim(dim)
+        ->Ids(ids.data())
+        ->Float32Vectors(datas.data())
+        ->Owner(false);
+
+    /******************* Create IVF Index *****************/
+    std::string ivf_build_params = R"(
+    {
+        "dtype": "float32",
+        "metric_type": "l2",
+        "dim": 128,
+        "index_param": {
+            "base_quantization_type": "fp32",
+            "partition_strategy_type": "gno_imi",
+            "ivf_train_type": "kmeans",
+            "first_order_buckets_count": 10,
+            "second_order_buckets_count": 10
+        }
+    })";
+    auto index = vsag::Factory::CreateIndex("ivf", ivf_build_params).value();
+
+    /******************* Build IVF Index *****************/
+
+    auto path = "example_gno_imi.index";
+    if (auto build_result = index->Build(base); build_result.has_value()) {
+        std::ofstream outfile(path, std::ios::out | std::ios::binary);
+        auto result = index->Serialize(outfile);
+        outfile.close();
+        std::cout << "After Build(), Index IVF contains: " << index->GetNumElements() << std::endl;
+    } else if (build_result.error().type == vsag::ErrorType::INTERNAL_ERROR) {
+        std::cerr << "Failed to build index: internalError" << std::endl;
+        exit(-1);
+    }
+
+    std::ifstream infile(path, std::ios::binary);
+    auto index2 = vsag::Factory::CreateIndex("ivf", ivf_build_params).value();
+    index2->Deserialize(infile);
+    infile.close();
+
+    /******************* Prepare Query Dataset *****************/
+    std::vector<float> query_vector(dim);
+    for (int64_t i = 0; i < dim; ++i) {
+        query_vector[i] = distrib_real(rng);
+    }
+
+    auto query = vsag::Dataset::Make();
+    query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector.data())->Owner(false);
+
+    /******************* KnnSearch For IVF Index *****************/
+    auto ivf_search_parameters = R"(
+    {
+        "ivf": {
+            "scan_buckets_count": 20,
+            "first_order_scan_ratio": 0.8
+        }
+    })";
+    int64_t topk = 10;
+    auto result = index->KnnSearch(query, topk, ivf_search_parameters).value();
+    auto result2 = index2->KnnSearch(query, topk, ivf_search_parameters).value();
+    /******************* Print Search Result *****************/
+    std::cout << "results: " << std::endl;
+    for (int64_t i = 0; i < result->GetDim(); ++i) {
+        std::cout << result->GetIds()[i] << ": " << result->GetDistances()[i] << std::endl;
+        std::cout << result2->GetIds()[i] << ": " << result2->GetDistances()[i] << std::endl;
+    }
+    return 0;
+}

examples/cpp/CMakeLists.txt

@@ -19,6 +19,9 @@ target_link_libraries (106_index_ivf vsag)
 add_executable (107_index_pyramid 107_index_pyramid.cpp)
 target_link_libraries (107_index_pyramid vsag)
 
+add_executable (108_index_gno_imi 108_index_gno_imi.cpp)
+target_link_libraries (108_index_gno_imi vsag)
+
 add_executable (201_custom_allocator 201_custom_allocator.cpp)
 target_link_libraries (201_custom_allocator vsag)
 

include/vsag/constants.h

@@ -166,4 +166,6 @@ extern const char* const IVF_PRECISE_QUANTIZATION_TYPE;
 extern const char* const IVF_PRECISE_IO_TYPE;
 extern const char* const IVF_PRECISE_FILE_PATH;
 
+extern const char* const GNO_IMI_FIRST_ORDER_BUCKETS_COUNT;
+extern const char* const GNO_IMI_SECOND_ORDER_BUCKETS_COUNT;
 }  // namespace vsag

src/algorithm/ivf.cpp

@@ -15,9 +15,12 @@
 
 #include "ivf.h"
 
+#include <set>
+
 #include "impl/basic_searcher.h"
 #include "index/index_impl.h"
 #include "inner_string_params.h"
+#include "ivf_partition/gno_imi_partition.h"
 #include "ivf_partition/ivf_nearest_partition.h"
 #include "utils/standard_heap.h"
 #include "utils/util_functions.h"
@@ -28,7 +31,6 @@ static constexpr const char* IVF_PARAMS_TEMPLATE =
     R"(
     {
         "type": "{INDEX_TYPE_IVF}",
-        "{IVF_TRAIN_TYPE_KEY}": "{IVF_TRAIN_TYPE_KMEANS}",
         "{BUCKET_PARAMS_KEY}": {
             "{IO_PARAMS_KEY}": {
                 "{IO_TYPE_KEY}": "{IO_TYPE_VALUE_BLOCK_MEMORY_IO}"
@@ -43,6 +45,15 @@ static constexpr const char* IVF_PARAMS_TEMPLATE =
             "{BUCKETS_COUNT_KEY}": 10,
             "{BUCKET_USE_RESIDUAL}": false
         },
+        "{IVF_PARTITION_STRATEGY_PARAMS_KEY}": {
+            "{IVF_PARTITION_STRATEGY_TYPE_KEY}": "{IVF_PARTITION_STRATEGY_TYPE_NEAREST}",
+            "{IVF_TRAIN_TYPE_KEY}": "{IVF_TRAIN_TYPE_KMEANS}",
+            "{IVF_PARTITION_STRATEGY_TYPE_GNO_IMI}": {
+                "{GNO_IMI_FIRST_ORDER_BUCKETS_COUNT_KEY}": 10,
+                "{GNO_IMI_SECOND_ORDER_BUCKETS_COUNT_KEY}": 10
+            }
+        },
+        "{BUCKET_PER_DATA_KEY}": 1,
         "{IVF_USE_REORDER_KEY}": false,
         "{IVF_PRECISE_CODES_KEY}": {
             "{IO_PARAMS_KEY}": {
@@ -83,7 +94,27 @@ IVF::CheckAndMappingExternalParam(const JsonType& external_param,
         },
         {
             IVF_TRAIN_TYPE,
-            {IVF_TRAIN_TYPE_KEY},
+            {IVF_PARTITION_STRATEGY_PARAMS_KEY, IVF_TRAIN_TYPE_KEY},
+        },
+        {
+            IVF_PARTITION_STRATEGY_TYPE_KEY,
+            {IVF_PARTITION_STRATEGY_PARAMS_KEY, IVF_PARTITION_STRATEGY_TYPE_KEY},
+        },
+        {
+            GNO_IMI_FIRST_ORDER_BUCKETS_COUNT,
+            {IVF_PARTITION_STRATEGY_PARAMS_KEY,
+             IVF_PARTITION_STRATEGY_TYPE_GNO_IMI,
+             GNO_IMI_FIRST_ORDER_BUCKETS_COUNT_KEY},
+        },
+        {
+            GNO_IMI_SECOND_ORDER_BUCKETS_COUNT,
+            {IVF_PARTITION_STRATEGY_PARAMS_KEY,
+             IVF_PARTITION_STRATEGY_TYPE_GNO_IMI,
+             GNO_IMI_SECOND_ORDER_BUCKETS_COUNT_KEY},
+        },
+        {
+            BUCKET_PER_DATA_KEY,
+            {BUCKET_PER_DATA_KEY},
         },
         {
             IVF_USE_REORDER,
@@ -116,13 +147,20 @@ IVF::CheckAndMappingExternalParam(const JsonType& external_param,
 }
 
 IVF::IVF(const IVFParameterPtr& param, const IndexCommonParam& common_param)
-    : InnerIndexInterface(param, common_param) {
+    : InnerIndexInterface(param, common_param), buckets_per_data_(param->buckets_per_data) {
     this->bucket_ = BucketInterface::MakeInstance(param->bucket_param, common_param);
     if (this->bucket_ == nullptr) {
         throw VsagException(ErrorType::INTERNAL_ERROR, "bucket init error");
     }
-    this->partition_strategy_ = std::make_shared<IVFNearestPartition>(
-        bucket_->bucket_count_, common_param, param->partition_train_type);
+    if (param->ivf_partition_strategy_parameter->partition_strategy_type ==
+        IVFPartitionStrategyType::IVF) {
+        this->partition_strategy_ = std::make_shared<IVFNearestPartition>(
+            bucket_->bucket_count_, common_param, param->ivf_partition_strategy_parameter);
+    } else if (param->ivf_partition_strategy_parameter->partition_strategy_type ==
+               IVFPartitionStrategyType::GNO_IMI) {
+        this->partition_strategy_ = std::make_shared<GNOIMIPartition>(
+            common_param, param->ivf_partition_strategy_parameter);
+    }
     this->use_reorder_ = param->use_reorder;
     if (this->use_reorder_) {
         this->reorder_codes_ = FlattenInterface::MakeInstance(param->flatten_param, common_param);
@@ -228,26 +266,34 @@ IVF::Add(const DatasetPtr& base) {
     auto num_element = base->GetNumElements();
     const auto* ids = base->GetIds();
     const auto* vectors = base->GetFloat32Vectors();
-    auto buckets = partition_strategy_->ClassifyDatas(vectors, num_element, 1);
+    auto buckets = partition_strategy_->ClassifyDatas(vectors, num_element, buckets_per_data_);
     Vector<float> normalize_data(dim_, allocator_);
     Vector<float> residual_data(dim_, allocator_);
     Vector<float> centroid(dim_, allocator_);
     for (int64_t i = 0; i < num_element; ++i) {
         const auto* data_ptr = vectors + i * dim_;
-        if (use_residual_) {
-            partition_strategy_->GetCentroid(buckets[i], centroid);
-            if (metric_ == MetricType::METRIC_TYPE_COSINE) {
-                Normalize(data_ptr, normalize_data.data(), dim_);
-                data_ptr = normalize_data.data();
+        for (int64_t j = 0; j < buckets_per_data_; ++j) {
+            auto idx = i * buckets_per_data_ + j;
+
+            if (use_residual_) {
+                partition_strategy_->GetCentroid(buckets[idx], centroid);
+                if (metric_ == MetricType::METRIC_TYPE_COSINE) {
+                    Normalize(data_ptr, normalize_data.data(), dim_);
+                    data_ptr = normalize_data.data();
+                }
+                FP32Sub(data_ptr, centroid.data(), residual_data.data(), dim_);
+                bucket_->InsertVector(residual_data.data(),
+                                      buckets[idx],
+                                      idx + total_elements_ * buckets_per_data_,
+                                      centroid.data());
+            } else {
+                bucket_->InsertVector(
+                    data_ptr, buckets[idx], idx + total_elements_ * buckets_per_data_);
             }
-            FP32Sub(data_ptr, centroid.data(), residual_data.data(), dim_);
-            bucket_->InsertVector(
-                residual_data.data(), buckets[i], i + total_elements_, centroid.data());
-        } else {
-            bucket_->InsertVector(data_ptr, buckets[i], i + total_elements_);
         }
         this->label_table_->Insert(i + total_elements_, ids[i]);
     }
+
     this->bucket_->Package();
     if (use_reorder_) {
         this->reorder_codes_->BatchInsertVector(base->GetFloat32Vectors(), base->GetNumElements());
@@ -351,6 +397,7 @@ IVF::Deserialize(StreamReader& reader) {
         this->reorder_codes_->Deserialize(reader);
     }
 }
+
 InnerSearchParam
 IVF::create_search_param(const std::string& parameters, const FilterPtr& filter) const {
     InnerSearchParam param;
@@ -363,6 +410,7 @@ IVF::create_search_param(const std::string& parameters, const FilterPtr& filter)
     param.scan_bucket_size = std::min(static_cast<BucketIdType>(search_param.scan_buckets_count),
                                       bucket_->bucket_count_);
     param.factor = search_param.topk_factor;
+    param.first_order_scan_ratio = search_param.first_order_scan_ratio;
     return param;
 }
 
@@ -407,9 +455,9 @@ IVF::search(const DatasetPtr& query, const InnerSearchParam& param) const {
         Normalize(query_data, normalize_data.data(), dim_);
         query_data = normalize_data.data();
     }
-    auto candidate_buckets =
-        partition_strategy_->ClassifyDatas(query_data, 1, param.scan_bucket_size);
+    auto candidate_buckets = partition_strategy_->ClassifyDatasForSearch(query_data, 1, param);
     auto computer = bucket_->FactoryComputer(query_data);
+
     Vector<float> dist(allocator_);
     auto cur_heap_top = std::numeric_limits<float>::max();
     int64_t topk = param.topk;
@@ -419,9 +467,23 @@ IVF::search(const DatasetPtr& query, const InnerSearchParam& param) const {
             topk = std::numeric_limits<int64_t>::max();
         }
     }
+    // Scale topk to ensure sufficient candidates after deduplication when buckets_per_data_ > 1
+    int64_t origin_topk = topk;
+    if (buckets_per_data_ > 1) {
+        if (topk <= std::numeric_limits<int64_t>::max() / buckets_per_data_) {
+            topk *= buckets_per_data_;
+        } else {
+            topk = std::numeric_limits<int64_t>::max();
+        }
+    }
+
     const auto& ft = param.is_inner_id_allowed;
     Vector<float> centroid(dim_, allocator_);
+
     for (auto& bucket_id : candidate_buckets) {
+        if (bucket_id == -1) {
+            break;
+        }
         auto bucket_size = bucket_->GetBucketSize(bucket_id);
         const auto* ids = bucket_->GetInnerIds(bucket_id);
         if (bucket_size > dist.size()) {
@@ -438,7 +500,8 @@ IVF::search(const DatasetPtr& query, const InnerSearchParam& param) const {
 
         bucket_->ScanBucketById(dist.data(), computer, bucket_id);
         for (int j = 0; j < bucket_size; ++j) {
-            if (ft == nullptr or ft->CheckValid(ids[j])) {
+            auto origin_id = ids[j] / buckets_per_data_;
+            if (ft == nullptr or ft->CheckValid(origin_id)) {
                 dist[j] -= ip_distance;
                 if constexpr (mode == KNN_SEARCH) {
                     if (search_result->Size() < topk or dist[j] < cur_heap_top) {
@@ -458,6 +521,35 @@ IVF::search(const DatasetPtr& query, const InnerSearchParam& param) const {
             }
         }
     }
+
+    // Deduplicate ids when buckets_per_data_ > 1
+    if (buckets_per_data_ > 1) {
+        std::unordered_map<InnerIdType, float> id_to_min_dist;
+        while (!search_result->Empty()) {
+            const auto& [dist_val, id] = search_result->Top();
+            auto origin_id = id / buckets_per_data_;
+            // Keep the smallest distance for each id
+            if (id_to_min_dist.find(origin_id) == id_to_min_dist.end() ||
+                dist_val < id_to_min_dist[origin_id]) {
+                id_to_min_dist[origin_id] = dist_val;
+            }
+            search_result->Pop();
+        }
+
+        auto cur_heap_top = std::numeric_limits<float>::max();
+        for (const auto& [origin_id, dist_val] : id_to_min_dist) {
+            if (dist_val < cur_heap_top) {
+                search_result->Push(dist_val, origin_id);
+            }
+            if (search_result->Size() > origin_topk) {
+                search_result->Pop();
+            }
+            if (not search_result->Empty() and search_result->Size() == origin_topk) {
+                cur_heap_top = search_result->Top().first;
+            }
+        }
+    }
+
     return search_result;
 }
 

src/algorithm/ivf.h

@@ -114,6 +114,7 @@ class IVF : public InnerIndexInterface {
     BucketInterfacePtr bucket_{nullptr};
 
     IVFPartitionStrategyPtr partition_strategy_{nullptr};
+    BucketIdType buckets_per_data_;
 
     int64_t total_elements_{0};