*: re-implement partition pruning for better performance

expression/builtin_time_test.go

@@ -14,7 +14,6 @@
 package expression
 
 import (
-	"fmt"
 	"math"
 	"strings"
 	"time"
@@ -1713,7 +1712,6 @@ func (s *testEvaluatorSuite) TestUnixTimestamp(c *C) {
 	}
 
 	for _, test := range tests {
-		fmt.Printf("Begin Test %v\n", test)
 		expr := s.datumsToConstants([]types.Datum{test.input})
 		expr[0].GetType().Decimal = test.inputDecimal
 		resetStmtContext(s.ctx)

go.sum

@@ -92,7 +92,6 @@ github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfb
 github.com/golang/protobuf v0.0.0-20180814211427-aa810b61a9c7/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
-github.com/golang/protobuf v1.3.2 h1:6nsPYzhq5kReh6QImI3k5qWzO4PEbvbIW2cwSfR/6xs=
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.3 h1:gyjaxf+svBWX08ZjK86iN9geUJF0H6gp2IRKX6Nf6/I=
 github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw=
@@ -211,15 +210,13 @@ github.com/pingcap/kvproto v0.0.0-20200108025604-a4dc183d2af5 h1:RUxQExD5yubAjWG
 github.com/pingcap/kvproto v0.0.0-20200108025604-a4dc183d2af5/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
 github.com/pingcap/kvproto v0.0.0-20200210234432-a965739f8162 h1:lsoIoCoXMpcHvW6jHcqP/prA4I6duAp1DVyG2ULz4bM=
 github.com/pingcap/kvproto v0.0.0-20200210234432-a965739f8162/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
-github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9 h1:AJD9pZYm72vMgPcQDww9rkZ1DnWfl0pXV3BOWlkYIjA=
 github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20200117041106-d28c14d3b1cd h1:CV3VsP3Z02MVtdpTMfEgRJ4T9NGgGTxdHpJerent7rM=
 github.com/pingcap/log v0.0.0-20200117041106-d28c14d3b1cd/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/parser v0.0.0-20200212063918-0829643f461c h1:QbFj6Ng/PvHeQNN7aPWpulXIzoo+j/J8odEM7ERUt7g=
 github.com/pingcap/parser v0.0.0-20200212063918-0829643f461c/go.mod h1:9v0Edh8IbgjGYW2ArJr19E+bvL8zKahsFp+ixWeId+4=
 github.com/pingcap/pd v1.1.0-beta.0.20200106144140-f5a7aa985497 h1:FzLErYtcXnSxtC469OuVDlgBbh0trJZzNxw0mNKzyls=
 github.com/pingcap/pd v1.1.0-beta.0.20200106144140-f5a7aa985497/go.mod h1:cfT/xu4Zz+Tkq95QrLgEBZ9ikRcgzy4alHqqoaTftqI=
-github.com/pingcap/sysutil v0.0.0-20191216090214-5f9620d22b3b h1:EEyo/SCRswLGuSk+7SB86Ak1p8bS6HL1Mi4Dhyuv6zg=
 github.com/pingcap/sysutil v0.0.0-20191216090214-5f9620d22b3b/go.mod h1:EB/852NMQ+aRKioCpToQ94Wl7fktV+FNnxf3CX/TTXI=
 github.com/pingcap/sysutil v0.0.0-20200206130906-2bfa6dc40bcd h1:k7CIHMFVKjHsda3PKkiN4zv++NEnexlUwiJEhryWpG0=
 github.com/pingcap/sysutil v0.0.0-20200206130906-2bfa6dc40bcd/go.mod h1:EB/852NMQ+aRKioCpToQ94Wl7fktV+FNnxf3CX/TTXI=

planner/core/partition_pruning_test.go

@@ -19,6 +19,8 @@ import (
 	"github.com/pingcap/parser/model"
 	"github.com/pingcap/tidb/ddl"
 	"github.com/pingcap/tidb/expression"
+	"github.com/pingcap/tidb/sessionctx"
+	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/mock"
 )
 
@@ -29,12 +31,6 @@ type testPartitionPruningSuite struct {
 }
 
 func (s *testPartitionPruningSuite) TestCanBePrune(c *C) {
-	p := parser.New()
-	stmt, err := p.ParseOneStmt("create table t (d datetime not null)", "", "")
-	c.Assert(err, IsNil)
-	tblInfo, err := ddl.BuildTableInfoFromAST(stmt.(*ast.CreateTableStmt))
-	c.Assert(err, IsNil)
-
 	// For the following case:
 	// CREATE TABLE t1 ( recdate  DATETIME NOT NULL )
 	// PARTITION BY RANGE( TO_DAYS(recdate) ) (
@@ -44,20 +40,18 @@ func (s *testPartitionPruningSuite) TestCanBePrune(c *C) {
 	// SELECT * FROM t1 WHERE recdate < '2007-03-08 00:00:00';
 	// SELECT * FROM t1 WHERE recdate > '2018-03-08 00:00:00';
 
-	ctx := mock.NewContext()
-	columns, names := expression.ColumnInfos2ColumnsAndNames(ctx, model.NewCIStr("t"), tblInfo.Name, tblInfo.Columns)
-	schema := expression.NewSchema(columns...)
-	partitionExpr, err := expression.ParseSimpleExprsWithNames(ctx, "to_days(d) < to_days('2007-03-08') and to_days(d) >= to_days('2007-03-07')", schema, names)
-	c.Assert(err, IsNil)
-	queryExpr, err := expression.ParseSimpleExprsWithNames(ctx, "d < '2000-03-08 00:00:00'", schema, names)
-	c.Assert(err, IsNil)
-	succ, err := s.canBePruned(ctx, nil, partitionExpr[0], queryExpr)
+	tc := prepareTestCtx(c,
+		"create table t (d datetime not null)",
+		"to_days(d)",
+	)
+	partitionExpr := tc.expr("to_days(d) < to_days('2007-03-08') and to_days(d) >= to_days('2007-03-07')")
+	queryExpr := tc.expr("d < '2000-03-08 00:00:00'")
+	succ, err := s.canBePruned(tc.sctx, nil, partitionExpr[0], queryExpr)
 	c.Assert(err, IsNil)
 	c.Assert(succ, IsTrue)
 
-	queryExpr, err = expression.ParseSimpleExprsWithNames(ctx, "d > '2018-03-08 00:00:00'", schema, names)
-	c.Assert(err, IsNil)
-	succ, err = s.canBePruned(ctx, nil, partitionExpr[0], queryExpr)
+	queryExpr = tc.expr("d > '2018-03-08 00:00:00'")
+	succ, err = s.canBePruned(tc.sctx, nil, partitionExpr[0], queryExpr)
 	c.Assert(err, IsNil)
 	c.Assert(succ, IsTrue)
 
@@ -72,24 +66,19 @@ func (s *testPartitionPruningSuite) TestCanBePrune(c *C) {
 	// 	PARTITION p2 VALUES LESS THAN (UNIX_TIMESTAMP('2010-01-01 00:00:00')),
 	// 	PARTITION p3 VALUES LESS THAN (MAXVALUE)
 	// );
-	stmt, err = p.ParseOneStmt("create table t (report_updated timestamp)", "", "")
-	c.Assert(err, IsNil)
-	tblInfo, err = ddl.BuildTableInfoFromAST(stmt.(*ast.CreateTableStmt))
-	c.Assert(err, IsNil)
-	columns, names = expression.ColumnInfos2ColumnsAndNames(ctx, model.NewCIStr("t"), tblInfo.Name, tblInfo.Columns)
-	schema = expression.NewSchema(columns...)
+	tc = prepareTestCtx(c,
+		"create table t (report_updated timestamp)",
+		"unix_timestamp(report_updated)",
+	)
 
-	partitionExpr, err = expression.ParseSimpleExprsWithNames(ctx, "unix_timestamp(report_updated) < unix_timestamp('2008-04-01') and unix_timestamp(report_updated) >= unix_timestamp('2008-01-01')", schema, names)
-	c.Assert(err, IsNil)
-	queryExpr, err = expression.ParseSimpleExprsWithNames(ctx, "report_updated > '2008-05-01 00:00:00'", schema, names)
-	c.Assert(err, IsNil)
-	succ, err = s.canBePruned(ctx, nil, partitionExpr[0], queryExpr)
+	partitionExpr = tc.expr("unix_timestamp(report_updated) < unix_timestamp('2008-04-01') and unix_timestamp(report_updated) >= unix_timestamp('2008-01-01')")
+	queryExpr = tc.expr("report_updated > '2008-05-01 00:00:00'")
+	succ, err = s.canBePruned(tc.sctx, nil, partitionExpr[0], queryExpr)
 	c.Assert(err, IsNil)
 	c.Assert(succ, IsTrue)
 
-	queryExpr, err = expression.ParseSimpleExprsWithNames(ctx, "report_updated > unix_timestamp('2008-05-01 00:00:00')", schema, names)
-	c.Assert(err, IsNil)
-	succ, err = s.canBePruned(ctx, nil, partitionExpr[0], queryExpr)
+	queryExpr = tc.expr("report_updated > unix_timestamp('2008-05-01 00:00:00')")
+	succ, err = s.canBePruned(tc.sctx, nil, partitionExpr[0], queryExpr)
 	c.Assert(err, IsNil)
 	_ = succ
 	// c.Assert(succ, IsTrue)
@@ -98,3 +87,181 @@ func (s *testPartitionPruningSuite) TestCanBePrune(c *C) {
 	// Because unix_timestamp('2008-05-01 00:00:00') is fold to constant int 1564761600, and compare it with timestamp (report_updated)
 	// need to convert 1564761600 to a timestamp, during that step, an error happen and the result is set to <nil>
 }
+
+func (s *testPartitionPruningSuite) TestPruneUseBinarySearch(c *C) {
+	lessThan := lessThanData{data: []int64{4, 7, 11, 14, 17, 0}, maxvalue: true}
+	cases := []struct {
+		input  dataForPrune
+		result partitionRange
+	}{
+		{dataForPrune{ast.EQ, 66}, partitionRange{5, 6}},
+		{dataForPrune{ast.EQ, 14}, partitionRange{4, 5}},
+		{dataForPrune{ast.EQ, 10}, partitionRange{2, 3}},
+		{dataForPrune{ast.EQ, 3}, partitionRange{0, 1}},
+		{dataForPrune{ast.LT, 66}, partitionRange{0, 6}},
+		{dataForPrune{ast.LT, 14}, partitionRange{0, 4}},
+		{dataForPrune{ast.LT, 10}, partitionRange{0, 3}},
+		{dataForPrune{ast.LT, 3}, partitionRange{0, 1}},
+		{dataForPrune{ast.GE, 66}, partitionRange{5, 6}},
+		{dataForPrune{ast.GE, 14}, partitionRange{4, 6}},
+		{dataForPrune{ast.GE, 10}, partitionRange{2, 6}},
+		{dataForPrune{ast.GE, 3}, partitionRange{0, 6}},
+		{dataForPrune{ast.GT, 66}, partitionRange{5, 6}},
+		{dataForPrune{ast.GT, 14}, partitionRange{4, 6}},
+		{dataForPrune{ast.GT, 10}, partitionRange{3, 6}},
+		{dataForPrune{ast.GT, 3}, partitionRange{1, 6}},
+		{dataForPrune{ast.GT, 2}, partitionRange{0, 6}},
+		{dataForPrune{ast.LE, 66}, partitionRange{0, 6}},
+		{dataForPrune{ast.LE, 14}, partitionRange{0, 5}},
+		{dataForPrune{ast.LE, 10}, partitionRange{0, 3}},
+		{dataForPrune{ast.LE, 3}, partitionRange{0, 1}},
+		{dataForPrune{ast.IsNull, 0}, partitionRange{0, 1}},
+		{dataForPrune{"illegal", 0}, partitionRange{0, 6}},
+	}
+
+	for i, ca := range cases {
+		start, end := pruneUseBinarySearch(lessThan, ca.input)
+		c.Assert(ca.result.start, Equals, start, Commentf("fail = %d", i))
+		c.Assert(ca.result.end, Equals, end, Commentf("fail = %d", i))
+	}
+}
+
+type testCtx struct {
+	c        *C
+	sctx     sessionctx.Context
+	schema   *expression.Schema
+	columns  []*expression.Column
+	names    types.NameSlice
+	lessThan lessThanData
+}
+
+func prepareTestCtx(c *C, createTable string, partitionExpr string) *testCtx {
+	p := parser.New()
+	stmt, err := p.ParseOneStmt(createTable, "", "")
+	c.Assert(err, IsNil)
+	sctx := mock.NewContext()
+	tblInfo, err := ddl.BuildTableInfoFromAST(stmt.(*ast.CreateTableStmt))
+	c.Assert(err, IsNil)
+	columns, names := expression.ColumnInfos2ColumnsAndNames(sctx, model.NewCIStr("t"), tblInfo.Name, tblInfo.Columns)
+	schema := expression.NewSchema(columns...)
+	return &testCtx{
+		c:       c,
+		sctx:    sctx,
+		schema:  schema,
+		columns: columns,
+		names:   names,
+	}
+}
+
+func (tc *testCtx) expr(expr string) []expression.Expression {
+	res, err := expression.ParseSimpleExprsWithNames(tc.sctx, expr, tc.schema, tc.names)
+	tc.c.Assert(err, IsNil)
+	return res
+}
+
+func (s *testPartitionPruningSuite) TestPartitionRangeForExpr(c *C) {
+	tc := prepareTestCtx(c,
+		"create table t (a int)",
+		"a",
+	)
+	lessThan := lessThanData{data: []int64{4, 7, 11, 14, 17, 0}, maxvalue: true}
+	cases := []struct {
+		input  string
+		result partitionRangeOR
+	}{
+		{"a > 3", partitionRangeOR{{1, 6}}},
+		{"a < 3", partitionRangeOR{{0, 1}}},
+		{"a >= 11", partitionRangeOR{{3, 6}}},
+		{"a > 11", partitionRangeOR{{3, 6}}},
+		{"a < 11", partitionRangeOR{{0, 3}}},
+		{"a = 16", partitionRangeOR{{4, 5}}},
+		{"a > 66", partitionRangeOR{{5, 6}}},
+		{"a > 2 and a < 10", partitionRangeOR{{0, 3}}},
+		{"a < 2 or a >= 15", partitionRangeOR{{0, 1}, {4, 6}}},
+		{"a is null", partitionRangeOR{{0, 1}}},
+		{"12 > a", partitionRangeOR{{0, 4}}},
+		{"4 <= a", partitionRangeOR{{1, 6}}},
+	}
+
+	for _, ca := range cases {
+		expr, err := expression.ParseSimpleExprsWithNames(tc.sctx, ca.input, tc.schema, tc.names)
+		c.Assert(err, IsNil)
+		result := fullRange(lessThan.length())
+		result = partitionRangeForExpr(tc.sctx, expr[0], lessThan, tc.columns[0], nil, result)
+		c.Assert(equalPartitionRangeOR(ca.result, result), IsTrue, Commentf("unexpected:", ca.input))
+	}
+}
+
+func equalPartitionRangeOR(x, y partitionRangeOR) bool {
+	if len(x) != len(y) {
+		return false
+	}
+	for i := 0; i < len(x); i++ {
+		if x[i] != y[i] {
+			return false
+		}
+	}
+	return true
+}
+
+func (s *testPartitionPruningSuite) TestPartitionRangeOperation(c *C) {
+	testIntersectionRange := []struct {
+		input1 partitionRangeOR
+		input2 partitionRange
+		result partitionRangeOR
+	}{
+		{input1: partitionRangeOR{{0, 3}, {6, 12}},
+			input2: partitionRange{4, 7},
+			result: partitionRangeOR{{6, 7}}},
+		{input1: partitionRangeOR{{0, 5}},
+			input2: partitionRange{6, 7},
+			result: partitionRangeOR{}},
+		{input1: partitionRangeOR{{0, 4}, {6, 7}, {8, 11}},
+			input2: partitionRange{3, 9},
+			result: partitionRangeOR{{3, 4}, {6, 7}, {8, 9}}},
+	}
+	for i, ca := range testIntersectionRange {
+		result := ca.input1.intersectionRange(ca.input2.start, ca.input2.end)
+		c.Assert(equalPartitionRangeOR(ca.result, result), IsTrue, Commentf("failed %d", i))
+	}
+
+	testIntersection := []struct {
+		input1 partitionRangeOR
+		input2 partitionRangeOR
+		result partitionRangeOR
+	}{
+		{input1: partitionRangeOR{{0, 3}, {6, 12}},
+			input2: partitionRangeOR{{4, 7}},
+			result: partitionRangeOR{{6, 7}}},
+		{input1: partitionRangeOR{{4, 7}},
+			input2: partitionRangeOR{{0, 3}, {6, 12}},
+			result: partitionRangeOR{{6, 7}}},
+		{input1: partitionRangeOR{{4, 7}, {8, 10}},
+			input2: partitionRangeOR{{0, 5}, {6, 12}},
+			result: partitionRangeOR{{4, 5}, {6, 7}, {8, 10}}},
+	}
+	for i, ca := range testIntersection {
+		result := ca.input1.intersection(ca.input2)
+		c.Assert(equalPartitionRangeOR(ca.result, result), IsTrue, Commentf("failed %d", i))
+	}
+
+	testUnion := []struct {
+		input1 partitionRangeOR
+		input2 partitionRangeOR
+		result partitionRangeOR
+	}{
+		{input1: partitionRangeOR{{0, 1}, {2, 7}},
+			input2: partitionRangeOR{{3, 5}},
+			result: partitionRangeOR{{0, 1}, {2, 7}}},
+		{input1: partitionRangeOR{{2, 7}},
+			input2: partitionRangeOR{{0, 3}, {4, 12}},
+			result: partitionRangeOR{{0, 12}}},
+		{input1: partitionRangeOR{{4, 7}, {8, 10}},
+			input2: partitionRangeOR{{0, 5}},
+			result: partitionRangeOR{{0, 7}, {8, 10}}},
+	}
+	for i, ca := range testUnion {
+		result := ca.input1.union(ca.input2)
+		c.Assert(equalPartitionRangeOR(ca.result, result), IsTrue, Commentf("failed %d", i))
+	}
+}