planner: support keep pruned conds

executor/set_test.go

@@ -520,6 +520,16 @@ func (s *testSerialSuite1) TestSetVar(c *C) {
 	tk.MustQuery(`select @@global.tidb_enable_ordered_result_mode`).Check(testkit.Rows("0"))
 	tk.MustQuery(`select @@tidb_enable_ordered_result_mode`).Check(testkit.Rows("1"))
 
+	// test for tidb_keep_pruned_conds
+	tk.MustQuery(`select @@tidb_keep_pruned_conds`).Check(testkit.Rows("0"))
+	tk.MustExec(`set global tidb_keep_pruned_conds = 1`)
+	tk.MustQuery(`select @@global.tidb_keep_pruned_conds`).Check(testkit.Rows("1"))
+	tk.MustExec(`set global tidb_keep_pruned_conds = 0`)
+	tk.MustQuery(`select @@global.tidb_keep_pruned_conds`).Check(testkit.Rows("0"))
+	tk.MustExec(`set tidb_keep_pruning_conds=1`)
+	tk.MustQuery(`select @@global.tidb_keep_pruned_conds`).Check(testkit.Rows("0"))
+	tk.MustQuery(`select @@tidb_keep_pruned_conds`).Check(testkit.Rows("1"))
+
 	// test for tidb_opt_enable_correlation_adjustment
 	tk.MustQuery(`select @@tidb_opt_enable_correlation_adjustment`).Check(testkit.Rows("1"))
 	tk.MustExec(`set global tidb_opt_enable_correlation_adjustment = 0`)

planner/core/integration_test.go

@@ -4406,3 +4406,22 @@ func (s *testIntegrationSerialSuite) TestIssue30271(c *C) {
 	tk.MustExec("set names utf8mb4 collate utf8mb4_general_ci;")
 	tk.MustQuery("select * from t where (a>'a' and b='a') or (b = 'A' and a < 'd') order by a,c;").Check(testkit.Rows("b a 1", "b A 2", "c a 3"))
 }
+
+func (s *testIntegrationSerialSuite) TestKeepPrunedConds(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("set @@session.tidb_partition_prune_mode = 'static'")
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t(a int, b int, c int, unique index idx(a, b, c)) PARTITION BY RANGE ( `a` ) (PARTITION `PART_202101` VALUES LESS THAN (202102), PARTITION `PART_202102` VALUES LESS THAN (202103), PARTITION `PART_202103` VALUES LESS THAN (202104));")
+	tk.MustExec("set @@tidb_keep_pruned_conds = 1")
+	tk.MustQuery("explain format = 'brief' select count(*) from t where a = 202103 and b = 1 and c = 1;").Check(testkit.Rows(""+
+		"StreamAgg 1.00 root  funcs:count(1)->Column#5",
+		"└─Point_Get 1.00 root table:t, partition:part_202103, index:idx(a, b, c) "))
+
+	tk.MustExec("set @@tidb_keep_pruned_conds = 0")
+	tk.MustQuery("explain format = 'brief' select count(*) from t where a = 202103 and b = 1 and c = 1;").Check(testkit.Rows(""+
+		"StreamAgg 1.00 root  funcs:count(1)->Column#5",
+		"└─IndexReader 0.01 root  index:Selection",
+		"  └─Selection 0.01 cop[tikv]  eq(test.t.b, 1), eq(test.t.c, 1)",
+		"    └─IndexFullScan 10000.00 cop[tikv] table:t, partition:part_202103, index:idx(a, b, c) keep order:false, stats:pseudo"))
+}

planner/core/rule_partition_processor.go

@@ -870,7 +870,7 @@ func (s *partitionProcessor) processRangePartition(ds *DataSource, pi *model.Par
 	if err != nil {
 		return nil, err
 	}
-	if prunedConds != nil {
+	if !ds.ctx.GetSessionVars().KeepPrunedConds && prunedConds != nil {
 		ds.pushedDownConds = prunedConds
 	}
 	return s.makeUnionAllChildren(ds, pi, used)

sessionctx/variable/session.go

@@ -871,6 +871,9 @@ type SessionVars struct {
 	// EnableStableResultMode if stabilize query results.
 	EnableStableResultMode bool
 
+	// KeepPrunedConds indicates whether to keep the pruning conditions after partition pruning.
+	KeepPrunedConds bool
+
 	// LocalTemporaryTables is *infoschema.LocalTemporaryTables, use interface to avoid circle dependency.
 	// It's nil if there is no local temporary table.
 	LocalTemporaryTables interface{}

sessionctx/variable/sysvar.go

@@ -1781,6 +1781,10 @@ var defaultSysVars = []*SysVar{
 		s.EnableStableResultMode = TiDBOptOn(val)
 		return nil
 	}},
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBKeepPrunedConds, Value: BoolToOnOff(DefTiDBKeepPrunedConds), Hidden: true, Type: TypeBool, SetSession: func(s *SessionVars, val string) error {
+		s.KeepPrunedConds = TiDBOptOn(val)
+		return nil
+	}},
 	{Scope: ScopeNone, Name: TiDBAllowFunctionForExpressionIndex, ReadOnly: true, Value: collectAllowFuncName4ExpressionIndex()},
 }
 

sessionctx/variable/tidb_vars.go

@@ -580,6 +580,9 @@ const (
 
 	// TiDBEnableOrderedResultMode indicates if stabilize query results.
 	TiDBEnableOrderedResultMode = "tidb_enable_ordered_result_mode"
+
+    // TiDBKeepPrunedConds indicates whether to keep the pruning conditions after partition pruning.
+	TiDBKeepPrunedConds = "tidb_keep_pruned_conds"
 )
 
 // TiDB vars that have only global scope
@@ -740,6 +743,7 @@ const (
 	DefTMPTableSize                       = 16777216
 	DefTiDBEnableLocalTxn                 = false
 	DefTiDBEnableOrderedResultMode        = false
+	DefTiDBKeepPrunedConds               = false
 )
 
 // Process global variables.