Skip to content

Commit 6d63367

Browse files
dmitry-chirkov-dremiodmitry-chirkov-dremio
committed
GH-838: [FlightSQL][JDBC] Fix timestamp unit conversion in PreparedStatement parameter binding
1 parent 27382cd commit 6d63367

2 files changed

Lines changed: 177 additions & 2 deletions

File tree

‎flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverter.java‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,31 @@
3333
/** AvaticaParameterConverter for Timestamp Arrow types. */
3434
public class TimestampAvaticaParameterConverter extends BaseAvaticaParameterConverter {
3535

36-
public TimestampAvaticaParameterConverter(ArrowType.Timestamp type) {}
36+
private final ArrowType.Timestamp type;
37+
38+
public TimestampAvaticaParameterConverter(ArrowType.Timestamp type) {
39+
this.type = type;
40+
}
41+
42+
/** Converts an epoch millisecond value from Avatica to the target time unit. */
43+
private long convertFromMillis(long epochMillis) {
44+
switch (type.getUnit()) {
45+
case SECOND:
46+
return epochMillis / 1_000L;
47+
case MILLISECOND:
48+
return epochMillis;
49+
case MICROSECOND:
50+
return epochMillis * 1_000L;
51+
case NANOSECOND:
52+
return epochMillis * 1_000_000L;
53+
default:
54+
throw new UnsupportedOperationException("Unsupported time unit: " + type.getUnit());
55+
}
56+
}
3757

3858
@Override
3959
public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
40-
long value = (long) typedValue.toLocal();
60+
long value = convertFromMillis((long) typedValue.toLocal());
4161
if (vector instanceof TimeStampSecVector) {
4262
((TimeStampSecVector) vector).setSafe(index, value);
4363
return true;

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/TimestampAvaticaParameterConverterTest.java‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.driver.jdbc.converter.impl;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertTrue;
21+
22+
import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension;
23+
import org.apache.arrow.memory.BufferAllocator;
24+
import org.apache.arrow.vector.TimeStampMicroTZVector;
25+
import org.apache.arrow.vector.TimeStampMicroVector;
26+
import org.apache.arrow.vector.TimeStampMilliVector;
27+
import org.apache.arrow.vector.TimeStampNanoTZVector;
28+
import org.apache.arrow.vector.TimeStampNanoVector;
29+
import org.apache.arrow.vector.TimeStampSecTZVector;
30+
import org.apache.arrow.vector.TimeStampSecVector;
31+
import org.apache.arrow.vector.types.TimeUnit;
32+
import org.apache.arrow.vector.types.pojo.ArrowType;
33+
import org.apache.calcite.avatica.ColumnMetaData;
34+
import org.apache.calcite.avatica.remote.TypedValue;
35+
import org.junit.jupiter.api.Test;
36+
import org.junit.jupiter.api.extension.RegisterExtension;
37+
38+
/** Tests for {@link TimestampAvaticaParameterConverter}. */
39+
public class TimestampAvaticaParameterConverterTest {
40+
41+
@RegisterExtension
42+
public static RootAllocatorTestExtension rootAllocatorTestExtension =
43+
new RootAllocatorTestExtension();
44+
45+
// A known epoch millis value: 2024-11-03 12:45:09.869 UTC
46+
private static final long EPOCH_MILLIS = 1730637909869L;
47+
48+
@Test
49+
public void testBindParameterMilliVector() {
50+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
51+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.MILLISECOND, null);
52+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
53+
54+
try (TimeStampMilliVector vector = new TimeStampMilliVector("ts", allocator)) {
55+
vector.allocateNew(1);
56+
TypedValue typedValue =
57+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
58+
assertTrue(converter.bindParameter(vector, typedValue, 0));
59+
assertEquals(EPOCH_MILLIS, vector.get(0));
60+
}
61+
}
62+
63+
@Test
64+
public void testBindParameterMicroVector() {
65+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
66+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.MICROSECOND, null);
67+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
68+
69+
try (TimeStampMicroVector vector = new TimeStampMicroVector("ts", allocator)) {
70+
vector.allocateNew(1);
71+
TypedValue typedValue =
72+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
73+
assertTrue(converter.bindParameter(vector, typedValue, 0));
74+
// Millis should be converted to micros
75+
assertEquals(EPOCH_MILLIS * 1_000L, vector.get(0));
76+
}
77+
}
78+
79+
@Test
80+
public void testBindParameterNanoVector() {
81+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
82+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.NANOSECOND, null);
83+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
84+
85+
try (TimeStampNanoVector vector = new TimeStampNanoVector("ts", allocator)) {
86+
vector.allocateNew(1);
87+
TypedValue typedValue =
88+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
89+
assertTrue(converter.bindParameter(vector, typedValue, 0));
90+
// Millis should be converted to nanos
91+
assertEquals(EPOCH_MILLIS * 1_000_000L, vector.get(0));
92+
}
93+
}
94+
95+
@Test
96+
public void testBindParameterSecVector() {
97+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
98+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.SECOND, null);
99+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
100+
101+
try (TimeStampSecVector vector = new TimeStampSecVector("ts", allocator)) {
102+
vector.allocateNew(1);
103+
TypedValue typedValue =
104+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
105+
assertTrue(converter.bindParameter(vector, typedValue, 0));
106+
// Millis should be converted to seconds
107+
assertEquals(EPOCH_MILLIS / 1_000L, vector.get(0));
108+
}
109+
}
110+
111+
@Test
112+
public void testBindParameterMicroTZVector() {
113+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
114+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC");
115+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
116+
117+
try (TimeStampMicroTZVector vector = new TimeStampMicroTZVector("ts", allocator, "UTC")) {
118+
vector.allocateNew(1);
119+
TypedValue typedValue =
120+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
121+
assertTrue(converter.bindParameter(vector, typedValue, 0));
122+
assertEquals(EPOCH_MILLIS * 1_000L, vector.get(0));
123+
}
124+
}
125+
126+
@Test
127+
public void testBindParameterNanoTZVector() {
128+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
129+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC");
130+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
131+
132+
try (TimeStampNanoTZVector vector = new TimeStampNanoTZVector("ts", allocator, "UTC")) {
133+
vector.allocateNew(1);
134+
TypedValue typedValue =
135+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
136+
assertTrue(converter.bindParameter(vector, typedValue, 0));
137+
assertEquals(EPOCH_MILLIS * 1_000_000L, vector.get(0));
138+
}
139+
}
140+
141+
@Test
142+
public void testBindParameterSecTZVector() {
143+
BufferAllocator allocator = rootAllocatorTestExtension.getRootAllocator();
144+
ArrowType.Timestamp type = new ArrowType.Timestamp(TimeUnit.SECOND, "UTC");
145+
TimestampAvaticaParameterConverter converter = new TimestampAvaticaParameterConverter(type);
146+
147+
try (TimeStampSecTZVector vector = new TimeStampSecTZVector("ts", allocator, "UTC")) {
148+
vector.allocateNew(1);
149+
TypedValue typedValue =
150+
TypedValue.ofLocal(ColumnMetaData.Rep.JAVA_SQL_TIMESTAMP, EPOCH_MILLIS);
151+
assertTrue(converter.bindParameter(vector, typedValue, 0));
152+
assertEquals(EPOCH_MILLIS / 1_000L, vector.get(0));
153+
}
154+
}
155+
}

0 commit comments

Comments
 (0)