faiss 0.1.3 → 0.1.4

data/vendor/faiss/faiss/{DirectMap.cpp → invlists/DirectMap.cpp}

@@ -7,7 +7,7 @@
 
 // -*- c++ -*-
 
-#include <faiss/DirectMap.h>
+#include <faiss/invlists/DirectMap.h>
 
 #include <cstdio>
 #include <cassert>

data/vendor/faiss/faiss/{DirectMap.h → invlists/DirectMap.h}

@@ -10,7 +10,7 @@
 #ifndef FAISS_DIRECT_MAP_H
 #define FAISS_DIRECT_MAP_H
 
-#include <faiss/InvertedLists.h>
+#include <faiss/invlists/InvertedLists.h>
 #include <unordered_map>
 
 

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

@@ -7,7 +7,7 @@
 
 // -*- c++ -*-
 
-#include <faiss/InvertedLists.h>
+#include <faiss/invlists/InvertedLists.h>
 
 #include <cstdio>
 
@@ -616,6 +616,77 @@ void MaskedInvertedLists::prefetch_lists (
     il1->prefetch_lists (list1.data(), list1.size());
 }
 
+/*****************************************
+ * MaskedInvertedLists implementation
+ ******************************************/
+
+
+StopWordsInvertedLists::StopWordsInvertedLists (
+        const InvertedLists *il0, size_t maxsize):
+    ReadOnlyInvertedLists (il0->nlist, il0->code_size),
+    il0 (il0), maxsize (maxsize)
+{
+
+}
+
+size_t StopWordsInvertedLists::list_size(size_t list_no) const
+{
+    size_t sz = il0->list_size(list_no);
+    return sz < maxsize ? sz : 0;
+}
+
+const uint8_t * StopWordsInvertedLists::get_codes (size_t list_no) const
+{
+    return il0->list_size(list_no) < maxsize ? il0->get_codes(list_no) : nullptr;
+}
+
+const idx_t * StopWordsInvertedLists::get_ids (size_t list_no) const
+{
+    return il0->list_size(list_no) < maxsize ? il0->get_ids(list_no) : nullptr;
+}
+
+void StopWordsInvertedLists::release_codes (
+      size_t list_no, const uint8_t *codes) const
+{
+    if (il0->list_size (list_no) < maxsize) {
+        il0->release_codes (list_no, codes);
+    }
+}
+
+void StopWordsInvertedLists::release_ids (size_t list_no, const idx_t *ids) const
+{
+    if (il0->list_size (list_no) < maxsize) {
+        il0->release_ids (list_no, ids);
+    }
+}
+
+idx_t StopWordsInvertedLists::get_single_id (size_t list_no, size_t offset) const
+{
+    FAISS_THROW_IF_NOT(il0->list_size (list_no) < maxsize);
+    return il0->get_single_id (list_no, offset);
+}
+
+const uint8_t * StopWordsInvertedLists::get_single_code (
+           size_t list_no, size_t offset) const
+{
+    FAISS_THROW_IF_NOT(il0->list_size (list_no) < maxsize);
+    return il0->get_single_code (list_no, offset);
+}
+
+void StopWordsInvertedLists::prefetch_lists (
+       const idx_t *list_nos, int nlist) const
+{
+    std::vector<idx_t> list0;
+    for (int i = 0; i < nlist; i++) {
+        idx_t list_no = list_nos[i];
+        if (list_no < 0) continue;
+        if (il0->list_size(list_no) < maxsize) {
+            list0.push_back(list_no);
+        }
+    }
+    il0->prefetch_lists (list0.data(), list0.size());
+}
+
 
 
 } // namespace faiss

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

