import copy
import functools
from . import __version__
from .types import IndexType, MetricType
from .check import check_pass_param
from .grpc_handler import GrpcHandler
from .http_handler import HttpHandler
from .exceptions import ParamError, NotConnectError
def check_connect(f):
@functools.wraps(f)
def wrapper(self, *args, **kwargs):
if not self.connected():
raise NotConnectError('Please connect to the server first')
return f(self, *args, **kwargs)
return wrapper
[docs]class Milvus:
def __init__(self, host=None, port=None, handler="GRPC", **kwargs):
if handler == "GRPC":
self._handler = GrpcHandler(host=host, port=port, **kwargs)
elif handler == "HTTP":
self._handler = HttpHandler(host=host, port=port, **kwargs)
else:
raise ParamError("Unknown handler options, please use \'GRPC\' or \'HTTP\'")
def __enter__(self):
self._handler.__enter__()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self._handler.__exit__(exc_type, exc_val, exc_tb)
[docs] def set_hook(self, **kwargs):
return self._handler.set_hook(**kwargs)
@property
def status(self):
return self._handler.status
@property
def handler(self):
if isinstance(self._handler, GrpcHandler):
return "GRPC"
if isinstance(self._handler, HttpHandler):
return "HTTP"
return "NULL"
[docs] def connect(self, host=None, port=None, uri=None, timeout=1):
"""
Connects to the Milvus server.
:type host: str
:type port: str
:type uri: str
:type timeout: float
:param host: (Optional) Ip address of the Milvus server. The default is `127.0.0.1`.
:param port: (Optional) Port of the Milvus server. The default is `19530`.
:param uri: (Optional) URI of the Milvus server. Only TCP is supported. The default is `tcp://127.0.0.1:19530`.
:param timeout: (Optional) Connection timeout in milliseconds. The default is 3000.
:return: Indicates whether the connection is successful.
:rtype: Status
:raises: NotConnectError
"""
return self._handler.connect(host, port, uri, timeout)
[docs] def connected(self):
"""
Checks whether the client is connected to the Milvus server.
:return: Whether the client is connected to the Milvus server.
:rtype: bool
"""
return self._handler.connected()
[docs] @check_connect
def disconnect(self):
"""
Disconnects from the Milvus server.
:return: Whether the client is disconnected from the Milvus server.
:rtype: Status
"""
return self._handler.disconnect()
[docs] def client_version(self):
"""
Returns the version of the client.
:return: Version of the client.
:rtype: (str)
"""
return __version__
[docs] @check_connect
def server_status(self, timeout=10):
"""
Returns the status of the Milvus server.
:return:
Status: Whether the operation is successful.
str : Status of the Milvus server.
:rtype: (Status, str)
"""
return self._cmd("status", timeout)
[docs] @check_connect
def server_version(self, timeout=10):
"""
Returns the version of the Milvus server.
:return:
Status: Whether the operation is successful.
str : Version of the Milvus server.
:rtype: (Status, str)
"""
return self._cmd("version", timeout)
@check_connect
def _cmd(self, cmd, timeout=10):
check_pass_param(cmd=cmd)
return self._handler._cmd(cmd, timeout)
[docs] @check_connect
def create_collection(self, param, timeout=10):
"""
Creates a collection.
:type param: dict
:param param: Information needed to create a collection.
`param={'collection_name': 'name',
'dimension': 16,
'index_file_size': 1024 (default)，
'metric_type': Metric_type.L2 (default)
}`
:param timeout: Timeout in seconds.
:type timeout: double
:return: Whether the operation is successful.
:rtype: Status
"""
if not isinstance(param, dict):
raise ParamError('Param type incorrect, expect {} but get {} instead'
.format(type(dict), type(param)))
collection_param = copy.deepcopy(param)
if 'collection_name' not in collection_param:
raise ParamError('collection_name is required')
collection_name = collection_param["collection_name"]
collection_param.pop('collection_name')
if 'dimension' not in collection_param:
raise ParamError('dimension is required')
dim = collection_param["dimension"]
collection_param.pop("dimension")
index_file_size = collection_param.get('index_file_size', 1024)
collection_param.pop('index_file_size', None)
metric_type = collection_param.get('metric_type', MetricType.L2)
collection_param.pop('metric_type', None)
check_pass_param(collection_name=collection_name, dimension=dim, index_file_size=index_file_size,
metric_type=metric_type)
return self._handler.create_table(collection_name, dim, index_file_size, metric_type, collection_param, timeout)
[docs] @check_connect
def has_collection(self, collection_name, timeout=10):
"""
Checks whether a collection exists.
:param collection_name: Name of the collection to check.
:type collection_name: str
:param timeout: Timeout in seconds.
:type timeout: int
:return:
Status: indicate whether the operation is successful.
bool if given collection_name exists
"""
check_pass_param(collection_name=collection_name)
return self._handler.has_table(collection_name, timeout)
[docs] @check_connect
def describe_collection(self, collection_name, timeout=10):
"""
Returns information of a collection.
:type collection_name: str
:param collection_name: Name of the collection to describe.
:returns: (Status, table_schema)
Status: indicate if query is successful
table_schema: return when operation is successful
:rtype: (Status, TableSchema)
"""
check_pass_param(collection_name=collection_name)
return self._handler.describe_table(collection_name, timeout)
[docs] @check_connect
def count_collection(self, collection_name, timeout=10):
"""
Returns the number of vectors in a collection.
:type collection_name: str
:param collection_name: target table name.
:returns:
Status: indicate if operation is successful
res: int, table row count
"""
check_pass_param(collection_name=collection_name)
return self._handler.count_table(collection_name, timeout)
[docs] @check_connect
def show_collections(self, timeout=10):
"""
Returns information of all collections.
:return:
Status: indicate if this operation is successful
collections: list of table names, return when operation
is successful
:rtype:
(Status, list[str])
"""
return self._handler.show_tables(timeout)
[docs] @check_connect
def collection_info(self, collection_name, timeout=10):
return self._handler.show_table_info(collection_name, timeout)
[docs] @check_connect
def preload_collection(self, collection_name, timeout=None):
"""
Loads a collection for caching.
:type collection_name: str
:param collection_name: table to preload
:returns:
Status: indicate if invoke is successful
"""
check_pass_param(collection_name=collection_name)
return self._handler.preload_table(collection_name, timeout)
[docs] @check_connect
def drop_collection(self, collection_name, timeout=10):
"""
Deletes a collection by name.
:type collection_name: str
:param collection_name: Name of the table being deleted
:return: Status, indicate if operation is successful
:rtype: Status
"""
check_pass_param(collection_name=collection_name)
return self._handler.drop_table(collection_name, timeout)
[docs] @check_connect
def insert(self, collection_name, records, ids=None, partition_tag=None, params=None, timeout=-1):
"""
Insert vectors to a collection.
:param ids: list of id
:type ids: list[int]
:type collection_name: str
:param collection_name: collection name been inserted
:type records: list[list[float]]
`example records: [[1.2345],[1.2345]]`
`OR using Prepare.records`
:param records: list of vectors been inserted
:type partition_tag: str or None.
If partition_tag is None, vectors will be inserted to the collection rather than partitions.
:param partition_tag: the tag string of a collection
:type
:type timeout: int
:param timeout: time waiting for server response
:returns:
Status: indicate whether vectors are inserted successfully.
ids: IDs of the inserted vectors.
:rtype: (Status, list(int))
"""
check_pass_param(collection_name=collection_name, records=records,
ids=ids, partition_tag=partition_tag)
if ids is not None and len(records) != len(ids):
raise ParamError("length of vectors do not match that of ids")
params = params or dict()
if not isinstance(params, dict):
raise ParamError("Params must be a dictionary type")
return self._handler.insert(collection_name, records, ids, partition_tag, params, timeout)
[docs] @check_connect
def get_vector_by_id(self, collection_name, vector_id, timeout=None):
check_pass_param(collection_name=collection_name, ids=[vector_id])
return self._handler.get_vector_by_id(collection_name, vector_id, timeout=timeout)
[docs] @check_connect
def get_vector_ids(self, collection_name, segment_name, timeout=None):
check_pass_param(collection_name=collection_name)
check_pass_param(collection_name=segment_name)
return self._handler.get_vector_ids(collection_name, segment_name, timeout)
[docs] @check_connect
def create_index(self, collection_name, index_type=None, params=None, timeout=None):
"""
Creates index for a collection.
:param collection_name: Collection used to create index.
:type collection_name: str
:param index: index params
:type index: dict
index_param can be None
`example (default) param={'index_type': IndexType.FLAT,
'nlist': 16384}`
:param timeout: grpc request timeout.
if `timeout` = -1, method invoke a synchronous call, waiting util grpc response
else method invoke a asynchronous call, timeout work here
:type timeout: int
:return: Whether the operation is successful.
"""
_index_type = IndexType.FLAT if index_type is None else index_type
check_pass_param(collection_name=collection_name, index_type=_index_type)
params = params or dict()
if not isinstance(params, dict):
raise ParamError("Params must be a dictionary type")
return self._handler.create_index(collection_name, _index_type, params, timeout)
[docs] @check_connect
def describe_index(self, collection_name, timeout=10):
"""
Show index information of a collection.
:type collection_name: str
:param collection_name: table name been queried
:returns:
Status: Whether the operation is successful.
IndexSchema:
"""
check_pass_param(collection_name=collection_name)
return self._handler.describe_index(collection_name, timeout)
[docs] @check_connect
def drop_index(self, collection_name, timeout=10):
"""
Removes an index.
:param collection_name: target collection name.
:type collection_name: str
:return:
Status: Whether the operation is successful.
：:rtype: Status
"""
check_pass_param(collection_name=collection_name)
return self._handler.drop_index(collection_name, timeout)
[docs] @check_connect
def create_partition(self, collection_name, partition_tag, timeout=10):
"""
create a partition for a collection.
:param collection_name: Name of the collection.
:type collection_name: str
:param partition_name: Name of the partition.
:type partition_name: str
:param partition_tag: Name of the partition tag.
:type partition_tag: str
:param timeout: time waiting for response.
:type timeout: int
:return:
Status: Whether the operation is successful.
"""
check_pass_param(collection_name=collection_name, partition_tag=partition_tag)
return self._handler.create_partition(collection_name, partition_tag, timeout)
[docs] @check_connect
def show_partitions(self, collection_name, timeout=10):
"""
Show all partitions in a collection.
:param collection_name: target table name.
:type collection_name: str
:param timeout: time waiting for response.
:type timeout: int
:return:
Status: Whether the operation is successful.
partition_list:
"""
check_pass_param(collection_name=collection_name)
return self._handler.show_partitions(collection_name, timeout)
[docs] @check_connect
def drop_partition(self, collection_name, partition_tag, timeout=10):
"""
Deletes a partition in a collection.
:param collection_name: Collection name.
:type collection_name: str
:param partition_tag: Partition name.
:type partition_tag: str
:param timeout: time waiting for response.
:type timeout: int
:return:
Status: Whether the operation is successful.
"""
check_pass_param(collection_name=collection_name, partition_tag=partition_tag)
return self._handler.drop_partition(collection_name, partition_tag, timeout)
[docs] @check_connect
def search(self, collection_name, top_k, query_records, partition_tags=None, params=None):
"""
Search vectors in a collection.
:param collection_name: Name of the collection.
:type collection_name: str
:param top_k: number of vertors which is most similar with query vectors
:type top_k: int
:param nprobe: cell number of probe
:type nprobe: int
:param query_records: vectors to query
:type query_records: list[list[float32]]
:param partition_tags: tags to search
:type partition_tags: list
:return
Status: Whether the operation is successful.
result: query result
:rtype: (Status, TopKQueryResult)
"""
check_pass_param(collection_name=collection_name, topk=top_k,
records=query_records, partition_tag_array=partition_tags)
params = params or dict()
if not isinstance(params, dict):
raise ParamError("Params must be a dictionary type")
return self._handler.search(collection_name, top_k, query_records, partition_tags, params)
[docs] @check_connect
def search_in_files(self, collection_name, file_ids, query_records, top_k, params=None):
"""
Searches for vectors in specific files of a collection.
The Milvus server stores vector data into multiple files. Searching for vectors in specific files is a
method used in Mishards. Obtain more detail about Mishards, see
<a href="https://github.com/milvus-io/milvus/tree/master/shards">
:type collection_name: str
:param collection_name: table name been queried
:type file_ids: list[str] or list[int]
:param file_ids: Specified files id array
:type query_records: list[list[float]]
:param query_records: all vectors going to be queried
:param query_ranges: Optional ranges for conditional search.
If not specified, search in the whole table
:type top_k: int
:param top_k: how many similar vectors will be searched
:returns:
Status: indicate if query is successful
results: query result
:rtype: (Status, TopKQueryResult)
"""
check_pass_param(collection_name=collection_name, topk=top_k, records=query_records)
params = params or dict()
if not isinstance(params, dict):
raise ParamError("Params must be a dictionary type")
return self._handler.search_in_files(collection_name, file_ids,
query_records, top_k, params)
[docs] @check_connect
def delete_by_id(self, collection_name, id_array, timeout=None):
"""
Deletes vectors in a collection by vector ID.
"""
check_pass_param(collection_name=collection_name, ids=id_array)
return self._handler.delete_by_id(collection_name, id_array, timeout)
[docs] @check_connect
def flush(self, collection_name_array=None):
"""
Flushes vector data in one collection or multiple collections to disk.
"""
if collection_name_array is None:
return self._handler.flush([])
if not isinstance(collection_name_array, list):
raise ParamError("Collection name array must be type of list")
if len(collection_name_array) <= 0:
raise ParamError("Collection name array is not allowed to be empty")
for name in collection_name_array:
check_pass_param(collection_name=name)
return self._handler.flush(collection_name_array)
[docs] @check_connect
def compact(self, collection_name, timeout=None):
"""
Compacts segments in a collection.
"""
check_pass_param(collection_name=collection_name)
return self._handler.compact(collection_name, timeout)
[docs] @check_connect
def get_config(self, parent_key, child_key):
"""
Gets Milvus configurations.
"""
cmd = "get_config {}.{}".format(parent_key, child_key)
return self._cmd(cmd)
[docs] @check_connect
def set_config(self, parent_key, child_key, value):
"""
Sets Milvus configurations.
"""
cmd = "set_config {}.{} {}".format(parent_key, child_key, value)
return self._cmd(cmd)
# In old version of pymilvus, some methods are different from the new.
# apply alternative method name for compatibility
# get_collection_row_count = count_collection
# delete_collection = drop_collection
add_vectors = insert
search_vectors = search
search_vectors_in_files = search_in_files