test_search_20.py 93.4 KB
Newer Older
1 2
import threading
import time
紫晴 已提交
3
import pytest
4 5 6
import random
import numpy as np

7
from base.client_base import TestcaseBase
紫晴 已提交
8
from utils.util_log import test_log as log
9 10
from common import common_func as cf
from common import common_type as ct
B
binbin 已提交
11
from common.common_type import CaseLabel, CheckTasks
紫晴 已提交
12

B
binbin 已提交
13
prefix = "search_collection"
14 15
search_num = 10
max_dim = ct.max_dim
16
epsilon = ct.epsilon
17
gracefulTime = ct.gracefulTime
18
default_nb = ct.default_nb
19
default_nb_medium = ct.default_nb_medium
20
default_nq = ct.default_nq
21
default_dim = ct.default_dim
22
default_limit = ct.default_limit
23 24
default_search_exp = "int64 >= 0"
default_search_field = ct.default_float_vec_field_name
25
default_search_params = ct.default_search_params
26 27
default_int64_field_name = ct.default_int64_field_name
default_float_field_name = ct.default_float_field_name
28
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
29

30
class TestCollectionSearchInvalid(TestcaseBase):
31
    """ Test case of search interface """
紫晴 已提交
32

33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
    @pytest.fixture(scope="function", params=ct.get_invalid_vectors)
    def get_invalid_vectors(self, request):
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_fields_type(self, request):
        if isinstance(request.param, str):
            pytest.skip("string is valid type for field")
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_fields_value(self, request):
        if not isinstance(request.param, str):
            pytest.skip("field value only support string")
        if request.param == "":
            pytest.skip("empty field is valid")
        yield request.param

51 52 53 54
    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_metric_type(self, request):
        yield request.param

55 56
    @pytest.fixture(scope="function", params=ct.get_invalid_ints)
    def get_invalid_limit(self, request):
57 58
        if isinstance(request.param, int) and request.param >= 0:
            pytest.skip("positive int is valid type for limit")
59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_expr_type(self, request):
        if isinstance(request.param, str):
            pytest.skip("string is valid type for expr")
        if request.param == None:
            pytest.skip("None is valid for expr")
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_expr_value(self, request):
        if not isinstance(request.param, str):
            pytest.skip("expression value only support string")
        if request.param == "":
            pytest.skip("empty field is valid")
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_partition(self, request):
        if request.param == []:
            pytest.skip("empty is valid for partition")
81 82
        if request.param == None:
            pytest.skip("None is valid for partition")
83 84 85 86
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_output_fields(self, request):
87 88 89 90
        if request.param == []:
            pytest.skip("empty is valid for output_fields")
        if request.param == None:
            pytest.skip("None is valid for output_fields")
91 92
        yield request.param

93

94 95
    """
    ******************************************************************
B
binbin 已提交
96
    #  The followings are invalid cases
97 98
    ******************************************************************
    """
99
    @pytest.mark.tags(CaseLabel.L1)
100
    def test_search_no_connection(self):
B
binbin 已提交
101
        """
102 103 104
        target: test search without connection
        method: create and delete connection, then search
        expected: raise exception and report the error
B
binbin 已提交
105
        """
106
        # 1. initialize with data
107 108
        collection_w = self.init_collection_general(prefix)[0]
        # 2. remove connection
109
        log.info("test_search_no_connection: removing connection")
110
        self.connection_wrap.remove_connection(alias='default')
111
        log.info("test_search_no_connection: removed connection")
112
        # 3. search without connection
113
        log.info("test_search_no_connection: searching without connection")
B
binbin 已提交
114 115 116 117
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit,
                            default_search_exp,
                            check_task=CheckTasks.err_res,
118 119
                            check_items={"err_code": 1,
                                         "err_msg": "should create connect first"})
120

121
    @pytest.mark.tags(CaseLabel.L1)
122
    def test_search_no_collection(self):
B
binbin 已提交
123
        """
124
        target: test the scenario which search the non-exist collection
B
binbin 已提交
125 126 127
        method: 1. create collection
                2. drop collection
                3. search the dropped collection
128
        expected: raise exception and report the error
B
binbin 已提交
129
        """
130
        # 1. initialize without data
B
binbin 已提交
131
        collection_w = self.init_collection_general(prefix)[0]
132
        # 2. Drop collection
B
binbin 已提交
133
        collection_w.drop()
134 135
        # 3. Search without collection
        log.info("test_search_no_collection: Searching without collection ")
B
binbin 已提交
136
        collection_w.search(vectors, default_search_field,
137
                            default_search_params, default_limit, default_search_exp,
B
binbin 已提交
138 139 140
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "collection %s doesn't exist!" % collection_w.name})
141

142
    @pytest.mark.tags(CaseLabel.L1)
143
    def test_search_param_missing(self):
B
binbin 已提交
144 145 146 147 148
        """
        target: test search with incomplete parameters
        method: search with incomplete parameters
        expected: raise exception and report the error
        """
149
        # 1. initialize without data
150
        collection_w = self.init_collection_general(prefix)[0]
151
        # 2. search collection with missing parameters
152 153
        log.info("test_search_param_missing: Searching collection %s "
                 "with missing parameters" % collection_w.name)
154
        try:
B
binbin 已提交
155
            collection_w.search()
156
        except TypeError as e:
157 158
            assert "missing 4 required positional arguments: 'data', " \
                   "'anns_field', 'param', and 'limit'" in str(e)
159

160 161 162 163
    @pytest.mark.tags(CaseLabel.L1)
    def test_search_param_invalid_vectors(self, get_invalid_vectors):
        """
        target: test search with invalid parameter values
164
        method: search with invalid data
165 166 167 168 169 170 171 172 173 174 175 176 177 178
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix)[0]
        # 2. search with invalid field
        invalid_vectors = get_invalid_vectors
        log.info("test_search_param_invalid_vectors: searching with "
                 "invalid vectors: {}".format(invalid_vectors))
        collection_w.search(invalid_vectors, default_search_field, default_search_params,
                            default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "`search_data` value {} is illegal".format(invalid_vectors)})

179
    @pytest.mark.tags(CaseLabel.L1)
180
    def test_search_param_invalid_dim(self):
B
binbin 已提交
181 182 183 184 185
        """
        target: test search with invalid parameter values
        method: search with invalid dim
        expected: raise exception and report the error
        """
186
        # 1. initialize with data
B
binbin 已提交
187
        collection_w = self.init_collection_general(prefix, True)[0]
188 189 190 191
        # 2. search with invalid dim
        log.info("test_search_param_invalid_dim: searching with invalid dim")
        wrong_dim = 129
        vectors = [[random.random() for _ in range(wrong_dim)] for _ in range(default_nq)]
B
binbin 已提交
192 193 194 195
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
196 197
                                         "err_msg": "The dimension of query entities "
                                                    "is different from schema"})
198

199
    @pytest.mark.tags(CaseLabel.L1)
200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220
    def test_search_param_invalid_field_type(self, get_invalid_fields_type):
        """
        target: test search with invalid parameter type
        method: search with invalid field
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix)[0]
        # 2. search with invalid field
        invalid_search_field = get_invalid_fields_type
        log.info("test_search_param_invalid_field_type: searching with "
                 "invalid field: %s" % invalid_search_field)
        collection_w.search(vectors[:default_nq], invalid_search_field, default_search_params,
                            default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items=
                            {"err_code": 1,
                             "err_msg": "`anns_field` value {} is illegal".format(invalid_search_field)})

    @pytest.mark.tags(CaseLabel.L1)
    def test_search_param_invalid_field_value(self, get_invalid_fields_value):
B
binbin 已提交
221 222
        """
        target: test search with invalid parameter values
223
        method: search with invalid field
B
binbin 已提交
224 225
        expected: raise exception and report the error
        """
226
        # 1. initialize with data
227
        collection_w = self.init_collection_general(prefix)[0]
