Support rerank

Signed-off-by: junjiejiangjjj <junjie.jiang@zilliz.com>

Author: junjiejiangjjj

examples/simple_rerank.py

@@ -0,0 +1,92 @@
+import time
+import numpy as np
+from pymilvus import (
+    MilvusClient,
+    DataType,
+    Function,
+    FunctionType,
+    AnnSearchRequest,
+)
+
+fmt = "\n=== {:30} ===\n"
+dim = 8
+collection_name = "hello_milvus"
+milvus_client = MilvusClient("http://localhost:19530")
+
+has_collection = milvus_client.has_collection(collection_name, timeout=5)
+if has_collection:
+    milvus_client.drop_collection(collection_name)
+
+schema = milvus_client.create_schema(enable_dynamic_field=False, auto_id=True)
+schema.add_field("id", DataType.INT64, is_primary=True)
+schema.add_field("embeddings", DataType.FLOAT_VECTOR, dim=dim)
+schema.add_field("ts", DataType.INT64)
+
+
+index_params = milvus_client.prepare_index_params()
+index_params.add_index(field_name = "embeddings", index_type="FLAT", metric_type="L2")
+milvus_client.create_collection(collection_name, schema=schema, index_params=index_params, consistency_level="Strong")
+
+print(fmt.format("    all collections    "))
+print(milvus_client.list_collections())
+
+print(fmt.format(f"schema of collection {collection_name}"))
+print(milvus_client.describe_collection(collection_name))
+
+rng = np.random.default_rng(seed=19530)
+rows = [
+        {"embeddings": rng.random((1, dim))[0], "ts": 100},
+        {"embeddings": rng.random((1, dim))[0], "ts": 200},
+        {"embeddings": rng.random((1, dim))[0], "ts": 300},
+        {"embeddings": rng.random((1, dim))[0], "ts": 400},
+        {"embeddings": rng.random((1, dim))[0], "ts": 500},
+        {"embeddings": rng.random((1, dim))[0], "ts": 600},
+]
+
+print(fmt.format("Start inserting entities"))
+insert_result = milvus_client.insert(collection_name, rows)
+print(fmt.format("Inserting entities done"))
+print(insert_result)
+
+
+print(fmt.format("Start load collection "))
+milvus_client.load_collection(collection_name)
+
+rng = np.random.default_rng(seed=19530)
+vectors_to_search = rng.random((1, dim))
+
+ranker = Function(
+    name="rerank_fn",
+    input_field_names=["ts"],
+    function_type=FunctionType.RERANK,
+    params={
+        "reranker": "decay",
+        "function": "exp",
+        "origin": 0,
+        "offset": 200,
+        "decay": 0.9,
+        "scale": 100
+    }
+)
+
+print(fmt.format(f"Start search with retrieve serveral fields."))
+result = milvus_client.search(collection_name, vectors_to_search, limit=3, output_fields=["*"], ranker=ranker)
+for hits in result:
+    for hit in hits:
+        print(f"hit: {hit}")
+
+vectors_to_search = rng.random((1, dim))
+search_param = {
+    "data": vectors_to_search,
+    "anns_field": "embeddings",
+    "param": {"metric_type": "L2"},
+    "limit": 3,
+}
+req = AnnSearchRequest(**search_param)
+
+hybrid_res = milvus_client.hybrid_search(collection_name, [req, req], ranker=ranker, limit=3, output_fields=["ts"])
+for hits in hybrid_res:
+    for hit in hits:
+        print(f" hybrid search hit: {hit}")
+
+milvus_client.drop_collection(collection_name)

pymilvus/client/async_grpc_handler.py

@@ -572,6 +572,7 @@ async def search(
         output_fields: Optional[List[str]] = None,
         round_decimal: int = -1,
         timeout: Optional[float] = None,
+        ranker: Optional["Function"] = None,
         **kwargs,
     ):
         await self.ensure_channel_ready()
