Preserve Order of Projections in GROUP BY Query

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

@@ -21,14 +21,18 @@ private sealed class EagerCosmosObject : CosmosObject
         {
             private readonly Dictionary<string, CosmosElement> dictionary;
 
-            public EagerCosmosObject(IDictionary<string, CosmosElement> dictionary)
+            public EagerCosmosObject(IReadOnlyDictionary<string, CosmosElement> dictionary)
             {
                 if (dictionary == null)
                 {
                     throw new ArgumentNullException($"{nameof(dictionary)}");
                 }
 
-                this.dictionary = new Dictionary<string, CosmosElement>(dictionary);
+                this.dictionary = new Dictionary<string, CosmosElement>();
+                foreach (KeyValuePair<string, CosmosElement> kvp in dictionary)
+                {
+                    this.dictionary[kvp.Key] = kvp.Value;
+                }
             }
 
             public override IEnumerable<string> Keys => this.dictionary.Keys;
@@ -37,7 +41,18 @@ public EagerCosmosObject(IDictionary<string, CosmosElement> dictionary)
 
             public override int Count => this.dictionary.Count;
 
-            public override CosmosElement this[string key] => this.dictionary[key];
+            public override CosmosElement this[string key]
+            {
+                get
+                {
+                    if (!this.TryGetValue(key, out CosmosElement value))
+                    {
+                        value = null;
+                    }
+
+                    return value;
+                }
+            }
 
             public override bool ContainsKey(string key) => this.dictionary.ContainsKey(key);
 

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

@@ -0,0 +1,111 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+namespace Microsoft.Azure.Cosmos.CosmosElements
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Linq;
+    using Microsoft.Azure.Cosmos.Json;
+    using Newtonsoft.Json;
+
+#if INTERNAL
+#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
+#pragma warning disable SA1601 // Partial elements should be documented
+    public
+#else
+    internal
+#endif
+    abstract partial class CosmosObject : CosmosElement, IReadOnlyDictionary<string, CosmosElement>
+    {
+        private sealed class OrderedCosmosObject : CosmosObject
+        {
+            private readonly List<KeyValuePair<string, CosmosElement>> properties;
+
+            public OrderedCosmosObject(IReadOnlyList<KeyValuePair<string, CosmosElement>> properties)
+            {
+                if (properties == null)
+                {
+                    throw new ArgumentNullException(nameof(properties));
+                }
+
+                this.properties = new List<KeyValuePair<string, CosmosElement>>(properties.Count);
+                foreach (KeyValuePair<string, CosmosElement> kvp in properties)
+                {
+                    if (this.properties.Any((property) => property.Key == kvp.Key))
+                    {
+                        throw new ArgumentException($"Duplicate key: {kvp.Key}");
+                    }
+
+                    this.properties.Add(kvp);
+                }
+            }
+
+            public override CosmosElement this[string key]
+            {
+                get
+                {
+                    if (!this.TryGetValue(key, out CosmosElement value))
+                    {
+                        value = null;
+                    }
+
+                    return value;
+                }
+            }
+
+            public override IEnumerable<string> Keys => this.properties.Select(kvp => kvp.Key);
+
+            public override IEnumerable<CosmosElement> Values => this.properties.Select(kvp => kvp.Value);
+
+            public override int Count => this.properties.Count;
+
+            public override bool ContainsKey(string key)
+            {
+                return this.TryGetValue(key, out CosmosElement unused);
+            }
+
+            public override IEnumerator<KeyValuePair<string, CosmosElement>> GetEnumerator()
+            {
+                return this.properties.GetEnumerator();
+            }
+
+            public override bool TryGetValue(string key, out CosmosElement 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)
+            {
+                if (jsonWriter == null)
+                {
+                    throw new ArgumentNullException($"{nameof(jsonWriter)}");
+                }
+
+                jsonWriter.WriteObjectStart();
+
+                foreach (KeyValuePair<string, CosmosElement> kvp in this.properties)
+                {
+                    jsonWriter.WriteFieldName(kvp.Key);
+                    kvp.Value.WriteTo(jsonWriter);
+                }
+
+                jsonWriter.WriteObjectEnd();
+            }
+        }
+    }
+#if INTERNAL
+#pragma warning restore SA1601 // Partial elements should be documented
+#pragma warning restore CS1591 // Missing XML comment for publicly visible type or member
+#endif
+}

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

