faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexFlat.h

@@ -74,7 +74,7 @@ struct IndexFlat : IndexFlatCodes {
 };
 
 struct IndexFlatIP : IndexFlat {
-    explicit IndexFlatIP(idx_t d) : IndexFlat(d, METRIC_INNER_PRODUCT) {}
+    explicit IndexFlatIP(idx_t d_in) : IndexFlat(d_in, METRIC_INNER_PRODUCT) {}
     IndexFlatIP() {}
 };
 
@@ -88,7 +88,7 @@ struct IndexFlatL2 : IndexFlat {
     /**
      * @param d dimensionality of the input vectors
      */
-    explicit IndexFlatL2(idx_t d) : IndexFlat(d, METRIC_L2) {}
+    explicit IndexFlatL2(idx_t d_in) : IndexFlat(d_in, METRIC_L2) {}
     IndexFlatL2() {}
 
     // override for l2 norms cache.
@@ -113,15 +113,16 @@ struct IndexFlatPanorama : IndexFlat {
      * @param batch_size batch size for Panorama storage
      */
     explicit IndexFlatPanorama(
-            idx_t d,
+            idx_t d_in,
             MetricType metric,
-            size_t n_levels,
-            size_t batch_size)
-            : IndexFlat(d, metric),
-              batch_size(batch_size),
-              n_levels(n_levels),
-              pano(code_size, n_levels, batch_size) {
-        FAISS_THROW_IF_NOT(metric == METRIC_L2);
+            size_t n_levels_in,
+            size_t batch_size_in)
+            : IndexFlat(d_in, metric),
+              batch_size(batch_size_in),
+              n_levels(n_levels_in),
+              pano(code_size, n_levels_in, batch_size_in) {
+        FAISS_THROW_IF_NOT(
+                metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT);
     }
 
     void add(idx_t n, const float* x) override;
@@ -173,10 +174,27 @@ struct IndexFlatL2Panorama : IndexFlatPanorama {
      * @param batch_size batch size for Panorama storage
      */
     explicit IndexFlatL2Panorama(
-            idx_t d,
-            size_t n_levels,
-            size_t batch_size = 512)
-            : IndexFlatPanorama(d, METRIC_L2, n_levels, batch_size) {}
+            idx_t d_in,
+            size_t n_levels_in,
+            size_t batch_size_in = 512)
+            : IndexFlatPanorama(d_in, METRIC_L2, n_levels_in, batch_size_in) {}
+};
+
+struct IndexFlatIPPanorama : IndexFlatPanorama {
+    /**
+     * @param d dimensionality of the input vectors
+     * @param n_levels number of Panorama levels
+     * @param batch_size batch size for Panorama storage
+     */
+    explicit IndexFlatIPPanorama(
+            idx_t d_in,
+            size_t n_levels_in,
+            size_t batch_size_in = 512)
+            : IndexFlatPanorama(
+                      d_in,
+                      METRIC_INNER_PRODUCT,
+                      n_levels_in,
+                      batch_size_in) {}
 };
 
 /// optimized version for 1D "vectors".

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -7,18 +7,21 @@
 
 #include <faiss/IndexFlatCodes.h>
 
+#include <atomic>
+
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/CodePacker.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 namespace faiss {
 
-IndexFlatCodes::IndexFlatCodes(size_t code_size, idx_t d, MetricType metric)
-        : Index(d, metric), code_size(code_size) {}
+IndexFlatCodes::IndexFlatCodes(size_t code_size_, idx_t d_, MetricType metric)
+        : Index(d_, metric), code_size(code_size_) {}
 
 IndexFlatCodes::IndexFlatCodes() : code_size(0) {}
 
@@ -131,12 +134,16 @@ struct GenericFlatCodesDistanceComputer : FlatCodesDistanceComputer {
     std::vector<float> vec_buffer;
     const float* query = nullptr;
 
-    GenericFlatCodesDistanceComputer(const IndexFlatCodes* codec, const VD& vd)
-            : FlatCodesDistanceComputer(codec->codes.data(), codec->code_size),
-              codec(*codec),
-              vd(vd),
-              code_buffer(codec->code_size * 4),
-              vec_buffer(codec->d * 4) {}
+    GenericFlatCodesDistanceComputer(
+            const IndexFlatCodes* codec_,
+            const VD& vd_)
+            : FlatCodesDistanceComputer(
+                      codec_->codes.data(),
+                      codec_->code_size),
+              codec(*codec_),
+              vd(vd_),
+              code_buffer(codec_->code_size * 4),
+              vec_buffer(codec_->d * 4) {}
 
     void set_query(const float* x) override {
         query = x;
@@ -182,15 +189,6 @@ struct GenericFlatCodesDistanceComputer : FlatCodesDistanceComputer {
     }
 };
 
-struct Run_get_distance_computer {
-    using T = FlatCodesDistanceComputer*;
-
-    template <class VD>
-    FlatCodesDistanceComputer* f(const VD& vd, const IndexFlatCodes* codec) {
-        return new GenericFlatCodesDistanceComputer<VD>(codec, vd);
-    }
-};
-
 template <class BlockResultHandler>
 struct Run_search_with_decompress {
     using T = void;
@@ -208,40 +206,54 @@ struct Run_search_with_decompress {
         using SingleResultHandler =
                 typename BlockResultHandler::SingleResultHandler;
         using DC = GenericFlatCodesDistanceComputer<VectorDistance>;
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel // if (res.nq > 100)
         {
-            std::unique_ptr<DC> dc(new DC(&index, vd));
-            SingleResultHandler resi(res);
+            std::unique_ptr<DC> dc;
+            std::unique_ptr<SingleResultHandler> resi;
+            try {
+                dc = std::make_unique<DC>(&index, vd);
+                resi = std::make_unique<SingleResultHandler>(res);
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 #pragma omp for
-            for (int64_t q = 0; q < res.nq; q++) {
-                resi.begin(q);
-                dc->set_query(xq + vd.d * q);
-                for (size_t i = 0; i < ntotal; i++) {
-                    if (res.is_in_selection(i)) {
-                        float dis = (*dc)(i);
-                        resi.add_result(dis, i);
+            for (int64_t q = 0; q < static_cast<int64_t>(res.nq); q++) {
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    resi->begin(q);
+                    dc->set_query(xq + vd.d * q);
+                    for (size_t i = 0; i < ntotal; i++) {
+                        if (res.is_in_selection(i)) {
+                            float dis = (*dc)(i);
+                            resi->add_result(dis, i);
+                        }
                     }
+                    resi->end();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
                 }
-                resi.end();
             }
         }
+        omp_rethrow_if_exception(ex);
     }
 };
 
 struct Run_search_with_decompress_res {
     using T = void;
 
-    template <class ResultHandler>
-    void f(ResultHandler& res, const IndexFlatCodes* index, const float* xq) {
-        Run_search_with_decompress<ResultHandler> r;
-        dispatch_VectorDistance(
-                index->d,
-                index->metric_type,
-                index->metric_arg,
-                r,
-                index,
-                xq,
-                res);
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res,
+           const IndexFlatCodes* index,
+           const float* xq) {
+        with_VectorDistance(
+                index->d, index->metric_type, index->metric_arg, [&](auto vd) {
+                    Run_search_with_decompress<BlockResultHandler> r;
+                    r.template f<decltype(vd)>(vd, index, xq, res);
+                });
     }
 };
 
@@ -249,8 +261,14 @@ struct Run_search_with_decompress_res {
 
 FlatCodesDistanceComputer* IndexFlatCodes::get_FlatCodesDistanceComputer()
         const {
-    Run_get_distance_computer r;
-    return dispatch_VectorDistance(d, metric_type, metric_arg, r, this);
+    return with_VectorDistance(
+            d,
+            metric_type,
+            metric_arg,
+            [&](auto vd) -> FlatCodesDistanceComputer* {
+                return new GenericFlatCodesDistanceComputer<decltype(vd)>(
+                        this, vd);
+            });
 }
 
 void IndexFlatCodes::search(
@@ -267,7 +285,7 @@ void IndexFlatCodes::search(
 }
 
 void IndexFlatCodes::range_search(
-        idx_t n,
+        idx_t /* n */,
         const float* x,
         float radius,
         RangeSearchResult* result,
@@ -277,4 +295,33 @@ void IndexFlatCodes::range_search(
     dispatch_range_ResultHandler(result, radius, metric_type, sel, r, this, x);
 }
 
+void IndexFlatCodes::search1(
+        const float* x,
+        ResultHandler& handler,
+        SearchParameters* params) const {
+    const IDSelector* sel = params ? params->sel : nullptr;
+    Run_search_with_decompress_res r;
+    if (sel) {
+        if (is_similarity_metric(metric_type)) {
+            SingleQueryBlockResultHandler<CMin<float, idx_t>, true> res(
+                    handler, sel);
+            r.f(res, this, x);
+        } else {
+            SingleQueryBlockResultHandler<CMax<float, idx_t>, true> res(
+                    handler, sel);
+            r.f(res, this, x);
+        }
+    } else {
+        if (is_similarity_metric(metric_type)) {
+            SingleQueryBlockResultHandler<CMin<float, idx_t>, false> res(
+                    handler);
+            r.f(res, this, x);
+        } else {
+            SingleQueryBlockResultHandler<CMax<float, idx_t>, false> res(
+                    handler);
+            r.f(res, this, x);
+        }
+    }
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.h

@@ -55,7 +55,8 @@ struct IndexFlatCodes : Index {
         return get_FlatCodesDistanceComputer();
     }
 
-    /** Search implemented by decoding */
+    /** Search implemented by decoding (most index types will have a faster
+     * implementation) */
     void search(
             idx_t n,
             const float* x,
@@ -71,6 +72,11 @@ struct IndexFlatCodes : Index {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) const override;
 
+    virtual void search1(
+            const float* x,
+            ResultHandler& handler,
+            SearchParameters* params = nullptr) const override;
+
     // returns a new instance of a CodePacker
     CodePacker* get_CodePacker() const;