plan,table: fix select null value for table partition (#7452)

plan/logical_plan_test.go

@@ -538,6 +538,17 @@ func (s *testPlanSuite) TestTablePartition(c *C) {
 			best: "Dual->Projection",
 			is:   is1,
 		},
+		{
+			// NULL will be located in the first partition.
+			sql:  "select * from t where t.h is null",
+			best: "Partition(41)->Projection",
+			is:   is,
+		},
+		{
+			sql:  "select * from t where t.h is null or t.h > 70",
+			best: "UnionAll{Partition(41)->Partition(44)}->Projection",
+			is:   is1,
+		},
 	}
 	for _, ca := range tests {
 		comment := Commentf("for %s", ca.sql)

plan/rule_partition_processor.go

@@ -14,12 +14,10 @@ package plan
 
 import (
 	"github.com/juju/errors"
-	"github.com/pingcap/tidb/ast"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/table/tables"
 	"github.com/pingcap/tidb/util/ranger"
-	log "github.com/sirupsen/logrus"
 )
 
 // partitionProcessor rewrites the ast for table partition.
@@ -76,8 +74,10 @@ func (s *partitionProcessor) prune(ds *DataSource) (LogicalPlan, error) {
 	}
 
 	var partitionExprs []expression.Expression
+	var col *expression.Column
 	if table, ok := ds.table.(partitionTable); ok {
 		partitionExprs = table.PartitionExpr().Ranges
+		col = table.PartitionExpr().Column
 	}
 	if len(partitionExprs) == 0 {
 		return nil, errors.New("partition expression missing")
@@ -86,8 +86,6 @@ func (s *partitionProcessor) prune(ds *DataSource) (LogicalPlan, error) {
 	// Rewrite data source to union all partitions, during which we may prune some
 	// partitions according to the filter conditions pushed to the DataSource.
 	children := make([]LogicalPlan, 0, len(pi.Definitions))
-
-	col := partitionExprAccessColumn(partitionExprs[0])
 	for i, expr := range partitionExprs {
 		if col != nil {
 			// If the selection condition would never be satisified, prune that partition.
@@ -143,21 +141,3 @@ func (s *partitionProcessor) canBePrune(ctx sessionctx.Context, col *expression.
 	}
 	return len(r) == 0, nil
 }
-
-// partitionExprAccessColumn extracts the column which is used by the partition expression.
-// If the partition expression is not a simple operation on one column, return nil.
-func partitionExprAccessColumn(expr expression.Expression) *expression.Column {
-	lt, ok := expr.(*expression.ScalarFunction)
-	if !ok || lt.FuncName.L != ast.LT {
-		// The partition expression is constructed by us, its format should always be
-		// expr < value, such as "id < 42", "timestamp(id) < maxvalue"
-		log.Warnf("illegal partition expr:%s", expr.String())
-		return nil
-	}
-	tmp := lt.GetArgs()[0]
-	col, ok := tmp.(*expression.Column)
-	if !ok {
-		return nil
-	}
-	return col
-}

table/tables/partition.go

@@ -98,14 +98,17 @@ func newPartitionedTable(tbl *Table, tblInfo *model.TableInfo) (table.Table, err
 //      p1 values less than (y1)
 //      p2 values less than (y2)
 //      p3 values less than (y3))
-// Ranges: (x < y1); (y1 <= x < y2); (y2 <= x < y3)
+// Ranges: (x < y1 or x is null); (y1 <= x < y2); (y2 <= x < y3)
 // UpperBounds: (x < y1); (x < y2); (x < y3)
 type PartitionExpr struct {
+	// Column is the column appeared in the by range expression, partition pruning need this to work.
+	Column      *expression.Column
 	Ranges      []expression.Expression
 	UpperBounds []expression.Expression
 }
 
 func generatePartitionExpr(tblInfo *model.TableInfo) (*PartitionExpr, error) {
+	var column *expression.Column
 	// The caller should assure partition info is not nil.
 	pi := tblInfo.GetPartitionInfo()
 	ctx := mock.NewContext()
@@ -129,6 +132,19 @@ func generatePartitionExpr(tblInfo *model.TableInfo) (*PartitionExpr, error) {
 
 		if i > 0 {
 			fmt.Fprintf(&buf, " and ((%s) >= (%s))", pi.Expr, pi.Definitions[i-1].LessThan[0])
+		} else {
+			// NULL will locate in the first partition, so its expression is (expr < value or expr is null).
+			fmt.Fprintf(&buf, " or ((%s) is null)", pi.Expr)
+
+			// Extracts the column of the partition expression, it will be used by partition prunning.
+			if tmp, err1 := expression.ParseSimpleExprWithTableInfo(ctx, pi.Expr, tblInfo); err1 == nil {
+				if col, ok := tmp.(*expression.Column); ok {
+					column = col
+				}
+			}
+			if column == nil {
+				log.Warnf("partition pruning won't work on this expr:%s", pi.Expr)
+			}
 		}
 
 		expr, err = expression.ParseSimpleExprWithTableInfo(ctx, buf.String(), tblInfo)
@@ -141,6 +157,7 @@ func generatePartitionExpr(tblInfo *model.TableInfo) (*PartitionExpr, error) {
 		buf.Reset()
 	}
 	return &PartitionExpr{
+		Column:      column,
 		Ranges:      partitionPruneExprs,
 		UpperBounds: locateExprs,
 	}, nil

table/tables/tables_test.go

@@ -341,6 +341,8 @@ PARTITION BY RANGE ( id ) (
 
 	_, err := ts.se.Execute(context.Background(), "set @@session.tidb_enable_table_partition=1")
 	c.Assert(err, IsNil)
+	_, err = ts.se.Execute(context.Background(), "drop table if exists t1;")
+	c.Assert(err, IsNil)
 	_, err = ts.se.Execute(context.Background(), createTable1)
 	c.Assert(err, IsNil)
 	tb, err := ts.dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
@@ -422,3 +424,44 @@ PARTITION BY RANGE ( id ) (
 		c.Assert(pd.ID, Equals, p.GetPhysicalID())
 	}
 }
+
+func (ts *testSuite) TestGeneratePartitionExpr(c *C) {
+	_, err := ts.se.Execute(context.Background(), "use test")
+	c.Assert(err, IsNil)
+	_, err = ts.se.Execute(context.Background(), "set @@session.tidb_enable_table_partition=1")
+	c.Assert(err, IsNil)
+	_, err = ts.se.Execute(context.Background(), "drop table if exists t1;")
+	c.Assert(err, IsNil)
+	_, err = ts.se.Execute(context.Background(), `create table t1 (id int)
+							partition by range (id) (
+							partition p0 values less than (4),
+							partition p1 values less than (7),
+							partition p3 values less than maxvalue)`)
+	c.Assert(err, IsNil)
+
+	tbl, err := ts.dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
+	c.Assert(err, IsNil)
+	type partitionExpr interface {
+		PartitionExpr() *tables.PartitionExpr
+	}
+	pe := tbl.(partitionExpr).PartitionExpr()
+	c.Assert(pe.Column.TblName.L, Equals, "t1")
+	c.Assert(pe.Column.ColName.L, Equals, "id")
+
+	ranges := []string{
+		"or(lt(t1.id, 4), isnull(t1.id))",
+		"and(lt(t1.id, 7), ge(t1.id, 4))",
+		"and(1, ge(t1.id, 7))",
+	}
+	upperBounds := []string{
+		"lt(t1.id, 4)",
+		"lt(t1.id, 7)",
+		"1",
+	}
+	for i, expr := range pe.Ranges {
+		c.Assert(expr.String(), Equals, ranges[i])
+	}
+	for i, expr := range pe.UpperBounds {
+		c.Assert(expr.String(), Equals, upperBounds[i])
+	}
+}