Skip to content

Commit c43b2e8

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

File tree

2 files changed

+189
-3
lines changed

2 files changed

+189
-3
lines changed

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

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,45 @@
3030
import org.apache.calcite.avatica.AvaticaParameter;
3131
import org.apache.calcite.avatica.remote.TypedValue;
3232

33-
/** AvaticaParameterConverter for Timestamp Arrow types. */
33+
/**
34+
* AvaticaParameterConverter for Timestamp Arrow types.
35+
*
36+
* <p>Avatica always provides timestamp values as epoch milliseconds via {@link TypedValue#toLocal}.
37+
* This converter scales the value to match the target vector's time unit (seconds, milliseconds,
38+
* microseconds, or nanoseconds).
39+
*/
3440
public class TimestampAvaticaParameterConverter extends BaseAvaticaParameterConverter {
3541

36-
public TimestampAvaticaParameterConverter(ArrowType.Timestamp type) {}
42+
private final ArrowType.Timestamp type;
43+
44+
public TimestampAvaticaParameterConverter(ArrowType.Timestamp type) {
45+
this.type = type;
46+
}
47+
48+
/**
49+
* Converts an epoch millisecond value to the target time unit.
50+
*
51+
* @param epochMillis the timestamp value in epoch milliseconds (from Avatica)
52+
* @return the timestamp value in the target unit
53+
*/
54+
private long convertFromMillis(long epochMillis) {
55+
switch (type.getUnit()) {
56+
case SECOND:
57+
return epochMillis / 1_000L;
58+
case MILLISECOND:
59+
return epochMillis;
60+
case MICROSECOND:
61+
return epochMillis * 1_000L;
62+
case NANOSECOND:
63+
return epochMillis * 1_000_000L;
64+
default:
65+
throw new UnsupportedOperationException("Unsupported time unit: " + type.getUnit());
66+
}
67+
}
3768

3869
@Override
3970
public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
40-
long value = (long) typedValue.toLocal();
71+
long value = convertFromMillis((long) typedValue.toLocal());
4172
if (vector instanceof TimeStampSecVector) {
4273
((TimeStampSecVector) vector).setSafe(index, value);
4374
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)