def do_write_callback(self, bsz):
d, n = 32, 1000
x = np.random.uniform(size=(n, d)).astype('float32')
index = faiss.IndexFlatL2(d)
index.add(x)
f = io.BytesIO()
# test with small block size
writer = faiss.PyCallbackIOWriter(f.write, 1234)
if bsz > 0:
writer = faiss.BufferedIOWriter(writer, bsz)
faiss.write_index(index, writer)
del writer # make sure all writes committed
if sys.version_info[0] < 3:
buf = f.getvalue()
else:
buf = f.getbuffer()
index2 = faiss.deserialize_index(np.frombuffer(buf, dtype='uint8'))
self.assertEqual(index.d, index2.d)
self.assertTrue(
np.all(
faiss.vector_to_array(index.xb) == faiss.vector_to_array(
index2.xb)))
# This is not a callable function: shoudl raise an exception
writer = faiss.PyCallbackIOWriter("blabla")
self.assertRaises(Exception, faiss.write_index, index, writer)
def test_transfer_pipe(self):
""" transfer an index through a Unix pipe """
d, n = 32, 1000
x = np.random.uniform(size=(n, d)).astype('float32')
index = faiss.IndexFlatL2(d)
index.add(x)
Dref, Iref = index.search(x, 10)
rf, wf = os.pipe()
# start thread that will decompress the index
def index_from_pipe():
reader = faiss.PyCallbackIOReader(lambda size: os.read(rf, size))
return faiss.read_index(reader)
with ThreadPool(1) as pool:
fut = pool.apply_async(index_from_pipe, ())
# write to pipe
writer = faiss.PyCallbackIOWriter(lambda b: os.write(wf, b))
faiss.write_index(index, writer)
index2 = fut.get()
# closing is not really useful but it does not hurt
os.close(wf)
os.close(rf)
Dnew, Inew = index2.search(x, 10)
np.testing.assert_array_equal(Iref, Inew)
np.testing.assert_array_equal(Dref, Dnew)
def output_faiss_index(self):
print('faiss index ntotal [%d]: %d' %
(self.rank, self.faiss_index.ntotal))
item_index_stream = _mindalpha.OutputStream(
self.item_index_output_path)
item_index_writer = faiss.PyCallbackIOWriter(item_index_stream.write)
faiss.write_index(self.faiss_index, item_index_writer)
self.output_index_meta()
self.item_ids_stream = None
