faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -18,21 +18,20 @@
 #include <algorithm>
 
 #include <faiss/utils/Heap.h>
-#include <faiss/utils/distances.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/utils.h>
 
 #include <faiss/Clustering.h>
 
 #include <faiss/utils/hamming.h>
 
-#include <faiss/impl/FaissAssert.h>
-
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
-
 #include <faiss/impl/ProductQuantizer.h>
-
-#include <faiss/impl/code_distance/code_distance.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
+#include <faiss/impl/simd_dispatch.h>
 
 namespace faiss {
 
@@ -427,7 +426,7 @@ void initialize_IVFPQ_precomputed_table(
     for (int m = 0; m < pq.M; m++)
         for (int j = 0; j < pq.ksub; j++)
             r_norms[m * pq.ksub + j] =
-                    fvec_norm_L2sqr(pq.get_centroids(m, j), pq.dsub);
+                    fvec_norm_L2sqr_dispatch(pq.get_centroids(m, j), pq.dsub);
 
     if (use_precomputed_table == 1) {
         precomputed_table.resize(nlist * pq.M * pq.ksub);
@@ -438,7 +437,7 @@ void initialize_IVFPQ_precomputed_table(
 
             float* tab = &precomputed_table[i * pq.M * pq.ksub];
             pq.compute_inner_prod_table(centroid.data(), tab);
-            fvec_madd(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd_dispatch(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
         }
     } else if (use_precomputed_table == 2) {
         const MultiIndexQuantizer* miq =
@@ -465,7 +464,7 @@ void initialize_IVFPQ_precomputed_table(
 
         for (size_t i = 0; i < cpq.ksub; i++) {
             float* tab = &precomputed_table[i * pq.M * pq.ksub];
-            fvec_madd(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd_dispatch(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
         }
     }
 }
@@ -628,7 +627,7 @@ struct QueryTables {
         // and dis0, the initial value
         ivfpq.quantizer->reconstruct(key, decoded_vec);
         // decoded_vec = centroid
-        float dis0 = fvec_inner_product(qi, decoded_vec, d);
+        float dis0 = fvec_inner_product_dispatch(qi, decoded_vec, d);
 
         if (polysemous_ht) {
             for (int i = 0; i < d; i++) {
@@ -657,7 +656,7 @@ struct QueryTables {
         } else if (use_precomputed_table == 1) {
             dis0 = coarse_dis;
 
-            fvec_madd(
+            fvec_madd_dispatch(
                     pq.M * pq.ksub,
                     ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
                     -2.0,
@@ -693,12 +692,12 @@ struct QueryTables {
 
                 if (polysemous_ht == 0) {
                     // sum up with query-specific table
-                    fvec_madd(Mf * pq.ksub, pc, -2.0, qtab, ltab);
+                    fvec_madd_dispatch(Mf * pq.ksub, pc, -2.0, qtab, ltab);
                     ltab += Mf * pq.ksub;
                     qtab += Mf * pq.ksub;
                 } else {
                     for (int m = cm * Mf; m < (cm + 1) * Mf; m++) {
-                        q_code[m] = fvec_madd_and_argmin(
+                        q_code[m] = fvec_madd_and_argmin_dispatch(
                                 pq.ksub, pc, -2, qtab, ltab);
                         pc += pq.ksub;
                         ltab += pq.ksub;
@@ -762,52 +761,31 @@ struct QueryTables {
     }
 };
 
-// This way of handling the selector is not optimal since all distances
-// are computed even if the id would filter it out.
 template <class C, bool use_sel>
-struct KnnSearchResults {
-    idx_t key;
-    const idx_t* ids;
-    const IDSelector* sel;
-
-    // heap params
-    size_t k;
-    float* heap_sim;
-    idx_t* heap_ids;
-
-    size_t nup;
-
-    inline bool skip_entry(idx_t j) {
-        return use_sel && !sel->is_member(ids[j]);
-    }
-
-    inline void add(idx_t j, float dis) {
-        if (C::cmp(heap_sim[0], dis)) {
-            idx_t id = ids ? ids[j] : lo_build(key, j);
-            heap_replace_top<C>(k, heap_sim, heap_ids, dis, id);
-            nup++;
-        }
-    }
-};
+struct WrappedSearchResult {
+    ResultHandler& res;
+    size_t nup = 0;
+    idx_t list_no;
 
-template <class C, bool use_sel>
-struct RangeSearchResults {
-    idx_t key;
     const idx_t* ids;
     const IDSelector* sel;
 
-    // wrapped result structure
-    float radius;
-    RangeQueryResult& rres;
+    WrappedSearchResult(
+            idx_t list_no,
+            const idx_t* ids,
+            const IDSelector* sel,
+            ResultHandler& res)
+            : res(res), list_no(list_no), ids(ids), sel(sel) {}
 
     inline bool skip_entry(idx_t j) {
         return use_sel && !sel->is_member(ids[j]);
     }
 
     inline void add(idx_t j, float dis) {
-        if (C::cmp(radius, dis)) {
-            idx_t id = ids ? ids[j] : lo_build(key, j);
-            rres.add(dis, id);
+        if (C::cmp(res.threshold, dis)) {
+            idx_t id = ids ? ids[j] : lo_build(this->list_no, j);
+            res.add_result(dis, id);
+            nup++;
         }
     }
 };
@@ -817,8 +795,9 @@ struct RangeSearchResults {
  * The scanning functions call their favorite precompute_*
  * function to precompute the tables they need.
  *****************************************************/
-template <typename IDType, MetricType METRIC_TYPE, class PQDecoder>
+template <typename IDType, MetricType METRIC_TYPE, class PQCodeDist>
 struct IVFPQScannerT : QueryTables {
+    using PQDecoder = typename PQCodeDist::PQDecoder;
     const uint8_t* list_codes;
     const IDType* list_ids;
     size_t list_size;
@@ -859,7 +838,7 @@ struct IVFPQScannerT : QueryTables {
     //         if (res.skip_entry(j)) {
     //             continue;
     //         }
-    //         float dis = dis0 + distance_single_code<PQDecoder>(
+    //         float dis = dis0 + PQCodeDist::distance_single_code(
     //             pq, sim_table, codes);
     //         res.add(j, dis);
     //     }
@@ -894,7 +873,7 @@ struct IVFPQScannerT : QueryTables {
                 float distance_1 = 0;
                 float distance_2 = 0;
                 float distance_3 = 0;
-                distance_four_codes<PQDecoder>(
+                PQCodeDist::distance_four_codes(
                         pq.M,
                         pq.nbits,
                         sim_table,
@@ -917,7 +896,7 @@ struct IVFPQScannerT : QueryTables {
 
         if (counter >= 1) {
             float dis = dis0 +
-                    distance_single_code<PQDecoder>(
+                    PQCodeDist::distance_single_code(
                                 pq.M,
                                 pq.nbits,
                                 sim_table,
@@ -926,7 +905,7 @@ struct IVFPQScannerT : QueryTables {
         }
         if (counter >= 2) {
             float dis = dis0 +
-                    distance_single_code<PQDecoder>(
+                    PQCodeDist::distance_single_code(
                                 pq.M,
                                 pq.nbits,
                                 sim_table,
@@ -935,7 +914,7 @@ struct IVFPQScannerT : QueryTables {
         }
         if (counter >= 3) {
             float dis = dis0 +
-                    distance_single_code<PQDecoder>(
+                    PQCodeDist::distance_single_code(
                                 pq.M,
                                 pq.nbits,
                                 sim_table,
@@ -979,7 +958,7 @@ struct IVFPQScannerT : QueryTables {
         if (by_residual) {
             if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
                 ivfpq.quantizer->reconstruct(key, residual_vec);
-                dis0 = fvec_inner_product(residual_vec, qi, d);
+                dis0 = fvec_inner_product_dispatch(residual_vec, qi, d);
             } else {
                 ivfpq.quantizer->compute_residual(qi, residual_vec, key);
             }
@@ -997,9 +976,9 @@ struct IVFPQScannerT : QueryTables {
 
             float dis;
             if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
-                dis = dis0 + fvec_inner_product(decoded_vec, qi, d);
+                dis = dis0 + fvec_inner_product_dispatch(decoded_vec, qi, d);
             } else {
-                dis = fvec_L2sqr(decoded_vec, dvec, d);
+                dis = fvec_L2sqr_dispatch(decoded_vec, dvec, d);
             }
             res.add(j, dis);
         }
@@ -1035,7 +1014,7 @@ struct IVFPQScannerT : QueryTables {
     //
     //                 float dis =
     //                         dis0 +
-    //                         distance_single_code<PQDecoder>(
+    //                         PQCodeDist::distance_single_code(
     //                             pq, sim_table, codes);
     //
     //                 res.add(j, dis);
@@ -1101,7 +1080,7 @@ struct IVFPQScannerT : QueryTables {
                 float distance_1 = dis0;
                 float distance_2 = dis0;
                 float distance_3 = dis0;
-                distance_four_codes<PQDecoder>(
+                PQCodeDist::distance_four_codes(
                         pq.M,
                         pq.nbits,
                         sim_table,
@@ -1132,7 +1111,7 @@ struct IVFPQScannerT : QueryTables {
             n_hamming_pass++;
 
             float dis = dis0 +
-                    distance_single_code<PQDecoder>(
+                    PQCodeDist::distance_single_code(
                                 pq.M,
                                 pq.nbits,
                                 sim_table,
@@ -1152,7 +1131,7 @@ struct IVFPQScannerT : QueryTables {
                 n_hamming_pass++;
 
                 float dis = dis0 +
-                        distance_single_code<PQDecoder>(
+                        PQCodeDist::distance_single_code(
                                     pq.M,
                                     pq.nbits,
                                     sim_table,
@@ -1199,8 +1178,8 @@ struct IVFPQScannerT : QueryTables {
  *
  * use_sel: store or ignore the IDSelector
  */
-template <MetricType METRIC_TYPE, class C, class PQDecoder, bool use_sel>
-struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
+template <MetricType METRIC_TYPE, class C, class PQCodeDist, bool use_sel>
+struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>,
                       InvertedListScanner {
     int precompute_mode;
     const IDSelector* sel;
@@ -1210,7 +1189,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
             bool store_pairs,
             int precompute_mode,
             const IDSelector* sel)
-            : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>(ivfpq, nullptr),
+            : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>(ivfpq, nullptr),
               precompute_mode(precompute_mode),
               sel(sel) {
         this->store_pairs = store_pairs;
@@ -1230,7 +1209,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
     float distance_to_code(const uint8_t* code) const override {
         assert(precompute_mode == 2);
         float dis = this->dis0 +
-                distance_single_code<PQDecoder>(
+                PQCodeDist::distance_single_code(
                             this->pq.M, this->pq.nbits, this->sim_table, code);
         return dis;
     }
@@ -1239,17 +1218,12 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
             size_t ncode,
             const uint8_t* codes,
             const idx_t* ids,
-            float* heap_sim,
-            idx_t* heap_ids,
-            size_t k) const override {
-        KnnSearchResults<C, use_sel> res = {
-                /* key */ this->key,
-                /* ids */ this->store_pairs ? nullptr : ids,
-                /* sel */ this->sel,
-                /* k */ k,
-                /* heap_sim */ heap_sim,
-                /* heap_ids */ heap_ids,
-                /* nup */ 0};
+            ResultHandler& handler) const override {
+        WrappedSearchResult<C, use_sel> res(
+                this->key,
+                this->store_pairs ? nullptr : ids,
+                this->sel,
+                handler);
 
         if (this->polysemous_ht > 0) {
             assert(precompute_mode == 2);
@@ -1265,85 +1239,53 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
         }
         return res.nup;
     }
-
-    void scan_codes_range(
-            size_t ncode,
-            const uint8_t* codes,
-            const idx_t* ids,
-            float radius,
-            RangeQueryResult& rres) const override {
-        RangeSearchResults<C, use_sel> res = {
-                /* key */ this->key,
-                /* ids */ this->store_pairs ? nullptr : ids,
-                /* sel */ this->sel,
-                /* radius */ radius,
-                /* rres */ rres};
-
-        if (this->polysemous_ht > 0) {
-            assert(precompute_mode == 2);
-            this->scan_list_polysemous(ncode, codes, res);
-        } else if (precompute_mode == 2) {
-            this->scan_list_with_table(ncode, codes, res);
-        } else if (precompute_mode == 1) {
-            this->scan_list_with_pointer(ncode, codes, res);
-        } else if (precompute_mode == 0) {
-            this->scan_on_the_fly_dist(ncode, codes, res);
-        } else {
-            FAISS_THROW_MSG("bad precomp mode");
-        }
-    }
 };
 
-template <class PQDecoder, bool use_sel>
-InvertedListScanner* get_InvertedListScanner1(
-        const IndexIVFPQ& index,
-        bool store_pairs,
-        const IDSelector* sel) {
-    if (index.metric_type == METRIC_INNER_PRODUCT) {
-        return new IVFPQScanner<
-                METRIC_INNER_PRODUCT,
-                CMin<float, idx_t>,
-                PQDecoder,
-                use_sel>(index, store_pairs, 2, sel);
-    } else if (index.metric_type == METRIC_L2) {
-        return new IVFPQScanner<
-                METRIC_L2,
-                CMax<float, idx_t>,
-                PQDecoder,
-                use_sel>(index, store_pairs, 2, sel);
-    }
-    return nullptr;
-}
-
-template <bool use_sel>
-InvertedListScanner* get_InvertedListScanner2(
-        const IndexIVFPQ& index,
-        bool store_pairs,
-        const IDSelector* sel) {
-    if (index.pq.nbits == 8) {
-        return get_InvertedListScanner1<PQDecoder8, use_sel>(
-                index, store_pairs, sel);
-    } else if (index.pq.nbits == 16) {
-        return get_InvertedListScanner1<PQDecoder16, use_sel>(
-                index, store_pairs, sel);
-    } else {
-        return get_InvertedListScanner1<PQDecoderGeneric, use_sel>(
-                index, store_pairs, sel);
-    }
-}
-
 } // anonymous namespace
 
 InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
         const IVFSearchParameters*) const {
-    if (sel) {
-        return get_InvertedListScanner2<true>(*this, store_pairs, sel);
-    } else {
-        return get_InvertedListScanner2<false>(*this, store_pairs, sel);
-    }
-    return nullptr;
+    return with_simd_level([&]<SIMDLevel SL>() -> InvertedListScanner* {
+        auto make =
+                [&]<class PQCodeDist, bool use_sel>() -> InvertedListScanner* {
+            if (metric_type == METRIC_INNER_PRODUCT) {
+                return new IVFPQScanner<
+                        METRIC_INNER_PRODUCT,
+                        CMin<float, idx_t>,
+                        PQCodeDist,
+                        use_sel>(*this, store_pairs, 2, sel);
+            } else if (metric_type == METRIC_L2) {
+                return new IVFPQScanner<
+                        METRIC_L2,
+                        CMax<float, idx_t>,
+                        PQCodeDist,
+                        use_sel>(*this, store_pairs, 2, sel);
+            } else {
+                FAISS_THROW_MSG("unsupported metric type");
+            }
+        };
+
+        auto with_decoder = [&]<bool use_sel>() -> InvertedListScanner* {
+            if (pq.nbits == 8) {
+                return make.template
+                operator()<PQCodeDistance<PQDecoder8, SL>, use_sel>();
+            } else if (pq.nbits == 16) {
+                return make.template
+                operator()<PQCodeDistance<PQDecoder16, SL>, use_sel>();
+            } else {
+                return make.template
+                operator()<PQCodeDistance<PQDecoderGeneric, SL>, use_sel>();
+            }
+        };
+
+        if (sel) {
+            return with_decoder.template operator()<true>();
+        } else {
+            return with_decoder.template operator()<false>();
+        }
+    });
 }
 
 IndexIVFPQStats indexIVFPQ_stats;

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -17,6 +17,7 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
 #include <faiss/utils/simdlib.h>
 
 #include <faiss/invlists/BlockInvertedLists.h>
@@ -307,6 +308,7 @@ struct IVFPQFastScanScanner : InvertedListScanner {
     const IndexIVFPQFastScan& index;
     AlignedTable<uint8_t> dis_tables;
     AlignedTable<uint16_t> biases;
+    std::vector<float> residual;
     std::array<float, 2> normalizers{};
     const float* xi = nullptr;
 
@@ -316,6 +318,7 @@ struct IVFPQFastScanScanner : InvertedListScanner {
             const IDSelector* sel)
             : InvertedListScanner(store_pairs, sel), index(index) {
         this->keep_max = is_similarity_metric(index.metric_type);
+        residual.resize(index.d);
     }
 
     void set_query(const float* query) override {
@@ -332,12 +335,40 @@ struct IVFPQFastScanScanner : InvertedListScanner {
         FastScanDistancePostProcessing empty_context{};
         index.compute_LUT_uint8(
                 1, xi, cq, dis_tables, biases, &normalizers[0], empty_context);
+        // used in distance_to_code
+        index.quantizer->compute_residual(
+                this->xi, residual.data(), this->list_no);
     }
 
-    float distance_to_code(const uint8_t* /* code */) const override {
-        // It's not really possible to implement a distance_to_code since codes
-        // for 32 database vectors are intermixed.
-        FAISS_THROW_MSG("not implemented");
+    float distance_to_code(const uint8_t* code) const override {
+        // directly use the PQ tables to compute the distance
+        const ProductQuantizer& pq = index.pq;
+        // when by_residual, codes are residuals so compare against query
+        // residual; otherwise codes are raw vectors so compare against raw
+        // query
+        const float* x = index.by_residual ? residual.data() : this->xi;
+        float accu = 0;
+        // implemented for all vector distances, although only L2 and IP are
+        // suppored by FastScan
+        with_VectorDistance(pq.dsub, index.metric_type, 0.0, [&](auto vd) {
+            int m;
+            for (m = 0; m + 1 < pq.M; m += 2) {
+                const float* cent;
+                uint8_t c = *code++;
+                cent = pq.get_centroids(m, c & 15);
+                accu += vd(cent, x);
+                x += pq.dsub;
+                cent = pq.get_centroids(m + 1, c >> 4);
+                accu += vd(cent, x);
+                x += pq.dsub;
+            }
+            if (m < pq.M) { // leftover
+                uint8_t c = *code++;
+                const float* cent = pq.get_centroids(m, c & 15);
+                accu += vd(cent, x);
+            }
+        });
+        return accu;
     }
 
     // Based on IVFFastScan search_implem_10, since it also deals with 1 query
@@ -388,7 +419,8 @@ struct IVFPQFastScanScanner : InvertedListScanner {
                 codes,
                 LUT,
                 *handler,
-                nullptr);
+                nullptr,
+                index.get_block_stride());
 
         // The handler is for the results of this iteration.
         // Then we need a second heap to combine across iterations.

data/vendor/faiss/faiss/IndexIVFPQR.cpp

@@ -11,6 +11,7 @@
 
 #include <cinttypes>
 
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
@@ -128,7 +129,7 @@ void IndexIVFPQR::search_preassigned(
         IndexIVFStats* stats) const {
     uint64_t t0;
     TIC;
-    size_t k_coarse = long(k * k_factor);
+    size_t k_coarse = long((size_t)k * k_factor);
     std::unique_ptr<idx_t[]> coarse_labels(new idx_t[k_coarse * n]);
     {
         // query with quantizer levels 1 and 2.

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -15,10 +15,15 @@
 #include <vector>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizer.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/expanded_scanners.h>
 
 namespace faiss {
 
+using rabitq_utils::SignBitFactorsWithError;
+
 IndexIVFRaBitQ::IndexIVFRaBitQ(
         Index* quantizer,
         const size_t d,
@@ -44,7 +49,7 @@ IndexIVFRaBitQ::IndexIVFRaBitQ() {
 void IndexIVFRaBitQ::train_encoder(
         idx_t n,
         const float* x,
-        const idx_t* assign) {
+        const idx_t* /*assign*/) {
     rabitq.train(n, x);
 }
 
@@ -182,7 +187,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     }
 
     /// following codes come from this inverted list
-    void set_list(idx_t list_no, float coarse_dis) override {
+    void set_list(idx_t list_no, float /*coarse_dis*/) override {
         this->list_no = list_no;
 
         reconstructed_centroid.resize(ivf_rabitq.d);
@@ -193,24 +198,31 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     }
 
     /// compute a single query-to-code distance
-    float distance_to_code(const uint8_t* code) const override {
+    float distance_to_code(const uint8_t* code) const final {
         return dc->distance_to_code(code);
     }
 
+    // redefiniing the scan_codes allows to inline the distance_to_code
+    // (this is unlikely to matter because it contains a virtual function call)
+    size_t scan_codes_1bit(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const {
+        return run_scan_codes(*this, list_size, codes, ids, handler);
+    }
+
     /// Override scan_codes to implement adaptive filtering for multi-bit codes
     size_t scan_codes(
             size_t list_size,
             const uint8_t* codes,
             const idx_t* ids,
-            float* simi,
-            idx_t* idxi,
-            size_t k) const override {
+            ResultHandler& handler) const override {
         size_t ex_bits = ivf_rabitq.rabitq.nb_bits - 1;
 
         // For 1-bit codes, use default implementation
         if (ex_bits == 0 || rabitq_dc == nullptr) {
-            return InvertedListScanner::scan_codes(
-                    list_size, codes, ids, simi, idxi, k);
+            return scan_codes_1bit(list_size, codes, ids, handler);
         }
 
         // Multi-bit: Two-stage search with adaptive filtering
@@ -233,33 +245,33 @@ struct RaBitInvertedListScanner : InvertedListScanner {
 
             local_1bit_evaluations++;
 
-            // Stage 1: Compute lower bound using 1-bit codes
-            float lower_bound = rabitq_dc->lower_bound_distance(codes);
-
-            // Stage 2: Adaptive filtering
-            // L2 (min-heap): filter if lower_bound < simi[0]
-            // IP (max-heap): filter if lower_bound > simi[0]
-            // Note: Using simi[0] directly (not cached) enables more aggressive
-            // filtering as the heap is updated with better candidates
-            bool should_refine = keep_max ? (lower_bound > simi[0])
-                                          : (lower_bound < simi[0]);
-
+            // Stage 1: Compute distance bound using 1-bit codes
+            // For L2 (min-heap): use lower_bound to safely skip if it's
+            //                    already worse than heap worst
+            // For IP (max-heap): use upper_bound because with a lower bound,
+            //                    we can't safely skip any candidate
+            float est_distance = rabitq_dc->distance_to_code_1bit(codes);
+
+            // Extract f_error and g_error for filtering
+            size_t code_size_base = (ivf_rabitq.d + 7) / 8;
+            const rabitq_utils::SignBitFactorsWithError* base_fac =
+                    reinterpret_cast<
+                            const rabitq_utils::SignBitFactorsWithError*>(
+                            codes + code_size_base);
+
+            bool should_refine = rabitq_utils::should_refine_candidate(
+                    est_distance,
+                    base_fac->f_error,
+                    rabitq_dc->g_error,
+                    handler.threshold,
+                    keep_max);
             if (should_refine) {
                 local_multibit_evaluations++;
                 // Lower bound is promising, compute full distance
                 float dis = distance_to_code(codes);
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
 
-                // Check if distance improves heap
-                bool improves_heap =
-                        keep_max ? (dis > simi[0]) : (dis < simi[0]);
-
-                if (improves_heap) {
-                    int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                    if (keep_max) {
-                        minheap_replace_top(k, simi, idxi, dis, id);
-                    } else {
-                        maxheap_replace_top(k, simi, idxi, dis, id);
-                    }
+                if (handler.add_result(dis, id)) {
                     nup++;
                 }
             }

data/vendor/faiss/faiss/IndexIVFRaBitQ.h

@@ -19,7 +19,7 @@
 namespace faiss {
 
 struct IVFRaBitQSearchParameters : IVFSearchParameters {
-    uint8_t qb = 0;
+    uint8_t qb = 4;
     bool centered = false;
 };
 
@@ -29,7 +29,7 @@ struct IndexIVFRaBitQ : IndexIVF {
 
     // the default number of bits to quantize a query with.
     // use '0' to disable quantization and use raw fp32 values.
-    uint8_t qb = 0;
+    uint8_t qb = 4;
 
     IndexIVFRaBitQ(
             Index* quantizer,