diff --git a/pkg/planner/optimize.go b/pkg/planner/optimize.go index 4200aa0698e9c..816c5d82deb0d 100644 --- a/pkg/planner/optimize.go +++ b/pkg/planner/optimize.go @@ -438,6 +438,110 @@ var planBuilderPool = sync.Pool{ }, } +type logicalPlanBuildCtx struct { + stmtCtxState stmtctx.LogicalPlanBuildState + plannerSelectBlockAsName *[]ast.HintTable + mapScalarSubQ []any + mapHashCode2UniqueID map[string]int + rewritePhaseInfo variable.RewritePhaseInfo +} + +func saveLogicalPlanBuildCtx(sessVars *variable.SessionVars) logicalPlanBuildCtx { + return logicalPlanBuildCtx{ + stmtCtxState: sessVars.StmtCtx.SaveLogicalPlanBuildState(), + plannerSelectBlockAsName: sessVars.PlannerSelectBlockAsName.Load(), + mapScalarSubQ: sessVars.MapScalarSubQ, + mapHashCode2UniqueID: sessVars.MapHashCode2UniqueID4ExtendedCol, + rewritePhaseInfo: sessVars.RewritePhaseInfo, + } +} + +func restoreLogicalPlanBuildCtx(sessVars *variable.SessionVars, logicalPlanCtx logicalPlanBuildCtx) { + sessVars.StmtCtx.RestoreLogicalPlanBuildState(logicalPlanCtx.stmtCtxState) + sessVars.PlannerSelectBlockAsName.Store(logicalPlanCtx.plannerSelectBlockAsName) + sessVars.MapScalarSubQ = logicalPlanCtx.mapScalarSubQ + sessVars.MapHashCode2UniqueID4ExtendedCol = logicalPlanCtx.mapHashCode2UniqueID + sessVars.RewritePhaseInfo = logicalPlanCtx.rewritePhaseInfo +} + +func buildAndOptimizeLogicalPlanRound( + ctx context.Context, + sctx planctx.PlanContext, + node *resolve.NodeW, + is infoschema.InfoSchema, + hintProcessor *hint.QBHintHandler, + checked *bool, + optimizeStarted *bool, + beginOpt *time.Time, + needRestoreLogicalPlanCtx bool, + bestPlan *base.PhysicalPlan, + bestNames *types.NameSlice, + bestCost *float64, + bestLogicalPlanCtx *logicalPlanBuildCtx, +) (base.Plan, types.NameSlice, bool, error) { + builder := planBuilderPool.Get().(*core.PlanBuilder) + defer planBuilderPool.Put(builder.ResetForReuse()) + // TODO: when buildRound > 1, only emit unused view-hint warnings for the winner build. + defer builder.HandleUnusedViewHints() + + builder.Init(sctx, is, hintProcessor) + + // todo: you can customize each round's special builder (like semi join rewrite or not by signal) + p, err := buildLogicalPlan(ctx, sctx, node, builder) + if err != nil { + return nil, nil, false, err + } + names := p.OutputNames() + + if !*checked { + // Keep privilege and lock checks fail-fast. These depend on visitInfo + // produced by the logical build, but not on the later cost winner. + if pm := privilege.GetPrivilegeManager(sctx); pm != nil { + visitInfo := core.VisitInfo4PrivCheck(ctx, is, node.Node, builder.GetVisitInfo()) + if err := core.CheckPrivilege(sctx.GetSessionVars().ActiveRoles, pm, visitInfo); err != nil { + return nil, nil, false, err + } + } + + if err := core.CheckTableLock(sctx, is, builder.GetVisitInfo()); err != nil { + return nil, nil, false, err + } + + if err := core.CheckTableMode(node); err != nil { + return nil, nil, false, err + } + *checked = true + } + + // Handle the non-logical plan statement. + logic, isLogicalPlan := p.(base.LogicalPlan) + if !isLogicalPlan { + return p, names, true, nil + } + + core.RecheckCTE(logic) + + // todo: also you can customize each round's special logical opt flag here (like decorrelate rule or not) + if !*optimizeStarted { + *optimizeStarted = true + *beginOpt = time.Now() + } + finalPlan, cost, err := core.DoOptimize(ctx, sctx, builder.GetOptFlag(), logic) + if err != nil { + return nil, nil, false, err + } + + if *bestPlan == nil || cost < *bestCost { + *bestCost = cost + *bestPlan = finalPlan + *bestNames = names + if needRestoreLogicalPlanCtx { + *bestLogicalPlanCtx = saveLogicalPlanBuildCtx(sctx.GetSessionVars()) + } + } + return p, names, false, nil +} + // optimizeCnt is a global variable only used for test. var optimizeCnt int @@ -456,54 +560,70 @@ func optimize(ctx context.Context, sctx planctx.PlanContext, node *resolve.NodeW topsql.MockHighCPULoad(sctx.GetSessionVars().StmtCtx.OriginalSQL, sqlPrefixes, 10) }) sessVars := sctx.GetSessionVars() + var ( + beginOpt time.Time + optimizeStarted bool + ) + defer func() { + if optimizeStarted { + sessVars.DurationOptimizer.Total = time.Since(beginOpt) + } + }() - // build logical plan + // Build the logical plan from the raw AST. The hint processor only keeps + // AST-derived metadata; per-build state is allocated inside PlanBuilder. hintProcessor := hint.NewQBHintHandler(sctx.GetSessionVars().StmtCtx) node.Node.Accept(hintProcessor) - 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 + + // build multi logical plan from raw AST. + var ( + buildRound = 1 + needRestoreLogicalPlanCtx = buildRound > 1 + bestCost = math.MaxFloat64 + bestPlan base.PhysicalPlan + bestNames types.NameSlice + bestLogicalPlanCtx logicalPlanBuildCtx + checked bool + ) + var initialLogicalPlanCtx logicalPlanBuildCtx + if needRestoreLogicalPlanCtx { + initialLogicalPlanCtx = saveLogicalPlanBuildCtx(sessVars) } + for i := range buildRound { + if needRestoreLogicalPlanCtx && i > 0 { + restoreLogicalPlanBuildCtx(sessVars, initialLogicalPlanCtx) + } - activeRoles := sessVars.ActiveRoles - // Check privilege. Maybe it's better to move this to the Preprocess, but - // we need the table information to check privilege, which is collected - // into the visitInfo in the logical plan builder. - if pm := privilege.GetPrivilegeManager(sctx); pm != nil { - visitInfo := core.VisitInfo4PrivCheck(ctx, is, node.Node, builder.GetVisitInfo()) - if err := core.CheckPrivilege(activeRoles, pm, visitInfo); err != nil { + p, names, nonLogical, err := buildAndOptimizeLogicalPlanRound( + ctx, + sctx, + node, + is, + hintProcessor, + &checked, + &optimizeStarted, + &beginOpt, + needRestoreLogicalPlanCtx, + &bestPlan, + &bestNames, + &bestCost, + &bestLogicalPlanCtx, + ) + if err != nil { return nil, nil, 0, err } + if nonLogical { + // keep compatible with the old. + return p, names, 0, nil + } } - - if err := core.CheckTableLock(sctx, is, builder.GetVisitInfo()); err != nil { - return nil, nil, 0, err + if bestPlan == nil { + return nil, nil, 0, errors.New("failed to build logical plan") } - - if err := core.CheckTableMode(node); err != nil { - return nil, nil, 0, err + if needRestoreLogicalPlanCtx { + restoreLogicalPlanBuildCtx(sessVars, bestLogicalPlanCtx) } - - names := p.OutputNames() - - // Handle the non-logical plan statement. - logic, isLogicalPlan := p.(base.LogicalPlan) - if !isLogicalPlan { - return p, names, 0, nil - } - - core.RecheckCTE(logic) - - beginOpt := time.Now() - finalPlan, cost, err := core.DoOptimize(ctx, sctx, builder.GetOptFlag(), logic) - // TODO: capture plan replayer here if it matches sql and plan digest - - sessVars.DurationOptimizer.Total = time.Since(beginOpt) - return finalPlan, names, cost, err + return bestPlan, bestNames, bestCost, nil } // OptimizeExecStmt to handle the "execute" statement