From b748d8af5dc9b6a3a700fc296ba2c26ffb6beebe Mon Sep 17 00:00:00 2001 From: nico <109071306+NicoYuan1986@users.noreply.github.com> Date: Fri, 14 Jun 2024 17:51:57 +0800 Subject: [PATCH] test: add different datatype for some functions (#33869) Signed-off-by: nico --- tests/python_client/testcases/test_query.py | 35 ++++++++++++ tests/python_client/testcases/test_search.py | 60 ++++++++++++++++---- 2 files changed, 84 insertions(+), 11 deletions(-) diff --git a/tests/python_client/testcases/test_query.py b/tests/python_client/testcases/test_query.py index efad358ec2..d1e25351c5 100644 --- a/tests/python_client/testcases/test_query.py +++ b/tests/python_client/testcases/test_query.py @@ -2363,6 +2363,20 @@ class TestQueryOperation(TestcaseBase): res, _ = collection_w.query(default_term_expr, output_fields=[ct.default_binary_vec_field_name]) assert res[0].keys() == set(fields) + @pytest.mark.tags(CaseLabel.L1) + @pytest.mark.parametrize("vector_data_type", ["FLOAT_VECTOR", "FLOAT16_VECTOR", "BFLOAT16_VECTOR"]) + def test_query_output_all_vector_type(self, vector_data_type): + """ + target: test query output different vector type + method: create index and specify vec field as output field + expected: return primary field and vec field + """ + collection_w, vectors = self.init_collection_general(prefix, True, + vector_data_type=vector_data_type)[0:2] + fields = [ct.default_int64_field_name, ct.default_float_vec_field_name] + res, _ = collection_w.query(default_term_expr, output_fields=[ct.default_float_vec_field_name]) + assert res[0].keys() == set(fields) + @pytest.mark.tags(CaseLabel.L2) def test_query_partition_repeatedly(self): """ @@ -3741,6 +3755,27 @@ class TestQueryIterator(TestcaseBase): check_items={"count": ct.default_nb - offset, "batch_size": batch_size}) + @pytest.mark.tags(CaseLabel.L2) + @pytest.mark.parametrize("vector_data_type", ["FLOAT_VECTOR", "FLOAT16_VECTOR", "BFLOAT16_VECTOR"]) + def test_query_iterator_output_different_vector_type(self, vector_data_type): + """ + target: test query iterator with output fields + method: 1. query iterator output different vector type + 2. check the result, expect pk + expected: query successfully + """ + # 1. initialize with data + batch_size = 400 + collection_w = self.init_collection_general(prefix, True, + vector_data_type=vector_data_type)[0] + # 2. query iterator + expr = "int64 >= 0" + collection_w.query_iterator(batch_size, expr=expr, + output_fields=[ct.default_float_vec_field_name], + check_task=CheckTasks.check_query_iterator, + check_items={"count": ct.default_nb, + "batch_size": batch_size}) + @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("batch_size", [10, 100, 777, 2000]) def test_query_iterator_with_different_batch_size(self, batch_size): diff --git a/tests/python_client/testcases/test_search.py b/tests/python_client/testcases/test_search.py index 829f3b7f86..e749d9fa8b 100644 --- a/tests/python_client/testcases/test_search.py +++ b/tests/python_client/testcases/test_search.py @@ -5074,6 +5074,8 @@ class TestSearchBase(TestcaseBase): # 4. check the search results for i in range(default_nq): assert res_ip[i].ids == res_cosine[i].ids + log.info(res_cosine[i].distances) + log.info(res_ip[i].distances) @pytest.mark.tags(CaseLabel.L2) def test_search_without_connect(self): @@ -5874,6 +5876,10 @@ class TestSearchPagination(TestcaseBase): def enable_dynamic_field(self, request): yield request.param + @pytest.fixture(scope="function", params=["FLOAT_VECTOR", "FLOAT16_VECTOR", "BFLOAT16_VECTOR"]) + def vector_data_type(self, request): + yield request.param + """ ****************************************************************** # The following are valid base cases @@ -5897,10 +5903,8 @@ class TestSearchPagination(TestcaseBase): collection_w = self.init_collection_general(prefix, True, auto_id=auto_id, dim=default_dim, enable_dynamic_field=enable_dynamic_field)[0] # 2. search pagination with offset - search_param = {"metric_type": "COSINE", - "params": {"nprobe": 10}, "offset": offset} - vectors = [[random.random() for _ in range(default_dim)] - for _ in range(default_nq)] + search_param = {"metric_type": "COSINE", "params": {"nprobe": 10}, "offset": offset} + vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)] search_res = collection_w.search(vectors[:default_nq], default_search_field, search_param, limit, default_search_exp, _async=_async, @@ -5937,10 +5941,8 @@ class TestSearchPagination(TestcaseBase): self.init_collection_general(prefix, True, auto_id=auto_id, dim=default_dim, enable_dynamic_field=enable_dynamic_field)[0:4] # 2. search - search_param = {"metric_type": "COSINE", - "params": {"nprobe": 10}, "offset": offset} - vectors = [[random.random() for _ in range(default_dim)] - for _ in range(default_nq)] + search_param = {"metric_type": "COSINE", "params": {"nprobe": 10}, "offset": offset} + vectors = [[random.random() for _ in range(default_dim)] for _ in range(default_nq)] output_fields = [default_string_field_name, default_float_field_name] search_res = collection_w.search(vectors[:default_nq], default_search_field, search_param, default_limit, @@ -5999,6 +6001,40 @@ class TestSearchPagination(TestcaseBase): assert sorted(search_res[0].distances, key=numpy.float32) == sorted( res[0].distances[offset:], key=numpy.float32) + @pytest.mark.tags(CaseLabel.L1) + def test_search_all_vector_type_with_pagination(self, vector_data_type): + """ + target: test search with pagination using different vector datatype + method: 1. connect and create a collection + 2. search pagination with offset + 3. search with offset+limit + 4. compare with the search results whose corresponding ids should be the same + expected: search successfully and ids is correct + """ + # 1. create a collection + auto_id = False + enable_dynamic_field = True + offset = 100 + limit = 20 + collection_w = self.init_collection_general(prefix, True, auto_id=auto_id, dim=default_dim, + enable_dynamic_field=enable_dynamic_field, + vector_data_type=vector_data_type)[0] + # 2. search pagination with offset + search_param = {"metric_type": "COSINE", "params": {"nprobe": 10}, "offset": offset} + vectors = cf.gen_vectors_based_on_vector_type(default_nq, default_dim, vector_data_type) + search_res = collection_w.search(vectors[:default_nq], default_search_field, + search_param, limit, + default_search_exp, + check_task=CheckTasks.check_search_results, + check_items={"nq": default_nq, + "limit": limit})[0] + # 3. search with offset+limit + res = collection_w.search(vectors[:default_nq], default_search_field, default_search_params, + limit + offset, default_search_exp)[0] + res_distance = res[0].distances[offset:] + # assert sorted(search_res[0].distances, key=numpy.float32) == sorted(res_distance, key=numpy.float32) + assert set(search_res[0].ids) == set(res[0].ids[offset:]) + @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("limit", [100, 3000, 10000]) def test_search_with_pagination_topK(self, limit, _async): @@ -9854,7 +9890,8 @@ class TestSearchIterator(TestcaseBase): """ Test case of search iterator """ @pytest.mark.tags(CaseLabel.L1) - def test_search_iterator_normal(self): + @pytest.mark.parametrize("vector_data_type", ["FLOAT_VECTOR", "FLOAT16_VECTOR", "BFLOAT16_VECTOR"]) + def test_search_iterator_normal(self, vector_data_type): """ target: test search iterator normal method: 1. search iterator @@ -9863,12 +9900,13 @@ class TestSearchIterator(TestcaseBase): """ # 1. initialize with data dim = 128 - collection_w = self.init_collection_general( - prefix, True, dim=dim, is_index=False)[0] + collection_w = self.init_collection_general(prefix, True, dim=dim, is_index=False, + vector_data_type=vector_data_type)[0] collection_w.create_index(field_name, {"metric_type": "L2"}) collection_w.load() # 2. search iterator search_params = {"metric_type": "L2"} + vectors = cf.gen_vectors_based_on_vector_type(1, dim, vector_data_type) batch_size = 200 collection_w.search_iterator(vectors[:1], field_name, search_params, batch_size, check_task=CheckTasks.check_search_iterator,