From 4e373cffeb159da9dd35f403feebd9b6b19be1ba Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Mon, 29 Jul 2024 20:44:48 +0800
Subject: [PATCH] planner: add newly created col for window projection (#52378)
 (#55003)

close pingcap/tidb#42734
---
 pkg/planner/core/casetest/windows/BUILD.bazel |  3 +-
 .../windows/widow_with_exist_subquery_test.go | 77 +++++++++++++++++++
 pkg/planner/core/logical_plan_builder.go      |  8 ++
 3 files changed, 87 insertions(+), 1 deletion(-)
 create mode 100644 pkg/planner/core/casetest/windows/widow_with_exist_subquery_test.go

diff --git a/pkg/planner/core/casetest/windows/BUILD.bazel b/pkg/planner/core/casetest/windows/BUILD.bazel
index bdac1460f976b..ae898167064ab 100644
--- a/pkg/planner/core/casetest/windows/BUILD.bazel
+++ b/pkg/planner/core/casetest/windows/BUILD.bazel
@@ -5,11 +5,12 @@ go_test(
     timeout = "short",
     srcs = [
         "main_test.go",
+        "widow_with_exist_subquery_test.go",
         "window_push_down_test.go",
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 4,
+    shard_count = 5,
     deps = [
         "//pkg/domain",
         "//pkg/planner/core/internal",
diff --git a/pkg/planner/core/casetest/windows/widow_with_exist_subquery_test.go b/pkg/planner/core/casetest/windows/widow_with_exist_subquery_test.go
new file mode 100644
index 0000000000000..e91c73f5274cd
--- /dev/null
+++ b/pkg/planner/core/casetest/windows/widow_with_exist_subquery_test.go
@@ -0,0 +1,77 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package windows
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/testkit"
+)
+
+func TestWindowWithCorrelatedSubQuery(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("CREATE TABLE temperature_data (temperature double);")
+	tk.MustExec("CREATE TABLE humidity_data (humidity double);")
+	tk.MustExec("CREATE TABLE weather_report (report_id double, report_date varchar(100));")
+
+	tk.MustExec("INSERT INTO temperature_data VALUES (1.0);")
+	tk.MustExec("INSERT INTO humidity_data VALUES (0.5);")
+	tk.MustExec("INSERT INTO weather_report VALUES (2.0, 'test');")
+
+	result := tk.MustQuery(`
+   SELECT
+     EXISTS (
+       SELECT
+         FIRST_VALUE(temp_data.temperature) OVER weather_window AS first_temperature,
+         MIN(report_data.report_id) OVER weather_window AS min_report_id
+       FROM
+         temperature_data AS temp_data
+       WINDOW weather_window AS (
+         PARTITION BY EXISTS (
+           SELECT
+             report_data.report_date AS report_date
+           FROM
+             humidity_data AS humidity_data
+           WHERE temp_data.temperature >= humidity_data.humidity
+         )
+       )
+     ) AS is_exist
+   FROM
+     weather_report AS report_data;
+ `)
+
+	result.Check(testkit.Rows("1"))
+
+	result = tk.MustQuery(`
+   SELECT
+     EXISTS (
+       SELECT
+         FIRST_VALUE(temp_data.temperature) OVER weather_window AS first_temperature,
+         MIN(report_data.report_id) OVER weather_window AS min_report_id
+       FROM
+         temperature_data AS temp_data
+       WINDOW weather_window AS (
+         PARTITION BY temp_data.temperature 
+       )
+     ) AS is_exist
+   FROM
+     weather_report AS report_data;
+ `)
+
+	result.Check(testkit.Rows("1"))
+}
diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go
index c08c0cff3ab3f..5f515f2f34ed6 100644
--- a/pkg/planner/core/logical_plan_builder.go
+++ b/pkg/planner/core/logical_plan_builder.go
@@ -6500,6 +6500,14 @@ func (b *PlanBuilder) buildByItemsForWindow(
 		}
 		if col, ok := it.(*expression.Column); ok {
 			retItems = append(retItems, property.SortItem{Col: col, Desc: item.Desc})
+			// We need to attempt to add this column because a subquery may be created during the expression rewrite process.
+			// Therefore, we need to ensure that the column from the newly created query plan is added.
+			// If the column is already in the schema, we don't need to add it again.
+			if !proj.schema.Contains(col) {
+				proj.Exprs = append(proj.Exprs, col)
+				proj.names = append(proj.names, types.EmptyName)
+				proj.schema.Append(col)
+			}
 			continue
 		}
 		proj.Exprs = append(proj.Exprs, it)