[Cosmos] Hybrid Search query pipeline

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 #### Features Added
 * Added full text policy and full text indexing policy. See [PR 37891](https://github.com/Azure/azure-sdk-for-python/pull/37891).
+* Added support for full text search and hybrid search queries. See [PR 38275](https://github.com/Azure/azure-sdk-for-python/pull/38275).
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/README.md

@@ -818,6 +818,40 @@ indexing_policy = {
 Modifying the index in a container is an asynchronous operation that can take a long time to finish. See [here][cosmos_index_policy_change] for more information.
 For more information on using full text policies and full text indexes, see [here][cosmos_fts].
 
+### Public Preview - Full Text Search and Hybrid Search
+
+With the addition of the full text indexing and full text policies, the SDK can now perform full text search and hybrid search queries.
+These queries can utilize the new query functions `FullTextContains()`, `FullTextContainsAll`, and `FullTextContainsAny` to efficiently
+search for the given terms within your item fields.
+
+Beyond these, you can also utilize the new `Order By RANK` and `Order By RANK RRF` along with `FullTextScore` to execute the [BM25][BM25] scoring algorithm
+or [Reciprocal Rank Fusion][RRF] (RRF) on your query, finding the items with the highest relevance to the terms you are looking for.
+All of these mentioned queries would look something like this:
+
+- `SELECT TOP 10 c.id, c.text FROM c WHERE FullTextContains(c.text, 'quantum')`
+
+
+- `SELECT TOP 10 c.id, c.text FROM c WHERE FullTextContainsAll(c.text, 'quantum', 'theory')`
+
+
+- `SELECT TOP 10 c.id, c.text FROM c WHERE FullTextContainsAny(c.text, 'quantum', 'theory')`
+
+
+- `SELECT TOP 10 c.id, c.text FROM c ORDER BY RANK FullTextScore(c.text, ['quantum', 'theory'])`
+
+
+- `SELECT TOP 10 c.id, c.text FROM c ORDER BY RANK RRF(FullTextScore(c.text, ['quantum', 'theory']), FullTextScore(c.text, ['model']))`
+
+
+- `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}))"`
+
+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
+may have additional latencies associated with searching in the service.
+
+You can find our sync samples [here][cosmos_index_sample] and our async samples [here][cosmos_index_sample_async] as well for additional guidance.
+
 ## Troubleshooting
 
 ### General
@@ -954,6 +988,8 @@ For more extensive documentation on the Cosmos DB service, see the [Azure Cosmos
 [cosmos_concurrency_sample]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/concurrency_sample.py
 [cosmos_index_sample]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/index_management.py
 [cosmos_index_sample_async]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/index_management_async.py
+[RRF]: https://learn.microsoft.com/azure/search/hybrid-search-ranking
+[BM25]: https://learn.microsoft.com/azure/search/index-similarity-and-scoring
 [cosmos_fts]: https://aka.ms/cosmosfulltextsearch
 [cosmos_index_policy_change]: https://learn.microsoft.com/azure/cosmos-db/index-policy#modifying-the-indexing-policy
 

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

@@ -3141,7 +3141,8 @@ def _GetQueryPlanThroughGateway(self, query: str, resource_link: str, **kwargs:
                                     documents._QueryFeature.OffsetAndLimit + "," +
                                     documents._QueryFeature.OrderBy + "," +
                                     documents._QueryFeature.Top + "," +
-                                    documents._QueryFeature.NonStreamingOrderBy)
+                                    documents._QueryFeature.NonStreamingOrderBy + "," +
+                                    documents._QueryFeature.HybridSearch)
         if os.environ.get('AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY', False):
             supported_query_features = (documents._QueryFeature.Aggregate + "," +
                                         documents._QueryFeature.CompositeAggregate + "," +

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

@@ -53,6 +53,7 @@ def __init__(self, partition_key_target_range, client, collection_link, query, d
         self._is_finished = False
         self._has_started = False
         self._cur_item = None
+        self._query = query
         # initiate execution context
 
         path = _base.GetPathFromLink(collection_link, "docs")

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

@@ -173,7 +173,7 @@ def __init__(self, execution_context, aggregate_operators):
         for operator in aggregate_operators:
             if operator == "Average":
                 self._local_aggregators.append(_AverageAggregator())
-            elif operator == "Count":
+            elif operator in ("Count", "CountIf"):
                 self._local_aggregators.append(_CountAggregator())
             elif operator == "Max":
                 self._local_aggregators.append(_MaxAggregator())

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

@@ -25,10 +25,10 @@
 
 import os
 from azure.cosmos._execution_context.aio import endpoint_component, multi_execution_aggregator
-from azure.cosmos._execution_context.aio import non_streaming_order_by_aggregator
+from azure.cosmos._execution_context.aio import non_streaming_order_by_aggregator, hybrid_search_aggregator
 from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.aio.base_execution_context import _DefaultQueryExecutionContext
-from azure.cosmos._execution_context.execution_dispatcher import _is_partitioned_execution_info
+from azure.cosmos._execution_context.execution_dispatcher import _is_partitioned_execution_info, _is_hybrid_search_query
 from azure.cosmos._execution_context.query_execution_info import _PartitionedQueryExecutionInfo
 from azure.cosmos.documents import _DistinctType
 from azure.cosmos.exceptions import CosmosHttpResponseError
@@ -89,7 +89,7 @@ async def fetch_next_block(self):
         try:
             return await self._execution_context.fetch_next_block()
         except CosmosHttpResponseError as e:
-            if _is_partitioned_execution_info(e): #cross partition query not servable
+            if _is_partitioned_execution_info(e) or _is_hybrid_search_query(self._query, e):
                 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))
@@ -126,6 +126,22 @@ async def _create_pipelined_execution_context(self, query_execution_info):
                                                                                         self._options,
                                                                                         query_execution_info)
             await execution_context_aggregator._configure_partition_ranges()
+        elif query_execution_info.has_hybrid_search_query_info():
+            hybrid_search_query_info = query_execution_info._query_execution_info['hybridSearchQueryInfo']
+            if not hybrid_search_query_info['take']:
+                raise ValueError("Executing a hybrid search query without TOP or LIMIT can consume many" +
+                                 " RUs very fast and have long runtimes. Please ensure you are using one" +
+                                 " of the two filters with your hybrid search query.")
+            if hybrid_search_query_info['take'] > os.environ.get('AZURE_COSMOS_HYBRID_SEARCH_MAX_ITEMS', 1000):
+                raise ValueError("Executing a hybrid search query with more items than the max is not allowed." +
+                                 "Please ensure you are using a limit smaller than the max, or change the max.")
+            execution_context_aggregator = \
+                hybrid_search_aggregator._HybridSearchContextAggregator(self._client,
+                                                                        self._resource_link,
+                                                                        self._options,
+                                                                        query_execution_info,
+                                                                        hybrid_search_query_info)
+            await execution_context_aggregator._run_hybrid_search()
         else:
             execution_context_aggregator = multi_execution_aggregator._MultiExecutionContextAggregator(self._client,
                                                                                                    self._resource_link,