faiss 0.1.1 → 0.1.2

data/vendor/faiss/impl/index_read.cpp

@@ -19,6 +19,7 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/io.h>
+#include <faiss/utils/hamming.h>
 
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
@@ -41,6 +42,7 @@
 #include <faiss/IndexBinaryFromFloat.h>
 #include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexBinaryHash.h>
 
 
 
@@ -364,6 +366,25 @@ ProductQuantizer * read_ProductQuantizer (IOReader *reader) {
   return pq;
 }
 
+static void read_direct_map (DirectMap *dm, IOReader *f) {
+    char maintain_direct_map;
+    READ1 (maintain_direct_map);
+    dm->type = (DirectMap::Type)maintain_direct_map;
+    READVECTOR (dm->array);
+    if (dm->type == DirectMap::Hashtable) {
+        using idx_t = Index::idx_t;
+        std::vector<std::pair<idx_t, idx_t>> v;
+        READVECTOR (v);
+        std::unordered_map<idx_t, idx_t> & map = dm->hashtable;
+        map.reserve (v.size());
+        for (auto it: v) {
+            map [it.first] = it.second;
+        }
+    }
+
+}
+
+
 static void read_ivf_header (
     IndexIVF *ivf, IOReader *f,
     std::vector<std::vector<Index::idx_t> > *ids = nullptr)
@@ -378,8 +399,7 @@ static void read_ivf_header (
         for (size_t i = 0; i < ivf->nlist; i++)
             READVECTOR ((*ids)[i]);
     }
-    READ1 (ivf->maintain_direct_map);
-    READVECTOR (ivf->direct_map);
+    read_direct_map (&ivf->direct_map, f);
 }
 
 // used for legacy formats
