#353 Rename config.h.in to version.h.in

CHANGELOG.md

@@ -8,13 +8,16 @@ Please mark all change in change log and use the ticket from JIRA.
 - \#246 - Exclude src/external folder from code coverage for jenkin ci
 - \#248 - Reside src/external in thirdparty
 - \#316 - Some files not merged after vectors added
+- \#327 - Search does not use GPU when index type is FLAT
 - \#340 - Test cases run failed on 0.6.0
 - \#353 - Rename config.h.in to version.h.in
+- \#374 - sdk_simple return empty result
 
 ## Feature
 - \#12 - Pure CPU version for Milvus
 - \#77 - Support table partition
 - \#226 - Experimental shards middleware for Milvus
+- \#127 - Support new Index type IVFPQ
 
 ## Improvement
 - \#275 - Rename C++ SDK IndexType
@@ -22,7 +25,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - \#260 - C++ SDK README
 - \#314 - add Find FAISS in CMake
 - \#310 - Add Q&A for 'protocol https not supported or disable in libcurl' issue
-- \#322 - Add option to enable / disable prometheus 
+- \#322 - Add option to enable / disable prometheus
 - \#358 - Add more information in build.sh and install.md
 
 ## Task

core/src/db/engine/ExecutionEngine.h

@@ -26,6 +26,7 @@
 namespace milvus {
 namespace engine {
 
+// TODO(linxj): replace with VecIndex::IndexType
 enum class EngineType {
     INVALID = 0,
     FAISS_IDMAP = 1,
@@ -33,7 +34,8 @@ enum class EngineType {
     FAISS_IVFSQ8,
     NSG_MIX,
     FAISS_IVFSQ8H,
-    MAX_VALUE = FAISS_IVFSQ8H,
+    FAISS_PQ,
+    MAX_VALUE = FAISS_PQ,
 };
 
 enum class MetricType {

core/src/db/engine/ExecutionEngineImpl.cpp

@@ -116,6 +116,14 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
             index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_HYBRID);
             break;
         }
+        case EngineType::FAISS_PQ: {
+#ifdef MILVUS_CPU_VERSION
+            index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_CPU);
+#else
+            index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_MIX);
+#endif
+            break;
+        }
         default: {
             ENGINE_LOG_ERROR << "Unsupported index type";
             return nullptr;

core/src/grpc/gen-milvus/milvus.pb.h

@@ -1321,7 +1321,8 @@ class RowRecord :
   void clear_vector_data();
   float vector_data(int index) const;
   void set_vector_data(int index, float value);
-  void add_vector_data(float value);
+//  void add_vector_data(float value);
+  void add_vector_data(std::vector<float>::const_iterator begin, std::vector<float>::const_iterator end);
   const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >&
       vector_data() const;
   ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >*
@@ -1473,7 +1474,9 @@ class InsertParam :
   void clear_row_id_array();
   ::PROTOBUF_NAMESPACE_ID::int64 row_id_array(int index) const;
   void set_row_id_array(int index, ::PROTOBUF_NAMESPACE_ID::int64 value);
-  void add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value);
+//  void add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value);
+  void add_row_id_array(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin,
+                        std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end);
   const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >&
       row_id_array() const;
   ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >*
@@ -2136,7 +2139,9 @@ class TopKQueryResult :
   void clear_ids();
   ::PROTOBUF_NAMESPACE_ID::int64 ids(int index) const;
   void set_ids(int index, ::PROTOBUF_NAMESPACE_ID::int64 value);
-  void add_ids(::PROTOBUF_NAMESPACE_ID::int64 value);
+//  void add_ids(::PROTOBUF_NAMESPACE_ID::int64 value);
+  void add_ids(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin,
+               std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end);
   const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >&
       ids() const;
   ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >*
@@ -2147,7 +2152,8 @@ class TopKQueryResult :
   void clear_distances();
   float distances(int index) const;
   void set_distances(int index, float value);
-  void add_distances(float value);
+//  void add_distances(float value);
+  void add_distances(std::vector<float>::const_iterator begin, std::vector<float>::const_iterator end);
   const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >&
       distances() const;
   ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >*
@@ -3922,9 +3928,14 @@ inline void RowRecord::set_vector_data(int index, float value) {
   vector_data_.Set(index, value);
   // @@protoc_insertion_point(field_set:milvus.grpc.RowRecord.vector_data)
 }
