From 2aad38831e0de638a923c5b731b73213f82af1e5 Mon Sep 17 00:00:00 2001
From: Jk Xu <54522439+Dousir9@users.noreply.github.com>
Date: Tue, 7 Feb 2023 18:45:58 -0600
Subject: [PATCH] planner: when pushing `topN` down below `proj`, check whether
 `topN.ByItems` contains a column(with ID=0) generated by `proj` (#40593)

close pingcap/tidb#40535
---
 planner/core/plan_test.go           | 13 +++++++++++++
 planner/core/rule_topn_push_down.go | 14 ++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/planner/core/plan_test.go b/planner/core/plan_test.go
index f1c65fcd9fb4d..a449cee57e69f 100644
--- a/planner/core/plan_test.go
+++ b/planner/core/plan_test.go
@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/expression/aggregation"
+	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/parser/mysql"
@@ -1142,3 +1143,15 @@ func TestJSONPlanInExplain(t *testing.T) {
 		}
 	}
 }
+
+func TestIssue40535(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	var cfg kv.InjectionConfig
+	tk := testkit.NewTestKit(t, kv.NewInjectedStore(store, &cfg))
+	tk.MustExec("use test;")
+	tk.MustExec("drop table if exists t1; drop table if exists t2;")
+	tk.MustExec("CREATE TABLE `t1`(`c1` bigint(20) NOT NULL DEFAULT '-2312745469307452950', `c2` datetime DEFAULT '5316-02-03 06:54:49', `c3` tinyblob DEFAULT NULL, PRIMARY KEY (`c1`) /*T![clustered_index] CLUSTERED */) ENGINE=InnoDB DEFAULT CHARSET=utf8 COLLATE=utf8_bin;")
+	tk.MustExec("CREATE TABLE `t2`(`c1` set('kn8pu','7et','vekx6','v3','liwrh','q14','1met','nnd5i','5o0','8cz','l') DEFAULT '7et,vekx6,liwrh,q14,1met', `c2` float DEFAULT '1.683167', KEY `k1` (`c2`,`c1`), KEY `k2` (`c2`)) ENGINE=InnoDB DEFAULT CHARSET=gbk COLLATE=gbk_chinese_ci;")
+	tk.MustExec("(select /*+ agg_to_cop()*/ locate(t1.c3, t1.c3) as r0, t1.c3 as r1 from t1 where not( IsNull(t1.c1)) order by r0,r1) union all (select concat_ws(',', t2.c2, t2.c1) as r0, t2.c1 as r1 from t2 order by r0, r1) order by 1 limit 273;")
+	require.Empty(t, tk.Session().LastMessage())
+}
diff --git a/planner/core/rule_topn_push_down.go b/planner/core/rule_topn_push_down.go
index 5dacac7579caa..2b4998049911a 100644
--- a/planner/core/rule_topn_push_down.go
+++ b/planner/core/rule_topn_push_down.go
@@ -136,6 +136,20 @@ func (p *LogicalProjection) pushDownTopN(topN *LogicalTopN, opt *logicalOptimize
 				topN.ByItems = append(topN.ByItems[:i], topN.ByItems[i+1:]...)
 			}
 		}
+
+		// if topN.ByItems contains a column(with ID=0) generated by projection, projection will prevent the optimizer from pushing topN down.
+		for _, by := range topN.ByItems {
+			cols := expression.ExtractColumns(by.Expr)
+			for _, col := range cols {
+				if col.ID == 0 && p.Schema().Contains(col) {
+					// check whether the column is generated by projection
+					if !p.children[0].Schema().Contains(col) {
+						p.children[0] = p.children[0].pushDownTopN(nil, opt)
+						return topN.setChild(p, opt)
+					}
+				}
+			}
+		}
 	}
 	p.children[0] = p.children[0].pushDownTopN(topN, opt)
 	return p