def scene_test(self, collection_name=None, vectors=None, ids=None):
"""
Scene test steps：
1.create collection with the specified collection name
2.insert data
3.flush data
4.create index
5.drop collection
"""
logger.debug("[scene_test] Start scene test : %s" % collection_name)
self.create_collection(dimension=128, collection_name=collection_name)
time.sleep(1)
collection_info = self.get_info(collection_name)
entities = utils.generate_entities(collection_info, vectors, ids)
logger.debug("[scene_test] Start insert : %s" % collection_name)
self.insert(entities, collection_name=collection_name)
logger.debug("[scene_test] Start flush : %s" % collection_name)
self.flush(collection_name=collection_name)
logger.debug("[scene_test] Start create index : %s" % collection_name)
self.create_index(field_name='float_vector',
index_type="ivf_sq8",
metric_type='l2',
collection_name=collection_name,
index_param={'nlist': 2048})
time.sleep(59)
logger.debug("[scene_test] Start drop : %s" % collection_name)
self.drop(collection_name=collection_name)
logger.debug("[scene_test]Scene test close : %s" % collection_name)
def prepare(self, **case_param):
""" According to the test case parameters, initialize the test """
collection_name = case_param["collection_name"]
metric_type = case_param["metric_type"]
dimension = case_param["dimension"]
vector_type = case_param["vector_type"]
index_type = case_param["index_type"]
index_param = case_param["index_param"]
index_field_name = case_param["index_field_name"]
self.milvus.set_collection(collection_name)
if self.milvus.exists_collection(collection_name):
logger.info("Re-create collection: %s" % collection_name)
self.milvus.drop()
dataset = case_param["dataset"]
self.milvus.create_collection(dimension, data_type=vector_type)
# Get the data set train for inserting into the collection
insert_vectors = utils.normalize(metric_type,
np.array(dataset["train"]))
if len(insert_vectors) != dataset["train"].shape[0]:
raise Exception(
"Row count of insert vectors: %d is not equal to dataset size: %d"
% (len(insert_vectors), dataset["train"].shape[0]))
logger.debug("The row count of entities to be inserted: %d" %
len(insert_vectors))
# Insert batch once
# milvus_instance.insert(insert_vectors)
info = self.milvus.get_info(collection_name)
loops = len(insert_vectors) // INSERT_INTERVAL + 1
for i in range(loops):
start = i * INSERT_INTERVAL
end = min((i + 1) * INSERT_INTERVAL, len(insert_vectors))
if start < end:
# Insert up to INSERT_INTERVAL=50000 at a time
tmp_vectors = insert_vectors[start:end]
ids = [i for i in range(start, end)]
if not isinstance(tmp_vectors, list):
entities = utils.generate_entities(info,
tmp_vectors.tolist(),
ids)
res_ids = self.milvus.insert(entities)
else:
entities = utils.generate_entities(tmp_vectors, ids)
res_ids = self.milvus.insert(entities)
assert res_ids == ids
logger.debug("End insert, start flush")
self.milvus.flush()
logger.debug("End flush")
res_count = self.milvus.count()
logger.info("Table: %s, row count: %d" % (collection_name, res_count))
if res_count != len(insert_vectors):
raise Exception("Table row count is not equal to insert vectors")
if self.milvus.describe_index(index_field_name):
self.milvus.drop_index(index_field_name)
logger.info("Re-create index: %s" % collection_name)
self.milvus.create_index(index_field_name,
index_type,
metric_type,
index_param=index_param)
logger.info(self.milvus.describe_index(index_field_name))
logger.info("Start load collection: %s" % collection_name)
# self.milvus.release_collection()
self.milvus.load_collection(timeout=600)
logger.info("End load collection: %s" % collection_name)