cherry pick pingcap#14921 to release-3.0

Signed-off-by: sre-bot <sre-bot@pingcap.com>

executor/builder.go

@@ -557,8 +557,15 @@ func (b *executorBuilder) buildSelectLock(v *plannercore.PhysicalLock) Executor
 		return src
 	}
 	e := &SelectLockExec{
+<<<<<<< HEAD
 		baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ExplainID(), src),
 		Lock:         v.Lock,
+=======
+		baseExecutor:     newBaseExecutor(b.ctx, v.Schema(), v.ExplainID(), src),
+		Lock:             v.Lock,
+		tblID2Handle:     v.TblID2Handle,
+		partitionedTable: v.PartitionedTable,
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 	}
 	return e
 }

executor/executor.go

@@ -742,6 +742,15 @@ type SelectLockExec struct {
 
 	Lock ast.SelectLockType
 	keys []kv.Key
+<<<<<<< HEAD
+=======
+
+	tblID2Handle     map[int64][]*expression.Column
+	partitionedTable []table.PartitionedTable
+
+	// tblID2Table is cached to reduce cost.
+	tblID2Table map[int64]table.PartitionedTable
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 }
 
 // Open implements the Executor Open interface.
@@ -755,6 +764,18 @@ func (e *SelectLockExec) Open(ctx context.Context) error {
 		// This operation is only for schema validator check.
 		txnCtx.UpdateDeltaForTable(id, 0, 0, map[int64]int64{})
 	}
