expression: refine performance of EXTRACT function

cmd/explaintest/r/tpch.result

@@ -448,7 +448,7 @@ id estRows task access object operator info
 Sort_23 769.96 root tpch.nation.n_name, tpch.nation.n_name, Column#50
 └─Projection_25 769.96 root tpch.nation.n_name, tpch.nation.n_name, Column#50, Column#52
  └─HashAgg_26 769.96 root group by:Column#50, tpch.nation.n_name, tpch.nation.n_name, funcs:sum(Column#51)->Column#52, funcs:firstrow(tpch.nation.n_name)->tpch.nation.n_name, funcs:firstrow(tpch.nation.n_name)->tpch.nation.n_name, funcs:firstrow(Column#50)->Column#50
- └─Projection_27 1957240.42 root tpch.nation.n_name, tpch.nation.n_name, extract(YEAR, cast(tpch.lineitem.l_shipdate, var_string(10)))->Column#50, mul(tpch.lineitem.l_extendedprice, minus(1, tpch.lineitem.l_discount))->Column#51
+ └─Projection_27 1957240.42 root tpch.nation.n_name, tpch.nation.n_name, extract(YEAR, tpch.lineitem.l_shipdate)->Column#50, mul(tpch.lineitem.l_extendedprice, minus(1, tpch.lineitem.l_discount))->Column#51
  └─Projection_28 1957240.42 root tpch.lineitem.l_extendedprice, tpch.lineitem.l_discount, tpch.lineitem.l_shipdate, tpch.nation.n_name, tpch.nation.n_name
  └─HashJoin_40 1957240.42 root inner join, equal:[eq(tpch.customer.c_nationkey, tpch.nation.n_nationkey)], other cond:or(and(eq(tpch.nation.n_name, "JAPAN"), eq(tpch.nation.n_name, "INDIA")), and(eq(tpch.nation.n_name, "INDIA"), eq(tpch.nation.n_name, "JAPAN")))
  ├─TableReader_94(Build) 2.00 root data:Selection_93
