planner, expr: eval readonly user var during plan phase | tidb-test=pr/2419

pkg/executor/executor.go

@@ -2085,6 +2085,7 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	vars.DurationWaitTS = 0
 	vars.CurrInsertBatchExtraCols = nil
 	vars.CurrInsertValues = chunk.Row{}
+	ctx.GetExprCtx().GetEvalCtx().ResetReadonlyVarMap()
 
 	return
 }

pkg/expression/builtin_other.go

@@ -950,11 +950,12 @@ func BuildGetVarFunction(ctx BuildContext, expr Expression, retType *types.Field
 	if builtinRetTp := f.getRetTp(); builtinRetTp.GetType() != mysql.TypeUnspecified || retType.GetType() == mysql.TypeUnspecified {
 		retType = builtinRetTp
 	}
-	return &ScalarFunction{
+	sf := &ScalarFunction{
 		FuncName: model.NewCIStr(ast.GetVar),
 		RetType:  retType,
 		Function: f,
-	}, nil
+	}
+	return FoldConstant(ctx, sf), nil
 }
 
 type getVarFunctionClass struct {

pkg/expression/constant_fold.go

@@ -15,6 +15,7 @@
 package expression
 
 import (
+	"github.com/pingcap/tidb/pkg/expression/contextopt"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/types"
@@ -33,6 +34,7 @@ func init() {
 		ast.Ifnull: ifNullFoldHandler,
 		ast.Case:   caseWhenHandler,
 		ast.IsNull: isNullHandler,
+		ast.GetVar: getVarHandler,
 	}
 }
 
@@ -159,6 +161,30 @@ func caseWhenHandler(ctx BuildContext, expr *ScalarFunction) (Expression, bool)
 	return expr, isDeferredConst
 }
 
