MS-288 update build stage

Former-commit-id: 8ee0976b88b0772d85d74f65b52e3484b4c76bd5

cpp/src/db/ExecutionEngine.h

@@ -19,6 +19,7 @@ enum class EngineType {
     FAISS_IDMAP = 1,
     FAISS_IVFFLAT,
     FAISS_IVFSQ8,
+    NSG_MIX,
     MAX_VALUE = FAISS_IVFSQ8,
 };
 

cpp/src/db/ExecutionEngineImpl.cpp

@@ -69,6 +69,10 @@ VecIndexPtr ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
             index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_MIX);
             break;
         }
+        case EngineType::NSG_MIX: {
+            index = GetVecIndexFactory(IndexType::NSG_MIX);
+            break;
+        }
         default: {
             ENGINE_LOG_ERROR << "Invalid engine type";
             return nullptr;

cpp/src/sdk/include/MilvusApi.h

@@ -19,6 +19,7 @@ enum class IndexType {
     cpu_idmap,
     gpu_ivfflat,
     gpu_ivfsq8,
+    mix_nsg,
 };
 
 /**

cpp/src/server/RequestTask.cpp

@@ -36,6 +36,7 @@ namespace {
                 {1, engine::EngineType::FAISS_IDMAP},
                 {2, engine::EngineType::FAISS_IVFFLAT},
                 {3, engine::EngineType::FAISS_IVFSQ8},
+                {4, engine::EngineType::NSG_MIX},
         };
 
         if(map_type.find(type) == map_type.end()) {
@@ -51,6 +52,7 @@ namespace {
                 {engine::EngineType::FAISS_IDMAP, 1},
                 {engine::EngineType::FAISS_IVFFLAT, 2},
                 {engine::EngineType::FAISS_IVFSQ8, 3},
+                {engine::EngineType::NSG_MIX, 4},
         };
 
         if(map_type.find(type) == map_type.end()) {

cpp/src/wrapper/knowhere/vec_index.cpp

@@ -7,6 +7,7 @@
 #include "knowhere/index/vector_index/idmap.h"
 #include "knowhere/index/vector_index/gpu_ivf.h"
 #include "knowhere/index/vector_index/cpu_kdt_rng.h"
+#include "knowhere/index/vector_index/nsg_index.h"
 #include "knowhere/common/exception.h"
 
 #include "vec_index.h"
@@ -103,10 +104,10 @@ VecIndexPtr GetVecIndexFactory(const IndexType &type) {
             index = std::make_shared<zilliz::knowhere::GPUIVFSQ>(0);
             return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
         }
-            //case IndexType::NSG: { // TODO(linxj): bug.
-            //    index = std::make_shared<zilliz::knowhere::NSG>();
-            //    break;
-            //}
+        case IndexType::NSG_MIX: { // TODO(linxj): bug.
+            index = std::make_shared<zilliz::knowhere::NSG>(0);
+            break;
+        }
         default: {
             return nullptr;
         }
@@ -199,6 +200,16 @@ void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Co
             if (!cfg.contains("nbits")) { cfg["nbits"] = int(8); }
             break;
         }
+        case IndexType::NSG_MIX: {
+            auto scale_factor = round(cfg["dim"].as<int>() / 128.0);
+            scale_factor = scale_factor >= 4 ? 4 : scale_factor;
+            if (!cfg.contains("nprobe")) { cfg["nprobe"] = 16 + 10 * scale_factor; }
+            if (!cfg.contains("knng")) { cfg["knng"] = 100 + 100 * scale_factor; }
+            if (!cfg.contains("search_length")) { cfg["search_length"] = 30 + 10 * scale_factor; }
+            if (!cfg.contains("out_degree")) { cfg["out_degree"] = 40 + 5 * scale_factor; }
+            if (!cfg.contains("candidate_pool_size")) { cfg["candidate_pool_size"] = 200 + 100 * scale_factor; }
+            break;
+        }
     }
 }
 

cpp/src/wrapper/knowhere/vec_index.h

@@ -32,7 +32,7 @@ enum class IndexType {
     FAISS_IVFPQ_GPU,
     SPTAG_KDT_RNT_CPU,
     FAISS_IVFSQ8_MIX,
-    //NSG,
+    NSG_MIX,
 };
 
 class VecIndex {

cpp/unittest/index_wrapper/knowhere_test.cpp

@@ -108,6 +108,12 @@ INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
                                             64, 100000, 10, 10,
                                             Config::object{{"dim", 64}, {"nlist", 1000}, {"nbits", 8}, {"metric_type", "L2"}},
                                             Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}}
+                            ),
+                            std::make_tuple(IndexType::NSG_MIX, "Default",
+                                            128, 100000, 10, 10,
+                                            Config::object{{"dim", 128}, {"nlist", 1638}, {"nprobe", 10}, {"metric_type", "L2"},
+                                                           {"knng", 50}, {"search_length", 40}, {"out_degree", 30}, {"candidate_pool_size", 200}},
+                                            Config::object{{"k", 10}, {"search_length", 20}}
                             )
                             //std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default",
                             //                64, 10000, 10, 10,