Adding COUNT and COUNT(DISTINCT) aggregation support

docs/cheatsheet.md

@@ -67,6 +67,18 @@ Foo.select.sumByOpt(_.myInt)                                        // Option[In
 Foo.select.size                                                     // Int
 // SELECT COUNT(1) FROM foo
 
+Foo.select.countBy(_.myInt)                                          // Int
+// SELECT COUNT(my_int) FROM foo
+
+Foo.select.countDistinctBy(_.myInt)                                  // Int
+// SELECT COUNT(DISTINCT my_int) FROM foo
+
+Foo.select.map(_.myInt).count                                        // Int
+// SELECT COUNT(my_int) FROM foo
+
+Foo.select.map(_.myInt).countDistinct                                // Int
+// SELECT COUNT(DISTINCT my_int) FROM foo
+
 Foo.select.aggregate(fs => (fs.sumBy(_.myInt), fs.maxBy(_.myInt)))  // (Int, Int)
 // SELECT SUM(my_int), MAX(my_int) FROM foo
 
@@ -200,14 +212,16 @@ to allow ScalaSql to work with it
 
 **Aggregate Functions**
 
-|                                      Scala |       SQL |
-|----------------------------------------------------:|------------------------:|
-|                                            `a.size` |              `COUNT(1)` |
-|                                   `a.mkString(sep)` |  `GROUP_CONCAT(a, sep)` |
-|  `a.sum`, `a.sumBy(_.myInt)`, `a.sumByOpt(_.myInt)` |           `SUM(my_int)` |
-|  `a.min`, `a.minBy(_.myInt)`, `a.minByOpt(_.myInt)` |           `MIN(my_int)` |
-|  `a.max`, `a.maxBy(_.myInt)`, `a.maxByOpt(_.myInt)` |           `MAX(my_int)` |
-|  `a.avg`, `a.avgBy(_.myInt)`, `a.avgByOpt(_.myInt)` |           `AVG(my_int)` |
+|                                                        Scala |                     SQL |
+|-------------------------------------------------------------:|------------------------:|
+|                                                     `a.size` |              `COUNT(1)` |
+|                                            `a.mkString(sep)` |  `GROUP_CONCAT(a, sep)` |
+|           `a.sum`, `a.sumBy(_.myInt)`, `a.sumByOpt(_.myInt)` |           `SUM(my_int)` |
+|           `a.min`, `a.minBy(_.myInt)`, `a.minByOpt(_.myInt)` |           `MIN(my_int)` |
+|           `a.max`, `a.maxBy(_.myInt)`, `a.maxByOpt(_.myInt)` |           `MAX(my_int)` |
+|           `a.avg`, `a.avgBy(_.myInt)`, `a.avgByOpt(_.myInt)` |           `AVG(my_int)` |
+|                 `a.countBy(_.myInt)`, `a.map(_.myInt).count` |         `COUNT(my_int)` |
+| `a.countDistinctBy(_.myInt)`, `a.map(_.myInt).countDistinct` |`COUNT(DISTINCT my_int)` |
 
 **Select Functions**
 

docs/reference.md

@@ -1133,6 +1133,91 @@ Purchase.select.sumBy(_.total)
 
 
 
+### Select.aggregate.countBy
+
+You can use `.countBy` to generate SQL `COUNT(column)` aggregates that count non-null values
+
+```scala
+Purchase.select.countBy(_.productId)
+```
+
+
+*
+    ```sql
+    SELECT COUNT(purchase0.product_id) AS res FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    7
+    ```
+
+### Select.aggregate.countDistinctBy
+
+You can use `.countDistinctBy` to generate SQL `COUNT(DISTINCT column)` aggregates
+that count unique non-null values
+
+```scala
+Purchase.select.countDistinctBy(_.productId)
+```
+
+
+*
+    ```sql
+    SELECT COUNT(DISTINCT purchase0.product_id) AS res FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    6
+    ```
+
+### Select.aggregate.count
+
+You can use `.count` on mapped expressions to generate SQL `COUNT(expr)` aggregates
+
+```scala
+Purchase.select.map(_.productId).count
+```
+
+
+*
+    ```sql
+    SELECT COUNT(purchase0.product_id) AS res FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    7
+    ```
+
+### Select.aggregate.countDistinct
+
+You can use `.countDistinct` on mapped expressions to generate SQL `COUNT(DISTINCT expr)` aggregates
+
+```scala
+Purchase.select.map(_.productId).countDistinct
+```
+
+
+*
+    ```sql
+    SELECT COUNT(DISTINCT purchase0.product_id) AS res FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    6
+    ```
+
 ### Select.aggregate.multiple
 
 If you want to perform multiple aggregates at once, you can use the `.aggregate` method
@@ -1155,6 +1240,27 @@ Purchase.select.aggregate(q => (q.sumBy(_.total), q.maxBy(_.total)))
     (12343.2, 10000.0)
     ```
 
+### Select.aggregate.multipleWithCount
+
+You can combine COUNT operations with other aggregates in a single query
+
+```scala
+Purchase.select.aggregate(q => (q.countBy(_.productId), q.countDistinctBy(_.productId), q.sumBy(_.total)))
+```
+
+
+*
+    ```sql
+    SELECT COUNT(purchase0.product_id) AS res_0, COUNT(DISTINCT purchase0.product_id) AS res_1, SUM(purchase0.total) AS res_2 FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    (7, 6, 12343.2)
+    ```
+
 
 
 ### Select.groupBy.simple
@@ -6320,6 +6426,88 @@ Purchase.select.mapAggregate((p, ps) =>
 
 
 
+### WindowFunction.aggregate.countBy
+
+Window functions can also use COUNT operations with partitioning and ordering.
+
+```scala
+Purchase.select.mapAggregate((p, ps) =>
+  (
+    p.shippingInfoId,
+    p.total,
+    ps.countBy(_.productId).over.partitionBy(p.shippingInfoId).sortBy(p.total).asc
+  )
+)
+```
+
+
+*
+    ```sql
+    SELECT
+      purchase0.shipping_info_id AS res_0,
+      purchase0.total AS res_1,
+      COUNT(purchase0.product_id) OVER (PARTITION BY purchase0.shipping_info_id ORDER BY purchase0.total ASC) AS res_2
+    FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    Seq(
+      (1, 15.7, 1),
+      (1, 888.0, 2),
+      (1, 900.0, 3),
+      (2, 493.8, 1),
+      (2, 10000.0, 2),
+      (3, 1.3, 1),
+      (3, 44.4, 2)
+    )
+    ```
+
+
+
+### WindowFunction.aggregate.countDistinctBy
+
+COUNT(DISTINCT) can also be used as a window function for running distinct counts.
+
+```scala
+Purchase.select.mapAggregate((p, ps) =>
+  (
+    p.shippingInfoId,
+    p.total,
+    ps.countDistinctBy(_.productId).over.partitionBy(p.shippingInfoId).sortBy(p.total).asc
+  )
+)
+```
+
+
+*
+    ```sql
+    SELECT
+      purchase0.shipping_info_id AS res_0,
+      purchase0.total AS res_1,
+      COUNT(DISTINCT purchase0.product_id) OVER (PARTITION BY purchase0.shipping_info_id ORDER BY purchase0.total ASC) AS res_2
+    FROM purchase purchase0
+    ```
+
+
+
+*
+    ```scala
+    Seq(
+      (1, 15.7, 1),
+      (1, 888.0, 2),
+      (1, 900.0, 3),
+      (2, 493.8, 1),
+      (2, 10000.0, 2),
+      (3, 1.3, 1),
+      (3, 44.4, 2)
+    )
+    ```
+
+
+
 ### WindowFunction.frames
 
 You can have further control over the window function call via `.frameStart`,

docs/tutorial.md

@@ -582,6 +582,34 @@ db.renderSql(query) ==>
 db.run(query) ==> 154
 ```
 
+ScalaSql also provides `.countBy` and `.countDistinctBy` for more specific counting
+operations. `.countBy` generates `COUNT(column)` and only counts non-null values,
+while `.countDistinctBy` generates `COUNT(DISTINCT column)` to count unique non-null values:
+```scala
+// Count non-null country codes
+val query1 = Country.select.countBy(_.code)
+db.renderSql(query1) ==> "SELECT COUNT(country0.code) AS res FROM country country0"
+
+// Count distinct continents
+val query2 = Country.select.countDistinctBy(_.continent)
+db.renderSql(query2) ==> "SELECT COUNT(DISTINCT country0.continent) AS res FROM country country0"
+
+db.run(query2) ==> 7  // 7 distinct continents in the database
+```
+
+You can also use `.count` and `.countDistinct` on mapped expressions:
+```scala
+// Count non-null population values after mapping
+val query3 = Country.select.map(_.population).count
+db.renderSql(query3) ==> "SELECT COUNT(country0.population) AS res FROM country country0"
+
+// Count distinct population density categories
+val query4 = Country.select
+  .map(c => c.population / c.surfaceArea)
+  .countDistinct
+db.renderSql(query4) ==> "SELECT COUNT(DISTINCT (country0.population / country0.surfacearea)) AS res FROM country country0"
+```
+
 If you want to perform multiple aggregates at once, you can use the `.aggregate`
 function. Below, we run a single query that returns the minimum, average, and
 maximum populations across all countries in our dataset
