# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import os import sys import numpy as np import faiss import unittest from common_faiss_tests import Randu10k ru = Randu10k() xb = ru.xb xt = ru.xt xq = ru.xq nb, d = xb.shape nq, d = xq.shape class IDRemap(unittest.TestCase): def test_id_remap_idmap(self): # reference: index without remapping index = faiss.IndexPQ(d, 8, 8) k = 10 index.train(xt) index.add(xb) _Dref, Iref = index.search(xq, k) # try a remapping ids = np.arange(nb)[::-1].copy().astype('int64') sub_index = faiss.IndexPQ(d, 8, 8) index2 = faiss.IndexIDMap(sub_index) index2.train(xt) index2.add_with_ids(xb, ids) _D, I = index2.search(xq, k) assert np.all(I == nb - 1 - Iref) def test_id_remap_ivf(self): # coarse quantizer in common coarse_quantizer = faiss.IndexFlatIP(d) ncentroids = 25 # reference: index without remapping index = faiss.IndexIVFPQ(coarse_quantizer, d, ncentroids, 8, 8) index.nprobe = 5 k = 10 index.train(xt) index.add(xb) _Dref, Iref = index.search(xq, k) # try a remapping ids = np.arange(nb)[::-1].copy().astype('int64') index2 = faiss.IndexIVFPQ(coarse_quantizer, d, ncentroids, 8, 8) index2.nprobe = 5 index2.train(xt) index2.add_with_ids(xb, ids) _D, I = index2.search(xq, k) assert np.all(I == nb - 1 - Iref) class Shards(unittest.TestCase): @unittest.skipIf(os.name == "posix" and os.uname().sysname == "Darwin", "There is a bug in the OpenMP implementation on OSX.") def test_shards(self): k = 32 ref_index = faiss.IndexFlatL2(d) print('ref search') ref_index.add(xb) _Dref, Iref = ref_index.search(xq, k) print(Iref[:5, :6]) shard_index = faiss.IndexShards(d) shard_index_2 = faiss.IndexShards(d, True, False) ni = 3 for i in range(ni): i0 = int(i * nb / ni) i1 = int((i + 1) * nb / ni) index = faiss.IndexFlatL2(d) index.add(xb[i0:i1]) shard_index.add_shard(index) index_2 = faiss.IndexFlatL2(d) irm = faiss.IndexIDMap(index_2) shard_index_2.add_shard(irm) # test parallel add shard_index_2.verbose = True shard_index_2.add(xb) for test_no in range(3): with_threads = test_no == 1 print('shard search test_no = %d' % test_no) if with_threads: remember_nt = faiss.omp_get_max_threads() faiss.omp_set_num_threads(1) shard_index.threaded = True else: shard_index.threaded = False if test_no != 2: _D, I = shard_index.search(xq, k) else: _D, I = shard_index_2.search(xq, k) print(I[:5, :6]) if with_threads: faiss.omp_set_num_threads(remember_nt) ndiff = (I != Iref).sum() print('%d / %d differences' % (ndiff, nq * k)) assert(ndiff < nq * k / 1000.) class Merge(unittest.TestCase): def make_index_for_merge(self, quant, index_type, master_index): ncent = 40 if index_type == 1: index = faiss.IndexIVFFlat(quant, d, ncent, faiss.METRIC_L2) if master_index: index.is_trained = True elif index_type == 2: index = faiss.IndexIVFPQ(quant, d, ncent, 4, 8) if master_index: index.pq = master_index.pq index.is_trained = True elif index_type == 3: index = faiss.IndexIVFPQR(quant, d, ncent, 4, 8, 8, 8) if master_index: index.pq = master_index.pq index.refine_pq = master_index.refine_pq index.is_trained = True elif index_type == 4: # quant used as the actual index index = faiss.IndexIDMap(quant) return index def do_test_merge(self, index_type): k = 16 quant = faiss.IndexFlatL2(d) ref_index = self.make_index_for_merge(quant, index_type, False) # trains the quantizer ref_index.train(xt) print('ref search') ref_index.add(xb) _Dref, Iref = ref_index.search(xq, k) print(Iref[:5, :6]) indexes = [] ni = 3 for i in range(ni): i0 = int(i * nb / ni) i1 = int((i + 1) * nb / ni) index = self.make_index_for_merge(quant, index_type, ref_index) index.is_trained = True index.add(xb[i0:i1]) indexes.append(index) index = indexes[0] for i in range(1, ni): print('merge ntotal=%d other.ntotal=%d ' % ( index.ntotal, indexes[i].ntotal)) index.merge_from(indexes[i], index.ntotal) _D, I = index.search(xq, k) print(I[:5, :6]) ndiff = (I != Iref).sum() print('%d / %d differences' % (ndiff, nq * k)) assert(ndiff < nq * k / 1000.) def test_merge(self): self.do_test_merge(1) self.do_test_merge(2) self.do_test_merge(3) def do_test_remove(self, index_type): k = 16 quant = faiss.IndexFlatL2(d) index = self.make_index_for_merge(quant, index_type, None) # trains the quantizer index.train(xt) if index_type < 4: index.add(xb) else: gen = np.random.RandomState(1234) id_list = gen.permutation(nb * 7)[:nb].astype('int64') index.add_with_ids(xb, id_list) print('ref search ntotal=%d' % index.ntotal) Dref, Iref = index.search(xq, k) toremove = np.zeros(nq * k, dtype='int64') nr = 0 for i in range(nq): for j in range(k): # remove all even results (it's ok if there are duplicates # in the list of ids) if Iref[i, j] % 2 == 0: nr = nr + 1 toremove[nr] = Iref[i, j] print('nr=', nr) idsel = faiss.IDSelectorBatch( nr, faiss.swig_ptr(toremove)) for i in range(nr): assert(idsel.is_member(int(toremove[i]))) nremoved = index.remove_ids(idsel) print('nremoved=%d ntotal=%d' % (nremoved, index.ntotal)) D, I = index.search(xq, k) # make sure results are in the same order with even ones removed ndiff = 0 for i in range(nq): j2 = 0 for j in range(k): if Iref[i, j] % 2 != 0: if I[i, j2] != Iref[i, j]: ndiff += 1 assert abs(D[i, j2] - Dref[i, j]) < 1e-5 j2 += 1 # draws are ordered arbitrarily assert ndiff < 5 def test_remove(self): self.do_test_remove(1) self.do_test_remove(2) self.do_test_remove(4) if __name__ == '__main__': unittest.main()