mirror of https://github.com/milvus-io/milvus.git
50 lines
1.2 KiB
Python
50 lines
1.2 KiB
Python
from ..search_index import *
|
|
|
|
import unittest
|
|
import numpy as np
|
|
|
|
|
|
class TestSearchSingleThread(unittest.TestCase):
|
|
def test_search_by_vectors(self):
|
|
d = 64
|
|
nb = 10000
|
|
nq = 100
|
|
_, xb, xq = get_dataset(d, nb, 500, nq)
|
|
|
|
index = faiss.IndexFlatL2(d)
|
|
index.add(xb)
|
|
|
|
# expect result
|
|
Dref, Iref = index.search(xq, 5)
|
|
|
|
searcher = FaissSearch(index)
|
|
result = searcher.search_by_vectors(xq, 5)
|
|
|
|
assert np.all(result.distance == Dref) \
|
|
and np.all(result.vectors == Iref)
|
|
pass
|
|
|
|
def test_top_k(selfs):
|
|
pass
|
|
|
|
|
|
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()
|