Skip to content

Commit 9202240

Browse files
xborderxborder
committed
working prototype
1 parent 82baf18 commit 9202240

File tree

11 files changed

+709
-3
lines changed

11 files changed

+709
-3
lines changed

‎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;
@@ -133,6 +135,9 @@ public static ArrowFlightJdbcAccessor createAccessor(
133135
} else if (vector instanceof LargeVarBinaryVector) {
134136
return new ArrowFlightJdbcBinaryVectorAccessor(
135137
(LargeVarBinaryVector) vector, getCurrentRow, setCursorWasNull);
138+
} else if (vector instanceof UuidVector) {
139+
return new ArrowFlightJdbcUuidVectorAccessor(
140+
(UuidVector) vector, getCurrentRow, setCursorWasNull);
136141
} else if (vector instanceof FixedSizeBinaryVector) {
137142
return new ArrowFlightJdbcBinaryVectorAccessor(
138143
(FixedSizeBinaryVector) vector, getCurrentRow, setCursorWasNull);

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

Lines changed: 88 additions & 0 deletions
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+
}

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

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
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 java.nio.ByteBuffer;
20+
import java.sql.Types;
21+
import java.util.UUID;
22+
import org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter;
23+
import org.apache.arrow.driver.jdbc.utils.SqlTypes;
24+
import org.apache.arrow.vector.FieldVector;
25+
import org.apache.arrow.vector.UuidVector;
26+
import org.apache.arrow.vector.types.pojo.Field;
27+
import org.apache.arrow.vector.util.UuidUtility;
28+
import org.apache.calcite.avatica.AvaticaParameter;
29+
import org.apache.calcite.avatica.remote.TypedValue;
30+
import org.apache.calcite.avatica.util.ByteString;
31+
32+
/**
33+
* AvaticaParameterConverter for UUID Arrow extension type.
34+
*
35+
* <p>Handles conversion of UUID values from JDBC parameters to Arrow's UUID extension type. Accepts
36+
* both {@link UUID} objects and String representations of UUIDs.
37+
*/
38+
public class UuidAvaticaParameterConverter implements AvaticaParameterConverter {
39+
40+
public UuidAvaticaParameterConverter() {}
41+
42+
@Override
43+
public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
44+
if (!(vector instanceof UuidVector)) {
45+
return false;
46+
}
47+
48+
UuidVector uuidVector = (UuidVector) vector;
49+
Object value = typedValue.toJdbc(null);
50+
51+
if (value == null) {
52+
uuidVector.setNull(index);
53+
return true;
54+
}
55+
56+
UUID uuid;
57+
if (value instanceof UUID) {
58+
uuid = (UUID) value;
59+
} else if (value instanceof String) {
60+
uuid = UUID.fromString((String) value);
61+
} else if (value instanceof byte[]) {
62+
byte[] bytes = (byte[]) value;
63+
if (bytes.length != 16) {
64+
throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length);
65+
}
66+
uuid = uuidFromBytes(bytes);
67+
} else if (value instanceof ByteString) {
68+
byte[] bytes = ((ByteString) value).getBytes();
69+
if (bytes.length != 16) {
70+
throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length);
71+
}
72+
uuid = uuidFromBytes(bytes);
73+
} else {
74+
throw new IllegalArgumentException(
75+
"Cannot convert " + value.getClass().getName() + " to UUID");
76+
}
77+
78+
uuidVector.setSafe(index, UuidUtility.getBytesFromUUID(uuid));
79+
return true;
80+
}
81+
82+
@Override
83+
public AvaticaParameter createParameter(Field field) {
84+
final String name = field.getName();
85+
final int jdbcType = Types.OTHER;
86+
final String typeName = SqlTypes.UUID_TYPE_NAME;
87+
final String className = UUID.class.getCanonicalName();
88+
return new AvaticaParameter(false, 0, 0, jdbcType, typeName, className, name);
89+
}
90+
91+
private static UUID uuidFromBytes(byte[] bytes) {
92+
// long msb = 0;
93+
// long lsb = 0;
94+
// for (int i = 0; i < 8; i++) {
95+
// msb = (msb << 8) | (bytes[i] & 0xff);
96+
// }
97+
// for (int i = 8; i < 16; i++) {
98+
// lsb = (lsb << 8) | (bytes[i] & 0xff);
99+
// }
100+
101+
final long mostSignificantBits;
102+
final long leastSignificantBits;
103+
ByteBuffer bb = ByteBuffer.wrap(bytes);
104+
mostSignificantBits = bb.getLong();
105+
leastSignificantBits = bb.getLong();
106+
107+
return new UUID(mostSignificantBits, leastSignificantBits);
108+
}
109+
}

‎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
@@ -39,10 +39,14 @@
3939
import org.apache.arrow.driver.jdbc.converter.impl.TimestampAvaticaParameterConverter;
4040
import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter;
4141
import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter;
42+
import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter;
4243
import org.apache.arrow.memory.BufferAllocator;
4344
import org.apache.arrow.vector.FieldVector;
4445
import org.apache.arrow.vector.VectorSchemaRoot;
46+
import org.apache.arrow.vector.extension.UuidType;
4547
import org.apache.arrow.vector.types.pojo.ArrowType;
48+
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor;
49+
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
4650
import org.apache.calcite.avatica.remote.TypedValue;
4751
import org.checkerframework.checker.nullness.qual.Nullable;
4852

