[Fix](nereids) fix rule SimplifyWindowExpression (#34099)

Co-authored-by: feiniaofeiafei <moailing@selectdb.com>
This commit is contained in:
feiniaofeiafei
2024-04-25 15:01:55 +08:00
committed by yiguolei
parent f34fe46bfa
commit 987f755206
3 changed files with 53 additions and 29 deletions

View File

@ -27,10 +27,12 @@ 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.BoundFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.trees.plans.logical.LogicalWindow;
import org.apache.doris.nereids.util.TypeCoercionUtils;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
@ -87,11 +89,13 @@ public class SimplifyWindowExpression extends OneRewriteRuleFactory {
if (function instanceof BoundFunction) {
BoundFunction boundFunction = (BoundFunction) function;
String name = ((BoundFunction) function).getName();
if ((name.equals(COUNT) && boundFunction.child(0).notNullable())
if ((name.equals(COUNT) && checkCount((Count) boundFunction))
|| REWRRITE_TO_CONST_WINDOW_FUNCTIONS.contains(name)) {
projectionsBuilder.add(new Alias(alias.getExprId(), new TinyIntLiteral((byte) 1), alias.getName()));
} else if (REWRRITE_TO_SLOT_WINDOW_FUNCTIONS.contains(name)) {
projectionsBuilder.add(new Alias(alias.getExprId(), boundFunction.child(0), alias.getName()));
projectionsBuilder.add(new Alias(alias.getExprId(),
TypeCoercionUtils.castIfNotSameType(boundFunction.child(0), boundFunction.getDataType()),
alias.getName()));
} else {
remainWindowExpression.add(expr);
}
@ -120,4 +124,8 @@ public class SimplifyWindowExpression extends OneRewriteRuleFactory {
window.child(0)));
}
}
private boolean checkCount(Count count) {
return count.isCountStar() || count.child(0).notNullable();
}
}

View File

@ -119,28 +119,28 @@
-- !select_avg --
\N \N \N
\N \N \N
1 1 1
1 1 1
2 2 2
3 3 3
3 3 3
4 4 4
5 5 5
5 5 5
7 7 7
1 1.0 1.0
1 1.0 1.0
2 2.0 2.0
3 3.0 3.0
3 3.0 3.0
4 4.0 4.0
5 5.0 5.0
5 5.0 5.0
7 7.0 7.0
-- !more_than_pk --
\N \N \N
\N \N \N
1 1 1
1 1 1
2 2 2
3 3 3
3 3 3
4 4 4
5 5 5
5 5 5
7 7 7
1 1.0 1.0
1 1.0 1.0
2 2.0 2.0
3 3.0 3.0
3 3.0 3.0
4 4.0 4.0
5 5.0 5.0
5 5.0 5.0
7 7.0 7.0
-- !select_last_value_shape --
PhysicalResultSink
@ -163,18 +163,31 @@ PhysicalResultSink
------filter((mal_test_simplify_window.__DORIS_DELETE_SIGN__ = 0))
--------PhysicalOlapScan[mal_test_simplify_window]
-- !select_count_star_col1 --
\N 1 1
1 1 1
1 1 1
2 1 1
2 1 1
2 1 1
3 1 1
3 1 1
4 1 1
6 1 1
6 1 1
-- !select_upper_plan_use_all_rewrite --
\N \N
\N \N
1 1
1 1
2 2
3 3
3 3
4 4
5 5
5 5
7 7
1 1.0
1 1.0
2 2.0
3 3.0
3 3.0
4 4.0
5 5.0
5 5.0
7 7.0
-- !select_upper_plan_use_rewrite_and_not_rewrite --
\N \N \N

View File

@ -78,6 +78,9 @@ suite("simplify_window_expression") {
explain shape plan
select b, avg(b) over (partition by a,b,c) c1, avg(b) over (partition by a,b,c order by b) c2
from mal_test_simplify_window"""
qt_select_count_star_col1 """
select a,count() over (partition by a,b) c1, count() over (partition by a,b order by a) c2
from mal_test_simplify_window order by 1,2,3;"""
qt_select_upper_plan_use_all_rewrite """
select b, c1 from (select b,avg(b) over (partition by a,b) c1