[Feat](nereids) support date function in partition prune (#38743) (#3…

…8898) cherry-pick #38743 to branch-2.1
apache · Aug 6, 2024 · ab3057b · ab3057b
1 parent bcea541
commit ab3057b
Show file tree

Hide file tree

Showing 3 changed files with 335 additions and 32 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
@@ -47,9 +47,9 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.DateTrunc;
 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;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.BooleanType;
-import org.apache.doris.nereids.types.DataType;
 import org.apache.doris.nereids.util.ExpressionUtils;
 import org.apache.doris.nereids.util.Utils;
 
@@ -650,36 +650,11 @@ public EvaluateRangeResult visitDate(Date date, EvaluateRangeInput context) {
         if (!(result.result instanceof Date)) {
             return result;
         }
-        date = (Date) result.result;
-        if (!(date.child() instanceof Slot) || !isPartitionSlot((Slot) date.child())) {
-            return result;
-        }
-        Slot partitionSlot = (Slot) date.child();
-        PartitionSlotType partitionSlotType = getPartitionSlotType(partitionSlot).get();
-        if (partitionSlotType != PartitionSlotType.RANGE || partitionSlotContainsNull.get(partitionSlot)) {
-            return result;
-        }
-        DataType childType = date.child().getDataType();
-        if (!childType.isDateTimeType() && !childType.isDateTimeV2Type()) {
-            return result;
-        }
-        ColumnRange dateTimeRange = result.childrenResult.get(0).columnRanges.get((Slot) date.child());
-        if (dateTimeRange.isEmptyRange()) {
-            return result;
+        Expression dateChild = date.child(0);
+        if (partitionSlotContainsNull.containsKey(dateChild)) {
+            partitionSlotContainsNull.put(date, true);
         }
-
-        Range<ColumnBound> span = dateTimeRange.span();
-        Literal lower = span.lowerEndpoint().getValue();
-        Literal upper = span.upperEndpoint().getValue();
-
-        Expression lowerDate = FoldConstantRuleOnFE.evaluate(new Date(lower), expressionRewriteContext);
-        Expression upperDate = FoldConstantRuleOnFE.evaluate(new Date(upper), expressionRewriteContext);
-
-        if (lowerDate instanceof Literal && upperDate instanceof Literal && lowerDate.equals(upperDate)) {
-            return new EvaluateRangeResult(lowerDate, result.columnRanges, result.childrenResult);
-        }
-
-        return result;
+        return computeMonotonicFunctionRange(result);
     }
 
     private boolean isPartitionSlot(Slot slot) {
@@ -875,7 +850,8 @@ private EvaluateRangeResult computeMonotonicFunctionRange(EvaluateRangeResult re
         }
         Range<ColumnBound> span = childRange.span();
         Literal lower = span.hasLowerBound() ? span.lowerEndpoint().getValue() : null;
-        Literal upper = span.hasUpperBound() ? span.upperEndpoint().getValue() : null;
+        Literal upper = span.hasUpperBound() && !(span.upperEndpoint().getValue() instanceof MaxLiteral)
+                ? span.upperEndpoint().getValue() : null;
         Expression lowerValue = lower != null ? FoldConstantRuleOnFE.evaluate(func.withConstantArgs(lower),
                 expressionRewriteContext) : null;
         Expression upperValue = upper != null ? FoldConstantRuleOnFE.evaluate(func.withConstantArgs(upper),

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
@@ -21,6 +21,8 @@
 import org.apache.doris.nereids.trees.expressions.Expression;
 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.shape.UnaryExpression;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.DateTimeType;
@@ -37,7 +39,7 @@
  * ScalarFunction 'date'. This class is generated by GenerateFunction.
  */
 public class Date extends ScalarFunction
-        implements UnaryExpression, ExplicitlyCastableSignature, AlwaysNullable {
+        implements UnaryExpression, ExplicitlyCastableSignature, AlwaysNullable, Monotonic {
 
     public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
             FunctionSignature.ret(DateV2Type.INSTANCE).args(DateTimeV2Type.SYSTEM_DEFAULT),
@@ -69,4 +71,19 @@ public List<FunctionSignature> getSignatures() {
     public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
         return visitor.visitDate(this, context);
     }
+
+    @Override
+    public boolean isPositive() {
+        return true;
+    }
+
+    @Override
+    public int getMonotonicFunctionChildIndex() {
+        return 0;
+    }
+
+    @Override
+    public Expression withConstantArgs(Literal literal) {
+        return new Date(literal);
+    }
 }