mem mapping and zero-copy python fixes (facebookresearch#4212)

Summary:
Pull Request resolved: facebookresearch#4212

Add files to TARGETS
fix python

Reviewed By: mengdilin

Differential Revision: D69984379

faiss/impl/mapped_io.cpp

@@ -1,3 +1,10 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #include <stdio.h>
 #include <string.h>
 
@@ -11,8 +18,8 @@
 
 #elif defined(_WIN32)
 
-#include <Windows.h>
-#include <io.h>
+#include <Windows.h> // @manual
+#include <io.h>      // @manual
 
 #endif
 
@@ -278,4 +285,4 @@ int MappedFileIOReader::filedescriptor() {
     return -1;
 }
 
-} // namespace faiss
+} // namespace faiss

faiss/impl/mapped_io.h

@@ -1,3 +1,10 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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 <cstddef>
@@ -41,4 +48,4 @@ struct MappedFileIOReader : IOReader {
     int filedescriptor() override;
 };
 
-} // namespace faiss
+} // namespace faiss

faiss/impl/maybe_owned_vector.h

@@ -1,3 +1,10 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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 <cstddef>
@@ -50,6 +57,13 @@ struct MaybeOwnedVector {
         c_size = owned_data.size();
     }
 
+    explicit MaybeOwnedVector(const std::vector<T>& vec)
+            : faiss::MaybeOwnedVector<T>(vec.size()) {
+        if (vec.size() > 0) {
+            memcpy(owned_data.data(), vec.data(), sizeof(T) * vec.size());
+        }
+    }
+
     MaybeOwnedVector(const MaybeOwnedVector& other) {
         is_owned = other.is_owned;
         owned_data = other.owned_data;
@@ -229,4 +243,4 @@ struct is_maybe_owned_vector<MaybeOwnedVector<T>> : std::true_type {};
 template <typename T>
 inline constexpr bool is_maybe_owned_vector_v = is_maybe_owned_vector<T>::value;
 
-} // namespace faiss
+} // namespace faiss

faiss/impl/zerocopy_io.cpp

@@ -1,3 +1,10 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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/impl/zerocopy_io.h>
 #include <cstring>
 
@@ -37,6 +44,9 @@ void ZeroCopyIOReader::reset() {
 }
 
 size_t ZeroCopyIOReader::operator()(void* ptr, size_t size, size_t nitems) {
+    if (size * nitems == 0) {
+        return 0;
+    }
     if (rp_ >= total_) {
         return 0;
     }
@@ -53,4 +63,4 @@ int ZeroCopyIOReader::filedescriptor() {
     return -1; // Indicating no file descriptor available for memory buffer
 }
 
-} // namespace faiss
+} // namespace faiss

faiss/impl/zerocopy_io.h

@@ -1,3 +1,10 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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 <cstdint>
@@ -22,4 +29,4 @@ struct ZeroCopyIOReader : public faiss::IOReader {
     int filedescriptor() override;
 };
 
-} // namespace faiss
+} // namespace faiss

faiss/python/swigfaiss.swig

@@ -32,6 +32,7 @@
 #pragma SWIG nowarn=341
 #pragma SWIG nowarn=512
 #pragma SWIG nowarn=362
+#pragma SWIG nowarn=509
 
 // we need explict control of these typedefs...
 // %include <stdint.i>

tests/test_fast_scan_ivf.py

@@ -270,8 +270,9 @@ def test_equiv_pq(self):
         index_pq = faiss.index_factory(32, "PQ16x4np")
         index_pq.pq = index.pq
         index_pq.is_trained = True
-        index_pq.codes = faiss. downcast_InvertedLists(
+        codevec = faiss.downcast_InvertedLists(
             index.invlists).codes.at(0)
+        index_pq.codes = faiss.MaybeOwnedVectorUInt8(codevec)
         index_pq.ntotal = index.ntotal
         Dnew, Inew = index_pq.search(xq, 4)
 

tests/test_io.py

@@ -481,3 +481,40 @@ def test_reader(self):
         finally:
             if os.path.exists(fname):
                 os.unlink(fname)
+
+
+class TestIOFlatMMap(unittest.TestCase):
+
+    def test_mmap(self):
+        xt, xb, xq = get_dataset_2(32, 0, 100, 50)
+        index = faiss.index_factory(32, "SQfp16", faiss.METRIC_L2)
+        # does not need training
+        index.add(xb)
+        Dref, Iref = index.search(xq, 10)
+
+        fd, fname = tempfile.mkstemp()
+        os.close(fd)
+        try:
+            faiss.write_index(index, fname)
+            index2 = faiss.read_index(fname, faiss.IO_FLAG_MMAP_IFC)
+            Dnew, Inew = index2.search(xq, 10)
+            np.testing.assert_array_equal(Iref, Inew)
+            np.testing.assert_array_equal(Dref, Dnew)
+        finally:
+            if os.path.exists(fname):
+                os.unlink(fname)
+
+    def test_zerocopy(self):
+        xt, xb, xq = get_dataset_2(32, 0, 100, 50)
+        index = faiss.index_factory(32, "SQfp16", faiss.METRIC_L2)
+        # does not need training
+        index.add(xb)
+        Dref, Iref = index.search(xq, 10)
+
+        serialized_index = faiss.serialize_index(index)
+        reader = faiss.ZeroCopyIOReader(
+            faiss.swig_ptr(serialized_index), serialized_index.size)
+        index2 = faiss.read_index(reader)
+        Dnew, Inew = index2.search(xq, 10)
+        np.testing.assert_array_equal(Iref, Inew)
+        np.testing.assert_array_equal(Dref, Dnew)