From c10cf53b4be579e1fe5bb75dba0527926ee1b3cd Mon Sep 17 00:00:00 2001 From: jiamingli-maker Date: Sun, 4 Jan 2026 18:57:22 +0800 Subject: [PATCH] test: Add HNSW_PRQ test cases and fix HNSW_PQ (#46680) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /kind improvement - Core invariant: index parameter validation and test expectations for the HNSW-family must be explicit, consistent, and deterministic — this PR enforces that by adding exhaustive parameter matrices for HNSW_PRQ (tests/python_client/testcases/indexes/{idx_hnsw_prq.py, test_hnsw_prq.py}) and normalizing expectations in idx_hnsw_pq.py via a shared success variable. - Logic removed / simplified: brittle, ad-hoc string expectations were consolidated — literal "success" occurrences were replaced with a single success variable and ambiguous short error messages were replaced by the canonical descriptive error text; this reduces duplicated assertion logic in tests and removes dependence on fragile, truncated messages. - Bug fix (tests): corrected HNSW_PQ test expectations to assert the full, authoritative error for invalid PQ m ("The dimension of the vector (dim) should be a multiple of the number of subquantizers (m).") and aligned HNSW_PRQ test matrices (idx_hnsw_prq.py) to the same explicit expectations — the change targets test assertions only and fixes false negatives caused by mismatched messages. - No data loss or behavior regression: only test code is added/modified (tests/python_client/testcases/indexes/*). Production code paths remain unmodified — collection creation, insert/flush, client.create_index, wait_for_index_ready, load_collection, search, and client.describe_index are invoked by tests but not changed; therefore persisted data, index artifacts, and runtime behavior are unaffected. Signed-off-by: zilliz --- .../testcases/indexes/idx_hnsw_pq.py | 6 +- .../testcases/indexes/idx_hnsw_prq.py | 591 ++++++++++++++++++ .../testcases/indexes/test_hnsw_prq.py | 258 ++++++++ 3 files changed, 852 insertions(+), 3 deletions(-) create mode 100644 tests/python_client/testcases/indexes/idx_hnsw_prq.py create mode 100644 tests/python_client/testcases/indexes/test_hnsw_prq.py diff --git a/tests/python_client/testcases/indexes/idx_hnsw_pq.py b/tests/python_client/testcases/indexes/idx_hnsw_pq.py index 39f210e87b..389948aaee 100644 --- a/tests/python_client/testcases/indexes/idx_hnsw_pq.py +++ b/tests/python_client/testcases/indexes/idx_hnsw_pq.py @@ -373,12 +373,12 @@ class HNSW_PQ: { "description": "Refine Disabled", "params": {"M": 16, "efConstruction": 200, "m": 32, "nbits": 8}, - "expected": "success" + "expected": success }, { "description": "Minimum Boundary Combination", "params": {"M": 2, "efConstruction": 1, "m": 1, "nbits": 1, "refine": True, "refine_type": "SQ8"}, - "expected": "success" + "expected": success }, { "description": "Maximum Boundary Combination", @@ -403,7 +403,7 @@ class HNSW_PQ: { "description": "Invalid PQ m (not divisor of dimension)", "params": {"M": 16,"efConstruction": 200,"m": 7, "nbits": 8, "refine": True, "refine_type": "FP32"}, - "expected": {"err_code": 999, "err_msg": "m"} + "expected": {"err_code": 999, "err_msg": "The dimension of the vector (dim) should be a multiple of the number of subquantizers (m)."} }, ] diff --git a/tests/python_client/testcases/indexes/idx_hnsw_prq.py b/tests/python_client/testcases/indexes/idx_hnsw_prq.py new file mode 100644 index 0000000000..107b410e23 --- /dev/null +++ b/tests/python_client/testcases/indexes/idx_hnsw_prq.py @@ -0,0 +1,591 @@ +from pymilvus import DataType + +success = "success" + + +class HNSW_PRQ: + supported_vector_types = [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + DataType.INT8_VECTOR + ] + + supported_metrics = ['L2', 'IP', 'COSINE'] + + build_params = [ + + # M params test + { + "description": "Minimum Boundary Test", + "params": {"M": 2}, + "expected": success + }, + { + "description": "Maximum Boundary Test", + "params": {"M": 2048}, + "expected": success + }, + { + "description": "Out of Range Test - Negative", + "params": {"M": -1}, + "expected": {"err_code": 1100, "err_msg": "param 'M' (-1) should be in range [2, 2048]"} + }, + { + "description": "Out of Range Test - Too Large", + "params": {"M": 2049}, + "expected": {"err_code": 1100, "err_msg": "param 'M' (2049) should be in range [2, 2048]"} + }, + { + "description": "String Type Test will ignore the wrong type", + "params": {"M": "16"}, + "expected": success + }, + { + "description": "Float Type Test", + "params": {"M": 16.0}, + "expected": {"err_code": 1100, "err_msg": "wrong data type in json"} + }, + { + "description": "Boolean Type Test", + "params": {"M": True}, + "expected": {"err_code": 1100, "err_msg": "invalid integer value, key: 'M', value: 'True': invalid parameter"} + }, + { + "description": "None Type Test, use default value", + "params": {"M": None}, + "expected": success + }, + { + "description": "List Type Test", + "params": {"M": [16]}, + "expected": {"err_code": 1100, "err_msg": "invalid integer value, key: 'M', value: '[16]': invalid parameter"} + }, + { + "description": "Nested dict in params", + "params": {"M": {"value": 16}}, + "expected": {"err_code": 1100, "err_msg": "invalid integer value"} + }, + # efConstruction params test + { + "description": "Minimum Boundary Test", + "params": {"efConstruction": 1}, + "expected": success + }, + { + "description": "Large Value Test", + "params": {"efConstruction": 10000}, + "expected": success + }, + { + "description": "Out of Range Test - Negative", + "params": {"efConstruction": -1}, + "expected": {"err_code": 1100, "err_msg": "param 'efConstruction' (-1) should be in range [1, 2147483647]"} + }, + { + "description": "String Type Test will ignore the wrong type", + "params": {"efConstruction": "100"}, + "expected": success + }, + { + "description": "Float Type Test", + "params": {"efConstruction": 100.0}, + "expected": {"err_code": 1100, "err_msg": "wrong data type in json"} + }, + { + "description": "Boolean Type Test", + "params": {"efConstruction": True}, + "expected": {"err_code": 1100, "err_msg": "invalid integer value, key: 'efConstruction', value: 'True': invalid parameter"} + }, + { + "description": "None Type Test, use default value", + "params": {"efConstruction": None}, + "expected": success + }, + { + "description": "List Type Test", + "params": {"efConstruction": [100]}, + "expected": {"err_code": 1100, "err_msg": "invalid integer value, key: 'efConstruction', value: '[100]': invalid parameter"} + }, + { + "description": "Nested List in Params", + "params": {"efConstruction": [[100]]}, + "expected": {"err_code": 1100, "err_msg": "invalid integer value"} + }, + # m params test + { + "description": "Minimum Boundary Test", + "params": {"m": 1}, + "expected": success + }, + # timeout + # { + # "description": "Half of Dimension Value Test", + # "params": {"m": 64}, + # "expected": success + # }, + # timeout + # { + # "description": "Maximum Boundary Test (Dimension)", + # "params": {"m": 128}, + # "expected": success + # }, + { + "description": "Negative Value Test", + "params": {"m": -1}, + "expected": { + "err_code": 1100, + "err_msg": "Out of range in json: param 'm' (-1) should be in range [1, 65536]" + } + }, + { + "description": "Larger Value Test", + "params": {"m": 256}, + "expected": { + "err_code": 1100, + "err_msg": "The dimension of a vector (dim) should be a multiple of the number of subquantizers (m)." + } + }, + { + "description": "Not Divisible by Dimension Value Test", + "params": {"m": 7}, + "expected": { + "err_code": 1100, + "err_msg": "The dimension of a vector (dim) should be a multiple of the number of subquantizers (m)." + } + }, + { + "description": "String Type Test", + "params": {"m": "16"}, + "expected": success + }, + { + "description": "Float Type Test", + "params": {"m": 16.0}, + "expected": { + "err_code": 1100, + "err_msg": "wrong data type in json, key: 'm', value: '16.0': invalid parameter" + } + }, + { + "description": "Boolean Type Test", + "params": {"m": True}, + "expected": { + "err_code": 1100, + "err_msg": "invalid integer value" + } + }, + { + "description": "List Type Test", + "params": {"m": [16]}, + "expected": { + "err_code": 1100, + "err_msg": "invalid integer value" + } + }, + { + "description": "None Type Test", + "params": {"m": None}, + "expected": success + }, + # nbits params test + { + "description": "Minimum Boundary Test", + "params": {"nbits": 1}, + "expected": success + }, + { + "description": "Default Value Test", + "params": {"nbits": 8}, + "expected": success + }, + { + "description": "Maximum Boundary Test", + "params": {"nbits": 10}, + "expected": success + }, + { + "description": "Negative Value Test", + "params": {"nbits": -1}, + "expected": { + "err_code": 1100, + "err_msg": "Out of range in json: param 'nbits' (-1) should be in range [1, 24]: invalid parameter" + } + }, + { + "description": "Large Value Test", + "params": {"nbits": 25}, + "expected": { + "err_code": 1100, + "err_msg": "Out of range in json: param 'nbits' (25) should be in range [1, 24]" + } + }, + { + "description": "String Type Test", + "params": {"nbits": "8"}, + "expected": success + }, + { + "description": "Float Type Test", + "params": {"nbits": 8.0}, + "expected": { + "err_code": 1100, + "err_msg": "wrong data type in json, key: 'nbits', value: '8.0'" + } + }, + { + "description": "Boolean Type Test", + "params": {"nbits": True}, + "expected": { + "err_code": 1100, + "err_msg": "invalid integer value" + } + }, + { + "description": "List Type Test", + "params": {"nbits": [8]}, + "expected": { + "err_code": 1100, + "err_msg": "invalid integer value" + } + }, + { + "description": "None Type Test", + "params": {"nbits": None}, + "expected": success + }, + # nrq params test + { + "description": "Minimum Boundary Test", + "params": {"nrq": 1}, + "expected": success + }, + { + "description": "Default Value Test", + "params": {"nrq": 2}, + "expected": success + }, + { + "description": "Maximum Boundary Test", + "params": {"nrq": 16}, + "expected": success + }, + { + "description": "Negative Value Test", + "params": {"nrq": -1}, + "expected": { + "err_code": 1100, + "err_msg": "Out of range in json: param 'nrq' (-1) should be in range [1, 16]" + } + }, + { + "description": "Larger Value Test", + "params": {"nrq": 17}, + "expected": { + "err_code": 1100, + "err_msg": "Out of range in json: param 'nrq' (17) should be in range [1, 16]" + } + }, + { + "description": "String Type Test", + "params": {"nrq": "4"}, + "expected": success + }, + { + "description": "Float Type Test", + "params": {"nrq": 4.0}, + "expected": { + "err_code": 1100, + "err_msg": "wrong data type in json, key: 'nrq', value: '4.0': invalid parameter" + } + }, + { + "description": "Boolean Type Test", + "params": {"nrq": True}, + "expected": { + "err_code": 1100, + "err_msg": "invalid integer value, key: 'nrq', value: 'True': invalid parameter" + } + }, + { + "description": "None Type Test", + "params": {"nrq": None}, + "expected": success + }, + { + "description": "List Type Test", + "params": {"nrq": [2]}, + "expected": { + "err_code": 1100, + "err_msg": "invalid integer value, key: 'nrq', value: '[2]': invalid parameter" + } + }, + + + # refine params test + { + "description": "refine = True", + "params": {"refine": True}, + "expected": success + }, + { + "description": "String Type Test", + "params": {"refine": "true"}, + "expected": success + }, + { + "description": "Invalid String Type Test", + "params": {"refine": "test"}, + "expected": {"err_code": 1100, "err_msg": "should be a boolean: invalid parameter"} + + }, + { + "description": "Integer Type Test", + "params": {"refine": 1}, + "expected": {"err_code": 1100, "err_msg": "should be a boolean: invalid parameter"} + }, + { + "description": "Float Type Test", + "params": {"refine": 1.0}, + "expected": {"err_code": 1100, "err_msg": "should be a boolean: invalid parameter"} + }, + { + "description": "List Type Test", + "params": {"refine": [True]}, + "expected": {"err_code": 1100, "err_msg": "should be a boolean: invalid parameter"} + }, + { + "description": "None Type Test, use default value", + "params": {"refine": None}, + "expected": success + }, + + # refine_type params test + { + "description": "Valid refine_type - SQ6", + "params": {"refine_type": "SQ6"}, + "expected": success + }, + { + "description": "Valid refine_type - SQ8", + "params": {"refine_type": "SQ8"}, + "expected": success + }, + { + "description": "Valid refine_type - BF16", + "params": {"refine_type": "BF16"}, + "expected": success + }, + { + "description": "Valid refine_type - FP16", + "params": {"refine_type": "FP16"}, + "expected": success + }, + { + "description": "Valid refine_type - FP32", + "params": {"refine_type": "FP32"}, + "expected": success + }, + { + "description": "Out of Range Test - unknown value", + "params": {"refine_type": "INT8"}, + "expected": {"err_code": 1100, "err_msg": "invalid refine type : INT8, optional types are [sq6, sq8, fp16, bf16, fp32, flat]: invalid parameter"} + }, + { + "description": "Integer Type Test", + "params": {"refine_type": 1}, + "expected": {"err_code": 1100, "err_msg": "invalid refine type : 1, optional types are [sq6, sq8, fp16, bf16, fp32, flat]: invalid parameter"} + }, + { + "description": "Float Type Test", + "params": {"refine_type": 1.0}, + "expected": {"err_code": 1100, "err_msg": "invalid refine type : 1.0, optional types are [sq6, sq8, fp16, bf16, fp32, flat]: invalid parameter"} + }, + { + "description": "List Type Test", + "params": {"refine_type": ["FP16"]}, + "expected": {"err_code": 1100, "err_msg": "['FP16'], optional types are [sq6, sq8, fp16, bf16, fp32, flat]: invalid parameter"} + }, + { + "description": "None Type Test, use default value", + "params": {"refine_type": None}, + "expected": success + }, + { + "description": "refine_type lower precision than sq_type but refine disabled", + "params": {"sq_type": "FP16", "refine_type": "SQ8"}, + "expected": success + }, + { + "description": "refine_type lower than sq_type", + "params": {"sq_type": "FP16", "refine_type": "SQ8", "refine": True}, + "expected": success + }, + # combination params test + { + "description": "empty dict params", + "params": {}, + "expected": success + }, + { + "description": "All optional parameters None", + "params": {"M": None, "efConstruction": None, "m": None, "nbits":None, "nrq":None,"refine": None, "refine_type": None}, + "expected": success + }, + { + "description": "Typical valid combination", + "params": {"M": 16, "efConstruction": 200, "m": 32, "nbits": 8, "nrq":1,"refine": True, "refine_type": "FP16"}, + "expected": success + }, + { + "description": "Refine Disabled", + "params": {"M": 16, "efConstruction": 200, "m": 32, "nbits": 8,"nrq": 1}, + "expected": success + }, + { + "description": "Minimum Boundary Combination", + "params": {"M": 2, "efConstruction": 1, "m": 1, "nbits": 1, "nrq":1, "refine": True, "refine_type": "SQ8"}, + "expected": success + }, + { + "description": "Maximum Boundary Combination", + "params": {"M": 2048, "efConstruction": 10000, "m": 128, "nbits": 8, "nrq":1, "refine": True, "refine_type": "FP32"}, + "expected": success + }, + { + "description": "Unknown extra parameter in combination", + "params": {"M": 16, "efConstruction": 200, "m": 32, "nbits": 8, "nrq":1, "refine": True, "refine_type": "FP16", "unknown_param": "nothing"}, + "expected": success + }, + { + "description": "Partial parameters set (M + m only)", + "params": {"M": 32, "m": 32}, + "expected": success + }, + { + "description": "Partial parameters set (efConstruction + nbits only)", + "params": {"efConstruction": 500,"nbits": 8}, + "expected": success + }, + { + "description": "Invalid m (not divisor of dimension)", + "params": {"M": 16,"efConstruction": 200,"m": 7, "nbits": 8, "refine": True, "refine_type": "FP32"}, + "expected": {"err_code": 1100, "err_msg": "The dimension of a vector (dim) should be a multiple of the number of subquantizers (m)."} + }, + + ] + + search_params = [ + # ef params test + { + "description": "Boundary Test - ef equals k", + "params": {"ef": 10}, + "expected": success + }, + { + "description": "Minimum Boundary Test", + "params": {"ef": 1}, + "expected": success + }, + { + "description": "Large Value Test", + "params": {"ef": 10000}, + "expected": success + }, + { + "description": "Out of Range Test - Negative", + "params": {"ef": -1}, + "expected": {"err_code": 65535, "err_msg": "param 'ef' (-1) should be in range [1, 2147483647]"} + }, + + { + "description": "String Type Test, not check data type", + "params": {"ef": "32"}, + "expected": success + }, + { + "description": "Float Type Test", + "params": {"ef": 32.0}, + "expected": {"err_code": 65535, "err_msg": "Type conflict in json: param 'ef' (32.0) should be integer"} + }, + { + "description": "Boolean Type Test", + "params": {"ef": True}, + "expected": {"err_code": 65535, "err_msg": "Type conflict in json: param 'ef' (true) should be integer"} + }, + { + "description": "None Type Test", + "params": {"ef": None}, + "expected": {"err_code": 65535, "err_msg": "Type conflict in json: param 'ef' (null) should be integer"} + }, + { + "description": "List Type Test", + "params": {"ef": [32]}, + "expected": {"err_code": 65535, "err_msg": "param 'ef' ([32]) should be integer"} + }, + + # refine_k params test + { + "description": "refine_k default boundary", + "params": {"refine_k": 1}, + "expected": success + }, + { + "description": "refine_k valid float", + "params": {"refine_k": 2.5}, + "expected": success + }, + { + "description": "refine_k out of range", + "params": {"refine_k": 0}, + "expected": {"err_code": 65535, "err_msg": "Out of range in json"} + }, + { + "description": "refine_k integer type", + "params": {"refine_k": 20}, + "expected": success + }, + { + "description": "String Type Test, not check data type", + "params": {"refine_k": "2.5"}, + "expected": success + }, + { + "description": "empty string type", + "params": {"refine_k": ""}, + "expected": {"err_code": 65535, "err_msg": "invalid float value"} + }, + { + "description": "refine_k boolean type", + "params": {"refine_k": True}, + "expected": {"err_code": 65535, "err_msg": "Type conflict in json: param 'refine_k' (true) should be a number"} + }, + { + "description": "None Type Test", + "params": {"refine_k": None}, + "expected": {"err_code": 65535, "err_msg": "Type conflict in json"} + }, + { + "description": "List Type Test", + "params": {"refine_k": [15]}, + "expected": {"err_code": 65535, "err_msg":"Type conflict in json"} + }, + + # combination params test + { + "description": "HNSW ef + SQ refine_k combination", + "params": {"ef": 64, "refine_k": 2}, + "expected": success + }, + { + "description": "Valid ef with invalid refine_k", + "params": {"ef": 64, "refine_k": 0}, + "expected": {"err_code": 65535, "err_msg":"Out of range in json"} + }, + { + "description": "empty dict params", + "params": {}, + "expected": success + }, + + ] diff --git a/tests/python_client/testcases/indexes/test_hnsw_prq.py b/tests/python_client/testcases/indexes/test_hnsw_prq.py new file mode 100644 index 0000000000..083b355f62 --- /dev/null +++ b/tests/python_client/testcases/indexes/test_hnsw_prq.py @@ -0,0 +1,258 @@ +import logging +from utils.util_pymilvus import * +from common.common_type import CaseLabel, CheckTasks +from common import common_type as ct +from common import common_func as cf +from base.client_v2_base import TestMilvusClientV2Base +import pytest +from idx_hnsw_prq import HNSW_PRQ + +index_type = "HNSW_PRQ" +success = "success" +pk_field_name = 'id' +vector_field_name = 'vector' +dim = ct.default_dim +default_nb = 2000 +default_build_params = {"M": 16, "efConstruction": 200, "m": 64, "nbits": 8, "nrq":1} +default_search_params = {"ef": 64, "refine_k": 1} + + +class TestHnswPRQBuildParams(TestMilvusClientV2Base): + @pytest.mark.tags(CaseLabel.L1) + @pytest.mark.parametrize("params", HNSW_PRQ.build_params) + def test_hnsw_prq_build_params(self, params): + """ + Test the build params of HNSW_PRQ index + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + schema, _ = self.create_schema(client) + schema.add_field(pk_field_name, datatype=DataType.INT64, is_primary=True, auto_id=False) + schema.add_field(vector_field_name, datatype=DataType.FLOAT_VECTOR, dim=dim) + self.create_collection(client, collection_name, schema=schema) + + all_rows = cf.gen_row_data_by_schema( + nb=default_nb, + schema=schema, + start=0, + random_pk=False + ) + + self.insert(client, collection_name, all_rows) + self.flush(client, collection_name) + + # create index + build_params = params.get("params", None) + index_params = self.prepare_index_params(client)[0] + index_params.add_index(field_name=vector_field_name, + metric_type=cf.get_default_metric_for_vector_type(vector_type=DataType.FLOAT_VECTOR), + index_type=index_type, + params=build_params) + # build index + if params.get("expected", None) != success: + self.create_index(client, collection_name, index_params, + check_task=CheckTasks.err_res, + check_items=params.get("expected")) + else: + self.create_index(client, collection_name, index_params) + self.wait_for_index_ready(client, collection_name, index_name=vector_field_name) + + # load collection + self.load_collection(client, collection_name) + + # search + nq = 2 + search_vectors = cf.gen_vectors(nq, dim=dim, vector_data_type=DataType.FLOAT_VECTOR) + self.search(client, collection_name, search_vectors, + search_params=default_search_params, + limit=ct.default_limit, + check_task=CheckTasks.check_search_results, + check_items={"enable_milvus_client_api": True, + "nq": nq, + "limit": ct.default_limit, + "pk_name": pk_field_name}) + + # verify the index params are persisted + idx_info = client.describe_index(collection_name, vector_field_name) + if build_params is not None: + for key, value in build_params.items(): + if value is not None: + assert key in idx_info.keys() + assert str(value) in idx_info.values() + + @pytest.mark.tags(CaseLabel.L2) + @pytest.mark.parametrize("vector_data_type", ct.all_vector_types) + def test_hnsw_prq_on_all_vector_types(self, vector_data_type): + """ + Test HNSW_PRQ index on all the vector types and metrics + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + schema, _ = self.create_schema(client) + schema.add_field(pk_field_name, datatype=DataType.INT64, is_primary=True, auto_id=False) + if vector_data_type == DataType.SPARSE_FLOAT_VECTOR: + schema.add_field(vector_field_name, datatype=vector_data_type) + else: + schema.add_field(vector_field_name, datatype=vector_data_type, dim=dim) + self.create_collection(client, collection_name, schema=schema) + + all_rows = cf.gen_row_data_by_schema( + nb=default_nb, + schema=schema, + start=0, + random_pk=False + ) + + self.insert(client, collection_name, all_rows) + self.flush(client, collection_name) + + # create index + index_params = self.prepare_index_params(client)[0] + metric_type = cf.get_default_metric_for_vector_type(vector_data_type) + index_params.add_index(field_name=vector_field_name, + metric_type=metric_type, + index_type=index_type, + params=default_build_params) + if vector_data_type not in HNSW_PRQ.supported_vector_types: + self.create_index(client, collection_name, index_params, + check_task=CheckTasks.err_res, + check_items={"err_code": 999, + "err_msg": f"can't build with this index HNSW_PRQ: invalid parameter"}) + + else: + self.create_index(client, collection_name, index_params) + self.wait_for_index_ready(client, collection_name, index_name=vector_field_name) + # load collection + self.load_collection(client, collection_name) + # search + nq = 2 + search_vectors = cf.gen_vectors(nq, dim=dim, vector_data_type=vector_data_type) + self.search(client, collection_name, search_vectors, + search_params=default_search_params, + limit=ct.default_limit, + check_task=CheckTasks.check_search_results, + check_items={"enable_milvus_client_api": True, + "nq": nq, + "limit": ct.default_limit, + "pk_name": pk_field_name}) + + @pytest.mark.tags(CaseLabel.L2) + @pytest.mark.parametrize("metric", HNSW_PRQ.supported_metrics) + def test_hnsw_prq_on_all_metrics(self, metric): + """ + Test the search params of HNSW_PRQ index + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + schema, _ = self.create_schema(client) + schema.add_field(pk_field_name, datatype=DataType.INT64, is_primary=True, auto_id=False) + schema.add_field(vector_field_name, datatype=DataType.FLOAT_VECTOR, dim=dim) + self.create_collection(client, collection_name, schema=schema) + + all_rows = cf.gen_row_data_by_schema( + nb=default_nb, + schema=schema, + start=0, + random_pk=False + ) + + self.insert(client, collection_name, all_rows) + self.flush(client, collection_name) + + # create index + index_params = self.prepare_index_params(client)[0] + index_params.add_index(field_name=vector_field_name, + metric_type=metric, + index_type=index_type, + params=default_build_params) + self.create_index(client, collection_name, index_params) + self.wait_for_index_ready(client, collection_name, index_name=vector_field_name) + # load collection + self.load_collection(client, collection_name) + # search + nq = 2 + search_vectors = cf.gen_vectors(nq, dim=dim, vector_data_type=DataType.FLOAT_VECTOR) + self.search(client, collection_name, search_vectors, + search_params=default_search_params, + limit=ct.default_limit, + check_task=CheckTasks.check_search_results, + check_items={"enable_milvus_client_api": True, + "nq": nq, + "limit": ct.default_limit, + "pk_name": pk_field_name}) + + +@pytest.mark.xdist_group("TestHnswPRQSearchParams") +class TestHnswPRQSearchParams(TestMilvusClientV2Base): + """Test search with pagination functionality for HNSW_PRQ index""" + + def setup_class(self): + super().setup_class(self) + self.collection_name = "TestHnswPRQSearchParams" + cf.gen_unique_str("_") + self.float_vector_field_name = vector_field_name + self.float_vector_dim = dim + self.primary_keys = [] + self.enable_dynamic_field = False + self.datas = [] + + @pytest.fixture(scope="class", autouse=True) + def prepare_collection(self, request): + """ + Initialize collection before test class runs + """ + client = self._client() + collection_schema = self.create_schema(client)[0] + collection_schema.add_field(pk_field_name, DataType.INT64, is_primary=True, auto_id=False) + collection_schema.add_field(self.float_vector_field_name, DataType.FLOAT_VECTOR, dim=128) + self.create_collection(client, self.collection_name, schema=collection_schema, + enable_dynamic_field=self.enable_dynamic_field, force_teardown=False) + all_data = cf.gen_row_data_by_schema( + nb=default_nb, + schema=collection_schema, + start=0, + random_pk=False + ) + self.insert(client, self.collection_name, data=all_data) + self.primary_keys.extend([i for i in range(default_nb)]) + + self.flush(client, self.collection_name) + # Create HNSW_PRQ index + index_params = self.prepare_index_params(client)[0] + index_params.add_index(field_name=self.float_vector_field_name, + metric_type="COSINE", + index_type=index_type, + params=default_build_params) + self.create_index(client, self.collection_name, index_params=index_params) + self.wait_for_index_ready(client, self.collection_name, index_name=self.float_vector_field_name) + self.load_collection(client, self.collection_name) + + def teardown(): + self.drop_collection(self._client(), self.collection_name) + request.addfinalizer(teardown) + + @pytest.mark.tags(CaseLabel.L1) + @pytest.mark.parametrize("params", HNSW_PRQ.search_params) + def test_hnsw_prq_search_params(self, params): + """ + Test the search params of HNSW_PRQ index + """ + client = self._client() + collection_name = self.collection_name + nq = 2 + search_vectors = cf.gen_vectors(nq, dim=self.float_vector_dim, vector_data_type=DataType.FLOAT_VECTOR) + search_params = params.get("params", None) + if params.get("expected", None) != success: + self.search(client, collection_name, search_vectors, + search_params=search_params, + limit=ct.default_limit, + check_task=CheckTasks.err_res, + check_items=params.get("expected")) + else: + self.search(client, collection_name, search_vectors, + search_params=search_params, + limit=ct.default_limit, + check_task=CheckTasks.check_search_results, + check_items={"enable_milvus_client_api": True, + "nq": nq, + "limit": ct.default_limit, + "pk_name": pk_field_name}) \ No newline at end of file