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,18 +24,23 @@
 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;
+
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
 import org.apache.arrow.driver.jdbc.utils.DateTimeUtils;
 import org.apache.arrow.vector.DateDayVector;
 import org.apache.arrow.vector.DateMilliVector;
 import org.apache.arrow.vector.ValueVector;
 
-/** Accessor for the Arrow types: {@link DateDayVector} and {@link DateMilliVector}. */
+/**
+ * Accessor for the Arrow types: {@link DateDayVector} and {@link DateMilliVector}.
+ */
 public class ArrowFlightJdbcDateVectorAccessor extends ArrowFlightJdbcAccessor {
 
   private final Getter getter;
@@ -45,9 +50,9 @@ public class ArrowFlightJdbcDateVectorAccessor extends ArrowFlightJdbcAccessor {
   /**
    * Instantiate an accessor for a {@link DateDayVector}.
    *
-   * @param vector an instance of a DateDayVector.
+   * @param vector             an instance of a DateDayVector.
    * @param currentRowSupplier the supplier to track the lines.
-   * @param setCursorWasNull the consumer to set if value was null.
+   * @param setCursorWasNull   the consumer to set if value was null.
    */
   public ArrowFlightJdbcDateVectorAccessor(
       DateDayVector vector,
@@ -62,7 +67,7 @@ public ArrowFlightJdbcDateVectorAccessor(
   /**
    * Instantiate an accessor for a {@link DateMilliVector}.
    *
-   * @param vector an instance of a DateMilliVector.
+   * @param vector             an instance of a DateMilliVector.
    * @param currentRowSupplier the supplier to track the lines.
    */
   public ArrowFlightJdbcDateVectorAccessor(
@@ -85,6 +90,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("invalid class");
+    }
+    return !type.isPrimitive() && wasNull ? null : type.cast(value);
+  }
+
   @Override
   public Date getDate(Calendar calendar) {
     fillHolder();
@@ -134,4 +152,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,35 +21,48 @@
 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.TimeZone;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntSupplier;
+
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
 import org.apache.arrow.vector.TimeStampVector;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.util.DateUtility;
 
-/** Accessor for the Arrow types extending from {@link TimeStampVector}. */
+/**
+ * Accessor for the Arrow types extending from {@link TimeStampVector}.
+ */
 public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAccessor {
 
   private final TimeZone timeZone;
   private final Getter getter;
   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. */
+  /**
+   * Functional interface used to convert a number (in any time resolution) to LocalDateTime.
+   */
   interface LongToLocalDateTime {
     LocalDateTime fromLong(long value);
   }
 
-  /** Instantiate a ArrowFlightJdbcTimeStampVectorAccessor for given vector. */
+  /**
+   * Instantiate a ArrowFlightJdbcTimeStampVectorAccessor for given vector.
+   */
   public ArrowFlightJdbcTimeStampVectorAccessor(
       TimeStampVector vector,
       IntSupplier currentRowSupplier,
@@ -58,6 +71,7 @@ public ArrowFlightJdbcTimeStampVectorAccessor(
     this.holder = new Holder();
     this.getter = createGetter(vector);
 
+    this.isZoned = getVectorIsZoned(vector);
     this.timeZone = getTimeZoneForVector(vector);
     this.timeUnit = getTimeUnitForVector(vector);
     this.longToLocalDateTime = getLongToLocalDateTimeForVector(vector, this.timeZone);
@@ -68,11 +82,57 @@ public Class<?> getObjectClass() {
     return Timestamp.class;
   }
 
+  @Override
+  public <T> T getObject(final Class<T> type) throws SQLException {
+    final Object value;
+    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("invalid class");
+    }
+    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 +145,8 @@ 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 =
@@ -177,4 +238,11 @@ protected static TimeZone getTimeZoneForVector(TimeStampVector vector) {
 
     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,11 +20,14 @@
 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;
+
 import org.apache.arrow.driver.jdbc.ArrowFlightJdbcTime;
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
@@ -48,9 +51,9 @@ public class ArrowFlightJdbcTimeVectorAccessor extends ArrowFlightJdbcAccessor {
   /**
    * Instantiate an accessor for a {@link TimeNanoVector}.
    *
-   * @param vector an instance of a TimeNanoVector.
+   * @param vector             an instance of a TimeNanoVector.
    * @param currentRowSupplier the supplier to track the lines.
-   * @param setCursorWasNull the consumer to set if value was null.
+   * @param setCursorWasNull   the consumer to set if value was null.
    */
   public ArrowFlightJdbcTimeVectorAccessor(
       TimeNanoVector vector,
@@ -65,9 +68,9 @@ public ArrowFlightJdbcTimeVectorAccessor(
   /**
    * Instantiate an accessor for a {@link TimeMicroVector}.
    *
-   * @param vector an instance of a TimeMicroVector.
+   * @param vector             an instance of a TimeMicroVector.
    * @param currentRowSupplier the supplier to track the lines.
-   * @param setCursorWasNull the consumer to set if value was null.
+   * @param setCursorWasNull   the consumer to set if value was null.
    */
   public ArrowFlightJdbcTimeVectorAccessor(
       TimeMicroVector vector,
@@ -82,7 +85,7 @@ public ArrowFlightJdbcTimeVectorAccessor(
   /**
    * Instantiate an accessor for a {@link TimeMilliVector}.
    *
-   * @param vector an instance of a TimeMilliVector.
+   * @param vector             an instance of a TimeMilliVector.
    * @param currentRowSupplier the supplier to track the lines.
    */
   public ArrowFlightJdbcTimeVectorAccessor(
@@ -98,7 +101,7 @@ public ArrowFlightJdbcTimeVectorAccessor(
   /**
    * Instantiate an accessor for a {@link TimeSecVector}.
    *
-   * @param vector an instance of a TimeSecVector.
+   * @param vector             an instance of a TimeSecVector.
    * @param currentRowSupplier the supplier to track the lines.
    */
   public ArrowFlightJdbcTimeVectorAccessor(
@@ -121,6 +124,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("invalid class");
+    }
+    return !type.isPrimitive() && wasNull ? null : type.cast(value);
+  }
+
   @Override
   public Time getTime(Calendar calendar) {
     fillHolder();
@@ -134,6 +150,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

@@ -120,7 +120,12 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
       case Time:
         return Types.TIME;
       case Timestamp:
-        return Types.TIMESTAMP;
+        String tz = ((ArrowType.Timestamp) arrowType).getTimezone();
+        if (tz != null){
+          return Types.TIMESTAMP_WITH_TIMEZONE;
+        } else {
+          return Types.TIMESTAMP;
+        }
       case Bool:
         return Types.BOOLEAN;
       case Decimal: