faiss 0.1.1 → 0.1.2

data/vendor/faiss/IndexBinaryHash.h

@@ -0,0 +1,116 @@
+/**
+ * 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++ -*-
+
+#ifndef FAISS_BINARY_HASH_H
+#define FAISS_BINARY_HASH_H
+
+
+
+#include <vector>
+#include <unordered_map>
+
+#include <faiss/IndexBinary.h>
+#include <faiss/IndexBinaryFlat.h>
+#include <faiss/utils/Heap.h>
+
+
+namespace faiss {
+
+struct RangeSearchResult;
+
+
+/** just uses the b first bits as a hash value */
+struct IndexBinaryHash : IndexBinary {
+
+    struct InvertedList {
+        std::vector<idx_t> ids;
+        std::vector<uint8_t> vecs;
+
+        void add (idx_t id, size_t code_size, const uint8_t *code);
+    };
+
+    using InvertedListMap = std::unordered_map<idx_t, InvertedList>;
+    InvertedListMap invlists;
+
+    int b, nflip;
+
+    IndexBinaryHash(int d, int b);
+
+    IndexBinaryHash();
+
+    void reset() override;
+
+    void add(idx_t n, const uint8_t *x) override;
+
+    void add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids) override;
+
+    void range_search(idx_t n, const uint8_t *x, int radius,
+                      RangeSearchResult *result) const override;
+
+    void search(idx_t n, const uint8_t *x, idx_t k,
+                int32_t *distances, idx_t *labels) const override;
+
+    void display() const;
+    size_t hashtable_size() const;
+
+};
+
+struct IndexBinaryHashStats {
+    size_t nq;       // nb of queries run
+    size_t n0;       // nb of empty lists
+    size_t nlist;    // nb of non-empty inverted lists scanned
+    size_t ndis;     // nb of distancs computed
+
+    IndexBinaryHashStats () {reset (); }
+    void reset ();
+};
+
+extern IndexBinaryHashStats indexBinaryHash_stats;
+
+
+/** just uses the b first bits as a hash value */
+struct IndexBinaryMultiHash: IndexBinary {
+
+    // where the vectors are actually stored
+    IndexBinaryFlat *storage;
+    bool own_fields;
+
+    // maps hash values to the ids that hash to them
+    using Map = std::unordered_map<idx_t, std::vector<idx_t> >;
+
+    // the different hashes, size nhash
+    std::vector<Map> maps;
+
+    int nhash; ///< nb of hash maps
+    int b; ///< nb bits per hash map
+    int nflip; ///< nb bit flips to use at search time
+
+    IndexBinaryMultiHash(int d, int nhash, int b);
+
+    IndexBinaryMultiHash();
+
+    ~IndexBinaryMultiHash();
+
+    void reset() override;
+
+    void add(idx_t n, const uint8_t *x) override;
+
+    void range_search(idx_t n, const uint8_t *x, int radius,
+                      RangeSearchResult *result) const override;
+
+     void search(idx_t n, const uint8_t *x, idx_t k,
+                int32_t *distances, idx_t *labels) const override;
+
+    size_t hashtable_size() const;
+
+};
+
+}
+
+#endif

data/vendor/faiss/IndexBinaryIVF.cpp

@@ -11,14 +11,17 @@
 #include <faiss/IndexBinaryIVF.h>
 
 #include <cstdio>
+#include <omp.h>
+
 #include <memory>
 
