branch-4.1: [fix](fe) Preserve CTE column order for repeated slots (branch-4.1)

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

@@ -426,15 +426,13 @@ public Plan visitLogicalPartitionTopN(LogicalPartitionTopN<? extends Plan> parti
     public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, Map<ExprId, Slot> replaceMap) {
         Map<Slot, Slot> consumerToProducerOutputMap = new LinkedHashMap<>();
         Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create();
-        for (Slot producerOutputSlot : cteConsumer.getConsumerToProducerOutputMap().values()) {
-            Optional<Slot> 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<Slot, Slot> 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);
     }

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<Slot, Slot> consumerToProducerOutputMap = Maps.newHashMap();
+            Map<Slot, Slot> consumerToProducerOutputMap = Maps.newLinkedHashMap();
             Multimap<Slot, Slot> producerToConsumerOutputMap = LinkedHashMultimap.create();
             Map<Slot, Map<List<String>, SlotReference>> oriSlotToSubPathToSlot = Maps.newHashMap();
             for (Map.Entry<Slot, Slot> consumerToProducer : cteConsumer.getConsumerToProducerOutputMap().entrySet()) {

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumer.java

@@ -53,6 +53,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<Slot, Slot> consumerToProducerOutputMap;
     private final Multimap<Slot, Slot> producerToConsumerOutputMap;
 
@@ -155,7 +156,7 @@ public LogicalCTEConsumer withRelationId(RelationId relationId) {
 
     @Override
     public List<Slot> computeOutput() {
-        return ImmutableList.copyOf(producerToConsumerOutputMap.values());
+        return ImmutableList.copyOf(consumerToProducerOutputMap.keySet());
     }
 
     @Override

fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AnalyzeCTETest.java

@@ -32,8 +32,10 @@
 import org.apache.doris.nereids.rules.rewrite.InApplyToJoin;
 import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderApply;
 import org.apache.doris.nereids.rules.rewrite.UnCorrelatedApplyFilter;
+import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Nullable;
+import org.apache.doris.nereids.trees.plans.ComputeResultSet;
 import org.apache.doris.nereids.trees.plans.commands.ExplainCommand;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan;
@@ -51,6 +53,8 @@
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
 public class AnalyzeCTETest extends TestWithFeService implements MemoPatternMatchSupported {
 
@@ -556,4 +560,20 @@ public void testRecCteInSubqueryException() {
         Assertions.assertTrue(
                 exception.getMessage().contains("Table [xx] does not exist in database"));
     }
+
+    @Test
+    public void testRepeatedProducerColumnOrder() {
+        String sql = "WITH q(x,y,z) AS (SELECT a,b,a FROM (SELECT 1 AS a,10 AS b) t) SELECT * FROM q";
+        List<Slot> output = PlanChecker.from(connectContext).analyze(sql).getPlan().getOutput();
+        Assertions.assertEquals(ImmutableList.of("x", "y", "z"), output.stream()
+                .map(Slot::getName).collect(Collectors.toList()));
+        PlanChecker.from(connectContext).checkPlannerResult(sql, planner -> {
+            PhysicalPlan plan = planner.getPhysicalPlan();
+            Assertions.assertTrue(plan instanceof ComputeResultSet);
+            Assertions.assertEquals(ImmutableList.of(ImmutableList.of("1", "10", "1")),
+                    ((ComputeResultSet) plan).computeResultInFe(
+                            planner.getCascadesContext(), Optional.empty(), plan.getOutput())
+                            .orElseThrow(() -> new AssertionError("Expected FE-computed CTE result")).getResultRows());
+        });
+    }
 }

fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalCTEConsumerTest.java

@@ -0,0 +1,89 @@
+// 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.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<Slot> output = consumer.getOutput();
+        Assertions.assertEquals(ImmutableList.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<Slot> output = ImmutableList.copyOf(consumer.getConsumerToProducerOutputMap().keySet());
+        LogicalCTEConsumer pruned = (LogicalCTEConsumer) consumer.pruneOutputs(
+                ImmutableList.of(output.get(1), output.get(2)));
+        Assertions.assertEquals(ImmutableList.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<ExprId, Slot> 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());
+    }
+}