mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-28 14:35:27 +08:00
test: add HNSW_SQ test cases (#46428)
/kind improvement /assign @yanliang567 Signed-off-by: zilliz <jiaming.li@zilliz.com>
This commit is contained in:
jiamingli-maker
GitHub
parent
89a002e12a
commit
b9fe8e9f9e
413
tests/python_client/testcases/indexes/idx_hnsw_sq.py
Normal file
413
tests/python_client/testcases/indexes/idx_hnsw_sq.py
Normal file
@ -0,0 +1,413 @@
|
||||
from pymilvus import DataType
|
||||
|
||||
success = "success"
|
||||
|
||||
|
||||
class HNSW_SQ:
|
||||
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"}
|
||||
},
|
||||
|
||||
# sq_type params test
|
||||
{
|
||||
"description": "Valid sq_type - SQ6",
|
||||
"params": {"sq_type": "SQ6"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Valid sq_type - SQ8",
|
||||
"params": {"sq_type": "SQ8"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Valid sq_type - BF16",
|
||||
"params": {"sq_type": "BF16"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Valid sq_type - FP16",
|
||||
"params": {"sq_type": "FP16"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Out of Range Test - Unknown String",
|
||||
"params": {"sq_type": "FP32"},
|
||||
"expected": {"err_code": 1100, "err_msg": "invalid scalar quantizer type: invalid parameter"}
|
||||
},
|
||||
{
|
||||
"description": "Integer Type Test",
|
||||
"params": {"sq_type": 8},
|
||||
"expected": {"err_code": 1100, "err_msg": "invalid scalar quantizer type: invalid parameter"}
|
||||
},
|
||||
{
|
||||
"description": "Float Type Test",
|
||||
"params": {"sq_type": 8.0},
|
||||
"expected": {"err_code": 1100, "err_msg": "invalid scalar quantizer type: invalid parameter"}
|
||||
},
|
||||
{
|
||||
"description": "Boolean Type Test",
|
||||
"params": {"sq_type": True},
|
||||
"expected": {"err_code": 1100, "err_msg": "invalid scalar quantizer type: invalid parameter"}
|
||||
},
|
||||
{
|
||||
"description": "None Type Test, use default value",
|
||||
"params": {"sq_type": None},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "List Type Test",
|
||||
"params": {"sq_type": ["SQ8"]},
|
||||
"expected": {"err_code": 1100, "err_msg": "invalid scalar quantizer type: invalid parameter"}
|
||||
},
|
||||
|
||||
# refine params test
|
||||
{
|
||||
"description": "refine = True",
|
||||
"params": {"refine": True},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "String Type Test",
|
||||
"params": {"refine": "true"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "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
|
||||
},
|
||||
# combination params test
|
||||
{
|
||||
"description": "empty dict params",
|
||||
"params": {},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "All optional parameters None",
|
||||
"params": {"M": None, "efConstruction": None, "sq_type": None, "refine": None, "refine_type": None},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Typical valid combination",
|
||||
"params": {"M": 16, "efConstruction": 200, "sq_type": "SQ8", "refine": True, "refine_type": "FP16"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Minimum boundary combination",
|
||||
"params": {"M": 2, "efConstruction": 1, "sq_type": "SQ6"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Maximum boundary combination",
|
||||
"params": {"M": 2048, "efConstruction": 10000, "sq_type": "FP16", "refine": True, "refine_type": "FP32"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Unknown extra parameter in combination",
|
||||
"params": {"M": 16, "efConstruction": 200, "sq_type": "SQ8", "refine": True, "refine_type": "FP16", "unknown_param": "nothing"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Partial parameters set (M + sq_type only)",
|
||||
"params": {"M": 32, "sq_type": "BF16"},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Partial parameters set (efConstruction + refine only)",
|
||||
"params": {"efConstruction": 500,"refine": True},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Invalid refine_type using vector data type",
|
||||
"params": {"sq_type": "SQ8", "refine": True, "refine_type": "INT8"},
|
||||
"expected": {"err_code": 1100, "err_msg": "invalid refine type"}
|
||||
}
|
||||
|
||||
]
|
||||
|
||||
search_params = [
|
||||
# ef params test
|
||||
{
|
||||
"description": "Boundary Test - ef equals k",
|
||||
"params": {"ef": 10},
|
||||
"expected": success
|
||||
},
|
||||
{
|
||||
"description": "Minimum Boundary Test",
|
||||
"params": {"ef": 1},
|
||||
"expected": {"err_code": 65535, "err_msg": "ef(1) should be larger than k(10)"} # assume default limit=10
|
||||
},
|
||||
{
|
||||
"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
|
||||
},
|
||||
|
||||
]
|
||||
272
tests/python_client/testcases/indexes/test_hnsw_sq.py
Normal file
272
tests/python_client/testcases/indexes/test_hnsw_sq.py
Normal file
@ -0,0 +1,272 @@
|
||||
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_sq import HNSW_SQ
|
||||
|
||||
index_type = "HNSW_SQ"
|
||||
success = "success"
|
||||
pk_field_name = 'id'
|
||||
vector_field_name = 'vector'
|
||||
dim = ct.default_dim
|
||||
default_nb = 2000
|
||||
default_build_params = {"M": 16, "efConstruction": 200, "sq_type": "SQ8"}
|
||||
default_search_params = {"ef": 64, "refine_k": 1}
|
||||
|
||||
|
||||
class TestHnswSQBuildParams(TestMilvusClientV2Base):
|
||||
@pytest.mark.tags(CaseLabel.L1)
|
||||
@pytest.mark.parametrize("params", HNSW_SQ.build_params)
|
||||
def test_hnsw_sq_build_params(self, params):
|
||||
"""
|
||||
Test the build params of HNSW_SQ 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)
|
||||
|
||||
# Insert data in 2 batches with unique primary keys
|
||||
insert_times = 2
|
||||
random_vectors = list(cf.gen_vectors(default_nb * insert_times, dim, vector_data_type=DataType.FLOAT_VECTOR))
|
||||
for j in range(insert_times):
|
||||
start_pk = j * default_nb
|
||||
rows = [{
|
||||
pk_field_name: i + start_pk,
|
||||
vector_field_name: random_vectors[i + start_pk]
|
||||
} for i in range(default_nb)]
|
||||
self.insert(client, collection_name, 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_sq_on_all_vector_types(self, vector_data_type):
|
||||
"""
|
||||
Test HNSW_SQ 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)
|
||||
|
||||
# Insert data in 2 batches with unique primary keys
|
||||
insert_times = 2
|
||||
random_vectors = list(cf.gen_vectors(default_nb*insert_times, dim, vector_data_type=vector_data_type)) \
|
||||
if vector_data_type == DataType.FLOAT_VECTOR \
|
||||
else cf.gen_vectors(default_nb*insert_times, dim, vector_data_type=vector_data_type)
|
||||
for j in range(insert_times):
|
||||
start_pk = j * default_nb
|
||||
rows = [{
|
||||
pk_field_name: i + start_pk,
|
||||
vector_field_name: random_vectors[i + start_pk]
|
||||
} for i in range(default_nb)]
|
||||
self.insert(client, collection_name, 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_SQ.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_SQ: 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_SQ.supported_metrics)
|
||||
def test_hnsw_sq_on_all_metrics(self, metric):
|
||||
"""
|
||||
Test the search params of HNSW_SQ 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)
|
||||
|
||||
# insert data
|
||||
insert_times = 2
|
||||
random_vectors = list(cf.gen_vectors(default_nb*insert_times, dim, vector_data_type=DataType.FLOAT_VECTOR))
|
||||
for j in range(insert_times):
|
||||
start_pk = j * default_nb
|
||||
rows = [{
|
||||
pk_field_name: i + start_pk,
|
||||
vector_field_name: random_vectors[i + start_pk]
|
||||
} for i in range(default_nb)]
|
||||
self.insert(client, collection_name, 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("TestHnswSQSearchParams")
|
||||
class TestHnswSQSearchParams(TestMilvusClientV2Base):
|
||||
"""Test search with pagination functionality for HNSW_SQ index"""
|
||||
|
||||
def setup_class(self):
|
||||
super().setup_class(self)
|
||||
self.collection_name = "TestHnswSQSearchParams" + 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)
|
||||
insert_times = 2
|
||||
float_vectors = cf.gen_vectors(default_nb * insert_times, dim=self.float_vector_dim,
|
||||
vector_data_type=DataType.FLOAT_VECTOR)
|
||||
for j in range(insert_times):
|
||||
rows = []
|
||||
for i in range(default_nb):
|
||||
pk = i + j * default_nb
|
||||
row = {
|
||||
pk_field_name: pk,
|
||||
self.float_vector_field_name: list(float_vectors[pk])
|
||||
}
|
||||
self.datas.append(row)
|
||||
rows.append(row)
|
||||
self.insert(client, self.collection_name, data=rows)
|
||||
self.primary_keys.extend([i + j * default_nb for i in range(default_nb)])
|
||||
self.flush(client, self.collection_name)
|
||||
# Create HNSW_SQ 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_SQ.search_params)
|
||||
def test_hnsw_sq_search_params(self, params):
|
||||
"""
|
||||
Test the search params of HNSW_SQ 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})
|
||||
Loading…
x
Reference in New Issue
Block a user