1
|
|
|
""" |
2
|
|
|
For testing index operations, including `create_index`, `get_index_info` and `drop_index` interfaces |
3
|
|
|
""" |
4
|
|
|
import logging |
5
|
|
|
import pytest |
6
|
|
|
import time |
7
|
|
|
import pdb |
8
|
|
|
import threading |
9
|
|
|
from multiprocessing import Pool, Process |
10
|
|
|
import numpy |
11
|
|
|
import sklearn.preprocessing |
12
|
|
|
from milvus import IndexType, MetricType |
13
|
|
|
from utils import * |
14
|
|
|
|
15
|
|
|
nb = 6000 |
16
|
|
|
dim = 128 |
17
|
|
|
index_file_size = 10 |
18
|
|
|
vectors = gen_vectors(nb, dim) |
19
|
|
|
vectors = sklearn.preprocessing.normalize(vectors, axis=1, norm='l2') |
20
|
|
|
vectors = vectors.tolist() |
21
|
|
|
BUILD_TIMEOUT = 300 |
22
|
|
|
nprobe = 1 |
23
|
|
|
tag = "1970-01-01" |
24
|
|
|
NLIST = 4046 |
25
|
|
|
INVALID_NLIST = 100000000 |
26
|
|
|
|
27
|
|
|
|
28
|
|
|
class TestIndexBase: |
29
|
|
|
@pytest.fixture( |
30
|
|
|
scope="function", |
31
|
|
|
params=gen_index() |
32
|
|
|
) |
33
|
|
|
def get_index(self, request, connect): |
34
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
35
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
36
|
|
|
pytest.skip("sq8h not support in CPU mode") |
37
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
38
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
39
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
40
|
|
|
return request.param |
41
|
|
|
|
42
|
|
View Code Duplication |
@pytest.fixture( |
|
|
|
|
43
|
|
|
scope="function", |
44
|
|
|
params=gen_simple_index() |
45
|
|
|
) |
46
|
|
|
def get_simple_index(self, request, connect): |
47
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
48
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
49
|
|
|
pytest.skip("sq8h not support in CPU mode") |
50
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
51
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
52
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
53
|
|
|
return request.param |
54
|
|
|
|
55
|
|
|
""" |
56
|
|
|
****************************************************************** |
57
|
|
|
The following cases are used to test `create_index` function |
58
|
|
|
****************************************************************** |
59
|
|
|
""" |
60
|
|
|
|
61
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
62
|
|
|
def test_create_index(self, connect, collection, get_simple_index): |
63
|
|
|
''' |
64
|
|
|
target: test create index interface |
65
|
|
|
method: create collection and add vectors in it, create index |
66
|
|
|
expected: return code equals to 0, and search success |
67
|
|
|
''' |
68
|
|
|
index_param = get_simple_index["index_param"] |
69
|
|
|
index_type = get_simple_index["index_type"] |
70
|
|
|
logging.getLogger().info(get_simple_index) |
71
|
|
|
status, ids = connect.insert(collection, vectors) |
72
|
|
|
status = connect.create_index(collection, index_type, index_param) |
73
|
|
|
assert status.OK() |
74
|
|
|
|
75
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
76
|
|
|
def test_create_index_no_vectors(self, connect, collection, get_simple_index): |
77
|
|
|
''' |
78
|
|
|
target: test create index interface |
79
|
|
|
method: create collection and add vectors in it, create index |
80
|
|
|
expected: return code equals to 0, and search success |
81
|
|
|
''' |
82
|
|
|
index_param = get_simple_index["index_param"] |
83
|
|
|
index_type = get_simple_index["index_type"] |
84
|
|
|
logging.getLogger().info(get_simple_index) |
85
|
|
|
status = connect.create_index(collection, index_type, index_param) |
86
|
|
|
assert status.OK() |
87
|
|
|
|
88
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
89
|
|
|
def test_create_index_partition(self, connect, collection, get_simple_index): |
90
|
|
|
''' |
91
|
|
|
target: test create index interface |
92
|
|
|
method: create collection, create partition, and add vectors in it, create index |
93
|
|
|
expected: return code equals to 0, and search success |
94
|
|
|
''' |
95
|
|
|
index_param = get_simple_index["index_param"] |
96
|
|
|
index_type = get_simple_index["index_type"] |
97
|
|
|
logging.getLogger().info(get_simple_index) |
98
|
|
|
status = connect.create_partition(collection, tag) |
99
|
|
|
status, ids = connect.insert(collection, vectors, partition_tag=tag) |
100
|
|
|
status = connect.create_index(collection, index_type, index_param) |
101
|
|
|
assert status.OK() |
102
|
|
|
|
103
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
104
|
|
|
def test_create_index_partition_flush(self, connect, collection, get_simple_index): |
105
|
|
|
''' |
106
|
|
|
target: test create index interface |
107
|
|
|
method: create collection, create partition, and add vectors in it, create index |
108
|
|
|
expected: return code equals to 0, and search success |
109
|
|
|
''' |
110
|
|
|
index_param = get_simple_index["index_param"] |
111
|
|
|
index_type = get_simple_index["index_type"] |
112
|
|
|
logging.getLogger().info(get_simple_index) |
113
|
|
|
status = connect.create_partition(collection, tag) |
114
|
|
|
status, ids = connect.insert(collection, vectors, partition_tag=tag) |
115
|
|
|
connect.flush() |
116
|
|
|
status = connect.create_index(collection, index_type, index_param) |
117
|
|
|
assert status.OK() |
118
|
|
|
|
119
|
|
|
# @pytest.mark.level(2) |
120
|
|
|
# def test_create_index_without_connect(self, dis_connect, collection): |
121
|
|
|
# ''' |
122
|
|
|
# target: test create index without connection |
123
|
|
|
# method: create collection and add vectors in it, check if added successfully |
124
|
|
|
# expected: raise exception |
125
|
|
|
# ''' |
126
|
|
|
# nlist = NLIST |
127
|
|
|
# index_type = IndexType.IVF_SQ8 |
128
|
|
|
# index_param = {"nlist": nlist} |
129
|
|
|
# with pytest.raises(Exception) as e: |
130
|
|
|
# status = dis_connect.create_index(collection, index_type, index_param) |
131
|
|
|
|
132
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
133
|
|
|
def test_create_index_search_with_query_vectors(self, connect, collection, get_simple_index): |
134
|
|
|
''' |
135
|
|
|
target: test create index interface, search with more query vectors |
136
|
|
|
method: create collection and add vectors in it, create index |
137
|
|
|
expected: return code equals to 0, and search success |
138
|
|
|
''' |
139
|
|
|
index_param = get_simple_index["index_param"] |
140
|
|
|
index_type = get_simple_index["index_type"] |
141
|
|
|
logging.getLogger().info(get_simple_index) |
142
|
|
|
status, ids = connect.insert(collection, vectors) |
143
|
|
|
status = connect.create_index(collection, index_type, index_param) |
144
|
|
|
logging.getLogger().info(connect.get_index_info(collection)) |
145
|
|
|
query_vecs = [vectors[0], vectors[1], vectors[2]] |
146
|
|
|
top_k = 5 |
147
|
|
|
search_param = get_search_param(index_type) |
148
|
|
|
status, result = connect.search(collection, top_k, query_vecs, params=search_param) |
149
|
|
|
assert status.OK() |
150
|
|
|
assert len(result) == len(query_vecs) |
151
|
|
|
logging.getLogger().info(result) |
152
|
|
|
|
153
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
154
|
|
|
@pytest.mark.level(2) |
155
|
|
|
def test_create_index_multithread(self, connect, collection, args): |
156
|
|
|
''' |
157
|
|
|
target: test create index interface with multiprocess |
158
|
|
|
method: create collection and add vectors in it, create index |
159
|
|
|
expected: return code equals to 0, and search success |
160
|
|
|
''' |
161
|
|
|
status, ids = connect.insert(collection, vectors) |
162
|
|
|
|
163
|
|
|
def build(connect): |
164
|
|
|
status = connect.create_index(collection, IndexType.IVFLAT, {"nlist": NLIST}) |
165
|
|
|
assert status.OK() |
166
|
|
|
|
167
|
|
|
threads_num = 8 |
168
|
|
|
threads = [] |
169
|
|
|
for i in range(threads_num): |
170
|
|
|
m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"]) |
171
|
|
|
t = threading.Thread(target=build, args=(m,)) |
172
|
|
|
threads.append(t) |
173
|
|
|
t.start() |
174
|
|
|
time.sleep(0.2) |
175
|
|
|
for t in threads: |
176
|
|
|
t.join() |
177
|
|
|
|
178
|
|
|
query_vec = [vectors[0]] |
179
|
|
|
top_k = 1 |
180
|
|
|
search_param = {"nprobe": nprobe} |
181
|
|
|
status, result = connect.search(collection, top_k, query_vec, params=search_param) |
182
|
|
|
assert len(result) == 1 |
183
|
|
|
assert len(result[0]) == top_k |
184
|
|
|
assert result[0][0].distance == 0.0 |
185
|
|
|
|
186
|
|
View Code Duplication |
@pytest.mark.level(2) |
|
|
|
|
187
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
188
|
|
|
def test_create_index_multithread_multicollection(self, connect, args): |
189
|
|
|
''' |
190
|
|
|
target: test create index interface with multiprocess |
191
|
|
|
method: create collection and add vectors in it, create index |
192
|
|
|
expected: return code equals to 0, and search success |
193
|
|
|
''' |
194
|
|
|
threads_num = 8 |
195
|
|
|
loop_num = 8 |
196
|
|
|
threads = [] |
197
|
|
|
collection = [] |
198
|
|
|
j = 0 |
199
|
|
|
while j < (threads_num*loop_num): |
200
|
|
|
collection_name = gen_unique_str("test_create_index_multiprocessing") |
201
|
|
|
collection.append(collection_name) |
202
|
|
|
param = {'collection_name': collection_name, |
203
|
|
|
'dimension': dim, |
204
|
|
|
'index_type': IndexType.FLAT, |
205
|
|
|
'store_raw_vector': False} |
206
|
|
|
connect.create_collection(param) |
207
|
|
|
j = j + 1 |
208
|
|
|
|
209
|
|
|
def create_index(): |
210
|
|
|
i = 0 |
211
|
|
|
while i < loop_num: |
212
|
|
|
# assert connect.has_collection(collection[ids*process_num+i]) |
213
|
|
|
status, ids = connect.insert(collection[ids*threads_num+i], vectors) |
|
|
|
|
214
|
|
|
status = connect.create_index(collection[ids*threads_num+i], IndexType.IVFLAT, {"nlist": NLIST}) |
215
|
|
|
assert status.OK() |
216
|
|
|
query_vec = [vectors[0]] |
217
|
|
|
top_k = 1 |
218
|
|
|
search_param = {"nprobe": nprobe} |
219
|
|
|
status, result = connect.search(collection[ids*threads_num+i], top_k, query_vec, params=search_param) |
220
|
|
|
assert len(result) == 1 |
221
|
|
|
assert len(result[0]) == top_k |
222
|
|
|
assert result[0][0].distance == 0.0 |
223
|
|
|
i = i + 1 |
224
|
|
|
for i in range(threads_num): |
225
|
|
|
m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"]) |
226
|
|
|
ids = i |
227
|
|
|
t = threading.Thread(target=create_index, args=(m, ids)) |
228
|
|
|
threads.append(t) |
229
|
|
|
t.start() |
230
|
|
|
time.sleep(0.2) |
231
|
|
|
for t in threads: |
232
|
|
|
t.join() |
233
|
|
|
|
234
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
235
|
|
|
@pytest.mark.level(2) |
236
|
|
|
def test_create_index_a_multithreads(self, connect, collection, args): |
237
|
|
|
status, ids = connect.insert(collection, vectors) |
238
|
|
|
def build(connect): |
239
|
|
|
status = connect.create_index(collection, IndexType.IVFLAT, {"nlist": NLIST}) |
240
|
|
|
assert status.OK() |
241
|
|
|
def count(connect): |
242
|
|
|
status, count = connect.count_entities(collection) |
243
|
|
|
assert status.OK() |
244
|
|
|
assert count == nb |
245
|
|
|
|
246
|
|
|
threads_num = 8 |
247
|
|
|
threads = [] |
248
|
|
|
uri = "tcp://%s:%s" % (args["ip"], args["port"]) |
249
|
|
|
for i in range(threads_num): |
250
|
|
|
m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"]) |
251
|
|
|
if(i % 2 == 0): |
252
|
|
|
p = threading.Thread(target=build, args=(m,)) |
253
|
|
|
else: |
254
|
|
|
p = threading.Thread(target=count, args=(m,)) |
255
|
|
|
threads.append(p) |
256
|
|
|
p.start() |
257
|
|
|
time.sleep(0.2) |
258
|
|
|
for p in threads: |
259
|
|
|
p.join() |
260
|
|
|
|
261
|
|
|
|
262
|
|
|
# TODO: enable |
263
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
264
|
|
|
@pytest.mark.level(2) |
265
|
|
|
def _test_create_index_multiprocessing(self, connect, collection, args): |
266
|
|
|
''' |
267
|
|
|
target: test create index interface with multiprocess |
268
|
|
|
method: create collection and add vectors in it, create index |
269
|
|
|
expected: return code equals to 0, and search success |
270
|
|
|
''' |
271
|
|
|
status, ids = connect.insert(collection, vectors) |
272
|
|
|
|
273
|
|
|
def build(connect): |
274
|
|
|
status = connect.create_index(collection, IndexType.IVFLAT, {"nlist": NLIST}) |
275
|
|
|
assert status.OK() |
276
|
|
|
|
277
|
|
|
process_num = 8 |
278
|
|
|
processes = [] |
279
|
|
|
for i in range(process_num): |
280
|
|
|
m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"]) |
281
|
|
|
p = Process(target=build, args=(m,)) |
282
|
|
|
processes.append(p) |
283
|
|
|
p.start() |
284
|
|
|
time.sleep(0.2) |
285
|
|
|
for p in processes: |
286
|
|
|
p.join() |
287
|
|
|
|
288
|
|
|
query_vec = [vectors[0]] |
289
|
|
|
top_k = 1 |
290
|
|
|
search_param = {"nprobe": nprobe} |
291
|
|
|
status, result = connect.search(collection, top_k, query_vec, params=search_param) |
292
|
|
|
assert len(result) == 1 |
293
|
|
|
assert len(result[0]) == top_k |
294
|
|
|
assert result[0][0].distance == 0.0 |
295
|
|
|
|
296
|
|
|
# TODO: enable |
297
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
298
|
|
|
def _test_create_index_multiprocessing_multicollection(self, connect, args): |
299
|
|
|
''' |
300
|
|
|
target: test create index interface with multiprocess |
301
|
|
|
method: create collection and add vectors in it, create index |
302
|
|
|
expected: return code equals to 0, and search success |
303
|
|
|
''' |
304
|
|
|
process_num = 8 |
305
|
|
|
loop_num = 8 |
306
|
|
|
processes = [] |
307
|
|
|
|
308
|
|
|
collection = [] |
309
|
|
|
j = 0 |
310
|
|
|
while j < (process_num*loop_num): |
311
|
|
|
collection_name = gen_unique_str("test_create_index_multiprocessing") |
312
|
|
|
collection.append(collection_name) |
313
|
|
|
param = {'collection_name': collection_name, |
314
|
|
|
'dimension': dim, |
315
|
|
|
'index_type': IndexType.FLAT, |
316
|
|
|
'store_raw_vector': False} |
317
|
|
|
connect.create_collection(param) |
318
|
|
|
j = j + 1 |
319
|
|
|
|
320
|
|
|
def create_index(): |
321
|
|
|
i = 0 |
322
|
|
|
while i < loop_num: |
323
|
|
|
# assert connect.has_collection(collection[ids*process_num+i]) |
324
|
|
|
status, ids = connect.insert(collection[ids*process_num+i], vectors) |
|
|
|
|
325
|
|
|
|
326
|
|
|
status = connect.create_index(collection[ids*process_num+i], IndexType.IVFLAT, {"nlist": NLIST}) |
327
|
|
|
assert status.OK() |
328
|
|
|
query_vec = [vectors[0]] |
329
|
|
|
top_k = 1 |
330
|
|
|
search_param = {"nprobe": nprobe} |
331
|
|
|
status, result = connect.search(collection[ids*process_num+i], top_k, query_vec, params=search_param) |
332
|
|
|
assert len(result) == 1 |
333
|
|
|
assert len(result[0]) == top_k |
334
|
|
|
assert result[0][0].distance == 0.0 |
335
|
|
|
i = i + 1 |
336
|
|
|
|
337
|
|
|
for i in range(process_num): |
338
|
|
|
m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"]) |
339
|
|
|
ids = i |
340
|
|
|
p = Process(target=create_index, args=(m,ids)) |
341
|
|
|
processes.append(p) |
342
|
|
|
p.start() |
343
|
|
|
time.sleep(0.2) |
344
|
|
|
for p in processes: |
345
|
|
|
p.join() |
346
|
|
|
|
347
|
|
|
def test_create_index_collection_not_existed(self, connect): |
348
|
|
|
''' |
349
|
|
|
target: test create index interface when collection name not existed |
350
|
|
|
method: create collection and add vectors in it, create index |
351
|
|
|
, make sure the collection name not in index |
352
|
|
|
expected: return code not equals to 0, create index failed |
353
|
|
|
''' |
354
|
|
|
collection_name = gen_unique_str(self.__class__.__name__) |
355
|
|
|
nlist = NLIST |
356
|
|
|
index_type = IndexType.IVF_SQ8 |
357
|
|
|
index_param = {"nlist": nlist} |
358
|
|
|
status = connect.create_index(collection_name, index_type, index_param) |
359
|
|
|
assert not status.OK() |
360
|
|
|
|
361
|
|
|
def test_create_index_collection_None(self, connect): |
362
|
|
|
''' |
363
|
|
|
target: test create index interface when collection name is None |
364
|
|
|
method: create collection and add vectors in it, create index with an collection_name: None |
365
|
|
|
expected: return code not equals to 0, create index failed |
366
|
|
|
''' |
367
|
|
|
collection_name = None |
368
|
|
|
nlist = NLIST |
369
|
|
|
index_type = IndexType.IVF_SQ8 |
370
|
|
|
index_param = {"nlist": nlist} |
371
|
|
|
with pytest.raises(Exception) as e: |
372
|
|
|
status = connect.create_index(collection_name, index_type, index_param) |
373
|
|
|
|
374
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
375
|
|
|
def test_create_index_no_vectors_then_insert(self, connect, collection, get_simple_index): |
376
|
|
|
''' |
377
|
|
|
target: test create index interface when there is no vectors in collection, and does not affect the subsequent process |
378
|
|
|
method: create collection and add no vectors in it, and then create index, add vectors in it |
379
|
|
|
expected: return code equals to 0 |
380
|
|
|
''' |
381
|
|
|
index_param = get_simple_index["index_param"] |
382
|
|
|
index_type = get_simple_index["index_type"] |
383
|
|
|
status = connect.create_index(collection, index_type, index_param) |
384
|
|
|
status, ids = connect.insert(collection, vectors) |
385
|
|
|
assert status.OK() |
386
|
|
|
|
387
|
|
|
@pytest.mark.level(2) |
388
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
389
|
|
|
def test_create_same_index_repeatedly(self, connect, collection, get_simple_index): |
390
|
|
|
''' |
391
|
|
|
target: check if index can be created repeatedly, with the same create_index params |
392
|
|
|
method: create index after index have been built |
393
|
|
|
expected: return code success, and search ok |
394
|
|
|
''' |
395
|
|
|
index_param = get_simple_index["index_param"] |
396
|
|
|
index_type = get_simple_index["index_type"] |
397
|
|
|
status = connect.create_index(collection, index_type, index_param) |
398
|
|
|
status = connect.create_index(collection, index_type, index_param) |
399
|
|
|
assert status.OK() |
400
|
|
|
|
401
|
|
View Code Duplication |
@pytest.mark.level(2) |
|
|
|
|
402
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
403
|
|
|
def test_create_different_index_repeatedly(self, connect, collection): |
404
|
|
|
''' |
405
|
|
|
target: check if index can be created repeatedly, with the different create_index params |
406
|
|
|
method: create another index with different index_params after index have been built |
407
|
|
|
expected: return code 0, and describe index result equals with the second index params |
408
|
|
|
''' |
409
|
|
|
nlist = NLIST |
410
|
|
|
status, ids = connect.insert(collection, vectors) |
411
|
|
|
index_type_1 = IndexType.IVF_SQ8 |
412
|
|
|
index_type_2 = IndexType.IVFLAT |
413
|
|
|
indexs = [{"index_type": index_type_1, "index_param": {"nlist": nlist}}, {"index_type": index_type_2, "index_param": {"nlist": nlist}}] |
414
|
|
|
logging.getLogger().info(indexs) |
415
|
|
|
for index in indexs: |
416
|
|
|
status = connect.create_index(collection, index["index_type"], index["index_param"]) |
417
|
|
|
assert status.OK() |
418
|
|
|
status, result = connect.get_index_info(collection) |
419
|
|
|
assert result._params["nlist"] == nlist |
420
|
|
|
assert result._collection_name == collection |
421
|
|
|
assert result._index_type == index_type_2 |
422
|
|
|
|
423
|
|
|
""" |
424
|
|
|
****************************************************************** |
425
|
|
|
The following cases are used to test `get_index_info` function |
426
|
|
|
****************************************************************** |
427
|
|
|
""" |
428
|
|
|
|
429
|
|
|
def test_get_index_info(self, connect, collection, get_index): |
430
|
|
|
''' |
431
|
|
|
target: test describe index interface |
432
|
|
|
method: create collection and add vectors in it, create index, call describe index |
433
|
|
|
expected: return code 0, and index instructure |
434
|
|
|
''' |
435
|
|
|
index_param = get_index["index_param"] |
436
|
|
|
index_type = get_index["index_type"] |
437
|
|
|
logging.getLogger().info(get_index) |
438
|
|
|
# status, ids = connect.insert(collection, vectors) |
439
|
|
|
status = connect.create_index(collection, index_type, index_param) |
440
|
|
|
if status.OK(): |
441
|
|
|
status, result = connect.get_index_info(collection) |
442
|
|
|
logging.getLogger().info(result) |
443
|
|
|
assert result._params == index_param |
444
|
|
|
assert result._collection_name == collection |
445
|
|
|
assert result._index_type == index_type |
446
|
|
|
|
447
|
|
View Code Duplication |
def test_describe_and_drop_index_multi_collections(self, connect, get_simple_index): |
|
|
|
|
448
|
|
|
''' |
449
|
|
|
target: test create, describe and drop index interface with multiple collections of L2 |
450
|
|
|
method: create collections and add vectors in it, create index, call describe index |
451
|
|
|
expected: return code 0, and index instructure |
452
|
|
|
''' |
453
|
|
|
nq = 100 |
454
|
|
|
vectors = gen_vectors(nq, dim) |
455
|
|
|
collection_list = [] |
456
|
|
|
for i in range(10): |
457
|
|
|
collection_name = gen_unique_str() |
458
|
|
|
collection_list.append(collection_name) |
459
|
|
|
param = {'collection_name': collection_name, |
460
|
|
|
'dimension': dim, |
461
|
|
|
'index_file_size': index_file_size, |
462
|
|
|
'metric_type': MetricType.L2} |
463
|
|
|
connect.create_collection(param) |
464
|
|
|
index_param = get_simple_index["index_param"] |
465
|
|
|
index_type = get_simple_index["index_type"] |
466
|
|
|
logging.getLogger().info(get_simple_index) |
467
|
|
|
status, ids = connect.insert(collection_name=collection_name, records=vectors) |
468
|
|
|
status = connect.create_index(collection_name, index_type, index_param) |
469
|
|
|
assert status.OK() |
470
|
|
|
|
471
|
|
|
for i in range(10): |
472
|
|
|
status, result = connect.get_index_info(collection_list[i]) |
473
|
|
|
logging.getLogger().info(result) |
474
|
|
|
assert result._params == index_param |
|
|
|
|
475
|
|
|
assert result._collection_name == collection_list[i] |
476
|
|
|
assert result._index_type == index_type |
|
|
|
|
477
|
|
|
|
478
|
|
|
for i in range(10): |
479
|
|
|
status = connect.drop_index(collection_list[i]) |
480
|
|
|
assert status.OK() |
481
|
|
|
status, result = connect.get_index_info(collection_list[i]) |
482
|
|
|
logging.getLogger().info(result) |
483
|
|
|
assert result._collection_name == collection_list[i] |
484
|
|
|
assert result._index_type == IndexType.FLAT |
485
|
|
|
|
486
|
|
|
# @pytest.mark.level(2) |
487
|
|
|
# def test_get_index_info_without_connect(self, dis_connect, collection): |
488
|
|
|
# ''' |
489
|
|
|
# target: test describe index without connection |
490
|
|
|
# method: describe index, and check if describe successfully |
491
|
|
|
# expected: raise exception |
492
|
|
|
# ''' |
493
|
|
|
# with pytest.raises(Exception) as e: |
494
|
|
|
# status = dis_connect.get_index_info(collection) |
495
|
|
|
|
496
|
|
|
def test_get_index_info_collection_not_existed(self, connect): |
497
|
|
|
''' |
498
|
|
|
target: test describe index interface when collection name not existed |
499
|
|
|
method: create collection and add vectors in it, create index |
500
|
|
|
, make sure the collection name not in index |
501
|
|
|
expected: return code not equals to 0, describe index failed |
502
|
|
|
''' |
503
|
|
|
collection_name = gen_unique_str(self.__class__.__name__) |
504
|
|
|
status, result = connect.get_index_info(collection_name) |
505
|
|
|
assert not status.OK() |
506
|
|
|
|
507
|
|
|
def test_get_index_info_collection_None(self, connect): |
508
|
|
|
''' |
509
|
|
|
target: test describe index interface when collection name is None |
510
|
|
|
method: create collection and add vectors in it, create index with an collection_name: None |
511
|
|
|
expected: return code not equals to 0, describe index failed |
512
|
|
|
''' |
513
|
|
|
collection_name = None |
514
|
|
|
with pytest.raises(Exception) as e: |
515
|
|
|
status = connect.get_index_info(collection_name) |
516
|
|
|
|
517
|
|
|
def test_get_index_info_not_create(self, connect, collection): |
518
|
|
|
''' |
519
|
|
|
target: test describe index interface when index not created |
520
|
|
|
method: create collection and add vectors in it, create index |
521
|
|
|
, make sure the collection name not in index |
522
|
|
|
expected: return code not equals to 0, describe index failed |
523
|
|
|
''' |
524
|
|
|
status, ids = connect.insert(collection, vectors) |
525
|
|
|
status, result = connect.get_index_info(collection) |
526
|
|
|
logging.getLogger().info(result) |
527
|
|
|
assert status.OK() |
528
|
|
|
# assert result._params["nlist"] == index_params["nlist"] |
529
|
|
|
# assert result._collection_name == collection |
530
|
|
|
# assert result._index_type == index_params["index_type"] |
531
|
|
|
|
532
|
|
|
""" |
533
|
|
|
****************************************************************** |
534
|
|
|
The following cases are used to test `drop_index` function |
535
|
|
|
****************************************************************** |
536
|
|
|
""" |
537
|
|
|
|
538
|
|
|
def test_drop_index(self, connect, collection, get_simple_index): |
539
|
|
|
''' |
540
|
|
|
target: test drop index interface |
541
|
|
|
method: create collection and add vectors in it, create index, call drop index |
542
|
|
|
expected: return code 0, and default index param |
543
|
|
|
''' |
544
|
|
|
index_param = get_simple_index["index_param"] |
545
|
|
|
index_type = get_simple_index["index_type"] |
546
|
|
|
# status, ids = connect.insert(collection, vectors) |
547
|
|
|
status = connect.create_index(collection, index_type, index_param) |
548
|
|
|
assert status.OK() |
549
|
|
|
status, result = connect.get_index_info(collection) |
550
|
|
|
logging.getLogger().info(result) |
551
|
|
|
status = connect.drop_index(collection) |
552
|
|
|
assert status.OK() |
553
|
|
|
status, result = connect.get_index_info(collection) |
554
|
|
|
logging.getLogger().info(result) |
555
|
|
|
assert result._collection_name == collection |
556
|
|
|
assert result._index_type == IndexType.FLAT |
557
|
|
|
|
558
|
|
View Code Duplication |
@pytest.mark.level(2) |
|
|
|
|
559
|
|
|
def test_drop_index_repeatly(self, connect, collection, get_simple_index): |
560
|
|
|
''' |
561
|
|
|
target: test drop index repeatly |
562
|
|
|
method: create index, call drop index, and drop again |
563
|
|
|
expected: return code 0 |
564
|
|
|
''' |
565
|
|
|
index_param = get_simple_index["index_param"] |
566
|
|
|
index_type = get_simple_index["index_type"] |
567
|
|
|
# status, ids = connect.insert(collection, vectors) |
568
|
|
|
status = connect.create_index(collection, index_type, index_param) |
569
|
|
|
assert status.OK() |
570
|
|
|
status, result = connect.get_index_info(collection) |
571
|
|
|
logging.getLogger().info(result) |
572
|
|
|
status = connect.drop_index(collection) |
573
|
|
|
assert status.OK() |
574
|
|
|
status = connect.drop_index(collection) |
575
|
|
|
assert status.OK() |
576
|
|
|
status, result = connect.get_index_info(collection) |
577
|
|
|
logging.getLogger().info(result) |
578
|
|
|
assert result._collection_name == collection |
579
|
|
|
assert result._index_type == IndexType.FLAT |
580
|
|
|
|
581
|
|
|
# @pytest.mark.level(2) |
582
|
|
|
# def test_drop_index_without_connect(self, dis_connect, collection): |
583
|
|
|
# ''' |
584
|
|
|
# target: test drop index without connection |
585
|
|
|
# method: drop index, and check if drop successfully |
586
|
|
|
# expected: raise exception |
587
|
|
|
# ''' |
588
|
|
|
# with pytest.raises(Exception) as e: |
589
|
|
|
# status = dis_connect.drop_index(collection) |
590
|
|
|
|
591
|
|
|
def test_drop_index_collection_not_existed(self, connect): |
592
|
|
|
''' |
593
|
|
|
target: test drop index interface when collection name not existed |
594
|
|
|
method: create collection and add vectors in it, create index |
595
|
|
|
, make sure the collection name not in index, and then drop it |
596
|
|
|
expected: return code not equals to 0, drop index failed |
597
|
|
|
''' |
598
|
|
|
collection_name = gen_unique_str(self.__class__.__name__) |
599
|
|
|
status = connect.drop_index(collection_name) |
600
|
|
|
assert not status.OK() |
601
|
|
|
|
602
|
|
|
def test_drop_index_collection_None(self, connect): |
603
|
|
|
''' |
604
|
|
|
target: test drop index interface when collection name is None |
605
|
|
|
method: create collection and add vectors in it, create index with an collection_name: None |
606
|
|
|
expected: return code not equals to 0, drop index failed |
607
|
|
|
''' |
608
|
|
|
collection_name = None |
609
|
|
|
with pytest.raises(Exception) as e: |
610
|
|
|
status = connect.drop_index(collection_name) |
611
|
|
|
|
612
|
|
|
def test_drop_index_collection_not_create(self, connect, collection): |
613
|
|
|
''' |
614
|
|
|
target: test drop index interface when index not created |
615
|
|
|
method: create collection and add vectors in it, create index |
616
|
|
|
expected: return code not equals to 0, drop index failed |
617
|
|
|
''' |
618
|
|
|
status, ids = connect.insert(collection, vectors) |
619
|
|
|
status, result = connect.get_index_info(collection) |
620
|
|
|
logging.getLogger().info(result) |
621
|
|
|
# no create index |
622
|
|
|
status = connect.drop_index(collection) |
623
|
|
|
logging.getLogger().info(status) |
624
|
|
|
assert status.OK() |
625
|
|
|
|
626
|
|
View Code Duplication |
@pytest.mark.level(2) |
|
|
|
|
627
|
|
|
def test_create_drop_index_repeatly(self, connect, collection, get_simple_index): |
628
|
|
|
''' |
629
|
|
|
target: test create / drop index repeatly, use the same index params |
630
|
|
|
method: create index, drop index, four times |
631
|
|
|
expected: return code 0 |
632
|
|
|
''' |
633
|
|
|
index_param = get_simple_index["index_param"] |
634
|
|
|
index_type = get_simple_index["index_type"] |
635
|
|
|
# status, ids = connect.insert(collection, vectors) |
636
|
|
|
for i in range(2): |
637
|
|
|
status = connect.create_index(collection, index_type, index_param) |
638
|
|
|
assert status.OK() |
639
|
|
|
status, result = connect.get_index_info(collection) |
640
|
|
|
logging.getLogger().info(result) |
641
|
|
|
status = connect.drop_index(collection) |
642
|
|
|
assert status.OK() |
643
|
|
|
status, result = connect.get_index_info(collection) |
644
|
|
|
logging.getLogger().info(result) |
645
|
|
|
assert result._collection_name == collection |
646
|
|
|
assert result._index_type == IndexType.FLAT |
647
|
|
|
|
648
|
|
|
def test_create_drop_index_repeatly_different_index_params(self, connect, collection): |
649
|
|
|
''' |
650
|
|
|
target: test create / drop index repeatly, use the different index params |
651
|
|
|
method: create index, drop index, four times, each tme use different index_params to create index |
652
|
|
|
expected: return code 0 |
653
|
|
|
''' |
654
|
|
|
nlist = NLIST |
655
|
|
|
indexs = [{"index_type": IndexType.IVFLAT, "index_param": {"nlist": nlist}}, {"index_type": IndexType.IVF_SQ8, "index_param": {"nlist": nlist}}] |
656
|
|
|
# status, ids = connect.insert(collection, vectors) |
657
|
|
|
for i in range(2): |
658
|
|
|
status = connect.create_index(collection, indexs[i]["index_type"], indexs[i]["index_param"]) |
659
|
|
|
assert status.OK() |
660
|
|
|
status, result = connect.get_index_info(collection) |
661
|
|
|
logging.getLogger().info(result) |
662
|
|
|
status = connect.drop_index(collection) |
663
|
|
|
assert status.OK() |
664
|
|
|
status, result = connect.get_index_info(collection) |
665
|
|
|
logging.getLogger().info(result) |
666
|
|
|
assert result._collection_name == collection |
667
|
|
|
assert result._index_type == IndexType.FLAT |
668
|
|
|
|
669
|
|
|
|
670
|
|
|
class TestIndexIP: |
671
|
|
|
@pytest.fixture( |
672
|
|
|
scope="function", |
673
|
|
|
params=gen_index() |
674
|
|
|
) |
675
|
|
|
def get_index(self, request, connect): |
676
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
677
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
678
|
|
|
pytest.skip("sq8h not support in CPU mode") |
679
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
680
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
681
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
682
|
|
|
if request.param["index_type"] == IndexType.RNSG: |
683
|
|
|
pytest.skip("rnsg not support in ip") |
684
|
|
|
return request.param |
685
|
|
|
|
686
|
|
|
@pytest.fixture( |
687
|
|
|
scope="function", |
688
|
|
|
params=gen_simple_index() |
689
|
|
|
) |
690
|
|
|
def get_simple_index(self, request, connect): |
691
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
692
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
693
|
|
|
pytest.skip("sq8h not support in CPU mode") |
694
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
695
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
696
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
697
|
|
|
if request.param["index_type"] == IndexType.RNSG: |
698
|
|
|
pytest.skip("rnsg not support in ip") |
699
|
|
|
return request.param |
700
|
|
|
""" |
701
|
|
|
****************************************************************** |
702
|
|
|
The following cases are used to test `create_index` function |
703
|
|
|
****************************************************************** |
704
|
|
|
""" |
705
|
|
|
@pytest.mark.level(2) |
706
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
707
|
|
|
def test_create_index(self, connect, ip_collection, get_simple_index): |
708
|
|
|
''' |
709
|
|
|
target: test create index interface |
710
|
|
|
method: create collection and add vectors in it, create index |
711
|
|
|
expected: return code equals to 0, and search success |
712
|
|
|
''' |
713
|
|
|
index_param = get_simple_index["index_param"] |
714
|
|
|
index_type = get_simple_index["index_type"] |
715
|
|
|
logging.getLogger().info(get_simple_index) |
716
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
717
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
718
|
|
|
assert status.OK() |
719
|
|
|
|
720
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
721
|
|
|
def test_create_index_collection(self, connect, ip_collection, get_simple_index): |
722
|
|
|
''' |
723
|
|
|
target: test create index interface |
724
|
|
|
method: create collection, create partition, and add vectors in it, create index on collection |
725
|
|
|
expected: return code equals to 0, and search success |
726
|
|
|
''' |
727
|
|
|
index_param = get_simple_index["index_param"] |
728
|
|
|
index_type = get_simple_index["index_type"] |
729
|
|
|
logging.getLogger().info(get_simple_index) |
730
|
|
|
status = connect.create_partition(ip_collection, tag) |
731
|
|
|
status, ids = connect.insert(ip_collection, vectors, partition_tag=tag) |
732
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
733
|
|
|
assert status.OK() |
734
|
|
|
|
735
|
|
|
# @pytest.mark.level(2) |
736
|
|
|
# def test_create_index_without_connect(self, dis_connect, ip_collection): |
737
|
|
|
# ''' |
738
|
|
|
# target: test create index without connection |
739
|
|
|
# method: create collection and add vectors in it, check if added successfully |
740
|
|
|
# expected: raise exception |
741
|
|
|
# ''' |
742
|
|
|
# nlist = NLIST |
743
|
|
|
# index_type = IndexType.IVF_SQ8 |
744
|
|
|
# index_param = {"nlist": nlist} |
745
|
|
|
# with pytest.raises(Exception) as e: |
746
|
|
|
# status = dis_connect.create_index(ip_collection, index_type, index_param) |
747
|
|
|
|
748
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
749
|
|
|
def test_create_index_search_with_query_vectors(self, connect, ip_collection, get_simple_index): |
750
|
|
|
''' |
751
|
|
|
target: test create index interface, search with more query vectors |
752
|
|
|
method: create collection and add vectors in it, create index, with no manual flush |
753
|
|
|
expected: return code equals to 0, and search success |
754
|
|
|
''' |
755
|
|
|
index_param = get_simple_index["index_param"] |
756
|
|
|
index_type = get_simple_index["index_type"] |
757
|
|
|
logging.getLogger().info(get_simple_index) |
758
|
|
|
logging.getLogger().info(connect.get_collection_info(ip_collection)) |
759
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
760
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
761
|
|
|
logging.getLogger().info(connect.get_index_info(ip_collection)) |
762
|
|
|
logging.getLogger().info(connect.get_collection_stats(ip_collection)) |
763
|
|
|
query_vecs = [vectors[0], vectors[1], vectors[2]] |
764
|
|
|
top_k = 5 |
765
|
|
|
search_param = get_search_param(index_type) |
766
|
|
|
status, result = connect.search(ip_collection, top_k, query_vecs, params=search_param) |
767
|
|
|
logging.getLogger().info(result) |
768
|
|
|
assert status.OK() |
769
|
|
|
assert len(result) == len(query_vecs) |
770
|
|
|
|
771
|
|
|
# TODO: enable |
772
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
773
|
|
|
@pytest.mark.level(2) |
774
|
|
|
def _test_create_index_multiprocessing(self, connect, ip_collection, args): |
775
|
|
|
''' |
776
|
|
|
target: test create index interface with multiprocess |
777
|
|
|
method: create collection and add vectors in it, create index |
778
|
|
|
expected: return code equals to 0, and search success |
779
|
|
|
''' |
780
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
781
|
|
|
def build(connect): |
782
|
|
|
status = connect.create_index(ip_collection, IndexType.IVFLAT, {"nlist": NLIST}) |
783
|
|
|
assert status.OK() |
784
|
|
|
|
785
|
|
|
process_num = 8 |
786
|
|
|
processes = [] |
787
|
|
|
|
788
|
|
|
for i in range(process_num): |
789
|
|
|
m = get_milvus(args["ip"], args["port"], handler=args["handler"]) |
790
|
|
|
p = Process(target=build, args=(m,)) |
791
|
|
|
processes.append(p) |
792
|
|
|
p.start() |
793
|
|
|
time.sleep(0.2) |
794
|
|
|
for p in processes: |
795
|
|
|
p.join() |
796
|
|
|
|
797
|
|
|
query_vec = [vectors[0]] |
798
|
|
|
top_k = 1 |
799
|
|
|
search_param = {"nprobe": nprobe} |
800
|
|
|
status, result = connect.search(ip_collection, top_k, query_vec, params=search_param) |
801
|
|
|
assert len(result) == 1 |
802
|
|
|
assert len(result[0]) == top_k |
803
|
|
|
assert result[0][0].distance == 0.0 |
804
|
|
|
|
805
|
|
|
# TODO: enable |
806
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
807
|
|
|
def _test_create_index_multiprocessing_multicollection(self, connect, args): |
808
|
|
|
''' |
809
|
|
|
target: test create index interface with multiprocess |
810
|
|
|
method: create collection and add vectors in it, create index |
811
|
|
|
expected: return code equals to 0, and search success |
812
|
|
|
''' |
813
|
|
|
process_num = 8 |
814
|
|
|
loop_num = 8 |
815
|
|
|
processes = [] |
816
|
|
|
|
817
|
|
|
collection = [] |
818
|
|
|
j = 0 |
819
|
|
|
while j < (process_num*loop_num): |
820
|
|
|
collection_name = gen_unique_str("test_create_index_multiprocessing") |
821
|
|
|
collection.append(collection_name) |
822
|
|
|
param = {'collection_name': collection_name, |
823
|
|
|
'dimension': dim} |
824
|
|
|
connect.create_collection(param) |
825
|
|
|
j = j + 1 |
826
|
|
|
|
827
|
|
|
def create_index(): |
828
|
|
|
i = 0 |
829
|
|
|
while i < loop_num: |
830
|
|
|
# assert connect.has_collection(collection[ids*process_num+i]) |
831
|
|
|
status, ids = connect.insert(collection[ids*process_num+i], vectors) |
|
|
|
|
832
|
|
|
|
833
|
|
|
status = connect.create_index(collection[ids*process_num+i], IndexType.IVFLAT, {"nlist": NLIST}) |
834
|
|
|
assert status.OK() |
835
|
|
|
query_vec = [vectors[0]] |
836
|
|
|
top_k = 1 |
837
|
|
|
search_param = {"nprobe": nprobe} |
838
|
|
|
status, result = connect.search(collection[ids*process_num+i], top_k, query_vec, params=search_param) |
839
|
|
|
assert len(result) == 1 |
840
|
|
|
assert len(result[0]) == top_k |
841
|
|
|
assert result[0][0].distance == 0.0 |
842
|
|
|
i = i + 1 |
843
|
|
|
|
844
|
|
|
for i in range(process_num): |
845
|
|
|
m = get_milvus(args["ip"], args["port"], handler=args["handler"]) |
846
|
|
|
ids = i |
847
|
|
|
p = Process(target=create_index, args=(m,ids)) |
848
|
|
|
processes.append(p) |
849
|
|
|
p.start() |
850
|
|
|
time.sleep(0.2) |
851
|
|
|
for p in processes: |
852
|
|
|
p.join() |
853
|
|
|
|
854
|
|
|
def test_create_index_no_vectors(self, connect, ip_collection): |
855
|
|
|
''' |
856
|
|
|
target: test create index interface when there is no vectors in collection |
857
|
|
|
method: create collection and add no vectors in it, and then create index |
858
|
|
|
expected: return code equals to 0 |
859
|
|
|
''' |
860
|
|
|
nlist = NLIST |
861
|
|
|
index_type = IndexType.IVF_SQ8 |
862
|
|
|
index_param = {"nlist": nlist} |
863
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
864
|
|
|
assert status.OK() |
865
|
|
|
|
866
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
867
|
|
|
def test_create_index_no_vectors_then_insert(self, connect, ip_collection, get_simple_index): |
868
|
|
|
''' |
869
|
|
|
target: test create index interface when there is no vectors in collection, and does not affect the subsequent process |
870
|
|
|
method: create collection and add no vectors in it, and then create index, add vectors in it |
871
|
|
|
expected: return code equals to 0 |
872
|
|
|
''' |
873
|
|
|
index_param = get_simple_index["index_param"] |
874
|
|
|
index_type = get_simple_index["index_type"] |
875
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
876
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
877
|
|
|
assert status.OK() |
878
|
|
|
|
879
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
880
|
|
|
def test_create_same_index_repeatedly(self, connect, ip_collection): |
881
|
|
|
''' |
882
|
|
|
target: check if index can be created repeatedly, with the same create_index params |
883
|
|
|
method: create index after index have been built |
884
|
|
|
expected: return code success, and search ok |
885
|
|
|
''' |
886
|
|
|
nlist = NLIST |
887
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
888
|
|
|
index_type = IndexType.IVF_SQ8 |
889
|
|
|
index_param = {"nlist": nlist} |
890
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
891
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
892
|
|
|
assert status.OK() |
893
|
|
|
query_vec = [vectors[0]] |
894
|
|
|
top_k = 1 |
895
|
|
|
search_param = {"nprobe": nprobe} |
896
|
|
|
status, result = connect.search(ip_collection, top_k, query_vec, params=search_param) |
897
|
|
|
assert len(result) == 1 |
898
|
|
|
assert len(result[0]) == top_k |
899
|
|
|
|
900
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
901
|
|
|
def test_create_different_index_repeatedly(self, connect, ip_collection): |
902
|
|
|
''' |
903
|
|
|
target: check if index can be created repeatedly, with the different create_index params |
904
|
|
|
method: create another index with different index_params after index have been built |
905
|
|
|
expected: return code 0, and describe index result equals with the second index params |
906
|
|
|
''' |
907
|
|
|
nlist = NLIST |
908
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
909
|
|
|
index_type_1 = IndexType.IVF_SQ8 |
910
|
|
|
index_type_2 = IndexType.IVFLAT |
911
|
|
|
indexs = [{"index_type": index_type_1, "index_param": {"nlist": nlist}}, {"index_type": index_type_2, "index_param": {"nlist": nlist}}] |
912
|
|
|
logging.getLogger().info(indexs) |
913
|
|
|
for index in indexs: |
914
|
|
|
status = connect.create_index(ip_collection, index["index_type"], index["index_param"]) |
915
|
|
|
assert status.OK() |
916
|
|
|
status, result = connect.get_index_info(ip_collection) |
917
|
|
|
assert result._params["nlist"] == nlist |
918
|
|
|
assert result._collection_name == ip_collection |
919
|
|
|
assert result._index_type == index_type_2 |
920
|
|
|
|
921
|
|
|
""" |
922
|
|
|
****************************************************************** |
923
|
|
|
The following cases are used to test `get_index_info` function |
924
|
|
|
****************************************************************** |
925
|
|
|
""" |
926
|
|
|
|
927
|
|
|
def test_get_index_info(self, connect, ip_collection, get_simple_index): |
928
|
|
|
''' |
929
|
|
|
target: test describe index interface |
930
|
|
|
method: create collection and add vectors in it, create index, call describe index |
931
|
|
|
expected: return code 0, and index instructure |
932
|
|
|
''' |
933
|
|
|
index_param = get_simple_index["index_param"] |
934
|
|
|
index_type = get_simple_index["index_type"] |
935
|
|
|
logging.getLogger().info(get_simple_index) |
936
|
|
|
# status, ids = connect.insert(ip_collection, vectors[:5000]) |
937
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
938
|
|
|
status, result = connect.get_index_info(ip_collection) |
939
|
|
|
logging.getLogger().info(result) |
940
|
|
|
assert result._collection_name == ip_collection |
941
|
|
|
status, mode = connect._cmd("mode") |
942
|
|
|
if str(mode) == "GPU" and index_type == IndexType.IVF_PQ: |
943
|
|
|
assert result._index_type == IndexType.FLAT |
944
|
|
|
assert result._params["nlist"] == NLIST |
945
|
|
|
else: |
946
|
|
|
assert result._index_type == index_type |
947
|
|
|
assert result._params == index_param |
948
|
|
|
|
949
|
|
View Code Duplication |
def test_get_index_info_partition(self, connect, ip_collection, get_simple_index): |
|
|
|
|
950
|
|
|
''' |
951
|
|
|
target: test describe index interface |
952
|
|
|
method: create collection, create partition and add vectors in it, create index, call describe index |
953
|
|
|
expected: return code 0, and index instructure |
954
|
|
|
''' |
955
|
|
|
index_param = get_simple_index["index_param"] |
956
|
|
|
index_type = get_simple_index["index_type"] |
957
|
|
|
logging.getLogger().info(get_simple_index) |
958
|
|
|
status = connect.create_partition(ip_collection, tag) |
959
|
|
|
status, ids = connect.insert(ip_collection, vectors, partition_tag=tag) |
960
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
961
|
|
|
status, result = connect.get_index_info(ip_collection) |
962
|
|
|
logging.getLogger().info(result) |
963
|
|
|
assert result._params == index_param |
964
|
|
|
assert result._collection_name == ip_collection |
965
|
|
|
assert result._index_type == index_type |
966
|
|
|
|
967
|
|
|
def test_get_index_info_partition_A(self, connect, ip_collection, get_simple_index): |
968
|
|
|
''' |
969
|
|
|
target: test describe index interface |
970
|
|
|
method: create collection, create partitions and add vectors in it, create index on partitions, call describe index |
971
|
|
|
expected: return code 0, and index instructure |
972
|
|
|
''' |
973
|
|
|
new_tag = "new_tag" |
974
|
|
|
index_param = get_simple_index["index_param"] |
975
|
|
|
index_type = get_simple_index["index_type"] |
976
|
|
|
logging.getLogger().info(get_simple_index) |
977
|
|
|
status = connect.create_partition(ip_collection, tag) |
978
|
|
|
status = connect.create_partition(ip_collection, new_tag) |
979
|
|
|
# status, ids = connect.insert(ip_collection, vectors, partition_tag=tag) |
980
|
|
|
# status, ids = connect.insert(ip_collection, vectors, partition_tag=new_tag) |
981
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
982
|
|
|
status, result = connect.get_index_info(ip_collection) |
983
|
|
|
logging.getLogger().info(result) |
984
|
|
|
assert result._params == index_param |
985
|
|
|
assert result._collection_name == ip_collection |
986
|
|
|
assert result._index_type == index_type |
987
|
|
|
|
988
|
|
View Code Duplication |
def test_describe_and_drop_index_multi_collections(self, connect, get_simple_index): |
|
|
|
|
989
|
|
|
''' |
990
|
|
|
target: test create, describe and drop index interface with multiple collections of IP |
991
|
|
|
method: create collections and add vectors in it, create index, call describe index |
992
|
|
|
expected: return code 0, and index instructure |
993
|
|
|
''' |
994
|
|
|
nq = 100 |
995
|
|
|
vectors = gen_vectors(nq, dim) |
996
|
|
|
collection_list = [] |
997
|
|
|
for i in range(10): |
998
|
|
|
collection_name = gen_unique_str() |
999
|
|
|
collection_list.append(collection_name) |
1000
|
|
|
param = {'collection_name': collection_name, |
1001
|
|
|
'dimension': dim, |
1002
|
|
|
'index_file_size': index_file_size, |
1003
|
|
|
'metric_type': MetricType.IP} |
1004
|
|
|
connect.create_collection(param) |
1005
|
|
|
index_param = get_simple_index["index_param"] |
1006
|
|
|
index_type = get_simple_index["index_type"] |
1007
|
|
|
logging.getLogger().info(get_simple_index) |
1008
|
|
|
status, ids = connect.insert(collection_name=collection_name, records=vectors) |
1009
|
|
|
status = connect.create_index(collection_name, index_type, index_param) |
1010
|
|
|
assert status.OK() |
1011
|
|
|
for i in range(10): |
1012
|
|
|
status, result = connect.get_index_info(collection_list[i]) |
1013
|
|
|
logging.getLogger().info(result) |
1014
|
|
|
assert result._params == index_param |
|
|
|
|
1015
|
|
|
assert result._collection_name == collection_list[i] |
1016
|
|
|
assert result._index_type == index_type |
|
|
|
|
1017
|
|
|
for i in range(10): |
1018
|
|
|
status = connect.drop_index(collection_list[i]) |
1019
|
|
|
assert status.OK() |
1020
|
|
|
status, result = connect.get_index_info(collection_list[i]) |
1021
|
|
|
logging.getLogger().info(result) |
1022
|
|
|
assert result._collection_name == collection_list[i] |
1023
|
|
|
assert result._index_type == IndexType.FLAT |
1024
|
|
|
|
1025
|
|
|
# @pytest.mark.level(2) |
1026
|
|
|
# def test_get_index_info_without_connect(self, dis_connect, ip_collection): |
1027
|
|
|
# ''' |
1028
|
|
|
# target: test describe index without connection |
1029
|
|
|
# method: describe index, and check if describe successfully |
1030
|
|
|
# expected: raise exception |
1031
|
|
|
# ''' |
1032
|
|
|
# with pytest.raises(Exception) as e: |
1033
|
|
|
# status = dis_connect.get_index_info(ip_collection) |
1034
|
|
|
|
1035
|
|
|
def test_get_index_info_not_create(self, connect, ip_collection): |
1036
|
|
|
''' |
1037
|
|
|
target: test describe index interface when index not created |
1038
|
|
|
method: create collection and add vectors in it, create index |
1039
|
|
|
, make sure the collection name not in index |
1040
|
|
|
expected: return code not equals to 0, describe index failed |
1041
|
|
|
''' |
1042
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
1043
|
|
|
status, result = connect.get_index_info(ip_collection) |
1044
|
|
|
logging.getLogger().info(result) |
1045
|
|
|
assert status.OK() |
1046
|
|
|
# assert result._params["nlist"] == index_params["nlist"] |
1047
|
|
|
# assert result._collection_name == collection |
1048
|
|
|
# assert result._index_type == index_params["index_type"] |
1049
|
|
|
|
1050
|
|
|
""" |
1051
|
|
|
****************************************************************** |
1052
|
|
|
The following cases are used to test `drop_index` function |
1053
|
|
|
****************************************************************** |
1054
|
|
|
""" |
1055
|
|
|
|
1056
|
|
View Code Duplication |
def test_drop_index(self, connect, ip_collection, get_simple_index): |
|
|
|
|
1057
|
|
|
''' |
1058
|
|
|
target: test drop index interface |
1059
|
|
|
method: create collection and add vectors in it, create index, call drop index |
1060
|
|
|
expected: return code 0, and default index param |
1061
|
|
|
''' |
1062
|
|
|
index_param = get_simple_index["index_param"] |
1063
|
|
|
index_type = get_simple_index["index_type"] |
1064
|
|
|
status, mode = connect._cmd("mode") |
1065
|
|
|
assert status.OK() |
1066
|
|
|
# status, ids = connect.insert(ip_collection, vectors) |
1067
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
1068
|
|
|
if str(mode) == "GPU" and (index_type == IndexType.IVF_PQ): |
1069
|
|
|
assert not status.OK() |
1070
|
|
|
else: |
1071
|
|
|
assert status.OK() |
1072
|
|
|
status, result = connect.get_index_info(ip_collection) |
1073
|
|
|
logging.getLogger().info(result) |
1074
|
|
|
status = connect.drop_index(ip_collection) |
1075
|
|
|
assert status.OK() |
1076
|
|
|
status, result = connect.get_index_info(ip_collection) |
1077
|
|
|
logging.getLogger().info(result) |
1078
|
|
|
assert result._collection_name == ip_collection |
1079
|
|
|
assert result._index_type == IndexType.FLAT |
1080
|
|
|
|
1081
|
|
|
def test_drop_index_partition(self, connect, ip_collection, get_simple_index): |
1082
|
|
|
''' |
1083
|
|
|
target: test drop index interface |
1084
|
|
|
method: create collection, create partition and add vectors in it, create index on collection, call drop collection index |
1085
|
|
|
expected: return code 0, and default index param |
1086
|
|
|
''' |
1087
|
|
|
index_param = get_simple_index["index_param"] |
1088
|
|
|
index_type = get_simple_index["index_type"] |
1089
|
|
|
status = connect.create_partition(ip_collection, tag) |
1090
|
|
|
status, ids = connect.insert(ip_collection, vectors, partition_tag=tag) |
1091
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
1092
|
|
|
assert status.OK() |
1093
|
|
|
status, result = connect.get_index_info(ip_collection) |
1094
|
|
|
logging.getLogger().info(result) |
1095
|
|
|
status = connect.drop_index(ip_collection) |
1096
|
|
|
assert status.OK() |
1097
|
|
|
status, result = connect.get_index_info(ip_collection) |
1098
|
|
|
logging.getLogger().info(result) |
1099
|
|
|
assert result._collection_name == ip_collection |
1100
|
|
|
assert result._index_type == IndexType.FLAT |
1101
|
|
|
|
1102
|
|
|
def test_drop_index_partition_C(self, connect, ip_collection, get_simple_index): |
1103
|
|
|
''' |
1104
|
|
|
target: test drop index interface |
1105
|
|
|
method: create collection, create partitions and add vectors in it, create index on partitions, call drop partition index |
1106
|
|
|
expected: return code 0, and default index param |
1107
|
|
|
''' |
1108
|
|
|
new_tag = "new_tag" |
1109
|
|
|
index_param = get_simple_index["index_param"] |
1110
|
|
|
index_type = get_simple_index["index_type"] |
1111
|
|
|
status = connect.create_partition(ip_collection, tag) |
1112
|
|
|
status = connect.create_partition(ip_collection, new_tag) |
1113
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
1114
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
1115
|
|
|
assert status.OK() |
1116
|
|
|
status = connect.drop_index(ip_collection) |
1117
|
|
|
assert status.OK() |
1118
|
|
|
status, result = connect.get_index_info(ip_collection) |
1119
|
|
|
logging.getLogger().info(result) |
1120
|
|
|
assert result._collection_name == ip_collection |
1121
|
|
|
assert result._index_type == IndexType.FLAT |
1122
|
|
|
|
1123
|
|
View Code Duplication |
@pytest.mark.level(2) |
|
|
|
|
1124
|
|
|
def test_drop_index_repeatly(self, connect, ip_collection, get_simple_index): |
1125
|
|
|
''' |
1126
|
|
|
target: test drop index repeatly |
1127
|
|
|
method: create index, call drop index, and drop again |
1128
|
|
|
expected: return code 0 |
1129
|
|
|
''' |
1130
|
|
|
index_param = get_simple_index["index_param"] |
1131
|
|
|
index_type = get_simple_index["index_type"] |
1132
|
|
|
# status, ids = connect.insert(ip_collection, vectors) |
1133
|
|
|
status, mode = connect._cmd("mode") |
1134
|
|
|
assert status.OK() |
1135
|
|
|
# status, ids = connect.insert(ip_collection, vectors) |
1136
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
1137
|
|
|
if str(mode) == "GPU" and (index_type == IndexType.IVF_PQ): |
1138
|
|
|
assert not status.OK() |
1139
|
|
|
else: |
1140
|
|
|
assert status.OK() |
1141
|
|
|
status, result = connect.get_index_info(ip_collection) |
1142
|
|
|
logging.getLogger().info(result) |
1143
|
|
|
status = connect.drop_index(ip_collection) |
1144
|
|
|
assert status.OK() |
1145
|
|
|
status = connect.drop_index(ip_collection) |
1146
|
|
|
assert status.OK() |
1147
|
|
|
status, result = connect.get_index_info(ip_collection) |
1148
|
|
|
logging.getLogger().info(result) |
1149
|
|
|
assert result._collection_name == ip_collection |
1150
|
|
|
assert result._index_type == IndexType.FLAT |
1151
|
|
|
|
1152
|
|
|
# @pytest.mark.level(2) |
1153
|
|
|
# def test_drop_index_without_connect(self, dis_connect, ip_collection): |
1154
|
|
|
# ''' |
1155
|
|
|
# target: test drop index without connection |
1156
|
|
|
# method: drop index, and check if drop successfully |
1157
|
|
|
# expected: raise exception |
1158
|
|
|
# ''' |
1159
|
|
|
# nlist = NLIST |
1160
|
|
|
# index_type = IndexType.IVFLAT |
1161
|
|
|
# index_param = {"nlist": nlist} |
1162
|
|
|
# with pytest.raises(Exception) as e: |
1163
|
|
|
# status = dis_connect.drop_index(ip_collection, index_type, index_param) |
1164
|
|
|
|
1165
|
|
|
def test_drop_index_collection_not_create(self, connect, ip_collection): |
1166
|
|
|
''' |
1167
|
|
|
target: test drop index interface when index not created |
1168
|
|
|
method: create collection and add vectors in it, create index |
1169
|
|
|
expected: return code not equals to 0, drop index failed |
1170
|
|
|
''' |
1171
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
1172
|
|
|
status, result = connect.get_index_info(ip_collection) |
1173
|
|
|
logging.getLogger().info(result) |
1174
|
|
|
# no create index |
1175
|
|
|
status = connect.drop_index(ip_collection) |
1176
|
|
|
logging.getLogger().info(status) |
1177
|
|
|
assert status.OK() |
1178
|
|
|
|
1179
|
|
View Code Duplication |
@pytest.mark.level(2) |
|
|
|
|
1180
|
|
|
def test_create_drop_index_repeatly(self, connect, ip_collection, get_simple_index): |
1181
|
|
|
''' |
1182
|
|
|
target: test create / drop index repeatly, use the same index params |
1183
|
|
|
method: create index, drop index, four times |
1184
|
|
|
expected: return code 0 |
1185
|
|
|
''' |
1186
|
|
|
index_param = get_simple_index["index_param"] |
1187
|
|
|
index_type = get_simple_index["index_type"] |
1188
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
1189
|
|
|
for i in range(2): |
1190
|
|
|
status = connect.create_index(ip_collection, index_type, index_param) |
1191
|
|
|
assert status.OK() |
1192
|
|
|
status, result = connect.get_index_info(ip_collection) |
1193
|
|
|
logging.getLogger().info(result) |
1194
|
|
|
status = connect.drop_index(ip_collection) |
1195
|
|
|
assert status.OK() |
1196
|
|
|
status, result = connect.get_index_info(ip_collection) |
1197
|
|
|
logging.getLogger().info(result) |
1198
|
|
|
assert result._collection_name == ip_collection |
1199
|
|
|
assert result._index_type == IndexType.FLAT |
1200
|
|
|
|
1201
|
|
|
def test_create_drop_index_repeatly_different_index_params(self, connect, ip_collection): |
1202
|
|
|
''' |
1203
|
|
|
target: test create / drop index repeatly, use the different index params |
1204
|
|
|
method: create index, drop index, four times, each tme use different index_params to create index |
1205
|
|
|
expected: return code 0 |
1206
|
|
|
''' |
1207
|
|
|
nlist = NLIST |
1208
|
|
|
indexs = [{"index_type": IndexType.IVFLAT, "index_param": {"nlist": nlist}}, {"index_type": IndexType.IVF_SQ8, "index_param": {"nlist": nlist}}] |
1209
|
|
|
status, ids = connect.insert(ip_collection, vectors) |
1210
|
|
|
for i in range(2): |
1211
|
|
|
status = connect.create_index(ip_collection, indexs[i]["index_type"], indexs[i]["index_param"]) |
1212
|
|
|
assert status.OK() |
1213
|
|
|
status, result = connect.get_index_info(ip_collection) |
1214
|
|
|
assert result._params == indexs[i]["index_param"] |
1215
|
|
|
assert result._collection_name == ip_collection |
1216
|
|
|
assert result._index_type == indexs[i]["index_type"] |
1217
|
|
|
status, result = connect.get_index_info(ip_collection) |
1218
|
|
|
logging.getLogger().info(result) |
1219
|
|
|
status = connect.drop_index(ip_collection) |
1220
|
|
|
assert status.OK() |
1221
|
|
|
status, result = connect.get_index_info(ip_collection) |
1222
|
|
|
logging.getLogger().info(result) |
1223
|
|
|
assert result._collection_name == ip_collection |
1224
|
|
|
assert result._index_type == IndexType.FLAT |
1225
|
|
|
|
1226
|
|
|
|
1227
|
|
|
class TestIndexJAC: |
1228
|
|
|
tmp, vectors = gen_binary_vectors(nb, dim) |
1229
|
|
|
|
1230
|
|
|
@pytest.fixture( |
1231
|
|
|
scope="function", |
1232
|
|
|
params=gen_index() |
1233
|
|
|
) |
1234
|
|
|
def get_index(self, request, connect): |
1235
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
1236
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
1237
|
|
|
pytest.skip("sq8h not support in CPU mode") |
1238
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
1239
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
1240
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
1241
|
|
|
return request.param |
1242
|
|
|
|
1243
|
|
View Code Duplication |
@pytest.fixture( |
|
|
|
|
1244
|
|
|
scope="function", |
1245
|
|
|
params=gen_simple_index() |
1246
|
|
|
) |
1247
|
|
|
def get_simple_index(self, request, connect): |
1248
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
1249
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
1250
|
|
|
pytest.skip("sq8h not support in CPU mode") |
1251
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
1252
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
1253
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
1254
|
|
|
return request.param |
1255
|
|
|
|
1256
|
|
View Code Duplication |
@pytest.fixture( |
|
|
|
|
1257
|
|
|
scope="function", |
1258
|
|
|
params=gen_simple_index() |
1259
|
|
|
) |
1260
|
|
|
def get_jaccard_index(self, request, connect): |
1261
|
|
|
logging.getLogger().info(request.param) |
1262
|
|
|
if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT: |
1263
|
|
|
return request.param |
1264
|
|
|
else: |
1265
|
|
|
pytest.skip("Skip index Temporary") |
1266
|
|
|
|
1267
|
|
|
""" |
1268
|
|
|
****************************************************************** |
1269
|
|
|
The following cases are used to test `create_index` function |
1270
|
|
|
****************************************************************** |
1271
|
|
|
""" |
1272
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
1273
|
|
|
def test_create_index(self, connect, jac_collection, get_jaccard_index): |
1274
|
|
|
''' |
1275
|
|
|
target: test create index interface |
1276
|
|
|
method: create collection and add vectors in it, create index |
1277
|
|
|
expected: return code equals to 0, and search success |
1278
|
|
|
''' |
1279
|
|
|
index_param = get_jaccard_index["index_param"] |
1280
|
|
|
index_type = get_jaccard_index["index_type"] |
1281
|
|
|
logging.getLogger().info(get_jaccard_index) |
1282
|
|
|
status, ids = connect.insert(jac_collection, self.vectors) |
1283
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1284
|
|
|
if index_type != IndexType.FLAT and index_type != IndexType.IVFLAT: |
1285
|
|
|
assert not status.OK() |
1286
|
|
|
else: |
1287
|
|
|
assert status.OK() |
1288
|
|
|
|
1289
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
1290
|
|
|
def test_create_index_partition(self, connect, jac_collection, get_jaccard_index): |
1291
|
|
|
''' |
1292
|
|
|
target: test create index interface |
1293
|
|
|
method: create collection, create partition, and add vectors in it, create index |
1294
|
|
|
expected: return code equals to 0, and search success |
1295
|
|
|
''' |
1296
|
|
|
index_param = get_jaccard_index["index_param"] |
1297
|
|
|
index_type = get_jaccard_index["index_type"] |
1298
|
|
|
logging.getLogger().info(get_jaccard_index) |
1299
|
|
|
status = connect.create_partition(jac_collection, tag) |
1300
|
|
|
status, ids = connect.insert(jac_collection, self.vectors, partition_tag=tag) |
1301
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1302
|
|
|
assert status.OK() |
1303
|
|
|
|
1304
|
|
|
# @pytest.mark.level(2) |
1305
|
|
|
# def test_create_index_without_connect(self, dis_connect, jac_collection): |
1306
|
|
|
# ''' |
1307
|
|
|
# target: test create index without connection |
1308
|
|
|
# method: create collection and add vectors in it, check if added successfully |
1309
|
|
|
# expected: raise exception |
1310
|
|
|
# ''' |
1311
|
|
|
# nlist = NLIST |
1312
|
|
|
# index_param = {"nlist": nlist} |
1313
|
|
|
# with pytest.raises(Exception) as e: |
1314
|
|
|
# status = dis_connect.create_index(jac_collection, IndexType.IVF_SQ8, index_param) |
1315
|
|
|
|
1316
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
1317
|
|
|
def test_create_index_search_with_query_vectors(self, connect, jac_collection, get_jaccard_index): |
1318
|
|
|
''' |
1319
|
|
|
target: test create index interface, search with more query vectors |
1320
|
|
|
method: create collection and add vectors in it, create index |
1321
|
|
|
expected: return code equals to 0, and search success |
1322
|
|
|
''' |
1323
|
|
|
index_param = get_jaccard_index["index_param"] |
1324
|
|
|
index_type = get_jaccard_index["index_type"] |
1325
|
|
|
logging.getLogger().info(get_jaccard_index) |
1326
|
|
|
status, ids = connect.insert(jac_collection, self.vectors) |
1327
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1328
|
|
|
logging.getLogger().info(connect.get_index_info(jac_collection)) |
1329
|
|
|
query_vecs = [self.vectors[0], self.vectors[1], self.vectors[2]] |
1330
|
|
|
top_k = 5 |
1331
|
|
|
search_param = get_search_param(index_type) |
1332
|
|
|
status, result = connect.search(jac_collection, top_k, query_vecs, params=search_param) |
1333
|
|
|
logging.getLogger().info(result) |
1334
|
|
|
assert status.OK() |
1335
|
|
|
assert len(result) == len(query_vecs) |
1336
|
|
|
|
1337
|
|
|
""" |
1338
|
|
|
****************************************************************** |
1339
|
|
|
The following cases are used to test `get_index_info` function |
1340
|
|
|
****************************************************************** |
1341
|
|
|
""" |
1342
|
|
|
|
1343
|
|
|
def test_get_index_info(self, connect, jac_collection, get_jaccard_index): |
1344
|
|
|
''' |
1345
|
|
|
target: test describe index interface |
1346
|
|
|
method: create collection and add vectors in it, create index, call describe index |
1347
|
|
|
expected: return code 0, and index instructure |
1348
|
|
|
''' |
1349
|
|
|
index_param = get_jaccard_index["index_param"] |
1350
|
|
|
index_type = get_jaccard_index["index_type"] |
1351
|
|
|
logging.getLogger().info(get_jaccard_index) |
1352
|
|
|
# status, ids = connect.insert(jac_collection, vectors[:5000]) |
1353
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1354
|
|
|
status, result = connect.get_index_info(jac_collection) |
1355
|
|
|
logging.getLogger().info(result) |
1356
|
|
|
assert result._collection_name == jac_collection |
1357
|
|
|
assert result._index_type == index_type |
1358
|
|
|
assert result._params == index_param |
1359
|
|
|
|
1360
|
|
View Code Duplication |
def test_get_index_info_partition(self, connect, jac_collection, get_jaccard_index): |
|
|
|
|
1361
|
|
|
''' |
1362
|
|
|
target: test describe index interface |
1363
|
|
|
method: create collection, create partition and add vectors in it, create index, call describe index |
1364
|
|
|
expected: return code 0, and index instructure |
1365
|
|
|
''' |
1366
|
|
|
index_param = get_jaccard_index["index_param"] |
1367
|
|
|
index_type = get_jaccard_index["index_type"] |
1368
|
|
|
logging.getLogger().info(get_jaccard_index) |
1369
|
|
|
status = connect.create_partition(jac_collection, tag) |
1370
|
|
|
status, ids = connect.insert(jac_collection, vectors, partition_tag=tag) |
1371
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1372
|
|
|
status, result = connect.get_index_info(jac_collection) |
1373
|
|
|
logging.getLogger().info(result) |
1374
|
|
|
assert result._params == index_param |
1375
|
|
|
assert result._collection_name == jac_collection |
1376
|
|
|
assert result._index_type == index_type |
1377
|
|
|
|
1378
|
|
|
""" |
1379
|
|
|
****************************************************************** |
1380
|
|
|
The following cases are used to test `drop_index` function |
1381
|
|
|
****************************************************************** |
1382
|
|
|
""" |
1383
|
|
|
|
1384
|
|
View Code Duplication |
def test_drop_index(self, connect, jac_collection, get_jaccard_index): |
|
|
|
|
1385
|
|
|
''' |
1386
|
|
|
target: test drop index interface |
1387
|
|
|
method: create collection and add vectors in it, create index, call drop index |
1388
|
|
|
expected: return code 0, and default index param |
1389
|
|
|
''' |
1390
|
|
|
index_param = get_jaccard_index["index_param"] |
1391
|
|
|
index_type = get_jaccard_index["index_type"] |
1392
|
|
|
status, mode = connect._cmd("mode") |
1393
|
|
|
assert status.OK() |
1394
|
|
|
# status, ids = connect.insert(ip_collection, vectors) |
1395
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1396
|
|
|
assert status.OK() |
1397
|
|
|
status, result = connect.get_index_info(jac_collection) |
1398
|
|
|
logging.getLogger().info(result) |
1399
|
|
|
status = connect.drop_index(jac_collection) |
1400
|
|
|
assert status.OK() |
1401
|
|
|
status, result = connect.get_index_info(jac_collection) |
1402
|
|
|
logging.getLogger().info(result) |
1403
|
|
|
assert result._collection_name == jac_collection |
1404
|
|
|
assert result._index_type == IndexType.FLAT |
1405
|
|
|
|
1406
|
|
|
def test_drop_index_partition(self, connect, jac_collection, get_jaccard_index): |
1407
|
|
|
''' |
1408
|
|
|
target: test drop index interface |
1409
|
|
|
method: create collection, create partition and add vectors in it, create index on collection, call drop collection index |
1410
|
|
|
expected: return code 0, and default index param |
1411
|
|
|
''' |
1412
|
|
|
index_param = get_jaccard_index["index_param"] |
1413
|
|
|
index_type = get_jaccard_index["index_type"] |
1414
|
|
|
status = connect.create_partition(jac_collection, tag) |
1415
|
|
|
status, ids = connect.insert(jac_collection, vectors, partition_tag=tag) |
1416
|
|
|
status = connect.create_index(jac_collection, index_type, index_param) |
1417
|
|
|
assert status.OK() |
1418
|
|
|
status, result = connect.get_index_info(jac_collection) |
1419
|
|
|
logging.getLogger().info(result) |
1420
|
|
|
status = connect.drop_index(jac_collection) |
1421
|
|
|
assert status.OK() |
1422
|
|
|
status, result = connect.get_index_info(jac_collection) |
1423
|
|
|
logging.getLogger().info(result) |
1424
|
|
|
assert result._collection_name == jac_collection |
1425
|
|
|
assert result._index_type == IndexType.FLAT |
1426
|
|
|
|
1427
|
|
|
|
1428
|
|
|
class TestIndexBinary: |
1429
|
|
|
tmp, vectors = gen_binary_vectors(nb, dim) |
1430
|
|
|
|
1431
|
|
|
@pytest.fixture( |
1432
|
|
|
scope="function", |
1433
|
|
|
params=gen_index() |
1434
|
|
|
) |
1435
|
|
|
def get_index(self, request, connect): |
1436
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
1437
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
1438
|
|
|
pytest.skip("sq8h not support in CPU mode") |
1439
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ or request.param["index_type"] == IndexType.HNSW: |
1440
|
|
|
pytest.skip("Skip PQ Temporary") |
1441
|
|
|
return request.param |
1442
|
|
|
|
1443
|
|
View Code Duplication |
@pytest.fixture( |
|
|
|
|
1444
|
|
|
scope="function", |
1445
|
|
|
params=gen_simple_index() |
1446
|
|
|
) |
1447
|
|
|
def get_simple_index(self, request, connect): |
1448
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
1449
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
1450
|
|
|
pytest.skip("sq8h not support in CPU mode") |
1451
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ or request.param["index_type"] == IndexType.HNSW: |
1452
|
|
|
pytest.skip("Skip PQ Temporary") |
1453
|
|
|
return request.param |
1454
|
|
|
|
1455
|
|
View Code Duplication |
@pytest.fixture( |
|
|
|
|
1456
|
|
|
scope="function", |
1457
|
|
|
params=gen_simple_index() |
1458
|
|
|
) |
1459
|
|
|
def get_hamming_index(self, request, connect): |
1460
|
|
|
logging.getLogger().info(request.param) |
1461
|
|
|
if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT: |
1462
|
|
|
return request.param |
1463
|
|
|
else: |
1464
|
|
|
pytest.skip("Skip index Temporary") |
1465
|
|
|
|
1466
|
|
|
@pytest.fixture( |
1467
|
|
|
scope="function", |
1468
|
|
|
params=gen_simple_index() |
1469
|
|
|
) |
1470
|
|
|
def get_substructure_index(self, request, connect): |
1471
|
|
|
logging.getLogger().info(request.param) |
1472
|
|
|
if request.param["index_type"] == IndexType.FLAT: |
1473
|
|
|
return request.param |
1474
|
|
|
else: |
1475
|
|
|
pytest.skip("Skip index Temporary") |
1476
|
|
|
|
1477
|
|
|
@pytest.fixture( |
1478
|
|
|
scope="function", |
1479
|
|
|
params=gen_simple_index() |
1480
|
|
|
) |
1481
|
|
|
def get_superstructure_index(self, request, connect): |
1482
|
|
|
logging.getLogger().info(request.param) |
1483
|
|
|
if request.param["index_type"] == IndexType.FLAT: |
1484
|
|
|
return request.param |
1485
|
|
|
else: |
1486
|
|
|
pytest.skip("Skip index Temporary") |
1487
|
|
|
|
1488
|
|
|
""" |
1489
|
|
|
****************************************************************** |
1490
|
|
|
The following cases are used to test `create_index` function |
1491
|
|
|
****************************************************************** |
1492
|
|
|
""" |
1493
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
1494
|
|
|
def test_create_index(self, connect, ham_collection, get_hamming_index): |
1495
|
|
|
''' |
1496
|
|
|
target: test create index interface |
1497
|
|
|
method: create collection and add vectors in it, create index |
1498
|
|
|
expected: return code equals to 0, and search success |
1499
|
|
|
''' |
1500
|
|
|
index_param = get_hamming_index["index_param"] |
1501
|
|
|
index_type = get_hamming_index["index_type"] |
1502
|
|
|
logging.getLogger().info(get_hamming_index) |
1503
|
|
|
status, ids = connect.insert(ham_collection, self.vectors) |
1504
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1505
|
|
|
if index_type != IndexType.FLAT and index_type != IndexType.IVFLAT: |
1506
|
|
|
assert not status.OK() |
1507
|
|
|
else: |
1508
|
|
|
assert status.OK() |
1509
|
|
|
|
1510
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
1511
|
|
|
def test_create_index_partition(self, connect, ham_collection, get_hamming_index): |
1512
|
|
|
''' |
1513
|
|
|
target: test create index interface |
1514
|
|
|
method: create collection, create partition, and add vectors in it, create index |
1515
|
|
|
expected: return code equals to 0, and search success |
1516
|
|
|
''' |
1517
|
|
|
index_param = get_hamming_index["index_param"] |
1518
|
|
|
index_type = get_hamming_index["index_type"] |
1519
|
|
|
logging.getLogger().info(get_hamming_index) |
1520
|
|
|
status = connect.create_partition(ham_collection, tag) |
1521
|
|
|
status, ids = connect.insert(ham_collection, self.vectors, partition_tag=tag) |
1522
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1523
|
|
|
assert status.OK() |
1524
|
|
|
status, res = connect.count_entities(ham_collection) |
1525
|
|
|
assert res == len(self.vectors) |
1526
|
|
|
|
1527
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
1528
|
|
|
def test_create_index_partition_structure(self, connect, substructure_collection, get_substructure_index): |
1529
|
|
|
''' |
1530
|
|
|
target: test create index interface |
1531
|
|
|
method: create collection, create partition, and add vectors in it, create index |
1532
|
|
|
expected: return code equals to 0, and search success |
1533
|
|
|
''' |
1534
|
|
|
index_param = get_substructure_index["index_param"] |
1535
|
|
|
index_type = get_substructure_index["index_type"] |
1536
|
|
|
logging.getLogger().info(get_substructure_index) |
1537
|
|
|
status = connect.create_partition(substructure_collection, tag) |
1538
|
|
|
status, ids = connect.insert(substructure_collection, self.vectors, partition_tag=tag) |
1539
|
|
|
status = connect.create_index(substructure_collection, index_type, index_param) |
1540
|
|
|
assert status.OK() |
1541
|
|
|
status, res = connect.count_entities(substructure_collection,) |
1542
|
|
|
assert res == len(self.vectors) |
1543
|
|
|
|
1544
|
|
|
# @pytest.mark.level(2) |
1545
|
|
|
# def test_create_index_without_connect(self, dis_connect, ham_collection): |
1546
|
|
|
# ''' |
1547
|
|
|
# target: test create index without connection |
1548
|
|
|
# method: create collection and add vectors in it, check if added successfully |
1549
|
|
|
# expected: raise exception |
1550
|
|
|
# ''' |
1551
|
|
|
# nlist = NLIST |
1552
|
|
|
# index_param = {"nlist": nlist} |
1553
|
|
|
# with pytest.raises(Exception) as e: |
1554
|
|
|
# status = dis_connect.create_index(ham_collection, IndexType.IVF_SQ8, index_param) |
1555
|
|
|
|
1556
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
1557
|
|
|
def test_create_index_search_with_query_vectors(self, connect, ham_collection, get_hamming_index): |
1558
|
|
|
''' |
1559
|
|
|
target: test create index interface, search with more query vectors |
1560
|
|
|
method: create collection and add vectors in it, create index |
1561
|
|
|
expected: return code equals to 0, and search success |
1562
|
|
|
''' |
1563
|
|
|
index_param = get_hamming_index["index_param"] |
1564
|
|
|
index_type = get_hamming_index["index_type"] |
1565
|
|
|
logging.getLogger().info(get_hamming_index) |
1566
|
|
|
status, ids = connect.insert(ham_collection, self.vectors) |
1567
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1568
|
|
|
logging.getLogger().info(connect.get_index_info(ham_collection)) |
1569
|
|
|
query_vecs = [self.vectors[0], self.vectors[1], self.vectors[2]] |
1570
|
|
|
top_k = 5 |
1571
|
|
|
search_param = get_search_param(index_type) |
1572
|
|
|
status, result = connect.search(ham_collection, top_k, query_vecs, params=search_param) |
1573
|
|
|
logging.getLogger().info(result) |
1574
|
|
|
assert status.OK() |
1575
|
|
|
assert len(result) == len(query_vecs) |
1576
|
|
|
|
1577
|
|
View Code Duplication |
@pytest.mark.timeout(BUILD_TIMEOUT) |
|
|
|
|
1578
|
|
|
def test_create_index_search_with_query_vectors_superstructure(self, connect, superstructure_collection, get_superstructure_index): |
1579
|
|
|
''' |
1580
|
|
|
target: test create index interface, search with more query vectors |
1581
|
|
|
method: create collection and add vectors in it, create index |
1582
|
|
|
expected: return code equals to 0, and search success |
1583
|
|
|
''' |
1584
|
|
|
index_param = get_superstructure_index["index_param"] |
1585
|
|
|
index_type = get_superstructure_index["index_type"] |
1586
|
|
|
logging.getLogger().info(get_superstructure_index) |
1587
|
|
|
status, ids = connect.insert(superstructure_collection, self.vectors) |
1588
|
|
|
status = connect.create_index(superstructure_collection, index_type, index_param) |
1589
|
|
|
logging.getLogger().info(connect.get_index_info(superstructure_collection)) |
1590
|
|
|
query_vecs = [self.vectors[0], self.vectors[1], self.vectors[2]] |
1591
|
|
|
top_k = 5 |
1592
|
|
|
search_param = get_search_param(index_type) |
1593
|
|
|
status, result = connect.search(superstructure_collection, top_k, query_vecs, params=search_param) |
1594
|
|
|
logging.getLogger().info(result) |
1595
|
|
|
assert status.OK() |
1596
|
|
|
assert len(result) == len(query_vecs) |
1597
|
|
|
|
1598
|
|
|
""" |
1599
|
|
|
****************************************************************** |
1600
|
|
|
The following cases are used to test `get_index_info` function |
1601
|
|
|
****************************************************************** |
1602
|
|
|
""" |
1603
|
|
|
|
1604
|
|
|
def test_get_index_info(self, connect, ham_collection, get_hamming_index): |
1605
|
|
|
''' |
1606
|
|
|
target: test describe index interface |
1607
|
|
|
method: create collection and add vectors in it, create index, call describe index |
1608
|
|
|
expected: return code 0, and index instructure |
1609
|
|
|
''' |
1610
|
|
|
index_param = get_hamming_index["index_param"] |
1611
|
|
|
index_type = get_hamming_index["index_type"] |
1612
|
|
|
logging.getLogger().info(get_hamming_index) |
1613
|
|
|
# status, ids = connect.insert(jac_collection, vectors[:5000]) |
1614
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1615
|
|
|
status, result = connect.get_index_info(ham_collection) |
1616
|
|
|
logging.getLogger().info(result) |
1617
|
|
|
assert result._collection_name == ham_collection |
1618
|
|
|
assert result._index_type == index_type |
1619
|
|
|
assert result._params == index_param |
1620
|
|
|
|
1621
|
|
View Code Duplication |
def test_get_index_info_partition(self, connect, ham_collection, get_hamming_index): |
|
|
|
|
1622
|
|
|
''' |
1623
|
|
|
target: test describe index interface |
1624
|
|
|
method: create collection, create partition and add vectors in it, create index, call describe index |
1625
|
|
|
expected: return code 0, and index instructure |
1626
|
|
|
''' |
1627
|
|
|
index_param = get_hamming_index["index_param"] |
1628
|
|
|
index_type = get_hamming_index["index_type"] |
1629
|
|
|
logging.getLogger().info(get_hamming_index) |
1630
|
|
|
status = connect.create_partition(ham_collection, tag) |
1631
|
|
|
status, ids = connect.insert(ham_collection, vectors, partition_tag=tag) |
1632
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1633
|
|
|
status, result = connect.get_index_info(ham_collection) |
1634
|
|
|
logging.getLogger().info(result) |
1635
|
|
|
assert result._params == index_param |
1636
|
|
|
assert result._collection_name == ham_collection |
1637
|
|
|
assert result._index_type == index_type |
1638
|
|
|
|
1639
|
|
View Code Duplication |
def test_get_index_info_partition_superstructrue(self, connect, superstructure_collection, get_superstructure_index): |
|
|
|
|
1640
|
|
|
''' |
1641
|
|
|
target: test describe index interface |
1642
|
|
|
method: create collection, create partition and add vectors in it, create index, call describe index |
1643
|
|
|
expected: return code 0, and index instructure |
1644
|
|
|
''' |
1645
|
|
|
index_param = get_superstructure_index["index_param"] |
1646
|
|
|
index_type = get_superstructure_index["index_type"] |
1647
|
|
|
logging.getLogger().info(get_superstructure_index) |
1648
|
|
|
status = connect.create_partition(superstructure_collection, tag) |
1649
|
|
|
status, ids = connect.insert(superstructure_collection, vectors, partition_tag=tag) |
1650
|
|
|
status = connect.create_index(superstructure_collection, index_type, index_param) |
1651
|
|
|
status, result = connect.get_index_info(superstructure_collection) |
1652
|
|
|
logging.getLogger().info(result) |
1653
|
|
|
assert result._params == index_param |
1654
|
|
|
assert result._collection_name == superstructure_collection |
1655
|
|
|
assert result._index_type == index_type |
1656
|
|
|
|
1657
|
|
|
""" |
1658
|
|
|
****************************************************************** |
1659
|
|
|
The following cases are used to test `drop_index` function |
1660
|
|
|
****************************************************************** |
1661
|
|
|
""" |
1662
|
|
|
|
1663
|
|
View Code Duplication |
def test_drop_index(self, connect, ham_collection, get_hamming_index): |
|
|
|
|
1664
|
|
|
''' |
1665
|
|
|
target: test drop index interface |
1666
|
|
|
method: create collection and add vectors in it, create index, call drop index |
1667
|
|
|
expected: return code 0, and default index param |
1668
|
|
|
''' |
1669
|
|
|
index_param = get_hamming_index["index_param"] |
1670
|
|
|
index_type = get_hamming_index["index_type"] |
1671
|
|
|
status, mode = connect._cmd("mode") |
1672
|
|
|
assert status.OK() |
1673
|
|
|
# status, ids = connect.insert(ip_collection, vectors) |
1674
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1675
|
|
|
assert status.OK() |
1676
|
|
|
status, result = connect.get_index_info(ham_collection) |
1677
|
|
|
logging.getLogger().info(result) |
1678
|
|
|
status = connect.drop_index(ham_collection) |
1679
|
|
|
assert status.OK() |
1680
|
|
|
status, result = connect.get_index_info(ham_collection) |
1681
|
|
|
logging.getLogger().info(result) |
1682
|
|
|
assert result._collection_name == ham_collection |
1683
|
|
|
assert result._index_type == IndexType.FLAT |
1684
|
|
|
|
1685
|
|
View Code Duplication |
def test_drop_index_substructure(self, connect, substructure_collection, get_substructure_index): |
|
|
|
|
1686
|
|
|
''' |
1687
|
|
|
target: test drop index interface |
1688
|
|
|
method: create collection and add vectors in it, create index, call drop index |
1689
|
|
|
expected: return code 0, and default index param |
1690
|
|
|
''' |
1691
|
|
|
index_param = get_substructure_index["index_param"] |
1692
|
|
|
index_type = get_substructure_index["index_type"] |
1693
|
|
|
status, mode = connect._cmd("mode") |
1694
|
|
|
assert status.OK() |
1695
|
|
|
status = connect.create_index(substructure_collection, index_type, index_param) |
1696
|
|
|
assert status.OK() |
1697
|
|
|
status, result = connect.get_index_info(substructure_collection) |
1698
|
|
|
logging.getLogger().info(result) |
1699
|
|
|
status = connect.drop_index(substructure_collection) |
1700
|
|
|
assert status.OK() |
1701
|
|
|
status, result = connect.get_index_info(substructure_collection) |
1702
|
|
|
logging.getLogger().info(result) |
1703
|
|
|
assert result._collection_name == substructure_collection |
1704
|
|
|
assert result._index_type == IndexType.FLAT |
1705
|
|
|
|
1706
|
|
|
def test_drop_index_partition(self, connect, ham_collection, get_hamming_index): |
1707
|
|
|
''' |
1708
|
|
|
target: test drop index interface |
1709
|
|
|
method: create collection, create partition and add vectors in it, create index on collection, call drop collection index |
1710
|
|
|
expected: return code 0, and default index param |
1711
|
|
|
''' |
1712
|
|
|
index_param = get_hamming_index["index_param"] |
1713
|
|
|
index_type = get_hamming_index["index_type"] |
1714
|
|
|
status = connect.create_partition(ham_collection, tag) |
1715
|
|
|
status, ids = connect.insert(ham_collection, vectors, partition_tag=tag) |
1716
|
|
|
status = connect.create_index(ham_collection, index_type, index_param) |
1717
|
|
|
assert status.OK() |
1718
|
|
|
status, result = connect.get_index_info(ham_collection) |
1719
|
|
|
logging.getLogger().info(result) |
1720
|
|
|
status = connect.drop_index(ham_collection) |
1721
|
|
|
assert status.OK() |
1722
|
|
|
status, result = connect.get_index_info(ham_collection) |
1723
|
|
|
logging.getLogger().info(result) |
1724
|
|
|
assert result._collection_name == ham_collection |
1725
|
|
|
assert result._index_type == IndexType.FLAT |
1726
|
|
|
|
1727
|
|
|
class TestIndexCollectionInvalid(object): |
1728
|
|
|
""" |
1729
|
|
|
Test create / describe / drop index interfaces with invalid collection names |
1730
|
|
|
""" |
1731
|
|
|
@pytest.fixture( |
1732
|
|
|
scope="function", |
1733
|
|
|
params=gen_invalid_collection_names() |
1734
|
|
|
) |
1735
|
|
|
def get_collection_name(self, request): |
1736
|
|
|
yield request.param |
1737
|
|
|
|
1738
|
|
|
@pytest.mark.level(1) |
1739
|
|
|
def test_create_index_with_invalid_collectionname(self, connect, get_collection_name): |
1740
|
|
|
collection_name = get_collection_name |
1741
|
|
|
nlist = NLIST |
1742
|
|
|
index_param = {"nlist": nlist} |
1743
|
|
|
status = connect.create_index(collection_name, IndexType.IVF_SQ8, index_param) |
1744
|
|
|
assert not status.OK() |
1745
|
|
|
|
1746
|
|
|
@pytest.mark.level(1) |
1747
|
|
|
def test_get_index_info_with_invalid_collectionname(self, connect, get_collection_name): |
1748
|
|
|
collection_name = get_collection_name |
1749
|
|
|
status, result = connect.get_index_info(collection_name) |
1750
|
|
|
assert not status.OK() |
1751
|
|
|
|
1752
|
|
|
@pytest.mark.level(1) |
1753
|
|
|
def test_drop_index_with_invalid_collectionname(self, connect, get_collection_name): |
1754
|
|
|
collection_name = get_collection_name |
1755
|
|
|
status = connect.drop_index(collection_name) |
1756
|
|
|
assert not status.OK() |
1757
|
|
|
|
1758
|
|
|
|
1759
|
|
|
class TestCreateIndexParamsInvalid(object): |
1760
|
|
|
""" |
1761
|
|
|
Test Building index with invalid collection names, collection names not in db |
1762
|
|
|
""" |
1763
|
|
|
@pytest.fixture( |
1764
|
|
|
scope="function", |
1765
|
|
|
params=gen_invalid_index() |
1766
|
|
|
) |
1767
|
|
|
def get_index(self, request): |
1768
|
|
|
yield request.param |
1769
|
|
|
|
1770
|
|
|
@pytest.mark.level(1) |
1771
|
|
|
def test_create_index_with_invalid_index_params(self, connect, collection, get_index): |
1772
|
|
|
index_param = get_index["index_param"] |
1773
|
|
|
index_type = get_index["index_type"] |
1774
|
|
|
logging.getLogger().info(get_index) |
1775
|
|
|
# status, ids = connect.insert(collection, vectors) |
1776
|
|
|
if (not index_type) or (not isinstance(index_type, IndexType)): |
1777
|
|
|
with pytest.raises(Exception) as e: |
1778
|
|
|
status = connect.create_index(collection, index_type, index_param) |
1779
|
|
|
else: |
1780
|
|
|
status = connect.create_index(collection, index_type, index_param) |
1781
|
|
|
assert not status.OK() |
1782
|
|
|
|
1783
|
|
|
""" |
1784
|
|
|
Test Building index with invalid nlist |
1785
|
|
|
""" |
1786
|
|
|
@pytest.fixture( |
1787
|
|
|
scope="function", |
1788
|
|
|
params=[IndexType.FLAT,IndexType.IVFLAT,IndexType.IVF_SQ8,IndexType.IVF_SQ8H] |
1789
|
|
|
) |
1790
|
|
|
def get_index_type(self, request): |
1791
|
|
|
yield request.param |
1792
|
|
|
|
1793
|
|
|
def test_create_index_with_invalid_nlist(self, connect, collection, get_index_type): |
1794
|
|
|
status, ids = connect.insert(collection, vectors) |
1795
|
|
|
status = connect.create_index(collection, get_index_type, {"nlist": INVALID_NLIST}) |
1796
|
|
|
if get_index_type != IndexType.FLAT: |
1797
|
|
|
assert not status.OK() |
1798
|
|
|
|
1799
|
|
|
''' |
1800
|
|
|
Test Building index with empty params |
1801
|
|
|
''' |
1802
|
|
|
def test_create_index_with_empty_param(self, connect, collection, get_index_type): |
1803
|
|
|
logging.getLogger().info(get_index_type) |
1804
|
|
|
status = connect.create_index(collection, get_index_type, {}) |
1805
|
|
|
if get_index_type != IndexType.FLAT : |
1806
|
|
|
assert not status.OK() |
1807
|
|
|
status, result = connect.get_index_info(collection) |
1808
|
|
|
logging.getLogger().info(result) |
1809
|
|
|
assert result._collection_name == collection |
1810
|
|
|
assert result._index_type == IndexType.FLAT |
1811
|
|
|
|
1812
|
|
|
class TestIndexAsync: |
1813
|
|
|
@pytest.fixture(scope="function", autouse=True) |
1814
|
|
|
def skip_http_check(self, args): |
1815
|
|
|
if args["handler"] == "HTTP": |
1816
|
|
|
pytest.skip("skip in http mode") |
1817
|
|
|
|
1818
|
|
|
""" |
1819
|
|
|
****************************************************************** |
1820
|
|
|
The following cases are used to test `create_index` function |
1821
|
|
|
****************************************************************** |
1822
|
|
|
""" |
1823
|
|
|
@pytest.fixture( |
1824
|
|
|
scope="function", |
1825
|
|
|
params=gen_index() |
1826
|
|
|
) |
1827
|
|
|
def get_index(self, request, connect): |
1828
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
1829
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
1830
|
|
|
pytest.skip("sq8h not support in CPU mode") |
1831
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
1832
|
|
|
if request.param["index_type"] == IndexType.IVF_PQ: |
1833
|
|
|
pytest.skip("ivfpq not support in GPU mode") |
1834
|
|
|
return request.param |
1835
|
|
|
|
1836
|
|
View Code Duplication |
@pytest.fixture( |
|
|
|
|
1837
|
|
|
scope="function", |
1838
|
|
|
params=gen_simple_index() |
1839
|
|
|
) |
1840
|
|
|
def get_simple_index(self, request, connect): |
1841
|
|
|
if str(connect._cmd("mode")[1]) == "CPU": |
1842
|
|
|
if request.param["index_type"] == IndexType.IVF_SQ8H: |
1843
|
|
|
pytest.skip("sq8h not support in CPU mode") |
1844
|
|
|
if str(connect._cmd("mode")[1]) == "GPU": |
1845
|
|
|
# if request.param["index_type"] == IndexType.IVF_PQ: |
1846
|
|
|
if request.param["index_type"] not in [IndexType.IVF_FLAT]: |
1847
|
|
|
# pytest.skip("ivfpq not support in GPU mode") |
1848
|
|
|
pytest.skip("debug ivf_flat in GPU mode") |
1849
|
|
|
return request.param |
1850
|
|
|
|
1851
|
|
|
def check_status(self, status): |
1852
|
|
|
logging.getLogger().info("In callback check status") |
1853
|
|
|
assert status.OK() |
1854
|
|
|
|
1855
|
|
|
""" |
1856
|
|
|
****************************************************************** |
1857
|
|
|
The following cases are used to test `create_index` function |
1858
|
|
|
****************************************************************** |
1859
|
|
|
""" |
1860
|
|
|
|
1861
|
|
|
@pytest.mark.timeout(BUILD_TIMEOUT) |
1862
|
|
|
def test_create_index(self, connect, collection, get_simple_index): |
1863
|
|
|
''' |
1864
|
|
|
target: test create index interface |
1865
|
|
|
method: create collection and add vectors in it, create index |
1866
|
|
|
expected: return code equals to 0, and search success |
1867
|
|
|
''' |
1868
|
|
|
index_param = get_simple_index["index_param"] |
1869
|
|
|
index_type = get_simple_index["index_type"] |
1870
|
|
|
logging.getLogger().info(get_simple_index) |
1871
|
|
|
vectors = gen_vectors(nb, dim) |
1872
|
|
|
status, ids = connect.insert(collection, vectors) |
1873
|
|
|
logging.getLogger().info("start index") |
1874
|
|
|
# future = connect.create_index(collection, index_type, index_param, _async=True, _callback=self.check_status) |
1875
|
|
|
future = connect.create_index(collection, index_type, index_param, _async=True) |
1876
|
|
|
logging.getLogger().info("before result") |
1877
|
|
|
status = future.result() |
1878
|
|
|
assert status.OK() |
1879
|
|
|
|
1880
|
|
|
def test_create_index_with_invalid_collectionname(self, connect): |
1881
|
|
|
collection_name = " " |
1882
|
|
|
nlist = NLIST |
1883
|
|
|
index_param = {"nlist": nlist} |
1884
|
|
|
future = connect.create_index(collection_name, IndexType.IVF_SQ8, index_param, _async=True) |
1885
|
|
|
status = future.result() |
1886
|
|
|
assert not status.OK() |
1887
|
|
|
|
1888
|
|
|
|