228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247
        # 2. search with invalid field
        invalid_search_field = get_invalid_fields_value
        log.info("test_search_param_invalid_field_value: searching with "
                 "invalid field: %s" % invalid_search_field)
        collection_w.search(vectors[:default_nq], invalid_search_field, default_search_params,
                            default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "Field %s doesn't exist in schema"
                                                    % invalid_search_field})

    @pytest.mark.tags(CaseLabel.L1)
    def test_search_param_invalid_metric_type(self, get_invalid_metric_type):
        """
        target: test search with invalid parameter values
        method: search with invalid metric type
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix, True, 10)[0]
248
        # 2. search with invalid metric_type
249
        log.info("test_search_param_invalid_metric_type: searching with invalid metric_type")
250 251
        invalid_metric = get_invalid_metric_type
        search_params = {"metric_type": invalid_metric, "params": {"nprobe": 10}}
B
binbin 已提交
252 253 254 255 256
        collection_w.search(vectors[:default_nq], default_search_field, search_params,
                            default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "metric type not found"})
257

258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
    @pytest.mark.tags(CaseLabel.L2)
    @pytest.mark.xfail(reason="issue 6727")
    @pytest.mark.parametrize("index, params",
                             zip(ct.all_index_types[:9],
                                 ct.default_index_params[:9]))
    def test_search_invalid_params_type(self, index, params):
        """
        target: test search with invalid search params
        method: test search with invalid params type
        expected: raise exception and report the error
        """
        if index == "FLAT":
            pytest.skip("skip in FLAT index")
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, 5000,
                                                                      is_index=True)
        # 2. create index and load
        default_index = {"index_type": index, "params": params, "metric_type": "L2"}
        collection_w.create_index("float_vector", default_index)
        collection_w.load()
        # 3. search
        invalid_search_params = cf.gen_invaild_search_params_type()
        for invalid_search_param in invalid_search_params:
            if index == invalid_search_param["index_type"]:
                search_params = {"metric_type": "L2", "params": invalid_search_param["search_params"]}
                collection_w.search(vectors[:default_nq], default_search_field,
                                    search_params, default_limit,
                                    default_search_exp,
                                    check_task=CheckTasks.err_res,
                                    check_items={"err_code": 0,
                                                 "err_msg": "metric type not found"})

290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308
    @pytest.mark.tags(CaseLabel.L1)
    def test_search_param_invalid_limit_type(self, get_invalid_limit):
        """
        target: test search with invalid limit type
        method: search with invalid limit
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix)[0]
        # 2. search with invalid field
        invalid_limit = get_invalid_limit
        log.info("test_search_param_invalid_limit_type: searching with "
                 "invalid limit: %s" % invalid_limit)
        collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
                            invalid_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "`limit` value %s is illegal" % invalid_limit})

309
    @pytest.mark.tags(CaseLabel.L1)
310
    @pytest.mark.parametrize("limit", [0, 16385])
311
    def test_search_param_invalid_limit_value(self, limit):
B
binbin 已提交
312
        """
313
        target: test search with invalid limit value
B
binbin 已提交
314 315 316
        method: search with invalid limit: 0 and maximum
        expected: raise exception and report the error
        """
317
        # 1. initialize with data
318
        collection_w = self.init_collection_general(prefix)[0]
319 320
        # 2. search with invalid limit (topK)
        log.info("test_search_param_invalid_limit: searching with "
321
                 "invalid limit (topK) = %s" % limit)        
322 323 324
        err_msg = "limit %d is too large!" % limit
        if limit == 0:
            err_msg = "`limit` value 0 is illegal"
B
binbin 已提交
325 326 327 328
        collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
                            limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
329
                                         "err_msg": err_msg})
330

331
    @pytest.mark.tags(CaseLabel.L1)
332
    def test_search_param_invalid_expr_type(self, get_invalid_expr_type):
B
binbin 已提交
333
        """
334
        target: test search with invalid parameter type
B
binbin 已提交
335 336 337
        method: search with invalid search expressions
        expected: raise exception and report the error
        """
338
        # 1. initialize with data
339 340 341 342 343
        collection_w = self.init_collection_general(prefix)[0]
        # 2 search with invalid expr
        invalid_search_expr = get_invalid_expr_type
        log.info("test_search_param_invalid_expr_type: searching with "
                 "invalid expr: {}".format(invalid_search_expr))
344
        
345 346 347 348 349 350 351 352 353 354 355 356 357 358 359
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit, invalid_search_expr,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "The type of expr must be string ,"
                                                    "but {} is given".format(type(invalid_search_expr))})
    @pytest.mark.tags(CaseLabel.L1)
    def test_search_param_invalid_expr_value(self, get_invalid_expr_value):
        """
        target: test search with invalid parameter values
        method: search with invalid search expressions
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix)[0]
360
        # 2 search with invalid expr
361 362
        invalid_search_expr = get_invalid_expr_value
        log.info("test_search_param_invalid_expr_value: searching with "
363
                 "invalid expr: %s" % invalid_search_expr)
B
binbin 已提交
364
        collection_w.search(vectors[:default_nq], default_search_field,
365
                            default_search_params, default_limit, invalid_search_expr,
B
binbin 已提交
366 367
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
368 369
                                         "err_msg": "invalid expression %s"
                                                    % invalid_search_expr})
370

371
    @pytest.mark.tags(CaseLabel.L2)
372
    def test_search_partition_invalid_type(self, get_invalid_partition):
373 374 375 376 377 378 379
        """
        target: test search invalid partition
        method: search with invalid partition type
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix)[0]
380
        # 2. search the invalid partition
381
        partition_name = get_invalid_partition
382
        err_msg = "`partition_name_array` value {} is illegal".format(partition_name)
383 384 385 386
        collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
                            default_limit, default_search_exp, partition_name,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
387
                                         "err_msg": err_msg})
388 389 390 391 392

    @pytest.mark.tags(CaseLabel.L2)
    def test_search_with_output_fields_invalid_type(self, get_invalid_output_fields):
        """
        target: test search with output fields
393 394
        method: search with invalid output_field
        expected: raise exception and report the error
395 396 397 398
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix)[0]
        # 2. search
399 400 401
        log.info("test_search_with_output_fields_invalid_type: Searching collection %s" % collection_w.name)
        output_fields = get_invalid_output_fields
        err_msg = "`output_fields` value {} is illegal".format(output_fields)
402 403
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit,
404
                            default_search_exp, output_fields=output_fields,
405 406
                            check_task=CheckTasks.err_res,
                            check_items={ct.err_code: 1,
407
                                         ct.err_msg: err_msg})
408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484

    @pytest.mark.tags(CaseLabel.L1)
    def test_search_release_collection(self):
        """
        target: test the scenario which search the released collection
        method: 1. create collection
                2. release collection
                3. search the released collection
        expected: raise exception and report the error
        """
        # 1. initialize without data
        collection_w = self.init_collection_general(prefix, True, 10)[0]
        # 2. release collection
        collection_w.release()
        # 3. Search the released collection
        log.info("test_search_release_collection: Searching without collection ")
        collection_w.search(vectors, default_search_field,
                            default_search_params, default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "collection %s was not loaded "
                                                    "into memory" % collection_w.name})

    @pytest.mark.tags(CaseLabel.L2)
    def test_search_release_partition(self):
        """
        target: test the scenario which search the released collection
        method: 1. create collection
                2. release partition
                3. search with specifying the released partition
        expected: raise exception and report the error
        """
        # 1. initialize with data
        partition_num = 1
        collection_w = self.init_collection_general(prefix, True, 10, partition_num)[0]
        par = collection_w.partitions
        par_name = par[partition_num].name
        # 2. release partition
        conn = self.connection_wrap.get_connection()[0]
        conn.release_partitions(collection_w.name, [par_name])
        # 3. Search the released partition
        log.info("test_search_release_partition: Searching specifying the released partition")
        limit = 10
        collection_w.search(vectors, default_search_field,
                            default_search_params, limit, default_search_exp,
                            [par_name],
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "partition has been released"})

    @pytest.mark.tags(CaseLabel.L1)
    def test_search_with_empty_collection(self):
        """
        target: test search with empty connection
        method: search the empty collection
        expected: raise exception and report the error
        """
        # 1. initialize without data
        collection_w = self.init_collection_general(prefix)[0]
        # 2. search collection without data before load
        log.info("test_search_with_empty_collection: Searching empty collection %s"
                 % collection_w.name)
        err_msg = "collection" + collection_w.name + "was not loaded into memory"
        collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
                            default_limit, default_search_exp, timeout=1,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": err_msg})
        # 3. search collection without data after load
        collection_w.load()
        collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
                            default_limit, default_search_exp,
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": default_nq,
                                         "ids": [],
                                         "limit": 0})

485 486 487 488 489 490 491 492 493 494 495 496
    @pytest.mark.tags(CaseLabel.L1)
    def test_search_partition_deleted(self):
        """
        target: test search deleted partition
        method: 1. search the collection
                2. delete a partition
                3. search the deleted partition
        expected: raise exception and report the error
        """
        # 1. initialize with data
        partition_num = 1
        collection_w = self.init_collection_general(prefix, True, 1000, partition_num)[0]
497
        # 2. delete partitions
498 499 500 501 502 503
        log.info("test_search_partition_deleted: deleting a partition")
        par = collection_w.partitions
        deleted_par_name = par[partition_num].name
        collection_w.drop_partition(deleted_par_name)
        log.info("test_search_partition_deleted: deleted a partition")
        collection_w.load()
504
        # 3. search after delete partitions
505 506 507 508 509 510 511 512
        log.info("test_search_partition_deleted: searching deleted partition")
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit, default_search_exp,
                            [deleted_par_name],
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "PartitonName: %s not found" % deleted_par_name})

513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549
    @pytest.mark.tags(CaseLabel.L2)
    @pytest.mark.xfail(reason="issue 6731")
    @pytest.mark.parametrize("index, params",
                             zip(ct.all_index_types[:9],
                                 ct.default_index_params[:9]))
    def test_search_different_index_invalid_params(self, nq, dim, index, params, auto_id, _async):
        """
        target: test search with different index
        method: test search with different index
        expected: searched successfully
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, 5000,
                                                                      partition_num=1,
                                                                      auto_id=auto_id,
                                                                      dim=dim, is_index=True)
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        # 2. create different index
        if params.get("m"):
            if (dim % params["m"]) != 0:
                params["m"] = dim//4
        log.info("test_search_different_index_invalid_params: Creating index-%s" % index)
        default_index = {"index_type": index, "params": params, "metric_type": "L2"}
        collection_w.create_index("float_vector", default_index)
        log.info("test_search_different_index_invalid_params: Created index-%s" % index)
        collection_w.load()
        # 3. search
        log.info("test_search_different_index_invalid_params: Searching after creating index-%s" % index)
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, default_limit,
                            default_search_exp, _async=_async,
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": nq,
                                         "ids": insert_ids,
                                         "limit": default_limit,
                                         "_async": _async})

