def do_read_callback(self, bsz):
d, n = 32, 1000
x = np.random.uniform(size=(n, d)).astype('float32')
index = faiss.IndexFlatL2(d)
index.add(x)
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
with open(fname, 'rb') as f:
reader = faiss.PyCallbackIOReader(f.read, 1234)
if bsz > 0:
reader = faiss.BufferedIOReader(reader, bsz)
index2 = faiss.read_index(reader)
self.assertEqual(index.d, index2.d)
np.testing.assert_array_equal(faiss.vector_to_array(index.xb),
faiss.vector_to_array(index2.xb))
# This is not a callable function: should raise an exception
reader = faiss.PyCallbackIOReader("blabla")
self.assertRaises(Exception, faiss.read_index, reader)
finally:
if os.path.exists(fname):
os.unlink(fname)
def test_buf_read(self):
x = np.random.uniform(size=20)
_, fname = tempfile.mkstemp()
try:
x.tofile(fname)
f = open(fname, 'rb')
reader = faiss.PyCallbackIOReader(f.read, 1234)
bsz = 123
reader = faiss.BufferedIOReader(reader, bsz)
y = np.zeros_like(x)
print('nbytes=', y.nbytes)
reader(faiss.swig_ptr(y), y.nbytes, 1)
np.testing.assert_array_equal(x, y)
finally:
if os.path.exists(fname):
os.unlink(fname)
def test_read_buffer(self):
d, n = 32, 1000
x = np.random.uniform(size=(n, d)).astype('float32')
index = faiss.IndexFlatL2(d)
index.add(x)
_, fname = tempfile.mkstemp()
try:
faiss.write_index(index, fname)
reader = faiss.BufferedIOReader(faiss.FileIOReader(fname), 1234)
index2 = faiss.read_index(reader)
self.assertEqual(index.d, index2.d)
np.testing.assert_array_equal(faiss.vector_to_array(index.xb),
faiss.vector_to_array(index2.xb))
finally:
if os.path.exists(fname):
os.unlink(fname)