mirror of https://github.com/milvus-io/milvus.git
130 lines
3.5 KiB
Python
130 lines
3.5 KiB
Python
import ujson
|
|
from milvus import Status
|
|
from functools import wraps
|
|
|
|
|
|
def error_status(func):
|
|
@wraps(func)
|
|
def inner(*args, **kwargs):
|
|
try:
|
|
results = func(*args, **kwargs)
|
|
except Exception as e:
|
|
return Status(code=Status.UNEXPECTED_ERROR, message=str(e)), None
|
|
|
|
return Status(code=0, message="Success"), results
|
|
|
|
return inner
|
|
|
|
|
|
class GrpcArgsParser(object):
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_CollectionSchema(cls, param):
|
|
_collection_schema = {
|
|
'collection_name': param.collection_name,
|
|
'dimension': param.dimension,
|
|
'index_file_size': param.index_file_size,
|
|
'metric_type': param.metric_type
|
|
}
|
|
|
|
return param.status, _collection_schema
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_CollectionName(cls, param):
|
|
return param.collection_name
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_FlushParam(cls, param):
|
|
return list(param.collection_name_array)
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_Index(cls, param):
|
|
_index = {
|
|
'index_type': param.index_type,
|
|
'params': param.extra_params[0].value
|
|
}
|
|
|
|
return _index
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_IndexParam(cls, param):
|
|
_collection_name = param.collection_name
|
|
_index_type = param.index_type
|
|
_index_param = {}
|
|
|
|
for params in param.extra_params:
|
|
if params.key == 'params':
|
|
_index_param = ujson.loads(str(params.value))
|
|
|
|
return _collection_name, _index_type, _index_param
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_Command(cls, param):
|
|
_cmd = param.cmd
|
|
|
|
return _cmd
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_RowRecord(cls, param):
|
|
return list(param.vector_data)
|
|
|
|
@classmethod
|
|
def parse_proto_PartitionParam(cls, param):
|
|
_collection_name = param.collection_name
|
|
_tag = param.tag
|
|
|
|
return _collection_name, _tag
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_SearchParam(cls, param):
|
|
_collection_name = param.collection_name
|
|
_topk = param.topk
|
|
|
|
if len(param.extra_params) == 0:
|
|
raise Exception("Search param loss")
|
|
_params = ujson.loads(str(param.extra_params[0].value))
|
|
|
|
_query_record_array = []
|
|
if param.query_record_array:
|
|
for record in param.query_record_array:
|
|
if record.float_data:
|
|
_query_record_array.append(list(record.float_data))
|
|
else:
|
|
_query_record_array.append(bytes(record.binary_data))
|
|
else:
|
|
raise Exception("Search argument parse error: record array is empty")
|
|
|
|
return _collection_name, _query_record_array, _topk, _params
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_DeleteByIDParam(cls, param):
|
|
_collection_name = param.collection_name
|
|
_id_array = list(param.id_array)
|
|
|
|
return _collection_name, _id_array
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_VectorIdentity(cls, param):
|
|
_collection_name = param.collection_name
|
|
_ids = list(param.id_array)
|
|
|
|
return _collection_name, _ids
|
|
|
|
@classmethod
|
|
@error_status
|
|
def parse_proto_GetVectorIDsParam(cls, param):
|
|
_collection__name = param.collection_name
|
|
_segment_name = param.segment_name
|
|
|
|
return _collection__name, _segment_name
|