@@ -36,6 +36,10 @@ struct InvertedLists {
 
     InvertedLists (size_t nlist, size_t code_size);
 
+    /// used for BlockInvertedLists, where the codes are packed into groups
+    /// and the individual code size is meaningless
+    static const size_t INVALID_CODE_SIZE = static_cast<size_t>(-1);
+
     /*************************
      *  Read only functions */
 
@@ -198,7 +202,7 @@ struct ArrayInvertedLists: InvertedLists {
 
     void resize (size_t list_no, size_t new_size) override;
 
-    virtual ~ArrayInvertedLists ();
+     ~ArrayInvertedLists () override;
 };
 
 /*****************************************************************
@@ -329,6 +333,33 @@ struct MaskedInvertedLists: ReadOnlyInvertedLists {
 
 };
 
+
+/** if the inverted list in il is smaller than maxsize then return it,
+ *  otherwise return an empty invlist */
+struct StopWordsInvertedLists: ReadOnlyInvertedLists {
+
+    const InvertedLists *il0;
+    size_t maxsize;
+
+    StopWordsInvertedLists (const InvertedLists *il, size_t maxsize);
+
+    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;
+
+    void release_codes (size_t list_no, const uint8_t *codes) const override;
+    void release_ids (size_t list_no, const idx_t *ids) const override;
+
+    idx_t get_single_id (size_t list_no, size_t offset) const override;
+
+    const uint8_t * get_single_code (
+           size_t list_no, size_t offset) const override;
+
+    void prefetch_lists (const idx_t *list_nos, int nlist) const override;
+
+};
+
+
 } // namespace faiss
 
 

data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp

@@ -0,0 +1,107 @@
+/**
+ * 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/InvertedListsIOHook.h>
+
+#include <faiss/impl/io.h>
+#include <faiss/impl/io_macros.h>
+#include <faiss/impl/FaissAssert.h>
+
+#include <faiss/invlists/BlockInvertedLists.h>
+
+#ifndef _MSC_VER
+#include <faiss/invlists/OnDiskInvertedLists.h>
+#endif // !_MSC_VER
+
+
+namespace faiss {
+
+
+/**********************************************************
+ * InvertedListIOHook's
+ **********************************************************/
+
+InvertedListsIOHook::InvertedListsIOHook(
+        const std::string & key, const std::string & classname):
+    key(key), classname(classname)
+{}
+
+namespace {
+
+/// std::vector that deletes its contents
+struct IOHookTable: std::vector<InvertedListsIOHook*> {
+
+    IOHookTable() {
+#ifndef _MSC_VER
+        push_back(new OnDiskInvertedListsIOHook());
+#endif
+        push_back(new BlockInvertedListsIOHook());
+    }
+
+    ~IOHookTable() {
+        for (auto x: *this) {
+            delete x;
+        }
+    }
+};
+
+static IOHookTable InvertedListsIOHook_table;
+
+} // anonymous namepsace
+
+InvertedListsIOHook* InvertedListsIOHook::lookup(int h)
+{
+    for(const auto & callback: InvertedListsIOHook_table) {
+        if (h == fourcc(callback->key)) {
+            return callback;
+        }
+    }
+    FAISS_THROW_FMT (
+        "read_InvertedLists: could not load ArrayInvertedLists as "
+        "%08x (\"%s\")", h, fourcc_inv_printable(h).c_str()
+    );
+}
+
+InvertedListsIOHook* InvertedListsIOHook::lookup_classname(const std::string & classname)
+{
+    for(const auto & callback: InvertedListsIOHook_table) {
+        if (callback->classname == classname) {
+            return callback;
+        }
+    }
+    FAISS_THROW_FMT (
+            "read_InvertedLists: could not find classname %s",
+            classname.c_str()
+    );
+}
+
+void InvertedListsIOHook::add_callback(InvertedListsIOHook *cb)
+{
+    InvertedListsIOHook_table.push_back(cb);
+}
+
+void InvertedListsIOHook::print_callbacks()
+{
+    printf("registered %zd InvertedListsIOHooks:\n",
+          InvertedListsIOHook_table.size());
+    for(const auto & cb: InvertedListsIOHook_table) {
+        printf("%08x %s %s\n",
+               fourcc(cb->key.c_str()),
+               cb->key.c_str(),
+               cb->classname.c_str());
+    }
+}
+
+InvertedLists * InvertedListsIOHook::read_ArrayInvertedLists(
+        IOReader *, int ,
+        size_t , size_t ,
+        const std::vector<size_t> &) const
+{
+    FAISS_THROW_FMT("read to array not implemented for %s", classname.c_str());
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h

@@ -0,0 +1,63 @@
+/**
+ * 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 <string>
+#include <faiss/invlists/InvertedLists.h>
+#include <faiss/impl/io.h>
+
+
+namespace faiss {
+
+
+/** 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)
+     *
+     * (default implementation fails)
+     */
+    virtual InvertedLists * read_ArrayInvertedLists(
+            IOReader *f, int io_flags,
+            size_t nlist, size_t code_size,
+            const std::vector<size_t> &sizes) const;
+
+    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);
+
+};
+
+} // namespace faiss

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

@@ -7,7 +7,7 @@
 
 // -*- c++ -*-
 
-#include <faiss/OnDiskInvertedLists.h>
+#include <faiss/invlists/OnDiskInvertedLists.h>
 
 #include <pthread.h>
 
@@ -130,7 +130,8 @@ struct LockLevels {
 
     void print () {
         pthread_mutex_lock(&mutex1);
-        printf("State: level3_in_use=%d n_level2=%d level1_holders: [", level3_in_use, n_level2);
+        printf("State: level3_in_use=%d n_level2=%d level1_holders: [",
+               int(level3_in_use), n_level2);
         for (int k : level1_holders) {
             printf("%d ", k);
         }
@@ -299,8 +300,7 @@ void OnDiskInvertedLists::update_totsize (size_t new_size)
     // unmap file
     if (ptr != nullptr) {
         int err = munmap (ptr, totsize);
-        FAISS_THROW_IF_NOT_FMT (err == 0, "munmap error: %s",
-                                strerror(errno));
+        FAISS_THROW_IF_NOT_FMT (err == 0, "munmap error: %s", strerror(errno));
     }
     if (totsize == 0) {
         // must create file before truncating it
@@ -516,8 +516,9 @@ size_t OnDiskInvertedLists::allocate_slot (size_t capacity) {
     if (it == slots.end()) {
         // not enough capacity
         size_t new_size = totsize == 0 ? 32 : totsize * 2;
-        while (new_size - totsize < capacity)
+        while (new_size - totsize < capacity) {
             new_size *= 2;
+        }
         locks->lock_3 ();
         update_totsize(new_size);
         locks->unlock_3 ();
@@ -678,6 +679,18 @@ void OnDiskInvertedLists::crop_invlists(size_t l0, size_t l1)
     nlist = l1 - l0;
 }
 
+
+void OnDiskInvertedLists::set_all_lists_sizes(const size_t *sizes)
+{
+    size_t ofs = 0;
+    for (size_t i = 0; i < nlist; i++) {
+        lists[i].offset = ofs;
+        lists[i].capacity = lists[i].size = sizes[i];
+        ofs += sizes[i] * (sizeof(idx_t) + code_size);
+    }
+
+}
+
 /*******************************************************
  * I/O support via callbacks
  *******************************************************/
@@ -755,7 +768,9 @@ InvertedLists * OnDiskInvertedListsIOHook::read(IOReader *f, int io_flags) const
 
     }
     READ1(od->totsize);
-    od->do_mmap();
+    if (!(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
+        od->do_mmap();
+    }
     return od;
 }
 

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

@@ -15,7 +15,7 @@
 #include <typeinfo>
 
 #include <faiss/IndexIVF.h>
-
+#include <faiss/invlists/InvertedListsIOHook.h>
 #include <faiss/index_io.h>
 
 namespace faiss {
@@ -109,7 +109,7 @@ struct OnDiskInvertedLists: InvertedLists {
 
     void prefetch_lists (const idx_t *list_nos, int nlist) const override;
 
-    virtual ~OnDiskInvertedLists ();
+    ~OnDiskInvertedLists () override;
 
     // private
 
@@ -126,6 +126,9 @@ struct OnDiskInvertedLists: InvertedLists {
     size_t allocate_slot (size_t capacity);
     void free_slot (size_t offset, size_t capacity);
 
+    /// override all list sizes and make a packed storage
+    void set_all_lists_sizes(const size_t *sizes);
+
     // empty constructor for the I/O functions
     OnDiskInvertedLists ();
 };

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

@@ -1,8 +1,15 @@
+/**
+ * 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 "Python.h"
 #include <faiss/impl/io.h>
-#include <faiss/InvertedLists.h>
+#include <faiss/invlists/InvertedLists.h>
 
 //  all callbacks have to acquire the GIL on input