branch-2.1: [fix](nereids) fix merge_percentile_to_array when has same agg function #44783 (#44879)

Cherry-picked from #44783

Co-authored-by: feiniaofeiafei <moailing@selectdb.com>
This commit is contained in:
github-actions[bot]
2024-12-04 22:02:16 +08:00
committed by GitHub
parent 60dde858b9
commit 191c86bb8b
3 changed files with 28 additions and 14 deletions

View File

@ -152,10 +152,10 @@ public class MergePercentileToArray extends OneRewriteRuleFactory {
(List<Expression>) (List) newPercentileArrays);
ImmutableList.Builder<NamedExpression> newProjectOutputExpressions = ImmutableList.builder();
newProjectOutputExpressions.addAll((List<NamedExpression>) (List) notChangeForProject);
Map<Expression, Alias> existsAliasMap = Maps.newHashMap();
Map<Expression, List<Alias>> existsAliasMap = Maps.newHashMap();
// existsAliasMap is used to keep upper plan refer the same expr
for (Alias alias : existsAliases) {
existsAliasMap.put(alias.child(), alias);
existsAliasMap.computeIfAbsent(alias.child(), k -> new ArrayList<>()).add(alias);
}
Map<DistinctAndExpr, Slot> slotMap = Maps.newHashMap();
// slotMap is used to find the correspondence
@ -169,20 +169,22 @@ public class MergePercentileToArray extends OneRewriteRuleFactory {
for (Map.Entry<DistinctAndExpr, List<AggregateFunction>> entry : funcMap.entrySet()) {
for (int i = 0; i < entry.getValue().size(); i++) {
AggregateFunction aggFunc = entry.getValue().get(i);
Alias originAlias = existsAliasMap.get(aggFunc);
DistinctAndExpr distinctAndExpr = new DistinctAndExpr(aggFunc.child(0), aggFunc.isDistinct());
Alias newAlias = new Alias(originAlias.getExprId(), new ElementAt(slotMap.get(distinctAndExpr),
new IntegerLiteral(i + 1)), originAlias.getName());
newProjectOutputExpressions.add(newAlias);
List<Alias> originAliases = existsAliasMap.get(aggFunc);
for (Alias originAlias : originAliases) {
DistinctAndExpr distinctAndExpr = new DistinctAndExpr(aggFunc.child(0), aggFunc.isDistinct());
Alias newAlias = new Alias(originAlias.getExprId(), new ElementAt(slotMap.get(distinctAndExpr),
new IntegerLiteral(i + 1)), originAlias.getName());
newProjectOutputExpressions.add(newAlias);
}
}
}
newProjectOutputExpressions.addAll(groupBySlots);
return new LogicalProject(newProjectOutputExpressions.build(), newAggregate);
return new LogicalProject<>(newProjectOutputExpressions.build(), newAggregate);
}
private static class DistinctAndExpr {
private Expression expression;
private boolean isDistinct;
private final Expression expression;
private final boolean isDistinct;
public DistinctAndExpr(Expression expression, boolean isDistinct) {
this.expression = expression;
@ -193,10 +195,6 @@ public class MergePercentileToArray extends OneRewriteRuleFactory {
return expression;
}
public boolean isDistinct() {
return isDistinct;
}
@Override
public boolean equals(Object o) {
if (this == o) {

View File

@ -41,3 +41,15 @@
7.0 \N \N
7.0 7.0 7
-- !same_percentile --
52 1.0 1.0 2.0
-- !same_percentile_group_by --
\N 6.0 6.0 6.0
2 3.0 3.0 3.0
25 3.0 3.0 3.0
4 2.0 2.0 2.0
5 1.0 1.0 1.6
7 6.0 6.0 6.0
9 1.2 1.2 1.8

View File

@ -57,4 +57,8 @@ suite("merge_percentile_to_array") {
percentile(abs(a), 0.55) as c2 from test_merge_percentile group by a) t;
"""
order_qt_same_percentile """select sum(a),percentile(pk, 0.1) as c1 , percentile(pk, 0.1) as c2 ,
percentile(pk, 0.4) as c2 from test_merge_percentile;"""
order_qt_same_percentile_group_by """select sum(a),percentile(pk, 0.1) as c1 , percentile(pk, 0.1) as c2 ,
percentile(pk, 0.4) as c2 from test_merge_percentile group by a;"""
}