fix: BitsetViewIDSelectorHelper miss id_offset (zilliztech#948)

Signed-off-by: cqy123456 <[email protected]>
weiliu1031 · Nov 14, 2024 · 0bb86e3 · 0bb86e3
1 parent bdb98cb
commit 0bb86e3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/tests/ut/test_bruteforce.cc b/tests/ut/test_bruteforce.cc
@@ -201,15 +201,15 @@ TEST_CASE("Test Brute Force", "[binary vector]") {
 TEST_CASE("Test Brute Force with input ids", "[float vector]") {
     using Catch::Approx;
     const int64_t nb = 1000;
-    const int64_t nq = 1;
+    const int64_t nq = 10;
     const int64_t dim = 128;
     const int64_t k = 10;
     const knowhere::Json conf = {
         {knowhere::meta::DIM, dim},
         {knowhere::meta::METRIC_TYPE, "L2"},
         {knowhere::meta::TOPK, k},
     };
-    std::vector<int64_t> block_prefix = {0, 333, 500, 555, 1000};
+    std::vector<int64_t> block_prefix = {0, 111, 333, 500, 555, 666, 888, 1000};
 
     // generate filter id and data
     auto filter_bits = GenerateBitsetWithRandomTbitsSet(nb, 100);

diff --git a/thirdparty/faiss/faiss/utils/distances.cpp b/thirdparty/faiss/faiss/utils/distances.cpp
@@ -151,9 +151,10 @@ struct IDSelectorHelper {
 struct BitsetViewSelectorHelper {
     // todo aguzhva: use avx gather instruction
     const knowhere::BitsetView bitset;
+    const size_t id_offset = 0;
 
     inline bool is_member(const size_t idx) const {
-        return !bitset.test(idx);
+        return !bitset.test(idx + id_offset);
     }
 };
 
@@ -254,8 +255,9 @@ void exhaustive_inner_product_seq(
     if (const auto* bitsetview_sel = dynamic_cast<const knowhere::BitsetViewIDSelector*>(sel)) {
         // A specialized case for Knowhere
         auto bitset = bitsetview_sel->bitset_view;
+        auto id_offset = bitsetview_sel->id_offset;
         if (!bitset.empty()) {
-            BitsetViewSelectorHelper bitset_helper{bitset};
+            BitsetViewSelectorHelper bitset_helper{bitset, id_offset};
             exhaustive_inner_product_seq_impl<BlockResultHandler, BitsetViewSelectorHelper>(
                 x, y, d, nx, ny, res, bitset_helper);
             return;
@@ -368,8 +370,9 @@ void exhaustive_L2sqr_seq(
     if (const auto* bitsetview_sel = dynamic_cast<const knowhere::BitsetViewIDSelector*>(sel)) {
         // A specialized case for Knowhere
         auto bitset = bitsetview_sel->bitset_view;
+        auto id_offset = bitsetview_sel->id_offset;
         if (!bitset.empty()) {
-            BitsetViewSelectorHelper bitset_helper{bitset};
+            BitsetViewSelectorHelper bitset_helper{bitset, id_offset};
             exhaustive_L2sqr_seq_impl<BlockResultHandler, BitsetViewSelectorHelper>(
                 x, y, d, nx, ny, res, bitset_helper);
             return;
@@ -494,8 +497,9 @@ void exhaustive_cosine_seq(
     if (const auto* bitsetview_sel = dynamic_cast<const knowhere::BitsetViewIDSelector*>(sel)) {
         // A specialized case for Knowhere
         auto bitset = bitsetview_sel->bitset_view;
+        auto id_offset = bitsetview_sel->id_offset;
         if (!bitset.empty()) {
-            BitsetViewSelectorHelper bitset_helper{bitset};
+            BitsetViewSelectorHelper bitset_helper{bitset, id_offset};
             exhaustive_cosine_seq_impl<BlockResultHandler, BitsetViewSelectorHelper>(
                 x, y, y_norms, d, nx, ny, res, bitset_helper);
             return;