diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java index ebe4016209..cdafeffc32 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java +++ b/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 getObject(final Class 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(); + } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java index debdd0fcb4..813fbc7cfd 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java +++ b/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 getObject(final Class 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; + } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java index 2c03ee631e..d525c2fdd2 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java +++ b/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 getObject(final Class 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; diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java index 96cb056db2..1b76ca0c95 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java +++ b/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: diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java index 88a172e4f2..70d3bcbd33 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java +++ b/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 expectedGetColumnsDataTypes = Arrays.asList(3, 93, 4); - List expectedGetColumnsTypeName = Arrays.asList("DECIMAL", "TIMESTAMP", "INTEGER"); + List expectedGetColumnsDataTypes = Arrays.asList(3, 2014, 4); + List expectedGetColumnsTypeName = + Arrays.asList("DECIMAL", "TIMESTAMP_WITH_TIMEZONE", "INTEGER"); List expectedGetColumnsRadix = Arrays.asList(10, null, 10); List expectedGetColumnsColumnSize = Arrays.asList(5, 29, 10); List expectedGetColumnsDecimalDigits = Arrays.asList(2, 9, 0); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java index 2e329f148e..e4863bd80e 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java @@ -16,17 +16,23 @@ */ package org.apache.arrow.driver.jdbc.accessor.impl.calendar; -import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor.getTimeUnitForVector; -import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor.getTimeZoneForVector; +import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor.*; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; 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.ZoneId; +import java.time.ZonedDateTime; import java.util.Calendar; +import java.util.Objects; import java.util.TimeZone; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; @@ -199,6 +205,99 @@ public void testShouldGetTimestampReturnValidTimestampWithCalendar( }); } + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidLocalDateTime( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + final LocalDateTime value = accessor.getObject(LocalDateTime.class); + final LocalDateTime expectedValue = + getZonedDateTime(currentRow, expectedTimeZone).toLocalDateTime(); + + assertThat(value, equalTo(expectedValue)); + assertThat(accessor.wasNull(), is(false)); + }); + } + + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidInstant( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + final boolean vectorHasTz = timeZone != null; + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + if (vectorHasTz) { + final Instant value = accessor.getObject(Instant.class); + final Instant expectedValue = + getZonedDateTime(currentRow, expectedTimeZone).toInstant(); + + assertThat(value, equalTo(expectedValue)); + assertThat(accessor.wasNull(), is(false)); + } else { + assertThrows(SQLException.class, () -> accessor.getObject(Instant.class)); + } + }); + } + + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidOffsetDateTime( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + final boolean vectorHasTz = timeZone != null; + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + if (vectorHasTz) { + final OffsetDateTime value = accessor.getObject(OffsetDateTime.class); + final OffsetDateTime expectedValue = + getZonedDateTime(currentRow, expectedTimeZone).toOffsetDateTime(); + + assertThat(value, equalTo(expectedValue)); + assertThat(value.getOffset(), equalTo(expectedValue.getOffset())); + assertThat(accessor.wasNull(), is(false)); + } else { + assertThrows(SQLException.class, () -> accessor.getObject(OffsetDateTime.class)); + } + }); + } + + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidZonedDateTime( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + final boolean vectorHasTz = timeZone != null; + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + if (vectorHasTz) { + final ZonedDateTime value = accessor.getObject(ZonedDateTime.class); + final ZonedDateTime expectedValue = getZonedDateTime(currentRow, expectedTimeZone); + + assertThat(value, equalTo(expectedValue)); + assertThat(value.getZone(), equalTo(ZoneId.of(expectedTimeZone))); + assertThat(accessor.wasNull(), is(false)); + } else { + assertThrows(SQLException.class, () -> accessor.getObject(ZonedDateTime.class)); + } + }); + } + @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnNull(Supplier vectorSupplier) { @@ -320,6 +419,21 @@ private Timestamp getTimestampForVector(int currentRow, String timeZone) { return expectedTimestamp; } + /** ZonedDateTime contains all necessary information to generate any java.time object. */ + private ZonedDateTime getZonedDateTime(int currentRow, String timeZone) { + Object object = vector.getObject(currentRow); + TimeZone tz = TimeZone.getTimeZone(timeZone); + ZonedDateTime expectedTimestamp = null; + if (object instanceof LocalDateTime) { + expectedTimestamp = ((LocalDateTime) object).atZone(tz.toZoneId()); + } else if (object instanceof Long) { + TimeUnit timeUnit = getTimeUnitForVector(vector); + Instant instant = Instant.ofEpochMilli(timeUnit.toMillis((Long) object)); + expectedTimestamp = ZonedDateTime.ofInstant(instant, tz.toZoneId()); + } + return expectedTimestamp; + } + @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClass(Supplier vectorSupplier) throws Exception { diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java index 00af3c96ba..a6dd6b3275 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java @@ -48,9 +48,15 @@ public void testGetSqlTypeIdFromArrowType() { assertEquals(Types.DATE, getSqlTypeIdFromArrowType(new ArrowType.Date(DateUnit.MILLISECOND))); assertEquals( Types.TIME, getSqlTypeIdFromArrowType(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); + assertEquals( + Types.TIMESTAMP, + getSqlTypeIdFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))); assertEquals( Types.TIMESTAMP, getSqlTypeIdFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, ""))); + assertEquals( + Types.TIMESTAMP_WITH_TIMEZONE, + getSqlTypeIdFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))); assertEquals(Types.BOOLEAN, getSqlTypeIdFromArrowType(new ArrowType.Bool())); @@ -95,9 +101,15 @@ public void testGetSqlTypeNameFromArrowType() { assertEquals("DATE", getSqlTypeNameFromArrowType(new ArrowType.Date(DateUnit.MILLISECOND))); assertEquals("TIME", getSqlTypeNameFromArrowType(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); + assertEquals( + "TIMESTAMP", + getSqlTypeNameFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))); assertEquals( "TIMESTAMP", getSqlTypeNameFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, ""))); + assertEquals( + "TIMESTAMP_WITH_TIMEZONE", + getSqlTypeNameFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))); assertEquals("BOOLEAN", getSqlTypeNameFromArrowType(new ArrowType.Bool()));