Fix scann filter search performance and cleanup ensure_topk_full config (#1179)

Signed-off-by: chasingegg <chao.gao@zilliz.com>

Author: chasingegg

src/index/ivf/ivf.cc

@@ -810,7 +810,7 @@ IvfIndexNode<DataType, IndexType>::Search(const DataSetPtr dataset, std::unique_
                     faiss::IVFSearchParameters base_search_params;
                     base_search_params.sel = id_selector;
                     base_search_params.nprobe = nprobe;
-                    base_search_params.ensure_topk_full = ivf_cfg.ensure_topk_full.value();
+                    base_search_params.ensure_topk_full = scann_cfg.ensure_topk_full.value();
                     if (base_search_params.ensure_topk_full) {
                         if (auto base_index_ptr = reinterpret_cast<faiss::IndexIVFPQFastScan*>(index_->base_index)) {
                             auto nlist = base_index_ptr->nlist;

src/index/ivf/ivf_config.h

@@ -27,7 +27,7 @@ class IvfConfig : public BaseConfig {
     CFG_INT nlist;
     CFG_INT nprobe;
     CFG_BOOL use_elkan;
-    CFG_BOOL ensure_topk_full;  // only take affect on temp index(IVF_FLAT_CC) now
+    CFG_BOOL ensure_topk_full;  // internal config, used for temp index
     CFG_INT max_empty_result_buckets;
     KNOHWERE_DECLARE_CONFIG(IvfConfig) {
         KNOWHERE_CONFIG_DECLARE_FIELD(nlist)
@@ -109,6 +109,7 @@ class ScannConfig : public IvfFlatConfig {
     CFG_INT reorder_k;
     CFG_BOOL with_raw_data;
     CFG_INT sub_dim;
+    CFG_BOOL ensure_topk_full;
     KNOHWERE_DECLARE_CONFIG(ScannConfig) {
         KNOWHERE_CONFIG_DECLARE_FIELD(reorder_k)
             .description("reorder k used for refining")
@@ -125,6 +126,10 @@ class ScannConfig : public IvfFlatConfig {
             .set_default(2)
             .for_train()
             .set_range(1, 65536);
+        KNOWHERE_CONFIG_DECLARE_FIELD(ensure_topk_full)
+            .set_default(false)
+            .description("whether to make sure topk results full")
+            .for_search();
     }
 
     Status

thirdparty/faiss/faiss/impl/pq4_fast_scan_search_1.cpp

@@ -125,12 +125,29 @@ void accumulate_fixed_blocks(
         ResultHandler& res,
         const Scaler& scaler) {
     constexpr int bbs = 32 * BB;
-    for (size_t j0 = 0; j0 < nb; j0 += bbs) {
+    for (size_t j0 = 0; j0 < nb; j0 += bbs, codes += bbs * nsq / 2) {
+        res.set_block_origin(0, j0);
+        // skip computing distances if all vectors inside a block are filtered out
+        if constexpr(has_sel_member_v<ResultHandler>) {
+            if (res.sel != nullptr) {  // we have filter here
+                bool skip_flag = true;
+                for (size_t jj = 0; jj < std::min<size_t>(bbs, res.ntotal - j0);
+                    jj++) {
+                    auto real_idx = res.adjust_id(0, jj);
+                    if (res.sel->is_member(real_idx)) {  // id is not filtered out, can not skip computing
+                        skip_flag = false;
+                        break;
+                    }
+                }
+
+                if (skip_flag) {
+                    continue;
+                }
+            }
+        }
         FixedStorageHandler<NQ, 2 * BB> res2;
         kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
-        res.set_block_origin(0, j0);
         res2.to_other_handler(res);
-        codes += bbs * nsq / 2;
     }
 }
 

thirdparty/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp

@@ -113,18 +113,6 @@ void kernel_accumulate_block(
     }
 }
 
-namespace {
-
-// a helper that checks whether a ResultHandler has a .sel member
-template<typename T, typename = void>
-struct has_sel_member : std::false_type {};
-template<typename T>
-struct has_sel_member<T, std::void_t<decltype(T::sel)>> : std::true_type {};
-template<typename T>
-inline constexpr bool has_sel_member_v = has_sel_member<T>::value;
-
-}
-
 // handle at most 4 blocks of queries
 template <int QBS, class ResultHandler, class Scaler>
 void accumulate_q_4step(

thirdparty/faiss/faiss/impl/simd_result_handlers.h

@@ -26,6 +26,18 @@
 
 namespace faiss {
 
+namespace {
+
+    // a helper that checks whether a ResultHandler has a .sel member
+    template<typename T, typename = void>
+    struct has_sel_member : std::false_type {};
+    template<typename T>
+    struct has_sel_member<T, std::void_t<decltype(T::sel)>> : std::true_type {};
+    template<typename T>
+    inline constexpr bool has_sel_member_v = has_sel_member<T>::value;
+    
+}
+
 struct SIMDResultHandler {
     // used to dispatch templates
     bool is_CMax = false;