expression, planner: push cast down to control function with enum type.

expression/builtin_cast.go

@@ -1804,6 +1804,12 @@ func BuildCastFunction4Union(ctx sessionctx.Context, expr Expression, tp *types.
 
 // BuildCastFunction builds a CAST ScalarFunction from the Expression.
 func BuildCastFunction(ctx sessionctx.Context, expr Expression, tp *types.FieldType) (res Expression) {
+	if res, success := TryPushCastDownToControlFunctionForHybridType(ctx, expr, tp); success {
+		if tp.EvalType() != types.ETJson {
+			res = FoldConstant(res)
+		}
+		return res
+	}
 	var fc functionClass
 	switch tp.EvalType() {
 	case types.ETInt:
@@ -1983,3 +1989,85 @@ func WrapWithCastAsJSON(ctx sessionctx.Context, expr Expression) Expression {
 	}
 	return BuildCastFunction(ctx, expr, tp)
 }
+
+// TryPushCastDownToControlFunctionForHybridType try to push cast down to control function for Hybrid Type.
+// If necessary, it will rebuild control function using changed args.
+func TryPushCastDownToControlFunctionForHybridType(ctx sessionctx.Context, expr Expression, tp *types.FieldType) (res Expression, success bool) {
+	sf, ok := expr.(*ScalarFunction)
+	if !ok {
+		return expr, false
+	}
+
+	var wrapCastFunc func(ctx sessionctx.Context, expr Expression) Expression
+	switch tp.EvalType() {
+	case types.ETInt:
+		wrapCastFunc = WrapWithCastAsInt
+	case types.ETReal:
+		wrapCastFunc = WrapWithCastAsReal
+	case types.ETString:
+		wrapCastFunc = WrapWithCastAsString
+	default:
+		return expr, false
+	}
+
+	args := sf.GetArgs()
+	switch sf.FuncName.L {
+	case ast.If:
+		if args[1].GetType().Hybrid() || args[2].GetType().Hybrid() {
+			args[1] = wrapCastFunc(ctx, args[1])
+			args[2] = wrapCastFunc(ctx, args[2])
+			f, err := funcs[ast.If].getFunction(ctx, args)
+			if err != nil {
+				return expr, false
+			}
+			sf.RetType, sf.Function = f.getRetTp(), f
+			return sf, true
+		}
+	case ast.Case:
+		hasHybrid := false
+		for i := 0; i < len(args)-1; i += 2 {
+			hasHybrid = hasHybrid || args[i+1].GetType().Hybrid()
+		}
+		if len(args)%2 == 1 {
+			hasHybrid = hasHybrid || args[len(args)-1].GetType().Hybrid()
+		}
+		if !hasHybrid {
+			return expr, false
+		}
+
+		for i := 0; i < len(args)-1; i += 2 {
+			args[i+1] = wrapCastFunc(ctx, args[i+1])
+		}
+		if len(args)%2 == 1 {
+			args[len(args)-1] = wrapCastFunc(ctx, args[len(args)-1])
+		}
+		f, err := funcs[ast.Case].getFunction(ctx, args)
+		if err != nil {
+			return expr, false
+		}
+		sf.RetType, sf.Function = f.getRetTp(), f
+		return sf, true
+	case ast.Elt:
+		hasHybrid := false
+		for i := 1; i < len(args); i++ {
+			hasHybrid = hasHybrid || args[i].GetType().Hybrid()
+		}
+		if !hasHybrid {
+			return expr, false
+		}
+
+		for i := 1; i < len(args); i++ {
+			args[i] = wrapCastFunc(ctx, args[i])
+		}
+		f, err := funcs[ast.Elt].getFunction(ctx, args)
+		if err != nil {
+			return expr, false
+		}
+		sf.RetType, sf.Function = f.getRetTp(), f
+		// Elt only supports ETString, so we need add extra cast to keep retType right.
+		return wrapCastFunc(ctx, sf), true
+	default:
+		return expr, false
+	}
+	return expr, false
+}

expression/integration_test.go

@@ -9336,3 +9336,45 @@ func (s *testIntegrationSuite) TestEnumIndex(c *C) {
 			"OFJHCEKCQGT:MXI7P3[YO4N0DF=2XJWJ4Z9Z;HQ8TMUTZV8YLQAHWJ4BDZHR3A -30 <nil>",
 			"ZOHBSCRMZPOI`IVTSEZAIDAF7DS@1TT20AP9 -30 <nil>"))
 }
