distsql: support global index in memIndexLookUpReader (#52240)

close #52132

pkg/distsql/BUILD.bazel

@@ -67,7 +67,7 @@ go_test(
     embed = [":distsql"],
     flaky = True,
     race = "on",
-    shard_count = 26,
+    shard_count = 27,
     deps = [
         "//pkg/distsql/context",
         "//pkg/domain/resourcegroup",

pkg/distsql/request_builder.go

@@ -560,7 +560,15 @@ func TableHandlesToKVRanges(tid int64, handles []kv.Handle) ([]kv.KeyRange, []in
 	hints := make([]int, 0, len(handles))
 	i := 0
 	for i < len(handles) {
-		if commonHandle, ok := handles[i].(*kv.CommonHandle); ok {
+		var isCommonHandle bool
+		var commonHandle *kv.CommonHandle
+		if partitionHandle, ok := handles[i].(kv.PartitionHandle); ok {
+			tid = partitionHandle.PartitionID
+			commonHandle, isCommonHandle = partitionHandle.Handle.(*kv.CommonHandle)
+		} else {
+			commonHandle, isCommonHandle = handles[i].(*kv.CommonHandle)
+		}
+		if isCommonHandle {
 			ran := kv.KeyRange{
 				StartKey: tablecodec.EncodeRowKey(tid, commonHandle.Encoded()),
 				EndKey:   tablecodec.EncodeRowKey(tid, kv.Key(commonHandle.Encoded()).Next()),
@@ -572,6 +580,9 @@ func TableHandlesToKVRanges(tid int64, handles []kv.Handle) ([]kv.KeyRange, []in
 		}
 		j := i + 1
 		for ; j < len(handles) && handles[j-1].IntValue() != math.MaxInt64; j++ {
+			if p, ok := handles[j].(kv.PartitionHandle); ok && p.PartitionID != tid {
+				break
+			}
 			if handles[j].IntValue() != handles[j-1].IntValue()+1 {
 				break
 			}

pkg/distsql/request_builder_test.go

@@ -72,6 +72,50 @@ func TestTableHandlesToKVRanges(t *testing.T) {
 	}
 }
 
+func TestTablePartitionHandlesToKVRanges(t *testing.T) {
+	handles := []kv.Handle{
+		// Partition handles in different partitions
+		kv.NewPartitionHandle(1, kv.IntHandle(0)),
+		kv.NewPartitionHandle(2, kv.IntHandle(2)),
+		kv.NewPartitionHandle(2, kv.IntHandle(3)),
+		kv.NewPartitionHandle(2, kv.IntHandle(4)),
+		kv.NewPartitionHandle(3, kv.IntHandle(5)),
+		kv.NewPartitionHandle(1, kv.IntHandle(10)),
+		kv.NewPartitionHandle(2, kv.IntHandle(11)),
+		kv.NewPartitionHandle(3, kv.IntHandle(100)),
+		kv.NewPartitionHandle(1, kv.IntHandle(9223372036854775806)),
+		kv.NewPartitionHandle(1, kv.IntHandle(9223372036854775807)),
+	}
+
+	// Build expected key ranges.
+	hrs := make([]*handleRange, 0, len(handles))
+	hrs = append(hrs, &handleRange{start: 0, end: 0})
+	hrs = append(hrs, &handleRange{start: 2, end: 4})
+	hrs = append(hrs, &handleRange{start: 5, end: 5})
+	hrs = append(hrs, &handleRange{start: 10, end: 10})
+	hrs = append(hrs, &handleRange{start: 11, end: 11})
+	hrs = append(hrs, &handleRange{start: 100, end: 100})
+	hrs = append(hrs, &handleRange{start: 9223372036854775806, end: 9223372036854775807})
+
+	expect := append(getExpectedRanges(1, hrs[:1]), getExpectedRanges(2, hrs[1:2])...)
+	expect = append(expect, getExpectedRanges(3, hrs[2:3])...)
+	expect = append(expect, getExpectedRanges(1, hrs[3:4])...)
+	expect = append(expect, getExpectedRanges(2, hrs[4:5])...)
+	expect = append(expect, getExpectedRanges(3, hrs[5:6])...)
+	expect = append(expect, getExpectedRanges(1, hrs[6:])...)
+
+	// Build actual key ranges.
+	actual, hints := TableHandlesToKVRanges(0, handles)
+
+	// Compare key ranges and expected key ranges.
+	require.Equal(t, len(expect), len(actual))
+	require.Equal(t, hints, []int{1, 3, 1, 1, 1, 1, 2})
+	for i := range actual {
+		require.Equal(t, expect[i].StartKey, actual[i].StartKey)
+		require.Equal(t, expect[i].EndKey, actual[i].EndKey)
+	}
+}
+
 func TestTableRangesToKVRanges(t *testing.T) {
 	ranges := []*ranger.Range{
 		{

pkg/executor/distsql.go

@@ -437,7 +437,7 @@ type IndexLookUpExecutor struct {
 	avgRowSize     float64
 
 	// fields about accessing partition tables
-	partitionTableMode bool                  // if this executor is accessing a partition table
+	partitionTableMode bool                  // if this executor is accessing a local index with partition table
 	prunedPartitions   []table.PhysicalTable // partition tables need to access
 	partitionIDMap     map[int64]struct{}    // partitionIDs that global index access
 	partitionRangeMap  map[int64][]*ranger.Range

pkg/executor/mem_reader.go

@@ -656,7 +656,7 @@ type memIndexLookUpReader struct {
 	idxReader *memIndexReader
 
 	// partition mode
-	partitionMode     bool                  // if it is accessing a partition table
+	partitionMode     bool                  // if this executor is accessing a local index with partition table
 	partitionTables   []table.PhysicalTable // partition tables to access
 	partitionKVRanges [][]kv.KeyRange       // kv ranges for these partition tables
 
@@ -742,9 +742,7 @@ func (m *memIndexLookUpReader) getMemRows(ctx context.Context) ([][]types.Datum,
 	}
 
 	if m.desc {
-		for i, j := 0, len(tblKVRanges)-1; i < j; i, j = i+1, j-1 {
-			tblKVRanges[i], tblKVRanges[j] = tblKVRanges[j], tblKVRanges[i]
-		}
+		slices.Reverse(tblKVRanges)
 	}
 
 	colIDs, pkColIDs, rd := getColIDAndPkColIDs(m.ctx, m.table, m.columns)

pkg/tablecodec/tablecodec.go

@@ -975,20 +975,36 @@ func decodeHandleInIndexKey(keySuffix []byte) (kv.Handle, error) {
 	return kv.NewCommonHandle(keySuffix)
 }
 
-func decodeHandleInIndexValue(value []byte) (kv.Handle, error) {
-	if getIndexVersion(value) == 1 {
-		seg := SplitIndexValueForClusteredIndexVersion1(value)
-		return kv.NewCommonHandle(seg.CommonHandle)
+func decodeHandleInIndexValue(value []byte) (handle kv.Handle, err error) {
+	var seg IndexValueSegments
+	if getIndexVersion(value) == 0 {
+		// For Old Encoding (IntHandle without any others options)
+		if len(value) <= MaxOldEncodeValueLen {
+			return decodeIntHandleInIndexValue(value), nil
+		}
+		// For IndexValueVersion0
+		seg = SplitIndexValue(value)
+	} else {
+		// For IndexValueForClusteredIndexVersion1
+		seg = SplitIndexValueForClusteredIndexVersion1(value)
+	}
+	if len(seg.IntHandle) != 0 {
+		handle = decodeIntHandleInIndexValue(seg.IntHandle)
 	}
-	if len(value) > MaxOldEncodeValueLen {
-		tailLen := value[0]
-		if tailLen >= 8 {
-			return decodeIntHandleInIndexValue(value[len(value)-int(tailLen):]), nil
+	if len(seg.CommonHandle) != 0 {
+		handle, err = kv.NewCommonHandle(seg.CommonHandle)
+		if err != nil {
+			return nil, err
+		}
+	}
+	if len(seg.PartitionID) != 0 {
+		_, pid, err := codec.DecodeInt(seg.PartitionID)
+		if err != nil {
+			return nil, err
 		}
-		handleLen := uint16(value[2])<<8 + uint16(value[3])
-		return kv.NewCommonHandle(value[4 : 4+handleLen])
+		handle = kv.NewPartitionHandle(pid, handle)
 	}
-	return decodeIntHandleInIndexValue(value), nil
+	return handle, nil
 }
 
 // decodeIntHandleInIndexValue uses to decode index value as int handle id.

tests/integrationtest/r/globalindex/mem_index_lookup.result

@@ -0,0 +1,47 @@
+set tidb_enable_global_index=true;
+# IntHandle
+drop table if exists t;
+CREATE TABLE `t` (
+`a` int(11) DEFAULT NULL,
+`b` int(11) DEFAULT NULL,
+UNIQUE KEY `idx1` (`b`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
+PARTITION BY HASH (`a`) PARTITIONS 5;
+begin;
+insert into t values (1, 2), (2, 3), (3, 4), (4, 5), (5, 1);
+explain select * from t use index(idx1) where b > 2;
+id	estRows	task	access object	operator info
+Projection_5	3333.33	root		globalindex__mem_index_lookup.t.a, globalindex__mem_index_lookup.t.b
+└─UnionScan_6	3333.33	root		gt(globalindex__mem_index_lookup.t.b, 2)
+  └─IndexLookUp_9	3333.33	root	partition:all	
+    ├─IndexRangeScan_7(Build)	3333.33	cop[tikv]	table:t, index:idx1(b)	range:(2,+inf], keep order:false, stats:pseudo
+    └─TableRowIDScan_8(Probe)	3333.33	cop[tikv]	table:t	keep order:false, stats:pseudo
+select * from t use index(idx1) where b > 2;
+a	b
+2	3
+3	4
+4	5
+rollback;
+# CommonHandle
+drop table if exists t;
+CREATE TABLE `t` (
+`a` year(4) primary key,
+`b` int(11) DEFAULT NULL,
+UNIQUE KEY `idx1` (`b`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
+PARTITION BY HASH (`a`) PARTITIONS 5;
+begin;
+insert into t values (2001, 2), (2002, 3), (2003, 4), (2004, 5), (2005, 1);
+explain select * from t use index(idx1) where b > 2;
+id	estRows	task	access object	operator info
+Projection_5	3333.33	root		globalindex__mem_index_lookup.t.a, globalindex__mem_index_lookup.t.b
+└─UnionScan_6	3333.33	root		gt(globalindex__mem_index_lookup.t.b, 2)
+  └─IndexLookUp_9	3333.33	root	partition:all	
+    ├─IndexRangeScan_7(Build)	3333.33	cop[tikv]	table:t, index:idx1(b)	range:(2,+inf], keep order:false, stats:pseudo
+    └─TableRowIDScan_8(Probe)	3333.33	cop[tikv]	table:t	keep order:false, stats:pseudo
+select * from t use index(idx1) where b > 2;
+a	b
+2002	3
+2003	4
+2004	5
+rollback;

tests/integrationtest/t/globalindex/mem_index_lookup.test

@@ -0,0 +1,32 @@
+set tidb_enable_global_index=true;
+
+--echo # IntHandle
+drop table if exists t;
+CREATE TABLE `t` (
+  `a` int(11) DEFAULT NULL,
+  `b` int(11) DEFAULT NULL,
+  UNIQUE KEY `idx1` (`b`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
+PARTITION BY HASH (`a`) PARTITIONS 5;
+
+begin;
+insert into t values (1, 2), (2, 3), (3, 4), (4, 5), (5, 1);
+explain select * from t use index(idx1) where b > 2;
+select * from t use index(idx1) where b > 2;
+rollback;
+
+--echo # CommonHandle
+drop table if exists t;
+CREATE TABLE `t` (
+  `a` year(4) primary key,
+  `b` int(11) DEFAULT NULL,
+  UNIQUE KEY `idx1` (`b`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
+PARTITION BY HASH (`a`) PARTITIONS 5;
+
+begin;
+insert into t values (2001, 2), (2002, 3), (2003, 4), (2004, 5), (2005, 1);
+explain select * from t use index(idx1) where b > 2;
+select * from t use index(idx1) where b > 2;
+rollback;
+