[fix](Nerieds) column prune should retain at least one column for uni…

…on all (#41613) (#41771)

cherry-pick #41613 to branch-2.1

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

@@ -322,17 +322,6 @@ private LogicalUnion pruneUnionOutput(LogicalUnion union, PruneContext context)
                 extractColumnIndex.add(i);
             }
         }
-        int len = extractColumnIndex.size();
-        ImmutableList.Builder<List<NamedExpression>> prunedConstantExprsList
-                = ImmutableList.builderWithExpectedSize(constantExprsList.size());
-        for (List<NamedExpression> row : constantExprsList) {
-            ImmutableList.Builder<NamedExpression> newRow = ImmutableList.builderWithExpectedSize(len);
-            for (int idx : extractColumnIndex) {
-                newRow.add(row.get(idx));
-            }
-            prunedConstantExprsList.add(newRow.build());
-        }
-
         if (prunedOutputs.isEmpty()) {
             List<NamedExpression> candidates = Lists.newArrayList(originOutput);
             candidates.retainAll(keys);
@@ -341,8 +330,19 @@ private LogicalUnion pruneUnionOutput(LogicalUnion union, PruneContext context)
             }
             NamedExpression minimumColumn = ExpressionUtils.selectMinimumColumn(candidates);
             prunedOutputs = ImmutableList.of(minimumColumn);
+            extractColumnIndex.add(originOutput.indexOf(minimumColumn));
         }
 
+        int len = extractColumnIndex.size();
+        ImmutableList.Builder<List<NamedExpression>> prunedConstantExprsList
+                = ImmutableList.builderWithExpectedSize(constantExprsList.size());
+        for (List<NamedExpression> row : constantExprsList) {
+            ImmutableList.Builder<NamedExpression> newRow = ImmutableList.builderWithExpectedSize(len);
+            for (int idx : extractColumnIndex) {
+                newRow.add(row.get(idx));
+            }
+            prunedConstantExprsList.add(newRow.build());
+        }
         if (prunedOutputs.equals(originOutput)) {
             return union;
         } else {

regression-test/suites/nereids_rules_p0/column_pruning/union_const_expr_column_pruning.groovy

@@ -0,0 +1,23 @@
+// 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.
+
+suite("const_expr_column_pruning") {
+    sql """SET ignore_shape_nodes='PhysicalDistribute,PhysicalProject'"""
+    // should only keep one column in union
+    sql "select count(1) from(select 3, 6 union all select 1, 3) t"
+    sql "select count(a) from(select 3 a, 6 union all select 1, 3) t"
+}