From afdc6800d82ea831ad256930be161632f8b41d4c Mon Sep 17 00:00:00 2001
From: feiniaofeiafei <53502832+feiniaofeiafei@users.noreply.github.com>
Date: Tue, 15 Oct 2024 00:09:02 +0800
Subject: [PATCH] [fix](Nerieds) column prune should retain at least one column
 for union all  (#41613) (#41771)

cherry-pick #41613 to branch-2.1
---
 .../nereids/rules/rewrite/ColumnPruning.java  | 22 +++++++++---------
 .../union_const_expr_column_pruning.groovy    | 23 +++++++++++++++++++
 2 files changed, 34 insertions(+), 11 deletions(-)
 create mode 100644 regression-test/suites/nereids_rules_p0/column_pruning/union_const_expr_column_pruning.groovy

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java
index b10a87479f2c95..4a60e94f76f766 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java
+++ b/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 {
diff --git a/regression-test/suites/nereids_rules_p0/column_pruning/union_const_expr_column_pruning.groovy b/regression-test/suites/nereids_rules_p0/column_pruning/union_const_expr_column_pruning.groovy
new file mode 100644
index 00000000000000..77d62a2189960b
--- /dev/null
+++ b/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"
+}
\ No newline at end of file