# 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. # a few tests of the swig wrapper import unittest import faiss import numpy as np class TestSWIGWrap(unittest.TestCase): """various regressions with the SWIG wrapper""" def test_size_t_ptr(self): # issue 1064 index = faiss.IndexHNSWFlat(10, 32) hnsw = index.hnsw index.add(np.random.rand(100, 10).astype("float32")) be = np.empty(2, "uint64") hnsw.neighbor_range(23, 0, faiss.swig_ptr(be), faiss.swig_ptr(be[1:])) def test_id_map_at(self): # issue 1020 n_features = 100 feature_dims = 10 features = np.random.random((n_features, feature_dims)).astype( np.float32 ) idx = np.arange(n_features).astype(np.int64) index = faiss.IndexFlatL2(feature_dims) index = faiss.IndexIDMap2(index) index.add_with_ids(features, idx) [index.id_map.at(int(i)) for i in range(index.ntotal)] def test_downcast_Refine(self): index = faiss.IndexRefineFlat( faiss.IndexScalarQuantizer(10, faiss.ScalarQuantizer.QT_8bit) ) # serialize and deserialize index2 = faiss.deserialize_index(faiss.serialize_index(index)) assert isinstance(index2, faiss.IndexRefineFlat) # Verify deserialized index is serializable again index3 = faiss.deserialize_index(faiss.serialize_index(index2)) assert isinstance(index3, faiss.IndexRefineFlat) def do_test_array_type(self, dtype): """tests swig_ptr and rev_swig_ptr for this type of array""" a = np.arange(12).astype(dtype) ptr = faiss.swig_ptr(a) a2 = faiss.rev_swig_ptr(ptr, 12) np.testing.assert_array_equal(a, a2) def test_all_array_types(self): self.do_test_array_type("float32") self.do_test_array_type("float64") self.do_test_array_type("int8") self.do_test_array_type("uint8") self.do_test_array_type("int16") self.do_test_array_type("uint16") self.do_test_array_type("int32") self.do_test_array_type("uint32") self.do_test_array_type("int64") self.do_test_array_type("uint64") def test_int64(self): # see https://github.com/facebookresearch/faiss/issues/1529 v = faiss.Int64Vector() for i in range(10): v.push_back(i) a = faiss.vector_to_array(v) assert a.dtype == "int64" np.testing.assert_array_equal(a, np.arange(10, dtype="int64")) # check if it works in an IDMap idx = faiss.IndexIDMap(faiss.IndexFlatL2(32)) idx.add_with_ids( np.random.rand(10, 32).astype("float32"), np.random.randint(1000, size=10, dtype="int64"), ) faiss.vector_to_array(idx.id_map) def test_asan(self): # this test should fail with ASAN index = faiss.IndexFlatL2(32) index.this.own(False) # this is a mem leak, should be catched by ASAN def test_SWIG_version(self): self.assertLess(faiss.swig_version(), 0x050000) def test_attribute_validation(self): """Test that setting invalid attributes raises AttributeError""" # Test IndexPreTransform - the main use case from the issue index = faiss.index_factory(256, "OPQ64,IVF16384,PQ64") # Should raise AttributeError when trying to set nprobe directly on # wrapper with self.assertRaises(AttributeError) as cm: index.nprobe = 16 # Check that the error message contains basic information error_msg = str(cm.exception) self.assertIn("IndexPreTransform", error_msg) self.assertIn("nprobe", error_msg) # Test with other IVF parameters that should be blocked with self.assertRaises(AttributeError) as cm: index.nlist = 8192 # Valid attributes should still work index.verbose = True # This should be allowed self.assertEqual(index.verbose, True) def test_attribute_validation_other_indexes(self): """Test that other index types allow normal attribute setting""" # Test with a regular IndexFlat - should allow most attributes now index = faiss.IndexFlatL2(10) # Valid attributes should work index.verbose = False self.assertEqual(index.verbose, False) def test_rabitq(self): """Test that other index types allow normal attribute setting""" # Test with a regular IndexFlat - should allow most attributes now index = faiss.IndexRaBitQ(10) # Valid attributes should work index.qb = 4 index.centered = True with self.assertRaises(AttributeError): index.centered2 = False def test_ivfsq_turboq_search_parameters(self): params = faiss.IVFSQTurboQSearchParameters() self.assertEqual(params.qb, 0) self.assertFalse(params.int_qjl) params.qb = 4 params.int_qjl = True self.assertEqual(params.qb, 4) self.assertTrue(params.int_qjl) def test_scalar_quantizer_tq_enum_values(self): # QT_2bit_tq..QT_5bit_tq were added in the TurboQuant commit and must # be accessible as ScalarQuantizer class attributes (pyi stub regression # test: they were absent until this fix). sq = faiss.ScalarQuantizer self.assertEqual(sq.QT_2bit_tq, 15) self.assertEqual(sq.QT_3bit_tq, 16) self.assertEqual(sq.QT_4bit_tq, 17) self.assertEqual(sq.QT_5bit_tq, 18) def test_scalar_quantizer_0bit_enum_value(self): # QT_0bit (centroid-only SQ, for IVF) was likewise absent from the # pyi stub until this fix. self.assertEqual(faiss.ScalarQuantizer.QT_0bit, 9) class TestRevSwigPtr(unittest.TestCase): def test_rev_swig_ptr(self): index = faiss.IndexFlatL2(4) xb0 = np.vstack( [i * 10 + np.array([1, 2, 3, 4], dtype="float32") for i in range(5)] ) index.add(xb0) xb = faiss.rev_swig_ptr(index.get_xb(), 4 * 5).reshape(5, 4) self.assertEqual(np.abs(xb0 - xb).sum(), 0) class TestException(unittest.TestCase): def test_exception(self): index = faiss.IndexFlatL2(10) a = np.zeros((5, 10), dtype="float32") b = np.zeros(5, dtype="int64") # an unsupported operation for IndexFlat self.assertRaises(RuntimeError, index.add_with_ids, a, b) # assert 'add_with_ids not implemented' in str(e) def test_exception_2(self): self.assertRaises( RuntimeError, faiss.index_factory, 12, "IVF256,Flat,PQ8" ) # assert 'could not parse' in str(e) class TestAddSACodes(unittest.TestCase): """Regression tests for add_sa_codes and sa_encode wrapper correctness.""" D = 32 NB = 200 def setUp(self): rng = np.random.default_rng(42) self.xb = rng.random((self.NB, self.D), dtype=np.float32) self.xtrain = rng.random((500, self.D), dtype=np.float32) # IDMap2 wrapping PQ8x2: outer index stores (id → code) pairs; # inner PQ8x2 encodes vectors via sa_encode / add_sa_codes. self.index = faiss.index_factory(self.D, "IDMap2,PQ8x2") self.index.train(self.xtrain) def test_add_sa_codes_int32_ids_stored_correctly(self): """int32 ids must be coerced to int64 before the C++ boundary. Without coercion swig_ptr hands the C++ side a pointer to 4-byte int32 data; C++ interprets every 8 bytes as one idx_t, producing garbage ids with no exception.""" codes = self.index.index.sa_encode(self.xb) ids_int64 = np.arange(self.NB, dtype="int64") + 1000 ids_int32 = ids_int64.astype("int32") self.index.add_sa_codes(codes, ids_int32) stored = faiss.vector_to_array(self.index.id_map) np.testing.assert_array_equal(stored, ids_int64) def test_add_sa_codes_int64_ids_unchanged(self): """int64 ids must pass through as-is (np.ascontiguousarray is a no-op when the array is already int64 and contiguous).""" codes = self.index.index.sa_encode(self.xb) ids = np.arange(self.NB, dtype="int64") + 500 self.index.add_sa_codes(codes, ids) stored = faiss.vector_to_array(self.index.id_map) np.testing.assert_array_equal(stored, ids) def test_sa_encode_preallocated_matches_auto_allocated(self): """sa_encode with a pre-allocated buffer must return the same codes as the auto-allocated path, exercising both branches of the hoisted sa_code_size() call.""" inner = self.index.index codes_auto = inner.sa_encode(self.xb) buf = np.empty_like(codes_auto) codes_pre = inner.sa_encode(self.xb, buf) np.testing.assert_array_equal(codes_auto, codes_pre) self.assertIs(codes_pre, buf) class TestSearchPreassignedDqNone(unittest.TestCase): """Dq=None must behave like Dq=zeros for search_preassigned and range_search_preassigned (regression: swig_ptr(None) raised ValueError).""" def _check(self, index, xq, Iq, k, thresh, zero_dtype): Dq_zero = np.zeros(Iq.shape, dtype=zero_dtype) _, I_zero = index.search_preassigned(xq, k, Iq, Dq_zero) _, I_none = index.search_preassigned(xq, k, Iq, None) np.testing.assert_array_equal(I_zero, I_none) lz, _, Iz = index.range_search_preassigned(xq, thresh, Iq, Dq_zero) ln, _, In = index.range_search_preassigned(xq, thresh, Iq, None) np.testing.assert_array_equal(lz, ln) np.testing.assert_array_equal(Iz, In) def test_float_ivf(self): rng = np.random.default_rng(42) xb = rng.random((200, 16), dtype=np.float32) xq = rng.random((10, 16), dtype=np.float32) index = faiss.index_factory(16, "IVF4,Flat") index.train(xb) index.add(xb) index.nprobe = 2 _, Iq = index.quantizer.search(xq, 2) self._check(index, xq, Iq, k=5, thresh=2.0, zero_dtype=np.float32) def test_binary_ivf(self): rng = np.random.default_rng(42) xb = rng.integers(0, 256, size=(200, 8), dtype=np.uint8) xq = rng.integers(0, 256, size=(10, 8), dtype=np.uint8) index = faiss.IndexBinaryIVF(faiss.IndexBinaryFlat(64), 64, 4) index.train(xb) index.add(xb) index.nprobe = 2 _, Iq = index.quantizer.search(xq, 2) self._check(index, xq, Iq, k=5, thresh=20, zero_dtype=np.int32) @unittest.skipIf( faiss.swig_version() < 0x040000, "swig < 4 does not support Doxygen comments", ) class TestDoxygen(unittest.TestCase): def test_doxygen_comments(self): maxheap_array = faiss.float_maxheap_array_t() self.assertTrue( "a template structure for a set of [min|max]-heaps" in maxheap_array.__doc__ ) class TestMapLong2Long(unittest.TestCase): def test_add_coerces_keys_and_vals(self): for dtype in (np.int32, np.int64): with self.subTest(dtype=dtype): m = faiss.MapLong2Long() keys = np.array([0, 1, 2, 3, 4], dtype=dtype) vals = np.array([100, 101, 102, 103, 104], dtype=dtype) m.add(keys, vals) for k, v in zip(range(5), range(100, 105)): self.assertEqual(m.search(k), v) def test_search_multiple_coerces_int32_keys(self): m = faiss.MapLong2Long() m.add( np.arange(5, dtype=np.int64), np.arange(5, dtype=np.int64) + 300, ) result = m.search_multiple(np.array([0, 2, 4], dtype=np.int32)) np.testing.assert_array_equal(result, [300, 302, 304]) class TestInvertedListsDowncast(unittest.TestCase): def test_downcast_ArrayInvertedListsPanorama(self): # downcast_InvertedLists() triggers the typemap; raw .invlists access # does not (typemap applies to function returns, not member getters). d, nlist, n_levels = 32, 8, 2 index = faiss.IndexIVFFlatPanorama( faiss.IndexFlatL2(d), d, nlist, n_levels ) il = faiss.downcast_InvertedLists(index.invlists) self.assertIsInstance(il, faiss.ArrayInvertedListsPanorama) self.assertEqual(il.n_levels, n_levels) def test_downcast_SliceInvertedLists(self): d, nlist = 4, 4 # code_size = d * sizeof(float) backing = faiss.ArrayInvertedLists(nlist, d * 4) sil = faiss.SliceInvertedLists(backing, 0, nlist) index = faiss.IndexIVFFlat(faiss.IndexFlatL2(d), d, nlist) # own=False: index doesn't delete sil; Python locals keep both alive. index.replace_invlists(sil, False) inner = faiss.downcast_InvertedLists(index.invlists) self.assertIsInstance(inner, faiss.SliceInvertedLists) self.assertEqual((inner.i0, inner.i1), (0, nlist))