faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp

@@ -8,8 +8,11 @@
 #include <faiss/IndexIVFRaBitQFastScan.h>
 
 #include <algorithm>
+#include <array>
 #include <cstdio>
+#include <memory>
 
+#include <faiss/impl/CodePackerRaBitQ.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/RaBitQUtils.h>
@@ -79,8 +82,6 @@ IndexIVFRaBitQFastScan::IndexIVFRaBitQFastScan(
     if (own_invlists) {
         replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true);
     }
-
-    flat_storage.clear();
 }
 
 // Constructor that converts an existing IndexIVFRaBitQ to FastScan format
@@ -97,41 +98,52 @@ IndexIVFRaBitQFastScan::IndexIVFRaBitQFastScan(
           rabitq(orig.rabitq) {}
 
 size_t IndexIVFRaBitQFastScan::compute_per_vector_storage_size() const {
-    const size_t ex_bits = rabitq.nb_bits - 1;
+    return rabitq_utils::compute_per_vector_storage_size(rabitq.nb_bits, d);
+}
 
-    if (ex_bits == 0) {
-        // 1-bit: only SignBitFactors (8 bytes)
-        return sizeof(SignBitFactors);
-    } else {
-        // Multi-bit: SignBitFactorsWithError + ExtraBitsFactors + ex-codes
-        return sizeof(SignBitFactorsWithError) + sizeof(ExtraBitsFactors) +
-                (d * ex_bits + 7) / 8;
-    }
+size_t IndexIVFRaBitQFastScan::code_packing_stride() const {
+    // Use code_size as stride to skip embedded factor data during packing
+    return code_size;
 }
 
-void IndexIVFRaBitQFastScan::preprocess_code_metadata(
-        idx_t n,
-        const uint8_t* flat_codes,
-        idx_t start_global_idx) {
-    // Unified approach: always use flat_storage for both 1-bit and multi-bit
-    const size_t storage_size = compute_per_vector_storage_size();
-    flat_storage.resize((start_global_idx + n) * storage_size);
+CodePacker* IndexIVFRaBitQFastScan::get_CodePacker() const {
+    return new CodePackerRaBitQ(M2, bbs, compute_per_vector_storage_size());
+}
+
+/*********************************************************
+ * postprocess_packed_codes: write auxiliary data into blocks
+ *********************************************************/
+
+void IndexIVFRaBitQFastScan::postprocess_packed_codes(
+        idx_t list_no,
+        size_t list_offset,
+        size_t n_added,
+        const uint8_t* flat_codes) {
+    auto* bil = dynamic_cast<BlockInvertedLists*>(invlists);
+    FAISS_THROW_IF_NOT(bil);
 
-    // Copy factors data directly to flat storage (no reordering needed)
+    uint8_t* block_data = bil->codes[list_no].data();
+    const size_t storage_size = compute_per_vector_storage_size();
     const size_t bit_pattern_size = (d + 7) / 8;
-    for (idx_t i = 0; i < n; i++) {
-        const uint8_t* code = flat_codes + i * code_size;
-        const uint8_t* source_factors_ptr = code + bit_pattern_size;
-        uint8_t* storage =
-                flat_storage.data() + (start_global_idx + i) * storage_size;
-        memcpy(storage, source_factors_ptr, storage_size);
+    const size_t packed_block_size = ((M2 + 1) / 2) * bbs;
+    const size_t full_block_size = get_block_stride();
+
+    for (size_t i = 0; i < n_added; i++) {
+        const uint8_t* src = flat_codes + i * code_size + bit_pattern_size;
+        uint8_t* dst = rabitq_utils::get_block_aux_ptr(
+                block_data,
+                list_offset + i,
+                bbs,
+                packed_block_size,
+                full_block_size,
+                storage_size);
+        memcpy(dst, src, storage_size);
     }
 }
 
-size_t IndexIVFRaBitQFastScan::code_packing_stride() const {
-    // Use code_size as stride to skip embedded factor data during packing
-    return code_size;
-}
+/*********************************************************
+ * train_encoder
+ *********************************************************/
 
 void IndexIVFRaBitQFastScan::train_encoder(
         idx_t n,
@@ -271,10 +283,11 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
             rotated_q,
             rotated_qq);
 
-    // Override query norm for inner product if original query is provided
     if (metric_type == MetricType::METRIC_INNER_PRODUCT &&
         original_query != nullptr) {
         query_factors.qr_norm_L2sqr = fvec_norm_L2sqr(original_query, d);
+        query_factors.q_dot_c = query_factors.qr_norm_L2sqr -
+                fvec_inner_product(original_query, residual, d);
     }
 
     const size_t ex_bits = rabitq.nb_bits - 1;
@@ -441,23 +454,22 @@ void IndexIVFRaBitQFastScan::reconstruct_from_offset(
         }
     }
 
-    // Get dp_multiplier directly from flat_storage
-    InvertedLists::ScopedIds list_ids(invlists, list_no);
-    idx_t global_id = list_ids[offset];
-
-    float dp_multiplier = 1.0f;
-    if (global_id >= 0) {
-        const size_t storage_size = compute_per_vector_storage_size();
-        const size_t storage_capacity = flat_storage.size() / storage_size;
-
-        if (static_cast<size_t>(global_id) < storage_capacity) {
-            const uint8_t* base_ptr =
-                    flat_storage.data() + global_id * storage_size;
-            const auto& base_factors =
-                    *reinterpret_cast<const SignBitFactors*>(base_ptr);
-            dp_multiplier = base_factors.dp_multiplier;
-        }
-    }
+    const size_t storage_size = compute_per_vector_storage_size();
+    const size_t packed_block_size = ((M2 + 1) / 2) * bbs;
+    const size_t full_block_size = get_block_stride();
+
+    InvertedLists::ScopedCodes list_block_codes(invlists, list_no);
+    const uint8_t* aux_ptr = rabitq_utils::get_block_aux_ptr(
+            list_block_codes.get(),
+            offset,
+            bbs,
+            packed_block_size,
+            full_block_size,
+            storage_size);
+
+    const auto& base_factors =
+            *reinterpret_cast<const SignBitFactors*>(aux_ptr);
+    const float dp_multiplier = base_factors.dp_multiplier;
 
     // Decode residual directly using dp_multiplier
     std::vector<float> residual(d);
@@ -573,7 +585,11 @@ IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::IVFRaBitQHeapHandler(
           nq(nq_val),
           k(k_val),
           context(ctx),
-          is_multibit(multibit) {
+          is_multibit(multibit),
+          storage_size(idx->compute_per_vector_storage_size()),
+          packed_block_size(((idx->M2 + 1) / 2) * idx->bbs),
+          full_block_size(idx->get_block_stride()),
+          packer(idx->get_CodePacker()) {
     current_list_no = 0;
     probe_indices.clear();
 
@@ -649,10 +665,13 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
 
         const float normalized_distance = d32tab[j] * one_a + bias;
 
-        // Get database factors from flat_storage
-        const size_t storage_size = index->compute_per_vector_storage_size();
-        const uint8_t* base_ptr =
-                index->flat_storage.data() + result_id * storage_size;
+        const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
+                list_codes_ptr,
+                idx_base + j,
+                index->bbs,
+                packed_block_size,
+                full_block_size,
+                storage_size);
 
         if (is_multibit) {
             // Track candidates actually considered for two-stage filtering
@@ -671,17 +690,18 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
                     index->qb,
                     index->d);
 
-            // Compute lower bound using error bound
-            float lower_bound =
-                    compute_lower_bound(dist_1bit, result_id, local_q, q);
-
             // Adaptive filtering: decide whether to compute full distance
             const bool is_similarity =
                     index->metric_type == MetricType::METRIC_INNER_PRODUCT;
-            bool should_refine = is_similarity
-                    ? (lower_bound > heap_dis[0])  // IP: keep if better
-                    : (lower_bound < heap_dis[0]); // L2: keep if better
 
+            float g_error = query_factors.g_error;
+
+            bool should_refine = rabitq_utils::should_refine_candidate(
+                    dist_1bit,
+                    full_factors.f_error,
+                    g_error,
+                    heap_dis[0],
+                    is_similarity);
             if (should_refine) {
                 local_multibit_evaluations++;
 
@@ -696,6 +716,7 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
                 if (Cfloat::cmp(heap_dis[0], dist_full)) {
                     heap_replace_top<Cfloat>(
                             k, heap_dis, heap_ids, dist_full, result_id);
+                    nup++;
                 }
             }
         } else {
@@ -715,6 +736,7 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
             if (Cfloat::cmp(heap_dis[0], adjusted_distance)) {
                 heap_replace_top<Cfloat>(
                         k, heap_dis, heap_ids, adjusted_distance, result_id);
+                nup++;
             }
         }
     }
@@ -732,6 +754,7 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::set_list_context(
         const std::vector<int>& probe_map) {
     current_list_no = list_no;
     probe_indices = probe_map;
+    list_codes_ptr = index->invlists->get_codes(list_no);
 }
 
 template <class C>
@@ -750,49 +773,23 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::end() {
     }
 }
 
-template <class C>
-float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::compute_lower_bound(
-        float dist_1bit,
-        size_t db_idx,
-        size_t local_q,
-        size_t global_q) const {
-    // Access f_error from SignBitFactorsWithError in flat storage
-    const size_t storage_size = index->compute_per_vector_storage_size();
-    const uint8_t* base_ptr =
-            index->flat_storage.data() + db_idx * storage_size;
-    const SignBitFactorsWithError& db_factors =
-            *reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
-    float f_error = db_factors.f_error;
-
-    // Get g_error from query factors
-    // Use local_q to access probe_indices (batch-local), global_q for storage
-    float g_error = 0.0f;
-    if (context && context->query_factors) {
-        size_t probe_rank = probe_indices[local_q];
-        size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe;
-        size_t storage_idx = global_q * nprobe + probe_rank;
-        g_error = context->query_factors[storage_idx].g_error;
-    }
-
-    // Compute error adjustment: f_error * g_error
-    float error_adjustment = f_error * g_error;
-
-    return dist_1bit - error_adjustment;
-}
-
 template <class C>
 float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::
         compute_full_multibit_distance(
-                size_t db_idx,
+                size_t /*db_idx*/,
                 size_t local_q,
                 size_t global_q,
                 size_t local_offset) const {
     const size_t ex_bits = index->rabitq.nb_bits - 1;
     const size_t dim = index->d;
 
-    const size_t storage_size = index->compute_per_vector_storage_size();
-    const uint8_t* base_ptr =
-            index->flat_storage.data() + db_idx * storage_size;
+    const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
+            list_codes_ptr,
+            local_offset,
+            index->bbs,
+            packed_block_size,
+            full_block_size,
+            storage_size);
 
     const size_t ex_code_size = (dim * ex_bits + 7) / 8;
     const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError);
