[fix](nereids)only push down subquery in non-window agg functions (#2…

…6034)
apache · Nov 4, 2023 · b387ac2 · b387ac2
1 parent 071538c
commit b387ac2
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 9 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeAggregate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeAggregate.java
@@ -103,20 +103,20 @@ public Rule build() {
 
             List<NamedExpression> aggregateOutput = aggregate.getOutputExpressions();
             Set<Alias> existsAlias = ExpressionUtils.mutableCollect(aggregateOutput, Alias.class::isInstance);
-            // we need push down subquery exprs in side non-distinct agg functions
-            Set<SubqueryExpr> subqueryExprs = ExpressionUtils.mutableCollect(
-                    Lists.newArrayList(ExpressionUtils.mutableCollect(aggregateOutput,
-                            expr -> expr instanceof AggregateFunction
-                                    && !((AggregateFunction) expr).isDistinct())),
+
+            List<AggregateFunction> aggFuncs = Lists.newArrayList();
+            aggregateOutput.forEach(o -> o.accept(CollectNonWindowedAggFuncs.INSTANCE, aggFuncs));
+
+            // we need push down subquery exprs inside non-window and non-distinct agg functions
+            Set<SubqueryExpr> subqueryExprs = ExpressionUtils.mutableCollect(aggFuncs.stream()
+                    .filter(aggFunc -> !aggFunc.isDistinct()).collect(Collectors.toList()),
                     SubqueryExpr.class::isInstance);
             Set<Expression> groupingByExprs = ImmutableSet.copyOf(aggregate.getGroupByExpressions());
             NormalizeToSlotContext bottomSlotContext =
                     NormalizeToSlotContext.buildContext(existsAlias, Sets.union(groupingByExprs, subqueryExprs));
             Set<NamedExpression> bottomOutputs =
                     bottomSlotContext.pushDownToNamedExpression(Sets.union(groupingByExprs, subqueryExprs));
 
-            List<AggregateFunction> aggFuncs = Lists.newArrayList();
-            aggregateOutput.forEach(o -> o.accept(CollectNonWindowedAggFuncs.INSTANCE, aggFuncs));
             // use group by context to normalize agg functions to process
             //   sql like: select sum(a + 1) from t group by a + 1
             //

diff --git a/regression-test/data/nereids_p0/subquery/test_subquery_in_project.out b/regression-test/data/nereids_p0/subquery/test_subquery_in_project.out
@@ -54,3 +54,15 @@ true
 -- !sql16 --
 12
 
+-- !sql17 --
+12
+12
+
+-- !sql18 --
+12
+12
+
+-- !sql20 --
+5
+7
+
diff --git a/regression-test/suites/nereids_p0/subquery/test_subquery_in_project.groovy b/regression-test/suites/nereids_p0/subquery/test_subquery_in_project.groovy
@@ -117,11 +117,23 @@ suite("test_subquery_in_project") {
     """
 
     qt_sql15 """
-        select sum(age + (select sum(age) from test_sql)) from test_sql;
+        select sum(age + (select sum(age) from test_sql)) from test_sql order by 1;
     """
 
     qt_sql16 """
-        select sum(distinct age + (select sum(age) from test_sql)) from test_sql;
+        select sum(distinct age + (select sum(age) from test_sql)) from test_sql order by 1;
+    """
+
+    qt_sql17 """
+        select sum(age + (select sum(age) from test_sql)) over() from test_sql order by 1;
+    """
+
+    qt_sql18 """
+        select sum(age + (select sum(age) from test_sql)) over() from test_sql group by dt, age order by 1;
+    """
+
+    qt_sql20 """
+        select sum(age + (select sum(age) from test_sql)) from test_sql group by dt, age order by 1;
     """
 
     sql """drop table if exists test_sql;"""