rewrite large In to min/max for partition prune

englefly · Dec 11, 2023 · 0914d6e · 0914d6e
1 parent ae923f7
commit 0914d6e
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 1 deletion.
diff --git a/...a/org/apache/doris/nereids/rules/expression/rules/TryEliminateUninterestedPredicates.java b/...a/org/apache/doris/nereids/rules/expression/rules/TryEliminateUninterestedPredicates.java
@@ -20,9 +20,19 @@
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext;
 import org.apache.doris.nereids.rules.expression.rules.TryEliminateUninterestedPredicates.Context;
+import org.apache.doris.nereids.trees.expressions.And;
 import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.GreaterThanEqual;
+import org.apache.doris.nereids.trees.expressions.InPredicate;
+import org.apache.doris.nereids.trees.expressions.LessThanEqual;
 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.DateLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.DateTimeLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.DoubleLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.FloatLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.Literal;
 import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
 
 import java.util.Set;
@@ -50,13 +60,64 @@ private TryEliminateUninterestedPredicates(Set<Slot> interestedSlots, CascadesCo
         this.expressionRewriteContext = new ExpressionRewriteContext(cascadesContext);
     }
 
+    /**
+     * rewrite
+     */
     public static Expression rewrite(Expression expression, Set<Slot> interestedSlots,
             CascadesContext cascadesContext) {
         // before eliminate uninterested predicate, we must push down `Not` under CompoundPredicate
         expression = expression.accept(new SimplifyNotExprRule(), null);
         TryEliminateUninterestedPredicates rewriter = new TryEliminateUninterestedPredicates(
                 interestedSlots, cascadesContext);
-        return expression.accept(rewriter, new Context());
+        Expression eliminated = expression.accept(rewriter, new Context());
+        if (cascadesContext.getConnectContext().getSessionVariable().enablePartitionPredicateRewrite) {
+            Expression toIn = OrToIn.INSTANCE.rewrite(eliminated, null);
+            return InToMinMax.rewrite(toIn, null);
+        } else {
+            return eliminated;
+        }
+    }
+
+    static class InToMinMax extends DefaultExpressionRewriter<Context> {
+        public static Expression rewrite(Expression expression, Context ctx) {
+            return expression.accept(new InToMinMax(), new Context());
+        }
+
+        @Override
+        public Expression visitInPredicate(InPredicate inPredicate, Context ctx) {
+            if (inPredicate.getOptions().size() > 10) {
+                Expression opt0 = inPredicate.getOptions().get(0);
+                if (opt0 instanceof IntegerLikeLiteral
+                        || opt0 instanceof DoubleLiteral
+                        || opt0 instanceof FloatLiteral
+                        || opt0 instanceof DateLiteral
+                        || opt0 instanceof DateTimeLiteral) {
+                    Literal minOpt = (Literal) inPredicate.getOptions().get(0);
+                    Double minVal = minOpt.getDouble();
+                    Literal maxOpt = (Literal) inPredicate.getOptions().get(0);
+                    Double maxVal = maxOpt.getDouble();
+                    for (int i = 1; i < inPredicate.getOptions().size(); i++) {
+                        Expression opt = inPredicate.getOptions().get(i);
+                        if (!(opt instanceof Literal)) {
+                            return inPredicate;
+                        }
+                        double optValue = ((Literal) opt).getDouble();
+                        if (optValue < minVal) {
+                            minVal = optValue;
+                            minOpt = (Literal) opt;
+                        } else if (optValue > maxVal) {
+                            maxVal = optValue;
+                            maxOpt = (Literal) opt;
+                        }
+                    }
+                    return new And(
+                            new GreaterThanEqual(inPredicate.getCompareExpr(), minOpt),
+                            new LessThanEqual(inPredicate.getCompareExpr(), maxOpt));
+                }
+            }
+            return inPredicate;
+        }
+
     }
 
     @Override

diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java
@@ -1127,6 +1127,9 @@ public void setMaxJoinNumberOfReorder(int maxJoinNumberOfReorder) {
     @VariableMgr.VarAttr(name = ENABLE_UNIQUE_KEY_PARTIAL_UPDATE, needForward = true)
     public boolean enableUniqueKeyPartialUpdate = false;
 
+    @VariableMgr.VarAttr(name = "enable_partition_predicate_rewrite", needForward = true)
+    public boolean enablePartitionPredicateRewrite = false;
+
     // If this fe is in fuzzy mode, then will use initFuzzyModeVariables to generate some variables,
     // not the default value set in the code.
     public void initFuzzyModeVariables() {