faiss 0.4.3 → 0.5.1

data/vendor/faiss/faiss/IndexIVFFastScan.h

@@ -7,9 +7,8 @@
 
 #pragma once
 
-#include <memory>
-
 #include <faiss/IndexIVF.h>
+#include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/utils/AlignedTable.h>
 
 namespace faiss {
@@ -63,6 +62,15 @@ struct IndexIVFFastScan : IndexIVF {
     // quantizer used to pack the codes
     Quantizer* fine_quantizer = nullptr;
 
+    /** Constructor for IndexIVFFastScan
+     *
+     * @param quantizer     coarse quantizer for IVF clustering
+     * @param d             dimensionality of vectors
+     * @param nlist         number of inverted lists
+     * @param code_size     size of each code in bytes
+     * @param metric        distance metric to use
+     * @param own_invlists  whether to own the inverted lists
+     */
     IndexIVFFastScan(
             Index* quantizer,
             size_t d,
@@ -73,7 +81,16 @@ struct IndexIVFFastScan : IndexIVF {
 
     IndexIVFFastScan();
 
-    /// called by implementations
+    /** Initialize the fast scan functionality (called by implementations)
+     *
+     * @param fine_quantizer  fine quantizer for encoding
+     * @param M               number of subquantizers
+     * @param nbits           number of bits per subquantizer
+     * @param nlist           number of inverted lists
+     * @param metric          distance metric to use
+     * @param bbs             block size for SIMD processing
+     * @param own_invlists    whether to own the inverted lists
+     */
     void init_fastscan(
             Quantizer* fine_quantizer,
             size_t M,
@@ -91,34 +108,72 @@ struct IndexIVFFastScan : IndexIVF {
     /// orig's inverted lists (for debugging)
     InvertedLists* orig_invlists = nullptr;
 
+    /** Add vectors with specific IDs to the index
+     *
+     * @param n     number of vectors to add
+     * @param x     vectors to add (n * d)
+     * @param xids  IDs for the vectors (n)
+     */
     void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
-
     // prepare look-up tables
 
     virtual bool lookup_table_is_3d() const = 0;
 
     // compact way of conveying coarse quantization results
     struct CoarseQuantized {
-        size_t nprobe;
+        size_t nprobe = 0;
         const float* dis = nullptr;
         const idx_t* ids = nullptr;
     };
 
+    /* Compute distance table for query set, given a list of coarse
+     * quantizers.
+     *
+     * @param n             number of queries
+     * @param x             query vectors (n, d)
+     * @param cq            coarse quantization results
+     * @param dis_tables    output distance tables
+     * @param biases        output bias values
+     * @param context       processing context containing query factors
+    processor
+     */
     virtual void compute_LUT(
             size_t n,
             const float* x,
             const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
-            AlignedTable<float>& biases) const = 0;
+            AlignedTable<float>& biases,
+            const FastScanDistancePostProcessing& context) const = 0;
 
+    /** Compute quantized lookup tables for distance computation
+     *
+     * @param n             number of query vectors
+     * @param x             query vectors (n * d)
+     * @param cq            coarse quantization results
+     * @param dis_tables    output quantized distance tables
+     * @param biases        output quantized bias values
+     * @param normalizers   output normalization factors
+     * @param context       processing context containing query factors
+     * processor
+     */
     void compute_LUT_uint8(
             size_t n,
             const float* x,
             const CoarseQuantized& cq,
             AlignedTable<uint8_t>& dis_tables,
             AlignedTable<uint16_t>& biases,
-            float* normalizers) const;
+            float* normalizers,
+            const FastScanDistancePostProcessing& context) const;
 
+    /** Search for k nearest neighbors
+     *
+     * @param n          number of query vectors
+     * @param x          query vectors (n * d)
+     * @param k          number of nearest neighbors to find
+     * @param distances  output distances (n * k)
+     * @param labels     output labels/indices (n * k)
+     * @param params     optional search parameters
+     */
     void search(
             idx_t n,
             const float* x,
@@ -127,6 +182,19 @@ struct IndexIVFFastScan : IndexIVF {
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
+    /** Search with pre-assigned coarse quantization
+     *
+     * @param n             number of query vectors
+     * @param x             query vectors (n * d)
+     * @param k             number of nearest neighbors to find
+     * @param assign        coarse cluster assignments (n * nprobe)
+     * @param centroid_dis  distances to centroids (n * nprobe)
+     * @param distances     output distances (n * k)
+     * @param labels        output labels/indices (n * k)
+     * @param store_pairs   whether to store cluster-relative pairs
+     * @param params        optional IVF search parameters
+     * @param stats         optional search statistics
+     */
     void search_preassigned(
             idx_t n,
             const float* x,
@@ -139,6 +207,14 @@ struct IndexIVFFastScan : IndexIVF {
             const IVFSearchParameters* params = nullptr,
             IndexIVFStats* stats = nullptr) const override;
 
+    /** Range search for all neighbors within radius
+     *
+     * @param n       number of query vectors
+     * @param x       query vectors (n * d)
+     * @param radius  search radius
+     * @param result  output range search results
+     * @param params  optional search parameters
+     */
     void range_search(
             idx_t n,
             const float* x,
@@ -146,7 +222,45 @@ struct IndexIVFFastScan : IndexIVF {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) const override;
 
-    // internal search funcs
+    /** Create a KNN handler for this index type
+     *
+     * This method can be overridden by derived classes to provide
+     * specialized handlers (e.g., IVFRaBitQHeapHandler for RaBitQ indexes).
+     * Base implementation creates standard handlers based on k and impl.
+     *
+     * @param is_max        true for max-heap (inner product), false for
+     *                      min-heap (L2 distance)
+     * @param impl          implementation number:
+     *                      - even (10, 12, 14): use heap for top-k
+     *                      - odd (11, 13, 15): use reservoir sampling
+     * @param n             number of queries
+     * @param k             number of neighbors to find per query
+     * @param distances     output array for distances (n * k), will be
+     *                      populated by handler
+     * @param labels        output array for result IDs (n * k), will be
+     *                      populated by handler
+     * @param sel           optional ID selector to filter results (nullptr =
+     *                      no filtering)
+     * @param context       processing context containing additional data
+     * @param normalizers   optional array of size 2*n for converting quantized
+     *                      uint16 distances to float.
+     *
+     * @return Allocated result handler (caller owns and must delete).
+     *         Handler processes SIMD batches and populates distances/labels.
+     *
+     * @note The returned handler must be deleted by caller after use.
+     *       Typical usage: handler->begin() → process batches → handler->end()
+     */
+    virtual SIMDResultHandlerToFloat* make_knn_handler(
+            bool is_max,
+            int impl,
+            idx_t n,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const IDSelector* sel,
+            const FastScanDistancePostProcessing& context,
+            const float* normalizers = nullptr) const;
 
     // dispatch to implementations and parallelize
     void search_dispatch_implem(
@@ -156,7 +270,7 @@ struct IndexIVFFastScan : IndexIVF {
             float* distances,
             idx_t* labels,
             const CoarseQuantized& cq,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     void range_search_dispatch_implem(
@@ -165,7 +279,7 @@ struct IndexIVFFastScan : IndexIVF {
             float radius,
             RangeSearchResult& rres,
             const CoarseQuantized& cq_in,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     // impl 1 and 2 are just for verification
@@ -177,7 +291,7 @@ struct IndexIVFFastScan : IndexIVF {
             float* distances,
             idx_t* labels,
             const CoarseQuantized& cq,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     template <class C>
@@ -188,7 +302,7 @@ struct IndexIVFFastScan : IndexIVF {
             float* distances,
             idx_t* labels,
             const CoarseQuantized& cq,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     // implem 10 and 12 are not multithreaded internally, so
@@ -200,7 +314,7 @@ struct IndexIVFFastScan : IndexIVF {
             const CoarseQuantized& cq,
             size_t* ndis_out,
             size_t* nlist_out,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     void search_implem_12(
@@ -210,7 +324,7 @@ struct IndexIVFFastScan : IndexIVF {
             const CoarseQuantized& cq,
             size_t* ndis_out,
             size_t* nlist_out,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     // implem 14 is multithreaded internally across nprobes and queries
@@ -222,7 +336,7 @@ struct IndexIVFFastScan : IndexIVF {
             idx_t* labels,
             const CoarseQuantized& cq,
             int impl,
-            const NormTableScaler* scaler,
+            const FastScanDistancePostProcessing& context,
             const IVFSearchParameters* params = nullptr) const;
 
     // reconstruct vectors from packed invlists
@@ -234,16 +348,57 @@ struct IndexIVFFastScan : IndexIVF {
     // reconstruct orig invlists (for debugging)
     void reconstruct_orig_invlists();
 
-    /** Decode a set of vectors.
+    /** Decode a set of vectors
      *
-     *  NOTE: The codes in the IndexFastScan object are non-contiguous.
-     *        But this method requires a contiguous representation.
+     * NOTE: The codes in the IndexFastScan object are non-contiguous.
+     *       But this method requires a contiguous representation.
      *
      * @param n       number of vectors
      * @param bytes   input encoded vectors, size n * code_size
      * @param x       output vectors, size n * d
      */
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+   protected:
+    /** Preprocess metadata from encoded vectors before packing.
+     *
+     * Called during add_with_ids after encode_vectors but before codes
+     * are packed into SIMD-friendly blocks. Subclasses can override to
+     * extract and store metadata embedded in codes or perform other
+     * pre-packing operations.
+     *
+     * Default implementation: no-op
+     *
+     * Example use case:
+     * - IndexIVFRaBitQFastScan extracts factor data from codes for use
+     *   during search-time distance corrections
+     *
+     * @param n                  number of vectors encoded
+     * @param flat_codes         encoded vectors (n * code_size bytes)
+     * @param start_global_idx   starting global index (ntotal before add)
+     */
+    virtual void preprocess_code_metadata(
+            idx_t n,
+            const uint8_t* flat_codes,
+            idx_t start_global_idx);
+
+    /** Get stride for interpreting codes during SIMD packing.
+     *
+     * The stride determines how to read codes when packing them into
+     * SIMD-friendly block format. This is needed when codes contain
+     * embedded metadata that should be skipped during packing.
+     *
+     * Default implementation: returns 0 (use standard M-byte stride)
+     *
+     * Example use case:
+     * - IndexIVFRaBitQFastScan returns code_size because codes contain
+     *   embedded factor data after the quantized bits
+     *
+     * @return stride in bytes:
+     *         - 0: use default stride (M bytes, standard PQ/AQ codes)
+     *         - >0: use custom stride (e.g., code_size for embedded metadata)
+     */
+    virtual size_t code_packing_stride() const;
 };
 
 struct IVFFastScanStats {

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -13,6 +13,7 @@
 
 #include <cinttypes>
 #include <cstdio>
+#include <numeric>
 
 #include <faiss/IndexFlat.h>
 

data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp

@@ -0,0 +1,251 @@
+/*
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IndexIVFFlatPanorama.h>
+
+#include <cstdio>
+
+#include <faiss/IndexFlat.h>
+
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/IDSelector.h>
+#include <faiss/impl/PanoramaStats.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+IndexIVFFlatPanorama::IndexIVFFlatPanorama(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        int n_levels,
+        MetricType metric,
+        bool own_invlists)
+        : IndexIVFFlat(quantizer, d, nlist, metric, false), n_levels(n_levels) {
+    // For now, we only support L2 distance.
+    // Supporting dot product and cosine distance is a trivial addition
+    // left for future work.
+    FAISS_THROW_IF_NOT(metric == METRIC_L2);
+
+    // We construct the inverted lists here so that we can use the
+    // level-oriented storage. This does not cause a leak as we constructed
+    // IndexIVF first, with own_invlists set to false.
+    this->invlists = new ArrayInvertedListsPanorama(nlist, code_size, n_levels);
+    this->own_invlists = own_invlists;
+}
+
+IndexIVFFlatPanorama::IndexIVFFlatPanorama() : n_levels(0) {}
+
+namespace {
+
+template <typename VectorDistance, bool use_sel>
+struct IVFFlatScannerPanorama : InvertedListScanner {
+    VectorDistance vd;
+    const ArrayInvertedListsPanorama* storage;
+    using C = typename VectorDistance::C;
+
+    IVFFlatScannerPanorama(
+            const VectorDistance& vd,
+            const ArrayInvertedListsPanorama* storage,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), vd(vd), storage(storage) {
+        keep_max = vd.is_similarity;
+        code_size = vd.d * sizeof(float);
+        cum_sums.resize(storage->n_levels + 1);
+    }
+
+    const float* xi = nullptr;
+    std::vector<float> cum_sums;
+    float q_norm = 0.0f;
+    void set_query(const float* query) override {
+        this->xi = query;
+        this->storage->pano.compute_query_cum_sums(query, cum_sums.data());
+        q_norm = cum_sums[0] * cum_sums[0];
+    }
+
+    void set_list(idx_t list_no, float /* coarse_dis */) override {
+        this->list_no = list_no;
+    }
+
+    /// This function is unreachable as `IndexIVF` only calls this within
+    /// iterators, which are not supported by `IndexIVFFlatPanorama`.
+    /// To avoid undefined behavior, we throw an error here.
+    float distance_to_code(const uint8_t* /* code */) const override {
+        FAISS_THROW_MSG(
+                "IndexIVFFlatPanorama does not support distance_to_code");
+    }
+
+    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 {
+        size_t nup = 0;
+
+        const size_t n_batches =
+                (list_size + storage->kBatchSize - 1) / storage->kBatchSize;
+
+        const float* cum_sums_data = storage->get_cum_sums(list_no);
+
+        std::vector<float> exact_distances(storage->kBatchSize);
+        std::vector<uint32_t> active_indices(storage->kBatchSize);
+
+        PanoramaStats local_stats;
+        local_stats.reset();
+
+        for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
+            size_t batch_start = batch_no * storage->kBatchSize;
+
+            size_t num_active =
+                    storage->pano
+                            .progressive_filter_batch<CMax<float, int64_t>>(
+                                    codes,
+                                    cum_sums_data,
+                                    xi,
+                                    cum_sums.data(),
+                                    batch_no,
+                                    list_size,
+                                    sel,
+                                    ids,
+                                    use_sel,
+                                    active_indices,
+                                    exact_distances,
+                                    simi[0],
+                                    local_stats);
+
+            // Add batch survivors to heap.
+            for (size_t i = 0; i < num_active; i++) {
+                uint32_t idx = active_indices[i];
+                size_t global_idx = batch_start + idx;
+                float dis = exact_distances[idx];
+
+                if (C::cmp(simi[0], dis)) {
+                    int64_t id = store_pairs ? lo_build(list_no, global_idx)
+                                             : ids[global_idx];
+                    heap_replace_top<C>(k, simi, idxi, dis, id);
+                    nup++;
+                }
+            }
+        }
+
+        indexPanorama_stats.add(local_stats);
+        return nup;
+    }
+
+    void scan_codes_range(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float radius,
+            RangeQueryResult& res) const override {
+        const size_t n_batches =
+                (list_size + storage->kBatchSize - 1) / storage->kBatchSize;
+
+        const float* cum_sums_data = storage->get_cum_sums(list_no);
+
+        std::vector<float> exact_distances(storage->kBatchSize);
+        std::vector<uint32_t> active_indices(storage->kBatchSize);
+
+        PanoramaStats local_stats;
+        local_stats.reset();
+
+        // Same progressive filtering as scan_codes, but with fixed radius
+        // threshold instead of dynamic heap threshold.
+        for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
+            size_t batch_start = batch_no * storage->kBatchSize;
+
+            size_t num_active =
+                    storage->pano
+                            .progressive_filter_batch<CMax<float, int64_t>>(
+                                    codes,
+                                    cum_sums_data,
+                                    xi,
+                                    cum_sums.data(),
+                                    batch_no,
+                                    list_size,
+                                    sel,
+                                    ids,
+                                    use_sel,
+                                    active_indices,
+                                    exact_distances,
+                                    radius,
+                                    local_stats);
+
+            // Add batch survivors to range result.
+            for (size_t i = 0; i < num_active; i++) {
+                uint32_t idx = active_indices[i];
+                size_t global_idx = batch_start + idx;
+                float dis = exact_distances[idx];
+
+                if (C::cmp(radius, dis)) {
+                    int64_t id = store_pairs ? lo_build(list_no, global_idx)
+                                             : ids[global_idx];
+                    res.add(dis, id);
+                }
+            }
+        }
+
+        indexPanorama_stats.add(local_stats);
+    }
+};
+
+struct Run_get_InvertedListScanner {
+    using T = InvertedListScanner*;
+
+    template <class VD>
+    InvertedListScanner* f(
+            VD& vd,
+            const IndexIVFFlatPanorama* ivf,
+            bool store_pairs,
+            const IDSelector* sel) {
+        // Safely cast to ArrayInvertedListsPanorama to access cumulative sums.
+        const ArrayInvertedListsPanorama* storage =
+                dynamic_cast<const ArrayInvertedListsPanorama*>(ivf->invlists);
+        FAISS_THROW_IF_NOT_MSG(
+                storage,
+                "IndexIVFFlatPanorama requires ArrayInvertedListsPanorama");
+
+        if (sel) {
+            return new IVFFlatScannerPanorama<VD, true>(
+                    vd, storage, store_pairs, sel);
+        } else {
+            return new IVFFlatScannerPanorama<VD, false>(
+                    vd, storage, store_pairs, sel);
+        }
+    }
+};
+
+} // anonymous namespace
+
+InvertedListScanner* IndexIVFFlatPanorama::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
+    Run_get_InvertedListScanner run;
+    return dispatch_VectorDistance(
+            d, metric_type, metric_arg, run, this, store_pairs, sel);
+}
+
+void IndexIVFFlatPanorama::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    const uint8_t* code = invlists->get_single_code(list_no, offset);
+    memcpy(recons, code, code_size);
+    invlists->release_codes(list_no, code);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFFlatPanorama.h

@@ -0,0 +1,64 @@
+/*
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#ifndef FAISS_INDEX_IVF_FLAT_PANORAMA_H
+#define FAISS_INDEX_IVF_FLAT_PANORAMA_H
+
+#include <stdint.h>
+#include "faiss/IndexIVFFlat.h"
+
+namespace faiss {
+
+/// Panorama adaptation of IndexIVFFlat following
+/// https://www.arxiv.org/pdf/2510.00566.
+///
+/// IDEA:
+/// Panorama adapts the storage layout within each cluster and uses
+/// pruning with bounds to improve the search performance.
+/// Combined with orthogonal transforms upstream that concentrate the energy
+/// in the early dimensions (like PCA, Cayley, etc.), Panorama can prune up
+/// to 95% of the vectors in the cluster.
+///
+/// OVERHEAD:
+/// To be more efficient, we compute the residual energies at insertion time
+/// and store them along the vectors, which comes with an additional storage
+/// overhead of exactly (nlevels + 1) floats per vector. Add time is also
+/// slightly higher due to the overhead of transposing the vectors.
+///
+/// NOTE:
+/// We inherit from IndexIVFFlat instead of IndexIVF so we can keep the same
+/// insertion logic. The code responsible for level-oriented storage is in
+/// `ArrayInvertedListsPanorama`, which is a struct member of `IndexIVF`.
+struct IndexIVFFlatPanorama : IndexIVFFlat {
+    size_t n_levels;
+
+    std::vector<MaybeOwnedVector<float>> cum_sums;
+
+    explicit IndexIVFFlatPanorama(
+            Index* quantizer,
+            size_t d,
+            size_t nlist_,
+            int n_levels,
+            MetricType = METRIC_L2,
+            bool own_invlists = true);
+
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
+
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
+    IndexIVFFlatPanorama();
+};
+
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -1163,7 +1163,9 @@ struct IVFPQScannerT : QueryTables {
         }
 
 #pragma omp critical
-        { indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
+        {
+            indexIVFPQ_stats.n_hamming_pass += n_hamming_pass;
+        }
     }
 
     template <class SearchResultType>

data/vendor/faiss/faiss/IndexIVFPQ.h

@@ -113,7 +113,7 @@ struct IndexIVFPQ : IndexIVF {
      */
     size_t find_duplicates(idx_t* ids, size_t* lims) const;
 
-    // map a vector to a binary code knowning the index
+    // map a vector to a binary code knowing the index
     void encode(idx_t key, const float* x, uint8_t* code) const;
 
     /** Encode multiple vectors