github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/colexec/preinsert/preinsert.go (about)

     1  // Copyright 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 preinsert
    16  
    17  import (
    18  	"bytes"
    19  
    20  	"github.com/matrixorigin/matrixone/pkg/catalog"
    21  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    22  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    23  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    24  	"github.com/matrixorigin/matrixone/pkg/logutil"
    25  	pb "github.com/matrixorigin/matrixone/pkg/pb/plan"
    26  	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
    27  	"github.com/matrixorigin/matrixone/pkg/sql/util"
    28  	"github.com/matrixorigin/matrixone/pkg/vm"
    29  	"go.uber.org/zap"
    30  )
    31  
    32  const argName = "preinsert"
    33  
    34  func (arg *Argument) String(buf *bytes.Buffer) {
    35  	buf.WriteString(argName)
    36  	buf.WriteString(": pre processing insert")
    37  }
    38  
    39  func (arg *Argument) Prepare(_ *proc) error {
    40  	return nil
    41  }
    42  
    43  func (arg *Argument) Call(proc *proc) (vm.CallResult, error) {
    44  	if err, isCancel := vm.CancelCheck(proc); isCancel {
    45  		return vm.CancelResult, err
    46  	}
    47  
    48  	result, err := arg.GetChildren(0).Call(proc)
    49  	if err != nil {
    50  		return result, err
    51  	}
    52  	analy := proc.GetAnalyze(arg.GetIdx(), arg.GetParallelIdx(), arg.GetParallelMajor())
    53  	analy.Start()
    54  	defer analy.Stop()
    55  
    56  	if result.Batch == nil || result.Batch.IsEmpty() {
    57  		return result, nil
    58  	}
    59  	bat := result.Batch
    60  
    61  	if arg.buf != nil {
    62  		proc.PutBatch(arg.buf)
    63  		arg.buf = nil
    64  	}
    65  
    66  	arg.buf = batch.NewWithSize(len(arg.Attrs))
    67  	// keep shuffleIDX unchanged
    68  	arg.buf.ShuffleIDX = bat.ShuffleIDX
    69  	arg.buf.Attrs = make([]string, 0, len(arg.Attrs))
    70  	for idx := range arg.Attrs {
    71  		arg.buf.Attrs = append(arg.buf.Attrs, arg.Attrs[idx])
    72  		srcVec := bat.Vecs[idx]
    73  		vec := proc.GetVector(*srcVec.GetType())
    74  		if err := vector.GetUnionAllFunction(*srcVec.GetType(), proc.Mp())(vec, srcVec); err != nil {
    75  			vec.Free(proc.Mp())
    76  			return result, err
    77  		}
    78  		arg.buf.SetVector(int32(idx), vec)
    79  	}
    80  	arg.buf.AddRowCount(bat.RowCount())
    81  
    82  	if arg.HasAutoCol {
    83  		err := genAutoIncrCol(arg.buf, proc, arg)
    84  		if err != nil {
    85  			return result, err
    86  		}
    87  	}
    88  	// check new rows not null
    89  	err = colexec.BatchDataNotNullCheck(arg.buf, arg.TableDef, proc.Ctx)
    90  	if err != nil {
    91  		return result, err
    92  	}
    93  
    94  	// calculate the composite primary key column and append the result vector to batch
    95  	err = genCompositePrimaryKey(arg.buf, proc, arg.TableDef)
    96  	if err != nil {
    97  		return result, err
    98  	}
    99  	err = genClusterBy(arg.buf, proc, arg.TableDef)
   100  	if err != nil {
   101  		return result, err
   102  	}
   103  	if arg.IsUpdate {
   104  		idx := len(bat.Vecs) - 1
   105  		arg.buf.Attrs = append(arg.buf.Attrs, catalog.Row_ID)
   106  		rowIdVec := proc.GetVector(*bat.GetVector(int32(idx)).GetType())
   107  		err = rowIdVec.UnionBatch(bat.Vecs[idx], 0, bat.Vecs[idx].Length(), nil, proc.Mp())
   108  		if err != nil {
   109  			rowIdVec.Free(proc.Mp())
   110  			return result, err
   111  		}
   112  		arg.buf.Vecs = append(arg.buf.Vecs, rowIdVec)
   113  	}
   114  
   115  	result.Batch = arg.buf
   116  	return result, nil
   117  }
   118  
   119  func genAutoIncrCol(bat *batch.Batch, proc *proc, arg *Argument) error {
   120  	lastInsertValue, err := proc.IncrService.InsertValues(
   121  		proc.Ctx,
   122  		arg.TableDef.TblId,
   123  		bat,
   124  		arg.EstimatedRowCount,
   125  	)
   126  	if err != nil {
   127  		if moerr.IsMoErrCode(err, moerr.ErrNoSuchTable) {
   128  			logutil.Error("insert auto increment column failed", zap.Error(err))
   129  			return moerr.NewNoSuchTableNoCtx(arg.SchemaName, arg.TableDef.Name)
   130  		}
   131  		return err
   132  	}
   133  	proc.SetLastInsertID(lastInsertValue)
   134  	return nil
   135  }
   136  
   137  func genCompositePrimaryKey(bat *batch.Batch, proc *proc, tableDef *pb.TableDef) error {
   138  	// Check whether the composite primary key column is included
   139  	if tableDef.Pkey.CompPkeyCol == nil {
   140  		return nil
   141  	}
   142  
   143  	return util.FillCompositeKeyBatch(bat, catalog.CPrimaryKeyColName, tableDef.Pkey.Names, proc)
   144  }
   145  
   146  func genClusterBy(bat *batch.Batch, proc *proc, tableDef *pb.TableDef) error {
   147  	if tableDef.ClusterBy == nil {
   148  		return nil
   149  	}
   150  	clusterBy := tableDef.ClusterBy.Name
   151  	if clusterBy == "" || !util.JudgeIsCompositeClusterByColumn(clusterBy) {
   152  		return nil
   153  	}
   154  	return util.FillCompositeClusterByBatch(bat, clusterBy, proc)
   155  }