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);