Fix UNNEST bulk insert column mapping

server/ast/aliased_table_expr.go

@@ -15,6 +15,8 @@
 package ast
 
 import (
+	"strings"
+
 	"github.com/cockroachdb/errors"
 
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
@@ -24,13 +26,116 @@ import (
 )
 
 // nodeAliasedTableExpr handles *tree.AliasedTableExpr nodes.
-func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.AliasedTableExpr, error) {
-	if node.Ordinality {
-		return nil, errors.Errorf("ordinality is not yet supported")
-	}
+func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (vitess.TableExpr, error) {
 	if node.IndexFlags != nil {
 		return nil, errors.Errorf("index flags are not yet supported")
 	}
+
+	// Handle RowsFromExpr specially - it can have WITH ORDINALITY and column aliases
+	if rowsFrom, ok := node.Expr.(*tree.RowsFromExpr); ok {
+		// Handle multi-argument UNNEST specially: UNNEST(arr1, arr2, ...)
+		// is syntactic sugar for ROWS FROM(unnest(arr1), unnest(arr2), ...)
+		// We need to detect this case and expand it to use RowsFromExpr.
+		if len(rowsFrom.Items) == 1 {
+			if funcExpr, ok := rowsFrom.Items[0].(*tree.FuncExpr); ok {
+				funcName := funcExpr.Func.String()
+				if strings.EqualFold(funcName, "unnest") && len(funcExpr.Exprs) > 1 {
+					// Expand multi-arg UNNEST into separate unnest calls
+					selectExprs := make(vitess.SelectExprs, len(funcExpr.Exprs))
+					for i, arg := range funcExpr.Exprs {
+						argExpr, err := nodeExpr(ctx, arg)
+						if err != nil {
+							return nil, err
+						}
+						selectExprs[i] = &vitess.AliasedExpr{
+							Expr: &vitess.FuncExpr{
+								Name:  vitess.NewColIdent("unnest"),
+								Exprs: vitess.SelectExprs{&vitess.AliasedExpr{Expr: argExpr}},
+							},
+						}
+					}
+
+					var columns vitess.Columns
+					if len(node.As.Cols) > 0 {
+						columns = make(vitess.Columns, len(node.As.Cols))
+						for i := range node.As.Cols {
+							columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
+						}
+					}
+
+					return &vitess.RowsFromExpr{
+						Exprs:          selectExprs,
+						WithOrdinality: node.Ordinality,
+						Alias:          vitess.NewTableIdent(string(node.As.Alias)),
+						Columns:        columns,
+					}, nil
+				}
+			}
+		}
+
+		// Use RowsFromExpr for:
+		// 1. Multiple functions: ROWS FROM(func1(), func2()) AS alias
+		// 2. WITH ORDINALITY: ROWS FROM(func()) WITH ORDINALITY
+		if len(rowsFrom.Items) > 1 || node.Ordinality {
+			selectExprs := make(vitess.SelectExprs, len(rowsFrom.Items))
+			for i, item := range rowsFrom.Items {
+				expr, err := nodeExpr(ctx, item)
+				if err != nil {
+					return nil, err
+				}
+				selectExprs[i] = &vitess.AliasedExpr{Expr: expr}
+			}
+
+			var columns vitess.Columns
+			if len(node.As.Cols) > 0 {
+				columns = make(vitess.Columns, len(node.As.Cols))
+				for i := range node.As.Cols {
+					columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
+				}
+			}
+
+			return &vitess.RowsFromExpr{
+				Exprs:          selectExprs,
+				WithOrdinality: node.Ordinality,
+				Alias:          vitess.NewTableIdent(string(node.As.Alias)),
+				Columns:        columns,
+			}, nil
+		}
+
+		// For single function without ordinality, fall through to use the existing
+		// table function infrastructure via nodeTableExpr
+		tableExpr, err := nodeTableExpr(ctx, rowsFrom)
+		if err != nil {
+			return nil, err
+		}
+
+		// Wrap in a subquery as the original code did
+		subquery := &vitess.Subquery{
+			Select: &vitess.Select{
+				From: vitess.TableExprs{tableExpr},
+			},
+		}
+
+		if len(node.As.Cols) > 0 {
+			columns := make([]vitess.ColIdent, len(node.As.Cols))
+			for i := range node.As.Cols {
+				columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
+			}
+			subquery.Columns = columns
+		}
+
+		return &vitess.AliasedTableExpr{
+			Expr:    subquery,
+			As:      vitess.NewTableIdent(string(node.As.Alias)),
+			Lateral: node.Lateral,
+		}, nil
+	}
+
+	// For non-RowsFromExpr expressions, ordinality is not yet supported
+	if node.Ordinality {
+		return nil, errors.Errorf("ordinality is only supported for ROWS FROM expressions")
+	}
+
 	var aliasExpr vitess.SimpleTableExpr
 	var authInfo vitess.AuthInformation
 
@@ -92,27 +197,6 @@ func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.Al
 			Select: selectStmt,
 		}
 
-		if len(node.As.Cols) > 0 {
-			columns := make([]vitess.ColIdent, len(node.As.Cols))
-			for i := range node.As.Cols {
-				columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
-			}
-			subquery.Columns = columns
-		}
-		aliasExpr = subquery
-	case *tree.RowsFromExpr:
-		tableExpr, err := nodeTableExpr(ctx, expr)
-		if err != nil {
-			return nil, err
-		}
-
-		// TODO: this should be represented as a table function more directly
-		subquery := &vitess.Subquery{
-			Select: &vitess.Select{
-				From: vitess.TableExprs{tableExpr},
-			},
-		}
-
 		if len(node.As.Cols) > 0 {
 			columns := make([]vitess.ColIdent, len(node.As.Cols))
 			for i := range node.As.Cols {

server/ast/select_clause.go

@@ -15,8 +15,9 @@
 package ast
 
 import (
-	"github.com/dolthub/go-mysql-server/sql/expression"
+	"strings"
 
+	"github.com/dolthub/go-mysql-server/sql/expression"
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
 
 	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
@@ -158,6 +159,32 @@ PostJoinRewrite:
 				}
 			}
 		}
+		// Handle multi-argument UNNEST: UNNEST(arr1, arr2, ...) produces a table with one column per array,
+		// where corresponding elements are "zipped" together. PostgreSQL pads shorter arrays with NULLs.
+		// We transform: SELECT * FROM UNNEST(arr1, arr2)
+		// Into:         SELECT * FROM ROWS FROM(unnest(arr1), unnest(arr2)) AS unnest
+		// This uses the native ROWS FROM table function which properly zips SRFs together.
+		if tableFuncExpr, ok := from[i].(*vitess.TableFuncExpr); ok {
+			if strings.EqualFold(tableFuncExpr.Name, "unnest") && len(tableFuncExpr.Exprs) > 1 {
+				selectExprs := make(vitess.SelectExprs, 0, len(tableFuncExpr.Exprs))
+				for _, argExpr := range tableFuncExpr.Exprs {
+					selectExprs = append(selectExprs, &vitess.AliasedExpr{
+						Expr: &vitess.FuncExpr{
+							Name:  vitess.NewColIdent("unnest"),
+							Exprs: vitess.SelectExprs{argExpr},
+						},
+					})
+				}
+				alias := tableFuncExpr.Alias
+				if alias.IsEmpty() {
+					alias = vitess.NewTableIdent("unnest")
+				}
+				from[i] = &vitess.RowsFromExpr{
+					Exprs: selectExprs,
+					Alias: alias,
+				}
+			}
+		}
 	}
 	distinct := node.Distinct
 	var distinctOn vitess.Exprs

