faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/IndexNSG.cpp

@@ -78,9 +78,10 @@ void IndexNSG::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const
-
-{
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT_MSG(
             storage,
             "Please use IndexNSGFlat (or variants) instead of IndexNSG directly");
@@ -104,9 +105,7 @@ void IndexNSG::search(
                 float* simi = distances + i * k;
                 dis->set_query(x + i * d);
 
-                maxheap_heapify(k, simi, idxi);
                 nsg.search(*dis, k, idxi, simi, vt);
-                maxheap_reorder(k, simi, idxi);
 
                 vt.advance();
             }
@@ -300,4 +299,37 @@ IndexNSGFlat::IndexNSGFlat(int d, int R, MetricType metric)
     is_trained = true;
 }
 
+/**************************************************************
+ * IndexNSGPQ implementation
+ **************************************************************/
+
+IndexNSGPQ::IndexNSGPQ() {}
+
+IndexNSGPQ::IndexNSGPQ(int d, int pq_m, int M)
+        : IndexNSG(new IndexPQ(d, pq_m, 8), M) {
+    own_fields = true;
+    is_trained = false;
+}
+
+void IndexNSGPQ::train(idx_t n, const float* x) {
+    IndexNSG::train(n, x);
+    (dynamic_cast<IndexPQ*>(storage))->pq.compute_sdc_table();
+}
+
+/**************************************************************
+ * IndexNSGSQ implementation
+ **************************************************************/
+
+IndexNSGSQ::IndexNSGSQ(
+        int d,
+        ScalarQuantizer::QuantizerType qtype,
+        int M,
+        MetricType metric)
+        : IndexNSG(new IndexScalarQuantizer(d, qtype, metric), M) {
+    is_trained = false;
+    own_fields = true;
+}
+
+IndexNSGSQ::IndexNSGSQ() {}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexNSG.h

@@ -13,6 +13,8 @@
 
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexNNDescent.h>
+#include <faiss/IndexPQ.h>
+#include <faiss/IndexScalarQuantizer.h>
 #include <faiss/impl/NSG.h>
 #include <faiss/utils/utils.h>
 
