executor: record the index usage on a the whole table (pingcap#50816)

close pingcap#50814

pkg/executor/batch_point_get.go

@@ -168,10 +168,9 @@ func (e *BatchPointGetExec) Close() error {
 	}
 	if e.indexUsageReporter != nil && e.idxInfo != nil {
 		kvReqTotal := e.stats.GetCmdRPCCount(tikvrpc.CmdBatchGet)
-		// We cannot distinguish how many rows are coming from each partition. Here, we record all index usages on the
-		// table itself (but not for the partition physical table), which is different from how an index usage is
-		// treated for a local index.
-		e.indexUsageReporter.ReportPointGetIndexUsage(e.tblInfo.ID, e.idxInfo.ID, e.ID(), kvReqTotal)
+		// We cannot distinguish how many rows are coming from each partition. Here, we calculate all index usages
+		// percentage according to the row counts for the whole table.
+		e.indexUsageReporter.ReportPointGetIndexUsage(e.tblInfo.ID, e.tblInfo.ID, e.idxInfo.ID, e.ID(), kvReqTotal)
 	}
 	e.inited = 0
 	e.index = 0

pkg/executor/distsql.go

@@ -227,7 +227,7 @@ func (e *IndexReaderExecutor) setDummy() {
 // Close clears all resources hold by current object.
 func (e *IndexReaderExecutor) Close() (err error) {
 	if e.indexUsageReporter != nil {
-		e.indexUsageReporter.ReportCopIndexUsage(e.physicalTableID, e.index.ID, e.plans[0].ID())
+		e.indexUsageReporter.ReportCopIndexUsageForTable(e.table, e.index.ID, e.plans[0].ID())
 	}
 
 	if e.result != nil {
@@ -832,8 +832,8 @@ func (e *IndexLookUpExecutor) Close() error {
 		defer e.Ctx().GetSessionVars().StmtCtx.RuntimeStatsColl.RegisterStats(e.ID(), e.stats)
 	}
 	if e.indexUsageReporter != nil {
-		e.indexUsageReporter.ReportCopIndexUsage(
-			e.table.Meta().ID,
+		e.indexUsageReporter.ReportCopIndexUsageForTable(
+			e.table,
 			e.index.ID,
 			e.idxPlans[0].ID())
 	}

pkg/executor/index_merge_reader.go

@@ -912,14 +912,13 @@ func (e *IndexMergeReaderExecutor) Close() error {
 		defer e.Ctx().GetSessionVars().StmtCtx.RuntimeStatsColl.RegisterStats(e.ID(), e.stats)
 	}
 	if e.indexUsageReporter != nil {
-		tableID := e.table.Meta().ID
 		for _, p := range e.partialPlans {
 			is, ok := p[0].(*plannercore.PhysicalIndexScan)
 			if !ok {
 				continue
 			}
 
-			e.indexUsageReporter.ReportCopIndexUsage(tableID, is.Index.ID, is.ID())
+			e.indexUsageReporter.ReportCopIndexUsageForTable(e.table, is.Index.ID, is.ID())
 		}
 	}
 	if e.finished == nil {

pkg/executor/internal/exec/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "//pkg/sessionctx/variable",
         "//pkg/statistics",
         "//pkg/statistics/handle/usage/indexusage",
+        "//pkg/table",
         "//pkg/types",
         "//pkg/util",
         "//pkg/util/chunk",

pkg/executor/internal/exec/indexusage.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/pkg/statistics"
 	"github.com/pingcap/tidb/pkg/statistics/handle/usage/indexusage"
+	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/util/execdetails"
 )
 
@@ -39,9 +40,23 @@ func NewIndexUsageReporter(reporter *indexusage.StmtIndexUsageCollector,
 	}
 }
 
-// ReportCopIndexUsage reports the index usage to the inside collector
-func (e *IndexUsageReporter) ReportCopIndexUsage(tableID int64, indexID int64, planID int) {
-	tableRowCount, ok := e.getTableRowCount(tableID)
+// ReportCopIndexUsageForTable wraps around `ReportCopIndexUsage` to get `tableID` and `physicalTableID` from the
+// `table.Table`. If it's expected to calculate the percentage according to the size of partition, the `tbl` argument
+// should be a `table.PhysicalTable`, or the percentage will be calculated using the size of whole table.
+func (e *IndexUsageReporter) ReportCopIndexUsageForTable(tbl table.Table, indexID int64, planID int) {
+	tableID := tbl.Meta().ID
+	physicalTableID := tableID
+	if physicalTable, ok := tbl.(table.PhysicalTable); ok {
+		physicalTableID = physicalTable.GetPhysicalID()
+	}
+
+	e.ReportCopIndexUsage(tableID, physicalTableID, indexID, planID)
+}
+
+// ReportCopIndexUsage reports the index usage to the inside collector. The index usage will be recorded in the
+// `tableID+indexID`, but the percentage is calculated using the size of the table specified by `physicalTableID`.
+func (e *IndexUsageReporter) ReportCopIndexUsage(tableID int64, physicalTableID int64, indexID int64, planID int) {
+	tableRowCount, ok := e.getTableRowCount(physicalTableID)
 	if !ok {
 		// skip if the table is empty or the stats is not valid
 		return
@@ -61,8 +76,8 @@ func (e *IndexUsageReporter) ReportCopIndexUsage(tableID int64, indexID int64, p
 }
 
 // ReportPointGetIndexUsage reports the index usage of a point get or batch point get
-func (e *IndexUsageReporter) ReportPointGetIndexUsage(tableID int64, indexID int64, planID int, kvRequestTotal int64) {
-	tableRowCount, ok := e.getTableRowCount(tableID)
+func (e *IndexUsageReporter) ReportPointGetIndexUsage(tableID int64, physicalTableID int64, indexID int64, planID int, kvRequestTotal int64) {
+	tableRowCount, ok := e.getTableRowCount(physicalTableID)
 	if !ok {
 		// skip if the table is empty or the stats is not valid
 		return

pkg/executor/internal/exec/indexusage_test.go

@@ -53,7 +53,7 @@ func TestIndexUsageReporter(t *testing.T) {
 	// For PointGet and BatchPointGet
 	planID := 3
 	runtimeStatsColl.GetBasicRuntimeStats(planID).Record(time.Second, 2024)
-	reporter.ReportPointGetIndexUsage(tableID, indexID, planID, 1)
+	reporter.ReportPointGetIndexUsage(tableID, tableID, indexID, planID, 1)
 
 	require.Eventually(t, func() bool {
 		tk.Session().ReportUsageStats()
@@ -73,7 +73,7 @@ func TestIndexUsageReporter(t *testing.T) {
 		ExecutorId:      &executorID,
 		Concurrency:     &zero,
 	})
-	reporter.ReportCopIndexUsage(tableID, indexID, planID)
+	reporter.ReportCopIndexUsage(tableID, tableID, indexID, planID)
 
 	require.Eventually(t, func() bool {
 		tk.Session().ReportUsageStats()
@@ -88,7 +88,7 @@ func TestIndexUsageReporter(t *testing.T) {
 	}
 	planID = 4
 	runtimeStatsColl.GetBasicRuntimeStats(planID).Record(time.Second, 2024)
-	reporter.ReportPointGetIndexUsage(tableID, indexID, planID, 1)
+	reporter.ReportPointGetIndexUsage(tableID, tableID, indexID, planID, 1)
 
 	require.Eventually(t, func() bool {
 		tk.Session().ReportUsageStats()
@@ -299,10 +299,6 @@ partition p3 values less than MAXVALUE)`)
 
 	table, err := dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t"))
 	require.NoError(t, err)
-	var pids []int64
-	for i := 0; i < 4; i++ {
-		pids = append(pids, table.Meta().Partition.GetPartitionIDByName(fmt.Sprintf("p%d", i)))
-	}
 	idx1ID := int64(0)
 	for _, idx := range table.Indices() {
 		if idx.Meta().Name.L == "idx_1" {
@@ -324,11 +320,11 @@ partition p3 values less than MAXVALUE)`)
 			"select id_1 from t where id_1 >= 30",
 			"PartitionUnion",
 			[]indexStatsExpect{
-				{pids[2], idx1ID, []indexusage.Sample{
+				{table.Meta().ID, idx1ID, []indexusage.Sample{
 					indexusage.NewSample(1, 1, 20, 30),
 				}},
-				{pids[3], idx1ID, []indexusage.Sample{
-					indexusage.NewSample(1, 1, 50, 50),
+				{table.Meta().ID, idx1ID, []indexusage.Sample{
+					indexusage.NewSample(0, 1, 50, 50),
 				}},
 			},
 		},
@@ -337,15 +333,15 @@ partition p3 values less than MAXVALUE)`)
 			"select id_1 from t where id_1 - 95 >= 0 and id_1 >= 90",
 			"IndexReader",
 			[]indexStatsExpect{
-				{pids[3], idx1ID, []indexusage.Sample{indexusage.NewSample(1, 1, 10, 50)}},
+				{table.Meta().ID, idx1ID, []indexusage.Sample{indexusage.NewSample(1, 1, 10, 50)}},
 			},
 		},
 		// PointGet in a partition
 		{
 			"select * from t where id_1 = 1",
 			"Point_Get",
 			[]indexStatsExpect{
-				{pids[0], idx1ID, []indexusage.Sample{indexusage.NewSample(1, 1, 1, 10)}},
+				{table.Meta().ID, idx1ID, []indexusage.Sample{indexusage.NewSample(1, 1, 1, 10)}},
 			},
 		},
 		// BatchPointGet in a partition

pkg/executor/point_get.go

@@ -205,11 +205,12 @@ func (e *PointGetExecutor) Close() error {
 	}
 	if e.indexUsageReporter != nil && e.idxInfo != nil {
 		tableID := e.tblInfo.ID
+		physicalTableID := tableID
 		if e.partitionDef != nil {
-			tableID = e.partitionDef.ID
+			physicalTableID = e.partitionDef.ID
 		}
 		kvReqTotal := e.stats.SnapshotRuntimeStats.GetCmdRPCCount(tikvrpc.CmdGet)
-		e.indexUsageReporter.ReportPointGetIndexUsage(tableID, e.idxInfo.ID, e.ID(), kvReqTotal)
+		e.indexUsageReporter.ReportPointGetIndexUsage(tableID, physicalTableID, e.idxInfo.ID, e.ID(), kvReqTotal)
 	}
 	e.done = false
 	return nil

pkg/statistics/handle/usage/indexusage/collector.go

@@ -257,6 +257,8 @@ func (s *StmtIndexUsageCollector) Update(tableID int64, indexID int64, sample Sa
 	s.Lock()
 	defer s.Unlock()
 
+	// If the usage of the table/index has been recorded in the statement, it'll not update the `QueryTotal`. Before the
+	// execution of each statement, the `StmtIndexUsageCollector` and internal map will be re-created.
 	idxID := GlobalIndexID{IndexID: indexID, TableID: tableID}
 	if _, ok := s.recordedIndex[idxID]; !ok {
 		sample.QueryTotal = 1