faiss 0.1.3 → 0.1.4

data/vendor/faiss/faiss/index_factory.cpp

@@ -34,6 +34,10 @@
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexLattice.h>
+#include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexIVFPQFastScan.h>
+#include <faiss/IndexRefine.h>
+
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryHNSW.h>
@@ -62,29 +66,90 @@ struct VTChain {
 /// what kind of training does this coarse quantizer require?
 char get_trains_alone(const Index *coarse_quantizer) {
     return
+        dynamic_cast<const IndexFlat*>(coarse_quantizer) ? 0 :
+        // multi index just needs to be quantized
         dynamic_cast<const MultiIndexQuantizer*>(coarse_quantizer) ? 1 :
         dynamic_cast<const IndexHNSWFlat*>(coarse_quantizer) ? 2 :
-        0;
+        2; // for complicated indexes, we assume they can't be used as a kmeans index
 }
 
+bool str_ends_with(const std::string& s, const std::string& suffix)
+{
+    return s.rfind(suffix) == std::abs(int(s.size()-suffix.size()));
+}
+
+// check if ends with suffix followed by digits
+bool str_ends_with_digits(const std::string& s, const std::string& suffix)
+{
+    int i;
+    for(i = s.length() - 1; i >= 0; i--) {
+        if (!isdigit(s[i])) break;
+    }
+    return str_ends_with(s.substr(0, i + 1), suffix);
+}
+
+void find_matching_parentheses(const std::string &s, int & i0, int & i1) {
+    int st = 0;
+    for (int i = 0; i < s.length(); i++) {
+        if (s[i] == '(') {
+            if (st == 0) {
+                i0 = i;
+            }
+            st++;
+        }
+
+        if (s[i] == ')') {
+            st--;
+            if (st == 0) {
+                i1 = i;
+                return;
+            }
+            if (st < 0) {
+                FAISS_THROW_FMT("factory string %s: unbalanced parentheses", s.c_str());
+            }
+        }
+
+    }
+    FAISS_THROW_FMT("factory string %s: unbalanced parentheses st=%d", s.c_str(), st);
 
 }
 
+} // anonymous namespace
+
 Index *index_factory (int d, const char *description_in, MetricType metric)
 {
     FAISS_THROW_IF_NOT(metric == METRIC_L2 ||
                        metric == METRIC_INNER_PRODUCT);
     VTChain vts;
     Index *coarse_quantizer = nullptr;
+    std::string parenthesis_ivf, parenthesis_refine;
     Index *index = nullptr;
     bool add_idmap = false;
-    bool make_IndexRefineFlat = false;
+    int d_in = d;
 
     ScopeDeleter1<Index> del_coarse_quantizer, del_index;
 
     std::string description(description_in);
     char *ptr;
 
+    // handle indexes in parentheses
+    while (description.find('(') != std::string::npos) {
+        // then we make a sub-index and remove the () from the description
+        int i0, i1;
+        find_matching_parentheses(description, i0, i1);
+
+        std::string sub_description = description.substr(i0 + 1, i1 - i0 - 1);
+
+        if (str_ends_with_digits(description.substr(0, i0), "IVF")) {
+            parenthesis_ivf = sub_description;
+        } else if (str_ends_with(description.substr(0, i0), "Refine")) {
+            parenthesis_refine = sub_description;
+        } else {
+            FAISS_THROW_MSG("don't know what to do with parenthesis index");
+        }
+        description = description.erase(i0, i1 - i0 + 1);
+    }
+
     int64_t ncentroids = -1;
     bool use_2layer = false;
     int hnsw_M = -1;
@@ -95,6 +160,8 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
         int d_out, opq_M, nbit, M, M2, pq_m, ncent, r2;
         std::string stok(tok);
         nbit = 8;
+        int bbs = -1;
+        char c;
 
         // to avoid mem leaks with exceptions:
         // do all tests before any instanciation
@@ -140,15 +207,20 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
         // coarse quantizers
         } else if (!coarse_quantizer &&
                    sscanf (tok, "IVF%" PRId64 "_HNSW%d", &ncentroids, &M) == 2) {
-            coarse_quantizer_1 = new IndexHNSWFlat (d, M);
+            coarse_quantizer_1 = new IndexHNSWFlat (d, M, metric);
 
         } else if (!coarse_quantizer &&
                    sscanf (tok, "IVF%" PRId64, &ncentroids) == 1) {
-            if (metric == METRIC_L2) {
+            if (!parenthesis_ivf.empty()) {
+                coarse_quantizer_1 =
+                    index_factory(d, parenthesis_ivf.c_str(), metric);
+
+            } else if (metric == METRIC_L2) {
                 coarse_quantizer_1 = new IndexFlatL2 (d);
             } else {
                 coarse_quantizer_1 = new IndexFlatIP (d);
             }
+
         } else if (!coarse_quantizer && sscanf (tok, "IMI2x%d", &nbit) == 1) {
             FAISS_THROW_IF_NOT_MSG (metric == METRIC_L2,
                              "MultiIndex not implemented for inner prod search");
@@ -228,6 +300,32 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
             del_coarse_quantizer.release ();
             index_ivf->own_fields = true;
             index_1 = index_ivf;
+        } else if (!index && (
+             sscanf (tok, "PQ%dx4fs_%d", &M, &bbs) == 2 ||
+            (sscanf (tok, "PQ%dx4f%c", &M, &c) == 2 && c == 's') ||
+            (sscanf (tok, "PQ%dx4fs%c", &M, &c) == 2 && c == 'r'))) {
+            if (bbs == -1) {
+                bbs = 32;
+            }
+            bool by_residual = str_ends_with(stok, "fsr");
+            if (coarse_quantizer) {
+                IndexIVFPQFastScan *index_ivf = new IndexIVFPQFastScan(
+                    coarse_quantizer, d, ncentroids, M, 4, metric, bbs
+                );
+                index_ivf->quantizer_trains_alone =
+                    get_trains_alone (coarse_quantizer);
+                index_ivf->metric_type = metric;
+                index_ivf->by_residual = by_residual;
+                index_ivf->cp.spherical = metric == METRIC_INNER_PRODUCT;
+                del_coarse_quantizer.release ();
+                index_ivf->own_fields = true;
+                index_1 = index_ivf;
+            } else {
+                IndexPQFastScan *index_pq = new IndexPQFastScan (
+                    d, M, 4, metric, bbs
+                );
+                index_1 = index_pq;
+            }
         } else if (!index && (sscanf (tok, "PQ%dx%d", &M, &nbit) == 2 ||
                               sscanf (tok, "PQ%d", &M) == 1 ||
                               sscanf (tok, "PQ%dnp", &M) == 1)) {
@@ -299,7 +397,12 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
             FAISS_THROW_IF_NOT(!coarse_quantizer);
             index_1 = new IndexLattice(d, M, nbit, r2);
         } else if (stok == "RFlat") {
-            make_IndexRefineFlat = true;
+            parenthesis_refine = "Flat";
+        } else if (stok == "Refine") {
+            FAISS_THROW_IF_NOT_MSG(
+                    !parenthesis_refine.empty(),
+                    "Refine index should be provided in parentheses"
+            );
         } else {
             FAISS_THROW_FMT( "could not parse token \"%s\" in %s\n",
                              tok, description_in);
@@ -356,8 +459,10 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
         index = index_pt;
     }
 
-    if (make_IndexRefineFlat) {
-        IndexRefineFlat *index_rf = new IndexRefineFlat (index);
+    if (!parenthesis_refine.empty()) {
+        Index *refine_index = index_factory(d_in, parenthesis_refine.c_str(), metric);
+        IndexRefine *index_rf = new IndexRefine(index, refine_index);
+        index_rf->own_refine_index = true;
         index_rf->own_fields = true;
         index = index_rf;
     }

data/vendor/faiss/faiss/index_io.h

@@ -51,7 +51,7 @@ const int IO_FLAG_READ_ONLY = 2;
 const int IO_FLAG_ONDISK_SAME_DIR = 4;
 // don't load IVF data to RAM, only list sizes
 const int IO_FLAG_SKIP_IVF_DATA = 8;
-// try to memmap data (useful for OnDiskInvertedLists)
+// try to memmap data (useful to load an ArrayInvertedLists as an OnDiskInvertedLists)
 const int IO_FLAG_MMAP = IO_FLAG_SKIP_IVF_DATA | 0x646f0000;
 
 
@@ -76,53 +76,6 @@ void write_InvertedLists (const InvertedLists *ils, IOWriter *f);
 InvertedLists *read_InvertedLists (IOReader *reader, int io_flags = 0);
 
 
-#ifndef _MSC_VER
-/** Callbacks to handle other types of InvertedList objects.
- *
- * The callbacks should be registered with add_callback before calling
- * read_index or read_InvertedLists. The callbacks for
- * OnDiskInvertedLists are registrered by default. The invlist type is
- * identified by:
- *
- * - the key (a fourcc) at read time
- * - the class name (as given by typeid.name) at write time
- */
-struct InvertedListsIOHook {
-    const std::string key; ///< string version of the fourcc
-    const std::string classname; ///< typeid.name
-
-    InvertedListsIOHook(const std::string & key, const std::string & classname);
-
-    /// write the index to the IOWriter (including the fourcc)
-    virtual void write(const InvertedLists *ils, IOWriter *f) const = 0;
-
-    /// called when the fourcc matches this class's fourcc
-    virtual InvertedLists * read(IOReader *f, int io_flags) const = 0;
-
-    /** read from a ArrayInvertedLists into this invertedlist type.
-     * For this to work, the callback has to be enabled and the io_flag has to be set to
-     * IO_FLAG_SKIP_IVF_DATA | (16 upper bits of the fourcc)
-     */
-    virtual InvertedLists * read_ArrayInvertedLists(
-            IOReader *f, int io_flags,
-            size_t nlist, size_t code_size,
-            const std::vector<size_t> &sizes) const = 0;
-
-    virtual ~InvertedListsIOHook() {}
-
-    /**************************** Manage the set of callbacks ******/
-
-    // transfers ownership
-    static void add_callback(InvertedListsIOHook *);
-    static void print_callbacks();
-    static InvertedListsIOHook* lookup(int h);
-    static InvertedListsIOHook* lookup_classname(const std::string & classname);
-
-};
-
-#endif // !_MSC_VER
-
-
 } // namespace faiss
 
 

data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp

@@ -0,0 +1,151 @@
+/**
+ * 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.
+ */
+
+#include <faiss/invlists/BlockInvertedLists.h>
+
+#include <faiss/impl/FaissAssert.h>
+
+#include <faiss/impl/io.h>
+#include <faiss/impl/io_macros.h>
+
+
+namespace faiss {
+
+BlockInvertedLists::BlockInvertedLists (
+        size_t nlist, size_t n_per_block,
+        size_t block_size):
+    InvertedLists (nlist, InvertedLists::INVALID_CODE_SIZE),
+    n_per_block(n_per_block), block_size(block_size)
+{
+    ids.resize (nlist);
+    codes.resize (nlist);
+}
+
+BlockInvertedLists::BlockInvertedLists ():
+    InvertedLists (0, InvertedLists::INVALID_CODE_SIZE),
+    n_per_block(0), block_size(0)
+{}
+
+
+size_t BlockInvertedLists::add_entries (
+           size_t list_no, size_t n_entry,
+           const idx_t* ids_in, const uint8_t *code)
+{
+    if (n_entry == 0) return 0;
+    FAISS_THROW_IF_NOT (list_no < nlist);
+    size_t o = ids [list_no].size();
+    FAISS_THROW_IF_NOT (o == 0); // not clear how we should handle subsequent adds
+    ids [list_no].resize (o + n_entry);
+    memcpy (&ids[list_no][o], ids_in, sizeof (ids_in[0]) * n_entry);
+
+    // copy whole blocks
+    size_t n_block = (n_entry + n_per_block - 1) / n_per_block;
+    codes [list_no].resize (n_block * block_size);
+    memcpy (&codes[list_no][o * code_size], code, n_block * block_size);
+    return o;
+}
+
+size_t BlockInvertedLists::list_size(size_t list_no) const
+{
+    assert (list_no < nlist);
+    return ids[list_no].size();
+}
+
+const uint8_t * BlockInvertedLists::get_codes (size_t list_no) const
+{
+    assert (list_no < nlist);
+    return codes[list_no].get();
+}
+
+const InvertedLists::idx_t * BlockInvertedLists::get_ids (size_t list_no) const
+{
+    assert (list_no < nlist);
+    return ids[list_no].data();
+}
+
+void BlockInvertedLists::resize (size_t list_no, size_t new_size)
+{
+    ids[list_no].resize (new_size);
+    size_t prev_nbytes = codes[list_no].size();
+    size_t n_block = (new_size + n_per_block - 1) / n_per_block;
+    size_t new_nbytes = n_block * block_size;
+    codes[list_no].resize (new_nbytes);
+    if (prev_nbytes < new_nbytes) {
+        // set new elements to 0
+        memset(
+            codes[list_no].data() + prev_nbytes, 0,
+            new_nbytes - prev_nbytes
+        );
+    }
+}
+
+void BlockInvertedLists::update_entries (
+      size_t , size_t , size_t ,
+      const idx_t *, const uint8_t *)
+{
+    FAISS_THROW_MSG("not impemented");
+    /*
+    assert (list_no < nlist);
+    assert (n_entry + offset <= ids[list_no].size());
+    memcpy (&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry);
+    memcpy (&codes[list_no][offset * code_size], codes_in, code_size * n_entry);
+    */
+}
+
+
+BlockInvertedLists::~BlockInvertedLists ()
+{}
+
+/**************************************************
+ * IO hook implementation
+ **************************************************/
+
+BlockInvertedListsIOHook::BlockInvertedListsIOHook():
+    InvertedListsIOHook("ilbl", typeid(BlockInvertedLists).name())
+{}
+
+
+void BlockInvertedListsIOHook::write(const InvertedLists *ils_in, IOWriter *f) const
+{
+    uint32_t h = fourcc ("ilbl");
+    WRITE1 (h);
+    const BlockInvertedLists *il =
+        dynamic_cast<const BlockInvertedLists*> (ils_in);
+    WRITE1 (il->nlist);
+    WRITE1 (il->code_size);
+    WRITE1 (il->n_per_block);
+    WRITE1 (il->block_size);
+
+    for (size_t i = 0; i < il->nlist; i++) {
+        WRITEVECTOR(il->ids[i]);
+        WRITEVECTOR(il->codes[i]);
+    }
+}
+
+InvertedLists * BlockInvertedListsIOHook::read(IOReader *f, int /* io_flags */) const
+{
+    BlockInvertedLists *il = new BlockInvertedLists();
+    READ1 (il->nlist);
+    READ1 (il->code_size);
+    READ1 (il->n_per_block);
+    READ1 (il->block_size);
+
+    il->ids.resize(il->nlist);
+    il->codes.resize(il->nlist);
+
+    for (size_t i = 0; i < il->nlist; i++) {
+        READVECTOR(il->ids[i]);
+        READVECTOR(il->codes[i]);
+    }
+
+    return il;
+}
+
+
+
+
+} // namespace faiss

data/vendor/faiss/faiss/invlists/BlockInvertedLists.h

@@ -0,0 +1,76 @@
+/**
+ * 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/invlists/InvertedLists.h>
+#include <faiss/invlists/InvertedListsIOHook.h>
+#include <faiss/utils/AlignedTable.h>
+#include <faiss/index_io.h>
+
+namespace faiss {
+
+/** Inverted Lists that are organized by blocks.
+ *
+ * Different from the regular inverted lists, the codes are organized by blocks
+ * of size block_size bytes that reprsent a set of n_per_block. Therefore, code
+ * allocations are always rounded up to block_size bytes. The codes are also
+ * aligned on 32-byte boundaries for use with SIMD.
+ *
+ * To avoid misinterpretations, the code_size is set to (size_t)(-1), even if
+ * arguably the amount of memory consumed by code is block_size / n_per_block.
+ *
+ * The writing functions add_entries and update_entries operate on block-aligned
+ * data.
+ */
+struct BlockInvertedLists: InvertedLists {
+
+    size_t n_per_block;   // nb of vectors stored per block
+    size_t block_size;    // nb bytes per block
+
+    std::vector<AlignedTable<uint8_t>> codes;
+    std::vector<std::vector<idx_t>> ids;
+
+
+    BlockInvertedLists (
+        size_t nlist, size_t vec_per_block,
+        size_t block_size
+    );
+
+    BlockInvertedLists();
+
+    size_t list_size(size_t list_no) const override;
+    const uint8_t * get_codes (size_t list_no) const override;
+    const idx_t * get_ids (size_t list_no) const override;
+
+    // works only on empty BlockInvertedLists
+    // the codes should be of size ceil(n_entry / n_per_block) * block_size
+    // and padded with 0s
+    size_t add_entries (
+           size_t list_no, size_t n_entry,
+           const idx_t* ids, const uint8_t *code) override;
+
+    /// not implemented
+    void update_entries (size_t list_no, size_t offset, size_t n_entry,
+                         const idx_t *ids, const uint8_t *code) override;
+
+    // also pads new data with 0s
+    void resize (size_t list_no, size_t new_size) override;
+
+    ~BlockInvertedLists () override;
+
+};
+
+struct BlockInvertedListsIOHook : InvertedListsIOHook {
+    BlockInvertedListsIOHook();
+    void write(const InvertedLists *ils, IOWriter *f) const override;
+    InvertedLists * read(IOReader *f, int io_flags) const override;
+};
+
+
+} // namespace faiss