Skip to content

Commit a1d8317

Browse files
xborderdependabot[bot]jhrotko
authored
GH-929: Add UUID support in JDBC driver (#930)
## What's Changed This PR adds UUID support to the Arrow Flight SQL JDBC driver, enabling JDBC applications to work with UUID data types when connecting to Flight SQL servers that use Arrow's canonical `arrow.uuid` extension type. ### Key Implementation Details - Added `ArrowFlightJdbcUuidVectorAccessor` to handle reading UUID values from `UuidVector` - `getObject()` returns `java.util.UUID` directly - `getString()` returns the standard hyphenated UUID format (e.g., "550e8400-e29b-41d4-a716-446655440000") - `getBytes()` returns the 16-byte binary representation - Added `UuidAvaticaParameterConverter` to handle parameter binding for UUID columns - Supports binding `java.util.UUID` objects directly via `setObject()` - Supports binding UUID string representations via `setString()` - Supports binding 16-byte arrays via `setBytes()` - UUID extension type maps to `java.sql.Types.OTHER`. - Updated `SqlTypes` to recognize `UuidType` and return appropriate SQL type ID **Examples** ``` java try (Statement stmt = conn.createStatement(); ResultSet rs = stmt.executeQuery("SELECT id, session_id FROM sessions")) { while (rs.next()) { int id = rs.getInt("id"); // getObject() returns java.util.UUID directly UUID sessionId = rs.getObject("session_id", UUID.class); // getString() returns hyphenated format: "550e8400-e29b-41d4-a716-..." String sessionIdStr = rs.getString("session_id"); System.out.printf("ID: %d, UUID: %s%n", id, sessionId); } } // Use PreparedStatement to bind UUID parameters String sql = "SELECT * FROM sessions WHERE session_id = ?"; try (PreparedStatement pstmt = conn.prepareStatement(sql)) { UUID targetId = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); // Bind UUID directly with setObject() pstmt.setObject(1, targetId); // Or bind as string: pstmt.setString(1, targetId.toString()); try (ResultSet rs = pstmt.executeQuery()) { if (rs.next()) { System.out.println("Found: " + rs.getObject("session_id")); } } } ``` Closes #929. --------- Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Joana Hrotko <[email protected]>
1 parent 71c418c commit a1d8317

File tree

15 files changed

+927
-1
lines changed

15 files changed

+927
-1
lines changed

docs/source/jdbc.rst

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,8 @@ Type Mapping
213213
------------
214214

215215
The Arrow to JDBC type mapping can be obtained at runtime via
216-
a method on ColumnBinder.
216+
a method on ColumnBinder. The Flight SQL JDBC driver follows the same
217+
mapping, with additional support for the UUID extension type noted below.
217218

218219
+----------------------------+----------------------------+-------+
219220
| Arrow Type | JDBC Type | Notes |
@@ -232,6 +233,8 @@ a method on ColumnBinder.
232233
+----------------------------+----------------------------+-------+
233234
| FixedSizeBinary | BINARY (setBytes) | |
234235
+----------------------------+----------------------------+-------+
236+
| Uuid (extension) | OTHER (setObject) | \(3) |
237+
+----------------------------+----------------------------+-------+
235238
| Float32 | REAL (setFloat) | |
236239
+----------------------------+----------------------------+-------+
237240
| Int8 | TINYINT (setByte) | |
@@ -276,3 +279,6 @@ a method on ColumnBinder.
276279
<https://docs.oracle.com/en/java/javase/11/docs/api/java.sql/java/sql/PreparedStatement.html#setTimestamp(int,java.sql.Timestamp)>`_,
277280
which will lead to the driver using the "default timezone" (that of
278281
the Java VM).
282+
* \(3) For the Flight SQL JDBC driver, the Arrow UUID extension type
283+
(``arrow.uuid``) maps to JDBC ``OTHER`` and is surfaced as
284+
``java.util.UUID`` values.

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,10 +75,12 @@
7575
import org.apache.arrow.vector.VarBinaryVector;
7676
import org.apache.arrow.vector.VarCharVector;
7777
import org.apache.arrow.vector.VectorSchemaRoot;
78+
import org.apache.arrow.vector.extension.UuidType;
7879
import org.apache.arrow.vector.ipc.ReadChannel;
7980
import org.apache.arrow.vector.ipc.message.MessageSerializer;
8081
import org.apache.arrow.vector.types.Types;
8182
import org.apache.arrow.vector.types.pojo.ArrowType;
83+
import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
8284
import org.apache.arrow.vector.types.pojo.Field;
8385
import org.apache.arrow.vector.types.pojo.Schema;
8486
import org.apache.arrow.vector.util.Text;
@@ -164,6 +166,9 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
164166
LONGNVARCHAR, SqlSupportsConvert.SQL_CONVERT_LONGVARCHAR_VALUE);
165167
sqlTypesToFlightEnumConvertTypes.put(DATE, SqlSupportsConvert.SQL_CONVERT_DATE_VALUE);
166168
sqlTypesToFlightEnumConvertTypes.put(TIMESTAMP, SqlSupportsConvert.SQL_CONVERT_TIMESTAMP_VALUE);
169+
170+
// Register the UUID extension type so it is always available for the driver
171+
ExtensionTypeRegistry.register(UuidType.INSTANCE);
167172
}
168173

