Skip to content

Commit c679c83

Browse files
authored
[fix](nereids) fix fold constant return wrong scale of datetime type (#50142)
fix fold constant return wrong scale of datetime type for example, this sql should return data type `datatime(6)`, but the fold constant rule will direct return the first argument, so return the wrong data type `datetime(0)` ```sql select nvl(cast('2025-04-17 01:02:03' as datetime(0)), cast('2025-04-17 01:02:03.123456' as datetime(6))) ```
1 parent 4572b59 commit c679c83

5 files changed

Lines changed: 127 additions & 30 deletions

File tree

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
import org.apache.doris.nereids.types.DataType;
9696
import org.apache.doris.nereids.types.coercion.DateLikeType;
9797
import org.apache.doris.nereids.util.ExpressionUtils;
98+
import org.apache.doris.nereids.util.TypeCoercionUtils;
9899
import org.apache.doris.qe.ConnectContext;
99100
import org.apache.doris.qe.GlobalVariable;
100101
import org.apache.doris.thrift.TUniqueId;
@@ -539,6 +540,7 @@ public Expression visitBinaryArithmetic(BinaryArithmetic binaryArithmetic, Expre
539540

540541
@Override
541542
public Expression visitCaseWhen(CaseWhen caseWhen, ExpressionRewriteContext context) {
543+
CaseWhen originCaseWhen = caseWhen;
542544
caseWhen = rewriteChildren(caseWhen, context);
543545
Expression newDefault = null;
544546
boolean foundNewDefault = false;
@@ -564,29 +566,35 @@ public Expression visitCaseWhen(CaseWhen caseWhen, ExpressionRewriteContext cont
564566
defaultResult = newDefault;
565567
}
566568
if (whenClauses.isEmpty()) {
567-
return defaultResult == null ? new NullLiteral(caseWhen.getDataType()) : defaultResult;
569+
return TypeCoercionUtils.ensureSameResultType(
570+
originCaseWhen, defaultResult == null ? new NullLiteral(caseWhen.getDataType()) : defaultResult,
571+
context
572+
);
568573
}
569574
if (defaultResult == null) {
570575
if (caseWhen.getDataType().isNullType()) {
571576
// if caseWhen's type is NULL_TYPE, means all possible return values are nulls
572577
// it's safe to return null literal here
573578
return new NullLiteral();
574579
} else {
575-
return new CaseWhen(whenClauses);
580+
return TypeCoercionUtils.ensureSameResultType(originCaseWhen, new CaseWhen(whenClauses), context);
576581
}
577582
}
578-
return new CaseWhen(whenClauses, defaultResult);
583+
return TypeCoercionUtils.ensureSameResultType(
584+
originCaseWhen, new CaseWhen(whenClauses, defaultResult), context
585+
);
579586
}
580587

581588
@Override
582589
public Expression visitIf(If ifExpr, ExpressionRewriteContext context) {
590+
If originIf = ifExpr;
583591
ifExpr = rewriteChildren(ifExpr, context);
584592
if (ifExpr.child(0) instanceof NullLiteral || ifExpr.child(0).equals(BooleanLiteral.FALSE)) {
585-
return ifExpr.child(2);
593+
return TypeCoercionUtils.ensureSameResultType(originIf, ifExpr.child(2), context);
586594
} else if (ifExpr.child(0).equals(BooleanLiteral.TRUE)) {
587-
return ifExpr.child(1);
595+
return TypeCoercionUtils.ensureSameResultType(originIf, ifExpr.child(1), context);
588596
}
589-
return ifExpr;
597+
return TypeCoercionUtils.ensureSameResultType(originIf, ifExpr, context);
590598
}
591599

592600
@Override
@@ -684,17 +692,20 @@ public Expression visitVersion(Version version, ExpressionRewriteContext context
684692

685693
@Override
686694
public Expression visitNvl(Nvl nvl, ExpressionRewriteContext context) {
695+
Nvl originNvl = nvl;
696+
nvl = rewriteChildren(nvl, context);
697+
687698
for (Expression expr : nvl.children()) {
688699
if (expr.isLiteral()) {
689700
if (!expr.isNullLiteral()) {
690-
return expr;
701+
return TypeCoercionUtils.ensureSameResultType(originNvl, expr, context);
691702
}
692703
} else {
693-
return nvl;
704+
return TypeCoercionUtils.ensureSameResultType(originNvl, nvl, context);
694705
}
695706
}
696707
// all nulls
697-
return nvl.child(0);
708+
return TypeCoercionUtils.ensureSameResultType(originNvl, nvl.child(0), context);
698709
}
699710

700711
private <E extends Expression> E rewriteChildren(E expr, ExpressionRewriteContext context) {

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/SimplifyConditionalFunction.java

Lines changed: 31 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
package org.apache.doris.nereids.rules.expression.rules;
1919

20+
import org.apache.doris.nereids.rules.expression.ExpressionMatchingContext;
2021
import org.apache.doris.nereids.rules.expression.ExpressionPatternMatcher;
2122
import org.apache.doris.nereids.rules.expression.ExpressionPatternRuleFactory;
2223
import org.apache.doris.nereids.rules.expression.ExpressionRuleType;
@@ -26,6 +27,7 @@
2627
import org.apache.doris.nereids.trees.expressions.functions.scalar.Nullable;
2728
import org.apache.doris.nereids.trees.expressions.functions.scalar.Nvl;
2829
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
30+
import org.apache.doris.nereids.util.TypeCoercionUtils;
2931

3032
import com.google.common.collect.ImmutableList;
3133

@@ -38,11 +40,11 @@ public class SimplifyConditionalFunction implements ExpressionPatternRuleFactory
3840
@Override
3941
public List<ExpressionPatternMatcher<? extends Expression>> buildRules() {
4042
return ImmutableList.of(
41-
matchesType(Coalesce.class).then(SimplifyConditionalFunction::rewriteCoalesce)
43+
matchesType(Coalesce.class).thenApply(SimplifyConditionalFunction::rewriteCoalesce)
4244
.toRule(ExpressionRuleType.SIMPLIFY_CONDITIONAL_FUNCTION),
43-
matchesType(Nvl.class).then(SimplifyConditionalFunction::rewriteNvl)
45+
matchesType(Nvl.class).thenApply(SimplifyConditionalFunction::rewriteNvl)
4446
.toRule(ExpressionRuleType.SIMPLIFY_CONDITIONAL_FUNCTION),
45-
matchesType(NullIf.class).then(SimplifyConditionalFunction::rewriteNullIf)
47+
matchesType(NullIf.class).thenApply(SimplifyConditionalFunction::rewriteNullIf)
4648
.toRule(ExpressionRuleType.SIMPLIFY_CONDITIONAL_FUNCTION)
4749
);
4850
}
@@ -53,46 +55,52 @@ public List<ExpressionPatternMatcher<? extends Expression>> buildRules() {
5355
* coalesce(null,null) => null
5456
* coalesce(expr1) => expr1
5557
* */
56-
private static Expression rewriteCoalesce(Coalesce expression) {
57-
if (1 == expression.arity()) {
58-
return expression.child(0);
58+
private static Expression rewriteCoalesce(ExpressionMatchingContext<Coalesce> ctx) {
59+
Coalesce coalesce = ctx.expr;
60+
if (1 == coalesce.arity()) {
61+
return TypeCoercionUtils.ensureSameResultType(coalesce, coalesce.child(0), ctx.rewriteContext);
5962
}
60-
if (!(expression.child(0) instanceof NullLiteral) && expression.child(0).nullable()) {
61-
return expression;
63+
if (!(coalesce.child(0) instanceof NullLiteral) && coalesce.child(0).nullable()) {
64+
return TypeCoercionUtils.ensureSameResultType(coalesce, coalesce, ctx.rewriteContext);
6265
}
6366
ImmutableList.Builder<Expression> childBuilder = ImmutableList.builder();
64-
for (int i = 0; i < expression.arity(); i++) {
65-
Expression child = expression.children().get(i);
67+
for (int i = 0; i < coalesce.arity(); i++) {
68+
Expression child = coalesce.children().get(i);
6669
if (child instanceof NullLiteral) {
6770
continue;
6871
}
6972
if (!child.nullable()) {
70-
return child;
73+
return TypeCoercionUtils.ensureSameResultType(coalesce, child, ctx.rewriteContext);
7174
} else {
72-
for (int j = i; j < expression.arity(); j++) {
73-
childBuilder.add(expression.children().get(j));
75+
for (int j = i; j < coalesce.arity(); j++) {
76+
childBuilder.add(coalesce.children().get(j));
7477
}
7578
break;
7679
}
7780
}
7881
List<Expression> newChildren = childBuilder.build();
7982
if (newChildren.isEmpty()) {
80-
return new NullLiteral(expression.getDataType());
83+
return TypeCoercionUtils.ensureSameResultType(
84+
coalesce, new NullLiteral(coalesce.getDataType()), ctx.rewriteContext
85+
);
8186
} else {
82-
return expression.withChildren(newChildren);
87+
return TypeCoercionUtils.ensureSameResultType(
88+
coalesce, coalesce.withChildren(newChildren), ctx.rewriteContext
89+
);
8390
}
8491
}
8592

8693
/*
8794
* nvl(null,R) => R
8895
* nvl(L(not-nullable ),R) => L
8996
* */
90-
private static Expression rewriteNvl(Nvl nvl) {
97+
private static Expression rewriteNvl(ExpressionMatchingContext<Nvl> ctx) {
98+
Nvl nvl = ctx.expr;
9199
if (nvl.child(0) instanceof NullLiteral) {
92-
return nvl.child(1);
100+
return TypeCoercionUtils.ensureSameResultType(nvl, nvl.child(1), ctx.rewriteContext);
93101
}
94102
if (!nvl.child(0).nullable()) {
95-
return nvl.child(0);
103+
return TypeCoercionUtils.ensureSameResultType(nvl, nvl.child(0), ctx.rewriteContext);
96104
}
97105
return nvl;
98106
}
@@ -101,9 +109,12 @@ private static Expression rewriteNvl(Nvl nvl) {
101109
* nullif(null, R) => Null
102110
* nullif(L, null) => Null
103111
*/
104-
private static Expression rewriteNullIf(NullIf nullIf) {
112+
private static Expression rewriteNullIf(ExpressionMatchingContext<NullIf> ctx) {
113+
NullIf nullIf = ctx.expr;
105114
if (nullIf.child(0) instanceof NullLiteral || nullIf.child(1) instanceof NullLiteral) {
106-
return new Nullable(nullIf.child(0));
115+
return TypeCoercionUtils.ensureSameResultType(
116+
nullIf, new Nullable(nullIf.child(0)), ctx.rewriteContext
117+
);
107118
} else {
108119
return nullIf;
109120
}

fe/fe-core/src/main/java/org/apache/doris/nereids/util/TypeCoercionUtils.java

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
import org.apache.doris.common.Pair;
2525
import org.apache.doris.nereids.annotation.Developing;
2626
import org.apache.doris.nereids.exceptions.AnalysisException;
27+
import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext;
28+
import org.apache.doris.nereids.rules.expression.rules.FoldConstantRuleOnFE;
2729
import org.apache.doris.nereids.trees.expressions.Add;
2830
import org.apache.doris.nereids.trees.expressions.BinaryArithmetic;
2931
import org.apache.doris.nereids.trees.expressions.BinaryOperator;
@@ -153,6 +155,26 @@ public class TypeCoercionUtils {
153155

154156
private static final Logger LOG = LogManager.getLogger(TypeCoercionUtils.class);
155157

158+
/**
159+
* ensure the result's data type equals to the originExpr's dataType,
160+
* ATTN: this method usually used in fold constant rule
161+
*/
162+
public static Expression ensureSameResultType(
163+
Expression originExpr, Expression result, ExpressionRewriteContext context) {
164+
DataType originDataType = originExpr.getDataType();
165+
DataType newDataType = result.getDataType();
166+
if (originDataType.equals(newDataType)) {
167+
return result;
168+
}
169+
// backend can direct use all string like type without cast
170+
if (originDataType.isStringLikeType() && newDataType.isStringLikeType()) {
171+
return result;
172+
}
173+
return FoldConstantRuleOnFE.PATTERN_MATCH_INSTANCE.visitCast(
174+
new Cast(result, originDataType), context
175+
);
176+
}
177+
156178
/**
157179
* Return Optional.empty() if we cannot do implicit cast.
158180
*/

fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/FoldConstantTest.java

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import org.apache.doris.nereids.rules.analysis.ExpressionAnalyzer;
2626
import org.apache.doris.nereids.rules.expression.rules.FoldConstantRule;
2727
import org.apache.doris.nereids.rules.expression.rules.FoldConstantRuleOnFE;
28+
import org.apache.doris.nereids.rules.expression.rules.SimplifyConditionalFunction;
29+
import org.apache.doris.nereids.trees.expressions.CaseWhen;
2830
import org.apache.doris.nereids.trees.expressions.Cast;
2931
import org.apache.doris.nereids.trees.expressions.Expression;
3032
import org.apache.doris.nereids.trees.expressions.GreaterThan;
@@ -122,6 +124,15 @@ void testCaseWhenFold() {
122124
assertRewriteAfterTypeCoercion("case when null = 2 then 1 else 4 end", "4");
123125
assertRewriteAfterTypeCoercion("case when null = 2 then 1 end", "null");
124126
assertRewriteAfterTypeCoercion("case when TA = TB then 1 when TC is null then 2 end", "CASE WHEN (TA = TB) THEN 1 WHEN TC IS NULL THEN 2 END");
127+
128+
// make sure the case when return datetime(6)
129+
Expression analyzedCaseWhen = ExpressionAnalyzer.analyzeFunction(null, null, PARSER.parseExpression(
130+
"case when true then cast('2025-04-17' as datetime(0)) else cast('2025-04-18 01:02:03.123456' as datetime(6)) end"));
131+
Assertions.assertEquals(DateTimeV2Type.of(6), analyzedCaseWhen.getDataType());
132+
Assertions.assertEquals(DateTimeV2Type.of(6), ((CaseWhen) analyzedCaseWhen).getWhenClauses().get(0).getResult().getDataType());
133+
Assertions.assertEquals(DateTimeV2Type.of(6), ((CaseWhen) analyzedCaseWhen).getDefaultValue().get().getDataType());
134+
Expression foldCaseWhen = executor.rewrite(analyzedCaseWhen, context);
135+
Assertions.assertEquals(new DateTimeV2Literal(DateTimeV2Type.of(6), "2025-04-17"), foldCaseWhen);
125136
}
126137

127138
@Test
@@ -1175,14 +1186,21 @@ void testFoldNvl() {
11751186
executor = new ExpressionRuleExecutor(ImmutableList.of(
11761187
ExpressionAnalyzer.FUNCTION_ANALYZER_RULE,
11771188
bottomUp(
1178-
FoldConstantRule.INSTANCE
1189+
FoldConstantRule.INSTANCE,
1190+
SimplifyConditionalFunction.INSTANCE
11791191
)
11801192
));
11811193

11821194
assertRewriteExpression("nvl(NULL, 1)", "1");
11831195
assertRewriteExpression("nvl(NULL, NULL)", "NULL");
11841196
assertRewriteAfterTypeCoercion("nvl(IA, NULL)", "ifnull(IA, NULL)");
11851197
assertRewriteAfterTypeCoercion("nvl(IA, 1)", "ifnull(IA, 1)");
1198+
1199+
Expression foldNvl = executor.rewrite(
1200+
PARSER.parseExpression("nvl(cast('2025-04-17' as datetime(0)), cast('2025-04-18 01:02:03.123456' as datetime(6)))"),
1201+
context
1202+
);
1203+
Assertions.assertEquals(new DateTimeV2Literal(DateTimeV2Type.of(6), "2025-04-17"), foldNvl);
11861204
}
11871205

11881206
private void assertRewriteExpression(String actualExpression, String expectedExpression) {

fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/SimplifyConditionalFunctionTest.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,15 @@
1919

2020
import org.apache.doris.nereids.rules.expression.ExpressionRewriteTestHelper;
2121
import org.apache.doris.nereids.rules.expression.ExpressionRuleExecutor;
22+
import org.apache.doris.nereids.trees.expressions.Cast;
2223
import org.apache.doris.nereids.trees.expressions.SlotReference;
2324
import org.apache.doris.nereids.trees.expressions.functions.scalar.Coalesce;
2425
import org.apache.doris.nereids.trees.expressions.functions.scalar.NullIf;
2526
import org.apache.doris.nereids.trees.expressions.functions.scalar.Nullable;
2627
import org.apache.doris.nereids.trees.expressions.functions.scalar.Nvl;
2728
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
2829
import org.apache.doris.nereids.types.BooleanType;
30+
import org.apache.doris.nereids.types.DateTimeV2Type;
2931
import org.apache.doris.nereids.types.StringType;
3032
import org.apache.doris.nereids.types.VarcharType;
3133

@@ -71,6 +73,18 @@ public void testCoalesce() {
7173

7274
// coalesce(null, nullable_slot, literal) -> coalesce(nullable_slot, slot, literal)
7375
assertRewrite(new Coalesce(slot, nonNullableSlot), new Coalesce(slot, nonNullableSlot));
76+
77+
SlotReference datetimeSlot = new SlotReference("dt", DateTimeV2Type.of(0), false);
78+
// coalesce(null_datetime(0), non-nullable_slot_datetime(6))
79+
assertRewrite(
80+
new Coalesce(new NullLiteral(DateTimeV2Type.of(6)), datetimeSlot),
81+
new Cast(datetimeSlot, DateTimeV2Type.of(6))
82+
);
83+
// coalesce(non-nullable_slot_datetime(6), null_datetime(0))
84+
assertRewrite(
85+
new Coalesce(datetimeSlot, new NullLiteral(DateTimeV2Type.of(6))),
86+
new Cast(datetimeSlot, DateTimeV2Type.of(6))
87+
);
7488
}
7589

7690
@Test
@@ -92,6 +106,18 @@ public void testNvl() {
92106

93107
// nvl(null, null) -> null
94108
assertRewrite(new Nvl(NullLiteral.INSTANCE, NullLiteral.INSTANCE), new NullLiteral(BooleanType.INSTANCE));
109+
110+
SlotReference datetimeSlot = new SlotReference("dt", DateTimeV2Type.of(0), false);
111+
// nvl(null_datetime(0), non-nullable_slot_datetime(6))
112+
assertRewrite(
113+
new Nvl(new NullLiteral(DateTimeV2Type.of(6)), datetimeSlot),
114+
new Cast(datetimeSlot, DateTimeV2Type.of(6))
115+
);
116+
// nvl(non-nullable_slot_datetime(6), null_datetime(0))
117+
assertRewrite(
118+
new Nvl(datetimeSlot, new NullLiteral(DateTimeV2Type.of(6))),
119+
new Cast(datetimeSlot, DateTimeV2Type.of(6))
120+
);
95121
}
96122

97123
@Test
@@ -108,6 +134,15 @@ public void testNullIf() {
108134

109135
// nullif(non-nullable_slot, null) -> non-nullable_slot
110136
assertRewrite(new NullIf(nonNullableSlot, NullLiteral.INSTANCE), new Nullable(nonNullableSlot));
137+
138+
// nullif(null_datetime(0), null_datetime(6)) -> null_datetime(6)
139+
assertRewrite(
140+
new NullIf(
141+
new NullLiteral(DateTimeV2Type.of(0)),
142+
new NullLiteral(DateTimeV2Type.of(6))
143+
),
144+
new Cast(new Nullable(new NullLiteral(DateTimeV2Type.of(0))), DateTimeV2Type.of(6))
145+
);
111146
}
112147

113148
}

0 commit comments

Comments
 (0)