cherry pick pingcap#19331 to release-3.0

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>

expression/expression.go

@@ -387,6 +387,355 @@ func IsBinaryLiteral(expr Expression) bool {
 	return ok && con.Value.Kind() == types.KindBinaryLiteral
 }
 
+<<<<<<< HEAD
+=======
+func canFuncBePushed(sf *ScalarFunction, storeType kv.StoreType) bool {
+	// Use the failpoint to control whether to push down an expression in the integration test.
+	// Push down all expression if the `failpoint expression` is `all`, otherwise, check
+	// whether scalar function's name is contained in the enabled expression list (e.g.`ne,eq,lt`).
+	// If neither of the above is true, switch to original logic.
+	failpoint.Inject("PushDownTestSwitcher", func(val failpoint.Value) {
+		enabled := val.(string)
+		if enabled == "all" {
+			failpoint.Return(true)
+		}
+		exprs := strings.Split(enabled, ",")
+		for _, expr := range exprs {
+			if strings.ToLower(strings.TrimSpace(expr)) == sf.FuncName.L {
+				failpoint.Return(true)
+			}
+		}
+	})
+
+	ret := false
+	switch sf.FuncName.L {
+	case
+		// op functions.
+		ast.LogicAnd,
+		ast.LogicOr,
+		ast.LogicXor,
+		ast.UnaryNot,
+		ast.And,
+		ast.Or,
+		ast.Xor,
+		ast.BitNeg,
+		ast.LeftShift,
+		ast.RightShift,
+		ast.UnaryMinus,
+
+		// compare functions.
+		ast.LT,
+		ast.LE,
+		ast.EQ,
+		ast.NE,
+		ast.GE,
+		ast.GT,
+		ast.NullEQ,
+		ast.In,
+		ast.IsNull,
+		ast.Like,
+		ast.IsTruth,
+		ast.IsFalsity,
+
+		// arithmetical functions.
+		ast.Plus,
+		ast.Minus,
+		ast.Mul,
+		ast.Div,
+		ast.Abs,
+
+		// math functions.
+		ast.Ceil,
+		ast.Ceiling,
+		ast.Floor,
+		ast.Sqrt,
+		ast.Sign,
+		ast.Ln,
+		ast.Log,
+		ast.Log2,
+		ast.Log10,
+		ast.Exp,
+		ast.Pow,
+		// Rust use the llvm math functions, which have different precision with Golang/MySQL(cmath)
+		// open the following switchers if we implement them in coprocessor via `cmath`
+		// ast.Sin,
+		// ast.Asin,
+		// ast.Cos,
+		// ast.Acos,
+		// ast.Tan,
+		// ast.Atan,
+		// ast.Atan2,
+		// ast.Cot,
+		ast.Radians,
+		ast.Degrees,
+		ast.Conv,
+		ast.CRC32,
+
+		// control flow functions.
+		ast.Case,
+		ast.If,
+		ast.Ifnull,
+		ast.Coalesce,
+
+		// string functions.
+		ast.Length,
+		ast.BitLength,
+		ast.Concat,
+		ast.ConcatWS,
+		// ast.Locate,
+		ast.Replace,
+		ast.ASCII,
+		ast.Hex,
+		ast.Reverse,
+		ast.LTrim,
+		ast.RTrim,
+		// ast.Left,
+		ast.Strcmp,
+		ast.Space,
+		ast.Elt,
+		ast.Field,
+
+		// json functions.
+		ast.JSONType,
+		ast.JSONExtract,
+		// FIXME: JSONUnquote is incompatible with Coprocessor
+		// ast.JSONUnquote,
+		ast.JSONObject,
+		ast.JSONArray,
+		ast.JSONMerge,
+		ast.JSONSet,
+		ast.JSONInsert,
+		// ast.JSONReplace,
+		ast.JSONRemove,
+		ast.JSONLength,
+
+		// date functions.
+		ast.DateFormat,
+		ast.FromDays,
+		// ast.ToDays,
+		ast.DayOfYear,
+		ast.DayOfMonth,
+		ast.Year,
+		ast.Month,
+		// FIXME: the coprocessor cannot keep the same behavior with TiDB in current compute framework
+		// ast.Hour,
+		// ast.Minute,
+		// ast.Second,
+		// ast.MicroSecond,
+		// ast.DayName,
+		ast.PeriodAdd,
+		ast.PeriodDiff,
+		ast.TimestampDiff,
+		ast.DateAdd,
+		ast.FromUnixTime,
+
+		// encryption functions.
+		ast.MD5,
+		ast.SHA1,
+		ast.UncompressedLength,
+
+		ast.Cast,
+
+		// misc functions.
+		ast.InetNtoa,
+		ast.InetAton,
+		ast.Inet6Ntoa,
+		ast.Inet6Aton,
+		ast.IsIPv4,
+		ast.IsIPv4Compat,
+		ast.IsIPv4Mapped,
+		ast.IsIPv6:
+		ret = true
+
+	// A special case: Only push down Round by signature
+	case ast.Round:
+		switch sf.Function.PbCode() {
+		case
+			tipb.ScalarFuncSig_RoundReal,
+			tipb.ScalarFuncSig_RoundInt,
+			tipb.ScalarFuncSig_RoundDec:
+			ret = true
+		}
+	case
+		ast.Substring,
+		ast.Substr:
+		switch sf.Function.PbCode() {
+		case
+			tipb.ScalarFuncSig_Substring2ArgsUTF8,
+			tipb.ScalarFuncSig_Substring3ArgsUTF8:
+			ret = true
+		}
+	case ast.Rand:
+		switch sf.Function.PbCode() {
+		case
+			tipb.ScalarFuncSig_RandWithSeedFirstGen:
+			ret = true
+		}
+	}
+	if ret {
+		switch storeType {
+		case kv.TiFlash:
+			ret = scalarExprSupportedByFlash(sf)
+		case kv.TiKV:
+			ret = scalarExprSupportedByTiKV(sf)
+		case kv.TiDB:
+			ret = scalarExprSupportedByTiDB(sf)
+		}
+	}
+	if ret {
+		ret = IsPushDownEnabled(sf.FuncName.L, storeType)
+	}
+	return ret
+}
+
+func storeTypeMask(storeType kv.StoreType) uint32 {
+	if storeType == kv.UnSpecified {
+		return 1<<kv.TiKV | 1<<kv.TiFlash | 1<<kv.TiDB
+	}
+	return 1 << storeType
+}
+
+// IsPushDownEnabled returns true if the input expr is not in the expr_pushdown_blacklist
+func IsPushDownEnabled(name string, storeType kv.StoreType) bool {
+	value, exists := DefaultExprPushDownBlacklist.Load().(map[string]uint32)[name]
+	if exists {
+		mask := storeTypeMask(storeType)
+		return !(value&mask == mask)
+	}
+
+	if storeType != kv.TiFlash && name == ast.AggFuncApproxCountDistinct {
+		// Can not push down approx_count_distinct to other store except tiflash by now.
+		return false
+	}
+
+	return true
+}
+
+// DefaultExprPushDownBlacklist indicates the expressions which can not be pushed down to TiKV.
+var DefaultExprPushDownBlacklist *atomic.Value
+
+func init() {
+	DefaultExprPushDownBlacklist = new(atomic.Value)
+	DefaultExprPushDownBlacklist.Store(make(map[string]uint32))
+}
+
+func canScalarFuncPushDown(scalarFunc *ScalarFunction, pc PbConverter, storeType kv.StoreType) bool {
+	pbCode := scalarFunc.Function.PbCode()
+	if pbCode <= tipb.ScalarFuncSig_Unspecified {
+		failpoint.Inject("PanicIfPbCodeUnspecified", func() {
+			panic(errors.Errorf("unspecified PbCode: %T", scalarFunc.Function))
+		})
+		return false
+	}
+
+	// Check whether this function can be pushed.
+	if !canFuncBePushed(scalarFunc, storeType) {
+		return false
+	}
+
+	// Check whether all of its parameters can be pushed.
+	for _, arg := range scalarFunc.GetArgs() {
+		if !canExprPushDown(arg, pc, storeType) {
+			return false
+		}
+	}
+
+	if metadata := scalarFunc.Function.metadata(); metadata != nil {
+		var err error
+		_, err = proto.Marshal(metadata)
+		if err != nil {
+			logutil.BgLogger().Error("encode metadata", zap.Any("metadata", metadata), zap.Error(err))
+			return false
+		}
+	}
+	return true
+}
+
+func canExprPushDown(expr Expression, pc PbConverter, storeType kv.StoreType) bool {
+	if storeType == kv.TiFlash && expr.GetType().Tp == mysql.TypeDuration {
+		return false
+	}
+	switch x := expr.(type) {
+	case *CorrelatedColumn:
+		return pc.conOrCorColToPBExpr(expr) != nil && pc.columnToPBExpr(&x.Column) != nil
+	case *Constant:
+		return pc.conOrCorColToPBExpr(expr) != nil
+	case *Column:
+		return pc.columnToPBExpr(x) != nil
+	case *ScalarFunction:
+		return canScalarFuncPushDown(x, pc, storeType)
+	}
+	return false
+}
+
+// PushDownExprs split the input exprs into pushed and remained, pushed include all the exprs that can be pushed down
+func PushDownExprs(sc *stmtctx.StatementContext, exprs []Expression, client kv.Client, storeType kv.StoreType) (pushed []Expression, remained []Expression) {
+	pc := PbConverter{sc: sc, client: client}
+	for _, expr := range exprs {
+		if canExprPushDown(expr, pc, storeType) {
+			pushed = append(pushed, expr)
+		} else {
+			remained = append(remained, expr)
+		}
+	}
+	return
+}
+
+// CanExprsPushDown return true if all the expr in exprs can be pushed down
+func CanExprsPushDown(sc *stmtctx.StatementContext, exprs []Expression, client kv.Client, storeType kv.StoreType) bool {
+	_, remained := PushDownExprs(sc, exprs, client, storeType)
+	return len(remained) == 0
+}
+
+func scalarExprSupportedByTiKV(function *ScalarFunction) bool {
+	switch function.FuncName.L {
+	case ast.Substr, ast.Substring, ast.DateAdd, ast.TimestampDiff,
+		ast.FromUnixTime:
+		return false
+	default:
+		return true
+	}
+}
+
+func scalarExprSupportedByTiDB(function *ScalarFunction) bool {
+	switch function.FuncName.L {
+	case ast.Substr, ast.Substring, ast.DateAdd, ast.TimestampDiff,
+		ast.FromUnixTime:
+		return false
+	default:
+		return true
+	}
+}
+
+func scalarExprSupportedByFlash(function *ScalarFunction) bool {
+	switch function.FuncName.L {
+	case ast.Plus, ast.Minus, ast.Div, ast.Mul, ast.GE, ast.LE,
+		ast.EQ, ast.NE, ast.LT, ast.GT, ast.Ifnull, ast.IsNull,
+		ast.Or, ast.In, ast.Mod, ast.And, ast.LogicOr, ast.LogicAnd,
+		ast.Like, ast.UnaryNot, ast.Case, ast.Month, ast.Substr,
+		ast.Substring, ast.TimestampDiff, ast.DateFormat, ast.FromUnixTime,
+		ast.JSONLength, ast.If, ast.BitNeg, ast.Xor:
+		return true
+	case ast.Cast:
+		switch function.Function.PbCode() {
+		case tipb.ScalarFuncSig_CastIntAsDecimal:
+			return true
+		default:
+			return false
+		}
+	case ast.DateAdd:
+		switch function.Function.PbCode() {
+		case tipb.ScalarFuncSig_AddDateDatetimeInt, tipb.ScalarFuncSig_AddDateStringInt:
+			return true
+		default:
+			return false
+		}
+	default:
+		return false
+	}
+}
+
+>>>>>>> 853e0b4... expression: fix errors in Apply when the type of correlation column is `bit` (#19331)
 // wrapWithIsTrue wraps `arg` with istrue function if the return type of expr is not
 // type int, otherwise, returns `arg` directly.
 // The `keepNull` controls what the istrue function will return when `arg` is null: