[Cosmos] Adds support for non streaming ORDER BY

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,8 @@
 ### 4.6.1 (Unreleased)
 
 #### Features Added
+* Adds vector embedding policy and vector indexing policy. See [PR 34882](https://github.com/Azure/azure-sdk-for-python/pull/34882).
+* Adds support for vector search non-streaming order by queries. See [PR 10101](https://github.com/Azure/azure-sdk-for-python/pull/10101).
 * Added support for using the start time option for change feed query API. See [PR 35090](https://github.com/Azure/azure-sdk-for-python/pull/35090)
 
 #### Breaking Changes
@@ -11,7 +13,6 @@
 * Fixed a bug where change feed query in Async client was not returning all pages due to case-sensitive response headers. See [PR 35090](https://github.com/Azure/azure-sdk-for-python/pull/35090)
 * Fixed a bug when a retryable exception occurs in the first page of a query execution causing query to return 0 results. See [PR 35090](https://github.com/Azure/azure-sdk-for-python/pull/35090).
 
-
 #### Other Changes
 
 ### 4.6.0 (2024-03-14)

sdk/cosmos/azure-cosmos/README.md

@@ -628,6 +628,72 @@ as well as containing the list of failed responses for the failed request.
 
 For more information on Transactional Batch, see [Azure Cosmos DB Transactional Batch][cosmos_transactional_batch].
 
+### Private Preview - Vector Embeddings and Vector Indexes
+We have added new capabilities to utilize vector embeddings and vector indexing for users to leverage vector
+search utilizing our Cosmos SDK. These two container-level configurations have to be turned on at the account-level
+before you can use them.
+
+Each vector embedding should have a path to the relevant vector field in your items being stored, a supported data type
+(float32, int8, uint8), the vector's dimensions (positive int <=1536), and the distance function being used for that embedding.
+A sample vector embedding policy would look like this:
+```python
+vector_embedding_policy = {
+    "vectorEmbeddings": [
+        {
+            "path": "/vector1",
+            "dataType": "float32",
+            "dimensions": 1000,
+            "distanceFunction": "euclidean"
+        },
+        {
+            "path": "/vector2",
+            "dataType": "int8",
+            "dimensions": 200,
+            "distanceFunction": "dotproduct"
+        },
+        {
+            "path": "/vector3",
+            "dataType": "uint8",
+            "dimensions": 400,
+            "distanceFunction": "cosine"
+        }
+    ]
+}
+```
+
+Separately, vector indexes have been added to the already existing indexing_policy and only require two fields per index:
+the path to the relevant field to be used, and the type of index from the possible options (flat, quantizedFlat, or diskANN).
+A sample indexing policy with vector indexes would look like this:
+```python
+indexing_policy = {
+        "automatic": True,
+        "indexingMode": "consistent",
+        "compositeIndexes": [
+            [
+                {"path": "/numberField", "order": "ascending"},
+                {"path": "/stringField", "order": "descending"}
+            ]
+        ],
+        "spatialIndexes": [
+            {"path": "/location/*", "types": [
+                "Point",
+                "Polygon"]}
+        ],
+        "vectorIndexes": [
+            {"path": "/vector1", "type": "flat"},
+            {"path": "/vector2", "type": "quantizedFlat"},
+            {"path": "/vector3", "type": "diskANN"}
+        ]
+    }
+```
+You would then pass in the relevant policies to your container creation method to ensure these configurations are used by it.
+The operation will fail if you pass new vector indexes to your indexing policy but forget to pass in an embedding policy.
+```python
+database.create_container(id=container_id, partition_key=PartitionKey(path="/id"),
+                          indexing_policy=indexing_policy, vector_embedding_policy=vector_embedding_policy)
+```
+***Note: vector embeddings and vector indexes CANNOT be edited by container replace operations. They are only available directly through creation.***
+
 ## Troubleshooting
 
 ### General

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

@@ -3107,7 +3107,8 @@ def _GetQueryPlanThroughGateway(self, query: str, resource_link: str, **kwargs:
                                     documents._QueryFeature.MultipleOrderBy + "," +
                                     documents._QueryFeature.OffsetAndLimit + "," +
                                     documents._QueryFeature.OrderBy + "," +
-                                    documents._QueryFeature.Top)
+                                    documents._QueryFeature.Top + "," +
+                                    documents._QueryFeature.NonStreamingOrderBy)
 
         options = {
             "contentType": runtime_constants.MediaTypes.Json,

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

@@ -271,3 +271,59 @@ def _validate_orderby_items(self, res1, res2):
             type2 = _OrderByHelper.getTypeStr(elt2)
             if type1 != type2:
                 raise ValueError("Expected {}, but got {}.".format(type1, type2))
+
+class _NonStreamingDocumentProducer(object):
+    """This class takes care of handling of the items to be sorted in a non-streaming context.
+    One instance of this document producer goes attached to every item coming in for the priority queue to be able
+    to properly sort items as they get inserted.
+    """
+
+    def __init__(self, item_result, sort_order):
+        """
+        Constructor
+        """
+        self._item_result = item_result
+        self._doc_producer_comp = _NonStreamingOrderByComparator(sort_order)
+
+
+
+class _NonStreamingOrderByComparator(object):
+    """Provide a Comparator for item results which respects orderby sort order.
+    """
+
+    def __init__(self, sort_order):  # pylint: disable=super-init-not-called
+        """Instantiates this class
+        :param list sort_order:
+            List of sort orders (i.e., Ascending, Descending)
+        :ivar list sort_order:
+            List of sort orders (i.e., Ascending, Descending)
+        """
+        self._sort_order = sort_order
+
+    async def compare(self, doc_producer1, doc_producer2):
+        """Compares the given two instances of DocumentProducers.
+        Based on the orderby query items and whether the sort order is Ascending
+        or Descending compares the peek result of the two DocumentProducers.
+        :param _DocumentProducer doc_producer1: first instance to be compared
+        :param _DocumentProducer doc_producer2: second instance to be compared
+        :return:
+            Integer value of compare result.
+                positive integer if doc_producers1 > doc_producers2
+                negative integer if doc_producers1 < doc_producers2
+        :rtype: int
+        """
+        order1 = doc_producer1._item_result["orderByItems"][0]
+        order2 = doc_producer2._item_result["orderByItems"][0]
+        type1_ord = _OrderByHelper.getTypeOrd(order1)
+        type2_ord = _OrderByHelper.getTypeOrd(order2)
+
+        type_ord_diff = type1_ord - type2_ord
+
+        if type_ord_diff:
+            return type_ord_diff
+
+        # the same type,
+        if type1_ord == 0:
+            return 0
+
+        return _compare_helper(order1['item'], order2['item'])

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

@@ -55,8 +55,15 @@ class _QueryExecutionOrderByEndpointComponent(_QueryExecutionEndpointComponent):
     """
     async def __anext__(self):
         payload = await self._execution_context.__anext__()
-        return payload["payload"]
+        return payload._item_result["payload"]
 
+class _QueryExecutionNonStreamingEndpointComponent(_QueryExecutionEndpointComponent):
+    """Represents an endpoint in handling a non-streaming order by query results.
+    For each processed orderby result it returns the item result.
+    """
+    async def __anext__(self):
+        payload = await self._execution_context.__anext__()
+        return payload._item_result["payload"]
 
 class _QueryExecutionTopEndpointComponent(_QueryExecutionEndpointComponent):
     """Represents an endpoint in handling top query.

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

@@ -23,8 +23,8 @@
 Cosmos database service.
 """
 
-from azure.cosmos._execution_context.aio import endpoint_component
-from azure.cosmos._execution_context.aio import multi_execution_aggregator
+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.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
@@ -106,13 +106,28 @@ async def _create_pipelined_execution_context(self, query_execution_info):
                                   and self._options["enableCrossPartitionQuery"]):
                 raise CosmosHttpResponseError(StatusCodes.BAD_REQUEST,
                                   "Cross partition query only supports 'VALUE <AggregateFunc>' for aggregates")
-
-        execution_context_aggregator = multi_execution_aggregator._MultiExecutionContextAggregator(self._client,
+        # throw exception here for vector search query without limit filter
+        if query_execution_info.get_has_non_streaming_order_by():
+            if query_execution_info.get_top() is None and query_execution_info.get_limit() is None:
+                # TODO: missing one last if statement here to check for the system variable bypass - need name
+                raise CosmosHttpResponseError(StatusCodes.BAD_REQUEST,
+                                              "Executing a vector 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 vector search query.")
+            execution_context_aggregator =\
+                non_streaming_order_by_aggregator._NonStreamingOrderByContextAggregator(self._client,
+                                                                                        self._resource_link,
+                                                                                        self._query,
+                                                                                        self._options,
+                                                                                        query_execution_info)
+            await execution_context_aggregator._configure_partition_ranges()
+        else:
+            execution_context_aggregator = multi_execution_aggregator._MultiExecutionContextAggregator(self._client,
                                                                                                    self._resource_link,
                                                                                                    self._query,
                                                                                                    self._options,
                                                                                                    query_execution_info)
-        await execution_context_aggregator._configure_partition_ranges()
+            await execution_context_aggregator._configure_partition_ranges()
         return _PipelineExecutionContext(self._client, self._options, execution_context_aggregator,
                                          query_execution_info)
 
@@ -134,7 +149,9 @@ def __init__(self, client, options, execution_context, query_execution_info):
         self._endpoint = endpoint_component._QueryExecutionEndpointComponent(execution_context)
 
         order_by = query_execution_info.get_order_by()
-        if order_by:
+        if query_execution_info.get_has_non_streaming_order_by():
+            self._endpoint = endpoint_component._QueryExecutionNonStreamingEndpointComponent(self._endpoint)
+        elif order_by:
             self._endpoint = endpoint_component._QueryExecutionOrderByEndpointComponent(self._endpoint)
 
         aggregates = query_execution_info.get_aggregates()