From 3a0cc24d57cd2d05d8b70de2c5a3b9b3f674873e Mon Sep 17 00:00:00 2001
From: "yudong.cai" <yudong.cai@zilliz.com>
Date: Fri, 29 Nov 2019 14:31:53 +0800
Subject: [PATCH] #331 update exception handle when search fail

---
 .../index/vector_index/IndexGPUIVF.cpp        | 11 +--
 .../knowhere/index/vector_index/IndexIVF.cpp  | 71 ++++++++++---------
 2 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp
index d496d8085c..872f05aa67 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp
+++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp
@@ -131,14 +131,9 @@ GPUIVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, i
     if (auto device_index = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) {
         auto search_cfg = std::dynamic_pointer_cast<IVFCfg>(cfg);
         device_index->nprobe = search_cfg->nprobe;
-        //        assert(device_index->getNumProbes() == search_cfg->nprobe);
-
-        try {
-            ResScope rs(res_, gpu_id_);
-            device_index->search(n, (float*)data, k, distances, labels);
-        } catch (faiss::FaissException& e) {
-            KNOWHERE_THROW_MSG(e.what());
-        }
+        //assert(device_index->getNumProbes() == search_cfg->nprobe);
+        ResScope rs(res_, gpu_id_);
+        device_index->search(n, (float*)data, k, distances, labels);
     } else {
         KNOWHERE_THROW_MSG("Not a GpuIndexIVF type.");
     }
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp
index f99c6e3972..884aa8c31d 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp
+++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp
@@ -119,42 +119,49 @@ IVF::Search(const DatasetPtr& dataset, const Config& config) {
 
     GETTENSOR(dataset)
 
-    auto elems = rows * search_cfg->k;
-    auto res_ids = (int64_t*)malloc(sizeof(int64_t) * elems);
-    auto res_dis = (float*)malloc(sizeof(float) * elems);
+    try {
+        auto elems = rows * search_cfg->k;
+        auto res_ids = (int64_t *) malloc(sizeof(int64_t) * elems);
+        auto res_dis = (float *) malloc(sizeof(float) * elems);
 
-    search_impl(rows, (float*)p_data, search_cfg->k, res_dis, res_ids, config);
+        search_impl(rows, (float *) p_data, search_cfg->k, res_dis, res_ids, config);
 
-    //    std::stringstream ss_res_id, ss_res_dist;
-    //    for (int i = 0; i < 10; ++i) {
-    //        printf("%llu", res_ids[i]);
-    //        printf("\n");
-    //        printf("%.6f", res_dis[i]);
-    //        printf("\n");
-    //        ss_res_id << res_ids[i] << " ";
-    //        ss_res_dist << res_dis[i] << " ";
-    //    }
-    //    std::cout << std::endl << "after search: " << std::endl;
-    //    std::cout << ss_res_id.str() << std::endl;
-    //    std::cout << ss_res_dist.str() << std::endl << std::endl;
+        //    std::stringstream ss_res_id, ss_res_dist;
+        //    for (int i = 0; i < 10; ++i) {
+        //        printf("%llu", res_ids[i]);
+        //        printf("\n");
+        //        printf("%.6f", res_dis[i]);
+        //        printf("\n");
+        //        ss_res_id << res_ids[i] << " ";
+        //        ss_res_dist << res_dis[i] << " ";
+        //    }
+        //    std::cout << std::endl << "after search: " << std::endl;
+        //    std::cout << ss_res_id.str() << std::endl;
+        //    std::cout << ss_res_dist.str() << std::endl << std::endl;
 
-    //    auto id_buf = MakeMutableBufferSmart((uint8_t*)res_ids, sizeof(int64_t) * elems);
-    //    auto dist_buf = MakeMutableBufferSmart((uint8_t*)res_dis, sizeof(float) * elems);
-    //
-    //    std::vector<BufferPtr> id_bufs{nullptr, id_buf};
-    //    std::vector<BufferPtr> dist_bufs{nullptr, dist_buf};
-    //
-    //    auto int64_type = std::make_shared<arrow::Int64Type>();
-    //    auto float_type = std::make_shared<arrow::FloatType>();
-    //
-    //    auto id_array_data = arrow::ArrayData::Make(int64_type, elems, id_bufs);
-    //    auto dist_array_data = arrow::ArrayData::Make(float_type, elems, dist_bufs);
-    //
-    //    auto ids = std::make_shared<NumericArray<arrow::Int64Type>>(id_array_data);
-    //    auto dists = std::make_shared<NumericArray<arrow::FloatType>>(dist_array_data);
-    //    std::vector<ArrayPtr> array{ids, dists};
+        //    auto id_buf = MakeMutableBufferSmart((uint8_t*)res_ids, sizeof(int64_t) * elems);
+        //    auto dist_buf = MakeMutableBufferSmart((uint8_t*)res_dis, sizeof(float) * elems);
+        //
+        //    std::vector<BufferPtr> id_bufs{nullptr, id_buf};
+        //    std::vector<BufferPtr> dist_bufs{nullptr, dist_buf};
+        //
+        //    auto int64_type = std::make_shared<arrow::Int64Type>();
+        //    auto float_type = std::make_shared<arrow::FloatType>();
+        //
+        //    auto id_array_data = arrow::ArrayData::Make(int64_type, elems, id_bufs);
+        //    auto dist_array_data = arrow::ArrayData::Make(float_type, elems, dist_bufs);
+        //
+        //    auto ids = std::make_shared<NumericArray<arrow::Int64Type>>(id_array_data);
+        //    auto dists = std::make_shared<NumericArray<arrow::FloatType>>(dist_array_data);
+        //    std::vector<ArrayPtr> array{ids, dists};
 
-    return std::make_shared<Dataset>((void*)res_ids, (void*)res_dis);
+        return std::make_shared<Dataset>((void *) res_ids, (void *) res_dis);
+
+    } catch (faiss::FaissException& e) {
+        KNOWHERE_THROW_MSG(e.what());
+    } catch (std::exception& e) {
+        KNOWHERE_THROW_MSG(e.what());
+    }
 }
 
 void