@@ -49,11 +49,16 @@ public static CosmosObject Create(
             return new LazyCosmosObject(jsonNavigator, jsonNavigatorNode);
         }
 
-        public static CosmosObject Create(IDictionary<string, CosmosElement> dictionary)
+        public static CosmosObject Create(IReadOnlyDictionary<string, CosmosElement> dictionary)
         {
             return new EagerCosmosObject(dictionary);
         }
 
+        public static CosmosObject Create(IReadOnlyList<KeyValuePair<string, CosmosElement>> properties)
+        {
+            return new OrderedCosmosObject(properties);
+        }
+
         public abstract bool ContainsKey(string key);
 
         public abstract bool TryGetValue(string key, out CosmosElement value);

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

@@ -29,6 +29,7 @@ internal abstract class SingleGroupAggregator
         public static SingleGroupAggregator Create(
             AggregateOperator[] aggregates,
             IReadOnlyDictionary<string, AggregateOperator?> aggregateAliasToAggregateType,
+            IReadOnlyList<string> orderedAliases,
             bool hasSelectValue)
         {
             SingleGroupAggregator aggregateValues;
@@ -47,7 +48,7 @@ public static SingleGroupAggregator Create(
             }
             else
             {
-                aggregateValues = SelectListAggregateValues.Create(aggregateAliasToAggregateType);
+                aggregateValues = SelectListAggregateValues.Create(aggregateAliasToAggregateType, orderedAliases);
             }
 
             return aggregateValues;
@@ -101,29 +102,45 @@ 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);
                     }
                 }
 
                 return CosmosObject.Create(aliasToElement);
             }
 
-            public static SelectListAggregateValues Create(IReadOnlyDictionary<string, AggregateOperator?> aggregateAliasToAggregateType)
+            public static SelectListAggregateValues Create(
+                IReadOnlyDictionary<string, AggregateOperator?> aggregateAliasToAggregateType,
+                IReadOnlyList<string> orderedAliases)
             {
                 Dictionary<string, AggregateValue> groupingTable = new Dictionary<string, AggregateValue>();
                 foreach (KeyValuePair<string, AggregateOperator?> aliasToAggregate in aggregateAliasToAggregateType)
@@ -133,7 +150,7 @@ public static SelectListAggregateValues Create(IReadOnlyDictionary<string, Aggre
                     groupingTable[alias] = AggregateValue.Create(aggregateOperator);
                 }
 
-                return new SelectListAggregateValues(groupingTable);
+                return new SelectListAggregateValues(groupingTable, orderedAliases);
             }
 
             public override void AddValues(CosmosElement values)

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

@@ -66,20 +66,22 @@ private AggregateDocumentQueryExecutionComponent(
         /// </summary>
         /// <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>
         /// <returns>The AggregateDocumentQueryExecutionComponent.</returns>
         public static async Task<AggregateDocumentQueryExecutionComponent> CreateAsync(
             AggregateOperator[] aggregates,
             IReadOnlyDictionary<string, AggregateOperator?> aliasToAggregateType,
+            IReadOnlyList<string> orderedAliases,
             bool hasSelectValue,
             string requestContinuation,
             Func<string, Task<IDocumentQueryExecutionComponent>> createSourceCallback)
         {
             return new AggregateDocumentQueryExecutionComponent(
                 await createSourceCallback(requestContinuation),
-                SingleGroupAggregator.Create(aggregates, aliasToAggregateType, hasSelectValue),
+                SingleGroupAggregator.Create(aggregates, aliasToAggregateType, orderedAliases, hasSelectValue),
                 aggregates != null && aggregates.Count() == 1);
         }
 

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

@@ -45,6 +45,7 @@ internal sealed class GroupByDocumentQueryExecutionComponent : DocumentQueryExec
         private static readonly AggregateOperator[] EmptyAggregateOperators = new AggregateOperator[] { };
 
         private readonly IReadOnlyDictionary<string, AggregateOperator?> groupByAliasToAggregateType;
+        private readonly IReadOnlyList<string> orderedAliases;
         private readonly Dictionary<UInt192, SingleGroupAggregator> groupingTable;
         private readonly bool hasSelectValue;
 
@@ -53,6 +54,7 @@ internal sealed class GroupByDocumentQueryExecutionComponent : DocumentQueryExec
 
         private GroupByDocumentQueryExecutionComponent(
             IReadOnlyDictionary<string, AggregateOperator?> groupByAliasToAggregateType,
+            IReadOnlyList<string> orderedAliases,
             bool hasSelectValue,
             IDocumentQueryExecutionComponent source)
             : base(source)
@@ -62,8 +64,14 @@ private GroupByDocumentQueryExecutionComponent(
                 throw new ArgumentNullException(nameof(groupByAliasToAggregateType));
             }
 
+            if (orderedAliases == null)
+            {
+                throw new ArgumentNullException(nameof(orderedAliases));
+            }
+
             this.groupingTable = new Dictionary<UInt192, SingleGroupAggregator>();
             this.groupByAliasToAggregateType = groupByAliasToAggregateType;
+            this.orderedAliases = orderedAliases;
             this.hasSelectValue = hasSelectValue;
         }
 
@@ -73,11 +81,13 @@ public static async Task<IDocumentQueryExecutionComponent> CreateAsync(
             string requestContinuation,
             Func<string, Task<IDocumentQueryExecutionComponent>> createSourceCallback,
             IReadOnlyDictionary<string, AggregateOperator?> groupByAliasToAggregateType,
+            IReadOnlyList<string> orderedAliases,
             bool hasSelectValue)
         {
             // We do not support continuation tokens for GROUP BY.
             return new GroupByDocumentQueryExecutionComponent(
                 groupByAliasToAggregateType,
+                orderedAliases,
                 hasSelectValue,
                 await createSourceCallback(requestContinuation));
         }
@@ -111,6 +121,7 @@ public override async Task<QueryResponseCore> DrainAsync(
                         singleGroupAggregator = SingleGroupAggregator.Create(
                             EmptyAggregateOperators,
                             this.groupByAliasToAggregateType,
+                            this.orderedAliases,
                             this.hasSelectValue);
                         this.groupingTable[groupByKeysHash] = singleGroupAggregator;
                     }

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs

@@ -220,6 +220,7 @@ private static async Task<PipelinedDocumentQueryExecutionContext> CreateHelperAs
                     return await AggregateDocumentQueryExecutionComponent.CreateAsync(
                         queryInfo.Aggregates,
                         queryInfo.GroupByAliasToAggregateType,
+                        queryInfo.GroupByAliases,
                         queryInfo.HasSelectValue,
                         continuationToken,
                         createSourceCallback);
@@ -248,6 +249,7 @@ private static async Task<PipelinedDocumentQueryExecutionContext> CreateHelperAs
                         continuationToken,
                         createSourceCallback,
                         queryInfo.GroupByAliasToAggregateType,
+                        queryInfo.GroupByAliases,
                         queryInfo.HasSelectValue);
                 };
             }

Microsoft.Azure.Cosmos/src/Query/Core/QueryInfo.cs

@@ -62,6 +62,13 @@ public string[] GroupByExpressions
             set;
         }
 
+        [JsonProperty("groupByAliases")]
+        public string[] GroupByAliases
+        {
+            get;
+            set;
+        }
+
         [JsonProperty("aggregates", ItemConverterType = typeof(StringEnumConverter))]
         public AggregateOperator[] Aggregates
         {