1515package analyzer
1616
1717import (
18+ "github.com/cockroachdb/errors"
1819 "github.com/dolthub/go-mysql-server/sql"
1920 "github.com/dolthub/go-mysql-server/sql/analyzer"
2021 "github.com/dolthub/go-mysql-server/sql/plan"
@@ -25,37 +26,74 @@ import (
2526)
2627
2728// OptimizeFunctions replaces all functions that fit specific criteria with their optimized variants. Also handles
28- // SRFs (set-returning functions) by setting the `IncludesNestedIters` flag on the Project node if any SRF is found.
29+ // SRFs (set-returning functions) by setting the `IncludesNestedIters` flag on the Project node if any SRF is found
30+ // inside projection expressions.
2931func OptimizeFunctions (ctx * sql.Context , a * analyzer.Analyzer , node sql.Node , scope * plan.Scope , selector analyzer.RuleSelector , qFlags * sql.QueryFlags ) (sql.Node , transform.TreeIdentity , error ) {
3032 // This is supposed to be one of the last rules to run. Subqueries break that assumption, so we skip this rule in such cases.
3133 if scope != nil && scope .CurrentNodeIsFromSubqueryExpression {
3234 return node , transform .SameTree , nil
3335 }
3436
37+ _ , isInsertNode := node .(* plan.InsertInto )
3538 return pgtransform .NodeWithOpaque (node , func (n sql.Node ) (sql.Node , transform.TreeIdentity , error ) {
36- _ , ok := n .(* plan.Project )
39+ projectNode , ok := n .(* plan.Project )
3740 if ! ok {
3841 return n , transform .SameTree , nil
3942 }
4043
44+ hasMultipleExpressionTuples := false
4145 hasSRF := false
42- n , same , err := transform .NodeExprs (n , func (expr sql.Expression ) (sql.Expression , transform.TreeIdentity , error ) {
46+ // Check if there is set returning function in the source node (e.g. SELECT * FROM unnest())
47+ n , sameNode , err := transform .NodeExprsWithNode (projectNode .Child , func (in sql.Node , expr sql.Expression ) (sql.Expression , transform.TreeIdentity , error ) {
4348 if compiledFunction , ok := expr .(* framework.CompiledFunction ); ok {
4449 hasSRF = hasSRF || compiledFunction .IsSRF ()
4550 if quickFunction := compiledFunction .GetQuickFunction (); quickFunction != nil {
4651 return quickFunction , transform .NewTree , nil
4752 }
4853 }
54+ if v , ok := in .(* plan.Values ); ok {
55+ hasMultipleExpressionTuples = len (v .ExpressionTuples ) > 1
56+ }
4957 return expr , transform .SameTree , nil
5058 })
59+ if err != nil {
60+ return nil , transform .SameTree , err
61+ }
62+ if ! sameNode {
63+ projectNode .Child = n
64+ }
65+
66+ // insert node cannot have more than 1 row value if it has set returning function
67+ if isInsertNode && hasMultipleExpressionTuples && hasSRF {
68+ return nil , false , errors .Errorf ("set-returning functions are not allowed in VALUES" )
69+ }
70+
71+ // Check if there is set returning function in the projection expressions (e.g. SELECT unnest() [FROM table/srf])
72+ hasSRFInProjection := false
73+ exprs , sameExprs , err := transform .Exprs (projectNode .Projections , func (expr sql.Expression ) (sql.Expression , transform.TreeIdentity , error ) {
74+ if compiledFunction , ok := expr .(* framework.CompiledFunction ); ok {
75+ hasSRFInProjection = hasSRFInProjection || compiledFunction .IsSRF ()
76+ if quickFunction := compiledFunction .GetQuickFunction (); quickFunction != nil {
77+ return quickFunction , transform .NewTree , nil
78+ }
79+ }
80+ return expr , transform .SameTree , nil
81+ })
82+ if err != nil {
83+ return nil , transform .SameTree , err
84+ }
85+ if ! sameExprs {
86+ projectNode .Projections = exprs
87+ }
5188
52- if hasSRF {
89+ // nested iter is used for set returning functions in the projections only
90+ if hasSRFInProjection {
5391 // Under some conditions, there will be no quick-function replacement, but changing the Project node to include
5492 // nested iterators is still a change we need to tell the transform functions about.
55- same = transform .NewTree
56- n = n .( * plan. Project ) .WithIncludesNestedIters (true )
93+ sameExprs = transform .NewTree
94+ projectNode = projectNode .WithIncludesNestedIters (true )
5795 }
5896
59- return n , same , err
97+ return projectNode , sameNode && sameExprs , err
6098 })
6199}
0 commit comments