github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/plan/build_update.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/pb/plan"
    19  	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
    20  )
    21  
    22  func buildTableUpdate(stmt *tree.Update, ctx CompilerContext) (p *Plan, err error) {
    23  	tblInfo, err := getUpdateTableInfo(ctx, stmt)
    24  	if err != nil {
    25  		return nil, err
    26  	}
    27  	rewriteInfo := &dmlSelectInfo{
    28  		typ:     "update",
    29  		rootId:  -1,
    30  		tblInfo: tblInfo,
    31  	}
    32  	builder := NewQueryBuilder(plan.Query_SELECT, ctx)
    33  	bindCtx := NewBindContext(builder, nil)
    34  	bindCtx.groupTag = builder.genNewTag()
    35  	bindCtx.aggregateTag = builder.genNewTag()
    36  	bindCtx.projectTag = builder.genNewTag()
    37  
    38  	err = initUpdateStmt(builder, bindCtx, rewriteInfo, stmt)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  
    43  	if tblInfo.haveConstraint {
    44  		for i, tableDef := range tblInfo.tableDefs {
    45  			err = rewriteDmlSelectInfo(builder, bindCtx, rewriteInfo, tableDef, rewriteInfo.derivedTableId, i)
    46  			if err != nil {
    47  				return nil, err
    48  			}
    49  		}
    50  	}
    51  
    52  	// append ProjectNode
    53  	rewriteInfo.rootId = builder.appendNode(&plan.Node{
    54  		NodeType:    plan.Node_PROJECT,
    55  		ProjectList: rewriteInfo.projectList,
    56  		Children:    []int32{rewriteInfo.rootId},
    57  		BindingTags: []int32{bindCtx.projectTag},
    58  	}, bindCtx)
    59  	bindCtx.results = rewriteInfo.projectList
    60  
    61  	builder.qry.Steps = append(builder.qry.Steps, rewriteInfo.rootId)
    62  	query, err := builder.createQuery()
    63  	if err != nil {
    64  		return nil, err
    65  	}
    66  
    67  	// append delete node
    68  	updateCtx := &plan.UpdateCtx{
    69  		Ref:       rewriteInfo.tblInfo.objRef,
    70  		TableDefs: rewriteInfo.tblInfo.tableDefs,
    71  		Idx:       make([]*plan.IdList, len(rewriteInfo.tblInfo.objRef)),
    72  		UpdateCol: make([]*plan.ColPosMap, len(rewriteInfo.tblInfo.updateCol)),
    73  
    74  		IdxRef: rewriteInfo.onIdxTbl,
    75  		IdxIdx: rewriteInfo.onIdx,
    76  
    77  		OnRestrictRef: rewriteInfo.onRestrictTbl,
    78  		OnRestrictIdx: rewriteInfo.onRestrict,
    79  
    80  		OnCascadeRef:       rewriteInfo.onCascadeRef,
    81  		OnCascadeDef:       rewriteInfo.onCascadeTableDef,
    82  		OnCascadeIdx:       make([]*plan.IdList, len(rewriteInfo.onCascade)),
    83  		OnCascadeUpdateCol: make([]*plan.ColPosMap, len(rewriteInfo.onCascadeUpdateCol)),
    84  
    85  		OnSetRef:       rewriteInfo.onSetRef,
    86  		OnSetDef:       rewriteInfo.onSetTableDef,
    87  		OnSetIdx:       make([]*plan.IdList, len(rewriteInfo.onSet)),
    88  		OnSetUpdateCol: make([]*plan.ColPosMap, len(rewriteInfo.onSetUpdateCol)),
    89  
    90  		ParentIdx: make([]*plan.ColPosMap, len(rewriteInfo.parentIdx)),
    91  	}
    92  	idx := int64(0)
    93  	for i, tableDef := range rewriteInfo.tblInfo.tableDefs {
    94  		updateCtx.TableDefs[i] = tableDef
    95  		idxList := make([]int64, len(tableDef.Cols))
    96  		for j := range tableDef.Cols {
    97  			idxList[j] = idx
    98  			idx++
    99  		}
   100  		updateCtx.Idx[i] = &plan.IdList{
   101  			List: idxList,
   102  		}
   103  	}
   104  	for i, idxMap := range rewriteInfo.tblInfo.updateCol {
   105  		updateCtx.UpdateCol[i] = &plan.ColPosMap{
   106  			Map: idxMap,
   107  		}
   108  	}
   109  	for i, idxList := range rewriteInfo.onCascade {
   110  		updateCtx.OnCascadeIdx[i] = &plan.IdList{
   111  			List: idxList,
   112  		}
   113  	}
   114  	for i, idxMap := range rewriteInfo.onCascadeUpdateCol {
   115  		updateCtx.OnCascadeUpdateCol[i] = &plan.ColPosMap{
   116  			Map: idxMap,
   117  		}
   118  	}
   119  	for i, idxList := range rewriteInfo.onSet {
   120  		updateCtx.OnSetIdx[i] = &plan.IdList{
   121  			List: idxList,
   122  		}
   123  	}
   124  	for i, idxMap := range rewriteInfo.onSetUpdateCol {
   125  		updateCtx.OnSetUpdateCol[i] = &plan.ColPosMap{
   126  			Map: idxMap,
   127  		}
   128  	}
   129  	for i, idxMap := range rewriteInfo.parentIdx {
   130  		updateCtx.ParentIdx[i] = &plan.ColPosMap{
   131  			Map: idxMap,
   132  		}
   133  	}
   134  
   135  	node := &Node{
   136  		NodeType:  plan.Node_UPDATE,
   137  		ObjRef:    nil,
   138  		TableDef:  nil,
   139  		Children:  []int32{query.Steps[len(query.Steps)-1]},
   140  		NodeId:    int32(len(query.Nodes)),
   141  		UpdateCtx: updateCtx,
   142  	}
   143  	query.Nodes = append(query.Nodes, node)
   144  	query.Steps[len(query.Steps)-1] = node.NodeId
   145  	query.StmtType = plan.Query_UPDATE
   146  
   147  	return &Plan{
   148  		Plan: &plan.Plan_Query{
   149  			Query: query,
   150  		},
   151  	}, err
   152  }
   153  
   154  func isSameColumnType(t1 *Type, t2 *Type) bool {
   155  	if t1.Id != t2.Id {
   156  		return false
   157  	}
   158  	if t1.Width == t2.Width && t1.Precision == t2.Precision && t1.Size == t2.Size && t1.Scale == t2.Scale {
   159  		return true
   160  	}
   161  	return true
   162  }