def ShowTableInfo(self, request, context):
_status, _table_name = Parser.parse_proto_TableName(request)
if not _status.OK():
return milvus_pb2.TableInfo(status=status_pb2.Status(
error_code=_status.code, reason=_status.message), )
metadata = {'resp_class': milvus_pb2.TableInfo}
logger.info('ShowTableInfo {}'.format(_table_name))
_status, _info = self._table_info(metadata=metadata,
table_name=_table_name)
if _status.OK():
_table_info = milvus_pb2.TableInfo(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
total_row_count=_info.count)
for par_stat in _info.partitions_stat:
_par = milvus_pb2.PartitionStat(tag=par_stat.tag,
total_row_count=par_stat.count)
for seg_stat in par_stat.segments_stat:
_par.segments_stat.add(
segment_name=seg_stat.segment_name,
row_count=seg_stat.count,
index_name=seg_stat.index_name,
data_size=seg_stat.data_size,
)
_table_info.partitions_stat.append(_par)
return _table_info
return milvus_pb2.TableInfo(status=status_pb2.Status(
error_code=_status.code, reason=_status.message), )
def DescribeCollection(self, request, context):
_status, _collection_name = Parser.parse_proto_CollectionName(request)
if not _status.OK():
return milvus_pb2.CollectionSchema(status=status_pb2.Status(
error_code=_status.code, reason=_status.message), )
metadata = {'resp_class': milvus_pb2.CollectionSchema}
logger.info('DescribeCollection {}'.format(_collection_name))
_status, _collection = self._describe_collection(metadata=metadata,
collection_name=_collection_name)
if _status.OK():
return milvus_pb2.CollectionSchema(
collection_name=_collection_name,
index_file_size=_collection.index_file_size,
dimension=_collection.dimension,
metric_type=_collection.metric_type,
status=status_pb2.Status(error_code=_status.code,
reason=_status.message),
)
return milvus_pb2.CollectionSchema(
collection_name=_collection_name,
status=status_pb2.Status(error_code=_status.code,
reason=_status.message),
)
def GetVectorsByID(self, request, context):
_status, unpacks = Parser.parse_proto_VectorIdentity(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
metadata = {'resp_class': milvus_pb2.VectorsData}
_collection_name, _ids = unpacks
logger.info('GetVectorByID {}'.format(_collection_name))
_status, vectors = self._get_vectors_by_id(_collection_name, _ids,
metadata)
_rpc_status = status_pb2.Status(error_code=_status.code,
reason=_status.message)
if not vectors:
return milvus_pb2.VectorsData(status=_rpc_status, )
if len(vectors) == 0:
return milvus_pb2.VectorsData(status=_rpc_status, vectors_data=[])
if isinstance(vectors[0], bytes):
records = [milvus_pb2.RowRecord(binary_data=v) for v in vectors]
else:
records = [milvus_pb2.RowRecord(float_data=v) for v in vectors]
response = milvus_pb2.VectorsData(status=_rpc_status)
response.vectors_data.extend(records)
return response
def Cmd(self, request, context):
_status, _cmd = Parser.parse_proto_Command(request)
logger.info('Cmd: {}'.format(_cmd))
if not _status.OK():
return milvus_pb2.StringReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message))
metadata = {'resp_class': milvus_pb2.StringReply}
if _cmd == 'conn_stats':
stats = self.router.readonly_topo.stats()
return milvus_pb2.StringReply(status=status_pb2.Status(
error_code=status_pb2.SUCCESS),
string_reply=json.dumps(stats, indent=2))
# if _cmd == 'version':
# _status, _reply = self._get_server_version(metadata=metadata)
# else:
# _status, _reply = self.router.connection(
# metadata=metadata).server_status()
_status, _reply = self._cmd(_cmd, metadata=metadata)
return milvus_pb2.StringReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
string_reply=_reply)
def DescribeIndex(self, request, context):
_status, _collection_name = Parser.parse_proto_CollectionName(request)
if not _status.OK():
return milvus_pb2.IndexParam(status=status_pb2.Status(
error_code=_status.code, reason=_status.message))
metadata = {'resp_class': milvus_pb2.IndexParam}
logger.info('DescribeIndex {}'.format(_collection_name))
_status, _index_param = self._describe_index(collection_name=_collection_name,
metadata=metadata)
if not _index_param:
return milvus_pb2.IndexParam(status=status_pb2.Status(
error_code=_status.code, reason=_status.message))
_index_type = _index_param._index_type
grpc_index = milvus_pb2.IndexParam(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
collection_name=_collection_name, index_type=_index_type)
grpc_index.extra_params.add(key='params', value=ujson.dumps(_index_param._params))
return grpc_index
def GetVectorByID(self, request, context):
_status, unpacks = Parser.parse_proto_VectorIdentity(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
metadata = {'resp_class': milvus_pb2.VectorData}
_collection_name, _id = unpacks
logger.info('GetVectorByID {}'.format(_collection_name))
_status, vector = self._get_vector_by_id(_collection_name, _id, metadata)
if not vector:
return milvus_pb2.VectorData(status=status_pb2.Status(
error_code=_status.code, reason=_status.message), )
if isinstance(vector, bytes):
records = milvus_pb2.RowRecord(binary_data=vector)
else:
records = milvus_pb2.RowRecord(float_data=vector)
return milvus_pb2.VectorData(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
vector_data=records
)
def DescribeTable(self, request, context):
_status, _table_name = Parser.parse_proto_TableName(request)
if not _status.OK():
return milvus_pb2.TableSchema(status=status_pb2.Status(
error_code=_status.code, reason=_status.message), )
metadata = {'resp_class': milvus_pb2.TableSchema}
logger.info('DescribeTable {}'.format(_table_name))
_status, _table = self._describe_table(metadata=metadata,
table_name=_table_name)
if _status.OK():
return milvus_pb2.TableSchema(
table_name=_table_name,
index_file_size=_table.index_file_size,
dimension=_table.dimension,
metric_type=_table.metric_type,
status=status_pb2.Status(error_code=_status.code,
reason=_status.message),
)
return milvus_pb2.TableSchema(
table_name=_table_name,
status=status_pb2.Status(error_code=_status.code,
reason=_status.message),
)
def DropIndex(self, request, context):
_status, _table_name = Parser.parse_proto_TableName(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
logger.info('DropIndex {}'.format(_table_name))
_status = self._drop_index(_table_name)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def Compact(self, request, context):
_status, _collection_name = Parser.parse_proto_CollectionName(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
logger.info('Compact {}'.format(_collection_name))
_status = self._compact(_collection_name)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def PreloadTable(self, request, context):
_status, _table_name = Parser.parse_proto_TableName(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
logger.info('PreloadTable {}'.format(_table_name))
_status = self._preload_table(_table_name)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def Flush(self, request, context):
_status, _collection_names = Parser.parse_proto_FlushParam(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
logger.info('Flush {}'.format(_collection_names))
_status = self._flush(_collection_names)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def PreloadCollection(self, request, context):
_status, _pack = Parser.parse_proto_PreloadCollectionParam(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
_collection_name, _partition_tags = _pack
logger.info('PreloadCollection {} | {}'.format(_collection_name, _partition_tags))
_status = self._preload_collection(_collection_name, _partition_tags)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def CreateTable(self, request, context):
_status, _table_schema = Parser.parse_proto_TableSchema(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
logger.info('CreateTable {}'.format(_table_schema['table_name']))
_status = self._create_table(_table_schema)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def ShowPartitions(self, request, context):
_status, _collection_name = Parser.parse_proto_CollectionName(request)
if not _status.OK():
return milvus_pb2.PartitionList(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
partition_array=[])
logger.info('ShowPartitions {}'.format(_collection_name))
_status, partition_array = self.router.connection().show_partitions(_collection_name)
return milvus_pb2.PartitionList(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
partition_tag_array=[param.tag for param in partition_array])
def DeleteByID(self, request, context):
_status, unpacks = Parser.parse_proto_DeleteByIDParam(request)
if not _status.OK():
logging.error('DeleteByID {}'.format(_status.message))
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
_collection_name, _ids = unpacks
logger.info('DeleteByID {}'.format(_collection_name))
_status = self._delete_by_id(_collection_name, _ids)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def DeleteByRange(self, request, context):
_status, unpacks = \
Parser.parse_proto_DeleteByRangeParam(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
_table_name, _start_date, _end_date = unpacks
logger.info('DeleteByRange {}: {} {}'.format(_table_name, _start_date,
_end_date))
_status = self._delete_by_range(_table_name, _start_date, _end_date)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def HasCollection(self, request, context):
_status, _collection_name = Parser.parse_proto_CollectionName(request)
if not _status.OK():
return milvus_pb2.BoolReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
bool_reply=False)
logger.info('HasCollection {}'.format(_collection_name))
_status, _bool = self._has_collection(_collection_name,
metadata={'resp_class': milvus_pb2.BoolReply})
return milvus_pb2.BoolReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
bool_reply=_bool)
def _do_merge(self, files_n_topk_results, topk, reverse=False, **kwargs):
status = status_pb2.Status(error_code=status_pb2.SUCCESS,
reason="Success")
if not files_n_topk_results:
return status, []
request_results = defaultdict(list)
calc_time = time.time()
for files_collection in files_n_topk_results:
if isinstance(files_collection, tuple):
status, _ = files_collection
return status, []
for request_pos, each_request_results in enumerate(
files_collection.topk_query_result):
request_results[request_pos].extend(
each_request_results.query_result_arrays)
request_results[request_pos] = sorted(
request_results[request_pos],
key=lambda x: x.distance,
reverse=reverse)[:topk]
calc_time = time.time() - calc_time
logger.info('Merge takes {}'.format(calc_time))
results = sorted(request_results.items())
topk_query_result = []
for result in results:
query_result = TopKQueryResult(query_result_arrays=result[1])
topk_query_result.append(query_result)
return status, topk_query_result
def CreateIndex(self, request, context):
_status, unpacks = Parser.parse_proto_IndexParam(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
_collection_name, _index_type, _index_param = unpacks
logger.info('CreateIndex {}'.format(_collection_name))
# TODO: interface create_collection incompleted
_status = self._create_index(_collection_name, _index_type, _index_param)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def test_span_deco(self):
request = 'request'
OK = status_pb2.Status(error_code=status_pb2.SUCCESS, reason='Success')
response = OK
rpc_info = FakeRpcInfo(request=request, response=response)
span = FakeSpan(context=None, tracer=FakeTracer())
span_deco = GrpcSpanDecorator()
span_deco(span, rpc_info)
assert len(span.logs) == 0
assert len(span.tags) == 0
response = milvus_pb2.BoolReply(status=OK, bool_reply=False)
rpc_info = FakeRpcInfo(request=request, response=response)
span = FakeSpan(context=None, tracer=FakeTracer())
span_deco = GrpcSpanDecorator()
span_deco(span, rpc_info)
assert len(span.logs) == 0
assert len(span.tags) == 0
response = 1
rpc_info = FakeRpcInfo(request=request, response=response)
span = FakeSpan(context=None, tracer=FakeTracer())
span_deco = GrpcSpanDecorator()
span_deco(span, rpc_info)
assert len(span.logs) == 1
assert len(span.tags) == 1
response = 0
rpc_info = FakeRpcInfo(request=request, response=response)
span = FakeSpan(context=None, tracer=FakeTracer())
span_deco = GrpcSpanDecorator()
span_deco(span, rpc_info)
assert len(span.logs) == 0
assert len(span.tags) == 0
def DropPartition(self, request, context):
_collection_name, _tag = Parser.parse_proto_PartitionParam(request)
_status = self.router.connection().drop_partition(
_collection_name, _tag)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def HasPartition(self, request, context):
_collection_name, _tag = Parser.parse_proto_PartitionParam(request)
_status, _ok = self.router.connection().has_partition(
_collection_name, _tag)
return milvus_pb2.BoolReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
bool_reply=_ok)
def _do_merge(self, files_n_topk_results, topk, reverse=False, **kwargs):
status = status_pb2.Status(error_code=status_pb2.SUCCESS,
reason="Success")
if not files_n_topk_results:
return status, [], []
merge_id_results = []
merge_dis_results = []
calc_time = time.time()
for files_collection in files_n_topk_results:
if isinstance(files_collection, tuple):
status, _ = files_collection
return status, [], []
if files_collection.status.error_code != 0:
return files_collection.status, [], []
row_num = files_collection.row_num
# row_num is equal to 0, result is empty
if not row_num:
continue
ids = files_collection.ids
diss = files_collection.distances # distance collections
# TODO: batch_len is equal to topk, may need to compare with topk
batch_len = len(ids) // row_num
for row_index in range(row_num):
id_batch = ids[row_index * batch_len:(row_index + 1) *
batch_len]
dis_batch = diss[row_index * batch_len:(row_index + 1) *
batch_len]
if len(merge_id_results) < row_index:
raise ValueError("merge error")
elif len(merge_id_results) == row_index:
# TODO: may bug here
merge_id_results.append(id_batch)
merge_dis_results.append(dis_batch)
else:
merge_id_results[row_index], merge_dis_results[row_index] = \
self._reduce(merge_id_results[row_index], id_batch,
merge_dis_results[row_index], dis_batch,
batch_len,
reverse)
calc_time = time.time() - calc_time
logger.info('Merge takes {}'.format(calc_time))
id_mrege_list = []
dis_mrege_list = []
for id_results, dis_results in zip(merge_id_results,
merge_dis_results):
id_mrege_list.extend(id_results)
dis_mrege_list.extend(dis_results)
return status, id_mrege_list, dis_mrege_list
def ShowCollections(self, request, context):
logger.info('ShowCollections')
metadata = {'resp_class': milvus_pb2.CollectionName}
_status, _results = self._show_collections(metadata=metadata)
return milvus_pb2.CollectionNameList(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
collection_names=_results)
def Insert(self, request, context):
logger.info('Insert')
# TODO: Ths SDK interface add_vectors() could update, add a key 'row_id_array'
_status, _ids = self._add_vectors(
metadata={'resp_class': milvus_pb2.VectorIds}, param=request)
return milvus_pb2.VectorIds(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
vector_id_array=_ids)
def resp_handler(err, error_code):
if not isinstance(err, exceptions.BaseException):
return status_pb2.Status(error_code=error_code, reason=str(err))
status = status_pb2.Status(error_code=error_code, reason=err.message)
if err.metadata is None:
return status
resp_class = err.metadata.get('resp_class', None)
if not resp_class:
return status
if resp_class == milvus_pb2.BoolReply:
return resp_class(status=status, bool_reply=False)
if resp_class == milvus_pb2.VectorIds:
return resp_class(status=status, vector_id_array=[])
if resp_class == milvus_pb2.TopKQueryResultList:
return resp_class(status=status, topk_query_result=[])
if resp_class == milvus_pb2.TableRowCount:
return resp_class(status=status, table_row_count=-1)
if resp_class == milvus_pb2.TableName:
return resp_class(status=status, table_name=[])
if resp_class == milvus_pb2.StringReply:
return resp_class(status=status, string_reply='')
if resp_class == milvus_pb2.TableSchema:
return milvus_pb2.TableSchema(
status=status
)
if resp_class == milvus_pb2.IndexParam:
return milvus_pb2.IndexParam(
table_name=milvus_pb2.TableName(
status=status
)
)
status.error_code = status_pb2.UNEXPECTED_ERROR
return status
def CreateCollection(self, request, context):
_status, unpacks = Parser.parse_proto_CollectionSchema(request)
if not _status.OK():
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
_status, _collection_schema = unpacks
# if _status.error_code != 0:
# logging.warning('[CreateCollection] collection schema error occurred: {}'.format(_status))
# return _status
logger.info('CreateCollection {}'.format(_collection_schema['collection_name']))
_status = self._create_collection(_collection_schema)
return status_pb2.Status(error_code=_status.code,
reason=_status.message)
def CountCollection(self, request, context):
_status, _collection_name = Parser.parse_proto_CollectionName(request)
if not _status.OK():
status = status_pb2.Status(error_code=_status.code,
reason=_status.message)
return milvus_pb2.CollectionRowCount(status=status)
logger.info('CountCollection {}'.format(_collection_name))
metadata = {'resp_class': milvus_pb2.CollectionRowCount}
_status, _count = self._count_collection(_collection_name, metadata=metadata)
return milvus_pb2.CollectionRowCount(
status=status_pb2.Status(error_code=_status.code,
reason=_status.message),
collection_row_count=_count if isinstance(_count, int) else -1)
def test_search(self, started_app):
collection_name = inspect.currentframe().f_code.co_name
to_index_cnt = random.randint(10, 20)
collection = TablesFactory(collection_id=collection_name,
state=Tables.NORMAL)
to_index_files = TableFilesFactory.create_batch(
to_index_cnt,
collection=collection,
file_type=TableFiles.FILE_TYPE_TO_INDEX)
topk = random.randint(5, 10)
nq = random.randint(5, 10)
param = {
'collection_name': collection_name,
'query_records': self.random_data(nq, collection.dimension),
'top_k': topk,
'params': {
'nprobe': 2049
}
}
result = [
milvus_pb2.TopKQueryResult(query_result_arrays=[
milvus_pb2.QueryResult(id=i, distance=random.random())
for i in range(topk)
]) for i in range(nq)
]
mock_results = milvus_pb2.TopKQueryResultList(status=status_pb2.Status(
error_code=status_pb2.SUCCESS, reason="Success"),
topk_query_result=result)
collection_schema = CollectionSchema(
collection_name=collection_name,
index_file_size=collection.index_file_size,
metric_type=collection.metric_type,
dimension=collection.dimension)
status, _ = self.client.search_vectors(**param)
assert status.code == Status.ILLEGAL_ARGUMENT
param['params']['nprobe'] = 2048
RouterMixin.connection = mock.MagicMock(return_value=Milvus())
RouterMixin.query_conn.conn = mock.MagicMock(return_value=Milvus())
Milvus.describe_collection = mock.MagicMock(
return_value=(BAD, collection_schema))
status, ret = self.client.search_vectors(**param)
assert status.code == Status.COLLECTION_NOT_EXISTS
Milvus.describe_collection = mock.MagicMock(
return_value=(OK, collection_schema))
Milvus.search_vectors_in_files = mock.MagicMock(
return_value=mock_results)
status, ret = self.client.search_vectors(**param)
assert status.OK()
assert len(ret) == nq
def Cmd(self, request, context):
_status, _cmd = Parser.parse_proto_Command(request)
logger.info('Cmd: {}'.format(_cmd))
if not _status.OK():
return milvus_pb2.StringReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message))
metadata = {'resp_class': milvus_pb2.StringReply}
if _cmd == 'version':
_status, _reply = self._get_server_version(metadata=metadata)
else:
_status, _reply = self.router.connection(
metadata=metadata).server_status()
return milvus_pb2.StringReply(status=status_pb2.Status(
error_code=_status.code, reason=_status.message),
string_reply=_reply)
