From b363c380034deea8506f7f0f836eac10e020182a Mon Sep 17 00:00:00 2001
From: Kenan Yao <cauchy1992@gmail.com>
Date: Thu, 17 Jun 2021 01:16:38 +0800
Subject: [PATCH] cherry pick #25345 to release-5.1

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 planner/core/integration_test.go    |  9 +++++++++
 planner/core/rule_column_pruning.go | 10 ++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index 394d5c02b0c1e..9b80c58b61d18 100644
--- a/planner/core/integration_test.go
+++ b/planner/core/integration_test.go
@@ -3772,6 +3772,15 @@ func (s *testIntegrationSuite) TestIssue24281(c *C) {
 		"UNION select 1 as v1, 2 as v2")
 }
 
+func (s *testIntegrationSuite) TestLimitWindowColPrune(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int)")
+	tk.MustExec("insert into t values(1)")
+	tk.MustQuery("select count(a) f1, row_number() over (order by count(a)) as f2 from t limit 1").Check(testkit.Rows("1 1"))
+}
+
 func (s *testIntegrationSuite) TestIncrementalAnalyzeStatsVer2(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("use test")
diff --git a/planner/core/rule_column_pruning.go b/planner/core/rule_column_pruning.go
index a6b16603b1585..212e9b6f4d242 100644
--- a/planner/core/rule_column_pruning.go
+++ b/planner/core/rule_column_pruning.go
@@ -443,8 +443,14 @@ func (p *LogicalLimit) PruneColumns(parentUsedCols []*expression.Column) error {
 		return nil
 	}
 
-	p.inlineProjection(parentUsedCols)
-	return p.children[0].PruneColumns(parentUsedCols)
+	savedUsedCols := make([]*expression.Column, len(parentUsedCols))
+	copy(savedUsedCols, parentUsedCols)
+	if err := p.children[0].PruneColumns(parentUsedCols); err != nil {
+		return err
+	}
+	p.schema = nil
+	p.inlineProjection(savedUsedCols)
+	return nil
 }
 
 func (*columnPruner) name() string {