@@ -522,7 +522,7 @@ Sort_30 719.02 root Column#62
 └─Projection_32 719.02 root Column#62, div(Column#64, Column#65)->Column#66
  └─HashAgg_33 719.02 root group by:Column#78, funcs:sum(Column#75)->Column#64, funcs:sum(Column#76)->Column#65, funcs:firstrow(Column#77)->Column#62
  └─Projection_123 563136.02 root case(eq(tpch.nation.n_name, INDIA), Column#63, 0)->Column#75, Column#63, Column#62, Column#62
- └─Projection_34 563136.02 root extract(YEAR, cast(tpch.orders.o_orderdate, var_string(10)))->Column#62, mul(tpch.lineitem.l_extendedprice, minus(1, tpch.lineitem.l_discount))->Column#63, tpch.nation.n_name
+ └─Projection_34 563136.02 root extract(YEAR, tpch.orders.o_orderdate)->Column#62, mul(tpch.lineitem.l_extendedprice, minus(1, tpch.lineitem.l_discount))->Column#63, tpch.nation.n_name
  └─Projection_35 563136.02 root tpch.lineitem.l_extendedprice, tpch.lineitem.l_discount, tpch.orders.o_orderdate, tpch.nation.n_name
  └─HashJoin_45 563136.02 root inner join, equal:[eq(tpch.supplier.s_nationkey, tpch.nation.n_nationkey)]
  ├─TableReader_121(Build) 25.00 root data:TableFullScan_120
@@ -598,7 +598,7 @@ id estRows task access object operator info
 Sort_26 2406.00 root tpch.nation.n_name, Column#53:desc
 └─Projection_28 2406.00 root tpch.nation.n_name, Column#53, Column#55
  └─HashAgg_29 2406.00 root group by:Column#53, tpch.nation.n_name, funcs:sum(Column#54)->Column#55, funcs:firstrow(tpch.nation.n_name)->tpch.nation.n_name, funcs:firstrow(Column#53)->Column#53
- └─Projection_30 241379546.70 root tpch.nation.n_name, extract(YEAR, cast(tpch.orders.o_orderdate, var_string(10)))->Column#53, minus(mul(tpch.lineitem.l_extendedprice, minus(1, tpch.lineitem.l_discount)), mul(tpch.partsupp.ps_supplycost, tpch.lineitem.l_quantity))->Column#54
+ └─Projection_30 241379546.70 root tpch.nation.n_name, extract(YEAR, tpch.orders.o_orderdate)->Column#53, minus(mul(tpch.lineitem.l_extendedprice, minus(1, tpch.lineitem.l_discount)), mul(tpch.partsupp.ps_supplycost, tpch.lineitem.l_quantity))->Column#54
  └─Projection_31 241379546.70 root tpch.lineitem.l_quantity, tpch.lineitem.l_extendedprice, tpch.lineitem.l_discount, tpch.partsupp.ps_supplycost, tpch.orders.o_orderdate, tpch.nation.n_name
  └─HashJoin_42 241379546.70 root inner join, equal:[eq(tpch.lineitem.l_orderkey, tpch.orders.o_orderkey)]
  ├─TableReader_117(Build) 75000000.00 root data:TableFullScan_116

expression/builtin_time.go

@@ -217,6 +217,7 @@ var (
  _ builtinFunc = &builtinStrToDateDurationSig{}
  _ builtinFunc = &builtinFromUnixTime1ArgSig{}
  _ builtinFunc = &builtinFromUnixTime2ArgSig{}
+ _ builtinFunc = &builtinExtractDatetimeFromStringSig{}
  _ builtinFunc = &builtinExtractDatetimeSig{}
  _ builtinFunc = &builtinExtractDurationSig{}
  _ builtinFunc = &builtinAddDateStringStringSig{}
@@ -2631,17 +2632,22 @@ func (c *extractFunctionClass) getFunction(ctx sessionctx.Context, args []Expres
  }
  var bf baseBuiltinFunc
  if isDatetimeUnit {
- decimalArg1 := int(types.MaxFsp)
  if args[1].GetType().EvalType() != types.ETString {
- decimalArg1 = 0
- }
- bf, err = newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETInt, types.ETString, types.ETString)
- if err != nil {
- return nil, err
+ bf, err = newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETInt, types.ETString, types.ETDatetime)
+ if err != nil {
+ return nil, err
+ }
+ sig = &builtinExtractDatetimeSig{bf}
+ sig.setPbCode(tipb.ScalarFuncSig_ExtractDatetime)
+ } else {
+ bf, err = newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETInt, types.ETString, types.ETString)
+ if err != nil {
+ return nil, err
+ }
+ bf.args[1].GetType().Decimal = int(types.MaxFsp)
+ sig = &builtinExtractDatetimeFromStringSig{bf}
+ sig.setPbCode(tipb.ScalarFuncSig_ExtractDatetimeFromString)
  }
