[BugFix] fix list partition with multi value issue (#53746)

Signed-off-by: stephen <[email protected]>
StarRocks · Dec 10, 2024 · 5509031 · 5509031
1 parent dc3f6b0
commit 5509031
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 3 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/ListPartitionInfo.java b/fe/fe-core/src/main/java/com/starrocks/catalog/ListPartitionInfo.java
@@ -213,6 +213,18 @@ public Map<Long, List<String>> getIdToValues() {
         return idToValues;
     }
 
+    public boolean isSingleValuePartition(long id) {
+        List<String> values = getIdToValues().get(id);
+        if (values != null && values.size() == 1) {
+            return true;
+        }
+        List<List<String>> multiValues = getIdToMultiValues().get(id);
+        if (multiValues != null && multiValues.size() == 1) {
+            return true;
+        }
+        return false;
+    }
+
     /**
      * serialize data to log
      *

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/OptOlapPartitionPruner.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/OptOlapPartitionPruner.java
@@ -307,7 +307,7 @@ private static List<Long> listPartitionPrune(OlapTable olapTable, ListPartitionI
 
         List<ScalarOperator> scalarOperatorList = Utils.extractConjuncts(operator.getPredicate());
         PartitionPruner partitionPruner = new ListPartitionPruner(columnToPartitionValuesMap,
-                columnToNullPartitions, scalarOperatorList, specifyPartitionIds);
+                columnToNullPartitions, scalarOperatorList, specifyPartitionIds, listPartitionInfo);
         try {
             List<Long> prune = partitionPruner.prune();
             if (prune == null && isTemporaryPartitionPrune) {

diff --git a/...re/src/main/java/com/starrocks/sql/optimizer/rule/transformation/ListPartitionPruner.java b/...re/src/main/java/com/starrocks/sql/optimizer/rule/transformation/ListPartitionPruner.java
@@ -20,6 +20,7 @@
 import com.starrocks.analysis.BinaryType;
 import com.starrocks.analysis.Expr;
 import com.starrocks.analysis.LiteralExpr;
+import com.starrocks.catalog.ListPartitionInfo;
 import com.starrocks.catalog.Type;
 import com.starrocks.common.AnalysisException;
 import com.starrocks.common.Pair;
@@ -83,17 +84,27 @@ public class ListPartitionPruner implements PartitionPruner {
     private final Set<Long> allPartitions;
     private final List<ColumnRefOperator> partitionColumnRefs;
     private final List<Long> specifyPartitionIds;
+    private final ListPartitionInfo listPartitionInfo;
 
     public ListPartitionPruner(
             Map<ColumnRefOperator, ConcurrentNavigableMap<LiteralExpr, Set<Long>>> columnToPartitionValuesMap,
             Map<ColumnRefOperator, Set<Long>> columnToNullPartitions,
             List<ScalarOperator> partitionConjuncts, List<Long> specifyPartitionIds) {
+        this(columnToPartitionValuesMap, columnToNullPartitions, partitionConjuncts, specifyPartitionIds, null);
+    }
+
+    public ListPartitionPruner(
+            Map<ColumnRefOperator, ConcurrentNavigableMap<LiteralExpr, Set<Long>>> columnToPartitionValuesMap,
+            Map<ColumnRefOperator, Set<Long>> columnToNullPartitions,
+            List<ScalarOperator> partitionConjuncts, List<Long> specifyPartitionIds,
+            ListPartitionInfo listPartitionInfo) {
         this.columnToPartitionValuesMap = columnToPartitionValuesMap;
         this.columnToNullPartitions = columnToNullPartitions;
         this.partitionConjuncts = partitionConjuncts;
         this.allPartitions = getAllPartitions();
         this.partitionColumnRefs = getPartitionColumnRefs();
         this.specifyPartitionIds = specifyPartitionIds;
+        this.listPartitionInfo = listPartitionInfo;
     }
 
     private Set<Long> getAllPartitions() {
@@ -335,7 +346,17 @@ private Set<Long> evalBinaryPredicate(BinaryPredicateOperator binaryPredicate) {
                 matches.removeAll(nullPartitions);
                 // remove partition matches literal
                 if (partitionValueMap.containsKey(literal)) {
-                    matches.removeAll(partitionValueMap.get(literal));
+                    if (listPartitionInfo == null) {
+                        // external table
+                        matches.removeAll(partitionValueMap.get(literal));
+                    } else {
+                        Set<Long> partitionIds = partitionValueMap.get(literal);
+                        for (Long id : partitionIds) {
+                            if (listPartitionInfo.isSingleValuePartition(id)) {
+                                matches.remove(id);
+                            }
+                        }
+                    }
                 }
                 return matches;
             case LE:
@@ -448,7 +469,18 @@ private Set<Long> evalInPredicate(InPredicateOperator inPredicate) {
             Set<Long> partitions = partitionValueMap.get(literal);
             if (partitions != null) {
                 if (inPredicate.isNotIn()) {
-                    matches.removeAll(partitions);
+                    // external table, one partition column for one partition can only have one value
+                    if (listPartitionInfo == null) {
+                        matches.removeAll(partitions);
+                    } else {
+                        // for olap table, if one partition is multi value partition like PARTITION pCalifornia VALUES IN ("Los Angeles","San Francisco","San Diego")
+                        // and we have a not in predicate like city not in ("Los Angeles"), it's not safe to remove this partition
+                        for (Long id : partitions) {
+                            if (listPartitionInfo.isSingleValuePartition(id)) {
+                                matches.remove(id);
+                            }
+                        }
+                    }
                 } else {
                     matches.addAll(partitions);
                 }