Merge #70699

70699: sql: make builtin generators with GeneratorWithExprs safer r=mgartner a=mgartner

The `Overload.GeneratorWithExprs` field was added in #70115. This commit
makes usage of this field safer by introducing the `IsGenerator` method
which can be used to determine if an overload is an SRF. This commit
also adds guardrails to prevent future issues with
`GeneratorWithExprs`in the `ConvertZipArraysToValues` normalization rule.

Release note: None

Co-authored-by: Marcus Gartner <marcus@cockroachlabs.com>

pkg/sql/opt/memo/statistics_builder.go

@@ -2257,7 +2257,7 @@ func (sb *statisticsBuilder) buildProjectSet(
 	var zipRowCount float64
 	for i := range projectSet.Zip {
 		if fn, ok := projectSet.Zip[i].Fn.(*FunctionExpr); ok {
-			if fn.Overload.Generator != nil {
+			if fn.Overload.IsGenerator() {
 				// TODO(rytaft): We may want to estimate the number of rows based on
 				// the type of generator function and its parameters.
 				zipRowCount = unknownGeneratorRowCount
@@ -2313,7 +2313,7 @@ func (sb *statisticsBuilder) colStatProjectSet(
 		for i := range projectSet.Zip {
 			item := &projectSet.Zip[i]
 			if item.Cols.ToSet().Intersects(reqZipCols) {
-				if fn, ok := item.Fn.(*FunctionExpr); ok && fn.Overload.Generator != nil {
+				if fn, ok := item.Fn.(*FunctionExpr); ok && fn.Overload.IsGenerator() {
 					// The columns(s) contain a generator function.
 					// TODO(rytaft): We may want to determine which generator function the
 					// requested columns correspond to, and estimate the distinct count and

pkg/sql/opt/norm/project_set_funcs.go

@@ -50,6 +50,10 @@ func (c *CustomFuncs) CanConstructValuesFromZips(zip memo.ZipExpr) bool {
 			// Argument is not an ArrayExpr or ConstExpr wrapping a DArray or DJSON.
 			return false
 		}
+		if fn.Overload.GeneratorWithExprs != nil {
+			// ConstructValuesFromZips does not handle GeneratorWithExprs.
+			return false
+		}
 	}
 	return true
 }
@@ -113,6 +117,9 @@ func (c *CustomFuncs) ConstructValuesFromZips(zip memo.ZipExpr) memo.RelExpr {
 			// Use a ValueGenerator to retrieve values from the datums wrapped
 			// in the ConstExpr. These generators are used at runtime to unnest
 			// values from regular and JSON arrays.
+			if function.Overload.GeneratorWithExprs != nil {
+				panic(errors.AssertionFailedf("unexpected GeneratorWithExprs"))
+			}
 			generator, err := function.Overload.Generator(c.f.evalCtx, tree.Datums{t.Value})
 			if err != nil {
 				panic(errors.AssertionFailedf("generator retrieval failed: %v", err))

pkg/sql/pg_catalog.go

@@ -2182,7 +2182,7 @@ https://www.postgresql.org/docs/9.5/catalog-pg-proc.html`,
 						isRetSet := false
 						if fixedRetType := builtin.FixedReturnType(); fixedRetType != nil {
 							var retOid oid.Oid
-							if fixedRetType.Family() == types.TupleFamily && builtin.Generator != nil {
+							if fixedRetType.Family() == types.TupleFamily && builtin.IsGenerator() {
 								isRetSet = true
 								// Functions returning tables with zero, or more than one
 								// columns are marked to return "anyelement"

pkg/sql/sem/tree/overload.go

@@ -151,6 +151,11 @@ func (b Overload) InferReturnTypeFromInputArgTypes(inputTypes []*types.T) *types
 	return retTyp
 }
 
+// IsGenerator returns true if the function is a set returning function (SRF).
+func (b Overload) IsGenerator() bool {
+	return b.Generator != nil || b.GeneratorWithExprs != nil
+}
+
 // Signature returns a human-readable signature.
 // If simplify is bool, tuple-returning functions with just
 // 1 tuple element unwrap the return type in the signature.