planner: when pushing topN down below proj, check whether `topN.B…

…yItems` contains a column(with ID=0) generated by `proj` (#40593) (#41178)

close #40535

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"
@@ -1151,3 +1152,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())
+}

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