[opt](nereids) improve prune partition with lots of in (xxx) (#46261)

improve prune partition with lots of `in (xxx)` this can reduce from 35s to 300ms when there have 4000 options
apache · Jan 2, 2025 · 5fc52ba · 5fc52ba
1 parent 518faa3
commit 5fc52ba
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/...main/java/org/apache/doris/nereids/rules/expression/rules/OneRangePartitionEvaluator.java b/...main/java/org/apache/doris/nereids/rules/expression/rules/OneRangePartitionEvaluator.java
@@ -63,6 +63,8 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Range;
+import com.google.common.collect.RangeSet;
+import com.google.common.collect.TreeRangeSet;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -363,13 +365,12 @@ public EvaluateRangeResult visitInPredicate(InPredicate inPredicate, EvaluateRan
         if (exprRanges.containsKey(inPredicate.getCompareExpr())
                 && inPredicate.getOptions().stream().allMatch(Literal.class::isInstance)) {
             Expression compareExpr = inPredicate.getCompareExpr();
-            ColumnRange unionLiteralRange = ColumnRange.empty();
             ColumnRange compareExprRange = result.childrenResult.get(0).columnRanges.get(compareExpr);
+            RangeSet<ColumnBound> union = TreeRangeSet.create();
             for (Expression expr : inPredicate.getOptions()) {
-                unionLiteralRange = unionLiteralRange.union(
-                        compareExprRange.intersect(ColumnRange.singleton((Literal) expr)));
+                union.addAll(compareExprRange.intersect(ColumnRange.singleton((Literal) expr)).asRanges());
             }
-            result = intersectSlotRange(result, exprRanges, compareExpr, unionLiteralRange);
+            result = intersectSlotRange(result, exprRanges, compareExpr, new ColumnRange(union));
         }
         result = result.withRejectNot(false);
         return result;