remove getRelationById in some case

pkg/sql/colexec/lockop/lock_op.go

@@ -87,6 +87,18 @@ func (lockOp *LockOp) Prepare(proc *process.Process) error {
 				GetFetchRowsFunc(lockOp.targets[idx].primaryColumnType))
 		}
 	}
+	if len(lockOp.ctr.relations) == 0 {
+		lockOp.ctr.relations = make([]engine.Relation, len(lockOp.targets))
+		for i, target := range lockOp.targets {
+			if target.objRef != nil {
+				rel, _, err := colexec.GetRelAndPartitionRelsByObjRef(proc.Ctx, proc, lockOp.engine, target.objRef, nil)
+				if err != nil {
+					return err
+				}
+				lockOp.ctr.relations[i] = rel
+			}
+		}
+	}
 	lockOp.ctr.parker = types.NewPacker()
 	return nil
 }
@@ -176,7 +188,7 @@ func performLock(
 			proc.Ctx,
 			lockOp.engine,
 			analyzer,
-			nil,
+			lockOp.ctr.relations[idx],
 			target.tableID,
 			proc,
 			priVec,
@@ -772,6 +784,7 @@ func (lockOp *LockOp) CopyToPipelineTarget() []*pipeline.LockTarget {
 			Mode:               target.mode,
 			LockRows:           plan.DeepCopyExpr(target.lockRows),
 			LockTableAtTheEnd:  target.lockTableAtTheEnd,
+			ObjRef:             plan.DeepCopyObjectRef(target.objRef),
 		}
 	}
 	return targets
