planner: fix order by sub-query couldn't find outer correlated columns (#33640)

close pingcap/tidb#26945

cmd/explaintest/r/explain_easy.result

@@ -822,12 +822,12 @@ Projection	1.00	root		Column#7
       └─TableDual(Probe)	1.00	root		rows:1
 explain format = 'brief' select count(a) from t group by b order by (select count(a));
 id	estRows	task	access object	operator info
-Sort	8000.00	root		Column#4
+Sort	8000.00	root		Column#5
 └─HashJoin	8000.00	root		CARTESIAN left outer join
   ├─TableDual(Build)	1.00	root		rows:1
-  └─HashAgg(Probe)	8000.00	root		group by:test.t.b, funcs:count(Column#8)->Column#4
+  └─HashAgg(Probe)	8000.00	root		group by:test.t.b, funcs:count(Column#9)->Column#5
     └─TableReader	8000.00	root		data:HashAgg
-      └─HashAgg	8000.00	cop[tikv]		group by:test.t.b, funcs:count(test.t.a)->Column#8
+      └─HashAgg	8000.00	cop[tikv]		group by:test.t.b, funcs:count(test.t.a)->Column#9
         └─TableFullScan	10000.00	cop[tikv]	table:t	keep order:false, stats:pseudo
 explain format = 'brief' select (select sum(count(a))) from t;
 id	estRows	task	access object	operator info
@@ -841,15 +841,15 @@ Projection	1.00	root		Column#5
       └─TableDual	1.00	root		rows:1
 explain format = 'brief' select sum(a), (select sum(a)), count(a) from t group by b order by (select count(a));
 id	estRows	task	access object	operator info
-Projection	8000.00	root		Column#4, Column#4, Column#5
-└─Sort	8000.00	root		Column#5
+Projection	8000.00	root		Column#5, Column#5, Column#6
+└─Sort	8000.00	root		Column#6
   └─HashJoin	8000.00	root		CARTESIAN left outer join
     ├─TableDual(Build)	1.00	root		rows:1
     └─HashJoin(Probe)	8000.00	root		CARTESIAN left outer join
       ├─TableDual(Build)	1.00	root		rows:1
-      └─HashAgg(Probe)	8000.00	root		group by:test.t.b, funcs:sum(Column#13)->Column#4, funcs:count(Column#14)->Column#5
+      └─HashAgg(Probe)	8000.00	root		group by:test.t.b, funcs:sum(Column#14)->Column#5, funcs:count(Column#15)->Column#6
         └─TableReader	8000.00	root		data:HashAgg
-          └─HashAgg	8000.00	cop[tikv]		group by:test.t.b, funcs:sum(test.t.a)->Column#13, funcs:count(test.t.a)->Column#14
+          └─HashAgg	8000.00	cop[tikv]		group by:test.t.b, funcs:sum(test.t.a)->Column#14, funcs:count(test.t.a)->Column#15
             └─TableFullScan	10000.00	cop[tikv]	table:t	keep order:false, stats:pseudo
 drop table if exists t;
 create table t(a tinyint, b smallint, c mediumint, d int, e bigint);

planner/core/integration_test.go

@@ -6449,3 +6449,27 @@ func TestIssue33042(t *testing.T) {
 		),
 	)
 }
+
+func TestIssue29663(t *testing.T) {
+	store, _, clean := testkit.CreateMockStoreAndDomain(t)
+	defer clean()
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t1")
+	tk.MustExec("drop table if exists t2")
+	tk.MustExec("create table t1 (a int, b int)")
+	tk.MustExec("create table t2 (c int, d int)")
+	tk.MustExec("insert into t1 values(1, 1), (1,2),(2,1),(2,2)")
+	tk.MustExec("insert into t2 values(1, 3), (1,4),(2,5),(2,6)")
+
+	tk.MustQuery("explain select one.a from t1 one order by (select two.d from t2 two where two.c = one.b)").Check(testkit.Rows(
+		"Projection_16 10000.00 root  test.t1.a",
+		"└─Sort_17 10000.00 root  test.t2.d",
+		"  └─Apply_20 10000.00 root  CARTESIAN left outer join",
+		"    ├─TableReader_22(Build) 10000.00 root  data:TableFullScan_21",
+		"    │ └─TableFullScan_21 10000.00 cop[tikv] table:one keep order:false, stats:pseudo",
+		"    └─MaxOneRow_23(Probe) 1.00 root  ",
+		"      └─TableReader_26 2.00 root  data:Selection_25",
+		"        └─Selection_25 2.00 cop[tikv]  eq(test.t2.c, test.t1.b)",
+		"          └─TableFullScan_24 2000.00 cop[tikv] table:two keep order:false, stats:pseudo"))
+}

planner/core/logical_plan_builder.go

@@ -2179,7 +2179,7 @@ func (a *havingWindowAndOrderbyExprResolver) Leave(n ast.Node) (node ast.Node, o
 // resolveHavingAndOrderBy will process aggregate functions and resolve the columns that don't exist in select fields.
 // If we found some columns that are not in select fields, we will append it to select fields and update the colMapper.
 // When we rewrite the order by / having expression, we will find column in map at first.
-func (b *PlanBuilder) resolveHavingAndOrderBy(sel *ast.SelectStmt, p LogicalPlan) (
+func (b *PlanBuilder) resolveHavingAndOrderBy(ctx context.Context, sel *ast.SelectStmt, p LogicalPlan) (
 	map[*ast.AggregateFuncExpr]int, map[*ast.AggregateFuncExpr]int, error) {
 	extractor := &havingWindowAndOrderbyExprResolver{
 		p:            p,
@@ -2219,6 +2219,50 @@ func (b *PlanBuilder) resolveHavingAndOrderBy(sel *ast.SelectStmt, p LogicalPlan
 		}
 	}
 	sel.Fields.Fields = extractor.selectFields
+	// this part is used to fetch correlated column from sub-query item in order-by clause, and append the origin
+	// auxiliary select filed in select list, otherwise, sub-query itself won't get the name resolved in outer schema.
+	if sel.OrderBy != nil {
+		for _, byItem := range sel.OrderBy.Items {
+			if _, ok := byItem.Expr.(*ast.SubqueryExpr); ok {
+				// correlated agg will be extracted completely latter.
+				_, np, err := b.rewrite(ctx, byItem.Expr, p, nil, true)
+				if err != nil {
+					return nil, nil, errors.Trace(err)
+				}
+				correlatedCols := ExtractCorrelatedCols4LogicalPlan(np)
+				for _, cone := range correlatedCols {
+					var colName *ast.ColumnName
+					for idx, pone := range p.Schema().Columns {
+						if cone.UniqueID == pone.UniqueID {
+							pname := p.OutputNames()[idx]
+							colName = &ast.ColumnName{
+								Schema: pname.DBName,
+								Table:  pname.TblName,
+								Name:   pname.ColName,
+							}
+							break
+						}
+					}
+					if colName != nil {
+						columnNameExpr := &ast.ColumnNameExpr{Name: colName}
+						for _, field := range sel.Fields.Fields {
+							if c, ok := field.Expr.(*ast.ColumnNameExpr); ok && colMatch(c.Name, columnNameExpr.Name) {
+								// deduplicate select fields: don't append it once it already has one.
+								columnNameExpr = nil
+								break
+							}
+						}
+						if columnNameExpr != nil {
+							sel.Fields.Fields = append(sel.Fields.Fields, &ast.SelectField{
+								Auxiliary: true,
+								Expr:      columnNameExpr,
+							})
+						}
+					}
+				}
+			}
+		}
+	}
 	return havingAggMapper, extractor.aggMapper, nil
 }
 
@@ -2402,7 +2446,7 @@ func (r *correlatedAggregateResolver) resolveSelect(sel *ast.SelectStmt) (err er
 		}
 	}
 
-	_, _, err = r.b.resolveHavingAndOrderBy(sel, p)
+	_, _, err = r.b.resolveHavingAndOrderBy(r.ctx, sel, p)
 	if err != nil {
 		return err
 	}
@@ -3627,7 +3671,7 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p L
 	// We must resolve having and order by clause before build projection,
 	// because when the query is "select a+1 as b from t having sum(b) < 0", we must replace sum(b) to sum(a+1),
 	// which only can be done before building projection and extracting Agg functions.
-	havingMap, orderMap, err = b.resolveHavingAndOrderBy(sel, p)
+	havingMap, orderMap, err = b.resolveHavingAndOrderBy(ctx, sel, p)
 	if err != nil {
 		return nil, err
 	}

planner/core/logical_plan_test.go

@@ -1507,7 +1507,7 @@ func TestUnion(t *testing.T) {
 			require.Error(t, err)
 			continue
 		}
-		require.NoError(t, err)
+		require.NoError(t, err, comment)
 		p := plan.(LogicalPlan)
 		p, err = logicalOptimize(ctx, builder.optFlag, p)
 		testdata.OnRecord(func() {