sql/parser: Define overload returnType interface and cleanup builtins

pkg/sql/distsqlrun/aggregator.go

@@ -43,7 +43,10 @@ func GetAggregateInfo(
 		for _, t := range b.Types.Types() {
 			if inputDatumType.Equivalent(t) {
 				// Found!
-				return b.AggregateFunc, sqlbase.DatumTypeToColumnType(b.ReturnType), nil
+				constructAgg := func() parser.AggregateFunc {
+					return b.AggregateFunc([]parser.Type{inputDatumType})
+				}
+				return constructAgg, sqlbase.DatumTypeToColumnType(b.FixedReturnType()), nil
 			}
 		}
 	}

pkg/sql/executor.go

@@ -1482,6 +1482,7 @@ func checkResultType(typ parser.Type) error {
 	case parser.TypeTimestampTZ:
 	case parser.TypeInterval:
 	case parser.TypeStringArray:
+	case parser.TypeNameArray:
 	case parser.TypeIntArray:
 	default:
 		// Compare all types that cannot rely on == equality.

pkg/sql/parser/aggregate_builtins.go

@@ -81,11 +81,17 @@ type AggregateFunc interface {
 // Exported for use in documentation.
 var Aggregates = map[string][]Builtin{
 	"array_agg": {
-		makeAggBuiltin(TypeInt, TypeIntArray, newIntArrayAggregate,
-			"Aggregates the selected values into an array."),
-		makeAggBuiltin(
-			TypeString, TypeStringArray, newStringArrayAggregate,
-			"Aggregates the selected values into an array."),
+		makeAggBuiltinWithReturnType(
+			TypeAny,
+			func(args []TypedExpr) Type {
+				if len(args) == 0 {
+					return unknownReturnType
+				}
+				return tArray{args[0].ResolvedType()}
+			},
+			newArrayAggregate,
+			"Aggregates the selected values into an array.",
+		),
 	},
 
 	"avg": {
@@ -117,14 +123,19 @@ var Aggregates = map[string][]Builtin{
 			"Concatenates all selected values."),
 	},
 
-	"count": countImpls(),
+	"count": {
+		makeAggBuiltin(TypeAny, TypeInt, newCountAggregate,
+			"Calculates the number of selected elements."),
+	},
 
-	"max": makeAggBuiltins(newMaxAggregate, "Identifies the maximum selected value.",
-		TypeBool, TypeInt, TypeFloat, TypeDecimal, TypeString, TypeBytes,
-		TypeDate, TypeTimestamp, TypeTimestampTZ, TypeInterval),
-	"min": makeAggBuiltins(newMinAggregate, "Identifies the minimum selected value.",
-		TypeBool, TypeInt, TypeFloat, TypeDecimal, TypeString, TypeBytes,
-		TypeDate, TypeTimestamp, TypeTimestampTZ, TypeInterval),
+	"max": collectBuiltins(func(t Type) Builtin {
+		return makeAggBuiltin(t, t, newMaxAggregate,
+			"Identifies the maximum selected value.")
+	}, TypesAnyNonArray...),
+	"min": collectBuiltins(func(t Type) Builtin {
+		return makeAggBuiltin(t, t, newMinAggregate,
+			"Identifies the minimum selected value.")
+	}, TypesAnyNonArray...),
 
 	"sum_int": {
 		makeAggBuiltin(TypeInt, TypeInt, newSmallIntSumAggregate,
@@ -161,41 +172,28 @@ var Aggregates = map[string][]Builtin{
 	},
 }
 
-func makeAggBuiltin(in, ret Type, f func() AggregateFunc, info string) Builtin {
+func makeAggBuiltin(in, ret Type, f func([]Type) AggregateFunc, info string) Builtin {
+	return makeAggBuiltinWithReturnType(in, fixedReturnType(ret), f, info)
+}
+
+func makeAggBuiltinWithReturnType(
+	in Type, retType returnTyper, f func([]Type) AggregateFunc, info string,
+) Builtin {
 	return Builtin{
 		// See the comment about aggregate functions in the definitions
 		// of the Builtins array above.
 		impure:        true,
 		class:         AggregateClass,
 		Types:         ArgTypes{{"arg", in}},
-		ReturnType:    ret,
+		ReturnType:    retType,
 		AggregateFunc: f,
-		WindowFunc: func() WindowFunc {
-			return newAggregateWindow(f())
+		WindowFunc: func(params []Type) WindowFunc {
+			return newAggregateWindow(f(params))
 		},
 		Info: info,
 	}
 }
 
-func makeAggBuiltins(f func() AggregateFunc, info string, types ...Type) []Builtin {
-	ret := make([]Builtin, len(types))
-	for i := range types {
-		ret[i] = makeAggBuiltin(types[i], types[i], f, info)
-	}
-	return ret
-}
-
-func countImpls() []Builtin {
-	types := []Type{TypeBool, TypeInt, TypeFloat, TypeDecimal, TypeString, TypeBytes,
-		TypeDate, TypeTimestamp, TypeTimestampTZ, TypeInterval, TypeTuple}
-	r := make([]Builtin, len(types))
-	for i := range types {
-		r[i] = makeAggBuiltin(types[i], TypeInt, newCountAggregate,
-			"Calculates the number of selected elements.")
-	}
-	return r
-}
-
 var _ AggregateFunc = &arrayAggregate{}
 var _ AggregateFunc = &avgAggregate{}
 var _ AggregateFunc = &countAggregate{}
@@ -248,12 +246,10 @@ type arrayAggregate struct {
 	arr *DArray
 }
 
-func newIntArrayAggregate() AggregateFunc {
-	return &arrayAggregate{arr: NewDArray(TypeInt)}
-}
-
-func newStringArrayAggregate() AggregateFunc {
-	return &arrayAggregate{arr: NewDArray(TypeString)}
+func newArrayAggregate(params []Type) AggregateFunc {
+	return &arrayAggregate{
+		arr: NewDArray(params[0]),
+	}
 }
 
 // Add accumulates the passed datum into the array.
@@ -276,14 +272,14 @@ type avgAggregate struct {
 	count int
 }
 
-func newIntAvgAggregate() AggregateFunc {
-	return &avgAggregate{agg: newIntSumAggregate()}
+func newIntAvgAggregate(params []Type) AggregateFunc {
+	return &avgAggregate{agg: newIntSumAggregate(params)}
 }
-func newFloatAvgAggregate() AggregateFunc {
-	return &avgAggregate{agg: newFloatSumAggregate()}
+func newFloatAvgAggregate(params []Type) AggregateFunc {
+	return &avgAggregate{agg: newFloatSumAggregate(params)}
 }
-func newDecimalAvgAggregate() AggregateFunc {
-	return &avgAggregate{agg: newDecimalSumAggregate()}
+func newDecimalAvgAggregate(params []Type) AggregateFunc {
+	return &avgAggregate{agg: newDecimalSumAggregate(params)}
 }
 
 // Add accumulates the passed datum into the average.
@@ -319,10 +315,10 @@ type concatAggregate struct {
 	result     bytes.Buffer
 }
 
-func newBytesConcatAggregate() AggregateFunc {
+func newBytesConcatAggregate(_ []Type) AggregateFunc {
 	return &concatAggregate{forBytes: true}
 }
-func newStringConcatAggregate() AggregateFunc {
+func newStringConcatAggregate(_ []Type) AggregateFunc {
 	return &concatAggregate{}
 }
 
@@ -357,7 +353,7 @@ type boolAndAggregate struct {
 	result     bool
 }
 
-func newBoolAndAggregate() AggregateFunc {
+func newBoolAndAggregate(_ []Type) AggregateFunc {
 	return &boolAndAggregate{}
 }
 
@@ -384,7 +380,7 @@ type boolOrAggregate struct {
 	result     bool
 }
 
-func newBoolOrAggregate() AggregateFunc {
+func newBoolOrAggregate(_ []Type) AggregateFunc {
 	return &boolOrAggregate{}
 }
 
@@ -407,7 +403,7 @@ type countAggregate struct {
 	count int
 }
 
-func newCountAggregate() AggregateFunc {
+func newCountAggregate(_ []Type) AggregateFunc {
 	return &countAggregate{}
 }
 
@@ -428,7 +424,7 @@ type MaxAggregate struct {
 	max Datum
 }
 
-func newMaxAggregate() AggregateFunc {
+func newMaxAggregate(_ []Type) AggregateFunc {
 	return &MaxAggregate{}
 }
 
@@ -460,7 +456,7 @@ type MinAggregate struct {
 	min Datum
 }
 
-func newMinAggregate() AggregateFunc {
+func newMinAggregate(_ []Type) AggregateFunc {
 	return &MinAggregate{}
 }
 
@@ -492,7 +488,7 @@ type smallIntSumAggregate struct {
 	seenNonNull bool
 }
 
-func newSmallIntSumAggregate() AggregateFunc {
+func newSmallIntSumAggregate(_ []Type) AggregateFunc {
 	return &smallIntSumAggregate{}
 }
 
@@ -525,7 +521,7 @@ type intSumAggregate struct {
 	seenNonNull bool
 }
 
-func newIntSumAggregate() AggregateFunc {
+func newIntSumAggregate(_ []Type) AggregateFunc {
 	return &intSumAggregate{}
 }
 
@@ -579,7 +575,7 @@ type decimalSumAggregate struct {
 	sawNonNull bool
 }
 
-func newDecimalSumAggregate() AggregateFunc {
+func newDecimalSumAggregate(_ []Type) AggregateFunc {
 	return &decimalSumAggregate{}
 }
 
@@ -608,7 +604,7 @@ type floatSumAggregate struct {
 	sawNonNull bool
 }
 
-func newFloatSumAggregate() AggregateFunc {
+func newFloatSumAggregate(_ []Type) AggregateFunc {
 	return &floatSumAggregate{}
 }
 
@@ -635,7 +631,7 @@ type intervalSumAggregate struct {
 	sawNonNull bool
 }
 
-func newIntervalSumAggregate() AggregateFunc {
+func newIntervalSumAggregate(_ []Type) AggregateFunc {
 	return &intervalSumAggregate{}
 }
 
@@ -663,7 +659,7 @@ type intVarianceAggregate struct {
 	tmpDec DDecimal
 }
 
-func newIntVarianceAggregate() AggregateFunc {
+func newIntVarianceAggregate(_ []Type) AggregateFunc {
 	return &intVarianceAggregate{}
 }
 
@@ -686,7 +682,7 @@ type floatVarianceAggregate struct {
 	sqrDiff float64
 }
 
-func newFloatVarianceAggregate() AggregateFunc {
+func newFloatVarianceAggregate(_ []Type) AggregateFunc {
 	return &floatVarianceAggregate{}
 }
 
@@ -723,7 +719,7 @@ type decimalVarianceAggregate struct {
 	tmp   inf.Dec
 }
 
-func newDecimalVarianceAggregate() AggregateFunc {
+func newDecimalVarianceAggregate(_ []Type) AggregateFunc {
 	return &decimalVarianceAggregate{}
 }
 
@@ -764,14 +760,14 @@ type stdDevAggregate struct {
 	agg AggregateFunc
 }
 
-func newIntStdDevAggregate() AggregateFunc {
-	return &stdDevAggregate{agg: newIntVarianceAggregate()}
+func newIntStdDevAggregate(params []Type) AggregateFunc {
+	return &stdDevAggregate{agg: newIntVarianceAggregate(params)}
 }
-func newFloatStdDevAggregate() AggregateFunc {
-	return &stdDevAggregate{agg: newFloatVarianceAggregate()}
+func newFloatStdDevAggregate(params []Type) AggregateFunc {
+	return &stdDevAggregate{agg: newFloatVarianceAggregate(params)}
 }
-func newDecimalStdDevAggregate() AggregateFunc {
-	return &stdDevAggregate{agg: newDecimalVarianceAggregate()}
+func newDecimalStdDevAggregate(params []Type) AggregateFunc {
+	return &stdDevAggregate{agg: newDecimalVarianceAggregate(params)}
 }
 
 // Add implements the AggregateFunc interface.

pkg/sql/parser/aggregate_builtins_test.go

@@ -29,8 +29,8 @@ import (
 // printing all values to strings once the accumulation has finished. If the string
 // slices are not equal, it means that the result Datums were modified during later
 // accumulation, which violates the "deep copy of any internal state" condition.
-func testAggregateResultDeepCopy(t *testing.T, aggFunc func() AggregateFunc, vals []Datum) {
-	aggImpl := aggFunc()
+func testAggregateResultDeepCopy(t *testing.T, aggFunc func([]Type) AggregateFunc, vals []Datum) {
+	aggImpl := aggFunc([]Type{vals[0].ResolvedType()})
 	runningDatums := make([]Datum, len(vals))
 	runningStrings := make([]string, len(vals))
 	for i := range vals {
@@ -223,10 +223,11 @@ func makeIntervalTestDatum(count int) []Datum {
 	return vals
 }
 
-func runBenchmarkAggregate(b *testing.B, aggFunc func() AggregateFunc, vals []Datum) {
+func runBenchmarkAggregate(b *testing.B, aggFunc func([]Type) AggregateFunc, vals []Datum) {
+	params := []Type{vals[0].ResolvedType()}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		aggImpl := aggFunc()
+		aggImpl := aggFunc(params)
 		for i := range vals {
 			aggImpl.Add(vals[i])
 		}