Improve check id loic for search test cases (#6340)

Signed-off-by: Binbin Lv <binbin.lv@zilliz.com>
This commit is contained in:
binbin 2021-07-07 17:50:00 +08:00 committed by GitHub
parent 6286d2b951
commit 8f1ec25dbb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 125 additions and 94 deletions

View File

@ -159,6 +159,7 @@ class TestcaseBase(Base):
collection_name = cf.gen_unique_str(prefix)
vectors = []
binary_raw_vectors = []
insert_ids = []
# 1 create collection
default_schema = cf.gen_default_collection_schema()
if is_binary:
@ -173,10 +174,10 @@ class TestcaseBase(Base):
cf.gen_partitions(collection_w, partition_num)
# 3 insert data if specified
if insert_data:
collection_w, vectors, binary_raw_vectors = \
collection_w, vectors, binary_raw_vectors, insert_ids = \
cf.insert_data(collection_w, nb, is_binary, is_all_data_type)
assert collection_w.is_empty is False
assert collection_w.num_entities == nb
collection_w.load()
return collection_w, vectors, binary_raw_vectors
return collection_w, vectors, binary_raw_vectors, insert_ids

View File

@ -151,7 +151,7 @@ class ResponseChecker:
"""
target: check the search results
method: 1. check the query number
2. check the limit(topK)
2. check the limit(topK) and ids
3. check the distance
expected: check the search is ok
"""
@ -176,11 +176,14 @@ class ResponseChecker:
assert len(hits) == check_items["limit"]
assert len(hits.ids) == check_items["limit"]
else:
for i in hits.ids:
assert i in range(check_items["nb"])
log.info("search_results_check: limit (topK) "
"searched for each query is correct")
log.info("search_results_check: search_results_check: checked the searching results")
ids_match = pc.list_contain_check(list(hits.ids),
list(check_items["ids"]))
if ids_match:
log.info("search_results_check: limit (topK) and "
"ids searched for each query are correct")
else:
log.error("search_results_check: ids searched not match")
assert ids_match
return True
@staticmethod

View File

@ -90,6 +90,24 @@ def list_equal_check(param1, param2):
return check_result
def list_contain_check(sublist, superlist):
if not isinstance(sublist, list):
raise Exception("%s isn't list type" % sublist)
if not isinstance(superlist, list):
raise Exception("%s isn't list type" % superlist)
check_result = True
for i in sublist:
if i not in superlist:
check_result = False
break
else:
superlist.remove(i)
if not check_result:
log.error("list_contain_check: List(%s) does not contain list(%s)"
% (str(superlist), str(sublist)))
return check_result
def get_connect_object_name(_list):
""" get the name of the objects that returned by the connection """

View File

@ -350,6 +350,7 @@ def insert_data(collection_w, nb=3000, is_binary=False, is_all_data_type=False):
num = len(par)
vectors = []
binary_raw_vectors = []
insert_ids = []
log.info("insert_data: inserting data into collection %s (num_entities: %s)"
% (collection_w.name, nb))
for i in range(num):
@ -359,11 +360,12 @@ def insert_data(collection_w, nb=3000, is_binary=False, is_all_data_type=False):
binary_raw_vectors.extend(binary_raw_data)
if is_all_data_type:
default_data = gen_dataframe_all_data_type(nb // num)
collection_w.insert(default_data, par[i].name)
insert_res = collection_w.insert(default_data, par[i].name)[0]
insert_ids.extend(insert_res.primary_keys)
vectors.append(default_data)
log.info("insert_data: inserted data into collection %s (num_entities: %s)"
% (collection_w.name, nb))
return collection_w, vectors, binary_raw_vectors
return collection_w, vectors, binary_raw_vectors, insert_ids
def _check_primary_keys(primary_keys, nb):

View File

@ -29,7 +29,7 @@ class TestQueryBase(TestcaseBase):
expected: verify query result
"""
# create collection, insert default_nb, load collection
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
pos = 5
term_expr = f'{ct.default_int64_field_name} in {int_values[:pos]}'
@ -92,7 +92,7 @@ class TestQueryBase(TestcaseBase):
method: query with expr None
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
error = {ct.err_code: 0, ct.err_msg: "The type of expr must be string"}
collection_w.query(None, check_task=CheckTasks.err_res, check_items=error)
@ -104,7 +104,7 @@ class TestQueryBase(TestcaseBase):
method: query with non-string expr, eg 1, [] ..
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
error = {ct.err_code: 0, ct.err_msg: "The type of expr must be string"}
collection_w.query(expr, check_task=CheckTasks.err_res, check_items=error)
@ -116,7 +116,7 @@ class TestQueryBase(TestcaseBase):
method: query with invalid string expr
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
error = {ct.err_code: 1, ct.err_msg: "Invalid expression!"}
collection_w.query(expr, check_task=CheckTasks.err_res, check_items=error)
@ -127,7 +127,7 @@ class TestQueryBase(TestcaseBase):
method: query with TermExpr
expected: query result is correct
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
res = vectors[0].iloc[:2, :2].to_dict('records')
collection_w.query(default_term_expr, check_task=CheckTasks.check_query_results, check_items={exp_res: res})
@ -139,7 +139,7 @@ class TestQueryBase(TestcaseBase):
method: query by term expr with fake field
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
term_expr = 'field in [1, 2]'
error = {ct.err_code: 1, ct.err_msg: "fieldName(field) not found"}
collection_w.query(term_expr, check_task=CheckTasks.err_res, check_items=error)
@ -152,7 +152,7 @@ class TestQueryBase(TestcaseBase):
method: query on float field
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
term_expr = f'{ct.default_float_field_name} in [1., 2.]'
error = {ct.err_code: 1, ct.err_msg: "column is not int64"}
collection_w.query(term_expr, check_task=CheckTasks.err_res, check_items=error)
@ -184,7 +184,7 @@ class TestQueryBase(TestcaseBase):
method: query with wrong keyword term expr
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
expr_1 = f'{ct.default_int64_field_name} inn [1, 2]'
error_1 = {ct.err_code: 1, ct.err_msg: f'unexpected token Identifier("inn")'}
collection_w.query(expr_1, check_task=CheckTasks.err_res, check_items=error_1)
@ -208,7 +208,7 @@ class TestQueryBase(TestcaseBase):
method: query with non-array term expr
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
error = {ct.err_code: 1, ct.err_msg: "right operand of the InExpr must be array"}
collection_w.query(expr, check_task=CheckTasks.err_res, check_items=error)
@ -220,7 +220,7 @@ class TestQueryBase(TestcaseBase):
expected: empty result
"""
term_expr = f'{ct.default_int64_field_name} in []'
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
res, _ = collection_w.query(term_expr)
assert len(res) == 0
@ -232,7 +232,7 @@ class TestQueryBase(TestcaseBase):
method: query with int field and float values
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
int_values = [1., 2.]
term_expr = f'{ct.default_int64_field_name} in {int_values}'
error = {ct.err_code: 1, ct.err_msg: "type mismatch"}
@ -246,7 +246,7 @@ class TestQueryBase(TestcaseBase):
method: query with term expr that has int and float type value
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
int_values = [1, 2.]
term_expr = f'{ct.default_int64_field_name} in {int_values}'
error = {ct.err_code: 1, ct.err_msg: "type mismatch"}
@ -261,7 +261,7 @@ class TestQueryBase(TestcaseBase):
method: query with non-constant array expr
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
term_expr = f'{ct.default_int64_field_name} in [{constant}]'
log.debug(term_expr)
error = {ct.err_code: 1, ct.err_msg: "unsupported leaf node"}
@ -274,7 +274,7 @@ class TestQueryBase(TestcaseBase):
method: query with output field=None
expected: return all fields
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
res, _ = collection_w.query(default_term_expr, output_fields=None)
fields = [ct.default_int64_field_name, ct.default_float_field_name]
assert set(res[0].keys()) == set(fields)
@ -286,7 +286,7 @@ class TestQueryBase(TestcaseBase):
method: query with output one field
expected: return one field
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
res, _ = collection_w.query(default_term_expr, output_fields=[ct.default_int64_field_name])
assert set(res[0].keys()) == set([ct.default_int64_field_name])
@ -297,7 +297,7 @@ class TestQueryBase(TestcaseBase):
method: query with output field=None
expected: return all fields
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
fields = [ct.default_int64_field_name, ct.default_float_field_name]
res, _ = collection_w.query(default_term_expr, output_fields=fields)
assert set(res[0].keys()) == set(fields)
@ -311,7 +311,7 @@ class TestQueryBase(TestcaseBase):
method: specify vec field as output field
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
fields = [[ct.default_float_vec_field_name], [ct.default_int64_field_name, ct.default_float_vec_field_name]]
error = {ct.err_code: 1, ct.err_msg: "Query does not support vector field currently"}
for output_fields in fields:
@ -325,7 +325,7 @@ class TestQueryBase(TestcaseBase):
method: specify int64 primary field as output field
expected: return int64 field
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
res, _ = collection_w.query(default_term_expr, output_fields=[ct.default_int64_field_name])
assert list(res[0].keys()) == [ct.default_int64_field_name]
@ -338,7 +338,7 @@ class TestQueryBase(TestcaseBase):
method: query with not existed output field
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
error = {ct.err_code: 1, ct.err_msg: 'Field int not exist'}
collection_w.query(default_term_expr, output_fields=output_fields, check_task=CheckTasks.err_res, check_items=error)
@ -349,7 +349,7 @@ class TestQueryBase(TestcaseBase):
method: query with empty output fields
expected: return all fields
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
query_res, _ = collection_w.query(default_term_expr, output_fields=[])
fields = [ct.default_int64_field_name, ct.default_float_field_name]
assert list(query_res[0].keys()) == fields
@ -363,7 +363,7 @@ class TestQueryBase(TestcaseBase):
method: query with invalid field fields
expected: raise exception
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
error = {ct.err_code: 0, ct.err_msg: f'Invalid query format. \'output_fields\' must be a list'}
collection_w.query(default_term_expr, output_fields=output_fields, check_task=CheckTasks.err_res, check_items=error)
@ -407,7 +407,7 @@ class TestQueryBase(TestcaseBase):
method: query on default partition
expected: verify query result
"""
collection_w, vectors, _, = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
res = vectors[0].iloc[:2, :2].to_dict('records')
collection_w.query(default_term_expr, partition_names=[ct.default_partition_name],
check_task=CheckTasks.check_query_results, check_items={exp_res: res})
@ -503,7 +503,7 @@ class TestQueryOperation(TestcaseBase):
"""
# init a collection and insert data
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)[0:3]
# query the first row of data
check_vec = vectors[0].iloc[:, [0, 1]][0:1].to_dict('records')
@ -520,7 +520,7 @@ class TestQueryOperation(TestcaseBase):
# init a collection and insert data
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True,
is_binary=True)
is_binary=True)[0:3]
# query the first row of data
check_vec = vectors[0].iloc[:, [0, 1]][0:1].to_dict('records')
@ -535,7 +535,7 @@ class TestQueryOperation(TestcaseBase):
"""
# init a collection and insert data
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)[0:3]
# data preparation
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
@ -569,7 +569,7 @@ class TestQueryOperation(TestcaseBase):
method: query with repeated array value
expected: verify query result
"""
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)[0:3]
int_values = [0, 0, 0, 0]
term_expr = f'{ct.default_int64_field_name} in {int_values}'
res, _ = collection_w.query(term_expr)
@ -583,7 +583,7 @@ class TestQueryOperation(TestcaseBase):
method: query after index
expected: query result is correct
"""
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)[0:3]
default_field_name = ct.default_float_vec_field_name
default_index_params = {"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}}
@ -607,14 +607,15 @@ class TestQueryOperation(TestcaseBase):
limit = 1000
nb_old = 500
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, True, nb_old)
collection_w, vectors, binary_raw_vectors, insert_ids = \
self.init_collection_general(prefix, True, nb_old)
# 2. search for original data after load
vectors_s = [[random.random() for _ in range(ct.default_dim)] for _ in range(ct.default_nq)]
collection_w.search(vectors_s[:ct.default_nq], ct.default_float_vec_field_name,
ct.default_search_params, limit, "int64 >= 0",
check_task=CheckTasks.check_search_results,
check_items={"nq": ct.default_nq, "limit": nb_old, "nb": nb_old})
check_items={"nq": ct.default_nq, "limit": nb_old, "ids": insert_ids})
# check number of entities and that method calls the flush interface
assert collection_w.num_entities == nb_old

View File

@ -157,6 +157,7 @@ class TestCollectionSearch(TestcaseBase):
default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"ids": [],
"limit": 0})
@pytest.mark.tags(CaseLabel.L1)
@ -179,7 +180,6 @@ class TestCollectionSearch(TestcaseBase):
"'anns_field', 'param', and 'limit'" in str(e)
@pytest.mark.tags(CaseLabel.L1)
@pytest.mark.xfail(reason="issue 6257")
def test_search_param_invalid_dim(self):
"""
target: test search with invalid parameter values
@ -197,7 +197,8 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, default_limit, default_search_exp,
check_task=CheckTasks.err_res,
check_items={"err_code": 1,
"err_msg": "UnexpectedError"})
"err_msg": "The dimension of query entities "
"is different from schema"})
@pytest.mark.tags(CaseLabel.L1)
def test_search_param_invalid_metric_type(self):
@ -417,7 +418,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_normal")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
# 2. search
log.info("test_search_normal: searching collection %s" % collection_w.name)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -426,7 +427,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L1)
@ -461,7 +462,7 @@ class TestCollectionSearch(TestcaseBase):
partition_num = 1
nb = 1000
limit = 1000
collection_w = self.init_collection_general(prefix, True, nb, partition_num)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb, partition_num)
# 2. search all the partitions before partition deletion
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
log.info("test_search_before_after_delete: searching before deleting partitions")
@ -469,12 +470,13 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb,
"ids": insert_ids,
"limit": limit})
# 3. delete partitions
log.info("test_search_before_after_delete: deleting a partition")
par = collection_w.partitions
deleted_entity_num = par[partition_num].num_entities
entity_num = nb - deleted_entity_num
collection_w.drop_partition(par[partition_num].name)
log.info("test_search_before_after_delete: deleted a partition")
collection_w.load()
@ -484,7 +486,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb-deleted_entity_num,
"ids": insert_ids[:entity_num],
"limit": limit-deleted_entity_num})
@pytest.mark.tags(CaseLabel.L2)
@ -501,7 +503,7 @@ class TestCollectionSearch(TestcaseBase):
nb = 1000
limit = 1000
partition_num = 1
collection_w = self.init_collection_general(prefix, True, nb, partition_num)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb, partition_num)
# 2. search all the partitions before partition deletion
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
log.info("test_search_partition_after_release_one: searching before deleting partitions")
@ -509,12 +511,13 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb,
"ids": insert_ids,
"limit": limit})
# 3. release one partition
log.info("test_search_partition_after_release_one: releasing a partition")
par = collection_w.partitions
deleted_entity_num = par[partition_num].num_entities
entity_num = nb - deleted_entity_num
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")
@ -524,7 +527,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb-deleted_entity_num,
"ids": insert_ids[:entity_num],
"limit": limit - deleted_entity_num})
@pytest.mark.tags(CaseLabel.L2)
@ -540,7 +543,7 @@ class TestCollectionSearch(TestcaseBase):
# 1. initialize with data
nb = 1000
limit = 1000
collection_w = self.init_collection_general(prefix, True, nb, 1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb, 1)
# 2. search all the partitions before partition deletion
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
log.info("test_search_partition_after_release_all: searching before deleting partitions")
@ -548,7 +551,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb,
"ids": insert_ids,
"limit": limit})
# 3. release all partitions
log.info("test_search_partition_after_release_all: releasing a partition")
@ -561,6 +564,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"ids": [],
"limit": 0})
@pytest.mark.tags(CaseLabel.L2)
@ -575,7 +579,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_collection_after_release_load")
# 1. initialize without data
collection_w = self.init_collection_general(prefix, True, default_nb, 1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, default_nb, 1)
# 2. release collection
collection_w.release()
# 3. Search the pre-released collection after load
@ -586,7 +590,7 @@ class TestCollectionSearch(TestcaseBase):
default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L2)
@ -601,7 +605,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_partition_after_release_load")
# 1. initialize without data
collection_w = self.init_collection_general(prefix, True, default_nb, 1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, default_nb, 1)
# 2. release collection
log.info("test_search_partition_after_release_load: releasing a partition")
par = collection_w.partitions
@ -617,7 +621,7 @@ class TestCollectionSearch(TestcaseBase):
limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": limit})
# 4. Search the pre-released partition after load
collection_w.search(vectors[:default_nq], default_search_field, default_search_params,
@ -625,7 +629,7 @@ class TestCollectionSearch(TestcaseBase):
[par[1].name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": par[1].num_entities,
"ids": insert_ids[par[0].num_entities:],
"limit": par[1].num_entities})
@pytest.mark.tags(CaseLabel.L2)
@ -641,7 +645,7 @@ class TestCollectionSearch(TestcaseBase):
# 1. initialize with data
collection_w = self.init_collection_general(prefix)[0]
# 2. insert data
cf.insert_data(collection_w, default_nb)
insert_ids = cf.insert_data(collection_w, default_nb)[3]
# 3. load data
collection_w.load()
# 4. flush and load
@ -653,7 +657,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L1)
@ -670,7 +674,7 @@ class TestCollectionSearch(TestcaseBase):
# 1. initialize with data
limit = 1000
nb_old = 500
collection_w = self.init_collection_general(prefix, True, nb_old)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb_old)
# 2. search for original data after load
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
log.info("test_search_new_data: searching for original data after load")
@ -678,17 +682,18 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb_old,
"ids": insert_ids,
"limit": nb_old})
# 3. insert new data
nb_new = 300
cf.insert_data(collection_w, nb_new, False)
insert_ids_new = cf.insert_data(collection_w, nb_new)[3]
insert_ids.extend(insert_ids_new)
# 4. search for new data without load
collection_w.search(vectors[:default_nq], default_search_field,
default_search_params, limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb_old+nb_new,
"ids": insert_ids,
"limit": nb_old+nb_new})
@pytest.mark.tags(CaseLabel.L2)
@ -703,7 +708,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_after_different_index")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, partition_num=1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, partition_num=1)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
# 2. create different index
log.info("test_search_after_different_index: Creating index-%s" % index)
@ -716,7 +721,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L2)
@ -731,7 +736,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_after_index_different_metric_type")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, partition_num=1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, partition_num=1)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
# 2. create different index
log.info("test_search_after_different_index: Creating index-%s" % index)
@ -744,7 +749,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L1)
@ -757,7 +762,7 @@ class TestCollectionSearch(TestcaseBase):
log.info("Test case of search interface: test_search_collection_multiple_times")
# 1. initialize with data
search_num = 5
collection_w = self.init_collection_general(prefix, True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
# 2. search for multiple times
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
for i in range(search_num):
@ -766,7 +771,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L1)
@ -778,7 +783,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_index_one_partition")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, partition_num=1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, partition_num=1)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
# 2. create index
default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
@ -792,7 +797,7 @@ class TestCollectionSearch(TestcaseBase):
[par[1].name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": par[1].num_entities,
"ids": insert_ids[par[0].num_entities:],
"limit": par[1].num_entities})
@pytest.mark.tags(CaseLabel.L1)
@ -804,7 +809,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_index_partitions")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, partition_num=1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, partition_num=1)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
# 2. create index
default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
@ -819,7 +824,7 @@ class TestCollectionSearch(TestcaseBase):
[par[0].name, par[1].name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": limit})
@pytest.mark.tags(CaseLabel.L2)
@ -834,7 +839,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_index_partitions_fuzzy")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, partition_num=1)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, partition_num=1)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
# 2. create index
default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": "L2"}
@ -845,14 +850,14 @@ class TestCollectionSearch(TestcaseBase):
nb = default_nb
par = collection_w.partitions
if partition_names == ["search(.*)"]:
nb = par[1].num_entities
insert_ids = insert_ids[par[0].num_entities:]
limit = par[1].num_entities
collection_w.search(vectors[:default_nq], default_search_field,
default_search_params, limit, default_search_exp,
partition_names,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb,
"ids": insert_ids,
"limit": limit})
@pytest.mark.tags(CaseLabel.L2)
@ -882,6 +887,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_exp, [partition_name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"ids": [],
"limit": 0})
@pytest.mark.tags(CaseLabel.L2)
@ -893,7 +899,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_binary_jaccard_flat_index")
# 1. initialize with binary data
collection_w, _, binary_raw_vector = \
collection_w, _, binary_raw_vector, _ = \
self.init_collection_general(prefix, True, 2, is_binary=True)
# 2. create index
default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "JACCARD"}
@ -917,7 +923,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_binary_hamming_flat_index")
# 1. initialize with binary data
collection_w, _, binary_raw_vector = \
collection_w, _, binary_raw_vector, _ = \
self.init_collection_general(prefix, True, 2, is_binary=True)
# 2. create index
default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "HAMMING"}
@ -940,7 +946,7 @@ class TestCollectionSearch(TestcaseBase):
expected: the return distance equals to the computed value
"""
# 1. initialize with binary data
collection_w, _, binary_raw_vector = \
collection_w, _, binary_raw_vector, _ = \
self.init_collection_general(prefix, True, 2, is_binary=True)
# 2. create index
default_index = {"index_type": "BIN_IVF_FLAT", "params": {"nlist": 128}, "metric_type": "TANIMOTO"}
@ -968,7 +974,7 @@ class TestCollectionSearch(TestcaseBase):
log.info("Test case of search interface: test_search_with_expression")
# 1. initialize with data
nb = 1000
collection_w = self.init_collection_general(prefix, True, nb)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, nb)
# 2. create index
index_param = {"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 100}}
collection_w.create_index("float_vector", index_param)
@ -979,7 +985,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, nb, expression,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": nb,
"ids": insert_ids,
"limit": limit})
@pytest.mark.tags(CaseLabel.L1)
@ -991,7 +997,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_with_output_fields_empty")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
# 2. search
log.info("test_search_with_output_fields_empty: Searching collection %s" % collection_w.name)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -1000,7 +1006,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_exp, output_fields=[],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})[0]
assert len(res[0][0].entity._row_data) == 0
@ -1013,7 +1019,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_with_output_fields_not_exist")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
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)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -1022,7 +1028,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_exp, output_fields=["int63"],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})[0]
assert len(res[0][0].entity._row_data) == 0
@ -1035,7 +1041,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_with_output_field")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
# 2. search
log.info("test_search_with_output_field: Searching collection %s" % collection_w.name)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -1044,7 +1050,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_exp, output_fields=[default_int64_field_name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})[0]
assert len(res[0][0].entity._row_data) != 0
assert default_int64_field_name in res[0][0].entity._row_data
@ -1058,7 +1064,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_with_output_fields")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, is_all_data_type=True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, is_all_data_type=True)
# 2. search
log.info("test_search_with_output_fields: Searching collection %s" % collection_w.name)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -1069,7 +1075,7 @@ class TestCollectionSearch(TestcaseBase):
default_float_field_name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})[0]
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
@ -1083,7 +1089,7 @@ class TestCollectionSearch(TestcaseBase):
"""
log.info("Test case of search interface: test_search_expression_all_data_type")
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True, is_all_data_type=True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True, is_all_data_type=True)
# 2. search
log.info("test_search_expression_all_data_type: Searching collection %s" % collection_w.name)
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -1095,7 +1101,7 @@ class TestCollectionSearch(TestcaseBase):
default_float_field_name],
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})[0]
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
@ -1113,7 +1119,7 @@ class TestCollectionSearch(TestcaseBase):
for i in range(collection_num):
# 1. initialize with data
log.info("test_search_multi_collections: search round %d" % (i + 1))
collection_w = self.init_collection_general(prefix, True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
# 2. search
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
log.info("test_search_multi_collections: searching %s entities (nq = %s) from collection %s" %
@ -1123,7 +1129,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
@pytest.mark.tags(CaseLabel.L2)
@ -1137,7 +1143,7 @@ class TestCollectionSearch(TestcaseBase):
threads_num = 10
threads = []
# 1. initialize with data
collection_w = self.init_collection_general(prefix, True)[0]
collection_w, _, _, insert_ids = self.init_collection_general(prefix, True)
def search(collection_w):
vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)]
@ -1145,7 +1151,7 @@ class TestCollectionSearch(TestcaseBase):
default_search_params, default_limit, default_search_exp,
check_task=CheckTasks.check_search_results,
check_items={"nq": default_nq,
"nb": default_nb,
"ids": insert_ids,
"limit": default_limit})
# 2. search with multi-processes