planner: correlate subquery rule

pkg/bindinfo/binding_auto_test.go

@@ -16,12 +16,14 @@ package bindinfo_test
 
 import (
 	"fmt"
+	"slices"
 	"strings"
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/bindinfo"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/auth"
+	"github.com/pingcap/tidb/pkg/sessionctx/vardef"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/testkit/testdata"
 	"github.com/stretchr/testify/require"
@@ -211,6 +213,30 @@ func TestRelevantOptVarsAndFixes(t *testing.T) {
 	}
 }
 
+func TestRelevantOptVarsCorrelateSubquery(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`create table t1 (a int, b int, key(a))`)
+	tk.MustExec(`create table t2 (a int, b int, key(a))`)
+
+	p := parser.New()
+	sql := "select * from t1 where a in (select a from t2)"
+
+	// The variable is recorded as relevant regardless of whether it is ON or OFF,
+	// because the code path where the variable affects plan choice was reached.
+	for _, enabled := range []string{"OFF", "ON"} {
+		tk.MustExec("set tidb_opt_enable_correlate_subquery = " + enabled)
+		p.Reset()
+		stmt, err := p.ParseOneStmt(sql, "", "")
+		require.NoError(t, err)
+		vars, _, err := bindinfo.RecordRelevantOptVarsAndFixes(tk.Session(), stmt)
+		require.NoError(t, err)
+		require.True(t, slices.Contains(vars, vardef.TiDBOptEnableCorrelateSubquery),
+			"enabled=%s: expected %s in recorded vars %v", enabled, vardef.TiDBOptEnableCorrelateSubquery, vars)
+	}
+}
+
 func TestExplainExploreAnalyze(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

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

@@ -7,6 +7,7 @@ go_test(
         "dual_test.go",
         "main_test.go",
         "rule_cdc_join_reorder_test.go",
+        "rule_correlate_test.go",
         "rule_derive_topn_from_window_test.go",
         "rule_eliminate_empty_selection_test.go",
         "rule_eliminate_projection_test.go",
@@ -19,7 +20,7 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 19,
+    shard_count = 20,
     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)
 	testDataMap.LoadTestSuiteData("testdata", "cdc_join_reorder_suite", true)
 
 	opts := []goleak.Option{
@@ -79,6 +80,10 @@ func GetOuterToSemiJoinSuiteData() testdata.TestData {
 	return testDataMap["outer_to_semi_join_suite"]
 }
 
+func GetCorrelateSuiteData() testdata.TestData {
+	return testDataMap["correlate_suite"]
+}
+
 func GetCDCJoinReorderSuiteData() testdata.TestData {
 	return testDataMap["cdc_join_reorder_suite"]
 }

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

@@ -0,0 +1,56 @@
+// 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, t3")
+		tk.MustExec("create table t1 (a int, b int, key(a))")
+		tk.MustExec("create table t2 (a int, b int, key(a))")
+		tk.MustExec("create table t3 (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)")
+		tk.MustExec("insert into t3 values (10,1),(20,2)")
+
+		// 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,17 @@
+[
+  {
+    "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)",
+      "select * from t1 where a in (select a from t2)",
+      "select * from t1 where exists (select 1 from t2)",
+      "select * from t1 where a not in (select a from t2)",
+      "select * from t1 where exists (select 1 from t2 where t2.a > t1.a)",
+      "select * from t1 where exists (select 1 from t2 where t2.a = t1.a and t2.b > t1.b)",
+      "select * from t1 where exists (select /*+ NO_DECORRELATE() */ 1 from t2 where t2.a = t1.a)",
+      "select * from t1 where a in (select t2.a from t2 inner join t3 on t3.a = t2.b where t3.b > 0)",
+      "select * from t1 where a in (select a from t2) order by a limit 10"
+    ]
+  }
+]

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

@@ -0,0 +1,168 @@
+[
+  {
+    "Name": "TestCorrelate",
+    "Cases": [
+      {
+        "SQL": "select * from t1 where exists (select 1 from t2 where t2.a = t1.a)",
+        "Plan": [
+          "Apply 10000.00 root  CARTESIAN semi join, left side:TableReader",
+          "├─TableReader(Build) 10000.00 root  data:TableFullScan",
+          "│ └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─Limit(Probe) 10000.00 root  offset:0, count:1",
+          "  └─IndexReader 10000.00 root  index:Limit",
+          "    └─Limit 10000.00 cop[tikv]  offset:0, count:1",
+          "      └─IndexRangeScan 10000.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",
+          "└─Limit(Probe) 10000.00 root  offset:0, count:1",
+          "  └─IndexReader 10000.00 root  index:Limit",
+          "    └─Limit 10000.00 cop[tikv]  offset:0, count:1",
+          "      └─IndexRangeScan 10000.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"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where a in (select a from t2)",
+        "Plan": [
+          "HashJoin 7992.00 root  semi join, left side:TableReader, equal:[eq(test.t1.a, test.t2.a)]",
+          "├─IndexReader(Build) 9990.00 root  index:IndexFullScan",
+          "│ └─IndexFullScan 9990.00 cop[tikv] table:t2, index:a(a) keep order:false, stats:pseudo",
+          "└─TableReader(Probe) 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"
+        ],
+        "Result": [
+          "1 1",
+          "2 2"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where exists (select 1 from t2)",
+        "Plan": [
+          "TableReader 10000.00 root  data:TableFullScan",
+          "└─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#10, ScalarQueryCol#11, ScalarQueryCol#12, ScalarQueryCol#13",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 1",
+          "2 2",
+          "3 3"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where a not in (select a from t2)",
+        "Plan": [
+          "HashJoin 8000.00 root  Null-aware anti semi join, left side:TableReader, equal:[eq(test.t1.a, test.t2.a)]",
+          "├─IndexReader(Build) 10000.00 root  index:IndexFullScan",
+          "│ └─IndexFullScan 10000.00 cop[tikv] table:t2, index:a(a) keep order:false, stats:pseudo",
+          "└─TableReader(Probe) 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "3 3"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where exists (select 1 from t2 where t2.a > t1.a)",
+        "Plan": [
+          "Apply 10000.00 root  CARTESIAN semi join, left side:TableReader",
+          "├─TableReader(Build) 10000.00 root  data:TableFullScan",
+          "│ └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─Limit(Probe) 10000.00 root  offset:0, count:1",
+          "  └─IndexReader 10000.00 root  index:Limit",
+          "    └─Limit 10000.00 cop[tikv]  offset:0, count:1",
+          "      └─Selection 809900.00 cop[tikv]  gt(test.t2.a, test.t1.a)",
+          "        └─IndexFullScan 1012375.00 cop[tikv] table:t2, index:a(a) keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 1"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where exists (select 1 from t2 where t2.a = t1.a and t2.b > t1.b)",
+        "Plan": [
+          "Apply 10000.00 root  CARTESIAN semi join, left side:TableReader",
+          "├─TableReader(Build) 10000.00 root  data:TableFullScan",
+          "│ └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─Limit(Probe) 10000.00 root  offset:0, count:1",
+          "  └─IndexLookUp 10000.00 root  ",
+          "    ├─IndexRangeScan(Build) 13375.00 cop[tikv] table:t2, index:a(a) range: decided by [eq(test.t2.a, test.t1.a)], keep order:false, stats:pseudo",
+          "    └─Limit(Probe) 10000.00 cop[tikv]  offset:0, count:1",
+          "      └─Selection 10000.00 cop[tikv]  gt(test.t2.b, test.t1.b)",
+          "        └─TableRowIDScan 13375.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 1",
+          "2 2"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where exists (select /*+ NO_DECORRELATE() */ 1 from t2 where t2.a = t1.a)",
+        "Plan": [
+          "Apply 10000.00 root  CARTESIAN semi join, left side:TableReader",
+          "├─TableReader(Build) 10000.00 root  data:TableFullScan",
+          "│ └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "└─Limit(Probe) 10000.00 root  offset:0, count:1",
+          "  └─IndexReader 10000.00 root  index:Limit",
+          "    └─Limit 10000.00 cop[tikv]  offset:0, count:1",
+          "      └─IndexRangeScan 10000.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 a in (select t2.a from t2 inner join t3 on t3.a = t2.b where t3.b > 0)",
+        "Plan": [
+          "HashJoin 7992.00 root  semi join, left side:TableReader, equal:[eq(test.t1.a, test.t2.a)]",
+          "├─HashJoin(Build) 4162.50 root  inner join, equal:[eq(test.t3.a, test.t2.b)]",
+          "│ ├─TableReader(Build) 3330.00 root  data:Selection",
+          "│ │ └─Selection 3330.00 cop[tikv]  gt(test.t3.b, 0), not(isnull(test.t3.a))",
+          "│ │   └─TableFullScan 10000.00 cop[tikv] table:t3 keep order:false, stats:pseudo",
+          "│ └─TableReader(Probe) 9980.01 root  data:Selection",
+          "│   └─Selection 9980.01 cop[tikv]  not(isnull(test.t2.a)), not(isnull(test.t2.b))",
+          "│     └─TableFullScan 10000.00 cop[tikv] table:t2 keep order:false, stats:pseudo",
+          "└─TableReader(Probe) 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"
+        ],
+        "Result": [
+          "1 1",
+          "2 2"
+        ]
+      },
+      {
+        "SQL": "select * from t1 where a in (select a from t2) order by a limit 10",
+        "Plan": [
+          "Limit 10.00 root  offset:0, count:10",
+          "└─MergeJoin 10.00 root  semi join, left side:Projection, left key:test.t1.a, right key:test.t2.a",
+          "  ├─IndexReader(Build) 12.50 root  index:IndexFullScan",
+          "  │ └─IndexFullScan 12.50 cop[tikv] table:t2, index:a(a) keep order:true, stats:pseudo",
+          "  └─Projection(Probe) 12.50 root  test.t1.a, test.t1.b",
+          "    └─IndexLookUp 12.50 root  ",
+          "      ├─IndexFullScan(Build) 12.50 cop[tikv] table:t1, index:a(a) keep order:true, stats:pseudo",
+          "      └─TableRowIDScan(Probe) 12.50 cop[tikv] table:t1 keep order:false, stats:pseudo"
+        ],
+        "Result": [
+          "1 1",
+          "2 2"
+        ]
+      }
+    ]
+  }
+]