expression: check timezone when encoding timestamp datum (pingcap#10303)

Author: eurekaka

expression/distsql_builtin.go

@@ -557,7 +557,7 @@ func convertTime(data []byte, ftPB *tipb.FieldType, tz *time.Location) (*Constan
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	if ft.Tp == mysql.TypeTimestamp && !t.IsZero() {
+	if ft.Tp == mysql.TypeTimestamp && tz != time.UTC {
 		err = t.ConvertTimeZone(time.UTC, tz)
 		if err != nil {
 			return nil, errors.Trace(err)

expression/expr_to_pb.go

@@ -106,7 +106,7 @@ func (pc PbConverter) conOrCorColToPBExpr(expr Expression) *tipb.Expr {
 		logutil.Logger(context.Background()).Error("eval constant or correlated column", zap.String("expression", expr.ExplainInfo()), zap.Error(err))
 		return nil
 	}
-	tp, val, ok := pc.encodeDatum(d)
+	tp, val, ok := pc.encodeDatum(ft, d)
 	if !ok {
 		return nil
 	}
@@ -117,7 +117,7 @@ func (pc PbConverter) conOrCorColToPBExpr(expr Expression) *tipb.Expr {
 	return &tipb.Expr{Tp: tp, Val: val, FieldType: ToPBFieldType(ft)}
 }
 
-func (pc *PbConverter) encodeDatum(d types.Datum) (tipb.ExprType, []byte, bool) {
+func (pc *PbConverter) encodeDatum(ft *types.FieldType, d types.Datum) (tipb.ExprType, []byte, bool) {
 	var (
 		tp  tipb.ExprType
 		val []byte
@@ -157,13 +157,11 @@ func (pc *PbConverter) encodeDatum(d types.Datum) (tipb.ExprType, []byte, bool)
 	case types.KindMysqlTime:
 		if pc.client.IsRequestTypeSupported(kv.ReqTypeDAG, int64(tipb.ExprType_MysqlTime)) {
 			tp = tipb.ExprType_MysqlTime
-			t := d.GetMysqlTime()
-			v, err := t.ToPackedUint()
+			val, err := codec.EncodeMySQLTime(pc.sc, d, ft.Tp, nil)
 			if err != nil {
 				logutil.Logger(context.Background()).Error("encode mysql time", zap.Error(err))
 				return tp, nil, false
 			}
-			val = codec.EncodeUint(nil, v)
 			return tp, val, true
 		}
 		return tp, nil, false

expression/integration_test.go

@@ -3897,3 +3897,17 @@ where
     datediff(b.date8, date(from_unixtime(a.starttime))) >= 0`
 	tk.MustQuery(q)
 }
+
+func (s *testIntegrationSuite) TestTimestampDatumEncode(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec(`drop table if exists t;`)
+	tk.MustExec(`create table t (a bigint primary key, b timestamp)`)
+	tk.MustExec(`insert into t values (1, "2019-04-29 11:56:12")`)
+	tk.MustQuery(`explain select * from t where b = (select max(b) from t)`).Check(testkit.Rows(
+		"TableReader_43 10.00 root data:Selection_42",
+		"└─Selection_42 10.00 cop eq(test.t.b, 2019-04-29 11:56:12)",
+		"  └─TableScan_41 10000.00 cop table:t, range:[-inf,+inf], keep order:false, stats:pseudo",
+	))
+	tk.MustQuery(`select * from t where b = (select max(b) from t)`).Check(testkit.Rows(`1 2019-04-29 11:56:12`))
+}

util/codec/codec.go

@@ -67,21 +67,10 @@ func encode(sc *stmtctx.StatementContext, b []byte, vals []types.Datum, comparab
 			b = encodeBytes(b, vals[i].GetBytes(), comparable)
 		case types.KindMysqlTime:
 			b = append(b, uintFlag)
-			t := vals[i].GetMysqlTime()
-			// Encoding timestamp need to consider timezone.
-			// If it's not in UTC, transform to UTC first.
-			if t.Type == mysql.TypeTimestamp && sc.TimeZone != time.UTC {
-				err = t.ConvertTimeZone(sc.TimeZone, time.UTC)
-				if err != nil {
-					return nil, errors.Trace(err)
-				}
-			}
-			var v uint64
-			v, err = t.ToPackedUint()
+			b, err = EncodeMySQLTime(sc, vals[i], mysql.TypeUnspecified, b)
 			if err != nil {
-				return nil, errors.Trace(err)
+				return nil, err
 			}
-			b = EncodeUint(b, v)
 		case types.KindMysqlDuration:
 			// duration may have negative value, so we cannot use String to encode directly.
 			b = append(b, durationFlag)
@@ -134,6 +123,29 @@ func encode(sc *stmtctx.StatementContext, b []byte, vals []types.Datum, comparab
 	return b, errors.Trace(err)
 }
 
+// EncodeMySQLTime encodes datum of `KindMysqlTime` to []byte.
+func EncodeMySQLTime(sc *stmtctx.StatementContext, d types.Datum, tp byte, b []byte) (_ []byte, err error) {
+	t := d.GetMysqlTime()
+	// Encoding timestamp need to consider timezone. If it's not in UTC, transform to UTC first.
+	// This is compatible with `PBToExpr > convertTime`, and coprocessor assumes the passed timestamp is in UTC as well.
+	if tp == mysql.TypeUnspecified {
+		tp = t.Type
+	}
+	if tp == mysql.TypeTimestamp && sc.TimeZone != time.UTC {
+		err = t.ConvertTimeZone(sc.TimeZone, time.UTC)
+		if err != nil {
+			return nil, err
+		}
+	}
+	var v uint64
+	v, err = t.ToPackedUint()
+	if err != nil {
+		return nil, err
+	}
+	b = EncodeUint(b, v)
+	return b, nil
+}
+
 func encodeBytes(b []byte, v []byte, comparable bool) []byte {
 	if comparable {
 		b = append(b, bytesFlag)