From 87a2f78388bb7ba64cc911656c85845fe14c5156 Mon Sep 17 00:00:00 2001
From: leiysky <leiysky@outlook.com>
Date: Sat, 20 Feb 2021 08:30:50 +0800
Subject: [PATCH] push down extract to tiflash

---
 expression/expr_to_pb_test.go | 15 +++++++++++++--
 expression/expression.go      | 10 +++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/expression/expr_to_pb_test.go b/expression/expr_to_pb_test.go
index 58725db674a72..b4d086a563386 100644
--- a/expression/expr_to_pb_test.go
+++ b/expression/expr_to_pb_test.go
@@ -628,6 +628,11 @@ func (s *testEvaluatorSuite) TestExprPushDownToFlash(c *C) {
 	c.Assert(err, IsNil)
 	exprs = append(exprs, function)
 
+	// ExtractDatetime: can be pushed
+	function, err = NewFunction(mock.NewContext(), ast.Extract, types.NewFieldType(mysql.TypeLonglong), stringColumn, datetimeColumn)
+	c.Assert(err, IsNil)
+	exprs = append(exprs, function)
+
 	// CastIntAsInt
 	function, err = NewFunction(mock.NewContext(), ast.Cast, types.NewFieldType(mysql.TypeLonglong), intColumn)
 	c.Assert(err, IsNil)
@@ -724,9 +729,15 @@ func (s *testEvaluatorSuite) TestExprPushDownToFlash(c *C) {
 	function, err = NewFunction(mock.NewContext(), ast.JSONDepth, types.NewFieldType(mysql.TypeLonglong), jsonColumn)
 	c.Assert(err, IsNil)
 	exprs = append(exprs, function)
+
+	// ExtractDatetimeFromString: can not be pushed
+	function, err = NewFunction(mock.NewContext(), ast.Extract, types.NewFieldType(mysql.TypeLonglong), stringColumn, stringColumn)
+	c.Assert(err, IsNil)
+	exprs = append(exprs, function)
+
 	pushed, remained := PushDownExprs(sc, exprs, client, kv.TiFlash)
-	c.Assert(len(pushed), Equals, len(exprs)-1)
-	c.Assert(len(remained), Equals, 1)
+	c.Assert(len(pushed), Equals, len(exprs)-2)
+	c.Assert(len(remained), Equals, 2)
 }
 
 func (s *testEvaluatorSuite) TestExprOnlyPushDownToFlash(c *C) {
diff --git a/expression/expression.go b/expression/expression.go
index 8eefc711530c7..03e81123ae07f 100644
--- a/expression/expression.go
+++ b/expression/expression.go
@@ -1070,6 +1070,7 @@ func canFuncBePushed(sf *ScalarFunction, storeType kv.StoreType) bool {
 		ast.TimestampDiff,
 		ast.DateAdd,
 		ast.FromUnixTime,
+		ast.Extract,
 
 		// encryption functions.
 		ast.MD5,
@@ -1233,7 +1234,7 @@ func CanExprsPushDown(sc *stmtctx.StatementContext, exprs []Expression, client k
 func scalarExprSupportedByTiKV(function *ScalarFunction) bool {
 	switch function.FuncName.L {
 	case ast.Substr, ast.Substring, ast.DateAdd, ast.TimestampDiff,
-		ast.FromUnixTime:
+		ast.FromUnixTime, ast.Extract:
 		return false
 	default:
 		return true
@@ -1285,6 +1286,13 @@ func scalarExprSupportedByFlash(function *ScalarFunction) bool {
 		default:
 			return false
 		}
+	case ast.Extract:
+		switch function.Function.PbCode() {
+		case tipb.ScalarFuncSig_ExtractDatetime:
+			return true
+		default:
+			return false
+		}
 	default:
 		return false
 	}