From 20a11660025a3b7a156db7990650ec30a42deddd Mon Sep 17 00:00:00 2001
From: DrewKimball <drewk@cockroachlabs.com>
Date: Mon, 29 Aug 2022 17:16:08 -0700
Subject: [PATCH] smith: add all argument columns to ORDER BY

This commit ensures that sqlsmith adds all function argument columns to
the `ORDER BY` clause when nondetermistic functions are disabled. This is
necessary for functions like `string_agg`, for which there may be multiple
possible values for each argument when the order is not specified.

As a special case, `lead` and `lag` require all input columns to be added
to the `ORDER BY` clause. This is because while they respect the specified
order, they can still access any value in the partition. This means that we
need to uniquely specify which rows fall on the boundaries of each peer group.
We can ensure this is the case by ordering on a key, and we can trivially
order on a key by ordering on all columns.

Partially fixes #85318
Fixes #87023

Release justification: testing-only change

Release note: None
---
 pkg/internal/sqlsmith/scalar.go | 50 +++++++++++++++++----------------
 1 file changed, 26 insertions(+), 24 deletions(-)

diff --git a/pkg/internal/sqlsmith/scalar.go b/pkg/internal/sqlsmith/scalar.go
index 3b61f149489f..99ee035d4812 100644
--- a/pkg/internal/sqlsmith/scalar.go
+++ b/pkg/internal/sqlsmith/scalar.go
@@ -410,10 +410,6 @@ func makeFunc(s *Smither, ctx Context, typ *types.T, refs colRefs) (tree.TypedEx
 		}
 	}
 
-	// Some aggregation functions benefit from an order by clause.
-	var orderExpr tree.Expr
-	var orderType *types.T
-
 	args := make(tree.TypedExprs, 0)
 	for _, argTyp := range fn.overload.Types.Types() {
 		// Postgres is picky about having Int4 arguments instead of Int8.
@@ -425,10 +421,6 @@ func makeFunc(s *Smither, ctx Context, typ *types.T, refs colRefs) (tree.TypedEx
 		if class == tree.AggregateClass || class == tree.WindowClass {
 			var ok bool
 			arg, ok = makeColRef(s, argTyp, refs)
-			if ok && len(args) == 0 {
-				orderExpr = arg
-				orderType = argTyp
-			}
 			if !ok {
 				// If we can't find a col ref for our aggregate function, just use a
 				// constant.
@@ -458,20 +450,30 @@ func makeFunc(s *Smither, ctx Context, typ *types.T, refs colRefs) (tree.TypedEx
 			order      tree.OrderBy
 			orderTypes []*types.T
 		)
+		addOrdCol := func(expr tree.Expr, typ *types.T) {
+			order = append(order, &tree.Order{Expr: expr, Direction: s.randDirection()})
+			orderTypes = append(orderTypes, typ)
+		}
 		s.sample(len(refs)-len(parts), 2, func(i int) {
 			ref := refs[i+len(parts)]
-			order = append(order, &tree.Order{
-				Expr:      ref.item,
-				Direction: s.randDirection(),
-			})
-			orderTypes = append(orderTypes, ref.typ)
+			addOrdCol(ref.item, ref.typ)
 		})
-		if s.disableNondeterministicFns && orderExpr != nil {
-			order = append(order, &tree.Order{
-				Expr:      orderExpr,
-				Direction: s.randDirection(),
-			})
-			orderTypes = append(orderTypes, orderType)
+		if s.disableNondeterministicFns {
+			switch fn.def.Name {
+			case "lead", "lag":
+				// Lead and lag need all columns to be added to the ORDER BY because
+				// while they respect the order, they can return any value within the
+				// partition. This means we need to be consistent about which rows fall
+				// on the boundary between peer groups.
+				for _, ref := range refs {
+					addOrdCol(ref.typedExpr(), ref.typ)
+				}
+			default:
+				// Other window functions only need to order by the argument columns.
+				for _, arg := range args {
+					addOrdCol(arg, arg.ResolvedType())
+				}
+			}
 		}
 		var frame *tree.WindowFrame
 		if s.coin() {
@@ -496,7 +498,7 @@ func makeFunc(s *Smither, ctx Context, typ *types.T, refs colRefs) (tree.TypedEx
 	)
 
 	// Some aggregation functions need an order by clause to be deterministic.
-	if s.disableNondeterministicFns && orderExpr != nil {
+	if s.disableNondeterministicFns && len(args) > 0 {
 		switch fn.def.Name {
 		case "array_agg",
 			"concat_agg",
@@ -508,10 +510,10 @@ func makeFunc(s *Smither, ctx Context, typ *types.T, refs colRefs) (tree.TypedEx
 			"string_agg",
 			"xmlagg":
 			funcExpr.AggType = tree.GeneralAgg
-			funcExpr.OrderBy = tree.OrderBy{{
-				Expr:      orderExpr,
-				Direction: s.randDirection(),
-			}}
+			funcExpr.OrderBy = make(tree.OrderBy, len(args))
+			for i := range funcExpr.OrderBy {
+				funcExpr.OrderBy[i] = &tree.Order{Expr: args[i], Direction: s.randDirection()}
+			}
 		}
 	}