planner: eliminate useless scalar subqueries in some scenarios of aggregate queries

pkg/planner/core/casetest/physicalplantest/physical_plan_test.go

@@ -167,6 +167,35 @@ func TestAggEliminator(t *testing.T) {
 	}
 }
 
+func TestIssue45822(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	var input []string
+	var output []struct {
+		SQL  string
+		Best string
+	}
+	planSuiteData := GetPlanSuiteData()
+	planSuiteData.LoadTestCases(t, &input, &output)
+	p := parser.New()
+	is := infoschema.MockInfoSchema([]*model.TableInfo{core.MockSignedTable(), core.MockUnsignedTable()})
+	tk.MustExec("use test")
+	for i, tt := range input {
+		comment := fmt.Sprintf("input: %s", tt)
+		stmt, err := p.ParseOneStmt(tt, "", "")
+		require.NoError(t, err, comment)
+		sc := tk.Session().GetSessionVars().StmtCtx
+		sc.IgnoreTruncate.Store(false)
+		p, _, err := planner.Optimize(context.TODO(), tk.Session(), stmt, is)
+		require.NoError(t, err)
+		testdata.OnRecord(func() {
+			output[i].SQL = tt
+			output[i].Best = core.ToString(p)
+		})
+		require.Equal(t, output[i].Best, core.ToString(p), fmt.Sprintf("input: %s", tt))
+	}
+}
+
 func TestINMJHint(t *testing.T) {
 	var (
 		input  []string

pkg/planner/core/casetest/physicalplantest/testdata/plan_suite_in.json

@@ -354,6 +354,14 @@
       "select max(a) from (select t1.a from t t1 join t t2 on t1.a=t2.a) t"
     ]
   },
+  {
+    "name": "TestIssue45822",
+    "cases": [
+      "select count(*) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t t1) t",
+      "select count(a) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t t1) t",
+      "select count(t) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t t1) t"
+    ]
+  },
   {
     "name": "TestUnmatchedTableInHint",
     "cases": [

pkg/planner/core/casetest/physicalplantest/testdata/plan_suite_out.json

@@ -2320,6 +2320,23 @@
       }
     ]
   },
+  {
+    "Name": "TestIssue45822",
+    "Cases": [
+      {
+        "SQL": "select count(*) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t t1) t;",
+        "Best": "IndexReader(Index(t.f)[[NULL,+inf]]->HashAgg)->HashAgg"
+      },
+      {
+        "SQL": "select count(a) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t t1) t;",
+        "Best": "IndexReader(Index(t.f)[[NULL,+inf]]->HashAgg)->HashAgg"
+      },
+      {
+        "SQL": "select count(t) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t t1) t;",
+        "Best": "Apply{IndexReader(Index(t.f)[[NULL,+inf]])->MergeInnerJoin{TableReader(Table(t))->TableReader(Table(t))}(test.t.a,test.t.a)->Limit}->HashAgg"
+      }
+    ]
+  },
   {
     "Name": "TestUnmatchedTableInHint",
     "Cases": [

pkg/planner/core/rule_column_pruning.go

@@ -26,6 +26,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/planner/util"
+	"github.com/pingcap/tidb/pkg/planner/util/fixcontrol"
 )
 
 type columnPruner struct {
@@ -205,6 +206,55 @@ func (la *LogicalAggregation) PruneColumns(parentUsedCols []*expression.Column,
 			la.SetChildren(childOfChild)
 		}
 	}
+	// Check for Apply elimination
+	if err := la.tryEliminateApply(opt); err != nil {
+		return err
+	}
+	return nil
+}
+
+// tryEliminateApply tries to eliminate apply operator.
+// For SQL like `select count(*) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t1) t;`, we could optimize it to `select count(*) from t1;`.
+// For SQL like `select count(a) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t1) t;`, we could optimize it to `select count(a) from t1;`.
+// For SQL like `select count(t) from (select a,(select t2.b from t t2,t t3 where t2.a=t3.a and t1.a=t2.a limit 1) t from t1) t;`, we couldn't optimize it.
+func (la *LogicalAggregation) tryEliminateApply(opt *logicalOptimizeOp) error {
+	allowEliminateApply := fixcontrol.GetBoolWithDefault(la.SCtx().GetSessionVars().GetOptimizerFixControlMap(), fixcontrol.Fix45822, false)
+	if !allowEliminateApply {
+		return nil
+	}
+	aggChild := la.Children()[0]
+	if proj, isProj := aggChild.(*LogicalProjection); isProj {
+		// If the child node is a LogicalProjection, we need to check whether the child node contains a LogicalApply.
+		projChild := proj.Children()[0]
+		if apply, isApply := projChild.(*LogicalApply); isApply {
+			usedlist := expression.GetUsedList(proj.Schema().Columns, apply.Children()[1].Schema())
+			used := false
+			for i := len(usedlist) - 1; i >= 0; i-- {
+				if usedlist[i] {
+					used = true
+					break
+				}
+			}
+			if !used {
+				applyEliminateTraceStep(projChild, opt)
+				proj.Children()[0] = apply.Children()[0]
+			}
+		}
+	} else if _, isApply := aggChild.(*LogicalApply); isApply {
+		// If the child node is a LogicalApply, we can check if the column is used by the parent node.
+		usedlist := expression.GetUsedList(la.Schema().Columns, aggChild.Children()[1].Schema())
+		used := false
+		for i := len(usedlist) - 1; i >= 0; i-- {
+			if usedlist[i] {
+				used = true
+				break
+			}
+		}
+		if !used {
+			applyEliminateTraceStep(la, opt)
+			la.Children()[0] = aggChild.Children()[0]
+		}
+	}
 	return nil
 }
 
@@ -714,3 +764,15 @@ func (*LogicalCTE) PruneColumns(_ []*expression.Column, _ *logicalOptimizeOp) er
 func (p *LogicalSequence) PruneColumns(parentUsedCols []*expression.Column, opt *logicalOptimizeOp) error {
 	return p.children[len(p.children)-1].PruneColumns(parentUsedCols, opt)
 }
+
+func applyEliminateTraceStep(lp LogicalPlan, opt *logicalOptimizeOp) {
+	action := func() string {
+		buffer := bytes.NewBufferString(
+			fmt.Sprintf("%v_%v is eliminated.", lp.TP(), lp.ID()))
+		return buffer.String()
+	}
+	reason := func() string {
+		return fmt.Sprintf("%v_%v can be eliminated because it hasn't been used by it's parent.", lp.TP(), lp.ID())
+	}
+	opt.appendStepToCurrent(lp.ID(), lp.TP(), reason, action)
+}

pkg/planner/util/fixcontrol/get.go

@@ -34,6 +34,8 @@ const (
 	Fix44855 uint64 = 44855
 	// Fix45132 controls whether to use access range row count to determine access path on the Skyline pruning.
 	Fix45132 uint64 = 45132
+	// Fix45822 controls whether to eliminate apply operator.
+	Fix45822 uint64 = 45822
 )
 
 // GetStr fetches the given key from the fix control map as a string type.