Preserve Order of Projections in GROUP BY Query

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.EagerCosmosObject.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.CosmosElements
 {
     using System;
     using System.Collections.Generic;
+    using System.Linq;
     using Microsoft.Azure.Cosmos.Json;
 
 #if INTERNAL
@@ -18,33 +19,60 @@ abstract partial class CosmosObject : CosmosElement, IReadOnlyDictionary<string,
     {
         private sealed class EagerCosmosObject : CosmosObject
         {
-            private readonly Dictionary<string, CosmosElement> dictionary;
+            private readonly List<KeyValuePair<string, CosmosElement>> properties;
 
-            public EagerCosmosObject(IDictionary<string, CosmosElement> dictionary)
+            public EagerCosmosObject(IReadOnlyList<KeyValuePair<string, CosmosElement>> properties)
             {
-                if (dictionary == null)
+                if (properties == null)
                 {
-                    throw new ArgumentNullException($"{nameof(dictionary)}");
+                    throw new ArgumentNullException(nameof(properties));
                 }
 
-                this.dictionary = new Dictionary<string, CosmosElement>(dictionary);
+                this.properties = new List<KeyValuePair<string, CosmosElement>>(properties);
             }
 
-            public override IEnumerable<string> Keys => this.dictionary.Keys;
+            public override CosmosElement this[string key]
+            {
+                get
+                {
+                    if (!this.TryGetValue(key, out CosmosElement value))
+                    {
+                        throw new KeyNotFoundException($"Failed to find key: {key}");
+                    }
+
+                    return value;
+                }
+            }
 
-            public override IEnumerable<CosmosElement> Values => this.dictionary.Values;
+            public override IEnumerable<string> Keys => this.properties.Select(kvp => kvp.Key);
 
-            public override int Count => this.dictionary.Count;
+            public override IEnumerable<CosmosElement> Values => this.properties.Select(kvp => kvp.Value);
 
-            public override CosmosElement this[string key] => this.dictionary[key];
+            public override int Count => this.properties.Count;
 
-            public override bool ContainsKey(string key) => this.dictionary.ContainsKey(key);
+            public override bool ContainsKey(string key)
+            {
+                return this.TryGetValue(key, out CosmosElement unused);
+            }
 
-            public override IEnumerator<KeyValuePair<string, CosmosElement>> GetEnumerator() => this.dictionary.GetEnumerator();
+            public override IEnumerator<KeyValuePair<string, CosmosElement>> GetEnumerator()
+            {
+                return this.properties.GetEnumerator();
+            }
 
             public override bool TryGetValue(string key, out CosmosElement value)
             {
-                return this.dictionary.TryGetValue(key, out value);
+                foreach (KeyValuePair<string, CosmosElement> property in this.properties)
+                {
+                    if (property.Key == key)
+                    {
+                        value = property.Value;
+                        return true;
+                    }
+                }
+
+                value = null;
+                return false;
             }
 
             public override void WriteTo(IJsonWriter jsonWriter)
@@ -56,7 +84,7 @@ public override void WriteTo(IJsonWriter jsonWriter)
 
                 jsonWriter.WriteObjectStart();
 
-                foreach (KeyValuePair<string, CosmosElement> kvp in this)
+                foreach (KeyValuePair<string, CosmosElement> kvp in this.properties)
                 {
                     jsonWriter.WriteFieldName(kvp.Key);
                     kvp.Value.WriteTo(jsonWriter);

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.LazyCosmosObject.cs

@@ -64,7 +64,7 @@ public override CosmosElement this[string key]
                 {
                     if (!this.TryGetValue(key, out CosmosElement value))
                     {
-                        value = null;
+                        throw new KeyNotFoundException($"Failed to find key: {key}");
                     }
 
                     return value;

Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.CosmosElements
 {
     using System.Collections;
     using System.Collections.Generic;
+    using System.Linq;
     using Microsoft.Azure.Cosmos.Json;
 
 #if INTERNAL
@@ -49,31 +50,36 @@ public static CosmosObject Create(
             return new LazyCosmosObject(jsonNavigator, jsonNavigatorNode);
         }
 
-        public static CosmosObject Create(IDictionary<string, CosmosElement> dictionary)
+        public static CosmosObject Create(Dictionary<string, CosmosElement> dictionary)
         {
-            return new EagerCosmosObject(dictionary);
+            return new EagerCosmosObject(dictionary.ToList());
+        }
+
+        public static CosmosObject Create(IReadOnlyList<KeyValuePair<string, CosmosElement>> properties)
+        {
+            return new EagerCosmosObject(properties);
         }
 
         public abstract bool ContainsKey(string key);
 
         public abstract bool TryGetValue(string key, out CosmosElement value);
 
-        public bool TryGetValue<TCosmosElement>(string key, out TCosmosElement typedCosmosElement)
+        public bool TryGetValue<TCosmosElement>(string key, out TCosmosElement value)
             where TCosmosElement : CosmosElement
         {
-            if (!this.TryGetValue(key, out CosmosElement cosmosElement))
+            if (!this.TryGetValue(key, out CosmosElement untypedCosmosElement))
             {
-                typedCosmosElement = default(TCosmosElement);
+                value = default;
                 return false;
             }
 
-            if (!(cosmosElement is TCosmosElement tCosmosElement))
+            if (!(untypedCosmosElement is TCosmosElement typedCosmosElement))
             {
-                typedCosmosElement = default(TCosmosElement);
+                value = default;
                 return false;
             }
 
-            typedCosmosElement = tCosmosElement;
+            value = typedCosmosElement;
             return true;
         }
 

Microsoft.Azure.Cosmos/src/Query/Core/Aggregation/MinMaxAggregator.cs

@@ -53,7 +53,7 @@ public void Aggregate(CosmosElement localMinMax)
             // and unwrap the object to get the actual item of interest
             if (localMinMax is CosmosObject cosmosObject)
             {
-                if (cosmosObject["count"] is CosmosNumber countToken)
+                if (cosmosObject.TryGetValue("count", out CosmosNumber countToken))
                 {
                     // We know the object looks like: {"min": MIN(c.blah), "count": COUNT(c.blah)}
                     long count;
@@ -72,11 +72,16 @@ public void Aggregate(CosmosElement localMinMax)
                         return;
                     }
 
-                    CosmosElement min = cosmosObject["min"];
-                    CosmosElement max = cosmosObject["max"];
+                    if (!cosmosObject.TryGetValue("min", out CosmosElement min))
+                    {
+                        min = null;
+                    }
 
-                    // Note that JToken won't equal null as long as a value is there
-                    // even if that value is a JSON null.
+                    if (!cosmosObject.TryGetValue("max", out CosmosElement max))
+                    {
+                        max = null;
+                    }
+
                     if (min != null)
                     {
                         localMinMax = min;

Microsoft.Azure.Cosmos/src/Query/Core/Aggregation/SingleGroupAggregator.cs

@@ -32,6 +32,7 @@ public static SingleGroupAggregator Create(
             CosmosQueryClient queryClient,
             AggregateOperator[] aggregates,
             IReadOnlyDictionary<string, AggregateOperator?> aggregateAliasToAggregateType,
+            IReadOnlyList<string> orderedAliases,
             bool hasSelectValue,
             string continuationToken)
         {
@@ -46,12 +47,12 @@ public static SingleGroupAggregator Create(
                 else
                 {
                     // SELECT VALUE <NON AGGREGATE>
-                    aggregateValues = SelectValueAggregateValues.Create(aggregateOperator: null, continuationToken: null);
+                    aggregateValues = SelectValueAggregateValues.Create(aggregateOperator: null, continuationToken: continuationToken);
                 }
             }
             else
             {
-                aggregateValues = SelectListAggregateValues.Create(queryClient, aggregateAliasToAggregateType, continuationToken);
+                aggregateValues = SelectListAggregateValues.Create(queryClient, aggregateAliasToAggregateType, orderedAliases, continuationToken);
             }
 
             return aggregateValues;
@@ -110,22 +111,36 @@ public override string ToString()
         private sealed class SelectListAggregateValues : SingleGroupAggregator
         {
             private readonly IReadOnlyDictionary<string, AggregateValue> aliasToValue;
+            private readonly IReadOnlyList<string> orderedAliases;
 
-            private SelectListAggregateValues(IReadOnlyDictionary<string, AggregateValue> aliasToValue)
+            private SelectListAggregateValues(
+                IReadOnlyDictionary<string, AggregateValue> aliasToValue,
+                IReadOnlyList<string> orderedAliases)
             {
+                if (aliasToValue == null)
+                {
+                    throw new ArgumentNullException(nameof(aliasToValue));
+                }
+
+                if (orderedAliases == null)
+                {
+                    throw new ArgumentNullException(nameof(orderedAliases));
+                }
+
                 this.aliasToValue = aliasToValue;
+                this.orderedAliases = orderedAliases;
             }
 
             public override CosmosElement GetResult()
             {
-                Dictionary<string, CosmosElement> aliasToElement = new Dictionary<string, CosmosElement>();
-                foreach (KeyValuePair<string, AggregateValue> aliasAndValue in this.aliasToValue)
+                List<KeyValuePair<string, CosmosElement>> aliasToElement = new List<KeyValuePair<string, CosmosElement>>();
+                foreach (string alias in this.orderedAliases)
                 {
-                    string alias = aliasAndValue.Key;
-                    AggregateValue aggregateValue = aliasAndValue.Value;
+                    AggregateValue aggregateValue = this.aliasToValue[alias];
                     if (aggregateValue.Result != null)
                     {
-                        aliasToElement[alias] = aggregateValue.Result;
+                        KeyValuePair<string, CosmosElement> kvp = new KeyValuePair<string, CosmosElement>(alias, aggregateValue.Result);
+                        aliasToElement.Add(kvp);
                     }
                 }
 
@@ -147,12 +162,13 @@ public override string GetContinuationToken()
             public static SelectListAggregateValues Create(
                 CosmosQueryClient cosmosQueryClient,
                 IReadOnlyDictionary<string, AggregateOperator?> aggregateAliasToAggregateType,
+                IReadOnlyList<string> orderedAliases,
                 string continuationToken)
             {
                 CosmosObject aliasToContinuationToken;
                 if (continuationToken != null)
                 {
-                    if (!CosmosElement.TryParse<CosmosObject>(continuationToken, out aliasToContinuationToken))
+                    if (!CosmosElement.TryParse(continuationToken, out aliasToContinuationToken))
                     {
                         throw cosmosQueryClient.CreateBadRequestException(
                             $"{nameof(SelectListAggregateValues)} continuation token is malformed: {continuationToken}.");
@@ -187,7 +203,7 @@ public static SelectListAggregateValues Create(
                     groupingTable[alias] = AggregateValue.Create(aggregateOperator, aliasContinuationToken);
                 }
 
-                return new SelectListAggregateValues(groupingTable);
+                return new SelectListAggregateValues(groupingTable, orderedAliases);
             }
 
             public override void AddValues(CosmosElement values)
@@ -201,7 +217,12 @@ public override void AddValues(CosmosElement values)
                 {
                     string alias = aliasAndValue.Key;
                     AggregateValue aggregateValue = aliasAndValue.Value;
-                    aggregateValue.AddValue(payload[alias]);
+                    if (!payload.TryGetValue(alias, out CosmosElement payloadValue))
+                    {
+                        payloadValue = null;
+                    }
+
+                    aggregateValue.AddValue(payloadValue);
                 }
             }
 

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/AggregateDocumentQueryExecutionComponent.Client.cs

@@ -26,6 +26,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
                 CosmosQueryClient queryClient,
                 AggregateOperator[] aggregates,
                 IReadOnlyDictionary<string, AggregateOperator?> aliasToAggregateType,
+                IReadOnlyList<string> orderedAliases,
                 bool hasSelectValue,
                 string requestContinuation,
                 Func<string, Task<IDocumentQueryExecutionComponent>> createSourceCallback)
@@ -35,6 +36,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
                     queryClient,
                     aggregates,
                     aliasToAggregateType,
+                    orderedAliases,
                     hasSelectValue,
                     continuationToken: null);
 

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/AggregateDocumentQueryExecutionComponent.Compute.cs

@@ -29,6 +29,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
                 CosmosQueryClient queryClient,
                 AggregateOperator[] aggregates,
                 IReadOnlyDictionary<string, AggregateOperator?> aliasToAggregateType,
+                IReadOnlyList<string> orderedAliases,
                 bool hasSelectValue,
                 string requestContinuation,
                 Func<string, Task<IDocumentQueryExecutionComponent>> createSourceCallback)
@@ -56,6 +57,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
                     queryClient,
                     aggregates,
                     aliasToAggregateType,
+                    orderedAliases,
                     hasSelectValue,
                     singleGroupAggregatorContinuationToken);
 

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/AggregateDocumentQueryExecutionComponent.cs

@@ -67,6 +67,7 @@ protected AggregateDocumentQueryExecutionComponent(
         /// <param name="queryClient">The query client.</param>
         /// <param name="aggregates">The aggregates.</param>
         /// <param name="aliasToAggregateType">The alias to aggregate type.</param>
+        /// <param name="orderedAliases">The ordering of the aliases.</param>
         /// <param name="hasSelectValue">Whether or not the query has the 'VALUE' keyword.</param>
         /// <param name="requestContinuation">The continuation token to resume from.</param>
         /// <param name="createSourceCallback">The callback to create the source component that supplies the local aggregates.</param>
@@ -76,6 +77,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
             CosmosQueryClient queryClient,
             AggregateOperator[] aggregates,
             IReadOnlyDictionary<string, AggregateOperator?> aliasToAggregateType,
+            IReadOnlyList<string> orderedAliases,
             bool hasSelectValue,
             string requestContinuation,
             Func<string, Task<IDocumentQueryExecutionComponent>> createSourceCallback)
@@ -88,6 +90,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
                         queryClient,
                         aggregates,
                         aliasToAggregateType,
+                        orderedAliases,
                         hasSelectValue,
                         requestContinuation,
                         createSourceCallback);
@@ -98,6 +101,7 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
                         queryClient,
                         aggregates,
                         aliasToAggregateType,
+                        orderedAliases,
                         hasSelectValue,
                         requestContinuation,
                         createSourceCallback);