From c884fa6d328bd75d72a3ea259e8c08ce46db9624 Mon Sep 17 00:00:00 2001
From: chasingegg <chao.gao@zilliz.com>
Date: Wed, 23 Apr 2025 01:54:06 +0800
Subject: [PATCH] Fix scann filter search performance and cleanup
 ensure_topk_full config

Signed-off-by: chasingegg <chao.gao@zilliz.com>
---
 src/index/ivf/ivf.cc                          |  2 +-
 src/index/ivf/ivf_config.h                    |  7 +++++-
 .../faiss/impl/pq4_fast_scan_search_1.cpp     | 23 ++++++++++++++++---
 .../faiss/impl/pq4_fast_scan_search_qbs.cpp   | 12 ----------
 .../faiss/faiss/impl/simd_result_handlers.h   | 12 ++++++++++
 5 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc
index baf01df88..88108436c 100644
--- a/src/index/ivf/ivf.cc
+++ b/src/index/ivf/ivf.cc
@@ -805,7 +805,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;
diff --git a/src/index/ivf/ivf_config.h b/src/index/ivf/ivf_config.h
index 93ef77642..112f3f708 100644
--- a/src/index/ivf/ivf_config.h
+++ b/src/index/ivf/ivf_config.h
@@ -21,7 +21,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)
@@ -103,6 +103,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")
@@ -119,6 +120,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
diff --git a/thirdparty/faiss/faiss/impl/pq4_fast_scan_search_1.cpp b/thirdparty/faiss/faiss/impl/pq4_fast_scan_search_1.cpp
index ca41f287f..5dcd6140c 100644
--- a/thirdparty/faiss/faiss/impl/pq4_fast_scan_search_1.cpp
+++ b/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;
     }
 }
 
diff --git a/thirdparty/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp b/thirdparty/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp
index d2ab08f4c..e88e44145 100644
--- a/thirdparty/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp
+++ b/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(
diff --git a/thirdparty/faiss/faiss/impl/simd_result_handlers.h b/thirdparty/faiss/faiss/impl/simd_result_handlers.h
index 88c7ce554..efc77d628 100644
--- a/thirdparty/faiss/faiss/impl/simd_result_handlers.h
+++ b/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;