planner: Add control flag to keep or remove ORDER BY in subquery (pin…

…gcap#33173)

close pingcap#32900

executor/set_test.go

@@ -593,6 +593,14 @@ func TestSetVar(t *testing.T) {
 	tk.MustExec("set global tidb_ignore_prepared_cache_close_stmt=0")
 	tk.MustQuery("select @@global.tidb_ignore_prepared_cache_close_stmt").Check(testkit.Rows("0"))
 	tk.MustQuery("show global variables like 'tidb_ignore_prepared_cache_close_stmt'").Check(testkit.Rows("tidb_ignore_prepared_cache_close_stmt OFF"))
+
+	// test for tidb_remove_orderby_in_subquery
+	tk.MustQuery("select @@session.tidb_remove_orderby_in_subquery").Check(testkit.Rows("0")) // default value is 0
+	tk.MustExec("set session tidb_remove_orderby_in_subquery=1")
+	tk.MustQuery("select @@session.tidb_remove_orderby_in_subquery").Check(testkit.Rows("1"))
+	tk.MustQuery("select @@global.tidb_remove_orderby_in_subquery").Check(testkit.Rows("0")) // default value is 0
+	tk.MustExec("set global tidb_remove_orderby_in_subquery=1")
+	tk.MustQuery("select @@global.tidb_remove_orderby_in_subquery").Check(testkit.Rows("1"))
 }
 
 func TestTruncateIncorrectIntSessionVar(t *testing.T) {

planner/core/logical_plan_builder.go

@@ -3759,13 +3759,19 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p L
 	}
 
 	if sel.OrderBy != nil {
-		if b.ctx.GetSessionVars().SQLMode.HasOnlyFullGroupBy() {
-			p, err = b.buildSortWithCheck(ctx, p, sel.OrderBy.Items, orderMap, windowMapper, projExprs, oldLen, sel.Distinct)
-		} else {
-			p, err = b.buildSort(ctx, p, sel.OrderBy.Items, orderMap, windowMapper)
-		}
-		if err != nil {
-			return nil, err
+		// We need to keep the ORDER BY clause for the following cases:
+		// 1. The select is top level query, order should be honored
+		// 2. The query has LIMIT clause
+		// 3. The control flag requires keeping ORDER BY explicitly
+		if len(b.selectOffset) == 1 || sel.Limit != nil || !b.ctx.GetSessionVars().RemoveOrderbyInSubquery {
+			if b.ctx.GetSessionVars().SQLMode.HasOnlyFullGroupBy() {
+				p, err = b.buildSortWithCheck(ctx, p, sel.OrderBy.Items, orderMap, windowMapper, projExprs, oldLen, sel.Distinct)
+			} else {
+				p, err = b.buildSort(ctx, p, sel.OrderBy.Items, orderMap, windowMapper)
+			}
+			if err != nil {
+				return nil, err
+			}
 		}
 	}
 

planner/core/logical_plan_test.go

@@ -2123,3 +2123,39 @@ func TestWindowLogicalPlanAmbiguous(t *testing.T) {
 		}
 	}
 }
+
+func TestRemoveOrderbyInSubquery(t *testing.T) {
+	tests := []struct {
+		sql  string
+		best string
+	}{
+		{
+			sql:  "select * from t order by a",
+			best: "DataScan(t)->Projection->Sort",
+		},
+		{
+			sql:  "select (select 1) from t order by a",
+			best: "DataScan(t)->Projection->Sort->Projection",
+		},
+		{
+			sql:  "select count(*) from (select b from t order by a) n",
+			best: "DataScan(t)->Projection->Projection->Aggr(count(1),firstrow(test.t.b))->Projection",
+		},
+		{
+			sql:  "select count(1) from (select b from t order by a limit 1) n",
+			best: "DataScan(t)->Projection->Sort->Limit->Projection->Aggr(count(1),firstrow(test.t.b))->Projection",
+		},
+	}
+
+	s := createPlannerSuite()
+	s.ctx.GetSessionVars().RemoveOrderbyInSubquery = true
+	ctx := context.TODO()
+	for i, tt := range tests {
+		comment := fmt.Sprintf("case:%v sql:%s", i, tt.sql)
+		stmt, err := s.p.ParseOneStmt(tt.sql, "", "")
+		require.NoError(t, err, comment)
+		p, _, err := BuildLogicalPlanForTest(ctx, s.ctx, stmt, s.is)
+		require.NoError(t, err, comment)
+		require.Equal(t, tt.best, ToString(p), comment)
+	}
+}

sessionctx/variable/session.go

@@ -1024,6 +1024,8 @@ type SessionVars struct {
 	BatchPendingTiFlashCount int
 	// RcReadCheckTS indicates if ts check optimization is enabled for current session.
 	RcReadCheckTS bool
+	// RemoveOrderbyInSubquery indicates whether to remove ORDER BY in subquery.
+	RemoveOrderbyInSubquery bool
 }
 
 // InitStatementContext initializes a StatementContext, the object is reused to reduce allocation.
@@ -1259,6 +1261,7 @@ func NewSessionVars() *SessionVars {
 		Rng:                         utilMath.NewWithTime(),
 		StatsLoadSyncWait:           StatsLoadSyncWait.Load(),
 		EnableLegacyInstanceScope:   DefEnableLegacyInstanceScope,
+		RemoveOrderbyInSubquery:     DefTiDBRemoveOrderbyInSubquery,
 	}
 	vars.KVVars = tikvstore.NewVariables(&vars.Killed)
 	vars.Concurrency = Concurrency{

sessionctx/variable/sysvar.go

@@ -1387,6 +1387,10 @@ var defaultSysVars = []*SysVar{
 		s.RcReadCheckTS = TiDBOptOn(val)
 		return nil
 	}},
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBRemoveOrderbyInSubquery, Value: BoolToOnOff(DefTiDBRemoveOrderbyInSubquery), Type: TypeBool, SetSession: func(s *SessionVars, val string) error {
+		s.RemoveOrderbyInSubquery = TiDBOptOn(val)
+		return nil
+	}},
 }
 
 // FeedbackProbability points to the FeedbackProbability in statistics package.

sessionctx/variable/tidb_vars.go

@@ -594,6 +594,9 @@ const (
 	// TiDBEnableOrderedResultMode indicates if stabilize query results.
 	TiDBEnableOrderedResultMode = "tidb_enable_ordered_result_mode"
 
+	// TiDBRemoveOrderbyInSubquery indicates whether to remove ORDER BY in subquery.
+	TiDBRemoveOrderbyInSubquery = "tidb_remove_orderby_in_subquery"
+
 	// TiDBEnablePseudoForOutdatedStats indicates whether use pseudo for outdated stats
 	TiDBEnablePseudoForOutdatedStats = "tidb_enable_pseudo_for_outdated_stats"
 
@@ -823,6 +826,7 @@ const (
 	DefTiDBIgnorePreparedCacheCloseStmt   = false
 	DefTiDBBatchPendingTiFlashCount       = 4000
 	DefRCReadCheckTS                      = false
+	DefTiDBRemoveOrderbyInSubquery        = false
 )
 
 // Process global variables.