GH-463: Improve TZ support for JDBC driver

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java

@@ -24,7 +24,9 @@
 import static org.apache.calcite.avatica.util.DateTimeUtils.unixDateToString;
 
 import java.sql.Date;
+import java.sql.SQLException;
 import java.sql.Timestamp;
+import java.time.LocalDate;
 import java.util.Calendar;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntSupplier;
@@ -85,6 +87,19 @@ public Object getObject() {
     return this.getDate(null);
   }
 
+  @Override
+  public <T> T getObject(final Class<T> type) throws SQLException {
+    final Object value;
+    if (type == LocalDate.class) {
+      value = getLocalDate();
+    } else if (type == Date.class) {
+      value = getObject();
+    } else {
+      throw new SQLException("Object type not supported for Date Vector");
+    }
+    return !type.isPrimitive() && wasNull ? null : type.cast(value);
+  }
+
   @Override
   public Date getDate(Calendar calendar) {
     fillHolder();
@@ -134,4 +149,8 @@ protected static TimeUnit getTimeUnitForVector(ValueVector vector) {
 
     throw new IllegalArgumentException("Invalid Arrow vector");
   }
+
+  private LocalDate getLocalDate() {
+    return getDate(null).toLocalDate();
+  }
 }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java

@@ -21,11 +21,18 @@
 import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorGetter.createGetter;
 
 import java.sql.Date;
+import java.sql.SQLException;
 import java.sql.Time;
 import java.sql.Timestamp;
+import java.time.Instant;
 import java.time.LocalDateTime;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
 import java.time.temporal.ChronoUnit;
 import java.util.Calendar;
+import java.util.Objects;
+import java.util.Set;
 import java.util.TimeZone;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntSupplier;
@@ -43,6 +50,7 @@ public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAcces
   private final TimeUnit timeUnit;
   private final LongToLocalDateTime longToLocalDateTime;
   private final Holder holder;
