From 98571d5f44031c3dc2ea11ac4c1a82f3ff936f82 Mon Sep 17 00:00:00 2001
From: "zirui.chen" <zirui.chen@zilliz.com>
Date: Thu, 14 Nov 2019 21:04:10 +0800
Subject: [PATCH] #127 makeNlist & unittest change

---
 core/src/wrapper/ConfAdapter.cpp       | 12 ++++++++++++
 core/src/wrapper/ConfAdapter.h         |  4 ++++
 core/unittest/wrapper/test_wrapper.cpp |  8 ++++----
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/wrapper/ConfAdapter.cpp
index 90d8c6523b..a94fdac8b2 100644
--- a/core/src/wrapper/ConfAdapter.cpp
+++ b/core/src/wrapper/ConfAdapter.cpp
@@ -154,6 +154,18 @@ IVFPQConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& typ
     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>();
diff --git a/core/src/wrapper/ConfAdapter.h b/core/src/wrapper/ConfAdapter.h
index 4723afb097..5ec3d52486 100644
--- a/core/src/wrapper/ConfAdapter.h
+++ b/core/src/wrapper/ConfAdapter.h
@@ -82,6 +82,10 @@ class IVFPQConfAdapter : public IVFConfAdapter {
 
     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 {
diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp
index 531d4a0024..c12bd5ea42 100644
--- a/core/unittest/wrapper/test_wrapper.cpp
+++ b/core/unittest/wrapper/test_wrapper.cpp
@@ -80,18 +80,18 @@ INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
 
 #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", DIM, 10000, 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) {