From d5a21de79674d9516fa97cadcd9b64ad3f4923e9 Mon Sep 17 00:00:00 2001
From: mch_ucchi <41606806+sohardforaname@users.noreply.github.com>
Date: Fri, 4 Aug 2023 13:32:48 +0800
Subject: [PATCH] [Enhancement](planner)support fold constant for date_trunc()
 (#22122)

---
 .../org/apache/doris/rewrite/FEFunctions.java | 77 +++++++++++++++++++
 .../test_date_function.groovy                 |  3 +
 2 files changed, 80 insertions(+)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/rewrite/FEFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/rewrite/FEFunctions.java
index a5cf6cb926..bce2dd7dfa 100755
--- a/fe/fe-core/src/main/java/org/apache/doris/rewrite/FEFunctions.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/rewrite/FEFunctions.java
@@ -350,6 +350,83 @@ public class FEFunctions {
         return null;
     }
 
+    @FEFunction(name = "date_trunc", argTypes = {"DATETIME", "VARCHAR"}, returnType = "DATETIME")
+    public static DateLiteral dateTrunc(LiteralExpr date, LiteralExpr truncate) {
+        if (date.getType().isDateLike()) {
+            DateLiteral dateLiteral = ((DateLiteral) date);
+            LocalDateTime localDate = dateTruncHelper(LocalDateTime.of(
+                            (int) dateLiteral.getYear(), (int) dateLiteral.getMonth(), (int) dateLiteral.getDay(),
+                            (int) dateLiteral.getHour(), (int) dateLiteral.getMinute(), (int) dateLiteral.getSecond()),
+                    truncate.getStringValue());
+
+            return new DateLiteral(localDate.getYear(), localDate.getMonthValue(), localDate.getDayOfMonth(),
+                    localDate.getHour(), localDate.getMinute(), localDate.getSecond(), date.getType());
+        }
+        return null;
+    }
+
+    @FEFunction(name = "date_trunc", argTypes = {"DATETIMEV2", "VARCHAR"}, returnType = "DATETIMEV2")
+    public static DateLiteral dateTruncV2(LiteralExpr date, LiteralExpr truncate) {
+        if (date.getType().isDateLike()) {
+            DateLiteral dateLiteral = ((DateLiteral) date);
+            LocalDateTime localDate = dateTruncHelper(LocalDateTime.of(
+                            (int) dateLiteral.getYear(), (int) dateLiteral.getMonth(), (int) dateLiteral.getDay(),
+                            (int) dateLiteral.getHour(), (int) dateLiteral.getMinute(), (int) dateLiteral.getSecond()),
+                    truncate.getStringValue());
+
+            return new DateLiteral(localDate.getYear(), localDate.getMonthValue(), localDate.getDayOfMonth(),
+                    localDate.getHour(), localDate.getMinute(), localDate.getSecond(), date.getType());
+        }
+        return null;
+    }
+
+    private static LocalDateTime dateTruncHelper(LocalDateTime dateTime, String trunc) {
+        int year = dateTime.getYear();
+        int month = dateTime.getMonthValue();
+        int day = dateTime.getDayOfMonth();
+        int hour = dateTime.getHour();
+        int minute = dateTime.getMinute();
+        int second = dateTime.getSecond();
+        switch (trunc.toLowerCase()) {
+            case "year":
+                month = 0;
+            case "quarter": // CHECKSTYLE IGNORE THIS LINE
+                month = ((month - 1) / 3) * 3 + 1;
+            case "month": // CHECKSTYLE IGNORE THIS LINE
+                day = 1;
+                break;
+            case "week":
+                LocalDateTime firstDayOfWeek = firstDayOfWeek(dateTime);
+                year = firstDayOfWeek.getYear();
+                month = firstDayOfWeek.getMonthValue();
+                day = firstDayOfWeek.getDayOfMonth();
+            default: // CHECKSTYLE IGNORE THIS LINE
+                break;
+        }
+        switch (trunc.toLowerCase()) {
+            case "year":
+            case "quarter":
+            case "month":
+            case "week":
+            case "day": // CHECKSTYLE IGNORE THIS LINE
+                hour = 0;
+            case "hour": // CHECKSTYLE IGNORE THIS LINE
+                minute = 0;
+            case "minute": // CHECKSTYLE IGNORE THIS LINE
+                second = 0;
+            default: // CHECKSTYLE IGNORE THIS LINE
+        }
+        return LocalDateTime.of(year, month, day, hour, minute, second);
+    }
+
+    private static int distanceToFirstDayOfWeek(LocalDateTime dateTime) {
+        return dateTime.getDayOfWeek().getValue() - 1;
+    }
+
+    private static LocalDateTime firstDayOfWeek(LocalDateTime dateTime) {
+        return dateTime.plusDays(-distanceToFirstDayOfWeek(dateTime));
+    }
+
     private static LocalDateTime toMonday(LocalDateTime dateTime) {
         LocalDateTime specialUpperBound = LocalDateTime.of(1970, 1, 4, 0, 0, 0);
         LocalDateTime specialLowerBound = LocalDateTime.of(1970, 1, 1, 0, 0, 0);
diff --git a/regression-test/suites/query_p0/sql_functions/datetime_functions/test_date_function.groovy b/regression-test/suites/query_p0/sql_functions/datetime_functions/test_date_function.groovy
index f3d775408e..02bbd80cd5 100644
--- a/regression-test/suites/query_p0/sql_functions/datetime_functions/test_date_function.groovy
+++ b/regression-test/suites/query_p0/sql_functions/datetime_functions/test_date_function.groovy
@@ -714,4 +714,7 @@ suite("test_date_function") {
         sql "select cast('20230631' as date), cast('20230632' as date)"
         result([[null, null]])
     }
+    
+    res = sql "explain select date_trunc('2022-04-24', 'day'), date_trunc('1999-03-12 00:31:23', 'hour')"
+    assertFalse(res.contains("date_trunc"))
 }