[Fix](nereids) fix NormalizeAgg, change the upper project projections…

… rewrite logic (apache#36623)

cherry-pick apache#36161 to branch-2.0

NormalizeAggregate rewrite logic has a bug, for sql like this:

SELECT
	CASE
		1 WHEN CAST( NULL AS SIGNED ) THEN NULL
		WHEN COUNT( DISTINCT CAST( NULL AS SIGNED ) ) THEN NULL
		ELSE null
	END ;

This is the plan after NormalizeAggregate, the LogicalAggregate only
output `count(DISTINCT cast(NULL as SIGNED))`#3, do not output cast(NULL
as SIGNED)#2, but the upper project use cast(NULL as SIGNED)#2, so Doris
report error "cast(NULL as SIGNED) not in aggregate's output".

LogicalResultSink[29] ( outputExprs=[__case_when_0#1] ) +--LogicalProject[26] ( distinct=false, projects=[CASE WHEN (1 = cast(NULL as SIGNED)#2) THEN NULL WHEN (1 = count(DISTINCT cast(NULL as SIGNED))#3) THEN NULL ELSE NULL END AS `CASE WHEN (1 = cast(NULL as SIGNED)) THEN NULL WHEN (1 = count(DISTINCT cast(NULL as SIGNED))) THEN NULL ELSE NULL END`#1], excepts=[] )
   +--LogicalAggregate[25] ( groupByExpr=[], outputExpr=[count(DISTINCT cast(NULL as SIGNED)#2) AS `count(DISTINCT cast(NULL as SIGNED))`#3], hasRepeat=false )
      +--LogicalProject[24] ( distinct=false, projects=[cast(NULL as SIGNED) AS `cast(NULL as SIGNED)`#2], excepts=[] )
         +--LogicalOneRowRelation ( projects=[0 AS `0`#0] )

The problem is that the cast(NULL as SIGNED)#2 should not outputted by
LogicalAggregate, cast(NULL as SIGNED) should be computed in
LogicalProject.
This pr change the upper project projections rewrite logic:
aggregateOutputs is rewritten and become the upper-level LogicalProject
projections. During the rewriting process, the expressions inside the
agg function can be rewritten with expressions in aggregate function
arguments and group by expressions, but the ones outside the agg
function can only be rewritten with group by expressions.

---------

Co-authored-by: moailing <[email protected]>

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeAggregate.java

@@ -193,12 +193,16 @@ private LogicalPlan normalizeAgg(LogicalAggregate<Plan> aggregate, LogicalHaving
 
         // push down 3 kinds of exprs, these pushed exprs will be used to normalize agg output later
         // 1. group by exprs
-        // 2. trivalAgg children
-        // 3. trivalAgg input slots
-        Set<Expression> allPushDownExprs =
-                Sets.union(groupingByExprs, Sets.union(needPushSelf, needPushInputSlots));
-        NormalizeToSlotContext bottomSlotContext =
-                NormalizeToSlotContext.buildContext(existsAlias, allPushDownExprs);
+        // 2. trivialAgg children
+        // 3. trivialAgg input slots
+        // We need to distinguish between expressions in aggregate function arguments and group by expressions.
+        NormalizeToSlotContext groupByExprContext = NormalizeToSlotContext.buildContext(existsAlias, groupingByExprs);
+        Set<Alias> existsAliasAndGroupByAlias = getExistsAlias(existsAlias, groupByExprContext.getNormalizeToSlotMap());
+        Set<Expression> argsOfAggFuncNeedPushDown = Sets.union(needPushSelf, needPushInputSlots);
+        NormalizeToSlotContext argsOfAggFuncNeedPushDownContext = NormalizeToSlotContext
+                .buildContext(existsAliasAndGroupByAlias, argsOfAggFuncNeedPushDown);
+        NormalizeToSlotContext bottomSlotContext = argsOfAggFuncNeedPushDownContext.mergeContext(groupByExprContext);
+
         Set<NamedExpression> pushedGroupByExprs =
                 bottomSlotContext.pushDownToNamedExpression(groupingByExprs);
         Set<NamedExpression> pushedTrivalAggChildren =
@@ -258,8 +262,12 @@ private LogicalPlan normalizeAgg(LogicalAggregate<Plan> aggregate, LogicalHaving
                 aggregate.withNormalized(normalizedGroupExprs, normalizedAggOutput, bottomPlan);
 
         // create upper projects by normalize all output exprs in old LogicalAggregate
+        // In aggregateOutput, the expressions inside the agg function can be rewritten
+        // with expressions in aggregate function arguments and group by expressions,
+        // but the ones outside the agg function can only be rewritten with group by expressions.
+        // After the above two rewrites are completed, use aggregate output agg functions to rewrite.
         List<NamedExpression> upperProjects = normalizeOutput(aggregateOutput,
-                bottomSlotContext, normalizedAggFuncsToSlotContext);
+                groupByExprContext, argsOfAggFuncNeedPushDownContext, normalizedAggFuncsToSlotContext);
 
         // create a parent project node
         LogicalProject<Plan> project = new LogicalProject(upperProjects, newAggregate);
@@ -304,11 +312,18 @@ private LogicalPlan normalizeAgg(LogicalAggregate<Plan> aggregate, LogicalHaving
     }
 
     private List<NamedExpression> normalizeOutput(List<NamedExpression> aggregateOutput,
-            NormalizeToSlotContext groupByToSlotContext, NormalizeToSlotContext normalizedAggFuncsToSlotContext) {
+            NormalizeToSlotContext groupByToSlotContext, NormalizeToSlotContext argsOfAggFuncNeedPushDownContext,
+            NormalizeToSlotContext normalizedAggFuncsToSlotContext) {
         // build upper project, use two context to do pop up, because agg output maybe contain two part:
-        //   group by keys and agg expressions
-        List<NamedExpression> upperProjects = groupByToSlotContext
-                .normalizeToUseSlotRefWithoutWindowFunction(aggregateOutput);
+        // group by keys and agg expressions
+        List<NamedExpression> upperProjects = new ArrayList<>();
+        for (Expression expr : aggregateOutput) {
+            Expression rewrittenExpr = expr.rewriteDownShortCircuit(
+                    e -> normalizeAggFuncChildren(
+                            argsOfAggFuncNeedPushDownContext, e));
+            upperProjects.add((NamedExpression) rewrittenExpr);
+        }
+        upperProjects = groupByToSlotContext.normalizeToUseSlotRefWithoutWindowFunction(upperProjects);
         upperProjects = normalizedAggFuncsToSlotContext.normalizeToUseSlotRefWithoutWindowFunction(upperProjects);
 
         Builder<NamedExpression> builder = new ImmutableList.Builder<>();
@@ -340,4 +355,28 @@ private List<Slot> collectAllUsedSlots(List<NamedExpression> expressions) {
         slots.addAll(ExpressionUtils.getInputSlotSet(expressions));
         return slots;
     }
+
+    private Set<Alias> getExistsAlias(Set<Alias> originAliases,
+            Map<Expression, NormalizeToSlotTriplet> groupingExprMap) {
+        Set<Alias> existsAlias = Sets.newHashSet();
+        existsAlias.addAll(originAliases);
+        for (NormalizeToSlotTriplet triplet : groupingExprMap.values()) {
+            if (triplet.pushedExpr instanceof Alias) {
+                Alias alias = (Alias) triplet.pushedExpr;
+                existsAlias.add(alias);
+            }
+        }
+        return existsAlias;
+    }
+
+    private Expression normalizeAggFuncChildren(NormalizeToSlotContext context, Expression expr) {
+        if (expr instanceof AggregateFunction) {
+            AggregateFunction function = (AggregateFunction) expr;
+            List<Expression> normalizedRealExpressions = context.normalizeToUseSlotRef(function.getArguments());
+            function = function.withChildren(normalizedRealExpressions);
+            return function;
+        } else {
+            return expr;
+        }
+    }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/NormalizeToSlot.java

@@ -49,6 +49,17 @@ public NormalizeToSlotContext(Map<Expression, NormalizeToSlotTriplet> normalizeT
             this.normalizeToSlotMap = normalizeToSlotMap;
         }
 
+        public Map<Expression, NormalizeToSlotTriplet> getNormalizeToSlotMap() {
+            return normalizeToSlotMap;
+        }
+
+        public NormalizeToSlotContext mergeContext(NormalizeToSlotContext context) {
+            Map<Expression, NormalizeToSlotTriplet> newMap = Maps.newHashMap();
+            newMap.putAll(this.normalizeToSlotMap);
+            newMap.putAll(context.getNormalizeToSlotMap());
+            return new NormalizeToSlotContext(newMap);
+        }
+
         /**
          * build normalization context by follow step.
          *   1. collect all exists alias by input parameters existsAliases build a reverted map: expr -> alias

regression-test/data/nereids_rules_p0/normalize_aggregate/normalize_aggregate_test.out

@@ -0,0 +1,3 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !test_upper_project_projections_rewrite2 --
+

regression-test/suites/nereids_rules_p0/normalize_aggregate/normalize_aggregate_test.groovy

@@ -0,0 +1,27 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+suite("normalize_aggregate") {
+    sql "SET enable_nereids_planner=true"
+    sql "SET enable_fallback_to_original_planner=false"
+
+    sql "drop table if exists normalize_aggregate_tab"
+    sql """CREATE TABLE normalize_aggregate_tab(col0 INTEGER, col1 INTEGER, col2 INTEGER) distributed by hash(col0) buckets 10
+        properties('replication_num' = '1'); """
+    qt_test_upper_project_projections_rewrite2 """
+    SELECT - + AVG ( DISTINCT - col0 ) * - col0 FROM
+    normalize_aggregate_tab WHERE + - col0 IS NULL GROUP BY col0 HAVING NULL IS NULL;"""
+}