faiss 0.5.3 → 0.6.0

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -21,7 +21,7 @@
 #include <faiss/impl/IDSelector.h>
 
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/distances.h>
+#include <faiss/impl/expanded_scanners.h>
 #include <faiss/utils/extra_distances.h>
 #include <faiss/utils/utils.h>
 
@@ -147,7 +147,7 @@ void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
 
 namespace {
 
-template <typename VectorDistance, bool use_sel>
+template <typename VectorDistance>
 struct IVFFlatScanner : InvertedListScanner {
     VectorDistance vd;
     using C = typename VectorDistance::C;
@@ -170,70 +170,18 @@ struct IVFFlatScanner : InvertedListScanner {
         this->list_no = list_no;
     }
 
-    float distance_to_code(const uint8_t* code) const override {
+    float distance_to_code(const uint8_t* code) const final {
         const float* yj = (float*)code;
         return vd(xi, yj);
     }
 
+    // redefining the scan_codes allows to inline the 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 {
-        const float* list_vecs = (const float*)codes;
-        size_t nup = 0;
-        for (size_t j = 0; j < list_size; j++) {
-            const float* yj = list_vecs + vd.d * j;
-            if (use_sel && !sel->is_member(ids[j])) {
-                continue;
-            }
-            float dis = vd(xi, yj);
-            if (C::cmp(simi[0], dis)) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                heap_replace_top<C>(k, simi, idxi, dis, id);
-                nup++;
-            }
-        }
-        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 float* list_vecs = (const float*)codes;
-        for (size_t j = 0; j < list_size; j++) {
-            const float* yj = list_vecs + vd.d * j;
-            if (use_sel && !sel->is_member(ids[j])) {
-                continue;
-            }
-            float dis = vd(xi, yj);
-            if (C::cmp(radius, dis)) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                res.add(dis, id);
-            }
-        }
-    }
-};
-
-struct Run_get_InvertedListScanner {
-    using T = InvertedListScanner*;
-
-    template <class VD>
-    InvertedListScanner* f(
-            VD& vd,
-            const IndexIVFFlat* ivf,
-            bool store_pairs,
-            const IDSelector* sel) {
-        if (sel) {
-            return new IVFFlatScanner<VD, true>(vd, store_pairs, sel);
-        } else {
-            return new IVFFlatScanner<VD, false>(vd, store_pairs, sel);
-        }
+            ResultHandler& handler) const {
+        return run_scan_codes_fix_C<C>(*this, list_size, codes, ids, handler);
     }
 };
 
@@ -243,9 +191,10 @@ InvertedListScanner* IndexIVFFlat::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);
+    return with_VectorDistance(
+            d, metric_type, metric_arg, [&](auto vd) -> InvertedListScanner* {
+                return new IVFFlatScanner<decltype(vd)>(vd, store_pairs, sel);
+            });
 }
 
 void IndexIVFFlat::reconstruct_from_offset(

data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp

@@ -12,13 +12,14 @@
 #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/extra_distances.h>
 #include <faiss/utils/utils.h>
 
@@ -32,10 +33,7 @@ IndexIVFFlatPanorama::IndexIVFFlatPanorama(
         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);
+    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
@@ -53,6 +51,7 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
     VectorDistance vd;
     const ArrayInvertedListsPanorama* storage;
     using C = typename VectorDistance::C;
+    static constexpr MetricType metric = VectorDistance::metric;
 
     IVFFlatScannerPanorama(
             const VectorDistance& vd,
@@ -90,9 +89,7 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
             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 =
@@ -109,22 +106,22 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         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 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,
+                        exact_distances,
+                        handler.threshold,
+                        local_stats);
+            });
 
             // Add batch survivors to heap.
             for (size_t i = 0; i < num_active; i++) {
@@ -132,10 +129,10 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
                 size_t global_idx = batch_start + 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);
+                    handler.add_result(dis, id);
                     nup++;
                 }
             }
@@ -144,88 +141,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 +149,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(