planner/core: support range partition pruning for 'in' expression (#17210)

cmd/explaintest/r/partition_pruning.result

@@ -3595,6 +3595,11 @@ id	estRows	task	access object	operator info
 TableReader_8	10.00	root		data:Selection_7
 └─Selection_7	10.00	cop[tikv]		eq(test.t2.a, 833)
   └─TableFullScan_6	10000.00	cop[tikv]	table:t2, partition:p4	keep order:false, stats:pseudo
+explain select * from t2 where a in (10,20,30);
+id	estRows	task	access object	operator info
+TableReader_8	30.00	root		data:Selection_7
+└─Selection_7	30.00	cop[tikv]		in(test.t2.a, 10, 20, 30)
+  └─TableFullScan_6	10000.00	cop[tikv]	table:t2, partition:p0	keep order:false, stats:pseudo
 explain select * from t2 where (a = 100 OR a = 900);
 id	estRows	task	access object	operator info
 Union_8	40.00	root		

cmd/explaintest/t/partition_pruning.test

@@ -766,6 +766,7 @@ explain select * from t2;
 explain select * from t2 where a = 101;
 explain select * from t2 where a = 550;
 explain select * from t2 where a = 833;
+explain select * from t2 where a in (10,20,30);
 explain select * from t2 where (a = 100 OR a = 900);
 explain select * from t2 where (a > 100 AND a < 600);
 explain select * from t2 where b = 4;
@@ -1010,4 +1011,3 @@ explain select * from t where ts <= '2020-04-14 23:59:59.123' -- p0,p1;
 explain select * from t where ts <= '2020-04-25 00:00:00' -- p0,p1,p2;
 explain select * from t where ts > '2020-04-25 00:00:00' or ts < '2020-01-02 00:00:00' -- p0;
 explain select * from t where ts > '2020-04-02 00:00:00' and ts < '2020-04-07 00:00:00' -- p0,p1;
-

planner/core/integration_test.go

@@ -556,6 +556,29 @@ func (s *testIntegrationSuite) TestPartitionTableStats(c *C) {
 	}
 }
 
+func (s *testIntegrationSuite) TestPartitionPruningForInExpr(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int(11), b int) partition by range (a) (partition p0 values less than (4), partition p1 values less than(10), partition p2 values less than maxvalue);")
+	tk.MustExec("insert into t values (1, 1),(10, 10),(11, 11)")
+
+	var input []string
+	var output []struct {
+		SQL  string
+		Plan []string
+	}
+	s.testData.GetTestCases(c, &input, &output)
+	for i, tt := range input {
+		s.testData.OnRecord(func() {
+			output[i].SQL = tt
+			output[i].Plan = s.testData.ConvertRowsToStrings(tk.MustQuery(tt).Rows())
+		})
+		tk.MustQuery(tt).Check(testkit.Rows(output[i].Plan...))
+	}
+}
+
 func (s *testIntegrationSuite) TestErrNoDB(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("create user test")

planner/core/rule_partition_processor.go

@@ -423,6 +423,12 @@ func partitionRangeForExpr(sctx sessionctx.Context, expr expression.Expression,
 			args := op.GetArgs()
 			newRange := partitionRangeForOrExpr(sctx, args[0], args[1], pruner)
 			return result.intersection(newRange)
+		} else if op.FuncName.L == ast.In {
+			if p, ok := pruner.(*rangePruner); ok {
+				newRange := partitionRangeForInExpr(sctx, op.GetArgs(), p)
+				return result.intersection(newRange)
+			}
+			return result
 		}
 	}
 
@@ -450,6 +456,12 @@ type rangePruner struct {
 }
 
 func (p *rangePruner) partitionRangeForExpr(sctx sessionctx.Context, expr expression.Expression) (int, int, bool) {
+	if constExpr, ok := expr.(*expression.Constant); ok {
+		if b, err := constExpr.Value.ToBool(sctx.GetSessionVars().StmtCtx); err == nil && b == 0 {
+			// A constant false expression.
+			return 0, 0, true
+		}
+	}
 	dataForPrune, ok := p.extractDataForPrune(sctx, expr)
 	if !ok {
 		return 0, 0, false
@@ -472,6 +484,39 @@ func partitionRangeForOrExpr(sctx sessionctx.Context, expr1, expr2 expression.Ex
 	return tmp1.union(tmp2)
 }
 
+func partitionRangeForInExpr(sctx sessionctx.Context, args []expression.Expression,
+	pruner *rangePruner) partitionRangeOR {
+	col, ok := args[0].(*expression.Column)
+	if !ok || col.ID != pruner.col.ID {
+		return pruner.fullRange()
+	}
+
+	var result partitionRangeOR
+	unsigned := mysql.HasUnsignedFlag(col.RetType.Flag)
+	for i := 1; i < len(args); i++ {
+		constExpr, ok := args[i].(*expression.Constant)
+		if !ok {
+			return pruner.fullRange()
+		}
+		switch constExpr.Value.Kind() {
+		case types.KindInt64, types.KindUint64:
+		case types.KindNull:
+			result = append(result, partitionRange{0, 1})
+			continue
+		default:
+			return pruner.fullRange()
+		}
+		val, err := constExpr.Value.ToInt64(sctx.GetSessionVars().StmtCtx)
+		if err != nil {
+			return pruner.fullRange()
+		}
+
+		start, end := pruneUseBinarySearch(pruner.lessThan, dataForPrune{op: ast.EQ, c: val}, unsigned)
+		result = append(result, partitionRange{start, end})
+	}
+	return result.simplify()
+}
+
 // monotoneIncFuncs are those functions that for any x y, if x > y => f(x) > f(y)
 var monotoneIncFuncs = map[string]struct{}{
 	ast.ToDays:        {},

planner/core/testdata/integration_suite_in.json

@@ -115,5 +115,17 @@
       "select /*+ IGNORE_INDEX() */ * from t1, t2 where t1.a=t2.a",
       "select /*+ USE_INDEX_MERGE() */ * from t1, t2 where t1.a=t2.a"
     ]
+  },
+  {
+    "name": "TestPartitionPruningForInExpr",
+    "cases": [
+      "explain select * from t where a in (1, 2,'11')",
+      "explain select * from t where a in (17, null)",
+      "explain select * from t where a in (16, 'abc')",
+      "explain select * from t where a in (15, 0.12, 3.47)",
+      "explain select * from t where a in (0.12, 3.47)",
+      "explain select * from t where a in (14, floor(3.47))",
+      "explain select * from t where b in (3, 4)"
+    ]
   }
 ]