-inline void RowRecord::add_vector_data(float value) {
-  vector_data_.Add(value);
-  // @@protoc_insertion_point(field_add:milvus.grpc.RowRecord.vector_data)
+//inline void RowRecord::add_vector_data(float value) {
+//  vector_data_.Add(value);
+//  // @@protoc_insertion_point(field_add:milvus.grpc.RowRecord.vector_data)
+//}
+inline void RowRecord::add_vector_data(std::vector<float>::const_iterator begin,
+                                       std::vector<float>::const_iterator end) {
+    vector_data_.Add(begin, end);
+    // @@protoc_insertion_point(field_add:milvus.grpc.RowRecord.vector_data)
 }
 inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >&
 RowRecord::vector_data() const {
@@ -4037,9 +4048,14 @@ inline void InsertParam::set_row_id_array(int index, ::PROTOBUF_NAMESPACE_ID::in
   row_id_array_.Set(index, value);
   // @@protoc_insertion_point(field_set:milvus.grpc.InsertParam.row_id_array)
 }
-inline void InsertParam::add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value) {
-  row_id_array_.Add(value);
-  // @@protoc_insertion_point(field_add:milvus.grpc.InsertParam.row_id_array)
+//inline void InsertParam::add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value) {
+//  row_id_array_.Add(value);
+//  // @@protoc_insertion_point(field_add:milvus.grpc.InsertParam.row_id_array)
+//}
+inline void InsertParam::add_row_id_array(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin,
+                                          std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end) {
+    row_id_array_.Add(begin, end);
+    // @@protoc_insertion_point(field_add:milvus.grpc.InsertParam.row_id_array)
 }
 inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >&
 InsertParam::row_id_array() const {
@@ -4588,9 +4604,14 @@ inline void TopKQueryResult::set_ids(int index, ::PROTOBUF_NAMESPACE_ID::int64 v
   ids_.Set(index, value);
   // @@protoc_insertion_point(field_set:milvus.grpc.TopKQueryResult.ids)
 }
-inline void TopKQueryResult::add_ids(::PROTOBUF_NAMESPACE_ID::int64 value) {
-  ids_.Add(value);
-  // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.ids)
+//inline void TopKQueryResult::add_ids(::PROTOBUF_NAMESPACE_ID::int64 value) {
+//  ids_.Add(value);
+//  // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.ids)
+//}
+inline void TopKQueryResult::add_ids(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin,
+                                     std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end) {
+    ids_.Add(begin,end);
+    // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.ids)
 }
 inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >&
 TopKQueryResult::ids() const {
@@ -4618,9 +4639,13 @@ inline void TopKQueryResult::set_distances(int index, float value) {
   distances_.Set(index, value);
   // @@protoc_insertion_point(field_set:milvus.grpc.TopKQueryResult.distances)
 }
-inline void TopKQueryResult::add_distances(float value) {
-  distances_.Add(value);
-  // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.distances)
+//inline void TopKQueryResult::add_distances(float value) {
+//  distances_.Add(value);
+//  // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.distances)
+//}
+inline void TopKQueryResult::add_distances(std::vector<float>::const_iterator begin, std::vector<float>::const_iterator end) {
+    distances_.Add(begin, end);
+    // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.distances)
 }
 inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >&
 TopKQueryResult::distances() const {

core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp

@@ -39,17 +39,19 @@ GPUIVFPQ::Train(const DatasetPtr& dataset, const Config& config) {
 
     GETTENSOR(dataset)
 
-    // TODO(linxj): set device here.
-    // TODO(linxj): set gpu resource here.
-    faiss::gpu::StandardGpuResources res;
-    faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, build_cfg->nlist, build_cfg->m, build_cfg->nbits,
-                                           GetMetricType(build_cfg->metric_type));  // IP not support
-    device_index.train(rows, (float*)p_data);
-
-    std::shared_ptr<faiss::Index> host_index = nullptr;
-    host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index));
-
-    return std::make_shared<IVFIndexModel>(host_index);
+    auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
+    if (temp_resource != nullptr) {
+        ResScope rs(temp_resource, gpu_id_, true);
+        auto device_index = new faiss::gpu::GpuIndexIVFPQ(temp_resource->faiss_res.get(), dim, build_cfg->nlist,
+                                                          build_cfg->m, build_cfg->nbits,
+                                                          GetMetricType(build_cfg->metric_type));  // IP not support
+        device_index->train(rows, (float*)p_data);
+        std::shared_ptr<faiss::Index> host_index = nullptr;
+        host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index));
+        return std::make_shared<IVFIndexModel>(host_index);
+    } else {
+        KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");
+    }
 }
 
 std::shared_ptr<faiss::IVFSearchParameters>
@@ -66,7 +68,14 @@ GPUIVFPQ::GenParams(const Config& config) {
 
 VectorIndexPtr
 GPUIVFPQ::CopyGpuToCpu(const Config& config) {
-    KNOWHERE_THROW_MSG("not support yet");
+    std::lock_guard<std::mutex> lk(mutex_);
+
+    faiss::Index* device_index = index_.get();
+    faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
+
+    std::shared_ptr<faiss::Index> new_index;
+    new_index.reset(host_index);
+    return std::make_shared<IVFPQ>(new_index);
 }
 
 }  // namespace knowhere

core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <memory>
+#include <utility>
 
 #include "IndexGPUIVF.h"
 
@@ -28,6 +29,10 @@ class GPUIVFPQ : public GPUIVF {
     explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) {
     }
 
+    GPUIVFPQ(std::shared_ptr<faiss::Index> index, const int64_t& device_id, ResPtr& resource)
+        : GPUIVF(std::move(index), device_id, resource) {
+    }
+
     IndexModelPtr
     Train(const DatasetPtr& dataset, const Config& config) override;
 

core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp

@@ -17,11 +17,19 @@
 
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFPQ.h>
+#ifdef MILVUS_GPU_VERSION
+#include <faiss/gpu/GpuCloner.h>
+#endif
+
 #include <memory>
 #include <utility>
 
 #include "knowhere/adapter/VectorAdapter.h"
 #include "knowhere/common/Exception.h"
+#ifdef MILVUS_GPU_VERSION
+#include "knowhere/index/vector_index/IndexGPUIVF.h"
+#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
+#endif
 #include "knowhere/index/vector_index/IndexIVFPQ.h"
 
 namespace knowhere {
@@ -60,4 +68,22 @@ IVFPQ::Clone_impl(const std::shared_ptr<faiss::Index>& index) {
     return std::make_shared<IVFPQ>(index);
 }
 
+VectorIndexPtr
+IVFPQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
+#ifdef MILVUS_GPU_VERSION
+    if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
+        ResScope rs(res, device_id, false);
+        auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
+
+        std::shared_ptr<faiss::Index> device_index;
+        device_index.reset(gpu_index);
+        return std::make_shared<GPUIVFPQ>(device_index, device_id, res);
+    } else {
+        KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
+    }
+#else
+    KNOWHERE_THROW_MSG("Calling IVFPQ::CopyCpuToGpu when we are using CPU version");
+#endif
+}
+
 }  // namespace knowhere