550
    @pytest.mark.tags(CaseLabel.L1)
551
    def test_search_index_partition_not_existed(self):
B
binbin 已提交
552 553 554 555 556
        """
        target: test search not existed partition
        method: search with not existed partition
        expected: raise exception and report the error
        """
557
        # 1. initialize with data
B
binbin 已提交
558
        collection_w = self.init_collection_general(prefix, True)[0]
559 560
        # 2. create index
        default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
B
binbin 已提交
561
        collection_w.create_index("float_vector", default_index)
562
        # 3. search the non exist partition
563
        partition_name = "search_non_exist"
B
binbin 已提交
564 565 566 567 568
        collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
                            default_limit, default_search_exp, [partition_name],
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "PartitonName: %s not found" % partition_name})
569

570
    @pytest.mark.tags(CaseLabel.L1)
571
    def test_search_param_invalid_binary(self):
B
binbin 已提交
572 573 574 575 576
        """
        target: test search within binary data (invalid parameter)
        method: search with wrong metric type
        expected: raise exception and report the error
        """
577
        # 1. initialize with binary data
B
binbin 已提交
578
        collection_w = self.init_collection_general(prefix, True, is_binary=True)[0]
579 580
        # 2. create index
        default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "JACCARD"}
B
binbin 已提交
581
        collection_w.create_index("binary_vector", default_index)
582
        # 3. search with exception
B
binbin 已提交
583
        binary_vectors = cf.gen_binary_vectors(3000, default_dim)[1]
584
        wrong_search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
B
binbin 已提交
585 586 587 588 589
        collection_w.search(binary_vectors[:default_nq], "binary_vector", wrong_search_params,
                            default_limit, default_search_exp,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "unsupported"})
590

591
    @pytest.mark.tags(CaseLabel.L2)
592 593 594 595 596 597 598 599 600
    def test_search_binary_flat_with_L2(self):
        """
        target: search binary collection using FlAT with L2
        method: search binary collection using FLAT with L2
        expected: raise exception and report error
        """
        # 1. initialize with binary data
        collection_w = self.init_collection_general(prefix, True, is_binary=True)[0]
        # 2. search and assert
601
        query_raw_vector, binary_vectors = cf.gen_binary_vectors(2, default_dim)
602 603 604 605 606 607 608
        search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
        collection_w.search(binary_vectors[:default_nq], "binary_vector",
                            search_params, default_limit, "int64 >= 0",
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "Search failed"})

609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626
    @pytest.mark.tags(CaseLabel.L1)
    def test_search_with_output_fields_not_exist(self):
        """
        target: test search with output fields
        method: search with non-exist output_field
        expected: search success
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
        # 2. search
        log.info("test_search_with_output_fields_not_exist: Searching collection %s" % collection_w.name)
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit,
                            default_search_exp, output_fields=["int63"],
                            check_task=CheckTasks.err_res,
                            check_items={ct.err_code: 1,
                                         ct.err_msg: 'Field int63 not exist'})

627
    @pytest.mark.tags(CaseLabel.L1)
628 629
    @pytest.mark.parametrize("output_fields", [[default_search_field], ["%"]])
    def test_search_output_field_vector(self, output_fields):
630 631
        """
        target: test search with vector as output field
632 633
        method: search with one vector output_field or
                wildcard for vector
634 635 636 637 638 639
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix, True)[0]
        # 2. search
        log.info("test_search_output_field_vector: Searching collection %s" % collection_w.name)
640 641
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit,
642
                            default_search_exp, output_fields=output_fields,
643 644 645 646
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": "Search doesn't support "
                                                    "vector field as output_fields"})
647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664
    @pytest.mark.tags(CaseLabel.L2)
    @pytest.mark.parametrize("output_fields", [["*%"], ["**"], ["*", "@"]])
    def test_search_output_field_invalid_wildcard(self, output_fields):
        """
        target: test search with invalid output wildcard
        method: search with invalid output_field wildcard
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w = self.init_collection_general(prefix, True)[0]
        # 2. search
        log.info("test_search_output_field_vector: Searching collection %s" % collection_w.name)
        collection_w.search(vectors[:default_nq], default_search_field,
                            default_search_params, default_limit,
                            default_search_exp, output_fields=output_fields,
                            check_task=CheckTasks.err_res,
                            check_items={"err_code": 1,
                                         "err_msg": f"Field {output_fields[-1]} not exist"})
665

666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689
class TestCollectionSearch(TestcaseBase):
    """ Test case of search interface """

    @pytest.fixture(scope="function",
                    params=[default_nb, default_nb_medium])
    def nb(self, request):
        yield request.param

    @pytest.fixture(scope="function", params=[2, 500])
    def nq(self, request):
        yield request.param

    @pytest.fixture(scope="function", params=[8, 128])
    def dim(self, request):
        yield request.param

    @pytest.fixture(scope="function", params=[False, True])
    def auto_id(self, request):
        yield request.param

    @pytest.fixture(scope="function", params=[False, True])
    def _async(self, request):
        yield request.param

690 691 692 693 694
    """
    ******************************************************************
    #  The following are valid base cases
    ******************************************************************
    """
695
    @pytest.mark.tags(CaseLabel.L0)
696
    def test_search_normal(self, nq, dim, auto_id):
697 698 699 700 701 702
        """
        target: test search normal case
        method: create connection, collection, insert and search
        expected: search successfully with limit(topK)
        """
        # 1. initialize with data
703 704
        collection_w, _, _, insert_ids = \
            self.init_collection_general(prefix, True, auto_id=auto_id, dim=dim)
705 706
        # 2. search
        log.info("test_search_normal: searching collection %s" % collection_w.name)
707 708
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        collection_w.search(vectors[:nq], default_search_field,
709 710 711
                            default_search_params, default_limit,
                            default_search_exp,
                            check_task=CheckTasks.check_search_results,
712
                            check_items={"nq": nq,
713
                                         "ids": insert_ids,
714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738
                                         "limit": default_limit}) 
                                    
    @pytest.mark.tag(CaseLabel.L0)
    def test_search_with_hit_vectors(self, nq, dim, auto_id):
        """
        target: test search with vectors in collections
        method: create connections,collection insert and search vectors in collections
        expected: search successfully with limit(topK) and can be hit at top 1 (min distance is 0)
        """
        collection_w, _vectors, _, insert_ids = \
            self.init_collection_general(prefix, True, auto_id=auto_id, dim=dim)
        # get vectors that inserted into collection
        vectors = np.array(_vectors[0]).tolist()
        vectors = [vectors[i][-1] for i in range(nq)]
        search_res, _ = collection_w.search(vectors[:nq], default_search_field,
                                            default_search_params, default_limit,
                                            default_search_exp,
                                            check_task=CheckTasks.check_search_results,
                                            check_items={"nq": nq,
                                                         "ids": insert_ids,
                                                         "limit": default_limit})
        for hits in search_res:
            # verify that top 1 hit is itself,so min distance is 0
            assert hits.distances[0] == 0.0
        
739
    @pytest.mark.tags(CaseLabel.L1)
740
    def test_search_with_empty_vectors(self, dim, auto_id, _async):
741 742 743 744 745
        """
        target: test search with empty query vector
        method: search using empty query vector
        expected: search successfully with 0 results
        """
746
        # 1. initialize without data
747
        collection_w = self.init_collection_general(prefix, True,
748
                                                    auto_id=auto_id, dim=dim)[0]
749 750 751 752
        # 2. search collection without data
        log.info("test_search_with_empty_vectors: Searching collection %s "
                 "using empty vector" % collection_w.name)
        collection_w.search([], default_search_field, default_search_params,
753
                            default_limit, default_search_exp, _async=_async,
754
                            check_task=CheckTasks.check_search_results,
755 756
                            check_items={"nq": 0,
                                         "_async": _async})
757

758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780
    @pytest.mark.tags(CaseLabel.L2)
    @pytest.mark.parametrize("search_params", [{}, {"params": {}}, {"params": {"nprobe": 10}}])
    def test_search_normal_default_params(self, dim, auto_id, search_params, _async):
        """
        target: test search normal case
        method: create connection, collection, insert and search
        expected: search successfully with limit(topK)
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = \
            self.init_collection_general(prefix, True, auto_id=auto_id, dim=dim)
        # 2. search
        log.info("test_search_normal: searching collection %s" % collection_w.name)
        vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)]
        collection_w.search(vectors[:default_nq], default_search_field,
                            search_params, default_limit,
                            default_search_exp, _async=_async,
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": default_nq,
                                         "ids": insert_ids,
                                         "limit": default_limit,
                                         "_async": _async})

