Memory-mapping and Zero-copy deserializers

faiss/CMakeLists.txt

@@ -69,12 +69,12 @@ set(FAISS_SRC
   impl/io.cpp
   impl/kmeans1d.cpp
   impl/lattice_Zn.cpp
+  impl/mapped_io.cpp
   impl/pq4_fast_scan.cpp
   impl/pq4_fast_scan_search_1.cpp
   impl/pq4_fast_scan_search_qbs.cpp
   impl/residual_quantizer_encode_steps.cpp
-  impl/io.cpp
-  impl/lattice_Zn.cpp
+  impl/zerocopy_io.cpp
   impl/NNDescent.cpp
   invlists/BlockInvertedLists.cpp
   invlists/DirectMap.cpp

faiss/IVFlib.cpp

@@ -200,12 +200,32 @@ static void shift_and_add(
     memcpy(dst.data() + insert_point, src.data(), src.size() * sizeof(T));
 }
 
+template <class T>
+static void shift_and_add(
+        MaybeOwnedVector<T>& dst,
+        size_t remove,
+        const MaybeOwnedVector<T>& src) {
+    if (remove > 0)
+        memmove(dst.data(),
+                dst.data() + remove,
+                (dst.size() - remove) * sizeof(T));
+    size_t insert_point = dst.size() - remove;
+    dst.resize(insert_point + src.size());
+    memcpy(dst.data() + insert_point, src.data(), src.size() * sizeof(T));
+}
+
 template <class T>
 static void remove_from_begin(std::vector<T>& v, size_t remove) {
     if (remove > 0)
         v.erase(v.begin(), v.begin() + remove);
 }
 
+template <class T>
+static void remove_from_begin(MaybeOwnedVector<T>& v, size_t remove) {
+    if (remove > 0)
+        v.erase(v.begin(), v.begin() + remove);
+}
+
 void SlidingIndexWindow::step(const Index* sub_index, bool remove_oldest) {
     FAISS_THROW_IF_NOT_MSG(
             !remove_oldest || n_slice > 0,

faiss/IndexBinaryFlat.h

@@ -14,14 +14,15 @@
 
 #include <faiss/IndexBinary.h>
 
+#include <faiss/impl/maybe_owned_vector.h>
 #include <faiss/utils/approx_topk/mode.h>
 
 namespace faiss {
 
 /** Index that stores the full vectors and performs exhaustive search. */
 struct IndexBinaryFlat : IndexBinary {
     /// database vectors, size ntotal * d / 8
-    std::vector<uint8_t> xb;
+    MaybeOwnedVector<uint8_t> xb;
 
     /** Select between using a heap or counting to select the k smallest values
      * when scanning inverted lists.

faiss/IndexFlatCodes.cpp

@@ -110,7 +110,7 @@ CodePacker* IndexFlatCodes::get_CodePacker() const {
 }
 
 void IndexFlatCodes::permute_entries(const idx_t* perm) {
-    std::vector<uint8_t> new_codes(codes.size());
+    MaybeOwnedVector<uint8_t> new_codes(codes.size());
 
     for (idx_t i = 0; i < ntotal; i++) {
         memcpy(new_codes.data() + i * code_size,

faiss/IndexFlatCodes.h

@@ -7,9 +7,11 @@
 
 #pragma once
 
+#include <vector>
+
 #include <faiss/Index.h>
 #include <faiss/impl/DistanceComputer.h>
-#include <vector>
+#include <faiss/impl/maybe_owned_vector.h>
 
 namespace faiss {
 
@@ -21,7 +23,7 @@ struct IndexFlatCodes : Index {
     size_t code_size;
 
     /// encoded dataset, size ntotal * code_size
-    std::vector<uint8_t> codes;
+    MaybeOwnedVector<uint8_t> codes;
 
     IndexFlatCodes();
 

faiss/impl/HNSW.cpp

@@ -1085,7 +1085,7 @@ void HNSW::permute_entries(const idx_t* map) {
     // swap everyone
     std::swap(levels, new_levels);
     std::swap(offsets, new_offsets);
-    std::swap(neighbors, new_neighbors);
+    neighbors = std::move(new_neighbors);
 }
 
 /**************************************************************

faiss/impl/HNSW.h

@@ -15,6 +15,7 @@
 
 #include <faiss/Index.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/maybe_owned_vector.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
@@ -121,7 +122,7 @@ struct HNSW {
 
     /// neighbors[offsets[i]:offsets[i+1]] is the list of neighbors of vector i
     /// for all levels. this is where all storage goes.
-    std::vector<storage_idx_t> neighbors;
+    MaybeOwnedVector<storage_idx_t> neighbors;
 
     /// entry point in the search structure (one of the points with maximum
     /// level