MB-65809 Handle subquery plans in prepared statements

For a prepared statement with SAVE option, or if plan stability is
turned on (PREPARED_ONLY, AD_HOC, AD_HOC_READ_ONLY), subquery plans
is part of the prepared statement plan. To properly unmarshal such
plans, we now put a couple of extra fields in planContext, which is
part of each operator. The extra fields keep track of the saved
subquery plans from the prepared statement (with the subquery text
as key), as well as unmarshalled subquery plan (with the
*algebra.Select as key).

When we unmarshal plans for a prepared statement, we first look for
the subquery plans, and if found, we generate a map of the saved
subquery plan (subqMap), and also create a new subquery plans
place holder (subqPlans). These two fields are put as part of the
plan context. We then proceed to unmarshal the main query plan. As
the main query plan is unmarshalled, when we need to parse an
expression, and there is a possiblity of encountering a subquery
in the expression, we use the new parsing function (introduced in
the previous checkin) for parsing, which will parse and generate
an expression, but then checks whether the subqMap is available,
if so, gather subqueries from the expression just parsed, and
look for saved subquery plan for each subquery. When found, it'll
then unmarshal the subquery plan, and save the subquery plan with
the new subquery pointer (*algebra.Select) into subqPlans. At the
end of this process when the main query plan is unmarshalled, all
the subquery expressions in the main plan should hopefully have
a plan that has been unmarshalled and saved in subqPlans. Then
subqPlans is put in the prepared plan, which then allows subquery
execution to find the appropriate plans (logic for looking for
subquery plans in a prepared statement already exists).

Change-Id: I6ebd5b8dc6288d3d0607596123962fa5db70c1a7
Reviewed-on: https://review.couchbase.org/c/query/+/242212
Reviewed-by: Sitaram Vemulapalli <sitaram.vemulapalli@couchbase.com>
Tested-by: Bingjie Miao <bingjie.miao@couchbase.com>

Author: miaobingjie

algebra/subquery.go

