planner: add extractor for tikv_region_peers

ref: pingcap#28330 Signed-off-by: IcePigZDB <icepigzdb@gmail.com>
IcePigZDB · Dec 13, 2021 · 27f38dc · 27f38dc
1 parent 7555536
commit 27f38dc
Show file tree

Hide file tree

Showing 3 changed files with 178 additions and 0 deletions.
diff --git a/planner/core/logical_plan_builder.go b/planner/core/logical_plan_builder.go
@@ -4352,6 +4352,8 @@ func (b *PlanBuilder) buildMemTable(_ context.Context, dbName model.CIStr, table
 			p.Extractor = &TiFlashSystemTableExtractor{}
 		case infoschema.TableStatementsSummary, infoschema.TableStatementsSummaryHistory:
 			p.Extractor = &StatementsSummaryExtractor{}
+		case infoschema.TableTiKVRegionPeers:
+			p.Extractor = &TikvRegionPeersExtractor{}
 		}
 	}
 	return p, nil

diff --git a/planner/core/memtable_predicate_extractor.go b/planner/core/memtable_predicate_extractor.go
@@ -1415,3 +1415,56 @@ func (e *StatementsSummaryExtractor) explainInfo(p *PhysicalMemTable) string {
 	}
 	return fmt.Sprintf("digests: [%s]", extractStringFromStringSet(e.Digests))
 }
+
+// TikvRegionPeersExtractor is used to extract some predicates of cluster table.
+type TikvRegionPeersExtractor struct {
+	extractHelper
+
+	// SkipRequest means the where clause always false, we don't need to request any component
+	SkipRequest bool
+
+	// RegionIDs/StoreIDs represents all region/store ids we should filter in PD to reduce network IO.
+	// e.g:
+	// 1. SELECT * FROM tikv_region_peers WHERE region_id=1
+	// 2. SELECT * FROM tikv_region_peers WHERE table_id in (11, 22)
+	RegionIDs []uint64
+	StoreIDs  []uint64
+}
+
+// Extract implements the MemTablePredicateExtractor Extract interface
+func (e *TikvRegionPeersExtractor) Extract(_ sessionctx.Context,
+	schema *expression.Schema,
+	names []*types.FieldName,
+	predicates []expression.Expression,
+) []expression.Expression {
+	// Extract the `region_id/store_id` columns.
+	remained, regionIDSkipRequest, regionIDs := e.extractCol(schema, names, predicates, "region_id", false)
+	remained, storeIDSkipRequest, storeIDs := e.extractCol(schema, names, remained, "store_id", false)
+	e.RegionIDs, e.StoreIDs = e.parseUint64(regionIDs), e.parseUint64(storeIDs)
+
+	e.SkipRequest = regionIDSkipRequest || storeIDSkipRequest
+	if e.SkipRequest {
+		return nil
+	}
+
+	return remained
+}
+
+func (e *TikvRegionPeersExtractor) explainInfo(p *PhysicalMemTable) string {
+	if e.SkipRequest {
+		return "skip_request:true"
+	}
+	r := new(bytes.Buffer)
+	if len(e.RegionIDs) > 0 {
+		r.WriteString(fmt.Sprintf("region_ids:[%s], ", extractStringFromUint64Slice(e.RegionIDs)))
+	}
+	if len(e.StoreIDs) > 0 {
+		r.WriteString(fmt.Sprintf("store_ids:[%s], ", extractStringFromUint64Slice(e.StoreIDs)))
+	}
+	// remove the last ", " in the message info
+	s := r.String()
+	if len(s) > 2 {
+		return s[:len(s)-2]
+	}
+	return s
+}
diff --git a/planner/core/memtable_predicate_extractor_test.go b/planner/core/memtable_predicate_extractor_test.go
@@ -1420,3 +1420,126 @@ func (s *extractorSuite) TestTiDBHotRegionsHistoryTableExtractor(c *C) {
 		}
 	}
 }
