Fix failure when aggregation function is wrapped in GROUP BY AUTO

core/trino-main/src/main/java/io/trino/sql/analyzer/AggregationAnalyzer.java

@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import io.trino.Session;
 import io.trino.metadata.FunctionResolver;
+import io.trino.metadata.ResolvedFunction;
 import io.trino.security.AccessControl;
 import io.trino.spi.StandardErrorCode;
 import io.trino.sql.PlannerContext;
@@ -32,6 +33,7 @@
 import io.trino.sql.tree.CurrentDate;
 import io.trino.sql.tree.CurrentTime;
 import io.trino.sql.tree.CurrentTimestamp;
+import io.trino.sql.tree.DefaultExpressionTraversalVisitor;
 import io.trino.sql.tree.DereferenceExpression;
 import io.trino.sql.tree.ExistsPredicate;
 import io.trino.sql.tree.Expression;
@@ -88,6 +90,8 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
@@ -100,6 +104,7 @@
 import static io.trino.spi.StandardErrorCode.INVALID_ARGUMENTS;
 import static io.trino.spi.StandardErrorCode.NESTED_AGGREGATION;
 import static io.trino.spi.StandardErrorCode.NESTED_WINDOW;
+import static io.trino.spi.function.FunctionKind.AGGREGATE;
 import static io.trino.sql.NodeUtils.getSortItemsFromOrderBy;
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractAggregateFunctions;
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractWindowExpressions;
@@ -810,4 +815,35 @@ private void verifyNoOrderByReferencesToOutputColumns(Node node, StandardErrorCo
                     throw semanticException(errorCode, expression, "%s", errorString);
                 });
     }
+
+    public static boolean containsAggregation(Expression expression, Function<FunctionCall, ResolvedFunction> functionResolver)
+    {
+        requireNonNull(functionResolver, "functionResolver is null");
+        requireNonNull(expression, "expression is null");
+
+        AtomicBoolean containsAggregation = new AtomicBoolean(false);
+        new AggregationFunctionVisitor(functionResolver).process(expression, containsAggregation);
+        return containsAggregation.get();
+    }
+
+    private static class AggregationFunctionVisitor
+            extends DefaultExpressionTraversalVisitor<AtomicBoolean>
+    {
+        private final Function<FunctionCall, ResolvedFunction> functionResolver;
+
+        public AggregationFunctionVisitor(Function<FunctionCall, ResolvedFunction> functionResolver)
+        {
+            this.functionResolver = functionResolver;
+        }
+
+        @Override
+        protected Void visitFunctionCall(FunctionCall node, AtomicBoolean containsAggregation)
+        {
+            if (functionResolver.apply(node).functionKind() == AGGREGATE) {
+                containsAggregation.set(true);
+                return null;
+            }
+            return super.visitFunctionCall(node, containsAggregation);
+        }
+    }
 }

core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java

@@ -382,6 +382,7 @@
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static io.trino.spi.type.VarcharType.createUnboundedVarcharType;
 import static io.trino.sql.NodeUtils.getSortItemsFromOrderBy;
+import static io.trino.sql.analyzer.AggregationAnalyzer.containsAggregation;
 import static io.trino.sql.analyzer.AggregationAnalyzer.verifyOrderByAggregations;
 import static io.trino.sql.analyzer.AggregationAnalyzer.verifySourceAggregations;
 import static io.trino.sql.analyzer.Analyzer.verifyNoAggregateWindowOrGroupingFunctions;
@@ -4557,16 +4558,11 @@ private GroupingSetAnalysis analyzeGroupBy(QuerySpecification node, Scope scope,
                     else if (groupingElement instanceof AutoGroupBy) {
                         // Analyze non-aggregation outputs
                         for (Expression column : outputExpressions) {
-                            if (column instanceof FunctionCall functionCall) {
-                                ResolvedFunction function = getResolvedFunction(functionCall);
-                                if (function.functionKind() == AGGREGATE) {
-                                    continue;
-                                }
-                            }
-                            else {
-                                verifyNoAggregateWindowOrGroupingFunctions(session, functionResolver, accessControl, column, "GROUP BY clause");
-                                analyzeExpression(column, scope);
+                            if (containsAggregation(column, this::getResolvedFunction)) {
+                                continue;
                             }
+                            verifyNoAggregateWindowOrGroupingFunctions(session, functionResolver, accessControl, column, "GROUP BY clause");
+                            analyzeExpression(column, scope);
 
                             ResolvedField field = analysis.getColumnReferenceFields().get(NodeRef.of(column));
                             if (field != null) {

core/trino-main/src/test/java/io/trino/sql/query/TestGroupBy.java

@@ -196,6 +196,22 @@ SELECT a, sum(b)
                 """))
                 .matches("VALUES (1, BIGINT '30')");
 
+        assertThat(assertions.query(
+                """
+                SELECT a, -sum(b)
+                FROM (VALUES (1, 10), (1, 20)) t(a, b)
+                GROUP BY AUTO
+                """))
+                .matches("VALUES (1, BIGINT '-30')");
+
+        assertThat(assertions.query(
+                """
+                SELECT a, abs(sum(b))
+                FROM (VALUES (1, -10), (1, -20)) t(a, b)
+                GROUP BY AUTO
+                """))
+                .matches("VALUES (1, BIGINT '30')");
+
         assertThat(assertions.query(
                 """
                 SELECT a AS new_a, sum(b) AS sum_b