Caiyd 1885 optimize knowhere ut (#1889)

* optimize knowhere unittest
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* code clean
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* optimize knowhere util
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix clang-format
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* update changelog
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix unittest build error
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* retry CI
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>
Co-authored-by: NJin Hai <hai.jin@zilliz.com>

CHANGELOG.md

@@ -24,10 +24,12 @@ Please mark all change in change log and use the issue from GitHub
 -   \#1784 Add Substructure and Superstructure in http module
 -   \#1858 Disable S3 build
 -   \#1882 Add index annoy into http module
+-   \#1885 Optimize knowhere unittest
 -   \#1886 Refactor log on search and insert request
 
 ## Task
 
+
 # Milvus 0.7.1 (2020-03-29)
 
 ## Bug

core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.h

@@ -27,7 +27,7 @@ struct MemoryIOWriter : public faiss::IOWriter {
     template <typename T>
     size_t
     write(T* ptr, size_t size, size_t nitems = 1) {
-        operator()((const void*)ptr, size, nitems);
+        return operator()((const void*)ptr, size, nitems);
     }
 };
 
@@ -42,7 +42,7 @@ struct MemoryIOReader : public faiss::IOReader {
     template <typename T>
     size_t
     read(T* ptr, size_t size, size_t nitems = 1) {
-        operator()((void*)ptr, size, nitems);
+        return operator()((void*)ptr, size, nitems);
     }
 };
 

core/src/index/unittest/CMakeLists.txt

@@ -28,6 +28,7 @@ set(basic_libs
 
 set(util_srcs
         ${MILVUS_THIRDPARTY_SRC}/easyloggingpp/easylogging++.cc
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.cpp
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexType.cpp
@@ -51,8 +52,7 @@ if (KNOWHERE_GPU_VERSION)
             )
 endif ()
 
-#<IVF-TEST>
-set(ivf_srcs
+set(faiss_srcs
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.cpp
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp
@@ -63,58 +63,113 @@ set(ivf_srcs
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
         )
 if (KNOWHERE_GPU_VERSION)
-    set(ivf_srcs ${ivf_srcs}
-            ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
-            ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp
-            ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp
-            ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp
-            ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp
-            ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp
-            )
+set(faiss_srcs ${faiss_srcs}
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp
+        )
+endif ()
+
+################################################################################
+#<INSTRUCTIONSET-TEST>
+if (NOT TARGET test_instructionset)
+    add_executable(test_instructionset test_instructionset.cpp)
+endif ()
+target_link_libraries(test_instructionset ${depend_libs} ${unittest_libs})
+
+################################################################################
+#<KNOWHERE-COMMON-TEST>
+if (NOT TARGET test_knowhere_common)
+    add_executable(test_knowhere_common test_common.cpp ${util_srcs})
+endif ()
+target_link_libraries(test_knowhere_common ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_knowhere_common DESTINATION unittest)
+
+if (KNOWHERE_GPU_VERSION)
+################################################################################
+#<GPU-RESOURCE-TEST>
+add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${faiss_srcs})
+target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_gpuresource DESTINATION unittest)
+
+################################################################################
+#<CUSTOMIZED-INDEX-TEST>
+add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${faiss_srcs})
+target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_customized_index DESTINATION unittest)
+endif ()
+
+################################################################################
+#<IDMAP-TEST>
+if (NOT TARGET test_idmap)
+    add_executable(test_idmap test_idmap.cpp ${faiss_srcs} ${util_srcs})
 endif ()
+target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_idmap DESTINATION unittest)
+
+################################################################################
+#<IVF-TEST>
 if (NOT TARGET test_ivf)
-    add_executable(test_ivf test_ivf.cpp ${ivf_srcs} ${util_srcs})
+    add_executable(test_ivf test_ivf.cpp ${faiss_srcs} ${util_srcs})
 endif ()
 target_link_libraries(test_ivf ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_ivf DESTINATION unittest)
+
+################################################################################
+#<BinaryIDMAP-TEST>
+if (NOT TARGET test_binaryidmap)
+    add_executable(test_binaryidmap test_binaryidmap.cpp ${faiss_srcs} ${util_srcs})
+endif ()
+target_link_libraries(test_binaryidmap ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_binaryidmap DESTINATION unittest)
 
+################################################################################
+#<BinaryIVF-TEST>
 if (NOT TARGET test_binaryivf)
