mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-04 12:59:23 +08:00
63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
import json
|
|
from utils.util_log import test_log as log
|
|
|
|
all_index_types = ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "ANNOY", "RHNSW_FLAT", "RHNSW_PQ", "RHNSW_SQ",
|
|
"BIN_FLAT", "BIN_IVF_FLAT"]
|
|
|
|
default_index_params = [{"nlist": 128}, {"nlist": 128}, {"nlist": 128}, {"nlist": 128, "m": 16, "nbits": 8},
|
|
{"M": 48, "efConstruction": 500}, {"n_trees": 50}, {"M": 48, "efConstruction": 500},
|
|
{"M": 48, "efConstruction": 500, "PQM": 8}, {"M": 48, "efConstruction": 500}, {"nlist": 128},
|
|
{"nlist": 128}]
|
|
|
|
index_params_map = dict(zip(all_index_types, default_index_params))
|
|
|
|
def gen_index_param(index_type):
|
|
metric_type = "L2"
|
|
if "BIN" in index_type:
|
|
metric_type = "HAMMING"
|
|
index_param = {
|
|
"index_type": index_type,
|
|
"params": index_params_map[index_type],
|
|
"metric_type": metric_type
|
|
}
|
|
return index_param
|
|
|
|
|
|
def gen_search_param(index_type, metric_type="L2"):
|
|
search_params = []
|
|
if index_type in ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_SQ8H", "IVF_PQ"]:
|
|
for nprobe in [10]:
|
|
ivf_search_params = {"metric_type": metric_type, "params": {"nprobe": nprobe}}
|
|
search_params.append(ivf_search_params)
|
|
elif index_type in ["BIN_FLAT", "BIN_IVF_FLAT"]:
|
|
for nprobe in [10]:
|
|
bin_search_params = {"metric_type": "HAMMING", "params": {"nprobe": nprobe}}
|
|
search_params.append(bin_search_params)
|
|
elif index_type in ["HNSW", "RHNSW_FLAT", "RHNSW_PQ", "RHNSW_SQ"]:
|
|
for ef in [64]:
|
|
hnsw_search_param = {"metric_type": metric_type, "params": {"ef": ef}}
|
|
search_params.append(hnsw_search_param)
|
|
elif index_type in ["NSG", "RNSG"]:
|
|
for search_length in [100]:
|
|
nsg_search_param = {"metric_type": metric_type, "params": {"search_length": search_length}}
|
|
search_params.append(nsg_search_param)
|
|
elif index_type == "ANNOY":
|
|
for search_k in [1000]:
|
|
annoy_search_param = {"metric_type": metric_type, "params": {"search_k": search_k}}
|
|
search_params.append(annoy_search_param)
|
|
else:
|
|
print("Invalid index_type.")
|
|
raise Exception("Invalid index_type.")
|
|
return search_params
|
|
|
|
|
|
|
|
def get_all_collections():
|
|
try:
|
|
with open("/tmp/ci_logs/all_collections.json", "r") as f:
|
|
data = json.load(f)
|
|
all_collections = data["all"]
|
|
except Exception as e:
|
|
log.error(f"get_all_collections error: {e}")
|
|
return []
|
|
return all_collections |