nested iter should be used for set returning functions as projections only

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/dolthub/dolt/go v0.40.5-0.20250919235907-2afe4e108385
 	github.com/dolthub/eventsapi_schema v0.0.0-20250915094920-eadfd39051ca
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
-	github.com/dolthub/go-mysql-server v0.20.1-0.20250919010705-b0670cfce79c
+	github.com/dolthub/go-mysql-server v0.20.1-0.20250922225301-7b1b2ec8fb3a
 	github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216
 	github.com/dolthub/vitess v0.0.0-20250918181259-ed0e1c5cb192
 	github.com/fatih/color v1.13.0

go.sum

@@ -258,8 +258,8 @@ github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20250916051405-78a38d478790 h1:zxMsH7RLiG+dlZ/y0LgJHTV26XoiSJcuWq+em6t6VVc=
 github.com/dolthub/go-icu-regex v0.0.0-20250916051405-78a38d478790/go.mod h1:F3cnm+vMRK1HaU6+rNqQrOCyR03HHhR1GWG2gnPOqaE=
-github.com/dolthub/go-mysql-server v0.20.1-0.20250919010705-b0670cfce79c h1:0vvuSeatbCpxh7NgCti9mhZTDCgrgKwIh9YP39EgEBU=
-github.com/dolthub/go-mysql-server v0.20.1-0.20250919010705-b0670cfce79c/go.mod h1:vYpsrVFnvdhCn3bD7Jsy7lcZsknq9HrqfNpi/DZOg18=
+github.com/dolthub/go-mysql-server v0.20.1-0.20250922225301-7b1b2ec8fb3a h1:qTp2CM5LveSYVic/sTAuI8LCtIBeIgQwt+bMnjQKjVg=
+github.com/dolthub/go-mysql-server v0.20.1-0.20250922225301-7b1b2ec8fb3a/go.mod h1:vYpsrVFnvdhCn3bD7Jsy7lcZsknq9HrqfNpi/DZOg18=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q=
 github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=

server/analyzer/optimize_functions.go

@@ -15,6 +15,7 @@
 package analyzer
 
 import (
+	"github.com/cockroachdb/errors"
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/analyzer"
 	"github.com/dolthub/go-mysql-server/sql/plan"
@@ -25,37 +26,74 @@ import (
 )
 
 // OptimizeFunctions replaces all functions that fit specific criteria with their optimized variants. Also handles
-// SRFs (set-returning functions) by setting the `IncludesNestedIters` flag on the Project node if any SRF is found.
+// SRFs (set-returning functions) by setting the `IncludesNestedIters` flag on the Project node if any SRF is found
+// inside projection expressions.
 func OptimizeFunctions(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope *plan.Scope, selector analyzer.RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
 	// This is supposed to be one of the last rules to run. Subqueries break that assumption, so we skip this rule in such cases.
 	if scope != nil && scope.CurrentNodeIsFromSubqueryExpression {
 		return node, transform.SameTree, nil
 	}
 
+	_, isInsertNode := node.(*plan.InsertInto)
 	return pgtransform.NodeWithOpaque(node, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
-		_, ok := n.(*plan.Project)
+		projectNode, ok := n.(*plan.Project)
 		if !ok {
 			return n, transform.SameTree, nil
 		}
 
+		hasMultipleExpressionTuples := false
 		hasSRF := false
-		n, same, err := transform.NodeExprs(n, func(expr sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
+		// Check if there is set returning function in the source node (e.g. SELECT * FROM unnest())
+		n, sameNode, err := transform.NodeExprsWithNode(projectNode.Child, func(in sql.Node, expr sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
 			if compiledFunction, ok := expr.(*framework.CompiledFunction); ok {
 				hasSRF = hasSRF || compiledFunction.IsSRF()
 				if quickFunction := compiledFunction.GetQuickFunction(); quickFunction != nil {
 					return quickFunction, transform.NewTree, nil
 				}
 			}
+			if v, ok := in.(*plan.Values); ok {
+				hasMultipleExpressionTuples = len(v.ExpressionTuples) > 1
+			}
 			return expr, transform.SameTree, nil
 		})
+		if err != nil {
+			return nil, transform.SameTree, err
+		}
+		if !sameNode {
+			projectNode.Child = n
+		}
+
+		// insert node cannot have more than 1 row value if it has set returning function
+		if isInsertNode && hasMultipleExpressionTuples && hasSRF {
+			return nil, false, errors.Errorf("set-returning functions are not allowed in VALUES")
+		}
+
+		// Check if there is set returning function in the projection expressions (e.g. SELECT unnest() [FROM table/srf])
+		hasSRFInProjection := false
+		exprs, sameExprs, err := transform.Exprs(projectNode.Projections, func(expr sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
+			if compiledFunction, ok := expr.(*framework.CompiledFunction); ok {
+				hasSRFInProjection = hasSRFInProjection || compiledFunction.IsSRF()
+				if quickFunction := compiledFunction.GetQuickFunction(); quickFunction != nil {
+					return quickFunction, transform.NewTree, nil
+				}
+			}
+			return expr, transform.SameTree, nil
+		})
+		if err != nil {
+			return nil, transform.SameTree, err
+		}
+		if !sameExprs {
+			projectNode.Projections = exprs
+		}
 
-		if hasSRF {
+		// nested iter is used for set returning functions in the projections only
+		if hasSRFInProjection {
 			// Under some conditions, there will be no quick-function replacement, but changing the Project node to include
 			// nested iterators is still a change we need to tell the transform functions about.
-			same = transform.NewTree
-			n = n.(*plan.Project).WithIncludesNestedIters(true)
+			sameExprs = transform.NewTree
+			projectNode = projectNode.WithIncludesNestedIters(true)
 		}
 
-		return n, same, err
+		return projectNode, sameNode && sameExprs, err
 	})
 }

testing/go/functions_test.go

@@ -3945,6 +3945,35 @@ func TestSetReturningFunctions(t *testing.T) {
 					},
 				},
 			},
+			{
+				Name: "generate_series as table function and projection",
+				Assertions: []ScriptTestAssertion{
+					{
+						Query:    `SELECT *, unnest(ARRAY['cat', 'dog', 'bird']) AS animal FROM generate_series(1, 3);`,
+						Expected: []sql.Row{{1, "cat"}, {1, "dog"}, {1, "bird"}, {2, "cat"}, {2, "dog"}, {2, "bird"}, {3, "cat"}, {3, "dog"}, {3, "bird"}},
+					},
+				},
+			},
+			{
+				Name: "insert with set returning function",
+				SetUpScript: []string{
+					"create table hash_parted (a int, b int, c int);",
+				},
+				Assertions: []ScriptTestAssertion{
+					{
+						Query:    `insert into hash_parted values(0, generate_series(1,3), generate_series(5,8));`,
+						Expected: []sql.Row{},
+					},
+					{
+						Query:       `insert into hash_parted values(0, generate_series(11,12), generate_series(51,54)), (1, generate_series(1,3), generate_series(5,8));`,
+						ExpectedErr: `set-returning functions are not allowed in VALUES`,
+					},
+					{
+						Query:    `select * from hash_parted;`,
+						Expected: []sql.Row{{0, 1, 5}, {0, 2, 6}, {0, 3, 7}, {0, nil, 8}},
+					},
+				},
+			},
 		},
 	)
 }