From 4c48cc17531d0006c53e10faf1fa02bbcb41f6ed Mon Sep 17 00:00:00 2001 From: 924060929 Date: Wed, 21 Oct 2026 16:43:24 +0800 Subject: [PATCH 1/2] [fix](fe) Preserve column identity for forwarded USING keys ### What problem does this PR solve? Issue Number: N/A Related PR: #64718 A qualified reference to a preserved-side USING key can be rejected even when the same key appears unqualified in GROUP BY. For example: ```sql SELECT t1.a, COUNT(*) FROM t1 LEFT JOIN t2 USING(a) GROUP BY a; ``` The USING-key projection creates a fresh Alias for a column forwarded directly from the preserved side. SELECT t1.a then references the original ExprId while GROUP BY a references the new one. Aggregate analysis rejects the query. The same mismatch affects HAVING and ORDER BY. Keep the input ExprId for forwarded USING keys while clearing the qualifier of the unqualified view. Preserve the original qualified outputs and the explicit asterisk output list. FULL OUTER JOIN still builds a new COALESCE expression with its own ExprId, and aggregate validation remains unchanged. ### Release note Fix spurious GROUP BY errors when qualified and unqualified references to a forwarded USING join key are mixed in SELECT, GROUP BY, HAVING, or ORDER BY. ### Check List (For Author) - Test: Unit Test - Complete BindUsingJoinTest suite passed through ./run-fe-ut.sh: 29 tests, 0 failures/errors. - Added complete-planner cases for both reference directions, HAVING/ORDER BY, compound grouping expressions, multiple USING keys, and derived inputs across INNER, LEFT/RIGHT, and SEMI/ANTI joins. - Existing coverage checks star expansion, qualified stars, chained joins, and FULL OUTER COALESCE keys; negative cases keep nullable-side and FULL OUTER key validation intact. - Added an EXPLAIN regression case to test_using_join_merge_key; its SQL is also exercised by the FE full-planner unit test. The regression suite itself has not been run locally. - Checkstyle and git diff --check. - Behavior changed: Yes, valid grouping queries now plan successfully. - Does this need documentation: No. --- .../rules/analysis/BindExpression.java | 8 +- .../rules/analysis/BindUsingJoinTest.java | 84 ++++++++++++++----- .../join/test_using_join_merge_key.groovy | 15 ++++ 3 files changed, 84 insertions(+), 23 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java index 456b7ca1892e8c..6b79a728f7de6d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java @@ -1053,7 +1053,9 @@ private LogicalPlan bindUsingJoin(MatchingContext> ImmutableList.Builder projectExprs = ImmutableList.builder(); ImmutableList.Builder asteriskOutputs = ImmutableList.builder(); - // 1. Add merged key columns (without qualifier) + // 1. Add merged key columns (without qualifier). A forwarded key keeps the preserved side's + // ExprId so qualified and unqualified references denote the same column in GROUP BY/HAVING/ORDER BY. + // Only FULL OUTER JOIN computes a new key and therefore needs a new Alias. for (int i = 0; i < leftConjunctsSlots.size(); i++) { Slot leftSlot = leftConjunctsSlots.get(i); Slot rightSlot = rightConjunctsSlots.get(i); @@ -1064,10 +1066,10 @@ private LogicalPlan bindUsingJoin(MatchingContext> mergeExpr = new Alias(new Coalesce(leftSlot, rightSlot), colName); } else if (joinType.isRightJoin() || joinType.isRightSemiOrAntiJoin()) { // RIGHT OUTER / RIGHT SEMI / RIGHT ANTI: use right side (preserved side) - mergeExpr = new Alias(rightSlot, colName); + mergeExpr = rightSlot.withQualifier(ImmutableList.of()); } else { // LEFT OUTER / INNER / LEFT SEMI / LEFT ANTI / CROSS: use left side - mergeExpr = new Alias(leftSlot, colName); + mergeExpr = leftSlot.withQualifier(ImmutableList.of()); } projectExprs.add(mergeExpr); asteriskOutputs.add(mergeExpr); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java index 28eef6bb1b4885..6ece9158a18eda 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java @@ -17,6 +17,7 @@ package org.apache.doris.nereids.rules.analysis; +import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.pattern.GeneratedPlanPatterns; import org.apache.doris.nereids.rules.RulePromise; import org.apache.doris.nereids.trees.expressions.Alias; @@ -119,13 +120,12 @@ void testUsingJoinProducesProjectOnTop() { // ---- 3. Merge key expression checks ---- @Test - void testInnerJoinUsingMergeKeyIsLeftAlias() { - String sql = "SELECT * FROM t1 JOIN t2 USING(a)"; - PlanChecker.from(connectContext) - .analyze(sql) - .matches(logicalProject().when(project -> project.getProjects().stream() - .filter(p -> p instanceof Alias && p.getName().equals("a")) - .anyMatch(alias -> alias.child(0) instanceof Slot))); + void testInnerJoinUsingMergeKeyKeepsLeftExprId() { + List output = PlanChecker.from(connectContext) + .analyze("SELECT a, t1.a FROM t1 JOIN t2 USING(a)").getPlan().getOutput(); + Assertions.assertEquals(output.get(1).getExprId(), output.get(0).getExprId()); + Assertions.assertTrue(output.get(0).getQualifier().isEmpty()); + Assertions.assertTrue(output.get(1).getQualifier().contains("t1")); } @Test @@ -187,19 +187,11 @@ void testRightJoinUsingAsteriskOutputHasSingleKey() { @Test void testRightJoinMergeKeyFromRightSide() { - String sql = "SELECT * FROM t1 RIGHT JOIN t2 USING(a)"; - PlanChecker.from(connectContext) - .analyze(sql) - .matches(logicalProject().when(project -> project.getProjects().stream() - .filter(p -> p instanceof Alias && p.getName().equals("a")) - .map(alias -> alias.child(0)) - .filter(Slot.class::isInstance) - .map(Slot.class::cast) - .anyMatch(slot -> { - // Merge key should come from right table (t2) - return slot.getQualifier().stream() - .anyMatch(q -> q.contains("t2")); - }))); + List output = PlanChecker.from(connectContext) + .analyze("SELECT a, t2.a FROM t1 RIGHT JOIN t2 USING(a)").getPlan().getOutput(); + Assertions.assertEquals(output.get(1).getExprId(), output.get(0).getExprId()); + Assertions.assertTrue(output.get(0).getQualifier().isEmpty()); + Assertions.assertTrue(output.get(1).getQualifier().contains("t2")); } // ---- 6. RIGHT SEMI JOIN merge key is visible ---- @@ -404,4 +396,56 @@ void testChainedFullOuterUsingJoinAsteriskOutput() { return outputNames.equals(List.of("ak", "pk", "lv", "rv", "nv")); })); } + + @Test + void testUsingJoinGroupByPreservedKey() { + String[] joinTypes = { + "JOIN", "LEFT JOIN", "LEFT SEMI JOIN", "LEFT ANTI JOIN", + "RIGHT JOIN", "RIGHT SEMI JOIN", "RIGHT ANTI JOIN" + }; + for (String joinType : joinTypes) { + String key = joinType.startsWith("RIGHT") ? "t2.a" : "t1.a"; + String from = " FROM t1 " + joinType + " t2 USING(a)"; + PlanChecker.from(connectContext).checkPlannerResult( + "SELECT " + key + ", COUNT(*)" + from + " GROUP BY a"); + PlanChecker.from(connectContext).checkPlannerResult( + "SELECT a, COUNT(*)" + from + " GROUP BY " + key); + PlanChecker.from(connectContext).checkPlannerResult( + "SELECT COUNT(*)" + from + " GROUP BY a HAVING " + key + " > 0 ORDER BY " + key); + PlanChecker.from(connectContext).checkPlannerResult( + "SELECT " + key + " + 1, COUNT(*)" + from + " GROUP BY a + 1"); + } + } + + @Test + void testUsingJoinGroupByMultipleKeys() { + PlanChecker.from(connectContext).checkPlannerResult( + "SELECT l.a, l.b, COUNT(*) FROM t1 l LEFT JOIN t1 r USING(a,b) GROUP BY a,b"); + PlanChecker.from(connectContext).checkPlannerResult( + "SELECT a, b, COUNT(*) FROM t1 l LEFT JOIN t1 r USING(a,b) GROUP BY l.a,l.b"); + } + + @Test + void testUsingJoinGroupByDoesNotMergeNullableSide() { + String[] queries = { + "SELECT t2.a, COUNT(*) FROM t1 LEFT JOIN t2 USING(a) GROUP BY a", + "SELECT t1.a, COUNT(*) FROM t1 RIGHT JOIN t2 USING(a) GROUP BY a", + "SELECT t1.a, COUNT(*) FROM t1 FULL JOIN t2 USING(a) GROUP BY a", + "SELECT a, COUNT(*) FROM t1 FULL JOIN t2 USING(a) GROUP BY t1.a" + }; + for (String sql : queries) { + AnalysisException error = Assertions.assertThrows(AnalysisException.class, + () -> PlanChecker.from(connectContext).analyze(sql), sql); + Assertions.assertTrue(error.getMessage().contains("must appear in the GROUP BY"), error.getMessage()); + } + } + + @Test + void testUsingJoinGroupByDerivedInput() { + PlanChecker.from(connectContext).checkPlannerResult( + "WITH l AS (SELECT 1 AS k, 10 AS v UNION ALL SELECT 2, 20), " + + "r AS (SELECT 1 AS k UNION ALL SELECT 3) " + + "SELECT l.k, SUM(l.v) FROM l LEFT JOIN r USING(k) " + + "GROUP BY k HAVING l.k > 0 ORDER BY l.k"); + } } diff --git a/regression-test/suites/query_p0/join/test_using_join_merge_key.groovy b/regression-test/suites/query_p0/join/test_using_join_merge_key.groovy index 1247afe8f7972a..7f61ac8e28e1a1 100644 --- a/regression-test/suites/query_p0/join/test_using_join_merge_key.groovy +++ b/regression-test/suites/query_p0/join/test_using_join_merge_key.groovy @@ -136,4 +136,19 @@ suite("test_using_join_merge_key", "query_p0") { FROM l FULL OUTER JOIN r USING(pk) FULL OUTER JOIN n USING(ak) ORDER BY l.pk NULLS LAST, r.pk NULLS LAST, n.ak NULLS LAST; """ + + // Qualified and unqualified references to a forwarded USING key must bind to the same slot. + sql """ + EXPLAIN + WITH l AS ( + SELECT 1 AS k, 10 AS v UNION ALL SELECT 2, 20 + ), r AS ( + SELECT 1 AS k UNION ALL SELECT 3 + ) + SELECT l.k, SUM(l.v) + FROM l LEFT JOIN r USING(k) + GROUP BY k + HAVING l.k > 0 + ORDER BY l.k + """ } From 363d34b70578aee08231e8082b6a0151ce199ac0 Mon Sep 17 00:00:00 2001 From: 924060929 Date: Wed, 21 Oct 2026 17:19:56 +0800 Subject: [PATCH 2/2] [fix](fe) Preserve CTE output order for repeated producer slots ### What problem does this PR solve? Issue Number: N/A Related PR: #69388 Problem Summary: A CTE selecting a forwarded USING key, a payload column, and the qualified key can expose columns in the wrong order because the reverse producer-to-consumer multimap groups equal producer ExprIds. Use the ordered consumer-to-producer map for output positions and preserve that order when rebuilding mappings for nullability, producer ExprId replacement and variant subpath pruning. Cover inner/left/right USING CTEs, FE-computed result values, consumer copies, pruning and mapping rewrites. ### Release note Preserve CTE result column order when the same source column appears more than once in its output. ### Check List (For Author) - Test: Unit Test - Official run-fe-ut.sh: 73 tests passed across BindUsingJoinTest, LogicalCTEConsumerTest, AnalyzeCTETest and EliminateGroupByKeyTest. - Reproduced the column-order failure before the repair. - Integrated Checkstyle and git diff --check passed. - Behavior changed: Yes, repeated CTE source columns retain their positions. - Does this need documentation: No --- .../nereids/rules/rewrite/AdjustNullable.java | 16 ++- .../rules/rewrite/EliminateGroupByKey.java | 10 +- .../rules/rewrite/VariantSubPathPruning.java | 2 +- .../plans/logical/LogicalCTEConsumer.java | 3 +- .../rules/analysis/BindUsingJoinTest.java | 30 +++++ .../plans/logical/LogicalCTEConsumerTest.java | 104 ++++++++++++++++++ 6 files changed, 149 insertions(+), 16 deletions(-) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumerTest.java diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java index d0d3d33822db01..73c2fe24c7d296 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java @@ -421,15 +421,13 @@ public Plan visitLogicalPartitionTopN(LogicalPartitionTopN parti public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Map replaceMap) { Map consumerToProducerOutputMap = new LinkedHashMap<>(); Multimap producerToConsumerOutputMap = LinkedHashMultimap.create(); - for (Slot producerOutputSlot : cteConsumer.getConsumerToProducerOutputMap().values()) { - Optional newProducerOutputSlot = updateExpression(producerOutputSlot, replaceMap, true); - for (Slot consumerOutputSlot : cteConsumer.getProducerToConsumerOutputMap().get(producerOutputSlot)) { - Slot slot = newProducerOutputSlot.orElse(producerOutputSlot); - Slot newConsumerOutputSlot = consumerOutputSlot.withNullable(slot.nullable()); - producerToConsumerOutputMap.put(slot, newConsumerOutputSlot); - consumerToProducerOutputMap.put(newConsumerOutputSlot, slot); - replaceMap.put(newConsumerOutputSlot.getExprId(), newConsumerOutputSlot); - } + for (Map.Entry entry : cteConsumer.getConsumerToProducerOutputMap().entrySet()) { + Slot producerOutputSlot = entry.getValue(); + Slot slot = updateExpression(producerOutputSlot, replaceMap, true).orElse(producerOutputSlot); + Slot newConsumerOutputSlot = entry.getKey().withNullable(slot.nullable()); + producerToConsumerOutputMap.put(slot, newConsumerOutputSlot); + consumerToProducerOutputMap.put(newConsumerOutputSlot, slot); + replaceMap.put(newConsumerOutputSlot.getExprId(), newConsumerOutputSlot); } return cteConsumer.withTwoMaps(consumerToProducerOutputMap, producerToConsumerOutputMap); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByKey.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByKey.java index cff93b15b2b58b..7b9d3b7dd324a6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByKey.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByKey.java @@ -143,15 +143,15 @@ public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Map newConsumerToProducer = new LinkedHashMap<>(); Multimap newProducerToConsumer = LinkedHashMultimap.create(); - for (Slot producerSlot : cteConsumer.getConsumerToProducerOutputMap().values()) { + for (Map.Entry entry : cteConsumer.getConsumerToProducerOutputMap().entrySet()) { + Slot producerSlot = entry.getValue(); ExprId newExprId = resolveExprIdChain(producerSlot.getExprId(), replaceMap); Slot effectiveProducerSlot = newExprId != null ? (Slot) producerSlot.withExprId(newExprId) : producerSlot; - for (Slot consumerSlot : cteConsumer.getProducerToConsumerOutputMap().get(producerSlot)) { - newProducerToConsumer.put(effectiveProducerSlot, consumerSlot); - newConsumerToProducer.put(consumerSlot, effectiveProducerSlot); - } + Slot consumerSlot = entry.getKey(); + newProducerToConsumer.put(effectiveProducerSlot, consumerSlot); + newConsumerToProducer.put(consumerSlot, effectiveProducerSlot); } return cteConsumer.withTwoMaps(newConsumerToProducer, newProducerToConsumer); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/VariantSubPathPruning.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/VariantSubPathPruning.java index 3036a2ba4d98af..6ce648f95d8fff 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/VariantSubPathPruning.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/VariantSubPathPruning.java @@ -422,7 +422,7 @@ public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Context cont .map(ExpressionTrait::getDataType).noneMatch(VariantType.class::isInstance)) { return cteConsumer; } - Map consumerToProducerOutputMap = Maps.newHashMap(); + Map consumerToProducerOutputMap = Maps.newLinkedHashMap(); Multimap producerToConsumerOutputMap = LinkedHashMultimap.create(); Map, SlotReference>> oriSlotToSubPathToSlot = Maps.newHashMap(); for (Map.Entry consumerToProducer : cteConsumer.getConsumerToProducerOutputMap().entrySet()) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumer.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumer.java index df4ffb7581e22e..e12301f7fdeb0f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumer.java @@ -54,6 +54,7 @@ public class LogicalCTEConsumer extends LogicalRelation implements BlockFuncDeps private final String name; private final CTEId cteId; + // Insertion order is the consumer column order, including repeated producer slots. private final Map consumerToProducerOutputMap; private final Multimap producerToConsumerOutputMap; @@ -159,7 +160,7 @@ public LogicalCTEConsumer withRelationId(RelationId relationId) { @Override public List computeOutput() { - return ImmutableList.copyOf(producerToConsumerOutputMap.values()); + return ImmutableList.copyOf(consumerToProducerOutputMap.keySet()); } @Override diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java index 6ece9158a18eda..d3f485fda57b89 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/BindUsingJoinTest.java @@ -25,6 +25,8 @@ import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.functions.scalar.Coalesce; +import org.apache.doris.nereids.trees.plans.ComputeResultSet; +import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan; import org.apache.doris.nereids.util.PlanChecker; import org.apache.doris.utframe.TestWithFeService; @@ -448,4 +450,32 @@ void testUsingJoinGroupByDerivedInput() { + "SELECT l.k, SUM(l.v) FROM l LEFT JOIN r USING(k) " + "GROUP BY k HAVING l.k > 0 ORDER BY l.k"); } + + @Test + void testUsingJoinCteColumnOrder() { + String[] joins = {"JOIN", "LEFT JOIN", "RIGHT JOIN"}; + for (String join : joins) { + String source = join.startsWith("RIGHT") ? "t2.a" : "t1.a"; + String payload = join.startsWith("RIGHT") ? "t2.c" : "t1.b"; + String sql = "WITH q(x,y,z) AS (SELECT a," + payload + "," + source + + " FROM t1 " + join + " t2 USING(a)) SELECT * FROM q"; + List output = PlanChecker.from(connectContext).analyze(sql).getPlan().getOutput(); + Assertions.assertEquals(List.of("x", "y", "z"), output.stream() + .map(Slot::getName).collect(Collectors.toList()), sql); + PlanChecker.from(connectContext).checkPlannerResult(sql); + } + } + + @Test + void testCteRepeatedColumnResults() { + String sql = "WITH q(x,y,z) AS (SELECT a,b,a FROM (SELECT 1 AS a,10 AS b) t) SELECT * FROM q"; + PlanChecker.from(connectContext).checkPlannerResult(sql, planner -> { + PhysicalPlan plan = planner.getPhysicalPlan(); + Assertions.assertTrue(plan instanceof ComputeResultSet); + Assertions.assertEquals(List.of(List.of("1", "10", "1")), + ((ComputeResultSet) plan).computeResultInFe(planner.getCascadesContext(), plan.getOutput()) + .orElseThrow().getResultRows()); + }); + } + } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumerTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumerTest.java new file mode 100644 index 00000000000000..6a20f2074cb596 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumerTest.java @@ -0,0 +1,104 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.trees.plans.logical; + +import org.apache.doris.nereids.rules.rewrite.AdjustNullable; +import org.apache.doris.nereids.rules.rewrite.EliminateGroupByKey; +import org.apache.doris.nereids.trees.expressions.CTEId; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.plans.RelationId; +import org.apache.doris.nereids.types.IntegerType; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Collectors; + +class LogicalCTEConsumerTest { + private LogicalCTEConsumer newConsumer() { + Slot key = new SlotReference("x", IntegerType.INSTANCE, false); + Slot payload = new SlotReference("y", IntegerType.INSTANCE, false); + LogicalEmptyRelation producer = new LogicalEmptyRelation(new RelationId(0), + ImmutableList.of(key, payload, key.withName("z"))); + return new LogicalCTEConsumer(new RelationId(1), new CTEId(0), "q", producer); + } + + private void assertColumnOrder(LogicalCTEConsumer consumer) { + List output = consumer.getOutput(); + Assertions.assertEquals(List.of("x", "y", "z"), output.stream() + .map(Slot::getName).collect(Collectors.toList())); + Assertions.assertEquals(consumer.getProducerSlot(output.get(0)).getExprId(), + consumer.getProducerSlot(output.get(2)).getExprId()); + Assertions.assertNotEquals(consumer.getProducerSlot(output.get(0)).getExprId(), + consumer.getProducerSlot(output.get(1)).getExprId()); + } + + @Test + void testRepeatedProducerSlotKeepsColumnOrder() { + LogicalCTEConsumer consumer = newConsumer(); + assertColumnOrder(consumer); + assertColumnOrder((LogicalCTEConsumer) consumer.withGroupExpression(Optional.empty())); + assertColumnOrder((LogicalCTEConsumer) consumer.withGroupExprLogicalPropChildren( + Optional.empty(), Optional.empty(), ImmutableList.of())); + } + + @Test + void testPruneOutputsKeepsColumnOrder() { + LogicalCTEConsumer consumer = newConsumer(); + List output = ImmutableList.copyOf(consumer.getConsumerToProducerOutputMap().keySet()); + LogicalCTEConsumer pruned = (LogicalCTEConsumer) consumer.pruneOutputs( + ImmutableList.of(output.get(1), output.get(2))); + Assertions.assertEquals(List.of("y", "z"), pruned.getOutput().stream() + .map(Slot::getName).collect(Collectors.toList())); + } + + @Test + void testAdjustNullableKeepsColumnOrder() { + LogicalCTEConsumer consumer = newConsumer(); + Slot producer = consumer.getConsumerToProducerOutputMap().values().iterator().next(); + Map replaceMap = new HashMap<>(); + replaceMap.put(producer.getExprId(), producer.withNullable(true)); + LogicalCTEConsumer updated = (LogicalCTEConsumer) new AdjustNullable(false) + .visitLogicalCTEConsumer(consumer, replaceMap); + assertColumnOrder(updated); + Assertions.assertTrue(updated.getOutput().get(0).nullable()); + Assertions.assertTrue(updated.getOutput().get(2).nullable()); + Assertions.assertFalse(updated.getOutput().get(1).nullable()); + } + + @Test + void testReplaceProducerIdKeepsColumnOrder() { + LogicalCTEConsumer consumer = newConsumer(); + Slot producer = consumer.getConsumerToProducerOutputMap().values().iterator().next(); + ExprId newId = new ExprId(999); + Map replaceMap = new HashMap<>(); + replaceMap.put(producer.getExprId(), newId); + LogicalCTEConsumer updated = (LogicalCTEConsumer) new EliminateGroupByKey() + .visitLogicalCTEConsumer(consumer, replaceMap); + assertColumnOrder(updated); + Assertions.assertEquals(newId, updated.getProducerSlot(updated.getOutput().get(0)).getExprId()); + Assertions.assertEquals(newId, updated.getProducerSlot(updated.getOutput().get(2)).getExprId()); + } +}