python读写milvus

目录

安装milvus客户端

创建milvus数据表

数据表插入向量

查询数据表

查向量数量

封装成类

---

python3.6版本

安装milvus客户端

pip install pymilvus==1.1.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

创建milvus数据表

from milvus import Milvus, MetricType

milvus = Milvus(host='milvus机器ip', port='19530')

params = {'collection_name':'test01',
          'dimension':512,
          'index_file_size':1024,
          'metric_type':MetricType.IP}

status = milvus.create_collection(params)
print(status)

Status(code=0, message='Create collection successfully!')

数据表插入向量

模拟10条512维的特征向量

import random

vectors = [[random.random() for _ in range(512)] for _ in range(10)]
print(len(vectors),len(vectors[0]))

10 512

模拟20条向量入库

from milvus import Milvus, MetricType

milvus = Milvus(host='milvus机器ip', port='19530')

params = {'collection_name':'test01',
          'dimension':512,
          'index_file_size':1024,
          'metric_type':MetricType.IP}

status = milvus.create_collection(params)
print(status)

vectors = [[1.0*i for _ in range(512)] for i in range(20)]
print(vectors[0][0],vectors[-1][0])

status, ids = milvus.insert(collection_name='test01', records=vectors)
print(status,ids)

milvus.flush(['test01'])
milvus.close()

Status(code=0, message='Create collection successfully!')
0.0 9.0
Status(code=0, message='Add vectors successfully!') [168846, ..., 1688460]

查询数据表

from milvus import Milvus, MetricType
from pprint import pprint

milvus = Milvus(host='milvus机器ip', port='19530')

status,info = milvus.get_collection_stats('test01')
print(status,info)

print('-'*20)
pprint(info)

Status(code=0, message='Success') {'partitions': [{'row_count': 20, 'segments': [{'data_size': 41120, 'index_name': 'IDMAP', 'name': '168846070', 'row_count': 20}], 'tag': '_default'}], 'row_count': 20}
--------------------
{'partitions': [{'row_count': 20,
                 'segments': [{'data_size': 41120,
                               'index_name': 'IDMAP',
                               'name': '1688460770',
                               'row_count': 20}],
                 'tag': '_default'}],
 'row_count': 20}

查向量数量

from milvus import Milvus, MetricType

milvus = Milvus(host='milvus机器ip', port='19530')

status,n = milvus.count_entities('test01')
print(status,n)

Status(code=0, message='Success!') 20

封装成类

class MilvusVector:
    def __init__(self):
        self.milvus = Milvus(host='milvus机器ip', port='19530')
        self.db_name = 'test01'
        params = {'collection_name': self.db_name, 'dimension': 512, 'index_file_size': 1024, 'metric_type': MetricType.IP}
        _, bool = self.milvus.has_collection(self.db_name)
        if bool == True:
            self.milvus.drop_collection(collection_name=self.db_name)
            self.milvus.create_collection(params)
            print('delete test01 ---> create test01')
        else:
            self.milvus.create_collection(params)
            print('create test01')

    def vectorNum(self):
        status, n = self.milvus.count_entities(self.db_name)
        return n

    def close(self):
        self.milvus.close()

    def insert(self,vector):
        status, ids = self.milvus.insert(collection_name=self.db_name,records=vector)
        self.milvus.flush([self.db_name])
        return ids[0]

    def search(self,vector):
        status, ips = self.milvus.search(collection_name=self.db_name, query_records=vector, top_k=1)
        id, ip = ips.id_array[0][0], ips.distance_array[0][0]
        return [id,ip]