plan, expression: support is null in hash partition pruning (#17281)

expression/partition_pruner.go

@@ -79,7 +79,18 @@ func (p *hashPartitionPruner) reduceColumnEQ() bool {
 
 func (p *hashPartitionPruner) reduceConstantEQ() bool {
 	for _, con := range p.conditions {
-		col, cond := validEqualCond(con)
+		var col *Column
+		var cond *Constant
+		if fn, ok := con.(*ScalarFunction); ok {
+			if fn.FuncName.L == ast.IsNull {
+				col, ok = fn.GetArgs()[0].(*Column)
+				if ok {
+					cond = Null.Clone().(*Constant)
+				}
+			} else {
+				col, cond = validEqualCond(con)
+			}
+		}
 		if col != nil {
 			id := p.getColID(col)
 			if p.constantMap[id] != nil {
@@ -94,18 +105,18 @@ func (p *hashPartitionPruner) reduceConstantEQ() bool {
 	return false
 }
 
-func (p *hashPartitionPruner) tryEvalPartitionExpr(piExpr Expression) (val int64, success bool, isNil bool) {
+func (p *hashPartitionPruner) tryEvalPartitionExpr(piExpr Expression) (val int64, success bool, isNull bool) {
 	switch pi := piExpr.(type) {
 	case *ScalarFunction:
 		if pi.FuncName.L == ast.Plus || pi.FuncName.L == ast.Minus || pi.FuncName.L == ast.Mul || pi.FuncName.L == ast.Div {
 			left, right := pi.GetArgs()[0], pi.GetArgs()[1]
-			leftVal, ok, isNil := p.tryEvalPartitionExpr(left)
-			if !ok {
-				return 0, ok, isNil
+			leftVal, ok, isNull := p.tryEvalPartitionExpr(left)
+			if !ok || isNull {
+				return 0, ok, isNull
 			}
-			rightVal, ok, isNil := p.tryEvalPartitionExpr(right)
-			if !ok {
-				return 0, ok, isNil
+			rightVal, ok, isNull := p.tryEvalPartitionExpr(right)
+			if !ok || isNull {
+				return 0, ok, isNull
 			}
 			switch pi.FuncName.L {
 			case ast.Plus:
@@ -123,11 +134,11 @@ func (p *hashPartitionPruner) tryEvalPartitionExpr(piExpr Expression) (val int64
 			val := p.constantMap[idx]
 			if val != nil {
 				pi.GetArgs()[0] = val
-				ret, _, err := pi.EvalInt(p.ctx, chunk.Row{})
+				ret, isNull, err := pi.EvalInt(p.ctx, chunk.Row{})
 				if err != nil {
 					return 0, false, false
 				}
-				return ret, true, false
+				return ret, true, isNull
 			}
 			return 0, false, false
 		}
@@ -137,7 +148,7 @@ func (p *hashPartitionPruner) tryEvalPartitionExpr(piExpr Expression) (val int64
 			return 0, false, false
 		}
 		if val.IsNull() {
-			return 0, false, true
+			return 0, true, true
 		}
 		if val.Kind() == types.KindInt64 {
 			return val.GetInt64(), true, false
@@ -164,8 +175,8 @@ func newHashPartitionPruner() *hashPartitionPruner {
 }
 
 // solve eval the hash partition expression, the first return value represent the result of partition expression. The second
-// return value is whether eval success. The third return value represent whether the eval result of partition value is null.
-func (p *hashPartitionPruner) solve(ctx sessionctx.Context, conds []Expression, piExpr Expression) (val int64, ok bool, isNil bool) {
+// return value is whether eval success. The third return value represent whether the query conditions is always false.
+func (p *hashPartitionPruner) solve(ctx sessionctx.Context, conds []Expression, piExpr Expression) (val int64, ok bool, isAlwaysFalse bool) {
 	p.ctx = ctx
 	for _, cond := range conds {
 		p.conditions = append(p.conditions, SplitCNFItems(cond)...)
@@ -176,17 +187,20 @@ func (p *hashPartitionPruner) solve(ctx sessionctx.Context, conds []Expression,
 	for _, col := range ExtractColumns(piExpr) {
 		p.insertCol(col)
 	}
-	p.constantMap = make([]*Constant, p.numColumn, p.numColumn)
-	conflict := p.reduceConstantEQ()
-	if conflict {
-		return 0, false, conflict
+	p.constantMap = make([]*Constant, p.numColumn)
+	isAlwaysFalse = p.reduceConstantEQ()
+	if isAlwaysFalse {
+		return 0, false, isAlwaysFalse
+	}
+	isAlwaysFalse = p.reduceColumnEQ()
+	if isAlwaysFalse {
+		return 0, false, isAlwaysFalse
 	}
-	conflict = p.reduceColumnEQ()
-	if conflict {
-		return 0, false, conflict
+	res, ok, isNull := p.tryEvalPartitionExpr(piExpr)
+	if isNull && ok {
+		return 0, ok, false
 	}
-	res, ok, isNil := p.tryEvalPartitionExpr(piExpr)
-	return res, ok, isNil
+	return res, ok, false
 }
 
 // FastLocateHashPartition is used to get hash partition quickly.

expression/partition_pruner_test.go

@@ -69,6 +69,8 @@ func (s *testSuite2) TestHashPartitionPruner(c *C) {
 	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;")
+	tk.MustExec("create table t6(a int, b int) partition by hash(a) partitions 3;")
+	tk.MustExec("create table t7(a int, b int) partition by hash(a + b) partitions 10;")
 
 	var input []string
 	var output []struct {

expression/testdata/partition_pruner_in.json

@@ -14,7 +14,11 @@
       "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 t4 where d = '2019-10-07 10:40:00' and a = 1",
-      "explain select * from t5 where d = '2019-10-07'"
+      "explain select * from t5 where d = '2019-10-07'",
+      "explain select * from t6 where a is null",
+      "explain select * from t6 where b is null",
+      "explain select * from t5 where d is null",
+      "explain select * from t7 where b = -3 and a is null"
     ]
   }
 ]

expression/testdata/partition_pruner_out.json

@@ -107,6 +107,43 @@
           "├─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"
         ]
+      },
+      {
+        "SQL": "explain select * from t6 where a is null",
+        "Result": [
+          "TableReader_8 10.00 root data:Selection_7",
+          "└─Selection_7 10.00 cop isnull(test_partition.t6.a)",
+          "  └─TableScan_6 10000.00 cop table:t6, partition:p0, range:[-inf,+inf], keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t6 where b is null",
+        "Result": [
+          "Union_9 30.00 root ",
+          "├─TableReader_12 10.00 root data:Selection_11",
+          "│ └─Selection_11 10.00 cop isnull(test_partition.t6.b)",
+          "│   └─TableScan_10 10000.00 cop table:t6, partition:p0, range:[-inf,+inf], keep order:false, stats:pseudo",
+          "├─TableReader_15 10.00 root data:Selection_14",
+          "│ └─Selection_14 10.00 cop isnull(test_partition.t6.b)",
+          "│   └─TableScan_13 10000.00 cop table:t6, partition:p1, range:[-inf,+inf], keep order:false, stats:pseudo",
+          "└─TableReader_18 10.00 root data:Selection_17",
+          "  └─Selection_17 10.00 cop isnull(test_partition.t6.b)",
+          "    └─TableScan_16 10000.00 cop table:t6, partition:p2, range:[-inf,+inf], keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t5 where d is null",
+        "Result": [
+          "TableDual_6 0.00 root rows:0"
+        ]
+      },
+      {
+        "SQL": "explain select * from t7 where b = -3 and a is null",
+        "Result": [
+          "TableReader_8 0.01 root data:Selection_7",
+          "└─Selection_7 0.01 cop eq(test_partition.t7.b, -3), isnull(test_partition.t7.a)",
+          "  └─TableScan_6 10000.00 cop table:t7, partition:p0, range:[-inf,+inf], keep order:false, stats:pseudo"
+        ]
       }
     ]
   }