@@ -437,10 +457,15 @@ Index *read_index (IOReader *f, int io_flags) {
     Index * idx = nullptr;
     uint32_t h;
     READ1 (h);
-    if (h == fourcc ("IxFI") || h == fourcc ("IxF2")) {
+    if (h == fourcc ("IxFI") || h == fourcc ("IxF2") || h == fourcc("IxFl")) {
         IndexFlat *idxf;
-        if (h == fourcc ("IxFI")) idxf = new IndexFlatIP ();
-        else                      idxf = new IndexFlatL2 ();
+        if (h == fourcc ("IxFI")) {
+            idxf = new IndexFlatIP ();
+        } else if (h == fourcc("IxF2")) {
+            idxf = new IndexFlatL2 ();
+        } else {
+            idxf = new IndexFlat ();
+        }
         read_index_header (idxf, f);
         READVECTOR (idxf->xb);
         FAISS_THROW_IF_NOT (idxf->xb.size() == idxf->ntotal * idxf->d);
@@ -726,10 +751,59 @@ static void read_binary_ivf_header (
         for (size_t i = 0; i < ivf->nlist; i++)
             READVECTOR ((*ids)[i]);
     }
-    READ1 (ivf->maintain_direct_map);
-    READVECTOR (ivf->direct_map);
+    read_direct_map (&ivf->direct_map, f);
+}
+
+static void read_binary_hash_invlists (
+        IndexBinaryHash::InvertedListMap &invlists,
+        int b, IOReader *f)
+{
+    size_t sz;
+    READ1 (sz);
+    int il_nbit = 0;
+    READ1 (il_nbit);
+    // buffer for bitstrings
+    std::vector<uint8_t> buf((b + il_nbit) * sz);
+    READVECTOR (buf);
+    BitstringReader rd (buf.data(), buf.size());
+    invlists.reserve (sz);
+    for (size_t i = 0; i < sz; i++) {
+        uint64_t hash = rd.read(b);
+        uint64_t ilsz = rd.read(il_nbit);
+        auto & il = invlists[hash];
+        READVECTOR (il.ids);
+        FAISS_THROW_IF_NOT (il.ids.size() == ilsz);
+        READVECTOR (il.vecs);
+    }
+}
+
+static void read_binary_multi_hash_map(
+        IndexBinaryMultiHash::Map &map,
+        int b, size_t ntotal,
+        IOReader *f)
+{
+    int id_bits;
+    size_t sz;
+    READ1 (id_bits);
+    READ1 (sz);
+    std::vector<uint8_t> buf;
+    READVECTOR (buf);
+    size_t nbit = (b + id_bits) * sz + ntotal * id_bits;
+    FAISS_THROW_IF_NOT (buf.size() == (nbit + 7) / 8);
+    BitstringReader rd (buf.data(), buf.size());
+    map.reserve (sz);
+    for (size_t i = 0; i < sz; i++) {
+        uint64_t hash = rd.read(b);
+        uint64_t ilsz = rd.read(id_bits);
+        auto & il = map[hash];
+        for (size_t j = 0; j < ilsz; j++) {
+            il.push_back (rd.read (id_bits));
+        }
+    }
 }
 
+
+
 IndexBinary *read_index_binary (IOReader *f, int io_flags) {
     IndexBinary * idx = nullptr;
     uint32_t h;
@@ -771,6 +845,28 @@ IndexBinary *read_index_binary (IOReader *f, int io_flags) {
             static_cast<IndexBinaryIDMap2*>(idxmap)->construct_rev_map ();
         }
         idx = idxmap;
+    } else if(h == fourcc("IBHh")) {
+        IndexBinaryHash *idxh = new IndexBinaryHash ();
+        read_index_binary_header (idxh, f);
+        READ1 (idxh->b);
+        READ1 (idxh->nflip);
+        read_binary_hash_invlists(idxh->invlists, idxh->b, f);
+        idx = idxh;
+    } else if(h == fourcc("IBHm")) {
+        IndexBinaryMultiHash* idxmh = new IndexBinaryMultiHash ();
+        read_index_binary_header (idxmh, f);
+        idxmh->storage = dynamic_cast<IndexBinaryFlat*> (read_index_binary (f));
+        FAISS_THROW_IF_NOT(idxmh->storage && idxmh->storage->ntotal == idxmh->ntotal);
+        idxmh->own_fields = true;
+        READ1 (idxmh->b);
+        READ1 (idxmh->nhash);
+        READ1 (idxmh->nflip);
+        idxmh->maps.resize (idxmh->nhash);
+        for (int i = 0; i < idxmh->nhash; i++) {
+            read_binary_multi_hash_map(
+                    idxmh->maps[i], idxmh->b, idxmh->ntotal, f);
+        }
+        idx = idxmh;
     } else {
         FAISS_THROW_FMT("Index type 0x%08x not supported\n", h);
         idx = nullptr;

data/vendor/faiss/impl/index_write.cpp

@@ -19,6 +19,7 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/io.h>
+#include <faiss/utils/hamming.h>
 
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
@@ -41,6 +42,7 @@
 #include <faiss/IndexBinaryFromFloat.h>
 #include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexBinaryHash.h>
 
 
 
@@ -286,20 +288,33 @@ static void write_HNSW (const HNSW *hnsw, IOWriter *f) {
     WRITE1 (hnsw->upper_beam);
 }
 
+static void write_direct_map (const DirectMap *dm, IOWriter *f) {
+    char maintain_direct_map = (char)dm->type; // for backwards compatibility with bool
+    WRITE1 (maintain_direct_map);
+    WRITEVECTOR (dm->array);
+    if (dm->type == DirectMap::Hashtable) {
+        using idx_t = Index::idx_t;
+        std::vector<std::pair<idx_t, idx_t>> v;
+        const std::unordered_map<idx_t, idx_t> & map = dm->hashtable;
+        v.resize (map.size());
+        std::copy(map.begin(), map.end(), v.begin());
+        WRITEVECTOR (v);
+    }
+}
+
 static void write_ivf_header (const IndexIVF *ivf, IOWriter *f) {
     write_index_header (ivf, f);
     WRITE1 (ivf->nlist);
     WRITE1 (ivf->nprobe);
     write_index (ivf->quantizer, f);
-    WRITE1 (ivf->maintain_direct_map);
-    WRITEVECTOR (ivf->direct_map);
+    write_direct_map (&ivf->direct_map, f);
 }
 
 void write_index (const Index *idx, IOWriter *f) {
     if (const IndexFlat * idxf = dynamic_cast<const IndexFlat *> (idx)) {
         uint32_t h = fourcc (
               idxf->metric_type == METRIC_INNER_PRODUCT ? "IxFI" :
-              idxf->metric_type == METRIC_L2 ? "IxF2" : nullptr);
+              idxf->metric_type == METRIC_L2 ? "IxF2" : "IxFl");
         WRITE1 (h);
         write_index_header (idx, f);
         WRITEVECTOR (idxf->xb);
@@ -499,8 +514,68 @@ static void write_binary_ivf_header (const IndexBinaryIVF *ivf, IOWriter *f) {
     WRITE1 (ivf->nlist);
     WRITE1 (ivf->nprobe);
     write_index_binary (ivf->quantizer, f);
-    WRITE1 (ivf->maintain_direct_map);
-    WRITEVECTOR (ivf->direct_map);
+    write_direct_map (&ivf->direct_map, f);
+}
+
+static void write_binary_hash_invlists (
+        const IndexBinaryHash::InvertedListMap &invlists,
+        int b, IOWriter *f)
+{
+    size_t sz = invlists.size();
+    WRITE1 (sz);
+    size_t maxil = 0;
+    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
+        if(it->second.ids.size() > maxil) {
+            maxil = it->second.ids.size();
+        }
+    }
+    int il_nbit = 0;
+    while(maxil >= ((uint64_t)1 << il_nbit)) {
+        il_nbit++;
+    }
+    WRITE1(il_nbit);
+
+    // first write sizes then data, may be useful if we want to
+    // memmap it at some point
+
+    // buffer for bitstrings
+    std::vector<uint8_t> buf (((b + il_nbit) * sz + 7) / 8);
+    BitstringWriter wr (buf.data(), buf.size());
+    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
+        wr.write (it->first, b);
+        wr.write (it->second.ids.size(), il_nbit);
+    }
+    WRITEVECTOR (buf);
+
+    for (auto it = invlists.begin(); it != invlists.end(); ++it) {
+        WRITEVECTOR (it->second.ids);
+        WRITEVECTOR (it->second.vecs);
+    }
+}
+
+static void write_binary_multi_hash_map(
+        const IndexBinaryMultiHash::Map &map,
+        int b, size_t ntotal,
+        IOWriter *f)
+{
+    int id_bits = 0;
+    while ((ntotal > ((Index::idx_t)1 << id_bits))) {
+        id_bits++;
+    }
+    WRITE1(id_bits);
+    size_t sz = map.size();
+    WRITE1(sz);
+    size_t nbit = (b + id_bits) * sz + ntotal * id_bits;
+    std::vector<uint8_t> buf((nbit + 7) / 8);
+    BitstringWriter wr (buf.data(), buf.size());
+    for (auto it = map.begin(); it != map.end(); ++it) {
+        wr.write(it->first, b);
+        wr.write(it->second.size(), id_bits);
+        for (auto id : it->second) {
+            wr.write(id, id_bits);
+        }
+    }
+    WRITEVECTOR (buf);
 }
 
 void write_index_binary (const IndexBinary *idx, IOWriter *f) {
@@ -539,6 +614,27 @@ void write_index_binary (const IndexBinary *idx, IOWriter *f) {
         write_index_binary_header (idxmap, f);
         write_index_binary (idxmap->index, f);
         WRITEVECTOR (idxmap->id_map);
+    } else if (const IndexBinaryHash *idxh =
+               dynamic_cast<const IndexBinaryHash *> (idx)) {
+        uint32_t h = fourcc ("IBHh");
+        WRITE1 (h);
+        write_index_binary_header (idxh, f);
+        WRITE1 (idxh->b);
+        WRITE1 (idxh->nflip);
+        write_binary_hash_invlists(idxh->invlists, idxh->b, f);
+    } else if (const IndexBinaryMultiHash *idxmh =
+               dynamic_cast<const IndexBinaryMultiHash *> (idx)) {
+        uint32_t h = fourcc ("IBHm");
+        WRITE1 (h);
+        write_index_binary_header (idxmh, f);
+        write_index_binary (idxmh->storage, f);
+        WRITE1 (idxmh->b);
+        WRITE1 (idxmh->nhash);
+        WRITE1 (idxmh->nflip);
+        for (int i = 0; i < idxmh->nhash; i++) {
+            write_binary_multi_hash_map(
+                    idxmh->maps[i], idxmh->b, idxmh->ntotal, f);
+        }
     } else {
         FAISS_THROW_MSG ("don't know how to serialize this type of index");
     }

data/vendor/faiss/impl/io.cpp

@@ -37,7 +37,6 @@ int IOWriter::fileno ()
  ***********************************************************************/
 
 
-
 size_t VectorIOWriter::operator()(
                 const void *ptr, size_t size, size_t nitems)
 {
@@ -132,6 +131,117 @@ int FileIOWriter::fileno()  {
     return ::fileno (f);
 }
 
+/***********************************************************************
+ * IO buffer
+ ***********************************************************************/
+
+BufferedIOReader::BufferedIOReader(IOReader *reader, size_t bsz, size_t totsz):
+    reader(reader), bsz(bsz), totsz(totsz), ofs(0), b0(0), b1(0), buffer(bsz)
+{
+}
+
+
+size_t BufferedIOReader::operator()(void *ptr, size_t unitsize, size_t nitems)
+{
+    size_t size = unitsize * nitems;
+    if (size == 0) return 0;
+    char * dst = (char*)ptr;
+    size_t nb;
+
+    { // first copy available bytes
+        nb = std::min(b1 - b0, size);
+        memcpy (dst, buffer.data() + b0, nb);
+        b0 += nb;
+        dst += nb;
+        size -= nb;
+    }
+
+    if (size > totsz - ofs) {
+        size = totsz - ofs;
+    }
+    // while we would like to have more data
+    while (size > 0) {
+        assert (b0 == b1); // buffer empty on input
+        // try to read from main reader
+        b0 = 0;
+        b1 = (*reader)(buffer.data(), 1, std::min(bsz, size));
+
+        if (b1 == 0) {
+            // no more bytes available
+            break;
+        }
+        ofs += b1;
+
+        // copy remaining bytes
+        size_t nb2 = std::min(b1, size);
+        memcpy (dst, buffer.data(), nb2);
+        b0 = nb2;
+        nb += nb2;
+        dst += nb2;
+        size -= nb2;
+    }
+    return nb / unitsize;
+}
+
+
+BufferedIOWriter::BufferedIOWriter(IOWriter *writer, size_t bsz):
+    writer(writer), bsz(bsz), b0(0), buffer(bsz)
+{
+}
+
+size_t BufferedIOWriter::operator()(const void *ptr, size_t unitsize, size_t nitems)
+{
+    size_t size = unitsize * nitems;
+    if (size == 0) return 0;
+    const char * src = (const char*)ptr;
+    size_t nb;
+
+    { // copy as many bytes as possible to buffer
+        nb = std::min(bsz - b0, size);
+        memcpy (buffer.data() + b0, src, nb);
+        b0 += nb;
+        src += nb;
+        size -= nb;
+    }
+    while (size > 0) {
+        assert(b0 == bsz);
+        // now we need to flush to add more bytes
+        size_t ofs = 0;
+        do {
+            assert (ofs < 10000000);
+            size_t written = (*writer)(buffer.data() + ofs, 1, bsz - ofs);
+            FAISS_THROW_IF_NOT(written > 0);
+            ofs += written;
+        } while(ofs != bsz);
+
+        // copy src to buffer
+        size_t nb1 = std::min(bsz, size);
+        memcpy (buffer.data(), src, nb1);
+        b0 = nb1;
+        nb += nb1;
+        src += nb1;
+        size -= nb1;
+    }
+
+    return nb / unitsize;
+}
+
+BufferedIOWriter::~BufferedIOWriter()
+{
+    size_t ofs = 0;
+    while(ofs != b0) {
+        printf("Destructor write %ld \n", b0 - ofs);
+        size_t written = (*writer)(buffer.data() + ofs, 1, b0 - ofs);
+        FAISS_THROW_IF_NOT(written > 0);
+        ofs += written;
+    }
+
+}
+
+
+
+
+
 uint32_t fourcc (const char sx[4]) {
     assert(4 == strlen(sx));
     const unsigned char *x = (unsigned char*)sx;

data/vendor/faiss/impl/io.h

@@ -9,6 +9,9 @@
 
 /***********************************************************
  * Abstract I/O objects
+ *
+ * I/O is always sequential, seek does not need to be supported
+ * (indexes could be read or written to a pipe).
  ***********************************************************/
 
 #pragma once
@@ -92,6 +95,41 @@ struct FileIOWriter: IOWriter {
     int fileno() override;
 };
 
+/*******************************************************
+ * Buffered reader + writer
+ *******************************************************/
+
+
+
+/** wraps an ioreader to make buffered reads to avoid too small reads */
+struct BufferedIOReader: IOReader {
+
+    IOReader *reader;
+    size_t bsz, totsz, ofs;
+    size_t b0, b1; ///< range of available bytes in the buffer
+    std::vector<char> buffer;
+
+    BufferedIOReader(IOReader *reader, size_t bsz,
+                     size_t totsz=(size_t)(-1));
+
+    size_t operator()(void *ptr, size_t size, size_t nitems) override;
+};
+
+struct BufferedIOWriter: IOWriter {
+
+    IOWriter *writer;
+    size_t bsz, ofs;
+    size_t b0; ///< amount of data in buffer
+    std::vector<char> buffer;
+
+    BufferedIOWriter(IOWriter *writer, size_t bsz);
+
+    size_t operator()(const void *ptr, size_t size, size_t nitems) override;
+
+    // flushes
+    ~BufferedIOWriter();
+};
+
 /// cast a 4-character string to a uint32_t that can be written and read easily
 uint32_t fourcc (const char sx[4]);