Fulltext query parameters async test fix (#42871)

* 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

* Update test_query_hybrid_search_async.py

Removed enable_cross_partition_query param from async tests as async cosmos db queries does not use this parameter.

* Update how we grab parameters for hybrid search queries.

Fixed issue that was causing original fix to using parameters with hybrid text queries to not work.

* additional test fixes

* test string change

* update edge case of parameterized hybrid search queries

* Fix pylint issues

* Update test_query_hybrid_search_async.py

---------

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

Author: bambriz

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

@@ -7,7 +7,7 @@
 from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.aio import document_producer
 from azure.cosmos._execution_context.hybrid_search_aggregator import _retrieve_component_scores, _rewrite_query_infos, \
-    _compute_rrf_scores, _compute_ranks, _coalesce_duplicate_rids
+    _compute_rrf_scores, _compute_ranks, _coalesce_duplicate_rids, _attach_parameters
 from azure.cosmos._routing import routing_range
 from azure.cosmos import exceptions
 
@@ -53,9 +53,15 @@ 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
+        self._parameters = None
         # 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
+        if query_execution_info:
+            self._parameters = (
+                query_execution_info.get("parameters")
+                if isinstance(query_execution_info, dict)
+                else getattr(query_execution_info, "parameters", None)
+            )
         self._hybrid_search_query_info = hybrid_search_query_info
         self._final_results = []
         self._aggregated_global_statistics = None
@@ -68,13 +74,8 @@ async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-ma
         if self._hybrid_search_query_info['requiresGlobalStatistics']:
             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
-                }
+            global_statistics_query = self._attach_parameters(self._hybrid_search_query_info['globalStatisticsQuery'])
+
             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,7 +114,7 @@ async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-ma
         component_query_infos = self._hybrid_search_query_info['componentQueryInfos']
         if self._aggregated_global_statistics:
             rewritten_query_infos = _rewrite_query_infos(self._hybrid_search_query_info,
-                                                         self._aggregated_global_statistics)
+                                                         self._aggregated_global_statistics, self._parameters)
         else:
             rewritten_query_infos = component_query_infos
 
@@ -123,10 +124,8 @@ async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-ma
         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