+
+func (s *extractorSuite) TestTikvRegionPeersExtractor(c *C) {
+	se, err := session.CreateSession4Test(s.store)
+	c.Assert(err, IsNil)
+
+	var cases = []struct {
+		sql                 string
+		regionIDs, storeIDs []uint64
+		skipRequest         bool
+	}{
+		// Test `region_id`, `store_id` columns.
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id=100",
+			regionIDs: []uint64{100},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where 100=region_id",
+			regionIDs: []uint64{100},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where 100=region_id or region_id=101",
+			regionIDs: []uint64{100, 101},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where 100=region_id or region_id=101 or region_id=102 or 103 = region_id",
+			regionIDs: []uint64{100, 101, 102, 103},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where (region_id=100 or region_id=101) and (store_id=200 or store_id=201)",
+			regionIDs: []uint64{100, 101},
+			storeIDs:  []uint64{200, 201},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id in (100, 101)",
+			regionIDs: []uint64{100, 101},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id in (100, 101) and store_id=200",
+			regionIDs: []uint64{100, 101},
+			storeIDs:  []uint64{200},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id in (100, 101) and store_id in (200, 201)",
+			regionIDs: []uint64{100, 101},
+			storeIDs:  []uint64{200, 201},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id=100 and store_id in (200, 201)",
+			regionIDs: []uint64{100},
+			storeIDs:  []uint64{200, 201},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id=100 and store_id=200",
+			regionIDs: []uint64{100},
+			storeIDs:  []uint64{200},
+		},
+		{
+			sql:         "select * from information_schema.tikv_region_peers where region_id=100 and region_id=101",
+			skipRequest: true,
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id=100 and region_id in (100,101)",
+			regionIDs: []uint64{100},
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id=100 and region_id in (100,101) and store_id=200 and store_id in (200,201)",
+			regionIDs: []uint64{100},
+			storeIDs:  []uint64{200},
+		},
+		{
+			sql:         "select * from information_schema.tikv_region_peers where region_id=100 and region_id in (101,102)",
+			skipRequest: true,
+		},
+		{
+			sql:         "select * from information_schema.tikv_region_peers where region_id=100 and region_id in (101,102) and store_id=200 and store_id in (200,201)",
+			skipRequest: true,
+		},
+		{
+			sql:         "select * from information_schema.tikv_region_peers where region_id=100 and region_id in (100,101) and store_id=200 and store_id in (201,202)",
+			skipRequest: true,
+		},
+		{
+			sql: `select * from information_schema.tikv_region_peers 
+								where region_id=100 and region_id in (100,101) 
+								and store_id=200 and store_id in (201,202)`,
+			skipRequest: true,
+		},
+		{
+			sql:       "select * from information_schema.tikv_region_peers where region_id in (100,101) and region_id in (101,102)",
+			regionIDs: []uint64{101},
+		},
+		{
+			sql: `select * from information_schema.tikv_region_peers 
+							where region_id in (100,101) 
+							and region_id in (101,102) 
+							and store_id in (200,201) 
+							and store_id in (201,202)`,
+			regionIDs: []uint64{101},
+			storeIDs:  []uint64{201},
+		},
+		{
+			sql: `select * from information_schema.tikv_region_peers 
+							where region_id in (100,101) 
+							and region_id in (100,102) 
+							and region_id in (102,103) 
+							and region_id in (103,104)`,
+			skipRequest: true,
+		},
+	}
+	parser := parser.New()
+	for _, ca := range cases {
+		logicalMemTable := s.getLogicalMemTable(c, se, parser, ca.sql)
+		c.Assert(logicalMemTable.Extractor, NotNil)
+
+		tikvRegionPeersExtractor := logicalMemTable.Extractor.(*plannercore.TikvRegionPeersExtractor)
+		if len(ca.regionIDs) > 0 {
+			c.Assert(tikvRegionPeersExtractor.RegionIDs, DeepEquals, ca.regionIDs, Commentf("SQL: %v", ca.sql))
+		}
+		if len(ca.storeIDs) > 0 {
+			c.Assert(tikvRegionPeersExtractor.StoreIDs, DeepEquals, ca.storeIDs, Commentf("SQL: %v", ca.sql))
+		}
+	}
+}