fix bugs

Signed-off-by: shuming.li <[email protected]>
LiShuMing · Dec 10, 2024 · ee21b3a · ee21b3a
1 parent dc535c3
commit ee21b3a
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 8 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/OlapTable.java b/fe/fe-core/src/main/java/com/starrocks/catalog/OlapTable.java
@@ -3680,19 +3680,24 @@ public void updateLastCollectProfileTime() {
         this.lastCollectProfileTime = System.currentTimeMillis();
     }
 
-    public Map<String, PCell> getPartitionCells() {
+    /**
+     * Return partition name and associate partition cell with specific partition columns.
+     * If partitionColumnsOpt is empty, return partition cell with all partition columns.
+     */
+    public Map<String, PCell> getPartitionCells(Optional<List<Column>> partitionColumnsOpt) {
         PartitionInfo partitionInfo = this.getPartitionInfo();
         if (partitionInfo.isUnPartitioned()) {
             return null;
         }
         if (partitionInfo.isRangePartition()) {
+            Preconditions.checkArgument(partitionColumnsOpt.isEmpty() || partitionColumnsOpt.get().size() == 1);
             Map<String, Range<PartitionKey>> rangeMap = getRangePartitionMap();
             if (rangeMap == null) {
                 return null;
             }
             return rangeMap.entrySet().stream().collect(Collectors.toMap(x -> x.getKey(), x -> new PRangeCell(x.getValue())));
         } else if (partitionInfo.isListPartition()) {
-            Map<String, PListCell> listMap = getListPartitionItems();
+            Map<String, PListCell> listMap = getListPartitionItems(partitionColumnsOpt);
             if (listMap == null) {
                 return null;
             }

diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/mv/MVTimelinessListPartitionArbiter.java b/fe/fe-core/src/main/java/com/starrocks/catalog/mv/MVTimelinessListPartitionArbiter.java
@@ -30,6 +30,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -89,7 +90,7 @@ public MvUpdateInfo getMVTimelinessUpdateInfoInChecked() throws AnalysisExceptio
 
         // remove ref base table's deleted partitions from `mvPartitionMap`
         // refresh ref base table's new added partitions
-        Map<String, PCell> mvPartitionNameToListMap = mv.getPartitionCells();
+        Map<String, PCell> mvPartitionNameToListMap = mv.getPartitionCells(Optional.empty());
         diff.getDeletes().keySet().forEach(mvPartitionNameToListMap::remove);
         mvPartitionNameToListMap.putAll(diff.getAdds());
 

diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/mv/MVTimelinessRangePartitionArbiter.java b/fe/fe-core/src/main/java/com/starrocks/catalog/mv/MVTimelinessRangePartitionArbiter.java
@@ -101,7 +101,7 @@ protected MvUpdateInfo getMVTimelinessUpdateInfoInChecked() throws AnalysisExcep
 
         // no needs to refresh the deleted partitions, because the deleted partitions are not in the mv's partition map.
         Set<String> mvToRefreshPartitionNames = Sets.newHashSet();
-        Map<String, PCell> mvPartitionToCells = mv.getPartitionCells();
+        Map<String, PCell> mvPartitionToCells = mv.getPartitionCells(Optional.empty());
 
         // remove ref base table's deleted partitions from `mvPartitionMap`
         mvToRefreshPartitionNames.addAll(diff.getDeletes().keySet());

diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/partitiontraits/OlapPartitionTraits.java b/fe/fe-core/src/main/java/com/starrocks/connector/partitiontraits/OlapPartitionTraits.java
@@ -63,7 +63,7 @@ public Map<String, Range<PartitionKey>> getPartitionKeyRange(Column partitionCol
 
     @Override
     public Map<String, PCell> getPartitionCells(List<Column> partitionColumns) {
-        return ((OlapTable) table).getPartitionCells();
+        return ((OlapTable) table).getPartitionCells(Optional.of(partitionColumns));
     }
 
     @Override

diff --git a/fe/fe-core/src/main/java/com/starrocks/scheduler/mv/MVPCTRefreshListPartitioner.java b/fe/fe-core/src/main/java/com/starrocks/scheduler/mv/MVPCTRefreshListPartitioner.java
@@ -63,6 +63,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -262,7 +263,7 @@ private static Expr getRefBaseTablePartitionPredicateExpr(Table table,
 
     @Override
     public Set<String> getMVPartitionsToRefreshWithForce() {
-        Map<String, PCell> mvValidListPartitionMapMap = mv.getPartitionCells();
+        Map<String, PCell> mvValidListPartitionMapMap = mv.getPartitionCells(Optional.empty());
         filterPartitionsByTTL(mvValidListPartitionMapMap, false);
         return mvValidListPartitionMapMap.keySet();
     }
@@ -348,7 +349,7 @@ public Set<String> getMVPartitionNamesWithTTL(MaterializedView mv,
                 }
             }
         } else {
-            mvListPartitionMap = mv.getPartitionCells();
+            mvListPartitionMap = mv.getPartitionCells(Optional.empty());
         }
         // filter all valid partitions by partition_retention_condition
         filterPartitionsByTTL(mvListPartitionMap, false);

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/common/ListPartitionDiffer.java b/fe/fe-core/src/main/java/com/starrocks/sql/common/ListPartitionDiffer.java
@@ -30,6 +30,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -231,7 +232,7 @@ public PartitionDiffResult computePartitionDiff(Range<PartitionKey> rangeToInclu
                                                     Map<Table, Map<String, PCell>> refBaseTablePartitionMap) {
         // generate the reference map between the base table and the mv
         // TODO: prune the partitions based on ttl
-        Map<String, PCell> mvPartitionNameToListMap = mv.getPartitionCells();
+        Map<String, PCell> mvPartitionNameToListMap = mv.getPartitionCells(Optional.empty());
 
         Map<String, PCell> allBasePartitionItems = collectBasePartitionCells(refBaseTablePartitionMap);
         PartitionDiff diff = ListPartitionDiffer.getListPartitionDiff(allBasePartitionItems, mvPartitionNameToListMap);