Skip to content

Commit 26630c3

Browse files
aiguoferaiguofer
authored
GH-463: Improve TZ support for JDBC driver (#464)
This PR adds support for natively fetching `java.time.*` objects through the JDBC driver. DateVector - getObject(LocalDate.class) DateTimeVector - getObject(OffsetDateTime.class) - getObject(LocalDateTime.class) - getObject(ZonedDateTime.class) - getObject(Instant.class) TimeVector - getObject(LocalTime.class) This PR also changes the behavior for vectors that include TZ info. These will now return as `TIMESTAMP_WITH_TIMEZONE`. The behavior for different ways to access a TimeStampVector are as follows: | | Vector with TZ | Vector W/O TZ | |---------------------------------|------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------| | getTimestamp() | Get timestamp in the vector TZ | Get timestamp in UTC | | getTimestamp(calendar) | Get timestamp by adjusting from the vector TZ to the desired calendar TZ | Get timestamp by adjusting from UTC to the desired calendar TZ (a bug, IMO) | | getObject(LocalDateTime.class) | Get LocalDateTime by taking the timestamp at the vector TZ and taking the "wall-clock" time at that moment | Treat the epoch value in the vector as the "wall-clock" time | | getObject(Instant.class) | Get Instant represented by the value in the vector TZ | Get Instant represented by the value in UTC | | getObject(OffsetDateTime.class) | Get OffsetDateTime represented by the value in the vector TZ (will account for daylight adjustment) | Get OffsetDateTime represented by the value in UTC (will account for daylight adjustment) | | getObject(ZonedDateTime.class) | Get ZonedDateTime represented by the value in the vector at in its TZ | Get ZonedDateTime represented by the value in the vector at in UTC | Closes #463
1 parent 0d296df commit 26630c3

File tree

7 files changed

+269
-13
lines changed

7 files changed

+269
-13
lines changed

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

+19
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424
import static org.apache.calcite.avatica.util.DateTimeUtils.unixDateToString;
2525

2626
import java.sql.Date;
27+
import java.sql.SQLException;
2728
import java.sql.Timestamp;
29+
import java.time.LocalDate;
2830
import java.util.Calendar;
2931
import java.util.concurrent.TimeUnit;
3032
import java.util.function.IntSupplier;
@@ -85,6 +87,19 @@ public Object getObject() {
8587
return this.getDate(null);
8688
}
8789

90+
@Override
91+
public <T> T getObject(final Class<T> type) throws SQLException {
92+
final Object value;
93+
if (type == LocalDate.class) {
94+
value = getLocalDate();
95+
} else if (type == Date.class) {
96+
value = getObject();
97+
} else {
98+
throw new SQLException("Object type not supported for Date Vector");
99+
}
100+
return !type.isPrimitive() && wasNull ? null : type.cast(value);
101+
}
102+
88103
@Override
89104
public Date getDate(Calendar calendar) {
90105
fillHolder();
@@ -134,4 +149,8 @@ protected static TimeUnit getTimeUnitForVector(ValueVector vector) {
134149

135150
throw new IllegalArgumentException("Invalid Arrow vector");
136151
}
152+
153+
private LocalDate getLocalDate() {
154+
return getDate(null).toLocalDate();
155+
}
137156
}

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

+93-8
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,18 @@
2121
import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorGetter.createGetter;
2222

2323
import java.sql.Date;
24+
import java.sql.SQLException;
2425
import java.sql.Time;
2526
import java.sql.Timestamp;
27+
import java.time.Instant;
2628
import java.time.LocalDateTime;
29+
import java.time.OffsetDateTime;
30+
import java.time.ZoneOffset;
31+
import java.time.ZonedDateTime;
2732
import java.time.temporal.ChronoUnit;
2833
import java.util.Calendar;
34+
import java.util.Objects;
35+
import java.util.Set;
2936
import java.util.TimeZone;
3037
import java.util.concurrent.TimeUnit;
3138
import java.util.function.IntSupplier;
@@ -43,6 +50,7 @@ public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAcces
4350
private final TimeUnit timeUnit;
4451
private final LongToLocalDateTime longToLocalDateTime;
4552
private final Holder holder;
53+
private final boolean isZoned;
4654

4755
/** Functional interface used to convert a number (in any time resolution) to LocalDateTime. */
4856
interface LongToLocalDateTime {
@@ -58,6 +66,9 @@ public ArrowFlightJdbcTimeStampVectorAccessor(
5866
this.holder = new Holder();
5967
this.getter = createGetter(vector);
6068

69+
// whether the vector included TZ info
70+
this.isZoned = getVectorIsZoned(vector);
71+
// non-null, either the vector TZ or default to UTC
6172
this.timeZone = getTimeZoneForVector(vector);
6273
this.timeUnit = getTimeUnitForVector(vector);
6374
this.longToLocalDateTime = getLongToLocalDateTimeForVector(vector, this.timeZone);
@@ -68,11 +79,62 @@ public Class<?> getObjectClass() {
6879
return Timestamp.class;
6980
}
7081

82+
@Override
83+
public <T> T getObject(final Class<T> type) throws SQLException {
84+
final Object value;
85+
if (!this.isZoned
86+
& Set.of(OffsetDateTime.class, ZonedDateTime.class, Instant.class).contains(type)) {
87+
throw new SQLException(
88+
"Vectors without timezones can't be converted to objects with offset/tz info.");
89+
} else if (type == OffsetDateTime.class) {
90+
value = getOffsetDateTime();
91+
} else if (type == LocalDateTime.class) {
92+
value = getLocalDateTime(null);
93+
} else if (type == ZonedDateTime.class) {
94+
value = getZonedDateTime();
95+
} else if (type == Instant.class) {
96+
value = getInstant();
97+
} else if (type == Timestamp.class) {
98+
value = getObject();
99+
} else {
100+
throw new SQLException("Object type not supported for TimeStamp Vector");
101+
}
102+
103+
return !type.isPrimitive() && wasNull ? null : type.cast(value);
104+
}
105+
71106
@Override
72107
public Object getObject() {
73108
return this.getTimestamp(null);
74109
}
75110

111+
private ZonedDateTime getZonedDateTime() {
112+
LocalDateTime localDateTime = getLocalDateTime(null);
113+
if (localDateTime == null) {
114+
return null;
115+
}
116+
117+
return localDateTime.atZone(this.timeZone.toZoneId());
118+
}
119+
120+
private OffsetDateTime getOffsetDateTime() {
121+
LocalDateTime localDateTime = getLocalDateTime(null);
122+
if (localDateTime == null) {
123+
return null;
124+
}
125+
ZoneOffset offset = this.timeZone.toZoneId().getRules().getOffset(localDateTime);
126+
return localDateTime.atOffset(offset);
127+
}
128+
129+
private Instant getInstant() {
130+
LocalDateTime localDateTime = getLocalDateTime(null);
131+
if (localDateTime == null) {
132+
return null;
133+
}
134+
ZoneOffset offset = this.timeZone.toZoneId().getRules().getOffset(localDateTime);
135+
return localDateTime.toInstant(offset);
136+
}
137+
76138
private LocalDateTime getLocalDateTime(Calendar calendar) {
77139
getter.get(getCurrentRow(), holder);
78140
this.wasNull = holder.isSet == 0;
@@ -85,7 +147,9 @@ private LocalDateTime getLocalDateTime(Calendar calendar) {
85147

86148
LocalDateTime localDateTime = this.longToLocalDateTime.fromLong(value);
87149

88-
if (calendar != null) {
150+
// Adjust timestamp to desired calendar (if provided) only if the column includes TZ info,
151+
// otherwise treat as wall-clock time
152+
if (calendar != null && this.isZoned) {
89153
TimeZone timeZone = calendar.getTimeZone();
90154
long millis = this.timeUnit.toMillis(value);
91155
localDateTime =
@@ -102,7 +166,7 @@ public Date getDate(Calendar calendar) {
102166
return null;
103167
}
104168

105-
return new Date(Timestamp.valueOf(localDateTime).getTime());
169+
return new Date(getTimestampWithOffset(calendar, localDateTime).getTime());
106170
}
107171

108172
@Override
@@ -112,7 +176,7 @@ public Time getTime(Calendar calendar) {
112176
return null;
113177
}
114178

115-
return new Time(Timestamp.valueOf(localDateTime).getTime());
179+
return new Time(getTimestampWithOffset(calendar, localDateTime).getTime());
116180
}
117181

118182
@Override
@@ -122,6 +186,24 @@ public Timestamp getTimestamp(Calendar calendar) {
122186
return null;
123187
}
124188

189+
return getTimestampWithOffset(calendar, localDateTime);
190+
}
191+
192+
/**
193+
* Apply offset to LocalDateTime to get a Timestamp with legacy behavior. Previously we applied
194+
* the offset to the LocalDateTime even if the underlying Vector did not have a TZ. In order to
195+
* support java.time.* accessors, we fixed this so we only apply the offset if the underlying
196+
* vector includes TZ info. In order to maintain backward compatibility, we apply the offset if
197+
* needed for getDate, getTime, and getTimestamp.
198+
*/
199+
private Timestamp getTimestampWithOffset(Calendar calendar, LocalDateTime localDateTime) {
200+
if (calendar != null && !isZoned) {
201+
TimeZone timeZone = calendar.getTimeZone();
202+
long millis = Timestamp.valueOf(localDateTime).getTime();
203+
localDateTime =
204+
localDateTime.minus(
205+
timeZone.getOffset(millis) - this.timeZone.getOffset(millis), ChronoUnit.MILLIS);
206+
}
125207
return Timestamp.valueOf(localDateTime);
126208
}
127209

@@ -170,11 +252,14 @@ protected static TimeZone getTimeZoneForVector(TimeStampVector vector) {
170252
ArrowType.Timestamp arrowType =
171253
(ArrowType.Timestamp) vector.getField().getFieldType().getType();
172254

173-
String timezoneName = arrowType.getTimezone();
174-
if (timezoneName == null) {
175-
return TimeZone.getTimeZone("UTC");
176-
}
177-
255+
String timezoneName = Objects.requireNonNullElse(arrowType.getTimezone(), "UTC");
178256
return TimeZone.getTimeZone(timezoneName);
179257
}
258+
259+
protected static boolean getVectorIsZoned(TimeStampVector vector) {
260+
ArrowType.Timestamp arrowType =
261+
(ArrowType.Timestamp) vector.getField().getFieldType().getType();
262+
263+
return arrowType.getTimezone() != null;
264+
}
180265
}

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

+19
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@
2020
import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.Holder;
2121
import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.createGetter;
2222

23+
import java.sql.SQLException;
2324
import java.sql.Time;
2425
import java.sql.Timestamp;
26+
import java.time.LocalTime;
2527
import java.util.Calendar;
2628
import java.util.concurrent.TimeUnit;
2729
import java.util.function.IntSupplier;
@@ -121,6 +123,19 @@ public Object getObject() {
121123
return this.getTime(null);
122124
}
123125

126+
@Override
127+
public <T> T getObject(final Class<T> type) throws SQLException {
128+
final Object value;
129+
if (type == LocalTime.class) {
130+
value = getLocalTime();
131+
} else if (type == Time.class) {
132+
value = getObject();
133+
} else {
134+
throw new SQLException("Object type not supported for Time Vector");
135+
}
136+
return !type.isPrimitive() && wasNull ? null : type.cast(value);
137+
}
138+
124139
@Override
125140
public Time getTime(Calendar calendar) {
126141
fillHolder();
@@ -134,6 +149,10 @@ public Time getTime(Calendar calendar) {
134149
return new ArrowFlightJdbcTime(DateTimeUtils.applyCalendarOffset(milliseconds, calendar));
135150
}
136151

152+
private LocalTime getLocalTime() {
153+
return getTime(null).toLocalTime();
154+
}
155+
137156
private void fillHolder() {
138157
getter.get(getCurrentRow(), holder);
139158
this.wasNull = holder.isSet == 0;

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

+7-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package org.apache.arrow.driver.jdbc.utils;
1818

19+
import com.google.common.base.Strings;
1920
import java.sql.Types;
2021
import java.util.HashMap;
2122
import java.util.Map;
@@ -120,7 +121,12 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
120121
case Time:
121122
return Types.TIME;
122123
case Timestamp:
123-
return Types.TIMESTAMP;
124+
String tz = ((ArrowType.Timestamp) arrowType).getTimezone();
125+
if (Strings.isNullOrEmpty(tz)) {
126+
return Types.TIMESTAMP;
127+
} else {
128+
return Types.TIMESTAMP_WITH_TIMEZONE;
129+
}
124130
case Bool:
125131
return Types.BOOLEAN;
126132
case Decimal:

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -299,8 +299,9 @@ public class ArrowDatabaseMetadataTest {
299299
private static Connection connection;
300300

301301
static {
302-
List<Integer> expectedGetColumnsDataTypes = Arrays.asList(3, 93, 4);
303-
List<String> expectedGetColumnsTypeName = Arrays.asList("DECIMAL", "TIMESTAMP", "INTEGER");
302+
List<Integer> expectedGetColumnsDataTypes = Arrays.asList(3, 2014, 4);
303+
List<String> expectedGetColumnsTypeName =
304+
Arrays.asList("DECIMAL", "TIMESTAMP_WITH_TIMEZONE", "INTEGER");
304305
List<Integer> expectedGetColumnsRadix = Arrays.asList(10, null, 10);
305306
List<Integer> expectedGetColumnsColumnSize = Arrays.asList(5, 29, 10);
306307
List<Integer> expectedGetColumnsDecimalDigits = Arrays.asList(2, 9, 0);

0 commit comments

Comments
 (0)