expression: Move more methods from SessionVars to BuildContext (#52440)

ref pingcap/tidb#52366

pkg/expression/contextimpl/sessionctx.go

@@ -17,6 +17,7 @@ package contextimpl
 import (
 	"context"
 	"math"
+	"sync/atomic"
 	"time"
 
 	"github.com/pingcap/tidb/pkg/errctx"
@@ -51,6 +52,8 @@ var _ exprctx.ExprContext = struct {
 type ExprCtxExtendedImpl struct {
 	sctx sessionctx.Context
 	*SessionEvalContext
+	inNullRejectCheck atomic.Bool
+	inUnionCast       atomic.Bool
 }
 
 // NewExprExtendedImpl creates a new ExprCtxExtendedImpl.
@@ -109,6 +112,31 @@ func (ctx *ExprCtxExtendedImpl) SetSkipPlanCache(reason error) {
 	ctx.sctx.GetSessionVars().StmtCtx.SetSkipPlanCache(reason)
 }
 
+// AllocPlanColumnID allocates column id for plan.
+func (ctx *ExprCtxExtendedImpl) AllocPlanColumnID() int64 {
+	return ctx.sctx.GetSessionVars().AllocPlanColumnID()
+}
+
+// SetInNullRejectCheck sets whether the expression is in null reject check.
+func (ctx *ExprCtxExtendedImpl) SetInNullRejectCheck(in bool) {
+	ctx.inNullRejectCheck.Store(in)
+}
+
+// IsInNullRejectCheck returns whether the expression is in null reject check.
+func (ctx *ExprCtxExtendedImpl) IsInNullRejectCheck() bool {
+	return ctx.inNullRejectCheck.Load()
+}
+
+// SetInUnionCast sets the flag to indicate whether the expression is in union cast.
+func (ctx *ExprCtxExtendedImpl) SetInUnionCast(in bool) {
+	ctx.inUnionCast.Store(in)
+}
+
+// IsInUnionCast indicates whether executing in special cast context that negative unsigned num will be zero.
+func (ctx *ExprCtxExtendedImpl) IsInUnionCast() bool {
+	return ctx.inUnionCast.Load()
+}
+
 // GetWindowingUseHighPrecision determines whether to compute window operations without loss of precision.
 // see https://dev.mysql.com/doc/refman/8.0/en/window-function-optimization.html for more details.
 func (ctx *ExprCtxExtendedImpl) GetWindowingUseHighPrecision() bool {

pkg/expression/contextimpl/sessionctx_test.go

@@ -253,6 +253,7 @@ func TestSessionBuildContext(t *testing.T) {
 	require.True(t, evalCtx.GetOptionalPropSet().IsFull())
 	require.Same(t, ctx, evalCtx.Sctx())
 
+	// charset and collation
 	vars := ctx.GetSessionVars()
 	err := vars.SetSystemVar("character_set_connection", "gbk")
 	require.NoError(t, err)
@@ -265,17 +266,45 @@ func TestSessionBuildContext(t *testing.T) {
 	require.Equal(t, "gbk_chinese_ci", collate)
 	require.Equal(t, "utf8mb4_0900_ai_ci", impl.GetDefaultCollationForUTF8MB4())
 
+	// SysdateIsNow
 	vars.SysdateIsNow = true
 	require.True(t, impl.GetSysdateIsNow())
 
+	// NoopFuncsMode
 	vars.NoopFuncsMode = 2
 	require.Equal(t, 2, impl.GetNoopFuncsMode())
 
+	// Rng
 	vars.Rng = mathutil.NewWithSeed(123)
 	require.Same(t, vars.Rng, impl.Rng())
 
+	// PlanCache
 	vars.StmtCtx.UseCache = true
 	require.True(t, impl.IsUseCache())
 	impl.SetSkipPlanCache(errors.New("mockReason"))
 	require.False(t, impl.IsUseCache())
+
+	// Alloc column id
+	prevID := vars.PlanColumnID.Load()
+	colID := impl.AllocPlanColumnID()
+	require.Equal(t, colID, prevID+1)
+	colID = impl.AllocPlanColumnID()
+	require.Equal(t, colID, prevID+2)
+	vars.AllocPlanColumnID()
+	colID = impl.AllocPlanColumnID()
+	require.Equal(t, colID, prevID+4)
+
+	// InNullRejectCheck
+	require.False(t, impl.IsInNullRejectCheck())
+	impl.SetInNullRejectCheck(true)
+	require.True(t, impl.IsInNullRejectCheck())
+	impl.SetInNullRejectCheck(false)
+	require.False(t, impl.IsInNullRejectCheck())
+
+	// InUnionCast
+	require.False(t, impl.IsInUnionCast())
+	impl.SetInUnionCast(true)
+	require.True(t, impl.IsInUnionCast())
+	impl.SetInUnionCast(false)
+	require.False(t, impl.IsInUnionCast())
 }