def SearchByID(self, request, context):
metadata = {'resp_class': milvus_pb2.TopKQueryResult}
collection_name = request.collection_name
topk = request.topk
if len(request.extra_params) == 0:
raise exceptions.SearchParamError(message="Search param loss",
metadata=metadata)
params = ujson.loads(str(request.extra_params[0].value))
logger.info('Search {}: topk={} params={}'.format(
collection_name, topk, params))
if topk > self.MAX_TOPK or topk <= 0:
raise exceptions.InvalidTopKError(
message='Invalid topk: {}'.format(topk), metadata=metadata)
collection_meta = self.collection_meta.get(collection_name, None)
if not collection_meta:
status, info = self.router.connection(
metadata=metadata).describe_collection(collection_name)
if not status.OK():
raise exceptions.CollectionNotFoundError(collection_name,
metadata=metadata)
self.collection_meta[collection_name] = info
collection_meta = info
start = time.time()
query_record_array = []
if int(collection_meta.metric_type) >= MetricType.HAMMING.value:
for query_record in request.query_record_array:
query_record_array.append(bytes(query_record.binary_data))
else:
for query_record in request.query_record_array:
query_record_array.append(list(query_record.float_data))
partition_tags = getattr(request, "partition_tag_array", [])
ids = getattr(request, "id_array", [])
search_result = self.router.connection(
metadata=metadata).search_by_ids(collection_name, ids, topk,
partition_tags, params)
# status, id_results, dis_results = self._do_query(context,
# collection_name,
# collection_meta,
# query_record_array,
# topk,
# params,
# partition_tags=getattr(request, "partition_tag_array", []),
# metadata=metadata)
now = time.time()
logger.info('SearchVector takes: {}'.format(now - start))
return search_result
def Search(self, request, context):
metadata = {'resp_class': milvus_pb2.TopKQueryResult}
collection_name = request.collection_name
topk = request.topk
if len(request.extra_params) == 0:
raise exceptions.SearchParamError(message="Search parma loss", metadata=metadata)
params = ujson.loads(str(request.extra_params[0].value))
logger.info('Search {}: topk={} params={}'.format(
collection_name, topk, params))
# if nprobe > self.MAX_NPROBE or nprobe <= 0:
# raise exceptions.InvalidArgumentError(
# message='Invalid nprobe: {}'.format(nprobe), metadata=metadata)
if topk > self.MAX_TOPK or topk <= 0:
raise exceptions.InvalidTopKError(
message='Invalid topk: {}'.format(topk), metadata=metadata)
collection_meta = self.collection_meta.get(collection_name, None)
if not collection_meta:
status, info = self.router.connection(
metadata=metadata).describe_collection(collection_name)
if not status.OK():
raise exceptions.CollectionNotFoundError(collection_name,
metadata=metadata)
self.collection_meta[collection_name] = info
collection_meta = info
start = time.time()
query_record_array = []
if int(collection_meta.metric_type) >= MetricType.HAMMING.value:
for query_record in request.query_record_array:
query_record_array.append(bytes(query_record.binary_data))
else:
for query_record in request.query_record_array:
query_record_array.append(list(query_record.float_data))
status, id_results, dis_results = self._do_query(context,
collection_name,
collection_meta,
query_record_array,
topk,
params,
partition_tags=getattr(request, "partition_tag_array", []),
metadata=metadata)
now = time.time()
logger.info('SearchVector takes: {}'.format(now - start))
topk_result_list = milvus_pb2.TopKQueryResult(
status=status_pb2.Status(error_code=status.error_code,
reason=status.reason),
row_num=len(request.query_record_array) if len(id_results) else 0,
ids=id_results,
distances=dis_results)
return topk_result_list