diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 683969b36..f35b719d4 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -1275,10 +1275,13 @@ def handle_MapLong2Long(the_class): def replacement_map_add(self, keys, vals): n, = keys.shape assert (n,) == vals.shape + keys = np.ascontiguousarray(keys, dtype='int64') + vals = np.ascontiguousarray(vals, dtype='int64') self.add_c(n, swig_ptr(keys), swig_ptr(vals)) def replacement_map_search_multiple(self, keys): n, = keys.shape + keys = np.ascontiguousarray(keys, dtype='int64') vals = np.empty(n, dtype='int64') self.search_multiple_c(n, swig_ptr(keys), swig_ptr(vals)) return vals diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 12be4c343..71f5665f8 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -947,6 +947,7 @@ typedef int cudaDataType_t; } %typemap(out) faiss::InvertedLists * { + DOWNCAST (ArrayInvertedListsPanorama) // subclass before parent DOWNCAST (ArrayInvertedLists) DOWNCAST (BlockInvertedLists) #ifndef SWIGWIN @@ -954,6 +955,7 @@ typedef int cudaDataType_t; #endif // !SWIGWIN DOWNCAST (VStackInvertedLists) DOWNCAST (HStackInvertedLists) + DOWNCAST (SliceInvertedLists) DOWNCAST (MaskedInvertedLists) DOWNCAST (StopWordsInvertedLists) DOWNCAST (CappedInvertedLists) diff --git a/tests/test_swig_wrapper.py b/tests/test_swig_wrapper.py index 7e147f121..ed451f15e 100644 --- a/tests/test_swig_wrapper.py +++ b/tests/test_swig_wrapper.py @@ -286,3 +286,49 @@ class TestDoxygen(unittest.TestCase): 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))