@@ -64,7 +66,8 @@ struct IndexNSG : Index {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
 
@@ -82,4 +85,25 @@ struct IndexNSGFlat : IndexNSG {
     IndexNSGFlat(int d, int R, MetricType metric = METRIC_L2);
 };
 
+/** PQ index topped with with a NSG structure to access elements
+ *  more efficiently.
+ */
+struct IndexNSGPQ : IndexNSG {
+    IndexNSGPQ();
+    IndexNSGPQ(int d, int pq_m, int M);
+    void train(idx_t n, const float* x) override;
+};
+
+/** SQ index topped with with a NSG structure to access elements
+ *  more efficiently.
+ */
+struct IndexNSGSQ : IndexNSG {
+    IndexNSGSQ();
+    IndexNSGSQ(
+            int d,
+            ScalarQuantizer::QuantizerType qtype,
+            int M,
+            MetricType metric = METRIC_L2);
+};
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexPQ.cpp

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #include <faiss/IndexPQ.h>
 
 #include <cinttypes>
@@ -17,7 +15,7 @@
 
 #include <algorithm>
 
-#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/hamming.h>
 
@@ -28,12 +26,13 @@ namespace faiss {
  ********************************************************/
 
 IndexPQ::IndexPQ(int d, size_t M, size_t nbits, MetricType metric)
-        : Index(d, metric), pq(d, M, nbits) {
+        : IndexFlatCodes(0, d, metric), pq(d, M, nbits) {
     is_trained = false;
     do_polysemous_training = false;
     polysemous_ht = nbits * M + 1;
     search_type = ST_PQ;
     encode_signs = false;
+    code_size = pq.code_size;
 }
 
 IndexPQ::IndexPQ() {
@@ -69,69 +68,19 @@ void IndexPQ::train(idx_t n, const float* x) {
     is_trained = true;
 }
 
-void IndexPQ::add(idx_t n, const float* x) {
-    FAISS_THROW_IF_NOT(is_trained);
-    codes.resize((n + ntotal) * pq.code_size);
-    pq.compute_codes(x, &codes[ntotal * pq.code_size], n);
-    ntotal += n;
-}
-
-size_t IndexPQ::remove_ids(const IDSelector& sel) {
-    idx_t j = 0;
-    for (idx_t i = 0; i < ntotal; i++) {
-        if (sel.is_member(i)) {
-            // should be removed
-        } else {
-            if (i > j) {
-                memmove(&codes[pq.code_size * j],
-                        &codes[pq.code_size * i],
-                        pq.code_size);
-            }
-            j++;
-        }
-    }
-    size_t nremove = ntotal - j;
-    if (nremove > 0) {
-        ntotal = j;
-        codes.resize(ntotal * pq.code_size);
-    }
-    return nremove;
-}
-
-void IndexPQ::reset() {
-    codes.clear();
-    ntotal = 0;
-}
-
-void IndexPQ::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
-    FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
-    for (idx_t i = 0; i < ni; i++) {
-        const uint8_t* code = &codes[(i0 + i) * pq.code_size];
-        pq.decode(code, recons + i * d);
-    }
-}
-
-void IndexPQ::reconstruct(idx_t key, float* recons) const {
-    FAISS_THROW_IF_NOT(key >= 0 && key < ntotal);
-    pq.decode(&codes[key * pq.code_size], recons);
-}
-
 namespace {
 
 template <class PQDecoder>
-struct PQDistanceComputer : DistanceComputer {
+struct PQDistanceComputer : FlatCodesDistanceComputer {
     size_t d;
     MetricType metric;
     Index::idx_t nb;
-    const uint8_t* codes;
-    size_t code_size;
     const ProductQuantizer& pq;
     const float* sdc;
     std::vector<float> precomputed_table;
     size_t ndis;
 
-    float operator()(idx_t i) override {
-        const uint8_t* code = codes + i * code_size;
+    float distance_to_code(const uint8_t* code) final {
         const float* dt = precomputed_table.data();
         PQDecoder decoder(code, pq.nbits);
         float accu = 0;
@@ -158,13 +107,15 @@ struct PQDistanceComputer : DistanceComputer {
         return accu;
     }
 
-    explicit PQDistanceComputer(const IndexPQ& storage) : pq(storage.pq) {
+    explicit PQDistanceComputer(const IndexPQ& storage)
+            : FlatCodesDistanceComputer(
+                      storage.codes.data(),
+                      storage.code_size),
+              pq(storage.pq) {
         precomputed_table.resize(pq.M * pq.ksub);
         nb = storage.ntotal;
         d = storage.d;
         metric = storage.metric_type;
-        codes = storage.codes.data();
-        code_size = pq.code_size;
         if (pq.sdc_table.size() == pq.ksub * pq.ksub * pq.M) {
             sdc = pq.sdc_table.data();
         } else {
@@ -184,7 +135,7 @@ struct PQDistanceComputer : DistanceComputer {
 
 } // namespace
 
-DistanceComputer* IndexPQ::get_distance_computer() const {
+FlatCodesDistanceComputer* IndexPQ::get_FlatCodesDistanceComputer() const {
     if (pq.nbits == 8) {
         return new PQDistanceComputer<PQDecoder8>(*this);
     } else if (pq.nbits == 16) {
@@ -203,10 +154,21 @@ void IndexPQ::search(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        const SearchParameters* iparams) const {
     FAISS_THROW_IF_NOT(k > 0);
-
     FAISS_THROW_IF_NOT(is_trained);
+
+    const SearchParametersPQ* params = nullptr;
+    Search_type_t search_type = this->search_type;
+
+    if (iparams) {
+        params = dynamic_cast<const SearchParametersPQ*>(iparams);
+        FAISS_THROW_IF_NOT_MSG(params, "invalid search params");
+        FAISS_THROW_IF_NOT_MSG(!params->sel, "selector not supported");
+        search_type = params->search_type;
+    }
+
     if (search_type == ST_PQ) { // Simple PQ search
 
         if (metric_type == METRIC_L2) {
@@ -225,8 +187,16 @@ void IndexPQ::search(
             search_type == ST_polysemous ||
             search_type == ST_polysemous_generalize) {
         FAISS_THROW_IF_NOT(metric_type == METRIC_L2);
-
-        search_core_polysemous(n, x, k, distances, labels);
+        int polysemous_ht =
+                params ? params->polysemous_ht : this->polysemous_ht;
+        search_core_polysemous(
+                n,
+                x,
+                k,
+                distances,
+                labels,
+                polysemous_ht,
+                search_type == ST_polysemous_generalize);
 
     } else { // code-to-code distances
 
@@ -302,12 +272,12 @@ static size_t polysemous_inner_loop(
         const uint8_t* q_code,
         size_t k,
         float* heap_dis,
-        int64_t* heap_ids) {
+        int64_t* heap_ids,
+        int ht) {
     int M = index.pq.M;
     int code_size = index.pq.code_size;
     int ksub = index.pq.ksub;
     size_t ntotal = index.ntotal;
-    int ht = index.polysemous_ht;
 
     const uint8_t* b_code = index.codes.data();
 
@@ -342,11 +312,16 @@ void IndexPQ::search_core_polysemous(
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
+        idx_t* labels,
+        int polysemous_ht,
+        bool generalized_hamming) const {
     FAISS_THROW_IF_NOT(k > 0);
-
     FAISS_THROW_IF_NOT(pq.nbits == 8);
 
+    if (polysemous_ht == 0) {
+        polysemous_ht = pq.nbits * pq.M + 1;
+    }
+
     // PQ distance tables
     float* dis_tables = new float[n * pq.ksub * pq.M];
     ScopeDeleter<float> del(dis_tables);
@@ -369,7 +344,9 @@ void IndexPQ::search_core_polysemous(
 
     size_t n_pass = 0;
 
-#pragma omp parallel for reduction(+ : n_pass)
+    int bad_code_size = 0;
+
+#pragma omp parallel for reduction(+ : n_pass, bad_code_size)
     for (idx_t qi = 0; qi < n; qi++) {
         const uint8_t* q_code = q_codes + qi * pq.code_size;
 
@@ -379,28 +356,24 @@ void IndexPQ::search_core_polysemous(
         float* heap_dis = distances + qi * k;
         maxheap_heapify(k, heap_dis, heap_ids);
 
-        if (search_type == ST_polysemous) {
+        if (!generalized_hamming) {
             switch (pq.code_size) {
-                case 4:
-                    n_pass += polysemous_inner_loop<HammingComputer4>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
-                case 8:
-                    n_pass += polysemous_inner_loop<HammingComputer8>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
-                case 16:
-                    n_pass += polysemous_inner_loop<HammingComputer16>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
-                case 32:
-                    n_pass += polysemous_inner_loop<HammingComputer32>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
-                case 20:
-                    n_pass += polysemous_inner_loop<HammingComputer20>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
+#define DISPATCH(cs)                                          \
+    case cs:                                                  \
+        n_pass += polysemous_inner_loop<HammingComputer##cs>( \
+                *this,                                        \
+                dis_table_qi,                                 \
+                q_code,                                       \
+                k,                                            \
+                heap_dis,                                     \
+                heap_ids,                                     \
+                polysemous_ht);                               \
+        break;
+                DISPATCH(4)
+                DISPATCH(8)
+                DISPATCH(16)
+                DISPATCH(32)
+                DISPATCH(20)
                 default:
                     if (pq.code_size % 4 == 0) {
                         n_pass += polysemous_inner_loop<HammingComputerDefault>(
@@ -409,28 +382,30 @@ void IndexPQ::search_core_polysemous(
                                 q_code,
                                 k,
                                 heap_dis,
-                                heap_ids);
+                                heap_ids,
+                                polysemous_ht);
                     } else {
-                        FAISS_THROW_FMT(
-                                "code size %zd not supported for polysemous",
-                                pq.code_size);
+                        bad_code_size++;
                     }
                     break;
             }
-        } else {
+#undef DISPATCH
+        } else { // generalized hamming
             switch (pq.code_size) {
-                case 8:
-                    n_pass += polysemous_inner_loop<GenHammingComputer8>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
-                case 16:
-                    n_pass += polysemous_inner_loop<GenHammingComputer16>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
-                case 32:
-                    n_pass += polysemous_inner_loop<GenHammingComputer32>(
-                            *this, dis_table_qi, q_code, k, heap_dis, heap_ids);
-                    break;
+#define DISPATCH(cs)                                             \
+    case cs:                                                     \
+        n_pass += polysemous_inner_loop<GenHammingComputer##cs>( \
+                *this,                                           \
+                dis_table_qi,                                    \
+                q_code,                                          \
+                k,                                               \
+                heap_dis,                                        \
+                heap_ids,                                        \
+                polysemous_ht);                                  \
+        break;
+                DISPATCH(8)
+                DISPATCH(16)
+                DISPATCH(32)
                 default:
                     if (pq.code_size % 8 == 0) {
                         n_pass += polysemous_inner_loop<GenHammingComputerM8>(
@@ -439,27 +414,29 @@ void IndexPQ::search_core_polysemous(
                                 q_code,
                                 k,
                                 heap_dis,
-                                heap_ids);
+                                heap_ids,
+                                polysemous_ht);
                     } else {
-                        FAISS_THROW_FMT(
-                                "code size %zd not supported for polysemous",
-                                pq.code_size);
+                        bad_code_size++;
                     }
                     break;
+#undef DISPATCH
             }
         }
         maxheap_reorder(k, heap_dis, heap_ids);
     }
 
+    if (bad_code_size) {
+        FAISS_THROW_FMT(
+                "code size %zd not supported for polysemous", pq.code_size);
+    }
+
     indexPQ_stats.nq += n;
     indexPQ_stats.ncode += n * ntotal;
     indexPQ_stats.n_hamming_pass += n_pass;
 }
 
 /* The standalone codec interface (just remaps to the PQ functions) */
-size_t IndexPQ::sa_code_size() const {
-    return pq.code_size;
-}
 
 void IndexPQ::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
     pq.compute_codes(x, bytes, n);
@@ -914,19 +891,25 @@ void MultiIndexQuantizer::train(idx_t n, const float* x) {
         ntotal *= pq.ksub;
 }
 
+// block size used in MultiIndexQuantizer::search
+int multi_index_quantizer_search_bs = 32768;
+
 void MultiIndexQuantizer::search(
         idx_t n,
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const {
-    if (n == 0)
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
+    if (n == 0) {
         return;
-
+    }
     FAISS_THROW_IF_NOT(k > 0);
 
     // the allocation just below can be severe...
-    idx_t bs = 32768;
+    idx_t bs = multi_index_quantizer_search_bs;
     if (n > bs) {
         for (idx_t i0 = 0; i0 < n; i0 += bs) {
             idx_t i1 = std::min(i0 + bs, n);
@@ -1061,9 +1044,14 @@ void MultiIndexQuantizer2::search(
         const float* x,
         idx_t K,
         float* distances,
-        idx_t* labels) const {
-    if (n == 0)
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !params, "search params not supported for this index");
+
+    if (n == 0) {
         return;
+    }
 
     int k2 = std::min(K, int64_t(pq.ksub));
     FAISS_THROW_IF_NOT(k2);

data/vendor/faiss/faiss/IndexPQ.h

@@ -12,7 +12,7 @@
 
 #include <vector>
 
-#include <faiss/Index.h>
+#include <faiss/IndexFlatCodes.h>
 #include <faiss/impl/PolysemousTraining.h>
 #include <faiss/impl/ProductQuantizer.h>
 #include <faiss/impl/platform_macros.h>
@@ -21,13 +21,10 @@ namespace faiss {
 
 /** Index based on a product quantizer. Stored vectors are
  * approximated by PQ codes. */
-struct IndexPQ : Index {
+struct IndexPQ : IndexFlatCodes {
     /// The product quantizer used to encode the vectors
     ProductQuantizer pq;
 
-    /// Codes. Size ntotal * pq.code_size
-    std::vector<uint8_t> codes;
-
     /** Constructor.
      *
      * @param d      dimensionality of the input vectors
@@ -43,31 +40,20 @@ struct IndexPQ : Index {
 
     void train(idx_t n, const float* x) override;
 
-    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 override;
-
-    void reset() override;
-
-    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
-
-    void reconstruct(idx_t key, float* recons) const override;
-
-    size_t remove_ids(const IDSelector& sel) override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     /* The standalone codec interface */
-    size_t sa_code_size() const override;
-
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
 
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 
-    DistanceComputer* get_distance_computer() const override;
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
 
     /******************************************************
      * Polysemous codes implementation
@@ -102,7 +88,9 @@ struct IndexPQ : Index {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const;
+            idx_t* labels,
+            int polysemous_ht,
+            bool generalized_hamming) const;
 
     /// prepare query for a polysemous search, but instead of
     /// computing the result, just get the histogram of Hamming
@@ -124,6 +112,12 @@ struct IndexPQ : Index {
     void hamming_distance_table(idx_t n, const float* x, int32_t* dis) const;
 };
 
+/// override search parameters from the class
+struct SearchParametersPQ : SearchParameters {
+    IndexPQ::Search_type_t search_type;
+    int polysemous_ht;
+};
+
 /// statistics are robust to internal threading, but not if
 /// IndexPQ::search is called by multiple threads
 struct IndexPQStats {
@@ -157,7 +151,8 @@ struct MultiIndexQuantizer : Index {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     /// add and reset will crash at runtime
     void add(idx_t n, const float* x) override;
@@ -168,6 +163,9 @@ struct MultiIndexQuantizer : Index {
     void reconstruct(idx_t key, float* recons) const override;
 };
 
+// block size used in MultiIndexQuantizer::search
+FAISS_API extern int multi_index_quantizer_search_bs;
+
 /** MultiIndexQuantizer where the PQ assignmnet is performed by sub-indexes
  */
 struct MultiIndexQuantizer2 : MultiIndexQuantizer {
@@ -190,7 +188,8 @@ struct MultiIndexQuantizer2 : MultiIndexQuantizer {
             const float* x,
             idx_t k,
             float* distances,
-            idx_t* labels) const override;
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 };
 
 } // namespace faiss