*: table partition dynamic prune mode with Physical Table ID from unistore

cmd/explaintest/r/generated_columns.result

@@ -1,4 +1,3 @@
-set @@tidb_partition_prune_mode='dynamic';
 DROP TABLE IF EXISTS person;
 CREATE TABLE person (
 id INT NOT NULL AUTO_INCREMENT PRIMARY KEY,
@@ -90,6 +89,8 @@ Projection	5.00	root		test.sgc1.j1, test.sgc1.j2, test.sgc1.a, test.sgc1.b, test
   └─TableReader(Probe)	5.00	root		data:Selection
     └─Selection	5.00	cop[tikv]		not(isnull(test.sgc1.a))
       └─TableFullScan	5.00	cop[tikv]	table:sgc1	keep order:false
+set @old_prune_mode = @@tidb_partition_prune_mode;
+set @@tidb_partition_prune_mode='static';
 DROP TABLE IF EXISTS sgc3;
 CREATE TABLE sgc3 (
 j JSON,
@@ -136,6 +137,31 @@ PartitionUnion	23263.33	root
 └─TableReader	3323.33	root		data:Selection
   └─Selection	3323.33	cop[tikv]		lt(test.sgc3.a, 7)
     └─TableFullScan	10000.00	cop[tikv]	table:sgc3, partition:max	keep order:false, stats:pseudo
+set @@tidb_partition_prune_mode='dynamic';
+DROP TABLE sgc3;
+CREATE TABLE sgc3 (
+j JSON,
+a INT AS (JSON_EXTRACT(j, "$.a")) STORED
+)
+PARTITION BY RANGE (a) (
+PARTITION p0 VALUES LESS THAN (1),
+PARTITION p1 VALUES LESS THAN (2),
+PARTITION p2 VALUES LESS THAN (3),
+PARTITION p3 VALUES LESS THAN (4),
+PARTITION p4 VALUES LESS THAN (5),
+PARTITION p5 VALUES LESS THAN (6),
+PARTITION max VALUES LESS THAN MAXVALUE);
+EXPLAIN format = 'brief' SELECT * FROM sgc3 WHERE a <= 1;
+id	estRows	task	access object	operator info
+TableReader	3323.33	root	partition:p0,p1	data:Selection
+└─Selection	3323.33	cop[tikv]		le(test.sgc3.a, 1)
+  └─TableFullScan	10000.00	cop[tikv]	table:sgc3	keep order:false, stats:pseudo
+EXPLAIN format = 'brief' SELECT * FROM sgc3 WHERE a < 7;
+id	estRows	task	access object	operator info
+TableReader	3323.33	root	partition:all	data:Selection
+└─Selection	3323.33	cop[tikv]		lt(test.sgc3.a, 7)
+  └─TableFullScan	10000.00	cop[tikv]	table:sgc3	keep order:false, stats:pseudo
+set @@tidb_partition_prune_mode = @old_prune_mode;
 DROP TABLE IF EXISTS t1;
 CREATE TABLE t1(a INT, b INT AS (a+1) VIRTUAL, c INT AS (b+1) VIRTUAL, d INT AS (c+1) VIRTUAL, KEY(b), INDEX IDX(c, d));
 INSERT INTO t1 (a) VALUES (0);

cmd/explaintest/r/select.result

@@ -354,6 +354,7 @@ Projection_3	1.00	root		sysdate()->Column#1, sleep(1)->Column#2, sysdate()->Colu
 └─TableDual_4	1.00	root		rows:1
 drop table if exists th;
 set @@session.tidb_enable_table_partition = '1';
+set @@session.tidb_partition_prune_mode = 'static';
 create table th (a int, b int) partition by hash(a) partitions 3;
 insert into th values (0,0),(1,1),(2,2),(3,3),(4,4),(5,5),(6,6),(7,7),(8,8);
 insert into th values (-1,-1),(-2,-2),(-3,-3),(-4,-4),(-5,-5),(-6,-6),(-7,-7),(-8,-8);
@@ -378,6 +379,21 @@ PartitionUnion_8	20000.00	root
 │ └─TableFullScan_9	10000.00	cop[tikv]	table:th, partition:p1	keep order:false, stats:pseudo
 └─TableReader_12	10000.00	root		data:TableFullScan_11
   └─TableFullScan_11	10000.00	cop[tikv]	table:th, partition:p2	keep order:false, stats:pseudo
+set @@session.tidb_partition_prune_mode = 'dynamic';
+desc select * from th where a=-2;
+id	estRows	task	access object	operator info
+TableReader_7	10.00	root	partition:p2	data:Selection_6
+└─Selection_6	10.00	cop[tikv]		eq(test.th.a, -2)
+  └─TableFullScan_5	10000.00	cop[tikv]	table:th	keep order:false, stats:pseudo
+desc select * from th;
+id	estRows	task	access object	operator info
+TableReader_5	10000.00	root	partition:all	data:TableFullScan_4
+└─TableFullScan_4	10000.00	cop[tikv]	table:th	keep order:false, stats:pseudo
+desc select * from th partition (p2,p1);
+id	estRows	task	access object	operator info
+TableReader_5	10000.00	root	partition:p1,p2	data:TableFullScan_4
+└─TableFullScan_4	10000.00	cop[tikv]	table:th	keep order:false, stats:pseudo
+set @@session.tidb_partition_prune_mode = DEFAULT;
 drop table if exists t;
 create table t(a int, b int);
 explain format = 'brief' select a != any (select a from t t2) from t t1;

cmd/explaintest/t/generated_columns.test

@@ -2,7 +2,6 @@
 -- Most of the cases are ported from other tests to make sure generated columns behaves the same.
 
 -- Stored generated columns as indices
-set @@tidb_partition_prune_mode='dynamic';
 
 DROP TABLE IF EXISTS person;
 CREATE TABLE person (
@@ -74,6 +73,8 @@ EXPLAIN format = 'brief' SELECT * from sgc1 join sgc2 on sgc1.a=sgc2.a;
 
 -- Stored generated columns as partition columns
 
+set @old_prune_mode = @@tidb_partition_prune_mode;
+set @@tidb_partition_prune_mode='static';
 DROP TABLE IF EXISTS sgc3;
 CREATE TABLE sgc3 (
 j JSON,
@@ -91,6 +92,26 @@ PARTITION max VALUES LESS THAN MAXVALUE);
 EXPLAIN format = 'brief' SELECT * FROM sgc3 WHERE a <= 1;
 EXPLAIN format = 'brief' SELECT * FROM sgc3 WHERE a < 7;
 
+set @@tidb_partition_prune_mode='dynamic';
+DROP TABLE sgc3;
+CREATE TABLE sgc3 (
+j JSON,
+a INT AS (JSON_EXTRACT(j, "$.a")) STORED
+)
+PARTITION BY RANGE (a) (
+PARTITION p0 VALUES LESS THAN (1),
+PARTITION p1 VALUES LESS THAN (2),
+PARTITION p2 VALUES LESS THAN (3),
+PARTITION p3 VALUES LESS THAN (4),
+PARTITION p4 VALUES LESS THAN (5),
+PARTITION p5 VALUES LESS THAN (6),
+PARTITION max VALUES LESS THAN MAXVALUE);
+
+EXPLAIN format = 'brief' SELECT * FROM sgc3 WHERE a <= 1;
+EXPLAIN format = 'brief' SELECT * FROM sgc3 WHERE a < 7;
+
+set @@tidb_partition_prune_mode = @old_prune_mode;
+
 -- Virtual generated columns as indices
 
 DROP TABLE IF EXISTS t1;

cmd/explaintest/t/select.test

@@ -181,12 +181,18 @@ desc select sysdate(), sleep(1), sysdate();
 # test select partition table
 drop table if exists th;
 set @@session.tidb_enable_table_partition = '1';
+set @@session.tidb_partition_prune_mode = 'static';
 create table th (a int, b int) partition by hash(a) partitions 3;
 insert into th values (0,0),(1,1),(2,2),(3,3),(4,4),(5,5),(6,6),(7,7),(8,8);
 insert into th values (-1,-1),(-2,-2),(-3,-3),(-4,-4),(-5,-5),(-6,-6),(-7,-7),(-8,-8);
 desc select * from th where a=-2;
 desc select * from th;
 desc select * from th partition (p2,p1);
+set @@session.tidb_partition_prune_mode = 'dynamic';
+desc select * from th where a=-2;
+desc select * from th;
+desc select * from th partition (p2,p1);
+set @@session.tidb_partition_prune_mode = DEFAULT;
 
 # test != any(subq) and = all(subq)
 drop table if exists t;

executor/analyze_test.go

@@ -663,6 +663,7 @@ func TestFailedAnalyzeRequest(t *testing.T) {
 	tk.MustExec("set @@tidb_analyze_version = 1")
 	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/executor/buildStatsFromResult", `return(true)`))
 	_, err := tk.Exec("analyze table t")
+	require.NotNil(t, err)
 	require.Equal(t, "mock buildStatsFromResult error", err.Error())
 	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/executor/buildStatsFromResult"))
 }

executor/builder.go

@@ -3096,7 +3096,8 @@ func buildNoRangeTableReader(b *executorBuilder, v *plannercore.PhysicalTableRea
 
 func extraPIDColumnIndex(schema *expression.Schema) offsetOptional {
 	for idx, col := range schema.Columns {
-		if col.ID == model.ExtraPidColID {
+		// TODO: Handle partitioned global index, i.e. both ExtraPidColID and ExtraPhysTblID is used.
+		if col.ID == model.ExtraPidColID || col.ID == model.ExtraPhysTblID {
 			return newOffset(idx)
 		}
 	}
@@ -3169,6 +3170,7 @@ func (b *executorBuilder) buildTableReader(v *plannercore.PhysicalTableReader) E
 
 	tmp, _ := b.is.TableByID(ts.Table.ID)
 	tbl := tmp.(table.PartitionedTable)
+	// TODO: Should this be done before checking dirty / ongoing transactions for avoiding adding PID columns?
 	partitions, err := partitionPruning(b.ctx, tbl, v.PartitionInfo.PruningConds, v.PartitionInfo.PartitionNames, v.PartitionInfo.Columns, v.PartitionInfo.ColumnNames)
 	if err != nil {
 		b.err = err
@@ -3649,6 +3651,9 @@ func buildNoRangeIndexMergeReader(b *executorBuilder, v *plannercore.PhysicalInd
 	}
 	collectTable := false
 	e.tableRequest.CollectRangeCounts = &collectTable
+	if tblInfo.Meta().Partition != nil {
+		e.extraPIDColumnIndex = extraPIDColumnIndex(v.Schema())
+	}
 	return e, nil
 }
 

executor/delete.go

@@ -122,7 +122,7 @@ func (e *DeleteExec) deleteSingleTableByChunk(ctx context.Context) error {
 
 			datumRow := make([]types.Datum, 0, len(fields))
 			for i, field := range fields {
-				if columns[i].ID == model.ExtraPidColID {
+				if columns[i].ID == model.ExtraPidColID || columns[i].ID == model.ExtraPhysTblID {
 					continue
 				}
 

executor/executor.go

@@ -949,6 +949,35 @@ func (e *SelectLockExec) Next(ctx context.Context, req *chunk.Chunk) error {
 					if err != nil {
 						return err
 					}
+
+					// Debug print.
+					if physicalID == 0 || (handle.IsInt() && handle.IntValue() == 0) || (handle.Len() == 0) {
+						txn, _ := e.ctx.Txn(false)
+						sql := e.ctx.GetSessionVars().StmtCtx.OriginalSQL
+						tblIDs := make([]int64, 0, len(e.tblID2Handle))
+						for tmpTblID := range e.tblID2Handle {
+							tblIDs = append(tblIDs, tmpTblID)
+						}
+						partTblLen := len(e.partitionedTable)
+						tblid2ColIdxMapKeys := make([]int64, 0, len(e.tblID2PIDColumnIndex))
+						tblid2ColIdxMapVals := make([]int, 0, len(e.tblID2PIDColumnIndex))
+						for key, val := range e.tblID2PIDColumnIndex {
+							tblid2ColIdxMapKeys = append(tblid2ColIdxMapKeys, key)
+							tblid2ColIdxMapVals = append(tblid2ColIdxMapVals, val)
+						}
+						logutil.Logger(ctx).Error("[for debug] the physicalID or handle value is unexpected",
+							zap.Uint64("ts", txn.StartTS()),
+							zap.Int64("id", id),
+							zap.Int64("physicalID", physicalID),
+							zap.Stringer("handle", handle),
+							zap.String("sql", sql),
+							zap.Int64s("tblIDs", tblIDs),
+							zap.Int("partTblLen", partTblLen),
+							zap.Int64s("tblID2PIDColIdxMapKeys", tblid2ColIdxMapKeys),
+							zap.Ints("tblid2ColIdxMapVals", tblid2ColIdxMapVals))
+						panic("unexpected lock key, check the tidb log with for debug")
+					}
+
 					e.keys = append(e.keys, tablecodec.EncodeRowKeyWithHandle(physicalID, handle))
 				}
 			}

executor/index_merge_reader.go

@@ -116,6 +116,10 @@ type IndexMergeReaderExecutor struct {
 	isCorColInPartialFilters []bool
 	isCorColInTableFilter    bool
 	isCorColInPartialAccess  []bool
+
+	// extraPIDColumnIndex is used for partition reader to add an extra partition ID column
+	// in static tidb_partition_prune_mode
+	extraPIDColumnIndex offsetOptional
 }
 
 // Open implements the Executor Open interface
@@ -599,16 +603,17 @@ func (e *IndexMergeReaderExecutor) startIndexMergeTableScanWorker(ctx context.Co
 
 func (e *IndexMergeReaderExecutor) buildFinalTableReader(ctx context.Context, tbl table.Table, handles []kv.Handle) (_ Executor, err error) {
 	tableReaderExec := &TableReaderExecutor{
-		baseExecutor:     newBaseExecutor(e.ctx, e.schema, e.getTablePlanRootID()),
-		table:            tbl,
-		dagPB:            e.tableRequest,
-		startTS:          e.startTS,
-		readReplicaScope: e.readReplicaScope,
-		isStaleness:      e.isStaleness,
-		streaming:        e.tableStreaming,
-		columns:          e.columns,
-		feedback:         statistics.NewQueryFeedback(0, nil, 0, false),
-		plans:            e.tblPlans,
+		baseExecutor:        newBaseExecutor(e.ctx, e.schema, e.getTablePlanRootID()),
+		table:               tbl,
+		dagPB:               e.tableRequest,
+		startTS:             e.startTS,
+		readReplicaScope:    e.readReplicaScope,
+		isStaleness:         e.isStaleness,
+		streaming:           e.tableStreaming,
+		columns:             e.columns,
+		feedback:            statistics.NewQueryFeedback(0, nil, 0, false),
+		plans:               e.tblPlans,
+		extraPIDColumnIndex: e.extraPIDColumnIndex,
 	}
 	if e.isCorColInTableFilter {
 		if tableReaderExec.dagPB.Executors, _, err = constructDistExec(e.ctx, e.tblPlans); err != nil {

executor/index_merge_reader_test.go

@@ -20,6 +20,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb/util"
@@ -474,3 +475,80 @@ func (test *testSerialSuite2) TestIndexMergeSplitTable(c *C) {
 	tk.MustExec("SPLIT TABLE tab2 BY (5);")
 	tk.MustQuery("SELECT /*+ use_index_merge(tab2) */ pk FROM tab2 WHERE (col4 > 565.89 OR col0 > 68 ) and col0 > 10 order by 1;").Check(testkit.Rows("0", "1", "2", "3", "4", "5", "6", "7"))
 }
+
+func (test *testSerialSuite2) TestPessimisticLockOnPartitionForIndexMerge(c *C) {
+	// Same purpose with TestPessimisticLockOnPartition, but test IndexMergeReader.
+	tk := testkit.NewTestKit(c, test.store)
+	tk.MustExec("use test")
+
+	tk.MustExec("drop table if exists t1, t2")
+	tk.MustExec(`create table t1 (c_datetime datetime, c1 int, c2 int, primary key (c_datetime), key(c1), key(c2))
+			partition by range (to_days(c_datetime)) (
+				partition p0 values less than (to_days('2020-02-01')),
+				partition p1 values less than (to_days('2020-04-01')),
+				partition p2 values less than (to_days('2020-06-01')),
+				partition p3 values less than maxvalue)`)
+	tk.MustExec("create table t2 (c_datetime datetime, unique key(c_datetime))")
+	tk.MustExec("insert into t1 values ('2020-06-26 03:24:00', 1, 1), ('2020-02-21 07:15:33', 2, 2), ('2020-04-27 13:50:58', 3, 3)")
+	tk.MustExec("insert into t2 values ('2020-01-10 09:36:00'), ('2020-02-04 06:00:00'), ('2020-06-12 03:45:18')")
+	tk.MustExec("analyze table t1")
+	tk.MustExec("analyze table t2")
+
+	tk1 := testkit.NewTestKit(c, test.store)
+	tk1.MustExec("use test")
+	tk1.MustExec("set @@tidb_partition_prune_mode = 'static'")
+
+	tk.MustExec("set @@tidb_partition_prune_mode = 'static'")
+	tk.MustExec("begin pessimistic")
+	tk.MustQuery(`explain format='brief' select /*+ use_index_merge(t1) */ c1 from t1 join t2
+			on t1.c_datetime >= t2.c_datetime
+			where t1.c1 < 10 or t1.c2 < 10 for update`).Check(testkit.Rows(
+		"Projection 16635.64 root  test.t1.c1",
+		"└─SelectLock 16635.64 root  for update 0",
+		"  └─Projection 16635.64 root  test.t1.c1, test.t1._tidb_rowid, test.t1._tidb_pid, test.t2._tidb_rowid",
+		"    └─HashJoin 16635.64 root  CARTESIAN inner join, other cond:ge(test.t1.c_datetime, test.t2.c_datetime)",
+		"      ├─IndexReader(Build) 3.00 root  index:IndexFullScan",
+		"      │ └─IndexFullScan 3.00 cop[tikv] table:t2, index:c_datetime(c_datetime) keep order:false",
+		"      └─PartitionUnion(Probe) 5545.21 root  ",
+		"        ├─IndexMerge 5542.21 root  ",
+		"        │ ├─IndexRangeScan(Build) 3323.33 cop[tikv] table:t1, partition:p0, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo",
+		"        │ ├─IndexRangeScan(Build) 3323.33 cop[tikv] table:t1, partition:p0, index:c2(c2) range:[-inf,10), keep order:false, stats:pseudo",
+		"        │ └─TableRowIDScan(Probe) 5542.21 cop[tikv] table:t1, partition:p0 keep order:false, stats:pseudo",
+		"        ├─IndexMerge 1.00 root  ",
+		"        │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p1, index:c1(c1) range:[-inf,10), keep order:false",
+		"        │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p1, index:c2(c2) range:[-inf,10), keep order:false",
+		"        │ └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p1 keep order:false",
+		"        ├─IndexMerge 1.00 root  ",
+		"        │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p2, index:c1(c1) range:[-inf,10), keep order:false",
+		"        │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p2, index:c2(c2) range:[-inf,10), keep order:false",
+		"        │ └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p2 keep order:false",
+		"        └─IndexMerge 1.00 root  ",
+		"          ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p3, index:c1(c1) range:[-inf,10), keep order:false",
+		"          ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p3, index:c2(c2) range:[-inf,10), keep order:false",
+		"          └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p3 keep order:false",
+	))
+	tk.MustQuery(`select /*+ use_index_merge(t1) */ c1 from t1 join t2
+			on t1.c_datetime >= t2.c_datetime
+			where t1.c1 < 10 or t1.c2 < 10 for update`).Sort().Check(testkit.Rows("1", "1", "1", "2", "2", "3", "3"))
+	tk1.MustExec("begin pessimistic")
+
+	ch := make(chan int32, 5)
+	go func() {
+		tk1.MustExec("update t1 set c_datetime = '2020-06-26 03:24:00' where c1 = 1")
+		ch <- 0
+		tk1.MustExec("rollback")
+		ch <- 0
+	}()
+
+	// Leave 50ms for tk1 to run, tk1 should be blocked at the update operation.
+	time.Sleep(50 * time.Millisecond)
+	ch <- 1
+
+	tk.MustExec("commit")
+	// tk1 should be blocked until tk commit, check the order.
+	c.Assert(<-ch, Equals, int32(1))
+	c.Assert(<-ch, Equals, int32(0))
+	<-ch // wait for goroutine to quit.
+
+	// TODO: add support for index merge reader in dynamic tidb_partition_prune_mode
+}

executor/table_reader.go

@@ -244,9 +244,12 @@ func (e *TableReaderExecutor) Next(ctx context.Context, req *chunk.Chunk) error
 	// When 'select ... for update' work on a partitioned table, the table reader should
 	// add the partition ID as an extra column. The SelectLockExec need this information
 	// to construct the lock key.
-	physicalID := getPhysicalTableID(e.table)
 	if e.extraPIDColumnIndex.valid() {
-		fillExtraPIDColumn(req, e.extraPIDColumnIndex.value(), physicalID)
+		physicalID := getPhysicalTableID(e.table)
+		if physicalID != e.table.Meta().ID {
+			// table partition in static prune mode (one TableReaderExecutor per partition)
+			fillExtraPIDColumn(req, e.extraPIDColumnIndex.value(), physicalID)
+		}
 	}
 
 	return nil