@@ -780,13 +793,15 @@ func (lockOp *LockOp) CopyToPipelineTarget() []*pipeline.LockTarget {
 // AddLockTarget add lock target, LockMode_Exclusive will used
 func (lockOp *LockOp) AddLockTarget(
 	tableID uint64,
+	objRef *plan.ObjectRef,
 	primaryColumnIndexInBatch int32,
 	primaryColumnType types.Type,
 	refreshTimestampIndexInBatch int32,
 	lockRows *plan.Expr,
 	lockTableAtTheEnd bool) *LockOp {
 	return lockOp.AddLockTargetWithMode(
 		tableID,
+		objRef,
 		lock.LockMode_Exclusive,
 		primaryColumnIndexInBatch,
 		primaryColumnType,
@@ -798,6 +813,7 @@ func (lockOp *LockOp) AddLockTarget(
 // AddLockTargetWithMode add lock target with lock mode
 func (lockOp *LockOp) AddLockTargetWithMode(
 	tableID uint64,
+	objRef *plan.ObjectRef,
 	mode lock.LockMode,
 	primaryColumnIndexInBatch int32,
 	primaryColumnType types.Type,
@@ -806,6 +822,7 @@ func (lockOp *LockOp) AddLockTargetWithMode(
 	lockTableAtTheEnd bool) *LockOp {
 	lockOp.targets = append(lockOp.targets, lockTarget{
 		tableID:                      tableID,
+		objRef:                       objRef,
 		primaryColumnIndexInBatch:    primaryColumnIndexInBatch,
 		primaryColumnType:            primaryColumnType,
 		refreshTimestampIndexInBatch: refreshTimestampIndexInBatch,
@@ -889,6 +906,7 @@ func (lockOp *LockOp) AddLockTargetWithPartitionAndMode(
 	// only one partition table, process as normal table
 	if len(tableIDs) == 1 {
 		return lockOp.AddLockTarget(tableIDs[0],
+			nil,
 			primaryColumnIndexInBatch,
 			primaryColumnType,
 			refreshTimestampIndexInBatch,
@@ -917,6 +935,7 @@ func (lockOp *LockOp) Reset(proc *process.Process, pipelineFailed bool, err erro
 	lockOp.resetParker()
 	lockOp.ctr.retryError = nil
 	lockOp.ctr.defChanged = false
+	lockOp.ctr.relations = nil
 }
 
 // Free free mem

pkg/sql/colexec/lockop/lock_op_test.go

@@ -381,7 +381,7 @@ func runLockNonBlockingOpTest(
 				IsLast:  false,
 			}
 			for idx, table := range tables {
-				arg.AddLockTarget(table, offset, pkType, offset+1, nil, true)
+				arg.AddLockTarget(table, nil, offset, pkType, offset+1, nil, true)
 
 				vec := vector.NewVec(pkType)
 				vector.AppendFixedList(vec, values[idx], nil, proc.Mp())

pkg/sql/colexec/lockop/types.go

@@ -104,6 +104,7 @@ func (lockOp *LockOp) Release() {
 
 type lockTarget struct {
 	tableID                      uint64
+	objRef                       *plan.ObjectRef
 	primaryColumnIndexInBatch    int32
 	refreshTimestampIndexInBatch int32
 	primaryColumnType            types.Type
@@ -133,6 +134,7 @@ type state struct {
 	retryError           error
 	defChanged           bool
 	fetchers             []FetchLockRowsFunc
+	relations            []engine.Relation
 	hasNewVersionInRange hasNewVersionInRangeFunc
 	lockCount            int64
 }

pkg/sql/colexec/preinsert/preinsert.go

@@ -276,26 +276,31 @@ func genAutoIncrCol(bat *batch.Batch, proc *proc, preInsert *PreInsert) error {
 			currentTxn.AddWorkspace(ws)
 			ws.BindTxnOp(currentTxn)
 		}
-		if _, _, rel, err := eng.GetRelationById(proc.Ctx, currentTxn, tableID); err == nil {
-			for col, idx := range needReCheck {
-				vec := bat.GetVector(int32(idx))
-				from, err := proc.GetIncrService().GetLastAllocateTS(proc.Ctx, tableID, col)
-				if err != nil {
-					return err
-				}
-				fromTs := types.TimestampToTS(from)
-				toTs := types.TimestampToTS(proc.Base.TxnOperator.SnapshotTS())
-				if mayChanged, err := rel.PrimaryKeysMayBeUpserted(proc.Ctx, fromTs, toTs, vec); err == nil {
-					if mayChanged {
-						logutil.Debugf("user may have manually specified the value to be inserted into the auto pk col before this transaction.")
-						return moerr.NewTxnNeedRetry(proc.Ctx)
-					}
-				} else {
-					return err
+		db, err := eng.Database(proc.Ctx, preInsert.SchemaName, currentTxn)
+		if err != nil {
+			return err
+		}
+		rel, err := db.Relation(proc.Ctx, preInsert.TableDef.Name, nil)
+		if err != nil {
+			return err
+		}
+
+		for col, idx := range needReCheck {
+			vec := bat.GetVector(int32(idx))
+			from, err := proc.GetIncrService().GetLastAllocateTS(proc.Ctx, tableID, col)
+			if err != nil {
+				return err
+			}
+			fromTs := types.TimestampToTS(from)
+			toTs := types.TimestampToTS(proc.Base.TxnOperator.SnapshotTS())
+			if mayChanged, err := rel.PrimaryKeysMayBeUpserted(proc.Ctx, fromTs, toTs, vec); err == nil {
+				if mayChanged {
+					logutil.Debugf("user may have manually specified the value to be inserted into the auto pk col before this transaction.")
+					return moerr.NewTxnNeedRetry(proc.Ctx)
 				}
+			} else {
+				return err
 			}
-		} else {
-			return err
 		}
 	}
 

pkg/sql/compile/operator.go

@@ -772,7 +772,7 @@ func constructLockOp(n *plan.Node, eng engine.Engine) (*lockop.LockOp, error) {
 		if target.IsPartitionTable {
 			arg.AddLockTargetWithPartition(target.GetPartitionTableIds(), target.GetPrimaryColIdxInBat(), typ, target.GetRefreshTsIdxInBat(), target.GetLockRows(), target.GetLockTableAtTheEnd(), target.GetFilterColIdxInBat())
 		} else {
-			arg.AddLockTarget(target.GetTableId(), target.GetPrimaryColIdxInBat(), typ, target.GetRefreshTsIdxInBat(), target.GetLockRows(), target.GetLockTableAtTheEnd())
+			arg.AddLockTarget(target.GetTableId(), target.GetObjRef(), target.GetPrimaryColIdxInBat(), typ, target.GetRefreshTsIdxInBat(), target.GetLockRows(), target.GetLockTableAtTheEnd())
 		}
 
 	}

pkg/sql/compile/remoterun.go

@@ -920,7 +920,7 @@ func convertToVmOperator(opr *pipeline.Instruction, ctx *scopeContext, eng engin
 		lockArg := lockop.NewArgumentByEngine(eng)
 		for _, target := range t.Targets {
 			typ := plan2.MakeTypeByPlan2Type(target.PrimaryColTyp)
-			lockArg.AddLockTarget(target.GetTableId(), target.GetPrimaryColIdxInBat(), typ, target.GetRefreshTsIdxInBat(), target.GetLockRows(), target.GetLockTableAtTheEnd())
+			lockArg.AddLockTarget(target.GetTableId(), target.GetObjRef(), target.GetPrimaryColIdxInBat(), typ, target.GetRefreshTsIdxInBat(), target.GetLockRows(), target.GetLockTableAtTheEnd())
 		}
 		for _, target := range t.Targets {
 			if target.LockTable {

pkg/sql/plan/bind_delete.go

@@ -269,6 +269,7 @@ func (builder *QueryBuilder) bindDelete(stmt *tree.Delete, bindCtx *BindContext)
 			if col.Name == tableDef.Pkey.PkeyColName {
 				lockTarget := &plan.LockTarget{
 					TableId:            tableDef.TblId,
+					ObjRef:             DeepCopyObjectRef(dmlCtx.objRefs[i]),
 					PrimaryColIdxInBat: int32(pkPos),
 					PrimaryColRelPos:   selectNodeTag,
 					PrimaryColTyp:      col.Typ,
@@ -310,6 +311,7 @@ func (builder *QueryBuilder) bindDelete(stmt *tree.Delete, bindCtx *BindContext)
 					if col.Name == idxNode.TableDef.Pkey.PkeyColName {
 						lockTargets = append(lockTargets, &plan.LockTarget{
 							TableId:            idxNode.TableDef.TblId,
+							ObjRef:             DeepCopyObjectRef(idxNode.ObjRef),
 							PrimaryColIdxInBat: int32(pkPos),
 							PrimaryColRelPos:   idxNode.BindingTags[0],
 							PrimaryColTyp:      col.Typ,

pkg/sql/plan/bind_insert.go

@@ -191,6 +191,7 @@ func (builder *QueryBuilder) appendDedupAndMultiUpdateNodesForBindInsert(
 		if col.Name == pkName && pkName != catalog.FakePrimaryKeyColName {
 			lockTarget := &plan.LockTarget{
 				TableId:            tableDef.TblId,
+				ObjRef:             DeepCopyObjectRef(objRef),
 				PrimaryColIdxInBat: int32(colName2Idx[tableDef.Name+"."+col.Name]),
 				PrimaryColRelPos:   selectTag,
 				PrimaryColTyp:      col.Typ,
@@ -221,6 +222,7 @@ func (builder *QueryBuilder) appendDedupAndMultiUpdateNodesForBindInsert(
 		}
 		lockTarget := &plan.LockTarget{
 			TableId:            idxTableDef.TblId,
+			ObjRef:             DeepCopyObjectRef(dmlCtx.objRefs[0]),
 			PrimaryColIdxInBat: pkIdxInBat,
 			PrimaryColRelPos:   selectTag,
 			PrimaryColTyp:      selectNode.ProjectList[int(pkIdxInBat)].Typ,

pkg/sql/plan/bind_update.go

@@ -326,6 +326,7 @@ func (builder *QueryBuilder) bindUpdate(stmt *tree.Update, bindCtx *BindContext)
 			if col.Name == tableDef.Pkey.PkeyColName {
 				lockTarget := &plan.LockTarget{
 					TableId:            tableDef.TblId,
+					ObjRef:             DeepCopyObjectRef(dmlCtx.objRefs[i]),
 					PrimaryColIdxInBat: int32(finalColIdx),
 					PrimaryColRelPos:   finalProjTag,
 					PrimaryColTyp:      col.Typ,
@@ -337,6 +338,7 @@ func (builder *QueryBuilder) bindUpdate(stmt *tree.Update, bindCtx *BindContext)
 					// need lock oldPk by old partition idx, lock new pk by new partition idx
 					lockTarget := &plan.LockTarget{
 						TableId:            tableDef.TblId,
+						ObjRef:             DeepCopyObjectRef(dmlCtx.objRefs[i]),
 						PrimaryColIdxInBat: int32(finalColIdx),
 						PrimaryColRelPos:   finalProjTag,
 						PrimaryColTyp:      col.Typ,

pkg/sql/plan/deepcopy.go

@@ -178,6 +178,7 @@ func DeepCopyLockTarget(target *plan.LockTarget) *plan.LockTarget {
 	}
 	return &plan.LockTarget{
 		TableId:            target.TableId,
+		ObjRef:             DeepCopyObjectRef(target.ObjRef),
 		PrimaryColIdxInBat: target.PrimaryColIdxInBat,
 		PrimaryColTyp:      target.PrimaryColTyp,
 		RefreshTsIdxInBat:  target.RefreshTsIdxInBat,

proto/pipeline.proto

@@ -201,6 +201,7 @@ message LockTarget {
   lock.LockMode Mode         = 8;
   plan.Expr lock_rows = 9;
   bool lock_table_at_the_end = 10;
+  plan.ObjectRef obj_ref = 11;
 }
 
 message LockOp {

proto/plan.proto

@@ -979,6 +979,7 @@ message LockTarget {
 	int32 filter_col_rel_pos  = 12;
 	Expr lock_rows = 13;
     bool lock_table_at_the_end = 14;
+    ObjectRef obj_ref = 15;
 }
 
 message PreInsertUkCtx {