planner: change the stage of count star rewriter rule from logical to…

… post (#39555)

close #37165

cmd/explaintest/r/access_path_selection.result

@@ -41,5 +41,5 @@ explain format = 'brief' select count(1) from access_path_selection;
 id estRows task access object operator info
 StreamAgg 1.00 root funcs:count(Column#18)->Column#4
 └─TableReader 1.00 root data:StreamAgg
- └─StreamAgg 1.00 cop[tikv] funcs:count(test.access_path_selection._tidb_rowid)->Column#18
+ └─StreamAgg 1.00 cop[tikv] funcs:count(1)->Column#18
  └─TableFullScan 10000.00 cop[tikv] table:access_path_selection keep order:false, stats:pseudo

cmd/explaintest/r/clustered_index.result

@@ -30,25 +30,25 @@ explain select count(*) from with_cluster_index.tbl_0 where col_0 < 5429 ;
 id estRows task access object operator info
 StreamAgg_17 1.00 root funcs:count(Column#8)->Column#6
 └─IndexReader_18 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(with_cluster_index.tbl_0.col_0)->Column#8
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#8
  └─IndexRangeScan_16 798.90 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[-inf,5429), keep order:false
 explain select count(*) from wout_cluster_index.tbl_0 where col_0 < 5429 ;
 id estRows task access object operator info
 StreamAgg_17 1.00 root funcs:count(Column#9)->Column#7
 └─IndexReader_18 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(wout_cluster_index.tbl_0.col_0)->Column#9
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#9
  └─IndexRangeScan_16 798.90 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[-inf,5429), keep order:false
 explain select count(*) from with_cluster_index.tbl_0 where col_0 < 41 ;
 id estRows task access object operator info
 StreamAgg_17 1.00 root funcs:count(Column#8)->Column#6
 └─IndexReader_18 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(with_cluster_index.tbl_0.col_0)->Column#8
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#8
  └─IndexRangeScan_16 41.00 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[-inf,41), keep order:false
 explain select count(*) from wout_cluster_index.tbl_0 where col_0 < 41 ;
 id estRows task access object operator info
 StreamAgg_17 1.00 root funcs:count(Column#9)->Column#7
 └─IndexReader_18 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(wout_cluster_index.tbl_0.col_0)->Column#9
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#9
  └─IndexRangeScan_16 41.00 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[-inf,41), keep order:false
 explain select col_14 from with_cluster_index.tbl_2 where col_11 <> '2013-11-01' ;
 id estRows task access object operator info
@@ -109,24 +109,24 @@ explain select count(*) from with_cluster_index.tbl_0 where col_0 <= 0 ;
 id estRows task access object operator info
 StreamAgg_16 1.00 root funcs:count(Column#8)->Column#6
 └─IndexReader_17 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(with_cluster_index.tbl_0.col_0)->Column#8
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#8
  └─IndexRangeScan_11 1.00 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[-inf,0], keep order:false
 explain select count(*) from wout_cluster_index.tbl_0 where col_0 <= 0 ;
 id estRows task access object operator info
 StreamAgg_16 1.00 root funcs:count(Column#9)->Column#7
 └─IndexReader_17 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(wout_cluster_index.tbl_0.col_0)->Column#9
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#9
  └─IndexRangeScan_11 1.00 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[-inf,0], keep order:false
 explain select count(*) from with_cluster_index.tbl_0 where col_0 >= 803163 ;
 id estRows task access object operator info
 StreamAgg_17 1.00 root funcs:count(Column#8)->Column#6
 └─IndexReader_18 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(with_cluster_index.tbl_0.col_0)->Column#8
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#8
  └─IndexRangeScan_16 109.70 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[803163,+inf], keep order:false
 explain select count(*) from wout_cluster_index.tbl_0 where col_0 >= 803163 ;
 id estRows task access object operator info
 StreamAgg_17 1.00 root funcs:count(Column#9)->Column#7
 └─IndexReader_18 1.00 root index:StreamAgg_9
- └─StreamAgg_9 1.00 cop[tikv] funcs:count(wout_cluster_index.tbl_0.col_0)->Column#9
+ └─StreamAgg_9 1.00 cop[tikv] funcs:count(1)->Column#9
  └─IndexRangeScan_16 109.70 cop[tikv] table:tbl_0, index:idx_3(col_0) range:[803163,+inf], keep order:false
 set @@tidb_enable_outer_join_reorder=false;

cmd/explaintest/r/explain_easy.result

@@ -218,14 +218,14 @@ id estRows task access object operator info
 Projection 1.00 root 1->Column#6
 └─StreamAgg 1.00 root funcs:count(Column#14)->Column#9
  └─TableReader 1.00 root data:StreamAgg
- └─StreamAgg 1.00 cop[tikv] funcs:count(test.t1.c1)->Column#14
+ └─StreamAgg 1.00 cop[tikv] funcs:count(1)->Column#14
  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo
 explain format = 'brief' select count(1) from (select max(c2), count(c3) as m from t1) k;
 id estRows task access object operator info
 StreamAgg 1.00 root funcs:count(1)->Column#6
 └─StreamAgg 1.00 root funcs:count(Column#13)->Column#8
  └─TableReader 1.00 root data:StreamAgg
- └─StreamAgg 1.00 cop[tikv] funcs:count(test.t1.c1)->Column#13
+ └─StreamAgg 1.00 cop[tikv] funcs:count(1)->Column#13
  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo
 explain format = 'brief' select count(1) from (select count(c2) from t1 group by c3) k;
 id estRows task access object operator info

cmd/explaintest/r/index_merge.result

@@ -126,7 +126,7 @@ Sort_11 5098.44 root test.t1.c1
  ├─StreamAgg_23(Build) 1.00 root funcs:min(Column#9)->Column#10, funcs:sum(0)->Column#11, funcs:count(1)->Column#12
  │ └─StreamAgg_43 1.00 root funcs:count(Column#25)->Column#9
  │ └─IndexReader_44 1.00 root index:StreamAgg_27
- │ └─StreamAgg_27 1.00 cop[tikv] funcs:count(test.t2._tidb_rowid)->Column#25
+ │ └─StreamAgg_27 1.00 cop[tikv] funcs:count(1)->Column#25
  │ └─IndexFullScan_41 10000.00 cop[tikv] table:t2, index:c1(c1) keep order:false, stats:pseudo
  └─IndexMerge_21(Probe) 2825.66 root type: union
  ├─IndexRangeScan_17(Build) 3323.33 cop[tikv] table:t1, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo
@@ -148,7 +148,7 @@ Sort_11 5098.44 root test.t1.c1
  ├─StreamAgg_23(Build) 1.00 root funcs:min(Column#9)->Column#10, funcs:sum(0)->Column#11, funcs:count(1)->Column#12
  │ └─StreamAgg_43 1.00 root funcs:count(Column#25)->Column#9
  │ └─IndexReader_44 1.00 root index:StreamAgg_27
- │ └─StreamAgg_27 1.00 cop[tikv] funcs:count(test.t2._tidb_rowid)->Column#25
+ │ └─StreamAgg_27 1.00 cop[tikv] funcs:count(1)->Column#25
  │ └─IndexFullScan_41 10000.00 cop[tikv] table:t2, index:c1(c1) keep order:false, stats:pseudo
  └─IndexMerge_21(Probe) 2825.66 root type: union
  ├─IndexRangeScan_17(Build) 3323.33 cop[tikv] table:t1, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo
@@ -170,7 +170,7 @@ Sort_11 5542.21 root test.t1.c1
  ├─StreamAgg_22(Build) 1.00 root funcs:max(Column#9)->Column#10, funcs:sum(0)->Column#11, funcs:count(1)->Column#12
  │ └─StreamAgg_42 1.00 root funcs:count(Column#25)->Column#9
  │ └─IndexReader_43 1.00 root index:StreamAgg_26
- │ └─StreamAgg_26 1.00 cop[tikv] funcs:count(test.t2._tidb_rowid)->Column#25
+ │ └─StreamAgg_26 1.00 cop[tikv] funcs:count(1)->Column#25
  │ └─IndexFullScan_40 10000.00 cop[tikv] table:t2, index:c1(c1) keep order:false, stats:pseudo
  └─IndexMerge_20(Probe) 5542.21 root type: union
  ├─IndexRangeScan_17(Build) 3323.33 cop[tikv] table:t1, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo
@@ -192,7 +192,7 @@ Sort_39 5542.21 root test.t1.c1
  ├─StreamAgg_50(Build) 1.00 root funcs:max(Column#13)->Column#14, funcs:sum(0)->Column#15, funcs:count(1)->Column#16
  │ └─StreamAgg_70 1.00 root funcs:count(Column#38)->Column#13
  │ └─IndexReader_71 1.00 root index:StreamAgg_54
- │ └─StreamAgg_54 1.00 cop[tikv] funcs:count(test.t2._tidb_rowid)->Column#38
+ │ └─StreamAgg_54 1.00 cop[tikv] funcs:count(1)->Column#38
  │ └─IndexFullScan_68 10000.00 cop[tikv] table:t2, index:c1(c1) keep order:false, stats:pseudo
  └─IndexMerge_48(Probe) 5542.21 root type: union
  ├─IndexRangeScan_45(Build) 3323.33 cop[tikv] table:t1, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo

planner/core/BUILD.bazel

@@ -48,7 +48,6 @@ go_library(
  "rule_aggregation_skew_rewrite.go",
  "rule_build_key_info.go",
  "rule_column_pruning.go",
- "rule_count_star_rewriter.go",
  "rule_decorrelate.go",
  "rule_eliminate_projection.go",
  "rule_generate_column_substitute.go",

planner/core/logical_plan_builder.go

@@ -204,7 +204,6 @@ func (b *PlanBuilder) buildAggregation(ctx context.Context, p LogicalPlan, aggFu
  b.optFlag |= flagPredicatePushDown
  b.optFlag |= flagEliminateAgg
  b.optFlag |= flagEliminateProjection
- b.optFlag |= flagCountStarRewriter
 
  if b.ctx.GetSessionVars().EnableSkewDistinctAgg {
  b.optFlag |= flagSkewDistinctAgg

planner/core/optimizer.go

@@ -16,11 +16,13 @@ package core
 
 import (
  "context"
+ "fmt"
  "math"
 
  "github.com/pingcap/errors"
  "github.com/pingcap/tidb/config"
  "github.com/pingcap/tidb/expression"
+ "github.com/pingcap/tidb/expression/aggregation"
  "github.com/pingcap/tidb/infoschema"
  "github.com/pingcap/tidb/kv"
  "github.com/pingcap/tidb/lock"
@@ -72,7 +74,6 @@ const (
  flagSyncWaitStatsLoadPoint
  flagJoinReOrder
  flagPrunColumnsAgain
- flagCountStarRewriter
 )
 
 var optRuleList = []logicalOptRule{
@@ -95,7 +96,6 @@ var optRuleList = []logicalOptRule{
  &syncWaitStatsLoadPoint{},
  &joinReOrderSolver{},
  &columnPruner{}, // column pruning again at last, note it will mess up the results of buildKeySolver
- &countStarRewriter{},
 }
 
 type logicalOptimizeOp struct {
@@ -392,6 +392,7 @@ func postOptimize(sctx sessionctx.Context, plan PhysicalPlan) (PhysicalPlan, err
  handleFineGrainedShuffle(sctx, plan)
  checkPlanCacheable(sctx, plan)
  propagateProbeParents(plan, nil)
+ countStarRewrite(plan)
  return plan, nil
 }
 
@@ -533,6 +534,120 @@ func prunePhysicalColumnsInternal(sctx sessionctx.Context, plan PhysicalPlan) er
  return nil
 }
 
+/*
+*
+The countStarRewriter is used to rewrite
+
+ count(*) -> count(not null column)
+
+**Only for TiFlash**
+Attention:
+Since count(*) is directly translated into count(1) during grammar parsing,
+the rewritten pattern actually matches count(constant)
+
+Pattern:
+PhysicalAggregation: count(constant)
+
+ |
+ TableFullScan: TiFlash
+
+Optimize:
+Table
+
+ <k1 bool not null, k2 int null, k3 bigint not null>
+
+Query: select count(*) from table
+ColumnPruningRule: datasource pick row_id
+countStarRewrite: datasource pick k1 instead of row_id
+
+ rewrite count(*) -> count(k1)
+
+Rewritten Query: select count(k1) from table
+*/
+func countStarRewrite(plan PhysicalPlan) {
+ countStarRewriteInternal(plan)
+ if tableReader, ok := plan.(*PhysicalTableReader); ok {
+ countStarRewrite(tableReader.tablePlan)
+ } else {
+ for _, child := range plan.Children() {
+ countStarRewrite(child)
+ }
+ }
+}
+
+func countStarRewriteInternal(plan PhysicalPlan) {
+ // match pattern any agg(count(constant)) -> tablefullscan(tiflash)
+ var physicalAgg *basePhysicalAgg
+ switch x := plan.(type) {
+ case *PhysicalHashAgg:
+ physicalAgg = x.getPointer()
+ case *PhysicalStreamAgg:
+ physicalAgg = x.getPointer()
+ default:
+ return
+ }
+ if len(physicalAgg.GroupByItems) > 0 || len(physicalAgg.children) != 1 {
+ return
+ }
+ for _, aggFunc := range physicalAgg.AggFuncs {
+ if aggFunc.Name != "count" || len(aggFunc.Args) != 1 || aggFunc.HasDistinct {
+ return
+ }
+ if _, ok := aggFunc.Args[0].(*expression.Constant); !ok {
+ return
+ }
+ }
+ physicalTableScan, ok := physicalAgg.Children()[0].(*PhysicalTableScan)
+ if !ok || !physicalTableScan.isFullScan() || physicalTableScan.StoreType != kv.TiFlash || len(physicalTableScan.schema.Columns) != 1 {
+ return
+ }
+ // rewrite datasource and agg args
+ rewriteTableScanAndAggArgs(physicalTableScan, physicalAgg.AggFuncs)
+}
+
+// rewriteTableScanAndAggArgs Pick the narrowest and not null column from table
+// If there is no not null column in Data Source, the row_id or pk column will be retained
+func rewriteTableScanAndAggArgs(physicalTableScan *PhysicalTableScan, aggFuncs []*aggregation.AggFuncDesc) {
+ var resultColumnInfo *model.ColumnInfo
+ var resultColumn *expression.Column
+
+ resultColumnInfo = physicalTableScan.Columns[0]
+ resultColumn = physicalTableScan.schema.Columns[0]
+ // prefer not null column from table
+ for _, columnInfo := range physicalTableScan.Table.Columns {
+ if columnInfo.FieldType.IsVarLengthType() {
+ continue
+ }
+ if mysql.HasNotNullFlag(columnInfo.GetFlag()) {
+ if columnInfo.GetFlen() < resultColumnInfo.GetFlen() {
+ resultColumnInfo = columnInfo
+ resultColumn = &expression.Column{
+ UniqueID: physicalTableScan.ctx.GetSessionVars().AllocPlanColumnID(),
+ ID: resultColumnInfo.ID,
+ RetType: resultColumnInfo.FieldType.Clone(),
+ OrigName: fmt.Sprintf("%s.%s.%s", physicalTableScan.DBName.L, physicalTableScan.Table.Name.L, resultColumnInfo.Name),
+ }
+ }
+ }
+ }
+ // table scan (row_id) -> (not null column)
+ physicalTableScan.Columns[0] = resultColumnInfo
+ physicalTableScan.schema.Columns[0] = resultColumn
+ // agg arg count(1) -> count(not null column)
+ arg := resultColumn.Clone()
+ for _, aggFunc := range aggFuncs {
+ constExpr, ok := aggFunc.Args[0].(*expression.Constant)
+ if !ok {
+ return
+ }
+ // count(null) shouldn't be rewritten
+ if constExpr.Value.IsNull() {
+ continue
+ }
+ aggFunc.Args[0] = arg
+ }
+}
+
 // Only for MPP(Window<-[Sort]<-ExchangeReceiver<-ExchangeSender).
 // TiFlashFineGrainedShuffleStreamCount:
 // < 0: fine grained shuffle is disabled.

planner/core/partition_pruner_test.go

@@ -91,12 +91,12 @@ func TestRangeColumnPartitionPruningForIn(t *testing.T) {
  "└─PartitionUnion 2.00 root ",
  " ├─HashAgg 1.00 root funcs:count(Column#7)->Column#5",
  " │ └─IndexReader 1.00 root index:HashAgg",
- " │ └─HashAgg 1.00 cop[tikv] funcs:count(test_range_col_in.t1.dt)->Column#7",
+ " │ └─HashAgg 1.00 cop[tikv] funcs:count(1)->Column#7",
  " │ └─Selection 20.00 cop[tikv] in(test_range_col_in.t1.dt, 2020-11-27 00:00:00.000000, 2020-11-28 00:00:00.000000)",
  " │ └─IndexFullScan 10000.00 cop[tikv] table:t1, partition:p20201127, index:PRIMARY(id, dt) keep order:false, stats:pseudo",
  " └─HashAgg 1.00 root funcs:count(Column#10)->Column#5",
  " └─IndexReader 1.00 root index:HashAgg",
- " └─HashAgg 1.00 cop[tikv] funcs:count(test_range_col_in.t1.dt)->Column#10",
+ " └─HashAgg 1.00 cop[tikv] funcs:count(1)->Column#10",
  " └─Selection 20.00 cop[tikv] in(test_range_col_in.t1.dt, 2020-11-27 00:00:00.000000, 2020-11-28 00:00:00.000000)",
  " └─IndexFullScan 10000.00 cop[tikv] table:t1, partition:p20201128, index:PRIMARY(id, dt) keep order:false, stats:pseudo"))
 

planner/core/physical_plans.go

@@ -1753,6 +1753,10 @@ type PhysicalHashAgg struct {
  basePhysicalAgg
 }
 
+func (p *PhysicalHashAgg) getPointer() *basePhysicalAgg {
+ return &p.basePhysicalAgg
+}
+
 // Clone implements PhysicalPlan interface.
 func (p *PhysicalHashAgg) Clone() (PhysicalPlan, error) {
  cloned := new(PhysicalHashAgg)
@@ -1787,6 +1791,10 @@ type PhysicalStreamAgg struct {
  basePhysicalAgg
 }
 
+func (p *PhysicalStreamAgg) getPointer() *basePhysicalAgg {
+ return &p.basePhysicalAgg
+}
+
 // Clone implements PhysicalPlan interface.
 func (p *PhysicalStreamAgg) Clone() (PhysicalPlan, error) {
  cloned := new(PhysicalStreamAgg)

planner/core/rule_column_pruning.go

@@ -336,24 +336,10 @@ func (ds *DataSource) PruneColumns(parentUsedCols []*expression.Column, opt *log
  appendColumnPruneTraceStep(ds, prunedColumns, opt)
  // For SQL like `select 1 from t`, tikv's response will be empty if no column is in schema.
  // So we'll force to push one if schema doesn't have any column.
- // There are two situations
- // case 1: tiflash. Select a non-empty and narrowest column.
- // The main reason is that tiflash is a column storage structure,
- // and choosing the narrowest column can reduce the amount of data read as much as possible,
- // making the reading efficiency the best.
- // case 2: tikv. Select row_id or pk column.
- // The main reason is that tikv is a kv store.
- // Select the key column for the best read efficiency.
  if ds.schema.Len() == 0 {
  var handleCol *expression.Column
  var handleColInfo *model.ColumnInfo
- // case 1: tiflash
- if ds.tableInfo.TiFlashReplica != nil {
- handleCol, handleColInfo = preferNotNullColumnFromTable(ds)
- } else {
- // case 2: tikv
- handleCol, handleColInfo = preferKeyColumnFromTable(ds, originSchemaColumns, originColumns)
- }
+ handleCol, handleColInfo = preferKeyColumnFromTable(ds, originSchemaColumns, originColumns)
  ds.Columns = append(ds.Columns, handleColInfo)
  ds.schema.Append(handleCol)
  }
@@ -661,38 +647,6 @@ func appendItemPruneTraceStep(p LogicalPlan, itemType string, prunedObjects []fm
  opt.appendStepToCurrent(p.ID(), p.TP(), reason, action)
 }
 
-// pick a not null and narrowest column from table
-func preferNotNullColumnFromTable(dataSource *DataSource) (*expression.Column, *model.ColumnInfo) {
- var resultColumnInfo *model.ColumnInfo
- var resultColumn *expression.Column
-
- if dataSource.handleCols != nil {
- resultColumn = dataSource.handleCols.GetCol(0)
- resultColumnInfo = resultColumn.ToInfo()
- } else {
- resultColumn = dataSource.newExtraHandleSchemaCol()
- resultColumnInfo = model.NewExtraHandleColInfo()
- }
-
- for _, columnInfo := range dataSource.tableInfo.Columns {
- if columnInfo.FieldType.IsVarLengthType() {
- continue
- }
- if mysql.HasNotNullFlag(columnInfo.GetFlag()) {
- if columnInfo.GetFlen() < resultColumnInfo.GetFlen() {
- resultColumnInfo = columnInfo
- resultColumn = &expression.Column{
- UniqueID: dataSource.ctx.GetSessionVars().AllocPlanColumnID(),
- ID: resultColumnInfo.ID,
- RetType: resultColumnInfo.FieldType.Clone(),
- OrigName: fmt.Sprintf("%s.%s.%s", dataSource.DBName.L, dataSource.tableInfo.Name.L, resultColumnInfo.Name),
- }
- }
- }
- }
- return resultColumn, resultColumnInfo
-}
-
 func preferKeyColumnFromTable(dataSource *DataSource, originColumns []*expression.Column,
  originSchemaColumns []*model.ColumnInfo) (*expression.Column, *model.ColumnInfo) {
  var resultColumnInfo *model.ColumnInfo