stats: support dump stats for partition table (#7753)

cmd/importer/stats.go

@@ -24,7 +24,6 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/codec"
-	"github.com/pingcap/tidb/util/mock"
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 )
@@ -39,8 +38,7 @@ func loadStats(tblInfo *model.TableInfo, path string) (*stats.Table, error) {
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	handle := stats.NewHandle(mock.NewContext(), 0)
-	return handle.LoadStatsFromJSONToTable(tblInfo, jsTable)
+	return stats.TableStatsFromJSON(tblInfo, tblInfo.ID, jsTable)
 }
 
 type histogram struct {

statistics/dump.go

@@ -33,6 +33,7 @@ type JSONTable struct {
 	Indices      map[string]*jsonColumn `json:"indices"`
 	Count        int64                  `json:"count"`
 	ModifyCount  int64                  `json:"modify_count"`
+	Partitions   map[string]*JSONTable  `json:"partitions"`
 }
 
 type jsonColumn struct {
@@ -58,7 +59,30 @@ func dumpJSONCol(hist *Histogram, CMSketch *CMSketch) *jsonColumn {
 
 // DumpStatsToJSON dumps statistic to json.
 func (h *Handle) DumpStatsToJSON(dbName string, tableInfo *model.TableInfo) (*JSONTable, error) {
-	tbl, err := h.tableStatsFromStorage(tableInfo, tableInfo.ID, true)
+	pi := tableInfo.GetPartitionInfo()
+	if pi == nil {
+		return h.tableStatsToJSON(dbName, tableInfo, tableInfo.ID)
+	}
+	jsonTbl := &JSONTable{
+		DatabaseName: dbName,
+		TableName:    tableInfo.Name.L,
+		Partitions:   make(map[string]*JSONTable, len(pi.Definitions)),
+	}
+	for _, def := range pi.Definitions {
+		tbl, err := h.tableStatsToJSON(dbName, tableInfo, def.ID)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		if tbl == nil {
+			continue
+		}
+		jsonTbl.Partitions[def.Name.L] = tbl
+	}
+	return jsonTbl, nil
+}
+
+func (h *Handle) tableStatsToJSON(dbName string, tableInfo *model.TableInfo, physicalID int64) (*JSONTable, error) {
+	tbl, err := h.tableStatsFromStorage(tableInfo, physicalID, true)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -91,11 +115,37 @@ func (h *Handle) DumpStatsToJSON(dbName string, tableInfo *model.TableInfo) (*JS
 
 // LoadStatsFromJSON will load statistic from JSONTable, and save it to the storage.
 func (h *Handle) LoadStatsFromJSON(is infoschema.InfoSchema, jsonTbl *JSONTable) error {
-	tableInfo, err := is.TableByName(model.NewCIStr(jsonTbl.DatabaseName), model.NewCIStr(jsonTbl.TableName))
+	table, err := is.TableByName(model.NewCIStr(jsonTbl.DatabaseName), model.NewCIStr(jsonTbl.TableName))
 	if err != nil {
 		return errors.Trace(err)
 	}
-	tbl, err := h.LoadStatsFromJSONToTable(tableInfo.Meta(), jsonTbl)
+	tableInfo := table.Meta()
+	pi := tableInfo.GetPartitionInfo()
+	if pi == nil {
+		err := h.loadStatsFromJSON(tableInfo, tableInfo.ID, jsonTbl)
+		if err != nil {
+			return errors.Trace(err)
+		}
+	} else {
+		if jsonTbl.Partitions == nil {
+			return errors.New("No partition statistics")
+		}
+		for _, def := range pi.Definitions {
+			tbl := jsonTbl.Partitions[def.Name.L]
+			if tbl == nil {
+				continue
+			}
+			err := h.loadStatsFromJSON(tableInfo, def.ID, tbl)
+			if err != nil {
+				return errors.Trace(err)
+			}
+		}
+	}
+	return errors.Trace(h.Update(is))
+}
+
+func (h *Handle) loadStatsFromJSON(tableInfo *model.TableInfo, physicalID int64, jsonTbl *JSONTable) error {
+	tbl, err := TableStatsFromJSON(tableInfo, physicalID, jsonTbl)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -116,13 +166,13 @@ func (h *Handle) LoadStatsFromJSON(is infoschema.InfoSchema, jsonTbl *JSONTable)
 	if err != nil {
 		return errors.Trace(err)
 	}
-	return errors.Trace(h.Update(is))
+	return nil
 }
 
-// LoadStatsFromJSONToTable load statistic from JSONTable and return the Table of statistic.
-func (h *Handle) LoadStatsFromJSONToTable(tableInfo *model.TableInfo, jsonTbl *JSONTable) (*Table, error) {
+// TableStatsFromJSON loads statistic from JSONTable and return the Table of statistic.
+func TableStatsFromJSON(tableInfo *model.TableInfo, physicalID int64, jsonTbl *JSONTable) (*Table, error) {
 	newHistColl := HistColl{
-		PhysicalID:     tableInfo.ID,
+		PhysicalID:     physicalID,
 		HavePhysicalID: true,
 		Count:          jsonTbl.Count,
 		ModifyCount:    jsonTbl.ModifyCount,

statistics/dump_test.go

@@ -14,6 +14,8 @@
 package statistics_test
 
 import (
+	"fmt"
+
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/kv"
@@ -62,7 +64,7 @@ func (s *testDumpStatsSuite) TestConversion(c *C) {
 	c.Assert(err, IsNil)
 	jsonTbl, err := h.DumpStatsToJSON("test", tableInfo.Meta())
 	c.Assert(err, IsNil)
-	loadTbl, err := h.LoadStatsFromJSONToTable(tableInfo.Meta(), jsonTbl)
+	loadTbl, err := statistics.TableStatsFromJSON(tableInfo.Meta(), tableInfo.Meta().ID, jsonTbl)
 	c.Assert(err, IsNil)
 
 	tbl := h.GetTableStats(tableInfo.Meta())
@@ -73,3 +75,49 @@ func (s *testDumpStatsSuite) TestConversion(c *C) {
 	loadTblInStorage := h.GetTableStats(tableInfo.Meta())
 	assertTableEqual(c, loadTblInStorage, tbl)
 }
+
+func (s *testDumpStatsSuite) TestDumpPartitions(c *C) {
+	defer cleanEnv(c, s.store, s.do)
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("set @@session.tidb_enable_table_partition=1")
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	createTable := `CREATE TABLE t (a int, b int, primary key(a), index idx(b))
+PARTITION BY RANGE ( a ) (
+		PARTITION p0 VALUES LESS THAN (6),
+		PARTITION p1 VALUES LESS THAN (11),
+		PARTITION p2 VALUES LESS THAN (16),
+		PARTITION p3 VALUES LESS THAN (21)
+)`
+	tk.MustExec(createTable)
+	for i := 1; i < 21; i++ {
+		tk.MustExec(fmt.Sprintf(`insert into t values (%d, %d)`, i, i))
+	}
+	tk.MustExec("analyze table t")
+	is := s.do.InfoSchema()
+	h := s.do.StatsHandle()
+	h.Update(is)
+
+	table, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("t"))
+	c.Assert(err, IsNil)
+	tableInfo := table.Meta()
+	jsonTbl, err := h.DumpStatsToJSON("test", tableInfo)
+	c.Assert(err, IsNil)
+	pi := tableInfo.GetPartitionInfo()
+	originTables := make([]*statistics.Table, 0, len(pi.Definitions))
+	for _, def := range pi.Definitions {
+		originTables = append(originTables, h.GetPartitionStats(tableInfo, def.ID))
+	}
+
+	tk.MustExec("truncate table mysql.stats_meta")
+	tk.MustExec("truncate table mysql.stats_histograms")
+	tk.MustExec("truncate table mysql.stats_buckets")
+	h.Clear()
+
+	err = h.LoadStatsFromJSON(s.do.InfoSchema(), jsonTbl)
+	c.Assert(err, IsNil)
+	for i, def := range pi.Definitions {
+		t := h.GetPartitionStats(tableInfo, def.ID)
+		assertTableEqual(c, originTables[i], t)
+	}
+}

statistics/table.go

@@ -229,7 +229,7 @@ func (h *Handle) columnStatsFromStorage(row chunk.Row, table *Table, tableInfo *
 
 // tableStatsFromStorage loads table stats info from storage.
 func (h *Handle) tableStatsFromStorage(tableInfo *model.TableInfo, physicalID int64, loadAll bool) (*Table, error) {
-	table, ok := h.statsCache.Load().(statsCache)[tableInfo.ID]
+	table, ok := h.statsCache.Load().(statsCache)[physicalID]
 	// If table stats is pseudo, we also need to copy it, since we will use the column stats when
 	// the average error rate of it is small.
 	if !ok {