faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/index_factory.h

@@ -19,6 +19,9 @@ Index* index_factory(
         const char* description,
         MetricType metric = METRIC_L2);
 
+/// set to > 0 to get more logs from index_factory
+FAISS_API extern int index_factory_verbose;
+
 IndexBinary* index_binary_factory(int d, const char* description);
 
 } // namespace faiss

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(