Merge cockroachdb#35711

35711: sql: permit subqueries within RHS of applyjoin r=jordanlewis a=jordanlewis

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.

Closes cockroachdb#35594.

Release note: None

Co-authored-by: Jordan Lewis <[email protected]>

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,14 +322,15 @@ 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(
 		params.ctx, &evalCtx, planCtx, params.p.Txn(), plan.plan, recv)
 	if recv.commErr != nil {
 		return recv.commErr
 	}
-	return nil
+	return rowResultWriter.err
 
 }
 

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

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')

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)