faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp

@@ -12,39 +12,43 @@
 #include <cstdio>
 
 #include <faiss/IndexFlat.h>
+#include <faiss/MetricType.h>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/PanoramaStats.h>
+#include <faiss/impl/ResultHandler.h>
 
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/distances.h>
+#include <faiss/utils/distances_dispatch.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,
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
+        int n_levels_in,
         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);
+        bool own_invlists_in,
+        size_t batch_size_in)
+        : IndexIVFFlat(quantizer_in, d_in, nlist_in, metric, false),
+          n_levels(n_levels_in),
+          batch_size(batch_size_in) {
+    FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT);
 
     // 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;
+    this->invlists = new ArrayInvertedListsPanorama(
+            nlist, code_size, n_levels, batch_size);
+    this->own_invlists = own_invlists_in;
 }
 
-IndexIVFFlatPanorama::IndexIVFFlatPanorama() : n_levels(0) {}
+IndexIVFFlatPanorama::IndexIVFFlatPanorama()
+        : n_levels(0), batch_size(Panorama::kDefaultBatchSize) {}
 
 namespace {
 
@@ -53,16 +57,28 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
     VectorDistance vd;
     const ArrayInvertedListsPanorama* storage;
     using C = typename VectorDistance::C;
+    static constexpr MetricType metric = VectorDistance::metric;
+
+    mutable std::vector<uint32_t> active_indices_;
+    mutable std::vector<uint8_t> active_byteset_;
+    mutable std::vector<float> exact_distances_;
+    mutable std::vector<float> dot_buffer_;
 
     IVFFlatScannerPanorama(
-            const VectorDistance& vd,
-            const ArrayInvertedListsPanorama* storage,
-            bool store_pairs,
-            const IDSelector* sel)
-            : InvertedListScanner(store_pairs, sel), vd(vd), storage(storage) {
+            const VectorDistance& vd_in,
+            const ArrayInvertedListsPanorama* storage_in,
+            bool store_pairs_in,
+            const IDSelector* sel_in)
+            : InvertedListScanner(store_pairs_in, sel_in),
+              vd(vd_in),
+              storage(storage_in) {
         keep_max = vd.is_similarity;
         code_size = vd.d * sizeof(float);
-        cum_sums.resize(storage->n_levels + 1);
+        cum_sums.resize(storage->pano.n_levels + 1);
+        active_indices_.resize(storage->pano.batch_size);
+        active_byteset_.resize(storage->pano.batch_size);
+        exact_distances_.resize(storage->pano.batch_size);
+        dot_buffer_.resize(storage->pano.batch_size);
     }
 
     const float* xi = nullptr;
@@ -74,8 +90,8 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         q_norm = cum_sums[0] * cum_sums[0];
     }
 
-    void set_list(idx_t list_no, float /* coarse_dis */) override {
-        this->list_no = list_no;
+    void set_list(idx_t list_no_in, float /* coarse_dis */) override {
+        this->list_no = list_no_in;
     }
 
     /// This function is unreachable as `IndexIVF` only calls this within
@@ -86,57 +102,61 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
                 "IndexIVFFlatPanorama does not support distance_to_code");
     }
 
+    using InvertedListScanner::scan_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 nup = 0;
 
-        const size_t n_batches =
-                (list_size + storage->kBatchSize - 1) / storage->kBatchSize;
+        const size_t bs = storage->pano.batch_size;
+        const size_t n_batches = (list_size + bs - 1) / bs;
 
         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);
