Skip to content

Commit 081b8a1

Browse files
aiguoferaiguofer
committed
Ensure local timezone is always respected
1 parent aee8a10 commit 081b8a1

22 files changed

+163
-100
lines changed

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

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.sql.SQLException;
2222
import java.sql.SQLFeatureNotSupportedException;
2323
import java.util.Arrays;
24+
import java.util.Calendar;
2425
import java.util.Map;
2526
import org.apache.arrow.driver.jdbc.accessor.impl.complex.AbstractArrowFlightJdbcListVectorAccessor;
2627
import org.apache.arrow.driver.jdbc.utils.SqlTypes;
@@ -42,6 +43,7 @@ public class ArrowFlightJdbcArray implements Array {
4243
private final FieldVector dataVector;
4344
private final long startOffset;
4445
private final long valuesCount;
46+
private final Calendar localCalendar;
4547

4648
/**
4749
* Instantiate an {@link Array} backed up by given {@link FieldVector}, limited by a start offset
@@ -51,10 +53,12 @@ public class ArrowFlightJdbcArray implements Array {
5153
* @param startOffset offset from FieldVector pointing to this Array's first value.
5254
* @param valuesCount how many items this Array contains.
5355
*/
54-
public ArrowFlightJdbcArray(FieldVector dataVector, long startOffset, long valuesCount) {
56+
public ArrowFlightJdbcArray(
57+
FieldVector dataVector, long startOffset, long valuesCount, Calendar localCalendar) {
5558
this.dataVector = dataVector;
5659
this.startOffset = startOffset;
5760
this.valuesCount = valuesCount;
61+
this.localCalendar = localCalendar;
5862
}
5963

6064
@Override
@@ -125,7 +129,8 @@ public ResultSet getResultSet(Map<String, Class<?>> map) throws SQLException {
125129
throw new SQLFeatureNotSupportedException();
126130
}
127131

128-
return getResultSetNoBoundariesCheck(this.dataVector, this.startOffset, this.valuesCount);
132+
return getResultSetNoBoundariesCheck(
133+
this.dataVector, this.startOffset, this.valuesCount, this.localCalendar);
129134
}
130135

131136
@Override
@@ -134,14 +139,15 @@ public ResultSet getResultSet(long index, int count) throws SQLException {
134139
}
135140

136141
private static ResultSet getResultSetNoBoundariesCheck(
137-
ValueVector dataVector, long start, long count) throws SQLException {
142+
ValueVector dataVector, long start, long count, Calendar localCalendar) throws SQLException {
138143
TransferPair transferPair = dataVector.getTransferPair(dataVector.getAllocator());
139144
transferPair.splitAndTransfer(
140145
LargeMemoryUtil.checkedCastToInt(start), LargeMemoryUtil.checkedCastToInt(count));
141146
FieldVector vectorSlice = (FieldVector) transferPair.getTo();
142147

143148
VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(vectorSlice);
144-
return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(vectorSchemaRoot);
149+
return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(
150+
vectorSchemaRoot, localCalendar);
145151
}
146152

147153
@Override
@@ -153,7 +159,10 @@ public ResultSet getResultSet(long index, int count, Map<String, Class<?>> map)
153159

154160
checkBoundaries(index, count);
155161
return getResultSetNoBoundariesCheck(
156-
this.dataVector, LargeMemoryUtil.checkedCastToInt(this.startOffset + index), count);
162+
this.dataVector,
163+
LargeMemoryUtil.checkedCastToInt(this.startOffset + index),
164+
count,
165+
this.localCalendar);
157166
}
158167

159168
@Override

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,18 +55,19 @@ public List<Accessor> createAccessors(
5555

5656
return IntStream.range(0, fieldVectors.size())
5757
.mapToObj(root::getVector)
58-
.map(this::createAccessor)
58+
.map(v -> this.createAccessor(v, localCalendar))
5959
.collect(Collectors.toCollection(() -> new ArrayList<>(fieldVectors.size())));
6060
}
6161

62-
private Accessor createAccessor(FieldVector vector) {
62+
private Accessor createAccessor(FieldVector vector, Calendar localCalendar) {
6363
return ArrowFlightJdbcAccessorFactory.createAccessor(
6464
vector,
6565
this::getCurrentRow,
6666
(boolean wasNull) -> {
6767
// AbstractCursor creates a boolean array of length 1 to hold the wasNull value
6868
this.wasNull[0] = wasNull;
69-
});
69+
},
70+
localCalendar);
7071
}
7172

7273
/**

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.sql.ResultSetMetaData;
2121
import java.sql.SQLException;
2222
import java.sql.Types;
23+
import java.util.Calendar;
2324
import java.util.HashSet;
2425
import java.util.List;
2526
import java.util.Objects;
@@ -68,10 +69,10 @@ public class ArrowFlightJdbcVectorSchemaRootResultSet extends AvaticaResultSet {
6869
* @return a ResultSet which accesses the given VectorSchemaRoot
6970
*/
7071
public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot(
71-
final VectorSchemaRoot vectorSchemaRoot) throws SQLException {
72+
final VectorSchemaRoot vectorSchemaRoot, Calendar localCalendar) throws SQLException {
7273
// Similar to how org.apache.calcite.avatica.util.ArrayFactoryImpl does
7374

74-
final TimeZone timeZone = TimeZone.getDefault();
75+
final TimeZone timeZone = localCalendar.getTimeZone();
7576
final QueryState state = new QueryState();
7677

7778
final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);

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

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package org.apache.arrow.driver.jdbc.accessor;
1818

