Query: Fixes non streaming OrderByCrossPartitionQueryPipelineStage to remove state and handle splits (#4493)

* simplify OrderByCrossPartitionQueryPipelineStage to remove state and be a static class.

* Add more test coverage for splits/merge and request charge

* Add more test coverage for splits

* incorporate code review feedback

Author: neildsh

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/BufferedOrderByResults.cs

@@ -0,0 +1,32 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy
+{
+    using System;
+    using System.Collections.Generic;
+
+    internal sealed class BufferedOrderByResults
+    {
+        public IEnumerator<OrderByQueryResult> Enumerator { get; }
+
+        public int Count { get; }
+
+        public double TotalRequestCharge { get; }
+
+        public QueryPageParameters QueryPageParameters { get; }
+
+        public BufferedOrderByResults(
+            IEnumerator<OrderByQueryResult> enumerator,
+            int itemCount,
+            double totalRequestCharge,
+            QueryPageParameters queryPageParameters)
+        {
+            this.Enumerator = enumerator ?? throw new ArgumentNullException(nameof(enumerator));
+            this.Count = itemCount;
+            this.TotalRequestCharge = totalRequestCharge;
+            this.QueryPageParameters = queryPageParameters ?? throw new ArgumentNullException(nameof(queryPageParameters));
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionEnumerator.cs

@@ -2,17 +2,17 @@
 // Copyright (c) Microsoft Corporation.  All rights reserved.
 // ------------------------------------------------------------
 
-namespace Microsoft.Azure.Cosmos.Pagination
+namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy
 {
     using System;
     using System.Collections;
     using System.Collections.Generic;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
+    using Microsoft.Azure.Cosmos.Pagination;
     using Microsoft.Azure.Cosmos.Query.Core.Collections;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
-    using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy;
     using Microsoft.Azure.Cosmos.Tracing;
 
     internal sealed class OrderByCrossPartitionEnumerator : IEnumerator<OrderByQueryResult>
@@ -25,69 +25,86 @@ internal sealed class OrderByCrossPartitionEnumerator : IEnumerator<OrderByQuery
 
         object IEnumerator.Current => this.Current;
 
-        public OrderByCrossPartitionEnumerator(PriorityQueue<IEnumerator<OrderByQueryResult>> queue)
+        private OrderByCrossPartitionEnumerator(PriorityQueue<IEnumerator<OrderByQueryResult>> queue)
         {
             this.queue = queue ?? throw new ArgumentNullException(nameof(queue));
         }
 
-        public static async Task<(IEnumerator<OrderByQueryResult> orderbyQueryResultEnumerator, double totalRequestCharge)> CreateAsync(
-            IEnumerable<OrderByQueryPartitionRangePageAsyncEnumerator> enumerators,
+        public static async Task<BufferedOrderByResults> CreateAsync(
+            ITracingAsyncEnumerator<TryCatch<OrderByQueryPage>> enumerator,
             IComparer<OrderByQueryResult> comparer,
             int levelSize,
             ITrace trace,
             CancellationToken cancellationToken)
         {
-            if (enumerators == null)
+            if (enumerator == null)
             {
-                throw new ArgumentNullException(nameof(enumerators));
+                throw new ArgumentNullException(nameof(enumerator));
             }
 
             if (comparer == null)
             {
                 throw new ArgumentNullException(nameof(comparer));
             }
 
+            QueryPageParameters queryPageParameters = null;
             double totalRequestCharge = 0;
+            int bufferedItemCount = 0;
             EnumeratorComparer enumeratorComparer = new EnumeratorComparer(comparer);
             PriorityQueue<IEnumerator<OrderByQueryResult>> queue = new PriorityQueue<IEnumerator<OrderByQueryResult>>(enumeratorComparer);
-            foreach (ITracingAsyncEnumerator<TryCatch<OrderByQueryPage>> enumerator in enumerators)
+            while (await enumerator.MoveNextAsync(trace, cancellationToken))
             {
-                while (await enumerator.MoveNextAsync(trace, cancellationToken))
+                TryCatch<OrderByQueryPage> currentPage = enumerator.Current;
+                if (currentPage.Failed)
                 {
-                    TryCatch<OrderByQueryPage> currentPage = enumerator.Current;
-                    if (currentPage.Failed)
-                    {
-                        throw currentPage.Exception;
-                    }
+                    throw currentPage.Exception;
+                }
 
-                    totalRequestCharge += currentPage.Result.RequestCharge;
-                    IReadOnlyList<CosmosElement> page = currentPage.Result.Page.Documents;
+                if (queryPageParameters == null)
+                {
+                    queryPageParameters = new QueryPageParameters(
+                        activityId: currentPage.Result.ActivityId,
+                        cosmosQueryExecutionInfo: currentPage.Result.Page.CosmosQueryExecutionInfo,
+                        distributionPlanSpec: currentPage.Result.Page.DistributionPlanSpec,
+                        additionalHeaders: currentPage.Result.AdditionalHeaders);
+                }
 
-                    if (page.Count > 0)
-                    {
-                        PageEnumerator pageEnumerator = new PageEnumerator(page);
-                        pageEnumerator.MoveNext();
+                totalRequestCharge += currentPage.Result.RequestCharge;
+                IReadOnlyList<CosmosElement> page = currentPage.Result.Page.Documents;
+                bufferedItemCount += page.Count;
+
+                if (page.Count > 0)
+                {
+                    PageEnumerator pageEnumerator = new PageEnumerator(page);
+                    pageEnumerator.MoveNext();
 
-                        queue.Enqueue(pageEnumerator);
+                    queue.Enqueue(pageEnumerator);
 
-                        if (queue.Count >= levelSize)
-                        {
-                            OrderByCrossPartitionEnumerator newEnumerator = new OrderByCrossPartitionEnumerator(queue);
-                            newEnumerator.MoveNext();
+                    if (queue.Count >= levelSize)
+                    {
+                        OrderByCrossPartitionEnumerator newEnumerator = new OrderByCrossPartitionEnumerator(queue);
+                        newEnumerator.MoveNext();
 
-                            queue = new PriorityQueue<IEnumerator<OrderByQueryResult>>(enumeratorComparer);
-                            queue.Enqueue(newEnumerator);
-                        }
+                        queue = new PriorityQueue<IEnumerator<OrderByQueryResult>>(enumeratorComparer);
+                        queue.Enqueue(newEnumerator);
                     }
                 }
             }
 
             if (queue.Count == 0)
             {
-                return (EmptyEnumerator.Instance, totalRequestCharge);
+                return new BufferedOrderByResults(
+                    EmptyEnumerator.Instance,
+                    itemCount: 0,
+                    totalRequestCharge,
+                    queryPageParameters);
             }
 
-            return (new OrderByCrossPartitionEnumerator(queue), totalRequestCharge);
+            return new BufferedOrderByResults(
+                new OrderByCrossPartitionEnumerator(queue),
+                bufferedItemCount,
+                totalRequestCharge,
+                queryPageParameters);
         }
 
         public bool MoveNext()