record index usage for primary keys

Signed-off-by: Yang Keao <yangkeao@chunibyo.icu>
pingcap · Aug 27, 2024 · ffc0814 · ffc0814
1 parent 99180a8
commit ffc0814
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 5 deletions.
diff --git a/pkg/executor/builder.go b/pkg/executor/builder.go
@@ -3482,6 +3482,7 @@ func buildNoRangeTableReader(b *executorBuilder, v *plannercore.PhysicalTableRea
 	e := &TableReaderExecutor{
 		BaseExecutorV2:             exec.NewBaseExecutorV2(b.ctx.GetSessionVars(), v.Schema(), v.ID()),
 		tableReaderExecutorContext: newTableReaderExecutorContext(b.ctx),
+		indexUsageReporter:         b.buildIndexUsageReporter(v),
 		dagPB:                      dagReq,
 		startTS:                    startTS,
 		txnScope:                   b.txnScope,

diff --git a/pkg/executor/index_merge_reader.go b/pkg/executor/index_merge_reader.go
@@ -928,12 +928,12 @@ func (e *IndexMergeReaderExecutor) Close() error {
 	}
 	if e.indexUsageReporter != nil {
 		for _, p := range e.partialPlans {
-			is, ok := p[0].(*plannercore.PhysicalIndexScan)
-			if !ok {
-				continue
+			switch p := p[0].(type) {
+			case *plannercore.PhysicalTableScan:
+				e.indexUsageReporter.ReportCopIndexUsageForHandle(e.table, p.ID())
+			case *plannercore.PhysicalIndexScan:
+				e.indexUsageReporter.ReportCopIndexUsageForTable(e.table, p.Index.ID, p.ID())
 			}
-
-			e.indexUsageReporter.ReportCopIndexUsageForTable(e.table, is.Index.ID, is.ID())
 		}
 	}
 	if e.finished == nil {

diff --git a/pkg/executor/internal/exec/indexusage.go b/pkg/executor/internal/exec/indexusage.go
@@ -40,6 +40,26 @@ func NewIndexUsageReporter(reporter *indexusage.StmtIndexUsageCollector,
 	}
 }
 
+// ReportCopIndexUsageForHandle wraps around `ReportCopIndexUsageForTable` to get the `indexID` automatically
+// from the `table.Table` if the table has a clustered index or integer primary key.
+func (e *IndexUsageReporter) ReportCopIndexUsageForHandle(tbl table.Table, planID int) {
+	var idxID int64
+	if tbl.Meta().PKIsHandle {
+		idxID = 0
+	} else if tbl.Meta().IsCommonHandle {
+		for _, idx := range tbl.Meta().Indices {
+			if idx.Primary {
+				idxID = idx.ID
+			}
+		}
+	} else {
+		// just ignore, this table is scanning through rowid.
+		return
+	}
+
+	e.ReportCopIndexUsageForTable(tbl, idxID, planID)
+}
+
 // 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.

diff --git a/pkg/executor/table_reader.go b/pkg/executor/table_reader.go
@@ -135,6 +135,7 @@ func newTableReaderExecutorContext(sctx sessionctx.Context) tableReaderExecutorC
 type TableReaderExecutor struct {
 	tableReaderExecutorContext
 	exec.BaseExecutorV2
+	indexUsageReporter *exec.IndexUsageReporter
 
 	table table.Table
 
@@ -341,6 +342,10 @@ func (e *TableReaderExecutor) Next(ctx context.Context, req *chunk.Chunk) error
 
 // Close implements the Executor Close interface.
 func (e *TableReaderExecutor) Close() error {
+	if e.indexUsageReporter != nil {
+		e.indexUsageReporter.ReportCopIndexUsageForHandle(e.table, e.plans[0].ID())
+	}
+
 	var err error
 	if e.resultHandler != nil {
 		err = e.resultHandler.Close()

diff --git a/pkg/planner/core/physical_plans.go b/pkg/planner/core/physical_plans.go
@@ -154,6 +154,12 @@ type PhysicalTableReader struct {
 	TableScanAndPartitionInfos []tableScanAndPartitionInfo `plan-cache-clone:"must-nil"`
 }
 
+// LoadTableStats loads the stats of the table read by this plan.
+func (p *PhysicalTableReader) LoadTableStats(ctx sessionctx.Context) {
+	ts := p.TablePlans[0].(*PhysicalTableScan)
+	loadTableStats(ctx, ts.Table, ts.physicalTableID)
+}
+
 // PhysPlanPartInfo indicates partition helper info in physical plan.
 type PhysPlanPartInfo struct {
 	PruningConds   []expression.Expression