From ffc1bcaaa7fc49e8ef173cd8d2dd0c2a1a5ebc4a Mon Sep 17 00:00:00 2001
From: Evan Zhou <coocood@gmail.com>
Date: Thu, 18 Jun 2020 21:11:31 +0800
Subject: [PATCH] unistore: fix clustered index point range check. (#18106)

---
 executor/point_get_test.go                    |  2 ++
 .../unistore/cophandler/closure_exec.go       | 23 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/executor/point_get_test.go b/executor/point_get_test.go
index e842c2734c..ed465155ca 100644
--- a/executor/point_get_test.go
+++ b/executor/point_get_test.go
@@ -550,8 +550,10 @@ func (s *testPointGetSuite) TestClusterIndexPointGet(c *C) {
 
 	tk.MustExec(`drop table if exists snp`)
 	tk.MustExec(`create table snp(id1 int, id2 int, v int, primary key(id1, id2))`)
+	tk.MustExec(`insert snp values (1, 1, 1), (2, 2, 2), (2, 3, 3)`)
 	tk.MustQuery(`explain select * from snp where id1 = 1`).Check(testkit.Rows("TableReader_6 10.00 root  data:TableRangeScan_5",
 		"└─TableRangeScan_5 10.00 cop[tikv] table:snp range:[1,1], keep order:false, stats:pseudo"))
 	tk.MustQuery(`explain select * from snp where id1 in (1, 100)`).Check(testkit.Rows("TableReader_6 20.00 root  data:TableRangeScan_5",
 		"└─TableRangeScan_5 20.00 cop[tikv] table:snp range:[1,1], [100,100], keep order:false, stats:pseudo"))
+	tk.MustQuery("select * from snp where id1 = 2").Check(testkit.Rows("2 2 2", "2 3 3"))
 }
diff --git a/store/mockstore/unistore/cophandler/closure_exec.go b/store/mockstore/unistore/cophandler/closure_exec.go
index 350c1eb0d8..cebe8eb3f0 100644
--- a/store/mockstore/unistore/cophandler/closure_exec.go
+++ b/store/mockstore/unistore/cophandler/closure_exec.go
@@ -351,7 +351,7 @@ func (e *closureExecutor) execute() ([]tipb.Chunk, error) {
 	}
 	dbReader := e.dbReader
 	for i, ran := range e.kvRanges {
-		if e.unique && ran.IsPoint() {
+		if e.isPointGetRange(ran) {
 			val, err := dbReader.Get(ran.StartKey, e.startTS)
 			if err != nil {
 				return nil, errors.Trace(err)
@@ -389,6 +389,27 @@ func (e *closureExecutor) execute() ([]tipb.Chunk, error) {
 	return e.oldChunks, err
 }
 
+func (e *closureExecutor) isPointGetRange(ran kv.KeyRange) bool {
+	if e.idxScanCtx != nil || len(e.primaryCols) == 0 {
+		return e.unique && ran.IsPoint()
+	}
+	// For common handle table scan, we also need to check if the column count of the start key equals.
+	if len(ran.StartKey) < tablecodec.RecordRowKeyLen {
+		return false
+	}
+	keyColValues := tablecodec.CutRowKeyPrefix(ran.StartKey)
+	var colCnt int
+	for len(keyColValues) > 0 {
+		var err error
+		_, keyColValues, err = codec.CutOne(keyColValues)
+		if err != nil {
+			return false
+		}
+		colCnt++
+	}
+	return colCnt == len(e.primaryCols) && ran.IsPoint()
+}
+
 func (e *closureExecutor) checkRangeLock() error {
 	if !e.ignoreLock && !e.lockChecked {
 		for _, ran := range e.kvRanges {