expression: reuse tipb.Expr 'val' to push down expression implicit parameters (#10790)

Author: lonng

expression/builtin.go

@@ -41,6 +41,15 @@ func (b *baseBuiltinFunc) PbCode() tipb.ScalarFuncSig {
 	return b.pbCode
 }
 
+// implicitArgs returns the implicit arguments of this function.
+// implicit arguments means some functions contain extra inner fields which will not
+// contain in `tipb.Expr.children` but must be pushed down to coprocessor
+func (b *baseBuiltinFunc) implicitArgs() []types.Datum {
+	// We will not use a field to store them because of only
+	// a few functions contain implicit parameters
+	return nil
+}
+
 func (b *baseBuiltinFunc) setPbCode(c tipb.ScalarFuncSig) {
 	b.pbCode = c
 }
@@ -244,6 +253,17 @@ type baseBuiltinCastFunc struct {
 	inUnion bool
 }
 
+// implicitArgs returns the implicit arguments of cast functions
+func (b *baseBuiltinCastFunc) implicitArgs() []types.Datum {
+	args := b.baseBuiltinFunc.implicitArgs()
+	if b.inUnion {
+		args = append(args, types.NewIntDatum(1))
+	} else {
+		args = append(args, types.NewIntDatum(0))
+	}
+	return args
+}
+
 func (b *baseBuiltinCastFunc) cloneFrom(from *baseBuiltinCastFunc) {
 	b.baseBuiltinFunc.cloneFrom(&from.baseBuiltinFunc)
 	b.inUnion = from.inUnion
@@ -284,6 +304,10 @@ type builtinFunc interface {
 	setPbCode(tipb.ScalarFuncSig)
 	// PbCode returns PbCode of this signature.
 	PbCode() tipb.ScalarFuncSig
+	// implicitArgs returns the implicit parameters of a function.
+	// implicit arguments means some functions contain extra inner fields which will not
+	// contain in `tipb.Expr.children` but must be pushed down to coprocessor
+	implicitArgs() []types.Datum
 	// Clone returns a copy of itself.
 	Clone() builtinFunc
 }

expression/expr_to_pb.go

@@ -243,9 +243,20 @@ func (pc PbConverter) scalarFuncToPBExpr(expr *ScalarFunction) *tipb.Expr {
 		children = append(children, pbArg)
 	}
 
+	var implicitArgs []byte
+	if args := expr.Function.implicitArgs(); len(args) > 0 {
+		encoded, err := codec.EncodeValue(pc.sc, nil, args...)
+		if err != nil {
+			logutil.Logger(context.Background()).Error("encode implicit parameters", zap.Any("datums", args), zap.Error(err))
+			return nil
+		}
+		implicitArgs = encoded
+	}
+
 	// construct expression ProtoBuf.
 	return &tipb.Expr{
 		Tp:        tipb.ExprType_ScalarFunc,
+		Val:       implicitArgs,
 		Sig:       pbCode,
 		Children:  children,
 		FieldType: ToPBFieldType(expr.RetType),

expression/expr_to_pb_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/mock"
 	"github.com/pingcap/tipb/go-tipb"
 )
@@ -593,3 +594,41 @@ func (s *testEvaluatorSuite) TestSortByItem2Pb(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(string(js), Equals, "{\"expr\":{\"tp\":201,\"val\":\"gAAAAAAAAAE=\",\"sig\":0,\"field_type\":{\"tp\":5,\"flag\":0,\"flen\":-1,\"decimal\":-1,\"collate\":46,\"charset\":\"\"}},\"desc\":true}")
 }
+
+func (s *testEvaluatorSuite) TestImplicitArgs(c *C) {
+	sc := new(stmtctx.StatementContext)
+	client := new(mock.Client)
+	dg := new(dataGen4Expr2PbTest)
+
+	c.Assert(failpoint.Enable("github.com/pingcap/tidb/expression/PushDownTestSwitcher", `return("all")`), IsNil)
+	defer func() { c.Assert(failpoint.Disable("github.com/pingcap/tidb/expression/PushDownTestSwitcher"), IsNil) }()
+
+	pc := PbConverter{client: client, sc: sc}
+
+	// InUnion flag is false in `BuildCastFunction` when `ScalarFuncSig_CastStringAsInt`
+	cast := BuildCastFunction(mock.NewContext(), dg.genColumn(mysql.TypeString, 1), types.NewFieldType(mysql.TypeLonglong))
+	c.Assert(cast.(*ScalarFunction).Function.implicitArgs(), DeepEquals, []types.Datum{types.NewIntDatum(0)})
+	expr := pc.ExprToPB(cast)
+	c.Assert(expr.Sig, Equals, tipb.ScalarFuncSig_CastStringAsInt)
+	c.Assert(len(expr.Val), Greater, 0)
+	datums, err := codec.Decode(expr.Val, 1)
+	c.Assert(err, IsNil)
+	c.Assert(datums, DeepEquals, []types.Datum{types.NewIntDatum(0)})
+
+	// InUnion flag is nil in `BuildCastFunction4Union` when `ScalarFuncSig_CastIntAsString`
+	castInUnion := BuildCastFunction4Union(mock.NewContext(), dg.genColumn(mysql.TypeLonglong, 1), types.NewFieldType(mysql.TypeString))
+	c.Assert(castInUnion.(*ScalarFunction).Function.implicitArgs(), IsNil)
+	expr = pc.ExprToPB(castInUnion)
+	c.Assert(expr.Sig, Equals, tipb.ScalarFuncSig_CastIntAsString)
+	c.Assert(len(expr.Val), Equals, 0)
+
+	// InUnion flag is true in `BuildCastFunction4Union` when `ScalarFuncSig_CastStringAsInt`
+	castInUnion = BuildCastFunction4Union(mock.NewContext(), dg.genColumn(mysql.TypeString, 1), types.NewFieldType(mysql.TypeLonglong))
+	c.Assert(castInUnion.(*ScalarFunction).Function.implicitArgs(), DeepEquals, []types.Datum{types.NewIntDatum(1)})
+	expr = pc.ExprToPB(castInUnion)
+	c.Assert(expr.Sig, Equals, tipb.ScalarFuncSig_CastStringAsInt)
+	c.Assert(len(expr.Val), Greater, 0)
+	datums, err = codec.Decode(expr.Val, 1)
+	c.Assert(err, IsNil)
+	c.Assert(datums, DeepEquals, []types.Datum{types.NewIntDatum(1)})
+}