169174
ArrowDatabaseMetadata(final AvaticaConnection connection) {

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.function.IntSupplier;
2020
import org.apache.arrow.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor;
2121
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor;
22+
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor;
2223
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor;
2324
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor;
2425
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor;
@@ -65,6 +66,7 @@
6566
import org.apache.arrow.vector.UInt2Vector;
6667
import org.apache.arrow.vector.UInt4Vector;
6768
import org.apache.arrow.vector.UInt8Vector;
69+
import org.apache.arrow.vector.UuidVector;
6870
import org.apache.arrow.vector.ValueVector;
6971
import org.apache.arrow.vector.VarBinaryVector;
7072
import org.apache.arrow.vector.VarCharVector;
@@ -138,6 +140,9 @@ public static ArrowFlightJdbcAccessor createAccessor(
138140
} else if (vector instanceof LargeVarBinaryVector) {
139141
return new ArrowFlightJdbcBinaryVectorAccessor(
140142
(LargeVarBinaryVector) vector, getCurrentRow, setCursorWasNull);
143+
} else if (vector instanceof UuidVector) {
144+
return new ArrowFlightJdbcUuidVectorAccessor(
145+
(UuidVector) vector, getCurrentRow, setCursorWasNull);
141146
} else if (vector instanceof FixedSizeBinaryVector) {
142147
return new ArrowFlightJdbcBinaryVectorAccessor(
143148
(FixedSizeBinaryVector) vector, getCurrentRow, setCursorWasNull);
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
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.accessor.impl.binary;
18+
19+
import java.util.UUID;
20+
import java.util.function.IntSupplier;
21+
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
22+
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
23+
import org.apache.arrow.vector.UuidVector;
24+
import org.apache.arrow.vector.util.UuidUtility;
25+
26+
/**
27+
* Accessor for the Arrow UUID extension type ({@link UuidVector}).
28+
*
29+
* <p>This accessor provides JDBC-compatible access to UUID values stored in Arrow's canonical UUID
30+
* extension type ('arrow.uuid'). It follows PostgreSQL JDBC driver conventions:
31+
*
32+
* <ul>
33+
* <li>{@link #getObject()} returns {@link java.util.UUID}
34+
* <li>{@link #getString()} returns the hyphenated string format (e.g.,
35+
* "550e8400-e29b-41d4-a716-446655440000")
36+
* <li>{@link #getBytes()} returns the 16-byte binary representation
37+
* </ul>
38+
*/
39+
public class ArrowFlightJdbcUuidVectorAccessor extends ArrowFlightJdbcAccessor {
40+
41+
private final UuidVector vector;
42+
43+
/**
44+
* Creates a new accessor for a UUID vector.
45+
*
46+
* @param vector the UUID vector to access
47+
* @param currentRowSupplier supplier for the current row index
48+
* @param setCursorWasNull consumer to set the wasNull flag
49+
*/
50+
public ArrowFlightJdbcUuidVectorAccessor(
51+
UuidVector vector,
52+
IntSupplier currentRowSupplier,
53+
ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
54+
super(currentRowSupplier, setCursorWasNull);
55+
this.vector = vector;
56+
}
57+
58+
@Override
59+
public Object getObject() {
60+
UUID uuid = vector.getObject(getCurrentRow());
61+
this.wasNull = uuid == null;
62+
this.wasNullConsumer.setWasNull(this.wasNull);
63+
return uuid;
64+
}
65+
66+
@Override
67+
public Class<?> getObjectClass() {
68+
return UUID.class;
69+
}
70+
71+
@Override
72+
public String getString() {
73+
UUID uuid = (UUID) getObject();
74+
if (uuid == null) {
75+
return null;
76+
}
77+
return uuid.toString();
78+
}
79+
80+
@Override
81+
public byte[] getBytes() {
82+
UUID uuid = (UUID) getObject();
83+
if (uuid == null) {
84+
return null;
85+
}
86+
return UuidUtility.getBytesFromUUID(uuid);
87+
}
88+
}
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
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.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeIdFromArrowType;
20+
import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeNameFromArrowType;
21+
22+
import java.nio.ByteBuffer;
23+
import java.util.UUID;
24+
import org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter;
25+
import org.apache.arrow.vector.FieldVector;
26+
import org.apache.arrow.vector.UuidVector;
27+
import org.apache.arrow.vector.types.pojo.Field;
28+
import org.apache.arrow.vector.util.UuidUtility;
29+
import org.apache.calcite.avatica.AvaticaParameter;
30+
import org.apache.calcite.avatica.remote.TypedValue;
31+
import org.apache.calcite.avatica.util.ByteString;
32+
33+
/**
34+
* AvaticaParameterConverter for UUID Arrow extension type.
35+
*
36+
* <p>Handles conversion of UUID values from JDBC parameters to Arrow's UUID extension type. Accepts
37+
* both {@link UUID} objects and String representations of UUIDs.
38+
*/
39+
public class UuidAvaticaParameterConverter implements AvaticaParameterConverter {
40+
41+
public UuidAvaticaParameterConverter() {}
42+
43+
@Override
44+
public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
45+
if (!(vector instanceof UuidVector)) {
46+
return false;
47+
}
48+
49+
UuidVector uuidVector = (UuidVector) vector;
50+
Object value = typedValue.toJdbc(null);
51+
52+
if (value == null) {
53+
uuidVector.setNull(index);
54+
return true;
55+
}
56+
57+
UUID uuid;
58+
if (value instanceof UUID) {
59+
uuid = (UUID) value;
60+
} else if (value instanceof String) {
61+
uuid = UUID.fromString((String) value);
62+
} else if (value instanceof byte[]) {
63+
byte[] bytes = (byte[]) value;
64+
if (bytes.length != 16) {
65+
throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length);
66+
}
67+
uuid = uuidFromBytes(bytes);
68+
} else if (value instanceof ByteString) {
69+
byte[] bytes = ((ByteString) value).getBytes();
70+
if (bytes.length != 16) {
71+
throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length);
72+
}
73+
uuid = uuidFromBytes(bytes);
74+
} else {
75+
throw new IllegalArgumentException(
76+
"Cannot convert " + value.getClass().getName() + " to UUID");
77+
}
78+
79+
uuidVector.setSafe(index, UuidUtility.getBytesFromUUID(uuid));
80+
return true;
81+
}
82+
83+
@Override
84+
public AvaticaParameter createParameter(Field field) {
85+
final String name = field.getName();
86+
final int jdbcType = getSqlTypeIdFromArrowType(field.getType());
87+
final String typeName = getSqlTypeNameFromArrowType(field.getType());
88+
final String className = UUID.class.getCanonicalName();
89+
return new AvaticaParameter(false, 0, 0, jdbcType, typeName, className, name);
90+
}
91+
92+
private static UUID uuidFromBytes(byte[] bytes) {
93+
final long mostSignificantBits;
94+
final long leastSignificantBits;
95+
ByteBuffer bb = ByteBuffer.wrap(bytes);
96+
// Reads the first eight bytes
97+
mostSignificantBits = bb.getLong();
98+
// Reads the first eight bytes at this buffer's current
99+
leastSignificantBits = bb.getLong();
100+
101+
return new UUID(mostSignificantBits, leastSignificantBits);
102+
}
103+
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,14 @@
4141
import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter;
4242
import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter;
4343
import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter;
44+
import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter;
4445
import org.apache.arrow.memory.BufferAllocator;
4546
import org.apache.arrow.vector.FieldVector;
4647
import org.apache.arrow.vector.VectorSchemaRoot;
48+
import org.apache.arrow.vector.extension.UuidType;
4749
import org.apache.arrow.vector.types.pojo.ArrowType;
50+
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor;
51+
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
4852
import org.apache.calcite.avatica.remote.TypedValue;
4953
import org.checkerframework.checker.nullness.qual.Nullable;
5054