core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h

@@ -34,6 +34,9 @@ class IVFPQ : public IVF {
     IndexModelPtr
     Train(const DatasetPtr& dataset, const Config& config) override;
 
+    VectorIndexPtr
+    CopyCpuToGpu(const int64_t& device_id, const Config& config) override;
+
  protected:
     std::shared_ptr<faiss::IVFSearchParameters>
     GenParams(const Config& config) override;

core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp

@@ -49,7 +49,7 @@ CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config
     if (auto cpu_index = std::dynamic_pointer_cast<IVFSQ>(index)) {
         return cpu_index->CopyCpuToGpu(device_id, config);
     } else if (auto cpu_index = std::dynamic_pointer_cast<IVFPQ>(index)) {
-        KNOWHERE_THROW_MSG("IVFPQ not support transfer to gpu");
+        return cpu_index->CopyCpuToGpu(device_id, config);
     } else if (auto cpu_index = std::dynamic_pointer_cast<IVF>(index)) {
         return cpu_index->CopyCpuToGpu(device_id, config);
     } else if (auto cpu_index = std::dynamic_pointer_cast<IDMAP>(index)) {

core/src/index/unittest/test_ivf.cpp

@@ -213,7 +213,7 @@ TEST_P(IVFTest, clone_test) {
 
     {
         // copy from gpu to cpu
-        std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"};
+        std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "GPUIVFPQ", "IVFSQHybrid"};
         auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
         if (finder != support_idx_vec.cend()) {
             EXPECT_NO_THROW({
@@ -238,7 +238,7 @@ TEST_P(IVFTest, clone_test) {
 
     {
         // copy to gpu
-        std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ"};
+        std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ", "IVFPQ", "GPUIVFPQ"};
         auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
         if (finder != support_idx_vec.cend()) {
             EXPECT_NO_THROW({

core/src/scheduler/interface/interfaces.h

@@ -33,8 +33,7 @@ namespace milvus {
 namespace interface {
 
 struct dumpable {
-    virtual ~dumpable() {
-    }
+    virtual ~dumpable() = default;
 
     virtual json
     Dump() const = 0;

core/src/scheduler/optimizer/OnlyGPUPass.cpp

@@ -34,7 +34,8 @@ OnlyGPUPass::Run(const TaskPtr& task) {
 
     auto search_task = std::static_pointer_cast<XSearchTask>(task);
     if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8 &&
-        search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT) {
+        search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT &&
+        search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IDMAP) {
         return false;
     }
 

core/src/sdk/grpc/ClientProxy.cpp

@@ -32,13 +32,6 @@ UriCheck(const std::string& uri) {
     return (index != std::string::npos);
 }
 
-void
-CopyRowRecord(::milvus::grpc::RowRecord* target, const RowRecord& src) {
-    auto vector_data = target->mutable_vector_data();
-    vector_data->Resize(static_cast<int>(src.data.size()), 0.0);
-    memcpy(vector_data->mutable_data(), src.data.data(), src.data.size() * sizeof(float));
-}
-
 Status
 ClientProxy::Connect(const ConnectParam& param) {
     std::string uri = param.ip_address + ":" + param.port;
@@ -196,17 +189,14 @@ ClientProxy::Insert(const std::string& table_name, const std::string& partition_
 
         for (auto& record : record_array) {
             ::milvus::grpc::RowRecord* grpc_record = insert_param.add_row_record_array();
-            CopyRowRecord(grpc_record, record);
+            grpc_record->add_vector_data(record.data.begin(), record.data.end());
         }
 
         // Single thread
         ::milvus::grpc::VectorIds vector_ids;
         if (!id_array.empty()) {
             /* set user's ids */
-            auto row_ids = insert_param.mutable_row_id_array();
-            row_ids->Reserve(static_cast<int>(id_array.size()));
-            memcpy(row_ids->mutable_data(), id_array.data(), id_array.size() * sizeof(int64_t));
-
+            insert_param.add_row_id_array(id_array.begin(), id_array.end());
             client_ptr_->Insert(vector_ids, insert_param, status);
         } else {
             client_ptr_->Insert(vector_ids, insert_param, status);
@@ -236,7 +226,7 @@ ClientProxy::Search(const std::string& table_name, const std::vector<std::string
         }
         for (auto& record : query_record_array) {
             ::milvus::grpc::RowRecord* row_record = search_param.add_query_record_array();
-            CopyRowRecord(row_record, record);
+            row_record->add_vector_data(record.data.begin(), record.data.end());
         }
 
         // step 2: convert range array

core/src/server/grpc_impl/GrpcRequestTask.cpp

@@ -682,11 +682,8 @@ SearchTask::OnExecute() {
 
         // step 7: construct result array
         topk_result_->set_row_num(record_count);
-        topk_result_->mutable_ids()->Resize(static_cast<int>(result_ids.size()), 0);
-        memcpy(topk_result_->mutable_ids()->mutable_data(), result_ids.data(), result_ids.size() * sizeof(int64_t));
-        topk_result_->mutable_distances()->Resize(static_cast<int>(result_distances.size()), 0.0);
-        memcpy(topk_result_->mutable_distances()->mutable_data(), result_distances.data(),
-               result_distances.size() * sizeof(float));
+        topk_result_->add_ids(result_ids.begin(), result_ids.end());
+        topk_result_->add_distances(result_distances.begin(), result_distances.end());
 
         // step 8: print time cost percent
         rc.RecordSection("construct result and send");

core/src/wrapper/ConfAdapter.cpp

@@ -16,6 +16,7 @@
 // under the License.
 
 #include "wrapper/ConfAdapter.h"
+#include "WrapperException.h"
 #include "knowhere/index/vector_index/helpers/IndexParameter.h"
 #include "utils/Log.h"
 
@@ -76,7 +77,7 @@ IVFConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) {
     if (size <= TYPICAL_COUNT / 16384 + 1) {
         // handle less row count, avoid nlist set to 0
         return 1;
-    } else if (int(size / TYPICAL_COUNT) * nlist == 0) {
+    } else if (int(size / TYPICAL_COUNT) * nlist <= 0) {
         // calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT
         return int(size / TYPICAL_COUNT * 16384);
     }
@@ -87,7 +88,11 @@ knowhere::Config
 IVFConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) {
     auto conf = std::make_shared<knowhere::IVFCfg>();
     conf->k = metaconf.k;
-    conf->nprobe = metaconf.nprobe;
+
+    if (metaconf.nprobe <= 0)
+        conf->nprobe = 16;  // hardcode here
+    else
+        conf->nprobe = metaconf.nprobe;
 
     switch (type) {
         case IndexType::FAISS_IVFFLAT_GPU:
@@ -123,11 +128,47 @@ IVFPQConfAdapter::Match(const TempMetaConf& metaconf) {
     conf->metric_type = metaconf.metric_type;
     conf->gpu_id = conf->gpu_id;
     conf->nbits = 8;
-    conf->m = 8;
+
+    if (!(conf->d % 4))
+        conf->m = conf->d / 4;  // compression radio = 16
+    else if (!(conf->d % 2))
+        conf->m = conf->d / 2;  // compression radio = 8
+    else if (!(conf->d % 3))
+        conf->m = conf->d / 3;  // compression radio = 12
+    else
+        conf->m = conf->d;  // same as SQ8, compression radio = 4
+
     MatchBase(conf);
     return conf;
 }
 
+knowhere::Config
+IVFPQConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) {
+    auto conf = std::make_shared<knowhere::IVFPQCfg>();
+    conf->k = metaconf.k;
+
+    if (metaconf.nprobe <= 0) {
+        WRAPPER_LOG_ERROR << "The nprobe of PQ is wrong!";
+        throw WrapperException("The nprobe of PQ is wrong!");
+    } else {
+        conf->nprobe = metaconf.nprobe;
+    }
+
+    return conf;
+}
+
+int64_t
+IVFPQConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) {
+    if (size <= TYPICAL_COUNT / 16384 + 1) {
+        // handle less row count, avoid nlist set to 0
+        return 1;
+    } else if (int(size / TYPICAL_COUNT) * nlist <= 0) {
+        // calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT
+        return int(size / TYPICAL_COUNT * 16384);
+    }
+    return nlist;
+}
+
 knowhere::Config
 NSGConfAdapter::Match(const TempMetaConf& metaconf) {
     auto conf = std::make_shared<knowhere::NSGCfg>();
@@ -136,13 +177,14 @@ NSGConfAdapter::Match(const TempMetaConf& metaconf) {
     conf->metric_type = metaconf.metric_type;
     conf->gpu_id = conf->gpu_id;
 
+    double factor = metaconf.size / TYPICAL_COUNT;
     auto scale_factor = round(metaconf.dim / 128.0);
     scale_factor = scale_factor >= 4 ? 4 : scale_factor;
-    conf->nprobe = 6 + 10 * scale_factor;
-    conf->knng = 100 + 100 * scale_factor;
-    conf->search_length = 40 + 5 * scale_factor;
-    conf->out_degree = 50 + 5 * scale_factor;
-    conf->candidate_pool_size = 200 + 100 * scale_factor;
+    conf->nprobe = conf->nlist > 10000 ? conf->nlist * 0.02 : conf->nlist * 0.1;
+    conf->knng = (100 + 100 * scale_factor) * factor;
+    conf->search_length = (40 + 5 * scale_factor) * factor;
+    conf->out_degree = (50 + 5 * scale_factor) * factor;
+    conf->candidate_pool_size = (200 + 100 * scale_factor) * factor;
     MatchBase(conf);
 
     //    WRAPPER_LOG_DEBUG << "nlist: " << conf->nlist
@@ -156,6 +198,9 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type)
     auto conf = std::make_shared<knowhere::NSGCfg>();
     conf->k = metaconf.k;
     conf->search_length = metaconf.search_length;
+    if (metaconf.search_length == TEMPMETA_DEFAULT_VALUE) {
+        conf->search_length = 30;  // TODO(linxj): hardcode here.
+    }
     return conf;
 }
 

core/src/wrapper/ConfAdapter.h

@@ -79,6 +79,13 @@ class IVFPQConfAdapter : public IVFConfAdapter {
  public:
     knowhere::Config
     Match(const TempMetaConf& metaconf) override;
+
+    knowhere::Config
+    MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override;
+
+ protected:
+    static int64_t
+    MatchNlist(const int64_t& size, const int64_t& nlist);
 };
 
 class NSGConfAdapter : public IVFConfAdapter {

core/src/wrapper/ConfAdapterMgr.cpp

@@ -53,6 +53,7 @@ AdapterMgr::RegisterAdapter() {
 
     REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_CPU, ivfpq_cpu);
     REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_GPU, ivfpq_gpu);
+    REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix);
 
     REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix);
 }

core/src/wrapper/VecIndex.cpp

@@ -145,6 +145,10 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) {
             index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
             break;
         }
+        case IndexType::FAISS_IVFPQ_MIX: {
+            index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
+            return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFPQ_MIX);
+        }
         case IndexType::FAISS_IVFSQ8_MIX: {
             index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
             return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
@@ -276,6 +280,10 @@ ConvertToCpuIndexType(const IndexType& type) {
         case IndexType::FAISS_IVFSQ8_MIX: {
             return IndexType::FAISS_IVFSQ8_CPU;
         }
+        case IndexType::FAISS_IVFPQ_GPU:
+        case IndexType::FAISS_IVFPQ_MIX: {
+            return IndexType::FAISS_IVFPQ_CPU;
+        }
         default: { return type; }
     }
 }
@@ -291,9 +299,12 @@ ConvertToGpuIndexType(const IndexType& type) {
         case IndexType::FAISS_IVFSQ8_CPU: {
             return IndexType::FAISS_IVFSQ8_GPU;
         }
+        case IndexType::FAISS_IVFPQ_MIX:
+        case IndexType::FAISS_IVFPQ_CPU: {
+            return IndexType::FAISS_IVFPQ_GPU;
+        }
         default: { return type; }
     }
 }
-
 }  // namespace engine
 }  // namespace milvus

core/src/wrapper/VecIndex.h

@@ -33,6 +33,7 @@ namespace engine {
 
 using Config = knowhere::Config;
 
+// TODO(linxj): replace with string, Do refactor serialization
 enum class IndexType {
     INVALID = 0,
     FAISS_IDMAP = 1,
@@ -47,6 +48,7 @@ enum class IndexType {
     FAISS_IVFSQ8_GPU,
     FAISS_IVFSQ8_HYBRID,  // only support build on gpu.
     NSG_MIX,
+    FAISS_IVFPQ_MIX,
 };
 
 class VecIndex;

core/unittest/wrapper/test_wrapper.cpp

@@ -29,33 +29,40 @@
 
 INITIALIZE_EASYLOGGINGPP
 
-using ::testing::Combine;
 using ::testing::TestWithParam;
 using ::testing::Values;
+using ::testing::Combine;
 
 class KnowhereWrapperTest
-    : public DataGenBase,
-      public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
+        : public DataGenBase,
+          public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
  protected:
-    void
-    SetUp() override {
+    void SetUp() override {
 #ifdef MILVUS_GPU_VERSION
         knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
 #endif
-
         std::string generator_type;
         std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam();
         GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
 
+        milvus::engine::TempMetaConf tempconf;
+        tempconf.metric_type = knowhere::METRICTYPE::L2;
+        tempconf.gpu_id = DEVICEID;
+        tempconf.size = nb;
+        tempconf.dim = dim;
+        tempconf.k = k;
+        tempconf.nprobe = 16;
+
         index_ = GetVecIndexFactory(index_type);
-        conf = ParamGenerator::GetInstance().Gen(index_type);
-        conf->k = k;
-        conf->d = dim;
-        conf->gpu_id = DEVICEID;
+        conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf);
+        searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf);
+
+//        conf->k = k;
+//        conf->d = dim;
+//        conf->gpu_id = DEVICEID;
     }
 
-    void
-    TearDown() override {
+    void TearDown() override {
 #ifdef MILVUS_GPU_VERSION
         knowhere::FaissGpuResourceMgr::GetInstance().Free();
 #endif
@@ -65,24 +72,27 @@ class KnowhereWrapperTest
     milvus::engine::IndexType index_type;
     milvus::engine::VecIndexPtr index_ = nullptr;
     knowhere::Config conf;
+    knowhere::Config searchconf;
 };
 
-INSTANTIATE_TEST_CASE_P(
-    WrapperParam, KnowhereWrapperTest,
-    Values(
-//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
+INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
+                        Values(
+                        //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
+
 #ifdef MILVUS_GPU_VERSION
         std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10),
-        std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 100000, 10, 10),
+        std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 1000, 10, 10),
         //                            std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB,
         //                            10, 10),
         std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10),
         std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10),
+        std::make_tuple(milvus::engine::IndexType::FAISS_IVFPQ_MIX, "Default", 64, 1000, 10, 10),
+
 //                            std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10),
 #endif
         //                            std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10),
-        std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10),
-        std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 100000, 10, 10),
+        std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 1000, 10, 10),
+        std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 1000, 10, 10),
         std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10)));
 
 TEST_P(KnowhereWrapperTest, BASE_TEST) {
@@ -93,12 +103,11 @@ TEST_P(KnowhereWrapperTest, BASE_TEST) {
     std::vector<float> res_dis(elems);
 
     index_->BuildAll(nb, xb.data(), ids.data(), conf);
-    index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+    index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
     AssertResult(res_ids, res_dis);
 }
 
 #ifdef MILVUS_GPU_VERSION
-
 TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
     EXPECT_EQ(index_->GetType(), index_type);
 
@@ -107,13 +116,13 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
     std::vector<float> res_dis(elems);
 
     index_->BuildAll(nb, xb.data(), ids.data(), conf);
-    index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+    index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
     AssertResult(res_ids, res_dis);
 
     {
         auto dev_idx = index_->CopyToGpu(DEVICEID);
         for (int i = 0; i < 10; ++i) {
-            dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+            dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
         }
         AssertResult(res_ids, res_dis);
     }
@@ -125,7 +134,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
 
         auto dev_idx = new_index->CopyToGpu(DEVICEID);
         for (int i = 0; i < 10; ++i) {
-            dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+            dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
         }
         AssertResult(res_ids, res_dis);
     }
@@ -139,7 +148,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
     std::vector<int64_t> res_ids(elems);
     std::vector<float> res_dis(elems);
     index_->BuildAll(nb, xb.data(), ids.data(), conf);
-    index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+    index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
     AssertResult(res_ids, res_dis);
 
     {
@@ -152,7 +161,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
 
         std::vector<int64_t> res_ids(elems);
         std::vector<float> res_dis(elems);
-        new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+        new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
         AssertResult(res_ids, res_dis);
     }
 
@@ -166,7 +175,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
 
         std::vector<int64_t> res_ids(elems);
         std::vector<float> res_dis(elems);
-        new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
+        new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
         AssertResult(res_ids, res_dis);
     }
 }

