Skip to content

planner: build multi alternative logical plan from shared AST (#66677) #68859

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ti-chi-bot wants to merge 1 commit into from
+257 −19
Closed

planner: build multi alternative logical plan from shared AST (#66677) #68859

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions pkg/planner/core/logical_plan_builder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
19 changes: 19 additions & 0 deletions pkg/planner/core/planbuilder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pkg/planner/optimize.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions pkg/sessionctx/stmtctx/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Comment on lines +48 to +52
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Jun 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Remove the leftover conflict markers from the Bazel rule.

Line 48 still contains merge-conflict markers, so Bazel cannot parse this file. Keep only the intended shard_count entry.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/sessionctx/stmtctx/BUILD.bazel` around lines 48 - 52, Remove the leftover
merge conflict markers in BUILD.bazel around the shard_count attribute (delete
the <<<<<<<, =======, and >>>>>>> lines) and leave a single shard_count =
<correct_value> entry (e.g., shard_count = 14 or shard_count = 17 as intended)
so the Bazel rule parses cleanly; ensure only one shard_count line remains in
the rule.

deps = [
"//pkg/errctx",
"//pkg/kv",
"//pkg/meta/model",
"//pkg/parser/ast",
"//pkg/sessionctx/variable",
"//pkg/testkit",
"//pkg/testkit/testfailpoint",
Expand Down
59 changes: 59 additions & 0 deletions pkg/sessionctx/stmtctx/stmtctx.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}
Comment on lines +68 to 93
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Jun 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Resolve the cherry-pick conflict before merge.

Line 68 still starts a <<<<<<< / ======= / >>>>>>> block, so this file does not parse. It also leaves jsonSQLWarn and LogicalPlanBuildState in an unresolved state.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/sessionctx/stmtctx/stmtctx.go` around lines 68 - 93, There is an
unresolved git conflict block left in stmtctx.go; remove the conflict markers
(<<<<<<<, =======, >>>>>>>) and reconcile the two definitions so the file
parses: either restore the jsonSQLWarn type if needed or keep
LogicalPlanBuildState (or combine them appropriately), ensuring only valid Go
types remain and references to jsonSQLWarn and LogicalPlanBuildState are
correctly defined and compiled; update imports if required and run go build to
verify the file parses.


// ReferenceCount indicates the reference count of StmtCtx.
Expand Down Expand Up @@ -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(),
Comment on lines +607 to +609
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Jun 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

ViewDepth is snapshotted/restored but not declared on StatementContext.

Lines 608 and 628 reference sc.ViewDepth, but there is no ViewDepth field on StatementContext anywhere in this file. This cherry-pick will still fail to compile after the conflict markers are fixed unless that field addition is brought over too.

Also applies to: 627-629

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/sessionctx/stmtctx/stmtctx.go` around lines 607 - 609, The code
references sc.ViewDepth but StatementContext lacks a ViewDepth field; add a
ViewDepth field to the StatementContext struct (matching the original type used
elsewhere in the project), and update any constructors/copy/snapshot/restore
logic so ViewDepth is properly preserved and copied (ensure sc.ViewDepth is
set/returned where StatementContext is initialized and that the
snapshot/Restore/Copy methods include ViewDepth alongside useDynamicPruneMode,
viewDepth, and colRefFromUpdatePlan). Use the existing symbols StatementContext,
sc.ViewDepth, Copy(), snapshot/restore (or the struct's constructor functions)
to locate where to add and wire this field.

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
Expand Down
86 changes: 86 additions & 0 deletions pkg/sessionctx/stmtctx/stmtctx_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions pkg/util/context/plancache.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading