faiss 0.2.6 → 0.2.7

data/vendor/faiss/faiss/IndexRowwiseMinMax.h

@@ -1,3 +1,10 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #pragma once
 
 #include <cstdint>

data/vendor/faiss/faiss/IndexShards.cpp

@@ -22,98 +22,17 @@ namespace faiss {
 // subroutines
 namespace {
 
-typedef Index::idx_t idx_t;
-
 // add translation to all valid labels
-void translate_labels(long n, idx_t* labels, long translation) {
+void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
     if (translation == 0)
         return;
-    for (long i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n; i++) {
         if (labels[i] < 0)
             continue;
         labels[i] += translation;
     }
 }
 
-/** merge result tables from several shards.
- * @param all_distances  size nshard * n * k
- * @param all_labels     idem
- * @param translartions  label translations to apply, size nshard
- */
-
-template <class IndexClass, class C>
-void merge_tables(
-        long n,
-        long k,
-        long nshard,
-        typename IndexClass::distance_t* distances,
-        idx_t* labels,
-        const std::vector<typename IndexClass::distance_t>& all_distances,
-        const std::vector<idx_t>& all_labels,
-        const std::vector<long>& translations) {
-    if (k == 0) {
-        return;
-    }
-    using distance_t = typename IndexClass::distance_t;
-
-    long stride = n * k;
-#pragma omp parallel
-    {
-        std::vector<int> buf(2 * nshard);
-        int* pointer = buf.data();
-        int* shard_ids = pointer + nshard;
-        std::vector<distance_t> buf2(nshard);
-        distance_t* heap_vals = buf2.data();
-#pragma omp for
-        for (long i = 0; i < n; i++) {
-            // the heap maps values to the shard where they are
-            // produced.
-            const distance_t* D_in = all_distances.data() + i * k;
-            const idx_t* I_in = all_labels.data() + i * k;
-            int heap_size = 0;
-
-            for (long s = 0; s < nshard; s++) {
-                pointer[s] = 0;
-                if (I_in[stride * s] >= 0) {
-                    heap_push<C>(
-                            ++heap_size,
-                            heap_vals,
-                            shard_ids,
-                            D_in[stride * s],
-                            s);
-                }
-            }
-
-            distance_t* D = distances + i * k;
-            idx_t* I = labels + i * k;
-
-            for (int j = 0; j < k; j++) {
-                if (heap_size == 0) {
-                    I[j] = -1;
-                    D[j] = C::neutral();
-                } else {
-                    // pop best element
-                    int s = shard_ids[0];
-                    int& p = pointer[s];
-                    D[j] = heap_vals[0];
-                    I[j] = I_in[stride * s + p] + translations[s];
-
-                    heap_pop<C>(heap_size--, heap_vals, shard_ids);
-                    p++;
-                    if (p < k && I_in[stride * s + p] >= 0) {
-                        heap_push<C>(
-                                ++heap_size,
-                                heap_vals,
-                                shard_ids,
-                                D_in[stride * s + p],
-                                s);
-                    }
-                }
-            }
-        }
-    }
-}
-
 } // anonymous namespace
 
 template <typename IndexT>
@@ -247,11 +166,9 @@ void IndexShardsTemplate<IndexT>::add_with_ids(
 
     if (!ids && !successive_ids) {
         aids.resize(n);
-
         for (idx_t i = 0; i < n; i++) {
             aids[i] = this->ntotal + i;
         }
-
         ids = aids.data();
     }
 
@@ -294,12 +211,23 @@ void IndexShardsTemplate<IndexT>::search(
             !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
-    long nshard = this->count();
+    int64_t nshard = this->count();
 
     std::vector<distance_t> all_distances(nshard * k * n);
     std::vector<idx_t> all_labels(nshard * k * n);
+    std::vector<int64_t> translations(nshard, 0);
+
+    // Because we just called runOnIndex above, it is safe to access the
+    // sub-index ntotal here
+    if (successive_ids) {
+        translations[0] = 0;
 
-    auto fn = [n, k, x, &all_distances, &all_labels](
+        for (int s = 0; s + 1 < nshard; s++) {
+            translations[s + 1] = translations[s] + this->at(s)->ntotal;
+        }
+    }
+
+    auto fn = [n, k, x, &all_distances, &all_labels, &translations](
                       int no, const IndexT* index) {
         if (index->verbose) {
             printf("begin query shard %d on %" PRId64 " points\n", no, n);
@@ -312,6 +240,9 @@ void IndexShardsTemplate<IndexT>::search(
                 all_distances.data() + no * k * n,
                 all_labels.data() + no * k * n);
 
+        translate_labels(
+                n * k, all_labels.data() + no * k * n, translations[no]);
+
         if (index->verbose) {
             printf("end query shard %d\n", no);
         }
@@ -319,38 +250,24 @@ void IndexShardsTemplate<IndexT>::search(
 
     this->runOnIndex(fn);
 
-    std::vector<long> translations(nshard, 0);
-
-    // Because we just called runOnIndex above, it is safe to access the
-    // sub-index ntotal here
-    if (successive_ids) {
-        translations[0] = 0;
-
-        for (int s = 0; s + 1 < nshard; s++) {
-            translations[s + 1] = translations[s] + this->at(s)->ntotal;
-        }
-    }
-
     if (this->metric_type == METRIC_L2) {
-        merge_tables<IndexT, CMin<distance_t, int>>(
+        merge_knn_results<idx_t, CMin<distance_t, int>>(
                 n,
                 k,
                 nshard,
+                all_distances.data(),
+                all_labels.data(),
                 distances,
-                labels,
-                all_distances,
-                all_labels,
-                translations);
+                labels);
     } else {
-        merge_tables<IndexT, CMax<distance_t, int>>(
+        merge_knn_results<idx_t, CMax<distance_t, int>>(
                 n,
                 k,
                 nshard,
+                all_distances.data(),
+                all_labels.data(),
                 distances,
-                labels,
-                all_distances,
-                all_labels,
-                translations);
+                labels);
     }
 }
 

data/vendor/faiss/faiss/IndexShards.h

@@ -18,7 +18,6 @@ namespace faiss {
  */
 template <typename IndexT>
 struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
-    using idx_t = typename IndexT::idx_t;
     using component_t = typename IndexT::component_t;
     using distance_t = typename IndexT::distance_t;
 
@@ -72,7 +71,7 @@ struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
      * Cases (successive_ids, xids):
      * - true, non-NULL       ERROR: it makes no sense to pass in ids and
      *                        request them to be shifted
-     * - true, NULL           OK, but should be called only once (calls add()
+     * - true, NULL           OK: but should be called only once (calls add()
      *                        on sub-indexes).
      * - false, non-NULL      OK: will call add_with_ids with passed in xids
      *                        distributed evenly over shards
@@ -96,7 +95,7 @@ struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
 
     /// Synchronize the top-level index (IndexShards) with data in the
     /// sub-indices
-    void syncWithSubIndexes();
+    virtual void syncWithSubIndexes();
 
    protected:
     /// Called just after an index is added

data/vendor/faiss/faiss/IndexShardsIVF.cpp

@@ -0,0 +1,246 @@
+/**
+ * Copyright (c) Facebook, Inc. and its 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/IndexShardsIVF.h>
+
+#include <cinttypes>
+#include <cstdio>
+#include <functional>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/WorkerThread.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+// subroutines
+namespace {
+
+// add translation to all valid labels
+void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
+    if (translation == 0) {
+        return;
+    }
+    for (int64_t i = 0; i < n; i++) {
+        if (labels[i] < 0) {
+            continue;
+        }
+        labels[i] += translation;
+    }
+}
+
+} // anonymous namespace
+
+/************************************************************
+ * IndexShardsIVF
+ ************************************************************/
+
+IndexShardsIVF::IndexShardsIVF(
+        Index* quantizer,
+        size_t nlist,
+        bool threaded,
+        bool successive_ids)
+        : IndexShardsTemplate<Index>(quantizer->d, threaded, successive_ids),
+          Level1Quantizer(quantizer, nlist) {
+    is_trained = quantizer->is_trained && quantizer->ntotal == nlist;
+}
+
+void IndexShardsIVF::addIndex(Index* index) {
+    auto index_ivf = dynamic_cast<IndexIVFInterface*>(index);
+    FAISS_THROW_IF_NOT_MSG(index_ivf, "can only add IndexIVFs");
+    FAISS_THROW_IF_NOT(index_ivf->nlist == nlist);
+    IndexShardsTemplate<Index>::addIndex(index);
+}
+
+void IndexShardsIVF::train(idx_t n, const component_t* x) {
+    if (verbose) {
+        printf("Training level-1 quantizer\n");
+    }
+    train_q1(n, x, verbose, metric_type);
+
+    // set the sub-quantizer codebooks
+    std::vector<float> centroids(nlist * d);
+    quantizer->reconstruct_n(0, nlist, centroids.data());
+
+    // probably not worth running in parallel
+    for (size_t i = 0; i < indices_.size(); i++) {
+        Index* index = indices_[i].first;
+        auto index_ivf = dynamic_cast<IndexIVFInterface*>(index);
+        Index* quantizer = index_ivf->quantizer;
+        if (!quantizer->is_trained) {
+            quantizer->train(nlist, centroids.data());
+        }
+        quantizer->add(nlist, centroids.data());
+        // finish training
+        index->train(n, x);
+    }
+
+    is_trained = true;
+}
+
+void IndexShardsIVF::add_with_ids(
+        idx_t n,
+        const component_t* x,
+        const idx_t* xids) {
+    // IndexIVF exposes add_core that we can use to factorize the
+    bool all_index_ivf = true;
+    for (size_t i = 0; i < indices_.size(); i++) {
+        Index* index = indices_[i].first;
+        all_index_ivf = all_index_ivf && dynamic_cast<IndexIVF*>(index);
+    }
+    if (!all_index_ivf) {
+        IndexShardsTemplate<Index>::add_with_ids(n, x, xids);
+        return;
+    }
+    FAISS_THROW_IF_NOT_MSG(
+            !(successive_ids && xids),
+            "It makes no sense to pass in ids and "
+            "request them to be shifted");
+
+    if (successive_ids) {
+        FAISS_THROW_IF_NOT_MSG(
+                !xids,
+                "It makes no sense to pass in ids and "
+                "request them to be shifted");
+        FAISS_THROW_IF_NOT_MSG(
+                this->ntotal == 0,
+                "when adding to IndexShards with sucessive_ids, "
+                "only add() in a single pass is supported");
+    }
+
+    // perform coarse quantization
+    std::vector<idx_t> Iq(n);
+    std::vector<float> Dq(n);
+    quantizer->search(n, x, 1, Dq.data(), Iq.data());
+
+    // possibly shift ids
+    idx_t nshard = this->count();
+    const idx_t* ids = xids;
+    std::vector<idx_t> aids;
+    if (!ids && !successive_ids) {
+        aids.resize(n);
+
+        for (idx_t i = 0; i < n; i++) {
+            aids[i] = this->ntotal + i;
+        }
+        ids = aids.data();
+    }
+    idx_t d = this->d;
+
+    auto fn = [n, ids, x, nshard, d, Iq](int no, Index* index) {
+        idx_t i0 = (idx_t)no * n / nshard;
+        idx_t i1 = ((idx_t)no + 1) * n / nshard;
+        const float* x0 = x + i0 * d;
+        auto index_ivf = dynamic_cast<IndexIVF*>(index);
+
+        if (index->verbose) {
+            printf("begin add shard %d on %" PRId64 " points\n", no, n);
+        }
+
+        index_ivf->add_core(
+                i1 - i0, x + i0 * d, ids ? ids + i0 : nullptr, Iq.data() + i0);
+
+        if (index->verbose) {
+            printf("end add shard %d on %" PRId64 " points\n", no, i1 - i0);
+        }
+    };
+
+    this->runOnIndex(fn);
+    syncWithSubIndexes();
+}
+
+void IndexShardsIVF::search(
+        idx_t n,
+        const component_t* x,
+        idx_t k,
+        distance_t* distances,
+        idx_t* labels,
+        const SearchParameters* params_in) const {
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(count() > 0);
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+    }
+
+    auto index0 = dynamic_cast<const IndexIVFInterface*>(at(0));
+    idx_t nprobe = params ? params->nprobe : index0->nprobe;
+
+    // coarse quantization (TODO: support tiling with search_precomputed)
+    std::vector<distance_t> Dq(n * nprobe);
+    std::vector<idx_t> Iq(n * nprobe);
+
+    quantizer->search(n, x, nprobe, Dq.data(), Iq.data());
+
+    int64_t nshard = this->count();
+
+    std::vector<distance_t> all_distances(nshard * k * n);
+    std::vector<idx_t> all_labels(nshard * k * n);
+    std::vector<int64_t> translations(nshard, 0);
+
+    if (successive_ids) {
+        translations[0] = 0;
+        for (int s = 0; s + 1 < nshard; s++) {
+            translations[s + 1] = translations[s] + this->at(s)->ntotal;
+        }
+    }
+
+    auto fn = [&](int no, const Index* indexIn) {
+        if (indexIn->verbose) {
+            printf("begin query shard %d on %" PRId64 " points\n", no, n);
+        }
+
+        auto index = dynamic_cast<const IndexIVFInterface*>(indexIn);
+
+        FAISS_THROW_IF_NOT_MSG(index->nprobe == nprobe, "inconsistent nprobe");
+
+        index->search_preassigned(
+                n,
+                x,
+                k,
+                Iq.data(),
+                Dq.data(),
+                all_distances.data() + no * k * n,
+                all_labels.data() + no * k * n,
+                false);
+
+        translate_labels(
+                n * k, all_labels.data() + no * k * n, translations[no]);
+
+        if (indexIn->verbose) {
+            printf("end query shard %d\n", no);
+        }
+    };
+
+    this->runOnIndex(fn);
+
+    if (this->metric_type == METRIC_L2) {
+        merge_knn_results<idx_t, CMin<distance_t, int>>(
+                n,
+                k,
+                nshard,
+                all_distances.data(),
+                all_labels.data(),
+                distances,
+                labels);
+    } else {
+        merge_knn_results<idx_t, CMax<distance_t, int>>(
+                n,
+                k,
+                nshard,
+                all_distances.data(),
+                all_labels.data(),
+                distances,
+                labels);
+    }
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexShardsIVF.h

@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/IndexIVF.h>
+#include <faiss/IndexShards.h>
+
+namespace faiss {
+
+/**
+ * IndexShards with a common coarse quantizer. All the indexes added should be
+ * IndexIVFInterface indexes so that the search_precomputed can be called.
+ */
+struct IndexShardsIVF : public IndexShards, Level1Quantizer {
+    explicit IndexShardsIVF(
+            Index* quantizer,
+            size_t nlist,
+            bool threaded = false,
+            bool successive_ids = true);
+
+    void addIndex(Index* index) override;
+
+    void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
+            override;
+
+    void train(idx_t n, const component_t* x) override;
+
+    void search(
+            idx_t n,
+            const component_t* x,
+            idx_t k,
+            distance_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/MetaIndexes.cpp

@@ -19,6 +19,8 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/WorkerThread.h>
+#include <faiss/utils/random.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
@@ -154,4 +156,88 @@ IndexSplitVectors::~IndexSplitVectors() {
     }
 }
 
+/********************************************************
+ * IndexRandom implementation
+ */
+
+IndexRandom::IndexRandom(
+        idx_t d,
+        idx_t ntotal,
+        int64_t seed,
+        MetricType metric_type)
+        : Index(d, metric_type), seed(seed) {
+    this->ntotal = ntotal;
+    is_trained = true;
+}
+
+void IndexRandom::add(idx_t n, const float*) {
+    ntotal += n;
+}
+
+void IndexRandom::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
+    FAISS_THROW_IF_NOT(k <= ntotal);
+#pragma omp parallel for if (n > 1000)
+    for (idx_t i = 0; i < n; i++) {
+        RandomGenerator rng(
+                seed + ivec_checksum(d, (const int32_t*)(x + i * d)));
+        idx_t* I = labels + i * k;
+        float* D = distances + i * k;
+        // assumes k << ntotal
+        if (k < 100 * ntotal) {
+            std::unordered_set<idx_t> map;
+            for (int j = 0; j < k; j++) {
+                idx_t ii;
+                for (;;) {
+                    // yes I know it's not strictly uniform...
+                    ii = rng.rand_int64() % ntotal;
+                    if (map.count(ii) == 0) {
+                        break;
+                    }
+                }
+                I[j] = ii;
+                map.insert(ii);
+            }
+        } else {
+            std::vector<idx_t> perm(ntotal);
+            for (idx_t j = 0; j < ntotal; j++) {
+                perm[j] = j;
+            }
+            for (int j = 0; j < k; j++) {
+                std::swap(perm[j], perm[rng.rand_int(ntotal)]);
+                I[j] = perm[j];
+            }
+        }
+        float dprev = 0;
+        for (int j = 0; j < k; j++) {
+            float step = rng.rand_float();
+            if (is_similarity_metric(metric_type)) {
+                step = -step;
+            }
+            dprev += step;
+            D[j] = dprev;
+        }
+    }
+}
+
+void IndexRandom::reconstruct(idx_t key, float* recons) const {
+    RandomGenerator rng(seed + 123332 + key);
+    for (size_t i = 0; i < d; i++) {
+        recons[i] = rng.rand_float();
+    }
+}
+
+void IndexRandom::reset() {
+    ntotal = 0;
+}
+
+IndexRandom::~IndexRandom() {}
+
 } // namespace faiss

data/vendor/faiss/faiss/MetaIndexes.h

@@ -49,6 +49,35 @@ struct IndexSplitVectors : Index {
     ~IndexSplitVectors() override;
 };
 
+/** index that returns random results.
+ * used mainly for time benchmarks
+ */
+struct IndexRandom : Index {
+    int64_t seed;
+
+    explicit IndexRandom(
+            idx_t d,
+            idx_t ntotal = 0,
+            int64_t seed = 1234,
+            MetricType mt = METRIC_L2);
+
+    void add(idx_t n, const float* x) override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    void reset() override;
+
+    ~IndexRandom() override;
+};
+
 } // namespace faiss
 
 #endif

data/vendor/faiss/faiss/MetricType.h

@@ -10,6 +10,8 @@
 #ifndef FAISS_METRIC_TYPE_H
 #define FAISS_METRIC_TYPE_H
 
+#include <faiss/impl/platform_macros.h>
+
 namespace faiss {
 
 /// The metric space for vector comparison for Faiss indices and algorithms.
@@ -29,8 +31,20 @@ enum MetricType {
     METRIC_Canberra = 20,
     METRIC_BrayCurtis,
     METRIC_JensenShannon,
+    METRIC_Jaccard, ///< defined as: sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i))
+                    ///< where a_i, b_i > 0
 };
 
+/// all vector indices are this type
+using idx_t = int64_t;
+
+/// this function is used to distinguish between min and max indexes since
+/// we need to support similarity and dis-similarity metrics in a flexible way
+constexpr bool is_similarity_metric(MetricType metric_type) {
+    return ((metric_type == METRIC_INNER_PRODUCT) ||
+            (metric_type == METRIC_Jaccard));
+}
+
 } // namespace faiss
 
 #endif