-                    }
+                    rewritten_query['rewrittenQuery'] = _attach_parameters(rewritten_query['rewrittenQuery'],
+                                                                           self._parameters)
                 component_query_execution_list.append(
                     document_producer._DocumentProducer(
                         pk_range,
@@ -195,6 +194,26 @@ async def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-ma
         drained_results.sort(key=lambda x: x['Score'], reverse=True)
         self._format_final_results(drained_results)
 
+    def _attach_parameters(self, query):
+        """Attach original query parameters (if any) without mutating the passed query object.
+
+        :param query: The original query (string or dict) to which saved parameters should be attached.
+        :type query: str or dict
+        :return: The query with parameters attached. Returns the original object if no parameters are stored.
+                 If the input was a string and parameters exist, a new dict is returned. If the input was a
+                 dict without "parameters", a shallow copied dict with "parameters" added is returned.
+        :rtype: str or dict
+        """
+        if not self._parameters:
+            return query
+        if isinstance(query, dict):
+            if "parameters" not in query:
+                new_query = dict(query)
+                new_query["parameters"] = self._parameters
+                return new_query
+            return query
+        return {"query": query, "parameters": self._parameters}
+
     def _format_final_results(self, results):
         skip = self._hybrid_search_query_info['skip'] or 0
         take = self._hybrid_search_query_info['take']

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

@@ -79,7 +79,7 @@ def _drain_and_coalesce_results(document_producers_to_drain):
     return all_results, is_singleton
 
 
-def _rewrite_query_infos(hybrid_search_query_info, global_statistics):
+def _rewrite_query_infos(hybrid_search_query_info, global_statistics, parameters=None):
     rewritten_query_infos = []
     for query_info in hybrid_search_query_info['componentQueryInfos']:
         assert query_info['orderBy']
@@ -90,6 +90,8 @@ def _rewrite_query_infos(hybrid_search_query_info, global_statistics):
                 _format_component_query_workaround(order_by_expression, global_statistics,
                                                    len(hybrid_search_query_info[
                                                            'componentQueryInfos'])))
+
+        query_info['rewrittenQuery'] = _attach_parameters(query_info['rewrittenQuery'], parameters)
         rewritten_query = _format_component_query_workaround(query_info['rewrittenQuery'],
                                                              global_statistics,
                                                              len(hybrid_search_query_info[
@@ -118,6 +120,10 @@ def _format_component_query(format_string, global_statistics):
 
 def _format_component_query_workaround(format_string, global_statistics, component_count):
     # TODO: remove this method once the fix is live and switch back to one above
+    parameters = None
+    if isinstance(format_string, dict):
+        parameters = format_string.get('parameters', None)
+        format_string = format_string['query']
     format_string = format_string.replace(_Placeholders.formattable_order_by, "true")
     query = format_string.replace(_Placeholders.total_document_count,
                                   str(global_statistics['documentCount']))
@@ -137,7 +143,29 @@ def _format_component_query_workaround(format_string, global_statistics, compone
 
         statistics_index += 1
 
-    return query
+    return _attach_parameters(query, parameters)
+
+
+def _attach_parameters(query, parameters=None):
+    """Attach original query parameters (if any) without mutating the passed query object.
+
+    :param query: The original query text or a query payload dict which may already contain parameters.
+    :type query: str or dict
+    :param parameters: Optional sequence of parameter definitions to attach.
+    :type parameters: list or None
+    :returns: The original query if no parameters to attach or already present, otherwise a new dict containing the
+     query and parameters.
+    :rtype: str or dict
+    """
+    if not parameters:
+        return query
+    if isinstance(query, dict):
+        if "parameters" not in query:
+            new_query = dict(query)
+            new_query["parameters"] = parameters
+            return new_query
+        return query
+    return {"query": query, "parameters": parameters}
 
 
 class _HybridSearchContextAggregator(_QueryExecutionContextBase):  # pylint: disable=too-many-instance-attributes
@@ -159,9 +187,15 @@ def __init__(self, client, resource_link, options,
         self._client = client
         self._resource_link = resource_link
         self._partitioned_query_ex_info = partitioned_query_execution_info
+        self._parameters = None
         # 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
+        if query_execution_info:
+            self._parameters = (
+                query_execution_info.get("parameters")
+                if isinstance(query_execution_info, dict)
+                else getattr(query_execution_info, "parameters", None)
+            )
         self._hybrid_search_query_info = hybrid_search_query_info
         self._final_results = []
         self._aggregated_global_statistics = None
@@ -174,13 +208,7 @@ def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-sta
         if self._hybrid_search_query_info['requiresGlobalStatistics']:
             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
-                }
+            global_statistics_query = self._attach_parameters(self._hybrid_search_query_info['globalStatisticsQuery'])
             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,7 +246,7 @@ def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-sta
         # re-write the component queries if needed
         if self._aggregated_global_statistics:
             rewritten_query_infos = _rewrite_query_infos(self._hybrid_search_query_info,
-                                                         self._aggregated_global_statistics)
+                                                         self._aggregated_global_statistics, self._parameters)
         else:
             rewritten_query_infos = self._hybrid_search_query_info['componentQueryInfos']
 
@@ -229,10 +257,7 @@ def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-sta
             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
-                    }
+                    rewritten_query['rewrittenQuery'] = self._attach_parameters(rewritten_query['rewrittenQuery'])
                 component_query_execution_list.append(
                     document_producer._DocumentProducer(
                         pk_range,
@@ -302,6 +327,25 @@ def _run_hybrid_search(self):  # pylint: disable=too-many-branches, too-many-sta
         drained_results.sort(key=lambda x: x['Score'], reverse=True)
         self._format_final_results(drained_results)
 
+    def _attach_parameters(self, query):
+        """Attach original query parameters (if any) without mutating the passed query object.
+
+        :param query: Query text or a query payload dict which may already contain parameters.
+        :type query: str or dict
+        :return: The original query if no parameters to attach or already present; otherwise a new dict containing the
+         query and parameters.
+        :rtype: str or dict
+        """
+        if not self._parameters:
+            return query
+        if isinstance(query, dict):
+            if "parameters" not in query:
+                new_query = dict(query)
+                new_query["parameters"] = self._parameters
+                return new_query
+            return query
+        return {"query": query, "parameters": self._parameters}
+
     def _format_final_results(self, results):
         skip = self._hybrid_search_query_info['skip'] or 0
         take = self._hybrid_search_query_info['take']
@@ -320,7 +364,7 @@ def _rewrite_query_infos(self):
                     _format_component_query_workaround(order_by_expression, self._aggregated_global_statistics,
                                                        len(self._hybrid_search_query_info[
                                                                'componentQueryInfos'])))
-
+            query_info['rewrittenQuery'] = _attach_parameters(query_info['rewrittenQuery'], self._parameters)
             rewritten_query = _format_component_query_workaround(query_info['rewrittenQuery'],
                                                                  self._aggregated_global_statistics,
                                                                  len(self._hybrid_search_query_info[

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

@@ -373,6 +373,7 @@ def test_hybrid_search_query_with_params_equivalence(self):
             "ORDER BY RANK FullTextScore(c.title, 'John')"
         )
         literal_results = self.test_container.query_items(literal_query, enable_cross_partition_query=True)
+        literal_results = list(literal_results)
         literal_indices = [res["index"] for res in literal_results]
 
         # Parameterized hybrid query (same as above, but using @term)
@@ -385,6 +386,7 @@ def test_hybrid_search_query_with_params_equivalence(self):
         param_results = self.test_container.query_items(
             param_query, parameters=params, enable_cross_partition_query=True
         )
+        param_results = list(param_results)
         param_indices = [res["index"] for res in param_results]
 
         # Checks: both forms produce the same results and match known expectation
@@ -398,6 +400,7 @@ def test_weighted_rrf_hybrid_search_with_params_and_response_hook(self):
             "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_results = list(literal_results)
         literal_indices = [res["index"] for res in literal_results]
 
         # Parameterized weighted RRF hybrid query (+ response hook)
@@ -414,6 +417,7 @@ def test_weighted_rrf_hybrid_search_with_params_and_response_hook(self):
         param_results = self.test_container.query_items(
             param_query, parameters=params, enable_cross_partition_query=True, response_hook=response_hook
         )
+        param_results = list(param_results)
         param_indices = [res["index"] for res in param_results]
 
         # Checks: number of results, equality against literal, and hook invoked
@@ -423,13 +427,13 @@ def test_weighted_rrf_hybrid_search_with_params_and_response_hook(self):
 
     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)
+        item = self.test_container.read_item('50', '1')
+        item_vector = item['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_results = list(literal_hybrid_results)
         literal_hybrid_indices = [res["index"] for res in literal_hybrid_results]
 
         param_hybrid = (
@@ -444,6 +448,7 @@ def test_hybrid_and_non_hybrid_param_queries_equivalence(self):
         param_hybrid_results = self.test_container.query_items(
             param_hybrid, parameters=params_hybrid, enable_cross_partition_query=True
         )
+        param_hybrid_results = list(param_hybrid_results)
         param_hybrid_indices = [res["index"] for res in param_hybrid_results]
 
         assert len(literal_hybrid_indices) == len(param_hybrid_indices) == 10
@@ -453,6 +458,7 @@ def test_hybrid_and_non_hybrid_param_queries_equivalence(self):
         # 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_results = list(literal_simple_results)
         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"