planner: fix mpp final agg couldn't co-exist with other non-final mode (#53455)

pkg/executor/test/tiflashtest/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
     ],
     flaky = True,
     race = "on",
-    shard_count = 42,
+    shard_count = 43,
     deps = [
         "//pkg/config",
         "//pkg/domain",

pkg/executor/test/tiflashtest/tiflash_test.go

@@ -1988,3 +1988,47 @@ func TestIssue50358(t *testing.T) {
 		tk.MustQuery("select 8 from t join t1").Check(testkit.Rows("8", "8"))
 	}
 }
+
+func TestMppAggShouldAlignFinalMode(t *testing.T) {
+	store := testkit.CreateMockStore(t, withMockTiFlash(1))
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t (" +
+		"  d date," +
+		"  v int," +
+		"  primary key(d, v)" +
+		") partition by range columns (d) (" +
+		"  partition p1 values less than ('2023-07-02')," +
+		"  partition p2 values less than ('2023-07-03')" +
+		");")
+	tk.MustExec("alter table t set tiflash replica 1")
+	tb := external.GetTableByName(t, tk, "test", "t")
+	err := domain.GetDomain(tk.Session()).DDL().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)
+	require.NoError(t, err)
+	tk.MustExec(`set tidb_partition_prune_mode='static';`)
+	err = failpoint.Enable("github.com/pingcap/tidb/pkg/expression/aggregation/show-agg-mode", "return(true)")
+	require.Nil(t, err)
+
+	tk.MustExec("set @@session.tidb_isolation_read_engines=\"tiflash\"")
+	tk.MustQuery("explain format='brief' select 1 from (" +
+		"  select /*+ read_from_storage(tiflash[t]) */ sum(1)" +
+		"  from t where d BETWEEN '2023-07-01' and '2023-07-03' group by d" +
+		") total;").Check(testkit.Rows("Projection 400.00 root  1->Column#4",
+		"└─HashAgg 400.00 root  group by:test.t.d, funcs:count(complete,1)->Column#8",
+		"  └─PartitionUnion 400.00 root  ",
+		"    ├─Projection 200.00 root  test.t.d",
+		"    │ └─HashAgg 200.00 root  group by:test.t.d, funcs:firstrow(partial2,test.t.d)->test.t.d, funcs:count(final,Column#12)->Column#9",
+		"    │   └─TableReader 200.00 root  MppVersion: 2, data:ExchangeSender",
+		"    │     └─ExchangeSender 200.00 mpp[tiflash]  ExchangeType: PassThrough",
+		"    │       └─HashAgg 200.00 mpp[tiflash]  group by:test.t.d, funcs:count(partial1,1)->Column#12",
+		"    │         └─TableRangeScan 250.00 mpp[tiflash] table:t, partition:p1 range:[2023-07-01,2023-07-03], keep order:false, stats:pseudo",
+		"    └─Projection 200.00 root  test.t.d",
+		"      └─HashAgg 200.00 root  group by:test.t.d, funcs:firstrow(partial2,test.t.d)->test.t.d, funcs:count(final,Column#14)->Column#10",
+		"        └─TableReader 200.00 root  MppVersion: 2, data:ExchangeSender",
+		"          └─ExchangeSender 200.00 mpp[tiflash]  ExchangeType: PassThrough",
+		"            └─HashAgg 200.00 mpp[tiflash]  group by:test.t.d, funcs:count(partial1,1)->Column#14",
+		"              └─TableRangeScan 250.00 mpp[tiflash] table:t, partition:p2 range:[2023-07-01,2023-07-03], keep order:false, stats:pseudo"))
+
+	err = failpoint.Disable("github.com/pingcap/tidb/pkg/expression/aggregation/show-agg-mode")
+	require.Nil(t, err)
+}

pkg/expression/aggregation/BUILD.bazel

@@ -50,6 +50,7 @@ go_library(
         "//pkg/util/mvmap",
         "//pkg/util/size",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_tipb//go-tipb",
     ],
 )

