Skip to content

Commit

Permalink
[Feat](nereids) add rewrite rule PushCountIntoUnionAll
Browse files Browse the repository at this point in the history
  • Loading branch information
feiniaofeiafei committed Apr 15, 2024
1 parent 9afdeda commit 1b24cf6
Showing 1 changed file with 18 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
* LogicalAggregate (groupByExpr=[c1#13], outputExpr=[c1#13, count(c1#13) AS `count(c1)`#15])
Expand All @@ -63,8 +64,8 @@
public class PushCountIntoUnionAll extends OneRewriteRuleFactory {
@Override
public Rule build() {
return logicalAggregate(logicalUnion().when(this::checkoutUnion))
.whenNot(this::hasUnsuportedAggFunc)
return logicalAggregate(logicalUnion().when(this::checkUnion))
.when(this::checkAgg)
.then(this::doPush)
.toRule(RuleType.PUSH_COUNT_INTO_UNION_ALL);
}
Expand Down Expand Up @@ -149,18 +150,30 @@ private <E extends Expression> List<E> replaceExpressionByUnionAll(List<E> expre
});
}

private boolean checkAgg(LogicalAggregate aggregate) {
Set<Count> res = ExpressionUtils.collect(aggregate.getOutputExpressions(), expr -> expr instanceof Count);
if (res.isEmpty()) {
return false;
}
return !hasUnsuportedAggFunc(aggregate);
}

private boolean hasUnsuportedAggFunc(LogicalAggregate aggregate) {
// only support count, and support count(distinct)
// only support count, not suport sum,min... and not support count(distinct)
return ExpressionUtils.deapAnyMatch(aggregate.getOutputExpressions(), expr -> {
if (expr instanceof AggregateFunction) {
return !(expr instanceof Count) || ((Count) expr).isDistinct();
if (!(expr instanceof Count)) {
return true;
} else {
return ((Count) expr).isDistinct();
}
} else {
return false;
}
});
}

private boolean checkoutUnion(LogicalUnion union) {
private boolean checkUnion(LogicalUnion union) {
if (union.getQualifier() != Qualifier.ALL) {
return false;
}
Expand Down

0 comments on commit 1b24cf6

Please sign in to comment.