Merge cockroachdb#26490

26490: opt/optbuilder: allow aggregation in order by r=rytaft a=rytaft

This commit fixes the optimizer so that it allows aggregations
in the `ORDER BY` clause, even if there are no aggregates in
the `SELECT` list and there is no `GROUP BY` clause. This change
is made for compatibility with Postgres.

Fixes cockroachdb#26457

Release note: None

Co-authored-by: Rebecca Taft <[email protected]>

pkg/sql/opt/optbuilder/groupby.go

@@ -87,17 +87,17 @@ type aggregateInfo struct {
 	args []memo.GroupID
 }
 
-func (b *Builder) needsAggregation(sel *tree.SelectClause) bool {
+func (b *Builder) needsAggregation(sel *tree.SelectClause, orderBy tree.OrderBy) bool {
 	// We have an aggregation if:
 	//  - we have a GROUP BY, or
 	//  - we have a HAVING clause, or
-	//  - we have aggregate functions in the select expressions.
-	return len(sel.GroupBy) > 0 || sel.Having != nil || b.hasAggregates(sel.Exprs)
+	//  - we have aggregate functions in the select and/or order by expressions.
+	return len(sel.GroupBy) > 0 || sel.Having != nil || b.hasAggregates(sel.Exprs, orderBy)
 }
 
 // hasAggregates determines if any of the given select expressions contain an
 // aggregate function.
-func (b *Builder) hasAggregates(selects tree.SelectExprs) bool {
+func (b *Builder) hasAggregates(selects tree.SelectExprs, orderBy tree.OrderBy) bool {
 	exprTransformCtx := transform.ExprTransformContext{}
 	for _, sel := range selects {
 		// TODO(rytaft): This function does not recurse into subqueries, so this
@@ -107,6 +107,12 @@ func (b *Builder) hasAggregates(selects tree.SelectExprs) bool {
 		}
 	}
 
+	for _, ob := range orderBy {
+		if exprTransformCtx.AggregateInExpr(ob.Expr, b.semaCtx.SearchPath) {
+			return true
+		}
+	}
+
 	return false
 }
 

pkg/sql/opt/optbuilder/scope.go

@@ -319,7 +319,7 @@ func (s *scope) startAggFunc() (aggInScope *scope, aggOutScope *scope) {
 		}
 	}
 
-	panic(unimplementedf("woops the aggregate semantic analysis is incomplete"))
+	panic(fmt.Errorf("aggregate function is not allowed in this context"))
 }
 
 // endAggFunc is called when the builder finishes building an aggregate

pkg/sql/opt/optbuilder/select.go

@@ -276,7 +276,7 @@ func (b *Builder) buildSelectClause(
 	outScope = fromScope
 
 	var projectionsScope *scope
-	if b.needsAggregation(sel) {
+	if b.needsAggregation(sel, orderBy) {
 		outScope, projectionsScope = b.buildAggregation(sel, orderBy, fromScope)
 	} else {
 		projectionsScope = fromScope.replace()

pkg/sql/opt/optbuilder/testdata/aggregate

@@ -2628,3 +2628,24 @@ project
       └── aggregations
            └── max [type=decimal]
                 └── variable: column6 [type=decimal]
+
+# Regression test for #26419
+build
+SELECT 123 FROM kv ORDER BY max(v)
+----
+sort
+ ├── columns: "123":5(int!null)
+ ├── ordering: +6
+ └── project
+      ├── columns: "123":5(int!null) column6:6(int)
+      ├── group-by
+      │    ├── columns: column6:6(int)
+      │    ├── project
+      │    │    ├── columns: v:2(int)
+      │    │    └── scan kv
+      │    │         └── columns: k:1(int!null) v:2(int) w:3(int) s:4(string)
+      │    └── aggregations
+      │         └── max [type=int]
+      │              └── variable: kv.v [type=int]
+      └── projections
+           └── const: 123 [type=int]