|
| 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