executor: batch point get support clustered index partition table. (p…

…ingcap#18829) Co-authored-by: ti-srebot <[email protected]>
m4ver1k · Jul 28, 2020 · 1e7454c · 1e7454c
1 parent 5973d14
commit 1e7454c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/executor/batch_point_get.go b/executor/batch_point_get.go
@@ -26,6 +26,7 @@ import (
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/math"
 	"github.com/pingcap/tidb/util/rowcodec"
@@ -152,7 +153,7 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
 		dedup := make(map[hack.MutableString]struct{})
 		keys := make([]kv.Key, 0, len(e.idxVals))
 		for _, idxVals := range e.idxVals {
-			physID := getPhysID(e.tblInfo, kv.IntHandle(idxVals[e.partPos].GetInt64()))
+			physID := getPhysID(e.tblInfo, idxVals[e.partPos].GetInt64())
 			idxKey, err1 := EncodeUniqueIndexKey(e.ctx, e.tblInfo, e.idxInfo, idxVals, physID)
 			if err1 != nil && !kv.ErrNotExist.Equal(err1) {
 				return err1
@@ -227,7 +228,15 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
 		if len(e.physIDs) > 0 {
 			tID = e.physIDs[i]
 		} else {
-			tID = getPhysID(e.tblInfo, handle)
+			if handle.IsInt() {
+				tID = getPhysID(e.tblInfo, handle.IntValue())
+			} else {
+				_, d, err1 := codec.DecodeOne(handle.EncodedCol(e.partPos))
+				if err1 != nil {
+					return err1
+				}
+				tID = getPhysID(e.tblInfo, d.GetInt64())
+			}
 		}
 		key := tablecodec.EncodeRowKeyWithHandle(tID, handle)
 		keys[i] = key
@@ -327,11 +336,11 @@ func (getter *PessimisticLockCacheGetter) Get(_ context.Context, key kv.Key) ([]
 	return nil, kv.ErrNotExist
 }
 
-func getPhysID(tblInfo *model.TableInfo, val kv.Handle) int64 {
+func getPhysID(tblInfo *model.TableInfo, intVal int64) int64 {
 	pi := tblInfo.Partition
 	if pi == nil {
 		return tblInfo.ID
 	}
-	partIdx := math.Abs(val.IntValue() % int64(pi.Num)) // TODO: fix me for table, partition on cluster index.
+	partIdx := math.Abs(intVal % int64(pi.Num))
 	return pi.Definitions[partIdx].ID
 }
diff --git a/executor/clustered_index_test.go b/executor/clustered_index_test.go
@@ -103,3 +103,12 @@ func (s *testClusteredSuite) TestClusteredHint(c *C) {
 	tk.MustExec("create table ht (a varchar(64) primary key, b int)")
 	tk.MustQuery("select * from ht use index (`PRIMARY`)")
 }
+
+func (s *testClusteredSuite) TestClusteredBatchPointGet(c *C) {
+	tk := s.newTK(c)
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("CREATE TABLE t (a int,b int,c int, PRIMARY KEY (a,b)) PARTITION BY HASH(a) PARTITIONS 3")
+	tk.MustExec("insert t values (1, 1, 1), (3, 3, 3), (5, 5, 5)")
+	tk.MustQuery("select * from t where (a, b) in ((1, 1), (3, 3), (5, 5))").Check(
+		testkit.Rows("1 1 1", "3 3 3", "5 5 5"))
+}