mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-29 06:55:27 +08:00
66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
from ..scheduler import *
|
|
|
|
import unittest
|
|
import faiss
|
|
import numpy as np
|
|
|
|
|
|
class TestScheduler(unittest.TestCase):
|
|
def test_schedule(self):
|
|
d = 64
|
|
nb = 10000
|
|
nq = 2
|
|
nt = 5000
|
|
xt, xb, xq = get_dataset(d, nb, nt, nq)
|
|
file_name = "/tmp/tempfile_1"
|
|
|
|
index = faiss.IndexFlatL2(d)
|
|
print(index.is_trained)
|
|
index.add(xb)
|
|
faiss.write_index(index, file_name)
|
|
Dref, Iref = index.search(xq, 5)
|
|
|
|
index2 = faiss.read_index(file_name)
|
|
|
|
scheduler_instance = Scheduler()
|
|
|
|
# query args 1
|
|
query_index = dict()
|
|
query_index['index'] = [file_name]
|
|
vectors = scheduler_instance.Search(query_index, vectors=xq, k=5)
|
|
assert np.all(vectors == Iref)
|
|
|
|
# query args 2
|
|
query_index = dict()
|
|
query_index['raw'] = xt
|
|
# Xiaojun TODO: 'raw_id' part
|
|
# query_index['raw_id'] =
|
|
query_index['dimension'] = d
|
|
query_index['index'] = [file_name]
|
|
|
|
# Xiaojun TODO: once 'raw_id' part added, open below
|
|
# vectors = scheduler_instance.Search(query_index, vectors=xq, k=5)
|
|
|
|
# print("success")
|
|
|
|
|
|
def get_dataset(d, nb, nt, nq):
|
|
"""A dataset that is not completely random but still challenging to
|
|
index
|
|
"""
|
|
d1 = 10 # intrinsic dimension (more or less)
|
|
n = nb + nt + nq
|
|
rs = np.random.RandomState(1338)
|
|
x = rs.normal(size=(n, d1))
|
|
x = np.dot(x, rs.rand(d1, d))
|
|
# now we have a d1-dim ellipsoid in d-dimensional space
|
|
# higher factor (>4) -> higher frequency -> less linear
|
|
x = x * (rs.rand(d) * 4 + 0.1)
|
|
x = np.sin(x)
|
|
x = x.astype('float32')
|
|
return x[:nt], x[nt:-nq], x[-nq:]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|