plan: fix a problem caused by union's schema

executor/aggregate_test.go

@@ -320,7 +320,7 @@ func (s *testSuite) TestAggregation(c *C) {
 	tk.MustExec("insert into t value(1), (0)")
 	tk.MustQuery("select a from t group by 1")
 	// This result is compatible with MySQL, the readable result is shown in the next case.
-	result = tk.MustQuery("select max(a) from t group by a")
+	result = tk.MustQuery("select max(a) from t group by a order by a")
 	result.Check(testkit.Rows(string([]byte{0x0}), string([]byte{0x1})))
 	result = tk.MustQuery("select cast(a as signed) as idx, cast(max(a) as signed),  cast(min(a) as signed) from t group by 1 order by idx")
 	result.Check(testkit.Rows("0 0 0", "1 1 1"))

plan/exhaust_physical_plans.go

@@ -881,6 +881,7 @@ func (p *LogicalUnionAll) exhaustPhysicalPlans(prop *property.PhysicalProperty)
 		chReqProps = append(chReqProps, &property.PhysicalProperty{ExpectedCnt: prop.ExpectedCnt})
 	}
 	ua := PhysicalUnionAll{}.init(p.ctx, p.stats.ScaleByExpectCnt(prop.ExpectedCnt), chReqProps...)
+	ua.SetSchema(p.Schema())
 	return []PhysicalPlan{ua}
 }
 

plan/logical_plan_builder.go

@@ -626,29 +626,32 @@ func unionJoinFieldType(a, b *types.FieldType) *types.FieldType {
 }
 
 func (b *planBuilder) buildProjection4Union(u *LogicalUnionAll) {
-	unionSchema := u.children[0].Schema().Clone()
+	unionCols := make([]*expression.Column, 0, u.children[0].Schema().Len())
 
 	// Infer union result types by its children's schema.
-	for i, col := range unionSchema.Columns {
-		var resultTp *types.FieldType
-		for j, child := range u.children {
-			childTp := child.Schema().Columns[i].RetType
-			if j == 0 {
-				resultTp = childTp
-			} else {
-				resultTp = unionJoinFieldType(resultTp, childTp)
-			}
-		}
-		unionSchema.Columns[i] = col.Clone().(*expression.Column)
-		unionSchema.Columns[i].RetType = resultTp
-		unionSchema.Columns[i].DBName = model.NewCIStr("")
+	for i, col := range u.children[0].Schema().Columns {
+		resultTp := col.RetType
+		for j := 1; j < len(u.children); j++ {
+			childTp := u.children[j].Schema().Columns[i].RetType
+			resultTp = unionJoinFieldType(resultTp, childTp)
+		}
+		unionCols = append(unionCols, &expression.Column{
+			ColName:  col.ColName,
+			TblName:  col.TblName,
+			RetType:  resultTp,
+			UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(),
+		})
 	}
-	// If the types of some child don't match the types of union, we add a projection with cast function.
+	u.schema = expression.NewSchema(unionCols...)
+	// Process each child
+	// 1. Check whether the `TypeField` is correct. If not, create a projection.
+	// 2. If we need to create projection, create it.
+	// 3. If we not need, but the child is not projection. We still need a projection. To make its schema the same with the union.
 	for childID, child := range u.children {
 		exprs := make([]expression.Expression, len(child.Schema().Columns))
 		needProjection := false
 		for i, srcCol := range child.Schema().Columns {
-			dstType := unionSchema.Columns[i].RetType
+			dstType := unionCols[i].RetType
 			srcType := srcCol.RetType
 			if !srcType.Equal(dstType) {
 				exprs[i] = expression.BuildCastFunction4Union(b.ctx, srcCol, dstType)
@@ -660,15 +663,10 @@ func (b *planBuilder) buildProjection4Union(u *LogicalUnionAll) {
 		if _, isProj := child.(*LogicalProjection); needProjection || !isProj {
 			b.optFlag |= flagEliminateProjection
 			proj := LogicalProjection{Exprs: exprs}.init(b.ctx)
-			if childID == 0 {
-				for _, col := range unionSchema.Columns {
-					col.UniqueID = b.ctx.GetSessionVars().AllocPlanColumnID()
-				}
-			}
 			proj.SetChildren(child)
 			u.children[childID] = proj
 		}
-		u.children[childID].(*LogicalProjection).SetSchema(unionSchema.Clone())
+		u.children[childID].(*LogicalProjection).SetSchema(expression.NewSchema(unionCols...))
 	}
 }
 
@@ -678,15 +676,15 @@ func (b *planBuilder) buildUnion(union *ast.UnionStmt) (LogicalPlan, error) {
 		return nil, errors.Trace(err)
 	}
 
-	unionDistinctPlan := b.buildSubUnion(distinctSelectPlans)
+	unionDistinctPlan := b.buildUnionAll(distinctSelectPlans)
 	if unionDistinctPlan != nil {
 		unionDistinctPlan = b.buildDistinct(unionDistinctPlan, unionDistinctPlan.Schema().Len())
 		if len(allSelectPlans) > 0 {
 			allSelectPlans = append(allSelectPlans, unionDistinctPlan)
 		}
 	}
 
-	unionAllPlan := b.buildSubUnion(allSelectPlans)
+	unionAllPlan := b.buildUnionAll(allSelectPlans)
 	unionPlan := unionDistinctPlan
 	if unionAllPlan != nil {
 		unionPlan = unionAllPlan
@@ -738,7 +736,7 @@ func (b *planBuilder) divideUnionSelectPlans(selects []*ast.SelectStmt) (distinc
 	return children[:firstUnionAllIdx], children[firstUnionAllIdx:], nil
 }
 
-func (b *planBuilder) buildSubUnion(subPlan []LogicalPlan) LogicalPlan {
+func (b *planBuilder) buildUnionAll(subPlan []LogicalPlan) LogicalPlan {
 	if len(subPlan) == 0 {
 		return nil
 	}

plan/logical_plan_test.go

@@ -869,7 +869,7 @@ func (s *testPlanSuite) TestEagerAggregation(c *C) {
 		},
 		{
 			sql:  "select max(c.b) from (select * from t a union all select * from t b) c group by c.a",
-			best: "UnionAll{DataScan(a)->Projection->DataScan(b)->Projection}->Projection->Projection",
+			best: "UnionAll{DataScan(a)->Projection->Projection->DataScan(b)->Projection->Projection}->Aggr(max(join_agg_0))->Projection",
 		},
 		{
 			sql:  "select max(a.c) from t a join t b on a.a=b.a and a.b=b.b group by a.b",

plan/logical_plans.go

@@ -560,7 +560,7 @@ func (ds *DataSource) TableInfo() *model.TableInfo {
 
 // LogicalUnionAll represents LogicalUnionAll plan.
 type LogicalUnionAll struct {
-	baseLogicalPlan
+	logicalSchemaProducer
 }
 
 // LogicalSort stands for the order by plan.

plan/physical_plans.go

@@ -267,7 +267,7 @@ type PhysicalLimit struct {
 
 // PhysicalUnionAll is the physical operator of UnionAll.
 type PhysicalUnionAll struct {
-	basePhysicalPlan
+	physicalSchemaProducer
 }
 
 // AggregationType stands for the mode of aggregation plan.

plan/rule_column_pruning.go

@@ -132,6 +132,7 @@ func (p *LogicalUnionAll) PruneColumns(parentUsedCols []*expression.Column) {
 	for _, child := range p.Children() {
 		child.PruneColumns(parentUsedCols)
 	}
+	p.schema.Columns = p.children[0].Schema().Columns
 }
 
 // PruneColumns implements LogicalPlan interface.

plan/rule_partition_processor.go

@@ -122,6 +122,7 @@ func (s *partitionProcessor) prune(ds *DataSource) (LogicalPlan, error) {
 	}
 	unionAll := LogicalUnionAll{}.init(ds.context())
 	unionAll.SetChildren(children...)
+	unionAll.SetSchema(ds.schema)
 	return unionAll, nil
 }