Support for grouping sets

docs/src/test/clauses.md

@@ -1131,6 +1131,52 @@ rendered.
     FROM "person"
     =#
 
+`GROUP` can accept the grouping mode or a vector of grouping sets.
+
+    c = FROM(:person) |> GROUP(:year_of_birth, sets = :ROLLUP)
+    #-> (…) |> GROUP(…, sets = :ROLLUP)
+
+    print(render(c |> SELECT(:year_of_birth, AGG(:count))))
+    #=>
+    SELECT
+      "year_of_birth",
+      count(*)
+    FROM "person"
+    GROUP BY ROLLUP("year_of_birth")
+    =#
+
+    c = FROM(:person) |> GROUP(:year_of_birth, sets = :CUBE)
+    #-> (…) |> GROUP(…, sets = :CUBE)
+
+    print(render(c |> SELECT(:year_of_birth, AGG(:count))))
+    #=>
+    SELECT
+      "year_of_birth",
+      count(*)
+    FROM "person"
+    GROUP BY CUBE("year_of_birth")
+    =#
+
+    c = FROM(:person) |> GROUP(:year_of_birth, sets = [[1], Int[]])
+    #-> (…) |> GROUP(…, sets = [[1], Int64[]])
+
+    print(render(c |> SELECT(:year_of_birth, AGG(:count))))
+    #=>
+    SELECT
+      "year_of_birth",
+      count(*)
+    FROM "person"
+    GROUP BY GROUPING SETS(("year_of_birth"), ())
+    =#
+
+`GROUP` raises an error when the vector of grouping sets is out of bounds.
+
+    FROM(:person) |> GROUP(:year_of_birth, sets = [[1, 2], [1], Int[]])
+    #=>
+    ERROR: DomainError with [[1, 2], [1], Int64[]]:
+    sets are out of bounds
+    =#
+
 
 ## `HAVING` Clause
 

docs/src/test/nodes.md

@@ -2502,6 +2502,95 @@ downstream.
     ) AS "person_2"
     =#
 
+`Group` allows specifying the grouping sets, either with grouping mode
+indicators `:cube` or `:rollup`, or by explicit enumeration.
+
+    q = From(person) |>
+        Group(Get.year_of_birth, sets = :cube)
+        Define(Agg.count())
+
+    display(q)
+    #=>
+    let person = SQLTable(:person, …),
+        q1 = From(person),
+        q2 = q1 |> Group(Get.year_of_birth, sets = :CUBE)
+        q2
+    end
+    =#
+
+    print(render(q))
+    #=>
+    SELECT "person_1"."year_of_birth"
+    FROM "person" AS "person_1"
+    GROUP BY CUBE("person_1"."year_of_birth")
+    =#
+
+    q = From(person) |>
+        Group(Get.year_of_birth, sets = [[1], Int[]])
+        Define(Agg.count())
+
+    display(q)
+    #=>
+    let person = SQLTable(:person, …),
+        q1 = From(person),
+        q2 = q1 |> Group(Get.year_of_birth, sets = [[1], []])
+        q2
+    end
+    =#
+
+    print(render(q))
+    #=>
+    SELECT "person_1"."year_of_birth"
+    FROM "person" AS "person_1"
+    GROUP BY GROUPING SETS(("person_1"."year_of_birth"), ())
+    =#
+
+`Group` allows specifying grouping sets using names of the grouping keys.
+
+    q = From(person) |>
+        Group(Get.year_of_birth, Get.gender_concept_id,
+              sets = ([:year_of_birth], ["gender_concept_id"]))
+        Define(Agg.count())
+
+    display(q)
+    #=>
+    let person = SQLTable(:person, …),
+        q1 = From(person),
+        q2 = q1 |>
+             Group(Get.year_of_birth, Get.gender_concept_id, sets = [[1], [2]])
+        q2
+    end
+    =#
+
+`Group` will report when a grouping set refers to an unknown key.
+
+    From(person) |>
+    Group(Get.year_of_birth, sets = [[:gender_concept_id], []])
+    #=>
+    ERROR: FunSQL.InvalidGroupingSetsError: `gender_concept_id` is not a valid key
+    =#
+
+`Group` complains about out-of-bound or incomplete grouping sets.
+
+    From(person) |>
+    Group(Get.year_of_birth, sets = [[1, 2], [1], []])
+    #=>
+    ERROR: FunSQL.InvalidGroupingSetsError: `2` is out of bounds in:
+    let q1 = Group(Get.year_of_birth, sets = [[1, 2], [1], []])
+        q1
+    end
+    =#
+
+    From(person) |>
+    Group(Get.year_of_birth, Get.gender_concept_id,
+          sets = [[1], []])
+    #=>
+    ERROR: FunSQL.InvalidGroupingSetsError: missing keys `[:year_of_birth]` in:
+    let q1 = Group(Get.year_of_birth, Get.gender_concept_id, sets = [[1], []])
+        q1
+    end
+    =#
+
 `Group` allows specifying the name of a group field.
 
     q = From(person) |>

docs/src/test/other.md

@@ -162,7 +162,7 @@ Any `Dict`-like object can serve as a query cache.
     #-> SQLCatalog(dialect = SQLDialect(), cache = Dict{Any, Any}())
 
     display(customcache_catalog)
-    #-> SQLCatalog(dialect = SQLDialect(), cache = Dict{Any, Any}())
+    #-> SQLCatalog(dialect = SQLDialect(), cache = (Dict{Any, Any})())
 
 The catalog behaves as a read-only `Dict` object.
 