+            size_t batch_start = batch_no * bs;
+            size_t num_active = with_metric_type(metric, [&]<MetricType M>() {
+                return storage->pano.progressive_filter_batch<C, M>(
+                        codes,
+                        cum_sums_data,
+                        xi,
+                        cum_sums.data(),
+                        batch_no,
+                        list_size,
+                        sel,
+                        ids,
+                        use_sel,
+                        active_indices_,
+                        active_byteset_,
+                        exact_distances_,
+                        dot_buffer_,
+                        handler.threshold,
+                        local_stats);
+            });
+
+            // num_active is the count of codes for which exact distance
+            // was computed in this batch (post-filter, post-pruning).
+            handler.stats.scan_cnt += num_active;
 
             // Add batch survivors to heap.
             for (size_t i = 0; i < num_active; i++) {
-                uint32_t idx = active_indices[i];
+                uint32_t idx = active_indices_[i];
                 size_t global_idx = batch_start + idx;
-                float dis = exact_distances[idx];
+                float dis = exact_distances_[idx];
 
-                if (C::cmp(simi[0], dis)) {
+                if (C::cmp(handler.threshold, 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++;
+                    if (handler.add_result(dis, id)) {
+                        handler.stats.nheap_updates++;
+                        nup++;
+                    }
                 }
             }
         }
@@ -144,88 +164,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         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
@@ -234,9 +172,22 @@ 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);
+    const ArrayInvertedListsPanorama* storage =
+            dynamic_cast<const ArrayInvertedListsPanorama*>(invlists);
+    FAISS_THROW_IF_NOT_MSG(
+            storage,
+            "IndexIVFFlatPanorama requires ArrayInvertedListsPanorama");
+
+    return with_VectorDistance(
+            d, metric_type, metric_arg, [&](auto vd) -> InvertedListScanner* {
+                if (sel) {
+                    return new IVFFlatScannerPanorama<decltype(vd), true>(
+                            vd, storage, store_pairs, sel);
+                } else {
+                    return new IVFFlatScannerPanorama<decltype(vd), false>(
+                            vd, storage, store_pairs, sel);
+                }
+            });
 }
 
 void IndexIVFFlatPanorama::reconstruct_from_offset(

data/vendor/faiss/faiss/IndexIVFFlatPanorama.h

@@ -37,6 +37,7 @@ namespace faiss {
 /// `ArrayInvertedListsPanorama`, which is a struct member of `IndexIVF`.
 struct IndexIVFFlatPanorama : IndexIVFFlat {
     size_t n_levels;
+    size_t batch_size;
 
     std::vector<MaybeOwnedVector<float>> cum_sums;
 
@@ -46,7 +47,8 @@ struct IndexIVFFlatPanorama : IndexIVFFlat {
             size_t nlist_,
             int n_levels,
             MetricType = METRIC_L2,
-            bool own_invlists = true);
+            bool own_invlists = true,
+            size_t batch_size = Panorama::kDefaultBatchSize);
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,

data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp

@@ -13,13 +13,13 @@
 namespace faiss {
 
 IndexIVFIndependentQuantizer::IndexIVFIndependentQuantizer(
-        Index* quantizer,
-        IndexIVF* index_ivf,
-        VectorTransform* vt)
-        : Index(quantizer->d, index_ivf->metric_type),
-          quantizer(quantizer),
-          vt(vt),
-          index_ivf(index_ivf) {
+        Index* quantizer_in,
+        IndexIVF* index_ivf_in,
+        VectorTransform* vt_in)
+        : Index(quantizer_in->d, index_ivf_in->metric_type),
+          quantizer(quantizer_in),
+          vt(vt_in),
+          index_ivf(index_ivf_in) {
     if (vt) {
         FAISS_THROW_IF_NOT_MSG(
                 vt->d_in == d && vt->d_out == index_ivf->d,
@@ -29,14 +29,16 @@ IndexIVFIndependentQuantizer::IndexIVFIndependentQuantizer(
     }
 
     if (quantizer->is_trained && quantizer->ntotal != 0) {
-        FAISS_THROW_IF_NOT(quantizer->ntotal == index_ivf->nlist);
+        FAISS_THROW_IF_NOT(
+                quantizer->ntotal == static_cast<idx_t>(index_ivf->nlist));
     }
     if (index_ivf->is_trained && vt) {
         FAISS_THROW_IF_NOT(vt->is_trained);
     }
     ntotal = index_ivf->ntotal;
     is_trained =
-            (quantizer->is_trained && quantizer->ntotal == index_ivf->nlist &&
+            (quantizer->is_trained &&
+             quantizer->ntotal == static_cast<idx_t>(index_ivf->nlist) &&
              (!vt || vt->is_trained) && index_ivf->is_trained);
 
     // disable precomputed tables because they use the distances that are
@@ -57,15 +59,16 @@ IndexIVFIndependentQuantizer::~IndexIVFIndependentQuantizer() {
 namespace {
 
 struct VTransformedVectors : TransformedVectors {
-    VTransformedVectors(const VectorTransform* vt, idx_t n, const float* x)
-            : TransformedVectors(x, vt ? vt->apply(n, x) : x) {}
+    VTransformedVectors(const VectorTransform* vt, idx_t n, const float* x_in)
+            : TransformedVectors(x_in, vt ? vt->apply(n, x_in) : x_in) {}
 };
 
 struct SubsampledVectors : TransformedVectors {
-    SubsampledVectors(int d, idx_t* n, idx_t max_n, const float* x)
+    SubsampledVectors(int d, idx_t* n, idx_t max_n, const float* x_in)
             : TransformedVectors(
-                      x,
-                      fvecs_maybe_subsample(d, (size_t*)n, max_n, x, true)) {}
+                      x_in,
+                      fvecs_maybe_subsample(d, (size_t*)n, max_n, x_in, true)) {
+    }
 };
 
 } // anonymous namespace
@@ -88,7 +91,7 @@ void IndexIVFIndependentQuantizer::search(
         idx_t* labels,
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT_MSG(!params, "search parameters not supported");
-    int nprobe = index_ivf->nprobe;
+    size_t nprobe = index_ivf->nprobe;
     std::vector<float> D(n * nprobe);
     std::vector<idx_t> I(n * nprobe);
     quantizer->search(n, x, nprobe, D.data(), I.data());