executor: fix pointGet and batchPointGet for global index

pkg/executor/batch_point_get.go

@@ -226,12 +226,16 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
 			}
 
 			var physID int64
-			if len(e.planPhysIDs) > 0 {
-				physID = e.planPhysIDs[i]
+			if e.partPos == core.GlobalWithoutColumnPos {
+				physID = e.tblInfo.ID
 			} else {
-				physID, err = core.GetPhysID(e.tblInfo, e.partExpr, e.partPos, idxVals[e.partPos])
-				if err != nil {
-					continue
+				if len(e.planPhysIDs) > 0 {
+					physID = e.planPhysIDs[i]
+				} else {
+					physID, err = core.GetPhysID(e.tblInfo, e.partExpr, e.partPos, idxVals[e.partPos])
+					if err != nil {
+						continue
+					}
 				}
 			}
 
@@ -300,8 +304,18 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
 				indexKeys = append(indexKeys, key)
 			}
 			if e.tblInfo.Partition != nil {
-				pid := tablecodec.DecodeTableID(key)
-				e.physIDs = append(e.physIDs, pid)
+				var pid int64
+				if e.idxInfo.Global {
+					segs := tablecodec.SplitIndexValue(handleVal)
+					_, pid, err = codec.DecodeInt(segs.PartitionID)
+					if err != nil {
+						return err
+					}
+					e.physIDs = append(e.physIDs, pid)
+				} else {
+					pid = tablecodec.DecodeTableID(key)
+					e.physIDs = append(e.physIDs, pid)
+				}
 				if e.lock {
 					e.UpdateDeltaForTableID(pid)
 				}

pkg/executor/point_get.go

@@ -303,6 +303,14 @@ func (e *PointGetExecutor) Next(ctx context.Context, req *chunk.Chunk) error {
 				// Wait `UPDATE` finished
 				failpoint.InjectContext(ctx, "pointGetRepeatableReadTest-step2", nil)
 			})
+			if e.idxInfo.Global {
+				segs := tablecodec.SplitIndexValue(e.handleVal)
+				_, pid, err := codec.DecodeInt(segs.PartitionID)
+				if err != nil {
+					return err
+				}
+				tblID = pid
+			}
 		}
 	}
 

pkg/executor/point_get_test.go

@@ -375,3 +375,39 @@ func TestWithTiDBSnapshot(t *testing.T) {
 
 	tk.MustQuery("select * from xx").Check(testkit.Rows("1", "7"))
 }
+
+func TestGlobalIndexPointGet(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("set tidb_enable_global_index=true")
+	defer func() {
+		tk.MustExec("set tidb_enable_global_index=default")
+	}()
+
+	tk.MustExec(`CREATE TABLE t ( a int, b int, c int default 0)
+						PARTITION BY RANGE (a) (
+		PARTITION p0 VALUES LESS THAN (10),
+		PARTITION p1 VALUES LESS THAN (20),
+		PARTITION p2 VALUES LESS THAN (30),
+		PARTITION p3 VALUES LESS THAN (40))`)
+	tk.MustExec("INSERT INTO t(a, b) values(1, 1), (2, 2), (3, 3), (15, 15), (25, 25), (35, 35)")
+	tk.MustExec("ALTER TABLE t ADD UNIQUE INDEX idx(b)")
+	tk.MustExec("analyze table t")
+
+	tk.MustQuery("select * from t use index(idx) where b in (15, 25, 35)").Sort().Check(testkit.Rows("15 15 0", "25 25 0", "35 35 0"))
+	tk.MustQuery("explain select * from t use index(idx) where b in (15, 25, 35)").Check(
+		testkit.Rows("Batch_Point_Get_1 3.00 root table:t, index:idx(b) keep order:false, desc:false"))
+
+	tk.MustQuery("select * from t use index(idx) where b in (select b from t use index(idx) where b>10)").Sort().Check(testkit.Rows("15 15 0", "25 25 0", "35 35 0"))
+
+	tk.MustQuery("select * from t use index(idx) where b = 15").Check(testkit.Rows("15 15 0"))
+	tk.MustQuery("explain select * from t use index(idx) where b = 15").Check(
+		testkit.Rows("Point_Get_1 1.00 root table:t, index:idx(b) "))
+
+	tk.MustQuery("select * from t use index(idx) where b in (select b from t use index(idx) where b > 10)").Sort().Check(
+		testkit.Rows("15 15 0", "25 25 0", "35 35 0"))
+
+	tk.MustQuery("select * from t partition(p1) use index(idx) where b = 3").Check(testkit.Rows())
+	tk.MustQuery("select * from t partition(p1) use index(idx) where b in (15, 25, 35)").Check(testkit.Rows("15 15 0"))
+}

pkg/planner/core/initialize.go

@@ -556,6 +556,10 @@ func (p *BatchPointGetPlan) Init(ctx sessionctx.Context, stats *property.StatsIn
 		d      types.Datum
 	)
 
