From 22fe728ae9572a07d6dd37d7e35ce081483fad5f Mon Sep 17 00:00:00 2001
From: "Jeremy D. Miller" <jeremydmiller@yahoo.com>
Date: Wed, 22 Apr 2026 17:15:32 -0500
Subject: [PATCH] Support Count/LongCount after GroupBy().Select()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

.GroupBy(...).Select(...).CountAsync() was throwing
"Marten does not know how to use result type int" because the
SingleValueMode set by CountAsync() on the grouping usage was never
transferred to the outer CollectionUsage, and the Count path didn't
know how to count groups instead of underlying rows.

Two small fixes in CompileGroupBy / ProcessSingleValueModeIfAny:

- Transfer SingleValueMode and IsAny from the grouping usage itself
  (not just from groupingUsage.Inner), so terminal aggregates applied
  directly after the Select projection are seen.
- When Count/LongCount is applied and the statement already has
  GROUP BY columns, wrap the query in a CTE and count its rows — the
  same pattern the Distinct path already uses. Replacing the SELECT
  with count(*) would have counted rows within each group instead of
  counting the groups themselves.

Fixes #4278.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 src/LinqTests/Operators/group_by_operator.cs  | 29 ++++++++++++
 .../Linq/CollectionUsage.Compilation.cs       | 44 +++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/src/LinqTests/Operators/group_by_operator.cs b/src/LinqTests/Operators/group_by_operator.cs
index 286f23e2b3..3d75f4e263 100644
--- a/src/LinqTests/Operators/group_by_operator.cs
+++ b/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);
+    }
 }
diff --git a/src/Marten/Linq/CollectionUsage.Compilation.cs b/src/Marten/Linq/CollectionUsage.Compilation.cs
index 2304466e43..9eb76e57bb 100644
--- a/src/Marten/Linq/CollectionUsage.Compilation.cs
+++ b/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;