Skip to content

Commit be59f2f

Browse files
aiguoferaiguofer
committed
Improve TZ support for JDBC driver
1 parent 3ef5450 commit be59f2f

File tree

4 files changed

+130
-15
lines changed

4 files changed

+130
-15
lines changed

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

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,18 +24,23 @@
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;
33+
3134
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
3235
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
3336
import org.apache.arrow.driver.jdbc.utils.DateTimeUtils;
3437
import org.apache.arrow.vector.DateDayVector;
3538
import org.apache.arrow.vector.DateMilliVector;
3639
import org.apache.arrow.vector.ValueVector;
3740

38-
/** Accessor for the Arrow types: {@link DateDayVector} and {@link DateMilliVector}. */
41+
/**
42+
* Accessor for the Arrow types: {@link DateDayVector} and {@link DateMilliVector}.
43+
*/
3944
public class ArrowFlightJdbcDateVectorAccessor extends ArrowFlightJdbcAccessor {
4045

4146
private final Getter getter;
@@ -45,9 +50,9 @@ public class ArrowFlightJdbcDateVectorAccessor extends ArrowFlightJdbcAccessor {
4550
/**
4651
* Instantiate an accessor for a {@link DateDayVector}.
4752
*
48-
* @param vector an instance of a DateDayVector.
53+
* @param vector an instance of a DateDayVector.
4954
* @param currentRowSupplier the supplier to track the lines.
50-
* @param setCursorWasNull the consumer to set if value was null.
55+
* @param setCursorWasNull the consumer to set if value was null.
5156
*/
5257
public ArrowFlightJdbcDateVectorAccessor(
5358
DateDayVector vector,
@@ -62,7 +67,7 @@ public ArrowFlightJdbcDateVectorAccessor(
6267
/**
6368
* Instantiate an accessor for a {@link DateMilliVector}.
6469
*
65-
* @param vector an instance of a DateMilliVector.
70+
* @param vector an instance of a DateMilliVector.
6671
* @param currentRowSupplier the supplier to track the lines.
6772
*/
6873
public ArrowFlightJdbcDateVectorAccessor(
@@ -85,6 +90,19 @@ public Object getObject() {
8590
return this.getDate(null);
8691
}
8792

93+
@Override
94+
public <T> T getObject(final Class<T> type) throws SQLException {
95+
final Object value;
96+
if (type == LocalDate.class) {
97+
value = getLocalDate();
98+
} else if (type == Date.class) {
99+
value = getObject();
100+
} else {
101+
throw new SQLException("invalid class");
102+
}
103+
return !type.isPrimitive() && wasNull ? null : type.cast(value);
104+
}
105+
88106
@Override
89107
public Date getDate(Calendar calendar) {
90108
fillHolder();
@@ -134,4 +152,8 @@ protected static TimeUnit getTimeUnitForVector(ValueVector vector) {
134152

135153
throw new IllegalArgumentException("Invalid Arrow vector");
136154
}
155+
156+
private LocalDate getLocalDate() {
157+
return getDate(null).toLocalDate();
158+
}
137159
}

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

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,35 +21,48 @@
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;
2934
import java.util.TimeZone;
3035
import java.util.concurrent.TimeUnit;
3136
import java.util.function.IntSupplier;
37+
3238
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
3339
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
3440
import org.apache.arrow.vector.TimeStampVector;
3541
import org.apache.arrow.vector.types.pojo.ArrowType;
3642
import org.apache.arrow.vector.util.DateUtility;
3743

38-
/** Accessor for the Arrow types extending from {@link TimeStampVector}. */
44+
/**
45+
* Accessor for the Arrow types extending from {@link TimeStampVector}.
46+
*/
3947
public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAccessor {
4048

4149
private final TimeZone timeZone;
4250
private final Getter getter;
4351
private final TimeUnit timeUnit;
4452
private final LongToLocalDateTime longToLocalDateTime;
4553
private final Holder holder;
54+
private final boolean isZoned;
4655

47-
/** Functional interface used to convert a number (in any time resolution) to LocalDateTime. */
56+
/**
57+
* Functional interface used to convert a number (in any time resolution) to LocalDateTime.
58+
*/
4859
interface LongToLocalDateTime {
4960
LocalDateTime fromLong(long value);
5061
}
5162

52-
/** Instantiate a ArrowFlightJdbcTimeStampVectorAccessor for given vector. */
63+
/**
64+
* Instantiate a ArrowFlightJdbcTimeStampVectorAccessor for given vector.
65+
*/
5366
public ArrowFlightJdbcTimeStampVectorAccessor(
5467
TimeStampVector vector,
5568
IntSupplier currentRowSupplier,
@@ -58,6 +71,7 @@ public ArrowFlightJdbcTimeStampVectorAccessor(
5871
this.holder = new Holder();
5972
this.getter = createGetter(vector);
6073

74+
this.isZoned = getVectorIsZoned(vector);
6175
this.timeZone = getTimeZoneForVector(vector);
6276
this.timeUnit = getTimeUnitForVector(vector);
6377
this.longToLocalDateTime = getLongToLocalDateTimeForVector(vector, this.timeZone);
@@ -68,11 +82,57 @@ public Class<?> getObjectClass() {
6882
return Timestamp.class;
6983
}
7084

85+
@Override
86+
public <T> T getObject(final Class<T> type) throws SQLException {
87+
final Object value;
88+
if (type == OffsetDateTime.class) {
89+
value = getOffsetDateTime();
90+
} else if (type == LocalDateTime.class) {
91+
value = getLocalDateTime(null);
92+
} else if (type == ZonedDateTime.class) {
93+
value = getZonedDateTime();
94+
} else if (type == Instant.class) {
95+
value = getInstant();
96+
} else if (type == Timestamp.class) {
97+
value = getObject();
98+
} else {
99+
throw new SQLException("invalid class");
100+
}
101+
return !type.isPrimitive() && wasNull ? null : type.cast(value);
102+
}
103+
71104
@Override
72105
public Object getObject() {
73106
return this.getTimestamp(null);
74107
}
75108

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

86146
LocalDateTime localDateTime = this.longToLocalDateTime.fromLong(value);
87147

88-
if (calendar != null) {
148+
// Adjust timestamp to desired calendar (if provided) only if the column includes TZ info, otherwise treat as wall-clock time
149+
if (calendar != null && this.isZoned) {
89150
TimeZone timeZone = calendar.getTimeZone();
90151
long millis = this.timeUnit.toMillis(value);
91152
localDateTime =
@@ -177,4 +238,11 @@ protected static TimeZone getTimeZoneForVector(TimeStampVector vector) {
177238

178239
return TimeZone.getTimeZone(timezoneName);
179240
}
241+
242+
protected static boolean getVectorIsZoned(TimeStampVector vector) {
243+
ArrowType.Timestamp arrowType =
244+
(ArrowType.Timestamp) vector.getField().getFieldType().getType();
245+
246+
return arrowType.getTimezone() != null;
247+
}
180248
}

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

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,14 @@
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;
30+
2831
import org.apache.arrow.driver.jdbc.ArrowFlightJdbcTime;
2932
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
3033
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
@@ -48,9 +51,9 @@ public class ArrowFlightJdbcTimeVectorAccessor extends ArrowFlightJdbcAccessor {
4851
/**
4952
* Instantiate an accessor for a {@link TimeNanoVector}.
5053
*
51-
* @param vector an instance of a TimeNanoVector.
54+
* @param vector an instance of a TimeNanoVector.
5255
* @param currentRowSupplier the supplier to track the lines.
53-
* @param setCursorWasNull the consumer to set if value was null.
56+
* @param setCursorWasNull the consumer to set if value was null.
5457
*/
5558
public ArrowFlightJdbcTimeVectorAccessor(
5659
TimeNanoVector vector,
@@ -65,9 +68,9 @@ public ArrowFlightJdbcTimeVectorAccessor(
6568
/**
6669
* Instantiate an accessor for a {@link TimeMicroVector}.
6770
*
68-
* @param vector an instance of a TimeMicroVector.
71+
* @param vector an instance of a TimeMicroVector.
6972
* @param currentRowSupplier the supplier to track the lines.
70-
* @param setCursorWasNull the consumer to set if value was null.
73+
* @param setCursorWasNull the consumer to set if value was null.
7174
*/
7275
public ArrowFlightJdbcTimeVectorAccessor(
7376
TimeMicroVector vector,
@@ -82,7 +85,7 @@ public ArrowFlightJdbcTimeVectorAccessor(
8285
/**
8386
* Instantiate an accessor for a {@link TimeMilliVector}.
8487
*
85-
* @param vector an instance of a TimeMilliVector.
88+
* @param vector an instance of a TimeMilliVector.
8689
* @param currentRowSupplier the supplier to track the lines.
8790
*/
8891
public ArrowFlightJdbcTimeVectorAccessor(
@@ -98,7 +101,7 @@ public ArrowFlightJdbcTimeVectorAccessor(
98101
/**
99102
* Instantiate an accessor for a {@link TimeSecVector}.
100103
*
101-
* @param vector an instance of a TimeSecVector.
104+
* @param vector an instance of a TimeSecVector.
102105
* @param currentRowSupplier the supplier to track the lines.
103106
*/
104107
public ArrowFlightJdbcTimeVectorAccessor(
@@ -121,6 +124,19 @@ public Object getObject() {
121124
return this.getTime(null);
122125
}
123126

127+
@Override
128+
public <T> T getObject(final Class<T> type) throws SQLException {
129+
final Object value;
130+
if (type == LocalTime.class) {
131+
value = getLocalTime();
132+
} else if (type == Time.class) {
133+
value = getObject();
134+
} else {
135+
throw new SQLException("invalid class");
136+
}
137+
return !type.isPrimitive() && wasNull ? null : type.cast(value);
138+
}
139+
124140
@Override
125141
public Time getTime(Calendar calendar) {
126142
fillHolder();
@@ -134,6 +150,10 @@ public Time getTime(Calendar calendar) {
134150
return new ArrowFlightJdbcTime(DateTimeUtils.applyCalendarOffset(milliseconds, calendar));
135151
}
136152

153+
private LocalTime getLocalTime() {
154+
return getTime(null).toLocalTime();
155+
}
156+
137157
private void fillHolder() {
138158
getter.get(getCurrentRow(), holder);
139159
this.wasNull = holder.isSet == 0;

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,12 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
120120
case Time:
121121
return Types.TIME;
122122
case Timestamp:
123-
return Types.TIMESTAMP;
123+
String tz = ((ArrowType.Timestamp) arrowType).getTimezone();
124+
if (tz != null){
125+
return Types.TIMESTAMP_WITH_TIMEZONE;
126+
} else {
127+
return Types.TIMESTAMP;
128+
}
124129
case Bool:
125130
return Types.BOOLEAN;
126131
case Decimal:

0 commit comments

Comments
 (0)