From 1984f5c665b67651e2ee1c0f63066fecd47a5c7f Mon Sep 17 00:00:00 2001
From: englefly <minghong.zhou@163.com>
Date: Fri, 1 Dec 2023 16:30:52 +0800
Subject: [PATCH] fix temp partition

---
 .../expression/rules/PartitionPruner.java      | 10 ----------
 .../rules/rewrite/PruneOlapScanPartition.java  | 18 ++++++++++++------
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/PartitionPruner.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/PartitionPruner.java
index a033629547619d..ae3a5add083486 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/PartitionPruner.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/PartitionPruner.java
@@ -18,7 +18,6 @@
 package org.apache.doris.nereids.rules.expression.rules;
 
 import org.apache.doris.catalog.ListPartitionItem;
-import org.apache.doris.catalog.PartitionInfo;
 import org.apache.doris.catalog.PartitionItem;
 import org.apache.doris.catalog.RangePartitionItem;
 import org.apache.doris.nereids.CascadesContext;
@@ -98,15 +97,6 @@ public List<Long> prune() {
                 .collect(ImmutableList.toImmutableList());
     }
 
-    /**
-     * prune partition with `partitionInfo` as parameter.
-     */
-    public static List<Long> prune(List<Slot> partitionSlots, Expression partitionPredicate,
-            PartitionInfo partitionInfo, CascadesContext cascadesContext, PartitionTableType partitionTableType) {
-        return prune(partitionSlots, partitionPredicate, partitionInfo.getAllPartitions(), cascadesContext,
-                partitionTableType);
-    }
-
     /**
      * prune partition with `idToPartitions` as parameter.
      */
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
index 1bdff7f6c10eba..655b8c7976b78e 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneOlapScanPartition.java
@@ -19,6 +19,7 @@
 
 import org.apache.doris.catalog.OlapTable;
 import org.apache.doris.catalog.PartitionInfo;
+import org.apache.doris.catalog.PartitionItem;
 import org.apache.doris.nereids.rules.Rule;
 import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.rules.expression.rules.PartitionPruner;
@@ -29,7 +30,6 @@
 import org.apache.doris.nereids.util.Utils;
 
 import com.google.common.collect.ImmutableList;
-import org.apache.commons.collections.CollectionUtils;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -64,15 +64,21 @@ public Rule build() {
                     .stream()
                     .map(column -> scanOutput.get(column.getName().toLowerCase()))
                     .collect(Collectors.toList());
+            List<Long> manuallySpecifiedPartitions = scan.getManuallySpecifiedPartitions();
 
+            Map<Long, PartitionItem> idToPartitions;
+            if (manuallySpecifiedPartitions.isEmpty()) {
+                idToPartitions = partitionInfo.getIdToItem(false);
+            } else {
+                Map<Long, PartitionItem> allPartitions = partitionInfo.getAllPartitions();
+                idToPartitions = allPartitions.keySet().stream()
+                        .filter(id -> manuallySpecifiedPartitions.contains(id))
+                        .collect(Collectors.toMap(Function.identity(), id -> allPartitions.get(id)));
+            }
             List<Long> prunedPartitions = new ArrayList<>(PartitionPruner.prune(
-                    partitionSlots, filter.getPredicate(), partitionInfo, ctx.cascadesContext,
+                    partitionSlots, filter.getPredicate(), idToPartitions, ctx.cascadesContext,
                     PartitionTableType.OLAP));
 
-            List<Long> manuallySpecifiedPartitions = scan.getManuallySpecifiedPartitions();
-            if (!CollectionUtils.isEmpty(manuallySpecifiedPartitions)) {
-                prunedPartitions.retainAll(manuallySpecifiedPartitions);
-            }
             LogicalOlapScan rewrittenScan = scan.withSelectedPartitionIds(ImmutableList.copyOf(prunedPartitions));
             return new LogicalFilter<>(filter.getConjuncts(), rewrittenScan);
         }).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE);