planner: refine PhysicalProjection.ResolveIndices during neighbouring…

… Proj (#8073)

Makefile

@@ -101,6 +101,7 @@ vet:
 clean:
 	$(GO) clean -i ./...
 	rm -rf *.out
+	rm -rf parser/
 
 todo:
 	@grep -n ^[[:space:]]*_[[:space:]]*=[[:space:]][[:alpha:]][[:alnum:]]* */*.go parser/parser.y || true

executor/executor_test.go

@@ -1083,6 +1083,24 @@ func (s *testSuite) TestUnion(c *C) {
 	tk.MustQuery("select a, b from (select a, 0 as d, b from t union all select a, 0 as d, b from t) test;").Check(testkit.Rows("1 2", "1 2"))
 }
 
+func (s *testSuite) TestNeighbouringProj(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+
+	tk.MustExec("drop table if exists t1, t2")
+	tk.MustExec("create table t1(a int, b int)")
+	tk.MustExec("create table t2(a int, b int)")
+	tk.MustExec("insert into t1 value(1, 1), (2, 2)")
+	tk.MustExec("insert into t2 value(1, 1), (2, 2)")
+	tk.MustQuery("select sum(c) from (select t1.a as a, t1.a as c, length(t1.b) from t1  union select a, b, b from t2) t;").Check(testkit.Rows("5"))
+
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a bigint, b bigint, c bigint);")
+	tk.MustExec("insert into t values(1, 1, 1), (2, 2, 2), (3, 3, 3);")
+	rs := tk.MustQuery("select cast(count(a) as signed), a as another, a from t group by a order by cast(count(a) as signed), a limit 10;")
+	rs.Check(testkit.Rows("1 1 1", "1 2 2", "1 3 3"))
+}
+
 func (s *testSuite) TestIn(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("use test")

planner/core/resolve_indices.go

@@ -15,6 +15,7 @@ package core
 
 import (
 	"github.com/pingcap/tidb/expression"
+	"github.com/pingcap/tidb/util/disjointset"
 )
 
 // ResolveIndices implements Plan interface.
@@ -23,6 +24,44 @@ func (p *PhysicalProjection) ResolveIndices() {
 	for i, expr := range p.Exprs {
 		p.Exprs[i] = expr.ResolveIndices(p.children[0].Schema())
 	}
+	childProj, isProj := p.children[0].(*PhysicalProjection)
+	if !isProj {
+		return
+	}
+	refine4NeighbourProj(p, childProj)
+}
+
+// refine4NeighbourProj refines the index for p.Exprs whose type is *Column when
+// there is two neighbouring Projections.
+// This function is introduced because that different childProj.Expr may refer
+// to the same index of childProj.Schema, so we need to keep this relation
+// between the specified expressions in the parent Projection.
+func refine4NeighbourProj(p, childProj *PhysicalProjection) {
+	inputIdx2OutputIdxes := make(map[int][]int)
+	for i, expr := range childProj.Exprs {
+		col, isCol := expr.(*expression.Column)
+		if !isCol {
+			continue
+		}
+		inputIdx2OutputIdxes[col.Index] = append(inputIdx2OutputIdxes[col.Index], i)
+	}
+	childSchemaUnionSet := disjointset.NewIntSet(childProj.schema.Len())
+	for _, outputIdxes := range inputIdx2OutputIdxes {
+		if len(outputIdxes) <= 1 {
+			continue
+		}
+		for i := 1; i < len(outputIdxes); i++ {
+			childSchemaUnionSet.Union(outputIdxes[0], outputIdxes[i])
+		}
+	}
+
+	for _, expr := range p.Exprs {
+		col, isCol := expr.(*expression.Column)
+		if !isCol {
+			continue
+		}
+		col.Index = childSchemaUnionSet.FindRoot(col.Index)
+	}
 }
 
 // ResolveIndices implements Plan interface.