planner: correlate subquery rule

pkg/planner/core/BUILD.bazel

@@ -49,6 +49,7 @@ go_library(
         "rule_aggregation_elimination.go",
         "rule_aggregation_push_down.go",
         "rule_aggregation_skew_rewrite.go",
+        "rule_correlate.go",
         "rule_decorrelate.go",
         "rule_derive_topn_from_window.go",
         "rule_eliminate_empty_selection.go",

pkg/planner/core/casetest/rule/BUILD.bazel

@@ -6,6 +6,7 @@ go_test(
     srcs = [
         "dual_test.go",
         "main_test.go",
+        "rule_correlate_test.go",
         "rule_derive_topn_from_window_test.go",
         "rule_eliminate_empty_selection_test.go",
         "rule_eliminate_projection_test.go",
@@ -18,7 +19,7 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 18,
+    shard_count = 19,
     deps = [
         "//pkg/config",
         "//pkg/domain",

pkg/planner/core/casetest/rule/main_test.go

@@ -35,6 +35,7 @@ func TestMain(m *testing.M) {
 	testDataMap.LoadTestSuiteData("testdata", "predicate_pushdown_suite", true)
 	testDataMap.LoadTestSuiteData("testdata", "predicate_simplification", true)
 	testDataMap.LoadTestSuiteData("testdata", "outer_to_semi_join_suite", true)
+	testDataMap.LoadTestSuiteData("testdata", "correlate_suite", true)
 
 	opts := []goleak.Option{
 		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
@@ -77,3 +78,7 @@ func GetPredicateSimplificationSuiteData() testdata.TestData {
 func GetOuterToSemiJoinSuiteData() testdata.TestData {
 	return testDataMap["outer_to_semi_join_suite"]
 }
+
+func GetCorrelateSuiteData() testdata.TestData {
+	return testDataMap["correlate_suite"]
+}

pkg/planner/core/casetest/rule/rule_correlate_test.go

@@ -0,0 +1,54 @@
+// Copyright 2025 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 rule
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/testdata"
+)
+
+func TestCorrelate(tt *testing.T) {
+	testkit.RunTestUnderCascades(tt, func(t *testing.T, tk *testkit.TestKit, cascades, caller string) {
+		tk.MustExec("use test")
+		tk.MustExec("drop table if exists t1, t2")
+		tk.MustExec("create table t1 (a int, b int, key(a))")
+		tk.MustExec("create table t2 (a int, b int, key(a))")
+		tk.MustExec("insert into t1 values (1,1),(2,2),(3,3)")
+		tk.MustExec("insert into t2 values (1,10),(2,20)")
+
+		// Enable the correlate rule.
+		tk.MustExec("set tidb_opt_enable_correlate_subquery = ON")
+
+		var input []string
+		var output []struct {
+			SQL    string
+			Plan   []string
+			Result []string
+		}
+		suite := GetCorrelateSuiteData()
+		suite.LoadTestCases(t, &input, &output, cascades, caller)
+		for i, sql := range input {
+			testdata.OnRecord(func() {
+				output[i].SQL = sql
+				output[i].Plan = testdata.ConvertRowsToStrings(tk.MustQuery("explain format = 'brief' " + sql).Rows())
+				output[i].Result = testdata.ConvertRowsToStrings(tk.MustQuery(sql).Rows())
+			})
+			tk.MustQuery("explain format = 'brief' " + sql).Check(testkit.Rows(output[i].Plan...))
+			tk.MustQuery(sql).Check(testkit.Rows(output[i].Result...))
+		}
+	})
+}

pkg/planner/core/casetest/rule/testdata/correlate_suite_in.json

@@ -0,0 +1,9 @@
+[
+  {
+    "name": "TestCorrelate",
+    "cases": [
+      "select * from t1 where exists (select 1 from t2 where t2.a = t1.a)",
+      "select * from t1 where not exists (select 1 from t2 where t2.a = t1.a)"
+    ]
+  }
+]

pkg/planner/core/casetest/rule/testdata/correlate_suite_out.json

@@ -0,0 +1,35 @@
+[
+  {
+    "Name": "TestCorrelate",
+    "Cases": [
+      {
+        "SQL": "select * from t1 where exists (select 1 from t2 where t2.a = t1.a)",
+        "Plan": [
+          "Apply 9990.00 root  CARTESIAN semi join, left side:TableReader",
+          "├─TableReader(Build) 9990.00 root  data:Selection",
+          "│ └─Selection 9990.00 cop[tikv]  not(isnull(test.t1.a))",
+          "│   └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─IndexReader(Probe) 99900.00 root  index:IndexRangeScan",
+          "  └─IndexRangeScan 99900.00 cop[tikv] table:t2, index:a(a) range: decided by [eq(test.t2.a, test.t1.a)], keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 1",
+          "2 2"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where not exists (select 1 from t2 where t2.a = t1.a)",
+        "Plan": [
+          "Apply 10000.00 root  CARTESIAN anti semi join, left side:TableReader",
+          "├─TableReader(Build) 10000.00 root  data:TableFullScan",
+          "│ └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─IndexReader(Probe) 100000.00 root  index:IndexRangeScan",
+          "  └─IndexRangeScan 100000.00 cop[tikv] table:t2, index:a(a) range: decided by [eq(test.t2.a, test.t1.a)], keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "3 3"
+        ]
+      }
+    ]
+  }
+]

pkg/planner/core/casetest/rule/testdata/correlate_suite_xut.json

@@ -0,0 +1,35 @@
+[
+  {
+    "Name": "TestCorrelate",
+    "Cases": [
+      {
+        "SQL": "select * from t1 where exists (select 1 from t2 where t2.a = t1.a)",
+        "Plan": [
+          "Apply 9990.00 root  CARTESIAN semi join, left side:TableReader",
+          "├─TableReader(Build) 9990.00 root  data:Selection",
+          "│ └─Selection 9990.00 cop[tikv]  not(isnull(test.t1.a))",
+          "│   └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─IndexReader(Probe) 99900.00 root  index:IndexRangeScan",
+          "  └─IndexRangeScan 99900.00 cop[tikv] table:t2, index:a(a) range: decided by [eq(test.t2.a, test.t1.a)], keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 1",
+          "2 2"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where not exists (select 1 from t2 where t2.a = t1.a)",
+        "Plan": [
+          "Apply 10000.00 root  CARTESIAN anti semi join, left side:TableReader",
+          "├─TableReader(Build) 10000.00 root  data:TableFullScan",
+          "│ └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─IndexReader(Probe) 100000.00 root  index:IndexRangeScan",
+          "  └─IndexRangeScan 100000.00 cop[tikv] table:t2, index:a(a) range: decided by [eq(test.t2.a, test.t1.a)], keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "3 3"
+        ]
+      }
+    ]
+  }
+]

pkg/planner/core/optimizer.go

@@ -110,6 +110,7 @@ var optRuleList = []base.LogicalOptRule{
 	&EliminateUnionAllDualItem{},
 	&EmptySelectionEliminator{},
 	&ResolveExpand{},
+	&CorrelateSolver{},
 }
 
 // Interaction Rule List
@@ -344,10 +345,6 @@ func VolcanoOptimize(ctx context.Context, sctx base.PlanContext, flag uint64, lo
 }
 
 func adjustOptimizationFlags(flag uint64, logic base.LogicalPlan) uint64 {
-	// If there is something after flagPrunColumns, do FlagPruneColumnsAgain.
-	if flag&rule.FlagPruneColumns > 0 && flag-rule.FlagPruneColumns > rule.FlagPruneColumns {
-		flag |= rule.FlagPruneColumnsAgain
-	}
 	if checkStableResultMode(logic.SCtx()) {
 		flag |= rule.FlagStabilizeResults
 	}
@@ -363,6 +360,20 @@ func adjustOptimizationFlags(flag uint64, logic base.LogicalPlan) uint64 {
 	if !logic.SCtx().GetSessionVars().StmtCtx.UseDynamicPruneMode {
 		flag |= rule.FlagPartitionProcessor // apply partition pruning under static mode
 	}
+	if logic.SCtx().GetSessionVars().EnableCorrelateSubquery {
+		flag |= rule.FlagCorrelate
+	}
+	// Recompute FlagPruneColumnsAgain after all conditional flag mutations so
+	// that conditionally-added flags (FlagCorrelate, FlagPartitionProcessor, …)
+	// are taken into account. A second column-prune pass is worthwhile when
+	// any rule above column pruning is enabled.
+	if flag&rule.FlagPruneColumns != 0 {
+		// Mask of all flag bits strictly above FlagPruneColumns.
+		const abovePruneColumns = ^(rule.FlagPruneColumns | (rule.FlagPruneColumns - 1))
+		if flag&abovePruneColumns != 0 {
+			flag |= rule.FlagPruneColumnsAgain
+		}
+	}
 	return flag
 }
 

pkg/planner/core/optimizer_test.go

@@ -15,6 +15,7 @@
 package core
 
 import (
+	"math/bits"
 	"reflect"
 	"strings"
 	"testing"
@@ -27,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/physicalop"
+	"github.com/pingcap/tidb/pkg/planner/core/rule"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util/coretestsdk"
 	"github.com/pingcap/tidb/pkg/types"
@@ -455,3 +457,18 @@ func TestCanTiFlashUseHashJoinV2(t *testing.T) {
 	// can not use hash join v2 due to null eq
 	require.False(t, hashJoin.CanTiFlashUseHashJoinV2(sctx))
 }
+
+func TestOptRuleListFlagAlignment(t *testing.T) {
+	// Each position i in optRuleList is gated by the flag bit 1<<i.
+	// The Flag* constants in rule/logical_rules.go are declared via iota in the
+	// same order. This test catches silent misalignment when a rule or flag is
+	// added/removed without updating the other.
+	//
+	// bits.Len64(lastFlag) == bit-position + 1 == expected list length.
+	numFlags := bits.Len64(rule.FlagCorrelate)
+	require.Equalf(t, numFlags, len(optRuleList),
+		"optRuleList length (%d) does not match Flag* count (%d); "+
+			"did you add a rule without a flag or vice versa? "+
+			"Update both optRuleList and the Flag* iota block in rule/logical_rules.go.",
+		len(optRuleList), numFlags)
+}

pkg/planner/core/rule/logical_rules.go

@@ -45,6 +45,7 @@ const (
 	FlagEliminateUnionAllDualItem
 	FlagEmptySelectionEliminator
 	FlagResolveExpand
+	FlagCorrelate
 )
 
 func setPredicatePushDownFlag(u uint64) uint64 {