Total Complexity | 194 |
Total Lines | 1887 |
Duplicated Lines | 40.43 % |
Changes | 0 |
Duplicate code is one of the most pungent code smells. A rule that is often used is to re-structure code once it is duplicated in three or more places.
Common duplication problems, and corresponding solutions are:
Complex classes like test_index often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.
Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.
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 |