Support Count/LongCount after GroupBy().Select() (#4278)

src/LinqTests/Operators/group_by_operator.cs

@@ -222,4 +222,33 @@ public async Task group_by_with_long_count()
         results.Single(x => x.Color == Colors.Blue).Count.ShouldBe(2L);
         results.Single(x => x.Color == Colors.Green).Count.ShouldBe(3L);
     }
+
+    // https://github.com/JasperFx/marten/issues/4278
+    [Fact]
+    public async Task group_by_count()
+    {
+        await SetupTargetData();
+
+        var count = await _session.Query<Target>()
+            .GroupBy(x => x.Color)
+            .Select(x => x.Key)      // Select must always follow GroupBy
+            .CountAsync();
+
+        // Blue, Green, Red -> three distinct groups
+        count.ShouldBe(3);
+    }
+
+    // https://github.com/JasperFx/marten/issues/4278
+    [Fact]
+    public async Task group_by_long_count()
+    {
+        await SetupTargetData();
+
+        var count = await _session.Query<Target>()
+            .GroupBy(x => x.Color)
+            .Select(x => x.Key)
+            .LongCountAsync();
+
+        count.ShouldBe(3L);
+    }
 }

src/Marten/Linq/CollectionUsage.Compilation.cs

@@ -486,6 +486,19 @@ public Statement CompileGroupBy(IMartenSession session,
             }
         }
 
+        // Transfer single-value operators applied directly on the grouping usage
+        // itself (e.g., .GroupBy(...).Select(...).CountAsync() / .AnyAsync()).
+        // See https://github.com/JasperFx/marten/issues/4278.
+        if (groupingUsage.SingleValueMode.HasValue)
+        {
+            SingleValueMode ??= groupingUsage.SingleValueMode;
+        }
+
+        if (groupingUsage.IsAny)
+        {
+            IsAny = true;
+        }
+
         // Transfer downstream operators from the grouping usage's Inner (if any)
         // e.g., OrderBy, Take, Skip after Select
         var downstream = groupingUsage.Inner;
@@ -630,6 +643,22 @@ internal void ProcessSingleValueModeIfAny(SelectorStatement statement, IMartenSe
                         return;
                     }
 
+                    if (statement.GroupByColumns.Count > 0)
+                    {
+                        // .GroupBy(...).Select(...).CountAsync() should return the
+                        // number of groups, not count(*) over the grouped rows.
+                        // Wrap the GROUP BY query in a CTE and count its rows.
+                        // See https://github.com/JasperFx/marten/issues/4278.
+                        statement.ConvertToCommonTableExpression(session);
+                        var groupCount = new SelectorStatement
+                        {
+                            SelectClause = new CountClause<int>(statement.ExportName)
+                        };
+
+                        statement.AddToEnd(groupCount);
+                        return;
+                    }
+
                     statement.SelectClause = new CountClause<int>(statement.SelectClause.FromObject);
 
                     break;
@@ -652,6 +681,21 @@ internal void ProcessSingleValueModeIfAny(SelectorStatement statement, IMartenSe
                         return;
                     }
 
+                    if (statement.GroupByColumns.Count > 0)
+                    {
+                        // .GroupBy(...).Select(...).LongCountAsync() should return
+                        // the number of groups. See
+                        // https://github.com/JasperFx/marten/issues/4278.
+                        statement.ConvertToCommonTableExpression(session);
+                        var groupLongCount = new SelectorStatement
+                        {
+                            SelectClause = new CountClause<long>(statement.ExportName)
+                        };
+
+                        statement.AddToEnd(groupLongCount);
+                        return;
+                    }
+
                     statement.SelectClause = new CountClause<long>(statement.SelectClause.FromObject);
                     break;