*: make partitioned table use statistics (pingcap#7340)

plan/cbo_test.go

@@ -411,6 +411,14 @@ func (s *testAnalyzeSuite) TestAnalyze(c *C) {
 
 	testKit.MustExec("create table t3 (a int, b int)")
 	testKit.MustExec("create index a on t3 (a)")
+
+	testKit.MustExec("set @@session.tidb_enable_table_partition=1")
+	testKit.MustExec("create table t4 (a int, b int) partition by range (a) (partition p1 values less than (2), partition p2 values less than (3))")
+	testKit.MustExec("create index a on t4 (a)")
+	testKit.MustExec("create index b on t4 (b)")
+	testKit.MustExec("insert into t4 (a,b) values (1,1),(1,2),(1,3),(1,4),(2,5),(2,6),(2,7),(2,8)")
+	testKit.MustExec("analyze table t4")
+
 	tests := []struct {
 		sql  string
 		best string
@@ -451,6 +459,19 @@ func (s *testAnalyzeSuite) TestAnalyze(c *C) {
 			sql:  "analyze table t2 index",
 			best: "Analyze{Index(a),Index(b)}",
 		},
+		// Test partitioned table.
+		{
+			sql:  "select * from t4 where t4.a <= 2",
+			best: "UnionAll{TableReader(Table(t4)->Sel([le(test.t4.a, 2)]))->TableReader(Table(t4)->Sel([le(test.t4.a, 2)]))}",
+		},
+		{
+			sql:  "select * from t4 where t4.b < 2",
+			best: "UnionAll{IndexLookUp(Index(t4.b)[[-inf,2)], Table(t4))->IndexLookUp(Index(t4.b)[[-inf,2)], Table(t4))}",
+		},
+		{
+			sql:  "select * from t4 where t4.a = 1 and t4.b <= 2",
+			best: "TableReader(Table(t4)->Sel([eq(test.t4.a, 1) le(test.t4.b, 2)]))",
+		},
 		// TODO: Refine these tests in the future.
 		//{
 		//	sql:  "select * from t2 where t2.a = 1 and t2.b <= 2",

plan/logical_plan_builder.go

@@ -33,6 +33,7 @@ import (
 	"github.com/pingcap/tidb/mysql"
 	"github.com/pingcap/tidb/parser"
 	"github.com/pingcap/tidb/parser/opcode"
+	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/table"
@@ -1719,15 +1720,20 @@ func (ds *DataSource) newExtraHandleSchemaCol() *expression.Column {
 // 1. tidb-server started and statistics handle has not been initialized.
 // 2. table row count from statistics is zero.
 // 3. statistics is outdated.
-func (b *planBuilder) getStatsTable(tblInfo *model.TableInfo) *statistics.Table {
-	statsHandle := domain.GetDomain(b.ctx).StatsHandle()
+func getStatsTable(ctx sessionctx.Context, tblInfo *model.TableInfo, pid int64) *statistics.Table {
+	statsHandle := domain.GetDomain(ctx).StatsHandle()
 
 	// 1. tidb-server started and statistics handle has not been initialized.
 	if statsHandle == nil {
 		return statistics.PseudoTable(tblInfo)
 	}
 
-	statsTbl := statsHandle.GetTableStats(tblInfo)
+	var statsTbl *statistics.Table
+	if pid != tblInfo.ID {
+		statsTbl = statsHandle.GetPartitionStats(tblInfo, pid)
+	} else {
+		statsTbl = statsHandle.GetTableStats(tblInfo)
+	}
 
 	// 2. table row count from statistics is zero.
 	if statsTbl.Count == 0 {
@@ -1773,12 +1779,16 @@ func (b *planBuilder) buildDataSource(tn *ast.TableName) (LogicalPlan, error) {
 	} else {
 		columns = tbl.Cols()
 	}
+	var statisticTable *statistics.Table
+	if _, ok := tbl.(table.PartitionedTable); !ok {
+		statisticTable = getStatsTable(b.ctx, tbl.Meta(), tbl.Meta().ID)
+	}
 
 	ds := DataSource{
 		DBName:              dbName,
 		table:               tbl,
 		tableInfo:           tableInfo,
-		statisticTable:      b.getStatsTable(tableInfo),
+		statisticTable:      statisticTable,
 		indexHints:          tn.IndexHints,
 		possibleAccessPaths: possiblePaths,
 		Columns:             make([]*model.ColumnInfo, 0, len(columns)),

plan/rule_partition_processor.go

@@ -109,7 +109,7 @@ func (s *partitionProcessor) prune(ds *DataSource) (LogicalPlan, error) {
 		// id as FromID. So we set the id of the newDataSource with the original one to
 		// avoid traversing the whole plan tree to update the references.
 		newDataSource.id = ds.id
-
+		newDataSource.statisticTable = getStatsTable(ds.context(), ds.table.Meta(), pi.Definitions[i].ID)
 		children = append(children, &newDataSource)
 	}
 	if len(children) == 0 {

table/tables/partition.go

@@ -245,3 +245,7 @@ func (t *partitionedTable) UpdateRecord(ctx sessionctx.Context, h int64, currDat
 	tbl := t.GetPartition(to)
 	return tbl.UpdateRecord(ctx, h, currData, newData, touched)
 }
+
+func (t *partitionedTable) GetID() int64 {
+	panic("GetID() should never be called on PartitionedTable")
+}