From 180f2767efb6a9a683aef6ca71fde5fcb0ec0d99 Mon Sep 17 00:00:00 2001
From: "REDMOND\\dmmelnik" <dmitri.melnikov@microsoft.com>
Date: Thu, 10 Sep 2020 08:52:22 -0700
Subject: [PATCH] Allocate on heap for a larger arrays

---
 .../Core/QueryPlan/QueryPartitionProvider.cs  | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs
index 4eca70d18c..5c5b009969 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs
+++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs
@@ -168,24 +168,23 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
             List<string> paths = new List<string>(partitionKeyDefinition.Paths);
             List<IReadOnlyList<string>> pathPartsList = new List<IReadOnlyList<string>>(paths.Count);
             uint[] partsLengths = new uint[paths.Count];
-            int i = 0;
             int allPartsLength = 0;
 
-            foreach (string path in paths)
+            for (int i = 0; i < paths.Count; i++)
             {
-                IReadOnlyList<string> pathParts = PathParser.GetPathParts(path);
-                partsLengths[i++] = (uint)pathParts.Count;
+                IReadOnlyList<string> pathParts = PathParser.GetPathParts(paths[i]);
+                partsLengths[i] = (uint)pathParts.Count;
                 pathPartsList.Add(pathParts);
                 allPartsLength += pathParts.Count;
             }
 
             string[] allParts = new string[allPartsLength];
-            i = 0;
+            int allPartsPointer = 0;
             foreach (IReadOnlyList<string> pathParts in pathPartsList)
             {
                 foreach (string part in pathParts)
                 {
-                    allParts[i++] = part;
+                    allParts[allPartsPointer++] = part;
                 }
             }
 
@@ -193,7 +192,7 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
 
             this.Initialize();
 
-            Span<byte> buffer = stackalloc byte[InitialBufferSize];
+            Span<byte> buffer = stackalloc byte[QueryPartitionProvider.InitialBufferSize];
             uint errorCode;
             uint serializedQueryExecutionInfoResultLength;
 
@@ -218,7 +217,11 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
 
                     if (errorCode == DISP_E_BUFFERTOOSMALL)
                     {
-                        buffer = stackalloc byte[(int)serializedQueryExecutionInfoResultLength];
+                        // Allocate on stack for smaller arrays, otherwise use heap.
+                        buffer = serializedQueryExecutionInfoResultLength < 4096
+                            ? stackalloc byte[(int)serializedQueryExecutionInfoResultLength]
+                            : new byte[serializedQueryExecutionInfoResultLength];
+
                         fixed (byte* bytePtr2 = buffer)
                         {
                             errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery(