diff --git a/faiss/impl/mapped_io.cpp b/faiss/impl/mapped_io.cpp index 2e23a8a47e..a62e22dee1 100644 --- a/faiss/impl/mapped_io.cpp +++ b/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 #include @@ -11,8 +18,8 @@ #elif defined(_WIN32) -#include -#include +#include // @manual +#include // @manual #endif @@ -278,4 +285,4 @@ int MappedFileIOReader::filedescriptor() { return -1; } -} // namespace faiss \ No newline at end of file +} // namespace faiss diff --git a/faiss/impl/mapped_io.h b/faiss/impl/mapped_io.h index efbde1e2e7..0e32df23d8 100644 --- a/faiss/impl/mapped_io.h +++ b/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 @@ -41,4 +48,4 @@ struct MappedFileIOReader : IOReader { int filedescriptor() override; }; -} // namespace faiss \ No newline at end of file +} // namespace faiss diff --git a/faiss/impl/maybe_owned_vector.h b/faiss/impl/maybe_owned_vector.h index bca98d19c2..2369c4ab4e 100644 --- a/faiss/impl/maybe_owned_vector.h +++ b/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 @@ -50,6 +57,13 @@ struct MaybeOwnedVector { c_size = owned_data.size(); } + explicit MaybeOwnedVector(const std::vector& vec) + : faiss::MaybeOwnedVector(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> : std::true_type {}; template inline constexpr bool is_maybe_owned_vector_v = is_maybe_owned_vector::value; -} // namespace faiss \ No newline at end of file +} // namespace faiss diff --git a/faiss/impl/zerocopy_io.cpp b/faiss/impl/zerocopy_io.cpp index 31214ce3b2..c754f1f07b 100644 --- a/faiss/impl/zerocopy_io.cpp +++ b/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 #include @@ -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 \ No newline at end of file +} // namespace faiss diff --git a/faiss/impl/zerocopy_io.h b/faiss/impl/zerocopy_io.h index 7033b5fc51..488b5d1e80 100644 --- a/faiss/impl/zerocopy_io.h +++ b/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 @@ -22,4 +29,4 @@ struct ZeroCopyIOReader : public faiss::IOReader { int filedescriptor() override; }; -} // namespace faiss \ No newline at end of file +} // namespace faiss diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index fc0e7c0df4..9cf81523d8 100644 --- a/faiss/python/swigfaiss.swig +++ b/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 diff --git a/tests/test_fast_scan_ivf.py b/tests/test_fast_scan_ivf.py index 75c9500f82..a1d6a21440 100644 --- a/tests/test_fast_scan_ivf.py +++ b/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) diff --git a/tests/test_io.py b/tests/test_io.py index 3cbd0a6e10..49dc3bf489 100644 --- a/tests/test_io.py +++ b/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)