async changes

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()

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

@@ -0,0 +1,170 @@
+# The MIT License (MIT)
+# Copyright (c) 2024 Microsoft Corporation
+
+"""Internal class for multi execution context aggregator implementation in the Azure Cosmos database service.
+"""
+
+from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase
+from azure.cosmos._execution_context.aio import document_producer, _queue_async_helper
+from azure.cosmos._routing import routing_range
+from azure.cosmos import exceptions
+
+# pylint: disable=protected-access
+
+
+class FixedSizePriorityQueue:
+    """Provides a Fixed Size Priority Queue abstraction data structure"""
+
+    def __init__(self, max_size):
+        self._heap = []
+        self.max_size = max_size
+
+    async def pop_async(self, document_producer_comparator):
+        return await _queue_async_helper.heap_pop(self._heap, document_producer_comparator)
+
+    async def push_async(self, item, document_producer_comparator):
+        await _queue_async_helper.heap_push(self._heap, item, document_producer_comparator)
+        if len(self._heap) > self.max_size:
+            await _queue_async_helper.heap_pop(self._heap, document_producer_comparator)
+
+    def peek(self):
+        return self._heap[0]
+
+    def size(self):
+        return len(self._heap)
+
+class _NonStreamingOrderByContextAggregator(_QueryExecutionContextBase):
+    """This class is a subclass of the query execution context base and serves for
+    non-streaming order by queries. It is very similar to the existing MultiExecutionContextAggregator,
+    but is needed since we're dealing with items and not document producers.
+
+    This class builds upon the multi-execution aggregator, building a document producer per partition
+    and draining their results entirely in order to create the result set relevant to the filters passed
+    by the user.
+    """
+
+    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info):
+        super(_NonStreamingOrderByContextAggregator, self).__init__(client, options)
+
+        # use the routing provider in the client
+        self._routing_provider = client._routing_map_provider
+        self._client = client
+        self._resource_link = resource_link
+        self._query = query
+        self._partitioned_query_ex_info = partitioned_query_ex_info
+        self._sort_orders = partitioned_query_ex_info.get_order_by()
+
+        pq_size = partitioned_query_ex_info.get_top() or partitioned_query_ex_info.get_limit()
+        self._orderByPQ = FixedSizePriorityQueue(pq_size)
+
+    async def __anext__(self):
+        """Returns the next result
+
+        :return: The next result.
+        :rtype: dict
+        :raises StopIteration: If no more result is left.
+        """
+        if self._orderByPQ.size() > 0:
+            res = await self._orderByPQ.pop_async(self._document_producer_comparator)
+            return res
+        raise StopAsyncIteration
+
+    async def fetch_next_block(self):
+
+        raise NotImplementedError("You should use pipeline's fetch_next_block.")
+
+    async def _repair_document_producer(self):
+        """Repairs the document producer context by using the re-initialized routing map provider in the client,
+        which loads in a refreshed partition key range cache to re-create the partition key ranges.
+        After loading this new cache, the document producers get re-created with the new valid ranges.
+        """
+        # refresh the routing provider to get the newly initialized one post-refresh
+        self._routing_provider = self._client._routing_map_provider
+        # will be a list of (partition_min, partition_max) tuples
+        targetPartitionRanges = await self._get_target_partition_key_range()
+
+        targetPartitionQueryExecutionContextList = []
+        for partitionTargetRange in targetPartitionRanges:
+            # create and add the child execution context for the target range
+            targetPartitionQueryExecutionContextList.append(
+                self._createTargetPartitionQueryExecutionContext(partitionTargetRange)
+            )
+
+        self._doc_producers = []
+        for targetQueryExContext in targetPartitionQueryExecutionContextList:
+            try:
+                await targetQueryExContext.peek()
+                # if there are matching results in the target ex range add it to the priority queue
+                self._doc_producers.append(targetQueryExContext)
+
+            except StopAsyncIteration:
+                continue
+
+    def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range):
+
+        rewritten_query = self._partitioned_query_ex_info.get_rewritten_query()
+        if rewritten_query:
+            if isinstance(self._query, dict):
+                # this is a parameterized query, collect all the parameters
+                query = dict(self._query)
+                query["query"] = rewritten_query
+            else:
+                query = rewritten_query
+        else:
+            query = self._query
+
+        return document_producer._DocumentProducer(
+            partition_key_target_range,
+            self._client,
+            self._resource_link,
+            query,
+            self._document_producer_comparator,
+            self._options,
+        )
+
+    async def _get_target_partition_key_range(self):
+
+        query_ranges = self._partitioned_query_ex_info.get_query_ranges()
+        return await self._routing_provider.get_overlapping_ranges(
+            self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
+        )
+
+    async def _configure_partition_ranges(self):
+        # will be a list of (partition_min, partition_max) tuples
+        targetPartitionRanges = await self._get_target_partition_key_range()
+
+        self._document_producer_comparator = document_producer._NonStreamingOrderByComparator(self._sort_orders)
+
+        targetPartitionQueryExecutionContextList = []
+        for partitionTargetRange in targetPartitionRanges:
+            # create and add the child execution context for the target range
+            targetPartitionQueryExecutionContextList.append(
+                self._createTargetPartitionQueryExecutionContext(partitionTargetRange)
+            )
+
+        self._doc_producers = []
+        for targetQueryExContext in targetPartitionQueryExecutionContextList:
+            try:
+                await targetQueryExContext.peek()
+                self._doc_producers.append(targetQueryExContext)
+            except exceptions.CosmosHttpResponseError as e:
+                if exceptions._partition_range_is_gone(e):
+                    # repairing document producer context on partition split
+                    await self._repair_document_producer()
+                else:
+                    raise
+
+            except StopAsyncIteration:
+                continue
+
+        pq_size = self._partitioned_query_ex_info.get_top() or self._partitioned_query_ex_info.get_limit()
+        self._orderByPQ = FixedSizePriorityQueue(pq_size)
+        for doc_producer in self._doc_producers:
+            while True:
+                try:
+                    result = await doc_producer.peek()
+                    item_result = document_producer._NonStreamingDocumentProducer(result, self._sort_orders)
+                    await self._orderByPQ.push_async(item_result, self._document_producer_comparator)
+                    await doc_producer.__anext__()
+                except StopAsyncIteration:
+                    break

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

@@ -326,7 +326,18 @@ def compare(self, doc_producer1, doc_producer2):
                 negative integer if doc_producers1 < doc_producers2
         :rtype: int
         """
-        # TODO: this is not fully safe - doesn't deal with scenario of having orderByItems of [{}]
-        rank1 = doc_producer1._item_result["orderByItems"][0]['item']
-        rank2 = doc_producer2._item_result["orderByItems"][0]['item']
-        return _compare_helper(rank1, rank2)
+        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/non_streaming_order_by_aggregator.py

@@ -131,7 +131,6 @@ def _repair_document_producer(self):
         self._doc_producers = []
         for targetQueryExContext in targetPartitionQueryExecutionContextList:
             try:
-                # TODO: we can also use more_itertools.peekable to be more python friendly
                 targetQueryExContext.peek()
                 # if there are matching results in the target ex range add it to the priority queue
                 self._doc_producers.append(targetQueryExContext)

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

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