@@ -584,6 +585,7 @@ async def search(
             output_fields=output_fields,
             guarantee_timestamp=kwargs.get("guarantee_timestamp"),
             timeout=timeout,
+            ranker=ranker,
         )
         request = Prepare.search_requests_with_expr(
             collection_name,
@@ -595,6 +597,7 @@ async def search(
             partition_names,
             output_fields,
             round_decimal,
+            ranker=ranker,
             **kwargs,
         )
         return await self._execute_search(request, timeout, round_decimal=round_decimal, **kwargs)
@@ -610,6 +613,7 @@ async def hybrid_search(
         output_fields: Optional[List[str]] = None,
         round_decimal: int = -1,
         timeout: Optional[float] = None,
+        ranker: Optional["Function"] = None,
         **kwargs,
     ):
         await self.ensure_channel_ready()
@@ -620,6 +624,7 @@ async def hybrid_search(
             output_fields=output_fields,
             guarantee_timestamp=kwargs.get("guarantee_timestamp"),
             timeout=timeout,
+            ranker=ranker,
         )
 
         requests = []
@@ -645,6 +650,7 @@ async def hybrid_search(
             partition_names,
             output_fields,
             round_decimal,
+            ranker=ranker,
             **kwargs,
         )
         return await self._execute_hybrid_search(

pymilvus/client/check.py

@@ -4,6 +4,7 @@
 
 from pymilvus.exceptions import ParamError
 from pymilvus.grpc_gen import milvus_pb2 as milvus_types
+from pymilvus.orm.schema import Function
 
 from . import entity_helper
 from .singleton_utils import Singleton
@@ -316,6 +317,14 @@ def is_legal_operate_privilege_group_type(operate_privilege_group_type: Any) ->
     )
 
 
+def is_legal_ranker(ranker: Any) -> bool:
+    if ranker is None:
+        return True
+    if not isinstance(ranker, Function):
+        return False
+    return True
+
+
 class ParamChecker(metaclass=Singleton):
     def __init__(self) -> None:
         self.check_dict = {
@@ -363,6 +372,7 @@ def __init__(self) -> None:
             "privilege_group": is_legal_privilege_group,
             "privileges": is_legal_privileges,
             "operate_privilege_group_type": is_legal_operate_privilege_group_type,
+            "ranker": is_legal_ranker,
         }
 
     def check(self, key: str, value: Callable):

pymilvus/client/grpc_handler.py

@@ -888,6 +888,7 @@ def search(
         output_fields: Optional[List[str]] = None,
         round_decimal: int = -1,
         timeout: Optional[float] = None,
+        ranker: Optional["Function"] = None,
         **kwargs,
     ):
         check_pass_param(
@@ -899,6 +900,7 @@ def search(
             output_fields=output_fields,
             guarantee_timestamp=kwargs.get("guarantee_timestamp"),
             timeout=timeout,
+            ranker=ranker,
         )
 
         request = Prepare.search_requests_with_expr(
@@ -911,6 +913,7 @@ def search(
             partition_names,
             output_fields,
             round_decimal,
+            ranker=ranker,
             **kwargs,
         )
         return self._execute_search(request, timeout, round_decimal=round_decimal, **kwargs)
@@ -920,7 +923,7 @@ def hybrid_search(
         self,
         collection_name: str,
         reqs: List[AnnSearchRequest],
-        rerank: BaseRanker,
+        rerank: Union[BaseRanker, "Function"],
         limit: int,
         partition_names: Optional[List[str]] = None,
         output_fields: Optional[List[str]] = None,
@@ -956,7 +959,7 @@ def hybrid_search(
         hybrid_search_request = Prepare.hybrid_search_request_with_ranker(
             collection_name,
             requests,
-            rerank.dict(),
+            rerank,
             limit,
             partition_names,
             output_fields,

pymilvus/client/prepare.py

@@ -9,7 +9,7 @@
 from pymilvus.grpc_gen import common_pb2 as common_types
 from pymilvus.grpc_gen import milvus_pb2 as milvus_types
 from pymilvus.grpc_gen import schema_pb2 as schema_types
-from pymilvus.orm.schema import CollectionSchema, FieldSchema
+from pymilvus.orm.schema import CollectionSchema, FieldSchema, Function
 from pymilvus.orm.types import infer_dtype_by_scalar_data
 
 from . import __version__, blob, check, entity_helper, ts_utils, utils
@@ -938,6 +938,7 @@ def search_requests_with_expr(
         partition_names: Optional[List[str]] = None,
         output_fields: Optional[List[str]] = None,
         round_decimal: int = -1,
+        ranker: Optional[Function] = None,
         **kwargs,
     ) -> milvus_types.SearchRequest:
         use_default_consistency = ts_utils.construct_guarantee_ts(collection_name, kwargs)
@@ -1049,24 +1050,31 @@ def search_requests_with_expr(
         if expr is not None:
             request.dsl = expr
 
+        if isinstance(ranker, Function):
+            request.function_score.CopyFrom(Prepare.ranker_to_function_score(ranker))
+
         return request
 
     @classmethod
     def hybrid_search_request_with_ranker(
         cls,
         collection_name: str,
         reqs: List,
-        rerank_param: Dict,
+        rerank: Union[Dict, Function],
         limit: int,
         partition_names: Optional[List[str]] = None,
         output_fields: Optional[List[str]] = None,
         round_decimal: int = -1,
         **kwargs,
     ) -> milvus_types.HybridSearchRequest:
         use_default_consistency = ts_utils.construct_guarantee_ts(collection_name, kwargs)
-        rerank_param["limit"] = limit
-        rerank_param["round_decimal"] = round_decimal
-        rerank_param["offset"] = kwargs.get("offset", 0)
+        rerank_param = {}
+        if isinstance(rerank, Dict):
+            rerank_param = rerank
+        else:
+            rerank_param["limit"] = limit
+            rerank_param["round_decimal"] = round_decimal
+            rerank_param["offset"] = kwargs.get("offset", 0)
 
         request = milvus_types.HybridSearchRequest(
             collection_name=collection_name,
@@ -1121,8 +1129,27 @@ def hybrid_search_request_with_ranker(
                 ]
             )
 
+        if isinstance(rerank, Function):
+            request.function_score.CopyFrom(Prepare.ranker_to_function_score(rerank))
         return request
 
+    @staticmethod
+    def ranker_to_function_score(ranker: Function) -> schema_types.FunctionScore:
+        function_score = schema_types.FunctionScore(
+            functions=[
+                schema_types.FunctionSchema(
+                    name=ranker.name,
+                    type=ranker.type,
+                    description=ranker.description,
+                    input_field_names=ranker.input_field_names,
+                )
+            ],
+        )
+        for k, v in ranker.params.items():
+            kv_pair = common_types.KeyValuePair(key=str(k), value=str(v))
+            function_score.functions[0].params.append(kv_pair)
+        return function_score
+
     @classmethod
     def create_alias_request(cls, collection_name: str, alias: str):
         return milvus_types.CreateAliasRequest(collection_name=collection_name, alias=alias)

pymilvus/client/types.py

@@ -121,6 +121,7 @@ class FunctionType(IntEnum):
     UNKNOWN = 0
     BM25 = 1
     TEXTEMBEDDING = 2
+    RERANK = 3
 
 
 class RangeType(IntEnum):

pymilvus/milvus_client/milvus_client.py

@@ -24,7 +24,7 @@
     ServerVersionIncompatibleException,
 )
 from pymilvus.orm import utility
-from pymilvus.orm.collection import CollectionSchema, FieldSchema
+from pymilvus.orm.collection import CollectionSchema, FieldSchema, Function
 from pymilvus.orm.connections import connections
 from pymilvus.orm.constants import FIELDS, METRIC_TYPE, TYPE, UNLIMITED
 from pymilvus.orm.iterator import QueryIterator, SearchIterator
@@ -297,7 +297,7 @@ def hybrid_search(
         self,
         collection_name: str,
         reqs: List[AnnSearchRequest],
-        ranker: BaseRanker,
+        ranker: Union[BaseRanker, Function],
         limit: int = 10,
         output_fields: Optional[List[str]] = None,
         timeout: Optional[float] = None,
@@ -309,7 +309,7 @@ def hybrid_search(
         Args:
             collection_name(``string``): The name of collection.
             reqs (``List[AnnSearchRequest]``): The vector search requests.
-            ranker (``BaseRanker``): The ranker for rearrange nummer of limit results.
+            ranker (``Union[BaseRanker, Function]``): The ranker for rearrange nummer of limit results.
             limit (``int``): The max number of returned record, also known as `topk`.
 
             partition_names (``List[str]``, optional): The names of partitions to search on.
@@ -375,6 +375,7 @@ def search(
         timeout: Optional[float] = None,
         partition_names: Optional[List[str]] = None,
         anns_field: Optional[str] = None,
+        ranker: Optional["Function"] = None,
         **kwargs,
     ) -> List[List[dict]]:
         """Search for a query vector/vectors.
@@ -389,6 +390,7 @@ def search(
             output_fields (List[str], optional): List of which field values to return. If None
                 specified, only primary fields including distances will be returned.
             search_params (dict, optional): The search params to use for the search.
+            ranker (Function, optional): The ranker to use for the search.
             timeout (float, optional): Timeout to use, overides the client level assigned at init.
                 Defaults to None.
 
@@ -412,6 +414,7 @@ def search(
                 partition_names=partition_names,
                 expr_params=kwargs.pop("filter_params", {}),
                 timeout=timeout,
+                ranker=ranker,
                 **kwargs,
             )
         except Exception as ex: