update comments

Author: Minh Le (from Dev Box)

Microsoft.Azure.Cosmos/src/Linq/ExpressionToSQL.cs

@@ -1763,7 +1763,7 @@ private static Collection VisitGroupBy(Type returnElementType, ReadOnlyCollectio
 
                 case ExpressionType.New:
                     // TODO: Multi Value Selector
-                    throw new DocumentQueryException(string.Format(CultureInfo.CurrentCulture, ClientResources.ExpressionTypeIsNotSupported, "Multiple value selector"));
+                    throw new DocumentQueryException(string.Format(CultureInfo.CurrentCulture, ClientResources.ExpressionTypeIsNotSupported, ExpressionType.New));
 
                 default:
                     throw new DocumentQueryException(string.Format(CultureInfo.CurrentCulture, ClientResources.ExpressionTypeIsNotSupported, valueSelectorExpression.NodeType));
@@ -1963,17 +1963,12 @@ private static SqlSelectClause VisitAggregateFunction(
                 context.PopParameter();
             }
             else if (arguments.Count == 2)
-            {
-                if (context.CurrentQuery.GroupByParameter != null)
-                {
-                    LambdaExpression lambda = Utilities.GetLambda(arguments[1]);
-                    aggregateExpression = ExpressionToSql.VisitNonSubqueryScalarLambda(lambda, context);
-                }
-                else
-                {
-                    LambdaExpression lambda = Utilities.GetLambda(arguments[1]);
-                    aggregateExpression = ExpressionToSql.VisitScalarExpression(lambda, context);
-                }
+            {
+                LambdaExpression lambda = Utilities.GetLambda(arguments[1]);
+                    
+                aggregateExpression = context.CurrentQuery.GroupByParameter != null
+                    ? ExpressionToSql.VisitNonSubqueryScalarLambda(lambda, context)
+                    : ExpressionToSql.VisitScalarExpression(lambda, context);
             }
             else
             {

Microsoft.Azure.Cosmos/src/Linq/QueryUnderConstruction.cs

@@ -251,7 +251,7 @@ public QueryUnderConstruction FlattenAsPossible()
                 }
 
                 // In case of Select -> Group by cases, the Select query should not be flattened and kept as a subquery
-                if (query.inputQuery != null && query.inputQuery.selectClause != null && query.groupByClause != null)
+                if ((query.inputQuery?.selectClause != null) && (query.groupByClause != null))
                 {
                     flattenQuery = this;
                     break;
@@ -302,7 +302,7 @@ private QueryUnderConstruction Flatten()
                 {
                     // If selectClause doesn't exists, use SELECT v0 where v0 is the input parameter, instead of SELECT *.
                     // If there is a groupby clause, the input parameter comes from the groupBy binding instead of the from clause binding
-                    string parameterName = this.GroupByParameter != null ? this.GroupByParameter.GetInputParameter().Name : this.FromParameters.GetInputParameter().Name;
+                    string parameterName = (this.GroupByParameter ?? this.FromParameters).GetInputParameter().Name;
                     SqlScalarExpression parameterExpression = SqlPropertyRefScalarExpression.Create(null, SqlIdentifier.Create(parameterName));
                     this.selectClause = SqlSelectClause.Create(SqlSelectValueSpec.Create(parameterExpression));
                 }
@@ -559,14 +559,16 @@ public bool ShouldBeOnNewQuery(string methodName, int argumentCount)
                     // New query is needed when there is already a Take or a non-distinct Select
                     shouldPackage = (this.topSpec != null) ||
                         (this.offsetSpec != null) ||
-                        (this.selectClause != null && !this.selectClause.HasDistinct) || this.groupByClause != null;
+                        (this.selectClause != null && !this.selectClause.HasDistinct) || 
+                        (this.groupByClause != null);
                     break;
 
                 case LinqMethods.GroupBy:
                     // New query is needed when there is already a Take or a Select or a Group by clause
                     shouldPackage = (this.topSpec != null) ||
                         (this.offsetSpec != null) ||
-                        (this.selectClause != null) || this.groupByClause != null;
+                        (this.selectClause != null) || 
+                        (this.groupByClause != null);
                     break;
 
                 case LinqMethods.Skip:

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Linq/LinqGeneralBaselineTests.cs

@@ -700,7 +700,9 @@ public void TestSubquery()
         [TestMethod]
         public void TestGroupByTranslation()
         {
-            List<LinqTestInput> inputs = new List<LinqTestInput>();
+            List<LinqTestInput> inputs = new List<LinqTestInput>();
+            inputs.Add(new LinqTestInput("GroupBy Single Value Select Key", b => getQuery(b).GroupBy(k => k /*keySelector*/,
+                                                                                (key, values) => key /*return the group by key */)));
             inputs.Add(new LinqTestInput("GroupBy Single Value Select Key", b => getQuery(b).GroupBy(k => k.Id /*keySelector*/,
                                                                                 (key, values) => key /*return the group by key */)));
             inputs.Add(new LinqTestInput("GroupBy Single Value Select Key Alias", b => getQuery(b).GroupBy(k => k.Id /*keySelector*/,
@@ -713,7 +715,7 @@ public void TestGroupByTranslation()
                                                                                 (key, values) => values.Max(value => value.Int) /*return the Max of each group */)));
             inputs.Add(new LinqTestInput("GroupBy Single Value With Count", b => getQuery(b).GroupBy(k => k.Id /*keySelector*/,
                                                                                 (key, values) => values.Count() /*return the Count of each group */)));
-            inputs.Add(new LinqTestInput("GroupBy Single Value With Count", b => getQuery(b).GroupBy(k => k.Id /*keySelector*/,
+            inputs.Add(new LinqTestInput("GroupBy Single Value With Average", b => getQuery(b).GroupBy(k => k.Id /*keySelector*/,
                                                                                 (key, values) => values.Average(value => value.Int) /*return the Count of each group */)));
 
             // Negative cases