planner: fix HashAgg cannot pushdown to tiflash_compute

planner/core/find_best_task.go

@@ -2012,13 +2012,14 @@ func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candid
 	}
 	// In disaggregated tiflash mode, only MPP is allowed, cop and batchCop is deprecated.
 	// So if prop.TaskTp is RootTaskType, have to use mppTask then convert to rootTask.
-	isDisaggregatedTiFlashPath := config.GetGlobalConfig().DisaggregatedTiFlash && ts.StoreType == kv.TiFlash
+	isDisaggregatedTiFlash := config.GetGlobalConfig().DisaggregatedTiFlash
+	isDisaggregatedTiFlashPath := isDisaggregatedTiFlash && ts.StoreType == kv.TiFlash
 	canMppConvertToRootForDisaggregatedTiFlash := isDisaggregatedTiFlashPath && prop.TaskTp == property.RootTaskType && ds.SCtx().GetSessionVars().IsMPPAllowed()
 	if prop.TaskTp == property.MppTaskType || canMppConvertToRootForDisaggregatedTiFlash {
 		if ts.KeepOrder {
 			return invalidTask, nil
 		}
-		if prop.MPPPartitionTp != property.AnyType || (ts.isPartition && !canMppConvertToRootForDisaggregatedTiFlash) {
+		if prop.MPPPartitionTp != property.AnyType || (ts.isPartition && !isDisaggregatedTiFlash) {
 			// If ts is a single partition, then this partition table is in static-only prune, then we should not choose mpp execution.
 			// But in disaggregated tiflash mode, we can only use mpp, so we add ExchangeSender and ExchangeReceiver above TableScan for static pruning partition table.
 			ds.SCtx().GetSessionVars().RaiseWarningWhenMPPEnforced("MPP mode may be blocked because table `" + ds.tableInfo.Name.O + "`is a partition table which is not supported when `@@tidb_partition_prune_mode=static`.")

planner/core/logical_plan_builder.go

@@ -50,7 +50,6 @@ import (
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/statistics/handle"
-	"github.com/pingcap/tidb/store/driver/backoff"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
 	"github.com/pingcap/tidb/table/temptable"
@@ -67,7 +66,6 @@ import (
 	"github.com/pingcap/tidb/util/plancodec"
 	"github.com/pingcap/tidb/util/set"
 	"github.com/pingcap/tidb/util/size"
-	"github.com/tikv/client-go/v2/tikv"
 )
 
 const (
@@ -717,11 +715,6 @@ func (ds *DataSource) setPreferredStoreType(hintInfo *tableHintInfo) {
 }
 
 func isTiFlashComputeNodeAvailable(ctx sessionctx.Context) bool {
-	bo := backoff.NewBackofferWithVars(context.Background(), 5000, nil)
-	stores, err := ctx.GetStore().(tikv.Storage).GetRegionCache().GetTiFlashComputeStores(bo.TiKVBackoffer())
-	if err != nil || len(stores) == 0 {
-		return false
-	}
 	return true
 }
 

planner/core/task.go

@@ -1820,6 +1820,9 @@ func (p *PhysicalHashAgg) attach2TaskForMpp(tasks ...task) task {
 	}
 	switch p.MppRunMode {
 	case Mpp1Phase:
+		if !aggFuncModeSame(p) {
+			return invalidTask
+		}
 		// 1-phase agg: when the partition columns can be satisfied, where the plan does not need to enforce Exchange
 		// only push down the original agg
 		proj := p.convertAvgForMPP()
@@ -1835,6 +1838,9 @@ func (p *PhysicalHashAgg) attach2TaskForMpp(tasks ...task) task {
 		if partialAgg == nil {
 			return invalidTask
 		}
+		if !aggFuncModeSame(finalAgg) {
+			return invalidTask
+		}
 		attachPlan2Task(partialAgg, mpp)
 		partitionCols := p.MppPartitionCols
 		if len(partitionCols) == 0 {
@@ -1882,6 +1888,9 @@ func (p *PhysicalHashAgg) attach2TaskForMpp(tasks ...task) task {
 		if finalAgg == nil {
 			return invalidTask
 		}
+		if !aggFuncModeSame(finalAgg) {
+			return invalidTask
+		}
 
 		// generate 3 stage aggregation for single count distinct if applicable.
 		//  select count(distinct a), count(b) from foo
@@ -1987,6 +1996,34 @@ func (p *PhysicalHashAgg) attach2TaskForMpp(tasks ...task) task {
 	}
 }
 
+func aggFuncModeSame(p PhysicalPlan) bool {
+	funcs := make([]*aggregation.AggFuncDesc, 0, 8)
+	sa, ok := p.(*PhysicalStreamAgg)
+	if ok {
+		for _, f := range sa.AggFuncs {
+			funcs = append(funcs, f)
+		}
+	} else {
+		ha, ok := p.(*PhysicalHashAgg)
+		if !ok {
+			return false
+		}
+		for _, f := range ha.AggFuncs {
+			funcs = append(funcs, f)
+		}
+	}
+	if len(funcs) == 0 {
+		return true
+	}
+	expFuncMode := funcs[0].Mode
+	for _, f := range funcs {
+		if f.Mode != expFuncMode {
+			return false
+		}
+	}
+	return true
+}
+
 func (p *PhysicalHashAgg) attach2Task(tasks ...task) task {
 	t := tasks[0].copy()
 	final := p