Skip to content

Commit

Permalink
fix for new-version of iterator(#1660) (#1658)
Browse files Browse the repository at this point in the history
Signed-off-by: MrPresent-Han <[email protected]>
  • Loading branch information
MrPresent-Han authored Aug 23, 2023
1 parent 58d42bb commit 95f7790
Showing 1 changed file with 20 additions and 19 deletions.
39 changes: 20 additions & 19 deletions pymilvus/orm/iterator.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,11 @@
from .constants import (
BATCH_SIZE,
CALC_DIST_COSINE,
CALC_DIST_HAMMING,
CALC_DIST_IP,
CALC_DIST_JACCARD,
CALC_DIST_L2,
CALC_DIST_TANIMOTO,
DEFAULT_SEARCH_EXTENSION_RATE,
FIELDS,
INT64_MAX,
Expand Down Expand Up @@ -192,31 +195,13 @@ def close(self) -> None:


def metrics_positive_related(metrics: str) -> bool:
if metrics is CALC_DIST_L2:
if metrics in [CALC_DIST_L2, CALC_DIST_JACCARD, CALC_DIST_HAMMING, CALC_DIST_TANIMOTO]:
return True
if metrics is CALC_DIST_IP or metrics is CALC_DIST_COSINE:
return False
raise MilvusException(message=f"unsupported metrics type for search iteration{metrics}")


class SearchHit:
def __init__(self, id: Any, distance: Any):
self._id = id
self._distance = distance


"""
class SearchHits:
def __init__(self, ids: list, distances: list):
self._ids = ids
self._distances = distances
def __next__(self):
return Hit()
"""


class SearchPage(LoopBase):
"""Since we only support nq=1 in search iteration, so search iteration response
should be different from raw response of search operation"""
Expand Down Expand Up @@ -257,6 +242,20 @@ def merge(self, others: List[Hits]):
for other in others:
self._results.append(other)

def ids(self):
ids = []
for res in self._results:
for hit in res:
ids.append(hit.id)
return ids

def distances(self):
distances = []
for res in self._results:
for hit in res:
distances.append(hit.distance)
return distances


class SearchIterator:
def __init__(
Expand Down Expand Up @@ -335,6 +334,8 @@ def __check_set_params(self, param: Dict):
self._param = {}
else:
self._param = deepcopy(param)
if PARAMS not in self._param:
self._param[PARAMS] = {}

def __setup__pk_prop(self):
fields = self._schema[FIELDS]
Expand Down

0 comments on commit 95f7790

Please sign in to comment.