[opt](nereids) Use a unified approach to deal with monotonic function…

… in partition prune (#42934) This pr use a unified approach to deal with monotonic function in partition prune. When adding a monotonic function in partition prune, we don't need to add visitXXX in OneRangePartitionEvaluator, only need to implement the Monotonic interface.
apache · Nov 28, 2024 · a4d9aaa · a4d9aaa
1 parent b761e4e
commit a4d9aaa
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 57 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
@@ -43,9 +43,8 @@
 import org.apache.doris.nereids.trees.expressions.Or;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.functions.Monotonic;
-import org.apache.doris.nereids.trees.expressions.functions.scalar.ConvertTz;
-import org.apache.doris.nereids.trees.expressions.functions.scalar.Date;
-import org.apache.doris.nereids.trees.expressions.functions.scalar.DateTrunc;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.NonNullable;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.Nullable;
 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.MaxLiteral;
@@ -91,6 +90,7 @@ public class OneRangePartitionEvaluator<K>
     private final List<Literal> lowers;
     private final List<Literal> uppers;
     private final List<List<Expression>> inputs;
+    // whether the Expression in partition range may be null.
     private final Map<Expression, Boolean> partitionSlotContainsNull;
     private final Map<Slot, PartitionSlotType> slotToType;
 
@@ -456,7 +456,12 @@ private EvaluateRangeResult evaluateChildrenThenThis(Expression expr, EvaluateRa
 
         for (int i = 0; i < children.size(); i++) {
             Expression child = children.get(i);
-            EvaluateRangeResult childResult = child.accept(this, context);
+            EvaluateRangeResult childResult;
+            if (child instanceof Monotonic) {
+                childResult = visitMonotonic(child, context);
+            } else {
+                childResult = child.accept(this, context);
+            }
             if (!childResult.result.equals(child)) {
                 hasNewChildren = true;
             }
@@ -615,45 +620,6 @@ private List<Literal> toMultiNereidsLiterals(PartitionKey partitionKey) {
         return literals;
     }
 
-    @Override
-    public EvaluateRangeResult visitDateTrunc(DateTrunc dateTrunc, EvaluateRangeInput context) {
-        EvaluateRangeResult result = super.visitDateTrunc(dateTrunc, context);
-        if (!(result.result instanceof DateTrunc)) {
-            return result;
-        }
-        Expression dateTruncChild = dateTrunc.child(0);
-        if (partitionSlotContainsNull.containsKey(dateTruncChild)) {
-            partitionSlotContainsNull.put(dateTrunc, true);
-        }
-        return computeMonotonicFunctionRange(result, context.rangeMap);
-    }
-
-    @Override
-    public EvaluateRangeResult visitDate(Date date, EvaluateRangeInput context) {
-        EvaluateRangeResult result = super.visitDate(date, context);
-        if (!(result.result instanceof Date)) {
-            return result;
-        }
-        Expression dateChild = date.child(0);
-        if (partitionSlotContainsNull.containsKey(dateChild)) {
-            partitionSlotContainsNull.put(date, true);
-        }
-        return computeMonotonicFunctionRange(result, context.rangeMap);
-    }
-
-    @Override
-    public EvaluateRangeResult visitConvertTz(ConvertTz convertTz, EvaluateRangeInput context) {
-        EvaluateRangeResult result = super.visitConvertTz(convertTz, context);
-        if (!(result.result instanceof ConvertTz)) {
-            return result;
-        }
-        Expression converTzChild = convertTz.child(0);
-        if (partitionSlotContainsNull.containsKey(converTzChild)) {
-            partitionSlotContainsNull.put(convertTz, true);
-        }
-        return computeMonotonicFunctionRange(result, context.rangeMap);
-    }
-
     private boolean isPartitionSlot(Slot slot) {
         return slotToType.containsKey(slot);
     }
@@ -820,15 +786,27 @@ private List<Map<Slot, PartitionSlotInput>> commonComputeOnePartitionInputs() {
         return onePartitionInputs;
     }
 
-    private EvaluateRangeResult computeMonotonicFunctionRange(EvaluateRangeResult result,
-            Map<Expression, ColumnRange> rangeMap) {
+    public EvaluateRangeResult visitMonotonic(Expression monotonic, EvaluateRangeInput context) {
+        EvaluateRangeResult rangeResult = evaluateChildrenThenThis(monotonic, context);
+        if (!rangeResult.result.getClass().equals(monotonic.getClass())) {
+            return rangeResult;
+        }
+        return computeMonotonicFunctionRange(rangeResult, context);
+    }
+
+    private EvaluateRangeResult computeMonotonicFunctionRange(EvaluateRangeResult result, EvaluateRangeInput context) {
         Monotonic func = (Monotonic) result.result;
-        if (rangeMap.containsKey(func)) {
+        if (context.rangeMap.containsKey(func)) {
             return new EvaluateRangeResult((Expression) func, ImmutableMap.of((Expression) func,
-                    rangeMap.get(func)), result.childrenResult);
+                    context.rangeMap.get(func)), result.childrenResult);
         }
         int childIndex = func.getMonotonicFunctionChildIndex();
         Expression funcChild = func.child(childIndex);
+        boolean isNullable = partitionSlotContainsNull.getOrDefault(funcChild, true);
+        Expression withNullable = func.withConstantArgs(isNullable ? new Nullable(funcChild)
+                : new NonNullable(funcChild));
+        partitionSlotContainsNull.put((Expression) func, withNullable.nullable());
+
         if (!result.childrenResult.get(0).columnRanges.containsKey(funcChild)) {
             return result;
         }
@@ -854,7 +832,7 @@ private EvaluateRangeResult computeMonotonicFunctionRange(EvaluateRangeResult re
         ColumnRange newRange = ColumnRange.all();
         if (lowerValue instanceof Literal && upperValue instanceof Literal && lowerValue.equals(upperValue)) {
             newRange = ColumnRange.singleton((Literal) lowerValue);
-            rangeMap.put((Expression) func, newRange);
+            context.rangeMap.put((Expression) func, newRange);
             newRanges.put((Expression) func, newRange);
             return new EvaluateRangeResult(lowerValue, newRanges, result.childrenResult);
         } else {
@@ -864,7 +842,7 @@ private EvaluateRangeResult computeMonotonicFunctionRange(EvaluateRangeResult re
             if (upperValue instanceof Literal) {
                 newRange = newRange.withUpperBound((Literal) upperValue);
             }
-            rangeMap.put((Expression) func, newRange);
+            context.rangeMap.put((Expression) func, newRange);
             newRanges.put((Expression) func, newRange);
             return new EvaluateRangeResult((Expression) func, newRanges, result.childrenResult);
         }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/Monotonic.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/Monotonic.java
@@ -18,7 +18,6 @@
 package org.apache.doris.nereids.trees.expressions.functions;
 
 import org.apache.doris.nereids.trees.expressions.Expression;
-import org.apache.doris.nereids.trees.expressions.literal.Literal;
 
 /** monotonicity of expressions */
 public interface Monotonic extends ExpressionTrait {
@@ -32,5 +31,5 @@ public interface Monotonic extends ExpressionTrait {
     // return the function with the arguments replaced by literal
     // e.g. date_trunc(dt, 'day'), dt in range ['2020-01-01 10:00:00', '2020-01-03 10:00:00']
     // return date_trunc('2020-01-01 10:00:00', 'day')
-    Expression withConstantArgs(Literal literal);
+    Expression withConstantArgs(Expression literal);
 }
diff --git a/.../src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/ConvertTz.java b/.../src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/ConvertTz.java
@@ -24,7 +24,6 @@
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
 import org.apache.doris.nereids.trees.expressions.functions.Monotonic;
 import org.apache.doris.nereids.trees.expressions.functions.PropagateNullLiteral;
-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.literal.StringLikeLiteral;
 import org.apache.doris.nereids.trees.expressions.shape.TernaryExpression;
@@ -99,7 +98,7 @@ public int getMonotonicFunctionChildIndex() {
     }
 
     @Override
-    public Expression withConstantArgs(Literal literal) {
+    public Expression withConstantArgs(Expression literal) {
         return new ConvertTz(literal, child(1), child(2));
     }
 }
diff --git a/...-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Date.java b/...-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Date.java
@@ -23,7 +23,6 @@
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
 import org.apache.doris.nereids.trees.expressions.functions.Monotonic;
 import org.apache.doris.nereids.trees.expressions.functions.PropagateNullLiteral;
-import org.apache.doris.nereids.trees.expressions.literal.Literal;
 import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.DateTimeType;
@@ -84,7 +83,7 @@ public int getMonotonicFunctionChildIndex() {
     }
 
     @Override
-    public Expression withConstantArgs(Literal literal) {
+    public Expression withConstantArgs(Expression literal) {
         return new Date(literal);
     }
 }
diff --git a/.../src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/DateTrunc.java b/.../src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/DateTrunc.java
@@ -23,7 +23,6 @@
 import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
 import org.apache.doris.nereids.trees.expressions.functions.Monotonic;
-import org.apache.doris.nereids.trees.expressions.literal.Literal;
 import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
 import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
@@ -105,7 +104,7 @@ public int getMonotonicFunctionChildIndex() {
     }
 
     @Override
-    public Expression withConstantArgs(Literal literal) {
+    public Expression withConstantArgs(Expression literal) {
         return new DateTrunc(literal, child(1));
     }
 }