Skip to content

Commit 6031d37

Browse files
authored
[fix](nereids)EliminateGroupByKeyByUniform should replace exprId for alias (#60020)
### What problem does this PR solve? Input plan has: ``` Sink(output #8) +->LogicalProject[projects=[200 AS `field1`#8, field2#9]] +-->Aggregate(gropuBy=[field1#8, field2], output=[field1#8, ...] ...) +->Project(a#7 as field1, ...) +->Filter(a=2) +->Scan(T) ``` `a` has unique value, so it could be removed from groupBy the output plan is ``` Sink(output #10) +->LogicalProject[projects=[200 AS `field1`#8, field2#9]] +-->Aggregate(gropuBy=[field2], output=[any_value(field1#8) as #10, ...]) +->Project(a#7 as field1, ...) +->Filter(a=2) +->Scan(T) ``` if `200 AS field1#8` is not rewritten by EliminateGroupByKeyByUniform, planner cannot find the input slot #10 for Sink from its child.
1 parent acc2a8f commit 6031d37

File tree

3 files changed

+336
-3
lines changed

3 files changed

+336
-3
lines changed

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByKeyByUniform.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ public Plan rewriteRoot(Plan plan, JobContext jobContext) {
6464
return plan;
6565
}
6666
Map<ExprId, ExprId> replaceMap = new HashMap<>();
67-
ExprIdRewriter.ReplaceRule replaceRule = new ExprIdRewriter.ReplaceRule(replaceMap);
67+
ExprIdRewriter.ReplaceRule replaceRule = new ExprIdRewriter.ReplaceRule(replaceMap, true);
6868
exprIdReplacer = new ExprIdRewriter(replaceRule, jobContext);
6969
return plan.accept(this, replaceMap);
7070
}

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ExprIdRewriter.java

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.apache.doris.nereids.rules.expression.ExpressionRewrite;
2626
import org.apache.doris.nereids.rules.expression.ExpressionRuleExecutor;
2727
import org.apache.doris.nereids.rules.expression.ExpressionRuleType;
28+
import org.apache.doris.nereids.trees.expressions.Alias;
2829
import org.apache.doris.nereids.trees.expressions.ExprId;
2930
import org.apache.doris.nereids.trees.expressions.Expression;
3031
import org.apache.doris.nereids.trees.expressions.Slot;
@@ -95,18 +96,72 @@ public Expression visitSlotReference(SlotReference slot, Map<ExprId, ExprId> rep
9596
}
9697
}
9798
};
99+
private static final DefaultExpressionRewriter<Map<ExprId, ExprId>> SLOT_ALIAS_REPLACER =
100+
new DefaultExpressionRewriter<Map<ExprId, ExprId>>() {
101+
@Override
102+
public Expression visitSlotReference(SlotReference slot, Map<ExprId, ExprId> replaceMap) {
103+
ExprId newId = replaceMap.get(slot.getExprId());
104+
if (newId == null) {
105+
return slot;
106+
}
107+
ExprId lastId = newId;
108+
while (true) {
109+
newId = replaceMap.get(lastId);
110+
if (newId == null) {
111+
return slot.withExprId(lastId);
112+
} else {
113+
lastId = newId;
114+
}
115+
}
116+
}
117+
118+
@Override
119+
public Expression visitAlias(Alias alias, Map<ExprId, ExprId> replaceMap) {
120+
ExprId newId = replaceMap.get(alias.getExprId());
121+
if (newId == null) {
122+
return alias;
123+
}
124+
ExprId lastId = newId;
125+
while (true) {
126+
newId = replaceMap.get(lastId);
127+
if (newId == null) {
128+
return alias.withExprId(lastId);
129+
} else {
130+
lastId = newId;
131+
}
132+
}
133+
}
134+
};
135+
98136
private final Map<ExprId, ExprId> replaceMap;
137+
private DefaultExpressionRewriter<Map<ExprId, ExprId>> replacer;
99138

100-
public ReplaceRule(Map<ExprId, ExprId> replaceMap) {
139+
/**
140+
* constructor
141+
*/
142+
public ReplaceRule(Map<ExprId, ExprId> replaceMap, boolean rewriteAlias) {
101143
this.replaceMap = replaceMap;
144+
if (rewriteAlias) {
145+
replacer = SLOT_ALIAS_REPLACER;
146+
} else {
147+
replacer = SLOT_REPLACER;
148+
}
149+
}
150+
151+
public ReplaceRule(Map<ExprId, ExprId> replaceMap) {
152+
this(replaceMap, false);
102153
}
103154

104155
@Override
105156
public List<ExpressionPatternMatcher<? extends Expression>> buildRules() {
106157
return ImmutableList.of(
107158
matchesType(SlotReference.class).thenApply(ctx -> {
108159
Slot slot = ctx.expr;
109-
return slot.accept(SLOT_REPLACER, replaceMap);
160+
return slot.accept(replacer, replaceMap);
161+
}).toRule(ExpressionRuleType.EXPR_ID_REWRITE_REPLACE),
162+
matchesType(Alias.class).thenApply(ctx -> {
163+
Alias alias = ctx.expr;
164+
return alias.accept(replacer, replaceMap);
110165
}).toRule(ExpressionRuleType.EXPR_ID_REWRITE_REPLACE)
111166
);
112167
}
Lines changed: 278 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,278 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
package org.apache.doris.nereids.rules.rewrite;
19+
20+
import org.apache.doris.nereids.trees.expressions.Alias;
21+
import org.apache.doris.nereids.trees.expressions.ExprId;
22+
import org.apache.doris.nereids.trees.expressions.Expression;
23+
import org.apache.doris.nereids.trees.expressions.SlotReference;
24+
import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral;
25+
import org.apache.doris.nereids.types.IntegerType;
26+
27+
import com.google.common.collect.ImmutableList;
28+
import org.junit.jupiter.api.Assertions;
29+
import org.junit.jupiter.api.Test;
30+
31+
import java.util.HashMap;
32+
import java.util.Map;
33+
34+
/**
35+
* Test for ExprIdRewriter, especially for visitAlias functionality.
36+
*
37+
* Background: In EliminateGroupByKeyByUniform rule, when a group by key is uniform,
38+
* it will be replaced with any_value() and get a new ExprId. For example:
39+
* - Original: field1#8 in aggregate output
40+
* - After: any_value(field1#8) AS `field1`#10
41+
* This creates a mapping {8 -> 10} in replaceMap.
42+
*
43+
* The problem: If upper LogicalProject has an Alias like `200 AS field1#8`, and
44+
* LogicalResultSink expects field1#10, the Alias's ExprId must be rewritten from #8 to #10.
45+
* Without visitAlias handling, the Alias ExprId won't be replaced, causing the
46+
* LogicalResultSink to fail finding the expected input field.
47+
*/
48+
public class ExprIdRewriterTest {
49+
50+
/**
51+
* Test that visitAlias correctly rewrites Alias ExprId.
52+
*
53+
* Scenario from EliminateGroupByKeyByUniform:
54+
* - Input plan has: LogicalProject[projects=[200 AS `field1`#8, field2#9]]
55+
* - After EliminateGroupByKeyByUniform, replaceMap contains {8 -> 10}
56+
* - Expected: Alias ExprId should be rewritten from #8 to #10
57+
*
58+
* Without visitAlias: Alias remains `200 AS field1#8`, but upper plan expects #10
59+
* With visitAlias: Alias becomes `200 AS field1#10`, matching upper plan expectation
60+
*/
61+
@Test
62+
void testVisitAliasRewritesExprId() {
63+
// Create an Alias with ExprId #8: 200 AS `field1`#8
64+
ExprId originalExprId = new ExprId(8);
65+
ExprId targetExprId = new ExprId(10);
66+
IntegerLiteral literal = new IntegerLiteral(200);
67+
Alias originalAlias = new Alias(originalExprId, literal, "field1");
68+
69+
// Create replaceMap: {8 -> 10}
70+
Map<ExprId, ExprId> replaceMap = new HashMap<>();
71+
replaceMap.put(originalExprId, targetExprId);
72+
73+
// Create ReplaceRule and apply it to the Alias
74+
Expression result = originalAlias.accept(
75+
new org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter<Map<ExprId, ExprId>>() {
76+
@Override
77+
public Expression visitAlias(Alias alias, Map<ExprId, ExprId> context) {
78+
ExprId newId = context.get(alias.getExprId());
79+
if (newId == null) {
80+
return alias;
81+
}
82+
ExprId lastId = newId;
83+
while (true) {
84+
newId = context.get(lastId);
85+
if (newId == null) {
86+
return alias.withExprId(lastId);
87+
} else {
88+
lastId = newId;
89+
}
90+
}
91+
}
92+
}, replaceMap);
93+
94+
// Verify: the result should be an Alias with ExprId #10
95+
Assertions.assertTrue(result instanceof Alias, "Result should be an Alias");
96+
Alias rewrittenAlias = (Alias) result;
97+
Assertions.assertEquals(targetExprId, rewrittenAlias.getExprId(),
98+
"Alias ExprId should be rewritten from #8 to #10");
99+
Assertions.assertEquals("field1", rewrittenAlias.getName(),
100+
"Alias name should remain unchanged");
101+
Assertions.assertEquals(literal, rewrittenAlias.child(),
102+
"Alias child expression should remain unchanged");
103+
}
104+
105+
/**
106+
* Test chained ExprId replacement for Alias.
107+
* e.g., replaceMap: {0 -> 3, 1 -> 6, 6 -> 7}
108+
* Alias with ExprId #1 should be rewritten to #7 (1 -> 6 -> 7)
109+
*/
110+
@Test
111+
void testVisitAliasChainedReplacement() {
112+
// Create an Alias with ExprId #1
113+
ExprId exprId1 = new ExprId(1);
114+
ExprId exprId6 = new ExprId(6);
115+
ExprId exprId7 = new ExprId(7);
116+
IntegerLiteral literal = new IntegerLiteral(100);
117+
Alias originalAlias = new Alias(exprId1, literal, "col");
118+
119+
// Create replaceMap: {0 -> 3, 1 -> 6, 6 -> 7}
120+
Map<ExprId, ExprId> replaceMap = new HashMap<>();
121+
replaceMap.put(new ExprId(0), new ExprId(3));
122+
replaceMap.put(exprId1, exprId6);
123+
replaceMap.put(exprId6, exprId7);
124+
125+
// Apply replacement using the same logic as ReplaceRule.visitAlias
126+
Expression result = originalAlias.accept(
127+
new org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter<Map<ExprId, ExprId>>() {
128+
@Override
129+
public Expression visitAlias(Alias alias, Map<ExprId, ExprId> context) {
130+
ExprId newId = context.get(alias.getExprId());
131+
if (newId == null) {
132+
return alias;
133+
}
134+
ExprId lastId = newId;
135+
while (true) {
136+
newId = context.get(lastId);
137+
if (newId == null) {
138+
return alias.withExprId(lastId);
139+
} else {
140+
lastId = newId;
141+
}
142+
}
143+
}
144+
}, replaceMap);
145+
146+
// Verify: ExprId should follow the chain 1 -> 6 -> 7
147+
Assertions.assertTrue(result instanceof Alias, "Result should be an Alias");
148+
Alias rewrittenAlias = (Alias) result;
149+
Assertions.assertEquals(exprId7, rewrittenAlias.getExprId(),
150+
"Alias ExprId should follow chain: 1 -> 6 -> 7, final ExprId should be #7");
151+
}
152+
153+
/**
154+
* Test that Alias without mapping in replaceMap remains unchanged.
155+
*/
156+
@Test
157+
void testVisitAliasNoMapping() {
158+
// Create an Alias with ExprId #5
159+
ExprId exprId5 = new ExprId(5);
160+
IntegerLiteral literal = new IntegerLiteral(300);
161+
Alias originalAlias = new Alias(exprId5, literal, "unmapped");
162+
163+
// Create replaceMap that doesn't contain #5
164+
Map<ExprId, ExprId> replaceMap = new HashMap<>();
165+
replaceMap.put(new ExprId(1), new ExprId(2));
166+
replaceMap.put(new ExprId(3), new ExprId(4));
167+
168+
// Apply replacement
169+
Expression result = originalAlias.accept(
170+
new org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter<Map<ExprId, ExprId>>() {
171+
@Override
172+
public Expression visitAlias(Alias alias, Map<ExprId, ExprId> context) {
173+
ExprId newId = context.get(alias.getExprId());
174+
if (newId == null) {
175+
return alias;
176+
}
177+
ExprId lastId = newId;
178+
while (true) {
179+
newId = context.get(lastId);
180+
if (newId == null) {
181+
return alias.withExprId(lastId);
182+
} else {
183+
lastId = newId;
184+
}
185+
}
186+
}
187+
}, replaceMap);
188+
189+
// Verify: Alias should remain unchanged
190+
Assertions.assertSame(originalAlias, result,
191+
"Alias without mapping should remain unchanged (same object reference)");
192+
Assertions.assertEquals(exprId5, ((Alias) result).getExprId(),
193+
"Alias ExprId should remain #5");
194+
}
195+
196+
/**
197+
* Test SlotReference ExprId rewriting for comparison.
198+
* This demonstrates that both SlotReference and Alias need proper ExprId rewriting.
199+
*/
200+
@Test
201+
void testVisitSlotReferenceRewritesExprId() {
202+
// Create a SlotReference with ExprId #8
203+
ExprId originalExprId = new ExprId(8);
204+
ExprId targetExprId = new ExprId(10);
205+
SlotReference originalSlot = new SlotReference(
206+
originalExprId, "field1", IntegerType.INSTANCE, true, ImmutableList.of());
207+
208+
// Create replaceMap: {8 -> 10}
209+
Map<ExprId, ExprId> replaceMap = new HashMap<>();
210+
replaceMap.put(originalExprId, targetExprId);
211+
212+
// Apply replacement using the same logic as ReplaceRule.visitSlotReference
213+
Expression result = originalSlot.accept(
214+
new org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter<Map<ExprId, ExprId>>() {
215+
@Override
216+
public Expression visitSlotReference(SlotReference slot, Map<ExprId, ExprId> context) {
217+
ExprId newId = context.get(slot.getExprId());
218+
if (newId == null) {
219+
return slot;
220+
}
221+
ExprId lastId = newId;
222+
while (true) {
223+
newId = context.get(lastId);
224+
if (newId == null) {
225+
return slot.withExprId(lastId);
226+
} else {
227+
lastId = newId;
228+
}
229+
}
230+
}
231+
}, replaceMap);
232+
233+
// Verify: SlotReference ExprId should be rewritten
234+
Assertions.assertTrue(result instanceof SlotReference, "Result should be a SlotReference");
235+
SlotReference rewrittenSlot = (SlotReference) result;
236+
Assertions.assertEquals(targetExprId, rewrittenSlot.getExprId(),
237+
"SlotReference ExprId should be rewritten from #8 to #10");
238+
Assertions.assertEquals("field1", rewrittenSlot.getName(),
239+
"SlotReference name should remain unchanged");
240+
}
241+
242+
/**
243+
* Integration test: Verify that without visitAlias, Alias ExprId won't be rewritten,
244+
* demonstrating the necessity of visitAlias in ExprIdRewriter.ReplaceRule.
245+
*
246+
* This test shows what happens when visitAlias is NOT implemented:
247+
* - The Alias ExprId remains unchanged
248+
* - This causes issues in EliminateGroupByKeyByUniform where upper plan expects new ExprId
249+
*/
250+
@Test
251+
void testWithoutVisitAliasExprIdNotRewritten() {
252+
// Create an Alias with ExprId #8
253+
ExprId originalExprId = new ExprId(8);
254+
ExprId targetExprId = new ExprId(10);
255+
IntegerLiteral literal = new IntegerLiteral(200);
256+
Alias originalAlias = new Alias(originalExprId, literal, "field1");
257+
258+
// Create replaceMap: {8 -> 10}
259+
Map<ExprId, ExprId> replaceMap = new HashMap<>();
260+
replaceMap.put(originalExprId, targetExprId);
261+
262+
// Apply replacement WITHOUT visitAlias override (default behavior)
263+
Expression result = originalAlias.accept(
264+
new org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter<Map<ExprId, ExprId>>() {
265+
// Note: No visitAlias override - uses default behavior
266+
}, replaceMap);
267+
268+
// Without visitAlias: Alias ExprId remains #8 (unchanged)
269+
Assertions.assertTrue(result instanceof Alias, "Result should be an Alias");
270+
Alias unchangedAlias = (Alias) result;
271+
Assertions.assertEquals(originalExprId, unchangedAlias.getExprId(),
272+
"Without visitAlias, Alias ExprId should remain #8 (not rewritten)");
273+
274+
// This demonstrates the bug: upper plan expects #10 but gets #8
275+
Assertions.assertNotEquals(targetExprId, unchangedAlias.getExprId(),
276+
"This shows the problem: ExprId #10 is expected but #8 is returned");
277+
}
278+
}

0 commit comments

Comments
 (0)