+func getVarHandler(ctx BuildContext, expr *ScalarFunction) (Expression, bool) {
+	arg0 := expr.GetArgs()[0]
+	c, ok := arg0.(*Constant)
+	if !ok {
+		return expr, false
+	}
+	if !ctx.GetEvalCtx().IsReadonlyVar(c.Value.GetString()) {
+		return expr, false
+	}
+	v, err := expr.Eval(ctx.GetEvalCtx(), chunk.Row{})
+	if err != nil {
+		return expr, false
+	}
+	sessVars, err := contextopt.SessionVarsPropReader{}.GetSessionVars(ctx.GetEvalCtx())
+	if err != nil {
+		return expr, false
+	}
+	d, ok := sessVars.GetUserVarVal(c.Value.GetString())
+	if ok && d.Kind() == types.KindBinaryLiteral {
+		v.SetBinaryLiteral(v.GetBytes())
+	}
+	return &Constant{Value: v, RetType: expr.RetType, DeferredExpr: expr}, false
+}
+
 func foldConstant(ctx BuildContext, expr Expression) (Expression, bool) {
 	switch x := expr.(type) {
 	case *ScalarFunction:

pkg/expression/context/context.go

@@ -83,6 +83,10 @@ type EvalContext interface {
 	GetOptionalPropSet() OptionalEvalPropKeySet
 	// GetOptionalPropProvider gets the optional property provider by key
 	GetOptionalPropProvider(OptionalEvalPropKey) (OptionalEvalPropProvider, bool)
+
+	ResetReadonlyVarMap()
+	SetReadonlyVarMap(map[string]struct{})
+	IsReadonlyVar(name string) bool
 }
 
 // BuildContext is used to build an expression

pkg/expression/contextsession/sessionctx.go

@@ -138,8 +138,9 @@ func (ctx *SessionExprContext) ConnectionID() uint64 {
 
 // SessionEvalContext implements the `expression.EvalContext` interface to provide evaluation context in session.
 type SessionEvalContext struct {
-	sctx  sessionctx.Context
-	props contextopt.OptionalEvalPropProviders
+	sctx         sessionctx.Context
+	props        contextopt.OptionalEvalPropProviders
+	readonlyVars map[string]struct{}
 }
 
 // NewSessionEvalContext creates a new SessionEvalContext.
@@ -156,9 +157,29 @@ func NewSessionEvalContext(sctx sessionctx.Context) *SessionEvalContext {
 	ctx.setOptionalProp(contextopt.DDLOwnerInfoProvider(sctx.IsDDLOwner))
 	// When EvalContext is created from a session, it should contain all the optional properties.
 	intest.Assert(ctx.props.PropKeySet().IsFull())
+	ctx.readonlyVars = make(map[string]struct{})
 	return ctx
 }
 
+// ResetReadonlyVarMap resets the readonly vars map.
+func (ctx *SessionEvalContext) ResetReadonlyVarMap() {
+	ctx.readonlyVars = nil
+}
+
+// SetReadonlyVarMap sets the readonly vars map.
+func (ctx *SessionEvalContext) SetReadonlyVarMap(vars map[string]struct{}) {
+	ctx.readonlyVars = vars
+}
+
+// IsReadonlyVar checks whether the variable is readonly.
+func (ctx *SessionEvalContext) IsReadonlyVar(name string) bool {
+	if ctx.readonlyVars == nil {
+		return false
+	}
+	_, ok := ctx.readonlyVars[name]
+	return ok
+}
+
 func (ctx *SessionEvalContext) setOptionalProp(prop exprctx.OptionalEvalPropProvider) {
 	intest.AssertFunc(func() bool {
 		return !ctx.props.Contains(prop.Desc().Key())

pkg/expression/contextstatic/evalctx.go

@@ -218,6 +218,19 @@ type StaticEvalContext struct {
 	staticEvalCtxState
 }
 
+// ResetReadonlyVarMap implements context.EvalContext.
+func (ctx *StaticEvalContext) ResetReadonlyVarMap() {
+}
+
+// SetReadonlyVarMap implements context.EvalContext.
+func (ctx *StaticEvalContext) SetReadonlyVarMap(map[string]struct{}) {
+}
+
+// IsReadonlyVar implements context.EvalContext.
+func (ctx *StaticEvalContext) IsReadonlyVar(name string) bool {
+	return false
+}
+
 // NewStaticEvalContext creates a new `StaticEvalContext` with the given options.
 func NewStaticEvalContext(opt ...StaticEvalCtxOption) *StaticEvalContext {
 	ctx := &StaticEvalContext{

pkg/expression/function_traits.go

@@ -55,7 +55,6 @@ var unFoldableFunctions = map[string]struct{}{
 	ast.RowFunc:   {},
 	ast.Values:    {},
 	ast.SetVar:    {},
-	ast.GetVar:    {},
 	ast.GetParam:  {},
 	ast.Benchmark: {},
 	ast.DayName:   {},
@@ -208,6 +207,7 @@ var AllowedPartition4UnaryOpMap = map[opcode.Op]struct{}{
 // inequalFunctions stores functions which cannot be propagated from column equal condition.
 var inequalFunctions = map[string]struct{}{
 	ast.IsNull: {},
+	ast.GetVar: {},
 }
 
 // mutableEffectsFunctions stores functions which are mutable or have side effects, specifically,

pkg/expression/util.go

@@ -1307,6 +1307,10 @@ func IsRuntimeConstExpr(expr Expression) bool {
 		if _, ok := unFoldableFunctions[x.FuncName.L]; ok {
 			return false
 		}
+		// If the GetVar is not folded, it's not readonly and can't be treated as a constant.
+		if x.FuncName.L == ast.GetVar {
+			return false
+		}
 		for _, arg := range x.GetArgs() {
 			if !IsRuntimeConstExpr(arg) {
 				return false

pkg/planner/core/plan_cache.go

@@ -57,16 +57,24 @@ func SetParameterValuesIntoSCtx(sctx base.PlanContext, isNonPrep bool, markers [
 	vars := sctx.GetSessionVars()
 	vars.PlanCacheParams.Reset()
 	for i, usingParam := range params {
-		val, err := usingParam.Eval(sctx.GetExprCtx().GetEvalCtx(), chunk.Row{})
-		if err != nil {
-			return err
-		}
-		if isGetVarBinaryLiteral(sctx, usingParam) {
-			binVal, convErr := val.ToBytes()
-			if convErr != nil {
-				return convErr
+		var (
+			val types.Datum
+			err error
+		)
+		if c, ok := usingParam.(*expression.Constant); ok {
+			val = c.Value
+		} else {
+			val, err = usingParam.Eval(sctx.GetExprCtx().GetEvalCtx(), chunk.Row{})
+			if err != nil {
+				return err
+			}
+			if isGetVarBinaryLiteral(sctx, usingParam) {
+				binVal, convErr := val.ToBytes()
+				if convErr != nil {
+					return convErr
+				}
+				val.SetBinaryLiteral(binVal)
 			}
-			val.SetBinaryLiteral(binVal)
 		}
 		if markers != nil {
 			param := markers[i].(*driver.ParamMarkerExpr)

pkg/planner/core/preprocess.go

@@ -127,6 +127,8 @@ func Preprocess(ctx context.Context, sctx sessionctx.Context, node ast.Node, pre
 		tableAliasInJoin:   make([]map[string]any, 0),
 		preprocessWith:     &preprocessWith{cteCanUsed: make([]string, 0), cteBeforeOffset: make([]int, 0)},
 		staleReadProcessor: staleread.NewStaleReadProcessor(ctx, sctx),
+		varsChanged:        make(map[string]struct{}),
+		varsReadonly:       make(map[string]struct{}),
 	}
 	for _, optFn := range preprocessOpt {
 		optFn(&v)
@@ -138,6 +140,10 @@ func Preprocess(ctx context.Context, sctx sessionctx.Context, node ast.Node, pre
 	node.Accept(&v)
 	// InfoSchema must be non-nil after preprocessing
 	v.ensureInfoSchema()
+	sctx.GetExprCtx().GetEvalCtx().SetReadonlyVarMap(v.varsReadonly)
+	if len(v.varsReadonly) > 0 {
+		sctx.GetSessionVars().StmtCtx.SetSkipPlanCache("read-only variables are used")
+	}
 	return errors.Trace(v.err)
 }
 
@@ -229,6 +235,9 @@ type preprocessor struct {
 
 	staleReadProcessor staleread.Processor
 
+	varsChanged  map[string]struct{}
+	varsReadonly map[string]struct{}
+
 	// values that may be returned
 	*PreprocessorReturn
 	err error
@@ -410,6 +419,16 @@ func (p *preprocessor) Enter(in ast.Node) (out ast.Node, skipChildren bool) {
 		}
 	case *ast.AnalyzeTableStmt:
 		p.flag |= inAnalyze
+	case *ast.VariableExpr:
+		if node.Value != nil {
+			p.varsChanged[node.Name] = struct{}{}
+			delete(p.varsReadonly, node.Name)
+		} else {
+			_, ok := p.varsChanged[node.Name]
+			if !ok {
+				p.varsReadonly[node.Name] = struct{}{}
+			}
+		}
 	default:
 		p.flag &= ^parentIsJoin
 	}

tests/integrationtest/r/index_merge.result

@@ -781,11 +781,11 @@ set @a = 1;
 explain select /*+ use_index_merge(t1) */ * from t1 where (c1 < 10 or c2 < 10) and length(substring(sqrt(c3), @a, 1)) = char_length(if(c1, c2, c3)) order by 1;
 id	estRows	task	access object	operator info
 Sort_5	4433.77	root		index_merge.t1.c1
-└─Selection_8	4433.77	root		eq(length(substring(cast(sqrt(cast(index_merge.t1.c3, double BINARY)), var_string(5)), getvar("a"), 1)), char_length(cast(if(index_merge.t1.c1, index_merge.t1.c2, index_merge.t1.c3), var_string(20))))
-  └─IndexMerge_12	5542.21	root		type: union
-    ├─IndexRangeScan_9(Build)	3323.33	cop[tikv]	table:t1, index:c1(c1)	range:[-inf,10), keep order:false, stats:pseudo
-    ├─IndexRangeScan_10(Build)	3323.33	cop[tikv]	table:t1, index:c2(c2)	range:[-inf,10), keep order:false, stats:pseudo
-    └─TableRowIDScan_11(Probe)	5542.21	cop[tikv]	table:t1	keep order:false, stats:pseudo
+└─IndexMerge_12	4433.77	root		type: union
+  ├─IndexRangeScan_8(Build)	3323.33	cop[tikv]	table:t1, index:c1(c1)	range:[-inf,10), keep order:false, stats:pseudo
+  ├─IndexRangeScan_9(Build)	3323.33	cop[tikv]	table:t1, index:c2(c2)	range:[-inf,10), keep order:false, stats:pseudo
+  └─Selection_11(Probe)	4433.77	cop[tikv]		eq(length(substring(cast(sqrt(cast(index_merge.t1.c3, double BINARY)), var_string(5)), 1, 1)), char_length(cast(if(index_merge.t1.c1, index_merge.t1.c2, index_merge.t1.c3), var_string(20))))
+    └─TableRowIDScan_10	5542.21	cop[tikv]	table:t1	keep order:false, stats:pseudo
 select /*+ use_index_merge(t1) */ * from t1 where (c1 < 10 or c2 < 10) and length(substring(sqrt(c3), @a, 1)) = char_length(if(c1, c2, c3)) order by 1;
 c1	c2	c3	c4	c5
 1	1	1	1	1