planner: support 3 stage aggregation for single scalar distinct agg (#…

…37203)

close #37202

executor/set_test.go

@@ -699,6 +699,14 @@ func TestSetVar(t *testing.T) {
  tk.MustExec("set global tidb_opt_skew_distinct_agg=1")
  tk.MustQuery("select @@global.tidb_opt_skew_distinct_agg").Check(testkit.Rows("1"))
 
+ // test for tidb_opt_three_stage_distinct_agg
+ tk.MustQuery("select @@session.tidb_opt_three_stage_distinct_agg").Check(testkit.Rows("1")) // default value is 1
+ tk.MustExec("set session tidb_opt_three_stage_distinct_agg=0")
+ tk.MustQuery("select @@session.tidb_opt_three_stage_distinct_agg").Check(testkit.Rows("0"))
+ tk.MustQuery("select @@global.tidb_opt_three_stage_distinct_agg").Check(testkit.Rows("1")) // default value is 1
+ tk.MustExec("set global tidb_opt_three_stage_distinct_agg=0")
+ tk.MustQuery("select @@global.tidb_opt_three_stage_distinct_agg").Check(testkit.Rows("0"))
+
  // the value of max_allowed_packet should be a multiple of 1024
  tk.MustExec("set @@global.max_allowed_packet=16385")
  tk.MustQuery("show warnings").Check(testkit.RowsWithSep("|", "Warning|1292|Truncated incorrect max_allowed_packet value: '16385'"))

expression/column.go

@@ -202,6 +202,18 @@ func (col *CorrelatedColumn) MemoryUsage() (sum int64) {
  return sum
 }
 
+// RemapColumn remaps columns with provided mapping and returns new expression
+func (col *CorrelatedColumn) RemapColumn(m map[int64]*Column) (Expression, error) {
+ mapped := m[(&col.Column).UniqueID]
+ if mapped == nil {
+ return nil, errors.Errorf("Can't remap column for %s", col)
+ }
+ return &CorrelatedColumn{
+ Column: *mapped,
+ Data: col.Data,
+ }, nil
+}
+
 // Column represents a column.
 type Column struct {
  RetType *types.FieldType
@@ -537,6 +549,15 @@ func (col *Column) resolveIndicesByVirtualExpr(schema *Schema) bool {
  return false
 }
 
+// RemapColumn remaps columns with provided mapping and returns new expression
+func (col *Column) RemapColumn(m map[int64]*Column) (Expression, error) {
+ mapped := m[col.UniqueID]
+ if mapped == nil {
+ return nil, errors.Errorf("Can't remap column for %s", col)
+ }
+ return mapped, nil
+}
+
 // Vectorized returns if this expression supports vectorized evaluation.
 func (col *Column) Vectorized() bool {
  return true

expression/constant.go

@@ -422,6 +422,11 @@ func (c *Constant) resolveIndicesByVirtualExpr(_ *Schema) bool {
  return true
 }
 
+// RemapColumn remaps columns with provided mapping and returns new expression
+func (c *Constant) RemapColumn(_ map[int64]*Column) (Expression, error) {
+ return c, nil
+}
+
 // Vectorized returns if this expression supports vectorized evaluation.
 func (c *Constant) Vectorized() bool {
  if c.DeferredExpr != nil {

expression/expression.go

@@ -164,6 +164,9 @@ type Expression interface {
  // resolveIndicesByVirtualExpr is called inside the `ResolveIndicesByVirtualExpr` It will perform on the expression itself.
  resolveIndicesByVirtualExpr(schema *Schema) bool
 
+ // RemapColumn remaps columns with provided mapping and returns new expression
+ RemapColumn(map[int64]*Column) (Expression, error)
+
  // ExplainInfo returns operator information to be explained.
  ExplainInfo() string
 

expression/scalar_function.go

@@ -488,6 +488,24 @@ func (sf *ScalarFunction) resolveIndicesByVirtualExpr(schema *Schema) bool {
  return true
 }
 
+// RemapColumn remaps columns with provided mapping and returns new expression
+func (sf *ScalarFunction) RemapColumn(m map[int64]*Column) (Expression, error) {
+ newSf, ok := sf.Clone().(*ScalarFunction)
+ if !ok {
+ return nil, errors.New("failed to cast to scalar function")
+ }
+ for i, arg := range sf.GetArgs() {
+ newArg, err := arg.RemapColumn(m)
+ if err != nil {
+ return nil, err
+ }
+ newSf.GetArgs()[i] = newArg
+ }
+ // clear hash code
+ newSf.hashcode = nil
+ return newSf, nil
+}
+
 // GetSingleColumn returns (Col, Desc) when the ScalarFunction is equivalent to (Col, Desc)
 // when used as a sort key, otherwise returns (nil, false).
 //

expression/util_test.go

@@ -579,6 +579,7 @@ func (m *MockExpr) ResolveIndices(schema *Schema) (Expression, error)
 func (m *MockExpr) resolveIndices(schema *Schema) error { return nil }
 func (m *MockExpr) ResolveIndicesByVirtualExpr(schema *Schema) (Expression, bool) { return m, true }
 func (m *MockExpr) resolveIndicesByVirtualExpr(schema *Schema) bool { return true }
+func (m *MockExpr) RemapColumn(_ map[int64]*Column) (Expression, error) { return m, nil }
 func (m *MockExpr) ExplainInfo() string { return "" }
 func (m *MockExpr) ExplainNormalizedInfo() string { return "" }
 func (m *MockExpr) HashCode(sc *stmtctx.StatementContext) []byte { return nil }

planner/core/enforce_mpp_test.go

@@ -484,3 +484,54 @@ func TestMPPSkewedGroupDistinctRewrite(t *testing.T) {
  require.Equal(t, output[i].Warn, testdata.ConvertSQLWarnToStrings(tk.Session().GetSessionVars().StmtCtx.GetWarnings()))
  }
 }
+
+// Test 3 stage aggregation for single count distinct
+func TestMPPSingleDistinct3Stage(t *testing.T) {
+ store := testkit.CreateMockStore(t)
+ tk := testkit.NewTestKit(t, store)
+
+ // test table
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t")
+ tk.MustExec("create table t(a int, b bigint not null, c bigint, d date, e varchar(20) collate utf8mb4_general_ci)")
+
+ // Create virtual tiflash replica info.
+ dom := domain.GetDomain(tk.Session())
+ is := dom.InfoSchema()
+ db, exists := is.SchemaByName(model.NewCIStr("test"))
+ require.True(t, exists)
+ for _, tblInfo := range db.Tables {
+ if tblInfo.Name.L == "t" {
+ tblInfo.TiFlashReplica = &model.TiFlashReplicaInfo{
+ Count: 1,
+ Available: true,
+ }
+ }
+ }
+
+ var input []string
+ var output []struct {
+ SQL string
+ Plan []string
+ Warn []string
+ }
+ enforceMPPSuiteData := plannercore.GetEnforceMPPSuiteData()
+ enforceMPPSuiteData.LoadTestCases(t, &input, &output)
+ for i, tt := range input {
+ testdata.OnRecord(func() {
+ output[i].SQL = tt
+ })
+ if strings.HasPrefix(tt, "set") || strings.HasPrefix(tt, "UPDATE") {
+ tk.MustExec(tt)
+ continue
+ }
+ testdata.OnRecord(func() {
+ output[i].SQL = tt
+ output[i].Plan = testdata.ConvertRowsToStrings(tk.MustQuery(tt).Rows())
+ output[i].Warn = testdata.ConvertSQLWarnToStrings(tk.Session().GetSessionVars().StmtCtx.GetWarnings())
+ })
+ res := tk.MustQuery(tt)
+ res.Check(testkit.Rows(output[i].Plan...))
+ require.Equal(t, output[i].Warn, testdata.ConvertSQLWarnToStrings(tk.Session().GetSessionVars().StmtCtx.GetWarnings()))
+ }
+}

planner/core/task.go

@@ -1436,6 +1436,35 @@ func (p *basePhysicalAgg) newPartialAggregate(copTaskType kv.StoreType, isMPPTas
  return partialAgg, finalAgg
 }
 
+// canUse3StageDistinctAgg returns true if this agg can use 3 stage for distinct aggregation
+func (p *basePhysicalAgg) canUse3StageDistinctAgg() bool {
+ num := 0
+ if !p.ctx.GetSessionVars().Enable3StageDistinctAgg || len(p.GroupByItems) > 0 {
+ return false
+ }
+ for _, fun := range p.AggFuncs {
+ if fun.HasDistinct {
+ num++
+ if num > 1 || fun.Name != ast.AggFuncCount {
+ return false
+ }
+ for _, arg := range fun.Args {
+ // bail out when args are not simple column, see GitHub issue #35417
+ if _, ok := arg.(*expression.Column); !ok {
+ return false
+ }
+ }
+ } else if len(fun.Args) > 1 {
+ return false
+ }
+
+ if len(fun.OrderByItems) > 0 {
+ return false
+ }
+ }
+ return num == 1
+}
+
 func genFirstRowAggForGroupBy(ctx sessionctx.Context, groupByItems []expression.Expression) ([]*aggregation.AggFuncDesc, error) {
  aggFuncs := make([]*aggregation.AggFuncDesc, 0, len(groupByItems))
  for _, groupBy := range groupByItems {
@@ -1642,15 +1671,98 @@ func (p *PhysicalHashAgg) attach2TaskForMpp(tasks ...task) task {
  if !mpp.needEnforceExchanger(prop) {
  return p.attach2TaskForMpp1Phase(mpp)
  }
+ // we have to check it before the content of p has been modified
+ canUse3StageAgg := p.canUse3StageDistinctAgg()
  proj := p.convertAvgForMPP()
  partialAgg, finalAgg := p.newPartialAggregate(kv.TiFlash, true)
  if finalAgg == nil {
  return invalidTask
  }
+
+ // generate 3 stage aggregation for single count distinct if applicable.
+ // select count(distinct a), count(b) from foo
+ // will generate plan:
+ // HashAgg sum(#1), sum(#2) -> final agg
+ // +- Exchange Passthrough
+ // +- HashAgg count(distinct a) #1, sum(#3) #2 -> middle agg
+ // +- Exchange HashPartition by a
+ // +- HashAgg count(b) #3, group by a -> partial agg
+ // +- TableScan foo
+ var middleAgg *PhysicalHashAgg = nil
+ if partialAgg != nil && canUse3StageAgg {
+ clonedAgg, err := finalAgg.Clone()
+ if err != nil {
+ return invalidTask
+ }
+ middleAgg = clonedAgg.(*PhysicalHashAgg)
+ distinctPos := 0
+ middleSchema := expression.NewSchema()
+ schemaMap := make(map[int64]*expression.Column, len(middleAgg.AggFuncs))
+ for i, fun := range middleAgg.AggFuncs {
+ col := &expression.Column{
+ UniqueID: p.ctx.GetSessionVars().AllocPlanColumnID(),
+ RetType: fun.RetTp,
+ }
+ if fun.HasDistinct {
+ distinctPos = i
+ } else {
+ fun.Mode = aggregation.Partial2Mode
+ originalCol := fun.Args[0].(*expression.Column)
+ schemaMap[originalCol.UniqueID] = col
+ }
+ middleSchema.Append(col)
+ }
+ middleAgg.schema = middleSchema
+
+ finalHashAgg := finalAgg.(*PhysicalHashAgg)
+ finalAggDescs := make([]*aggregation.AggFuncDesc, 0, len(finalHashAgg.AggFuncs))
+ for i, fun := range finalHashAgg.AggFuncs {
+ newArgs := make([]expression.Expression, 0, 1)
+ if distinctPos == i {
+ // change count(distinct) to sum()
+ fun.Name = ast.AggFuncSum
+ fun.HasDistinct = false
+ newArgs = append(newArgs, middleSchema.Columns[i])
+ } else {
+ for _, arg := range fun.Args {
+ newCol, err := arg.RemapColumn(schemaMap)
+ if err != nil {
+ return invalidTask
+ }
+ newArgs = append(newArgs, newCol)
+ }
+ }
+ fun.Args = newArgs
+ finalAggDescs = append(finalAggDescs, fun)
+ }
+ finalHashAgg.AggFuncs = finalAggDescs
+ }
+
  // partial agg would be null if one scalar agg cannot run in two-phase mode
  if partialAgg != nil {
  attachPlan2Task(partialAgg, mpp)
  }
+
+ if middleAgg != nil && canUse3StageAgg {
+ items := partialAgg.(*PhysicalHashAgg).GroupByItems
+ partitionCols := make([]*property.MPPPartitionColumn, 0, len(items))
+ for _, expr := range items {
+ col, ok := expr.(*expression.Column)
+ if !ok {
+ continue
+ }
+ partitionCols = append(partitionCols, &property.MPPPartitionColumn{
+ Col: col,
+ CollateID: property.GetCollateIDByNameForPartition(col.GetType().GetCollate()),
+ })
+ }
+
+ prop := &property.PhysicalProperty{TaskTp: property.MppTaskType, ExpectedCnt: math.MaxFloat64, MPPPartitionTp: property.HashType, MPPPartitionCols: partitionCols}
+ newMpp := mpp.enforceExchanger(prop)
+ attachPlan2Task(middleAgg, newMpp)
+ mpp = newMpp
+ }
+
  newMpp := mpp.enforceExchanger(prop)
  attachPlan2Task(finalAgg, newMpp)
  if proj == nil {

planner/core/testdata/enforce_mpp_suite_in.json

@@ -113,5 +113,24 @@
  "EXPLAIN select a, count(b), avg(distinct c), count(distinct c) from t group by a; -- multi distinct funcs, bail out",
  "EXPLAIN select count(b), count(distinct c) from t; -- single distinct func but no group key, bail out"
  ]
+ },
+ {
+ "name": "TestMPPSingleDistinct3Stage",
+ "cases": [
+ "set @@tidb_allow_mpp=1;set @@tidb_enforce_mpp=1;",
+ "EXPLAIN select count(distinct b) from t;",
+ "EXPLAIN select count(distinct c) from t;",
+ "EXPLAIN select count(distinct e) from t;",
+ "EXPLAIN select count(distinct a,b,c,e) from t;",
+ "EXPLAIN select count(distinct c), count(a), count(*) from t;",
+ "EXPLAIN select sum(b), count(a), count(*), count(distinct c) from t;",
+ "EXPLAIN select sum(b+a), count(*), count(distinct c), count(a) from t having count(distinct c) > 2;",
+ "EXPLAIN select sum(b+a), count(*), count(a) from t having count(distinct c) > 2;",
+ "EXPLAIN select sum(b+a), max(b), count(distinct c), count(*) from t having count(a) > 2;",
+ "EXPLAIN select sum(b), count(distinct a, b, e), count(a+b) from t;",
+ "EXPLAIN select count(distinct b), json_objectagg(d,c) from t;",
+ "EXPLAIN select count(distinct c+a), count(a) from t;",
+ "EXPLAIN select sum(b), count(distinct c+a, b, e), count(a+b) from t;"
+ ]
  }
 ]