enhance search iterator extension

milvus-io · Nov 24, 2023 · 51ebd74 · 51ebd74
1 parent 9075512
commit 51ebd74
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/pymilvus/orm/constants.py b/pymilvus/orm/constants.py
@@ -49,4 +49,4 @@
 MAX_BATCH_SIZE: int = 16384
 DEFAULT_SEARCH_EXTENSION_RATE: int = 10
 UNLIMITED: int = -1
-MAX_TRY_TIME: int = 10
+MAX_TRY_TIME: int = 20
diff --git a/pymilvus/orm/iterator.py b/pymilvus/orm/iterator.py
@@ -336,10 +336,16 @@ def __set_up_range_parameters(self, page: SearchPage):
         else:
             self._width = first_hit.distance - last_hit.distance
         self._tail_band = last_hit.distance
+        LOGGER.info(
+            f"set up init parameter for searchIterator width:{self._width} tail_band:{self._tail_band}"
+        )
 
     def __check_reached_limit(self) -> bool:
         if self._limit == UNLIMITED or self._returned_count < self._limit:
             return False
+        LOGGER.debug(
+            f"reached search limit:{self._limit}, returned_count:{self._returned_count}, directly return"
+        )
         return True
 
     def __check_set_params(self, param: Dict):
@@ -491,14 +497,15 @@ def __try_search_fill(self) -> SearchPage:
             if len(new_page) > 0:
                 final_page.merge(new_page.get_res())
                 self._tail_band = new_page[-1].distance
-                # if the current ring contains vectors, we always set coefficient back to 1
-                coefficient = 1
-            else:
-                # if there's a ring containing no vectors matched, then we need to extend
-                # the ring continually to avoid empty ring problem
-                coefficient += 1
-            if len(final_page) > self._iterator_params[BATCH_SIZE] or try_time > MAX_TRY_TIME:
+            if len(final_page) > self._iterator_params[BATCH_SIZE]:
+                break
+            if try_time > MAX_TRY_TIME:
+                LOGGER.warning(f"Search probe exceed max try times:{MAX_TRY_TIME} directly break")
                 break
+            # if there's a ring containing no vectors matched, then we need to extend
+            # the ring continually to avoid empty ring problem
+            coefficient += 1
+
         return final_page
 
     def __execute_next_search(self, next_params: dict, next_expr: str) -> SearchPage:
@@ -542,7 +549,7 @@ def __filtered_duplicated_result_expr(self, expr: str):
 
     def __next_params(self, coefficient: int):
         coefficient = max(1, coefficient)
-        next_params = self._param.copy()
+        next_params = deepcopy(self._param)
         if metrics_positive_related(self._param[METRIC_TYPE]):
             next_radius = self._tail_band + self._width * coefficient
             if RADIUS in self._param[PARAMS] and next_radius > self._param[PARAMS][RADIUS]:
@@ -556,6 +563,12 @@ def __next_params(self, coefficient: int):
             else:
                 next_params[PARAMS][RADIUS] = next_radius
         next_params[PARAMS][RANGE_FILTER] = self._tail_band
+        LOGGER.debug(
+            f"next round search iteration radius:{next_params[PARAMS][RADIUS]},"
+            f"range_filter:{next_params[PARAMS][RANGE_FILTER]},"
+            f"coefficient:{coefficient}"
+        )
+
         return next_params
 
     def close(self):