planner: optimize decorrelate by removing redundant LIMIT when join key is unique

pkg/planner/core/main_test.go

@@ -36,6 +36,7 @@ func TestMain(m *testing.M) {
 	testDataMap.LoadTestSuiteData("testdata", "index_merge_suite", true)
 	testDataMap.LoadTestSuiteData("testdata", "runtime_filter_generator_suite")
 	testDataMap.LoadTestSuiteData("testdata", "plan_cache_suite")
+	testDataMap.LoadTestSuiteData("testdata", "decorrelate_limit_suite", true)
 
 	indexMergeSuiteData = testDataMap["index_merge_suite"]
 	planSuiteUnexportedData = testDataMap["plan_suite_unexported"]
@@ -72,3 +73,11 @@ func GetIndexMergeSuiteData() testdata.TestData {
 func GetRuntimeFilterGeneratorData() testdata.TestData {
 	return testDataMap["runtime_filter_generator_suite"]
 }
+
+func GetDecorrelateLimitSuiteData() testdata.TestData {
+	return testDataMap["decorrelate_limit_suite"]
+}
+
+func GetCascadesSuiteData() testdata.TestData {
+	return testDataMap["cascades_suite"]
+}

pkg/planner/core/plan_test.go

@@ -36,6 +36,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/util/coretestsdk"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/testdata"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/dbterror/plannererrors"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
@@ -732,3 +733,26 @@ func TestImportIntoBuildPlan(t *testing.T) {
 	require.ErrorIs(t, tk.ExecToErr("IMPORT INTO t3 FROM select * from t2"),
 		infoschema.ErrTableNotExists)
 }
+
+func TestDecorrelateLimitOptimization(t *testing.T) {
+	testkit.RunTestUnderCascadesWithDomain(t, func(t *testing.T, testKit *testkit.TestKit, dom *domain.Domain, cascades, caller string) {
+		testKit.MustExec("use test")
+		testKit.MustExec("CREATE TABLE IF NOT EXISTS employees (\n    id INT PRIMARY KEY,\n    name VARCHAR(50),\n    dept_id INT,\n    salary DECIMAL(10, 2),\n    alias VARCHAR(50)\n)")
+		testKit.MustExec("CREATE TABLE IF NOT EXISTS employee_notes (\n    id INT PRIMARY KEY,\n    employee_id INT,\n    note TEXT,\n    created_at TIMESTAMP,\n    INDEX idx_employee_id (employee_id)\n)")
+		var input []string
+		var output []struct {
+			SQL  string
+			Plan []string
+		}
+		decorrelateLimitSuiteData := core.GetDecorrelateLimitSuiteData()
+		decorrelateLimitSuiteData.LoadTestCases(t, &input, &output, cascades, caller)
+		for i, sql := range input {
+			plan := testKit.MustQuery(sql)
+			testdata.OnRecord(func() {
+				output[i].SQL = sql
+				output[i].Plan = testdata.ConvertRowsToStrings(plan.Rows())
+			})
+			plan.Check(testkit.Rows(output[i].Plan...))
+		}
+	})
+}

pkg/planner/core/testdata/decorrelate_limit_suite_in.json

@@ -0,0 +1,22 @@
+[
+  {
+    "Name": "TestDecorrelateLimitOptimization",
+    "Cases": [
+      "EXPLAIN format = 'plan_tree'  SELECT e.name, e.salary , ( select salary FROM employees e2 WHERE e2.id = e.id LIMIT 1 OFFSET 0 ) AS avg_dept_salary FROM employees e WHERE e.dept_id > 1",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary , ( select salary FROM employees e2 WHERE e2.id = e.id LIMIT 1 OFFSET 1 ) AS avg_dept_salary FROM employees e WHERE e.dept_id > 1",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary, ( SELECT e2.salary FROM employees e2 WHERE e2.dept_id = e.dept_id LIMIT 1 OFFSET 0) AS avg_dept_salary FROM employees e WHERE e.dept_id = 1",
+      "EXPLAIN format = 'plan_tree' SELECT e.id, e.name, e.salary, (SELECT en.note FROM employees e2 JOIN employee_notes en ON en.employee_id = e2.id WHERE e2.id = e.id ORDER BY en.created_at DESC LIMIT 1) AS latest_note FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary,(SELECT en.note FROM employees e2 JOIN employee_notes en ON en.employee_id = e2.id LEFT JOIN employees e3 ON e3.id = e2.dept_id WHERE e2.id = e.id LIMIT 1) AS note_multi_join FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary,(SELECT e2.salary FROM employees e2 INNER JOIN employee_notes en ON en.employee_id = e2.id WHERE e2.id = e.id LIMIT 1) AS salary_inner_join FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary, (SELECT AVG(e2.salary)  FROM employees e2  WHERE e2.id = e.id  GROUP BY e2.dept_id  HAVING AVG(e2.salary) > 1000  LIMIT 1) AS avg_salary_having FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary, (SELECT count(e2.dept_id)  FROM employees e2  WHERE e2.id = e.id limit 1) AS distinct_dept_id FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary,(SELECT e2.salary FROM employees e2 WHERE e2.id = e.id AND e2.dept_id IN (     SELECT dept_id     FROM employees e3     WHERE e3.id = e.id     LIMIT 1 ) LIMIT 1) AS salary_nested FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary,(SELECT e2.salary FROM employees e2 WHERE e2.id = e.id AND EXISTS (     SELECT 1     FROM employee_notes en     WHERE en.employee_id = e2.id ) LIMIT 1) AS salary_exists FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary, (SELECT e2.salary  FROM employees e2  WHERE e2.id = e.id  ORDER BY e2.dept_id, e2.salary DESC  LIMIT 1) AS salary_order_multi FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary,  (SELECT DISTINCT e2.dept_id   FROM employees e2   WHERE e2.id = e.id   LIMIT 1) AS distinct_dept_id FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary, (SELECT ROW_NUMBER() OVER (ORDER BY e2.salary DESC)  FROM employees e2  WHERE e2.id = e.id  LIMIT 1) AS row_num FROM employees e",
+      "EXPLAIN format = 'plan_tree' SELECT e.name, e.salary, (SELECT DISTINCT e2.dept_id  FROM employees e2  WHERE e2.id = e.id  LIMIT 1) AS distinct_dept_id FROM employees e"
+    ]
+  }
+]
+