From 43cf9a2245b91c647b46816ad3d5424ef90f1070 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 27 Mar 2024 17:32:18 +0800
Subject: [PATCH] planner: apply rule_partition_pruning when optimizing CTE
 under static mode (#51903) (#52148)

close pingcap/tidb#51873
---
 pkg/planner/core/casetest/BUILD.bazel         |  2 +-
 pkg/planner/core/casetest/integration_test.go | 32 +++++++++++++++++++
 pkg/planner/core/optimizer.go                 |  4 ++-
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/pkg/planner/core/casetest/BUILD.bazel b/pkg/planner/core/casetest/BUILD.bazel
index a9cac84c4ba66..037ae24961c47 100644
--- a/pkg/planner/core/casetest/BUILD.bazel
+++ b/pkg/planner/core/casetest/BUILD.bazel
@@ -12,7 +12,7 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 20,
+    shard_count = 21,
     deps = [
         "//pkg/domain",
         "//pkg/parser",
diff --git a/pkg/planner/core/casetest/integration_test.go b/pkg/planner/core/casetest/integration_test.go
index c52f0aa939305..a6c6b0289db00 100644
--- a/pkg/planner/core/casetest/integration_test.go
+++ b/pkg/planner/core/casetest/integration_test.go
@@ -348,6 +348,38 @@ func TestTiFlashFineGrainedShuffle(t *testing.T) {
 	}
 }
 
+func TestIssue51873(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+	tk.MustExec(`CREATE TABLE h1 (
+  id bigint(20) NOT NULL AUTO_INCREMENT,
+  position_date date NOT NULL,
+  asset_id varchar(32) DEFAULT NULL,
+  portfolio_code varchar(50) DEFAULT NULL,
+  PRIMARY KEY (id,position_date) /*T![clustered_index] NONCLUSTERED */,
+  UNIQUE KEY uidx_posi_asset_balance_key (position_date,portfolio_code,asset_id)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin AUTO_INCREMENT=30002
+PARTITION BY RANGE COLUMNS(position_date)
+(PARTITION p202401 VALUES LESS THAN ('2024-02-01'))`)
+	tk.MustExec(`create table h2 like h1`)
+	tk.MustExec(`insert into h1 values(1,'2024-01-01',1,1)`)
+	tk.MustExec(`insert into h2 values(1,'2024-01-01',1,1)`)
+	tk.MustExec(`analyze table h1`)
+	tk.MustExec(`set @@tidb_skip_missing_partition_stats=0`)
+	tk.MustQuery(`with assetBalance AS
+    (SELECT asset_id, portfolio_code FROM h1 pab WHERE pab.position_date = '2024-01-01' ),
+cashBalance AS (SELECT portfolio_code, asset_id
+    FROM h2 pcb WHERE pcb.position_date = '2024-01-01' ),
+assetIdList AS (SELECT DISTINCT asset_id AS assetId
+    FROM assetBalance )
+SELECT main.portfolioCode
+FROM (SELECT DISTINCT balance.portfolio_code AS portfolioCode
+    FROM assetBalance balance
+    LEFT JOIN assetIdList
+        ON balance.asset_id = assetIdList.assetId ) main`).Check(testkit.Rows("1"))
+}
+
 func TestIssue50926(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)
diff --git a/pkg/planner/core/optimizer.go b/pkg/planner/core/optimizer.go
index fada1322343f6..ace88fa1526b9 100644
--- a/pkg/planner/core/optimizer.go
+++ b/pkg/planner/core/optimizer.go
@@ -362,7 +362,9 @@ func adjustOptimizationFlags(flag uint64, logic LogicalPlan) uint64 {
 	}
 	flag |= flagCollectPredicateColumnsPoint
 	flag |= flagSyncWaitStatsLoadPoint
-
+	if !logic.SCtx().GetSessionVars().StmtCtx.UseDynamicPruneMode {
+		flag |= flagPartitionProcessor // apply partition pruning under static mode
+	}
 	return flag
 }