+
+func (s *testIntegrationSuite) TestEnumControlFunction(c *C) {
+	defer s.cleanEnv(c)
+
+	// issue 23114
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists e;")
+	tk.MustExec("create table e(e enum('c', 'b', 'a'));")
+	tk.MustExec("insert into e values ('a'),('b'),('a'),('b');")
+	tk.MustQuery("select e from e where if(e>1, e, e);").Sort().Check(
+		testkit.Rows("a", "a", "b", "b"))
+	tk.MustQuery("select e from e where case e when 1 then e else e end;").Sort().Check(
+		testkit.Rows("a", "a", "b", "b"))
+	tk.MustQuery("select e from e where case 1 when e then e end;").Check(testkit.Rows())
+
+	tk.MustQuery("select if(e>1,e,e)='a' from e").Sort().Check(
+		testkit.Rows("0", "0", "1", "1"))
+	tk.MustQuery("select if(e>1,e,e)=1 from e").Sort().Check(
+		testkit.Rows("0", "0", "0", "0"))
+
+	// issue 24494
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t(a int,b enum(\"b\",\"y\",\"1\"));")
+	tk.MustExec("insert into t values(0,\"y\"),(1,\"b\"),(null,null),(2,\"1\");")
+	tk.MustQuery("SELECT count(*) FROM t where if(a,b ,null);").Check(testkit.Rows("2"))
+
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t(a int,b enum(\"b\"),c enum(\"c\"));")
+	tk.MustExec("insert into t values(1,1,1),(2,1,1),(1,1,1),(2,1,1);")
+	tk.MustQuery("select a from t where if(a=1,b,c)=\"b\";").Check(testkit.Rows("1", "1"))
+	tk.MustQuery("select a from t where if(a=1,b,c)=\"c\";").Check(testkit.Rows("2", "2"))
+	tk.MustQuery("select a from t where if(a=1,b,c)=1;").Sort().Check(testkit.Rows("1", "1", "2", "2"))
+	tk.MustQuery("select a from t where if(a=1,b,c);").Sort().Check(testkit.Rows("1", "1", "2", "2"))
+
+	tk.MustExec("drop table if exists e;")
+	tk.MustExec("create table e(e enum('c', 'b', 'a'));")
+	tk.MustExec("insert into e values(3)")
+	tk.MustQuery("select elt(1,e) = 'a' from e").Check(testkit.Rows("1"))
+	tk.MustQuery("select elt(1,e) = 3 from e").Check(testkit.Rows("1"))
+	tk.MustQuery("select e from e where elt(1,e)").Check(testkit.Rows("a"))
+}

planner/core/logical_plan_builder.go

@@ -975,6 +975,18 @@ func (b *PlanBuilder) buildSelection(ctx context.Context, p LogicalPlan, where a
 	if len(cnfExpres) == 0 {
 		return p, nil
 	}
+	// check expr field types.
+	for i, expr := range cnfExpres {
+		if expr.GetType().EvalType() == types.ETString {
+			tp := types.NewFieldType(mysql.TypeDouble)
+			tp.Flen, tp.Decimal = mysql.MaxRealWidth, types.UnspecifiedLength
+			types.SetBinChsClnFlag(tp)
+			tp.Flag |= expr.GetType().Flag & mysql.UnsignedFlag
+			if res, ok := expression.TryPushCastDownToControlFunctionForHybridType(b.ctx, expr, tp); ok {
+				cnfExpres[i] = res
+			}
+		}
+	}
 	selection.Conditions = cnfExpres
 	selection.SetChildren(p)
 	return selection, nil