planner: fix the issue that the optimizer terminates the optimization process for DataSource too early (#48186)

planner/core/find_best_task.go

@@ -30,6 +30,7 @@ import (
  "github.com/pingcap/tidb/planner/util"
  "github.com/pingcap/tidb/sessionctx"
  "github.com/pingcap/tidb/sessionctx/stmtctx"
+ "github.com/pingcap/tidb/sessionctx/variable"
  "github.com/pingcap/tidb/statistics"
  "github.com/pingcap/tidb/types"
  tidbutil "github.com/pingcap/tidb/util"
@@ -841,6 +842,14 @@ func (ds *DataSource) isPointGetConvertableSchema() bool {
  return true
 }
 
+// exploreEnforcedPlan determines whether to explore enforced plans for this DataSource if it has already found an unenforced plan.
+// See #46177 for more information.
+func (ds *DataSource) exploreEnforcedPlan() bool {
+ // default value is false to keep it compatible with previous versions.
+ fixValue, ok := ds.ctx.GetSessionVars().GetOptimizerFixControlValue(variable.TiDBOptFixControl46177)
+ return ok && variable.TiDBOptOn(fixValue)
+}
+
 // findBestTask implements the PhysicalPlan interface.
 // It will enumerate all the available indices and choose a plan with least cost.
 func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter *PlanCounterTp, opt *physicalOptimizeOp) (t task, cntPlan int64, err error) {
@@ -881,23 +890,25 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
  return
  }
  var cnt int64
+ var unenforcedTask task
  // If prop.CanAddEnforcer is true, the prop.SortItems need to be set nil for ds.findBestTask.
  // Before function return, reset it for enforcing task prop and storing map<prop,task>.
  oldProp := prop.CloneEssentialFields()
  if prop.CanAddEnforcer {
  // First, get the bestTask without enforced prop
  prop.CanAddEnforcer = false
- t, cnt, err = ds.findBestTask(prop, planCounter, opt)
+ unenforcedTask, cnt, err = ds.findBestTask(prop, planCounter, opt)
  if err != nil {
  return nil, 0, err
  }
- prop.CanAddEnforcer = true
- if t != invalidTask {
- ds.storeTask(prop, t)
- cntPlan = cnt
- return
+ if !unenforcedTask.invalid() && !ds.exploreEnforcedPlan() {
+ ds.storeTask(prop, unenforcedTask)
+ return unenforcedTask, cnt, nil
  }
- // Next, get the bestTask with enforced prop
+
+ // Then, explore the bestTask with enforced prop
+ prop.CanAddEnforcer = true
+ cntPlan += cnt
  prop.SortItems = []property.SortItem{}
  prop.MPPPartitionTp = property.AnyType
  } else if prop.MPPPartitionTp != property.AnyType {
@@ -912,6 +923,18 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
  t = enforceProperty(prop, t, ds.basePlan.ctx)
  prop.CanAddEnforcer = true
  }
+
+ if unenforcedTask != nil && !unenforcedTask.invalid() {
+ curIsBest, cerr := compareTaskCost(ds.SCtx(), unenforcedTask, t, opt)
+ if cerr != nil {
+ err = cerr
+ return
+ }
+ if curIsBest {
+ t = unenforcedTask
+ }
+ }
+
  ds.storeTask(prop, t)
  if ds.SampleInfo != nil && !t.invalid() {
  if _, ok := t.plan().(*PhysicalTableSample); !ok {

planner/core/integration_test.go

@@ -5278,6 +5278,53 @@ func TestVirtualExprPushDown(t *testing.T) {
  tk.MustQuery("explain select * from t where c2 > 1;").CheckAt([]int{0, 2, 4}, rows)
 }
 
+func TestIssue46177(t *testing.T) {
+ store := testkit.CreateMockStore(t)
+ tk := testkit.NewTestKit(t, store)
+ tk.MustExec(`use test`)
+ tk.MustExec(` CREATE TABLE sbtest (
+ id int(10) unsigned NOT NULL AUTO_INCREMENT,
+ k int(10) unsigned NOT NULL DEFAULT '0',
+ c char(120) NOT NULL DEFAULT '',
+ pad char(60) NOT NULL DEFAULT '',
+ PRIMARY KEY (id) /*T![clustered_index] CLUSTERED */,
+ KEY k (k)
+)`)
+
+ // cannot choose the best plan with RangeScan.
+ tk.MustExec(`set @@tidb_opt_fix_control = '46177:off'`)
+ tk.MustQuery(`explain format='brief' select row_number() over(order by a.k) from (select * from sbtest where id<10) a`).Check(testkit.Rows(
+ `Projection 10.00 root Column#6`,
+ `└─Window 10.00 root row_number()->Column#6 over(order by test.sbtest.k rows between current row and current row)`,
+ ` └─IndexReader 10.00 root index:Selection`,
+ ` └─Selection 10.00 cop[tikv] lt(test.sbtest.id, 10)`,
+ ` └─IndexFullScan 10000.00 cop[tikv] table:sbtest, index:k(k) keep order:true, stats:pseudo`))
+
+ tk.MustExec(`set @@tidb_opt_fix_control = '46177:on'`)
+ tk.MustQuery(`explain format='brief' select row_number() over(order by a.k) from (select * from sbtest where id<10) a`).Check(testkit.Rows(
+ `Projection 10.00 root Column#6`,
+ `└─Window 10.00 root row_number()->Column#6 over(order by test.sbtest.k rows between current row and current row)`,
+ ` └─Sort 10.00 root test.sbtest.k`,
+ ` └─TableReader 10.00 root data:TableRangeScan`,
+ ` └─TableRangeScan 10.00 cop[tikv] table:sbtest range:[0,10), keep order:false, stats:pseudo`))
+
+ // cannot choose the range scan plan.
+ tk.MustExec(`set @@tidb_opt_fix_control = '46177:off'`)
+ tk.MustQuery(`explain format='brief' select /*+ stream_agg() */ count(1) from sbtest where id<1 group by k`).Check(testkit.Rows(
+ `StreamAgg 1.00 root group by:test.sbtest.k, funcs:count(Column#6)->Column#5`,
+ `└─IndexReader 1.00 root index:StreamAgg`,
+ ` └─StreamAgg 1.00 cop[tikv] group by:test.sbtest.k, funcs:count(1)->Column#6`,
+ ` └─Selection 1.00 cop[tikv] lt(test.sbtest.id, 1)`,
+ ` └─IndexFullScan 10000.00 cop[tikv] table:sbtest, index:k(k) keep order:true, stats:pseudo`))
+
+ tk.MustExec(`set @@tidb_opt_fix_control = '46177:on'`)
+ tk.MustQuery(`explain format='brief' select /*+ stream_agg() */ count(1) from sbtest where id<1 group by k`).Check(testkit.Rows(
+ `StreamAgg 1.00 root group by:test.sbtest.k, funcs:count(1)->Column#5`,
+ `└─Sort 1.00 root test.sbtest.k`,
+ ` └─TableReader 1.00 root data:TableRangeScan`,
+ ` └─TableRangeScan 1.00 cop[tikv] table:sbtest range:[0,1), keep order:false, stats:pseudo`))
+}
+
 // https://github.com/pingcap/tidb/issues/41458
 func TestIssue41458(t *testing.T) {
  store := testkit.CreateMockStore(t)

sessionctx/variable/session.go

@@ -1498,6 +1498,8 @@ var (
  // TiDBOptFixControl44855 controls whether to use a more accurate upper bound when estimating row count of index
  // range scan under inner side of index join.
  TiDBOptFixControl44855 uint64 = 44855
+ // TiDBOptFixControl46177 controls whether to explore enforced plans for DataSource if it has already found an unenforced plan.
+ TiDBOptFixControl46177 uint64 = 46177
 )
 
 // GetOptimizerFixControlValue returns the specified value of the optimizer fix control.