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)