- bf.args[1].GetType().Decimal = decimalArg1
- sig = &builtinExtractDatetimeSig{bf}
- sig.setPbCode(tipb.ScalarFuncSig_ExtractDatetime)
  } else {
  bf, err = newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETInt, types.ETString, types.ETDuration)
  if err != nil {
@@ -2653,19 +2659,19 @@ func (c *extractFunctionClass) getFunction(ctx sessionctx.Context, args []Expres
  return sig, nil
 }
 
-type builtinExtractDatetimeSig struct {
+type builtinExtractDatetimeFromStringSig struct {
  baseBuiltinFunc
 }
 
-func (b *builtinExtractDatetimeSig) Clone() builtinFunc {
- newSig := &builtinExtractDatetimeSig{}
+func (b *builtinExtractDatetimeFromStringSig) Clone() builtinFunc {
+ newSig := &builtinExtractDatetimeFromStringSig{}
  newSig.cloneFrom(&b.baseBuiltinFunc)
  return newSig
 }
 
 // evalInt evals a builtinExtractDatetimeSig.
 // See https://dev.mysql.com/doc/refman/5.7/en/date-and-time-functions.html#function_extract
-func (b *builtinExtractDatetimeSig) evalInt(row chunk.Row) (int64, bool, error) {
+func (b *builtinExtractDatetimeFromStringSig) evalInt(row chunk.Row) (int64, bool, error) {
  unit, isNull, err := b.args[0].EvalString(b.ctx, row)
  if isNull || err != nil {
  return 0, isNull, err
@@ -2712,6 +2718,31 @@ func (b *builtinExtractDatetimeSig) evalInt(row chunk.Row) (int64, bool, error)
  return res, err != nil, err
 }
 
+type builtinExtractDatetimeSig struct {
+ baseBuiltinFunc
+}
+
+func (b *builtinExtractDatetimeSig) Clone() builtinFunc {
+ newSig := &builtinExtractDatetimeSig{}
+ newSig.cloneFrom(&b.baseBuiltinFunc)
+ return newSig
+}
+
+// evalInt evals a builtinExtractDatetimeSig.
+// See https://dev.mysql.com/doc/refman/5.7/en/date-and-time-functions.html#function_extract
+func (b *builtinExtractDatetimeSig) evalInt(row chunk.Row) (int64, bool, error) {
+ unit, isNull, err := b.args[0].EvalString(b.ctx, row)
+ if isNull || err != nil {
+ return 0, isNull, err
+ }
+ dt, isNull, err := b.args[1].EvalTime(b.ctx, row)
+ if isNull || err != nil {
+ return 0, isNull, err
+ }
+ res, err := types.ExtractDatetimeNum(&dt, unit)
+ return res, err != nil, err
+}
+
 type builtinExtractDurationSig struct {
  baseBuiltinFunc
 }

expression/builtin_time_vec.go

@@ -205,12 +205,12 @@ func (b *builtinSysDateWithoutFspSig) vecEvalTime(input *chunk.Chunk, result *ch
  return nil
 }
 
-func (b *builtinExtractDatetimeSig) vectorized() bool {
+func (b *builtinExtractDatetimeFromStringSig) vectorized() bool {
  // TODO: to fix https://github.com/pingcap/tidb/issues/9716 in vectorized evaluation.
  return false
 }
 
-func (b *builtinExtractDatetimeSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
+func (b *builtinExtractDatetimeFromStringSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
  n := input.NumRows()
  buf, err := b.bufAllocator.get(types.ETString, n)
  if err != nil {
@@ -1001,6 +1001,47 @@ func (b *builtinWeekWithModeSig) vecEvalInt(input *chunk.Chunk, result *chunk.Co
  return nil
 }
 
+func (b *builtinExtractDatetimeSig) vectorized() bool {
+ return true
+}
+
+func (b *builtinExtractDatetimeSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
+ n := input.NumRows()
+ unit, err := b.bufAllocator.get(types.ETString, n)
+ if err != nil {
+ return err
+ }
+ defer b.bufAllocator.put(unit)
+ if err := b.args[0].VecEvalString(b.ctx, input, unit); err != nil {
+ return err
+ }
+ dt, err := b.bufAllocator.get(types.ETDatetime, n)
+ if err != nil {
+ return err
+ }
+ defer b.bufAllocator.put(dt)
+ if err = b.args[1].VecEvalTime(b.ctx, input, dt); err != nil {
+ return err
+ }
+ result.ResizeInt64(n, false)
+ result.MergeNulls(unit, dt)
+ i64s := result.Int64s()
+ tmIs := dt.Times()
+ var t types.Time
+ for i := 0; i < n; i++ {
+ if result.IsNull(i) {
+ continue
+ }
+ unitI := unit.GetString(i)
+ t = tmIs[i]
+ i64s[i], err = types.ExtractDatetimeNum(&t, unitI)
+ if err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
 func (b *builtinExtractDurationSig) vectorized() bool {
  return true
 }

expression/distsql_builtin.go

@@ -875,6 +875,8 @@ func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *ti
  f = &builtinFromUnixTime1ArgSig{base}
  case tipb.ScalarFuncSig_FromUnixTime2Arg:
  f = &builtinFromUnixTime2ArgSig{base}
+ case tipb.ScalarFuncSig_ExtractDatetimeFromString:
+ f = &builtinExtractDatetimeFromStringSig{base}
  case tipb.ScalarFuncSig_ExtractDatetime:
  f = &builtinExtractDatetimeSig{base}
  case tipb.ScalarFuncSig_ExtractDuration:

expression/integration_test.go

@@ -2086,6 +2086,8 @@ func (s *testIntegrationSuite2) TestTimeBuiltin(c *C) {
  result.Check(testkit.Rows("11212 3612 1212 2400"))
  result = tk.MustQuery("select extract(day_hour from '2017-01-01 12:12:12'), extract(day_hour from '01 12:12:12'), extract(day_hour from '12:12:12'), extract(day_hour from '01 00:00:00.89')")
  result.Check(testkit.Rows("112 36 12 24"))
+ result = tk.MustQuery("select extract(day_microsecond from cast('2017-01-01 12:12:12' as datetime)), extract(day_second from cast('2017-01-01 12:12:12' as datetime)), extract(day_minute from cast('2017-01-01 12:12:12' as datetime)), extract(day_hour from cast('2017-01-01 12:12:12' as datetime))")
+ result.Check(testkit.Rows("1121212000000 1121212 11212 112"))
 
  // for adddate, subdate
  dateArithmeticalTests := []struct {

expression/util_test.go

@@ -158,7 +158,7 @@ func (s *testUtilSuite) TestClone(c *check.C) {
  &builtinQuarterSig{}, &builtinSecToTimeSig{}, &builtinTimeToSecSig{}, &builtinTimestampAddSig{}, &builtinToDaysSig{},
  &builtinToSecondsSig{}, &builtinUTCTimeWithArgSig{}, &builtinUTCTimeWithoutArgSig{}, &builtinTimestamp1ArgSig{}, &builtinTimestamp2ArgsSig{},
  &builtinTimestampLiteralSig{}, &builtinLastDaySig{}, &builtinStrToDateDateSig{}, &builtinStrToDateDatetimeSig{}, &builtinStrToDateDurationSig{},
- &builtinFromUnixTime1ArgSig{}, &builtinFromUnixTime2ArgSig{}, &builtinExtractDatetimeSig{}, &builtinExtractDurationSig{}, &builtinAddDateStringStringSig{},
+ &builtinFromUnixTime1ArgSig{}, &builtinFromUnixTime2ArgSig{}, &builtinExtractDatetimeFromStringSig{}, &builtinExtractDatetimeSig{}, &builtinExtractDurationSig{}, &builtinAddDateStringStringSig{},
  &builtinAddDateStringIntSig{}, &builtinAddDateStringRealSig{}, &builtinAddDateStringDecimalSig{}, &builtinAddDateIntStringSig{}, &builtinAddDateIntIntSig{},
  &builtinAddDateIntRealSig{}, &builtinAddDateIntDecimalSig{}, &builtinAddDateDatetimeStringSig{}, &builtinAddDateDatetimeIntSig{}, &builtinAddDateDatetimeRealSig{},
  &builtinAddDateDatetimeDecimalSig{}, &builtinSubDateStringStringSig{}, &builtinSubDateStringIntSig{}, &builtinSubDateStringRealSig{}, &builtinSubDateStringDecimalSig{},

go.mod

@@ -49,7 +49,7 @@ require (
  github.com/pingcap/parser v0.0.0-20210203141130-32ef3e0e2e13
  github.com/pingcap/sysutil v0.0.0-20201130064824-f0c8aa6a6966
  github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible
- github.com/pingcap/tipb v0.0.0-20201229060814-148bc717ce4c
+ github.com/pingcap/tipb v0.0.0-20210204051656-2870a0852037
  github.com/prometheus/client_golang v1.5.1
  github.com/prometheus/client_model v0.2.0
  github.com/prometheus/common v0.9.1
@@ -70,7 +70,7 @@ require (
  go.uber.org/multierr v1.6.0 // indirect
  go.uber.org/zap v1.16.0
  golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a // indirect
- golang.org/x/net v0.0.0-20201021035429-f5854403a974
+ golang.org/x/net v0.0.0-20201110031124-69a78807bb2b
  golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9
  golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4
  golang.org/x/text v0.3.5

go.sum

@@ -182,6 +182,7 @@ github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/me
 github.com/go-yaml/yaml v2.1.0+incompatible h1:RYi2hDdss1u4YE7GwixGzWwVo47T8UQwnTLB6vQiq+o=
 github.com/go-yaml/yaml v2.1.0+incompatible/go.mod h1:w2MrLa16VYP0jy6N7M5kHaCkaLENm+P+Tv+MfurjSw0=
 github.com/goccy/go-graphviz v0.0.5/go.mod h1:wXVsXxmyMQU6TN3zGRttjNn3h+iCAS7xQFC6TlNvLhk=
+github.com/gogo/protobuf v0.0.0-20171007142547-342cbe0a0415/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v0.0.0-20180717141946-636bf0302bc9/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
@@ -418,8 +419,8 @@ github.com/pingcap/sysutil v0.0.0-20201130064824-f0c8aa6a6966 h1:JI0wOAb8aQML0vA
 github.com/pingcap/sysutil v0.0.0-20201130064824-f0c8aa6a6966/go.mod h1:EB/852NMQ+aRKioCpToQ94Wl7fktV+FNnxf3CX/TTXI=
 github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible h1:ceznmu/lLseGHP/jKyOa/3u/5H3wtLLLqkH2V3ssSjg=
 github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible/go.mod h1:XGdcy9+yqlDSEMTpOXnwf3hiTeqrV6MN/u1se9N8yIM=
-github.com/pingcap/tipb v0.0.0-20201229060814-148bc717ce4c h1:kvrdp2hY+asgSvVXCj4eebA9DH4SSouRVQUZpa1Se/Y=
-github.com/pingcap/tipb v0.0.0-20201229060814-148bc717ce4c/go.mod h1:RtkHW8WbcNxj8lsbzjaILci01CtYnYbIkQhjyZWrWVI=
+github.com/pingcap/tipb v0.0.0-20210204051656-2870a0852037 h1:FVIyv52hHnkhWX7FIUCrfXC5BBDo+yaGX2+w5lV65Xs=
+github.com/pingcap/tipb v0.0.0-20210204051656-2870a0852037/go.mod h1:nsEhnMokcn7MRqd2J60yxpn/ac3ZH8A6GOJ9NslabUo=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
@@ -658,8 +659,9 @@ golang.org/x/net v0.0.0-20200202094626-16171245cfb2/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200301022130-244492dfa37a/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200520182314-0ba52f642ac2/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
-golang.org/x/net v0.0.0-20201021035429-f5854403a974 h1:IX6qOQeG5uLjB/hjjwjedwfjND0hgjPMMyO1RoIXQNI=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.0.0-20201110031124-69a78807bb2b h1:uwuIcX0g4Yl1NC5XAz37xsr2lTtcqevgzYNVt49waME=
+golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
@@ -749,6 +751,7 @@ golang.org/x/tools v0.0.0-20191216173652-a0e659d51361/go.mod h1:TB2adYChydJhpapK
 golang.org/x/tools v0.0.0-20191227053925-7b8e75db28f4/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
 golang.org/x/tools v0.0.0-20200225230052-807dcd883420/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
 golang.org/x/tools v0.0.0-20200527183253-8e7acdbce89d/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
+golang.org/x/tools v0.0.0-20201125231158-b5590deeca9b/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.0 h1:po9/4sTYwZU9lPhi1tOrb4hCv3qrhiQ77LZfGa2OjwY=
 golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=