From 6ba2f681afa84102d13cd3d58150a316b7255724 Mon Sep 17 00:00:00 2001
From: Zhang Wenxin <101034200+morrySnow@users.noreply.github.com>
Date: Wed, 17 May 2023 17:08:42 +0800
Subject: [PATCH] [fix](Nereids) result error when do agg with distinct under
 pipeline (#19735)

---
 .../translator/PhysicalPlanTranslator.java    |  2 +-
 .../implementation/AggregateStrategies.java   |  4 ++--
 .../functions/agg/AggregateParam.java         | 22 +++++++++++++++----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
index e0340ba20b..87a3cc4318 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
@@ -391,7 +391,7 @@ public class PhysicalPlanTranslator extends DefaultPlanVisitor<PlanFragment, Pla
         // TODO: nereids forbid all parallel scan under aggregate temporary, because nereids could generate
         //  so complex aggregate plan than legacy planner, and should add forbid parallel scan hint when
         //  generate physical aggregate plan.
-        if (leftMostNode instanceof OlapScanNode && aggregate.getAggMode() == AggMode.INPUT_TO_RESULT) {
+        if (leftMostNode instanceof OlapScanNode && aggregate.getAggregateParam().needColocateScan) {
             currentFragment.getPlanRoot().setShouldColoScan();
             currentFragment.setHasColocatePlanNode(!ConnectContext.get().getSessionVariable().enablePipelineEngine());
         }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java
index 1a1f085c5c..5efef3fda6 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java
@@ -788,7 +788,7 @@ public class AggregateStrategies implements ImplementationRuleFactory {
                 .addAll(distinctArguments)
                 .build();
 
-        AggregateParam inputToBufferParam = new AggregateParam(AggPhase.LOCAL, AggMode.INPUT_TO_BUFFER);
+        AggregateParam inputToBufferParam = new AggregateParam(AggPhase.LOCAL, AggMode.INPUT_TO_BUFFER, true);
 
         Map<AggregateFunction, Alias> nonDistinctAggFunctionToAliasPhase1 = aggregateFunctions.stream()
                 .filter(aggregateFunction -> !aggregateFunction.isDistinct())
@@ -1295,7 +1295,7 @@ public class AggregateStrategies implements ImplementationRuleFactory {
                 .addAll(distinctArguments)
                 .build();
 
-        AggregateParam inputToBufferParam = new AggregateParam(AggPhase.LOCAL, AggMode.INPUT_TO_BUFFER);
+        AggregateParam inputToBufferParam = new AggregateParam(AggPhase.LOCAL, AggMode.INPUT_TO_BUFFER, true);
 
         Map<AggregateFunction, Alias> nonDistinctAggFunctionToAliasPhase1 = aggregateFunctions.stream()
                 .filter(aggregateFunction -> !aggregateFunction.isDistinct())
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AggregateParam.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AggregateParam.java
index 97851151e6..2ff8eb262f 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AggregateParam.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AggregateParam.java
@@ -29,26 +29,39 @@ public class AggregateParam {
 
     public final AggMode aggMode;
 
+    // TODO remove this flag, and generate it in enforce and cost job
+    public boolean needColocateScan;
+
     /** AggregateParam */
     public AggregateParam(AggPhase aggPhase, AggMode aggMode) {
+        this(aggPhase, aggMode, false);
+    }
+
+    /** AggregateParam */
+    public AggregateParam(AggPhase aggPhase, AggMode aggMode, boolean needColocateScan) {
         this.aggMode = Objects.requireNonNull(aggMode, "aggMode cannot be null");
         this.aggPhase = Objects.requireNonNull(aggPhase, "aggPhase cannot be null");
+        this.needColocateScan = needColocateScan;
     }
 
     public static AggregateParam localResult() {
-        return new AggregateParam(AggPhase.LOCAL, AggMode.INPUT_TO_RESULT);
+        return new AggregateParam(AggPhase.LOCAL, AggMode.INPUT_TO_RESULT, true);
     }
 
     public AggregateParam withAggPhase(AggPhase aggPhase) {
-        return new AggregateParam(aggPhase, aggMode);
+        return new AggregateParam(aggPhase, aggMode, needColocateScan);
     }
 
     public AggregateParam withAggPhase(AggMode aggMode) {
-        return new AggregateParam(aggPhase, aggMode);
+        return new AggregateParam(aggPhase, aggMode, needColocateScan);
     }
 
     public AggregateParam withAppPhaseAndAppMode(AggPhase aggPhase, AggMode aggMode) {
-        return new AggregateParam(aggPhase, aggMode);
+        return new AggregateParam(aggPhase, aggMode, needColocateScan);
+    }
+
+    public AggregateParam withNeedColocateScan(boolean needColocateScan) {
+        return new AggregateParam(aggPhase, aggMode, needColocateScan);
     }
 
     @Override
@@ -74,6 +87,7 @@ public class AggregateParam {
         return "AggregateParam{"
                 + "aggPhase=" + aggPhase
                 + ", aggMode=" + aggMode
+                + ", needColocateScan=" + needColocateScan
                 + '}';
     }
 }