Fix UNNEST bulk insert column mapping

server/analyzer/init.go

@@ -15,13 +15,15 @@
 package analyzer
 
 import (
+	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/analyzer"
 	"github.com/dolthub/go-mysql-server/sql/expression"
 	"github.com/dolthub/go-mysql-server/sql/memo"
 	"github.com/dolthub/go-mysql-server/sql/plan"
 	"github.com/dolthub/go-mysql-server/sql/planbuilder"
 
 	pgexpression "github.com/dolthub/doltgresql/server/expression"
+	pgnodes "github.com/dolthub/doltgresql/server/node"
 )
 
 // IDs are basically arbitrary, we just need to ensure that they do not conflict with existing IDs
@@ -114,6 +116,13 @@ func initEngine() {
 
 	planbuilder.IsAggregateFunc = IsAggregateFunc
 
+	planbuilder.BuildMultiExprTableFunc = func(
+		exprs []sql.Expression, alias string,
+		withOrdinality bool, columnAliases []string,
+	) (sql.Node, error) {
+		return pgnodes.NewRowsFrom(exprs, alias, withOrdinality, columnAliases), nil
+	}
+
 	expression.DefaultExpressionFactory = pgexpression.PostgresExpressionFactory{}
 
 	// There are a couple places during analysis where SplitConjunction in GMS cannot correctly split up

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,115 @@
 )
 
 // 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), ...)
+		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.TableFuncExpr{
+						Exprs:          selectExprs,
+						WithOrdinality: node.Ordinality,
+						Alias:          vitess.NewTableIdent(string(node.As.Alias)),
+						Columns:        columns,
+					}, nil
+				}
+			}
+		}
+
+		// Use TableFuncExpr (nameless) 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.TableFuncExpr{
+				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 +196,6 @@
 			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.TableFuncExpr{
+					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,55 @@ 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), ...)
+		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.TableFuncExpr{
+						Exprs: selectExprs,
+					}, nil
+				}
+			}
+		}
+
+		// For explicit ROWS FROM with multiple functions, use TableFuncExpr (nameless)
+		// 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.TableFuncExpr{
+				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]}