expression: push down EXTRACT to TiFlash (pingcap#22832)

denimyftiu · Mar 12, 2021 · b128c3c · b128c3c
1 parent 0b68d8a
commit b128c3c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/expression/expr_to_pb_test.go b/expression/expr_to_pb_test.go
@@ -629,6 +629,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)
@@ -725,9 +730,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
@@ -1070,6 +1070,7 @@ func canFuncBePushed(sf *ScalarFunction, storeType kv.StoreType) bool {
 		ast.TimestampDiff,
 		ast.DateAdd,
 		ast.FromUnixTime,
+		ast.Extract,
 
 		// encryption functions.
 		ast.MD5,
@@ -1243,7 +1244,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
@@ -1295,6 +1296,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
 	}