[CALCITE-7491] Literals of type TIMESTAMP WITH TIME ZONE cause crashes

Signed-off-by: Mihai Budiu <mbudiu@feldera.com>

Author: mihaibudiu

babel/src/test/resources/sql/big-query.iq

@@ -30,6 +30,43 @@
 !use scott-big-query
 !set outputformat mysql
 
+# Test case for [CALCITE-7491] https://issues.apache.org/jira/browse/CALCITE-7491
+# Literals of type TIMESTAMP WITH TIME ZONE cause crashes
+# This will change once we fix [CALCITE-7494]
+# Avatica conversion to string of TIMESTAMP WITH TIME ZONE
+# does not include time zone
+select TIMESTAMP WITH TIME ZONE '2020-01-01 00:00:00 America/New_York';
++---------------------+
+| EXPR$0              |
++---------------------+
+| 2020-01-01 05:00:00 |
++---------------------+
+(1 row)
+
+!ok
+
+# Two timestamps with time zone are equal if they represent the same UTC time
+SELECT TIMESTAMP WITH TIME ZONE '2020-01-01 08:10:10 America/New_York' = TIMESTAMP WITH TIME ZONE '2020-01-01 05:10:10 America/Los_Angeles';
++--------+
+| EXPR$0 |
++--------+
+| true   |
++--------+
+(1 row)
+
+!ok
+
+# Two equal timestamps in different time zones are different if they represent different UTC times
+SELECT TIMESTAMP WITH TIME ZONE '2020-01-01 08:10:10 America/New_York' = TIMESTAMP WITH TIME ZONE '2020-01-01 08:10:10 America/Los_Angeles';
++--------+
+| EXPR$0 |
++--------+
+| false  |
++--------+
+(1 row)
+
+!ok
+
 # Two tests for [CALCITE-7094] Using a type alias as a constructor function
 # causes a validator assertion failure
 select int64();

core/src/main/java/org/apache/calcite/adapter/enumerable/RexToLixTranslator.java

@@ -1077,6 +1077,7 @@ public static Expression translateLiteral(
                   () -> "value for " + literal).toString()));
     case DATE:
     case TIME:
+    case TIME_TZ:
     case TIME_WITH_LOCAL_TIME_ZONE:
     case INTERVAL_YEAR:
     case INTERVAL_YEAR_MONTH:
@@ -1085,6 +1086,7 @@ public static Expression translateLiteral(
       javaClass = int.class;
       break;
     case TIMESTAMP:
+    case TIMESTAMP_TZ:
     case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
     case INTERVAL_DAY:
     case INTERVAL_DAY_HOUR:

core/src/main/java/org/apache/calcite/rex/RexLiteral.java

@@ -60,6 +60,11 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.text.SimpleDateFormat;
+import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
 import java.util.Calendar;
 import java.util.List;
 import java.util.Locale;
@@ -1183,9 +1188,18 @@ public boolean isNull() {
       break;
     case TIMESTAMP_TZ:
       if (clazz == Long.class) {
-        return clazz.cast(((TimestampWithTimeZoneString) value)
+        TimestampWithTimeZoneString tstz = (TimestampWithTimeZoneString) value;
+        long ms = tstz
             .getLocalTimestampString()
-            .getMillisSinceEpoch());
+            .getMillisSinceEpoch();
+        // Interpret the timestamp part as a UTC timestamp
+        LocalDateTime local = Instant.ofEpochMilli(ms).atZone(ZoneOffset.UTC).toLocalDateTime();
+        // Adjust for the time zone
+        ZoneId id = tstz.getTimeZone().toZoneId();
+        ZonedDateTime zoned = local.atZone(id);
+        ZonedDateTime utc = zoned.withZoneSameInstant(ZoneOffset.UTC);
+        ms = utc.toInstant().toEpochMilli();
+        return clazz.cast(ms);
       } else if (clazz == Calendar.class) {
         TimestampWithTimeZoneString ts = (TimestampWithTimeZoneString) value;
         return clazz.cast(ts.getLocalTimestampString().toCalendar(ts.getTimeZone()));

core/src/main/java/org/apache/calcite/sql2rel/SqlNodeToRexConverterImpl.java

@@ -26,6 +26,7 @@
 import org.apache.calcite.sql.SqlLiteral;
 import org.apache.calcite.sql.SqlTimeLiteral;
 import org.apache.calcite.sql.SqlTimestampLiteral;
+import org.apache.calcite.sql.SqlTimestampTzLiteral;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.BitString;
 import org.apache.calcite.util.DateString;
@@ -131,7 +132,7 @@ public class SqlNodeToRexConverterImpl implements SqlNodeToRexConverter {
     case TIMESTAMP_TZ:
       return rexBuilder.makeTimestampTzLiteral(
           literal.getValueAs(TimestampWithTimeZoneString.class),
-          ((SqlTimestampLiteral) literal).getPrec());
+          ((SqlTimestampTzLiteral) literal).getPrec());
     case TIME:
       return rexBuilder.makeTimeLiteral(
           literal.getValueAs(TimeString.class),

core/src/main/java/org/apache/calcite/util/TimestampWithTimeZoneString.java

@@ -173,7 +173,7 @@ public TimestampWithTimeZoneString withTimeZone(TimeZone timeZone) {
   }
 
   @Override public int compareTo(TimestampWithTimeZoneString o) {
-    return v.compareTo(o.v);
+    return this.pt.getCalendar().compareTo(o.pt.getCalendar());
   }
 
   public TimestampWithTimeZoneString round(int precision) {

site/_docs/reference.md

@@ -1682,6 +1682,7 @@ charSet:
 
 timeZone:
       WITHOUT TIME ZONE
+  |   WITH TIME ZONE
   |   WITH LOCAL TIME ZONE
 {% endhighlight %}