faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/IndexIVFSQFastScan.cpp

@@ -0,0 +1,855 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/IndexIVFSQFastScan.h>
+#include <faiss/impl/sq_fastscan_utils.h>
+
+#include <omp.h>
+
+#include <algorithm>
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/invlists/BlockInvertedLists.h>
+#include <faiss/invlists/DirectMap.h>
+#include <faiss/invlists/InvertedLists.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/hamming.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+using namespace faiss::sq_fastscan;
+
+// -----------------------------------------------------------------------
+// Constructors
+// -----------------------------------------------------------------------
+
+IndexIVFSQFastScan::IndexIVFSQFastScan(
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
+        ScalarQuantizer::QuantizerType qtype,
+        MetricType metric,
+        int bbs_in,
+        bool by_residual_in)
+        : IndexIVFFastScan(quantizer_in, d_in, nlist_in, 0, metric, false),
+          sq(d_in, qtype) {
+    if (is_fallback(qtype)) {
+        by_residual = by_residual_in;
+        if (ScalarQuantizer::TurboQuantRefine::is_turboq_full(qtype)) {
+            by_residual = false;
+        }
+    } else {
+        // FastScan path: by_residual=false because our 2D LUT cannot
+        // handle per-probe residuals (same rationale as IndexIVFPQFastScan).
+        by_residual = false;
+    }
+
+    if (is_fallback(qtype)) {
+        // Fallback: use standard ArrayInvertedLists with full SQ codes.
+        // Search uses IndexIVF's InvertedListScanner path.
+        code_size = sq.code_size;
+        replace_invlists(new ArrayInvertedLists(nlist_in, code_size), true);
+    } else {
+        size_t M_in = d_in;
+        M = M_in;
+        nbits = 4;
+        bbs = bbs_in;
+        ksub = 16;
+        M2 = roundup(M_in, 2);
+        code_size = M2 / 2;
+
+        replace_invlists(
+                new BlockInvertedLists(nlist_in, get_CodePacker()), true);
+
+        if (needs_rerank(qtype)) {
+            lo_codes_invlists =
+                    new ArrayInvertedLists(nlist_in, sq_rerank_size(d, qtype));
+            direct_map.set_type(DirectMap::Hashtable, invlists, 0);
+        }
+    }
+}
+
+IndexIVFSQFastScan::IndexIVFSQFastScan() = default;
+
+IndexIVFSQFastScan::IndexIVFSQFastScan(
+        const IndexIVFScalarQuantizer& orig,
+        int bbs_in)
+        : IndexIVFFastScan(
+                  orig.quantizer,
+                  orig.d,
+                  orig.nlist,
+                  0,
+                  orig.metric_type,
+                  false),
+          sq(orig.sq) {
+    ntotal = orig.ntotal;
+    is_trained = orig.is_trained;
+    nprobe = orig.nprobe;
+
+    if (is_fallback(sq.qtype)) {
+        by_residual = orig.by_residual;
+        code_size = sq.code_size;
+        replace_invlists(new ArrayInvertedLists(nlist, code_size), true);
+
+        for (size_t list_no = 0; list_no < nlist; list_no++) {
+            size_t list_size = orig.invlists->list_size(list_no);
+            if (list_size == 0) {
+                continue;
+            }
+            InvertedLists::ScopedCodes codes(orig.invlists, list_no);
+            InvertedLists::ScopedIds ids(orig.invlists, list_no);
+            invlists->add_entries(list_no, list_size, ids.get(), codes.get());
+        }
+        return;
+    }
+
+    FAISS_THROW_IF_NOT_MSG(
+            !orig.by_residual,
+            "IndexIVFSQFastScan: conversion from IndexIVFScalarQuantizer "
+            "with by_residual=true is not supported for native/rerank "
+            "types. Set orig.by_residual=false and retrain, or construct "
+            "IndexIVFSQFastScan directly.");
+
+    by_residual = false;
+
+    size_t M_in = d;
+    M = M_in;
+    nbits = 4;
+    bbs = bbs_in;
+    ksub = 16;
+    M2 = roundup(M_in, 2);
+    code_size = M2 / 2;
+
+    replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true);
+
+    if (needs_rerank(sq.qtype)) {
+        lo_codes_invlists = new ArrayInvertedLists(nlist, lo_code_size());
+        direct_map.set_type(DirectMap::Hashtable, invlists, 0);
+    }
+
+    for (size_t list_no = 0; list_no < nlist; list_no++) {
+        size_t list_size = orig.invlists->list_size(list_no);
+        if (list_size == 0) {
+            continue;
+        }
+
+        InvertedLists::ScopedCodes orig_codes(orig.invlists, list_no);
+        InvertedLists::ScopedIds orig_ids(orig.invlists, list_no);
+
+        // Decode to float and re-encode for the scan.
+        std::vector<float> recon(list_size * d);
+        sq.decode(orig_codes.get(), recon.data(), list_size);
+
+        std::vector<uint8_t> flat4(list_size * M2 / 2, 0);
+        if (is_native_4bit(sq.qtype)) {
+            sq.compute_codes(recon.data(), flat4.data(), list_size);
+        } else {
+            // Top 4 bits feed the scan, the rest are kept for rerank.
+            std::vector<uint8_t> lo(list_size * lo_code_size(), 0);
+            float_to_split_codes(
+                    recon.data(),
+                    flat4.data(),
+                    lo.data(),
+                    list_size,
+                    d,
+                    M2,
+                    sq);
+            lo_codes_invlists->add_entries(
+                    list_no, list_size, orig_ids.get(), lo.data());
+            for (size_t j = 0; j < list_size; j++) {
+                direct_map.add_single_id(orig_ids.get()[j], list_no, j);
+            }
+        }
+
+        BlockInvertedLists* bil = dynamic_cast<BlockInvertedLists*>(invlists);
+        FAISS_THROW_IF_NOT(bil);
+
+        size_t nb2 = roundup(list_size, bbs);
+        AlignedTable<uint8_t> packed(nb2 * M2 / 2);
+        pq4_pack_codes(flat4.data(), list_size, M, nb2, bbs, M2, packed.get());
+        bil->add_entries(list_no, list_size, orig_ids.get(), packed.get());
+    }
+}
+
+IndexIVFSQFastScan::~IndexIVFSQFastScan() {
+    delete lo_codes_invlists;
+}
+
+// -----------------------------------------------------------------------
+// Virtual method implementations
+// -----------------------------------------------------------------------
+
+size_t IndexIVFSQFastScan::fast_scan_code_size() const {
+    return M2 / 2;
+}
+
+void IndexIVFSQFastScan::train_encoder(
+        idx_t n,
+        const float* x,
+        const idx_t* /*assign*/) {
+    sq.train(n, x);
+}
+
+idx_t IndexIVFSQFastScan::train_encoder_num_vectors() const {
+    return 100000;
+}
+
+void IndexIVFSQFastScan::encode_vectors(
+        idx_t n,
+        const float* x,
+        const idx_t* list_nos,
+        uint8_t* codes_out,
+        bool include_listnos) const {
+    std::vector<float> residuals;
+    const float* to_encode = x;
+
+    if (by_residual) {
+        residuals.resize(n * d);
+        for (idx_t i = 0; i < n; i++) {
+            if (list_nos[i] >= 0) {
+                quantizer->compute_residual(
+                        x + i * d, residuals.data() + i * d, list_nos[i]);
+            } else {
+                memset(residuals.data() + i * d, 0, sizeof(float) * d);
+            }
+        }
+        to_encode = residuals.data();
+    }
+
+    if (is_fallback(sq.qtype)) {
+        sq.compute_codes(to_encode, codes_out, n);
+    } else if (is_native_4bit(sq.qtype)) {
+        sq.compute_codes(to_encode, codes_out, n);
+    } else {
+        // add_core keeps the leftover bits in lo_codes_invlists.
+        std::vector<uint8_t> lo_scratch(size_t(n) * lo_code_size());
+        float_to_split_codes(
+                to_encode, codes_out, lo_scratch.data(), n, d, M2, sq);
+    }
+
+    if (include_listnos) {
+        size_t coarse_size = coarse_code_size();
+        for (idx_t i = n - 1; i >= 0; i--) {
+            uint8_t* code = codes_out + i * (coarse_size + code_size);
+            memmove(code + coarse_size, codes_out + i * code_size, code_size);
+            encode_listno(list_nos[i], code);
+        }
+    }
+}
+
+bool IndexIVFSQFastScan::lookup_table_is_3d() const {
+    // For SQ, the LUT doesn't depend on which list we're scanning
+    // (unlike PQ with residuals and precomputed tables).
+    return false;
+}
+
+void IndexIVFSQFastScan::compute_LUT(
+        size_t n,
+        const float* x,
+        const CoarseQuantized& /*cq*/,
+        AlignedTable<float>& dis_tables,
+        AlignedTable<float>& /*biases*/,
+        const FastScanDistancePostProcessing&) const {
+    size_t dim = d;
+    size_t ksub_val = 16;
+    size_t dim12 = dim * ksub_val;
+
+    // Build reconstruction table: for each dimension, the 16 centroid values.
+    bool is_uniform = is_uniform_range(sq.qtype) ||
+            sq.qtype == ScalarQuantizer::QT_4bit_uniform;
+
+    std::vector<float> recon_table(dim * ksub_val);
+
+    if (is_uniform) {
+        float vmin, vdiff;
+        get_uniform_range(sq, vmin, vdiff);
+        for (size_t c = 0; c < ksub_val; c++) {
+            float recon = vmin + ((c + 0.5f) / 15.0f) * vdiff;
+            for (size_t m = 0; m < dim; m++) {
+                recon_table[m * ksub_val + c] = recon;
+            }
+        }
+    } else {
+        const float* vmin = sq.trained.data();
+        const float* vdiff = sq.trained.data() + dim;
+        for (size_t m = 0; m < dim; m++) {
+            for (size_t c = 0; c < ksub_val; c++) {
+                recon_table[m * ksub_val + c] =
+                        vmin[m] + ((c + 0.5f) / 15.0f) * vdiff[m];
+            }
+        }
+    }
+
+    // by_residual is always false for the fast-scan path (non-fallback
+    // types), so codes store raw quantized values and the LUT computes
+    // distances directly from the raw query.  No biases needed.
+
+    dis_tables.resize(n * dim12);
+
+    if (metric_type == METRIC_L2) {
+        for (size_t i = 0; i < n; i++) {
+            const float* xi = x + i * dim;
+            float* lut_i = dis_tables.get() + i * dim12;
+            for (size_t m = 0; m < dim; m++) {
+                float qi = xi[m];
+                const float* recon_m = recon_table.data() + m * ksub_val;
+                float* lut_m = lut_i + m * ksub_val;
+                for (size_t c = 0; c < ksub_val; c++) {
+                    float diff = qi - recon_m[c];
+                    lut_m[c] = diff * diff;
+                }
+            }
+        }
+    } else {
+        for (size_t i = 0; i < n; i++) {
+            const float* xi = x + i * dim;
+            float* lut_i = dis_tables.get() + i * dim12;
+            for (size_t m = 0; m < dim; m++) {
+                float qi = xi[m];
+                const float* recon_m = recon_table.data() + m * ksub_val;
+                float* lut_m = lut_i + m * ksub_val;
+                for (size_t c = 0; c < ksub_val; c++) {
+                    lut_m[c] = qi * recon_m[c];
+                }
+            }
+        }
+    }
+}
+
+void IndexIVFSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
+        const {
+    size_t coarse_size = coarse_code_size();
+
+#pragma omp parallel if (n > 1)
+    {
+        std::vector<float> residual(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = bytes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+            sq.decode(code + coarse_size, xi, 1);
+            if (by_residual) {
+                quantizer->reconstruct(list_no, residual.data());
+                for (int j = 0; j < d; j++) {
+                    xi[j] += residual[j];
+                }
+            }
+        }
+    }
+}
+
+void IndexIVFSQFastScan::recombine_code(
+        int64_t list_no,
+        int64_t offset,
+        uint8_t* scratch,
+        int* values,
+        uint8_t* code_out) const {
+    // Interleaved for SIMD, so read one sub-quantizer at a time.
+    memset(scratch, 0, size_t(M2) / 2);
+    InvertedLists::ScopedCodes list_codes(invlists, list_no);
+    BitstringWriter bsw(scratch, M2 / 2);
+    for (size_t m = 0; m < M; m++) {
+        bsw.write(
+                pq4_get_packed_element(list_codes.get(), bbs, M2, offset, m),
+                4);
+    }
+    InvertedLists::ScopedCodes lo(lo_codes_invlists, list_no);
+    split_to_codes(
+            scratch, lo.get() + offset * lo_code_size(), d, sq.qtype, values);
+    // BitstringWriter matches Codec6bit/Codec8bit, so distance_to_code reads
+    // it.
+    memset(code_out, 0, sq.code_size);
+    BitstringWriter out(code_out, sq.code_size);
+    const int bits = sq_bits(sq.qtype);
+    for (int m = 0; m < d; m++) {
+        out.write(values[m], bits);
+    }
+}
+
+void IndexIVFSQFastScan::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    if (is_fallback(sq.qtype)) {
+        // Fallback: codes stored directly in invlists.
+        InvertedLists::ScopedCodes codes(invlists, list_no);
+        sq.decode(codes.get() + offset * sq.code_size, recons, 1);
+    } else if (lo_codes_invlists) {
+        // Rerank: rebuild the original code from its two nibble halves.
+        std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
+        std::vector<int> values(d);
+        recombine_code(
+                list_no, offset, scratch.data(), values.data(), code.data());
+        sq.decode(code.data(), recons, 1);
+    } else {
+        // Native 4-bit: unpack from block inverted lists.
+        std::vector<uint8_t> code(M2 / 2, 0);
+        InvertedLists::ScopedCodes list_codes(invlists, list_no);
+        BitstringWriter bsw(code.data(), M2 / 2);
+        for (size_t m = 0; m < M; m++) {
+            uint8_t c = pq4_get_packed_element(
+                    list_codes.get(), bbs, M2, offset, m);
+            bsw.write(c, 4);
+        }
+        sq.decode(code.data(), recons, 1);
+    }
+
+    if (by_residual) {
+        std::vector<float> centroid(d);
+        quantizer->reconstruct(list_no, centroid.data());
+        for (int j = 0; j < d; j++) {
+            recons[j] += centroid[j];
+        }
+    }
+}
+
+InvertedListScanner* IndexIVFSQFastScan::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
+    return sq.select_InvertedListScanner(
+            metric_type, quantizer, store_pairs, sel, by_residual);
+}
+
+// -----------------------------------------------------------------------
+// search — fallback types use IndexIVF's scanner-based path
+// -----------------------------------------------------------------------
+
+void IndexIVFSQFastScan::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    if (is_fallback(sq.qtype)) {
+        IndexIVF::search(n, x, k, distances, labels, params);
+    } else if (needs_rerank(sq.qtype)) {
+        idx_t k_coarse = idx_t(double(k) * rerank_factor);
+        FAISS_THROW_IF_NOT(k_coarse >= k);
+        FAISS_THROW_IF_NOT_MSG(
+                n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
+
+        std::vector<float> coarse_dis(n * k_coarse);
+        std::vector<idx_t> coarse_ids(n * k_coarse);
+        IndexIVFFastScan::search(
+                n, x, k_coarse, coarse_dis.data(), coarse_ids.data(), params);
+
+        // Rerank using original full-precision codes
+#pragma omp parallel
+        {
+            std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
+                    sq.get_distance_computer(metric_type));
+            std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
+            std::vector<int> values(d);
+
+#pragma omp for
+            for (idx_t i = 0; i < n; i++) {
+                const float* xi = x + i * d;
+                dc->set_query(xi);
+
+                float* heap_dis = distances + i * k;
+                idx_t* heap_ids = labels + i * k;
+
+                if (metric_type == METRIC_L2) {
+                    maxheap_heapify(k, heap_dis, heap_ids);
+                } else {
+                    minheap_heapify(k, heap_dis, heap_ids);
+                }
+
+                for (idx_t j = 0; j < k_coarse; j++) {
+                    idx_t id = coarse_ids[i * k_coarse + j];
+                    if (id < 0) {
+                        continue;
+                    }
+
+                    idx_t lo = direct_map.get(id);
+                    idx_t list_no = lo_listno(lo);
+                    idx_t offset = lo_offset(lo);
+
+                    recombine_code(
+                            list_no,
+                            offset,
+                            scratch.data(),
+                            values.data(),
+                            code.data());
+                    float dis = dc->distance_to_code(code.data());
+
+                    if (metric_type == METRIC_L2) {
+                        if (dis < heap_dis[0]) {
+                            maxheap_replace_top(k, heap_dis, heap_ids, dis, id);
+                        }
+                    } else {
+                        if (dis > heap_dis[0]) {
+                            minheap_replace_top(k, heap_dis, heap_ids, dis, id);
+                        }
+                    }
+                }
+
+                if (metric_type == METRIC_L2) {
+                    maxheap_reorder(k, heap_dis, heap_ids);
+                } else {
+                    minheap_reorder(k, heap_dis, heap_ids);
+                }
+            }
+        }
+    } else {
+        IndexIVFFastScan::search(n, x, k, distances, labels, params);
+    }
+}
+
+void IndexIVFSQFastScan::search_preassigned(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        const idx_t* assign,
+        const float* centroid_dis,
+        float* distances,
+        idx_t* labels,
+        bool store_pairs,
+        const IVFSearchParameters* params,
+        IndexIVFStats* stats) const {
+    if (is_fallback(sq.qtype)) {
+        IndexIVF::search_preassigned(
+                n,
+                x,
+                k,
+                assign,
+                centroid_dis,
+                distances,
+                labels,
+                store_pairs,
+                params,
+                stats);
+    } else if (needs_rerank(sq.qtype)) {
+        idx_t k_coarse = idx_t(double(k) * rerank_factor);
+        FAISS_THROW_IF_NOT(k_coarse >= k);
+        FAISS_THROW_IF_NOT_MSG(
+                n <= INT64_MAX / k_coarse, "n * k_coarse would overflow int64");
+
+        std::vector<float> coarse_dis(n * k_coarse);
+        std::vector<idx_t> coarse_ids(n * k_coarse);
+        IndexIVFFastScan::search_preassigned(
+                n,
+                x,
+                k_coarse,
+                assign,
+                centroid_dis,
+                coarse_dis.data(),
+                coarse_ids.data(),
+                false,
+                params,
+                stats);
+
+#pragma omp parallel
+        {
+            std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
+                    sq.get_distance_computer(metric_type));
+            std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
+            std::vector<int> values(d);
+
+#pragma omp for
+            for (idx_t i = 0; i < n; i++) {
+                const float* xi = x + i * d;
+                dc->set_query(xi);
+
+                float* heap_dis = distances + i * k;
+                idx_t* heap_ids = labels + i * k;
+
+                if (metric_type == METRIC_L2) {
+                    maxheap_heapify(k, heap_dis, heap_ids);
+                } else {
+                    minheap_heapify(k, heap_dis, heap_ids);
+                }
+
+                for (idx_t j = 0; j < k_coarse; j++) {
+                    idx_t id = coarse_ids[i * k_coarse + j];
+                    if (id < 0) {
+                        continue;
+                    }
+
+                    idx_t lo = direct_map.get(id);
+                    idx_t list_no = lo_listno(lo);
+                    idx_t offset = lo_offset(lo);
+
+                    recombine_code(
+                            list_no,
+                            offset,
+                            scratch.data(),
+                            values.data(),
+                            code.data());
+                    float dis = dc->distance_to_code(code.data());
+
+                    if (metric_type == METRIC_L2) {
+                        if (dis < heap_dis[0]) {
+                            maxheap_replace_top(k, heap_dis, heap_ids, dis, id);
+                        }
+                    } else {
+                        if (dis > heap_dis[0]) {
+                            minheap_replace_top(k, heap_dis, heap_ids, dis, id);
+                        }
+                    }
+                }
+
+                if (metric_type == METRIC_L2) {
+                    maxheap_reorder(k, heap_dis, heap_ids);
+                } else {
+                    minheap_reorder(k, heap_dis, heap_ids);
+                }
+            }
+        }
+    } else {
+        IndexIVFFastScan::search_preassigned(
+                n,
+                x,
+                k,
+                assign,
+                centroid_dis,
+                distances,
+                labels,
+                store_pairs,
+                params,
+                stats);
+    }
+}
+
+void IndexIVFSQFastScan::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    if (is_fallback(sq.qtype)) {
+        IndexIVF::range_search(n, x, radius, result, params);
+    } else if (needs_rerank(sq.qtype)) {
+        const IVFSearchParameters* ivf_params = nullptr;
+        if (params) {
+            ivf_params = dynamic_cast<const IVFSearchParameters*>(params);
+        }
+        size_t cur_nprobe = ivf_params ? ivf_params->nprobe : nprobe;
+        const IDSelector* sel = params ? params->sel : nullptr;
+
+        std::vector<idx_t> assign(n * cur_nprobe);
+        std::vector<float> centroid_dis(n * cur_nprobe);
+        quantizer->search(n, x, cur_nprobe, centroid_dis.data(), assign.data());
+
+        std::vector<RangeSearchPartialResult*> partial_results(
+                omp_get_max_threads());
+
+#pragma omp parallel
+        {
+            int rank = omp_get_thread_num();
+            RangeSearchPartialResult* pres =
+                    new RangeSearchPartialResult(result);
+            partial_results[rank] = pres;
+
+            std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
+                    sq.get_distance_computer(metric_type));
+            std::vector<uint8_t> scratch(M2 / 2), code(sq.code_size);
+            std::vector<int> values(d);
+
+#pragma omp for
+            for (idx_t i = 0; i < n; i++) {
+                dc->set_query(x + i * d);
+                RangeQueryResult& qres = pres->new_result(i);
+
+                for (size_t p = 0; p < cur_nprobe; p++) {
+                    idx_t list_no = assign[i * cur_nprobe + p];
+                    if (list_no < 0) {
+                        continue;
+                    }
+                    size_t list_size = lo_codes_invlists->list_size(list_no);
+                    if (list_size == 0) {
+                        continue;
+                    }
+                    InvertedLists::ScopedIds ids(lo_codes_invlists, list_no);
+
+                    for (size_t j = 0; j < list_size; j++) {
+                        if (sel && !sel->is_member(ids[j])) {
+                            continue;
+                        }
+                        recombine_code(
+                                list_no,
+                                j,
+                                scratch.data(),
+                                values.data(),
+                                code.data());
+                        float dis = dc->distance_to_code(code.data());
+                        if (metric_type == METRIC_L2) {
+                            if (dis < radius) {
+                                qres.add(dis, ids[j]);
+                            }
+                        } else {
+                            if (dis > radius) {
+                                qres.add(dis, ids[j]);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        RangeSearchPartialResult::merge(partial_results);
+    } else {
+        IndexIVFFastScan::range_search(n, x, radius, result, params);
+    }
+}
+
+// -----------------------------------------------------------------------
+// add_with_ids — custom implementation to store both packed and original codes
+// -----------------------------------------------------------------------
+
+void IndexIVFSQFastScan::add_with_ids(
+        idx_t n,
+        const float* x,
+        const idx_t* xids) {
+    FAISS_THROW_IF_NOT(is_trained);
+
+    constexpr idx_t bs = 65536;
+    if (n > bs) {
+        for (idx_t i0 = 0; i0 < n; i0 += bs) {
+            idx_t i1 = std::min(n, i0 + bs);
+            add_with_ids(i1 - i0, x + i0 * d, xids ? xids + i0 : nullptr);
+        }
+        return;
+    }
+
+    // Assign to coarse clusters
+    std::unique_ptr<idx_t[]> idx(new idx_t[n]);
+    quantizer->assign(n, x, idx.get());
+
+    if (is_fallback(sq.qtype)) {
+        // Fallback: store SQ codes directly in invlists (same as
+        // IndexIVFScalarQuantizer).
+        std::vector<uint8_t> flat_codes(n * sq.code_size);
+        encode_vectors(n, x, idx.get(), flat_codes.data());
+
+        for (idx_t i = 0; i < n; i++) {
+            if (idx[i] < 0) {
+                continue;
+            }
+            idx_t id = xids ? xids[i] : ntotal + i;
+            invlists->add_entry(
+                    idx[i], id, flat_codes.data() + i * sq.code_size);
+        }
+        ntotal += n;
+        return;
+    }
+
+    // Produce 4-bit packed codes for the SIMD scan,
+    // and (for rerank types) also store original codes.
+    AlignedTable<uint8_t> flat_codes(n * code_size);
+    encode_vectors(n, x, idx.get(), flat_codes.get());
+
+    // Same quantisation as encode_vectors, or the halves disagree.
+    std::vector<uint8_t> lo_flat_codes;
+    if (needs_rerank(sq.qtype)) {
+        std::vector<float> to_encode(size_t(n) * d);
+        for (idx_t i = 0; i < n; i++) {
+            if (by_residual && idx[i] >= 0) {
+                quantizer->compute_residual(
+                        x + i * d, to_encode.data() + i * d, idx[i]);
+            } else {
+                memcpy(to_encode.data() + i * d, x + i * d, sizeof(float) * d);
+            }
+        }
+        lo_flat_codes.assign(size_t(n) * lo_code_size(), 0);
+        std::vector<uint8_t> hi_discard(size_t(n) * (M2 / 2));
+        float_to_split_codes(
+                to_encode.data(),
+                hi_discard.data(),
+                lo_flat_codes.data(),
+                n,
+                d,
+                M2,
+                sq);
+    }
+
+    // Sort by list assignment
+    std::vector<idx_t> order(n);
+    for (idx_t i = 0; i < n; i++) {
+        order[i] = i;
+    }
+    std::stable_sort(order.begin(), order.end(), [&idx](idx_t a, idx_t b) {
+        return idx[a] < idx[b];
+    });
+
+    BlockInvertedLists* bil = dynamic_cast<BlockInvertedLists*>(invlists);
+    FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported");
+
+    idx_t i0 = 0;
+    while (i0 < n) {
+        idx_t list_no = idx[order[i0]];
+        idx_t i1 = i0 + 1;
+        while (i1 < n && idx[order[i1]] == list_no) {
+            i1++;
+        }
+
+        if (list_no == -1) {
+            i0 = i1;
+            continue;
+        }
+
+        size_t list_size = bil->list_size(list_no);
+        bil->resize(list_no, list_size + i1 - i0);
+
+        AlignedTable<uint8_t> list_codes((i1 - i0) * code_size);
+        for (idx_t i = i0; i < i1; i++) {
+            size_t ofs = list_size + i - i0;
+            idx_t id = xids ? xids[order[i]] : ntotal + order[i];
+            bil->ids[list_no][ofs] = id;
+            memcpy(list_codes.data() + (i - i0) * code_size,
+                   flat_codes.data() + order[i] * code_size,
+                   code_size);
+
+            if (needs_rerank(sq.qtype)) {
+                lo_codes_invlists->add_entry(
+                        list_no,
+                        id,
+                        lo_flat_codes.data() + order[i] * lo_code_size());
+                direct_map.add_single_id(id, list_no, ofs);
+            }
+        }
+
+        pq4_pack_codes_range(
+                list_codes.data(),
+                M,
+                list_size,
+                list_size + i1 - i0,
+                bbs,
+                M2,
+                bil->codes[list_no].data(),
+                0,
+                get_block_stride());
+
+        i0 = i1;
+    }
+
+    ntotal += n;
+}
+
+void IndexIVFSQFastScan::reset() {
+    IndexIVFFastScan::reset();
+    if (lo_codes_invlists) {
+        for (size_t i = 0; i < nlist; i++) {
+            lo_codes_invlists->resize(i, 0);
+        }
+        direct_map.clear();
+    }
+}
+
+} // namespace faiss