planner/core/testdata/integration_suite_out.json

@@ -500,5 +500,87 @@
         ]
       }
     ]
+  },
+  {
+    "Name": "TestPartitionPruningForInExpr",
+    "Cases": [
+      {
+        "SQL": "explain select * from t where a in (1, 2,'11')",
+        "Plan": [
+          "Union_8 60.00 root  ",
+          "├─TableReader_11 30.00 root  data:Selection_10",
+          "│ └─Selection_10 30.00 cop[tikv]  in(test.t.a, 1, 2, 11)",
+          "│   └─TableFullScan_9 10000.00 cop[tikv] table:t, partition:p0 keep order:false, stats:pseudo",
+          "└─TableReader_14 30.00 root  data:Selection_13",
+          "  └─Selection_13 30.00 cop[tikv]  in(test.t.a, 1, 2, 11)",
+          "    └─TableFullScan_12 10000.00 cop[tikv] table:t, partition:p2 keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t where a in (17, null)",
+        "Plan": [
+          "Union_8 20.00 root  ",
+          "├─TableReader_11 10.00 root  data:Selection_10",
+          "│ └─Selection_10 10.00 cop[tikv]  in(test.t.a, 17, NULL)",
+          "│   └─TableFullScan_9 10000.00 cop[tikv] table:t, partition:p0 keep order:false, stats:pseudo",
+          "└─TableReader_14 10.00 root  data:Selection_13",
+          "  └─Selection_13 10.00 cop[tikv]  in(test.t.a, 17, NULL)",
+          "    └─TableFullScan_12 10000.00 cop[tikv] table:t, partition:p2 keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t where a in (16, 'abc')",
+        "Plan": [
+          "Union_8 40.00 root  ",
+          "├─TableReader_11 20.00 root  data:Selection_10",
+          "│ └─Selection_10 20.00 cop[tikv]  in(test.t.a, 16, 0)",
+          "│   └─TableFullScan_9 10000.00 cop[tikv] table:t, partition:p0 keep order:false, stats:pseudo",
+          "└─TableReader_14 20.00 root  data:Selection_13",
+          "  └─Selection_13 20.00 cop[tikv]  in(test.t.a, 16, 0)",
+          "    └─TableFullScan_12 10000.00 cop[tikv] table:t, partition:p2 keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t where a in (15, 0.12, 3.47)",
+        "Plan": [
+          "TableReader_8 10.00 root  data:Selection_7",
+          "└─Selection_7 10.00 cop[tikv]  or(eq(test.t.a, 15), 0)",
+          "  └─TableFullScan_6 10000.00 cop[tikv] table:t, partition:p2 keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t where a in (0.12, 3.47)",
+        "Plan": [
+          "TableDual_6 0.00 root  rows:0"
+        ]
+      },
+      {
+        "SQL": "explain select * from t where a in (14, floor(3.47))",
+        "Plan": [
+          "Union_8 40.00 root  ",
+          "├─TableReader_11 20.00 root  data:Selection_10",
+          "│ └─Selection_10 20.00 cop[tikv]  in(test.t.a, 14, 3)",
+          "│   └─TableFullScan_9 10000.00 cop[tikv] table:t, partition:p0 keep order:false, stats:pseudo",
+          "└─TableReader_14 20.00 root  data:Selection_13",
+          "  └─Selection_13 20.00 cop[tikv]  in(test.t.a, 14, 3)",
+          "    └─TableFullScan_12 10000.00 cop[tikv] table:t, partition:p2 keep order:false, stats:pseudo"
+        ]
+      },
+      {
+        "SQL": "explain select * from t where b in (3, 4)",
+        "Plan": [
+          "Union_9 60.00 root  ",
+          "├─TableReader_12 20.00 root  data:Selection_11",
+          "│ └─Selection_11 20.00 cop[tikv]  in(test.t.b, 3, 4)",
+          "│   └─TableFullScan_10 10000.00 cop[tikv] table:t, partition:p0 keep order:false, stats:pseudo",
+          "├─TableReader_15 20.00 root  data:Selection_14",
+          "│ └─Selection_14 20.00 cop[tikv]  in(test.t.b, 3, 4)",
+          "│   └─TableFullScan_13 10000.00 cop[tikv] table:t, partition:p1 keep order:false, stats:pseudo",
+          "└─TableReader_18 20.00 root  data:Selection_17",
+          "  └─Selection_17 20.00 cop[tikv]  in(test.t.b, 3, 4)",
+          "    └─TableFullScan_16 10000.00 cop[tikv] table:t, partition:p2 keep order:false, stats:pseudo"
+        ]
+      }
+    ]
   }
 ]