plan, expression: add date function support for hash partition (#15068)

expression/partition_pruner.go

@@ -117,6 +117,19 @@ func (p *hashPartitionPruner) tryEvalPartitionExpr(piExpr Expression) (val int64
 			case ast.Div:
 				return rightVal / leftVal, true, false
 			}
+		} else if pi.FuncName.L == ast.Year || pi.FuncName.L == ast.Month || pi.FuncName.L == ast.ToDays {
+			col := pi.GetArgs()[0].(*Column)
+			idx := p.getColID(col)
+			val := p.constantMap[idx]
+			if val != nil {
+				pi.GetArgs()[0] = val
+				ret, _, err := pi.EvalInt(p.ctx, chunk.Row{})
+				if err != nil {
+					return 0, false, false
+				}
+				return ret, true, false
+			}
+			return 0, false, false
 		}
 	case *Constant:
 		val, err := pi.Eval(chunk.Row{})

expression/partition_pruner_test.go

@@ -66,6 +66,9 @@ func (s *testSuite2) TestHashPartitionPruner(c *C) {
 	tk.MustExec("drop table if exists t1, t2;")
 	tk.MustExec("create table t2(id int, a int, b int, primary key(id, a)) partition by hash(id + a) partitions 10;")
 	tk.MustExec("create table t1(id int primary key, a int, b int) partition by hash(id) partitions 10;")
+	tk.MustExec("create table t3(id int, a int, b int, primary key(id, a)) partition by hash(id) partitions 10;")
+	tk.MustExec("create table t4(d datetime, a int, b int, primary key(d, a)) partition by hash(year(d)) partitions 10;")
+	tk.MustExec("create table t5(d date, a int, b int, primary key(d, a)) partition by hash(month(d)) partitions 10;")
 
 	var input []string
 	var output []struct {

expression/testdata/partition_pruner_in.json

@@ -16,8 +16,9 @@
       "explain select * from t1 left join t2 on true where t1.a = 1 and false",
       "explain select * from t1 left join t2 on true where t1.a = 1 and null",
       "explain select * from t1 left join t2 on true where t1.a = null",
-      "explain select * from t1 where t1.a > 7 and t1.a < 3",
-      "explain select * from t1 where t1.a between 7 and 3"
+      // Case with date.
+      "explain select * from t4 where d = '2019-10-07 10:40:00' and a = 1",
+      "explain select * from t5 where d = '2019-10-07'"
     ]
   }
 ]

expression/testdata/partition_pruner_out.json

@@ -93,15 +93,19 @@
         ]
       },
       {
-        "SQL": "explain select * from t1 where t1.a > 7 and t1.a < 3",
+        "SQL": "explain select * from t4 where d = '2019-10-07 10:40:00' and a = 1",
         "Result": [
-          "TableDual_6 0.00 root rows:0"
+          "IndexLookUp_8 1.00 root ",
+          "├─IndexScan_6 1.00 cop table:t4, partition:p9, index:d, a, range:[2019-10-07 10:40:00 1,2019-10-07 10:40:00 1], keep order:false, stats:pseudo",
+          "└─TableScan_7 1.00 cop table:t4, partition:p9, keep order:false, stats:pseudo"
         ]
       },
       {
-        "SQL": "explain select * from t1 where t1.a between 7 and 3",
+        "SQL": "explain select * from t5 where d = '2019-10-07'",
         "Result": [
-          "TableDual_6 0.00 root rows:0"
+          "IndexLookUp_11 10.00 root ",
+          "├─IndexScan_9 10.00 cop table:t5, partition:p0, index:d, a, range:[2019-10-07,2019-10-07], keep order:false, stats:pseudo",
+          "└─TableScan_10 10.00 cop table:t5, partition:p0, keep order:false, stats:pseudo"
         ]
       }
     ]