From df56205ca9ecbdfafa320f82ee962f933d3b7cbc Mon Sep 17 00:00:00 2001
From: Lingyu Song <songlingyu@pingcap.com>
Date: Wed, 25 Mar 2020 11:03:09 +0800
Subject: [PATCH] plan, expression: add date function support for hash
 partition (#15068) (#15619)

---
 expression/partition_pruner.go                | 13 +++++++++++++
 expression/partition_pruner_test.go           |  3 +++
 expression/testdata/partition_pruner_in.json  |  5 +++--
 expression/testdata/partition_pruner_out.json | 14 +++++++++-----
 4 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/expression/partition_pruner.go b/expression/partition_pruner.go
index f380844d9e7a8..4483af107498c 100644
--- a/expression/partition_pruner.go
+++ b/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{})
diff --git a/expression/partition_pruner_test.go b/expression/partition_pruner_test.go
index bf19567f831d3..8f181d821169a 100644
--- a/expression/partition_pruner_test.go
+++ b/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 {
diff --git a/expression/testdata/partition_pruner_in.json b/expression/testdata/partition_pruner_in.json
index 44934d894c96d..82b138f99874a 100644
--- a/expression/testdata/partition_pruner_in.json
+++ b/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'"
     ]
   }
 ]
diff --git a/expression/testdata/partition_pruner_out.json b/expression/testdata/partition_pruner_out.json
index b661fe23713bd..2998b08c7bc1e 100644
--- a/expression/testdata/partition_pruner_out.json
+++ b/expression/testdata/partition_pruner_out.json
@@ -93,17 +93,21 @@
         ]
       },
       {
-        "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[tikv] 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[tikv] 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[tikv] table:t5, partition:p0, index:d, a, range:[2019-10-07,2019-10-07], keep order:false, stats:pseudo",
+          "└─TableScan_10 10.00 cop[tikv] table:t5, partition:p0, keep order:false, stats:pseudo"
         ]
       }
     ]
   }
-]
\ No newline at end of file
+]