faiss 0.2.4 → 0.2.5

data/vendor/faiss/faiss/index_factory.cpp

@@ -26,10 +26,12 @@
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexIVFPQFastScan.h>
@@ -42,6 +44,7 @@
 #include <faiss/IndexPQFastScan.h>
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexRefine.h>
+#include <faiss/IndexRowwiseMinMax.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/VectorTransform.h>
@@ -149,11 +152,15 @@ std::map<std::string, AdditiveQuantizer::Search_type_t> aq_search_type = {
         {"_Nqint4", AdditiveQuantizer::ST_norm_qint4},
         {"_Ncqint8", AdditiveQuantizer::ST_norm_cqint8},
         {"_Ncqint4", AdditiveQuantizer::ST_norm_cqint4},
+        {"_Nlsq2x4", AdditiveQuantizer::ST_norm_lsq2x4},
+        {"_Nrq2x4", AdditiveQuantizer::ST_norm_rq2x4},
 };
 
 const std::string aq_def_pattern = "[0-9]+x[0-9]+(_[0-9]+x[0-9]+)*";
 const std::string aq_norm_pattern =
-        "(|_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4)";
+        "(|_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4|_Nlsq2x4|_Nrq2x4)";
+
+const std::string paq_def_pattern = "([0-9]+)x([0-9]+)x([0-9]+)";
 
 AdditiveQuantizer::Search_type_t aq_parse_search_type(
         std::string stok,
@@ -169,7 +176,7 @@ AdditiveQuantizer::Search_type_t aq_parse_search_type(
 std::vector<size_t> aq_parse_nbits(std::string stok) {
     std::vector<size_t> nbits;
     std::smatch sm;
-    while (std::regex_search(stok, sm, std::regex("([0-9]+)x([0-9]+)"))) {
+    while (std::regex_search(stok, sm, std::regex("[^q]([0-9]+)x([0-9]+)"))) {
         int M = std::stoi(sm[1].str());
         int nbit = std::stoi(sm[2].str());
         nbits.resize(nbits.size() + M, nbit);
@@ -341,6 +348,53 @@ IndexIVF* parse_IndexIVF(
         }
         return index_ivf;
     }
+    if (match("(PRQ|PLSQ)" + paq_def_pattern + aq_norm_pattern)) {
+        int nsplits = mres_to_int(sm[2]);
+        int Msub = mres_to_int(sm[3]);
+        int nbit = mres_to_int(sm[4]);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), mt);
+        IndexIVF* index_ivf;
+        if (sm[1].str() == "PRQ") {
+            index_ivf = new IndexIVFProductResidualQuantizer(
+                    get_q(), d, nlist, nsplits, Msub, nbit, mt, st);
+        } else {
+            index_ivf = new IndexIVFProductLocalSearchQuantizer(
+                    get_q(), d, nlist, nsplits, Msub, nbit, mt, st);
+        }
+        return index_ivf;
+    }
+    if (match("(RQ|LSQ)([0-9]+)x4fs(r?)(_[0-9]+)?" + aq_norm_pattern)) {
+        int M = std::stoi(sm[2].str());
+        int bbs = mres_to_int(sm[4], 32, 1);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), mt);
+        IndexIVFAdditiveQuantizerFastScan* index_ivf;
+        if (sm[1].str() == "RQ") {
+            index_ivf = new IndexIVFResidualQuantizerFastScan(
+                    get_q(), d, nlist, M, 4, mt, st, bbs);
+        } else {
+            index_ivf = new IndexIVFLocalSearchQuantizerFastScan(
+                    get_q(), d, nlist, M, 4, mt, st, bbs);
+        }
+        index_ivf->by_residual = (sm[3].str() == "r");
+        return index_ivf;
+    }
+    if (match("(PRQ|PLSQ)([0-9]+)x([0-9]+)x4fs(r?)(_[0-9]+)?" +
+              aq_norm_pattern)) {
+        int nsplits = std::stoi(sm[2].str());
+        int Msub = std::stoi(sm[3].str());
+        int bbs = mres_to_int(sm[5], 32, 1);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), mt);
+        IndexIVFAdditiveQuantizerFastScan* index_ivf;
+        if (sm[1].str() == "PRQ") {
+            index_ivf = new IndexIVFProductResidualQuantizerFastScan(
+                    get_q(), d, nlist, nsplits, Msub, 4, mt, st, bbs);
+        } else {
+            index_ivf = new IndexIVFProductLocalSearchQuantizerFastScan(
+                    get_q(), d, nlist, nsplits, Msub, 4, mt, st, bbs);
+        }
+        index_ivf->by_residual = (sm[4].str() == "r");
+        return index_ivf;
+    }
     if (match("(ITQ|PCA|PCAR)([0-9]+)?,SH([-0-9.e]+)?([gcm])?")) {
         int outdim = mres_to_int(sm[2], d); // is also the number of bits
         std::unique_ptr<VectorTransform> vt;
@@ -421,6 +475,37 @@ IndexHNSW* parse_IndexHNSW(
     return nullptr;
 }
 
+/***************************************************************
+ * Parse IndexNSG
+ */
+
+IndexNSG* parse_IndexNSG(
+        const std::string code_string,
+        int d,
+        MetricType mt,
+        int nsg_R) {
+    std::smatch sm;
+    auto match = [&sm, &code_string](const std::string& pattern) {
+        return re_match(code_string, pattern, sm);
+    };
+
+    if (match("Flat|")) {
+        return new IndexNSGFlat(d, nsg_R, mt);
+    }
+    if (match("PQ([0-9]+)(np)?")) {
+        int M = std::stoi(sm[1].str());
+        IndexNSGPQ* ipq = new IndexNSGPQ(d, M, nsg_R);
+        dynamic_cast<IndexPQ*>(ipq->storage)->do_polysemous_training =
+                sm[2].str() != "np";
+        return ipq;
+    }
+    if (match(sq_pattern)) {
+        return new IndexNSGSQ(d, sq_types[sm[1].str()], nsg_R, mt);
+    }
+
+    return nullptr;
+}
+
 /***************************************************************
  * Parse basic indexes
  */
@@ -454,11 +539,6 @@ Index* parse_other_indexes(
         return new IndexLattice(d, M, nbit, r2);
     }
 
-    // IndexNSGFlat
-    if (match("NSG([0-9]+)(,Flat)?")) {
-        return new IndexNSGFlat(d, std::stoi(sm[1].str()), metric);
-    }
-
     // IndexScalarQuantizer
     if (match(sq_pattern)) {
         return new IndexScalarQuantizer(d, sq_types[description], metric);
@@ -505,6 +585,60 @@ Index* parse_other_indexes(
         return new IndexLocalSearchQuantizer(d, M, nbit, metric, st);
     }
 
+    // IndexProductResidualQuantizer
+    if (match("PRQ" + paq_def_pattern + aq_norm_pattern)) {
+        int nsplits = mres_to_int(sm[1]);
+        int Msub = mres_to_int(sm[2]);
+        int nbit = mres_to_int(sm[3]);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), metric);
+        return new IndexProductResidualQuantizer(
+                d, nsplits, Msub, nbit, metric, st);
+    }
+
+    // IndexProductLocalSearchQuantizer
+    if (match("PLSQ" + paq_def_pattern + aq_norm_pattern)) {
+        int nsplits = mres_to_int(sm[1]);
+        int Msub = mres_to_int(sm[2]);
+        int nbit = mres_to_int(sm[3]);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), metric);
+        return new IndexProductLocalSearchQuantizer(
+                d, nsplits, Msub, nbit, metric, st);
+    }
+
+    // IndexAdditiveQuantizerFastScan
+    // RQ{M}x4fs_{bbs}_{search_type}
+    pattern = "(LSQ|RQ)([0-9]+)x4fs(_[0-9]+)?" + aq_norm_pattern;
+    if (match(pattern)) {
+        int M = std::stoi(sm[2].str());
+        int bbs = mres_to_int(sm[3], 32, 1);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), metric);
+
+        if (sm[1].str() == "RQ") {
+            return new IndexResidualQuantizerFastScan(d, M, 4, metric, st, bbs);
+        } else if (sm[1].str() == "LSQ") {
+            return new IndexLocalSearchQuantizerFastScan(
+                    d, M, 4, metric, st, bbs);
+        }
+    }
+
+    // IndexProductAdditiveQuantizerFastScan
+    // PRQ{nsplits}x{Msub}x4fs_{bbs}_{search_type}
+    pattern = "(PLSQ|PRQ)([0-9]+)x([0-9]+)x4fs(_[0-9]+)?" + aq_norm_pattern;
+    if (match(pattern)) {
+        int nsplits = std::stoi(sm[2].str());
+        int Msub = std::stoi(sm[3].str());
+        int bbs = mres_to_int(sm[4], 32, 1);
+        auto st = aq_parse_search_type(sm[sm.size() - 1].str(), metric);
+
+        if (sm[1].str() == "PRQ") {
+            return new IndexProductResidualQuantizerFastScan(
+                    d, nsplits, Msub, 4, metric, st, bbs);
+        } else if (sm[1].str() == "PLSQ") {
+            return new IndexProductLocalSearchQuantizerFastScan(
+                    d, nsplits, Msub, 4, metric, st, bbs);
+        }
+    }
+
     return nullptr;
 }
 