+	if p.PartitionColPos == GlobalWithoutColumnPos {
+		return p
+	}
+
 	if p.PartitionExpr != nil {
 		if len(p.Handles) > 0 {
 			for _, handle := range p.Handles {

pkg/planner/core/point_get_plan.go

@@ -59,6 +59,9 @@ import (
 	"go.uber.org/zap"
 )
 
+// GlobalWithoutColumnPos marks the index has no partition column.
+const GlobalWithoutColumnPos = -1
+
 // PointGetPlan is a fast plan for simple point get.
 // When we detect that the statement has a unique equal access condition, this plan is used.
 // This plan is much faster to build and to execute because it avoid the optimization and coprocessor cost.
@@ -1069,7 +1072,7 @@ func tryPointGetPlan(ctx sessionctx.Context, selStmt *ast.SelectStmt, check bool
 			return p
 		}
 		if partitionInfo == nil {
-			return nil
+			return checkTblIndexForPointPlan(ctx, tblName, schema, names, pairs, nil, pos, true, isTableDual, check)
 		}
 		// Take partition selection into consideration.
 		if len(tblName.PartitionNames) > 0 {
@@ -1100,16 +1103,45 @@ func tryPointGetPlan(ctx sessionctx.Context, selStmt *ast.SelectStmt, check bool
 		return nil
 	}
 
+	return checkTblIndexForPointPlan(ctx, tblName, schema, names, pairs, partitionInfo, pos, false, isTableDual, check)
+}
+
+func checkTblIndexForPointPlan(ctx sessionctx.Context, tblName *ast.TableName, schema *expression.Schema,
+	names []*types.FieldName, pairs []nameValuePair, partitionInfo *model.PartitionDefinition,
+	pos int, globalIndexCheck, isTableDual, check bool) *PointGetPlan {
+	if globalIndexCheck {
+		// when partitions are specified or some partition is in ddl, not use point get plan for global index.
+		// TODO: Add partition ID filter for Global Index Point Get.
+		// partitions are specified in select stmt.
+		if len(tblName.PartitionNames) > 0 {
+			return nil
+		}
+		tbl := tblName.TableInfo
+		// some partition is in ddl.
+		if tbl == nil ||
+			len(tbl.GetPartitionInfo().AddingDefinitions) > 0 ||
+			len(tbl.GetPartitionInfo().DroppingDefinitions) > 0 {
+			return nil
+		}
+	}
 	check = check || ctx.GetSessionVars().IsIsolation(ast.ReadCommitted)
 	check = check && ctx.GetSessionVars().ConnectionID > 0
 	var latestIndexes map[int64]*model.IndexInfo
 	var err error
 
+	tbl := tblName.TableInfo
+	dbName := tblName.Schema.L
+	if dbName == "" {
+		dbName = ctx.GetSessionVars().CurrentDB
+	}
 	for _, idxInfo := range tbl.Indices {
 		if !idxInfo.Unique || idxInfo.State != model.StatePublic || idxInfo.Invisible || idxInfo.MVIndex ||
 			!indexIsAvailableByHints(idxInfo, tblName.IndexHints) {
 			continue
 		}
+		if globalIndexCheck && !idxInfo.Global {
+			continue
+		}
 		if isTableDual {
 			if check && latestIndexes == nil {
 				latestIndexes, check, err = getLatestIndexInfo(ctx, tbl.ID, 0)
@@ -1998,6 +2030,9 @@ func getColumnPosInIndex(idx *model.IndexInfo, colName *model.CIStr) int {
 			return i
 		}
 	}
+	if idx.Global {
+		return GlobalWithoutColumnPos
+	}
 	panic("unique index must include all partition columns")
 }
 

tests/integrationtest/r/executor/partition/global_index.result

@@ -91,22 +91,14 @@ alter table test_t1 add unique  p_a (a);
 insert into test_t1 values (1,1,1);
 explain format='brief' select * from test_t1 where a = 1;
 id	estRows	task	access object	operator info
-PartitionUnion	2.00	root		
-├─TableReader	1.00	root		data:Selection
-│ └─Selection	1.00	cop[tikv]		eq(executor__partition__global_index.test_t1.a, 1)
-│   └─TableFullScan	10000.00	cop[tikv]	table:test_t1, partition:p0	keep order:false, stats:pseudo
-└─TableReader	1.00	root		data:Selection
-  └─Selection	1.00	cop[tikv]		eq(executor__partition__global_index.test_t1.a, 1)
-    └─TableFullScan	10000.00	cop[tikv]	table:test_t1, partition:p1	keep order:false, stats:pseudo
+Point_Get	1.00	root	table:test_t1, index:p_a(a)	
 select * from test_t1 where a = 1;
 a	b	c
 1	1	1
 analyze table test_t1;
 explain format='brief' select * from test_t1 where a = 1;
 id	estRows	task	access object	operator info
-IndexLookUp	1.00	root	partition:all	
-├─IndexRangeScan(Build)	1.00	cop[tikv]	table:test_t1, index:p_a(a)	range:[1,1], keep order:false
-└─TableRowIDScan(Probe)	1.00	cop[tikv]	table:test_t1	keep order:false
+Point_Get	1.00	root	table:test_t1, index:p_a(a)	
 select * from test_t1 where a = 1;
 a	b	c
 1	1	1