rescore with prefetch option

Author: generall

engine/clients/qdrant/search.py

@@ -4,7 +4,7 @@
 import httpx
 from qdrant_client import QdrantClient
 from qdrant_client._pydantic_compat import construct
-from qdrant_client.http import models as rest
+from qdrant_client import models
 
 from dataset_reader.base_reader import Query
 from engine.base_client.search import BaseSearcher
@@ -37,24 +37,35 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic
 
     @classmethod
     def search_one(cls, query: Query, top: int) -> List[Tuple[int, float]]:
+
         # Can query only one till we introduce re-ranking in the benchmarks
         if query.sparse_vector is None:
             query_vector = query.vector
         else:
             query_vector = construct(
-                rest.SparseVector,
+                models.SparseVector,
                 indices=query.sparse_vector.indices,
                 values=query.sparse_vector.values,
             )
 
+
+        prefetch = cls.search_params.get("prefetch")
+
+        if prefetch:
+            prefetch = models.Prefetch(
+                **prefetch,
+                query=query_vector,
+            )
+
         try:
             res = cls.client.query_points(
                 using="sparse" if query.sparse_vector else None,
                 collection_name=QDRANT_COLLECTION_NAME,
+                prefetch=prefetch,
                 query=query_vector,
                 query_filter=cls.parser.parse(query.meta_conditions),
                 limit=top,
-                search_params=rest.SearchParams(**cls.search_params.get("config", {})),
+                search_params=models.SearchParams(**cls.search_params.get("config", {})),
                 with_payload=cls.search_params.get("with_payload", False),
             )
         except Exception as ex:

experiments/configurations/qdrant-single-node.json

@@ -55,10 +55,18 @@
       {
         "parallel": 8,
         "config": {
-          "hnsw_ef": 128,
           "quantization": {
             "rescore": true
           }
+        },
+        "prefetch": {
+          "limit": 400,
+          "params": {
+            "hnsw_ef": 128,
+            "quantization": {
+              "rescore": false
+            }
+          }
         }
       }
     ],