From fbfc5989a34c58ab43da0cc9c3934c39f8a59fc9 Mon Sep 17 00:00:00 2001
From: wuwenchi <wuwenchihdu@hotmail.com>
Date: Wed, 23 Aug 2023 20:05:34 +0800
Subject: [PATCH] fix aggOp before split

---
 .../glue/translator/PhysicalPlanTranslator.java      | 12 +++++++++---
 .../glue/translator/PlanTranslatorContext.java       | 11 +++++++++++
 2 files changed, 20 insertions(+), 3 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 2551cb64411d87..f92e0372bb160a 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
@@ -91,6 +91,7 @@
 import org.apache.doris.nereids.trees.plans.physical.PhysicalCTEAnchor;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalCTEConsumer;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalCTEProducer;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalCatalogRelation;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalDeferMaterializeOlapScan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalDeferMaterializeResultSink;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalDeferMaterializeTopN;
@@ -434,7 +435,9 @@ public PlanFragment visitPhysicalFileScan(PhysicalFileScan fileScan, PlanTransla
         } else {
             throw new RuntimeException("do not support table type " + table.getType());
         }
+
         scanNode.addConjuncts(translateToLegacyConjuncts(fileScan.getConjuncts()));
+        scanNode.setPushDownAggNoGrouping(context.getTablePushAggOp(table.getId()));
 
         TableName tableName = new TableName(null, "", "");
         TableRef ref = new TableRef(tableName, null, null);
@@ -797,9 +800,7 @@ public PlanFragment visitPhysicalStorageLayerAggregate(
                         || storageLayerAggregate.getRelation() instanceof PhysicalFileScan),
                 "PhysicalStorageLayerAggregate only support PhysicalOlapScan and PhysicalFileScan: "
                         + storageLayerAggregate.getRelation().getClass().getName());
-        PlanFragment planFragment = storageLayerAggregate.getRelation().accept(this, context);
 
-        ScanNode scanNode = (ScanNode) planFragment.getPlanRoot();
         TPushAggOp pushAggOp;
         switch (storageLayerAggregate.getAggOp()) {
             case COUNT:
@@ -815,7 +816,12 @@ public PlanFragment visitPhysicalStorageLayerAggregate(
                 throw new AnalysisException("Unsupported storage layer aggregate: "
                         + storageLayerAggregate.getAggOp());
         }
-        scanNode.setPushDownAggNoGrouping(pushAggOp);
+
+        context.setTablePushAggOp(
+                ((PhysicalCatalogRelation) storageLayerAggregate.getRelation()).getTable().getId(), pushAggOp);
+
+        PlanFragment planFragment = storageLayerAggregate.getRelation().accept(this, context);
+
         updateLegacyPlanIdToPhysicalPlan(planFragment.getPlanRoot(), storageLayerAggregate);
         return planFragment;
     }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java
index 1c176da409e355..4e90d3eea251ae 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PlanTranslatorContext.java
@@ -41,6 +41,7 @@
 import org.apache.doris.planner.PlanNode;
 import org.apache.doris.planner.PlanNodeId;
 import org.apache.doris.planner.ScanNode;
+import org.apache.doris.thrift.TPushAggOp;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
@@ -93,6 +94,8 @@ public class PlanTranslatorContext {
 
     private final Map<PlanFragmentId, CTEScanNode> cteScanNodeMap = Maps.newHashMap();
 
+    private final Map<Long, TPushAggOp> tablePushAggOp = Maps.newHashMap();
+
     public PlanTranslatorContext(CascadesContext ctx) {
         this.translator = new RuntimeFilterTranslator(ctx.getRuntimeFilterContext());
     }
@@ -235,4 +238,12 @@ public TupleDescriptor getTupleDesc(TupleId tupleId) {
     public DescriptorTable getDescTable() {
         return descTable;
     }
+
+    public void setTablePushAggOp(Long tableId, TPushAggOp aggOp) {
+        tablePushAggOp.put(tableId, aggOp);
+    }
+
+    public TPushAggOp getTablePushAggOp(Long tableId) {
+        return tablePushAggOp.getOrDefault(tableId, TPushAggOp.NONE);
+    }
 }