def test_query_after_search(self, connect, collection):
"""
target: test query after search
method: query after search
expected: query result is correct
"""
entities, ids = init_data(connect, collection)
assert len(ids) == ut.default_nb
top_k = 10
nq = 2
query, _ = ut.gen_query_vectors(ut.default_float_vec_field_name,
entities,
top_k=top_k,
nq=nq)
connect.load_collection(collection)
search_res = connect.search(collection, query)
assert len(search_res) == nq
assert len(search_res[0]) == top_k
term_expr = f'{default_int_field_name} in {ids[:default_pos]}'
res = connect.query(collection, term_expr, output_fields=["*", "%"])
logging.getLogger().info(res)
assert len(res) == default_pos
for _id, index in enumerate(ids[:default_pos]):
if res[index][default_int_field_name] == entities[0]["values"][
index]:
assert res[index][default_float_field_name] == entities[1][
"values"][index]
ut.assert_equal_vector(
res[index][ut.default_float_vec_field_name],
entities[2]["values"][index])
def test_query_valid(self, connect, collection):
"""
target: test query
method: query with term expr
expected: verify query result
"""
entities, ids = init_data(connect, collection)
assert len(ids) == ut.default_nb
connect.load_collection(collection)
term_expr = f'{default_int_field_name} in {ids[:default_pos]}'
res = connect.query(collection, term_expr, output_fields=["*", "%"])
assert len(res) == default_pos
for _id, index in enumerate(ids[:default_pos]):
if res[index][default_int_field_name] == entities[0]["values"][
index]:
assert res[index][default_float_field_name] == entities[1][
"values"][index]
res = connect.query(collection,
term_expr,
output_fields=[ut.default_float_vec_field_name])
assert len(res) == default_pos
for _id, index in enumerate(ids[:default_pos]):
if res[index][default_int_field_name] == entities[0]["values"][
index]:
ut.assert_equal_vector(
res[index][ut.default_float_vec_field_name],
entities[2]["values"][index])
def test_query_default_partition(self, connect, collection):
"""
target: test query on default partition
method: query on default partition
expected: verify query result
"""
entities, ids = init_data(connect, collection)
assert len(ids) == ut.default_nb
connect.load_collection(collection)
res = connect.query(collection, default_term_expr, partition_names=[ut.default_partition_name], output_fields=["*", "%"])
for _id, index in enumerate(ids[:default_pos]):
if res[index][default_int_field_name] == entities[0]["values"][index]:
assert res[index][default_float_field_name] == entities[1]["values"][index]
ut.assert_equal_vector(res[index][ut.default_float_vec_field_name], entities[2]["values"][index])
def test_query_expr_single_term_array(self, connect, collection):
"""
target: test query with single array term expr
method: query with single array value
expected: query result is one entity
"""
entities, ids = init_data(connect, collection)
assert len(ids) == ut.default_nb
connect.load_collection(collection)
term_expr = f'{default_int_field_name} in [0]'
res = connect.query(collection, term_expr, output_fields=["*", "%"])
assert len(res) == 1
assert res[0][default_int_field_name] == entities[0]["values"][0]
assert res[0][default_float_field_name] == entities[1]["values"][0]
ut.assert_equal_vector(res[0][ut.default_float_vec_field_name], entities[2]["values"][0])
def test_query_after_index(self, connect, collection, get_simple_index):
"""
target: test query after creating index
method: query after index
expected: query result is correct
"""
entities, ids = init_data(connect, collection)
assert len(ids) == ut.default_nb
connect.create_index(collection, ut.default_float_vec_field_name, get_simple_index)
connect.load_collection(collection)
term_expr = f'{default_int_field_name} in {ids[:default_pos]}'
res = connect.query(collection, term_expr, output_fields=["*", "%"])
logging.getLogger().info(res)
assert len(res) == default_pos
for _id, index in enumerate(ids[:default_pos]):
if res[index][default_int_field_name] == entities[0]["values"][index]:
assert res[index][default_float_field_name] == entities[1]["values"][index]
ut.assert_equal_vector(res[index][ut.default_float_vec_field_name], entities[-1]["values"][index])