@@ -288,5 +292,15 @@ public Boolean visit(ArrowType.RunEndEncoded type) {
288292
throw new UnsupportedOperationException(
289293
"No Avatica parameter binder implemented for type " + type);
290294
}
295+
296+
@Override
297+
public Boolean visit(ExtensionType type) {
298+
if (type instanceof UuidType) {
299+
return new UuidAvaticaParameterConverter().bindParameter(vector, typedValue, index);
300+
}
301+
302+
// fallback to default implementation
303+
return ArrowTypeVisitor.super.visit(type);
304+
}
291305
}
292306
}

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

Lines changed: 16 additions & 3 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;
@@ -68,7 +72,6 @@ public static List<ColumnMetaData> convertArrowFieldsToColumnMetaDataList(
6872
.map(
6973
index -> {
7074
final Field field = fields.get(index);
71-
final ArrowType fieldType = field.getType();
7275

7376
final Common.ColumnMetaData.Builder builder =
7477
Common.ColumnMetaData.newBuilder()
@@ -80,8 +83,8 @@ public static List<ColumnMetaData> convertArrowFieldsToColumnMetaDataList(
8083

8184
builder.setType(
8285
Common.AvaticaType.newBuilder()
83-
.setId(SqlTypes.getSqlTypeIdFromArrowType(fieldType))
84-
.setName(SqlTypes.getSqlTypeNameFromArrowType(fieldType))
86+
.setId(SqlTypes.getSqlTypeIdFromField(field))
87+
.setName(SqlTypes.getSqlTypeNameFromField(field))
8588
.build());
8689

8790
return ColumnMetaData.fromProto(builder.build());
@@ -294,5 +297,15 @@ public AvaticaParameter visit(ArrowType.RunEndEncoded type) {
294297
throw new UnsupportedOperationException(
295298
"No Avatica parameter binder implemented for type " + type);
296299
}
300+
301+
@Override
302+
public AvaticaParameter visit(ExtensionType type) {
303+
if (type instanceof UuidType) {
304+
return new UuidAvaticaParameterConverter().createParameter(field);
305+
}
306+
307+
// fallback to default implementation
308+
return ArrowTypeVisitor.super.visit(type);
309+
}
297310
}
298311
}

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,16 @@
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;
26+
import org.apache.arrow.vector.types.pojo.Field;
2527

2628
/** SQL Types utility functions. */
2729
public class SqlTypes {
30+
/** SQL type name for UUID (matches PostgreSQL JDBC driver convention). */
31+
public static final String UUID_TYPE_NAME = "uuid";
32+
2833
private static final Map<Integer, String> typeIdToName = new HashMap<>();
2934

3035
static {
@@ -162,4 +167,44 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
162167

163168
throw new IllegalArgumentException("Unsupported ArrowType " + arrowType);
164169
}
170+
171+
/**
172+
* Convert given {@link Field} to its corresponding SQL type ID, handling extension types.
173+
*
174+
* @param field field to convert from
175+
* @return corresponding SQL type ID.
176+
* @see java.sql.Types
177+
*/
178+
public static int getSqlTypeIdFromField(Field field) {
179+
ArrowType arrowType = field.getType();
180+
if (arrowType instanceof UuidType) {
181+
return Types.OTHER;
182+
}
183+
return getSqlTypeIdFromArrowType(arrowType);
184+
}
185+
186+
/**
187+
* Convert given {@link Field} to its corresponding SQL type name, handling extension types.
188+
*
189+
* @param field field to convert from
190+
* @return corresponding SQL type name.
191+
* @see java.sql.Types
192+
*/
193+
public static String getSqlTypeNameFromField(Field field) {
194+
ArrowType arrowType = field.getType();
195+
if (arrowType instanceof UuidType) {
196+
return UUID_TYPE_NAME;
197+
}
198+
return getSqlTypeNameFromArrowType(arrowType);
199+
}
200+
201+
/**
202+
* Check if the given field represents a UUID type.
203+
*
204+
* @param field field to check
205+
* @return true if the field is a UUID extension type
206+
*/
207+
public static boolean isUuidField(Field field) {
208+
return field.getType() instanceof UuidType;
209+
}
165210
}

‎flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java‎

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

19+
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
1920
import static org.junit.jupiter.api.Assertions.assertTrue;
2021

2122
import java.util.function.IntSupplier;
2223
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor;
24+
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor;
2325
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor;
2426
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor;
2527
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor;
@@ -471,4 +473,15 @@ public void createAccessorForMapVector() {
471473
assertTrue(accessor instanceof ArrowFlightJdbcMapVectorAccessor);
472474
}
473475
}
476+
477+
@Test
478+
public void createAccessorForUuidVector() {
479+
try (ValueVector valueVector = rootAllocatorTestExtension.createUuidVector()) {
480+
ArrowFlightJdbcAccessor accessor =
481+
ArrowFlightJdbcAccessorFactory.createAccessor(
482+
valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {});
483+
484+
assertInstanceOf(ArrowFlightJdbcUuidVectorAccessor.class, accessor);
485+
}
486+
}
474487
}

0 commit comments

Comments
 (0)