core/unittest/wrapper/utils.h

@@ -15,18 +15,22 @@
 // specific language governing permissions and limitations
 // under the License.
 
+
 #pragma once
 
 #include <gtest/gtest.h>
-#include <cstdio>
-#include <cstdlib>
-#include <fstream>
 #include <memory>
 #include <vector>
+#include <cstdlib>
+#include <cstdio>
+#include <fstream>
+
 
-#include "knowhere/index/vector_index/helpers/IndexParameter.h"
 #include "wrapper/VecIndex.h"
 #include "wrapper/utils.h"
+#include "knowhere/index/vector_index/helpers/IndexParameter.h"
+#include "wrapper/ConfAdapterMgr.h"
+#include "wrapper/ConfAdapter.h"
 
 class DataGenBase;
 
@@ -40,29 +44,31 @@ constexpr int64_t PINMEM = 1024 * 1024 * 200;
 constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
 constexpr int64_t RESNUM = 2;
 
-static const char* CONFIG_PATH = "/tmp/milvus_test";
-static const char* CONFIG_FILE = "/server_config.yaml";
+static const char *CONFIG_PATH = "/tmp/milvus_test";
+static const char *CONFIG_FILE = "/server_config.yaml";
 
 class KnowhereTest : public ::testing::Test {
  protected:
-    void
-    SetUp() override;
-    void
-    TearDown() override;
+    void SetUp() override;
+    void TearDown() override;
 };
 
 class DataGenBase {
  public:
-    virtual void
-    GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids, const int& k,
-            int64_t* gt_ids, float* gt_dis);
+    virtual void GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids,
+                         const int& k, int64_t* gt_ids, float* gt_dis);
 