+  private final boolean isZoned;
 
   /** Functional interface used to convert a number (in any time resolution) to LocalDateTime. */
   interface LongToLocalDateTime {
@@ -58,6 +66,9 @@ public ArrowFlightJdbcTimeStampVectorAccessor(
     this.holder = new Holder();
     this.getter = createGetter(vector);
 
+    // whether the vector included TZ info
+    this.isZoned = getVectorIsZoned(vector);
+    // non-null, either the vector TZ or default to UTC
     this.timeZone = getTimeZoneForVector(vector);
     this.timeUnit = getTimeUnitForVector(vector);
     this.longToLocalDateTime = getLongToLocalDateTimeForVector(vector, this.timeZone);
@@ -68,11 +79,62 @@ public Class<?> getObjectClass() {
     return Timestamp.class;
   }
 
+  @Override
+  public <T> T getObject(final Class<T> type) throws SQLException {
+    final Object value;
+    if (!this.isZoned
+        & Set.of(OffsetDateTime.class, ZonedDateTime.class, Instant.class).contains(type)) {
+      throw new SQLException(
+          "Vectors without timezones can't be converted to objects with offset/tz info.");
+    } else if (type == OffsetDateTime.class) {
+      value = getOffsetDateTime();
+    } else if (type == LocalDateTime.class) {
+      value = getLocalDateTime(null);
+    } else if (type == ZonedDateTime.class) {
+      value = getZonedDateTime();
+    } else if (type == Instant.class) {
+      value = getInstant();
+    } else if (type == Timestamp.class) {
+      value = getObject();
+    } else {
+      throw new SQLException("Object type not supported for TimeStamp Vector");
+    }
+
+    return !type.isPrimitive() && wasNull ? null : type.cast(value);
+  }
+
   @Override
   public Object getObject() {
     return this.getTimestamp(null);
   }
 
+  private ZonedDateTime getZonedDateTime() {
+    LocalDateTime localDateTime = getLocalDateTime(null);
+    if (localDateTime == null) {
+      return null;
+    }
+
+    return localDateTime.atZone(this.timeZone.toZoneId());
+  }
+
+  private OffsetDateTime getOffsetDateTime() {
+    LocalDateTime localDateTime = getLocalDateTime(null);
+    if (localDateTime == null) {
+      return null;
+    }
+    ZoneOffset offset = this.timeZone.toZoneId().getRules().getOffset(localDateTime);
+    return localDateTime.atOffset(offset);
+  }
+
+  private Instant getInstant() {
+    LocalDateTime localDateTime = getLocalDateTime(null);
+    if (localDateTime == null) {
+      return null;
+    }
+    ZoneOffset offset = this.timeZone.toZoneId().getRules().getOffset(localDateTime);
+    return localDateTime.toInstant(offset);
+  }
+
   private LocalDateTime getLocalDateTime(Calendar calendar) {
     getter.get(getCurrentRow(), holder);
     this.wasNull = holder.isSet == 0;
@@ -85,7 +147,9 @@ private LocalDateTime getLocalDateTime(Calendar calendar) {
 
     LocalDateTime localDateTime = this.longToLocalDateTime.fromLong(value);
 
-    if (calendar != null) {
+    // Adjust timestamp to desired calendar (if provided) only if the column includes TZ info,
+    // otherwise treat as wall-clock time
+    if (calendar != null && this.isZoned) {
       TimeZone timeZone = calendar.getTimeZone();
       long millis = this.timeUnit.toMillis(value);
       localDateTime =
@@ -102,7 +166,7 @@ public Date getDate(Calendar calendar) {
       return null;
     }
 
-    return new Date(Timestamp.valueOf(localDateTime).getTime());
+    return new Date(getTimestampWithOffset(calendar, localDateTime).getTime());
   }
 
   @Override
@@ -112,7 +176,7 @@ public Time getTime(Calendar calendar) {
       return null;
     }
 
-    return new Time(Timestamp.valueOf(localDateTime).getTime());
+    return new Time(getTimestampWithOffset(calendar, localDateTime).getTime());
   }
 
   @Override
@@ -122,6 +186,24 @@ public Timestamp getTimestamp(Calendar calendar) {
       return null;
     }
 
+    return getTimestampWithOffset(calendar, localDateTime);
+  }
+
+  /**
+   * Apply offset to LocalDateTime to get a Timestamp with legacy behavior. Previously we applied
+   * the offset to the LocalDateTime even if the underlying Vector did not have a TZ. In order to
+   * support java.time.* accessors, we fixed this so we only apply the offset if the underlying
+   * vector includes TZ info. In order to maintain backward compatibility, we apply the offset if
+   * needed for getDate, getTime, and getTimestamp.
+   */
+  private Timestamp getTimestampWithOffset(Calendar calendar, LocalDateTime localDateTime) {
+    if (calendar != null && !isZoned) {
+      TimeZone timeZone = calendar.getTimeZone();
+      long millis = Timestamp.valueOf(localDateTime).getTime();
+      localDateTime =
+          localDateTime.minus(
+              timeZone.getOffset(millis) - this.timeZone.getOffset(millis), ChronoUnit.MILLIS);
+    }
     return Timestamp.valueOf(localDateTime);
   }
 
@@ -170,11 +252,14 @@ protected static TimeZone getTimeZoneForVector(TimeStampVector vector) {
     ArrowType.Timestamp arrowType =
         (ArrowType.Timestamp) vector.getField().getFieldType().getType();
 
-    String timezoneName = arrowType.getTimezone();
-    if (timezoneName == null) {
-      return TimeZone.getTimeZone("UTC");
-    }
-
+    String timezoneName = Objects.requireNonNullElse(arrowType.getTimezone(), "UTC");
     return TimeZone.getTimeZone(timezoneName);
   }
+
+  protected static boolean getVectorIsZoned(TimeStampVector vector) {
+    ArrowType.Timestamp arrowType =
+        (ArrowType.Timestamp) vector.getField().getFieldType().getType();
+
+    return arrowType.getTimezone() != null;
+  }
 }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java

@@ -20,8 +20,10 @@
 import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.Holder;
 import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.createGetter;
 
+import java.sql.SQLException;
 import java.sql.Time;
 import java.sql.Timestamp;
+import java.time.LocalTime;
 import java.util.Calendar;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntSupplier;
@@ -121,6 +123,19 @@ public Object getObject() {
     return this.getTime(null);
   }
 
+  @Override
+  public <T> T getObject(final Class<T> type) throws SQLException {
+    final Object value;
+    if (type == LocalTime.class) {
+      value = getLocalTime();
+    } else if (type == Time.class) {
+      value = getObject();
+    } else {
+      throw new SQLException("Object type not supported for Time Vector");
+    }
+    return !type.isPrimitive() && wasNull ? null : type.cast(value);
+  }
+
   @Override
   public Time getTime(Calendar calendar) {
     fillHolder();
@@ -134,6 +149,10 @@ public Time getTime(Calendar calendar) {
     return new ArrowFlightJdbcTime(DateTimeUtils.applyCalendarOffset(milliseconds, calendar));
   }
 
+  private LocalTime getLocalTime() {
+    return getTime(null).toLocalTime();
+  }
+
   private void fillHolder() {
     getter.get(getCurrentRow(), holder);
     this.wasNull = holder.isSet == 0;

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java

@@ -16,6 +16,7 @@
  */
 package org.apache.arrow.driver.jdbc.utils;
 
+import com.google.common.base.Strings;
 import java.sql.Types;
 import java.util.HashMap;
 import java.util.Map;
@@ -120,7 +121,12 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
       case Time:
         return Types.TIME;
       case Timestamp:
-        return Types.TIMESTAMP;
+        String tz = ((ArrowType.Timestamp) arrowType).getTimezone();
+        if (Strings.isNullOrEmpty(tz)) {
+          return Types.TIMESTAMP;
+        } else {
+          return Types.TIMESTAMP_WITH_TIMEZONE;
+        }
       case Bool:
         return Types.BOOLEAN;
       case Decimal:

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java

@@ -299,8 +299,9 @@ public class ArrowDatabaseMetadataTest {
   private static Connection connection;
 
   static {
-    List<Integer> expectedGetColumnsDataTypes = Arrays.asList(3, 93, 4);
-    List<String> expectedGetColumnsTypeName = Arrays.asList("DECIMAL", "TIMESTAMP", "INTEGER");
+    List<Integer> expectedGetColumnsDataTypes = Arrays.asList(3, 2014, 4);
+    List<String> expectedGetColumnsTypeName =
+        Arrays.asList("DECIMAL", "TIMESTAMP_WITH_TIMEZONE", "INTEGER");
     List<Integer> expectedGetColumnsRadix = Arrays.asList(10, null, 10);
     List<Integer> expectedGetColumnsColumnSize = Arrays.asList(5, 29, 10);
     List<Integer> expectedGetColumnsDecimalDigits = Arrays.asList(2, 9, 0);