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)