Repository: pingcap/tidb
File: pkg/statistics/handle/bootstrap.go (Go)

Code before:
		return nil, errors.Trace(err)
	}
	defer terror.Call(rc.Close)

	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
	req := rc.NewChunk(nil)
	iter := chunk.NewIterator4Chunk(req)
	tablesWithBuckets := make(map[int64]struct{})

	for {
		err := rc.Next(ctx, req)
		if err != nil {
			return nil, errors.Trace(err)
		}
		if req.NumRows() == 0 {
			break
		}
		for row := iter.Begin(); row != iter.End(); row = iter.Next() {
			tableID := row.GetInt64(0)
			tablesWithBuckets[tableID] = struct{}{}
		}
	}

	return tablesWithBuckets, nil
}

func (h *Handle) initStatsTopNByPaging(cache statstypes.StatsCache, task initstats.Task, totalMemory uint64) error {
	return h.Pool.SPool().WithSession(func(se *syssession.Session) error {
		return se.WithSessionContext(func(sctx sessionctx.Context) error {
			return h.initStatsTopNByPagingWithSCtx(sctx, cache, task, totalMemory)
		})
	})
}

// initStatsTopNByPagingWithSCtx contains the core business logic for initStatsTopNByPaging.
// This method preserves git blame history by keeping the original logic intact.
func (h *Handle) initStatsTopNByPagingWithSCtx(sctx sessionctx.Context, cache statstypes.StatsCache, task initstats.Task, totalMemory uint64) error {
	// First, get the tables with buckets in this range
	tablesWithBuckets, err := getTablesWithBucketsInRange(sctx, [2]int64{task.StartTid, task.EndTid})
	if err != nil {
		return errors.Trace(err)
	}

	sql := genInitStatsTopNSQLForIndexes(true, [2]int64{task.StartTid, task.EndTid})
	rc, err := util.Exec(sctx, sql)
	if err != nil {
		return errors.Trace(err)
	}
	defer terror.Call(rc.Close)
	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
	req := rc.NewChunk(nil)

Review comment (performance):
The new function adds an additional query to mysql.stats_buckets for each task range. While the hint USE_INDEX(stats_buckets, tbl) is specified, this introduces overhead as it must scan the index to find distinct table_ids before the main TopN loading begins. Consider whether this pre-query is necessary for all cases or if it could be optimized. For example, if the range contains many tables without buckets, this provides value, but if most tables have buckets, this becomes additional overhead. You might want to add metrics or logging to track how often this optimization is beneficial.

Code after:
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer terror.Call(rc.Close)

	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
	req := rc.NewChunk(nil)
	iter := chunk.NewIterator4Chunk(req)
	tablesWithBuckets := make(map[int64]struct{})

	for {
		err := rc.Next(ctx, req)
		if err != nil {
			return nil, errors.Trace(err)
		}
		if req.NumRows() == 0 {
			break
		}
		for row := iter.Begin(); row != iter.End(); row = iter.Next() {
			tableID := row.GetInt64(0)
			tablesWithBuckets[tableID] = struct{}{}
		}
	}

	return tablesWithBuckets, nil
}

func (h *Handle) initStatsTopNByPaging(cache statstypes.StatsCache, task initstats.Task, totalMemory uint64) error {
	return h.Pool.SPool().WithSession(func(se *syssession.Session) error {
		return se.WithSessionContext(func(sctx sessionctx.Context) error {
			return h.initStatsTopNByPagingWithSCtx(sctx, cache, task, totalMemory)
		})
	})
}

// initStatsTopNByPagingWithSCtx contains the core business logic for initStatsTopNByPaging.
// This method preserves git blame history by keeping the original logic intact.
func (h *Handle) initStatsTopNByPagingWithSCtx(sctx sessionctx.Context, cache statstypes.StatsCache, task initstats.Task, totalMemory uint64) error {
	// First, get the tables with buckets in this range
	tablesWithBuckets, err := getTablesWithBucketsInRange(sctx, [2]int64{task.StartTid, task.EndTid})
	if err != nil {
		return errors.Trace(err)
	}

	sql := genInitStatsTopNSQLForIndexes(true, [2]int64{task.StartTid, task.EndTid})
	rc, err := util.Exec(sctx, sql)
	if err != nil {
		return errors.Trace(err)
	}
	defer terror.Call(rc.Close)
	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)