Allow optional second argument denoting timezone for datetime extract…

…ion functions in the multi-stage engine (#14998)
apache · Feb 5, 2025 · e2575fc · e2575fc
1 parent 8c2e6c1
commit e2575fc
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 26 deletions.
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
@@ -126,18 +126,30 @@ public enum TransformFunctionType {
   DATE_TRUNC("dateTrunc", ReturnTypes.BIGINT, OperandTypes.family(
       List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
           SqlTypeFamily.CHARACTER), i -> i > 1)),
-  YEAR("year"),
-  YEAR_OF_WEEK("yearOfWeek", "yow"),
-  QUARTER("quarter"),
-  MONTH_OF_YEAR("monthOfYear", "month"),
-  WEEK_OF_YEAR("weekOfYear", "week"),
-  DAY_OF_YEAR("dayOfYear", "doy"),
-  DAY_OF_MONTH("dayOfMonth", "day"),
-  DAY_OF_WEEK("dayOfWeek", "dow"),
-  HOUR("hour"),
-  MINUTE("minute"),
-  SECOND("second"),
-  MILLISECOND("millisecond"),
+  YEAR("year", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1)),
+  YEAR_OF_WEEK("yearOfWeek", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1), "yow"),
+  QUARTER("quarter", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1)),
+  MONTH_OF_YEAR("monthOfYear", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1), "month"),
+  WEEK_OF_YEAR("weekOfYear", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1), "week"),
+  DAY_OF_YEAR("dayOfYear", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1), "doy"),
+  DAY_OF_MONTH("dayOfMonth", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1), "day"),
+  DAY_OF_WEEK("dayOfWeek", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1), "dow"),
+  HOUR("hour", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1)),
+  MINUTE("minute", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1)),
+  SECOND("second", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1)),
+  MILLISECOND("millisecond", ReturnTypes.BIGINT_NULLABLE, OperandTypes.family(
+      List.of(SqlTypeFamily.DATETIME, SqlTypeFamily.CHARACTER), i -> i == 1)),
   EXTRACT("extract"),
 
   // Array functions

diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java
@@ -234,21 +234,7 @@ public static PinotOperatorTable instance() {
       SqlStdOperatorTable.TIMESTAMP_ADD,
       SqlStdOperatorTable.TIMESTAMP_DIFF,
       SqlStdOperatorTable.CAST,
-
       SqlStdOperatorTable.EXTRACT,
-      // TODO: The following operators are all rewritten to EXTRACT. Consider removing them because they are all
-      //       supported without rewrite.
-      SqlStdOperatorTable.YEAR,
-      SqlStdOperatorTable.QUARTER,
-      SqlStdOperatorTable.MONTH,
-      SqlStdOperatorTable.WEEK,
-      SqlStdOperatorTable.DAYOFYEAR,
-      SqlStdOperatorTable.DAYOFMONTH,
-      SqlStdOperatorTable.DAYOFWEEK,
-      SqlStdOperatorTable.HOUR,
-      SqlStdOperatorTable.MINUTE,
-      SqlStdOperatorTable.SECOND,
-
       SqlStdOperatorTable.ITEM,
       SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR,
       SqlStdOperatorTable.LISTAGG

diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
@@ -262,6 +262,7 @@ protected Object[][] provideQueries() {
         new Object[]{"SELECT FREQUENT_STRINGS_SKETCH(col1) FROM a"},
         new Object[]{"SELECT FREQUENT_LONGS_SKETCH(col3, 1024) FROM a"},
         new Object[]{"SELECT FREQUENT_LONGS_SKETCH(col3) FROM a"},
+        new Object[]{"SELECT DAY_OF_WEEK(ts_timestamp, 'UTC') FROM a"}
     };
   }