This wrapper exists to isolate the parquet-variant dependency from Arrow's public API, + * allowing the vector module to expose variant functionality without requiring users to depend on + * parquet-variant directly. It also ensures that nested variant values (from arrays and objects) + * are consistently wrapped. + */ +public class Variant { + + private final org.apache.parquet.variant.Variant delegate; + + /** Creates a Variant from raw metadata and value byte arrays. */ + public Variant(byte[] metadata, byte[] value) { + this.delegate = new org.apache.parquet.variant.Variant(value, metadata); + } + + /** Creates a Variant by copying data from ArrowBuf instances. */ + public Variant( + ArrowBuf metadataBuffer, + int metadataStart, + int metadataEnd, + ArrowBuf valueBuffer, + int valueStart, + int valueEnd) { + byte[] metadata = new byte[metadataEnd - metadataStart]; + byte[] value = new byte[valueEnd - valueStart]; + metadataBuffer.getBytes(metadataStart, metadata); + valueBuffer.getBytes(valueStart, value); + this.delegate = new org.apache.parquet.variant.Variant(value, metadata); + } + + private Variant(org.apache.parquet.variant.Variant delegate) { + this.delegate = delegate; + } + + /** Constructs a Variant from a NullableVariantHolder. */ + public Variant(NullableVariantHolder holder) { + this( + holder.metadataBuffer, + holder.metadataStart, + holder.metadataEnd, + holder.valueBuffer, + holder.valueStart, + holder.valueEnd); + } + + public ByteBuffer getValueBuffer() { + return delegate.getValueBuffer(); + } + + public ByteBuffer getMetadataBuffer() { + return delegate.getMetadataBuffer(); + } + + public boolean getBoolean() { + return delegate.getBoolean(); + } + + public byte getByte() { + return delegate.getByte(); + } + + public short getShort() { + return delegate.getShort(); + } + + public int getInt() { + return delegate.getInt(); + } + + public long getLong() { + return delegate.getLong(); + } + + public double getDouble() { + return delegate.getDouble(); + } + + public BigDecimal getDecimal() { + return delegate.getDecimal(); + } + + public float getFloat() { + return delegate.getFloat(); + } + + public ByteBuffer getBinary() { + return delegate.getBinary(); + } + + public UUID getUUID() { + return delegate.getUUID(); + } + + public String getString() { + return delegate.getString(); + } + + public Type getType() { + return Type.fromParquet(delegate.getType()); + } + + public int numObjectElements() { + return delegate.numObjectElements(); + } + + public Variant getFieldByKey(String key) { + org.apache.parquet.variant.Variant result = delegate.getFieldByKey(key); + return result != null ? wrap(result) : null; + } + + public ObjectField getFieldAtIndex(int idx) { + org.apache.parquet.variant.Variant.ObjectField field = delegate.getFieldAtIndex(idx); + return new ObjectField(field.key, wrap(field.value)); + } + + public int numArrayElements() { + return delegate.numArrayElements(); + } + + public Variant getElementAtIndex(int index) { + org.apache.parquet.variant.Variant result = delegate.getElementAtIndex(index); + return result != null ? wrap(result) : null; + } + + private static Variant wrap(org.apache.parquet.variant.Variant parquetVariant) { + return new Variant(parquetVariant); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + Variant variant = (Variant) o; + return delegate.getMetadataBuffer().equals(variant.delegate.getMetadataBuffer()) + && delegate.getValueBuffer().equals(variant.delegate.getValueBuffer()); + } + + @Override + public int hashCode() { + return Objects.hash(delegate.getMetadataBuffer(), delegate.getValueBuffer()); + } + + @Override + public String toString() { + return "Variant{type=" + getType() + '}'; + } + + public enum Type { + OBJECT, + ARRAY, + NULL, + BOOLEAN, + BYTE, + SHORT, + INT, + LONG, + STRING, + DOUBLE, + DECIMAL4, + DECIMAL8, + DECIMAL16, + DATE, + TIMESTAMP_TZ, + TIMESTAMP_NTZ, + FLOAT, + BINARY, + TIME, + TIMESTAMP_NANOS_TZ, + TIMESTAMP_NANOS_NTZ, + UUID; + + static Type fromParquet(org.apache.parquet.variant.Variant.Type parquetType) { + return Type.valueOf(parquetType.name()); + } + } + + public static final class ObjectField { + public final String key; + public final Variant value; + + public ObjectField(String key, Variant value) { + this.key = key; + this.value = value; + } + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantType.java b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantType.java new file mode 100644 index 0000000000..3deb70cdc0 --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantType.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.variant.extension; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.variant.impl.VariantWriterImpl; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.FieldType; + +/** + * Arrow extension type for Parquet + * Variant binary encoding. The type itself does not support shredded variant data. + */ +public final class VariantType extends ExtensionType { + + public static final VariantType INSTANCE = new VariantType(); + + public static final String EXTENSION_NAME = "parquet.variant"; + + static { + ExtensionTypeRegistry.register(INSTANCE); + } + + private VariantType() {} + + @Override + public ArrowType storageType() { + return ArrowType.Struct.INSTANCE; + } + + @Override + public String extensionName() { + return EXTENSION_NAME; + } + + @Override + public boolean extensionEquals(ExtensionType other) { + return other instanceof VariantType; + } + + @Override + public String serialize() { + return ""; + } + + @Override + public ArrowType deserialize(ArrowType storageType, String serializedData) { + if (!storageType.equals(this.storageType())) { + throw new UnsupportedOperationException( + "Cannot construct VariantType from underlying type " + storageType); + } + return INSTANCE; + } + + @Override + public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { + return new VariantVector(name, allocator); + } + + @Override + public boolean isComplex() { + // The type itself is not complex meaning we need separate functions to convert/extract + // different types. + // Meanwhile, the containing vector is complex in terms of containing multiple values (metadata + // and value) + return false; + } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + return new VariantWriterImpl((VariantVector) vector); + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantVector.java b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantVector.java new file mode 100644 index 0000000000..1bbf1a6bdb --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantVector.java @@ -0,0 +1,348 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.variant.extension; + +import java.nio.ByteBuffer; +import java.util.List; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.variant.holders.VariantHolder; +import org.apache.arrow.vector.BitVectorHelper; +import org.apache.arrow.vector.ExtensionTypeVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.complex.AbstractStructVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.Binary; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.CallBack; +import org.apache.arrow.vector.util.TransferPair; + +/** + * Arrow vector for storing {@link VariantType} values. + * + *
Stores semi-structured data (like JSON) as metadata + value binary pairs, allowing
+ * type-flexible columnar storage within Arrow's type system.
+ */
+public class VariantVector extends ExtensionTypeVector This writer handles writing variant data to a {@link VariantVector}. It accepts both {@link
+ * VariantHolder} and {@link NullableVariantHolder} objects containing metadata and value buffers
+ * and writes them to the appropriate position in the vector.
+ */
+public class VariantWriterImpl extends AbstractExtensionTypeWriter This method handles {@link ExtensionHolder} by delegating to {@link #write(ExtensionHolder)}
+ * and {@link Variant} by delegating to {@link #writeVariant(Variant)}.
+ *
+ * @param object the object to write, must be an {@link ExtensionHolder} or {@link Variant}
+ * @throws IllegalArgumentException if the object is not an {@link ExtensionHolder} or {@link
+ * Variant}
+ */
+ @Override
+ public void writeExtension(Object object) {
+ if (object instanceof ExtensionHolder) {
+ write((ExtensionHolder) object);
+ } else if (object instanceof Variant) {
+ writeVariant((Variant) object);
+ } else {
+ throw new IllegalArgumentException(
+ String.format(UNSUPPORTED_TYPE_TEMPLATE, object.getClass().getName()));
+ }
+ }
+
+ private void writeVariant(Variant variant) {
+ java.nio.ByteBuffer metadataBuffer = variant.getMetadataBuffer();
+ java.nio.ByteBuffer valueBuffer = variant.getValueBuffer();
+ int metadataLength = metadataBuffer.remaining();
+ int valueLength = valueBuffer.remaining();
+ try (ArrowBuf metadataBuf = vector.getAllocator().buffer(metadataLength);
+ ArrowBuf valueBuf = vector.getAllocator().buffer(valueLength)) {
+ metadataBuf.setBytes(0, metadataBuffer.duplicate());
+ valueBuf.setBytes(0, valueBuffer.duplicate());
+ NullableVariantHolder holder = new NullableVariantHolder();
+ holder.isSet = 1;
+ holder.metadataBuffer = metadataBuf;
+ holder.metadataStart = 0;
+ holder.metadataEnd = metadataLength;
+ holder.valueBuffer = valueBuf;
+ holder.valueStart = 0;
+ holder.valueEnd = valueLength;
+ vector.setSafe(getPosition(), holder);
+ vector.setValueCount(getPosition() + 1);
+ }
+ }
+
+ @Override
+ public void writeExtension(Object value, ArrowType type) {
+ writeExtension(value);
+ }
+
+ /**
+ * Writes a variant holder to the vector at the current position.
+ *
+ * The holder can be either a {@link VariantHolder} (non-nullable, always set) or a {@link
+ * NullableVariantHolder} (nullable, may be null). The data is written using {@link
+ * VariantVector#setSafe(int, NullableVariantHolder)} which handles buffer allocation and copying.
+ *
+ * @param extensionHolder the variant holder to write, must be a {@link VariantHolder} or {@link
+ * NullableVariantHolder}
+ * @throws IllegalArgumentException if the holder is neither a {@link VariantHolder} nor a {@link
+ * NullableVariantHolder}
+ */
+ @Override
+ public void write(ExtensionHolder extensionHolder) {
+ if (extensionHolder instanceof VariantHolder) {
+ vector.setSafe(getPosition(), (VariantHolder) extensionHolder);
+ } else if (extensionHolder instanceof NullableVariantHolder) {
+ vector.setSafe(getPosition(), (NullableVariantHolder) extensionHolder);
+ } else {
+ throw new IllegalArgumentException(
+ String.format(UNSUPPORTED_TYPE_TEMPLATE, extensionHolder.getClass().getName()));
+ }
+ vector.setValueCount(getPosition() + 1);
+ }
+}
diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java b/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java
new file mode 100644
index 0000000000..bc46a68616
--- /dev/null
+++ b/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java
@@ -0,0 +1,439 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.variant;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.util.UUID;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.parquet.variant.VariantBuilder;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class TestVariant {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ void beforeEach() {
+ allocator = new RootAllocator();
+ }
+
+ @AfterEach
+ void afterEach() {
+ allocator.close();
+ }
+
+ static Variant buildVariant(VariantBuilder builder) {
+ org.apache.parquet.variant.Variant parquetVariant = builder.build();
+ ByteBuffer valueBuf = parquetVariant.getValueBuffer();
+ ByteBuffer metaBuf = parquetVariant.getMetadataBuffer();
+ byte[] valueBytes = new byte[valueBuf.remaining()];
+ byte[] metaBytes = new byte[metaBuf.remaining()];
+ valueBuf.get(valueBytes);
+ metaBuf.get(metaBytes);
+ return new Variant(metaBytes, valueBytes);
+ }
+
+ public static Variant variantString(String value) {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendString(value);
+ return buildVariant(builder);
+ }
+
+ @Test
+ void testConstructionWithArrowBuf() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendInt(42);
+ Variant source = buildVariant(builder);
+ int metaLen = source.getMetadataBuffer().remaining();
+ int valueLen = source.getValueBuffer().remaining();
+
+ try (ArrowBuf metadataArrowBuf = allocator.buffer(metaLen + 2);
+ ArrowBuf valueArrowBuf = allocator.buffer(valueLen + 3)) {
+ metadataArrowBuf.setBytes(2, source.getMetadataBuffer());
+ valueArrowBuf.setBytes(3, source.getValueBuffer());
+
+ Variant variant =
+ new Variant(metadataArrowBuf, 2, 2 + metaLen, valueArrowBuf, 3, 3 + valueLen);
+
+ assertEquals(Variant.Type.INT, variant.getType());
+ assertEquals(42, variant.getInt());
+ }
+ }
+
+ @Test
+ void testNullType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendNull();
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.NULL, variant.getType());
+ }
+
+ @Test
+ void testBooleanType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendBoolean(true);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.BOOLEAN, variant.getType());
+ assertTrue(variant.getBoolean());
+
+ builder = new VariantBuilder();
+ builder.appendBoolean(false);
+ variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.BOOLEAN, variant.getType());
+ assertFalse(variant.getBoolean());
+ }
+
+ @Test
+ void testByteType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendByte((byte) 42);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.BYTE, variant.getType());
+ assertEquals((byte) 42, variant.getByte());
+ }
+
+ @Test
+ void testShortType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendShort((short) 1234);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.SHORT, variant.getType());
+ assertEquals((short) 1234, variant.getShort());
+ }
+
+ @Test
+ void testIntType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendInt(123456);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.INT, variant.getType());
+ assertEquals(123456, variant.getInt());
+ }
+
+ @Test
+ void testLongType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendLong(9876543210L);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.LONG, variant.getType());
+ assertEquals(9876543210L, variant.getLong());
+ }
+
+ @Test
+ void testFloatType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendFloat(3.14f);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.FLOAT, variant.getType());
+ assertEquals(3.14f, variant.getFloat(), 0.001f);
+ }
+
+ @Test
+ void testDoubleType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendDouble(3.14159265359);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.DOUBLE, variant.getType());
+ assertEquals(3.14159265359, variant.getDouble(), 0.0000001);
+ }
+
+ @Test
+ void testStringType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendString("hello world");
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.STRING, variant.getType());
+ assertEquals("hello world", variant.getString());
+ }
+
+ @Test
+ void testDecimalType() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendDecimal(new BigDecimal("123.456"));
+ Variant variant = buildVariant(builder);
+
+ assertTrue(
+ variant.getType() == Variant.Type.DECIMAL4
+ || variant.getType() == Variant.Type.DECIMAL8
+ || variant.getType() == Variant.Type.DECIMAL16);
+ assertEquals(new BigDecimal("123.456"), variant.getDecimal());
+ }
+
+ @Test
+ void testBinaryType() {
+ VariantBuilder builder = new VariantBuilder();
+ byte[] data = new byte[] {1, 2, 3, 4, 5};
+ builder.appendBinary(ByteBuffer.wrap(data));
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.BINARY, variant.getType());
+ ByteBuffer result = variant.getBinary();
+ byte[] resultBytes = new byte[result.remaining()];
+ result.get(resultBytes);
+ assertArrayEquals(data, resultBytes);
+ }
+
+ @Test
+ void testUuidType() {
+ VariantBuilder builder = new VariantBuilder();
+ UUID uuid = UUID.randomUUID();
+ builder.appendUUID(uuid);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.UUID, variant.getType());
+ assertEquals(uuid, variant.getUUID());
+ }
+
+ @Test
+ void testDateType() {
+ VariantBuilder builder = new VariantBuilder();
+ int daysSinceEpoch = 19000;
+ builder.appendDate(daysSinceEpoch);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.DATE, variant.getType());
+ }
+
+ @Test
+ void testTimestampTzType() {
+ VariantBuilder builder = new VariantBuilder();
+ long micros = System.currentTimeMillis() * 1000;
+ builder.appendTimestampTz(micros);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.TIMESTAMP_TZ, variant.getType());
+ }
+
+ @Test
+ void testTimestampNtzType() {
+ VariantBuilder builder = new VariantBuilder();
+ long micros = System.currentTimeMillis() * 1000;
+ builder.appendTimestampNtz(micros);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.TIMESTAMP_NTZ, variant.getType());
+ }
+
+ @Test
+ void testTimeType() {
+ VariantBuilder builder = new VariantBuilder();
+ long micros = 12345678L;
+ builder.appendTime(micros);
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.TIME, variant.getType());
+ }
+
+ @Test
+ void testObjectType() {
+ VariantBuilder builder = new VariantBuilder();
+ var objBuilder = builder.startObject();
+ objBuilder.appendKey("name");
+ objBuilder.appendString("test");
+ objBuilder.appendKey("value");
+ objBuilder.appendInt(42);
+ builder.endObject();
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.OBJECT, variant.getType());
+ assertEquals(2, variant.numObjectElements());
+
+ Variant nameField = variant.getFieldByKey("name");
+ assertNotNull(nameField);
+ assertEquals(Variant.Type.STRING, nameField.getType());
+ assertEquals("test", nameField.getString());
+
+ Variant valueField = variant.getFieldByKey("value");
+ assertNotNull(valueField);
+ assertEquals(Variant.Type.INT, valueField.getType());
+ assertEquals(42, valueField.getInt());
+
+ assertNull(variant.getFieldByKey("nonexistent"));
+
+ // Empty object
+ builder = new VariantBuilder();
+ builder.startObject();
+ builder.endObject();
+ Variant emptyObj = buildVariant(builder);
+ assertEquals(Variant.Type.OBJECT, emptyObj.getType());
+ assertEquals(0, emptyObj.numObjectElements());
+ }
+
+ @Test
+ void testObjectFieldAtIndex() {
+ VariantBuilder builder = new VariantBuilder();
+ var objBuilder = builder.startObject();
+ objBuilder.appendKey("alpha");
+ objBuilder.appendInt(1);
+ objBuilder.appendKey("beta");
+ objBuilder.appendInt(2);
+ builder.endObject();
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.OBJECT, variant.getType());
+ assertEquals(2, variant.numObjectElements());
+
+ Variant.ObjectField field0 = variant.getFieldAtIndex(0);
+ assertNotNull(field0);
+ assertNotNull(field0.key);
+ assertNotNull(field0.value);
+
+ Variant.ObjectField field1 = variant.getFieldAtIndex(1);
+ assertNotNull(field1);
+ assertNotNull(field1.key);
+ assertNotNull(field1.value);
+ }
+
+ @Test
+ void testArrayType() {
+ VariantBuilder builder = new VariantBuilder();
+ var arrayBuilder = builder.startArray();
+ arrayBuilder.appendInt(1);
+ arrayBuilder.appendInt(2);
+ arrayBuilder.appendInt(3);
+ builder.endArray();
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.ARRAY, variant.getType());
+ assertEquals(3, variant.numArrayElements());
+
+ Variant elem0 = variant.getElementAtIndex(0);
+ assertNotNull(elem0);
+ assertEquals(Variant.Type.INT, elem0.getType());
+ assertEquals(1, elem0.getInt());
+
+ Variant elem1 = variant.getElementAtIndex(1);
+ assertEquals(2, elem1.getInt());
+
+ Variant elem2 = variant.getElementAtIndex(2);
+ assertEquals(3, elem2.getInt());
+
+ assertNull(variant.getElementAtIndex(-1));
+ assertNull(variant.getElementAtIndex(3));
+
+ // Empty array
+ builder = new VariantBuilder();
+ builder.startArray();
+ builder.endArray();
+ Variant emptyArr = buildVariant(builder);
+ assertEquals(Variant.Type.ARRAY, emptyArr.getType());
+ assertEquals(0, emptyArr.numArrayElements());
+ }
+
+ @Test
+ void testNestedStructure() {
+ VariantBuilder builder = new VariantBuilder();
+ var objBuilder = builder.startObject();
+ objBuilder.appendKey("items");
+ var arrayBuilder = objBuilder.startArray();
+ arrayBuilder.appendString("a");
+ arrayBuilder.appendString("b");
+ objBuilder.endArray();
+ builder.endObject();
+ Variant variant = buildVariant(builder);
+
+ assertEquals(Variant.Type.OBJECT, variant.getType());
+ Variant items = variant.getFieldByKey("items");
+ assertNotNull(items);
+ assertEquals(Variant.Type.ARRAY, items.getType());
+ assertEquals(2, items.numArrayElements());
+ assertEquals("a", items.getElementAtIndex(0).getString());
+ assertEquals("b", items.getElementAtIndex(1).getString());
+ }
+
+ @Test
+ void testEquals() {
+ VariantBuilder builder1 = new VariantBuilder();
+ builder1.appendString("test");
+ Variant variant1 = buildVariant(builder1);
+
+ VariantBuilder builder2 = new VariantBuilder();
+ builder2.appendString("test");
+ Variant variant2 = buildVariant(builder2);
+
+ VariantBuilder builder3 = new VariantBuilder();
+ builder3.appendString("different");
+ Variant variant3 = buildVariant(builder3);
+
+ assertEquals(variant1, variant1);
+ assertEquals(variant1, variant2);
+ assertNotEquals(variant1, variant3);
+ assertNotEquals(variant1, null);
+ assertNotEquals(variant1, "not a variant");
+ }
+
+ @Test
+ void testHashCode() {
+ VariantBuilder builder1 = new VariantBuilder();
+ builder1.appendInt(42);
+ Variant variant1 = buildVariant(builder1);
+
+ VariantBuilder builder2 = new VariantBuilder();
+ builder2.appendInt(42);
+ Variant variant2 = buildVariant(builder2);
+
+ assertEquals(variant1.hashCode(), variant2.hashCode());
+ }
+
+ @Test
+ void testToString() {
+ VariantBuilder builder = new VariantBuilder();
+ builder.appendString("test");
+ Variant variant = buildVariant(builder);
+
+ String str = variant.toString();
+ assertNotNull(str);
+ assertTrue(str.contains("type="));
+ }
+
+ @Test
+ void testTypeEnumsMatch() {
+ for (Variant.Type arrowType : Variant.Type.values()) {
+ org.apache.parquet.variant.Variant.Type parquetType =
+ org.apache.parquet.variant.Variant.Type.valueOf(arrowType.name());
+ assertEquals(arrowType, Variant.Type.fromParquet(parquetType));
+ }
+ for (org.apache.parquet.variant.Variant.Type parquetType :
+ org.apache.parquet.variant.Variant.Type.values()) {
+ Variant.Type arrowType = Variant.Type.valueOf(parquetType.name());
+ assertEquals(parquetType.name(), arrowType.name());
+ }
+ }
+}
diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantExtensionType.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantExtensionType.java
new file mode 100644
index 0000000000..f3213d523a
--- /dev/null
+++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantExtensionType.java
@@ -0,0 +1,249 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.variant.extension;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.channels.SeekableByteChannel;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+import java.util.Collections;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.variant.TestVariant;
+import org.apache.arrow.variant.Variant;
+import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.compare.Range;
+import org.apache.arrow.vector.compare.RangeEqualsVisitor;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.complex.writer.BaseWriter;
+import org.apache.arrow.vector.ipc.ArrowFileReader;
+import org.apache.arrow.vector.ipc.ArrowFileWriter;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
+import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.VectorBatchAppender;
+import org.apache.arrow.vector.validate.ValidateVectorVisitor;
+import org.junit.jupiter.api.Test;
+
+public class TestVariantExtensionType {
+
+ private static void ensureRegistered(ArrowType.ExtensionType type) {
+ if (ExtensionTypeRegistry.lookup(type.extensionName()) == null) {
+ ExtensionTypeRegistry.register(type);
+ }
+ }
+
+ @Test
+ public void roundtripVariant() throws IOException {
+ ensureRegistered(VariantType.INSTANCE);
+ final Schema schema =
+ new Schema(Collections.singletonList(Field.nullable("a", VariantType.INSTANCE)));
+ try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+ VariantVector vector = (VariantVector) root.getVector("a");
+ vector.allocateNew();
+
+ vector.setSafe(0, TestVariant.variantString("hello"));
+ vector.setSafe(1, TestVariant.variantString("world"));
+ vector.setValueCount(2);
+ root.setRowCount(2);
+
+ final File file = File.createTempFile("varianttest", ".arrow");
+ try (final WritableByteChannel channel =
+ FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE);
+ final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) {
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+ }
+
+ try (final SeekableByteChannel channel =
+ Files.newByteChannel(Paths.get(file.getAbsolutePath()));
+ final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
+ reader.loadNextBatch();
+ final VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot();
+ assertEquals(root.getSchema(), readerRoot.getSchema());
+
+ final Field field = readerRoot.getSchema().getFields().get(0);
+ final VariantType expectedType = VariantType.INSTANCE;
+ assertEquals(
+ field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
+ expectedType.extensionName());
+ assertEquals(
+ field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
+ expectedType.serialize());
+
+ final ExtensionTypeVector deserialized =
+ (ExtensionTypeVector) readerRoot.getFieldVectors().get(0);
+ assertEquals(vector.getValueCount(), deserialized.getValueCount());
+ for (int i = 0; i < vector.getValueCount(); i++) {
+ assertEquals(vector.isNull(i), deserialized.isNull(i));
+ if (!vector.isNull(i)) {
+ assertEquals(vector.getObject(i), deserialized.getObject(i));
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void readVariantAsUnderlyingType() throws IOException {
+ ensureRegistered(VariantType.INSTANCE);
+ final Schema schema =
+ new Schema(Collections.singletonList(VariantVector.createVariantField("a")));
+ try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+ VariantVector vector = (VariantVector) root.getVector("a");
+ vector.allocateNew();
+
+ vector.setSafe(0, TestVariant.variantString("hello"));
+ vector.setValueCount(1);
+ root.setRowCount(1);
+
+ final File file = File.createTempFile("varianttest", ".arrow");
+ try (final WritableByteChannel channel =
+ FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE);
+ final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) {
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+ }
+
+ ExtensionTypeRegistry.unregister(VariantType.INSTANCE);
+
+ try (final SeekableByteChannel channel =
+ Files.newByteChannel(Paths.get(file.getAbsolutePath()));
+ final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
+ reader.loadNextBatch();
+ VectorSchemaRoot readRoot = reader.getVectorSchemaRoot();
+
+ // Verify schema properties
+ assertEquals(1, readRoot.getSchema().getFields().size());
+ assertEquals("a", readRoot.getSchema().getFields().get(0).getName());
+ assertTrue(readRoot.getSchema().getFields().get(0).getType() instanceof ArrowType.Struct);
+
+ // Verify extension metadata is preserved
+ final Field field = readRoot.getSchema().getFields().get(0);
+ assertEquals(
+ VariantType.EXTENSION_NAME,
+ field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME));
+ assertEquals("", field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA));
+
+ // Verify vector type and row count
+ assertEquals(1, readRoot.getRowCount());
+ FieldVector readVector = readRoot.getVector("a");
+ assertEquals(StructVector.class, readVector.getClass());
+
+ // Verify value count matches
+ StructVector structVector = (StructVector) readVector;
+ assertEquals(vector.getValueCount(), structVector.getValueCount());
+
+ // Verify the underlying data can be accessed from child vectors
+ VarBinaryVector metadataVector =
+ structVector.getChild(VariantVector.METADATA_VECTOR_NAME, VarBinaryVector.class);
+ VarBinaryVector valueVector =
+ structVector.getChild(VariantVector.VALUE_VECTOR_NAME, VarBinaryVector.class);
+ assertNotNull(metadataVector);
+ assertNotNull(valueVector);
+ assertEquals(1, metadataVector.getValueCount());
+ assertEquals(1, valueVector.getValueCount());
+ }
+ }
+ }
+
+ @Test
+ public void testVariantVectorCompare() {
+ VariantType variantType = VariantType.INSTANCE;
+ ExtensionTypeRegistry.register(variantType);
+ Variant hello = TestVariant.variantString("hello");
+ Variant world = TestVariant.variantString("world");
+ try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+ VariantVector a1 =
+ (VariantVector)
+ variantType.getNewVector("a", FieldType.nullable(variantType), allocator);
+ VariantVector a2 =
+ (VariantVector)
+ variantType.getNewVector("a", FieldType.nullable(variantType), allocator);
+ VariantVector bb =
+ (VariantVector)
+ variantType.getNewVector("a", FieldType.nullable(variantType), allocator)) {
+
+ ValidateVectorVisitor validateVisitor = new ValidateVectorVisitor();
+ validateVisitor.visit(a1, null);
+
+ a1.allocateNew();
+ a2.allocateNew();
+ bb.allocateNew();
+
+ a1.setSafe(0, hello);
+ a1.setSafe(1, world);
+ a1.setValueCount(2);
+
+ a2.setSafe(0, hello);
+ a2.setSafe(1, world);
+ a2.setValueCount(2);
+
+ bb.setSafe(0, world);
+ bb.setSafe(1, hello);
+ bb.setValueCount(2);
+
+ Range range = new Range(0, 0, a1.getValueCount());
+ RangeEqualsVisitor visitor = new RangeEqualsVisitor(a1, a2);
+ assertTrue(visitor.rangeEquals(range));
+
+ visitor = new RangeEqualsVisitor(a1, bb);
+ assertFalse(visitor.rangeEquals(range));
+
+ VectorBatchAppender.batchAppend(a1, a2, bb);
+ assertEquals(6, a1.getValueCount());
+ validateVisitor.visit(a1, null);
+ }
+ }
+
+ @Test
+ public void testVariantCopyAsValueThrowsException() {
+ ensureRegistered(VariantType.INSTANCE);
+ try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+ VariantVector vector = new VariantVector("variant", allocator)) {
+ vector.allocateNew();
+ vector.setSafe(0, TestVariant.variantString("hello"));
+ vector.setValueCount(1);
+
+ var reader = vector.getReader();
+ reader.setPosition(0);
+
+ assertThrows(
+ IllegalArgumentException.class, () -> reader.copyAsValue((BaseWriter.StructWriter) null));
+ }
+ }
+}
diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInListVector.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInListVector.java
new file mode 100644
index 0000000000..8b6000bc46
--- /dev/null
+++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInListVector.java
@@ -0,0 +1,202 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.variant.extension;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.variant.TestVariant;
+import org.apache.arrow.variant.Variant;
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.impl.UnionListReader;
+import org.apache.arrow.vector.complex.impl.UnionListWriter;
+import org.apache.arrow.vector.complex.reader.FieldReader;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.util.TransferPair;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class TestVariantInListVector {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ public void init() {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void terminate() throws Exception {
+ allocator.close();
+ }
+
+ @Test
+ public void testListVectorWithVariantExtensionType() {
+ final FieldType type = FieldType.nullable(VariantType.INSTANCE);
+ try (ListVector inVector = new ListVector("input", allocator, type, null)) {
+ Variant variant1 = TestVariant.variantString("hello");
+ Variant variant2 = TestVariant.variantString("bye");
+
+ UnionListWriter writer = inVector.getWriter();
+ writer.allocate();
+
+ writer.setPosition(0);
+ writer.startList();
+ ExtensionWriter extensionWriter = writer.extension(VariantType.INSTANCE);
+ extensionWriter.writeExtension(variant1);
+ extensionWriter.writeExtension(variant2);
+ writer.endList();
+ inVector.setValueCount(1);
+
+ ArrayList
A "producer" allocator owns the exported batch; if the C Data release callback fires, the
+ * producer drains to zero. A too-small consumer allocator forces an OOM part-way through the
+ * import. The test asserts the producer drains, confirming the release callback fired despite the
+ * failure.
+ */
+final class ImportOutOfMemoryTest {
+ private static final int ROWS = 1024;
+ private static final int VALUE_BYTES = 256;
+ private static final int COLUMNS = 4;
+ // Far smaller than the exported batch, so the import OOMs part-way through the buffers.
+ private static final long TINY_LIMIT = 16 * 1024;
+
+ private RootAllocator root;
+
+ @BeforeEach
+ public void setUp() {
+ root = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ root.close();
+ }
+
+ @Test
+ public void importOomDoesNotLeakExportedArray() {
+ // "producer" owns only the exported batch buffers; the C Data struct containers live on a
+ // separate allocator (they are consumed/closed by import, which would otherwise muddy the
+ // producer's balance). So producer draining to zero is an exact signal that the array's release
+ // callback fired.
+ try (BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE);
+ BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE)) {
+ try (ArrowArray array = ArrowArray.allocateNew(structs);
+ ArrowSchema schema = ArrowSchema.allocateNew(structs)) {
+ exportBatch(producer, array, schema);
+ assertTrue(
+ producer.getAllocatedMemory() > 0, "producer holds the exported batch before import");
+
+ // A consumer allocator far too small to hold the batch: the import throws part-way through.
+ try (BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT);
+ CDataDictionaryProvider provider = new CDataDictionaryProvider()) {
+ Schema importSchema = Data.importSchema(consumer, schema, provider);
+ try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) {
+ Exception thrown =
+ assertThrows(
+ Exception.class,
+ () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider));
+ assertTrue(
+ hasOutOfMemoryCause(thrown),
+ "mid-import failure must be an allocator OOM: " + thrown);
+ }
+ }
+
+ // The array's release callback must have fired despite the mid-import OOM, freeing the
+ // whole exported batch. On the unfixed retain-before-wrap code the batch is stranded.
+ assertEquals(
+ 0L,
+ producer.getAllocatedMemory(),
+ "import OOM leaked the exported batch (producer not drained)");
+ }
+ }
+ }
+
+ /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */
+ private static boolean hasOutOfMemoryCause(Throwable t) {
+ for (Throwable cause = t; cause != null; cause = cause.getCause()) {
+ if (cause instanceof OutOfMemoryException) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Builds a wide multi-column VarChar batch on {@code alloc} and exports it into the C structs.
+ */
+ private void exportBatch(BufferAllocator alloc, ArrowArray array, ArrowSchema schema) {
+ byte[] value = new byte[VALUE_BYTES];
+ for (int i = 0; i < value.length; i++) {
+ value[i] = (byte) 'x';
+ }
+ List This accessor provides JDBC-compatible access to UUID values stored in Arrow's canonical UUID
+ * extension type ('arrow.uuid'). It follows PostgreSQL JDBC driver conventions:
+ *
+ * This provider handles service-to-service authentication where no user interaction is required.
+ * Tokens are cached and automatically refreshed before expiration.
+ */
+public class ClientCredentialsTokenProvider extends AbstractOAuthTokenProvider {
+
+ /**
+ * Creates a new ClientCredentialsTokenProvider.
+ *
+ * @param tokenUri the OAuth token endpoint URI
+ * @param clientId the OAuth client ID
+ * @param clientSecret the OAuth client secret
+ * @param scope optional OAuth scopes (space-separated)
+ */
+ ClientCredentialsTokenProvider(
+ URI tokenUri, String clientId, String clientSecret, @Nullable String scope) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ Objects.requireNonNull(clientId, "clientId cannot be null");
+ Objects.requireNonNull(clientSecret, "clientSecret cannot be null");
+ this.clientAuth = new ClientSecretBasic(new ClientID(clientId), new Secret(clientSecret));
+ this.scope = (scope != null && !scope.isEmpty()) ? Scope.parse(scope) : null;
+ }
+
+ @Override
+ protected TokenRequest buildTokenRequest() {
+ return new TokenRequest(tokenUri, clientAuth, new ClientCredentialsGrant(), scope);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java
new file mode 100644
index 0000000000..cba9d4c2e6
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java
@@ -0,0 +1,240 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.driver.jdbc.client.oauth;
+
+import com.nimbusds.oauth2.sdk.GrantType;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.sql.SQLException;
+import java.util.Locale;
+import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Configuration class for OAuth settings parsed from connection properties. */
+public class OAuthConfiguration {
+
+ private final GrantType grantType;
+ private final URI tokenUri;
+ private final @Nullable String clientId;
+ private final @Nullable String clientSecret;
+ private final @Nullable String scope;
+ private final @Nullable String subjectToken;
+ private final @Nullable String subjectTokenType;
+ private final @Nullable String actorToken;
+ private final @Nullable String actorTokenType;
+ private final @Nullable String audience;
+ private final @Nullable String resource;
+ private final @Nullable String requestedTokenType;
+
+ private OAuthConfiguration(Builder builder) throws SQLException {
+ this.grantType = builder.grantType;
+ this.tokenUri = builder.tokenUri;
+ this.clientId = builder.clientId;
+ this.clientSecret = builder.clientSecret;
+ this.scope = builder.scope;
+ this.subjectToken = builder.subjectToken;
+ this.subjectTokenType = builder.subjectTokenType;
+ this.actorToken = builder.actorToken;
+ this.actorTokenType = builder.actorTokenType;
+ this.audience = builder.audience;
+ this.resource = builder.resource;
+ this.requestedTokenType = builder.requestedTokenType;
+
+ validate();
+ }
+
+ private void validate() throws SQLException {
+ Objects.requireNonNull(grantType, "OAuth grant type is required");
+ Objects.requireNonNull(tokenUri, "Token URI is required");
+
+ if (GrantType.CLIENT_CREDENTIALS.equals(grantType)) {
+ if (clientId == null || clientId.isEmpty()) {
+ throw new SQLException("clientId is required for client_credentials flow");
+ }
+ if (clientSecret == null || clientSecret.isEmpty()) {
+ throw new SQLException("clientSecret is required for client_credentials flow");
+ }
+ } else if (GrantType.TOKEN_EXCHANGE.equals(grantType)) {
+ if (subjectToken == null || subjectToken.isEmpty()) {
+ throw new SQLException("subjectToken is required for token_exchange flow");
+ }
+ if (subjectTokenType == null || subjectTokenType.isEmpty()) {
+ throw new SQLException("subjectTokenType is required for token_exchange flow");
+ }
+ } else {
+ throw new SQLException("Unsupported OAuth grant type: " + grantType);
+ }
+ }
+
+ /**
+ * Creates an OAuthTokenProvider based on the configured grant type.
+ *
+ * @return the token provider
+ * @throws SQLException if the grant type is not supported or configuration is invalid
+ */
+ public OAuthTokenProvider createTokenProvider() throws SQLException {
+ if (GrantType.CLIENT_CREDENTIALS.equals(grantType)) {
+ return OAuthTokenProviders.clientCredentials()
+ .tokenUri(tokenUri)
+ .clientId(clientId)
+ .clientSecret(clientSecret)
+ .scope(scope)
+ .build();
+ } else if (GrantType.TOKEN_EXCHANGE.equals(grantType)) {
+ OAuthTokenProviders.TokenExchangeBuilder builder =
+ OAuthTokenProviders.tokenExchange()
+ .tokenUri(tokenUri)
+ .subjectToken(subjectToken)
+ .subjectTokenType(subjectTokenType)
+ .actorToken(actorToken)
+ .actorTokenType(actorTokenType)
+ .audience(audience)
+ .requestedTokenType(requestedTokenType)
+ .scope(scope)
+ .resource(resource);
+
+ if (clientId != null && clientSecret != null) {
+ builder.clientCredentials(clientId, clientSecret);
+ }
+
+ return builder.build();
+ } else {
+ throw new SQLException("Unsupported OAuth grant type: " + grantType);
+ }
+ }
+
+ /** Builder for OAuthConfiguration. */
+ public static class Builder {
+ private GrantType grantType;
+ private URI tokenUri;
+ private @Nullable String clientId;
+ private @Nullable String clientSecret;
+ private @Nullable String scope;
+ private @Nullable String subjectToken;
+ private @Nullable String subjectTokenType;
+ private @Nullable String actorToken;
+ private @Nullable String actorTokenType;
+ private @Nullable String audience;
+ private @Nullable String resource;
+ private @Nullable String requestedTokenType;
+
+ /**
+ * Sets the OAuth grant type from a string value.
+ *
+ * Accepts either user-friendly names ("client_credentials", "token_exchange") or the full
+ * URN format as defined in RFC 6749 and RFC 8693.
+ *
+ * @param flowStr the flow type string (e.g., "client_credentials", "token_exchange")
+ * @return this builder
+ * @throws SQLException if the flow string is invalid
+ */
+ public Builder flow(String flowStr) throws SQLException {
+ if (flowStr == null || flowStr.isEmpty()) {
+ throw new SQLException("OAuth flow cannot be null or empty");
+ }
+ try {
+ String normalized = flowStr.toLowerCase(Locale.ROOT);
+ // Map user-friendly names to URN format for token_exchange
+ if ("token_exchange".equals(normalized)) {
+ normalized = GrantType.TOKEN_EXCHANGE.getValue();
+ }
+ GrantType parsed = GrantType.parse(normalized);
+ if (!parsed.equals(GrantType.CLIENT_CREDENTIALS)
+ && !parsed.equals(GrantType.TOKEN_EXCHANGE)) {
+ throw new SQLException("Unsupported OAuth flow: " + flowStr);
+ }
+ this.grantType = parsed;
+ } catch (com.nimbusds.oauth2.sdk.ParseException e) {
+ throw new SQLException("Invalid OAuth flow: " + flowStr, e);
+ }
+ return this;
+ }
+
+ /**
+ * Sets the token URI.
+ *
+ * @param tokenUri the OAuth token endpoint URI
+ * @return this builder
+ * @throws SQLException if the URI is invalid
+ */
+ public Builder tokenUri(String tokenUri) throws SQLException {
+ if (tokenUri == null || tokenUri.isEmpty()) {
+ throw new SQLException("Token URI cannot be null or empty");
+ }
+ try {
+ this.tokenUri = new URI(tokenUri);
+ } catch (URISyntaxException e) {
+ throw new SQLException("Invalid token URI: " + tokenUri, e);
+ }
+ return this;
+ }
+
+ public Builder clientId(@Nullable String clientId) {
+ this.clientId = clientId;
+ return this;
+ }
+
+ public Builder clientSecret(@Nullable String clientSecret) {
+ this.clientSecret = clientSecret;
+ return this;
+ }
+
+ public Builder scope(@Nullable String scope) {
+ this.scope = scope;
+ return this;
+ }
+
+ public Builder subjectToken(@Nullable String subjectToken) {
+ this.subjectToken = subjectToken;
+ return this;
+ }
+
+ public Builder subjectTokenType(@Nullable String subjectTokenType) {
+ this.subjectTokenType = subjectTokenType;
+ return this;
+ }
+
+ public Builder actorToken(@Nullable String actorToken) {
+ this.actorToken = actorToken;
+ return this;
+ }
+
+ public Builder actorTokenType(@Nullable String actorTokenType) {
+ this.actorTokenType = actorTokenType;
+ return this;
+ }
+
+ public Builder audience(@Nullable String audience) {
+ this.audience = audience;
+ return this;
+ }
+
+ public Builder resource(@Nullable String resource) {
+ this.resource = resource;
+ return this;
+ }
+
+ public Builder requestedTokenType(@Nullable String requestedTokenType) {
+ this.requestedTokenType = requestedTokenType;
+ return this;
+ }
+
+ public OAuthConfiguration build() throws SQLException {
+ return new OAuthConfiguration(this);
+ }
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java
new file mode 100644
index 0000000000..0d4ad4689f
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.driver.jdbc.client.oauth;
+
+import java.sql.SQLException;
+import java.util.Objects;
+import java.util.function.Consumer;
+import org.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.auth2.Auth2Constants;
+
+/** Writes OAuth bearer tokens to Flight call headers. */
+public class OAuthCredentialWriter implements Consumer This class provides a single entry point for creating all OAuth token providers with a
+ * consistent builder API. It supports:
+ *
+ * Example usage:
+ *
+ * This provider exchanges one token for another, commonly used for federated authentication,
+ * delegation, or impersonation scenarios. Tokens are cached and automatically refreshed.
+ */
+public class TokenExchangeTokenProvider extends AbstractOAuthTokenProvider {
+
+ @VisibleForTesting TokenExchangeGrant grant;
+
+ @VisibleForTesting @Nullable List Handles conversion of UUID values from JDBC parameters to Arrow's UUID extension type. Accepts
+ * both {@link UUID} objects and String representations of UUIDs.
+ */
+public class UuidAvaticaParameterConverter implements AvaticaParameterConverter {
+
+ public UuidAvaticaParameterConverter() {}
+
+ @Override
+ public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
+ if (!(vector instanceof UuidVector)) {
+ return false;
+ }
+
+ UuidVector uuidVector = (UuidVector) vector;
+ Object value = typedValue.toJdbc(null);
+
+ if (value == null) {
+ uuidVector.setNull(index);
+ return true;
+ }
+
+ UUID uuid;
+ if (value instanceof UUID) {
+ uuid = (UUID) value;
+ } else if (value instanceof String) {
+ uuid = UUID.fromString((String) value);
+ } else if (value instanceof byte[]) {
+ byte[] bytes = (byte[]) value;
+ if (bytes.length != 16) {
+ throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length);
+ }
+ uuid = uuidFromBytes(bytes);
+ } else if (value instanceof ByteString) {
+ byte[] bytes = ((ByteString) value).getBytes();
+ if (bytes.length != 16) {
+ throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length);
+ }
+ uuid = uuidFromBytes(bytes);
+ } else {
+ throw new IllegalArgumentException(
+ "Cannot convert " + value.getClass().getName() + " to UUID");
+ }
+
+ uuidVector.setSafe(index, UuidUtility.getBytesFromUUID(uuid));
+ return true;
+ }
+
+ @Override
+ public AvaticaParameter createParameter(Field field) {
+ final String name = field.getName();
+ final int jdbcType = getSqlTypeIdFromArrowType(field.getType());
+ final String typeName = getSqlTypeNameFromArrowType(field.getType());
+ final String className = UUID.class.getCanonicalName();
+ return new AvaticaParameter(false, 0, 0, jdbcType, typeName, className, name);
+ }
+
+ private static UUID uuidFromBytes(byte[] bytes) {
+ final long mostSignificantBits;
+ final long leastSignificantBits;
+ ByteBuffer bb = ByteBuffer.wrap(bytes);
+ // Reads the first eight bytes
+ mostSignificantBits = bb.getLong();
+ // Reads the first eight bytes at this buffer's current
+ leastSignificantBits = bb.getLong();
+
+ return new UUID(mostSignificantBits, leastSignificantBits);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java
index 76ba964a53..d0ba74dbcc 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.driver.jdbc.utils;
+import java.sql.SQLException;
import java.time.Duration;
import java.util.Arrays;
import java.util.HashMap;
@@ -23,6 +24,7 @@
import java.util.Objects;
import java.util.Properties;
import org.apache.arrow.driver.jdbc.ArrowFlightConnection;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthConfiguration;
import org.apache.arrow.flight.CallHeaders;
import org.apache.arrow.flight.CallOption;
import org.apache.arrow.flight.FlightCallHeaders;
@@ -31,6 +33,7 @@
import org.apache.calcite.avatica.ConnectionConfig;
import org.apache.calcite.avatica.ConnectionConfigImpl;
import org.apache.calcite.avatica.ConnectionProperty;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** A {@link ConnectionConfig} for the {@link ArrowFlightConnection}. */
public final class ArrowFlightConnectionConfigImpl extends ConnectionConfigImpl {
@@ -211,6 +214,38 @@ public Map These tests verify that OAuth tokens obtained from an OAuth server are correctly used in
+ * Flight SQL requests.
+ */
+public class OAuthIntegrationTest {
+
+ private static final String VALID_ACCESS_TOKEN = "valid-oauth-access-token-12345";
+ private static final String CLIENT_ID = "test-client-id";
+ private static final String CLIENT_SECRET = "test-client-secret";
+ private static final String SUBJECT_TOKEN = "original-subject-token";
+ private static final String SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt";
+ private static final String TEST_SCOPE = "dremio.all";
+
+ private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer();
+
+ @RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION;
+
+ static {
+ FLIGHT_SERVER_TEST_EXTENSION =
+ new FlightServerTestExtension.Builder()
+ .authentication(new TokenAuthentication.Builder().token(VALID_ACCESS_TOKEN).build())
+ .producer(FLIGHT_SQL_PRODUCER)
+ .build();
+ }
+
+ @StartStop private final MockWebServer oauthServer = new MockWebServer();
+ private URI tokenEndpoint;
+
+ @BeforeAll
+ public static void setUpClass() {
+ // Register a simple catalog query handler
+ FLIGHT_SQL_PRODUCER.addCatalogQuery(
+ CommandGetCatalogs.getDefaultInstance(),
+ listener -> {
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(Schemas.GET_CATALOGS_SCHEMA, allocator)) {
+ root.setRowCount(0);
+ listener.start(root);
+ listener.putNext();
+ } catch (Throwable t) {
+ listener.error(t);
+ } finally {
+ listener.completed();
+ }
+ });
+
+ // Register a simple schema query handler for getSchemas()
+ FLIGHT_SQL_PRODUCER.addCatalogQuery(
+ CommandGetDbSchemas.getDefaultInstance(),
+ listener -> {
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(Schemas.GET_SCHEMAS_SCHEMA, allocator)) {
+ root.setRowCount(0);
+ listener.start(root);
+ listener.putNext();
+ } catch (Throwable t) {
+ listener.error(t);
+ } finally {
+ listener.completed();
+ }
+ });
+ }
+
+ @AfterAll
+ public static void tearDownClass() {
+ AutoCloseables.closeNoChecked(FLIGHT_SQL_PRODUCER);
+ }
+
+ @BeforeEach
+ public void setUp() {
+ tokenEndpoint = oauthServer.url("/oauth/token").uri();
+ }
+
+ @AfterEach
+ public void tearDown() {
+ oauthServer.close();
+ }
+
+ // Helper methods for mock OAuth responses
+
+ private void enqueueSuccessfulTokenResponse() {
+ enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 3600);
+ }
+
+ private void enqueueSuccessfulTokenResponse(String token, int expiresIn) {
+ String body =
+ String.format(
+ "{\"access_token\":\"%s\",\"token_type\":\"Bearer\",\"expires_in\":%d}",
+ token, expiresIn);
+ oauthServer.enqueue(
+ new MockResponse.Builder()
+ .code(200)
+ .setHeader("Content-Type", "application/json")
+ .body(body)
+ .build());
+ }
+
+ private void enqueueErrorResponse(String error, String description) {
+ String body =
+ String.format("{\"error\":\"%s\",\"error_description\":\"%s\"}", error, description);
+ oauthServer.enqueue(
+ new MockResponse.Builder()
+ .code(400)
+ .setHeader("Content-Type", "application/json")
+ .body(body)
+ .build());
+ }
+
+ private Properties createBaseProperties() {
+ Properties props = new Properties();
+ props.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost");
+ props.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ props.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false);
+ return props;
+ }
+
+ private String getJdbcUrl() {
+ return String.format(
+ "jdbc:arrow-flight-sql://localhost:%d", FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ }
+
+ // ==================== Client Credentials Flow Tests ====================
+
+ @Test
+ public void testClientCredentialsFlowSuccess() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+ props.put(ArrowFlightConnectionProperty.OAUTH_SCOPE.camelName(), TEST_SCOPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertFalse(conn.isClosed());
+ // Trigger a Flight call to force OAuth token retrieval
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ // Verify OAuth request was made
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ assertEquals("POST", oauthRequest.getMethod());
+ String body = oauthRequest.getBody().utf8();
+ assertTrue(body.contains("grant_type=client_credentials"));
+ assertTrue(body.contains("scope=" + TEST_SCOPE));
+ }
+
+ @Test
+ public void testClientCredentialsFlowWithUrlParameters() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ String url =
+ String.format(
+ "jdbc:arrow-flight-sql://localhost:%d?useEncryption=false"
+ + "&oauth.flow=client_credentials"
+ + "&oauth.tokenUri=%s"
+ + "&oauth.clientId=%s"
+ + "&oauth.clientSecret=%s",
+ FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ tokenEndpoint.toString(),
+ CLIENT_ID,
+ CLIENT_SECRET);
+
+ try (Connection conn = DriverManager.getConnection(url)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ assertEquals(1, oauthServer.getRequestCount());
+ }
+
+ @Test
+ public void testClientCredentialsFlowInvalidCredentials() throws Exception {
+ enqueueErrorResponse("invalid_client", "Client authentication failed");
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), "wrong-client");
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), "wrong-secret");
+
+ Exception ex =
+ assertThrows(
+ Exception.class,
+ () -> {
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+ });
+ // Verify the error message contains the OAuth error somewhere in the exception chain
+ assertTrue(
+ containsInExceptionChain(ex, "invalid_client"),
+ "Exception chain should contain 'invalid_client'");
+ }
+
+ private boolean containsInExceptionChain(Throwable t, String message) {
+ while (t != null) {
+ if (t.getMessage() != null && t.getMessage().contains(message)) {
+ return true;
+ }
+ t = t.getCause();
+ }
+ return false;
+ }
+
+ // ==================== Token Exchange Flow Tests ====================
+
+ @Test
+ public void testTokenExchangeFlowMinimalParameters() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ String body = oauthRequest.getBody().utf8();
+ assertTrue(
+ body.contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange"),
+ "Should contain token exchange grant type");
+ assertTrue(body.contains("subject_token=" + SUBJECT_TOKEN));
+ }
+
+ @Test
+ public void testTokenExchangeFlowWithAllParameters() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ String actorToken = "actor-token-value";
+ String actorTokenType = "urn:ietf:params:oauth:token-type:access_token";
+ String audience = "https://api.example.com";
+ String resource = "https://api.example.com/resource";
+ String requestedTokenType = "urn:ietf:params:oauth:token-type:access_token";
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+ props.put(ArrowFlightConnectionProperty.OAUTH_SCOPE.camelName(), TEST_SCOPE);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+ props.put(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN.camelName(), actorToken);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE.camelName(), actorTokenType);
+ props.put(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_AUDIENCE.camelName(), audience);
+ props.put(ArrowFlightConnectionProperty.OAUTH_RESOURCE.camelName(), resource);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.camelName(),
+ requestedTokenType);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ String body = oauthRequest.getBody().utf8();
+ assertTrue(body.contains("subject_token=" + SUBJECT_TOKEN));
+ assertTrue(body.contains("actor_token=" + actorToken));
+ }
+
+ @Test
+ public void testTokenExchangeFlowWithClientAuthentication() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS);
+ assertNotNull(oauthRequest, "OAuth request should have been made");
+ String authHeader = oauthRequest.getHeaders().get("Authorization");
+ assertNotNull(authHeader, "Should have Basic auth header for client authentication");
+ assertTrue(authHeader.startsWith("Basic "));
+ }
+
+ // ==================== Token Caching Tests ====================
+
+ @Test
+ public void testTokenCachingAcrossMultipleOperations() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ // Execute multiple operations
+ conn.isValid(5);
+ conn.getMetaData().getCatalogs().close();
+ conn.getMetaData().getSchemas().close();
+ }
+
+ // Should only have made one OAuth request due to caching
+ assertEquals(1, oauthServer.getRequestCount());
+ }
+
+ @Test
+ public void testTokenRefreshAfterExpiration() throws Exception {
+ enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 1);
+ enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 3600);
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN);
+ props.put(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(),
+ SUBJECT_TOKEN_TYPE);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ // First operation triggers initial token fetch
+ conn.getMetaData().getCatalogs().close();
+
+ // Token with 1s expiry is immediately considered expired (due to 30s buffer)
+ // so the next operation should trigger a refresh
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ // Should have made exactly 2 OAuth requests: initial + refresh
+ assertEquals(2, oauthServer.getRequestCount());
+ }
+
+ // ==================== Error Handling Tests ====================
+
+ @Test
+ public void testMissingRequiredParametersClientCredentials() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ // Missing client_id and client_secret
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ @Test
+ public void testMissingRequiredParametersTokenExchange() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ // Missing subject_token and subject_token_type
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ @Test
+ public void testInvalidOAuthFlow() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "invalid_flow");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ @Test
+ public void testMalformedTokenEndpoint() {
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), "not-a-valid-uri://");
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+
+ assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props));
+ }
+
+ // ==================== Authorization Header Verification ====================
+
+ @Test
+ public void testOAuthTokenSentAsBearer() throws Exception {
+ enqueueSuccessfulTokenResponse();
+
+ Properties props = createBaseProperties();
+ props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials");
+ props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString());
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID);
+ props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET);
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ conn.getMetaData().getCatalogs().close();
+ }
+
+ // Verify the Flight server received the bearer token
+ String authHeader =
+ FLIGHT_SERVER_TEST_EXTENSION
+ .getInterceptorFactory()
+ .getHeader(org.apache.arrow.flight.FlightMethod.GET_FLIGHT_INFO, "authorization");
+ assertNotNull(authHeader, "Authorization header should be present in Flight requests");
+ assertEquals("Bearer " + VALID_ACCESS_TOKEN, authHeader);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java
index 569b5495fe..3a5a39be3d 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java
@@ -22,8 +22,10 @@
import static org.hamcrest.CoreMatchers.allOf;
import static org.hamcrest.CoreMatchers.anyOf;
import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.*;
@@ -31,7 +33,9 @@
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
+import java.sql.PreparedStatement;
import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.SQLTimeoutException;
import java.sql.Statement;
@@ -42,6 +46,7 @@
import java.util.List;
import java.util.Random;
import java.util.Set;
+import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers;
import org.apache.arrow.driver.jdbc.utils.FallbackFlightSqlProducer;
@@ -61,6 +66,7 @@
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -795,4 +801,174 @@ public void testResultSetAppMetadata() throws Exception {
"foo".getBytes(StandardCharsets.UTF_8));
}
}
+
+ @Test
+ public void testSelectQueryWithUuidColumn() throws SQLException {
+ // Expectations
+ final int expectedRowCount = 4;
+ final UUID[] expectedUuids =
+ new UUID[] {
+ CoreMockedSqlProducers.UUID_1,
+ CoreMockedSqlProducers.UUID_2,
+ CoreMockedSqlProducers.UUID_3,
+ null
+ };
+
+ final Integer[] expectedIds = new Integer[] {1, 2, 3, 4};
+
+ final List Verifies that the accessor correctly handles UUID values from Arrow's UUID extension type,
+ * following PostgreSQL JDBC driver conventions.
+ */
+public class ArrowFlightJdbcUuidVectorAccessorTest {
+
+ @RegisterExtension
+ public static RootAllocatorTestExtension rootAllocatorTestExtension =
+ new RootAllocatorTestExtension();
+
+ private static final UUID UUID_1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000");
+ private static final UUID UUID_2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8");
+ private static final UUID UUID_3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479");
+
+ private UuidVector vector;
+ private ArrowFlightJdbcUuidVectorAccessor accessor;
+ private boolean wasNullCalled;
+ private boolean wasNullValue;
+
+ @BeforeEach
+ public void setUp() {
+ vector = rootAllocatorTestExtension.createUuidVector();
+ wasNullCalled = false;
+ wasNullValue = false;
+ ArrowFlightJdbcAccessorFactory.WasNullConsumer wasNullConsumer =
+ (wasNull) -> {
+ wasNullCalled = true;
+ wasNullValue = wasNull;
+ };
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, wasNullConsumer);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ vector.close();
+ }
+
+ @Test
+ public void testGetObjectReturnsUuid() {
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ Object result = accessor.getObject();
+ assertThat(result, is(UUID_1));
+ assertThat(accessor.wasNull(), is(false));
+ }
+
+ @Test
+ public void testGetObjectReturnsCorrectUuidForEachRow() {
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ assertThat(accessor.getObject(), is(UUID_1));
+
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 1, (wasNull) -> {});
+ assertThat(accessor.getObject(), is(UUID_2));
+
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 2, (wasNull) -> {});
+ assertThat(accessor.getObject(), is(UUID_3));
+ }
+
+ @Test
+ public void testGetObjectReturnsNullForNullValue() {
+ vector.reset();
+ vector.allocateNew(1);
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ Object result = accessor.getObject();
+ assertThat(result, nullValue());
+ assertThat(accessor.wasNull(), is(true));
+ }
+
+ @Test
+ public void testGetObjectClassReturnsUuidClass() {
+ assertThat(accessor.getObjectClass(), equalTo(UUID.class));
+ }
+
+ @Test
+ public void testGetStringReturnsHyphenatedFormat() {
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ String result = accessor.getString();
+ assertThat(result, is("550e8400-e29b-41d4-a716-446655440000"));
+ assertThat(accessor.wasNull(), is(false));
+ }
+
+ @Test
+ public void testGetStringReturnsNullForNullValue() {
+ vector.reset();
+ vector.allocateNew(1);
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ String result = accessor.getString();
+ assertThat(result, nullValue());
+ assertThat(accessor.wasNull(), is(true));
+ }
+
+ @Test
+ public void testGetBytesReturns16ByteArray() {
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ byte[] result = accessor.getBytes();
+ assertThat(result.length, is(16));
+ assertThat(result, is(UuidUtility.getBytesFromUUID(UUID_1)));
+ assertThat(accessor.wasNull(), is(false));
+ }
+
+ @Test
+ public void testGetBytesReturnsNullForNullValue() {
+ vector.reset();
+ vector.allocateNew(1);
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {});
+ byte[] result = accessor.getBytes();
+ assertThat(result, nullValue());
+ assertThat(accessor.wasNull(), is(true));
+ }
+
+ @Test
+ public void testWasNullConsumerIsCalled() {
+ accessor =
+ new ArrowFlightJdbcUuidVectorAccessor(
+ vector,
+ () -> 0,
+ (wasNull) -> {
+ wasNullCalled = true;
+ wasNullValue = wasNull;
+ });
+ accessor.getObject();
+ assertThat(wasNullCalled, is(true));
+ assertThat(wasNullValue, is(false));
+ }
+
+ @Test
+ public void testWasNullConsumerIsCalledWithTrueForNull() {
+ vector.reset();
+ vector.allocateNew(1);
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ accessor =
+ new ArrowFlightJdbcUuidVectorAccessor(
+ vector,
+ () -> 0,
+ (wasNull) -> {
+ wasNullCalled = true;
+ wasNullValue = wasNull;
+ });
+ accessor.getObject();
+ assertThat(wasNullCalled, is(true));
+ assertThat(wasNullValue, is(true));
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java
index ad689837e2..c5eb6e34ef 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java
@@ -191,7 +191,12 @@ public void testShouldGetArrayGetResultSetReturnValidResultSet(
try (ResultSet rs = array.getResultSet()) {
int count = 0;
while (rs.next()) {
- final int value = rs.getInt(1);
+ // Column 1: 1-based index (per JDBC spec)
+ final int index = rs.getInt(1);
+ assertThat(index, equalTo(count + 1));
+
+ // Column 2: actual value (per JDBC spec)
+ final int value = rs.getInt(2);
assertThat(value, equalTo(currentRow * count));
count++;
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java
index 696e5afb71..f2d1725fd8 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java
@@ -153,15 +153,15 @@ public void testShouldGetArrayReturnValidArray() throws SQLException {
try (ResultSet resultSet = array.getResultSet()) {
assertTrue(resultSet.next());
- Map entry = resultSet.getObject(1, Map.class);
+ Map entry = resultSet.getObject(2, Map.class);
assertEquals(1, entry.get("key"));
assertEquals(11, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(2, entry.get("key"));
assertEquals(22, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(3, entry.get("key"));
assertEquals(33, entry.get("value"));
assertFalse(resultSet.next());
@@ -173,7 +173,7 @@ public void testShouldGetArrayReturnValidArray() throws SQLException {
assertFalse(accessor.wasNull());
try (ResultSet resultSet = array.getResultSet()) {
assertTrue(resultSet.next());
- Map entry = resultSet.getObject(1, Map.class);
+ Map entry = resultSet.getObject(2, Map.class);
assertEquals(2, entry.get("key"));
assertNull(entry.get("value"));
assertFalse(resultSet.next());
@@ -185,19 +185,19 @@ public void testShouldGetArrayReturnValidArray() throws SQLException {
assertFalse(accessor.wasNull());
try (ResultSet resultSet = array.getResultSet()) {
assertTrue(resultSet.next());
- Map entry = resultSet.getObject(1, Map.class);
+ Map entry = resultSet.getObject(2, Map.class);
assertEquals(0, entry.get("key"));
assertEquals(2000, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(1, entry.get("key"));
assertEquals(2001, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(2, entry.get("key"));
assertEquals(2002, entry.get("value"));
assertTrue(resultSet.next());
- entry = resultSet.getObject(1, Map.class);
+ entry = resultSet.getObject(2, Map.class);
assertEquals(3, entry.get("key"));
assertEquals(2003, entry.get("value"));
assertFalse(resultSet.next());
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java
index a2f6fd586f..82876f4aa1 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java
@@ -24,6 +24,8 @@
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.when;
@@ -46,6 +48,8 @@
import org.apache.arrow.vector.DateMilliVector;
import org.apache.arrow.vector.TimeMilliVector;
import org.apache.arrow.vector.TimeStampVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarCharVector;
import org.apache.arrow.vector.util.Text;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -695,4 +699,26 @@ public void testShouldGetObjectClassReturnString() {
final Class clazz = accessor.getObjectClass();
assertThat(clazz, equalTo(String.class));
}
+
+ @Test
+ public void testViewVarcharVector() throws Exception {
+ try (VarCharVector varCharVector =
+ new VarCharVector("", rootAllocatorTestExtension.getRootAllocator());
+ ViewVarCharVector viewVarCharVector =
+ new ViewVarCharVector("", rootAllocatorTestExtension.getRootAllocator())) {
+ varCharVector.allocateNew(1);
+ viewVarCharVector.allocateNew(1);
+
+ ArrowFlightJdbcVarCharVectorAccessor varCharVectorAccessor =
+ new ArrowFlightJdbcVarCharVectorAccessor(varCharVector, () -> 0, (boolean wasNull) -> {});
+ ArrowFlightJdbcVarCharVectorAccessor viewVarcharVectorAccessor =
+ new ArrowFlightJdbcVarCharVectorAccessor(
+ viewVarCharVector, () -> 0, (boolean wasNull) -> {});
+ assertNull(viewVarcharVectorAccessor.getString());
+
+ varCharVector.set(0, new Text("looooong_string"));
+ viewVarCharVector.set(0, new Text("looooong_string"));
+ assertEquals(varCharVectorAccessor.getString(), viewVarcharVectorAccessor.getString());
+ }
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java
new file mode 100644
index 0000000000..d5973ab5d8
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.driver.jdbc.client;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Optional;
+import org.apache.arrow.flight.CallOption;
+import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.CloseSessionRequest;
+import org.apache.arrow.flight.FlightStatusCode;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+public class ArrowFlightSqlClientHandlerTest {
+
+ @ParameterizedTest
+ @MethodSource
+ public void testCloseHandlesFlightRuntimeException(
+ boolean throwFromCloseSession, CallStatus callStatus, boolean shouldSuppress)
+ throws Exception {
+ FlightSqlClient sqlClient = mock(FlightSqlClient.class);
+ String cacheKey = "cacheKey";
+ Optional Verifies that the converter correctly handles UUID parameter binding from JDBC to Arrow's UUID
+ * extension type.
+ */
+public class UuidAvaticaParameterConverterTest {
+
+ @RegisterExtension
+ public static RootAllocatorTestExtension rootAllocatorTestExtension =
+ new RootAllocatorTestExtension();
+
+ private static final UUID TEST_UUID = UUID.fromString("550e8400-e29b-41d4-a716-446655440000");
+
+ private UuidVector vector;
+ private UuidAvaticaParameterConverter converter;
+
+ @BeforeEach
+ public void setUp() {
+ vector = new UuidVector("uuid_param", rootAllocatorTestExtension.getRootAllocator());
+ vector.allocateNew(5);
+ converter = new UuidAvaticaParameterConverter();
+ }
+
+ @AfterEach
+ public void tearDown() {
+ vector.close();
+ }
+
+ @Test
+ public void testBindParameterWithUuidObject() {
+ TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, TEST_UUID);
+
+ boolean result = converter.bindParameter(vector, typedValue, 0);
+
+ assertTrue(result);
+ assertThat(vector.getObject(0), is(TEST_UUID));
+ }
+
+ @Test
+ public void testBindParameterWithUuidString() {
+ String uuidString = "550e8400-e29b-41d4-a716-446655440000";
+ TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.STRING, uuidString);
+
+ boolean result = converter.bindParameter(vector, typedValue, 0);
+
+ assertTrue(result);
+ assertThat(vector.getObject(0), is(TEST_UUID));
+ }
+
+ @Test
+ public void testBindParameterWithByteArray() {
+ byte[] uuidBytes = UuidUtility.getBytesFromUUID(TEST_UUID);
+ ByteString byteString = new ByteString(uuidBytes);
+ TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.BYTE_STRING, byteString);
+
+ boolean result = converter.bindParameter(vector, typedValue, 0);
+
+ assertTrue(result);
+ assertThat(vector.getObject(0), is(TEST_UUID));
+ }
+
+ @Test
+ public void testBindParameterWithNullValue() {
+ TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, null);
+
+ boolean result = converter.bindParameter(vector, typedValue, 0);
+
+ assertTrue(result);
+ assertTrue(vector.isNull(0));
+ assertThat(vector.getObject(0), nullValue());
+ }
+
+ @Test
+ public void testBindParameterWithInvalidByteArrayLength() {
+ byte[] invalidBytes = new byte[8]; // Should be 16 bytes
+ ByteString byteString = new ByteString(invalidBytes);
+ TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.BYTE_STRING, byteString);
+
+ assertThrows(
+ IllegalArgumentException.class, () -> converter.bindParameter(vector, typedValue, 0));
+ }
+
+ @Test
+ public void testBindParameterWithInvalidType() {
+ TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.INTEGER, 12345);
+
+ assertThrows(
+ IllegalArgumentException.class, () -> converter.bindParameter(vector, typedValue, 0));
+ }
+
+ @Test
+ public void testBindParameterMultipleValues() {
+ UUID uuid1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000");
+ UUID uuid2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8");
+ UUID uuid3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479");
+
+ converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid1), 0);
+ converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid2), 1);
+ converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid3), 2);
+
+ assertThat(vector.getObject(0), is(uuid1));
+ assertThat(vector.getObject(1), is(uuid2));
+ assertThat(vector.getObject(2), is(uuid3));
+ }
+
+ @Test
+ public void testCreateParameter() {
+ Field uuidField = new Field("uuid_col", new FieldType(true, UuidType.INSTANCE, null), null);
+
+ AvaticaParameter parameter = converter.createParameter(uuidField);
+
+ assertThat(parameter.name, is("uuid_col"));
+ assertThat(parameter.parameterType, is(Types.OTHER));
+ assertThat(parameter.typeName, is("OTHER"));
+ assertThat(parameter.className, equalTo(UUID.class.getCanonicalName()));
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java
index b6fdc99694..f128ca7c73 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java
@@ -69,6 +69,19 @@ public void testShouldConvertArrowFieldsToColumnMetaDataList() {
.tableName("table1")
.build()
.getMetadataMap()),
+ null),
+ new Field(
+ "col2",
+ new FieldType(
+ true,
+ ArrowType.Utf8View.INSTANCE,
+ null,
+ new FlightSqlColumnMetadata.Builder()
+ .catalogName("catalog1")
+ .schemaName("schema1")
+ .tableName("table1")
+ .build()
+ .getMetadataMap()),
null));
final List Returns {@link System#identityHashCode(Object)} which provides a unique identifier for this
+ * buffer without any per-instance memory overhead.
*
- * @return integer id
+ * @return the identity hash code for this buffer
*/
public long getId() {
- return id;
+ return System.identityHashCode(this);
}
/**
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
index 4f9d3c61c6..dbd6da3291 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
@@ -25,9 +25,10 @@
public interface BufferAllocator extends AutoCloseable {
/**
- * Allocate a new or reused buffer of the provided size. Note that the buffer may technically be
- * larger than the requested size for rounding purposes. However, the buffer's capacity will be
- * set to the configured size.
+ * Allocate a new or reused buffer of the provided size. The buffer may be larger than the
+ * requested size for rounding purposes (e.g. to a power of two), and the buffer's capacity will
+ * reflect the actual allocated size. Use {@link ArrowBuf#capacity(long)} to set the capacity to
+ * the requested size if needed.
*
* @param size The size in bytes.
* @return a new ArrowBuf, or null if the request can't be satisfied
@@ -36,9 +37,10 @@ public interface BufferAllocator extends AutoCloseable {
ArrowBuf buffer(long size);
/**
- * Allocate a new or reused buffer of the provided size. Note that the buffer may technically be
- * larger than the requested size for rounding purposes. However, the buffer's capacity will be
- * set to the configured size.
+ * Allocate a new or reused buffer of the provided size. The buffer may be larger than the
+ * requested size for rounding purposes (e.g. to a power of two), and the buffer's capacity will
+ * reflect the actual allocated size. Use {@link ArrowBuf#capacity(long)} to set the capacity to
+ * the requested size if needed.
*
* @param size The size in bytes.
* @param manager A buffer manager to manage reallocation.
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
index b562a421e7..eb90efcbb5 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
@@ -17,8 +17,7 @@
package org.apache.arrow.memory;
import java.util.IdentityHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import org.apache.arrow.memory.util.CommonUtil;
import org.apache.arrow.memory.util.HistoricalLog;
import org.apache.arrow.util.Preconditions;
@@ -32,12 +31,13 @@
public class BufferLedger implements ValueWithKeyIncluded This benchmark measures the heap memory overhead of creating many ArrowBuf instances. The
+ * optimizations using AtomicFieldUpdater instead of AtomicLong/AtomicInteger objects should reduce
+ * memory overhead significantly.
+ *
+ * Expected savings per instance: - ArrowBuf: 8 bytes (id field removed) - BufferLedger: 28 bytes
+ * (20 from AtomicInteger + 8 from ledgerId) - Accountant: 48 bytes (3 × 16 bytes from AtomicLong
+ * objects)
+ *
+ * For 1M ArrowBuf instances, this should save approximately 8 MB of heap memory.
+ */
+@State(Scope.Benchmark)
+@Fork(
+ value = 1,
+ jvmArgs = {"-Xms2g", "-Xmx2g"})
+@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+public class MemoryFootprintBenchmarks {
+
+ /** Number of ArrowBuf instances to create for memory footprint measurement. */
+ private static final int NUM_BUFFERS = 100_000;
+
+ /** Size in bytes of each buffer allocation. */
+ private static final int BUFFER_SIZE = 1024;
+
+ /** Root allocator used for all buffer allocations in the benchmark. */
+ private RootAllocator allocator;
+
+ /** Array to hold references to allocated buffers, preventing garbage collection. */
+ private ArrowBuf[] buffers;
+
+ /** JMX bean for querying heap memory usage statistics. */
+ private MemoryMXBean memoryBean;
+
+ /**
+ * Sets up the benchmark state before each trial.
+ *
+ * Initializes the memory monitoring bean, creates a root allocator with sufficient capacity,
+ * and allocates the buffer reference array.
+ */
+ @Setup(Level.Trial)
+ public void setup() {
+ memoryBean = ManagementFactory.getMemoryMXBean();
+ allocator = new RootAllocator((long) NUM_BUFFERS * BUFFER_SIZE);
+ buffers = new ArrowBuf[NUM_BUFFERS];
+ }
+
+ /**
+ * Cleans up buffers after each benchmark invocation.
+ *
+ * Closes all allocated buffers to prevent memory leaks and ensure each iteration starts with a
+ * clean slate. This is critical for the memory footprint benchmark which allocates many buffers
+ * that would otherwise accumulate across warmup and measurement iterations.
+ */
+ @TearDown(Level.Invocation)
+ public void tearDown() {
+ for (int i = 0; i < NUM_BUFFERS; i++) {
+ if (buffers[i] != null) {
+ buffers[i].close();
+ buffers[i] = null;
+ }
+ }
+ }
+
+ /**
+ * Cleans up the allocator after the trial completes.
+ *
+ * Closes the root allocator to release all resources after all warmup and measurement
+ * iterations are complete.
+ */
+ @TearDown(Level.Trial)
+ public void tearDownTrial() {
+ allocator.close();
+ }
+
+ /**
+ * Benchmark that measures heap memory usage when creating many ArrowBuf instances.
+ *
+ * This benchmark creates {@value #NUM_BUFFERS} ArrowBuf instances and measures the heap memory
+ * used. With the AtomicFieldUpdater optimizations, we expect to save approximately 800 KB of heap
+ * memory (8 bytes × 100,000 instances) just from removing the id field in ArrowBuf.
+ *
+ * The benchmark performs garbage collection before and after allocation to ensure accurate
+ * measurement of heap memory delta. Results are printed to stdout for analysis.
+ *
+ * @return the total heap memory used by the allocated buffers in bytes
+ */
+ @Benchmark
+ @BenchmarkMode(Mode.SingleShotTime)
+ @OutputTimeUnit(TimeUnit.MILLISECONDS)
+ public long measureArrowBufMemoryFootprint() {
+ // Force GC before measurement
+ System.gc();
+ System.gc();
+ System.gc();
+
+ MemoryUsage heapBefore = memoryBean.getHeapMemoryUsage();
+ long usedBefore = heapBefore.getUsed();
+
+ // Allocate buffers
+ for (int i = 0; i < NUM_BUFFERS; i++) {
+ buffers[i] = allocator.buffer(BUFFER_SIZE);
+ }
+
+ // Force GC to get accurate measurement
+ System.gc();
+ System.gc();
+ System.gc();
+
+ MemoryUsage heapAfter = memoryBean.getHeapMemoryUsage();
+ long usedAfter = heapAfter.getUsed();
+
+ long memoryUsed = usedAfter - usedBefore;
+
+ // Print memory usage for analysis
+ System.out.printf(
+ "Created %d ArrowBuf instances. Heap memory used: %d bytes (%.2f MB)%n",
+ NUM_BUFFERS, memoryUsed, memoryUsed / (1024.0 * 1024.0));
+ System.out.printf(
+ "Average memory per ArrowBuf: %.2f bytes%n", (double) memoryUsed / NUM_BUFFERS);
+
+ return memoryUsed;
+ }
+
+ /**
+ * Benchmark that measures allocation and deallocation performance.
+ *
+ * This complements the memory footprint benchmark by measuring the time it takes to allocate
+ * and deallocate 1,000 buffers in a tight loop. This helps identify any performance regressions
+ * introduced by memory optimizations.
+ *
+ * Uses a local buffer array to avoid interference with the shared {@link #buffers} array used
+ * by other benchmarks.
+ */
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void measureAllocationPerformance() {
+ ArrowBuf[] localBuffers = new ArrowBuf[1000];
+
+ for (int i = 0; i < 1000; i++) {
+ localBuffers[i] = allocator.buffer(BUFFER_SIZE);
+ }
+
+ for (int i = 0; i < 1000; i++) {
+ localBuffers[i].close();
+ }
+ }
+
+ /**
+ * Main entry point for running the benchmarks standalone.
+ *
+ * This allows running the benchmarks directly from the command line or IDE without using the
+ * Maven JMH plugin. Example usage:
+ *
+ *
+ *
+ */
+public class ArrowFlightJdbcUuidVectorAccessor extends ArrowFlightJdbcAccessor {
+
+ private final UuidVector vector;
+
+ /**
+ * Creates a new accessor for a UUID vector.
+ *
+ * @param vector the UUID vector to access
+ * @param currentRowSupplier supplier for the current row index
+ * @param setCursorWasNull consumer to set the wasNull flag
+ */
+ public ArrowFlightJdbcUuidVectorAccessor(
+ UuidVector vector,
+ IntSupplier currentRowSupplier,
+ ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
+ super(currentRowSupplier, setCursorWasNull);
+ this.vector = vector;
+ }
+
+ @Override
+ public Object getObject() {
+ UUID uuid = vector.getObject(getCurrentRow());
+ this.wasNull = uuid == null;
+ this.wasNullConsumer.setWasNull(this.wasNull);
+ return uuid;
+ }
+
+ @Override
+ public Class getObjectClass() {
+ return UUID.class;
+ }
+
+ @Override
+ public String getString() {
+ UUID uuid = (UUID) getObject();
+ if (uuid == null) {
+ return null;
+ }
+ return uuid.toString();
+ }
+
+ @Override
+ public byte[] getBytes() {
+ UUID uuid = (UUID) getObject();
+ if (uuid == null) {
+ return null;
+ }
+ return UuidUtility.getBytesFromUUID(uuid);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java
index ebebf6ca74..7b04e89346 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java
@@ -35,6 +35,7 @@
import org.apache.arrow.driver.jdbc.utils.DateTimeUtils;
import org.apache.arrow.vector.LargeVarCharVector;
import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarCharVector;
import org.apache.arrow.vector.util.Text;
/** Accessor for the Arrow types: {@link VarCharVector} and {@link LargeVarCharVector}. */
@@ -62,6 +63,13 @@ public ArrowFlightJdbcVarCharVectorAccessor(
this(vector::get, currentRowSupplier, setCursorWasNull);
}
+ public ArrowFlightJdbcVarCharVectorAccessor(
+ ViewVarCharVector vector,
+ IntSupplier currentRowSupplier,
+ ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
+ this(vector::get, currentRowSupplier, setCursorWasNull);
+ }
+
ArrowFlightJdbcVarCharVectorAccessor(
Getter getter,
IntSupplier currentRowSupplier,
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
index a3f6900373..719cc38a2b 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
@@ -32,6 +32,9 @@
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthConfiguration;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthCredentialWriter;
+import org.apache.arrow.driver.jdbc.client.oauth.OAuthTokenProvider;
import org.apache.arrow.driver.jdbc.client.utils.ClientAuthenticationUtils;
import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache;
import org.apache.arrow.driver.jdbc.client.utils.FlightLocationQueue;
@@ -47,7 +50,6 @@
import org.apache.arrow.flight.FlightStatusCode;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.LocationSchemes;
-import org.apache.arrow.flight.SessionOptionValue;
import org.apache.arrow.flight.SessionOptionValueFactory;
import org.apache.arrow.flight.SetSessionOptionsRequest;
import org.apache.arrow.flight.SetSessionOptionsResult;
@@ -147,20 +149,26 @@ public List
+ *
+ *
+ * {@code
+ * // Client Credentials flow
+ * OAuthTokenProvider provider = OAuthTokenProviders.clientCredentials()
+ * .tokenUri("https://auth.example.com/token")
+ * .clientId("my-client")
+ * .clientSecret("my-secret")
+ * .scope("read write")
+ * .build();
+ *
+ * // Token Exchange flow
+ * OAuthTokenProvider provider = OAuthTokenProviders.tokenExchange()
+ * .tokenUri("https://auth.example.com/token")
+ * .subjectToken("user-token")
+ * .subjectTokenType("urn:ietf:params:oauth:token-type:access_token")
+ * .build();
+ * }
+ */
+public final class OAuthTokenProviders {
+
+ private OAuthTokenProviders() {}
+
+ /**
+ * Creates a new builder for Client Credentials flow.
+ *
+ * @return a new ClientCredentialsBuilder instance
+ */
+ public static ClientCredentialsBuilder clientCredentials() {
+ return new ClientCredentialsBuilder();
+ }
+
+ /**
+ * Creates a new builder for Token Exchange flow.
+ *
+ * @return a new TokenExchangeBuilder instance
+ */
+ public static TokenExchangeBuilder tokenExchange() {
+ return new TokenExchangeBuilder();
+ }
+
+ /** Builder for creating {@link ClientCredentialsTokenProvider} instances. */
+ public static class ClientCredentialsBuilder {
+ private @Nullable URI tokenUri;
+ private @Nullable String clientId;
+ private @Nullable String clientSecret;
+ private @Nullable String scope;
+
+ ClientCredentialsBuilder() {}
+
+ /**
+ * Sets the OAuth token endpoint URI (required).
+ *
+ * @param tokenUri the token endpoint URI
+ * @return this builder
+ */
+ public ClientCredentialsBuilder tokenUri(URI tokenUri) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth token endpoint URI from a string (required).
+ *
+ * @param tokenUri the token endpoint URI string
+ * @return this builder
+ * @throws IllegalArgumentException if the URI is invalid
+ */
+ public ClientCredentialsBuilder tokenUri(String tokenUri) {
+ Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ try {
+ this.tokenUri = new URI(tokenUri);
+ } catch (URISyntaxException e) {
+ throw new IllegalArgumentException("Invalid token URI: " + tokenUri, e);
+ }
+ return this;
+ }
+
+ /**
+ * Sets the OAuth client ID (required).
+ *
+ * @param clientId the client ID
+ * @return this builder
+ */
+ public ClientCredentialsBuilder clientId(String clientId) {
+ this.clientId = Objects.requireNonNull(clientId, "clientId cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth client secret (required).
+ *
+ * @param clientSecret the client secret
+ * @return this builder
+ */
+ public ClientCredentialsBuilder clientSecret(String clientSecret) {
+ this.clientSecret = Objects.requireNonNull(clientSecret, "clientSecret cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the OAuth scopes (optional).
+ *
+ * @param scope the space-separated scope string
+ * @return this builder
+ */
+ public ClientCredentialsBuilder scope(@Nullable String scope) {
+ this.scope = scope;
+ return this;
+ }
+
+ /**
+ * Builds a new ClientCredentialsTokenProvider instance.
+ *
+ * @return the configured ClientCredentialsTokenProvider
+ * @throws IllegalStateException if required parameters are missing
+ */
+ public ClientCredentialsTokenProvider build() {
+ if (tokenUri == null) {
+ throw new IllegalStateException("tokenUri is required");
+ }
+ if (clientId == null) {
+ throw new IllegalStateException("clientId is required");
+ }
+ if (clientSecret == null) {
+ throw new IllegalStateException("clientSecret is required");
+ }
+ return new ClientCredentialsTokenProvider(tokenUri, clientId, clientSecret, scope);
+ }
+ }
+
+ /** Builder for creating {@link TokenExchangeTokenProvider} instances. */
+ public static class TokenExchangeBuilder {
+ private @Nullable URI tokenUri;
+ private @Nullable String subjectToken;
+ private @Nullable String subjectTokenType;
+ private @Nullable String actorToken;
+ private @Nullable String actorTokenType;
+ private @Nullable String audience;
+ private @Nullable String requestedTokenType;
+ private @Nullable Scope scope;
+ private @Nullable ListautoCloseables
* when close is called.
*/
- public static AutoCloseable all(final Collection autoCloseables) {
+ public static AutoCloseable all(
+ final @Nullable Collection autoCloseables) {
return new AutoCloseable() {
@Override
public void close() throws Exception {
@@ -48,7 +51,10 @@ public void close() throws Exception {
* @param t the throwable to add suppressed exception to
* @param autoCloseables the closeables to close
*/
- public static void close(Throwable t, AutoCloseable... autoCloseables) {
+ public static void close(Throwable t, @Nullable AutoCloseable... autoCloseables) {
+ if (autoCloseables == null) {
+ return;
+ }
close(t, Arrays.asList(autoCloseables));
}
@@ -58,7 +64,8 @@ public static void close(Throwable t, AutoCloseable... autoCloseables) {
* @param t the throwable to add suppressed exception to
* @param autoCloseables the closeables to close
*/
- public static void close(Throwable t, Iterable autoCloseables) {
+ public static void close(
+ Throwable t, @Nullable Iterable autoCloseables) {
try {
close(autoCloseables);
} catch (Exception e) {
@@ -71,7 +78,10 @@ public static void close(Throwable t, Iterable autoClos
*
* @param autoCloseables the closeables to close
*/
- public static void close(AutoCloseable... autoCloseables) throws Exception {
+ public static void close(@Nullable AutoCloseable... autoCloseables) throws Exception {
+ if (autoCloseables == null) {
+ return;
+ }
close(Arrays.asList(autoCloseables));
}
@@ -80,7 +90,8 @@ public static void close(AutoCloseable... autoCloseables) throws Exception {
*
* @param ac the closeables to close
*/
- public static void close(Iterable ac) throws Exception {
+ public static void close(@Nullable Iterable ac)
+ throws Exception {
// this method can be called on a single object if it implements Iterableac to a {@link Iterable} filtering out any null values. */
- public static Iterablelist to the rollback list. */
- public void addAll(AutoCloseable... list) {
+ public void addAll(@Nullable AutoCloseable... list) {
+ if (list == null) {
+ return;
+ }
closeables.addAll(Arrays.asList(list));
}
/** Add all of list to the rollback list. */
- public void addAll(Iterable list) {
+ public void addAll(@Nullable Iterable list) {
+ if (list == null) {
+ return;
+ }
for (AutoCloseable ac : list) {
closeables.add(ac);
}
@@ -189,7 +214,7 @@ public void close() throws Exception {
}
/** Creates an {@link RollbackCloseable} from the given closeables. */
- public static RollbackCloseable rollbackable(AutoCloseable... closeables) {
+ public static RollbackCloseable rollbackable(@Nullable AutoCloseable... closeables) {
return new RollbackCloseable(closeables);
}
@@ -203,7 +228,7 @@ public static RollbackCloseable rollbackable(AutoCloseable... closeables) {
* @throws RuntimeException if an Exception occurs; the Exception is wrapped by the
* RuntimeException
*/
- public static void closeNoChecked(final AutoCloseable autoCloseable) {
+ public static void closeNoChecked(final @Nullable AutoCloseable autoCloseable) {
if (autoCloseable != null) {
try {
autoCloseable.close();
diff --git a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java
index b5e0a71e7e..f74bf63f82 100644
--- a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java
+++ b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java
@@ -20,32 +20,27 @@
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.condition.JRE.JAVA_16;
+import org.apache.arrow.memory.util.MemoryUtil;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledForJreRange;
public class TestOpens {
- /** Instantiating the RootAllocator should poke MemoryUtil and fail. */
+ /** Accessing MemoryUtil.directBuffer should fail as add-opens is not configured. */
@Test
@EnabledForJreRange(min = JAVA_16)
public void testMemoryUtilFailsLoudly() {
// This test is configured by Maven to run WITHOUT add-opens. So this should fail on JDK16+
// (where JEP396 means that add-opens is required to access JDK internals).
// The test will likely fail in your IDE if it doesn't correctly pick this up.
- Throwable e =
- assertThrows(
- Throwable.class,
- () -> {
- BufferAllocator allocator = new RootAllocator();
- allocator.close();
- });
+ Throwable e = assertThrows(Throwable.class, () -> MemoryUtil.directBuffer(0, 10));
boolean found = false;
while (e != null) {
- e = e.getCause();
- if (e instanceof RuntimeException
- && e.getMessage().contains("Failed to initialize MemoryUtil")) {
+ if (e instanceof UnsupportedOperationException
+ && e.getMessage().contains("java.nio.DirectByteBuffer.{@code
+ * java -cp target/benchmarks.jar org.apache.arrow.memory.MemoryFootprintBenchmarks
+ * }
+ *
+ * @param args command line arguments (not used)
+ * @throws RunnerException if the benchmark runner encounters an error
+ */
+ public static void main(String[] args) throws RunnerException {
+ Options opt =
+ new OptionsBuilder()
+ .include(MemoryFootprintBenchmarks.class.getSimpleName())
+ .forks(1)
+ .build();
+
+ new Runner(opt).run();
+ }
+}
diff --git a/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java
new file mode 100644
index 0000000000..b5f87e7a75
--- /dev/null
+++ b/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector;
+
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.complex.impl.UuidWriterImpl;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.profile.GCProfiler;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+/** Benchmarks for {@link UuidVector}. */
+@State(Scope.Benchmark)
+public class UuidVectorBenchmarks {
+ // checkstyle:off: MissingJavadocMethod
+
+ private static final int VECTOR_LENGTH = 10_000;
+
+ private static final int ALLOCATOR_CAPACITY = 1024 * 1024;
+
+ private BufferAllocator allocator;
+
+ private UuidVector vector;
+
+ private UUID[] testUuids;
+
+ @Setup
+ public void prepare() {
+ allocator = new RootAllocator(ALLOCATOR_CAPACITY);
+ vector = new UuidVector("vector", allocator);
+ vector.allocateNew(VECTOR_LENGTH);
+ vector.setValueCount(VECTOR_LENGTH);
+
+ // Pre-generate UUIDs for consistent benchmarking
+ testUuids = new UUID[VECTOR_LENGTH];
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ testUuids[i] = new UUID(i, i * 2L);
+ }
+ }
+
+ @TearDown
+ public void tearDown() {
+ vector.close();
+ allocator.close();
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void setWithHolder() {
+ NullableUuidHolder holder = new NullableUuidHolder();
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ vector.get(i, holder);
+ vector.setSafe(i, holder);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void setUuidDirectly() {
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ vector.setSafe(i, testUuids[i]);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void setWithWriter() {
+ UuidWriterImpl writer = new UuidWriterImpl(vector);
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ writer.writeExtension(testUuids[i]);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void getWithUuidHolder() {
+ NullableUuidHolder holder = new NullableUuidHolder();
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ vector.get(i, holder);
+ }
+ }
+
+ @Benchmark
+ @BenchmarkMode(Mode.AverageTime)
+ @OutputTimeUnit(TimeUnit.MICROSECONDS)
+ public void getUuidDirectly() {
+ for (int i = 0; i < VECTOR_LENGTH; i++) {
+ UUID uuid = vector.getObject(i);
+ }
+ }
+
+ public static void main(String[] args) throws RunnerException {
+ Options opt =
+ new OptionsBuilder()
+ .include(UuidVectorBenchmarks.class.getSimpleName())
+ .forks(1)
+ .addProfiler(GCProfiler.class)
+ .build();
+
+ new Runner(opt).run();
+ }
+ // checkstyle:on: MissingJavadocMethod
+}
diff --git a/pom.xml b/pom.xml
index f6c9053db9..78137eb4e9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -23,12 +23,12 @@ under the License.
Supports setting and retrieving UUIDs with efficient storage and nullable value handling. + * + *
Usage: + * + *
{@code
+ * UuidVector vector = new UuidVector("uuid_col", allocator);
+ * vector.set(0, UUID.randomUUID());
+ * UUID value = vector.getObject(0);
+ * }
+ *
+ * @see UuidType
+ * @see UuidHolder
+ * @see NullableUuidHolder
+ */
+public class UuidVector extends ExtensionTypeVectorCreates a new underlying FixedSizeBinaryVector with 16-byte width.
+ *
+ * @param name the name of the vector
+ * @param allocator the buffer allocator
+ */
+ public UuidVector(String name, BufferAllocator allocator) {
+ super(name, allocator, new FixedSizeBinaryVector(name, allocator, UUID_BYTE_WIDTH));
+ this.field = new Field(name, FieldType.nullable(UuidType.INSTANCE), null);
+ }
+
+ /**
+ * Constructs a UUID vector from a field and allocator.
+ *
+ * @param field the field definition (should contain UuidType)
+ * @param allocator the buffer allocator
+ */
+ public UuidVector(Field field, BufferAllocator allocator) {
+ super(
+ field.getName(),
+ allocator,
+ new FixedSizeBinaryVector(field.getName(), allocator, UUID_BYTE_WIDTH));
+ this.field = field;
+ }
+
+ @Override
+ public UUID getObject(int index) {
+ if (isSet(index) == 0) {
+ return null;
+ }
+ final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index));
+ return new UUID(bb.getLong(), bb.getLong());
+ }
+
+ @Override
+ public int hashCode(int index) {
+ return hashCode(index, null);
+ }
+
+ @Override
+ public int hashCode(int index, ArrowBufHasher hasher) {
+ int start = this.getStartOffset(index);
+ return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, start + UUID_BYTE_WIDTH);
+ }
+
+ /**
+ * Checks if the value at the given index is set (non-null).
+ *
+ * @param index the index to check
+ * @return 1 if the value is set, 0 if null
+ */
+ public int isSet(int index) {
+ return getUnderlyingVector().isSet(index);
+ }
+
+ /**
+ * Reads the UUID value at the given index into a NullableUuidHolder.
+ *
+ * @param index the index to read from
+ * @param holder the holder to populate with the UUID data
+ */
+ public void get(int index, NullableUuidHolder holder) {
+ Preconditions.checkArgument(index >= 0, "Cannot get negative index in UUID vector.");
+ if (isSet(index) == 0) {
+ holder.isSet = 0;
+ return;
+ }
+ holder.isSet = 1;
+ holder.buffer = getDataBuffer();
+ holder.start = getStartOffset(index);
+ }
+
+ /**
+ * Calculates the byte offset for a given index in the data buffer.
+ *
+ * @param index the index of the UUID value
+ * @return the byte offset in the data buffer
+ */
+ public final int getStartOffset(int index) {
+ return index * UUID_BYTE_WIDTH;
+ }
+
+ /**
+ * Sets the UUID value at the given index.
+ *
+ * @param index the index to set
+ * @param value the UUID value to set, or null to set a null value
+ */
+ public void set(int index, UUID value) {
+ if (value != null) {
+ set(index, UuidUtility.getBytesFromUUID(value));
+ } else {
+ getUnderlyingVector().setNull(index);
+ }
+ }
+
+ /**
+ * Sets the UUID value at the given index from a UuidHolder.
+ *
+ * @param index the index to set
+ * @param holder the holder containing the UUID data
+ */
+ public void set(int index, UuidHolder holder) {
+ this.set(index, holder.buffer, holder.start);
+ }
+
+ /**
+ * Sets the UUID value at the given index from a NullableUuidHolder.
+ *
+ * @param index the index to set
+ * @param holder the holder containing the UUID data
+ */
+ public void set(int index, NullableUuidHolder holder) {
+ if (holder.isSet == 0) {
+ getUnderlyingVector().setNull(index);
+ } else {
+ this.set(index, holder.buffer, holder.start);
+ }
+ }
+
+ /**
+ * Sets the UUID value at the given index by copying from a source buffer.
+ *
+ * @param index the index to set
+ * @param source the source buffer to copy from
+ * @param sourceOffset the offset in the source buffer where the UUID data starts
+ */
+ public void set(int index, ArrowBuf source, int sourceOffset) {
+ Preconditions.checkNotNull(source, "Cannot set UUID vector, the source buffer is null.");
+
+ BitVectorHelper.setBit(getUnderlyingVector().getValidityBuffer(), index);
+ getUnderlyingVector()
+ .getDataBuffer()
+ .setBytes((long) index * UUID_BYTE_WIDTH, source, sourceOffset, UUID_BYTE_WIDTH);
+ }
+
+ /**
+ * Sets the UUID value at the given index from a byte array.
+ *
+ * @param index the index to set
+ * @param value the 16-byte array containing the UUID data
+ */
+ public void set(int index, byte[] value) {
+ getUnderlyingVector().set(index, value);
+ }
+
+ /**
+ * Sets the UUID value at the given index, expanding capacity if needed.
+ *
+ * @param index the index to set
+ * @param value the UUID value to set, or null to set a null value
+ */
+ public void setSafe(int index, UUID value) {
+ if (value != null) {
+ setSafe(index, UuidUtility.getBytesFromUUID(value));
+ } else {
+ getUnderlyingVector().setNull(index);
+ }
+ }
+
+ /**
+ * Sets the UUID value at the given index from a NullableUuidHolder, expanding capacity if needed.
+ *
+ * @param index the index to set
+ * @param holder the holder containing the UUID data, or null to set a null value
+ */
+ public void setSafe(int index, NullableUuidHolder holder) {
+ if (holder == null || holder.isSet == 0) {
+ getUnderlyingVector().setNull(index);
+ } else {
+ this.setSafe(index, holder.buffer, holder.start);
+ }
+ }
+
+ /**
+ * Sets the UUID value at the given index from a UuidHolder, expanding capacity if needed.
+ *
+ * @param index the index to set
+ * @param holder the holder containing the UUID data
+ */
+ public void setSafe(int index, UuidHolder holder) {
+ this.setSafe(index, holder.buffer, holder.start);
+ }
+
+ /**
+ * Sets the UUID value at the given index by copying from a source buffer, expanding capacity if
+ * needed.
+ *
+ * @param index the index to set
+ * @param buffer the source buffer to copy from
+ * @param start the offset in the source buffer where the UUID data starts
+ */
+ public void setSafe(int index, ArrowBuf buffer, int start) {
+ getUnderlyingVector().handleSafe(index);
+ this.set(index, buffer, start);
+ }
+
+ /**
+ * Sets the UUID value at the given index from a byte array, expanding capacity if needed.
+ *
+ * @param index the index to set
+ * @param value the 16-byte array containing the UUID data
+ */
+ public void setSafe(int index, byte[] value) {
+ getUnderlyingVector().setIndexDefined(index);
+ getUnderlyingVector().setSafe(index, value);
+ }
+
+ /**
+ * Sets the UUID value at the given index from an ArrowBuf, expanding capacity if needed.
+ *
+ * @param index the index to set
+ * @param value the buffer containing the 16-byte UUID data
+ */
+ public void setSafe(int index, ArrowBuf value) {
+ getUnderlyingVector().setSafe(index, value);
+ }
+
+ @Override
+ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
+ getUnderlyingVector()
+ .copyFromSafe(fromIndex, thisIndex, ((UuidVector) from).getUnderlyingVector());
+ }
+
+ @Override
+ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
+ getUnderlyingVector()
+ .copyFromSafe(fromIndex, thisIndex, ((UuidVector) from).getUnderlyingVector());
+ }
+
+ @Override
+ public Field getField() {
+ return field;
+ }
+
+ @Override
+ public ArrowBufPointer getDataPointer(int i) {
+ return getUnderlyingVector().getDataPointer(i);
+ }
+
+ @Override
+ public ArrowBufPointer getDataPointer(int i, ArrowBufPointer arrowBufPointer) {
+ return getUnderlyingVector().getDataPointer(i, arrowBufPointer);
+ }
+
+ @Override
+ public void allocateNew(int valueCount) {
+ getUnderlyingVector().allocateNew(valueCount);
+ }
+
+ @Override
+ public void zeroVector() {
+ getUnderlyingVector().zeroVector();
+ }
+
+ @Override
+ public TransferPair makeTransferPair(ValueVector to) {
+ return new TransferImpl((UuidVector) to);
+ }
+
+ @Override
+ protected FieldReader getReaderImpl() {
+ return new UuidReaderImpl(this);
+ }
+
+ @Override
+ public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+ return new TransferImpl(field, allocator);
+ }
+
+ @Override
+ public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) {
+ return getTransferPair(field, allocator);
+ }
+
+ @Override
+ public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
+ return new TransferImpl(ref, allocator);
+ }
+
+ @Override
+ public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
+ return getTransferPair(ref, allocator);
+ }
+
+ @Override
+ public TransferPair getTransferPair(BufferAllocator allocator) {
+ return getTransferPair(this.getField().getName(), allocator);
+ }
+
+ @Override
+ public int getTypeWidth() {
+ return UUID_BYTE_WIDTH;
+ }
+
+ /** {@link TransferPair} for {@link UuidVector}. */
+ public class TransferImpl implements TransferPair {
+ UuidVector to;
+
+ /**
+ * Constructs a transfer pair with the given target vector.
+ *
+ * @param to the target UUID vector
+ */
+ public TransferImpl(UuidVector to) {
+ this.to = to;
+ }
+
+ /**
+ * Constructs a transfer pair, creating a new target vector from the field and allocator.
+ *
+ * @param field the field definition for the target vector
+ * @param allocator the buffer allocator for the target vector
+ */
+ public TransferImpl(Field field, BufferAllocator allocator) {
+ this.to = new UuidVector(field, allocator);
+ }
+
+ /**
+ * Constructs a transfer pair, creating a new target vector with the given name and allocator.
+ *
+ * @param ref the name for the target vector
+ * @param allocator the buffer allocator for the target vector
+ */
+ public TransferImpl(String ref, BufferAllocator allocator) {
+ this.to = new UuidVector(ref, allocator);
+ }
+
+ /**
+ * Gets the target vector of this transfer pair.
+ *
+ * @return the target UUID vector
+ */
+ public UuidVector getTo() {
+ return this.to;
+ }
+
+ /** Transfers ownership of data from the source vector to the target vector. */
+ public void transfer() {
+ getUnderlyingVector().transferTo(to.getUnderlyingVector());
+ }
+
+ /**
+ * Splits and transfers a range of values from the source vector to the target vector.
+ *
+ * @param startIndex the starting index in the source vector
+ * @param length the number of values to transfer
+ */
+ public void splitAndTransfer(int startIndex, int length) {
+ getUnderlyingVector().splitAndTransferTo(startIndex, length, to.getUnderlyingVector());
+ }
+
+ /**
+ * Copies a value from the source vector to the target vector, expanding capacity if needed.
+ *
+ * @param fromIndex the index in the source vector
+ * @param toIndex the index in the target vector
+ */
+ public void copyValueSafe(int fromIndex, int toIndex) {
+ to.copyFromSafe(fromIndex, toIndex, (ValueVector) UuidVector.this);
+ }
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/ValueVector.java b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
index e0628c2ee1..3a5058256c 100644
--- a/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
@@ -22,7 +22,6 @@
import org.apache.arrow.memory.OutOfMemoryException;
import org.apache.arrow.memory.util.hash.ArrowBufHasher;
import org.apache.arrow.vector.compare.VectorVisitor;
-import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
import org.apache.arrow.vector.complex.reader.FieldReader;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.Field;
@@ -310,30 +309,6 @@ public interface ValueVector extends Closeable, Iterable