[opt](nereids) speedup huge InPredicate for partition pruning

fe/fe-core/src/main/java/org/apache/doris/analysis/LargeIntLiteral.java

@@ -134,7 +134,7 @@ public boolean isMinValue() {
     }
 
     @Override
-    public Object getRealValue() {
+    public BigInteger getRealValue() {
         return this.value;
     }
 

fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSortedPartitionsCacheManager.java

@@ -30,6 +30,7 @@
 import org.apache.doris.nereids.rules.expression.rules.SortedPartitionRanges.PartitionItemAndId;
 import org.apache.doris.nereids.rules.expression.rules.SortedPartitionRanges.PartitionItemAndRange;
 import org.apache.doris.nereids.trees.plans.algebra.CatalogRelation;
+import org.apache.doris.qe.ConnectContext;
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
@@ -65,6 +66,11 @@ public NereidsSortedPartitionsCacheManager() {
 
     public Optional<SortedPartitionRanges<?>> get(
             SupportBinarySearchFilteringPartitions table, CatalogRelation scan) {
+        ConnectContext connectContext = ConnectContext.get();
+        if (connectContext != null && !connectContext.getSessionVariable().enableBinarySearchFilteringPartitions) {
+            return Optional.empty();
+        }
+
         DatabaseIf<?> database = table.getDatabase();
         if (database == null) {
             return Optional.empty();

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PredicatesSplitter.java

@@ -86,7 +86,7 @@ public Void visitComparisonPredicate(ComparisonPredicate comparisonPredicate, Vo
         @Override
         public Void visitInPredicate(InPredicate inPredicate, Void context) {
             if (containOnlyColumnRef(inPredicate.getCompareExpr(), true)
-                    && (ExpressionUtils.isAllLiteral(inPredicate.getOptions()))) {
+                    && (inPredicate.optionsAreLiterals())) {
                 rangePredicates.put(inPredicate, ExpressionInfo.EMPTY);
             } else {
                 residualPredicates.put(inPredicate, ExpressionInfo.EMPTY);

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/ColumnRange.java

@@ -30,7 +30,7 @@
 
 /** ColumnRange */
 public class ColumnRange {
-    private final RangeSet<ColumnBound> rangeSet;
+    public final RangeSet<ColumnBound> rangeSet;
 
     public ColumnRange() {
         rangeSet = ImmutableRangeSet.of();
@@ -45,9 +45,25 @@ public ColumnRange(RangeSet<ColumnBound> rangeSet) {
     }
 
     public ColumnRange intersect(ColumnRange range) {
-        RangeSet<ColumnBound> newSet = TreeRangeSet.create();
-        range.rangeSet.asRanges().forEach(r -> newSet.addAll(rangeSet.subRangeSet(r)));
-        return new ColumnRange(newSet);
+        return new ColumnRange(intersect(rangeSet, range.rangeSet));
+    }
+
+    /** intersect */
+    public static <T extends Comparable<T>> RangeSet<T> intersect(RangeSet<T> r1, RangeSet<T> r2) {
+        RangeSet<T> bigRangeSet;
+        RangeSet<T> smallRangeSet;
+        if (r1.asRanges().size() < r2.asRanges().size()) {
+            bigRangeSet = r2;
+            smallRangeSet = r1;
+        } else {
+            bigRangeSet = r1;
+            smallRangeSet = r2;
+        }
+        RangeSet<T> newSet = TreeRangeSet.create();
+        for (Range<T> smallRange : smallRangeSet.asRanges()) {
+            newSet.addAll(bigRangeSet.subRangeSet(smallRange));
+        }
+        return newSet;
     }
 
     public ColumnRange union(ColumnRange range) {

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/OneListPartitionEvaluator.java

@@ -17,14 +17,20 @@
 
 package org.apache.doris.nereids.rules.expression.rules;
 
+import org.apache.doris.analysis.LiteralExpr;
 import org.apache.doris.catalog.ListPartitionItem;
+import org.apache.doris.catalog.PartitionKey;
 import org.apache.doris.common.Pair;
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext;
 import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.InPredicate;
 import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral;
 import org.apache.doris.nereids.trees.expressions.literal.Literal;
+import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
 import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
+import org.apache.doris.nereids.types.BooleanType;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -58,23 +64,74 @@ public K getPartitionIdent() {
 
     @Override
     public List<Map<Slot, PartitionSlotInput>> getOnePartitionInputs() {
+        if (partitionSlots.size() == 1) {
+            // fast path
+            return getInputsByOneSlot();
+        } else {
+            // slow path
+            return getInputsByMultiSlots();
+        }
+    }
+
+    private List<Map<Slot, PartitionSlotInput>> getInputsByOneSlot() {
+        ImmutableList.Builder<Map<Slot, PartitionSlotInput>> inputs
+                = ImmutableList.builderWithExpectedSize(partitionItem.getItems().size());
+        Slot slot = partitionSlots.get(0);
+        for (PartitionKey item : partitionItem.getItems()) {
+            LiteralExpr legacy = item.getKeys().get(0);
+            inputs.add(ImmutableMap.of(
+                    slot,
+                    new PartitionSlotInput(Literal.fromLegacyLiteral(legacy, legacy.getType()), ImmutableMap.of()))
+            );
+        }
+        return inputs.build();
+    }
+
+    private List<Map<Slot, PartitionSlotInput>> getInputsByMultiSlots() {
         return partitionItem.getItems().stream()
-            .map(keys -> {
-                List<Literal> literals = keys.getKeys()
-                        .stream()
-                        .map(literal -> Literal.fromLegacyLiteral(literal, literal.getType()))
-                        .collect(ImmutableList.toImmutableList());
-
-                return IntStream.range(0, partitionSlots.size())
-                        .mapToObj(index -> {
-                            Slot partitionSlot = partitionSlots.get(index);
-                            // partitionSlot will be replaced to this literal
-                            Literal literal = literals.get(index);
-                            // list partition don't need to compute the slot's range,
-                            // so we pass through an empty map
-                            return Pair.of(partitionSlot, new PartitionSlotInput(literal, ImmutableMap.of()));
-                        }).collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value));
-            }).collect(ImmutableList.toImmutableList());
+                .map(keys -> {
+                    List<Literal> literals = keys.getKeys()
+                            .stream()
+                            .map(literal -> Literal.fromLegacyLiteral(literal, literal.getType()))
+                            .collect(ImmutableList.toImmutableList());
+
+                    return IntStream.range(0, partitionSlots.size())
+                            .mapToObj(index -> {
+                                Slot partitionSlot = partitionSlots.get(index);
+                                // partitionSlot will be replaced to this literal
+                                Literal literal = literals.get(index);
+                                // list partition don't need to compute the slot's range,
+                                // so we pass through an empty map
+                                return Pair.of(partitionSlot, new PartitionSlotInput(literal, ImmutableMap.of()));
+                            }).collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value));
+                }).collect(ImmutableList.toImmutableList());
+    }
+
+    @Override
+    public Expression visitInPredicate(InPredicate inPredicate, Map<Slot, PartitionSlotInput> context) {
+        if (!inPredicate.optionsAreLiterals()) {
+            return super.visitInPredicate(inPredicate, context);
+        }
+
+        Expression newCompareExpr = inPredicate.getCompareExpr().accept(this, context);
+        if (newCompareExpr.isNullLiteral()) {
+            return new NullLiteral(BooleanType.INSTANCE);
+        }
+
+        try {
+            // fast path
+            boolean contains = inPredicate.getLiteralOptionSet().contains(newCompareExpr);
+            if (contains) {
+                return BooleanLiteral.TRUE;
+            }
+            if (inPredicate.optionsContainsNullLiteral()) {
+                return new NullLiteral(BooleanType.INSTANCE);
+            }
+            return BooleanLiteral.FALSE;
+        } catch (Throwable t) {
+            // slow path
+            return super.visitInPredicate(inPredicate, context);
+        }
     }
 
     @Override

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/OneRangePartitionEvaluator.java

@@ -19,7 +19,6 @@
 
 import org.apache.doris.analysis.LiteralExpr;
 import org.apache.doris.catalog.PartitionKey;
-import org.apache.doris.catalog.PrimitiveType;
 import org.apache.doris.catalog.RangePartitionItem;
 import org.apache.doris.catalog.Type;
 import org.apache.doris.common.Pair;
@@ -63,8 +62,6 @@
 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;
@@ -356,21 +353,38 @@ public EvaluateRangeResult visitNullSafeEqual(NullSafeEqual nullSafeEqual, Evalu
 
     @Override
     public EvaluateRangeResult visitInPredicate(InPredicate inPredicate, EvaluateRangeInput context) {
-        EvaluateRangeResult result = evaluateChildrenThenThis(inPredicate, context);
-        if (!(result.result instanceof InPredicate)) {
+        if (!inPredicate.optionsAreLiterals()) {
+            EvaluateRangeResult result = evaluateChildrenThenThis(inPredicate, context);
+            if (!(result.result instanceof InPredicate)) {
+                return result;
+            }
+            result = result.withRejectNot(false);
             return result;
         }
-        inPredicate = (InPredicate) result.result;
-        Map<Expression, ColumnRange> exprRanges = result.childrenResult.get(0).columnRanges;
-        if (exprRanges.containsKey(inPredicate.getCompareExpr())
-                && inPredicate.getOptions().stream().allMatch(Literal.class::isInstance)) {
-            Expression compareExpr = inPredicate.getCompareExpr();
-            ColumnRange compareExprRange = result.childrenResult.get(0).columnRanges.get(compareExpr);
-            RangeSet<ColumnBound> union = TreeRangeSet.create();
-            for (Expression expr : inPredicate.getOptions()) {
-                union.addAll(compareExprRange.intersect(ColumnRange.singleton((Literal) expr)).asRanges());
+
+        EvaluateRangeResult leftResult = inPredicate.getCompareExpr().accept(this, context);
+        Map<Expression, ColumnRange> exprRanges = leftResult.columnRanges;
+        if (!exprRanges.containsKey(inPredicate.getCompareExpr())) {
+            EvaluateRangeResult result = evaluateChildrenThenThis(inPredicate, context);
+            if (!(result.result instanceof InPredicate)) {
+                return result;
             }
-            result = intersectSlotRange(result, exprRanges, compareExpr, new ColumnRange(union));
+            result = result.withRejectNot(false);
+            return result;
+        }
+
+        Expression compareExpr = inPredicate.getCompareExpr();
+        ColumnRange compareExprRange = exprRanges.get(compareExpr);
+        ColumnRange optionsRange = inPredicate.getLiteralOptionsRangeSet();
+
+        ColumnRange intersect = optionsRange.intersect(compareExprRange);
+        Map<Expression, ColumnRange> newColumnRanges = replaceExprRange(exprRanges, compareExpr, intersect);
+
+        EvaluateRangeResult result;
+        if (intersect.isEmptyRange()) {
+            result = new EvaluateRangeResult(BooleanLiteral.FALSE, newColumnRanges, ImmutableList.of(leftResult));
+        } else {
+            result = new EvaluateRangeResult(inPredicate, newColumnRanges, ImmutableList.of(leftResult));
         }
         result = result.withRejectNot(false);
         return result;
@@ -399,7 +413,7 @@ public EvaluateRangeResult visitAnd(And and, EvaluateRangeInput context) {
         EvaluateRangeResult result = evaluateChildrenThenThis(and, context);
 
         EvaluateRangeResult andResult = result.childrenResult.get(0);
-        for (int i = 1; i < andResult.childrenResult.size(); i++) {
+        for (int i = 1; i < result.childrenResult.size(); i++) {
             andResult = mergeRanges(result.result, andResult, result.childrenResult.get(i),
                     context.rangeMap,
                     (leftRange, rightRange) -> leftRange.intersect(rightRange));
@@ -616,8 +630,7 @@ private List<Literal> toNereidsLiterals(PartitionKey partitionKey) {
     private List<Literal> toSingleNereidsLiteral(PartitionKey partitionKey) {
         List<LiteralExpr> keys = partitionKey.getKeys();
         LiteralExpr literalExpr = keys.get(0);
-        PrimitiveType primitiveType = partitionKey.getTypes().get(0);
-        Type type = Type.fromPrimitiveType(primitiveType);
+        Type type = literalExpr.getType();
         return ImmutableList.of(Literal.fromLegacyLiteral(literalExpr, type));
     }
 
@@ -626,8 +639,7 @@ private List<Literal> toMultiNereidsLiterals(PartitionKey partitionKey) {
         List<Literal> literals = Lists.newArrayListWithCapacity(keys.size());
         for (int i = 0; i < keys.size(); i++) {
             LiteralExpr literalExpr = keys.get(i);
-            PrimitiveType primitiveType = partitionKey.getTypes().get(i);
-            Type type = Type.fromPrimitiveType(primitiveType);
+            Type type = literalExpr.getType();
             literals.add(Literal.fromLegacyLiteral(literalExpr, type));
         }
         return literals;

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/OrToIn.java

@@ -80,7 +80,6 @@ public Expression rewriteTree(Expression expr, ExpressionRewriteContext context)
         ExpressionBottomUpRewriter simplify = ExpressionRewrite.bottomUp(SimplifyRange.INSTANCE);
         expr = simplify.rewrite(expr, context);
         return rewriteTree(expr);
-
     }
 
     /**

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/PartitionPredicateToRange.java

@@ -84,7 +84,7 @@ public RangeSet<MultiColumnBound> visitAnd(And and, Void context) {
                 intersects = childRanges;
                 continue;
             }
-            intersects = intersection(childRanges, intersects);
+            intersects = ColumnRange.intersect(childRanges, intersects);
             if (intersects.isEmpty()) {
                 break;
             }
@@ -165,12 +165,9 @@ public RangeSet<MultiColumnBound> visitInPredicate(InPredicate inPredicate, Void
         Expression compareExpr = inPredicate.getCompareExpr();
         if (compareExpr instanceof SlotReference) {
             SlotReference slot = (SlotReference) compareExpr;
-            if (slotIds.contains((slot).getExprId().asInt())) {
+            if (slotIds.contains((slot).getExprId().asInt()) && inPredicate.optionsAreLiterals()) {
                 RangeSet<MultiColumnBound> union = TreeRangeSet.create();
                 for (Expression option : inPredicate.getOptions()) {
-                    if (!(option instanceof Literal)) {
-                        return null;
-                    }
                     Literal literal = (Literal) option;
                     union.addAll(
                             toRangeSet(slot, literal, BoundType.CLOSED, literal, BoundType.CLOSED)
@@ -258,12 +255,4 @@ private RangeSet<MultiColumnBound> toRangeSet(SlotReference slotReference,
         Range<MultiColumnBound> range = Range.range(lowerBound, lowerBoundType, upperBound, upperBoundType);
         return ImmutableRangeSet.of(range);
     }
-
-    public static <T extends Comparable<?>> RangeSet<T> intersection(RangeSet<T> rangeSet1, RangeSet<T> rangeSet2) {
-        RangeSet<T> result = TreeRangeSet.create();
-        for (Range<T> range : rangeSet1.asRanges()) {
-            result.addAll(rangeSet2.subRangeSet(range));
-        }
-        return result;
-    }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java

@@ -406,6 +406,32 @@ public boolean isKeyColumnFromTable() {
                 && ((SlotReference) this).getColumn().get().isKey();
     }
 
+    /** containsNullLiteralChildren */
+    public boolean containsNullLiteralChildren() {
+        return getOrInitMutableState("CONTAINS_NULL_LITERAL_CHILDREN", () -> {
+            for (Expression child : children) {
+                if (child instanceof NullLiteral) {
+                    return true;
+                }
+            }
+            return false;
+        });
+    }
+
+    /** allChildrenAreLiteral */
+    public boolean allChildrenAreLiteral() {
+        return getOrInitMutableState("ALL_CHILDREN_ARE_LITERAL", () -> {
+            boolean allLiteral = true;
+            for (Expression child : getArguments()) {
+                if (!(child instanceof Literal)) {
+                    allLiteral = false;
+                    break;
+                }
+            }
+            return allLiteral;
+        });
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {