Add max_empty_result_buckerts for ScaNN range search (#467)

Signed-off-by: chasingegg <[email protected]>

Author: chasingegg

src/index/ivf/ivf.cc

@@ -769,14 +769,11 @@ IvfIndexNode<DataType, IndexType>::RangeSearch(const DataSet& dataset, const Con
                     }
 
                     // todo aguzhva: this is somewhat alogical. Refactor?
-                    faiss::IVFSearchParameters base_search_params;
-                    base_search_params.max_empty_result_buckets = ivf_cfg.max_empty_result_buckets.value();
-                    base_search_params.sel = id_selector;
-
-                    faiss::IndexScaNNSearchParameters scann_search_params;
-                    scann_search_params.base_index_params = &base_search_params;
+                    faiss::IVFSearchParameters search_params;
+                    search_params.max_empty_result_buckets = ivf_cfg.max_empty_result_buckets.value();
+                    search_params.sel = id_selector;
 
-                    index_->range_search(1, cur_query, radius, &res, &scann_search_params);
+                    index_->range_search(1, cur_query, radius, &res, &search_params);
                 } else {
                     auto cur_query = (const float*)xq + index * dim;
                     if (is_cosine) {

thirdparty/faiss/faiss/IndexIVFFastScan.cpp

@@ -1369,6 +1369,7 @@ void IndexIVFFastScan::range_search_implem_12(
 
     const size_t nprobe = params ? params->nprobe : this->nprobe;
     const size_t max_codes = params ? params->max_codes : this->max_codes;
+    const int max_empty_result_buckets = params ? params->max_empty_result_buckets : 1;
     const IDSelector* sel = params ? params->sel : nullptr;
     const SearchParameters* quantizer_params =
             params ? params->quantizer_params : nullptr;
@@ -1445,6 +1446,8 @@ void IndexIVFFastScan::range_search_implem_12(
 
     size_t i0 = 0;
     uint64_t t_copy_pack = 0, t_scan = 0;
+    handler->in_range_num = 0;
+    int ndup = 0;
     while (i0 < qcs.size()) {
         uint64_t tt0 = get_cy();
 
@@ -1498,7 +1501,7 @@ void IndexIVFFastScan::range_search_implem_12(
         handler->ntotal = list_size;
         handler->q_map = q_map.data();
         handler->id_map = ids.get();
-        handler->in_range_num = 0;
+        int prev_in_range_num = handler->in_range_num;
         uint64_t tt1 = get_cy();
 
         pq4_accumulate_loop_qbs(
@@ -1509,7 +1512,12 @@ void IndexIVFFastScan::range_search_implem_12(
                 LUT.get(),
                 *(handler.get()),
                 scaler);
-        if (handler->in_range_num <= 0) {
+        if (handler->in_range_num == prev_in_range_num) {
+            ndup++;
+        } else {
+            ndup = 0;
+        }
+        if (ndup == max_empty_result_buckets) {
             break;
         }
 

thirdparty/faiss/faiss/IndexScaNN.cpp

@@ -202,23 +202,21 @@ void IndexScaNN::range_search(
 
     FAISS_THROW_IF_NOT(is_trained);
 
-    const IndexScaNNSearchParameters* params = nullptr;
+    const IVFSearchParameters* params = nullptr;
     if (params_in) {
-        params = dynamic_cast<const IndexScaNNSearchParameters*>(params_in);
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
         FAISS_THROW_IF_NOT_MSG(params, "IndexScaNN params have incorrect type");
     }
 
-    SearchParameters* base_index_params = 
-        (params != nullptr) ? params->base_index_params : nullptr;
-
-
     auto base = dynamic_cast<const IndexIVFPQFastScan*>(base_index);
     FAISS_THROW_IF_NOT(base);
 
     IVFSearchParameters ivf_search_params;
     ivf_search_params.nprobe = base->nlist;
+    ivf_search_params.max_empty_result_buckets = params->max_empty_result_buckets;
     // todo aguzhva: this is somewhat hacky
-    ivf_search_params.sel = base_index_params->sel;
+    ivf_search_params.sel = params->sel;
+
     base->range_search(n, x, radius, result, &ivf_search_params);
 
     // nothing to refine, directly return the result