@@ -604,6 +738,14 @@ std::unique_ptr<Index> index_factory_sub(
 
     // IndexIDMap -- it turns out is was used both as a prefix and a suffix, so
     // support both
+    if (re_match(description, "(.+),IDMap2", sm) ||
+        re_match(description, "IDMap2,(.+)", sm)) {
+        IndexIDMap2* idmap2 = new IndexIDMap2(
+                index_factory_sub(d, sm[1].str(), metric).release());
+        idmap2->own_fields = true;
+        return std::unique_ptr<Index>(idmap2);
+    }
+
     if (re_match(description, "(.+),IDMap", sm) ||
         re_match(description, "IDMap,(.+)", sm)) {
         IndexIDMap* idmap = new IndexIDMap(
@@ -642,6 +784,53 @@ std::unique_ptr<Index> index_factory_sub(
         return std::unique_ptr<Index>(index);
     }
 
+    // NSG variants (it was unclear in the old version that the separator was a
+    // "," so we support both "_" and ",")
+    if (re_match(description, "NSG([0-9]*)([,_].*)?", sm)) {
+        int nsg_R = mres_to_int(sm[1], 32);
+        // We also accept empty code string (synonym of Flat)
+        std::string code_string =
+                sm[2].length() > 0 ? sm[2].str().substr(1) : "";
+        if (verbose) {
+            printf("parsing NSG string %s code_string=%s nsg_R=%d\n",
+                   description.c_str(),
+                   code_string.c_str(),
+                   nsg_R);
+        }
+
+        IndexNSG* index = parse_IndexNSG(code_string, d, metric, nsg_R);
+        FAISS_THROW_IF_NOT_FMT(
+                index,
+                "could not parse NSG code description %s in %s",
+                code_string.c_str(),
+                description.c_str());
+        return std::unique_ptr<Index>(index);
+    }
+
+    // IndexRowwiseMinMax, fp32 version
+    if (description.compare(0, 7, "MinMax,") == 0) {
+        size_t comma = description.find(",");
+        std::string sub_index_string = description.substr(comma + 1);
+        auto sub_index = index_factory_sub(d, sub_index_string, metric);
+
+        auto index = new IndexRowwiseMinMax(sub_index.release());
+        index->own_fields = true;
+
+        return std::unique_ptr<Index>(index);
+    }
+
+    // IndexRowwiseMinMax, fp16 version
+    if (description.compare(0, 11, "MinMaxFP16,") == 0) {
+        size_t comma = description.find(",");
+        std::string sub_index_string = description.substr(comma + 1);
+        auto sub_index = index_factory_sub(d, sub_index_string, metric);
+
+        auto index = new IndexRowwiseMinMaxFP16(sub_index.release());
+        index->own_fields = true;
+
+        return std::unique_ptr<Index>(index);
+    }
+
     // IndexIVF
     {
         size_t nlist;

data/vendor/faiss/faiss/index_io.h

@@ -50,6 +50,8 @@ 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;
+// don't initialize precomputed table after loading
+const int IO_FLAG_SKIP_PRECOMPUTE_TABLE = 16;
 // try to memmap data (useful to load an ArrayInvertedLists as an
 // OnDiskInvertedLists)
 const int IO_FLAG_MMAP = IO_FLAG_SKIP_IVF_DATA | 0x646f0000;
@@ -63,7 +65,10 @@ IndexBinary* read_index_binary(FILE* f, int io_flags = 0);
 IndexBinary* read_index_binary(IOReader* reader, int io_flags = 0);
 
 void write_VectorTransform(const VectorTransform* vt, const char* fname);
+void write_VectorTransform(const VectorTransform* vt, IOWriter* f);
+
 VectorTransform* read_VectorTransform(const char* fname);
+VectorTransform* read_VectorTransform(IOReader* f);
 
 ProductQuantizer* read_ProductQuantizer(const char* fname);
 ProductQuantizer* read_ProductQuantizer(IOReader* reader);

data/vendor/faiss/faiss/invlists/DirectMap.cpp

@@ -14,6 +14,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 
 namespace faiss {
 

data/vendor/faiss/faiss/invlists/InvertedLists.cpp

@@ -28,7 +28,10 @@ InvertedLists::~InvertedLists() {}
 InvertedLists::idx_t InvertedLists::get_single_id(size_t list_no, size_t offset)
         const {
     assert(offset < list_size(list_no));
-    return get_ids(list_no)[offset];
+    const idx_t* ids = get_ids(list_no);
+    idx_t id = ids[offset];
+    release_ids(list_no, ids);
+    return id;
 }
 
 void InvertedLists::release_codes(size_t, const uint8_t*) const {}

data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp

@@ -278,13 +278,14 @@ void OnDiskInvertedLists::do_mmap() {
     uint8_t* ptro =
             (uint8_t*)mmap(nullptr, totsize, prot, MAP_SHARED, fileno(f), 0);
 
+    fclose(f);
+
     FAISS_THROW_IF_NOT_FMT(
             ptro != MAP_FAILED,
             "could not mmap %s: %s",
             filename.c_str(),
             strerror(errno));
     ptr = ptro;
-    fclose(f);
 }
 
 void OnDiskInvertedLists::update_totsize(size_t new_size) {

data/vendor/faiss/faiss/python/python_callbacks.cpp

@@ -107,3 +107,30 @@ PyCallbackIOReader::~PyCallbackIOReader() {
     PyThreadLock gil;
     Py_DECREF(callback);
 }
+
+/***********************************************************
+ * Callbacks for IDSelector
+ ***********************************************************/
+
+PyCallbackIDSelector::PyCallbackIDSelector(PyObject* callback)
+        : callback(callback) {
+    PyThreadLock gil;
+    Py_INCREF(callback);
+}
+
+bool PyCallbackIDSelector::is_member(idx_t id) const {
+    FAISS_THROW_IF_NOT((id >> 32) == 0);
+    PyThreadLock gil;
+    PyObject* result = PyObject_CallFunction(callback, "(n)", int(id));
+    if (result == NULL) {
+        FAISS_THROW_MSG("propagate py error");
+    }
+    bool b = PyObject_IsTrue(result);
+    Py_DECREF(result);
+    return b;
+}
+
+PyCallbackIDSelector::~PyCallbackIDSelector() {
+    PyThreadLock gil;
+    Py_DECREF(callback);
+}

data/vendor/faiss/faiss/python/python_callbacks.h

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/io.h>
 #include <faiss/invlists/InvertedLists.h>
 #include "Python.h"
@@ -43,3 +44,17 @@ struct PyCallbackIOReader : faiss::IOReader {
 
     ~PyCallbackIOReader() override;
 };
+
+/***********************************************************
+ * Callbacks for IDSelector
+ ***********************************************************/
+
+struct PyCallbackIDSelector : faiss::IDSelector {
+    PyObject* callback;
+
+    explicit PyCallbackIDSelector(PyObject* callback);
+
+    bool is_member(idx_t id) const override;
+
+    ~PyCallbackIDSelector() override;
+};

data/vendor/faiss/faiss/utils/Heap.h

@@ -47,21 +47,25 @@ inline void heap_pop(size_t k, typename C::T* bh_val, typename C::TI* bh_ids) {
     bh_val--; /* Use 1-based indexing for easier node->child translation */
     bh_ids--;
     typename C::T val = bh_val[k];
+    typename C::TI id = bh_ids[k];
     size_t i = 1, i1, i2;
     while (1) {
         i1 = i << 1;
         i2 = i1 + 1;
         if (i1 > k)
             break;
-        if (i2 == k + 1 || C::cmp(bh_val[i1], bh_val[i2])) {
-            if (C::cmp(val, bh_val[i1]))
+        if ((i2 == k + 1) ||
+            C::cmp2(bh_val[i1], bh_val[i2], bh_ids[i1], bh_ids[i2])) {
+            if (C::cmp2(val, bh_val[i1], id, bh_ids[i1])) {
                 break;
+            }
             bh_val[i] = bh_val[i1];
             bh_ids[i] = bh_ids[i1];
             i = i1;
         } else {
-            if (C::cmp(val, bh_val[i2]))
+            if (C::cmp2(val, bh_val[i2], id, bh_ids[i2])) {
                 break;
+            }
             bh_val[i] = bh_val[i2];
             bh_ids[i] = bh_ids[i2];
             i = i2;
@@ -80,24 +84,28 @@ inline void heap_push(
         typename C::T* bh_val,
         typename C::TI* bh_ids,
         typename C::T val,
-        typename C::TI ids) {
+        typename C::TI id) {
     bh_val--; /* Use 1-based indexing for easier node->child translation */
     bh_ids--;
     size_t i = k, i_father;
     while (i > 1) {
         i_father = i >> 1;
-        if (!C::cmp(val, bh_val[i_father])) /* the heap structure is ok */
+        if (!C::cmp2(val, bh_val[i_father], id, bh_ids[i_father])) {
+            /* the heap structure is ok */
             break;
+        }
         bh_val[i] = bh_val[i_father];
         bh_ids[i] = bh_ids[i_father];
         i = i_father;
     }
     bh_val[i] = val;
-    bh_ids[i] = ids;
+    bh_ids[i] = id;
 }
 
-/** Replace the top element from the heap defined by bh_val[0..k-1] and
- * bh_ids[0..k-1].
+/**
+ * Replaces the top element from the heap defined by bh_val[0..k-1] and
+ * bh_ids[0..k-1], and for identical bh_val[] values also sorts by bh_ids[]
+ * values.
  */
 template <class C>
 inline void heap_replace_top(
@@ -105,31 +113,39 @@ inline void heap_replace_top(
         typename C::T* bh_val,
         typename C::TI* bh_ids,
         typename C::T val,
-        typename C::TI ids) {
+        typename C::TI id) {
     bh_val--; /* Use 1-based indexing for easier node->child translation */
     bh_ids--;
     size_t i = 1, i1, i2;
     while (1) {
         i1 = i << 1;
         i2 = i1 + 1;
-        if (i1 > k)
+        if (i1 > k) {
             break;
-        if (i2 == k + 1 || C::cmp(bh_val[i1], bh_val[i2])) {
-            if (C::cmp(val, bh_val[i1]))
+        }
+
+        // Note that C::cmp2() is a bool function answering
+        // `(a1 > b1) || ((a1 == b1) && (a2 > b2))` for max
+        // heap and same with the `<` sign for min heap.
+        if ((i2 == k + 1) ||
+            C::cmp2(bh_val[i1], bh_val[i2], bh_ids[i1], bh_ids[i2])) {
+            if (C::cmp2(val, bh_val[i1], id, bh_ids[i1])) {
                 break;
+            }
             bh_val[i] = bh_val[i1];
             bh_ids[i] = bh_ids[i1];
             i = i1;
         } else {
-            if (C::cmp(val, bh_val[i2]))
+            if (C::cmp2(val, bh_val[i2], id, bh_ids[i2])) {
                 break;
+            }
             bh_val[i] = bh_val[i2];
             bh_ids[i] = bh_ids[i2];
             i = i2;
         }
     }
     bh_val[i] = val;
-    bh_ids[i] = ids;
+    bh_ids[i] = id;
 }
 
 /* Partial instanciation for heaps with TI = int64_t */
@@ -294,7 +310,7 @@ inline void maxheap_addn(
  * Heap finalization (reorder elements)
  *******************************************************************/
 
-/* This function maps a binary heap into an sorted structure.
+/* This function maps a binary heap into a sorted structure.
    It returns the number  */
 template <typename C>
 inline size_t heap_reorder(