From 9319d173dd75adfa0108cb27d20edfafdaa5fe9a Mon Sep 17 00:00:00 2001
From: walter <w41ter.l@gmail.com>
Date: Fri, 17 Nov 2023 23:58:21 +0800
Subject: [PATCH] [refactor](planner) filter empty partitions in a unified
 location (#27190)

---
 .../org/apache/doris/catalog/OlapTable.java     | 12 ++++++++++++
 .../rules/rewrite/PruneEmptyPartition.java      |  6 +-----
 .../org/apache/doris/planner/OlapScanNode.java  | 17 ++---------------
 3 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/OlapTable.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/OlapTable.java
index 661a7893b0..a1ca88bf7c 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/catalog/OlapTable.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/OlapTable.java
@@ -965,6 +965,18 @@ public class OlapTable extends Table {
         return partition;
     }
 
+    // select the non-empty partition ids belonging to this table.
+    //
+    // ATTN: partitions not belonging to this table will be filtered.
+    public List<Long> selectNonEmptyPartitionIds(Collection<Long> partitionIds) {
+        return partitionIds.stream()
+                .map(this::getPartition)
+                .filter(p -> p != null)
+                .filter(Partition::hasData)
+                .map(Partition::getId)
+                .collect(Collectors.toList());
+    }
+
     public int getPartitionNum() {
         return idToPartition.size();
     }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java
index 5b3154e29f..a408ea95a2 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneEmptyPartition.java
@@ -22,8 +22,6 @@ import org.apache.doris.nereids.rules.Rule;
 import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
 
-import java.util.stream.Collectors;
-
 /**
  * Used to prune empty partition.
  */
@@ -34,9 +32,7 @@ public class PruneEmptyPartition extends OneRewriteRuleFactory {
         return logicalOlapScan().thenApply(ctx -> {
             LogicalOlapScan scan = ctx.root;
             OlapTable table = scan.getTable();
-            return scan.withSelectedPartitionIds(scan.getSelectedPartitionIds().stream()
-                    .filter(partitionId -> table.getPartition(partitionId).hasData())
-                    .collect(Collectors.toList()));
+            return scan.withSelectedPartitionIds(table.selectNonEmptyPartitionIds(scan.getSelectedPartitionIds()));
         }).toRule(RuleType.PRUNE_EMPTY_PARTITION);
     }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java
index 1164e84f7e..6208d7f39b 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java
@@ -873,22 +873,9 @@ public class OlapScanNode extends ScanNode {
         if (partitionInfo.getType() == PartitionType.RANGE || partitionInfo.getType() == PartitionType.LIST) {
             selectedPartitionIds = partitionPrune(partitionInfo, partitionNames);
         } else {
-            selectedPartitionIds = null;
-        }
-
-        if (selectedPartitionIds == null) {
-            selectedPartitionIds = Lists.newArrayList();
-            for (Partition partition : olapTable.getPartitions()) {
-                if (!partition.hasData()) {
-                    continue;
-                }
-                selectedPartitionIds.add(partition.getId());
-            }
-        } else {
-            selectedPartitionIds = selectedPartitionIds.stream()
-                    .filter(id -> olapTable.getPartition(id).hasData())
-                    .collect(Collectors.toList());
+            selectedPartitionIds = olapTable.getPartitionIds();
         }
+        selectedPartitionIds = olapTable.selectNonEmptyPartitionIds(selectedPartitionIds);
         selectedPartitionNum = selectedPartitionIds.size();
 
         for (long id : selectedPartitionIds) {