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/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 28eef6bb1b4885..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 @@ -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; @@ -24,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; @@ -119,13 +122,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 +189,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 +398,84 @@ 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"); + } + + @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()); + } +} 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 + """ }