faiss 0.1.2 → 0.1.3

data/vendor/faiss/{IndexReplicas.h → faiss/IndexReplicas.h}

@@ -65,9 +65,15 @@ class IndexReplicasTemplate : public ThreadedIndex<IndexT> {
   /// reconstructs from the first index
   void reconstruct(idx_t, component_t *v) const override;
 
+  /// Synchronize the top-level index (IndexShards) with data in the sub-indices
+  void syncWithSubIndexes();
+
  protected:
   /// Called just after an index is added
   void onAfterAddIndex(IndexT* index) override;
+
+  /// Called just after an index is removed
+  void onAfterRemoveIndex(IndexT* index) override;
 };
 
 using IndexReplicas = IndexReplicasTemplate<Index>;

data/vendor/faiss/{IndexScalarQuantizer.cpp → faiss/IndexScalarQuantizer.cpp}

@@ -77,7 +77,7 @@ void IndexScalarQuantizer::search(
         ScopeDeleter1<InvertedListScanner> del(scanner);
 
 #pragma omp for
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             float * D = distances + k * i;
             idx_t * I = labels + k * i;
             // re-order heap
@@ -197,7 +197,7 @@ void IndexIVFScalarQuantizer::encode_vectors(idx_t n, const float* x,
         std::vector<float> residual (d);
 
 #pragma omp for
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             int64_t list_no = list_nos [i];
             if (list_no >= 0) {
                 const float *xi = x + i * d;
@@ -227,7 +227,7 @@ void IndexIVFScalarQuantizer::sa_decode (idx_t n, const uint8_t *codes,
         std::vector<float> residual (d);
 
 #pragma omp for
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             const uint8_t *code = codes + i * (code_size + coarse_size);
             int64_t list_no = decode_listno (code);
             float *xi = x + i * d;

data/vendor/faiss/{IndexShards.cpp → faiss/IndexShards.cpp}

@@ -9,6 +9,7 @@
 
 #include <faiss/IndexShards.h>
 
+#include <cinttypes>
 #include <cstdio>
 #include <functional>
 
@@ -132,18 +133,20 @@ IndexShardsTemplate<IndexT>::IndexShardsTemplate(bool threaded,
 template <typename IndexT>
 void
 IndexShardsTemplate<IndexT>::onAfterAddIndex(IndexT* index /* unused */) {
-  sync_with_shard_indexes();
+  syncWithSubIndexes();
 }
 
 template <typename IndexT>
 void
 IndexShardsTemplate<IndexT>::onAfterRemoveIndex(IndexT* index /* unused */) {
-  sync_with_shard_indexes();
+  syncWithSubIndexes();
 }
 
+// FIXME: assumes that nothing is currently running on the sub-indexes, which is
+// true with the normal API, but should use the runOnIndex API instead
 template <typename IndexT>
 void
-IndexShardsTemplate<IndexT>::sync_with_shard_indexes() {
+IndexShardsTemplate<IndexT>::syncWithSubIndexes() {
   if (!this->count()) {
     this->is_trained = false;
     this->ntotal = 0;
@@ -160,6 +163,31 @@ IndexShardsTemplate<IndexT>::sync_with_shard_indexes() {
     auto index = this->at(i);
     FAISS_THROW_IF_NOT(this->metric_type == index->metric_type);
     FAISS_THROW_IF_NOT(this->d == index->d);
+    FAISS_THROW_IF_NOT(this->is_trained == index->is_trained);
+
+    this->ntotal += index->ntotal;
+  }
+}
+
+// No metric_type for IndexBinary
+template <>
+void
+IndexShardsTemplate<IndexBinary>::syncWithSubIndexes() {
+  if (!this->count()) {
+    this->is_trained = false;
+    this->ntotal = 0;
+
+    return;
+  }
+
+  auto firstIndex = this->at(0);
+  this->is_trained = firstIndex->is_trained;
+  this->ntotal = firstIndex->ntotal;
+
+  for (int i = 1; i < this->count(); ++i) {
+    auto index = this->at(i);
+    FAISS_THROW_IF_NOT(this->d == index->d);
+    FAISS_THROW_IF_NOT(this->is_trained == index->is_trained);
 
     this->ntotal += index->ntotal;
   }
@@ -172,7 +200,7 @@ IndexShardsTemplate<IndexT>::train(idx_t n,
   auto fn =
     [n, x](int no, IndexT *index) {
       if (index->verbose) {
-        printf("begin train shard %d on %ld points\n", no, n);
+        printf("begin train shard %d on %" PRId64 " points\n", no, n);
       }
 
       index->train(n, x);
@@ -183,7 +211,7 @@ IndexShardsTemplate<IndexT>::train(idx_t n,
     };
 
   this->runOnIndex(fn);
-  sync_with_shard_indexes();
+  syncWithSubIndexes();
 }
 
 template <typename IndexT>
@@ -237,7 +265,7 @@ IndexShardsTemplate<IndexT>::add_with_ids(idx_t n,
       auto x0 = x + i0 * components_per_vec;
 
       if (index->verbose) {
-        printf ("begin add shard %d on %ld points\n", no, n);
+        printf ("begin add shard %d on %" PRId64 " points\n", no, n);
       }
 
       if (ids) {
@@ -247,15 +275,12 @@ IndexShardsTemplate<IndexT>::add_with_ids(idx_t n,
       }
 
       if (index->verbose) {
-        printf ("end add shard %d on %ld points\n", no, i1 - i0);
+        printf ("end add shard %d on %" PRId64 " points\n", no, i1 - i0);
       }
     };
 
   this->runOnIndex(fn);
-
-  // This is safe to do here because the current thread controls execution in
-  // all threads, and nothing else is happening
-  this->ntotal += n;
+  syncWithSubIndexes();
 }
 
 template <typename IndexT>
@@ -273,7 +298,7 @@ IndexShardsTemplate<IndexT>::search(idx_t n,
   auto fn =
     [n, k, x, &all_distances, &all_labels](int no, const IndexT *index) {
       if (index->verbose) {
-        printf ("begin query shard %d on %ld points\n", no, n);
+        printf ("begin query shard %d on %" PRId64 " points\n", no, n);
       }
 
       index->search (n, x, k,

data/vendor/faiss/{IndexShards.h → faiss/IndexShards.h}

@@ -79,12 +79,11 @@ struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
 
   void train(idx_t n, const component_t* x) override;
 
-  // update metric_type and ntotal. Call if you changes something in
-  // the shard indexes.
-  void sync_with_shard_indexes();
-
   bool successive_ids;
 
+  /// Synchronize the top-level index (IndexShards) with data in the sub-indices
+  void syncWithSubIndexes();
+
  protected:
   /// Called just after an index is added
   void onAfterAddIndex(IndexT* index) override;

data/vendor/faiss/{InvertedLists.cpp → faiss/InvertedLists.cpp}

@@ -274,7 +274,7 @@ const uint8_t * HStackInvertedLists::get_single_code (
         }
         offset -= sz;
     }
-    FAISS_THROW_FMT ("offset %ld unknown", offset);
+    FAISS_THROW_FMT ("offset %zd unknown", offset);
 }
 
 
@@ -309,7 +309,7 @@ Index::idx_t HStackInvertedLists::get_single_id (
         }
         offset -= sz;
     }
-    FAISS_THROW_FMT ("offset %ld unknown", offset);
+    FAISS_THROW_FMT ("offset %zd unknown", offset);
 }
 
 

data/vendor/faiss/{InvertedLists.h → faiss/InvertedLists.h}

@@ -208,6 +208,7 @@ struct ArrayInvertedLists: InvertedLists {
  * that can be stacked horizontally, vertically and sliced.
  *****************************************************************/
 
+/// invlists that fail for all write functions
 struct ReadOnlyInvertedLists: InvertedLists {
 
     ReadOnlyInvertedLists (size_t nlist, size_t code_size):

data/vendor/faiss/{MetaIndexes.cpp → faiss/MetaIndexes.cpp}

@@ -9,7 +9,9 @@
 
 #include <faiss/MetaIndexes.h>
 
+#include <cinttypes>
 #include <cstdio>
+#include <limits>
 #include <stdint.h>
 
 #include <faiss/impl/FaissAssert.h>
@@ -199,7 +201,7 @@ void IndexIDMap2Template<IndexT>::reconstruct
     try {
         this->index->reconstruct (rev_map.at (key), recons);
     } catch (const std::out_of_range& e) {
-        FAISS_THROW_FMT ("key %ld not found", key);
+        FAISS_THROW_FMT ("key %" PRId64 " not found", key);
     }
 }
 
@@ -274,7 +276,7 @@ void IndexSplitVectors::search (
         float *distances1 = no == 0 ? distances : all_distances + no * k * n;
         idx_t *labels1 = no == 0 ? labels : all_labels + no * k * n;
         if (index->verbose)
-            printf ("begin query shard %d on %ld points\n", no, n);
+            printf ("begin query shard %d on %" PRId64 " points\n", no, n);
         const Index * sub_index = index->sub_indexes[no];
         int64_t sub_d = sub_index->d, d = index->d;
         idx_t ofs = 0;
@@ -319,7 +321,7 @@ void IndexSplitVectors::search (
                     distances[j] += distances_i[j];
                 } else {
                     labels[j] = -1;
-                    distances[j] = 0.0 / 0.0;
+                    distances[j] = std::numeric_limits<float>::quiet_NaN();
                 }
             }
         }

data/vendor/faiss/{OnDiskInvertedLists.cpp → faiss/OnDiskInvertedLists.cpp}

@@ -16,10 +16,14 @@
 #include <sys/mman.h>
 #include <unistd.h>
 #include <sys/types.h>
+#include <sys/stat.h>
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/utils.h>
 
+#include <faiss/impl/io.h>
+#include <faiss/impl/io_macros.h>
+
 
 namespace faiss {
 
@@ -320,7 +324,7 @@ void OnDiskInvertedLists::update_totsize (size_t new_size)
     totsize = new_size;
 
     // create file
-    printf ("resizing %s to %ld bytes\n", filename.c_str(), totsize);
+    printf ("resizing %s to %zd bytes\n", filename.c_str(), totsize);
 
     int err = truncate (filename.c_str(), totsize);
 
@@ -341,7 +345,7 @@ void OnDiskInvertedLists::update_totsize (size_t new_size)
 
 #define INVALID_OFFSET (size_t)(-1)
 
-OnDiskInvertedLists::List::List ():
+OnDiskOneList::OnDiskOneList ():
     size (0), capacity (0), offset (INVALID_OFFSET)
 {}
 
@@ -640,7 +644,7 @@ size_t OnDiskInvertedLists::merge_from (const InvertedLists **ils, int n_il,
                 nmerged++;
                 double t1 = getmillisecs();
                 if (t1 - last_t > 500) {
-                    printf("merged %ld lists in %.3f s\r",
+                    printf("merged %zd lists in %.3f s\r",
                            nmerged, (t1 - t0) / 1000.0);
                     fflush(stdout);
                     last_t = t1;
@@ -656,6 +660,12 @@ size_t OnDiskInvertedLists::merge_from (const InvertedLists **ils, int n_il,
 }
 
 
+size_t OnDiskInvertedLists::merge_from_1 (const InvertedLists *ils, bool verbose)
+{
+    return merge_from (&ils, 1, verbose);
+}
+
+
 void OnDiskInvertedLists::crop_invlists(size_t l0, size_t l1)
 {
     FAISS_THROW_IF_NOT(0 <= l0 && l0 <= l1 && l1 <= nlist);
@@ -668,6 +678,134 @@ void OnDiskInvertedLists::crop_invlists(size_t l0, size_t l1)
     nlist = l1 - l0;
 }
 
+/*******************************************************
+ * I/O support via callbacks
+ *******************************************************/
+
+
+
+
+OnDiskInvertedListsIOHook::OnDiskInvertedListsIOHook():
+    InvertedListsIOHook("ilod", typeid(OnDiskInvertedLists).name())
+{}
+
+
+void OnDiskInvertedListsIOHook::write(const InvertedLists *ils, IOWriter *f) const
+{
+    uint32_t h = fourcc ("ilod");
+    WRITE1 (h);
+    WRITE1 (ils->nlist);
+    WRITE1 (ils->code_size);
+    const OnDiskInvertedLists *od = dynamic_cast<const OnDiskInvertedLists*> (ils);
+    // this is a POD object
+    WRITEVECTOR (od->lists);
+
+    {
+        std::vector<OnDiskInvertedLists::Slot> v(
+                od->slots.begin(), od->slots.end());
+        WRITEVECTOR(v);
+    }
+    {
+        std::vector<char> x(od->filename.begin(), od->filename.end());
+        WRITEVECTOR(x);
+    }
+    WRITE1(od->totsize);
+}
+
+InvertedLists * OnDiskInvertedListsIOHook::read(IOReader *f, int io_flags) const
+{
+    OnDiskInvertedLists *od = new OnDiskInvertedLists();
+    od->read_only = io_flags & IO_FLAG_READ_ONLY;
+    READ1 (od->nlist);
+    READ1 (od->code_size);
+    // this is a POD object
+    READVECTOR (od->lists);
+    {
+        std::vector<OnDiskInvertedLists::Slot> v;
+        READVECTOR(v);
+        od->slots.assign(v.begin(), v.end());
+    }
+    {
+        std::vector<char> x;
+        READVECTOR(x);
+        od->filename.assign(x.begin(), x.end());
+
+        if (io_flags & IO_FLAG_ONDISK_SAME_DIR) {
+            FileIOReader *reader = dynamic_cast<FileIOReader*>(f);
+            FAISS_THROW_IF_NOT_MSG (
+                    reader, "IO_FLAG_ONDISK_SAME_DIR only supported "
+                    "when reading from file");
+            std::string indexname = reader->name;
+            std::string dirname = "./";
+            size_t slash = indexname.find_last_of('/');
+            if (slash != std::string::npos) {
+                dirname = indexname.substr(0, slash + 1);
+            }
+            std::string filename = od->filename;
+            slash = filename.find_last_of('/');
+            if (slash != std::string::npos) {
+                filename = filename.substr(slash + 1);
+            }
+            filename = dirname + filename;
+            printf("IO_FLAG_ONDISK_SAME_DIR: "
+                   "updating ondisk filename from %s to %s\n",
+                   od->filename.c_str(), filename.c_str());
+            od->filename = filename;
+        }
+
+    }
+    READ1(od->totsize);
+    od->do_mmap();
+    return od;
+}
+
+/** read from a ArrayInvertedLists into this invertedlist type */
+InvertedLists * OnDiskInvertedListsIOHook::read_ArrayInvertedLists(
+        IOReader *f, int /* io_flags */,
+        size_t nlist, size_t code_size,
+        const std::vector<size_t> &sizes) const
+{
+    auto ails = new OnDiskInvertedLists ();
+    ails->nlist = nlist;
+    ails->code_size = code_size;
+    ails->read_only = true;
+    ails->lists.resize (nlist);
+
+    FileIOReader *reader = dynamic_cast<FileIOReader*>(f);
+    FAISS_THROW_IF_NOT_MSG(reader, "mmap only supported for File objects");
+    FILE *fdesc = reader->f;
+    size_t o0 = ftell(fdesc);
+    size_t o = o0;
+    { // do the mmap
+        struct stat buf;
+        int ret = fstat (fileno(fdesc), &buf);
+        FAISS_THROW_IF_NOT_FMT (ret == 0,
+                                "fstat failed: %s", strerror(errno));
+        ails->totsize = buf.st_size;
+        ails->ptr = (uint8_t*)mmap (nullptr, ails->totsize,
+                                    PROT_READ, MAP_SHARED,
+                                    fileno(fdesc), 0);
+        FAISS_THROW_IF_NOT_FMT (ails->ptr != MAP_FAILED,
+                                "could not mmap: %s",
+                                strerror(errno));
+    }
+
+    FAISS_THROW_IF_NOT(o <= ails->totsize);
+
+    for (size_t i = 0; i < ails->nlist; i++) {
+        OnDiskInvertedLists::List & l = ails->lists[i];
+        l.size = l.capacity = sizes[i];
+        l.offset = o;
+        o += l.size * (sizeof(OnDiskInvertedLists::idx_t) +
+                       ails->code_size);
+    }
+    // resume normal reading of file
+    fseek (fdesc, o, SEEK_SET);
+
+    return ails;
+}
+
+
 
 
 

data/vendor/faiss/{OnDiskInvertedLists.h → faiss/OnDiskInvertedLists.h}

@@ -12,14 +12,26 @@
 
 #include <vector>
 #include <list>
+#include <typeinfo>
 
 #include <faiss/IndexIVF.h>
 
+#include <faiss/index_io.h>
+
 namespace faiss {
 
 
 struct LockLevels;
 
+
+struct OnDiskOneList {
+    size_t size;     // size of inverted list (entries)
+    size_t capacity; // allocated size (entries)
+    size_t offset;   // offset in buffer (bytes)
+    OnDiskOneList ();
+};
+
+
 /** On-disk storage of inverted lists.
  *
  * The data is stored in a mmapped chunk of memory (base ptointer ptr,
@@ -49,13 +61,7 @@ struct LockLevels;
  * lists in parallel.
  */
 struct OnDiskInvertedLists: InvertedLists {
-
-    struct List {
-        size_t size;     // size of inverted list (entries)
-        size_t capacity; // allocated size (entries)
-        size_t offset;   // offset in buffer (bytes)
-        List ();
-    };
+    using List = OnDiskOneList;
 
     // size nlist
     std::vector<List> lists;
@@ -95,6 +101,9 @@ struct OnDiskInvertedLists: InvertedLists {
     // allocating slots)
     size_t merge_from (const InvertedLists **ils, int n_il, bool verbose=false);
 
+    /// same as merge_from for a single invlist
+    size_t merge_from_1 (const InvertedLists *il, bool verbose=false);
+
     /// restrict the inverted lists to l0:l1 without touching the mmapped region
     void crop_invlists(size_t l0, size_t l1);
 
@@ -121,6 +130,17 @@ struct OnDiskInvertedLists: InvertedLists {
     OnDiskInvertedLists ();
 };
 
+struct OnDiskInvertedListsIOHook: InvertedListsIOHook {
+    OnDiskInvertedListsIOHook();
+    void write(const InvertedLists *ils, IOWriter *f) const override;
+    InvertedLists * read(IOReader *f, int io_flags) const override;
+    InvertedLists * read_ArrayInvertedLists(
+            IOReader *f, int io_flags,
+            size_t nlist, size_t code_size,
+            const std::vector<size_t> &sizes) const override;
+};
+
+
 
 } // namespace faiss