class Ingester: def __init__(self, host, port, collection, collection_param, partition=None, drop=False, batch_size=100, dtype=np.float32): self.collection = collection self.partition = partition self.client = Milvus(host, port) self.dtype = dtype self.batch_size = batch_size if drop and self.client.has_collection(collection): self.client.drop_collection(collection) if not self.client.has_collection(collection): self.client.create_collection(collection, collection_param) if partition and not self.client.has_partition(collection, partition): self.client.create_partition(collection, partition) def ingest(self, entities, ids): if self.partition: return self.client.insert(self.collection, entities, ids=ids, partition_tag=self.partition) else: return self.client.insert(self.collection, entities, ids=ids)
def main(): sentences = word2vec.Text8Corpus("text8") # 加载语料 model = word2vec.Word2Vec(sentences, size=200, window=5, min_count=5) # 训练模型 word_set = model.wv.index2word # 单词集合 word_vec = model.wv.vectors # word2vec结果向量集合 milvus = Milvus() milvus.connect(host='localhost', port='19530') param = { 'collection_name': 'word2vec', 'dimension': 200, 'index_file_size': 1024, 'metric_type': MetricType.L2 } milvus.create_collection(param) status, ids = milvus.insert(collection_name='word2vec', records=word_vec) # 单词分类 ivf_param = {'nlist': 100} # 分成100类 milvus.create_index('word2vec', IndexType.IVF_FLAT, ivf_param) # 增加索引 status, index = milvus.describe_index( 'word2vec') # 相当于将word分成100个类别 做了聚类算法 # 查找相似度最高的单词 res = milvus.search(collection_name='word2vec', query_records=[list(word_vec[word_set.index('king')])], top_k=10, params={'nprobe': 16}) for i in range(10): id = res[1][0][i].id print(word_set[ids.index(id)]) print(1)
def milvus_test(usr_features, IS_INFER, mov_features=None, ids=None): _HOST = '127.0.0.1' _PORT = '19530' # default value table_name = 'recommender_demo' milvus = Milvus() param = {'host': _HOST, 'port': _PORT} status = milvus.connect(**param) if status.OK(): print("Server connected.") else: print("Server connect fail.") sys.exit(1) if IS_INFER: status = milvus.drop_collection(table_name) time.sleep(3) status, ok = milvus.has_collection(table_name) if not ok: if mov_features is None: print("Insert vectors is none!") sys.exit(1) param = { 'collection_name': table_name, 'dimension': 200, 'index_file_size': 1024, # optional 'metric_type': MetricType.IP # optional } print(milvus.create_collection(param)) insert_vectors = normaliz_data(mov_features) status, ids = milvus.insert(collection_name=table_name, records=insert_vectors, ids=ids) time.sleep(1) status, result = milvus.count_collection(table_name) print("rows in table recommender_demo:", result) search_vectors = normaliz_data(usr_features) param = { 'collection_name': table_name, 'query_records': search_vectors, 'top_k': 5, 'params': { 'nprobe': 16 } } time1 = time.time() status, results = milvus.search(**param) time2 = time.time() print("Top\t", "Ids\t", "Title\t", "Score") for i, re in enumerate(results[0]): title = paddle.dataset.movielens.movie_info()[int(re.id)].title print(i, "\t", re.id, "\t", title, "\t", float(re.distance) * 5)
def predict(start_date_str, end_date_str): print("加载模型") model = gensim.models.doc2vec.Doc2Vec.load("./doc2vec.model") print("建立milvus链接") client = Milvus(host=milvus_ip, port='19530') print("读取数据ing") start_date = datetime.strptime(start_date_str, '%Y-%m-%d').timestamp() * 1000 end_date = datetime.strptime(end_date_str, '%Y-%m-%d').timestamp() * 1000 res = Paper.query_by_time_interval(start_date, end_date) num = 0 start = time.time() id_list = [] user_id_list = [] vecs = [] for i in res: paper_id = i.id paper_user_id = i.user_id paper_str = i.title + " . " + i.description vec = get_vector(model, [paper_str]) # 将词向量写入到Milvus id_list.append(paper_id) user_id_list.append(paper_user_id) vecs.append(list(vec)) # 将词向量写入数据库 paper_vec = str(vec).replace('\n', '').replace('[', '').replace( ']', '').replace(" ", " ").replace(" ", ",")[1:] paper_vec = paper_vec.replace(",,", ",0,") Paper.update_SQL('doc_vector', paper_vec, paper_user_id) num += 1 if num % 200 == 0: print("完成了", num, '篇', '--用时:', time.time() - start) start = time.time() # hybrid_entities = [ # {"name": "id", "values": id_list, "type": DataType.INT32}, # {"name": "Vec", "values": vecs, "type": DataType.FLOAT_VECTOR} # ] client.insert('ideaman', records=vecs, ids=id_list) client.flush(collection_name_array=["ideaman"]) user_id_list.clear() id_list.clear() vecs.clear()
def multi_thread_opr(table_name, utid): print("[{}] | T{} | Running .....".format(datetime.datetime.now(), utid)) client0 = Milvus(handler="HTTP") table_param = {'table_name': table_name, 'dimension': 64} vectors = [[random.random() for _ in range(64)] for _ in range(10000)] client0.connect() client0.create_table(table_param) # print("[{}] | T{} | O{} | Start insert data .....".format(datetime.datetime.now(), utid, i)) client0.insert(table_name, vectors) # print("[{}] | T{} | O{} | Stop insert data .....".format(datetime.datetime.now(), utid, i)) client0.disconnect()
def test_not_connect(self): client = Milvus() with pytest.raises(NotConnectError): client.create_collection({}) with pytest.raises(NotConnectError): client.has_collection("a") with pytest.raises(NotConnectError): client.describe_collection("a") with pytest.raises(NotConnectError): client.drop_collection("a") with pytest.raises(NotConnectError): client.create_index("a") with pytest.raises(NotConnectError): client.insert("a", [], None) with pytest.raises(NotConnectError): client.count_collection("a") with pytest.raises(NotConnectError): client.show_collections() with pytest.raises(NotConnectError): client.search("a", 1, 2, [], None) with pytest.raises(NotConnectError): client.search_in_files("a", [], [], 2, 1, None) with pytest.raises(NotConnectError): client._cmd("") with pytest.raises(NotConnectError): client.preload_collection("a") with pytest.raises(NotConnectError): client.describe_index("a") with pytest.raises(NotConnectError): client.drop_index("")
def _add(): milvus = Milvus(**server_config) vectors = _generate_vectors(128, 10000) print('\n\tPID: {}, insert {} vectors'.format(os.getpid(), 10000)) status, _ = milvus.insert(_collection_name, vectors) if not status.OK(): print("PID {} insert failed: {}".format(os.getpid(), status.message)) milvus.close()
def create(): _HOST = 'localhost' _PORT = '19530' _collection_name = 'chs_stars_faces_512' _DIM = 512 # dimension of vector _INDEX_FILE_SIZE = 256 # max file size of stored index milvus = Milvus(_HOST, _PORT) param = { 'collection_name': _collection_name, 'dimension': _DIM, 'index_file_size': _INDEX_FILE_SIZE, # optional 'metric_type': MetricType.IP # optional } milvus.create_collection(param) index_param = { 'nlist': 2048 # 推荐 4 * sqrt(n) } status = milvus.create_index(_collection_name, IndexType.IVF_SQ8, index_param) # with open("chs_stars_features_pca.pickle", "rb") as f: # pca = pickle.load(f) # # with open("../chs_stars_features_pca.csv", "w") as fw, open("../chs_stars_features.csv", "r") as fr: # reader = csv.reader(fr) # writer = csv.writer(fw) # for index, line in enumerate(tqdm(reader)): # star, fname, features = line # features = np.array(json.loads(features)) # features = np.resize(features, (1, 512)) # features = normalize(features) # features = pca.transform(features).squeeze() # status, ids = milvus.insert(collection_name=_collection_name, records=[features.tolist()], ids=[index]) # if not status.OK(): # print(status) # continue # writer.writerow([index, star, fname, features]) with open("../chs_stars_labels.csv", "w") as fw, open("../chs_stars_features.csv", "r") as fr: reader = csv.reader(fr) writer = csv.writer(fw) for index, line in enumerate(tqdm(reader)): star, fname, features = line # features = np.array(json.loads(features)) # features = np.resize(features, (1, 512)) #features = normalize(features) features = json.loads(features) status, ids = milvus.insert(collection_name=_collection_name, records=[features], ids=[index]) if not status.OK(): print(status) continue writer.writerow([index, star, fname])
def insert_vectors(name, vectors): milvus = Milvus() try: milvus.connect(MILVUS_ADDR, MILVUS_PORT) res, ids = milvus.insert(collection_name=name, records=vectors) if not res.OK(): raise MilvusError("There has some error when insert vectors", res) return ids except Exception as e: logger.error("There has some error when insert vectors", exc_info=True) raise MilvusError("There has some error when insert vectors", e)
def insert_vectors(name, vectors): try: milvus = Milvus(host=MILVUS_ADDR, port=MILVUS_PORT) res, ids = milvus.insert(collection_name=name, records=vectors) if not res.OK(): err_msg = "There was some error when insert vectors" logger.error(f"{err_msg} : {str(res)}", exc_info=True) raise MilvusError(err_msg, res) return ids except Exception as e: err_msg = "There was some error when insert vectors" logger.error(f"{err_msg} : {str(e)}", exc_info=True) raise MilvusError(err_msg, e)
def milvus_test(usr_features, mov_features, ids): _HOST = '127.0.0.1' _PORT = '19530' # default value milvus = Milvus() param = {'host': _HOST, 'port': _PORT} status = milvus.connect(**param) if status.OK(): print("\nServer connected.") else: print("\nServer connect fail.") sys.exit(1) table_name = 'paddle_demo1' status, ok = milvus.has_table(table_name) if not ok: param = { 'table_name': table_name, 'dimension': 200, 'index_file_size': 1024, # optional 'metric_type': MetricType.IP # optional } milvus.create_table(param) insert_vectors = normaliz_data([usr_features.tolist()]) status, ids = milvus.insert(table_name=table_name, records=insert_vectors, ids=ids) time.sleep(1) status, result = milvus.count_table(table_name) print("rows in table paddle_demo1:", result) status, table = milvus.describe_table(table_name) search_vectors = normaliz_data([mov_features.tolist()]) param = { 'table_name': table_name, 'query_records': search_vectors, 'top_k': 1, 'nprobe': 16 } status, results = milvus.search_vectors(**param) print("Searched ids:", results[0][0].id) print("Score:", float(results[0][0].distance) * 5) status = milvus.drop_table(table_name)
def _add_milvus_question(self, question_vector, collection: str, partition: str, milvus: mv.Milvus) -> int: """ 添加标准问题 @param {object} question_vector - 问题向量 @param {str} collection - 问题分类 @param {str} partition - 场景 @param {mv.Milvus} milvus - Milvus服务连接对象 @returns {int} - 返回milvus_id """ _status, _milvus_ids = milvus.insert( collection, [question_vector, ], partition_tag=partition) self.confirm_milvus_status(_status, 'insert') self._log_debug('insert _milvus_ids: %s' % str(_milvus_ids)) return _milvus_ids[0]
class MilvusConnection: def __init__(self, env, name="movies_L2", port="19530", param=None): if param is None: param = dict() param = { "collection_name": name, "dimension": 128, "index_file_size": 1024, "metric_type": MetricType.L2, **param, } self.name = name self.client = Milvus(host="localhost", port=port) self.statuses = {} if not self.client.has_collection(name)[1]: status_created_collection = self.client.create_collection(param) vectors = env.base.embeddings.detach().cpu().numpy().astype( "float32") target_ids = list(range(vectors.shape[0])) status_inserted, inserted_vector_ids = self.client.insert( collection_name=name, records=vectors, ids=target_ids) status_flushed = self.client.flush([name]) status_compacted = self.client.compact(collection_name=name) self.statuses["created_collection"] = status_created_collection self.statuses["inserted"] = status_inserted self.statuses["flushed"] = status_flushed self.statuses["compacted"] = status_compacted def search(self, search_vecs, topk=10, search_param=None): if search_param is None: search_param = dict() search_param = {"nprobe": 16, **search_param} status, results = self.client.search( collection_name=self.name, query_records=search_vecs, top_k=topk, params=search_param, ) self.statuses['last_search'] = status return torch.tensor(results.id_array) def get_log(self): return self.statuses
def milvus_test(usr_features, mov_features, ids): _HOST = '127.0.0.1' _PORT = '19530' # default value milvus = Milvus(_HOST, _PORT) table_name = 'paddle_demo1' status, ok = milvus.has_collection(table_name) if not ok: param = { 'collection_name': table_name, 'dimension': 200, 'index_file_size': 1024, # optional 'metric_type': MetricType.IP # optional } milvus.create_collection(param) insert_vectors = normaliz_data([usr_features.tolist()]) status, ids = milvus.insert(collection_name=table_name, records=insert_vectors, ids=ids) time.sleep(1) status, result = milvus.count_entities(table_name) print("rows in table paddle_demo1:", result) # status, table = milvus.count_entities(table_name) search_vectors = normaliz_data([mov_features.tolist()]) param = { 'collection_name': table_name, 'query_records': search_vectors, 'top_k': 1, 'params': { 'nprobe': 16 } } status, results = milvus.search(**param) print("Searched ids:", results[0][0].id) print("Score:", float(results[0][0].distance) * 5) status = milvus.drop_collection(table_name)
def main(): # connect_milvus_server() milvus = Milvus(host=SERVER_ADDR, port=SERVER_PORT) create_milvus_collection(milvus) build_collection(milvus) count = 0 while count < (VEC_NUM // BASE_LEN): vectors = load_bvecs_data(FILE_PATH, BASE_LEN, count) vectors_ids = [ id for id in range(count * BASE_LEN, (count + 1) * BASE_LEN) ] sex = [random.randint(0, 2) for _ in range(10000)] get_time = [random.randint(2017, 2020) for _ in range(10000)] is_glasses = [random.randint(10, 13) for _ in range(10000)] hybrid_entities = [{ "name": "sex", "values": sex, "type": DataType.INT32 }, { "name": "is_glasses", "values": is_glasses, "type": DataType.INT32 }, { "name": "get_time", "values": get_time, "type": DataType.INT32 }, { "name": "Vec", "values": vectors, "type": DataType.FLOAT_VECTOR }] time_start = time.time() result = milvus.insert('mixed06', hybrid_entities, ids=vectors_ids) time_end = time.time() print("insert milvue time: ", time_end - time_start) count = count + 1
def main(): milvus = Milvus(handler="HTTP") # Connect to Milvus server # You may need to change _HOST and _PORT accordingly param = {'host': _HOST, 'port': _PORT} status = milvus.connect(**param) if status.OK(): print("Server connected.") else: print("Server connect fail.") sys.exit(1) # Create table demo_table if it dosen't exist. table_name = 'demo_tables' status, ok = milvus.has_table(table_name) if not ok: param = { 'table_name': table_name, 'dimension': _DIM, 'index_file_size': _INDEX_FILE_SIZE, # optional 'metric_type': MetricType.L2 # optional } milvus.create_table(param) # Show tables in Milvus server _, tables = milvus.show_tables() # Describe demo_table _, table = milvus.describe_table(table_name) print(table) # 10000 vectors with 16 dimension # element per dimension is float32 type # vectors should be a 2-D array vectors = [[random.random() for _ in range(_DIM)] for _ in range(100000)] # You can also use numpy to generate random vectors: # `vectors = np.random.rand(10000, 16).astype(np.float32).tolist()` # Insert vectors into demo_table, return status and vectors id list status, ids = milvus.insert(table_name=table_name, records=vectors) # Wait for 6 seconds, until Milvus server persist vector data. time.sleep(6) # Get demo_table row count status, result = milvus.count_table(table_name) # create index of vectors, search more rapidly index_param = { 'index_type': IndexType.IVFLAT, # choice ivflat index 'nlist': 2048 } # Create ivflat index in demo_table # You can search vectors without creating index. however, Creating index help to # search faster status = milvus.create_index(table_name, index_param) # describe index, get information of index status, index = milvus.describe_index(table_name) print(index) # Use the top 10 vectors for similarity search query_vectors = vectors[0:10] # execute vector similarity search param = { 'table_name': table_name, 'query_records': query_vectors, 'top_k': 1, 'nprobe': 16 } status, results = milvus.search(**param) if status.OK(): # indicate search result # also use by: # `results.distance_array[0][0] == 0.0 or results.id_array[0][0] == ids[0]` if results[0][0].distance == 0.0 or results[0][0].id == ids[0]: print('Query result is correct') else: print('Query result isn\'t correct') # print results print(results) # Delete demo_table status = milvus.drop_table(table_name) # Disconnect from Milvus status = milvus.disconnect()
def main(): # Specify server addr when create milvus client instance milvus = Milvus(_HOST, _PORT) # Create collection demo_collection if it dosen't exist. collection_name = 'example_async_collection_' status, ok = milvus.has_collection(collection_name) if not ok: param = { 'collection_name': collection_name, 'dimension': _DIM, 'index_file_size': 128, # optional 'metric_type': MetricType.L2 # optional } status = milvus.create_collection(param) if not status.OK(): print("Create collection failed: {}".format(status.message), file=sys.stderr) print("exiting ...", file=sys.stderr) sys.exit(1) # Show collections in Milvus server _, collections = milvus.list_collections() # Describe demo_collection _, collection = milvus.get_collection_info(collection_name) print(collection) # 10000 vectors with 16 dimension # element per dimension is float32 type # vectors should be a 2-D array vectors = [[random.random() for _ in range(_DIM)] for _ in range(100000)] # You can also use numpy to generate random vectors: # `vectors = np.random.rand(10000, 16).astype(np.float32)` def _insert_callback(status, ids): if status.OK(): print("Insert successfully") else: print("Insert failed.", status.message) # Insert vectors into demo_collection, adding callback function insert_future = milvus.insert(collection_name=collection_name, records=vectors, _async=True, _callback=_insert_callback) # Or invoke result() to get results: # insert_future = milvus.insert(collection_name=collection_name, records=vectors, _async=True) # status, ids = insert_future.result() insert_future.done() # Flush collection inserted data to disk. def _flush_callback(status): if status.OK(): print("Flush successfully") else: print("Flush failed.", status.message) flush_future = milvus.flush([collection_name], _async=True, _callback=_flush_callback) # Or invoke result() to get results: # flush_future = milvus.flush([collection_name], _async=True) # status = flush_future.result() flush_future.done() def _compact_callback(status): if status.OK(): print("Compact successfully") else: print("Compact failed.", status.message) compact_furure = milvus.compact(collection_name, _async=True, _cakkback=_compact_callback) # Or invoke result() to get results: # compact_future = milvus.compact(collection_name, _async=True) # status = compact_future.result() compact_furure.done() # Get demo_collection row count status, result = milvus.count_entities(collection_name) # present collection info _, info = milvus.get_collection_stats(collection_name) print(info) # create index of vectors, search more rapidly index_param = {'nlist': 2048} def _index_callback(status): if status.OK(): print("Create index successfully") else: print("Create index failed.", status.message) # Create ivflat index in demo_collection # You can search vectors without creating index. however, Creating index help to # search faster print("Creating index: {}".format(index_param)) index_future = milvus.create_index(collection_name, IndexType.IVF_FLAT, index_param, _async=True, _callback=_index_callback) # Or invoke result() to get results: # index_future = milvus.create_index(collection_name, IndexType.IVF_FLAT, index_param, _async=True) # status = index_future.result() index_future.done() # describe index, get information of index status, index = milvus.get_index_info(collection_name) print(index) # Use the top 10 vectors for similarity search query_vectors = vectors[0:10] # execute vector similarity search search_param = {"nprobe": 16} print("Searching ... ") def _search_callback(status, results): # if status.OK(): # print("Search successfully") # else: # print("Search failed.", status.message) if status.OK(): # indicate search result # also use by: # `results.distance_array[0][0] == 0.0 or results.id_array[0][0] == ids[0]` if results[0][0].distance == 0.0: # or results[0][0].id == ids[0]: print('Query result is correct') else: print('Query result isn\'t correct') # print results print(results) else: print("Search failed. ", status) param = { 'collection_name': collection_name, 'query_records': query_vectors, 'top_k': 1, 'params': search_param, "_async": True, "_callback": _search_callback } search_future = milvus.search(**param) # Or invoke result() to get results: # # param = { # 'collection_name': collection_name, # 'query_records': query_vectors, # 'top_k': 1, # 'params': search_param, # "_async": True, # } # search_future = milvus.search(param) # status, results = index_future.result() search_future.done() # Delete demo_collection status = milvus.drop_collection(collection_name)
def main(): # Specify server addr when create milvus client instance # milvus client instance maintain a connection pool, param # `pool_size` specify the max connection num. # 获取服务端的连接 milvus = Milvus(_HOST, _PORT) # Create collection demo_collection if it dosen't exist. # 创建collection collection_name = 'example_collection_' # 看是否有这个collection status, ok = milvus.has_collection(collection_name) # 如果没有则创建 if not ok: param = { 'collection_name': collection_name, 'dimension': _DIM, 'index_file_size': _INDEX_FILE_SIZE, # optional 'metric_type': MetricType.L2 # optional } # 创建collection milvus.create_collection(param) # Show collections in Milvus server # 查看所有的collection _, collections = milvus.list_collections() print(collections) # Describe demo_collection # 得到当前的collection _, collection = milvus.get_collection_info(collection_name) print(collection) # 10000 vectors with 128 dimension # element per dimension is float32 type # vectors should be a 2-D array # 创建10个长度为8的向量 vectors = [[random.random() for _ in range(_DIM)] for _ in range(10)] print(vectors) # You can also use numpy to generate random vectors: # vectors = np.random.rand(10000, _DIM).astype(np.float32) # Insert vectors into demo_collection, return status and vectors id list # 把这10个向量都插入milvus status, ids = milvus.insert(collection_name=collection_name, records=vectors) if not status.OK(): print("Insert failed: {}".format(status)) print(ids) # Flush collection inserted data to disk. # 数据落盘 milvus.flush([collection_name]) # Get demo_collection row count # 得到当前row的数量 status, result = milvus.count_entities(collection_name) print(status) print(result) # present collection statistics info # 查看collection的统计数据 _, info = milvus.get_collection_stats(collection_name) print(info) # Obtain raw vectors by providing vector ids # 得到前十个数据 status, result_vectors = milvus.get_entity_by_id(collection_name, ids[:10]) print(result_vectors) # create index of vectors, search more rapidly # 创建索引 index_param = {'nlist': 2048} # Create ivflat index in demo_collection # You can search vectors without creating index. however, Creating index help to # search faster # 创建ivf_flat print("Creating index: {}".format(index_param)) status = milvus.create_index(collection_name, IndexType.IVF_FLAT, index_param) # describe index, get information of index # 得到索引的信息 status, index = milvus.get_index_info(collection_name) print(index) # Use the top 10 vectors for similarity search # 对前10个数据进行query query_vectors = vectors[0:10] # execute vector similarity search # 索引的搜索的中心点数量 search_param = {"nprobe": 16} print("Searching ... ") param = { 'collection_name': collection_name, 'query_records': query_vectors, 'top_k': 1, 'params': search_param, } # 进行搜索 status, results = milvus.search(**param) if status.OK(): print(results) # indicate search result # also use by: # `results.distance_array[0][0] == 0.0 or results.id_array[0][0] == ids[0]` if results[0][0].distance == 0.0 or results[0][0].id == ids[0]: print('Query result is correct') else: print('Query result isn\'t correct') # print results print(results) else: print("Search failed. ", status) # Delete demo_collection # 删除掉collection status = milvus.drop_collection(collection_name)
import numpy as np from pathlib import Path from milvus import Milvus, IndexType, MetricType, Status if __name__ == "__main__": milvus = Milvus(host='localhost', port='19530') collection_name = 'celeba' encodings = np.load('encodings.npy', allow_pickle=True).tolist() vectorEncoding = [] vectorId = [] for encoding in encodings: vectorEncoding.append(encoding['encondings'].tolist()) vectorId.append(encoding['id']) status, inserted_vector_ids = milvus.insert( collection_name=collection_name, records=vectorEncoding, ids=vectorId) milvus.flush([collection_name])
'collection_name': col_name, 'dimension': dim, 'index_file_size': 1024, 'metric_type': MetricType.L2 } milvus.create_collection(param=param) ivf_param = {'nlist': 16384} milvus.create_index(collection_name=col_name, index_type=IndexType.IVF_FLAT, params=ivf_param) vectors = [[random.random() for _ in range(dim)] for _ in range(2000)] vector_ids = list(range(2000)) _, ids = milvus.insert(collection_name=col_name, records=vectors, ids=vector_ids) # print(ids) time.sleep(1) search_param = {'nprobe': 16} q_records = [[random.random() for _ in range(dim)] for _ in range(5)] _, result = milvus.search(collection_name=col_name, query_records=q_records, top_k=2, params=search_param) # for r in result: # print(r) print(result.id_array) print(result)
import random import numpy as np #vectors = [[random.random() for _ in range(256)] for _ in range(3)] #print(np.shape(np.array(vectors))) #vector_ids = [1,2,3] vectors = [] vector_ids = [] id_query = {} with open("embed.txt", "r") as f: for line in f: line_lst = line.strip().split("\t") vectors.append(list(map(float, (line_lst[2].split())))) vector_ids.append(int(line_lst[1])) id_query[int(line_lst[1])] = line_lst[0] milvus.insert(collection_name='test01', records=vectors, ids=vector_ids) ivf_param = {'nlist': 16384} milvus.create_index('test01', IndexType.IVF_FLAT, ivf_param) search_param = {'nprobe': 16} q_records = vectors[:1] t = time.time() result = milvus.search(collection_name='test01', query_records=q_records, top_k=10, params=search_param)[1] print(result) print(time.time() - t) t = time.time()
class MilvusDocumentStore(SQLDocumentStore): """ Milvus (https://milvus.io/) is a highly reliable, scalable Document Store specialized on storing and processing vectors. Therefore, it is particularly suited for Haystack users that work with dense retrieval methods (like DPR). In contrast to FAISS, Milvus ... - runs as a separate service (e.g. a Docker container) and can scale easily in a distributed environment - allows dynamic data management (i.e. you can insert/delete vectors without recreating the whole index) - encapsulates multiple ANN libraries (FAISS, ANNOY ...) This class uses Milvus for all vector related storage, processing and querying. The meta-data (e.g. for filtering) and the document text are however stored in a separate SQL Database as Milvus does not allow these data types (yet). Usage: 1. Start a Milvus server (see https://milvus.io/docs/v1.0.0/install_milvus.md) 2. Init a MilvusDocumentStore in Haystack """ def __init__( self, sql_url: str = "sqlite:///", milvus_url: str = "tcp://localhost:19530", connection_pool: str = "SingletonThread", index: str = "document", vector_dim: int = 768, index_file_size: int = 1024, similarity: str = "dot_product", index_type: IndexType = IndexType.FLAT, index_param: Optional[Dict[str, Any]] = None, search_param: Optional[Dict[str, Any]] = None, update_existing_documents: bool = False, return_embedding: bool = False, embedding_field: str = "embedding", progress_bar: bool = True, **kwargs, ): """ :param sql_url: SQL connection URL for storing document texts and metadata. It defaults to a local, file based SQLite DB. For large scale deployment, Postgres is recommended. If using MySQL then same server can also be used for Milvus metadata. For more details see https://milvus.io/docs/v1.0.0/data_manage.md. :param milvus_url: Milvus server connection URL for storing and processing vectors. Protocol, host and port will automatically be inferred from the URL. See https://milvus.io/docs/v1.0.0/install_milvus.md for instructions to start a Milvus instance. :param connection_pool: Connection pool type to connect with Milvus server. Default: "SingletonThread". :param index: Index name for text, embedding and metadata (in Milvus terms, this is the "collection name"). :param vector_dim: The embedding vector size. Default: 768. :param index_file_size: Specifies the size of each segment file that is stored by Milvus and its default value is 1024 MB. When the size of newly inserted vectors reaches the specified volume, Milvus packs these vectors into a new segment. Milvus creates one index file for each segment. When conducting a vector search, Milvus searches all index files one by one. As a rule of thumb, we would see a 30% ~ 50% increase in the search performance after changing the value of index_file_size from 1024 to 2048. Note that an overly large index_file_size value may cause failure to load a segment into the memory or graphics memory. (From https://milvus.io/docs/v1.0.0/performance_faq.md#How-can-I-get-the-best-performance-from-Milvus-through-setting-index_file_size) :param similarity: The similarity function used to compare document vectors. 'dot_product' is the default and recommended for DPR embeddings. 'cosine' is recommended for Sentence Transformers, but is not directly supported by Milvus. However, you can normalize your embeddings and use `dot_product` to get the same results. See https://milvus.io/docs/v1.0.0/metric.md?Inner-product-(IP)#floating. :param index_type: Type of approximate nearest neighbour (ANN) index used. The choice here determines your tradeoff between speed and accuracy. Some popular options: - FLAT (default): Exact method, slow - IVF_FLAT, inverted file based heuristic, fast - HSNW: Graph based, fast - ANNOY: Tree based, fast See: https://milvus.io/docs/v1.0.0/index.md :param index_param: Configuration parameters for the chose index_type needed at indexing time. For example: {"nlist": 16384} as the number of cluster units to create for index_type IVF_FLAT. See https://milvus.io/docs/v1.0.0/index.md :param search_param: Configuration parameters for the chose index_type needed at query time For example: {"nprobe": 10} as the number of cluster units to query for index_type IVF_FLAT. See https://milvus.io/docs/v1.0.0/index.md :param update_existing_documents: Whether to update any existing documents with the same ID when adding documents. When set as True, any document with an existing ID gets updated. If set to False, an error is raised if the document ID of the document being added already exists. :param return_embedding: To return document embedding. :param embedding_field: Name of field containing an embedding vector. :param progress_bar: Whether to show a tqdm progress bar or not. Can be helpful to disable in production deployments to keep the logs clean. """ self.milvus_server = Milvus(uri=milvus_url, pool=connection_pool) self.vector_dim = vector_dim self.index_file_size = index_file_size if similarity == "dot_product": self.metric_type = MetricType.IP self.similarity = similarity else: raise ValueError( "The Milvus document store can currently only support dot_product similarity. " "Please set similarity=\"dot_product\"") self.index_type = index_type self.index_param = index_param or {"nlist": 16384} self.search_param = search_param or {"nprobe": 10} self.index = index self._create_collection_and_index_if_not_exist(self.index) self.return_embedding = return_embedding self.embedding_field = embedding_field self.progress_bar = progress_bar super().__init__(url=sql_url, update_existing_documents=update_existing_documents, index=index) def __del__(self): return self.milvus_server.close() def _create_collection_and_index_if_not_exist( self, index: Optional[str] = None, index_param: Optional[Dict[str, Any]] = None): index = index or self.index index_param = index_param or self.index_param status, ok = self.milvus_server.has_collection(collection_name=index) if not ok: collection_param = { 'collection_name': index, 'dimension': self.vector_dim, 'index_file_size': self.index_file_size, 'metric_type': self.metric_type } status = self.milvus_server.create_collection(collection_param) if status.code != Status.SUCCESS: raise RuntimeError( f'Collection creation on Milvus server failed: {status}') status = self.milvus_server.create_index(index, self.index_type, index_param) if status.code != Status.SUCCESS: raise RuntimeError( f'Index creation on Milvus server failed: {status}') def _create_document_field_map(self) -> Dict: return { self.index: self.embedding_field, } def write_documents(self, documents: Union[List[dict], List[Document]], index: Optional[str] = None, batch_size: int = 10_000): """ Add new documents to the DocumentStore. :param documents: List of `Dicts` or List of `Documents`. If they already contain the embeddings, we'll index them right away in Milvus. If not, you can later call update_embeddings() to create & index them. :param index: (SQL) index name for storing the docs and metadata :param batch_size: When working with large number of documents, batching can help reduce memory footprint. :return: """ index = index or self.index self._create_collection_and_index_if_not_exist(index) field_map = self._create_document_field_map() if len(documents) == 0: logger.warning( "Calling DocumentStore.write_documents() with empty list") return document_objects = [ Document.from_dict(d, field_map=field_map) if isinstance(d, dict) else d for d in documents ] add_vectors = False if document_objects[0].embedding is None else True batched_documents = get_batches_from_generator(document_objects, batch_size) with tqdm(total=len(document_objects), disable=not self.progress_bar) as progress_bar: for document_batch in batched_documents: vector_ids = [] if add_vectors: doc_ids = [] embeddings = [] for doc in document_batch: doc_ids.append(doc.id) if isinstance(doc.embedding, np.ndarray): embeddings.append(doc.embedding.tolist()) elif isinstance(doc.embedding, list): embeddings.append(doc.embedding) else: raise AttributeError( f'Format of supplied document embedding {type(doc.embedding)} is not ' f'supported. Please use list or numpy.ndarray') if self.update_existing_documents: existing_docs = super().get_documents_by_id( ids=doc_ids, index=index) self._delete_vector_ids_from_milvus( documents=existing_docs, index=index) status, vector_ids = self.milvus_server.insert( collection_name=index, records=embeddings) if status.code != Status.SUCCESS: raise RuntimeError( f'Vector embedding insertion failed: {status}') docs_to_write_in_sql = [] for idx, doc in enumerate(document_batch): meta = doc.meta if add_vectors: meta["vector_id"] = vector_ids[idx] docs_to_write_in_sql.append(doc) super().write_documents(docs_to_write_in_sql, index=index) progress_bar.update(batch_size) progress_bar.close() self.milvus_server.flush([index]) if self.update_existing_documents: self.milvus_server.compact(collection_name=index) def update_embeddings( self, retriever: BaseRetriever, index: Optional[str] = None, batch_size: int = 10_000, update_existing_embeddings: bool = True, filters: Optional[Dict[str, List[str]]] = None, ):
client.connect() collection_name = 'example_collection_vector' param = { 'collection_name': collection_name, 'dimension': _DIM, 'index_file_size': 10, # optional 'metric_type': MetricType.L2 # optional } client.create_collection(param) # randomly generate 100000 vectors and insert collection vectors = [[random.random() for _ in range(_DIM)] for _ in range(10000)] client.insert(collection_name, vectors) # flush data to disk client.flush([collection_name]) # query collection's statistical information. status, info = client.collection_info(collection_name) if not status.OK(): print("Query collection statistical information fail. exiting ....") sys.exit(1) # show collection information print("Total amount of vectors in collection {} is {}".format( collection_name, info.count)) for par in info.partitions_stat: print("\tpartition tag: {}, vector count: {}".format(
def main(): milvus = Milvus() # Connect to Milvus server # You may need to change _HOST and _PORT accordingly param = {'host': _HOST, 'port': _PORT} status = milvus.connect(**param) if status.OK(): print("Server connected.") else: print("Server connect fail.") sys.exit(1) # Create collection demo_collection if it dosen't exist. collection_name = 'example_collection' status, ok = milvus.has_collection(collection_name) if not ok: param = { 'collection_name': collection_name, 'dimension': _DIM, 'index_file_size': _INDEX_FILE_SIZE, # optional 'metric_type': MetricType.L2 # optional } milvus.create_collection(param) # Show collections in Milvus server _, collections = milvus.show_collections() # present collection info _, info = milvus.collection_info(collection_name) print(info) # Describe demo_collection _, collection = milvus.describe_collection(collection_name) print(collection) # 10000 vectors with 16 dimension # element per dimension is float32 type # vectors should be a 2-D array vectors = [[random.random() for _ in range(_DIM)] for _ in range(10000)] # You can also use numpy to generate random vectors: # `vectors = np.random.rand(10000, 16).astype(np.float32)` # Insert vectors into demo_collection, return status and vectors id list status, ids = milvus.insert(collection_name=collection_name, records=vectors) # Flush collection inserted data to disk. milvus.flush([collection_name]) # Get demo_collection row count status, result = milvus.count_collection(collection_name) # create index of vectors, search more rapidly index_param = { 'nlist': 2048 } # Create ivflat index in demo_collection # You can search vectors without creating index. however, Creating index help to # search faster print("Creating index: {}".format(index_param)) status = milvus.create_index(collection_name, IndexType.IVF_FLAT, index_param) # describe index, get information of index status, index = milvus.describe_index(collection_name) print(index) # Use the top 10 vectors for similarity search query_vectors = vectors[0:10] # execute vector similarity search search_param = { "nprobe": 16 } param = { 'collection_name': collection_name, 'query_records': query_vectors, 'top_k': 1, 'params': search_param } print("Searching ... ") status, results = milvus.search(**param) if status.OK(): # indicate search result # also use by: # `results.distance_array[0][0] == 0.0 or results.id_array[0][0] == ids[0]` if results[0][0].distance == 0.0 or results[0][0].id == ids[0]: print('Query result is correct') else: print('Query result isn\'t correct') # print results print(results) # Delete demo_collection status = milvus.drop_collection(collection_name) # Disconnect from Milvus status = milvus.disconnect()
class MilvusHelper: def __init__(self): try: self.client = Milvus(host=MILVUS_HOST, port=MILVUS_PORT) LOGGER.debug( "Successfully connect to Milvus with IP:{} and PORT:{}".format( MILVUS_HOST, MILVUS_PORT)) except Exception as e: LOGGER.error("Failed to connect Milvus: {}".format(e)) sys.exit(1) # Return if Milvus has the collection def has_collection(self, collection_name): try: status = self.client.has_collection(collection_name)[1] return status except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1) # Create milvus collection if not exists def create_colllection(self, collection_name): try: if not self.has_collection(collection_name): collection_param = { 'collection_name': collection_name, 'dimension': VECTOR_DIMENSION, 'index_file_size': INDEX_FILE_SIZE, 'metric_type': METRIC_TYPE } status = self.client.create_collection(collection_param) if status.code != 0: raise Exception(status.message) LOGGER.debug( "Create Milvus collection: {}".format(collection_name)) except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1) # Batch insert vectors to milvus collection def insert(self, collection_name, vectors): try: self.create_colllection(collection_name) status, ids = self.client.insert(collection_name=collection_name, records=vectors) if not status.code: LOGGER.debug( "Insert vectors to Milvus in collection: {} with {} rows". format(collection_name, len(vectors))) return ids else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1) # Create IVF_FLAT index on milvus collection def create_index(self, collection_name): try: index_param = {'nlist': 16384} status = self.client.create_index(collection_name, IndexType.IVF_FLAT, index_param) if not status.code: LOGGER.debug( "Successfully create index in collection:{} with param:{}". format(collection_name, index_param)) return status else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to create index: {}".format(e)) sys.exit(1) # Delete Milvus collection def delete_collection(self, collection_name): try: status = self.client.drop_collection( collection_name=collection_name) if not status.code: LOGGER.debug( "Successfully drop collection: {}".format(collection_name)) return status else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to drop collection: {}".format(e)) sys.exit(1) # Search vector in milvus collection def search_vectors(self, collection_name, vectors, top_k): try: search_param = {'nprobe': 16} status, result = self.client.search( collection_name=collection_name, query_records=vectors, top_k=top_k, params=search_param) if not status.code: LOGGER.debug("Successfully search in collection: {}".format( collection_name)) return result else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to search vectors in Milvus: {}".format(e)) sys.exit(1) # Get the number of milvus collection def count(self, collection_name): try: status, num = self.client.count_entities( collection_name=collection_name) if not status.code: LOGGER.debug( "Successfully get the num:{} of the collection:{}".format( num, collection_name)) return num else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to count vectors in Milvus: {}".format(e)) sys.exit(1)
def main(): # Connect to Milvus server # You may need to change _HOST and _PORT accordingly param = {'host': _HOST, 'port': _PORT} # You can create a instance specified server addr and # invoke rpc method directly client = Milvus(**param) # Create collection demo_collection if it dosen't exist. collection_name = 'demo_partition_collection' partition_tag = "random" status, ok = client.has_collection(collection_name) # if collection exists, then drop it if status.OK() and ok: client.drop_collection(collection_name) param = { 'collection_name': collection_name, 'dimension': _DIM, 'index_file_size': _INDEX_FILE_SIZE, # optional 'metric_type': MetricType.L2 # optional } client.create_collection(param) # Show collections in Milvus server _, collections = client.show_collections() # Describe collection _, collection = client.describe_collection(collection_name) print(collection) # create partition client.create_partition(collection_name, partition_tag=partition_tag) # display partitions _, partitions = client.show_partitions(collection_name) # 10000 vectors with 16 dimension # element per dimension is float32 type # vectors should be a 2-D array vectors = [[random.random() for _ in range(_DIM)] for _ in range(10000)] # You can also use numpy to generate random vectors: # `vectors = np.random.rand(10000, 16).astype(np.float32).tolist()` # Insert vectors into partition of collection, return status and vectors id list status, ids = client.insert(collection_name=collection_name, records=vectors, partition_tag=partition_tag) # Wait for 6 seconds, until Milvus server persist vector data. time.sleep(6) # Get demo_collection row count status, num = client.count_collection(collection_name) # create index of vectors, search more rapidly index_param = { 'nlist': 2048 } # Create ivflat index in demo_collection # You can search vectors without creating index. however, Creating index help to # search faster status = client.create_index(collection_name, IndexType.IVF_FLAT, index_param) # describe index, get information of index status, index = client.describe_index(collection_name) print(index) # Use the top 10 vectors for similarity search query_vectors = vectors[0:10] # execute vector similarity search, search range in partition `partition1` search_param = { "nprobe": 10 } param = { 'collection_name': collection_name, 'query_records': query_vectors, 'top_k': 1, 'partition_tags': ["random"], 'params': search_param } status, results = client.search(**param) if status.OK(): # indicate search result # also use by: # `results.distance_array[0][0] == 0.0 or results.id_array[0][0] == ids[0]` if results[0][0].distance == 0.0 or results[0][0].id == ids[0]: print('Query result is correct') else: print('Query result isn\'t correct') # print results print(results) # Delete partition. You can also invoke `drop_collection()`, so that all of partitions belongs to # designated collections will be deleted. # status = client.drop_partition(collection_name, partition_tag) # Delete collection. All of partitions of this collection will be dropped. status = client.drop_collection(collection_name)
"values": release_years, "type": DataType.INT32 }, { "name": "embedding", "values": embeddings, "type": DataType.FLOAT_VECTOR }, ] # ------ # Basic insert: # After preparing the data, we are going to insert them into our collection. # The number of films inserted should be 8657. # ------ ids = client.insert(collection_name, hybrid_entities, ids) client.flush([collection_name]) after_flush_counts = client.count_entities(collection_name) print(" > There are {} films in collection `{}` after flush".format( after_flush_counts, collection_name)) # ------ # Basic create index: # Now that we have inserted all the films into Milvus, we are going to build index with these data. # # While building index, we have to indicate which `field` to build index for, the `index_type`, # `metric_type` and params for the specific index type. In our case, we want to build a `IVF_FLAT` # index, so the specific params are "nlist". See pymilvus documentation # (https://milvus-io.github.io/milvus-sdk-python/pythondoc/v0.3.0/index.html) for `index_type` we # support and the params accordingly.
}, { "_id": 2, "duration": 226, "release_year": 2002, "embedding": random_bin_vector(128) }, { "_id": 3, "duration": 252, "release_year": 2003, "embedding": random_bin_vector(128) } ] ids = client.insert(collection_name, The_Lord_of_the_Rings, partition_tag="American") client.flush([collection_name]) print("\n----------insert----------") print("Films are inserted and the ids are: {}".format(ids)) # ------ # Basic collection stats: # We can get the detail of collection statistics info by `get_collection_stats` # ------ stats = client.get_collection_stats(collection_name) print("\n----------get collection stats----------") pprint(stats) # ------
class MilvusClient(object): def __init__(self, collection_name=None, host=None, port=None, timeout=180): self._collection_name = collection_name start_time = time.time() if not host: host = SERVER_HOST_DEFAULT if not port: port = SERVER_PORT_DEFAULT logger.debug(host) logger.debug(port) # retry connect remote server i = 0 while time.time() < start_time + timeout: try: self._milvus = Milvus(host=host, port=port, try_connect=False, pre_ping=False) break except Exception as e: logger.error(str(e)) logger.error("Milvus connect failed: %d times" % i) i = i + 1 time.sleep(i) if time.time() > start_time + timeout: raise Exception("Server connect timeout") # self._metric_type = None def __str__(self): return 'Milvus collection %s' % self._collection_name def check_status(self, status): if not status.OK(): logger.error(status.message) logger.error(self._milvus.server_status()) logger.error(self.count()) raise Exception("Status not ok") def check_result_ids(self, result): for index, item in enumerate(result): if item[0].distance >= epsilon: logger.error(index) logger.error(item[0].distance) raise Exception("Distance wrong") # only support the given field name def create_collection(self, dimension, data_type=DataType.FLOAT_VECTOR, auto_id=False, collection_name=None, other_fields=None): self._dimension = dimension if not collection_name: collection_name = self._collection_name vec_field_name = utils.get_default_field_name(data_type) fields = [{ "name": vec_field_name, "type": data_type, "params": { "dim": dimension } }] if other_fields: other_fields = other_fields.split(",") if "int" in other_fields: fields.append({ "name": utils.DEFAULT_INT_FIELD_NAME, "type": DataType.INT64 }) if "float" in other_fields: fields.append({ "name": utils.DEFAULT_FLOAT_FIELD_NAME, "type": DataType.FLOAT }) create_param = {"fields": fields, "auto_id": auto_id} try: self._milvus.create_collection(collection_name, create_param) logger.info("Create collection: <%s> successfully" % collection_name) except Exception as e: logger.error(str(e)) raise def create_partition(self, tag, collection_name=None): if not collection_name: collection_name = self._collection_name self._milvus.create_partition(collection_name, tag) def generate_values(self, data_type, vectors, ids): values = None if data_type in [DataType.INT32, DataType.INT64]: values = ids elif data_type in [DataType.FLOAT, DataType.DOUBLE]: values = [(i + 0.0) for i in ids] elif data_type in [DataType.FLOAT_VECTOR, DataType.BINARY_VECTOR]: values = vectors return values def generate_entities(self, vectors, ids=None, collection_name=None): entities = [] if collection_name is None: collection_name = self._collection_name info = self.get_info(collection_name) for field in info["fields"]: field_type = field["type"] entities.append({ "name": field["name"], "type": field_type, "values": self.generate_values(field_type, vectors, ids) }) return entities @time_wrapper def insert(self, entities, ids=None, collection_name=None): tmp_collection_name = self._collection_name if collection_name is None else collection_name try: insert_ids = self._milvus.insert(tmp_collection_name, entities, ids=ids) return insert_ids except Exception as e: logger.error(str(e)) def get_dimension(self): info = self.get_info() for field in info["fields"]: if field["type"] in [ DataType.FLOAT_VECTOR, DataType.BINARY_VECTOR ]: return field["params"]["dim"] def get_rand_ids(self, length): segment_ids = [] while True: stats = self.get_stats() segments = stats["partitions"][0]["segments"] # random choice one segment segment = random.choice(segments) try: segment_ids = self._milvus.list_id_in_segment( self._collection_name, segment["id"]) except Exception as e: logger.error(str(e)) if not len(segment_ids): continue elif len(segment_ids) > length: return random.sample(segment_ids, length) else: logger.debug("Reset length: %d" % len(segment_ids)) return segment_ids # def get_rand_ids_each_segment(self, length): # res = [] # status, stats = self._milvus.get_collection_stats(self._collection_name) # self.check_status(status) # segments = stats["partitions"][0]["segments"] # segments_num = len(segments) # # random choice from each segment # for segment in segments: # status, segment_ids = self._milvus.list_id_in_segment(self._collection_name, segment["name"]) # self.check_status(status) # res.extend(segment_ids[:length]) # return segments_num, res # def get_rand_entities(self, length): # ids = self.get_rand_ids(length) # status, get_res = self._milvus.get_entity_by_id(self._collection_name, ids) # self.check_status(status) # return ids, get_res def get(self): get_ids = random.randint(1, 1000000) self._milvus.get_entity_by_id(self._collection_name, [get_ids]) @time_wrapper def get_entities(self, get_ids): get_res = self._milvus.get_entity_by_id(self._collection_name, get_ids) return get_res @time_wrapper def delete(self, ids, collection_name=None): tmp_collection_name = self._collection_name if collection_name is None else collection_name self._milvus.delete_entity_by_id(tmp_collection_name, ids) def delete_rand(self): delete_id_length = random.randint(1, 100) count_before = self.count() logger.debug("%s: length to delete: %d" % (self._collection_name, delete_id_length)) delete_ids = self.get_rand_ids(delete_id_length) self.delete(delete_ids) self.flush() logger.info("%s: count after delete: %d" % (self._collection_name, self.count())) get_res = self._milvus.get_entity_by_id(self._collection_name, delete_ids) for item in get_res: assert not item # if count_before - len(delete_ids) < self.count(): # logger.error(delete_ids) # raise Exception("Error occured") @time_wrapper def flush(self, _async=False, collection_name=None): tmp_collection_name = self._collection_name if collection_name is None else collection_name self._milvus.flush([tmp_collection_name], _async=_async) @time_wrapper def compact(self, collection_name=None): tmp_collection_name = self._collection_name if collection_name is None else collection_name status = self._milvus.compact(tmp_collection_name) self.check_status(status) @time_wrapper def create_index(self, field_name, index_type, metric_type, _async=False, index_param=None): index_type = INDEX_MAP[index_type] metric_type = utils.metric_type_trans(metric_type) logger.info( "Building index start, collection_name: %s, index_type: %s, metric_type: %s" % (self._collection_name, index_type, metric_type)) if index_param: logger.info(index_param) index_params = { "index_type": index_type, "metric_type": metric_type, "params": index_param } self._milvus.create_index(self._collection_name, field_name, index_params, _async=_async) # TODO: need to check def describe_index(self, field_name): # stats = self.get_stats() info = self._milvus.describe_index(self._collection_name, field_name) index_info = {"index_type": "flat", "index_param": None} for field in info["fields"]: for index in field['indexes']: if not index or "index_type" not in index: continue else: for k, v in INDEX_MAP.items(): if index['index_type'] == v: index_info['index_type'] = k index_info['index_param'] = index['params'] return index_info return index_info def drop_index(self, field_name): logger.info("Drop index: %s" % self._collection_name) return self._milvus.drop_index(self._collection_name, field_name) @time_wrapper def query(self, vector_query, filter_query=None, collection_name=None): tmp_collection_name = self._collection_name if collection_name is None else collection_name must_params = [vector_query] if filter_query: must_params.extend(filter_query) query = {"bool": {"must": must_params}} result = self._milvus.search(tmp_collection_name, query) return result @time_wrapper def load_and_query(self, vector_query, filter_query=None, collection_name=None): tmp_collection_name = self._collection_name if collection_name is None else collection_name must_params = [vector_query] if filter_query: must_params.extend(filter_query) query = {"bool": {"must": must_params}} self.load_collection(tmp_collection_name) result = self._milvus.search(tmp_collection_name, query) return result def get_ids(self, result): idss = result._entities.ids ids = [] len_idss = len(idss) len_r = len(result) top_k = len_idss // len_r for offset in range(0, len_idss, top_k): ids.append(idss[offset:min(offset + top_k, len_idss)]) return ids def query_rand(self, nq_max=100): # for ivf search dimension = 128 top_k = random.randint(1, 100) nq = random.randint(1, nq_max) nprobe = random.randint(1, 100) search_param = {"nprobe": nprobe} query_vectors = [[random.random() for _ in range(dimension)] for _ in range(nq)] metric_type = random.choice(["l2", "ip"]) logger.info("%s, Search nq: %d, top_k: %d, nprobe: %d" % (self._collection_name, nq, top_k, nprobe)) vec_field_name = utils.get_default_field_name() vector_query = { "vector": { vec_field_name: { "topk": top_k, "query": query_vectors, "metric_type": utils.metric_type_trans(metric_type), "params": search_param } } } self.query(vector_query) def load_query_rand(self, nq_max=100): # for ivf search dimension = 128 top_k = random.randint(1, 100) nq = random.randint(1, nq_max) nprobe = random.randint(1, 100) search_param = {"nprobe": nprobe} query_vectors = [[random.random() for _ in range(dimension)] for _ in range(nq)] metric_type = random.choice(["l2", "ip"]) logger.info("%s, Search nq: %d, top_k: %d, nprobe: %d" % (self._collection_name, nq, top_k, nprobe)) vec_field_name = utils.get_default_field_name() vector_query = { "vector": { vec_field_name: { "topk": top_k, "query": query_vectors, "metric_type": utils.metric_type_trans(metric_type), "params": search_param } } } self.load_and_query(vector_query) # TODO: need to check def count(self, collection_name=None): if collection_name is None: collection_name = self._collection_name row_count = self._milvus.get_collection_stats( collection_name)["row_count"] logger.debug("Row count: %d in collection: <%s>" % (row_count, collection_name)) return row_count def drop(self, timeout=120, collection_name=None): timeout = int(timeout) if collection_name is None: collection_name = self._collection_name logger.info("Start delete collection: %s" % collection_name) self._milvus.drop_collection(collection_name) i = 0 while i < timeout: try: row_count = self.count(collection_name=collection_name) if row_count: time.sleep(1) i = i + 1 continue else: break except Exception as e: logger.debug(str(e)) break if i >= timeout: logger.error("Delete collection timeout") def get_stats(self): return self._milvus.get_collection_stats(self._collection_name) def get_info(self, collection_name=None): # pdb.set_trace() if collection_name is None: collection_name = self._collection_name return self._milvus.get_collection_info(collection_name) def show_collections(self): return self._milvus.list_collections() def exists_collection(self, collection_name=None): if collection_name is None: collection_name = self._collection_name res = self._milvus.has_collection(collection_name) return res def clean_db(self): collection_names = self.show_collections() for name in collection_names: self.drop(collection_name=name) @time_wrapper def load_collection(self, collection_name=None): if collection_name is None: collection_name = self._collection_name return self._milvus.load_collection(collection_name, timeout=3000) @time_wrapper def release_collection(self, collection_name=None): if collection_name is None: collection_name = self._collection_name return self._milvus.release_collection(collection_name, timeout=3000) @time_wrapper def load_partitions(self, tag_names, collection_name=None): if collection_name is None: collection_name = self._collection_name return self._milvus.load_partitions(collection_name, tag_names, timeout=3000) @time_wrapper def release_partitions(self, tag_names, collection_name=None): if collection_name is None: collection_name = self._collection_name return self._milvus.release_partitions(collection_name, tag_names, timeout=3000)
def main(): milvus = Milvus(_HOST, _PORT) # num = random.randint(1, 100000) num = 100000 # Create collection demo_collection if it dosen't exist. collection_name = 'example_hybrid_collections_{}'.format(num) if milvus.has_collection(collection_name): milvus.drop_collection(collection_name) collection_param = { "fields": [{ "field": "A", "type": DataType.INT32 }, { "field": "B", "type": DataType.INT32 }, { "field": "C", "type": DataType.INT64 }, { "field": "Vec", "type": DataType.FLOAT_VECTOR, "params": { "dim": 128, "metric_type": "L2" } }], "segment_size": 100 } milvus.create_collection(collection_name, collection_param) milvus.compact(collection_name) # milvus.create_partition(collection_name, "p_01", timeout=1800) # pars = milvus.list_partitions(collection_name) # ok = milvus.has_partition(collection_name, "p_01", timeout=1800) # assert ok # ok = milvus.has_partition(collection_name, "p_02") # assert not ok # for p in pars: # if p == "_default": # continue # milvus.drop_partition(collection_name, p) # milvus.drop_collection(collection_name) # sys.exit(0) A_list = [random.randint(0, 255) for _ in range(num)] vec = [[random.random() for _ in range(128)] for _ in range(num)] hybrid_entities = [{ "field": "A", "values": A_list, "type": DataType.INT32 }, { "field": "B", "values": A_list, "type": DataType.INT32 }, { "field": "C", "values": A_list, "type": DataType.INT64 }, { "field": "Vec", "values": vec, "type": DataType.FLOAT_VECTOR, "params": { "dim": 128 } }] for slice_e in utils.entities_slice(hybrid_entities): ids = milvus.insert(collection_name, slice_e) milvus.flush([collection_name]) print("Flush ... ") # time.sleep(3) count = milvus.count_entities(collection_name) milvus.delete_entity_by_id(collection_name, ids[:1]) milvus.flush([collection_name]) print("Get entity be id start ...... ") entities = milvus.get_entity_by_id(collection_name, ids[:1]) et = entities.dict() milvus.delete_entity_by_id(collection_name, ids[1:2]) milvus.flush([collection_name]) print("Create index ......") milvus.create_index(collection_name, "Vec", { "index_type": "IVF_FLAT", "metric_type": "L2", "params": { "nlist": 100 } }) print("Create index done.") info = milvus.get_collection_info(collection_name) print(info) stats = milvus.get_collection_stats(collection_name) print("\nstats\n") print(stats) query_hybrid = \ { "bool": { "must": [ { "term": { "A": [1, 2, 5] } }, { "range": { "B": {"GT": 1, "LT": 100} } }, { "vector": { "Vec": { "topk": 10, "query": vec[: 10000], "params": {"nprobe": 10} } } } ], }, } # print("Start searach ..", flush=True) # results = milvus.search(collection_name, query_hybrid) # print(results) # # for r in list(results): # print("ids", r.ids) # print("distances", r.distances) t0 = time.time() count = 0 results = milvus.search(collection_name, query_hybrid, fields=["B"]) for r in list(results): # print("ids", r.ids) # print("distances", r.distances) for rr in r: count += 1 # print(rr.entity.get("B")) print("Search cost {} s".format(time.time() - t0)) # for result in results: # for r in result: # print(f"{r}") # itertor entity id # for result in results: # for r in result: # # get distance # dis = r.distance # id_ = r.id # # obtain all field name # fields = r.entity.fields # for f in fields: # # get field value by field name # # fv = r.entity. # fv = r.entity.value_of_field(f) # print(fv) milvus.drop_collection(collection_name)