From 709b1713d912bb6777855c9f2b04414b51521adc Mon Sep 17 00:00:00 2001 From: Jordan Lewis Date: Wed, 13 Mar 2019 19:28:02 -0400 Subject: [PATCH] sql: permit subqueries within RHS of applyjoin Previously, the code assumed that subqueries would have already been promoted to top-level subqueries and run exactly once by the re-optimization phase of apply join. However, that promotion happens in execbuild, not in the optimizer, and therefore won't have already happened by the time that we go to execbuild the re-optimized RHS in apply join. For now, the solution is to re-run the subqueries in the RHS every time the RHS is run. This is suboptimal because subqueries need only be run once per query regardless of their position within the apply join tree. However, setting this up currently is difficult at the moment and would have required a more invasive change. When WITH support is available, the optimizer will promote subqueries into WITH clauses up front, at which point we will be able to safely remove the requirement that apply join rerun subqueries on the RHS. Release note: None --- pkg/sql/apply_join.go | 29 ++++++++----- .../logictest/testdata/logic_test/apply_join | 42 +++++++++++++++++++ .../testdata/logic_test/subquery_correlated | 4 +- .../opt/exec/execbuilder/scalar_builder.go | 6 ++- 4 files changed, 67 insertions(+), 14 deletions(-) diff --git a/pkg/sql/apply_join.go b/pkg/sql/apply_join.go index 4a16b8a7fda0..9ea780499f1b 100644 --- a/pkg/sql/apply_join.go +++ b/pkg/sql/apply_join.go @@ -23,7 +23,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/opt/memo" "github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical" "github.com/cockroachdb/cockroach/pkg/sql/opt/xform" - "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/rowcontainer" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/sql/sqlbase" @@ -257,8 +256,6 @@ func (a *applyJoinNode) Next(params runParams) (bool, error) { factory := a.optimizer.Factory() execbuilder.ReplaceVars(factory, a.right, a.rightProps, bindings) - a.optimizer.Memo().SetRoot(a.optimizer.Memo().RootExpr().(memo.RelExpr), a.rightProps) - a.optimizer.Optimize() newRightSide := a.optimizer.Memo().RootExpr() @@ -270,12 +267,6 @@ func (a *applyJoinNode) Next(params runParams) (bool, error) { } plan := p.(*planTop) - if len(plan.subqueryPlans) > 0 { - return false, pgerror.NewAssertionErrorf("apply join found an "+ - "unexpected re-optimized right-hand-side with non-zero subqueries:\n%s", - newRightSide) - } - if err := a.runRightSidePlan(params, plan); err != nil { return false, err } @@ -306,6 +297,23 @@ func (a *applyJoinNode) runRightSidePlan(params runParams, plan *planTop) error ) defer recv.Release() + if !params.p.extendedEvalCtx.ExecCfg.DistSQLPlanner.PlanAndRunSubqueries( + params.ctx, + params.p, + func() *extendedEvalContext { + ret := *params.extendedEvalCtx + return &ret + }, + plan.subqueryPlans, + recv, + true, + ) { + if err := rowResultWriter.Err(); err != nil { + return err + } + return recv.commErr + } + // Make a copy of the EvalContext so it can be safely modified. evalCtx := *params.p.ExtendedEvalContext() planCtx := params.p.extendedEvalCtx.ExecCfg.DistSQLPlanner.newLocalPlanningCtx(params.ctx, &evalCtx) @@ -314,6 +322,7 @@ func (a *applyJoinNode) runRightSidePlan(params runParams, plan *planTop) error plannerCopy := *params.p planCtx.planner = &plannerCopy planCtx.planner.curPlan = *plan + planCtx.ExtendedEvalCtx.Planner = &plannerCopy planCtx.stmtType = recv.stmtType params.p.extendedEvalCtx.ExecCfg.DistSQLPlanner.PlanAndRun( @@ -321,7 +330,7 @@ func (a *applyJoinNode) runRightSidePlan(params runParams, plan *planTop) error if recv.commErr != nil { return recv.commErr } - return nil + return rowResultWriter.err } diff --git a/pkg/sql/logictest/testdata/logic_test/apply_join b/pkg/sql/logictest/testdata/logic_test/apply_join index 4b9b674b5cfb..a1e747c354ba 100644 --- a/pkg/sql/logictest/testdata/logic_test/apply_join +++ b/pkg/sql/logictest/testdata/logic_test/apply_join @@ -149,3 +149,45 @@ EXECUTE e 3 three 3 three 3 three 4 four 4 four 4 four 5 five 5 five 5 five + +# Test subqueries within an apply join. + +statement ok +PREPARE f AS OPT PLAN ' +(Root + (InnerJoinApply + (Scan [(Table "t") (Cols "k,str") ]) + (Select + (Scan [(Table "u") (Cols "l,str2") ]) + [ (Eq (Plus (Var "k") + (Subquery (Values [(Tuple [(Const 1)] "tuple{int}") ] + [(Cols [(NewColumn "z" "int")] )]) + [])) + (Var "l") )] + ) + [] + [] + ) + (Presentation "k,str,l,str2") + (NoOrdering) +)' + +query ITIT rowsort +EXECUTE f +---- +1 one 2 two +2 two 3 three +3 three 4 four +4 four 5 five + +# Another test of subqueries within an apply join. + +query I +SELECT + (SELECT * FROM (VALUES ((SELECT x FROM (VALUES (1)) AS s (x)) + y))) +FROM + (VALUES (1), (2), (3)) AS t (y) +---- +2 +3 +4 diff --git a/pkg/sql/logictest/testdata/logic_test/subquery_correlated b/pkg/sql/logictest/testdata/logic_test/subquery_correlated index 93c77b763a76..007f1764227e 100644 --- a/pkg/sql/logictest/testdata/logic_test/subquery_correlated +++ b/pkg/sql/logictest/testdata/logic_test/subquery_correlated @@ -928,13 +928,11 @@ ON c.c_id=o.c_id AND EXISTS(SELECT * FROM o WHERE o.c_id=c.c_id AND ship IS NULL 4 70 4 80 -query II rowsort +query error more than one row returned by a subquery SELECT c.c_id, o.o_id FROM c INNER JOIN o ON c.c_id=o.c_id AND o.ship = (SELECT o.ship FROM o WHERE o.c_id=c.c_id); ----- -6 90 statement error AS OF SYSTEM TIME must be provided on a top-level statement SELECT (SELECT c_id FROM o AS OF SYSTEM TIME '-1us') diff --git a/pkg/sql/opt/exec/execbuilder/scalar_builder.go b/pkg/sql/opt/exec/execbuilder/scalar_builder.go index c318d194a104..d8fbcbe6496d 100644 --- a/pkg/sql/opt/exec/execbuilder/scalar_builder.go +++ b/pkg/sql/opt/exec/execbuilder/scalar_builder.go @@ -557,8 +557,12 @@ func (b *Builder) buildSubquery( func (b *Builder) addSubquery( mode exec.SubqueryMode, typ types.T, root exec.Node, originalExpr *tree.Subquery, ) *tree.Subquery { + var originalSelect tree.SelectStatement + if originalExpr != nil { + originalSelect = originalExpr.Select + } exprNode := &tree.Subquery{ - Select: originalExpr.Select, + Select: originalSelect, Exists: mode == exec.SubqueryExists, } exprNode.SetType(typ)