+
+	if len(e.tblID2Handle) > 0 && len(e.partitionedTable) > 0 {
+		e.tblID2Table = make(map[int64]table.PartitionedTable, len(e.partitionedTable))
+		for id := range e.tblID2Handle {
+			for _, p := range e.partitionedTable {
+				if id == p.Meta().ID {
+					e.tblID2Table[id] = p
+				}
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -774,12 +795,30 @@ func (e *SelectLockExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	if len(e.Schema().TblID2Handle) == 0 || e.Lock != ast.SelectLockForUpdate {
 		return nil
 	}
-	if req.NumRows() != 0 {
+
+	if req.NumRows() > 0 {
 		iter := chunk.NewIterator4Chunk(req)
+<<<<<<< HEAD
 		for id, cols := range e.Schema().TblID2Handle {
 			for _, col := range cols {
 				for row := iter.Begin(); row != iter.End(); row = iter.Next() {
 					e.keys = append(e.keys, tablecodec.EncodeRowKeyWithHandle(id, row.GetInt64(col.Index)))
+=======
+		for row := iter.Begin(); row != iter.End(); row = iter.Next() {
+			for id, cols := range e.tblID2Handle {
+				physicalID := id
+				if pt, ok := e.tblID2Table[id]; ok {
+					// On a partitioned table, we have to use physical ID to encode the lock key!
+					p, err := pt.GetPartitionByRow(e.ctx, row.GetDatumRow(e.base().retFieldTypes))
+					if err != nil {
+						return err
+					}
+					physicalID = p.GetPhysicalID()
+				}
+
+				for _, col := range cols {
+					e.keys = append(e.keys, tablecodec.EncodeRowKeyWithHandle(physicalID, row.GetInt64(col.Index)))
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 				}
 			}
 		}

executor/write.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/expression"
+	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
@@ -117,8 +118,23 @@ func updateRecord(ctx sessionctx.Context, h int64, oldData, newData []types.Datu
 		if ctx.GetSessionVars().ClientCapability&mysql.ClientFoundRows > 0 {
 			sc.AddAffectedRows(1)
 		}
+<<<<<<< HEAD
 		unchangedRowKey := tablecodec.EncodeRowKeyWithHandle(t.Meta().ID, h)
 		txnCtx := ctx.GetSessionVars().TxnCtx
+=======
+
+		physicalID := t.Meta().ID
+		if pt, ok := t.(table.PartitionedTable); ok {
+			p, err := pt.GetPartitionByRow(sctx, oldData)
+			if err != nil {
+				return false, false, 0, err
+			}
+			physicalID = p.GetPhysicalID()
+		}
+
+		unchangedRowKey := tablecodec.EncodeRowKeyWithHandle(physicalID, h)
+		txnCtx := sctx.GetSessionVars().TxnCtx
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 		if txnCtx.IsPessimistic {
 			txnCtx.AddUnchangedRowKey(unchangedRowKey)
 		}

planner/core/exhaust_physical_plans.go

@@ -1280,7 +1280,13 @@ func (p *LogicalLimit) exhaustPhysicalPlans(prop *property.PhysicalProperty) []P
 func (p *LogicalLock) exhaustPhysicalPlans(prop *property.PhysicalProperty) []PhysicalPlan {
 	childProp := prop.Clone()
 	lock := PhysicalLock{
+<<<<<<< HEAD
 		Lock: p.Lock,
+=======
+		Lock:             p.Lock,
+		TblID2Handle:     p.tblID2Handle,
+		PartitionedTable: p.partitionedTable,
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 	}.Init(p.ctx, p.stats.ScaleByExpectCnt(prop.ExpectedCnt), childProp)
 	return []PhysicalPlan{lock}
 }

planner/core/logical_plan_builder.go

@@ -2239,6 +2239,7 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName) (L
 
 	if tableInfo.GetPartitionInfo() != nil {
 		b.optFlag = b.optFlag | flagPartitionProcessor
+		b.partitionedTable = append(b.partitionedTable, tbl.(table.PartitionedTable))
 		// check partition by name.
 		for _, name := range tn.PartitionNames {
 			_, err = tables.FindPartitionByName(tableInfo, name.L)

planner/core/logical_plans.go

@@ -684,7 +684,13 @@ type LogicalLimit struct {
 type LogicalLock struct {
 	baseLogicalPlan
 
+<<<<<<< HEAD
 	Lock ast.SelectLockType
+=======
+	Lock             ast.SelectLockType
+	tblID2Handle     map[int64][]*expression.Column
+	partitionedTable []table.PartitionedTable
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 }
 
 // WindowFrame represents a window function frame.

planner/core/physical_plans.go

@@ -21,6 +21,7 @@ import (
 	"github.com/pingcap/tidb/planner/property"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/statistics"
+	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/ranger"
 )
@@ -280,6 +281,12 @@ type PhysicalLock struct {
 	basePhysicalPlan
 
 	Lock ast.SelectLockType
+<<<<<<< HEAD
+=======
+
+	TblID2Handle     map[int64][]*expression.Column
+	PartitionedTable []table.PartitionedTable
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 }
 
 // PhysicalLimit is the physical operator of Limit.

planner/core/planbuilder.go

@@ -199,7 +199,69 @@ type PlanBuilder struct {
 	// "STRAIGHT_JOIN" option.
 	inStraightJoin bool
 
+<<<<<<< HEAD
 	windowSpecs map[string]*ast.WindowSpec
+=======
+	// handleHelper records the handle column position for tables. Delete/Update/SelectLock/UnionScan may need this information.
+	// It collects the information by the following procedure:
+	//   Since we build the plan tree from bottom to top, we maintain a stack to record the current handle information.
+	//   If it's a dataSource/tableDual node, we create a new map.
+	//   If it's a aggregation, we pop the map and push a nil map since no handle information left.
+	//   If it's a union, we pop all children's and push a nil map.
+	//   If it's a join, we pop its children's out then merge them and push the new map to stack.
+	//   If we meet a subquery, it's clearly that it's a independent problem so we just pop one map out when we finish building the subquery.
+	handleHelper *handleColHelper
+
+	hintProcessor *BlockHintProcessor
+	// selectOffset is the offsets of current processing select stmts.
+	selectOffset []int
+
+	// SelectLock need this information to locate the lock on partitions.
+	partitionedTable []table.PartitionedTable
+}
+
+type handleColHelper struct {
+	id2HandleMapStack []map[int64][]*expression.Column
+	stackTail         int
+}
+
+func (hch *handleColHelper) appendColToLastMap(tblID int64, col *expression.Column) {
+	tailMap := hch.id2HandleMapStack[hch.stackTail-1]
+	tailMap[tblID] = append(tailMap[tblID], col)
+}
+
+func (hch *handleColHelper) popMap() map[int64][]*expression.Column {
+	ret := hch.id2HandleMapStack[hch.stackTail-1]
+	hch.stackTail--
+	hch.id2HandleMapStack = hch.id2HandleMapStack[:hch.stackTail]
+	return ret
+}
+
+func (hch *handleColHelper) pushMap(m map[int64][]*expression.Column) {
+	hch.id2HandleMapStack = append(hch.id2HandleMapStack, m)
+	hch.stackTail++
+}
+
+func (hch *handleColHelper) mergeAndPush(m1, m2 map[int64][]*expression.Column) {
+	newMap := make(map[int64][]*expression.Column, mathutil.Max(len(m1), len(m2)))
+	for k, v := range m1 {
+		newMap[k] = make([]*expression.Column, len(v))
+		copy(newMap[k], v)
+	}
+	for k, v := range m2 {
+		if _, ok := newMap[k]; ok {
+			newMap[k] = append(newMap[k], v...)
+		} else {
+			newMap[k] = make([]*expression.Column, len(v))
+			copy(newMap[k], v)
+		}
+	}
+	hch.pushMap(newMap)
+}
+
+func (hch *handleColHelper) tailMap() map[int64][]*expression.Column {
+	return hch.id2HandleMapStack[hch.stackTail-1]
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 }
 
 // GetVisitInfo gets the visitInfo of the PlanBuilder.
@@ -575,7 +637,15 @@ func removeIgnoredPaths(paths, ignoredPaths []*accessPath, tblInfo *model.TableI
 }
 
 func (b *PlanBuilder) buildSelectLock(src LogicalPlan, lock ast.SelectLockType) *LogicalLock {
+<<<<<<< HEAD
 	selectLock := LogicalLock{Lock: lock}.Init(b.ctx)
+=======
+	selectLock := LogicalLock{
+		Lock:             lock,
+		tblID2Handle:     b.handleHelper.tailMap(),
+		partitionedTable: b.partitionedTable,
+	}.Init(b.ctx)
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 	selectLock.SetChildren(src)
 	return selectLock
 }

planner/core/rule_column_pruning.go

@@ -368,7 +368,17 @@ func (p *LogicalLock) PruneColumns(parentUsedCols []*expression.Column) error {
 		return p.baseLogicalPlan.PruneColumns(parentUsedCols)
 	}
 
+<<<<<<< HEAD
 	for _, cols := range p.children[0].Schema().TblID2Handle {
+=======
+	if len(p.partitionedTable) > 0 {
+		// If the children include partitioned tables, do not prune columns.
+		// Because the executor needs the partitioned columns to calculate the lock key.
+		return p.children[0].PruneColumns(p.Schema().Columns)
+	}
+
+	for _, cols := range p.tblID2Handle {
+>>>>>>> b3469e7... *: fix a bug that the pessimistic lock doesn't work on a partition (#14921)
 		parentUsedCols = append(parentUsedCols, cols...)
 	}
 	return p.children[0].PruneColumns(parentUsedCols)