Add Weighted RRF to Azure Cosmos DB Python SDK (#40899)

* Add Weighted Reciprocal Rank Fusion to Python Cosmos DB

Adds weighted reciprocal rank fusion to the python sdk.

* Update hybrid_search_aggregator.py

* Update CHANGELOG.md

* fix pylint

* Update README.md

* added query optimizations as well

* add valid vector search test for wrrf

* update quick fixes

* pylint fix

* Update CHANGELOG.md

* update changes

* review requested changes

* Update CHANGELOG.md

* updates should fix deprecration of list in full text score

* Update test_query_hybrid_search.py

* Update CHANGELOG.md

* Update tests and Readme

* Update _version.py

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,15 +1,16 @@
 ## Release History
 
-### 4.11.0b2 (Unreleased)
+### 4.12.0b1 (unreleased)
 
 #### Features Added
-
-#### Breaking Changes
+* Added ability to use weighted RRF (Reciprocal Rank Fusion) for Hybrid full text search queries. See [PR 40899](https://github.com/Azure/azure-sdk-for-python/pull/40899/files).
 
 #### Bugs Fixed
 * Fixed Diagnostics Error Log Formatting to handle error messages from non-CosmosHttpResponseExceptions. See [PR 40889](https://github.com/Azure/azure-sdk-for-python/pull/40889/files)
 * Fixed bug where `multiple_write_locations` option in client was not being honored. See [PR 40999](https://github.com/Azure/azure-sdk-for-python/pull/40999).
 
+#### Breaking Changes
+
 #### Other Changes
 
 ### 4.11.0b1 (2025-04-30)

sdk/cosmos/azure-cosmos/README.md

@@ -867,6 +867,13 @@ All of these mentioned queries would look something like this:
 
 - `SELECT TOP 10 c.id, c.text FROM c ORDER BY RANK RRF(FullTextScore(c.text, ['quantum', 'theory']), FullTextScore(c.text, ['model']), VectorDistance(c.embedding, {item_embedding}))"`
 
+You can also use Weighted Reciprocal Rank Fusion to assign different weights to the different scores being used in the RRF function.
+This is done by passing in a list of weights to the RRF function in the query. **NOTE: If more weights are given than there are components of the RRF function, or if weights are missing a BAD REQUEST exception will occur.**
+- `SELECT TOP 10 c.id, c.text FROM c ORDER BY RANK RRF(FullTextScore(c.text, ['quantum', 'theory']), FullTextScore(c.text, ['model']), VectorDistance(c.embedding, {item_embedding}), [0.5, 0.3, 0.2])`
+
+
+- `SELECT TOP 10 c.id, c.text FROM c ORDER BY RANK RRF(FullTextScore(c.text, ['quantum', 'theory']), FullTextScore(c.text, ['model']), VectorDistance(c.embedding, {item_embedding}), [-0.5, 0.3, 0.2])`
+
 These queries must always use a TOP or LIMIT clause within the query since hybrid search queries have to look through a lot of data otherwise and may become too expensive or long-running.
 Since these queries are relatively expensive, the SDK sets a default limit of 1000 max items per query - if you'd like to raise that further, you
 can use the `AZURE_COSMOS_HYBRID_SEARCH_MAX_ITEMS` environment variable to do so. However, be advised that queries with too many vector results

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -3196,7 +3196,8 @@ def _GetQueryPlanThroughGateway(self, query: str, resource_link: str, **kwargs:
                                     documents._QueryFeature.Top + "," +
                                     documents._QueryFeature.NonStreamingOrderBy + "," +
                                     documents._QueryFeature.HybridSearch + "," +
-                                    documents._QueryFeature.CountIf)
+                                    documents._QueryFeature.CountIf + "," +
+                                    documents._QueryFeature.WeightedRankFusion)
         if os.environ.get(Constants.NON_STREAMING_ORDER_BY_DISABLED_CONFIG,
                           Constants.NON_STREAMING_ORDER_BY_DISABLED_CONFIG_DEFAULT) == "True":
             supported_query_features = (documents._QueryFeature.Aggregate + "," +

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py

@@ -59,7 +59,7 @@ def __init__(self, client, resource_link, options, partitioned_query_execution_i
         self._document_producer_comparator = None
         self._response_hook = response_hook
 
-    async def _run_hybrid_search(self):
+    async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-statements
         # Check if we need to run global statistics queries, and if so do for every partition in the container
         if self._hybrid_search_query_info['requiresGlobalStatistics']:
             target_partition_key_ranges = await self._get_target_partition_key_range(target_all_ranges=True)
@@ -147,21 +147,32 @@ async def _run_hybrid_search(self):
             self._format_final_results(drained_results)
             return
 
+        # Get the Components weights if any
+        if self._hybrid_search_query_info.get('componentWeights'):
+            component_weights = self._hybrid_search_query_info['componentWeights']
+        else:
+            # If no weights are provided, we default to 1.0 for all components
+            component_weights = [1.0] * len(self._hybrid_search_query_info['componentQueryInfos'])
+
         # Sort drained results by _rid
         drained_results.sort(key=lambda x: x['_rid'])
 
         # Compose component scores matrix, where each tuple is (score, index)
         component_scores = _retrieve_component_scores(drained_results)
 
-        # Sort by scores in descending order
-        for score_tuples in component_scores:
-            score_tuples.sort(key=lambda x: x[0], reverse=True)
+        # Sort by scores using component weights
+        for index, score_tuples in enumerate(component_scores):
+            # Negative Weights will change sorting from Descending to Ascending
+            ordering = self._hybrid_search_query_info['componentQueryInfos'][index]['orderBy'][0]
+            comparison_factor = not ordering.lower() == 'ascending'
+            #  pylint: disable=cell-var-from-loop
+            score_tuples.sort(key=lambda x: x[0], reverse=comparison_factor)
 
         # Compute the ranks
         ranks = _compute_ranks(component_scores)
 
         # Compute the RRF scores and add them to output
-        _compute_rrf_scores(ranks, drained_results)
+        _compute_rrf_scores(ranks, component_weights, drained_results)
 
         # Finally, sort on the RRF scores to build the final result to return
         drained_results.sort(key=lambda x: x['Score'], reverse=True)

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py

@@ -3,7 +3,7 @@
 
 """Internal class for multi execution context aggregator implementation in the Azure Cosmos database service.
 """
-
+from typing import List, Union
 from azure.cosmos._execution_context.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context import document_producer
 from azure.cosmos._routing import routing_range
@@ -36,12 +36,13 @@ def _retrieve_component_scores(drained_results):
     return component_scores_list
 
 
-def _compute_rrf_scores(ranks, query_results):
+def _compute_rrf_scores(ranks: List[List[int]], component_weights: List[Union[int, float]], query_results: List[dict]):
     component_count = len(ranks)
     for index, result in enumerate(query_results):
         rrf_score = 0.0
         for component_index in range(component_count):
-            rrf_score += 1.0 / (RRF_CONSTANT + ranks[component_index][index])
+            rrf_score += component_weights[component_index] / (RRF_CONSTANT + ranks[component_index][index])
+
         # Add the score to the item to be returned
         result['Score'] = rrf_score
 
@@ -54,7 +55,7 @@ def _compute_ranks(component_scores):
         rank = 1  # ranks are 1-based
         for index, score_tuple in enumerate(scores):
             # Identical scores should have the same rank
-            if index > 0 and score_tuple[0] < scores[index - 1][0]:
+            if index > 0 and score_tuple[0] != scores[index - 1][0]:
                 rank += 1
             ranks[component_index][score_tuple[1]] = rank
 
@@ -164,7 +165,7 @@ def __init__(self, client, resource_link, options,
         self._document_producer_comparator = None
         self._response_hook = response_hook
 
-    def _run_hybrid_search(self):
+    def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-statements
         # Check if we need to run global statistics queries, and if so do for every partition in the container
         if self._hybrid_search_query_info['requiresGlobalStatistics']:
             target_partition_key_ranges = self._get_target_partition_key_range(target_all_ranges=True)
@@ -251,21 +252,33 @@ def _run_hybrid_search(self):
             self._format_final_results(drained_results)
             return
 
+        # Get the Components weight if any
+        if self._hybrid_search_query_info.get('componentWeights'):
+            component_weights = self._hybrid_search_query_info['componentWeights']
+        else:
+            # If no weights are provided, we assume all components have equal weight
+            component_weights = [1.0] * len(self._hybrid_search_query_info['componentQueryInfos'])
+
         # Sort drained results by _rid
         drained_results.sort(key=lambda x: x['_rid'])
 
         # Compose component scores matrix, where each tuple is (score, index)
         component_scores = _retrieve_component_scores(drained_results)
 
-        # Sort by scores in descending order
-        for score_tuples in component_scores:
-            score_tuples.sort(key=lambda x: x[0], reverse=True)
+        # Sort by scores using component weights
+        for index, score_tuples in enumerate(component_scores):
+            # Ordering of the component query is based on if the weight is negative or positive
+            # A positive weight ordering means descending order, a negative weight ordering means ascending order
+            ordering = self._hybrid_search_query_info['componentQueryInfos'][index]['orderBy'][0]
+            comparison_factor = not ordering.lower() == 'ascending'
+            #  pylint: disable=cell-var-from-loop
+            score_tuples.sort(key=lambda x: x[0], reverse=comparison_factor)
 
         # Compute the ranks
         ranks = _compute_ranks(component_scores)
 
         # Compute the RRF scores and add them to output
-        _compute_rrf_scores(ranks, drained_results)
+        _compute_rrf_scores(ranks, component_weights, drained_results)
 
         # Finally, sort on the RRF scores to build the final result to return
         drained_results.sort(key=lambda x: x['Score'], reverse=True)

sdk/cosmos/azure-cosmos/azure/cosmos/_version.py

@@ -19,4 +19,4 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-VERSION = "4.11.0b2"
+VERSION = "4.12.0b1"

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py

@@ -3207,7 +3207,8 @@ async def _GetQueryPlanThroughGateway(self, query: str, resource_link: str, **kw
                                     documents._QueryFeature.Top + "," +
                                     documents._QueryFeature.NonStreamingOrderBy + "," +
                                     documents._QueryFeature.HybridSearch + "," +
-                                    documents._QueryFeature.CountIf)
+                                    documents._QueryFeature.CountIf + "," +
+                                    documents._QueryFeature.WeightedRankFusion)
         if os.environ.get(Constants.NON_STREAMING_ORDER_BY_DISABLED_CONFIG,
                           Constants.NON_STREAMING_ORDER_BY_DISABLED_CONFIG_DEFAULT) == "True":
             supported_query_features = (documents._QueryFeature.Aggregate + "," +

sdk/cosmos/azure-cosmos/azure/cosmos/documents.py

@@ -438,7 +438,7 @@ class _QueryFeature:
     NonStreamingOrderBy: Literal["NonStreamingOrderBy"] = "NonStreamingOrderBy"
     HybridSearch: Literal["HybridSearch"] = "HybridSearch"
     CountIf: Literal["CountIf"] = "CountIf"
-
+    WeightedRankFusion: Literal["WeightedRankFusion"] = "WeightedRankFusion"
 
 class _DistinctType:
     NoneType: Literal["None"] = "None"