@@ -599,6 +627,25 @@ FROM country country0
 db.run(query) ==> (0, 25434098, 1277558000)
 ```
 
+You can combine COUNT operations with other aggregates in the same query:
+```scala
+val query = Country.select
+  .aggregate(cs => (
+    cs.countBy(_.population),        // Count non-null populations
+    cs.countDistinctBy(_.continent), // Count distinct continents
+    cs.sumBy(_.population)           // Sum all populations
+  ))
+db.renderSql(query) ==> """
+SELECT
+  COUNT(country0.population) AS res_0,
+  COUNT(DISTINCT country0.continent) AS res_1,
+  SUM(country0.population) AS res_2
+FROM country country0
+"""
+
+db.run(query) ==> (239, 7, 6078749450)
+```
+
 ### Sort/Drop/Take
 
 You can use `.sortBy` to order the returned rows, and `.drop` and `.take`

scalasql/operations/src/AggAnyOps.scala

@@ -0,0 +1,20 @@
+package scalasql.operations
+
+import scalasql.core.{Aggregatable, Expr, Queryable, TypeMapper}
+import scalasql.core.SqlStr.SqlStringSyntax
+
+/**
+ * Aggregations that apply to any element type `T`, e.g. COUNT and COUNT(DISTINCT)
+ * over an aggregated `Expr[T]` sequence.
+ */
+class AggAnyOps[T](v: Aggregatable[Expr[T]])(
+    implicit tmInt: TypeMapper[Int],
+    qrInt: Queryable.Row[Expr[Int], Int]
+) {
+  /** Counts non-null values */
+  def count: Expr[Int] = v.aggregateExpr[Int](expr => implicit ctx => sql"COUNT($expr)")
+
+  /** Counts distinct non-null values */
+  def countDistinct: Expr[Int] =
+    v.aggregateExpr[Int](expr => implicit ctx => sql"COUNT(DISTINCT $expr)")
+}

scalasql/operations/src/AggOps.scala

@@ -11,6 +11,16 @@ class AggOps[T](v: Aggregatable[T])(implicit qr: Queryable.Row[T, ?], dialect: D
   /** Counts the rows */
   def size: Expr[Int] = v.aggregateExpr(_ => _ => sql"COUNT(1)")
 
+  /** Counts non-null values in the selected column */
+  def countBy[V](f: T => Expr[V])(implicit qrInt: Queryable.Row[Expr[Int], Int]): Expr[Int] =
+    v.aggregateExpr[Int](expr => implicit ctx => sql"COUNT(${f(expr)})")
+
+  /** Counts distinct non-null values in the selected column */
+  def countDistinctBy[V](
+      f: T => Expr[V]
+  )(implicit qrInt: Queryable.Row[Expr[Int], Int]): Expr[Int] =
+    v.aggregateExpr[Int](expr => implicit ctx => sql"COUNT(DISTINCT ${f(expr)})")
+
   /** Computes the sum of column values */
   def sumBy[V: TypeMapper](f: T => Expr[V])(
       implicit qr: Queryable.Row[Expr[V], V]

scalasql/src/dialects/Dialect.scala

@@ -302,6 +302,10 @@ trait Dialect extends DialectTypeMappers {
       implicit qr: Queryable.Row[T, ?]
   ): operations.AggOps[T] = new operations.AggOps(v)
 
+  implicit def AggAnyOpsConv[T: TypeMapper](v: Aggregatable[Expr[T]])(
+      implicit qrInt: Queryable.Row[Expr[Int], Int]
+  ): operations.AggAnyOps[T] = new operations.AggAnyOps(v)
+
   implicit def ExprAggOpsConv[T](v: Aggregatable[Expr[T]]): operations.ExprAggOps[T]
 
   implicit def TableOpsConv[V[_[_]]](t: Table[V]): TableOps[V] = new TableOps(t)