Add a new Optimizer Rule to Remove Distinct Aggregation for Queries containing Grouped Aggregation Operators

[pgandhi999]

Description

This optimization removes the redundant Aggregation operator introduced by Distinct clause when the query contains Grouped Aggregation operator.

For context, we noticed a query failing due to OOM error where the query plan contained a Grouped Aggregation Operator followed by the Distinct Aggregation Operator. The Distinct Aggregation Operator is redundant for Grouping Aggregations. After removing the DISTINCT clause, the query succeeded. Adding an Optimizer Rule for the same.

Additional context and related issues

Release notes

( ) This is not user-visible or is docs only, and no release notes are required.
( ) Release notes are required. Please propose a release note for me.
(x) Release notes are required, with the following suggested text:

# General
* Improve performance for queries with a redundant `DISTINCT` over already-distinct `GROUP BY` aggregation ({issue}`23087`)

[pgandhi999]

@pettyjamesm

[pettyjamesm]

Latest iteration looks good to me, @martint - can you take another look and @pgandhi999 please resolve the review comments that have already been addressed?

[martint]

This is still not quite correct. It's technically not legal to compare the grouping keys from the leaf aggregation with the grouping keys passed from above without performing the reverse translations defined by this projection, as the meaning of symbols is local to each pair of operations and it may have been changed by the projection.

[pettyjamesm]

It may not locally be correct to perform the check: groupingKeys.containsAll(aggregationNode.getGroupingSets().getGroupingKeys()) in isolation, since we don't know at that point whether the grouping key symbols from the parent and child aggregations have the same meaning. However, when we return the leaf aggregation symbols back up the call stack, we verfiy that any ProjectNodes encountered only perform identity projections for all leaf aggregation symbols- at which point we can say that they do have the same meaning by the time we reach the parent aggregation again and the transform should be considered to be valid.

[martint]

My concern is that it's too brittle and it relies on non-local assumptions. The call to isDistinctOverGroupingKeys(lookup.resolve(projectNode.getSource()), lookup, groupingKeys) is incorrect by construction, and whether childGroupingKeys.stream().allMatch(projectNode.getAssignments()::isIdentity) is sufficient to compensate depends on knowledge that this specific callsite doesn't have. For example, adding support of another type of node might break these assumptions.

Also, in the call to childGroupingKeys.stream().allMatch(projectNode.getAssignments()::isIdentity), it's incorrect to perform a lookup in the project node assignments using the symbols from the input scope as the key. The keys belong in the "output" scope. Today that works, but if in the future we make such lookups stricter by failing if the key is not from the proper scope, this will break.

Separately, the check is too specific. For example, it won't be able to catch plans like these ones:

- Aggregate [k]
  - Project k = x
    - Aggregate [x]

[pettyjamesm]

My concern is that it's too brittle and it relies on non-local assumptions

That’s fair, although I would point out that the assumptions are self contained to the single static private method, so they’re reasonably localized.

For example, adding support of another type of node might break these assumptions.

Agreed. Maybe we can add a comment to clarify the specific assumptions being made and warn against expanding the rule to accept new intervening node types?

The keys belong in the "output" scope. Today that works, but if in the future we make such lookups stricter by failing if the key is not from the proper scope, this will break.

It’s hard for me to guess what that might look like, but keys of interest by definition straddle the input and output scope, and the check for identity transforms guarantees that they are the same. (This is related to the next point, more on that there).

Separately, the check is too specific. For example, it won't be able to catch plans like these ones:

- Aggregate [k]
 - Project k = x
   - Aggregate [x]

It won’t detect that, true- but it’s a missed optimization opportunity at that point and not a correctness bug. I’m not sure whether there are any plan optimizers today that would create new symbols for an identity assignment, but if there are then I guess we should plan for that? It’s definitely going to be more expensive to find any identity mappings from all input symbols and consider the intersection of those with the parent grouping symbols, but should be doable at that higher cost. Is that the recommendation here?

[martint]

Ok, fine. Let's add that documentation and get it in as is. We can fix those issue later when they come up. Also, please add a comment explaining that the lookups are not "proper" and that we're taking a shortcut.

[pgandhi999]

Thank you for all the comments. I have rewritten the method to include translating the grouping set symbols for project nodes.

[pettyjamesm]

Merged, thanks @pgandhi999!

[sopel39]

This should use io.trino.sql.planner.optimizations.ActualProperties#localProperties instead of implementing derivations itself as part of RemoveRedundantDistinctAggregation. It would be more generic then too because it would work for aggregations that are distinct within stream, but not necessarily globally.