Register dateTimeConverter,timeConvert,dateTrunc to v2 functions (apa…

…che#11097)
s0nskar · Aug 10, 2023 · 23a292d · 23a292d
1 parent 111e396
commit 23a292d
Show file tree

Hide file tree

Showing 5 changed files with 443 additions and 7 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
@@ -35,6 +35,8 @@
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.type.SqlTypeTransforms;
+import org.apache.pinot.spi.data.DateTimeFieldSpec;
+import org.apache.pinot.spi.data.DateTimeFormatSpec;
 
 
 public enum TransformFunctionType {
@@ -98,9 +100,25 @@ public enum TransformFunctionType {
       OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)), "json_extract_key"),
 
   // date time functions
-  TIMECONVERT("timeConvert", "time_convert"),
-  DATETIMECONVERT("dateTimeConvert", "date_time_convert"),
-  DATETRUNC("dateTrunc", "datetrunc"),
+  TIMECONVERT("timeConvert",
+      ReturnTypes.BIGINT_FORCE_NULLABLE,
+      OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER)),
+      "time_convert"),
+
+  DATETIMECONVERT("dateTimeConvert",
+      ReturnTypes.cascade(
+          opBinding -> dateTimeConverterReturnTypeInference(opBinding),
+          SqlTypeTransforms.FORCE_NULLABLE),
+      OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
+          SqlTypeFamily.CHARACTER)), "date_time_convert"),
+
+  DATETRUNC("dateTrunc",
+      ReturnTypes.BIGINT_FORCE_NULLABLE,
+      OperandTypes.family(
+          ImmutableList.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
+              SqlTypeFamily.CHARACTER),
+          ordinal -> ordinal > 1)),
+
   YEAR("year"),
   YEAR_OF_WEEK("yearOfWeek", "yow"),
   QUARTER("quarter"),
@@ -134,7 +152,14 @@ public enum TransformFunctionType {
   CLPDECODE("clpDecode"),
 
   // Regexp functions
-  REGEXP_EXTRACT("regexpExtract"),
+  REGEXP_EXTRACT("regexpExtract", "regexp_extract"),
+  REGEXPREPLACE("regexpReplace",
+      ReturnTypes.VARCHAR_2000_NULLABLE,
+      OperandTypes.family(
+          ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
+              SqlTypeFamily.INTEGER, SqlTypeFamily.INTEGER, SqlTypeFamily.CHARACTER),
+          ordinal -> ordinal > 2),
+      "regexp_replace"),
 
   // Special type for annotation based scalar functions
   SCALAR("scalar"),
@@ -252,6 +277,20 @@ private static RelDataType positionalReturnTypeInferenceFromStringLiteral(SqlOpe
     return opBinding.getTypeFactory().createSqlType(defaultSqlType);
   }
 
+  private static RelDataType dateTimeConverterReturnTypeInference(SqlOperatorBinding opBinding) {
+    int outputFormatPos = 2;
+    if (opBinding.getOperandCount() > outputFormatPos
+        && opBinding.isOperandLiteral(outputFormatPos, false)) {
+      String outputFormatStr = opBinding.getOperandLiteralValue(outputFormatPos, String.class).toUpperCase();
+      DateTimeFormatSpec dateTimeFormatSpec = new DateTimeFormatSpec(outputFormatStr);
+      if ((dateTimeFormatSpec.getTimeFormat() == DateTimeFieldSpec.TimeFormat.EPOCH) || (
+          dateTimeFormatSpec.getTimeFormat() == DateTimeFieldSpec.TimeFormat.TIMESTAMP)) {
+        return opBinding.getTypeFactory().createSqlType(SqlTypeName.BIGINT);
+      }
+    }
+    return opBinding.getTypeFactory().createSqlType(SqlTypeName.VARCHAR);
+  }
+
   private static RelDataType inferTypeFromStringLiteral(String operandTypeStr, RelDataTypeFactory typeFactory) {
     switch (operandTypeStr) {
       case "INT":

diff --git a/...tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java b/...tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
@@ -145,6 +145,7 @@ private void testHardcodedQueriesCommon()
       throws Exception {
     String query;
     String h2Query;
+
     query = "SELECT COUNT(*) FROM mytable WHERE CarrierDelay=15 AND ArrDelay > CarrierDelay LIMIT 1";
     testQuery(query);
     query = "SELECT ArrDelay, CarrierDelay, (ArrDelay - CarrierDelay) AS diff FROM mytable WHERE CarrierDelay=15 AND "
@@ -252,6 +253,33 @@ private void testHardcodedQueriesCommon()
         + "'1:DAYS:SIMPLE_DATE_FORMAT:yyyy-MM-dd''T''HH:mm:ss.SSS''Z''', '1:DAYS') = '2014-09-05T00:00:00.000Z'";
     h2Query = "SELECT DistanceGroup FROM mytable WHERE DaysSinceEpoch = 16318 LIMIT 10000";
     testQuery(query, h2Query);
+
+    // DateTimeConverter
+    query = "SELECT dateTimeConvert(DaysSinceEpoch,'1:DAYS:EPOCH','1:HOURS:EPOCH','1:HOURS'), COUNT(*) FROM mytable "
+        + "GROUP BY dateTimeConvert(DaysSinceEpoch,'1:DAYS:EPOCH','1:HOURS:EPOCH','1:HOURS') "
+        + "ORDER BY COUNT(*), dateTimeConvert(DaysSinceEpoch,'1:DAYS:EPOCH','1:HOURS:EPOCH','1:HOURS') DESC";
+    h2Query = "SELECT DaysSinceEpoch * 24, COUNT(*) FROM mytable "
+        + "GROUP BY DaysSinceEpoch * 24 "
+        + "ORDER BY COUNT(*), DaysSinceEpoch DESC";
+    testQuery(query, h2Query);
+
+    // TimeConvert
+    query = "SELECT timeConvert(DaysSinceEpoch,'DAYS','SECONDS'), COUNT(*) FROM mytable "
+        + "GROUP BY timeConvert(DaysSinceEpoch,'DAYS','SECONDS') "
+        + "ORDER BY COUNT(*), timeConvert(DaysSinceEpoch,'DAYS','SECONDS') DESC";
+    h2Query = "SELECT DaysSinceEpoch * 86400, COUNT(*) FROM mytable "
+        + "GROUP BY DaysSinceEpoch * 86400"
+        + "ORDER BY COUNT(*), DaysSinceEpoch * 86400 DESC";
+    testQuery(query, h2Query);
+
+    // test arithmetic operations on date time columns
+    query = "SELECT sub(DaysSinceEpoch,25), COUNT(*) FROM mytable "
+        + "GROUP BY sub(DaysSinceEpoch,25) "
+        + "ORDER BY COUNT(*),sub(DaysSinceEpoch,25) DESC";
+    h2Query = "SELECT DaysSinceEpoch - 25, COUNT(*) FROM mytable "
+        + "GROUP BY DaysSinceEpoch "
+        + "ORDER BY COUNT(*), DaysSinceEpoch DESC";
+    testQuery(query, h2Query);
   }
 
   private void testHardCodedQueriesV1()