fix normalize agg

apache · Apr 28, 2024 · 9cef386 · 9cef386
1 parent d80e4ba
commit 9cef386
Showing 1 changed file with 18 additions and 16 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeRepeat.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeRepeat.java
@@ -25,7 +25,6 @@
 import org.apache.doris.nereids.trees.expressions.Alias;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.NamedExpression;
-import org.apache.doris.nereids.trees.expressions.OrderExpression;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
 import org.apache.doris.nereids.trees.expressions.VirtualSlotReference;
@@ -127,10 +126,9 @@ private LogicalAggregate<Plan> normalizeRepeat(LogicalRepeat<Plan> repeat) {
                 .collect(ImmutableList.toImmutableList());
 
         // replace the arguments of grouping scalar function to virtual slots
-        List<NamedExpression> normalizedAggOutput = repeat.getOutputExpressions().stream()
-                .map(expr -> (NamedExpression) expr.rewriteDownShortCircuit(
-                        e -> normalizeGroupingScalarFunction(context, e)))
-                .collect(Collectors.toList());
+        // replace some complex expression to slot, e.g. `a + 1`
+        List<NamedExpression> normalizedAggOutput = context.normalizeToUseSlotRef(
+                repeat.getOutputExpressions(), this::normalizeGroupingScalarFunction);
 
         Set<VirtualSlotReference> virtualSlotsInFunction =
                 ExpressionUtils.collect(normalizedAggOutput, VirtualSlotReference.class::isInstance);
@@ -185,17 +183,21 @@ private Set<Expression> collectNeedToSlotExpressions(LogicalRepeat<Plan> repeat)
                 .flatMap(function -> function.getArguments().stream())
                 .collect(ImmutableSet.toImmutableSet());
 
-        List<AggregateFunction> aggregateFunctions = CollectNonWindowedAggFuncs.collect(repeat.getOutputExpressions());
-
-        ImmutableSet<Expression> argumentsOfAggregateFunction = aggregateFunctions.stream()
-                .flatMap(function -> function.getArguments().stream().map(arg -> {
-                    if (arg instanceof OrderExpression) {
-                        return arg.child(0);
-                    } else {
-                        return arg;
-                    }
-                }))
-                .collect(ImmutableSet.toImmutableSet());
+        // List<AggregateFunction> aggregateFunctions = CollectNonWindowedAggFuncs
+        // .collect(repeat.getOutputExpressions());
+
+        // ImmutableSet<Expression> argumentsOfAggregateFunction = aggregateFunctions.stream()
+        //         .flatMap(function -> function.getArguments().stream().map(arg -> {
+        //             if (arg instanceof OrderExpression) {
+        //                 return arg.child(0);
+        //             } else {
+        //                 return arg;
+        //             }
+        //         }))
+        //         .collect(ImmutableSet.toImmutableSet());
+
+        Set<SlotReference> argumentsOfAggregateFunction = ExpressionUtils.collect(
+                repeat.getOutputExpressions(), expr -> expr.getClass().equals(SlotReference.class));
 
         ImmutableSet<Expression> needPushDown = ImmutableSet.<Expression>builder()
                 // grouping sets should be pushed down, e.g. grouping sets((k + 1)),