statistics: Update stats meta count and modify count for dropped table partitions

pkg/statistics/handle/ddl/BUILD.bazel

@@ -10,6 +10,8 @@ go_library(
         "//pkg/sessionctx/variable",
         "//pkg/statistics/handle/types",
         "//pkg/statistics/handle/util",
+        "//pkg/util/logutil",
+        "@org_uber_go_zap//:zap",
     ],
 )
 
@@ -18,10 +20,11 @@ go_test(
     timeout = "short",
     srcs = ["ddl_test.go"],
     flaky = True,
-    shard_count = 4,
+    shard_count = 5,
     deps = [
         "//pkg/parser/model",
         "//pkg/planner/cardinality",
+        "//pkg/statistics/handle/util",
         "//pkg/testkit",
         "//pkg/types",
         "//pkg/util/mock",

pkg/statistics/handle/ddl/ddl.go

@@ -19,6 +19,8 @@ import (
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/statistics/handle/types"
 	"github.com/pingcap/tidb/pkg/statistics/handle/util"
+	"github.com/pingcap/tidb/pkg/util/logutil"
+	"go.uber.org/zap"
 )
 
 type ddlHandlerImpl struct {
@@ -116,13 +118,36 @@ func (h *ddlHandlerImpl) HandleDDLEvent(t *util.DDLEvent) error {
 			}
 		}
 	case model.ActionDropTablePartition:
-		// TODO: Update the modify count and count for the global table.
-		_, droppedPartitionInfo := t.GetDropPartitionInfo()
+		globalTableInfo, droppedPartitionInfo := t.GetDropPartitionInfo()
+
+		delta := int64(0)
+		count := int64(0)
 		for _, def := range droppedPartitionInfo.Definitions {
+			// Get the count and modify count of the partition.
+			stats := h.statsHandler.GetPartitionStats(globalTableInfo, def.ID)
+			if stats.Pseudo {
+				logutil.BgLogger().Warn(
+					"drop partition with pseudo stats, "+
+						"usually it won't happen because we always load stats when initializing the handle",
+					zap.String("table", globalTableInfo.Name.O),
+					zap.String("partition", def.Name.O),
+				)
+			} else {
+				delta -= stats.RealtimeCount
+				count += stats.RealtimeCount
+			}
+			// Always reset the partition stats.
 			if err := h.statsWriter.ResetTableStats2KVForDrop(def.ID); err != nil {
 				return err
 			}
 		}
+		if count != 0 {
+			if err := h.statsWriter.UpdateStatsMetaDelta(
+				globalTableInfo.ID, count, delta,
+			); err != nil {
+				return err
+			}
+		}
 	case model.ActionReorganizePartition:
 		globalTableInfo, addedPartInfo, _ := t.GetReorganizePartitionInfo()
 		for _, def := range addedPartInfo.Definitions {

pkg/statistics/handle/ddl/ddl_test.go

@@ -15,10 +15,12 @@
 package ddl_test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/planner/cardinality"
+	"github.com/pingcap/tidb/pkg/statistics/handle/util"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/mock"
@@ -284,3 +286,60 @@ PARTITION BY RANGE ( a ) (
 		}
 	}
 }
+
+func TestDropAPartition(t *testing.T) {
+	store, do := testkit.CreateMockStoreAndDomain(t)
+	testKit := testkit.NewTestKit(t, store)
+	h := do.StatsHandle()
+	testKit.MustExec("use test")
+	testKit.MustExec("drop table if exists t")
+	testKit.MustExec(`
+		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)
+		)
+	`)
+	testKit.MustExec("insert into t values (1,2),(2,2),(6,2),(11,2),(16,2)")
+	testKit.MustExec("analyze table t")
+	is := do.InfoSchema()
+	tbl, err := is.TableByName(
+		model.NewCIStr("test"), model.NewCIStr("t"),
+	)
+	require.NoError(t, err)
+	tableInfo := tbl.Meta()
+	pi := tableInfo.GetPartitionInfo()
+	for _, def := range pi.Definitions {
+		statsTbl := h.GetPartitionStats(tableInfo, def.ID)
+		require.False(t, statsTbl.Pseudo)
+	}
+	err = h.Update(is)
+	require.NoError(t, err)
+
+	testKit.MustExec("alter table t drop partition p0")
+	// Find the drop partition event.
+	var dropPartitionEvent *util.DDLEvent
+	for {
+		event := <-h.DDLEventCh()
+		if event.GetType() == model.ActionDropTablePartition {
+			dropPartitionEvent = event
+			break
+		}
+	}
+	err = h.HandleDDLEvent(dropPartitionEvent)
+	require.NoError(t, err)
+	// Check the global stats meta.
+	// Because we have dropped a partition, the count should be 3 and the modify count should be 2.
+	testKit.MustQuery(
+		fmt.Sprintf("select count, modify_count from mysql.stats_meta where table_id = %d", tableInfo.ID),
+	).Check(
+		testkit.Rows("3 2"),
+	)
+}

pkg/statistics/handle/storage/save.go

@@ -325,9 +325,18 @@ func SaveTableStatsToStorage(sctx sessionctx.Context,
 // If count is negative, both count and modify count would not be used and not be written to the table. Unless, corresponding
 // fields in the stats_meta table will be updated.
 // TODO: refactor to reduce the number of parameters
-func SaveStatsToStorage(sctx sessionctx.Context,
-	tableID int64, count, modifyCount int64, isIndex int, hg *statistics.Histogram,
-	cms *statistics.CMSketch, topN *statistics.TopN, statsVersion int, isAnalyzed int64, updateAnalyzeTime bool) (statsVer uint64, err error) {
+func SaveStatsToStorage(
+	sctx sessionctx.Context,
+	tableID int64,
+	count, modifyCount int64,
+	isIndex int,
+	hg *statistics.Histogram,
+	cms *statistics.CMSketch,
+	topN *statistics.TopN,
+	statsVersion int,
+	isAnalyzed int64,
+	updateAnalyzeTime bool,
+) (statsVer uint64, err error) {
 	version, err := util.GetStartTS(sctx)
 	if err != nil {
 		return 0, errors.Trace(err)

pkg/statistics/handle/storage/stats_read_writer.go

@@ -215,6 +215,33 @@ func (s *statsReadWriter) StatsMetaCountAndModifyCount(tableID int64) (count, mo
 	return
 }
 
+// UpdateStatsMetaDelta updates the count and modify_count for the given table in mysql.stats_meta.
+func (s *statsReadWriter) UpdateStatsMetaDelta(tableID int64, count, delta int64) (err error) {
+	err = util.CallWithSCtx(s.statsHandler.SPool(), func(sctx sessionctx.Context) error {
+		lockedTables, err := s.statsHandler.GetLockedTables(tableID)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		isLocked := false
+		if len(lockedTables) > 0 {
+			isLocked = true
+		}
+		startTS, err := util.GetStartTS(sctx)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		err = UpdateStatsMeta(
+			sctx,
+			startTS,
+			variable.TableDelta{Count: count, Delta: delta},
+			tableID,
+			isLocked,
+		)
+		return err
+	}, util.FlagWrapTxn)
+	return
+}
+
 // TableStatsFromStorage loads table stats info from storage.
 func (s *statsReadWriter) TableStatsFromStorage(tableInfo *model.TableInfo, physicalID int64, loadAll bool, snapshot uint64) (statsTbl *statistics.Table, err error) {
 	err = util.CallWithSCtx(s.statsHandler.SPool(), func(sctx sessionctx.Context) error {
@@ -233,8 +260,18 @@ func (s *statsReadWriter) TableStatsFromStorage(tableInfo *model.TableInfo, phys
 // If count is negative, both count and modify count would not be used and not be written to the table. Unless, corresponding
 // fields in the stats_meta table will be updated.
 // TODO: refactor to reduce the number of parameters
-func (s *statsReadWriter) SaveStatsToStorage(tableID int64, count, modifyCount int64, isIndex int, hg *statistics.Histogram,
-	cms *statistics.CMSketch, topN *statistics.TopN, statsVersion int, isAnalyzed int64, updateAnalyzeTime bool, source string) (err error) {
+func (s *statsReadWriter) SaveStatsToStorage(
+	tableID int64,
+	count, modifyCount int64,
+	isIndex int,
+	hg *statistics.Histogram,
+	cms *statistics.CMSketch,
+	topN *statistics.TopN,
+	statsVersion int,
+	isAnalyzed int64,
+	updateAnalyzeTime bool,
+	source string,
+) (err error) {
 	var statsVer uint64
 	err = util.CallWithSCtx(s.statsHandler.SPool(), func(sctx sessionctx.Context) error {
 		statsVer, err = SaveStatsToStorage(sctx, tableID,

pkg/statistics/handle/types/interfaces.go

@@ -228,6 +228,10 @@ type StatsReadWriter interface {
 	// StatsMetaCountAndModifyCount reads count and modify_count for the given table from mysql.stats_meta.
 	StatsMetaCountAndModifyCount(tableID int64) (count, modifyCount int64, err error)
 
+	// UpdateStatsMetaDelta updates the count and modify_count for the given table in mysql.stats_meta.
+	// It will add the delta to the original count and modify_count. The delta can be positive or negative.
+	UpdateStatsMetaDelta(tableID int64, count, delta int64) (err error)
+
 	// LoadNeededHistograms will load histograms for those needed columns/indices and put them into the cache.
 	LoadNeededHistograms() (err error)