From dc79ec20268af87607a404e8c4c68a1a07d8c1fe Mon Sep 17 00:00:00 2001 From: Arenatlx <314806019@qq.com> Date: Fri, 6 Mar 2026 15:06:59 +0800 Subject: [PATCH] This is an automated cherry-pick of #66677 Signed-off-by: ti-chi-bot --- pkg/planner/core/logical_plan_builder.go | 13 ++-- pkg/planner/core/planbuilder.go | 19 ++++++ pkg/planner/optimize.go | 2 +- pkg/sessionctx/stmtctx/BUILD.bazel | 6 ++ pkg/sessionctx/stmtctx/stmtctx.go | 59 ++++++++++++++++ pkg/sessionctx/stmtctx/stmtctx_test.go | 86 ++++++++++++++++++++++++ pkg/util/context/plancache.go | 20 ++++++ pkg/util/hint/hint_query_block.go | 71 +++++++++++++++---- 8 files changed, 257 insertions(+), 19 deletions(-) diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index ead149d7bdf80..b4bb77558c860 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -3749,7 +3749,7 @@ func (b *PlanBuilder) pushHintWithoutTableWarning(hint *ast.TableOptimizerHint) } func (b *PlanBuilder) pushTableHints(hints []*ast.TableOptimizerHint, currentLevel int) { - hints = b.hintProcessor.GetCurrentStmtHints(hints, currentLevel) + hints = b.hintProcessor.GetCurrentStmtHints(hints, currentLevel, b.hintState) sessionVars := b.ctx.GetSessionVars() currentDB := sessionVars.CurrentDB warnHandler := sessionVars.StmtCtx @@ -4628,7 +4628,7 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as // Because of the nested views, so we should check the left table list in hint when build the data source from the view inside the current view. currentQBNameMap4View[qbName] = viewQBNameHintTable[1:] currentViewHints[qbName] = b.hintProcessor.ViewQBNameToHints[qbName] - b.hintProcessor.ViewQBNameUsed[qbName] = struct{}{} + b.hintProcessor.MarkViewQBNameUsed(qbName, b.hintState) } } return b.BuildDataSourceFromView(ctx, dbName, tableInfo, currentQBNameMap4View, currentViewHints) @@ -5109,18 +5109,21 @@ func (b *PlanBuilder) BuildDataSourceFromView(ctx context.Context, dbName pmodel hintProcessor.ViewQBNameToTable = qbNameMap4View hintProcessor.ViewQBNameToHints = viewHints - hintProcessor.ViewQBNameUsed = make(map[string]struct{}) - hintProcessor.QBOffsetToHints = currentQbHints hintProcessor.QBNameToSelOffset = currentQbNameMap + hintState := hintProcessor.NewBuildState() + hintState.QBOffsetToHints = currentQbHints originHintProcessor := b.hintProcessor + originHintState := b.hintState originPlannerSelectBlockAsName := b.ctx.GetSessionVars().PlannerSelectBlockAsName.Load() b.hintProcessor = hintProcessor + b.hintState = hintState newPlannerSelectBlockAsName := make([]ast.HintTable, hintProcessor.MaxSelectStmtOffset()+1) b.ctx.GetSessionVars().PlannerSelectBlockAsName.Store(&newPlannerSelectBlockAsName) defer func() { - b.hintProcessor.HandleUnusedViewHints() + b.hintProcessor.SetWarns(b.hintProcessor.HandleUnusedViewHints(b.hintState, nil)) b.hintProcessor = originHintProcessor + b.hintState = originHintState b.ctx.GetSessionVars().PlannerSelectBlockAsName.Store(originPlannerSelectBlockAsName) }() nodeW := resolve.NewNodeWWithCtx(selectNode, b.resolveCtx) diff --git a/pkg/planner/core/planbuilder.go b/pkg/planner/core/planbuilder.go index 09d433dbb9967..c909b6bab29c6 100644 --- a/pkg/planner/core/planbuilder.go +++ b/pkg/planner/core/planbuilder.go @@ -256,7 +256,10 @@ type PlanBuilder struct { // finish building the subquery or CTE. handleHelper *handleColHelper + // read-only meta derived from ast node. hintProcessor *hint.QBHintHandler + // mutable state of QBHint when building. + hintState *hint.QBHintBuildState // qbOffset is the offsets of current processing select stmts. qbOffset []int @@ -405,6 +408,11 @@ func GetDBTableInfo(visitInfo []visitInfo) []stmtctx.TableEntry { return tables } +// GetHintState gets the HintState from the PlanBuilder. +func (b *PlanBuilder) GetHintState() *hint.QBHintBuildState { + return b.hintState +} + // GetOptFlag gets the OptFlag of the PlanBuilder. func (b *PlanBuilder) GetOptFlag() uint64 { if b.isSampling { @@ -482,6 +490,9 @@ func (b *PlanBuilder) Init(sctx base.PlanContext, is infoschema.InfoSchema, proc b.ctx = sctx b.is = is b.hintProcessor = processor + if processor != nil { + b.hintState = processor.NewBuildState() + } b.isForUpdateRead = sctx.GetSessionVars().IsPessimisticReadConsistency() b.noDecorrelate = sctx.GetSessionVars().EnableNoDecorrelateInSelect if savedBlockNames == nil { @@ -523,6 +534,14 @@ func (b *PlanBuilder) ResetForReuse() *PlanBuilder { return b } +// HandleUnusedViewHints appends warnings for unused view hints in the current build. +func (b *PlanBuilder) HandleUnusedViewHints() { + if b.hintProcessor == nil { + return + } + b.hintProcessor.SetWarns(b.hintProcessor.HandleUnusedViewHints(b.hintState, nil)) +} + // Build builds the ast node to a Plan. func (b *PlanBuilder) Build(ctx context.Context, node *resolve.NodeW) (base.Plan, error) { // Build might be called recursively, right now they all share the same resolve diff --git a/pkg/planner/optimize.go b/pkg/planner/optimize.go index 8f19e77467d85..4da6ad5539ecc 100644 --- a/pkg/planner/optimize.go +++ b/pkg/planner/optimize.go @@ -492,10 +492,10 @@ func optimize(ctx context.Context, sctx planctx.PlanContext, node *resolve.NodeW // build logical plan hintProcessor := hint.NewQBHintHandler(sctx.GetSessionVars().StmtCtx) node.Node.Accept(hintProcessor) - defer hintProcessor.HandleUnusedViewHints() builder := planBuilderPool.Get().(*core.PlanBuilder) defer planBuilderPool.Put(builder.ResetForReuse()) builder.Init(sctx, is, hintProcessor) + defer builder.HandleUnusedViewHints() p, err := buildLogicalPlan(ctx, sctx, node, builder) if err != nil { return nil, nil, 0, err diff --git a/pkg/sessionctx/stmtctx/BUILD.bazel b/pkg/sessionctx/stmtctx/BUILD.bazel index 3f45049846d15..10c88e5d5f61d 100644 --- a/pkg/sessionctx/stmtctx/BUILD.bazel +++ b/pkg/sessionctx/stmtctx/BUILD.bazel @@ -45,10 +45,16 @@ go_test( ], embed = [":stmtctx"], flaky = True, +<<<<<<< HEAD shard_count = 14, +======= + shard_count = 17, +>>>>>>> d59e531fe61 (planner: build multi alternative logical plan from shared AST (#66677)) deps = [ "//pkg/errctx", "//pkg/kv", + "//pkg/meta/model", + "//pkg/parser/ast", "//pkg/sessionctx/variable", "//pkg/testkit", "//pkg/testkit/testfailpoint", diff --git a/pkg/sessionctx/stmtctx/stmtctx.go b/pkg/sessionctx/stmtctx/stmtctx.go index 7e3feeb3732f0..6fca84e217dbb 100644 --- a/pkg/sessionctx/stmtctx/stmtctx.go +++ b/pkg/sessionctx/stmtctx/stmtctx.go @@ -65,10 +65,31 @@ func AllocateTaskID() uint64 { // SQLWarn relates a sql warning and it's level. type SQLWarn = contextutil.SQLWarn +<<<<<<< HEAD type jsonSQLWarn struct { Level string `json:"level"` SQLErr *terror.Error `json:"err,omitempty"` Msg string `json:"msg,omitempty"` +======= +// LogicalPlanBuildState stores the statement-scoped planner state that is mutated while +// building a logical plan from AST. +type LogicalPlanBuildState struct { + warnings []SQLWarn + extraWarnings []SQLWarn + tables []TableEntry + tableStats map[int64]any + lockTableIDs map[int64]struct{} + tblInfo2UnionScan map[*model.TableInfo]bool + useDynamicPruneMode bool + viewDepth int32 + colRefFromUpdatePlan intset.FastIntSet + // plan cache related stuff + planCacheUseCache bool + planCacheType contextutil.PlanCacheType + planCacheUnqualified string + planCacheForce bool + planCacheAlwaysWarn bool +>>>>>>> d59e531fe61 (planner: build multi alternative logical plan from shared AST (#66677)) } // ReferenceCount indicates the reference count of StmtCtx. @@ -572,6 +593,44 @@ func (sc *StatementContext) Reset() bool { return true } +// SaveLogicalPlanBuildState captures the statement-scoped planner state before building +// another logical plan candidate from the same AST. +func (sc *StatementContext) SaveLogicalPlanBuildState() LogicalPlanBuildState { + planCacheUseCache, planCacheType, planCacheUnqualified, planCacheForce, planCacheAlwaysWarn := sc.PlanCacheTracker.Save() + return LogicalPlanBuildState{ + warnings: slices.Clone(sc.GetWarnings()), + extraWarnings: slices.Clone(sc.GetExtraWarnings()), + tables: slices.Clone(sc.Tables), + tableStats: maps.Clone(sc.TableStats), + lockTableIDs: maps.Clone(sc.LockTableIDs), + tblInfo2UnionScan: maps.Clone(sc.TblInfo2UnionScan), + useDynamicPruneMode: sc.UseDynamicPruneMode, + viewDepth: sc.ViewDepth, + colRefFromUpdatePlan: sc.ColRefFromUpdatePlan.Copy(), + planCacheUseCache: planCacheUseCache, + planCacheType: planCacheType, + planCacheUnqualified: planCacheUnqualified, + planCacheForce: planCacheForce, + planCacheAlwaysWarn: planCacheAlwaysWarn, + } +} + +// RestoreLogicalPlanBuildState restores the statement-scoped planner state after a +// discarded logical plan build attempt. +func (sc *StatementContext) RestoreLogicalPlanBuildState(state LogicalPlanBuildState) { + sc.SetWarnings(slices.Clone(state.warnings)) + sc.SetExtraWarnings(slices.Clone(state.extraWarnings)) + sc.Tables = slices.Clone(state.tables) + sc.TableStats = maps.Clone(state.tableStats) + sc.LockTableIDs = maps.Clone(state.lockTableIDs) + sc.TblInfo2UnionScan = maps.Clone(state.tblInfo2UnionScan) + sc.UseDynamicPruneMode = state.useDynamicPruneMode + sc.ViewDepth = state.viewDepth + sc.ColRefFromUpdatePlan.CopyFrom(state.colRefFromUpdatePlan) + sc.PlanCacheTracker.Restore(state.planCacheUseCache, state.planCacheType, state.planCacheUnqualified, state.planCacheForce, state.planCacheAlwaysWarn) + sc.RangeFallbackHandler = contextutil.NewRangeFallbackHandler(&sc.PlanCacheTracker, sc) +} + // CtxID returns the context id of the statement func (sc *StatementContext) CtxID() uint64 { return sc.ctxID diff --git a/pkg/sessionctx/stmtctx/stmtctx_test.go b/pkg/sessionctx/stmtctx/stmtctx_test.go index b5d61cfe5a036..a6d23c0f196c1 100644 --- a/pkg/sessionctx/stmtctx/stmtctx_test.go +++ b/pkg/sessionctx/stmtctx/stmtctx_test.go @@ -27,6 +27,8 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/tidb/pkg/errctx" "github.com/pingcap/tidb/pkg/kv" + "github.com/pingcap/tidb/pkg/meta/model" + "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/sessionctx/stmtctx" "github.com/pingcap/tidb/pkg/sessionctx/variable" "github.com/pingcap/tidb/pkg/testkit" @@ -218,6 +220,90 @@ func TestMarshalSQLWarn(t *testing.T) { tk.MustQuery("show warnings").Check(rows) } +func TestLogicalPlanBuildStateRestore(t *testing.T) { + sc := stmtctx.NewStmtCtx() + sc.AppendWarning(errors.New("baseline warning")) + sc.AppendExtraWarning(errors.New("baseline extra warning")) + sc.Tables = []stmtctx.TableEntry{{DB: "test", Table: "t"}} + sc.TableStats = map[int64]any{42: "baseline stats"} + sc.LockTableIDs = map[int64]struct{}{1: {}} + tblInfo := &model.TableInfo{ID: 42} + sc.TblInfo2UnionScan = map[*model.TableInfo]bool{tblInfo: true} + sc.UseDynamicPruneMode = true + sc.ViewDepth = 2 + sc.ColRefFromUpdatePlan.Insert(7) + sc.SetCacheType(contextutil.SessionNonPrepared) + sc.EnablePlanCache() + + state := sc.SaveLogicalPlanBuildState() + + sc.AppendWarning(errors.New("candidate warning")) + sc.AppendExtraWarning(errors.New("candidate extra warning")) + sc.Tables = []stmtctx.TableEntry{{DB: "candidate", Table: "t2"}} + sc.TableStats = map[int64]any{99: "candidate stats"} + sc.LockTableIDs[2] = struct{}{} + sc.TblInfo2UnionScan = map[*model.TableInfo]bool{{ID: 99}: false} + sc.UseDynamicPruneMode = false + sc.ViewDepth = 9 + sc.ColRefFromUpdatePlan.Insert(9) + sc.SetSkipPlanCache("candidate reason") + + sc.RestoreLogicalPlanBuildState(state) + + warnings := sc.GetWarnings() + require.Len(t, warnings, 1) + require.Equal(t, "baseline warning", warnings[0].Err.Error()) + + extraWarnings := sc.GetExtraWarnings() + require.Len(t, extraWarnings, 1) + require.Equal(t, "baseline extra warning", extraWarnings[0].Err.Error()) + + require.Equal(t, []stmtctx.TableEntry{{DB: "test", Table: "t"}}, sc.Tables) + require.Equal(t, map[int64]any{42: "baseline stats"}, sc.TableStats) + require.Equal(t, map[int64]struct{}{1: {}}, sc.LockTableIDs) + require.Equal(t, map[*model.TableInfo]bool{tblInfo: true}, sc.TblInfo2UnionScan) + require.True(t, sc.UseDynamicPartitionPrune()) + require.Equal(t, int32(2), sc.ViewDepth) + require.True(t, sc.ColRefFromUpdatePlan.Has(7)) + require.False(t, sc.ColRefFromUpdatePlan.Has(9)) + require.True(t, sc.UseCache()) + require.Empty(t, sc.PlanCacheUnqualified()) +} + +func TestQBHintHandlerBuildState(t *testing.T) { + handler := hint.NewQBHintHandler(nil) + handler.QBNameToSelOffset = map[string]int{"qb_1": 1} + handler.ViewQBNameToTable = map[string][]ast.HintTable{ + "view_qb": {{TableName: ast.NewCIStr("t")}}, + } + handler.ViewQBNameToHints = map[string][]*ast.TableOptimizerHint{ + "view_qb": {{HintName: ast.NewCIStr("merge_join")}}, + } + handler.Enter(&ast.SelectStmt{}) + handler.Enter(&ast.SelectStmt{}) + state := handler.NewBuildState() + hints := handler.GetCurrentStmtHints([]*ast.TableOptimizerHint{ + {HintName: ast.NewCIStr("use_index"), QBName: ast.NewCIStr("qb_1")}, + }, 1, state) + handler.MarkViewQBNameUsed("view_qb", state) + + require.Len(t, hints, 1) + require.Equal(t, "use_index", hints[0].HintName.L) + + require.Equal(t, 2, handler.MaxSelectStmtOffset()) + require.Equal(t, map[string]int{"qb_1": 1}, handler.QBNameToSelOffset) + require.Equal(t, map[string][]*ast.TableOptimizerHint{ + "view_qb": {{HintName: ast.NewCIStr("merge_join")}}, + }, handler.ViewQBNameToHints) + require.Equal(t, map[string][]ast.HintTable{ + "view_qb": {{TableName: ast.NewCIStr("t")}}, + }, handler.ViewQBNameToTable) + require.Equal(t, map[int][]*ast.TableOptimizerHint{ + 1: {{HintName: ast.NewCIStr("use_index"), QBName: ast.NewCIStr("qb_1")}}, + }, state.QBOffsetToHints) + require.Equal(t, map[string]struct{}{"view_qb": {}}, state.ViewQBNameUsed) +} + func TestApproxRuntimeInfo(t *testing.T) { var n = rand.Intn(19000) + 1000 var valRange = rand.Int31n(10000) + 1000 diff --git a/pkg/util/context/plancache.go b/pkg/util/context/plancache.go index 9cac5ec81012e..d6c46cd4db81e 100644 --- a/pkg/util/context/plancache.go +++ b/pkg/util/context/plancache.go @@ -122,6 +122,26 @@ func (h *PlanCacheTracker) EnablePlanCache() { h.useCache = true } +// Save captures the mutable planning-time state of the tracker. +func (h *PlanCacheTracker) Save() (useCache bool, cacheType PlanCacheType, planCacheUnqualified string, forcePlanCache bool, alwaysWarnSkipCache bool) { + h.mu.Lock() + defer h.mu.Unlock() + + return h.useCache, h.cacheType, h.planCacheUnqualified, h.forcePlanCache, h.alwaysWarnSkipCache +} + +// Restore restores the mutable planning-time state of the tracker. +func (h *PlanCacheTracker) Restore(useCache bool, cacheType PlanCacheType, planCacheUnqualified string, forcePlanCache bool, alwaysWarnSkipCache bool) { + h.mu.Lock() + defer h.mu.Unlock() + + h.useCache = useCache + h.cacheType = cacheType + h.planCacheUnqualified = planCacheUnqualified + h.forcePlanCache = forcePlanCache + h.alwaysWarnSkipCache = alwaysWarnSkipCache +} + // UseCache returns whether to use plan cache. func (h *PlanCacheTracker) UseCache() bool { h.mu.Lock() diff --git a/pkg/util/hint/hint_query_block.go b/pkg/util/hint/hint_query_block.go index 99bfc646837e6..3b98e93a69bd0 100644 --- a/pkg/util/hint/hint_query_block.go +++ b/pkg/util/hint/hint_query_block.go @@ -31,18 +31,23 @@ import ( // In both cases, the `use_index` hint doesn't take effect directly, since a specific qb_name is specified, and this // QBHintHandler is used to handle this cases. type QBHintHandler struct { - QBNameToSelOffset map[string]int // map[QBName]SelectOffset - QBOffsetToHints map[int][]*ast.TableOptimizerHint // map[QueryBlockOffset]Hints + QBNameToSelOffset map[string]int // map[QBName]SelectOffset // Used for the view's hint ViewQBNameToTable map[string][]ast.HintTable // map[QBName]HintedTable ViewQBNameToHints map[string][]*ast.TableOptimizerHint // map[QBName]Hints - ViewQBNameUsed map[string]struct{} // map[QBName]Used warnHandler hintWarnHandler selectStmtOffset int } +// QBHintBuildState stores the per-build runtime state for a QBHintHandler. +// The handler itself only keeps AST-derived metadata and can be shared safely. +type QBHintBuildState struct { + QBOffsetToHints map[int][]*ast.TableOptimizerHint // map[QueryBlockOffset]Hints + ViewQBNameUsed map[string]struct{} // map[QBName]Used +} + // hintWarnHandler is used to handle the warning when parsing hints. type hintWarnHandler interface { SetHintWarning(warn string) @@ -56,6 +61,20 @@ func NewQBHintHandler(warnHandler hintWarnHandler) *QBHintHandler { } } +// NewBuildState creates the per-build runtime state for the handler. +func (p *QBHintHandler) NewBuildState() *QBHintBuildState { + if p == nil { + return nil + } + + state := &QBHintBuildState{} + if len(p.ViewQBNameToTable) == 0 { + return state + } + state.ViewQBNameUsed = make(map[string]struct{}, len(p.ViewQBNameToTable)) + return state +} + // MaxSelectStmtOffset returns the current stmt offset. func (p *QBHintHandler) MaxSelectStmtOffset() int { return p.selectStmtOffset @@ -133,7 +152,6 @@ func (p *QBHintHandler) handleViewHints(hints []*ast.TableOptimizerHint, offset usedHints[i] = true if p.ViewQBNameToTable == nil { p.ViewQBNameToTable = make(map[string][]ast.HintTable) - p.ViewQBNameUsed = make(map[string]struct{}) } qbName := hint.QBName.L if qbName == "" { @@ -202,15 +220,21 @@ func (p *QBHintHandler) handleViewHints(hints []*ast.TableOptimizerHint, offset } // HandleUnusedViewHints handle the unused view hints. -func (p *QBHintHandler) HandleUnusedViewHints() { +func (p *QBHintHandler) HandleUnusedViewHints(state *QBHintBuildState, warn []string) []string { + if state == nil { + return warn + } + + warn = warn[:0] if p.ViewQBNameToTable != nil { for qbName := range p.ViewQBNameToTable { - _, ok := p.ViewQBNameUsed[qbName] + _, ok := state.ViewQBNameUsed[qbName] if !ok && p.warnHandler != nil { - p.warnHandler.SetHintWarning(fmt.Sprintf("The qb_name hint %s is unused, please check whether the table list in the qb_name hint %s is correct", qbName, qbName)) + warn = append(warn, fmt.Sprintf("The qb_name hint %s is unused, please check whether the table list in the qb_name hint %s is correct", qbName, qbName)) } } } + return warn } const ( @@ -243,6 +267,16 @@ func (p *QBHintHandler) getBlockOffset(blockName model.CIStr) int { return -1 } +// SetWarns set the warning from a list of strings. +func (p *QBHintHandler) SetWarns(warns []string) { + if p == nil || p.warnHandler == nil || len(warns) == 0 { + return + } + for _, one := range warns { + p.warnHandler.SetHintWarning(one) + } +} + // GetHintOffset gets the offset of stmt that the hints take effects. func (p *QBHintHandler) GetHintOffset(qbName model.CIStr, currentOffset int) int { if qbName.L != "" { @@ -280,9 +314,12 @@ func (p *QBHintHandler) isHint4View(hint *ast.TableOptimizerHint) bool { } // GetCurrentStmtHints extracts all hints that take effects at current stmt. -func (p *QBHintHandler) GetCurrentStmtHints(hints []*ast.TableOptimizerHint, currentOffset int) []*ast.TableOptimizerHint { - if p.QBOffsetToHints == nil { - p.QBOffsetToHints = make(map[int][]*ast.TableOptimizerHint) +func (p *QBHintHandler) GetCurrentStmtHints(hints []*ast.TableOptimizerHint, currentOffset int, state *QBHintBuildState) []*ast.TableOptimizerHint { + if state == nil { + state = &QBHintBuildState{} + } + if state.QBOffsetToHints == nil { + state.QBOffsetToHints = make(map[int][]*ast.TableOptimizerHint) } for _, hint := range hints { if hint.HintName.L == hintQBName { @@ -296,11 +333,19 @@ func (p *QBHintHandler) GetCurrentStmtHints(hints []*ast.TableOptimizerHint, cur } continue } - if !slices.Contains(p.QBOffsetToHints[offset], hint) { - p.QBOffsetToHints[offset] = append(p.QBOffsetToHints[offset], hint) + if !slices.Contains(state.QBOffsetToHints[offset], hint) { + state.QBOffsetToHints[offset] = append(state.QBOffsetToHints[offset], hint) } } - return p.QBOffsetToHints[currentOffset] + return state.QBOffsetToHints[currentOffset] +} + +// MarkViewQBNameUsed records that the named view qb hint was used in the current build. +func (*QBHintHandler) MarkViewQBNameUsed(qbName string, state *QBHintBuildState) { + if state == nil || state.ViewQBNameUsed == nil { + return + } + state.ViewQBNameUsed[qbName] = struct{}{} } // GenerateQBName builds QBName from offset.