-    add_executable(test_binaryivf test_binaryivf.cpp ${ivf_srcs} ${util_srcs})
+    add_executable(test_binaryivf test_binaryivf.cpp ${faiss_srcs} ${util_srcs})
 endif ()
 target_link_libraries(test_binaryivf ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_binaryivf DESTINATION unittest)
 
 
-#<IDMAP-TEST>
-if (NOT TARGET test_idmap)
-    add_executable(test_idmap test_idmap.cpp ${ivf_srcs} ${util_srcs})
+################################################################################
+#<NSG-TEST>
+add_definitions(-std=c++11 -O3 -lboost -march=native -Wall -DINFO)
+
+find_package(OpenMP REQUIRED)
+if (OpenMP_FOUND)
+    set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
+else ()
+    message(FATAL_ERROR "no OpenMP supprot")
 endif ()
-target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs})
 
-#<ANNOY-TEST>
-set(annoy_srcs
-        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexAnnoy.cpp
+include_directories(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/impl/nsg)
+aux_source_directory(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/impl/nsg nsg_src)
+set(interface_src
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexNSG.cpp
         )
-
-if (NOT TARGET test_annoy)
-    add_executable(test_annoy test_annoy.cpp ${annoy_srcs} ${util_srcs})
+if (NOT TARGET test_nsg)
+    add_executable(test_nsg test_nsg.cpp ${interface_src} ${nsg_src} ${util_srcs} ${faiss_srcs})
 endif ()
-target_link_libraries(test_annoy ${depend_libs} ${unittest_libs} ${basic_libs})
+target_link_libraries(test_nsg ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_nsg DESTINATION unittest)
 
+################################################################################
 #<HNSW-TEST>
 set(hnsw_srcs
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexHNSW.cpp
         )
-
 if (NOT TARGET test_hnsw)
     add_executable(test_hnsw test_hnsw.cpp ${hnsw_srcs} ${util_srcs})
 endif ()
 target_link_libraries(test_hnsw ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_hnsw DESTINATION unittest)
 
-#<BinaryIDMAP-TEST>
-if (NOT TARGET test_binaryidmap)
-    add_executable(test_binaryidmap test_binaryidmap.cpp ${ivf_srcs} ${util_srcs})
-endif ()
-target_link_libraries(test_binaryidmap ${depend_libs} ${unittest_libs} ${basic_libs})
-
+################################################################################
 #<SPTAG-TEST>
 set(sptag_srcs
         ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.cpp
@@ -127,40 +182,20 @@ endif ()
 target_link_libraries(test_sptag
         SPTAGLibStatic
         ${depend_libs} ${unittest_libs} ${basic_libs})
+install(TARGETS test_sptag DESTINATION unittest)
 
-if (KNOWHERE_GPU_VERSION)
-    add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs})
-    target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})
-
-    add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs})
-    target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})
-endif ()
-
-#<INSTRUCTIONSET-TEST>
-if (NOT TARGET test_instructionset)
-    add_executable(test_instructionset test_instructionset.cpp)
-endif ()
-target_link_libraries(test_instructionset ${depend_libs} ${unittest_libs})
-
-if (NOT TARGET test_knowhere_common)
-    add_executable(test_knowhere_common test_common.cpp ${util_srcs})
+################################################################################
+#<ANNOY-TEST>
+set(annoy_srcs
+        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexAnnoy.cpp
+        )
+if (NOT TARGET test_annoy)
+    add_executable(test_annoy test_annoy.cpp ${annoy_srcs} ${util_srcs})
 endif ()
-target_link_libraries(test_knowhere_common ${depend_libs} ${unittest_libs} ${basic_libs})
-
-install(TARGETS test_ivf DESTINATION unittest)
-install(TARGETS test_hnsw DESTINATION unittest)
-install(TARGETS test_binaryivf DESTINATION unittest)
-install(TARGETS test_idmap DESTINATION unittest)
-install(TARGETS test_binaryidmap DESTINATION unittest)
-install(TARGETS test_sptag DESTINATION unittest)
-install(TARGETS test_knowhere_common DESTINATION unittest)
+target_link_libraries(test_annoy ${depend_libs} ${unittest_libs} ${basic_libs})
 install(TARGETS test_annoy DESTINATION unittest)
 
-if (KNOWHERE_GPU_VERSION)
-    install(TARGETS test_gpuresource DESTINATION unittest)
-    install(TARGETS test_customized_index DESTINATION unittest)
-endif ()
+
 #add_subdirectory(faiss_ori)
 #add_subdirectory(faiss_benchmark)
-add_subdirectory(test_nsg)
 

core/src/index/unittest/Helper.h

@@ -63,6 +63,7 @@ IndexFactory(const milvus::knowhere::IndexType& type, const milvus::knowhere::In
         }
 #endif
     }
+    return nullptr;
 }
 
 class ParamGenerator {
@@ -109,6 +110,7 @@ class ParamGenerator {
         } else {
             std::cout << "Invalid index type " << type << std::endl;
         }
+        return milvus::knowhere::Config();
     }
 };
 

core/src/index/unittest/test_binaryidmap.cpp

@@ -21,11 +21,11 @@ using ::testing::Combine;
 using ::testing::TestWithParam;
 using ::testing::Values;
 
-class BinaryIDMAPTest : public BinaryDataGen, public TestWithParam<std::string> {
+class BinaryIDMAPTest : public DataGen, public TestWithParam<std::string> {
  protected:
     void
     SetUp() override {
-        Init_with_binary_default();
+        Init_with_default(true);
         index_ = std::make_shared<milvus::knowhere::BinaryIDMAP>();
     }
 
@@ -40,7 +40,7 @@ INSTANTIATE_TEST_CASE_P(METRICParameters, BinaryIDMAPTest,
                         Values(std::string("JACCARD"), std::string("TANIMOTO"), std::string("HAMMING")));
 
 TEST_P(BinaryIDMAPTest, binaryidmap_basic) {
-    ASSERT_TRUE(!xb.empty());
+    ASSERT_TRUE(!xb_bin.empty());
 
     std::string MetricType = GetParam();
     milvus::knowhere::Config conf{

core/src/index/unittest/test_binaryivf.cpp

@@ -24,12 +24,12 @@ using ::testing::Combine;
 using ::testing::TestWithParam;
 using ::testing::Values;
 
-class BinaryIVFTest : public BinaryDataGen, public TestWithParam<std::string> {
+class BinaryIVFTest : public DataGen, public TestWithParam<std::string> {
  protected:
     void
     SetUp() override {
         std::string MetricType = GetParam();
-        Init_with_binary_default();
+        Init_with_default(true);
         //        nb = 1000000;
         //        nq = 1000;
         //        k = 1000;
@@ -58,14 +58,9 @@ INSTANTIATE_TEST_CASE_P(METRICParameters, BinaryIVFTest,
                         Values(std::string("JACCARD"), std::string("TANIMOTO"), std::string("HAMMING")));
 
 TEST_P(BinaryIVFTest, binaryivf_basic) {
-    assert(!xb.empty());
-
-    //    auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
-    //    index_->set_preprocessor(preprocessor);
+    assert(!xb_bin.empty());
 
     index_->Train(base_dataset, conf);
-    //    index_->set_index_model(model);
-    //    index_->Add(base_dataset, conf);
     EXPECT_EQ(index_->Count(), nb);
     EXPECT_EQ(index_->Dim(), dim);
 

core/src/index/unittest/test_nsg/CMakeLists.txt

-##############################
-#include_directories(/usr/local/include/gperftools)
-#link_directories(/usr/local/lib)
-
-add_definitions(-std=c++11 -O3 -lboost -march=native -Wall -DINFO)
-
-find_package(OpenMP REQUIRED)
-if (OpenMP_FOUND)
-    set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
-else ()
-    message(FATAL_ERROR "no OpenMP supprot")
-endif ()
-
-include_directories(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/impl/nsg)
-aux_source_directory(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/impl/nsg nsg_src)
-
-set(interface_src
-        ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexNSG.cpp
-        )
-
-if (NOT TARGET test_nsg)
-    add_executable(test_nsg test_nsg.cpp ${interface_src} ${nsg_src} ${util_srcs} ${ivf_srcs})
-endif ()
-
-target_link_libraries(test_nsg ${depend_libs} ${unittest_libs} ${basic_libs})
-##############################
-
-install(TARGETS test_nsg DESTINATION unittest)
\ No newline at end of file

core/src/index/unittest/utils.cpp

@@ -29,121 +29,88 @@ InitLog() {
 }
 
 void
-DataGen::Init_with_default() {
-    Generate(dim, nb, nq);
+DataGen::Init_with_default(const bool is_binary) {
+    Generate(dim, nb, nq, is_binary);
 }
 
 void
-BinaryDataGen::Init_with_binary_default() {
-    Generate(dim, nb, nq);
-}
-
-void
-DataGen::Generate(const int& dim, const int& nb, const int& nq) {
-    this->nb = nb;
-    this->nq = nq;
+DataGen::Generate(const int dim, const int nb, const int nq, const bool is_binary) {
     this->dim = dim;
-
-    GenAll(dim, nb, xb, ids, xids, nq, xq);
-    assert(xb.size() == (size_t)dim * nb);
-    assert(xq.size() == (size_t)dim * nq);
-
-    base_dataset = generate_dataset(nb, dim, xb.data(), ids.data());
-    query_dataset = generate_query_dataset(nq, dim, xq.data());
-    id_dataset = generate_id_dataset(nq, ids.data());
-    xid_dataset = generate_id_dataset(nq, xids.data());
-    xid_dataset->Set(milvus::knowhere::meta::DIM, (int64_t)dim);
-}
-
-void
-BinaryDataGen::Generate(const int& dim, const int& nb, const int& nq) {
     this->nb = nb;
     this->nq = nq;
-    this->dim = dim;
 
-    int64_t dim_x = dim / 8;
-    GenBinaryAll(dim_x, nb, xb, ids, xids, nq, xq);
-    assert(xb.size() == (size_t)dim_x * nb);
-    assert(xq.size() == (size_t)dim_x * nq);
-
-    base_dataset = generate_dataset(nb, dim, xb.data(), ids.data());
-    query_dataset = generate_query_dataset(nq, dim, xq.data());
-    id_dataset = generate_id_dataset(nq, ids.data());
-    xid_dataset = generate_id_dataset(nq, xids.data());
-}
-
-// not used
-#if 0
-knowhere::DatasetPtr
-DataGen::GenQuery(const int& nq) {
-    xq.resize(nq * dim);
-    for (int i = 0; i < nq * dim; ++i) {
-        xq[i] = xb[i];
+    if (!is_binary) {
+        GenAll(dim, nb, xb, ids, xids, nq, xq);
+        assert(xb.size() == (size_t)dim * nb);
+        assert(xq.size() == (size_t)dim * nq);
+
+        base_dataset = milvus::knowhere::GenDatasetWithIds(nb, dim, xb.data(), ids.data());
+        query_dataset = milvus::knowhere::GenDataset(nq, dim, xq.data());
+    } else {
+        int64_t dim_x = dim / 8;
+        GenAll(dim_x, nb, xb_bin, ids, xids, nq, xq_bin);
+        assert(xb_bin.size() == (size_t)dim_x * nb);
+        assert(xq_bin.size() == (size_t)dim_x * nq);
+
+        base_dataset = milvus::knowhere::GenDatasetWithIds(nb, dim, xb_bin.data(), ids.data());
+        query_dataset = milvus::knowhere::GenDataset(nq, dim, xq_bin.data());
     }
-    return generate_query_dataset(nq, dim, xq.data());
+
+    id_dataset = milvus::knowhere::GenDatasetWithIds(nq, dim, nullptr, ids.data());
+    xid_dataset = milvus::knowhere::GenDatasetWithIds(nq, dim, nullptr, xids.data());
 }
-#endif
 
 void
-GenAll(const int64_t dim, const int64_t& nb, std::vector<float>& xb, std::vector<int64_t>& ids,
-       std::vector<int64_t>& xids, const int64_t& nq, std::vector<float>& xq) {
+GenAll(const int64_t dim, const int64_t nb, std::vector<float>& xb, std::vector<int64_t>& ids,
+       std::vector<int64_t>& xids, const int64_t nq, std::vector<float>& xq) {
     xb.resize(nb * dim);
     xq.resize(nq * dim);
     ids.resize(nb);
     xids.resize(1);
-    GenAll(dim, nb, xb.data(), ids.data(), xids.data(), nq, xq.data());
-}
-
-void
-GenAll(const int64_t& dim, const int64_t& nb, float* xb, int64_t* ids, int64_t* xids, const int64_t& nq, float* xq) {
-    GenBase(dim, nb, xb, ids);
-    for (int64_t i = 0; i < nq * dim; ++i) {
-        xq[i] = xb[i];
-    }
-    xids[0] = 3;  // pseudo random
+    GenBase(dim, nb, xb.data(), ids.data(), nq, xq.data(), xids.data(), false);
 }
 
 void
-GenBinaryAll(const int64_t dim, const int64_t& nb, std::vector<uint8_t>& xb, std::vector<int64_t>& ids,
-             std::vector<int64_t>& xids, const int64_t& nq, std::vector<uint8_t>& xq) {
+GenAll(const int64_t dim, const int64_t nb, std::vector<uint8_t>& xb, std::vector<int64_t>& ids,
+       std::vector<int64_t>& xids, const int64_t nq, std::vector<uint8_t>& xq) {
     xb.resize(nb * dim);
     xq.resize(nq * dim);
     ids.resize(nb);
     xids.resize(1);
-    GenBinaryAll(dim, nb, xb.data(), ids.data(), xids.data(), nq, xq.data());
+    GenBase(dim, nb, xb.data(), ids.data(), nq, xq.data(), xids.data(), true);
 }
 
 void
-GenBinaryAll(const int64_t& dim, const int64_t& nb, uint8_t* xb, int64_t* ids, int64_t* xids, const int64_t& nq,
-             uint8_t* xq) {
-    GenBinaryBase(dim, nb, xb, ids);
-    for (int64_t i = 0; i < nq * dim; ++i) {
-        xq[i] = xb[i];
-    }
-    xids[0] = 3;  // pseudo random
-}
-
-void
-GenBase(const int64_t& dim, const int64_t& nb, float* xb, int64_t* ids) {
-    for (auto i = 0; i < nb; ++i) {
-        for (auto j = 0; j < dim; ++j) {
-            // p_data[i * d + j] = float(base + i);
-            xb[i * dim + j] = drand48();
+GenBase(const int64_t dim, const int64_t nb, const void* xb, int64_t* ids, const int64_t nq, const void* xq,
+        int64_t* xids, bool is_binary) {
+    if (!is_binary) {
+        float* xb_f = (float*)xb;
+        float* xq_f = (float*)xq;
+        for (auto i = 0; i < nb; ++i) {
+            for (auto j = 0; j < dim; ++j) {
+                xb_f[i * dim + j] = drand48();
+            }
+            xb_f[dim * i] += i / 1000.;
+            ids[i] = i;
         }
-        xb[dim * i] += i / 1000.;
-        ids[i] = i;
-    }
-}
-
-void
-GenBinaryBase(const int64_t& dim, const int64_t& nb, uint8_t* xb, int64_t* ids) {
-    for (auto i = 0; i < nb; ++i) {
-        for (auto j = 0; j < dim; ++j) {
-            // p_data[i * d + j] = float(base + i);
-            xb[i * dim + j] = (uint8_t)lrand48();
+        for (int64_t i = 0; i < nq * dim; ++i) {
+            xq_f[i] = xb_f[i];
+        }
+    } else {
+        uint8_t* xb_u = (uint8_t*)xb;
+        uint8_t* xq_u = (uint8_t*)xq;
+        for (auto i = 0; i < nb; ++i) {
+            for (auto j = 0; j < dim; ++j) {
+                xb_u[i * dim + j] = (uint8_t)lrand48();
+            }
+            xb_u[dim * i] += i / 1000.;
+            ids[i] = i;
+        }
+        for (int64_t i = 0; i < nq * dim; ++i) {
+            xq_u[i] = xb_u[i];
         }
-        ids[i] = i;
     }
+    xids[0] = 3;  // pseudo random
 }
 
 FileIOReader::FileIOReader(const std::string& fname) {
@@ -176,33 +143,6 @@ FileIOWriter::operator()(void* ptr, size_t size) {
     return size;
 }
 
-milvus::knowhere::DatasetPtr
-generate_dataset(int64_t nb, int64_t dim, const void* xb, const int64_t* ids) {
-    auto ret_ds = std::make_shared<milvus::knowhere::Dataset>();
-    ret_ds->Set(milvus::knowhere::meta::ROWS, nb);
-    ret_ds->Set(milvus::knowhere::meta::DIM, dim);
-    ret_ds->Set(milvus::knowhere::meta::TENSOR, xb);
-    ret_ds->Set(milvus::knowhere::meta::IDS, ids);
-    return ret_ds;
-}
-
-milvus::knowhere::DatasetPtr
-generate_query_dataset(int64_t nb, int64_t dim, const void* xb) {
-    auto ret_ds = std::make_shared<milvus::knowhere::Dataset>();
-    ret_ds->Set(milvus::knowhere::meta::ROWS, nb);
-    ret_ds->Set(milvus::knowhere::meta::DIM, dim);
-    ret_ds->Set(milvus::knowhere::meta::TENSOR, xb);
-    return ret_ds;
-}
-
-milvus::knowhere::DatasetPtr
-generate_id_dataset(int64_t nb, const int64_t* ids) {
-    auto ret_ds = std::make_shared<milvus::knowhere::Dataset>();
-    ret_ds->Set(milvus::knowhere::meta::ROWS, nb);
-    ret_ds->Set(milvus::knowhere::meta::IDS, ids);
-    return ret_ds;
-}
-
 void
 AssertAnns(const milvus::knowhere::DatasetPtr& result, const int nq, const int k, const CheckMode check_mode) {
     auto ids = result->Get<int64_t*>(milvus::knowhere::meta::IDS);

core/src/index/unittest/utils.h

@@ -23,13 +23,10 @@
 class DataGen {
  protected:
     void
-    Init_with_default();
+    Init_with_default(const bool is_binary = false);
 
     void
-    Generate(const int& dim, const int& nb, const int& nq);
-
-    milvus::knowhere::DatasetPtr
-    GenQuery(const int& nq);
+    Generate(const int dim, const int nb, const int nq, const bool is_binary = false);
 
  protected:
     int nb = 10000;
@@ -38,6 +35,8 @@ class DataGen {
     int k = 10;
     std::vector<float> xb;
     std::vector<float> xq;
+    std::vector<uint8_t> xb_bin;
+    std::vector<uint8_t> xq_bin;
     std::vector<int64_t> ids;
     std::vector<int64_t> xids;
     milvus::knowhere::DatasetPtr base_dataset = nullptr;
@@ -46,65 +45,21 @@ class DataGen {
     milvus::knowhere::DatasetPtr xid_dataset = nullptr;
 };
 
-class BinaryDataGen {
- protected:
-    void
-    Init_with_binary_default();
-
-    void
-    Generate(const int& dim, const int& nb, const int& nq);
-
-    milvus::knowhere::DatasetPtr
-    GenQuery(const int& nq);
-
- protected:
-    int nb = 10000;
-    int nq = 10;
-    int dim = 512;
-    int k = 10;
-    std::vector<uint8_t> xb;
-    std::vector<uint8_t> xq;
-    std::vector<int64_t> ids;
-    std::vector<int64_t> xids;
-    milvus::knowhere::DatasetPtr base_dataset = nullptr;
-    milvus::knowhere::DatasetPtr query_dataset = nullptr;
-    milvus::knowhere::DatasetPtr id_dataset = nullptr;
-    milvus::knowhere::DatasetPtr xid_dataset = nullptr;
-};
-
-extern void
-GenAll(const int64_t dim, const int64_t& nb, std::vector<float>& xb, std::vector<int64_t>& ids,
-       std::vector<int64_t>& xids, const int64_t& nq, std::vector<float>& xq);
-
 extern void
-GenAll(const int64_t& dim, const int64_t& nb, float* xb, int64_t* ids, int64_t* xids, const int64_t& nq, float* xq);
+GenAll(const int64_t dim, const int64_t nb, std::vector<float>& xb, std::vector<int64_t>& ids,
+       std::vector<int64_t>& xids, const int64_t nq, std::vector<float>& xq);
 
 extern void
-GenBinaryAll(const int64_t dim, const int64_t& nb, std::vector<uint8_t>& xb, std::vector<int64_t>& ids,
-             std::vector<int64_t>& xids, const int64_t& nq, std::vector<uint8_t>& xq);
+GenAll(const int64_t dim, const int64_t nb, std::vector<uint8_t>& xb, std::vector<int64_t>& ids,
+       std::vector<int64_t>& xids, const int64_t nq, std::vector<uint8_t>& xq);
 
 extern void
-GenBinaryAll(const int64_t& dim, const int64_t& nb, uint8_t* xb, int64_t* ids, int64_t* xids, const int64_t& nq,
-             uint8_t* xq);
-
-extern void
-GenBase(const int64_t& dim, const int64_t& nb, float* xb, int64_t* ids);
-
-extern void
-GenBinaryBase(const int64_t& dim, const int64_t& nb, uint8_t* xb, int64_t* ids);
+GenBase(const int64_t dim, const int64_t nb, const void* xb, int64_t* ids, const int64_t nq, const void* xq,
+        int64_t* xids, const bool is_binary);
 
 extern void
 InitLog();
 
-milvus::knowhere::DatasetPtr
-generate_dataset(int64_t nb, int64_t dim, const void* xb, const int64_t* ids);
-
-milvus::knowhere::DatasetPtr
-generate_query_dataset(int64_t nb, int64_t dim, const void* xb);
-
-milvus::knowhere::DatasetPtr
-generate_id_dataset(int64_t nb, const int64_t* ids);
-
 enum class CheckMode {
     CHECK_EQUAL = 0,
     CHECK_NOT_EQUAL = 1,