plan: fix a problem caused by union's schema

cmd/explaintest/r/explain_complex.result

@@ -179,23 +179,23 @@ CREATE TABLE `tbl_008` (`a` int, `b` int);
 CREATE TABLE `tbl_009` (`a` int, `b` int);
 explain select sum(a) from (select * from tbl_001 union all select * from tbl_002 union all select * from tbl_003 union all select * from tbl_004 union all select * from tbl_005 union all select * from tbl_006 union all select * from tbl_007 union all select * from tbl_008 union all select * from tbl_009) x group by b;
 id count task operator info
-HashAgg_25 72000.00 root group by:x.b, funcs:sum(x.a)
-└─Union_26 90000.00 root 
- ├─TableReader_29 10000.00 root data:TableScan_28
- │ └─TableScan_28 10000.00 cop table:tbl_001, range:[-inf,+inf], keep order:false, stats:pseudo
- ├─TableReader_32 10000.00 root data:TableScan_31
- │ └─TableScan_31 10000.00 cop table:tbl_002, range:[-inf,+inf], keep order:false, stats:pseudo
- ├─TableReader_35 10000.00 root data:TableScan_34
- │ └─TableScan_34 10000.00 cop table:tbl_003, range:[-inf,+inf], keep order:false, stats:pseudo
+HashAgg_34 72000.00 root group by:x.b, funcs:sum(x.a)
+└─Union_35 90000.00 root 
  ├─TableReader_38 10000.00 root data:TableScan_37
- │ └─TableScan_37 10000.00 cop table:tbl_004, range:[-inf,+inf], keep order:false, stats:pseudo
+ │ └─TableScan_37 10000.00 cop table:tbl_001, range:[-inf,+inf], keep order:false, stats:pseudo
  ├─TableReader_41 10000.00 root data:TableScan_40
- │ └─TableScan_40 10000.00 cop table:tbl_005, range:[-inf,+inf], keep order:false, stats:pseudo
+ │ └─TableScan_40 10000.00 cop table:tbl_002, range:[-inf,+inf], keep order:false, stats:pseudo
  ├─TableReader_44 10000.00 root data:TableScan_43
- │ └─TableScan_43 10000.00 cop table:tbl_006, range:[-inf,+inf], keep order:false, stats:pseudo
+ │ └─TableScan_43 10000.00 cop table:tbl_003, range:[-inf,+inf], keep order:false, stats:pseudo
  ├─TableReader_47 10000.00 root data:TableScan_46
- │ └─TableScan_46 10000.00 cop table:tbl_007, range:[-inf,+inf], keep order:false, stats:pseudo
+ │ └─TableScan_46 10000.00 cop table:tbl_004, range:[-inf,+inf], keep order:false, stats:pseudo
  ├─TableReader_50 10000.00 root data:TableScan_49
- │ └─TableScan_49 10000.00 cop table:tbl_008, range:[-inf,+inf], keep order:false, stats:pseudo
- └─TableReader_53 10000.00 root data:TableScan_52
- └─TableScan_52 10000.00 cop table:tbl_009, range:[-inf,+inf], keep order:false, stats:pseudo
+ │ └─TableScan_49 10000.00 cop table:tbl_005, range:[-inf,+inf], keep order:false, stats:pseudo
+ ├─TableReader_53 10000.00 root data:TableScan_52
+ │ └─TableScan_52 10000.00 cop table:tbl_006, range:[-inf,+inf], keep order:false, stats:pseudo
+ ├─TableReader_56 10000.00 root data:TableScan_55
+ │ └─TableScan_55 10000.00 cop table:tbl_007, range:[-inf,+inf], keep order:false, stats:pseudo
+ ├─TableReader_59 10000.00 root data:TableScan_58
+ │ └─TableScan_58 10000.00 cop table:tbl_008, range:[-inf,+inf], keep order:false, stats:pseudo
+ └─TableReader_62 10000.00 root data:TableScan_61
+ └─TableScan_61 10000.00 cop table:tbl_009, range:[-inf,+inf], keep order:false, stats:pseudo

cmd/explaintest/r/explain_complex_stats.result

@@ -205,23 +205,23 @@ CREATE TABLE tbl_009 (a int, b int);
 load stats 's/explain_complex_stats_tbl_009.json';
 explain select sum(a) from (select * from tbl_001 union all select * from tbl_002 union all select * from tbl_003 union all select * from tbl_004 union all select * from tbl_005 union all select * from tbl_006 union all select * from tbl_007 union all select * from tbl_008 union all select * from tbl_009) x group by b;
 id count task operator info
-HashAgg_25 18000.00 root group by:x.b, funcs:sum(x.a)
-└─Union_26 18000.00 root 
- ├─TableReader_29 2000.00 root data:TableScan_28
- │ └─TableScan_28 2000.00 cop table:tbl_001, range:[-inf,+inf], keep order:false
- ├─TableReader_32 2000.00 root data:TableScan_31
- │ └─TableScan_31 2000.00 cop table:tbl_002, range:[-inf,+inf], keep order:false
- ├─TableReader_35 2000.00 root data:TableScan_34
- │ └─TableScan_34 2000.00 cop table:tbl_003, range:[-inf,+inf], keep order:false
+HashAgg_34 18000.00 root group by:x.b, funcs:sum(x.a)
+└─Union_35 18000.00 root 
  ├─TableReader_38 2000.00 root data:TableScan_37
- │ └─TableScan_37 2000.00 cop table:tbl_004, range:[-inf,+inf], keep order:false
+ │ └─TableScan_37 2000.00 cop table:tbl_001, range:[-inf,+inf], keep order:false
  ├─TableReader_41 2000.00 root data:TableScan_40
- │ └─TableScan_40 2000.00 cop table:tbl_005, range:[-inf,+inf], keep order:false
+ │ └─TableScan_40 2000.00 cop table:tbl_002, range:[-inf,+inf], keep order:false
  ├─TableReader_44 2000.00 root data:TableScan_43
- │ └─TableScan_43 2000.00 cop table:tbl_006, range:[-inf,+inf], keep order:false
+ │ └─TableScan_43 2000.00 cop table:tbl_003, range:[-inf,+inf], keep order:false
  ├─TableReader_47 2000.00 root data:TableScan_46
- │ └─TableScan_46 2000.00 cop table:tbl_007, range:[-inf,+inf], keep order:false
+ │ └─TableScan_46 2000.00 cop table:tbl_004, range:[-inf,+inf], keep order:false
  ├─TableReader_50 2000.00 root data:TableScan_49
- │ └─TableScan_49 2000.00 cop table:tbl_008, range:[-inf,+inf], keep order:false
- └─TableReader_53 2000.00 root data:TableScan_52
- └─TableScan_52 2000.00 cop table:tbl_009, range:[-inf,+inf], keep order:false
+ │ └─TableScan_49 2000.00 cop table:tbl_005, range:[-inf,+inf], keep order:false
+ ├─TableReader_53 2000.00 root data:TableScan_52
+ │ └─TableScan_52 2000.00 cop table:tbl_006, range:[-inf,+inf], keep order:false
+ ├─TableReader_56 2000.00 root data:TableScan_55
+ │ └─TableScan_55 2000.00 cop table:tbl_007, range:[-inf,+inf], keep order:false
+ ├─TableReader_59 2000.00 root data:TableScan_58
+ │ └─TableScan_58 2000.00 cop table:tbl_008, range:[-inf,+inf], keep order:false
+ └─TableReader_62 2000.00 root data:TableScan_61
+ └─TableScan_61 2000.00 cop table:tbl_009, range:[-inf,+inf], keep order:false

cmd/explaintest/r/explain_easy.result

@@ -151,35 +151,35 @@ TableReader_5 10000.00 root data:TableScan_4
 └─TableScan_4 10000.00 cop table:t1, range:[-inf,+inf], keep order:false, stats:pseudo
 explain select c1 from t2 union select c1 from t2 union all select c1 from t2;
 id count task operator info
-Union_14 26000.00 root 
-├─TableReader_17 10000.00 root data:TableScan_16
-│ └─TableScan_16 10000.00 cop table:t2, range:[-inf,+inf], keep order:false, stats:pseudo
-└─HashAgg_21 16000.00 root group by:t2.c1, funcs:firstrow(join_agg_0)
- └─Union_22 16000.00 root 
- ├─StreamAgg_35 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
- │ └─IndexReader_36 8000.00 root index:StreamAgg_26
- │ └─StreamAgg_26 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
- │ └─IndexScan_34 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
- └─StreamAgg_52 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
- └─IndexReader_53 8000.00 root index:StreamAgg_43
- └─StreamAgg_43 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
- └─IndexScan_51 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
+Union_17 26000.00 root 
+├─TableReader_20 10000.00 root data:TableScan_19
+│ └─TableScan_19 10000.00 cop table:t2, range:[-inf,+inf], keep order:false, stats:pseudo
+└─HashAgg_24 16000.00 root group by:t2.c1, funcs:firstrow(join_agg_0)
+ └─Union_25 16000.00 root 
+ ├─StreamAgg_38 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
+ │ └─IndexReader_39 8000.00 root index:StreamAgg_29
+ │ └─StreamAgg_29 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
+ │ └─IndexScan_37 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
+ └─StreamAgg_55 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
+ └─IndexReader_56 8000.00 root index:StreamAgg_46
+ └─StreamAgg_46 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
+ └─IndexScan_54 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
 explain select c1 from t2 union all select c1 from t2 union select c1 from t2;
 id count task operator info
-HashAgg_15 24000.00 root group by:t2.c1, funcs:firstrow(join_agg_0)
-└─Union_16 24000.00 root 
- ├─StreamAgg_29 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
- │ └─IndexReader_30 8000.00 root index:StreamAgg_20
- │ └─StreamAgg_20 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
- │ └─IndexScan_28 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
- ├─StreamAgg_46 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
- │ └─IndexReader_47 8000.00 root index:StreamAgg_37
- │ └─StreamAgg_37 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
- │ └─IndexScan_45 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
- └─StreamAgg_63 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
- └─IndexReader_64 8000.00 root index:StreamAgg_54
- └─StreamAgg_54 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
- └─IndexScan_62 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
+HashAgg_18 24000.00 root group by:t2.c1, funcs:firstrow(join_agg_0)
+└─Union_19 24000.00 root 
+ ├─StreamAgg_32 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
+ │ └─IndexReader_33 8000.00 root index:StreamAgg_23
+ │ └─StreamAgg_23 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
+ │ └─IndexScan_31 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
+ ├─StreamAgg_49 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
+ │ └─IndexReader_50 8000.00 root index:StreamAgg_40
+ │ └─StreamAgg_40 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
+ │ └─IndexScan_48 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
+ └─StreamAgg_66 8000.00 root group by:col_2, funcs:firstrow(col_0), firstrow(col_1)
+ └─IndexReader_67 8000.00 root index:StreamAgg_57
+ └─StreamAgg_57 8000.00 cop group by:test.t2.c1, funcs:firstrow(test.t2.c1), firstrow(test.t2.c1)
+ └─IndexScan_65 10000.00 cop table:t2, index:c1, range:[NULL,+inf], keep order:true, stats:pseudo
 set @@session.tidb_opt_insubquery_unfold = 0;
 explain select sum(t1.c1 in (select c1 from t2)) from t1;
 id count task operator info

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

@@ -646,49 +646,41 @@ 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 and add a projection above original child.
+ // So the schema of `UnionAll` can be the same with its children's.
  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)
- needProjection = true
  } else {
  exprs[i] = srcCol
  }
  }
- 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())
+ b.optFlag |= flagEliminateProjection
+ proj := LogicalProjection{Exprs: exprs}.init(b.ctx)
+ proj.SetSchema(expression.NewSchema(unionCols...))
+ proj.SetChildren(child)
+ u.children[childID] = proj
  }
 }
 
@@ -698,15 +690,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
@@ -758,7 +750,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

@@ -372,7 +372,7 @@ func (s *testPlanSuite) TestPredicatePushDown(c *C) {
  },
  {
  sql: "select a, d from (select * from t union all select * from t union all select * from t) z where a < 10",
- best: "UnionAll{DataScan(t)->Projection->DataScan(t)->Projection->DataScan(t)->Projection}->Projection",
+ best: "UnionAll{DataScan(t)->Projection->Projection->DataScan(t)->Projection->Projection->DataScan(t)->Projection->Projection}->Projection",
  },
  {
  sql: "select (select count(*) from t where t.a = k.a) from t k",
@@ -1011,7 +1011,7 @@ func (s *testPlanSuite) TestEagerAggregation(c *C) {
  },
  {
  sql: "select sum(c1) from (select c c1, d c2 from t a union all select a c1, b c2 from t b union all select b c1, e c2 from t c) x group by c2",
- best: "UnionAll{DataScan(a)->Aggr(sum(a.c),firstrow(a.d))->DataScan(b)->Aggr(sum(b.a),firstrow(b.b))->DataScan(c)->Aggr(sum(c.b),firstrow(c.e))}->Aggr(sum(join_agg_0))->Projection",
+ best: "UnionAll{DataScan(a)->Projection->Aggr(sum(a.c1),firstrow(a.c2))->DataScan(b)->Projection->Aggr(sum(b.c1),firstrow(b.c2))->DataScan(c)->Projection->Aggr(sum(c.c1),firstrow(c.c2))}->Aggr(sum(join_agg_0))->Projection",
  },
  {
  sql: "select max(a.b), max(b.b) from t a join t b on a.c = b.c group by a.a",
@@ -1023,7 +1023,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->Projection->DataScan(b)->Projection->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
 }