diff --git a/planner/core/logical_plan_builder.go b/planner/core/logical_plan_builder.go index c3b1239d4ffac..6d4e4238dd362 100644 --- 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 index 923d025ac7d13..57b53a75cc080 100644 --- 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 index 311b9ae4a8838..b63d042c03a7e 100644 --- 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)) + } + } +}