@@ -290,5 +294,15 @@ public Boolean visit(ArrowType.RunEndEncoded type) {
290294
throw new UnsupportedOperationException(
291295
"No Avatica parameter binder implemented for type " + type);
292296
}
297+
298+
@Override
299+
public Boolean visit(ExtensionType type) {
300+
if (type instanceof UuidType) {
301+
return new UuidAvaticaParameterConverter().bindParameter(vector, typedValue, index);
302+
}
303+
304+
// fallback to default implementation
305+
return ArrowTypeVisitor.super.visit(type);
306+
}
293307
}
294308
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,12 @@
4343
import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter;
4444
import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter;
4545
import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter;
46+
import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter;
4647
import org.apache.arrow.flight.sql.FlightSqlColumnMetadata;
48+
import org.apache.arrow.vector.extension.UuidType;
4749
import org.apache.arrow.vector.types.pojo.ArrowType;
50+
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor;
51+
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
4852
import org.apache.arrow.vector.types.pojo.Field;
4953
import org.apache.calcite.avatica.AvaticaParameter;
5054
import org.apache.calcite.avatica.ColumnMetaData;
@@ -294,5 +298,15 @@ public AvaticaParameter visit(ArrowType.RunEndEncoded type) {
294298
throw new UnsupportedOperationException(
295299
"No Avatica parameter binder implemented for type " + type);
296300
}
301+
302+
@Override
303+
public AvaticaParameter visit(ExtensionType type) {
304+
if (type instanceof UuidType) {
305+
return new UuidAvaticaParameterConverter().createParameter(field);
306+
}
307+
308+
// fallback to default implementation
309+
return ArrowTypeVisitor.super.visit(type);
310+
}
297311
}
298312
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,13 @@
2020
import java.sql.Types;
2121
import java.util.HashMap;
2222
import java.util.Map;
23+
import org.apache.arrow.vector.extension.UuidType;
2324
import org.apache.arrow.vector.types.FloatingPointPrecision;
2425
import org.apache.arrow.vector.types.pojo.ArrowType;
2526

2627
/** SQL Types utility functions. */
2728
public class SqlTypes {
29+
2830
private static final Map<Integer, String> typeIdToName = new HashMap<>();
2931

3032
static {
@@ -110,6 +112,9 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
110112
case BinaryView:
111113
return Types.VARBINARY;
112114
case FixedSizeBinary:
115+
if (arrowType instanceof UuidType) {
116+
return Types.OTHER;
117+
}
113118
return Types.BINARY;
114119
case LargeBinary:
115120
return Types.LONGVARBINARY;

0 commit comments

Comments
 (0)