pkg/expression/aggregation/aggregation.go

@@ -134,6 +134,23 @@ const (
 	DedupMode
 )
 
+// ToString show the agg mode.
+func (a AggFunctionMode) ToString() string {
+	switch a {
+	case CompleteMode:
+		return "complete"
+	case FinalMode:
+		return "final"
+	case Partial1Mode:
+		return "partial1"
+	case Partial2Mode:
+		return "partial2"
+	case DedupMode:
+		return "deduplicate"
+	}
+	return ""
+}
+
 type aggFunction struct {
 	*AggFuncDesc
 }

pkg/expression/aggregation/explain.go

@@ -18,14 +18,26 @@ import (
 	"bytes"
 	"fmt"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 )
 
 // ExplainAggFunc generates explain information for a aggregation function.
 func ExplainAggFunc(ctx expression.EvalContext, agg *AggFuncDesc, normalized bool) string {
 	var buffer bytes.Buffer
-	fmt.Fprintf(&buffer, "%s(", agg.Name)
+	showMode := false
+	failpoint.Inject("show-agg-mode", func(v failpoint.Value) {
+		if v.(bool) {
+			showMode = true
+		}
+	})
+	if showMode {
+		fmt.Fprintf(&buffer, "%s(%s,", agg.Name, agg.Mode.ToString())
+	} else {
+		fmt.Fprintf(&buffer, "%s(", agg.Name)
+	}
+
 	if agg.HasDistinct {
 		buffer.WriteString("distinct ")
 	}

pkg/planner/core/exhaust_physical_plans.go

@@ -3226,6 +3226,22 @@ func (la *LogicalAggregation) tryToGetMppHashAggs(prop *property.PhysicalPropert
 			}
 		}
 	}
+	// ref: https://github.com/pingcap/tiflash/blob/3ebb102fba17dce3d990d824a9df93d93f1ab
+	// 766/dbms/src/Flash/Coprocessor/AggregationInterpreterHelper.cpp#L26
+	validMppAgg := func(mppAgg *PhysicalHashAgg) bool {
+		isFinalAgg := true
+		if mppAgg.AggFuncs[0].Mode != aggregation.FinalMode && mppAgg.AggFuncs[0].Mode != aggregation.CompleteMode {
+			isFinalAgg = false
+		}
+		for _, one := range mppAgg.AggFuncs[1:] {
+			otherIsFinalAgg := one.Mode == aggregation.FinalMode || one.Mode == aggregation.CompleteMode
+			if isFinalAgg != otherIsFinalAgg {
+				// different agg mode detected in mpp side.
+				return false
+			}
+		}
+		return true
+	}
 
 	if len(la.GroupByItems) > 0 {
 		partitionCols := la.GetPotentialPartitionKeys()
@@ -3259,7 +3275,9 @@ func (la *LogicalAggregation) tryToGetMppHashAggs(prop *property.PhysicalPropert
 			agg.SetSchema(la.schema.Clone())
 			agg.MppRunMode = Mpp1Phase
 			finalAggAdjust(agg.AggFuncs)
-			hashAggs = append(hashAggs, agg)
+			if validMppAgg(agg) {
+				hashAggs = append(hashAggs, agg)
+			}
 		}
 
 		// Final agg can't be split into multi-stage aggregate, so exit early
@@ -3274,7 +3292,9 @@ func (la *LogicalAggregation) tryToGetMppHashAggs(prop *property.PhysicalPropert
 		agg.SetSchema(la.schema.Clone())
 		agg.MppRunMode = Mpp2Phase
 		agg.MppPartitionCols = partitionCols
-		hashAggs = append(hashAggs, agg)
+		if validMppAgg(agg) {
+			hashAggs = append(hashAggs, agg)
+		}
 
 		// agg runs on TiDB with a partial agg on TiFlash if possible
 		if prop.TaskTp == property.RootTaskType {