[Enhancement] support convert_tz as scalaroperatorfunction (#34684)

so that we can support partition prune like CONVERT_TZ(TIMESTAMP('2023-08-01'),'America/Los_Angeles','UTC') Signed-off-by: zombee0 <ewang2027@gmail.com>
StarRocks · Nov 10, 2023 · 35d5967 · 35d5967
1 parent 74600ce
commit 35d5967
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java
@@ -86,6 +86,7 @@
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
+import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
@@ -407,6 +408,19 @@ public static ConstantOperator timestamp(ConstantOperator arg) throws AnalysisEx
         return arg;
     }
 
+    @ConstantFunction.List(list = {
+            @ConstantFunction(name = "convert_tz", argTypes = {DATE, VARCHAR, VARCHAR}, returnType = DATETIME),
+            @ConstantFunction(name = "convert_tz", argTypes = {DATETIME, VARCHAR, VARCHAR}, returnType = DATETIME)
+    })
+    public static ConstantOperator convert_tz(ConstantOperator arg, ConstantOperator fromTz, ConstantOperator toTz) {
+        LocalDateTime dt = arg.getDatetime();
+        ZoneId oldZone = ZoneId.of(fromTz.getVarchar());
+
+        ZoneId newZone = ZoneId.of(toTz.getVarchar());
+        LocalDateTime newDateTime = dt.atZone(oldZone).withZoneSameInstant(newZone).toLocalDateTime();
+        return ConstantOperator.createDatetime(newDateTime);
+    }
+
     @ConstantFunction(name = "unix_timestamp", argTypes = {}, returnType = BIGINT)
     public static ConstantOperator unixTimestampNow() {
         return unixTimestamp(now());

diff --git a/...e-core/src/test/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctionsTest.java b/...e-core/src/test/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctionsTest.java
@@ -538,6 +538,21 @@ public void unixTimestamp() {
                 ScalarOperatorFunctions.unixTimestamp(O_DT_20150323_092355).getBigint());
     }
 
+    @Test
+    public void convert_tz() {
+        ConstantOperator olddt = ConstantOperator.createDatetime(LocalDateTime.of(2019, 8, 1, 13, 21, 3));
+        assertEquals("2019-07-31T22:21:03",
+                ScalarOperatorFunctions.convert_tz(olddt,
+                        ConstantOperator.createVarchar("Asia/Shanghai"),
+                        ConstantOperator.createVarchar("America/Los_Angeles")).getDatetime().toString());
+
+        ConstantOperator oldd = ConstantOperator.createDate(LocalDateTime.of(2019, 8, 1, 0, 0, 0));
+        assertEquals("2019-07-31T09:00",
+                ScalarOperatorFunctions.convert_tz(oldd,
+                        ConstantOperator.createVarchar("Asia/Shanghai"),
+                        ConstantOperator.createVarchar("America/Los_Angeles")).getDatetime().toString());
+    }
+
     @Test
     public void fromUnixTime() throws AnalysisException {
         assertEquals("1970-01-01 08:00:10",