Add with_timezone wrapper (linkedin#169)

wenruimeng · Oct 27, 2021 · e7ebd55 · e7ebd55
1 parent 92820e1
commit e7ebd55
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/coral-trino/src/main/java/com/linkedin/coral/trino/rel2trino/Calcite2TrinoUDFConverter.java b/coral-trino/src/main/java/com/linkedin/coral/trino/rel2trino/Calcite2TrinoUDFConverter.java
@@ -301,7 +301,9 @@ private Optional<RexNode> visitCast(RexCall call) {
 
       if (call.getType().getSqlTypeName() == DECIMAL && leftOperand.getType().getSqlTypeName() == TIMESTAMP) {
         SqlOperator trinoToUnixTime = createUDF("to_unixtime", explicit(DOUBLE));
-        return Optional.of(rexBuilder.makeCast(call.getType(), rexBuilder.makeCall(trinoToUnixTime, leftOperand)));
+        SqlOperator trinoWithTimeZone = createUDF("with_timezone", explicit(TIMESTAMP /* should be WITH TIME ZONE */));
+        return Optional.of(rexBuilder.makeCast(call.getType(), rexBuilder.makeCall(trinoToUnixTime,
+            rexBuilder.makeCall(trinoWithTimeZone, leftOperand, rexBuilder.makeLiteral("UTC")))));
       }
 
       return Optional.empty();

diff --git a/coral-trino/src/test/java/com/linkedin/coral/trino/rel2trino/HiveToTrinoConverterTest.java b/coral-trino/src/test/java/com/linkedin/coral/trino/rel2trino/HiveToTrinoConverterTest.java
@@ -338,7 +338,8 @@ public void testCastTimestampToDecimal() {
     RelNode relNode = hiveToRelConverter
         .convertSql("SELECT CAST(a_timestamp AS DECIMAL(10, 0)) AS d\nFROM test.table_from_utc_timestamp");
     String targetSql =
-        "SELECT CAST(\"to_unixtime\"(\"a_timestamp\") AS DECIMAL(10, 0)) AS \"d\"\nFROM \"test\".\"table_from_utc_timestamp\"";
+        "SELECT CAST(\"to_unixtime\"(\"with_timezone\"(\"a_timestamp\", 'UTC')) AS DECIMAL(10, 0)) AS \"d\"\n"
+            + "FROM \"test\".\"table_from_utc_timestamp\"";
     String expandedSql = relToTrinoConverter.convert(relNode);
     assertEquals(expandedSql, targetSql);
   }
@@ -350,14 +351,16 @@ public void testCastNestedTimestampToDecimal() {
     RelNode relNode = hiveToRelConverter.convertSql(
         "SELECT CAST(CAST(a_date AS TIMESTAMP) AS DECIMAL(10, 0)) AS d\nFROM test.table_from_utc_timestamp");
     String targetSql =
-        "SELECT CAST(\"to_unixtime\"(CAST(\"a_date\" AS TIMESTAMP)) AS DECIMAL(10, 0)) AS \"d\"\nFROM \"test\".\"table_from_utc_timestamp\"";
+        "SELECT CAST(\"to_unixtime\"(\"with_timezone\"(CAST(\"a_date\" AS TIMESTAMP), 'UTC')) AS DECIMAL(10, 0)) AS \"d\"\n"
+            + "FROM \"test\".\"table_from_utc_timestamp\"";
     String expandedSql = relToTrinoConverter.convert(relNode);
     assertEquals(expandedSql, targetSql);
 
     relNode = hiveToRelConverter.convertSql(
         "SELECT CAST(from_utc_timestamp(a_date, 'America/Los_Angeles') AS DECIMAL(10, 0)) AS d\nFROM test.table_from_utc_timestamp");
     targetSql =
-        "SELECT CAST(\"to_unixtime\"(CAST(\"at_timezone\"(\"from_unixtime\"(\"to_unixtime\"(\"with_timezone\"(\"a_date\", 'UTC'))), \"$canonicalize_hive_timezone_id\"('America/Los_Angeles')) AS TIMESTAMP(3))) AS DECIMAL(10, 0)) AS \"d\"\nFROM \"test\".\"table_from_utc_timestamp\"";
+        "SELECT CAST(\"to_unixtime\"(\"with_timezone\"(CAST(\"at_timezone\"(\"from_unixtime\"(\"to_unixtime\"(\"with_timezone\"(\"a_date\", 'UTC'))), \"$canonicalize_hive_timezone_id\"('America/Los_Angeles')) AS TIMESTAMP(3)), 'UTC')) AS DECIMAL(10, 0)) AS \"d\"\n"
+            + "FROM \"test\".\"table_from_utc_timestamp\"";
     expandedSql = relToTrinoConverter.convert(relNode);
     assertEquals(expandedSql, targetSql);
   }