@@ -809,8 +806,7 @@ float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::
     InvertedLists::ScopedCodes list_codes(index->invlists, list_no);
 
     std::vector<uint8_t> unpacked_code(index->code_size);
-    CodePackerPQ4 packer(index->M2, index->bbs);
-    packer.unpack_1(list_codes.get(), local_offset, unpacked_code.data());
+    packer->unpack_1(list_codes.get(), local_offset, unpacked_code.data());
     const uint8_t* sign_bits = unpacked_code.data();
 
     return rabitq_utils::compute_full_multibit_distance(
@@ -818,11 +814,164 @@ float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::
             ex_code,
             ex_fac,
             query_factors.rotated_q.data(),
-            query_factors.qr_to_c_L2sqr,
-            query_factors.qr_norm_L2sqr,
+            (index->metric_type == MetricType::METRIC_INNER_PRODUCT)
+                    ? query_factors.q_dot_c
+                    : query_factors.qr_to_c_L2sqr,
             dim,
             ex_bits,
             index->metric_type);
 }
 
+/*********************************************************
+ * IVFRaBitQFastScanScanner implementation
+ *********************************************************/
+
+namespace {
+
+/// Provides IVF scanner interface using FastScan's SIMD batch processing.
+struct IVFRaBitQFastScanScanner : InvertedListScanner {
+    static constexpr int impl = 10;
+    static constexpr size_t nq = 1;
+
+    const IndexIVFRaBitQFastScan& index;
+
+    AlignedTable<uint8_t> dis_tables;
+    AlignedTable<uint16_t> biases;
+    /// [scale, offset] for converting uint16 to float
+    std::array<float, 2> normalizers{};
+
+    const float* xi = nullptr;
+
+    QueryFactorsData query_factors;
+    FastScanDistancePostProcessing context;
+
+    std::unique_ptr<FlatCodesDistanceComputer> dc;
+    std::vector<float> centroid;
+
+    IVFRaBitQFastScanScanner(
+            const IndexIVFRaBitQFastScan& index,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), index(index) {
+        this->keep_max = is_similarity_metric(index.metric_type);
+    }
+
+    void set_query(const float* query) override {
+        this->xi = query;
+    }
+
+    void set_list(idx_t list_no, float coarse_dis) override {
+        this->list_no = list_no;
+
+        IndexIVFFastScan::CoarseQuantized cq{
+                .nprobe = 1,
+                .dis = &coarse_dis,
+                .ids = &list_no,
+        };
+
+        // Set up context for use in scan_codes
+        context = FastScanDistancePostProcessing{};
+        context.query_factors = &query_factors;
+        context.nprobe = 1;
+
+        index.compute_LUT_uint8(
+                1, xi, cq, dis_tables, biases, &normalizers[0], context);
+
+        // Set up distance computer for distance_to_code
+        centroid.resize(index.d);
+        index.quantizer->reconstruct(list_no, centroid.data());
+        dc.reset(index.rabitq.get_distance_computer(
+                index.qb, centroid.data(), index.centered));
+        dc->set_query(xi);
+    }
+
+    float distance_to_code(const uint8_t* code) const override {
+        FAISS_THROW_IF_NOT_MSG(
+                dc,
+                "set_query and set_list must be called before distance_to_code");
+        return dc->distance_to_code(code);
+    }
+
+   public:
+    size_t scan_codes(
+            size_t ntotal,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* distances,
+            idx_t* labels,
+            size_t k) const override {
+        // initialize the current iteration heap to the worst possible value of
+        // the prior loop
+        std::vector<float> curr_dists(k, distances[0]);
+        std::vector<idx_t> curr_labels(k, labels[0]);
+
+        std::unique_ptr<SIMDResultHandlerToFloat> handler(
+                index.make_knn_handler(
+                        !keep_max,
+                        impl,
+                        nq,
+                        k,
+                        curr_dists.data(),
+                        curr_labels.data(),
+                        sel,
+                        context,
+                        &normalizers[0]));
+
+        int qmap1[1] = {0};
+        handler->q_map = qmap1;
+        handler->begin(&normalizers[0]);
+
+        const uint8_t* LUT = dis_tables.get();
+        handler->dbias = biases.get();
+        handler->ntotal = ntotal;
+        handler->id_map = ids;
+
+        // RaBitQ needs list context for factor lookup
+        std::vector<int> probe_map = {0};
+        handler->set_list_context(list_no, probe_map);
+
+        pq4_accumulate_loop(
+                1,
+                roundup(ntotal, index.bbs),
+                index.bbs,
+                static_cast<int>(index.M2),
+                codes,
+                LUT,
+                *handler,
+                nullptr,
+                index.get_block_stride());
+
+        // Combine results across iterations
+        handler->end();
+        if (keep_max) {
+            minheap_addn(
+                    k,
+                    distances,
+                    labels,
+                    curr_dists.data(),
+                    curr_labels.data(),
+                    k);
+        } else {
+            maxheap_addn(
+                    k,
+                    distances,
+                    labels,
+                    curr_dists.data(),
+                    curr_labels.data(),
+                    k);
+        }
+
+        return handler->num_updates();
+    }
+};
+
+} // anonymous namespace
+
+InvertedListScanner* IndexIVFRaBitQFastScan::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
+    return new IVFRaBitQFastScanScanner(*this, store_pairs, sel);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <memory>
 #include <vector>
 
 #include <faiss/IndexIVFFastScan.h>