src/clauses/group.jl

@@ -1,21 +1,47 @@
 # GROUP BY clause.
 
+module GROUPING_MODE
+
+@enum GroupingMode::UInt8 begin
+    ROLLUP
+    CUBE
+end
+
+Base.convert(::Type{GroupingMode}, s::Symbol) =
+    s in (:rollup, :ROLLUP) ?
+        ROLLUP :
+    s in (:cube, :CUBE) ?
+        CUBE :
+    throw(DomainError(QuoteNode(s), "expected :rollup or :cube"))
+
+end
+
+import .GROUPING_MODE.GroupingMode
+
 mutable struct GroupClause <: AbstractSQLClause
     over::Union{SQLClause, Nothing}
     by::Vector{SQLClause}
+    sets::Union{Vector{Vector{Int}}, GroupingMode, Nothing}
 
-    GroupClause(;
+    function GroupClause(;
                 over = nothing,
-                by = SQLClause[]) =
-        new(over, by)
+                by = SQLClause[],
+                sets = nothing)
+        c = new(over, by, sets isa Symbol ? convert(GroupingMode, sets) : sets)
+        s = c.sets
+        if s isa Vector{Vector{Int}} && !checkbounds(Bool, c.by, s)
+            throw(DomainError(s, "sets are out of bounds"))
+        end
+        c
+    end
 end
 
-GroupClause(by...; over = nothing) =
-    GroupClause(over = over, by = SQLClause[by...])
+GroupClause(by...; over = nothing, sets = nothing) =
+    GroupClause(over = over, by = SQLClause[by...], sets = sets)
 
 """
-    GROUP(; over = nothing, by = [])
-    GROUP(by...; over = nothing)
+    GROUP(; over = nothing, by = [], sets = nothing)
+    GROUP(by...; over = nothing, sets = nothing)
 
 A `GROUP BY` clause.
 
@@ -43,12 +69,16 @@ dissect(scr::Symbol, ::typeof(GROUP), pats::Vector{Any}) =
 function PrettyPrinting.quoteof(c::GroupClause, ctx::QuoteContext)
     ex = Expr(:call, nameof(GROUP))
     append!(ex.args, quoteof(c.by, ctx))
+    s = c.sets
+    if s !== nothing
+        push!(ex.args, Expr(:kw, :sets, s isa GroupingMode ? QuoteNode(Symbol(s)) : s))
+    end
     if c.over !== nothing
         ex = Expr(:call, :|>, quoteof(c.over, ctx), ex)
     end
     ex
 end
 
 rebase(c::GroupClause, c′) =
-    GroupClause(over = rebase(c.over, c′), by = c.by)
+    GroupClause(over = rebase(c.over, c′), by = c.by, sets = c.sets)
 

src/link.jl

@@ -118,7 +118,7 @@ end
 function dismantle(n::GroupNode, ctx)
     over′ = dismantle(n.over, ctx)
     by′ = dismantle_scalar(n.by, ctx)
-    Group(over = over′, by = by′, name = n.name, label_map = n.label_map)
+    Group(over = over′, by = by′, sets = n.sets, name = n.name, label_map = n.label_map)
 end
 
 function dismantle(n::IterateNode, ctx)
@@ -290,6 +290,10 @@ function link(n::GroupNode, ctx)
     # To avoid duplicate SQL, they must be evaluated in a nested subquery.
     refs = SQLNode[]
     append!(refs, n.by)
+    if n.sets !== nothing
+        # Force evaluation in a nested subquery.
+        append!(refs, n.by)
+    end
     # Ignore `SELECT DISTINCT` case.
     if has_aggregates
         ctx′ = LinkContext(ctx, refs = refs)
@@ -311,7 +315,7 @@ function link(n::GroupNode, ctx)
         over = Padding(over = over)
     end
     over′ = Linked(refs, 0, over = link(over, ctx, refs))
-    Group(over = over′, by = n.by, name = n.name, label_map = n.label_map)
+    Group(over = over′, by = n.by, sets = n.sets, name = n.name, label_map = n.label_map)
 end
 
 function link(n::IterateNode, ctx)

src/nodes.jl

@@ -348,7 +348,7 @@ A node that cannot be rebased.
 struct RebaseError <: FunSQLError
     path::Vector{SQLNode}
 
-    RebaseError(; path = SQLNode[])=
+    RebaseError(; path = SQLNode[]) =
         new(path)
 end
 
@@ -357,6 +357,30 @@ function Base.showerror(io::IO, err::RebaseError)
     showpath(io, err.path)
 end
 
+"""
+Grouping sets are specified incorrectly.
+"""
+struct InvalidGroupingSetsError <: FunSQLError
+    value::Union{Int, Symbol, Vector{Symbol}}
+    path::Vector{SQLNode}
+
+    InvalidGroupingSetsError(value; path = SQLNode[]) =
+        new(value, path)
+end
+
+function Base.showerror(io::IO, err::InvalidGroupingSetsError)
+    print(io, "FunSQL.InvalidGroupingSetsError: ")
+    value = err.value
+    if value isa Int
+        print(io, "`$value` is out of bounds")
+    elseif value isa Symbol
+        print(io, "`$value` is not a valid key")
+    elseif value isa Vector{Symbol}
+        print(io, "missing keys `$value`")
+    end
+    showpath(io, err.path)
+end
+
 module REFERENCE_ERROR_TYPE
 
 @enum ReferenceErrorType::UInt8 begin