server/ast/table_expr.go

@@ -15,6 +15,8 @@
 package ast
 
 import (
+	"strings"
+
 	"github.com/cockroachdb/errors"
 
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
@@ -99,12 +101,56 @@ func nodeTableExpr(ctx *Context, node tree.TableExpr) (vitess.TableExpr, error)
 			Exprs: vitess.TableExprs{tableExpr},
 		}, nil
 	case *tree.RowsFromExpr:
+		// Handle multi-argument UNNEST specially: UNNEST(arr1, arr2, ...)
+		// is syntactic sugar for ROWS FROM(unnest(arr1), unnest(arr2), ...)
+		// We need to detect this case and expand it to use RowsFromExpr.
+		if len(node.Items) == 1 {
+			if funcExpr, ok := node.Items[0].(*tree.FuncExpr); ok {
+				funcName := funcExpr.Func.String()
+				if strings.EqualFold(funcName, "unnest") && len(funcExpr.Exprs) > 1 {
+					// Expand multi-arg UNNEST into separate unnest calls
+					selectExprs := make(vitess.SelectExprs, len(funcExpr.Exprs))
+					for i, arg := range funcExpr.Exprs {
+						argExpr, err := nodeExpr(ctx, arg)
+						if err != nil {
+							return nil, err
+						}
+						selectExprs[i] = &vitess.AliasedExpr{
+							Expr: &vitess.FuncExpr{
+								Name:  vitess.NewColIdent("unnest"),
+								Exprs: vitess.SelectExprs{&vitess.AliasedExpr{Expr: argExpr}},
+							},
+						}
+					}
+					return &vitess.RowsFromExpr{
+						Exprs: selectExprs,
+					}, nil
+				}
+			}
+		}
+
+		// For explicit ROWS FROM with multiple functions, use RowsFromExpr
+		// This handles: ROWS FROM(generate_series(1,3), generate_series(10,12))
+		if len(node.Items) > 1 {
+			selectExprs := make(vitess.SelectExprs, len(node.Items))
+			for i, item := range node.Items {
+				expr, err := nodeExpr(ctx, item)
+				if err != nil {
+					return nil, err
+				}
+				selectExprs[i] = &vitess.AliasedExpr{Expr: expr}
+			}
+			return &vitess.RowsFromExpr{
+				Exprs: selectExprs,
+			}, nil
+		}
+
+		// For single functions, use the original ValuesStatement approach
+		// which works with the existing table function infrastructure
 		exprs, err := nodeExprs(ctx, node.Items)
 		if err != nil {
 			return nil, err
 		}