@@ -369,6 +369,8 @@ func (this *SubqueryPlans) MarshalPlans(lock bool, prepName string,
 	for key := range this.plans {
 		r := make(map[string]interface{}, 4)
 		str := marshal(r, key, this.plans[key], this.isks[key], prepName)
+		// duplicates are removed, i.e. if the exact same subquery appears multiple times in
+		// the query, the subquery plan is only saved once
 		subqueries[str] = r
 	}
 	if lock {

execution/context.go

@@ -34,6 +34,7 @@ import (
 	"github.com/couchbase/query/plan"
 	"github.com/couchbase/query/planner"
 	"github.com/couchbase/query/sequences"
+	"github.com/couchbase/query/settings"
 	"github.com/couchbase/query/system"
 	"github.com/couchbase/query/tenant"
 	"github.com/couchbase/query/timestamp"
@@ -1466,6 +1467,10 @@ func (this *opContext) EvaluateSubquery(query *algebra.Select, parent value.Valu
 				}
 			}
 		}
+		if !planFound && (this.prepared.Persist() || settings.IsPlanStabilityEnabled()) {
+			logging.Debugf("Plan Stability (mode %d): subquery plan not found. Prepared name(%s) persist(%t) subquery (%s)",
+				settings.GetPlanStabilityMode(), this.prepared.Name(), this.prepared.Persist(), query.String())
+		}
 	}
 
 	if !planFound {

plan/context.go

@@ -9,6 +9,9 @@
 package plan
 
 import (
+	"encoding/json"
+
+	"github.com/couchbase/query/algebra"
 	"github.com/couchbase/query/expression"
 	"github.com/couchbase/query/value"
 )
@@ -19,20 +22,40 @@ type planContext struct {
 	withs     *value.ScopeValue
 	vars      *value.ScopeValue
 	keyspaces *value.ScopeValue
+
+	// for subquery plans in prepared
+	subqMap   map[string]*subqPlanInfo
+	subqPlans *algebra.SubqueryPlans
+}
+
+func newPlanContext(subqMap map[string]*subqPlanInfo, subqPlans *algebra.SubqueryPlans) *planContext {
+	rv := &planContext{
+		withs:     value.NewScopeValue(make(map[string]interface{}), nil),
+		vars:      value.NewScopeValue(make(map[string]interface{}), nil),
+		keyspaces: value.NewScopeValue(make(map[string]interface{}), nil),
+		subqMap:   subqMap,
+		subqPlans: subqPlans,
+	}
+
+	rv.SetMapper(rv)
+	return rv
 }
 
-func newPlanContext(parent *planContext) *planContext {
-	var wv, vv, kv value.Value
-	if parent != nil {
-		wv = parent.withs
-		vv = parent.vars
-		kv = parent.keyspaces
+func subqPlanContext(parent *planContext) *planContext {
+	// parent should not be nil
+	if parent == nil {
+		panic("subqPlanContext: parent pointer is nil")
 	}
 
+	// the subqMap and subqPlans should contain subqueries at all levels, and it's ok to share
+	// the pointers since this is currently only used for unmarshal of subquery plans from a
+	// prepared statement (no parallelism)
 	rv := &planContext{
-		withs:     value.NewScopeValue(make(map[string]interface{}), wv),
-		vars:      value.NewScopeValue(make(map[string]interface{}), vv),
-		keyspaces: value.NewScopeValue(make(map[string]interface{}), kv),
+		withs:     value.NewScopeValue(make(map[string]interface{}), parent.withs),
+		vars:      value.NewScopeValue(make(map[string]interface{}), parent.vars),
+		keyspaces: value.NewScopeValue(make(map[string]interface{}), parent.keyspaces),
+		subqMap:   parent.subqMap,
+		subqPlans: parent.subqPlans,
 	}
 
 	rv.SetMapper(rv)
@@ -222,7 +245,7 @@ func (this *planContext) VisitIdentifier(expr *expression.Identifier) (interface
 
 func (this *planContext) VisitSubquery(expr expression.Subquery) (interface{}, error) {
 	// use a new planContext for an extra scope
-	planContext := newPlanContext(this)
+	planContext := subqPlanContext(this)
 
 	err := expr.MapChildren(planContext)
 	if err != nil {
@@ -259,3 +282,54 @@ func (this *planContext) addSubqTermAlias(alias string) {
 		this.keyspaces.SetField(alias, value.NewValue(ident_flags))
 	}
 }
+
+// functionality for handling subquery map
+
+type subqPlanInfo struct {
+	op   []byte
+	isks map[string]bool
+}
+
+func (this *planContext) hasSubqueryMap() bool {
+	return this.subqMap != nil
+}
+
+// Look for subquery expressions in the expr just parsed, and for each subquery, match with the
+// subquery plans from prepared statement; when matched, unmarshal the subquery plan and put the
+// subquery plan into algebra.SubqueryPlans (for the prepared statement)
+func (this *planContext) checkSubqueryMap(expr expression.Expression) (expression.Expression, error) {
+	if len(this.subqMap) == 0 {
+		return expr, nil
+	}
+
+	subqs, err := expression.ListSubqueries(expression.Expressions{expr}, false)
+	if err != nil {
+		return expr, err
+	}
+
+	for _, subqExpr := range subqs {
+		subq, _ := subqExpr.(*algebra.Subquery)
+		str := subq.String()
+		if planInfo, ok := this.subqMap[str]; ok {
+			// unmarshal the subquery plan here (in case there are identical subqueries
+			// present we need separate plan for each subquery)
+			var op_type struct {
+				Operator string `json:"#operator"`
+			}
+
+			err = json.Unmarshal(planInfo.op, &op_type)
+			if err != nil {
+				return nil, err
+			}
+
+			subPlanContext := subqPlanContext(this)
+			subqPlan, err := MakeOperator(op_type.Operator, planInfo.op, subPlanContext)
+			if err != nil {
+				return nil, err
+			}
+			// no need for locking since this is only called under unmarshal (no parallelism)
+			this.subqPlans.Set(subq.Select(), nil, NewQueryPlan(subqPlan), planInfo.isks, false)
+		}
+	}
+	return expr, nil
+}

plan/parser.go

@@ -13,6 +13,16 @@ import (
 	"github.com/couchbase/query/expression/parser"
 )
 
+// This function is used to parse expressions, and if there are subquery plans available when
+// unmarshalling prepared statements, it handles the matching of parsed subquery expression and
+// unmarshalled subquery plan
 func parseWithContext(s string, plContext *planContext) (expression.Expression, error) {
-	return parser.Parse(s)
+	expr, err := parser.Parse(s)
+	if expr == nil || err != nil {
+		return expr, err
+	}
+	if plContext != nil && plContext.hasSubqueryMap() {
+		expr, err = plContext.checkSubqueryMap(expr)
+	}
+	return expr, err
 }

plan/prepared.go

@@ -23,7 +23,6 @@ import (
 	"github.com/couchbase/query/algebra"
 	"github.com/couchbase/query/datastore"
 	"github.com/couchbase/query/errors"
-	"github.com/couchbase/query/expression/parser"
 	"github.com/couchbase/query/util"
 	"github.com/couchbase/query/value"
 )
@@ -307,45 +306,33 @@ func (this *Prepared) unmarshalInternal(body []byte) error {
 		}
 	}
 
-	planContext := newPlanContext(nil)
-	this.Operator, err = MakeOperator(op_type.Operator, _unmarshalled.Operator, planContext)
-	if err != nil {
-		return err
-	}
-
+	var subqMap map[string]*subqPlanInfo
+	var subqPlans *algebra.SubqueryPlans
+	hasSubq := false
 	if len(_unmarshalled.SubqPlans) > 0 {
-		this.subqueryPlans = algebra.NewSubqueryPlans()
+		hasSubq = true
+		subqPlans = algebra.NewSubqueryPlans()
+		subqMap = make(map[string]*subqPlanInfo, len(_unmarshalled.SubqPlans))
 		for _, planP := range _unmarshalled.SubqPlans {
 			subqText := planP.SubqText
 			if subqText == "" {
 				return errors.NewPlanInternalError("prepared.unmarshalInternal: missing subquery text")
 			}
-			subqExpr, err := parser.Parse(subqText)
-			if err != nil {
-				return err
-			}
-			subquery, ok := subqExpr.(*algebra.Subquery)
-			if !ok {
-				return errors.NewPlanInternalError("prepared.unmarshalInternal: subquery expression is not a Subquery.")
-			}
-
-			var op_type1 struct {
-				Operator string `json:"#operator"`
-			}
-
-			err = json.Unmarshal(planP.PlanOp, &op_type1)
-			if err != nil {
-				return err
+			subqMap[subqText] = &subqPlanInfo{
+				op:   planP.PlanOp,
+				isks: planP.IndexScanKeyspaces,
 			}
+		}
+	}
 
-			subPlanContext := newPlanContext(planContext)
-			subqPlan, err := MakeOperator(op_type1.Operator, planP.PlanOp, subPlanContext)
-			if err != nil {
-				return err
-			}
+	planContext := newPlanContext(subqMap, subqPlans)
+	this.Operator, err = MakeOperator(op_type.Operator, _unmarshalled.Operator, planContext)
+	if err != nil {
+		return err
+	}
 
-			this.subqueryPlans.Set(subquery.Select(), nil, NewQueryPlan(subqPlan), planP.IndexScanKeyspaces, false)
-		}
+	if hasSubq {
+		this.subqueryPlans = subqPlans
 	}
 
 	return nil
@@ -894,7 +881,7 @@ func marshalSubqueryPlan(r map[string]interface{}, subselect *algebra.Select, pl
 	}
 	text := subselect.String()
 	if subselect.IsCorrelated() {
-		// needed for parsing to work properly
+		// this text needs to match what gets generated from Subquery.String()
 		text = "correlated (" + text + ")"
 	}
 	r["subquery"] = text