@@ -55,17 +56,6 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
     /// Use zero-centered scalar quantizer for queries
     bool centered = false;
 
-    /// Per-vector auxiliary data (1-bit codes stored separately in `codes`)
-    ///
-    /// 1-bit codes (sign bits) are stored in the inherited `codes` array from
-    /// IndexFastScan in packed FastScan format for SIMD processing.
-    ///
-    /// This flat_storage holds per-vector factors and refinement-bit codes:
-    /// Layout for 1-bit: [SignBitFactors (8 bytes)]
-    /// Layout for multi-bit: [SignBitFactorsWithError
-    /// (12B)][ref_codes][ExtraBitsFactors (8B)]
-    std::vector<uint8_t> flat_storage;
-
     // Constructors
 
     IndexIVFRaBitQFastScan();
@@ -94,16 +84,20 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             bool include_listnos = false) const override;
 
    protected:
-    /// Extract and store RaBitQ factors from encoded vectors
-    void preprocess_code_metadata(
-            idx_t n,
-            const uint8_t* flat_codes,
-            idx_t start_global_idx) override;
-
     /// Return code_size as stride to skip embedded factor data during packing
     size_t code_packing_stride() const override;
 
    public:
+    /// Return CodePackerRaBitQ with enlarged block size
+    CodePacker* get_CodePacker() const override;
+
+    /// Write per-vector auxiliary data into block auxiliary region
+    void postprocess_packed_codes(
+            idx_t list_no,
+            size_t list_offset,
+            size_t n_added,
+            const uint8_t* flat_codes) override;
+
     /// Reconstruct a single vector from an inverted list
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;
@@ -111,7 +105,7 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
     /// Override sa_decode to handle RaBitQ reconstruction
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 
-    /// Compute storage size per vector in flat_storage based on nb_bits
+    /// Compute per-vector auxiliary storage size based on nb_bits
     size_t compute_per_vector_storage_size() const;
 
    private:
@@ -166,6 +160,13 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             const FastScanDistancePostProcessing& context,
             const float* normalizers = nullptr) const override;
 
+    /// Get an InvertedListScanner for single-query scanning.
+    /// This provides compatibility with the standard IVF search interface
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs = false,
+            const IDSelector* sel = nullptr,
+            const IVFSearchParameters* params = nullptr) const override;
+
     /** SIMD result handler for IndexIVFRaBitQFastScan that applies
      * RaBitQ-specific distance corrections during batch processing.
      *
@@ -192,11 +193,19 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
         int64_t* heap_labels;  // [nq * k]
         const size_t nq, k;
         size_t current_list_no = 0;
+        const uint8_t* list_codes_ptr = nullptr; // raw block data for list
         std::vector<int>
                 probe_indices; // probe index for each query in current batch
         const FastScanDistancePostProcessing*
                 context;        // Processing context with query factors
         const bool is_multibit; // Whether to use multi-bit two-stage search
+        size_t nup = 0;         // Number of heap updates
+
+        // Cached block-layout constants (invariant for handler lifetime)
+        const size_t storage_size;
+        const size_t packed_block_size;
+        const size_t full_block_size;
+        std::unique_ptr<CodePacker> packer; // cached for unpack in hot path
 
         // Use float-based comparator for heap operations
         using Cfloat = typename std::conditional<
@@ -224,6 +233,10 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
 
         void end() override;
 
+        size_t num_updates() override {
+            return nup;
+        }
+
        private:
         /// Compute full multi-bit distance for a candidate vector (multi-bit
         /// only)
@@ -232,20 +245,10 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
         /// @param global_q Global query index (for storage indexing)
         /// @param local_offset Offset within the current inverted list
         float compute_full_multibit_distance(
-                size_t db_idx,
+                size_t /*db_idx*/,
                 size_t local_q,
                 size_t global_q,
                 size_t local_offset) const;
