[fix](nereids) window function with grouping sets work not well (#31475)

```sql        
select a, c, sum(sum(b)) over(partition by c order by c rows between unbounded preceding and current row) 
from test_window_table2 group by grouping sets((a),( c)) having a > 1 order by 1,2,3;
```
for this kind of case:
sum(sum(col)) over, nereids has cannot find slot problem.
the output slot of repeat and aggregate is computed wrongly.
Only collecting the trival-agg  in NormalizeRepeat can fix this problem.


Co-authored-by: feiniaofeiafei <moailing@selectdb.com>
This commit is contained in:
feiniaofeiafei
2024-02-29 16:02:43 +08:00
committed by yiguolei
parent a6ab6c1cb8
commit 686938f5db
5 changed files with 76 additions and 23 deletions

View File

@ -31,13 +31,13 @@ import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
import org.apache.doris.nereids.trees.expressions.WindowExpression;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionVisitor;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
import org.apache.doris.nereids.trees.plans.logical.LogicalHaving;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.util.ExpressionUtils;
import org.apache.doris.nereids.util.PlanUtils;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableList.Builder;
@ -145,7 +145,7 @@ public class NormalizeAggregate implements RewriteRuleFactory, NormalizeToSlot {
// collect all trival-agg
List<NamedExpression> aggregateOutput = aggregate.getOutputExpressions();
List<AggregateFunction> aggFuncs = Lists.newArrayList();
aggregateOutput.forEach(o -> o.accept(CollectNonWindowedAggFuncs.INSTANCE, aggFuncs));
aggregateOutput.forEach(o -> o.accept(PlanUtils.CollectNonWindowedAggFuncs.INSTANCE, aggFuncs));
// split non-distinct agg child as two part
// TRUE part 1: need push down itself, if it contains subqury or window expression
@ -291,23 +291,4 @@ public class NormalizeAggregate implements RewriteRuleFactory, NormalizeToSlot {
}
return builder.build();
}
private static class CollectNonWindowedAggFuncs extends DefaultExpressionVisitor<Void, List<AggregateFunction>> {
private static final CollectNonWindowedAggFuncs INSTANCE = new CollectNonWindowedAggFuncs();
@Override
public Void visitWindow(WindowExpression windowExpression, List<AggregateFunction> context) {
for (Expression child : windowExpression.getExpressionsInWindowSpec()) {
child.accept(this, context);
}
return null;
}
@Override
public Void visitAggregateFunction(AggregateFunction aggregateFunction, List<AggregateFunction> context) {
context.add(aggregateFunction);
return null;
}
}
}

View File

@ -37,9 +37,11 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.trees.plans.logical.LogicalRepeat;
import org.apache.doris.nereids.util.ExpressionUtils;
import org.apache.doris.nereids.util.PlanUtils;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.collect.Sets.SetView;
@ -169,8 +171,9 @@ public class NormalizeRepeat extends OneAnalysisRuleFactory {
.flatMap(function -> function.getArguments().stream())
.collect(ImmutableSet.toImmutableSet());
Set<AggregateFunction> aggregateFunctions = ExpressionUtils.collect(
repeat.getOutputExpressions(), AggregateFunction.class::isInstance);
List<AggregateFunction> aggregateFunctions = Lists.newArrayList();
repeat.getOutputExpressions().forEach(
o -> o.accept(PlanUtils.CollectNonWindowedAggFuncs.INSTANCE, aggregateFunctions));
ImmutableSet<Expression> argumentsOfAggregateFunction = aggregateFunctions.stream()
.flatMap(function -> function.getArguments().stream().map(arg -> {

View File

@ -22,6 +22,9 @@ import org.apache.doris.nereids.trees.expressions.ComparisonPredicate;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.WindowExpression;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionVisitor;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
@ -121,4 +124,26 @@ public class PlanUtils {
.collect(ImmutableSet.toImmutableSet());
return resultSet;
}
/**
* collect non_window_agg_func
*/
public static class CollectNonWindowedAggFuncs extends DefaultExpressionVisitor<Void, List<AggregateFunction>> {
public static final CollectNonWindowedAggFuncs INSTANCE = new CollectNonWindowedAggFuncs();
@Override
public Void visitWindow(WindowExpression windowExpression, List<AggregateFunction> context) {
for (Expression child : windowExpression.getExpressionsInWindowSpec()) {
child.accept(this, context);
}
return null;
}
@Override
public Void visitAggregateFunction(AggregateFunction aggregateFunction, List<AggregateFunction> context) {
context.add(aggregateFunction);
return null;
}
}
}