dj.U(...).aggr gives wrong results or fails when the group is a union, an aggregation, or a computed projection #1661

What happens

dj.U(...).aggr(group, ...) should aggregate over the rows of group. When group is a plain table, a join, or a restriction of these, it does. When group is a union, an aggregation, or a projection with a computed attribute, it isn't aggregating the rows that group returns:

• Union: the second operand is dropped without a warning.
• Aggregation: the inner grouping is lost. Depending on the computation, the result is a wrong count with no error, an SQL error, or a missing column. A restriction applied before the inner grouping is dropped too.
• Projection with a computed attribute, aggregated per attribute: SQL error.

A dj.U(...) & expr group, whose DISTINCT is dropped in the same place, is covered by #1660 together with the joins that drop it.

The regular aggr (Parent.aggr(group, ...)) handles the same groups correctly.

Reproduction

On master (3c912c4), the same on MySQL 8.0 and PostgreSQL:

import datajoint as dj

schema = dj.Schema("u_aggr_repro")

@schema
class Parent(dj.Manual):
    definition = """
    parent_id : int32
    ---
    label : varchar(8)
    """

@schema
class Child(dj.Manual):
    definition = """
    -> Parent
    child_id : int32
    ---
    color : varchar(8)
    """

@schema
class ParentInfo(dj.Manual):
    definition = """
    -> Parent
    ---
    info : varchar(8)
    """

Parent.insert([
    {"parent_id": 1, "label": "a"},
    {"parent_id": 2, "label": "b"},
    {"parent_id": 3, "label": "a"},
])
Child.insert([
    {"parent_id": 1, "child_id": 1, "color": "red"},
    {"parent_id": 1, "child_id": 2, "color": "blue"},
    {"parent_id": 2, "child_id": 1, "color": "blue"},
])
ParentInfo.insert([
    {"parent_id": 1, "info": "x"},
    {"parent_id": 2, "info": "y"},
    {"parent_id": 3, "info": "z"},
])

u = (Parent.proj() & "parent_id=1") + (ParentInfo.proj() & "parent_id > 1")
len(u)                                       # 3
dj.U().aggr(u, n="count(*)").fetch1("n")     # 1, expected 3

per_parent = Parent.aggr(Child, n="count(child_id)")     # n: parent 1 -> 2, 2 -> 1, 3 -> 0
dj.U().aggr(per_parent, total="sum(n)").fetch1("total")  # error, expected 3
dj.U().aggr(per_parent, k="count(*)").fetch1("k")        # 4, expected 3

Four of the eight return a wrong number with no error.

The SQL shows what is lost. For the union, the parent_info side is gone:

SELECT count(*) as `n` FROM (SELECT `parent_id` FROM `parent` WHERE ( (parent_id=1))) as `$0`

For the aggregation, the inner GROUP BY parent_id is gone, so count(*) counts the joined rows (4) instead of the parents (3), and sum(n) fails:

SELECT count(*) as `k` FROM `parent` LEFT JOIN `child` USING (`parent_id`)
SELECT sum(n) as `total` FROM (SELECT `parent_id`,count(child_id) as `n` FROM `parent` LEFT JOIN `child` USING (`parent_id`)) as `$1`

These cases work:

• dj.U().aggr(Child.proj(score="child_id * 10"), m="max(score)") returns 20 on both backends. With no grouping attribute, the inner query has no GROUP BY to break.
• Grouping by a computed or renamed attribute works: dj.U("score").aggr(Child.proj(score="child_id * 10"), n="count(*)") and dj.U("pid").aggr(Child.proj(pid="parent_id"), n="count(*)").
• A union that is restricted again, dj.U().aggr(u & "parent_id > 1", n="count(*)"), returns 2. Restricting a union already wraps it in a subquery.
• Parent.aggr(u, n="count(*)") and Parent.aggr(Child.proj(score="child_id * 10"), m="max(score)") are correct.

Cause

U.aggr (src/datajoint/expression.py:1532) builds the aggregation from three parts of the group and sets its own grouping (lines 1567–1573):

result._heading = group.heading.set_primary_key(list(self.primary_key))
result._support = group.support
result._joins = group._joins
result._left_restrict = group.restriction
result._grouping_attributes = list(self.primary_key)

These three parts describe a table, a join, or a restriction completely. Other groups lose something:

• Union: Union.create puts both operands in _support (line 1429) and has no _joins. from_clause() (lines 105–128) pairs each entry of _support after the first with an entry of _joins. With no _joins, only the first operand reaches the FROM clause.
• Aggregation: its _support and _joins are the joined tables. Its restriction holds the HAVING conditions, and its _left_restrict and _grouping_attributes hold the WHERE and GROUP BY. U.aggr keeps the tables and moves HAVING into WHERE, where n is not a column yet. It drops the inner WHERE and GROUP BY. n keeps its expression, count(child_id), in the heading.
• Computed attribute: score keeps its expression in the aggregation's heading. Because max(score) refers to a computed attribute, proj() wraps the aggregation in a subquery (line 613). The inner query then selects all of Child's columns under GROUP BY parent_id.

The regular aggr covers these groups. Aggregation.create wraps a group with computed attributes in a subquery (line 1336), and join() wraps aggregation operands and operands built from several tables (lines 404–421). Before 2.0, U.aggr called Aggregation.create. Commit ba81cb8 (first released in 2.0.0) replaced that call with the direct construction above, which also skips the computed-attribute check.

Expected

The values in the "Expected" column above. In U.aggr, wrap the group in a subquery unless it is a table, a join, or a restriction of these. That covers a Union, an Aggregation, and a group with computed attributes (group.heading.new_attributes); a DISTINCT group (group._distinct) is #1660. Calling group.make_subquery() by hand already gives the expected value for every row of the table, on both backends. Plain tables and joins should keep the current SQL with no subquery.

PR #1657 adds the same wrapping in this function for a group limited by dj.Top (if group._top is not None: group = group.make_subquery()) and lists this bug under "Not changed here". The fix can extend that condition. Neither change alone covers the other's cases.