-    virtual void
-    GenData(const int& dim, const int& nb, const int& nq, std::vector<float>& xb, std::vector<float>& xq,
-            std::vector<int64_t>& ids, const int& k, std::vector<int64_t>& gt_ids, std::vector<float>& gt_dis);
+    virtual void GenData(const int& dim,
+                         const int& nb,
+                         const int& nq,
+                         std::vector<float>& xb,
+                         std::vector<float>& xq,
+                         std::vector<int64_t>& ids,
+                         const int& k,
+                         std::vector<int64_t>& gt_ids,
+                         std::vector<float>& gt_dis);
 
-    void
-    AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis);
+    void AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis);
 
     int dim = DIM;
     int nb = NB;
@@ -79,14 +85,22 @@ class DataGenBase {
 
 class ParamGenerator {
  public:
-    static ParamGenerator&
-    GetInstance() {
+    static ParamGenerator& GetInstance() {
         static ParamGenerator instance;
         return instance;
     }
 
-    knowhere::Config
-    Gen(const milvus::engine::IndexType& type) {
+    knowhere::Config GenSearchConf(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) {
+        auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type);
+        return adapter->MatchSearch(conf, type);
+    }
+
+    knowhere::Config GenBuild(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) {
+        auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type);
+        return adapter->Match(conf);
+    }
+
+    knowhere::Config Gen(const milvus::engine::IndexType& type) {
         switch (type) {
             case milvus::engine::IndexType::FAISS_IDMAP: {
                 auto tempconf = std::make_shared<knowhere::Cfg>();
@@ -113,34 +127,37 @@ class ParamGenerator {
                 tempconf->metric_type = knowhere::METRICTYPE::L2;
                 return tempconf;
             }
-                //            case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
-                //            case milvus::engine::IndexType::FAISS_IVFPQ_GPU: {
-                //                auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
-                //                tempconf->nlist = 100;
-                //                tempconf->nprobe = 16;
-                //                tempconf->nbits = 8;
-                //                tempconf->m = 8;
-                //                tempconf->metric_type = knowhere::METRICTYPE::L2;
-                //                return tempconf;
-                //            }
-                //            case milvus::engine::IndexType::NSG_MIX: {
-                //                auto tempconf = std::make_shared<knowhere::NSGCfg>();
-                //                tempconf->nlist = 100;
-                //                tempconf->nprobe = 16;
-                //                tempconf->search_length = 8;
-                //                tempconf->knng = 200;
-                //                tempconf->search_length = 40; // TODO(linxj): be 20 when search
-                //                tempconf->out_degree = 60;
-                //                tempconf->candidate_pool_size = 200;
-                //                tempconf->metric_type = knowhere::METRICTYPE::L2;
-                //                return tempconf;
-                //            }
+            case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
+            case milvus::engine::IndexType::FAISS_IVFPQ_GPU:
+            case milvus::engine::IndexType::FAISS_IVFPQ_MIX: {
+                auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
+                tempconf->nlist = 100;
+                tempconf->nprobe = 16;
+                tempconf->nbits = 8;
+                tempconf->m = 8;
+                tempconf->metric_type = knowhere::METRICTYPE::L2;
+                return tempconf;
+            }
+            case milvus::engine::IndexType::NSG_MIX: {
+                auto tempconf = std::make_shared<knowhere::NSGCfg>();
+                tempconf->nlist = 100;
+                tempconf->nprobe = 16;
+                tempconf->search_length = 8;
+                tempconf->knng = 200;
+                tempconf->search_length = 40; // TODO(linxj): be 20 when search
+                tempconf->out_degree = 60;
+                tempconf->candidate_pool_size = 200;
+                tempconf->metric_type = knowhere::METRICTYPE::L2;
+                return tempconf;
+            }
         }
     }
 };
 
-// class SanityCheck : public DataGenBase {
+
+//class SanityCheck : public DataGenBase {
 // public:
 //    void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, long *ids,
 //                 const int &k, long *gt_ids, float *gt_dis) override;
 //};
+

tests/milvus_python_test/conftest.py

@@ -42,16 +42,22 @@ def connect(request):
     port = request.config.getoption("--port")
     milvus = Milvus()
     try:
-        milvus.connect(host=ip, port=port)
-    except:
+        status = milvus.connect(host=ip, port=port)
+        logging.getLogger().info(status)
+        if not status.OK():
+            # try again
+            logging.getLogger().info("------------------------------------")
+            logging.getLogger().info("Try to connect again")
+            logging.getLogger().info("------------------------------------")
+            res = milvus.connect(host=ip, port=port)
+    except Exception as e:
+        logging.getLogger().error(str(e))
         pytest.exit("Milvus server can not connected, exit pytest ...")
-
     def fin():
         try:
             milvus.disconnect()
         except:
             pass
-
     request.addfinalizer(fin)
     return milvus
 
@@ -129,4 +135,4 @@ def ip_table(request, connect):
 
     request.addfinalizer(teardown)
 
-    return table_name
\ No newline at end of file
+    return table_name

tests/milvus_python_test/test_connect.py

@@ -249,8 +249,8 @@ class TestConnect:
         '''
         milvus = Milvus()
         uri_value = "tcp://%s:19540" % args["ip"]
-        milvus.connect(host=args["ip"], port="", uri=uri_value)
-        assert milvus.connected()
+        with pytest.raises(Exception) as e:
+            milvus.connect(host=args["ip"], port="", uri=uri_value)
 
     def test_connect_param_priority_uri(self, args):
         '''