+
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
-
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/IndexFlat.h>
+#include <faiss/IndexLSH.h>
 
 
 namespace faiss {
@@ -29,7 +32,6 @@ IndexBinaryIVF::IndexBinaryIVF(IndexBinary *quantizer, size_t d, size_t nlist)
       own_invlists(true),
       nprobe(1),
       max_codes(0),
-      maintain_direct_map(false),
       quantizer(quantizer),
       nlist(nlist),
       own_fields(false),
@@ -46,7 +48,6 @@ IndexBinaryIVF::IndexBinaryIVF()
       own_invlists(false),
       nprobe(1),
       max_codes(0),
-      maintain_direct_map(false),
       quantizer(nullptr),
       nlist(0),
       own_fields(false),
@@ -65,8 +66,7 @@ void IndexBinaryIVF::add_core(idx_t n, const uint8_t *x, const idx_t *xids,
                               const idx_t *precomputed_idx) {
   FAISS_THROW_IF_NOT(is_trained);
   assert(invlists);
-  FAISS_THROW_IF_NOT_MSG(!(maintain_direct_map && xids),
-                         "cannot have direct map and add with ids");
+  direct_map.check_can_add (xids);
 
   const idx_t * idx;
 
@@ -85,13 +85,15 @@ void IndexBinaryIVF::add_core(idx_t n, const uint8_t *x, const idx_t *xids,
     idx_t id = xids ? xids[i] : ntotal + i;
     idx_t list_no = idx[i];
 
-    if (list_no < 0)
-      continue;
-    const uint8_t *xi = x + i * code_size;
-    size_t offset = invlists->add_entry(list_no, id, xi);
+    if (list_no < 0) {
+        direct_map.add_single_id (id, -1, 0);
+    } else {
+        const uint8_t *xi = x + i * code_size;
+        size_t offset = invlists->add_entry(list_no, id, xi);
+
+        direct_map.add_single_id (id, list_no, offset);
+    }
 
-    if (maintain_direct_map)
-      direct_map.push_back(list_no << 32 | offset);
     n_add++;
   }
   if (verbose) {
@@ -101,29 +103,21 @@ void IndexBinaryIVF::add_core(idx_t n, const uint8_t *x, const idx_t *xids,
   ntotal += n_add;
 }
 
-void IndexBinaryIVF::make_direct_map(bool new_maintain_direct_map) {
-  // nothing to do
-  if (new_maintain_direct_map == maintain_direct_map)
-    return;
-
-  if (new_maintain_direct_map) {
-    direct_map.resize(ntotal, -1);
-    for (size_t key = 0; key < nlist; key++) {
-      size_t list_size = invlists->list_size(key);
-      const idx_t *idlist = invlists->get_ids(key);
-
-      for (size_t ofs = 0; ofs < list_size; ofs++) {
-        FAISS_THROW_IF_NOT_MSG(0 <= idlist[ofs] && idlist[ofs] < ntotal,
-                               "direct map supported only for seuquential ids");
-        direct_map[idlist[ofs]] = key << 32 | ofs;
-      }
+void IndexBinaryIVF::make_direct_map (bool b)
+{
+    if (b) {
+        direct_map.set_type (DirectMap::Array, invlists, ntotal);
+    } else {
+        direct_map.set_type (DirectMap::NoMap, invlists, ntotal);
     }
-  } else {
-    direct_map.clear();
-  }
-  maintain_direct_map = new_maintain_direct_map;
 }
 
+void IndexBinaryIVF::set_direct_map_type (DirectMap::Type type)
+{
+    direct_map.set_type (type, invlists, ntotal);
+}
+
+
 void IndexBinaryIVF::search(idx_t n, const uint8_t *x, idx_t k,
                             int32_t *distances, idx_t *labels) const {
   std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
@@ -142,11 +136,8 @@ void IndexBinaryIVF::search(idx_t n, const uint8_t *x, idx_t k,
 }
 
 void IndexBinaryIVF::reconstruct(idx_t key, uint8_t *recons) const {
-  FAISS_THROW_IF_NOT_MSG(direct_map.size() == ntotal,
-                         "direct map is not initialized");
-  idx_t list_no = direct_map[key] >> 32;
-  idx_t offset = direct_map[key] & 0xffffffff;
-  reconstruct_from_offset(list_no, offset, recons);
+    idx_t lo = direct_map.get (key);
+    reconstruct_from_offset (lo_listno(lo), lo_offset(lo), recons);
 }
 
 void IndexBinaryIVF::reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const {
@@ -215,39 +206,9 @@ void IndexBinaryIVF::reset() {
 }
 
 size_t IndexBinaryIVF::remove_ids(const IDSelector& sel) {
-  FAISS_THROW_IF_NOT_MSG(!maintain_direct_map,
-                         "direct map remove not implemented");
-
-  std::vector<idx_t> toremove(nlist);
-
-#pragma omp parallel for
-  for (idx_t i = 0; i < nlist; i++) {
-    idx_t l0 = invlists->list_size (i), l = l0, j = 0;
-    const idx_t *idsi = invlists->get_ids(i);
-    while (j < l) {
-      if (sel.is_member(idsi[j])) {
-        l--;
-        invlists->update_entry(
-          i, j,
-          invlists->get_single_id(i, l),
-          invlists->get_single_code(i, l));
-      } else {
-        j++;
-      }
-    }
-    toremove[i] = l0 - l;
-  }
-  // this will not run well in parallel on ondisk because of possible shrinks
-  size_t nremove = 0;
-  for (idx_t i = 0; i < nlist; i++) {
-    if (toremove[i] > 0) {
-      nremove += toremove[i];
-      invlists->resize(
-        i, invlists->list_size(i) - toremove[i]);
-    }
-  }
-  ntotal -= nremove;
-  return nremove;
+    size_t nremove = direct_map.remove_ids (sel, invlists);
+    ntotal -= nremove;
+    return nremove;
 }
 
 void IndexBinaryIVF::train(idx_t n, const uint8_t *x) {
@@ -267,9 +228,6 @@ void IndexBinaryIVF::train(idx_t n, const uint8_t *x) {
     Clustering clus(d, nlist, cp);
     quantizer->reset();
 
-    std::unique_ptr<float[]> x_f(new float[n * d]);
-    binary_to_real(n * d, x, x_f.get());
-
     IndexFlatL2 index_tmp(d);
 
     if (clustering_index && verbose) {
@@ -277,8 +235,12 @@ void IndexBinaryIVF::train(idx_t n, const uint8_t *x) {
              clustering_index->d);
     }
 
-    clus.train(n, x_f.get(), clustering_index ? *clustering_index : index_tmp);
+    // LSH codec that is able to convert the binary vectors to floats.
+    IndexLSH codec(d, d, false, false);
+
+    clus.train_encoded (n, x, &codec, clustering_index ? *clustering_index : index_tmp);
 
+    // convert clusters to binary
     std::unique_ptr<uint8_t[]> x_b(new uint8_t[clus.k * code_size]);
     real_to_binary(d * clus.k, clus.centroids.data(), x_b.get());
 
@@ -294,8 +256,7 @@ void IndexBinaryIVF::merge_from(IndexBinaryIVF &other, idx_t add_id) {
   FAISS_THROW_IF_NOT(other.d == d);
   FAISS_THROW_IF_NOT(other.nlist == nlist);
   FAISS_THROW_IF_NOT(other.code_size == code_size);
-  FAISS_THROW_IF_NOT_MSG((!maintain_direct_map &&
-                          !other.maintain_direct_map),
+  FAISS_THROW_IF_NOT_MSG(direct_map.no() && other.direct_map.no(),
                          "direct map copy not implemented");
   FAISS_THROW_IF_NOT_MSG(typeid (*this) == typeid (other),
                          "can only merge indexes of the same type");
@@ -322,13 +283,15 @@ namespace {
 using idx_t = Index::idx_t;
 
 
-template<class HammingComputer, bool store_pairs>
+template<class HammingComputer>
 struct IVFBinaryScannerL2: BinaryInvertedListScanner {
 
     HammingComputer hc;
     size_t code_size;
+    bool store_pairs;
 
-    IVFBinaryScannerL2 (size_t code_size): code_size (code_size)
+    IVFBinaryScannerL2 (size_t code_size, bool store_pairs):
+        code_size (code_size), store_pairs(store_pairs)
     {}
 
     void set_query (const uint8_t *query_vector) override {
@@ -357,7 +320,7 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
             uint32_t dis = hc.hamming (codes);
             if (dis < simi[0]) {
                 heap_pop<C> (k, simi, idxi);
-                idx_t id = store_pairs ? (list_no << 32 | j) : ids[j];
+                idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
                 heap_push<C> (k, simi, idxi, dis, id);
                 nup++;
             }
@@ -366,6 +329,24 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
         return nup;
     }
 
+    void scan_codes_range (size_t n,
+                           const uint8_t *codes,
+                           const idx_t *ids,
+                           int radius,
+                           RangeQueryResult &result) const
+    {
+        size_t nup = 0;
+        for (size_t j = 0; j < n; j++) {
+            uint32_t dis = hc.hamming (codes);
+            if (dis < radius) {
+                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
+                result.add (dis, id);
+            }
+            codes += code_size;
+        }
+
+    }
+
 
 };
 
@@ -373,29 +354,6 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
 template <bool store_pairs>
 BinaryInvertedListScanner *select_IVFBinaryScannerL2 (size_t code_size) {
 
-    switch (code_size) {
-#define HANDLE_CS(cs)                                                  \
-    case cs:                                                            \
-        return new IVFBinaryScannerL2<HammingComputer ## cs, store_pairs> (cs);
-      HANDLE_CS(4);
-      HANDLE_CS(8);
-      HANDLE_CS(16);
-      HANDLE_CS(20);
-      HANDLE_CS(32);
-      HANDLE_CS(64);
-#undef HANDLE_CS
-    default:
-        if (code_size % 8 == 0) {
-            return new IVFBinaryScannerL2<HammingComputerM8,
-                store_pairs> (code_size);
-        } else if (code_size % 4 == 0) {
-            return new IVFBinaryScannerL2<HammingComputerM4,
-                store_pairs> (code_size);
-        } else {
-            return new IVFBinaryScannerL2<HammingComputerDefault,
-                store_pairs> (code_size);
-        }
-    }
 }
 
 
@@ -466,8 +424,10 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
                     ids = sids->get();
                 }
 
-                nheap += scanner->scan_codes (list_size, scodes.get(),
-                                              ids, simi, idxi, k);
+                nheap += scanner->scan_codes (
+                        list_size, scodes.get(),
+                        ids, simi, idxi, k
+                );
 
                 nscan += list_size;
                 if (max_codes && nscan >= max_codes)
@@ -627,11 +587,26 @@ void search_knn_hamming_count_1 (
 BinaryInvertedListScanner *IndexBinaryIVF::get_InvertedListScanner
       (bool store_pairs) const
 {
-    if (store_pairs) {
-        return select_IVFBinaryScannerL2<true> (code_size);
-    } else {
-        return select_IVFBinaryScannerL2<false> (code_size);
+
+#define HC(name) return new IVFBinaryScannerL2<name> (code_size, store_pairs)
+    switch (code_size) {
+    case 4: HC(HammingComputer4);
+    case 8: HC(HammingComputer8);
+    case 16: HC(HammingComputer16);
+    case 20: HC(HammingComputer20);
+    case 32: HC(HammingComputer32);
+    case 64: HC(HammingComputer64);
+    default:
+        if (code_size % 8 == 0) {
+            HC(HammingComputerM8);
+        } else if (code_size % 4 == 0) {
+            HC(HammingComputerM4);
+        } else {
+            HC(HammingComputerDefault);
+        }
     }
+#undef HC
+
 }
 
 void IndexBinaryIVF::search_preassigned(idx_t n, const uint8_t *x, idx_t k,
@@ -657,6 +632,84 @@ void IndexBinaryIVF::search_preassigned(idx_t n, const uint8_t *x, idx_t k,
     }
 }
 
+
+void IndexBinaryIVF::range_search(
+        idx_t n, const uint8_t *x, int radius,
+        RangeSearchResult *res) const
+{
+
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
+
+    double t0 = getmillisecs();
+    quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
+    indexIVF_stats.quantization_time += getmillisecs() - t0;
+
+    t0 = getmillisecs();
+    invlists->prefetch_lists(idx.get(), n * nprobe);
+
+    bool store_pairs = false;
+    size_t nlistv = 0, ndis = 0;
+
+    std::vector<RangeSearchPartialResult *> all_pres (omp_get_max_threads());
+
+#pragma omp parallel reduction(+: nlistv, ndis)
+    {
+        RangeSearchPartialResult pres(res);
+        std::unique_ptr<BinaryInvertedListScanner> scanner
+            (get_InvertedListScanner(store_pairs));
+        FAISS_THROW_IF_NOT (scanner.get ());
+
+        all_pres[omp_get_thread_num()] = &pres;
+
+        auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult &qres)
+        {
+
+            idx_t key = idx[i * nprobe + ik];  /* select the list  */
+            if (key < 0) return;
+            FAISS_THROW_IF_NOT_FMT (
+                    key < (idx_t) nlist,
+                    "Invalid key=%ld  at ik=%ld nlist=%ld\n",
+                    key, ik, nlist);
+            const size_t list_size = invlists->list_size(key);
+
+            if (list_size == 0) return;
+
+            InvertedLists::ScopedCodes scodes (invlists, key);
+            InvertedLists::ScopedIds ids (invlists, key);
+
+            scanner->set_list (key, coarse_dis[i * nprobe + ik]);
+            nlistv++;
+            ndis += list_size;
+            scanner->scan_codes_range (list_size, scodes.get(),
+                                       ids.get(), radius, qres);
+        };
+
+#pragma omp for
+        for (size_t i = 0; i < n; i++) {
+            scanner->set_query (x + i * code_size);
+
+            RangeQueryResult & qres = pres.new_result (i);
+
+            for (size_t ik = 0; ik < nprobe; ik++) {
+                scan_list_func (i, ik, qres);
+            }
+
+        }
+
+        pres.finalize();
+
+    }
+    indexIVF_stats.nq += n;
+    indexIVF_stats.nlist += nlistv;
+    indexIVF_stats.ndis += ndis;
+    indexIVF_stats.search_time += getmillisecs() - t0;
+
+}
+
+
+
+
 IndexBinaryIVF::~IndexBinaryIVF() {
   if (own_invlists) {
     delete invlists;