19+
import java.util.Calendar;
1920
import java.util.function.IntSupplier;
2021
import org.apache.arrow.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor;
2122
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor;
@@ -87,7 +88,10 @@ public class ArrowFlightJdbcAccessorFactory {
8788
* @return an instance of one of the accessors.
8889
*/
8990
public static ArrowFlightJdbcAccessor createAccessor(
90-
ValueVector vector, IntSupplier getCurrentRow, WasNullConsumer setCursorWasNull) {
91+
ValueVector vector,
92+
IntSupplier getCurrentRow,
93+
WasNullConsumer setCursorWasNull,
94+
Calendar localCalendar) {
9195
if (vector instanceof UInt1Vector) {
9296
return new ArrowFlightJdbcBaseIntVectorAccessor(
9397
(UInt1Vector) vector, getCurrentRow, setCursorWasNull);
@@ -138,7 +142,7 @@ public static ArrowFlightJdbcAccessor createAccessor(
138142
(FixedSizeBinaryVector) vector, getCurrentRow, setCursorWasNull);
139143
} else if (vector instanceof TimeStampVector) {
140144
return new ArrowFlightJdbcTimeStampVectorAccessor(
141-
(TimeStampVector) vector, getCurrentRow, setCursorWasNull);
145+
(TimeStampVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
142146
} else if (vector instanceof TimeNanoVector) {
143147
return new ArrowFlightJdbcTimeVectorAccessor(
144148
(TimeNanoVector) vector, getCurrentRow, setCursorWasNull);
@@ -180,22 +184,22 @@ public static ArrowFlightJdbcAccessor createAccessor(
180184
(StructVector) vector, getCurrentRow, setCursorWasNull);
181185
} else if (vector instanceof MapVector) {
182186
return new ArrowFlightJdbcMapVectorAccessor(
183-
(MapVector) vector, getCurrentRow, setCursorWasNull);
187+
(MapVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
184188
} else if (vector instanceof ListVector) {
185189
return new ArrowFlightJdbcListVectorAccessor(
186-
(ListVector) vector, getCurrentRow, setCursorWasNull);
190+
(ListVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
187191
} else if (vector instanceof LargeListVector) {
188192
return new ArrowFlightJdbcLargeListVectorAccessor(
189-
(LargeListVector) vector, getCurrentRow, setCursorWasNull);
193+
(LargeListVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
190194
} else if (vector instanceof FixedSizeListVector) {
191195
return new ArrowFlightJdbcFixedSizeListVectorAccessor(
192-
(FixedSizeListVector) vector, getCurrentRow, setCursorWasNull);
196+
(FixedSizeListVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
193197
} else if (vector instanceof UnionVector) {
194198
return new ArrowFlightJdbcUnionVectorAccessor(
195-
(UnionVector) vector, getCurrentRow, setCursorWasNull);
199+
(UnionVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
196200
} else if (vector instanceof DenseUnionVector) {
197201
return new ArrowFlightJdbcDenseUnionVectorAccessor(
198-
(DenseUnionVector) vector, getCurrentRow, setCursorWasNull);
202+
(DenseUnionVector) vector, getCurrentRow, setCursorWasNull, localCalendar);
199203
} else if (vector instanceof NullVector || vector == null) {
200204
return new ArrowFlightJdbcNullVectorAccessor(setCursorWasNull);
201205
}

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAcces
5151
private final LongToLocalDateTime longToLocalDateTime;
5252
private final Holder holder;
5353
private final boolean isZoned;
54+
private final Calendar localCalendar;
5455

5556
/** Functional interface used to convert a number (in any time resolution) to LocalDateTime. */
5657
interface LongToLocalDateTime {
@@ -61,9 +62,11 @@ interface LongToLocalDateTime {
6162
public ArrowFlightJdbcTimeStampVectorAccessor(
6263
TimeStampVector vector,
6364
IntSupplier currentRowSupplier,
64-
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
65+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull,
66+
Calendar localCalendar) {
6567
super(currentRowSupplier, setCursorWasNull);
6668
this.holder = new Holder();
69+
this.localCalendar = localCalendar;
6770
this.getter = createGetter(vector);
6871

6972
// whether the vector included TZ info
@@ -89,7 +92,7 @@ public <T> T getObject(final Class<T> type) throws SQLException {
8992
} else if (type == OffsetDateTime.class) {
9093
value = getOffsetDateTime();
9194
} else if (type == LocalDateTime.class) {
92-
value = getLocalDateTime(null);
95+
value = getLocalDateTime(localCalendar);
9396
} else if (type == ZonedDateTime.class) {
9497
value = getZonedDateTime();
9598
} else if (type == Instant.class) {
@@ -105,7 +108,7 @@ public <T> T getObject(final Class<T> type) throws SQLException {
105108

106109
@Override
107110
public Object getObject() {
108-
return this.getTimestamp(null);
111+
return this.getTimestamp(localCalendar);
109112
}
110113

111114
private ZonedDateTime getZonedDateTime() {
@@ -154,7 +157,7 @@ private LocalDateTime getLocalDateTime(Calendar calendar) {
154157
long millis = this.timeUnit.toMillis(value);
155158
localDateTime =
156159
localDateTime.minus(
157-
timeZone.getOffset(millis) - this.timeZone.getOffset(millis), ChronoUnit.MILLIS);
160+
this.timeZone.getOffset(millis) - timeZone.getOffset(millis), ChronoUnit.MILLIS);
158161
}
159162
return localDateTime;
160163
}
@@ -202,7 +205,7 @@ private Timestamp getTimestampWithOffset(Calendar calendar, LocalDateTime localD
202205
long millis = Timestamp.valueOf(localDateTime).getTime();
203206
localDateTime =
204207
localDateTime.minus(
205-
timeZone.getOffset(millis) - this.timeZone.getOffset(millis), ChronoUnit.MILLIS);
208+
this.timeZone.getOffset(millis) - timeZone.getOffset(millis), ChronoUnit.MILLIS);
206209
}
207210
return Timestamp.valueOf(localDateTime);
208211
}

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package org.apache.arrow.driver.jdbc.accessor.impl.complex;
1818

1919
import java.sql.Array;
20+
import java.util.Calendar;
2021
import java.util.List;
2122
import java.util.function.IntSupplier;
2223
import org.apache.arrow.driver.jdbc.ArrowFlightJdbcArray;
@@ -33,10 +34,14 @@
3334
*/
3435
public abstract class AbstractArrowFlightJdbcListVectorAccessor extends ArrowFlightJdbcAccessor {
3536

37+
private final Calendar localCalendar;
38+
3639
protected AbstractArrowFlightJdbcListVectorAccessor(
3740
IntSupplier currentRowSupplier,
38-
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
41+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull,
42+
Calendar localCalendar) {
3943
super(currentRowSupplier, setCursorWasNull);
44+
this.localCalendar = localCalendar;
4045
}
4146

4247
@Override
@@ -67,6 +72,6 @@ public final Array getArray() {
6772
long endOffset = getEndOffset(index);
6873

6974
long valuesCount = endOffset - startOffset;
70-
return new ArrowFlightJdbcArray(dataVector, startOffset, valuesCount);
75+
return new ArrowFlightJdbcArray(dataVector, startOffset, valuesCount, localCalendar);
7176
}
7277
}

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,17 @@ public abstract class AbstractArrowFlightJdbcUnionVectorAccessor extends ArrowFl
5050
*/
5151
private final ArrowFlightJdbcAccessor[] accessors = new ArrowFlightJdbcAccessor[128];
5252

53+
protected final Calendar localCalendar;
54+
5355
private final ArrowFlightJdbcNullVectorAccessor nullAccessor =
5456
new ArrowFlightJdbcNullVectorAccessor((boolean wasNull) -> {});
5557

5658
protected AbstractArrowFlightJdbcUnionVectorAccessor(
5759
IntSupplier currentRowSupplier,
58-
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
60+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull,
61+
Calendar localCalendar) {
5962
super(currentRowSupplier, setCursorWasNull);
63+
this.localCalendar = localCalendar;
6064
}
6165

6266
protected abstract ArrowFlightJdbcAccessor createAccessorForVector(ValueVector vector);

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package org.apache.arrow.driver.jdbc.accessor.impl.complex;
1818

19+
import java.util.Calendar;
1920
import java.util.function.IntSupplier;
2021
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
2122
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
@@ -38,15 +39,19 @@ public class ArrowFlightJdbcDenseUnionVectorAccessor
3839
public ArrowFlightJdbcDenseUnionVectorAccessor(
3940
DenseUnionVector vector,
4041
IntSupplier currentRowSupplier,
41-
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
42-
super(currentRowSupplier, setCursorWasNull);
42+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull,
43+
Calendar localCalendar) {
44+
super(currentRowSupplier, setCursorWasNull, localCalendar);
4345
this.vector = vector;
4446
}
4547

4648
@Override
4749
protected ArrowFlightJdbcAccessor createAccessorForVector(ValueVector vector) {
4850
return ArrowFlightJdbcAccessorFactory.createAccessor(
49-
vector, () -> this.vector.getOffset(this.getCurrentRow()), (boolean wasNull) -> {});
51+
vector,
52+
() -> this.vector.getOffset(this.getCurrentRow()),
53+
(boolean wasNull) -> {},
54+
localCalendar);
5055
}
5156

5257
@Override

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package org.apache.arrow.driver.jdbc.accessor.impl.complex;
1818

19+
import java.util.Calendar;
1920
import java.util.List;
2021
import java.util.function.IntSupplier;
2122
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
@@ -31,8 +32,9 @@ public class ArrowFlightJdbcFixedSizeListVectorAccessor
3132
public ArrowFlightJdbcFixedSizeListVectorAccessor(
3233
FixedSizeListVector vector,
3334
IntSupplier currentRowSupplier,
34-
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
35-
super(currentRowSupplier, setCursorWasNull);
35+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull,
36+
Calendar localCalendar) {
37+
super(currentRowSupplier, setCursorWasNull, localCalendar);
3638
this.vector = vector;
3739
}
3840

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package org.apache.arrow.driver.jdbc.accessor.impl.complex;
1818

19+
import java.util.Calendar;
1920
import java.util.List;
2021
import java.util.function.IntSupplier;
2122
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
@@ -31,8 +32,9 @@ public class ArrowFlightJdbcLargeListVectorAccessor
3132
public ArrowFlightJdbcLargeListVectorAccessor(
3233
LargeListVector vector,
3334
IntSupplier currentRowSupplier,
34-
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
35-
super(currentRowSupplier, setCursorWasNull);
35+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull,
36+
Calendar localCalendar) {
37+
super(currentRowSupplier, setCursorWasNull, localCalendar);
3638
this.vector = vector;
3739
}
3840

0 commit comments

Comments
 (0)