-		//TODO: not sure if this is correct at all. I think we want to return one result per row, but maybe not.
-		// This needs to be tested to verify.
 		rows := make([]vitess.ValTuple, len(exprs))
 		for i := range exprs {
 			rows[i] = vitess.ValTuple{exprs[i]}

testing/go/functions_test.go

@@ -1552,6 +1552,91 @@ func TestArrayFunctions(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name: "multi-argument unnest",
+			Assertions: []ScriptTestAssertion{
+				{
+					// Basic multi-argument UNNEST with equal-length arrays
+					Query:    `SELECT * FROM UNNEST(ARRAY['a','b','c'], ARRAY[1,2,3])`,
+					Expected: []sql.Row{{"a", int64(1)}, {"b", int64(2)}, {"c", int64(3)}},
+				},
+				{
+					// Multi-argument UNNEST with unequal-length arrays (shorter padded with NULL)
+					Query:    `SELECT * FROM UNNEST(ARRAY['a','b'], ARRAY[1,2,3])`,
+					Expected: []sql.Row{{"a", int64(1)}, {"b", int64(2)}, {nil, int64(3)}},
+				},
+				{
+					// Multi-argument UNNEST with empty array
+					Query:    `SELECT * FROM UNNEST(ARRAY['a','b'], ARRAY[]::int[])`,
+					Expected: []sql.Row{{"a", nil}, {"b", nil}},
+				},
+				{
+					// Multi-argument UNNEST with three arrays (booleans come as "t"/"f" strings from PostgreSQL wire protocol)
+					Query:    `SELECT * FROM UNNEST(ARRAY[1,2], ARRAY['x','y'], ARRAY[true,false])`,
+					Expected: []sql.Row{{int64(1), "x", "t"}, {int64(2), "y", "f"}},
+				},
+				{
+					// Multi-argument UNNEST with alias
+					Query:    `SELECT u.* FROM UNNEST(ARRAY['a','b'], ARRAY[1,2]) AS u`,
+					Expected: []sql.Row{{"a", int64(1)}, {"b", int64(2)}},
+				},
+			},
+		},
+		{
+			Name: "ROWS FROM with generate_series",
+			Assertions: []ScriptTestAssertion{
+				{
+					// Basic ROWS FROM with two generate_series calls
+					Query:    `SELECT * FROM ROWS FROM(generate_series(1,3), generate_series(10,12))`,
+					Expected: []sql.Row{{int64(1), int64(10)}, {int64(2), int64(11)}, {int64(3), int64(12)}},
+				},
+				{
+					// ROWS FROM with unequal-length series (shorter padded with NULL)
+					Query:    `SELECT * FROM ROWS FROM(generate_series(1,2), generate_series(10,13))`,
+					Expected: []sql.Row{{int64(1), int64(10)}, {int64(2), int64(11)}, {nil, int64(12)}, {nil, int64(13)}},
+				},
+				{
+					// ROWS FROM with table alias
+					Query:    `SELECT r.* FROM ROWS FROM(generate_series(1,2), generate_series(10,11)) AS r`,
+					Expected: []sql.Row{{int64(1), int64(10)}, {int64(2), int64(11)}},
+				},
+				{
+					// ROWS FROM with three functions
+					Query:    `SELECT * FROM ROWS FROM(generate_series(1,2), generate_series(10,11), generate_series(100,101))`,
+					Expected: []sql.Row{{int64(1), int64(10), int64(100)}, {int64(2), int64(11), int64(101)}},
+				},
+			},
+		},
+		{
+			Name: "ROWS FROM with unnest",
+			Assertions: []ScriptTestAssertion{
+				{
+					// ROWS FROM with explicit unnest calls
+					Query:    `SELECT * FROM ROWS FROM(unnest(ARRAY['a','b']), unnest(ARRAY[1,2]))`,
+					Expected: []sql.Row{{"a", int64(1)}, {"b", int64(2)}},
+				},
+				{
+					// ROWS FROM with unequal-length unnest
+					Query:    `SELECT * FROM ROWS FROM(unnest(ARRAY['a','b','c']), unnest(ARRAY[1,2]))`,
+					Expected: []sql.Row{{"a", int64(1)}, {"b", int64(2)}, {"c", nil}},
+				},
+			},
+		},
+		{
+			Name: "ROWS FROM mixed functions",
+			Assertions: []ScriptTestAssertion{
+				{
+					// Mix generate_series and unnest
+					Query:    `SELECT * FROM ROWS FROM(generate_series(1,3), unnest(ARRAY['x','y','z']))`,
+					Expected: []sql.Row{{int64(1), "x"}, {int64(2), "y"}, {int64(3), "z"}},
+				},
+				{
+					// Mix with unequal lengths
+					Query:    `SELECT * FROM ROWS FROM(generate_series(1,2), unnest(ARRAY['x','y','z']))`,
+					Expected: []sql.Row{{int64(1), "x"}, {int64(2), "y"}, {nil, "z"}},
+				},
+			},
+		},
 		{
 			Name:        "array_to_string",
 			SetUpScript: []string{},

testing/go/insert_test.go

@@ -293,5 +293,27 @@ ON CONFLICT (id) do update set c1 = $4`,
 				},
 			},
 		},
+		{
+			Name: "insert from unnest",
+			SetUpScript: []string{
+				`CREATE TABLE "django_content_type" (id serial primary key, app_label varchar, model varchar)`,
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query: `INSERT INTO "django_content_type" ("app_label", "model")
+SELECT * FROM UNNEST(('{debug_app,debug_app}')::varchar[],
+                     ('{debugmodel1,debugmodel2}')::varchar[])
+RETURNING "django_content_type"."id"`,
+					Expected: []sql.Row{{1}, {2}},
+				},
+				{
+					Query: `SELECT "app_label", "model" FROM "django_content_type" ORDER BY "id"`,
+					Expected: []sql.Row{
+						{"debug_app", "debugmodel1"},
+						{"debug_app", "debugmodel2"},
+					},
+				},
+			},
+		},
 	})
 }