executor: update memIndexLookUp to make it can read partition table correctly and add more cases about dynamic-mode with UnionScan

executor/mem_reader.go

@@ -379,6 +379,11 @@ type memIndexLookUpReader struct {
 	retFieldTypes []*types.FieldType
 
 	idxReader *memIndexReader
+
+	// partition mode
+	partitionMode     bool                  // if it is accessing a partition table
+	partitionTables   []table.PhysicalTable // partition tables to access
+	partitionKVRanges [][]kv.KeyRange       // kv ranges for these partition tables
 }
 
 func buildMemIndexLookUpReader(us *UnionScanExec, idxLookUpReader *IndexLookUpExecutor) *memIndexLookUpReader {
@@ -404,16 +409,43 @@ func buildMemIndexLookUpReader(us *UnionScanExec, idxLookUpReader *IndexLookUpEx
 		conditions:    us.conditions,
 		retFieldTypes: retTypes(us),
 		idxReader:     memIdxReader,
+
+		partitionMode:     idxLookUpReader.partitionTableMode,
+		partitionKVRanges: idxLookUpReader.partitionKVRanges,
+		partitionTables:   idxLookUpReader.prunedPartitions,
 	}
 }
 
 func (m *memIndexLookUpReader) getMemRows() ([][]types.Datum, error) {
-	handles, err := m.idxReader.getMemRowsHandle()
-	if err != nil || len(handles) == 0 {
-		return nil, err
+	kvRanges := [][]kv.KeyRange{m.idxReader.kvRanges}
+	tbls := []table.Table{m.table}
+	if m.partitionMode {
+		m.idxReader.desc = false // keep-order if always false for IndexLookUp reading partitions so this parameter makes no sense
+		kvRanges = m.partitionKVRanges
+		tbls = tbls[:0]
+		for _, p := range m.partitionTables {
+			tbls = append(tbls, p)
+		}
+	}
+
+	tblKVRanges := make([]kv.KeyRange, 0, 16)
+	numHandles := 0
+	for i, tbl := range tbls {
+		m.idxReader.kvRanges = kvRanges[i]
+		handles, err := m.idxReader.getMemRowsHandle()
+		if err != nil {
+			return nil, err
+		}
+		if len(handles) == 0 {
+			continue
+		}
+		numHandles += len(handles)
+		tblKVRanges = append(tblKVRanges, distsql.TableHandlesToKVRanges(getPhysicalTableID(tbl), handles)...)
+	}
+	if numHandles == 0 {
+		return nil, nil
 	}
 
-	tblKVRanges := distsql.TableHandlesToKVRanges(getPhysicalTableID(m.table), handles)
 	colIDs := make(map[int64]int, len(m.columns))
 	for i, col := range m.columns {
 		colIDs[col.ID] = i
@@ -440,7 +472,7 @@ func (m *memIndexLookUpReader) getMemRows() ([][]types.Datum, error) {
 		columns:       m.columns,
 		kvRanges:      tblKVRanges,
 		conditions:    m.conditions,
-		addedRows:     make([][]types.Datum, 0, len(handles)),
+		addedRows:     make([][]types.Datum, 0, numHandles),
 		retFieldTypes: m.retFieldTypes,
 		colIDs:        colIDs,
 		pkColIDs:      pkColIDs,

executor/partition_table_test.go

@@ -1737,6 +1737,81 @@ func (s *partitionTableSuite) TestParallelApply(c *C) {
 	}
 }
 
+func (s *partitionTableSuite) TestDirectReadingWithUnionScan(c *C) {
+	if israce.RaceEnabled {
+		c.Skip("exhaustive types test, skip race test")
+	}
+
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("create database test_unionscan")
+	defer tk.MustExec(`drop database test_unionscan`)
+	tk.MustExec("use test_unionscan")
+	tk.MustExec("set @@tidb_partition_prune_mode = 'dynamic'")
+
+	tk.MustExec(`create table trange(a int, b int, index idx_a(a)) partition by range(a) (
+		partition p0 values less than (10),
+		partition p1 values less than (30),
+		partition p2 values less than (50))`)
+	tk.MustExec(`create table thash(a int, b int, index idx_a(a)) partition by hash(a) partitions 4`)
+	tk.MustExec(`create table tnormal(a int, b int, index idx_a(a))`)
+
+	vals := make([]string, 0, 1000)
+	for i := 0; i < 1000; i++ {
+		vals = append(vals, fmt.Sprintf("(%v, %v)", rand.Intn(50), rand.Intn(50)))
+	}
+	for _, tb := range []string{`trange`, `tnormal`, `thash`} {
+		sql := fmt.Sprintf(`insert into %v values `+strings.Join(vals, ", "), tb)
+		tk.MustExec(sql)
+	}
+
+	randCond := func(col string) string {
+		la, ra := rand.Intn(50), rand.Intn(50)
+		if la > ra {
+			la, ra = ra, la
+		}
+		return fmt.Sprintf(`%v>=%v and %v<=%v`, col, la, col, ra)
+	}
+
+	tk.MustExec(`begin`)
+	for i := 0; i < 1000; i++ {
+		if i == 0 || rand.Intn(2) == 0 { // insert some inflight rows
+			val := fmt.Sprintf("(%v, %v)", rand.Intn(50), rand.Intn(50))
+			for _, tb := range []string{`trange`, `tnormal`, `thash`} {
+				sql := fmt.Sprintf(`insert into %v values `+val, tb)
+				tk.MustExec(sql)
+			}
+		} else {
+			var sql string
+			switch rand.Intn(3) {
+			case 0: // table scan
+				sql = `select * from %v ignore index(idx_a) where ` + randCond(`b`)
+			case 1: // index reader
+				sql = `select a from %v use index(idx_a) where ` + randCond(`a`)
+			case 2: // index lookup
+				sql = `select * from %v use index(idx_a) where ` + randCond(`a`) + ` and ` + randCond(`b`)
+			}
+			switch rand.Intn(2) {
+			case 0: // order by a
+				sql += ` order by a`
+			case 1: // order by b
+				sql += ` order by b`
+			}
+
+			var result [][]interface{}
+			for _, tb := range []string{`trange`, `tnormal`, `thash`} {
+				q := fmt.Sprintf(sql, tb)
+				tk.HasPlan(q, `UnionScan`)
+				if result == nil {
+					result = tk.MustQuery(q).Sort().Rows()
+				} else {
+					tk.MustQuery(q).Sort().Check(result)
+				}
+			}
+		}
+	}
+	tk.MustExec(`rollback`)
+}
+
 func (s *partitionTableSuite) TestDirectReadingWithAgg(c *C) {
 	if israce.RaceEnabled {
 		c.Skip("exhaustive types test, skip race test")