-
-        /// Compute lower bound using 1-bit distance and error bound (multi-bit
-        /// only)
-        /// @param local_q Batch-local query index (for probe_indices access)
-        /// @param global_q Global query index (for storage indexing)
-        float compute_lower_bound(
-                float dist_1bit,
-                size_t db_idx,
-                size_t local_q,
-                size_t global_q) const;
     };
 };
 

data/vendor/faiss/faiss/IndexLSH.cpp

@@ -86,12 +86,14 @@ void IndexLSH::train(idx_t n, const float* x) {
 
         for (idx_t i = 0; i < nbits; i++) {
             float* xi = transposed_x.get() + i * n;
-            // std::nth_element
-            std::sort(xi, xi + n);
-            if (n % 2 == 1)
-                thresholds[i] = xi[n / 2];
-            else
-                thresholds[i] = (xi[n / 2 - 1] + xi[n / 2]) / 2;
+            // Use nth_element (O(n)) instead of sort (O(n log n))
+            std::nth_element(xi, xi + n / 2, xi + n);
+            float median = xi[n / 2];
+            if (n % 2 == 0) {
+                std::nth_element(xi, xi + n / 2 - 1, xi + n);
+                median = (median + xi[n / 2 - 1]) / 2;
+            }
+            thresholds[i] = median;
         }
     }
     is_trained = true;