Skip to content

Commit 0396d70

Browse files
myersmyers
committed
Add ROWS FROM support for multi-argument UNNEST
Expands multi-argument UNNEST(arr1, arr2, ...) into ROWS FROM(unnest(arr1), unnest(arr2), ...) which properly zips results together with NULL padding for shorter arrays. Changes: - table_expr.go: Detect and expand multi-arg UNNEST to RowsFromExpr - aliased_table_expr.go: Support WITH ORDINALITY via RowsFromExpr - select_clause.go: Update multi-arg UNNEST rewrite to use RowsFromExpr Depends on: - dolthub/vitess#454 - dolthub/go-mysql-server#3412
1 parent 174bf4d commit 0396d70

File tree

3 files changed

+146
-36
lines changed

3 files changed

+146
-36
lines changed

‎server/ast/aliased_table_expr.go‎

Lines changed: 110 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
package ast
1616

1717
import (
18+
"strings"
19+
1820
"github.com/cockroachdb/errors"
1921

2022
vitess "github.com/dolthub/vitess/go/vt/sqlparser"
@@ -24,13 +26,117 @@ import (
2426
)
2527

2628
// nodeAliasedTableExpr handles *tree.AliasedTableExpr nodes.
27-
func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.AliasedTableExpr, error) {
28-
if node.Ordinality {
29-
return nil, errors.Errorf("ordinality is not yet supported")
30-
}
29+
func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (vitess.TableExpr, error) {
3130
if node.IndexFlags != nil {
3231
return nil, errors.Errorf("index flags are not yet supported")
3332
}
33+
34+
// Handle RowsFromExpr specially - it can have WITH ORDINALITY and column aliases
35+
if rowsFrom, ok := node.Expr.(*tree.RowsFromExpr); ok {
36+
// Handle multi-argument UNNEST specially: UNNEST(arr1, arr2, ...)
37+
// is syntactic sugar for ROWS FROM(unnest(arr1), unnest(arr2), ...)
38+
// We need to detect this case and expand it to use RowsFromExpr.
39+
if len(rowsFrom.Items) == 1 {
40+
if funcExpr, ok := rowsFrom.Items[0].(*tree.FuncExpr); ok {
41+
funcName := funcExpr.Func.String()
42+
if strings.EqualFold(funcName, "unnest") && len(funcExpr.Exprs) > 1 {
43+
// Expand multi-arg UNNEST into separate unnest calls
44+
selectExprs := make(vitess.SelectExprs, len(funcExpr.Exprs))
45+
for i, arg := range funcExpr.Exprs {
46+
argExpr, err := nodeExpr(ctx, arg)
47+
if err != nil {
48+
return nil, err
49+
}
50+
selectExprs[i] = &vitess.AliasedExpr{
51+
Expr: &vitess.FuncExpr{
52+
Name: vitess.NewColIdent("unnest"),
53+
Exprs: vitess.SelectExprs{&vitess.AliasedExpr{Expr: argExpr}},
54+
},
55+
}
56+
}
57+
58+
var columns vitess.Columns
59+
if len(node.As.Cols) > 0 {
60+
columns = make(vitess.Columns, len(node.As.Cols))
61+
for i := range node.As.Cols {
62+
columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
63+
}
64+
}
65+
66+
return &vitess.RowsFromExpr{
67+
Exprs: selectExprs,
68+
WithOrdinality: node.Ordinality,
69+
Alias: vitess.NewTableIdent(string(node.As.Alias)),
70+
Columns: columns,
71+
}, nil
72+
}
73+
}
74+
}
75+
76+
// For single functions or non-multi-arg-UNNEST cases, use the existing
77+
// subquery-based approach that works with the table function infrastructure.
78+
// Only WITH ORDINALITY requires the new RowsFromExpr approach.
79+
if node.Ordinality {
80+
// Use RowsFromExpr for WITH ORDINALITY support
81+
selectExprs := make(vitess.SelectExprs, len(rowsFrom.Items))
82+
for i, item := range rowsFrom.Items {
83+
expr, err := nodeExpr(ctx, item)
84+
if err != nil {
85+
return nil, err
86+
}
87+
selectExprs[i] = &vitess.AliasedExpr{Expr: expr}
88+
}
89+
90+
var columns vitess.Columns
91+
if len(node.As.Cols) > 0 {
92+
columns = make(vitess.Columns, len(node.As.Cols))
93+
for i := range node.As.Cols {
94+
columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
95+
}
96+
}
97+
98+
return &vitess.RowsFromExpr{
99+
Exprs: selectExprs,
100+
WithOrdinality: node.Ordinality,
101+
Alias: vitess.NewTableIdent(string(node.As.Alias)),
102+
Columns: columns,
103+
}, nil
104+
}
105+
106+
// For non-ordinality cases, fall through to use the existing
107+
// table function infrastructure via nodeTableExpr
108+
tableExpr, err := nodeTableExpr(ctx, rowsFrom)
109+
if err != nil {
110+
return nil, err
111+
}
112+
113+
// Wrap in a subquery as the original code did
114+
subquery := &vitess.Subquery{
115+
Select: &vitess.Select{
116+
From: vitess.TableExprs{tableExpr},
117+
},
118+
}
119+
120+
if len(node.As.Cols) > 0 {
121+
columns := make([]vitess.ColIdent, len(node.As.Cols))
122+
for i := range node.As.Cols {
123+
columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
124+
}
125+
subquery.Columns = columns
126+
}
127+
128+
return &vitess.AliasedTableExpr{
129+
Expr: subquery,
130+
As: vitess.NewTableIdent(string(node.As.Alias)),
131+
Lateral: node.Lateral,
132+
}, nil
133+
}
134+
135+
// For non-RowsFromExpr expressions, ordinality is not yet supported
136+
if node.Ordinality {
137+
return nil, errors.Errorf("ordinality is only supported for ROWS FROM expressions")
138+
}
139+
34140
var aliasExpr vitess.SimpleTableExpr
35141
var authInfo vitess.AuthInformation
36142

@@ -92,27 +198,6 @@ func nodeAliasedTableExpr(ctx *Context, node *tree.AliasedTableExpr) (*vitess.Al
92198
Select: selectStmt,
93199
}
94200

95-
if len(node.As.Cols) > 0 {
96-
columns := make([]vitess.ColIdent, len(node.As.Cols))
97-
for i := range node.As.Cols {
98-
columns[i] = vitess.NewColIdent(string(node.As.Cols[i]))
99-
}
100-
subquery.Columns = columns
101-
}
102-
aliasExpr = subquery
103-
case *tree.RowsFromExpr:
104-
tableExpr, err := nodeTableExpr(ctx, expr)
105-
if err != nil {
106-
return nil, err
107-
}
108-
109-
// TODO: this should be represented as a table function more directly
110-
subquery := &vitess.Subquery{
111-
Select: &vitess.Select{
112-
From: vitess.TableExprs{tableExpr},
113-
},
114-
}
115-
116201
if len(node.As.Cols) > 0 {
117202
columns := make([]vitess.ColIdent, len(node.As.Cols))
118203
for i := range node.As.Cols {

‎server/ast/select_clause.go‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,8 @@ PostJoinRewrite:
162162
// Handle multi-argument UNNEST: UNNEST(arr1, arr2, ...) produces a table with one column per array,
163163
// where corresponding elements are "zipped" together. PostgreSQL pads shorter arrays with NULLs.
164164
// We transform: SELECT * FROM UNNEST(arr1, arr2)
165-
// Into: SELECT * FROM (SELECT unnest(arr1), unnest(arr2)) AS unnest
166-
// GMS's ProjectRowWithNestedIters handles multiple SRFs by zipping them together correctly.
165+
// Into: SELECT * FROM ROWS FROM(unnest(arr1), unnest(arr2)) AS unnest
166+
// This uses the native ROWS FROM table function which properly zips SRFs together.
167167
if tableFuncExpr, ok := from[i].(*vitess.TableFuncExpr); ok {
168168
if strings.EqualFold(tableFuncExpr.Name, "unnest") && len(tableFuncExpr.Exprs) > 1 {
169169
selectExprs := make(vitess.SelectExprs, 0, len(tableFuncExpr.Exprs))
@@ -179,13 +179,9 @@ PostJoinRewrite:
179179
if alias.IsEmpty() {
180180
alias = vitess.NewTableIdent("unnest")
181181
}
182-
from[i] = &vitess.AliasedTableExpr{
183-
Expr: &vitess.Subquery{
184-
Select: &vitess.Select{
185-
SelectExprs: selectExprs,
186-
},
187-
},
188-
As: alias,
182+
from[i] = &vitess.RowsFromExpr{
183+
Exprs: selectExprs,
184+
Alias: alias,
189185
}
190186
}
191187
}

‎server/ast/table_expr.go‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
package ast
1616

1717
import (
18+
"strings"
19+
1820
"github.com/cockroachdb/errors"
1921

2022
vitess "github.com/dolthub/vitess/go/vt/sqlparser"
@@ -99,12 +101,39 @@ func nodeTableExpr(ctx *Context, node tree.TableExpr) (vitess.TableExpr, error)
99101
Exprs: vitess.TableExprs{tableExpr},
100102
}, nil
101103
case *tree.RowsFromExpr:
104+
// Handle multi-argument UNNEST specially: UNNEST(arr1, arr2, ...)
105+
// is syntactic sugar for ROWS FROM(unnest(arr1), unnest(arr2), ...)
106+
// We need to detect this case and expand it to use RowsFromExpr.
107+
if len(node.Items) == 1 {
108+
if funcExpr, ok := node.Items[0].(*tree.FuncExpr); ok {
109+
funcName := funcExpr.Func.String()
110+
if strings.EqualFold(funcName, "unnest") && len(funcExpr.Exprs) > 1 {
111+
// Expand multi-arg UNNEST into separate unnest calls
112+
selectExprs := make(vitess.SelectExprs, len(funcExpr.Exprs))
113+
for i, arg := range funcExpr.Exprs {
114+
argExpr, err := nodeExpr(ctx, arg)
115+
if err != nil {
116+
return nil, err
117+
}
118+
selectExprs[i] = &vitess.AliasedExpr{
119+
Expr: &vitess.FuncExpr{
120+
Name: vitess.NewColIdent("unnest"),
121+
Exprs: vitess.SelectExprs{&vitess.AliasedExpr{Expr: argExpr}},
122+
},
123+
}
124+
}
125+
return &vitess.RowsFromExpr{
126+
Exprs: selectExprs,
127+
}, nil
128+
}
129+
}
130+
}
131+
// For single functions or other cases, use the original ValuesStatement approach
132+
// which works with the existing table function infrastructure
102133
exprs, err := nodeExprs(ctx, node.Items)
103134
if err != nil {
104135
return nil, err
105136
}
106-
//TODO: not sure if this is correct at all. I think we want to return one result per row, but maybe not.
107-
// This needs to be tested to verify.
108137
rows := make([]vitess.ValTuple, len(exprs))
109138
for i := range exprs {
110139
rows[i] = vitess.ValTuple{exprs[i]}

0 commit comments

Comments
 (0)