From 4b34ecefbe1ff6c648e8bf99d6987b92830b133a Mon Sep 17 00:00:00 2001
From: starocean999 <40539150+starocean999@users.noreply.github.com>
Date: Fri, 22 Mar 2024 10:19:44 +0800
Subject: [PATCH] [fix](nereids)str_to_date function's signature for folding
 constant is wrong (#32474)

---
 .../main/java/org/apache/doris/mtmv/MTMVUtil.java | 15 +++++++++++----
 .../executable/DateTimeExtractAndTransform.java   | 11 ++++++++---
 .../test_date_function_prune.groovy               |  5 +++++
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVUtil.java b/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVUtil.java
index 16459fa130..01033a9615 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVUtil.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVUtil.java
@@ -33,6 +33,8 @@ import org.apache.doris.nereids.trees.expressions.functions.executable.DateTimeA
 import org.apache.doris.nereids.trees.expressions.functions.executable.DateTimeArithmetic;
 import org.apache.doris.nereids.trees.expressions.functions.executable.DateTimeExtractAndTransform;
 import org.apache.doris.nereids.trees.expressions.literal.DateTimeLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.DateTimeV2Literal;
+import org.apache.doris.nereids.trees.expressions.literal.DateV2Literal;
 import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral;
 import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
 
@@ -163,12 +165,17 @@ public class MTMVUtil {
         String dateFormat = dateFormatOptional.get();
         Expression strToDate = DateTimeExtractAndTransform
                 .strToDate(new VarcharLiteral(expr.getStringValue()), new VarcharLiteral(dateFormat));
-        if (!(strToDate instanceof DateTimeLiteral)) {
+        if (strToDate instanceof DateTimeV2Literal) {
+            return ((IntegerLiteral) DateTimeExtractAndTransform
+                    .unixTimestamp((DateTimeV2Literal) strToDate)).getValue();
+        } else if (strToDate instanceof DateV2Literal) {
+            return ((IntegerLiteral) DateTimeExtractAndTransform
+                    .unixTimestamp((DateV2Literal) strToDate)).getValue();
+        } else {
             throw new AnalysisException(
-                    String.format("strToDate failed, stringValue: %s, dateFormat: %s", expr.getStringValue(),
-                            dateFormat));
+                    String.format("strToDate failed, stringValue: %s, dateFormat: %s",
+                            expr.getStringValue(), dateFormat));
         }
-        return ((IntegerLiteral) DateTimeExtractAndTransform.unixTimestamp((DateTimeLiteral) strToDate)).getValue();
     }
 
     /**
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/DateTimeExtractAndTransform.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/DateTimeExtractAndTransform.java
index fa1e1ee2a7..2e53f45475 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/DateTimeExtractAndTransform.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/DateTimeExtractAndTransform.java
@@ -593,10 +593,15 @@ public class DateTimeExtractAndTransform {
     /**
      * date transformation function: str_to_date
      */
-    @ExecFunction(name = "str_to_date", argTypes = {"VARCHAR, VARCHAR"}, returnType = "DATETIME")
+    @ExecFunction(name = "str_to_date", argTypes = {"VARCHAR", "VARCHAR"}, returnType = "DATETIMEV2")
     public static Expression strToDate(VarcharLiteral str, VarcharLiteral format) {
-        return DateTimeLiteral.fromJavaDateType(DateUtils.getTime(DateUtils.formatBuilder(format.getValue())
-                        .toFormatter(), str.getValue()));
+        if (org.apache.doris.analysis.DateLiteral.hasTimePart(format.getStringValue())) {
+            return DateTimeV2Literal.fromJavaDateType(DateUtils.getTime(DateUtils.formatBuilder(format.getValue())
+                    .toFormatter(), str.getValue()));
+        } else {
+            return DateV2Literal.fromJavaDateType(DateUtils.getTime(DateUtils.formatBuilder(format.getValue())
+                    .toFormatter(), str.getValue()));
+        }
     }
 
     @ExecFunction(name = "timestamp", argTypes = {"DATETIME"}, returnType = "DATETIME")
diff --git a/regression-test/suites/nereids_rules_p0/partition_prune/test_date_function_prune.groovy b/regression-test/suites/nereids_rules_p0/partition_prune/test_date_function_prune.groovy
index 6f57067231..c126206eba 100644
--- a/regression-test/suites/nereids_rules_p0/partition_prune/test_date_function_prune.groovy
+++ b/regression-test/suites/nereids_rules_p0/partition_prune/test_date_function_prune.groovy
@@ -86,4 +86,9 @@ suite("test_date_function_prune") {
         sql "select * from dp where date(date_time) = null or date(date_time) = '2020-01-01'"
         contains("partitions=1/3 (p1)")
     }
+
+    explain {
+        sql "select * from dp where date_time > str_to_date('2020-01-02','%Y-%m-%d')"
+        contains("partitions=2/3 (p2,p3)")
+    }
 }
\ No newline at end of file