test: Add HNSW_PRQ test cases and fix HNSW_PQ (#46680)

/kind improvement

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
- 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: zilliz <jiaming.li@zilliz.com>
This commit is contained in:
jiamingli-maker 2026-01-04 18:57:22 +08:00 committed by GitHub
parent 146f173eea
commit c10cf53b4b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 852 additions and 3 deletions

View File

@ -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)."}
},
]

View File

@ -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
},
]

View File

@ -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})