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])