vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/subquery_op.go (about)

     1  /*
     2  Copyright 2022 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package planbuilder
    18  
    19  import (
    20  	"vitess.io/vitess/go/vt/sqlparser"
    21  	"vitess.io/vitess/go/vt/vtgate/engine"
    22  	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators"
    23  	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
    24  )
    25  
    26  func transformSubQueryPlan(ctx *plancontext.PlanningContext, op *operators.SubQueryOp) (logicalPlan, error) {
    27  	innerPlan, err := transformToLogicalPlan(ctx, op.Inner, false)
    28  	if err != nil {
    29  		return nil, err
    30  	}
    31  	innerPlan, err = planHorizon(ctx, innerPlan, op.Extracted.Subquery.Select, true)
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	argName := op.Extracted.GetArgName()
    37  	hasValuesArg := op.Extracted.GetHasValuesArg()
    38  	outerPlan, err := transformToLogicalPlan(ctx, op.Outer, false)
    39  
    40  	merged := mergeSubQueryOpPlan(ctx, innerPlan, outerPlan, op)
    41  	if merged != nil {
    42  		return merged, nil
    43  	}
    44  	plan := newPulloutSubquery(engine.PulloutOpcode(op.Extracted.OpCode), argName, hasValuesArg, innerPlan)
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  	plan.underlying = outerPlan
    49  	return plan, err
    50  }
    51  
    52  func transformCorrelatedSubQueryPlan(ctx *plancontext.PlanningContext, op *operators.CorrelatedSubQueryOp) (logicalPlan, error) {
    53  	outer, err := transformToLogicalPlan(ctx, op.Outer, false)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	inner, err := transformToLogicalPlan(ctx, op.Inner, false)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	return newSemiJoin(outer, inner, op.Vars, op.LHSColumns), nil
    62  }
    63  
    64  func mergeSubQueryOpPlan(ctx *plancontext.PlanningContext, inner, outer logicalPlan, n *operators.SubQueryOp) logicalPlan {
    65  	iroute, ok := inner.(*routeGen4)
    66  	if !ok {
    67  		return nil
    68  	}
    69  	oroute, ok := outer.(*routeGen4)
    70  	if !ok {
    71  		return nil
    72  	}
    73  
    74  	if canMergeSubqueryPlans(ctx, iroute, oroute) {
    75  		// n.extracted is an expression that lives in oroute.Select.
    76  		// Instead of looking for it in the AST, we have a copy in the subquery tree that we can update
    77  		n.Extracted.NeedsRewrite = true
    78  		replaceSubQuery(ctx, oroute.Select)
    79  		return mergeSystemTableInformation(oroute, iroute)
    80  	}
    81  	return nil
    82  }
    83  
    84  // mergeSystemTableInformation copies over information from the second route to the first and appends to it
    85  func mergeSystemTableInformation(a *routeGen4, b *routeGen4) logicalPlan {
    86  	// safe to append system table schema and system table names, since either the routing will match or either side would be throwing an error
    87  	// during run-time which we want to preserve. For example outer side has User in sys table schema and inner side has User and Main in sys table schema
    88  	// Inner might end up throwing an error at runtime, but if it doesn't then it is safe to merge.
    89  	a.eroute.SysTableTableSchema = append(a.eroute.SysTableTableSchema, b.eroute.SysTableTableSchema...)
    90  	for k, v := range b.eroute.SysTableTableName {
    91  		a.eroute.SysTableTableName[k] = v
    92  	}
    93  	return a
    94  }
    95  
    96  func canMergeSubqueryPlans(ctx *plancontext.PlanningContext, a, b *routeGen4) bool {
    97  	// this method should be close to tryMerge below. it does the same thing, but on logicalPlans instead of queryTrees
    98  	if a.eroute.Keyspace.Name != b.eroute.Keyspace.Name {
    99  		return false
   100  	}
   101  	switch a.eroute.Opcode {
   102  	case engine.Unsharded, engine.Reference:
   103  		return a.eroute.Opcode == b.eroute.Opcode
   104  	case engine.DBA:
   105  		return canSelectDBAMerge(a, b)
   106  	case engine.EqualUnique:
   107  		// Check if they target the same shard.
   108  		if b.eroute.Opcode == engine.EqualUnique &&
   109  			a.eroute.Vindex == b.eroute.Vindex &&
   110  			a.condition != nil &&
   111  			b.condition != nil &&
   112  			gen4ValuesEqual(ctx, []sqlparser.Expr{a.condition}, []sqlparser.Expr{b.condition}) {
   113  			return true
   114  		}
   115  	}
   116  	return false
   117  }