github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/plan/optimize.go (about)

     1  // Copyright 2021 - 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package plan
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    19  	"github.com/matrixorigin/matrixone/pkg/pb/plan"
    20  	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
    21  	"github.com/matrixorigin/matrixone/pkg/sql/plan/rule"
    22  )
    23  
    24  var defaultRules = []Rule{}
    25  var prepareRules = []Rule{}
    26  
    27  func init() {
    28  	defaultRules = []Rule{
    29  		rule.NewConstantFold(false),
    30  		// rule.NewPredicatePushdown(),
    31  	}
    32  	prepareRules = []Rule{
    33  		rule.NewConstantFold(true),
    34  	}
    35  }
    36  
    37  func NewBaseOptimizer(ctx CompilerContext) *BaseOptimizer {
    38  	return &BaseOptimizer{
    39  		ctx:   ctx,
    40  		rules: defaultRules,
    41  	}
    42  }
    43  
    44  func NewPrepareOptimizer(ctx CompilerContext) *BaseOptimizer {
    45  	return &BaseOptimizer{
    46  		ctx:   ctx,
    47  		rules: prepareRules,
    48  	}
    49  }
    50  
    51  func (opt *BaseOptimizer) CurrentContext() CompilerContext {
    52  	return opt.ctx
    53  }
    54  
    55  func (opt *BaseOptimizer) Optimize(stmt tree.Statement, isPrepareStmt bool) (*Query, error) {
    56  	pn, err := BuildPlan(opt.ctx, stmt, isPrepareStmt)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	qry, ok := pn.Plan.(*plan.Plan_Query)
    61  	if !ok {
    62  		panic(moerr.NewInternalError(opt.ctx.GetContext(), pn.String()))
    63  	}
    64  	opt.qry = qry.Query
    65  	return opt.optimize()
    66  }
    67  
    68  func (opt *BaseOptimizer) optimize() (*Query, error) {
    69  	if len(opt.qry.Steps) == 0 {
    70  		return opt.qry, nil
    71  	}
    72  	for _, step := range opt.qry.Steps {
    73  		opt.exploreNode(opt.qry.Nodes[step])
    74  	}
    75  
    76  	opt.pruneUsedNodes(opt.qry)
    77  
    78  	return opt.qry, nil
    79  }
    80  
    81  func (opt *BaseOptimizer) exploreNode(n *Node) {
    82  	for i := range n.Children {
    83  		opt.exploreNode(opt.qry.Nodes[n.Children[i]])
    84  	}
    85  	for _, rule := range opt.rules {
    86  		if rule.Match(n) {
    87  			rule.Apply(n, opt.qry, opt.ctx.GetProcess())
    88  		}
    89  	}
    90  }
    91  
    92  func (opt *BaseOptimizer) pruneUsedNodes(qry *plan.Query) {
    93  	newSteps := make([]int32, 0, len(qry.Steps))
    94  	newNodes := make([]*plan.Node, 0, len(qry.Nodes))
    95  
    96  	for _, step := range qry.Steps {
    97  		newSteps = append(newSteps, opt.compactPlanTree(qry, step, &newNodes))
    98  	}
    99  
   100  	qry.Steps = newSteps
   101  	qry.Nodes = newNodes
   102  }
   103  
   104  func (opt *BaseOptimizer) compactPlanTree(qry *plan.Query, nodeID int32, nodes *[]*plan.Node) int32 {
   105  	node := qry.Nodes[nodeID]
   106  
   107  	for i, childID := range node.Children {
   108  		node.Children[i] = opt.compactPlanTree(qry, childID, nodes)
   109  	}
   110  
   111  	node.NodeId = int32(len(*nodes))
   112  	*nodes = append(*nodes, node)
   113  
   114  	return node.NodeId
   115  }