Fix hybrid text query with parameters (#42787)

* Fix hybrid text query with parameters

* Update CHANGELOG.md

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

Co-authored-by: Copilot <[email protected]>

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

Co-authored-by: Copilot <[email protected]>

* pylint fix and formatting fix

---------

Co-authored-by: Copilot <[email protected]>

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -13,6 +13,7 @@
 * Fixed bug where during health checks read regions were marked as unavailable for write operations. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).
 * Fixed bug where containers named with spaces or special characters using session consistency would fall back to eventual consistency. See [PR 42608](https://github.com/Azure/azure-sdk-for-python/pull/42608)
 * Fixed bug where `excluded_locations` was not being honored for some metadata calls. See [PR 42266](https://github.com/Azure/azure-sdk-for-python/pull/42266).
+* Fixed bug where Hybrid Search queries using parameters were not working. See [PR 42787](https://github.com/Azure/azure-sdk-for-python/pull/42787)
 * Fixed partition scoping for per partition circuit breaker. See [PR 42751](https://github.com/Azure/azure-sdk-for-python/pull/42751)
 
 #### Other Changes

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

@@ -68,6 +68,11 @@ async def _create_execution_context_with_query_plan(self):
         query_to_use = self._query if self._query is not None else "Select * from root r"
         query_execution_info = _PartitionedQueryExecutionInfo(await self._client._GetQueryPlanThroughGateway
         (query_to_use, self._resource_link, self._options.get('excludedLocations')))
+        qe_info = getattr(query_execution_info, "_query_execution_info", None)
+        if isinstance(qe_info, dict) and isinstance(query_to_use, dict):
+            params = query_to_use.get("parameters")
+            if params is not None:
+                query_execution_info._query_execution_info['parameters'] = params
         self._execution_context = await self._create_pipelined_execution_context(query_execution_info)
 
     async def __anext__(self):

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

@@ -34,7 +34,7 @@ async def _drain_and_coalesce_results(document_producers_to_drain):
     return all_results, is_singleton
 
 
-class _HybridSearchContextAggregator(_QueryExecutionContextBase):
+class _HybridSearchContextAggregator(_QueryExecutionContextBase):  # pylint: disable=too-many-instance-attributes
     """This class is a subclass of the query execution context base and serves for
     full text search and hybrid search queries. It is very similar to the existing MultiExecutionContextAggregator,
     but is needed since we have a lot more additional client-side logic to take care of.
@@ -53,6 +53,9 @@ def __init__(self, client, resource_link, options, partitioned_query_execution_i
         self._client = client
         self._resource_link = resource_link
         self._partitioned_query_ex_info = partitioned_query_execution_info
+        # If the query uses parameters, we must save them to add them back to the component queries
+        query_execution_info = getattr(self._partitioned_query_ex_info, "_query_execution_info", None)
+        self._parameters = getattr(query_execution_info, "parameters", None) if query_execution_info else None
         self._hybrid_search_query_info = hybrid_search_query_info
         self._final_results = []
         self._aggregated_global_statistics = None
@@ -66,6 +69,12 @@ async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-ma
             target_partition_key_ranges = await self._get_target_partition_key_range(target_all_ranges=True)
             global_statistics_doc_producers = []
             global_statistics_query = self._hybrid_search_query_info['globalStatisticsQuery']
+            # If query was given parameters we must add them back in
+            if self._parameters:
+                global_statistics_query = {
+                    'query': global_statistics_query,
+                    'parameters': self._parameters
+                }
             partitioned_query_execution_context_list = []
             for partition_key_target_range in target_partition_key_ranges:
                 # create a document producer for each partition key range
@@ -113,6 +122,11 @@ async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-ma
         target_partition_key_ranges = await self._get_target_partition_key_range(target_all_ranges=False)
         for rewritten_query in rewritten_query_infos:
             for pk_range in target_partition_key_ranges:
+                if self._parameters:
+                    rewritten_query['rewrittenQuery'] = {
+                        'query': rewritten_query['rewrittenQuery'],
+                        'parameters': self._parameters
+                    }
                 component_query_execution_list.append(
                     document_producer._DocumentProducer(
                         pk_range,

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

@@ -98,6 +98,12 @@ def _create_execution_context_with_query_plan(self):
         query_to_use = self._query if self._query is not None else "Select * from root r"
         query_execution_info = _PartitionedQueryExecutionInfo(self._client._GetQueryPlanThroughGateway
         (query_to_use, self._resource_link, self._options.get('excludedLocations')))
+
+        qe_info = getattr(query_execution_info, "_query_execution_info", None)
+        if isinstance(qe_info, dict) and isinstance(query_to_use, dict):
+            params = query_to_use.get("parameters")
+            if params is not None:
+                query_execution_info._query_execution_info['parameters'] = params
         self._execution_context = self._create_pipelined_execution_context(query_execution_info)
 
     def __next__(self):

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

@@ -140,7 +140,7 @@ def _format_component_query_workaround(format_string, global_statistics, compone
     return query
 
 
-class _HybridSearchContextAggregator(_QueryExecutionContextBase):
+class _HybridSearchContextAggregator(_QueryExecutionContextBase):  # pylint: disable=too-many-instance-attributes
     """This class is a subclass of the query execution context base and serves for
     full text search and hybrid search queries. It is very similar to the existing MultiExecutionContextAggregator,
     but is needed since we have a lot more additional client-side logic to take care of.
@@ -159,6 +159,9 @@ def __init__(self, client, resource_link, options,
         self._client = client
         self._resource_link = resource_link
         self._partitioned_query_ex_info = partitioned_query_execution_info
+        # If the query uses parameters, we must save them to add them back to the component queries
+        query_execution_info = getattr(self._partitioned_query_ex_info, "_query_execution_info", None)
+        self._parameters = getattr(query_execution_info, "parameters", None) if query_execution_info else None
         self._hybrid_search_query_info = hybrid_search_query_info
         self._final_results = []
         self._aggregated_global_statistics = None
@@ -172,6 +175,12 @@ def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-sta
             target_partition_key_ranges = self._get_target_partition_key_range(target_all_ranges=True)
             global_statistics_doc_producers = []
             global_statistics_query = self._hybrid_search_query_info['globalStatisticsQuery']
+            # If query was given parameters we must add them back in
+            if self._parameters:
+                global_statistics_query = {
+                    'query': global_statistics_query,
+                    'parameters': self._parameters
+                }
             partitioned_query_execution_context_list = []
             for partition_key_target_range in target_partition_key_ranges:
                 # create a document producer for each partition key range
@@ -218,6 +227,12 @@ def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-sta
         target_partition_key_ranges = self._get_target_partition_key_range(target_all_ranges=False)
         for rewritten_query in rewritten_query_infos:
             for pk_range in target_partition_key_ranges:
+                # If query was given parameters we must add them back in
+                if self._parameters:
+                    rewritten_query['rewrittenQuery'] = {
+                        'query': rewritten_query['rewrittenQuery'],
+                        'parameters': self._parameters
+                    }
                 component_query_execution_list.append(
                     document_producer._DocumentProducer(
                         pk_range,

sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py

@@ -365,6 +365,105 @@ def test_weighted_reciprocal_rank_fusion_with_response_hook(self):
         assert len(result_list) == 10
         assert response_hook.count > 0  # Ensure the response hook was called
 
+    def test_hybrid_search_query_with_params_equivalence(self):
+        # Literal hybrid query
+        literal_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "WHERE FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') "
+            "ORDER BY RANK FullTextScore(c.title, 'John')"
+        )
+        literal_results = self.test_container.query_items(literal_query, enable_cross_partition_query=True)
+        literal_indices = [res["index"] for res in literal_results]
+
+        # Parameterized hybrid query (same as above, but using @term)
+        param_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "WHERE FullTextContains(c.title, @term) OR FullTextContains(c.text, @term) "
+            "ORDER BY RANK FullTextScore(c.title, @term)"
+        )
+        params = [{"name": "@term", "value": "John"}]
+        param_results = self.test_container.query_items(
+            param_query, parameters=params, enable_cross_partition_query=True
+        )
+        param_indices = [res["index"] for res in param_results]
+
+        # Checks: both forms produce the same results and match known expectation
+        assert len(literal_indices) == len(param_indices) == 3
+        assert set(literal_indices) == set(param_indices) == {2, 85, 57}
+
+    def test_weighted_rrf_hybrid_search_with_params_and_response_hook(self):
+        # Literal weighted RRF hybrid query
+        literal_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States'), [1, 0.5])"
+        )
+        literal_results = self.test_container.query_items(literal_query, enable_cross_partition_query=True)
+        literal_indices = [res["index"] for res in literal_results]
+
+        # Parameterized weighted RRF hybrid query (+ response hook)
+        response_hook = test_config.ResponseHookCaller()
+        param_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.title, @titleTerm), FullTextScore(c.text, @textTerm), @weights)"
+        )
+        params = [
+            {"name": "@titleTerm", "value": "John"},
+            {"name": "@textTerm", "value": "United States"},
+            {"name": "@weights", "value": [1, 0.5]},
+        ]
+        param_results = self.test_container.query_items(
+            param_query, parameters=params, enable_cross_partition_query=True, response_hook=response_hook
+        )
+        param_indices = [res["index"] for res in param_results]
+
+        # Checks: number of results, equality against literal, and hook invoked
+        assert len(literal_indices) == len(param_indices) == 10
+        assert set(literal_indices) == set(param_indices)
+        assert response_hook.count > 0
+
+    def test_hybrid_and_non_hybrid_param_queries_equivalence(self):
+        # Hybrid query with vector distance (literal vs param) and compare equality
+        item_vector = self.test_container.read_item("50", "1")["vector"]
+        literal_hybrid = (
+            "SELECT c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.text, 'United States'), VectorDistance(c.vector, {})) "
+            "OFFSET 0 LIMIT 10"
+        ).format(item_vector)
+        literal_hybrid_results = self.test_container.query_items(literal_hybrid, enable_cross_partition_query=True)
+        literal_hybrid_indices = [res["index"] for res in literal_hybrid_results]
+
+        param_hybrid = (
+            "SELECT c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.text, @country), VectorDistance(c.vector, @vec)) "
+            "OFFSET 0 LIMIT 10"
+        )
+        params_hybrid = [
+            {"name": "@country", "value": "United States"},
+            {"name": "@vec", "value": item_vector},
+        ]
+        param_hybrid_results = self.test_container.query_items(
+            param_hybrid, parameters=params_hybrid, enable_cross_partition_query=True
+        )
+        param_hybrid_indices = [res["index"] for res in param_hybrid_results]
+
+        assert len(literal_hybrid_indices) == len(param_hybrid_indices) == 10
+        # Compare ordered lists to ensure identical ranking
+        assert literal_hybrid_indices == param_hybrid_indices
+
+        # Non-hybrid parameterized query equivalence on same container
+        literal_simple = "SELECT TOP 5 c.index FROM c WHERE c.pk = '1' ORDER BY c.index"
+        literal_simple_results = self.test_container.query_items(literal_simple, enable_cross_partition_query=True)
+        literal_simple_indices = [res["index"] for res in literal_simple_results]
+
+        param_simple = "SELECT TOP 5 c.index FROM c WHERE c.pk = @pk ORDER BY c.index"
+        params_simple = [{"name": "@pk", "value": "1"}]
+        param_simple_results = self.test_container.query_items(
+            param_simple, parameters=params_simple, enable_cross_partition_query=True
+        )
+        param_simple_indices = [res["index"] for res in param_simple_results]
+
+        assert len(literal_simple_indices) == len(param_simple_indices) == 5
+        assert literal_simple_indices == param_simple_indices
 
 
 if __name__ == "__main__":

sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py

@@ -368,8 +368,105 @@ async def test_weighted_reciprocal_rank_fusion_with_response_hook_async(self):
         assert len(result_list) == 10
         assert response_hook.count > 0  # Ensure the response hook was called
 
-
-
+    async def test_hybrid_search_query_with_params_equivalence_async(self):
+        # Literal hybrid query
+        literal_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "WHERE FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') "
+            "ORDER BY RANK FullTextScore(c.title, 'John')"
+        )
+        literal_results = self.test_container.query_items(literal_query, enable_cross_partition_query=True)
+        literal_indices = [res["index"] async for res in literal_results]
+
+        # Parameterized hybrid query (same as above, but using @term)
+        param_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "WHERE FullTextContains(c.title, @term) OR FullTextContains(c.text, @term) "
+            "ORDER BY RANK FullTextScore(c.title, @term)"
+        )
+        params = [{"name": "@term", "value": "John"}]
+        param_results = self.test_container.query_items(
+            param_query, parameters=params, enable_cross_partition_query=True
+        )
+        param_indices = [res["index"] async for res in param_results]
+
+        # Checks: both forms produce the same results and match known expectation
+        assert len(literal_indices) == len(param_indices) == 3
+        assert set(literal_indices) == set(param_indices) == {2, 85, 57}
+
+    async def test_weighted_rrf_hybrid_search_with_params_and_response_hook_async(self):
+        # Literal weighted RRF hybrid query
+        literal_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States'), [1, 0.5])"
+        )
+        literal_results = self.test_container.query_items(literal_query, enable_cross_partition_query=True)
+        literal_indices = [res["index"] async for res in literal_results]
+
+        # Parameterized weighted RRF hybrid query (+ response hook)
+        response_hook = test_config.ResponseHookCaller()
+        param_query = (
+            "SELECT TOP 10 c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.title, @titleTerm), FullTextScore(c.text, @textTerm), @weights)"
+        )
+        params = [
+            {"name": "@titleTerm", "value": "John"},
+            {"name": "@textTerm", "value": "United States"},
+            {"name": "@weights", "value": [1, 0.5]},
+        ]
+        param_results = self.test_container.query_items(
+            param_query, parameters=params, enable_cross_partition_query=True, response_hook=response_hook
+        )
+        param_indices = [res["index"] async for res in param_results]
+
+        # Checks: number of results, equality against literal, and hook invoked
+        assert len(literal_indices) == len(param_indices) == 10
+        assert set(literal_indices) == set(param_indices)
+        assert response_hook.count > 0
+
+    async def test_hybrid_and_non_hybrid_param_queries_equivalence_async(self):
+        # Hybrid query with vector distance (literal vs param) and compare equality
+        item_vector = (await self.test_container.read_item("50", "1"))["vector"]
+        literal_hybrid = (
+            "SELECT c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.text, 'United States'), VectorDistance(c.vector, {})) "
+            "OFFSET 0 LIMIT 10"
+        ).format(item_vector)
+        literal_hybrid_results = self.test_container.query_items(literal_hybrid, enable_cross_partition_query=True)
+        literal_hybrid_indices = [res["index"] async for res in literal_hybrid_results]
+
+        param_hybrid = (
+            "SELECT c.index, c.title FROM c "
+            "ORDER BY RANK RRF(FullTextScore(c.text, @country), VectorDistance(c.vector, @vec)) "
+            "OFFSET 0 LIMIT 10"
+        )
+        params_hybrid = [
+            {"name": "@country", "value": "United States"},
+            {"name": "@vec", "value": item_vector},
+        ]
+        param_hybrid_results = self.test_container.query_items(
+            param_hybrid, parameters=params_hybrid, enable_cross_partition_query=True
+        )
+        param_hybrid_indices = [res["index"] async for res in param_hybrid_results]
+
+        assert len(literal_hybrid_indices) == len(param_hybrid_indices) == 10
+        # Compare ordered lists to ensure identical ranking
+        assert literal_hybrid_indices == param_hybrid_indices
+
+        # Non-hybrid parameterized query equivalence on same container
+        literal_simple = "SELECT TOP 5 c.index FROM c WHERE c.pk = '1' ORDER BY c.index"
+        literal_simple_results = self.test_container.query_items(literal_simple, enable_cross_partition_query=True)
+        literal_simple_indices = [res["index"] async for res in literal_simple_results]
+
+        param_simple = "SELECT TOP 5 c.index FROM c WHERE c.pk = @pk ORDER BY c.index"
+        params_simple = [{"name": "@pk", "value": "1"}]
+        param_simple_results = self.test_container.query_items(
+            param_simple, parameters=params_simple, enable_cross_partition_query=True
+        )
+        param_simple_indices = [res["index"] async for res in param_simple_results]
+
+        assert len(literal_simple_indices) == len(param_simple_indices) == 5
+        assert literal_simple_indices == param_simple_indices
 
 
 if __name__ == "__main__":