Y
yanliang567 已提交
781
    @pytest.mark.tags(CaseLabel.L1)
782
    def test_search_before_after_delete(self, nq, dim, auto_id, _async):
B
binbin 已提交
783 784
        """
        target: test search function before and after deletion
785 786 787
        method: 1. search the collection
                2. delete a partition
                3. search the collection
B
binbin 已提交
788 789
        expected: the deleted entities should not be searched
        """
790
        # 1. initialize with data
791
        nb = 1000
792
        limit = 1000
793 794 795 796 797
        partition_num = 1
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      partition_num,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
B
binbin 已提交
798
        # 2. search all the partitions before partition deletion
799
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
800
        log.info("test_search_before_after_delete: searching before deleting partitions")
801 802 803
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
804
                            check_task=CheckTasks.check_search_results,
805
                            check_items={"nq": nq,
806
                                         "ids": insert_ids,
807 808
                                         "limit": limit,
                                         "_async": _async})
809 810
        # 3. delete partitions
        log.info("test_search_before_after_delete: deleting a partition")
B
binbin 已提交
811
        par = collection_w.partitions
812
        deleted_entity_num = par[partition_num].num_entities
813
        entity_num = nb - deleted_entity_num
B
binbin 已提交
814
        collection_w.drop_partition(par[partition_num].name)
815
        log.info("test_search_before_after_delete: deleted a partition")
B
binbin 已提交
816
        collection_w.load()
817
        # 4. search non-deleted part after delete partitions
818
        log.info("test_search_before_after_delete: searching after deleting partitions")
819 820 821
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
822
                            check_task=CheckTasks.check_search_results,
823
                            check_items={"nq": nq,
824
                                         "ids": insert_ids[:entity_num],
825 826
                                         "limit": limit-deleted_entity_num,
                                         "_async": _async})
827 828

    @pytest.mark.tags(CaseLabel.L2)
829
    def test_search_partition_after_release_one(self, nq, dim, auto_id, _async):
830 831 832 833 834 835 836 837
        """
        target: test search function before and after release
        method: 1. search the collection
                2. release a partition
                3. search the collection
        expected: the deleted entities should not be searched
        """
        # 1. initialize with data
838
        nb = 1000
839
        limit = 1000
840
        partition_num = 1
841 842 843 844
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      partition_num,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
845
        # 2. search all the partitions before partition deletion
846
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
847
        log.info("test_search_partition_after_release_one: searching before deleting partitions")
848 849 850
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
851
                            check_task=CheckTasks.check_search_results,
852
                            check_items={"nq": nq,
853
                                         "ids": insert_ids,
854 855
                                         "limit": limit,
                                         "_async": _async})
856
        # 3. release one partition
857 858 859
        log.info("test_search_partition_after_release_one: releasing a partition")
        par = collection_w.partitions
        deleted_entity_num = par[partition_num].num_entities
860
        entity_num = nb - deleted_entity_num
861 862 863
        conn = self.connection_wrap.get_connection()[0]
        conn.release_partitions(collection_w.name, [par[partition_num].name])
        log.info("test_search_partition_after_release_one: released a partition")
864
        # 4. search collection after release one partition
865
        log.info("test_search_partition_after_release_one: searching after deleting partitions")
866 867 868
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
869
                            check_task=CheckTasks.check_search_results,
870
                            check_items={"nq": nq,
871
                                         "ids": insert_ids[:entity_num],
872 873
                                         "limit": limit - deleted_entity_num,
                                         "_async": _async})
874 875

    @pytest.mark.tags(CaseLabel.L2)
876
    def test_search_partition_after_release_all(self, nq, dim, auto_id, _async):
877 878 879 880 881 882 883 884
        """
        target: test search function before and after release
        method: 1. search the collection
                2. release a partition
                3. search the collection
        expected: the deleted entities should not be searched
        """
        # 1. initialize with data
885
        nb = 1000
886
        limit = 1000
887 888 889
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      1, auto_id=auto_id,
                                                                      dim=dim)
890
        # 2. search all the partitions before partition deletion
891
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
892
        log.info("test_search_partition_after_release_all: searching before deleting partitions")
893 894 895
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
896
                            check_task=CheckTasks.check_search_results,
897
                            check_items={"nq": nq,
898
                                         "ids": insert_ids,
899 900
                                         "limit": limit,
                                         "_async": _async})
901
        # 3. release all partitions
902 903 904 905 906
        log.info("test_search_partition_after_release_all: releasing a partition")
        par = collection_w.partitions
        conn = self.connection_wrap.get_connection()[0]
        conn.release_partitions(collection_w.name, [par[0].name, par[1].name])
        log.info("test_search_partition_after_release_all: released a partition")
907
        # 4. search collection after release all partitions
908 909 910
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
911
                            check_task=CheckTasks.check_search_results,
912
                            check_items={"nq": nq,
913
                                         "ids": [],
914 915
                                         "limit": 0,
                                         "_async": _async})
916

917
    @pytest.mark.tags(CaseLabel.L2)
918
    def test_search_collection_after_release_load(self, nb, nq, dim, auto_id, _async):
919 920 921 922 923 924 925 926 927
        """
        target: search the pre-released collection after load
        method: 1. create collection
                2. release collection
                3. load collection
                4. search the pre-released collection
        expected: search successfully
        """
        # 1. initialize without data
928 929 930
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      1, auto_id=auto_id,
                                                                      dim=dim)
931 932 933 934
        # 2. release collection
        collection_w.release()
        # 3. Search the pre-released collection after load
        collection_w.load()
935 936 937 938
        log.info("test_search_collection_after_release_load: searching after load")
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        collection_w.search(vectors[:nq], default_search_field, default_search_params,
                            default_limit, default_search_exp, _async=_async,
939
                            check_task=CheckTasks.check_search_results,
940
                            check_items={"nq": nq,
941
                                         "ids": insert_ids,
942 943
                                         "limit": default_limit,
                                         "_async": _async})
944 945

    @pytest.mark.tags(CaseLabel.L2)
946
    @pytest.mark.xfail(reason="issue 6997")
947
    def test_search_partition_after_release_load(self, nb, nq, dim, auto_id, _async):
948 949 950 951 952
        """
        target: search the pre-released collection after load
        method: 1. create collection
                2. release a partition
                3. load partition
953
                4. search the pre-released partition
954 955 956
        expected: search successfully
        """
        # 1. initialize without data
957 958 959
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      1, auto_id=auto_id,
                                                                      dim=dim)
960 961 962 963 964 965
        # 2. release collection
        log.info("test_search_partition_after_release_load: releasing a partition")
        par = collection_w.partitions
        conn = self.connection_wrap.get_connection()[0]
        conn.release_partitions(collection_w.name, [par[1].name])
        log.info("test_search_partition_after_release_load: released a partition")
966 967
        # 3. Search the collection after load
        limit = 1000
968 969
        collection_w.load()
        log.info("test_search_partition_after_release_load: searching after load")
970 971 972
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        collection_w.search(vectors[:nq], default_search_field, default_search_params,
                            limit, default_search_exp, _async=_async,
973
                            check_task=CheckTasks.check_search_results,
974
                            check_items={"nq": nq,
975
                                         "ids": insert_ids,
976 977
                                         "limit": limit,
                                         "_async": _async})
978
        # 4. Search the pre-released partition after load
