clean up serialization tests (#4700)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4700

A pass over the serialization / deserialization tests, so that they don't use temp files when not needed.

Reviewed By: pankajsingh88, junjieqi

Differential Revision: D87615610

fbshipit-source-id: 3e6366b938d41e083781745179ff19ecf5247ba2
This commit is contained in:
Matthijs Douze
2025-12-09 05:16:38 -08:00
committed by meta-codesync[bot]
parent ac2e3abe38
commit 1ea99d8073
8 changed files with 50 additions and 179 deletions
+24 -79
View File
@@ -10,8 +10,6 @@ from __future__ import absolute_import, division, print_function, unicode_litera
import numpy as np
import unittest
import faiss
import os
import tempfile
def make_binary_dataset(d, nb, nt, nq):
@@ -37,21 +35,12 @@ class TestBinaryFlat(unittest.TestCase):
index = faiss.IndexBinaryFlat(d)
index.add(self.xb)
D, I = index.search(self.xq, 3)
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
D2, I2 = index2.search(self.xq, 3)
index2 = faiss.read_index_binary(tmpnam)
D2, I2 = index2.search(self.xq, 3)
assert (I2 == I).all()
assert (D2 == D).all()
finally:
os.remove(tmpnam)
assert (I2 == I).all()
assert (D2 == D).all()
class TestBinaryIVF(unittest.TestCase):
@@ -76,20 +65,12 @@ class TestBinaryIVF(unittest.TestCase):
index.add(self.xb)
D, I = index.search(self.xq, 3)
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
index2 = faiss.read_index_binary(tmpnam)
D2, I2 = index2.search(self.xq, 3)
D2, I2 = index2.search(self.xq, 3)
assert (I2 == I).all()
assert (D2 == D).all()
finally:
os.remove(tmpnam)
assert (I2 == I).all()
assert (D2 == D).all()
class TestObjectOwnership(unittest.TestCase):
@@ -109,16 +90,10 @@ class TestObjectOwnership(unittest.TestCase):
index = faiss.IndexBinaryFlat(d)
index.add(self.xb)
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
# this is the output of read_index_binary (==> checks ownership)
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
index2 = faiss.read_index_binary(tmpnam)
assert index2.thisown
finally:
os.remove(tmpnam)
assert index2.thisown
class TestBinaryFromFloat(unittest.TestCase):
@@ -140,21 +115,11 @@ class TestBinaryFromFloat(unittest.TestCase):
index.add(self.xb)
D, I = index.search(self.xq, 3)
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
index2 = faiss.read_index_binary(tmpnam)
D2, I2 = index2.search(self.xq, 3)
assert (I2 == I).all()
assert (D2 == D).all()
finally:
os.remove(tmpnam)
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
D2, I2 = index2.search(self.xq, 3)
assert (I2 == I).all()
assert (D2 == D).all()
class TestBinaryHNSW(unittest.TestCase):
@@ -174,20 +139,12 @@ class TestBinaryHNSW(unittest.TestCase):
index.add(self.xb)
D, I = index.search(self.xq, 3)
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
index2 = faiss.read_index_binary(tmpnam)
D2, I2 = index2.search(self.xq, 3)
D2, I2 = index2.search(self.xq, 3)
assert (I2 == I).all()
assert (D2 == D).all()
finally:
os.remove(tmpnam)
assert (I2 == I).all()
assert (D2 == D).all()
def test_ivf_hnsw(self):
d = self.xq.shape[1] * 8
@@ -200,21 +157,9 @@ class TestBinaryHNSW(unittest.TestCase):
index.add(self.xb)
D, I = index.search(self.xq, 3)
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
D2, I2 = index2.search(self.xq, 3)
index2 = faiss.read_index_binary(tmpnam)
D2, I2 = index2.search(self.xq, 3)
assert (I2 == I).all()
assert (D2 == D).all()
finally:
os.remove(tmpnam)
if __name__ == '__main__':
unittest.main()
assert (I2 == I).all()
assert (D2 == D).all()
+6 -22
View File
@@ -6,8 +6,6 @@
import unittest
import time
import os
import tempfile
import numpy as np
import faiss
@@ -587,16 +585,9 @@ class TestAQFastScan(unittest.TestCase):
index.add(ds.get_database())
D1, I1 = index.search(ds.get_queries(), 1)
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
index2 = faiss.read_index(fname)
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
finally:
if os.path.exists(fname):
os.unlink(fname)
index2 = faiss.deserialize_index(faiss.serialize_index(index))
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
def test_io(self):
self.subtest_io('LSQ4x4fs_Nlsq2x4')
@@ -685,16 +676,9 @@ class TestPAQFastScan(unittest.TestCase):
index.add(ds.get_database())
D1, I1 = index.search(ds.get_queries(), 1)
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
index2 = faiss.read_index(fname)
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
finally:
if os.path.exists(fname):
os.unlink(fname)
index2 = faiss.deserialize_index(faiss.serialize_index(index))
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
def test_io(self):
self.subtest_io('PLSQ2x3x4fs_Nlsq2x4')
+7 -23
View File
@@ -4,9 +4,7 @@
# LICENSE file in the root directory of this source tree.
import os
import unittest
import tempfile
import numpy as np
import faiss
@@ -821,18 +819,11 @@ class TestIVFAQFastScan(unittest.TestCase):
index = faiss.index_factory(d, factory_str)
index.train(ds.get_train())
index.add(ds.get_database())
D1, I1 = index.search(ds.get_queries(), 1)
_, I1 = index.search(ds.get_queries(), 1)
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
index2 = faiss.read_index(fname)
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
finally:
if os.path.exists(fname):
os.unlink(fname)
index2 = faiss.deserialize_index(faiss.serialize_index(index))
_, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
def test_io(self):
self.subtest_io('IVF16,LSQ4x4fs_Nlsq2x4')
@@ -929,16 +920,9 @@ class TestIVFPAQFastScan(unittest.TestCase):
index.add(ds.get_database())
D1, I1 = index.search(ds.get_queries(), 1)
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
index2 = faiss.read_index(fname)
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
finally:
if os.path.exists(fname):
os.unlink(fname)
index2 = faiss.deserialize_index(faiss.serialize_index(index))
D2, I2 = index2.search(ds.get_queries(), 1)
np.testing.assert_array_equal(I1, I2)
def test_io(self):
self.subtest_io('IVF16,PLSQ2x3x4fsr_Nlsq2x4')
-1
View File
@@ -15,7 +15,6 @@ Paper: https://www.arxiv.org/pdf/2510.00566
"""
import unittest
import tempfile
import os
import faiss
+2 -26
View File
@@ -8,8 +8,6 @@
import numpy as np
import unittest
import faiss
import tempfile
import os
from common_faiss_tests import get_dataset_2
@@ -280,14 +278,7 @@ class TestNSG(unittest.TestCase):
return knn_graph
def subtest_io_and_clone(self, index, Dnsg, Insg):
fd, tmpfile = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, tmpfile)
index2 = faiss.read_index(tmpfile)
finally:
if os.path.exists(tmpfile):
os.unlink(tmpfile)
index2 = faiss.deserialize_index(faiss.serialize_index(index))
Dnsg2, Insg2 = index2.search(self.xq, 1)
np.testing.assert_array_equal(Dnsg2, Dnsg)
@@ -306,8 +297,6 @@ class TestNSG(unittest.TestCase):
def subtest_add(self, build_type, thresh, metric=faiss.METRIC_L2):
d = self.xq.shape[1]
metrics = {faiss.METRIC_L2: 'L2',
faiss.METRIC_INNER_PRODUCT: 'IP'}
flat_index = faiss.IndexFlat(d, metric)
flat_index.add(self.xb)
@@ -381,8 +370,6 @@ class TestNSG(unittest.TestCase):
def test_reset(self):
"""test IndexNSG.reset()"""
d = self.xq.shape[1]
metrics = {faiss.METRIC_L2: 'L2',
faiss.METRIC_INNER_PRODUCT: 'IP'}
metric = faiss.METRIC_L2
flat_index = faiss.IndexFlat(d, metric)
@@ -546,14 +533,7 @@ class TestNNDescent(unittest.TestCase):
self.assertGreaterEqual(recalls, 450) # 462
# do some IO tests
fd, tmpfile = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, tmpfile)
index2 = faiss.read_index(tmpfile)
finally:
if os.path.exists(tmpfile):
os.unlink(tmpfile)
index2 = faiss.deserialize_index(faiss.serialize_index(index))
D2, I2 = index2.search(self.xq, 1)
np.testing.assert_array_equal(D2, D)
@@ -592,10 +572,6 @@ class TestNNDescentKNNG(unittest.TestCase):
self.subtest(32, 10, faiss.METRIC_INNER_PRODUCT)
def subtest(self, d, K, metric):
metric_names = {faiss.METRIC_L1: 'L1',
faiss.METRIC_L2: 'L2',
faiss.METRIC_INNER_PRODUCT: 'IP'}
nb = 1000
_, xb, _ = get_dataset_2(d, 0, nb, 0)
+2 -16
View File
@@ -4,7 +4,6 @@
# LICENSE file in the root directory of this source tree.
""" more elaborate that test_index.py """
from __future__ import absolute_import, division, print_function
import numpy as np
import unittest
@@ -193,13 +192,7 @@ class TestRemove(unittest.TestCase):
assert False, 'should have raised an exception'
# while we are there, let's test I/O as well...
fd, tmpnam = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index_binary(index, tmpnam)
index = faiss.read_index_binary(tmpnam)
finally:
os.remove(tmpnam)
index = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
assert index.reconstruct(1004)[0] == 104
try:
@@ -469,14 +462,7 @@ class TestIVFFlatDedup(unittest.TestCase):
check_ref_knn_with_draws(Dref, Iref, Dnew, Inew)
# test I/O
fd, tmpfile = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index_new, tmpfile)
index_st = faiss.read_index(tmpfile)
finally:
if os.path.exists(tmpfile):
os.unlink(tmpfile)
index_st = faiss.deserialize_index(faiss.serialize_index(index_new))
Dst, Ist = index_st.search(xq, 20)
check_ref_knn_with_draws(Dnew, Inew, Dst, Ist)
+9 -11
View File
@@ -359,21 +359,19 @@ class Test_IO_PQ(unittest.TestCase):
try:
faiss.write_ProductQuantizer(index.pq, fname)
read_pq = faiss.read_ProductQuantizer(fname)
self.assertEqual(index.pq.M, read_pq.M)
self.assertEqual(index.pq.nbits, read_pq.nbits)
self.assertEqual(index.pq.dsub, read_pq.dsub)
self.assertEqual(index.pq.ksub, read_pq.ksub)
np.testing.assert_array_equal(
faiss.vector_to_array(index.pq.centroids),
faiss.vector_to_array(read_pq.centroids)
)
finally:
if os.path.exists(fname):
os.unlink(fname)
self.assertEqual(index.pq.M, read_pq.M)
self.assertEqual(index.pq.nbits, read_pq.nbits)
self.assertEqual(index.pq.dsub, read_pq.dsub)
self.assertEqual(index.pq.ksub, read_pq.ksub)
np.testing.assert_array_equal(
faiss.vector_to_array(index.pq.centroids),
faiss.vector_to_array(read_pq.centroids)
)
class Test_IO_IndexLSH(unittest.TestCase):
-1
View File
@@ -15,7 +15,6 @@ Paper: https://www.arxiv.org/pdf/2510.00566
"""
import unittest
import tempfile
import os
import faiss