Unified query executor that routes based on query count and reranker type.
| 103 | |
| 104 | |
| 105 | class QueryExecutor: |
| 106 | """Unified query executor that routes based on query count and reranker type.""" |
| 107 | |
| 108 | def __init__(self, schema: CollectionSchema): |
| 109 | self._schema = schema |
| 110 | |
| 111 | def _build_queries( |
| 112 | self, ctx: QueryContext, collection: _Collection |
| 113 | ) -> list[_SearchQuery]: |
| 114 | """Build query vector list (no validation, conversion only).""" |
| 115 | if not ctx.queries: |
| 116 | return [self._build_base_search_query(ctx)] |
| 117 | return [ |
| 118 | self._build_search_query(ctx, query, collection) for query in ctx.queries |
| 119 | ] |
| 120 | |
| 121 | def execute(self, ctx: QueryContext, collection: _Collection) -> DocList: |
| 122 | """Execute a query, routing by query count. |
| 123 | |
| 124 | A single (or vector-less) query is sent to C++ as a ``_SearchQuery``; |
| 125 | multiple queries are assembled into a ``_MultiQuery``. |
| 126 | """ |
| 127 | queries = self._build_queries(ctx, collection) |
| 128 | if not queries: |
| 129 | raise ValueError("No query to execute") |
| 130 | |
| 131 | if len(queries) == 1: |
| 132 | return self._execute_single_query(queries[0], collection) |
| 133 | return self._execute_multi_query(ctx, queries, collection) |
| 134 | |
| 135 | def _execute_single_query( |
| 136 | self, query: _SearchQuery, collection: _Collection |
| 137 | ) -> DocList: |
| 138 | """Single/vector-less query: send a ``_SearchQuery`` to C++.""" |
| 139 | docs = collection.Query(query) |
| 140 | return [convert_to_py_doc(doc, self._schema) for doc in docs] |
| 141 | |
| 142 | def _execute_multi_query( |
| 143 | self, ctx: QueryContext, queries: list[_SearchQuery], collection: _Collection |
| 144 | ) -> DocList: |
| 145 | """Multiple queries: send a ``_MultiQuery`` to C++. |
| 146 | |
| 147 | A Python-only reranker (e.g. a model/API-based one) cannot run inside |
| 148 | the C++ MultiQuery, so each route is executed individually and merged by |
| 149 | the reranker in Python. The built-in RRF/Weighted/Callback rerankers use |
| 150 | the C++ variant-based fast path. |
| 151 | """ |
| 152 | reranker = ctx.reranker |
| 153 | if reranker is None: |
| 154 | raise ValueError( |
| 155 | "A reranker is required to merge results from multiple queries; " |
| 156 | "specify the 'reranker' argument." |
| 157 | ) |
| 158 | if not isinstance(reranker, (RrfReRanker, WeightedReRanker, CallbackReRanker)): |
| 159 | docs_list = self._execute_python_pipeline(queries, collection) |
| 160 | return self._merge_and_rerank(ctx, docs_list) |
| 161 | |
| 162 | multi_query = self._build_multi_query(ctx, queries) |
no outgoing calls