979 980 981 982 983
        if limit > par[1].num_entities:
            limit_check = par[1].num_entities
        else:
            limit_check = limit
        collection_w.search(vectors[:nq], default_search_field, default_search_params,
984
                            limit, default_search_exp,
985
                            [par[1].name], _async=_async,
986
                            check_task=CheckTasks.check_search_results,
987
                            check_items={"nq": nq,
988
                                         "ids": insert_ids[par[0].num_entities:],
989 990
                                         "limit": limit_check,
                                         "_async": _async})
991 992

    @pytest.mark.tags(CaseLabel.L2)
993
    def test_search_load_flush_load(self, nb, nq, dim, auto_id, _async):
994 995 996 997 998 999 1000 1001
        """
        target: test search when load before flush
        method: 1. search the collection
                2. insert data and load
                3. flush, and load
        expected: search success with limit(topK)
        """
        # 1. initialize with data
1002
        collection_w = self.init_collection_general(prefix, auto_id=auto_id, dim=dim)[0]
1003
        # 2. insert data
1004
        insert_ids = cf.insert_data(collection_w, nb, auto_id=auto_id, dim=dim)[3]
1005 1006 1007
        # 3. load data
        collection_w.load()
        # 4. flush and load
Z
zhuwenxing 已提交
1008
        collection_w.num_entities
1009 1010
        collection_w.load()
        # 5. search for new data without load
1011 1012 1013 1014
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, default_limit,
                            default_search_exp, _async=_async,
1015
                            check_task=CheckTasks.check_search_results,
1016
                            check_items={"nq": nq,
1017
                                         "ids": insert_ids,
1018 1019
                                         "limit": default_limit,
                                         "_async": _async})
1020

1021 1022
    @pytest.mark.tags(CaseLabel.L2)
    def test_search_new_data(self, nq, dim, auto_id, _async):
1023 1024
        """
        target: test search new inserted data without load
1025 1026 1027
        method: 1. search the collection
                2. insert new data
                3. search the collection without load again
1028 1029 1030 1031 1032
        expected: new data should be searched
        """
        # 1. initialize with data
        limit = 1000
        nb_old = 500
1033 1034 1035
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb_old,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
1036
        # 2. search for original data after load
1037
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
1038
        log.info("test_search_new_data: searching for original data after load")
1039 1040 1041
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
1042
                            check_task=CheckTasks.check_search_results,
1043
                            check_items={"nq": nq,
1044
                                         "ids": insert_ids,
1045 1046
                                         "limit": nb_old,
                                         "_async": _async})
1047 1048
        # 3. insert new data
        nb_new = 300
1049 1050
        insert_ids_new = cf.insert_data(collection_w, nb_new,
                                        auto_id=auto_id, dim=dim)[3]
1051
        insert_ids.extend(insert_ids_new)
1052 1053 1054
        # gracefulTime is default as 1s which allows data
        # could not be searched instantly in gracefulTime
        time.sleep(gracefulTime)
1055
        # 4. search for new data without load
1056 1057 1058
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, limit,
                            default_search_exp, _async=_async,
1059
                            check_task=CheckTasks.check_search_results,
1060
                            check_items={"nq": nq,
1061
                                         "ids": insert_ids,
1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086
                                         "limit": nb_old+nb_new,
                                         "_async": _async})

    @pytest.mark.tags(CaseLabel.L2)
    def test_search_max_dim(self, nq, auto_id, _async):
        """
        target: test search normal case
        method: create connection, collection, insert and search
        expected: search successfully with limit(topK)
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, default_nb,
                                                                      auto_id=auto_id,
                                                                      dim=max_dim)
        # 2. search
        log.info("test_search_max_dim: searching collection %s" % collection_w.name)
        vectors = [[random.random() for _ in range(max_dim)] for _ in range(nq)]
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, 2,
                            default_search_exp, _async=_async,
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": nq,
                                         "ids": insert_ids,
                                         "limit": 2,
                                         "_async": _async})
1087

1088
    @pytest.mark.tags(CaseLabel.L2)
1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106
    @pytest.mark.parametrize("index, params",
                             zip(ct.all_index_types[:9],
                                 ct.default_index_params[:9]))
    def test_search_after_different_index_with_params(self, dim, index, params, auto_id, _async):
        """
        target: test search with invalid search params
        method: test search with invalid params type
        expected: raise exception and report the error
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, 5000,
                                                                      partition_num=1,
                                                                      auto_id=auto_id,
                                                                      dim=dim, is_index=True)
        # 2. create index and load
        if params.get("m"):
            if (dim % params["m"]) != 0:
                params["m"] = dim//4
1107 1108 1109
        if params.get("PQM"):
            if (dim % params["PQM"]) != 0:
                params["PQM"] = dim//4
1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127
        default_index = {"index_type": index, "params": params, "metric_type": "L2"}
        collection_w.create_index("float_vector", default_index)
        collection_w.load()
        # 3. search
        search_params = cf.gen_search_param(index)
        vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)]
        for search_param in search_params:
            log.info("Searching with search params: {}".format(search_param))
            collection_w.search(vectors[:default_nq], default_search_field,
                                search_param, default_limit,
                                default_search_exp, _async=_async,
                                check_task=CheckTasks.check_search_results,
                                check_items={"nq": default_nq,
                                             "ids": insert_ids,
                                             "limit": default_limit,
                                             "_async": _async})

    @pytest.mark.tags(CaseLabel.L2)
1128 1129 1130
    @pytest.mark.parametrize("index, params",
                             zip(ct.all_index_types[:9],
                                 ct.default_index_params[:9]))
1131
    def test_search_after_index_different_metric_type(self, dim, index, params, auto_id, _async):
B
binbin 已提交
1132 1133 1134 1135 1136
        """
        target: test search with different metric type
        method: test search with different metric type
        expected: searched successfully
        """
1137
        # 1. initialize with data
1138
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, 5000,
1139 1140
                                                                      partition_num=1,
                                                                      auto_id=auto_id,
1141
                                                                      dim=dim, is_index=True)
1142
        # 2. create different index
1143 1144 1145
        if params.get("m"):
            if (dim % params["m"]) != 0:
                params["m"] = dim//4
1146 1147 1148
        if params.get("PQM"):
            if (dim % params["PQM"]) != 0:
                params["PQM"] = dim//4
1149
        log.info("test_search_after_index_different_metric_type: Creating index-%s" % index)
1150 1151
        default_index = {"index_type": index, "params": params, "metric_type": "IP"}
        collection_w.create_index("float_vector", default_index)
1152 1153
        log.info("test_search_after_index_different_metric_type: Created index-%s" % index)
        collection_w.load()
1154
        # 3. search
1155
        search_params = cf.gen_search_param(index, "IP")
1156
        vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)]
1157 1158 1159 1160 1161 1162 1163 1164 1165 1166
        for search_param in search_params:
            log.info("Searching with search params: {}".format(search_param))
            collection_w.search(vectors[:default_nq], default_search_field,
                                search_param, default_limit,
                                default_search_exp, _async=_async,
                                check_task=CheckTasks.check_search_results,
                                check_items={"nq": default_nq,
                                             "ids": insert_ids,
                                             "limit": default_limit,
                                             "_async": _async})
1167

1168 1169
    @pytest.mark.tags(CaseLabel.L2)
    def test_search_collection_multiple_times(self, nb, nq, dim, auto_id, _async):
B
binbin 已提交
1170 1171 1172 1173 1174
        """
        target: test search for multiple times
        method: search for multiple times
        expected: searched successfully
        """
1175
        # 1. initialize with data
1176 1177 1178
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
1179
        # 2. search for multiple times
1180
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
1181
        for i in range(search_num):
B
binbin 已提交
1182
            log.info("test_search_collection_multiple_times: searching round %d" % (i+1))
1183 1184 1185
            collection_w.search(vectors[:nq], default_search_field,
                                default_search_params, default_limit,
                                default_search_exp, _async=_async,
1186
                                check_task=CheckTasks.check_search_results,
1187
                                check_items={"nq": nq,
1188
                                             "ids": insert_ids,
1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231
                                             "limit": default_limit,
                                             "_async": _async})

    @pytest.mark.tags(CaseLabel.L2)
    def test_search_sync_async_multiple_times(self, nb, nq, dim, auto_id):
        """
        target: test async search after sync search case
        method: create connection, collection, insert,
                sync search and async search
        expected: search successfully with limit(topK)
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
        # 2. search
        log.info("test_search_sync_async_multiple_times: searching collection %s" % collection_w.name)
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        for i in range(search_num):
            log.info("test_search_sync_async_multiple_times: searching round %d" % (i + 1))
            for _async in [False, True]:
                collection_w.search(vectors[:nq], default_search_field,
                                    default_search_params, default_limit,
                                    default_search_exp, _async=_async,
                                    check_task=CheckTasks.check_search_results,
                                    check_items={"nq": nq,
                                                 "ids": insert_ids,
                                                 "limit": default_limit,
                                                 "_async": _async})

    @pytest.mark.tags(CaseLabel.L2)
    def test_search_multiple_vectors(self, nb, nq, dim, auto_id, _async):
        """
        target: test search with multiple vectors
        method: create connection, collection with multiple
                vectors, insert and search
        expected: search successfully with limit(topK)
        """
        # 1. connect
        self._connect()
        # 2. create collection with multiple vectors
        c_name = cf.gen_unique_str(prefix)
        fields = [cf.gen_int64_field(is_primary=True), cf.gen_float_field(),
1232 1233
                  cf.gen_float_vec_field(dim=dim), cf.gen_float_vec_field(name="tmp", dim=dim)]
        schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id)
1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265
        collection_w = self.collection_wrap.init_collection(c_name, schema=schema,
                                                            check_task=CheckTasks.check_collection_property,
                                                            check_items={"name": c_name, "schema": schema})[0]
        # 3. insert
        vectors = [[random.random() for _ in range(dim)] for _ in range(nb)]
        vectors_tmp = [[random.random() for _ in range(dim)] for _ in range(nb)]
        data = [[i for i in range(nb)], [np.float32(i) for i in range(nb)], vectors, vectors_tmp]
        if auto_id:
            data = [[np.float32(i) for i in range(nb)], vectors, vectors_tmp]
        res = collection_w.insert(data)
        insert_ids = res.primary_keys
        assert collection_w.num_entities == nb
        # 4. load
        collection_w.load()
        # 5. search all the vectors
        log.info("test_search_multiple_vectors: searching collection %s" % collection_w.name)
        collection_w.search(vectors[:nq], default_search_field,
                            default_search_params, default_limit,
                            default_search_exp, _async=_async,
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": nq,
                                         "ids": insert_ids,
                                         "limit": default_limit,
                                         "_async": _async})
        collection_w.search(vectors[:nq], "tmp",
                            default_search_params, default_limit,
                            default_search_exp, _async=_async,
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": nq,
                                         "ids": insert_ids,
                                         "limit": default_limit,
                                         "_async": _async})
1266

1267
    @pytest.mark.tags(CaseLabel.L1)
1268
    def test_search_index_one_partition(self, nb, auto_id, _async):
B
binbin 已提交
1269 1270 1271 1272 1273
        """
        target: test search from partition
        method: search from one partition
        expected: searched successfully
        """
1274
        # 1. initialize with data
1275 1276
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      partition_num=1,
1277 1278
                                                                      auto_id=auto_id,
                                                                      is_index=True)
1279
        
1280 1281
        # 2. create index
        default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
B
binbin 已提交
1282
        collection_w.create_index("float_vector", default_index)
1283
        collection_w.load()
1284
        # 3. search in one partition
B
binbin 已提交
1285 1286
        log.info("test_search_index_one_partition: searching (1000 entities) through one partition")
        limit = 1000
1287
        par = collection_w.partitions
1288 1289 1290 1291
        if limit > par[1].num_entities:
            limit_check = par[1].num_entities
        else:
            limit_check = limit
1292
        search_params = {"metric_type": "L2", "params": {"nprobe": 128}}
1293
        collection_w.search(vectors[:default_nq], default_search_field,
1294
                            search_params, limit, default_search_exp,
1295
                            [par[1].name], _async=_async,
1296 1297
                            check_task=CheckTasks.check_search_results,
                            check_items={"nq": default_nq,
1298
                                         "ids": insert_ids[par[0].num_entities:],
1299 1300
                                         "limit": limit_check,
                                         "_async": _async})
1301

1302 1303
    @pytest.mark.tags(CaseLabel.L2)
    def test_search_index_partitions(self, nb, nq, dim, auto_id, _async):
B
binbin 已提交
1304 1305
        """
        target: test search from partitions
1306
        method: search from partitions
B
binbin 已提交
1307 1308
        expected: searched successfully
        """
1309
        # 1. initialize with data
1310 1311 1312
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      partition_num=1,
                                                                      auto_id=auto_id,
1313 1314
                                                                      dim=dim,
                                                                      is_index=True)
1315
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
1316 1317
        # 2. create index
        default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
B
binbin 已提交
1318
        collection_w.create_index("float_vector", default_index)
1319
        collection_w.load()
1320 1321
        # 3. search through partitions
        log.info("test_search_index_partitions: searching (1000 entities) through partitions")
B
binbin 已提交
1322
        par = collection_w.partitions
1323
        log.info("test_search_index_partitions: partitions: %s" % par)
B
binbin 已提交
1324
        limit = 1000
1325
        collection_w.search(vectors[:nq], default_search_field,
1326
                            default_search_params, limit, default_search_exp,
1327
                            [par[0].name, par[1].name], _async=_async,
1328
                            check_task=CheckTasks.check_search_results,
1329
                            check_items={"nq": nq,
1330
                                         "ids": insert_ids,
1331 1332
                                         "limit": limit,
                                         "_async": _async})
1333

1334
    @pytest.mark.tags(CaseLabel.L2)
1335 1336
    @pytest.mark.parametrize("partition_names",
                             [["(.*)"], ["search(.*)"]])
1337
    def test_search_index_partitions_fuzzy(self, nb, nq, dim, partition_names, auto_id, _async):
1338 1339 1340 1341 1342 1343 1344
        """
        target: test search from partitions
        method: search from partitions with fuzzy
                partition name
        expected: searched successfully
        """
        # 1. initialize with data
1345 1346 1347 1348 1349
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      partition_num=1,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
1350 1351 1352 1353 1354 1355
        # 2. create index
        default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
        collection_w.create_index("float_vector", default_index)
        # 3. search through partitions
        log.info("test_search_index_partitions_fuzzy: searching through partitions")
        limit = 1000
1356
        limit_check = limit
1357 1358
        par = collection_w.partitions
        if partition_names == ["search(.*)"]:
1359
            insert_ids = insert_ids[par[0].num_entities:]
1360 1361 1362
            if limit > par[1].num_entities:
                limit_check = par[1].num_entities
        collection_w.search(vectors[:nq], default_search_field,
1363
                            default_search_params, limit, default_search_exp,
1364
                            partition_names, _async=_async,
1365
                            check_task=CheckTasks.check_search_results,
1366
                            check_items={"nq": nq,
1367
                                         "ids": insert_ids,
1368 1369
                                         "limit": limit_check,
                                         "_async": _async})
1370

1371
    @pytest.mark.tags(CaseLabel.L2)
1372
    def test_search_index_partition_empty(self, nq, dim, auto_id, _async):
B
binbin 已提交
1373 1374 1375 1376 1377
        """
        target: test search the empty partition
        method: search from the empty partition
        expected: searched successfully with 0 results
        """
1378
        # 1. initialize with data
1379 1380
        collection_w = self.init_collection_general(prefix, True, auto_id=auto_id,
                                                    dim=dim, is_index=True)[0]
1381
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
1382 1383
        # 2. create empty partition
        partition_name = "search_partition_empty"
B
binbin 已提交
1384 1385
        collection_w.create_partition(partition_name=partition_name, description="search partition empty")
        par = collection_w.partitions
1386
        log.info("test_search_index_partition_empty: partitions: %s" % par)
B
binbin 已提交
1387
        collection_w.load()
1388 1389
        # 3. create index
        default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
B
binbin 已提交
1390
        collection_w.create_index("float_vector", default_index)
1391
        # 4. search the empty partition
1392 1393 1394
        log.info("test_search_index_partition_empty: searching %s "
                 "entities through empty partition" % default_limit)
        collection_w.search(vectors[:nq], default_search_field,
1395 1396
                            default_search_params, default_limit,
                            default_search_exp, [partition_name],
1397
                            _async=_async,
1398
                            check_task=CheckTasks.check_search_results,
1399
                            check_items={"nq": nq,
1400
                                         "ids": [],
1401 1402
                                         "limit": 0,
                                         "_async": _async})
1403

1404
    @pytest.mark.tags(CaseLabel.L2)
1405
    def test_search_binary_jaccard_flat_index(self, nq, dim, auto_id, _async):
B
binbin 已提交
1406
        """
1407
        target: search binary_collection, and check the result: distance
1408
        method: compare the return distance value with value computed with JACCARD
1409
        expected: the return distance equals to the computed value
B
binbin 已提交
1410
        """
1411
        # 1. initialize with binary data
1412 1413 1414
        collection_w, _, binary_raw_vector, insert_ids = self.init_collection_general(prefix, True, 2,
                                                                                      is_binary=True,
                                                                                      auto_id=auto_id,
1415 1416
                                                                                      dim=dim,
                                                                                      is_index=True)
1417 1418
        # 2. create index
        default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "JACCARD"}
B
binbin 已提交
1419
        collection_w.create_index("binary_vector", default_index)
1420
        collection_w.load()
1421
        # 3. compute the distance
1422
        query_raw_vector, binary_vectors = cf.gen_binary_vectors(3000, dim)
1423 1424 1425 1426
        distance_0 = cf.jaccard(query_raw_vector[0], binary_raw_vector[0])
        distance_1 = cf.jaccard(query_raw_vector[0], binary_raw_vector[1])
        # 4. search and compare the distance
        search_params = {"metric_type": "JACCARD", "params": {"nprobe": 10}}
1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437
        res = collection_w.search(binary_vectors[:nq], "binary_vector",
                                  search_params, default_limit, "int64 >= 0",
                                  _async=_async,
                                  check_task=CheckTasks.check_search_results,
                                  check_items={"nq": nq,
                                               "ids": insert_ids,
                                               "limit": 2,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1438 1439
        assert abs(res[0]._distances[0] - min(distance_0, distance_1)) <= epsilon

1440
    @pytest.mark.tags(CaseLabel.L2)
1441
    def test_search_binary_hamming_flat_index(self, nq, dim, auto_id, _async):
B
binbin 已提交
1442
        """
1443
        target: search binary_collection, and check the result: distance
1444
        method: compare the return distance value with value computed with HAMMING
1445
        expected: the return distance equals to the computed value
B
binbin 已提交
1446
        """
1447
        # 1. initialize with binary data
1448 1449 1450
        collection_w, _, binary_raw_vector, insert_ids = self.init_collection_general(prefix, True, 2,
                                                                                      is_binary=True,
                                                                                      auto_id=auto_id,
1451 1452
                                                                                      dim=dim,
                                                                                      is_index=True)
1453 1454
        # 2. create index
        default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "HAMMING"}
B
binbin 已提交
1455
        collection_w.create_index("binary_vector", default_index)
1456
        # 3. compute the distance
1457 1458
        collection_w.load()
        query_raw_vector, binary_vectors = cf.gen_binary_vectors(3000, dim)
1459 1460 1461 1462
        distance_0 = cf.hamming(query_raw_vector[0], binary_raw_vector[0])
        distance_1 = cf.hamming(query_raw_vector[0], binary_raw_vector[1])
        # 4. search and compare the distance
        search_params = {"metric_type": "HAMMING", "params": {"nprobe": 10}}
1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473
        res = collection_w.search(binary_vectors[:nq], "binary_vector",
                                  search_params, default_limit, "int64 >= 0",
                                  _async=_async,
                                  check_task=CheckTasks.check_search_results,
                                  check_items={"nq": nq,
                                               "ids": insert_ids,
                                               "limit": 2,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1474 1475
        assert abs(res[0]._distances[0] - min(distance_0, distance_1)) <= epsilon

1476
    @pytest.mark.tags(CaseLabel.L2)
B
binbin 已提交
1477
    @pytest.mark.xfail(reason="issue 6843")
1478
    def test_search_binary_tanimoto_flat_index(self, nq, dim, auto_id, _async):
B
binbin 已提交
1479
        """
1480
        target: search binary_collection, and check the result: distance
1481
        method: compare the return distance value with value computed with TANIMOTO
1482
        expected: the return distance equals to the computed value
B
binbin 已提交
1483
        """
1484
        # 1. initialize with binary data
1485 1486 1487
        collection_w, _, binary_raw_vector, insert_ids = self.init_collection_general(prefix, True, 2,
                                                                                      is_binary=True,
                                                                                      auto_id=auto_id,
1488 1489
                                                                                      dim=dim,
                                                                                      is_index=True)
1490
        log.info("auto_id= %s, _async= %s" % (auto_id, _async))
1491 1492
        # 2. create index
        default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "TANIMOTO"}
B
binbin 已提交
1493
        collection_w.create_index("binary_vector", default_index)
1494
        collection_w.load()
1495
        # 3. compute the distance
1496
        query_raw_vector, binary_vectors = cf.gen_binary_vectors(3000, dim)
1497 1498 1499 1500
        distance_0 = cf.tanimoto(query_raw_vector[0], binary_raw_vector[0])
        distance_1 = cf.tanimoto(query_raw_vector[0], binary_raw_vector[1])
        # 4. search and compare the distance
        search_params = {"metric_type": "TANIMOTO", "params": {"nprobe": 10}}
1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511
        res = collection_w.search(binary_vectors[:nq], "binary_vector",
                                  search_params, default_limit, "int64 >= 0",
                                  _async=_async,
                                  check_task=CheckTasks.check_search_results,
                                  check_items={"nq": nq,
                                               "ids": insert_ids,
                                               "limit": 2,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1512 1513
        assert abs(res[0]._distances[0] - min(distance_0, distance_1)) <= epsilon

1514
    @pytest.mark.tags(CaseLabel.L1)
1515 1516
    @pytest.mark.parametrize("expression", cf.gen_normal_expressions())
    def test_search_with_expression(self, dim, expression, _async):
1517 1518 1519 1520 1521 1522
        """
        target: test search with different expressions
        method: test search with different expressions
        expected: searched successfully with correct limit(topK)
        """
        # 1. initialize with data
1523
        nb = 1000
1524
        collection_w, _vectors, _, insert_ids = self.init_collection_general(prefix, True,
1525 1526
                                                                             nb, dim=dim,
                                                                             is_index=True)
1527

Z
zhuwenxing 已提交
1528
        # filter result with expression in collection
1529
        _vectors = _vectors[0]
Z
zhuwenxing 已提交
1530 1531
        expression = expression.replace("&&", "and").replace("||", "or")
        filter_ids = []
1532
        for i, _id in enumerate(insert_ids):
1533 1534 1535 1536 1537
            int64 = _vectors.int64[i]
            float = _vectors.float[i]
            if not expression or eval(expression):
                filter_ids.append(_id)
        
1538 1539 1540
        # 2. create index
        index_param = {"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 100}}
        collection_w.create_index("float_vector", index_param)
1541
        collection_w.load()
1542 1543

        # 3. search with expression
1544
        log.info("test_search_with_expression: searching with expression: %s" % expression)
1545
        vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)]
1546
        search_res, _ = collection_w.search(vectors[:default_nq], default_search_field,
1547 1548 1549 1550 1551 1552 1553
                                            default_search_params, nb, expression,
                                            _async=_async,
                                            check_task=CheckTasks.check_search_results,
                                            check_items={"nq": default_nq,
                                                         "ids": insert_ids,
                                                         "limit": min(nb, len(filter_ids)),
                                                         "_async": _async})
1554 1555 1556 1557 1558 1559 1560
        if _async:
            search_res.done()
            search_res = search_res.result()
        
        filter_ids_set = set(filter_ids)
        for hits in search_res:
            ids = hits.ids
Z
zhuwenxing 已提交
1561
            assert set(ids).issubset(filter_ids_set)
1562

1563
    @pytest.mark.tags(CaseLabel.L2)
1564
    @pytest.mark.parametrize("expression", cf.gen_normal_expressions_field(default_float_field_name))
1565
    def test_search_with_expression_auto_id(self, dim, expression, _async):
1566 1567 1568 1569 1570 1571 1572
        """
        target: test search with different expressions
        method: test search with different expressions
        expected: searched successfully with correct limit(topK)
        """
        # 1. initialize with data
        nb = 1000
1573
        collection_w, _vectors, _, insert_ids = self.init_collection_general(prefix, True, nb,
1574 1575 1576
                                                                             auto_id=True,
                                                                             dim=dim,
                                                                             is_index=True)
1577 1578
        
        
Z
zhuwenxing 已提交
1579
        # filter result with expression in collection
1580
        _vectors = _vectors[0]
1581
        expression = expression.replace("&&", "and").replace("||", "or")
1582 1583 1584 1585
        filter_ids = []
        for i, _id in enumerate(insert_ids):
            exec(f"{default_float_field_name} = _vectors.{default_float_field_name}[i]")
            if not expression or eval(expression):
Z
zhuwenxing 已提交
1586
                filter_ids.append(_id)
1587
        
1588 1589 1590
        # 2. create index
        index_param = {"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 100}}
        collection_w.create_index("float_vector", index_param)
1591
        collection_w.load()
1592 1593


1594 1595 1596
        # 3. search with different expressions
        log.info("test_search_with_expression: searching with expression: %s" % expression)
        vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)]
1597
        search_res, _ = collection_w.search(vectors[:default_nq], default_search_field,
1598 1599 1600 1601 1602 1603 1604
                                            default_search_params, nb, expression,
                                            _async=_async,
                                            check_task=CheckTasks.check_search_results,
                                            check_items={"nq": default_nq,
                                                         "ids": insert_ids,
                                                         "limit": min(nb, len(filter_ids)),
                                                         "_async": _async})
1605 1606 1607 1608 1609 1610 1611 1612
        if _async:
            search_res.done()
            search_res = search_res.result()
        
        filter_ids_set = set(filter_ids)
        for hits in search_res:
            ids = hits.ids
            assert set(ids).issubset(filter_ids_set)        
1613

1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646
    @pytest.mark.tags(CaseLabel.L2)
    def test_search_expression_all_data_type(self, nb, nq, dim, auto_id, _async):
        """
        target: test search using different supported data type
        method: search using different supported data type
        expected: search success
        """
        # 1. initialize with data
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      is_all_data_type=True,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
        # 2. search
        log.info("test_search_expression_all_data_type: Searching collection %s" % collection_w.name)
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        search_exp = "int64 >= 0 && int32 >= 0 && int16 >= 0 " \
                     "&& int8 >= 0 && float >= 0 && double >= 0"
        res = collection_w.search(vectors[:nq], default_search_field,
                                  default_search_params, default_limit,
                                  search_exp, _async=_async,
                                  output_fields=[default_int64_field_name,
                                                 default_float_field_name],
                                  check_task=CheckTasks.check_search_results,
                                  check_items={"nq": nq,
                                               "ids": insert_ids,
                                               "limit": default_limit,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
        assert len(res[0][0].entity._row_data) != 0
        assert (default_int64_field_name and default_float_field_name) in res[0][0].entity._row_data

1647 1648
    @pytest.mark.tags(CaseLabel.L2)
    def test_search_with_output_fields_empty(self, nb, nq, dim, auto_id, _async):
1649 1650 1651 1652 1653 1654
        """
        target: test search with output fields
        method: search with empty output_field
        expected: search success
        """
        # 1. initialize with data
1655 1656 1657
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
1658 1659
        # 2. search
        log.info("test_search_with_output_fields_empty: Searching collection %s" % collection_w.name)
1660 1661
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        res = collection_w.search(vectors[:nq], default_search_field,
1662
                                  default_search_params, default_limit,
1663 1664
                                  default_search_exp, _async=_async,
                                  output_fields=[],
1665
                                  check_task=CheckTasks.check_search_results,
1666
                                  check_items={"nq": nq,
1667
                                               "ids": insert_ids,
1668 1669 1670 1671 1672
                                               "limit": default_limit,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1673 1674 1675
        assert len(res[0][0].entity._row_data) == 0

    @pytest.mark.tags(CaseLabel.L1)
1676
    def test_search_with_output_field(self, auto_id, _async):
1677 1678 1679 1680 1681 1682
        """
        target: test search with output fields
        method: search with one output_field
        expected: search success
        """
        # 1. initialize with data
1683 1684
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True,
                                                                      auto_id=auto_id)
1685 1686
        # 2. search
        log.info("test_search_with_output_field: Searching collection %s" % collection_w.name)
1687
        
1688 1689
        res = collection_w.search(vectors[:default_nq], default_search_field,
                                  default_search_params, default_limit,
1690 1691
                                  default_search_exp, _async=_async,
                                  output_fields=[default_int64_field_name],
1692 1693
                                  check_task=CheckTasks.check_search_results,
                                  check_items={"nq": default_nq,
1694
                                               "ids": insert_ids,
1695 1696 1697 1698 1699
                                               "limit": default_limit,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1700 1701 1702 1703
        assert len(res[0][0].entity._row_data) != 0
        assert default_int64_field_name in res[0][0].entity._row_data

    @pytest.mark.tags(CaseLabel.L2)
1704
    def test_search_with_output_fields(self, nb, nq, dim, auto_id, _async):
1705 1706 1707 1708 1709 1710
        """
        target: test search with output fields
        method: search with multiple output_field
        expected: search success
        """
        # 1. initialize with data
1711 1712 1713 1714
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      is_all_data_type=True,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
1715 1716
        # 2. search
        log.info("test_search_with_output_fields: Searching collection %s" % collection_w.name)
1717 1718
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        res = collection_w.search(vectors[:nq], default_search_field,
1719
                                  default_search_params, default_limit,
1720
                                  default_search_exp, _async=_async,
1721 1722 1723
                                  output_fields=[default_int64_field_name,
                                                 default_float_field_name],
                                  check_task=CheckTasks.check_search_results,
1724
                                  check_items={"nq": nq,
1725
                                               "ids": insert_ids,
1726 1727 1728 1729 1730
                                               "limit": default_limit,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1731 1732 1733 1734
        assert len(res[0][0].entity._row_data) != 0
        assert (default_int64_field_name and default_float_field_name) in res[0][0].entity._row_data

    @pytest.mark.tags(CaseLabel.L2)
1735 1736
    @pytest.mark.parametrize("output_fields", [["*"], ["*", default_float_field_name]])
    def test_search_with_output_field_wildcard(self, output_fields, auto_id, _async):
1737
        """
1738 1739
        target: test search with output fields using wildcard
        method: search with one output_field (wildcard)
1740 1741 1742
        expected: search success
        """
        # 1. initialize with data
1743 1744
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True,
                                                                      auto_id=auto_id)
1745
        # 2. search
1746
        log.info("test_search_with_output_field_wildcard: Searching collection %s" % collection_w.name)
1747
        
1748
        res = collection_w.search(vectors[:default_nq], default_search_field,
1749
                                  default_search_params, default_limit,
1750 1751
                                  default_search_exp, _async=_async,
                                  output_fields=output_fields,
1752
                                  check_task=CheckTasks.check_search_results,
1753
                                  check_items={"nq": default_nq,
1754
                                               "ids": insert_ids,
1755 1756 1757 1758 1759
                                               "limit": default_limit,
                                               "_async": _async})[0]
        if _async:
            res.done()
            res = res.result()
1760 1761 1762 1763
        assert len(res[0][0].entity._row_data) != 0
        assert (default_int64_field_name and default_float_field_name) in res[0][0].entity._row_data

    @pytest.mark.tags(CaseLabel.L2)
1764
    def test_search_multi_collections(self, nb, nq, dim, auto_id, _async):
B
binbin 已提交
1765
        """
1766 1767 1768
        target: test search multi collections of L2
        method: add vectors into 10 collections, and search
        expected: search status ok, the length of result
B
binbin 已提交
1769
        """
1770
        self._connect()
1771 1772
        collection_num = 10
        for i in range(collection_num):
1773
            # 1. initialize with data
1774
            log.info("test_search_multi_collections: search round %d" % (i + 1))
1775 1776 1777
            collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                          auto_id=auto_id,
                                                                          dim=dim)
1778
            # 2. search
1779
            vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
1780
            log.info("test_search_multi_collections: searching %s entities (nq = %s) from collection %s" %
1781 1782
                     (default_limit, nq, collection_w.name))
            collection_w.search(vectors[:nq], default_search_field,
1783
                                default_search_params, default_limit,
1784
                                default_search_exp, _async=_async,
1785
                                check_task=CheckTasks.check_search_results,
1786
                                check_items={"nq": nq,
1787
                                             "ids": insert_ids,
1788 1789
                                             "limit": default_limit,
                                             "_async": _async})
1790

1791
    @pytest.mark.tags(CaseLabel.L2)
1792
    def test_search_concurrent_multi_threads(self, nb, nq, dim, auto_id, _async):
B
binbin 已提交
1793
        """
1794 1795 1796
        target: test concurrent search with multi-processes
        method: search with 10 processes, each process uses dependent connection
        expected: status ok and the returned vectors should be query_records
B
binbin 已提交
1797
        """
1798
        # 1. initialize with data
1799 1800
        threads_num = 10
        threads = []
1801 1802 1803
        collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb,
                                                                      auto_id=auto_id,
                                                                      dim=dim)
1804

B
binbin 已提交
1805
        def search(collection_w):
1806 1807 1808 1809
            vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
            collection_w.search(vectors[:nq], default_search_field,
                                default_search_params, default_limit,
                                default_search_exp, _async=_async,
1810
                                check_task=CheckTasks.check_search_results,
1811
                                check_items={"nq": nq,
1812
                                             "ids": insert_ids,
1813 1814
                                             "limit": default_limit,
                                             "_async": _async})
1815 1816

        # 2. search with multi-processes
B
binbin 已提交
1817
        log.info("test_search_concurrent_multi_threads: searching with %s processes" % threads_num)
1818
        for i in range(threads_num):
B
binbin 已提交
1819
            t = threading.Thread(target=search, args=(collection_w,))
1820 1821 1822 1823 1824
            threads.append(t)
            t.start()
            time.sleep(0.2)
        for t in threads:
            t.join()
1825

1826

1827 1828