```
-
-[1]: https://logback.qos.ch/manual/configuration.html
-[2]: https://github.com/apache/arrow/blob/main/cpp/README.md
-[3]: http://google.github.io/styleguide/javaguide.html
-[4]: https://maven.apache.org/surefire/maven-failsafe-plugin/
diff --git a/adapter/avro/pom.xml b/adapter/avro/pom.xml
index cf9d353330..18c48a0e8f 100644
--- a/adapter/avro/pom.xml
+++ b/adapter/avro/pom.xml
@@ -23,7 +23,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
../../pom.xml
diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java
index 87b594af9e..e09b99f670 100644
--- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java
+++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java
@@ -17,10 +17,14 @@
package org.apache.arrow.adapter.avro;
import java.util.ArrayList;
+import java.util.HashSet;
import java.util.List;
+import java.util.Set;
+import java.util.regex.Pattern;
import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer;
import org.apache.arrow.adapter.avro.producers.AvroBooleanProducer;
import org.apache.arrow.adapter.avro.producers.AvroBytesProducer;
+import org.apache.arrow.adapter.avro.producers.AvroEnumProducer;
import org.apache.arrow.adapter.avro.producers.AvroFixedSizeBinaryProducer;
import org.apache.arrow.adapter.avro.producers.AvroFixedSizeListProducer;
import org.apache.arrow.adapter.avro.producers.AvroFloat2Producer;
@@ -41,6 +45,7 @@
import org.apache.arrow.adapter.avro.producers.AvroUint8Producer;
import org.apache.arrow.adapter.avro.producers.BaseAvroProducer;
import org.apache.arrow.adapter.avro.producers.CompositeAvroProducer;
+import org.apache.arrow.adapter.avro.producers.DictionaryDecodingProducer;
import org.apache.arrow.adapter.avro.producers.Producer;
import org.apache.arrow.adapter.avro.producers.logical.AvroDateDayProducer;
import org.apache.arrow.adapter.avro.producers.logical.AvroDateMilliProducer;
@@ -59,6 +64,7 @@
import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampSecProducer;
import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampSecTzProducer;
import org.apache.arrow.util.Preconditions;
+import org.apache.arrow.vector.BaseIntVector;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.DateDayVector;
@@ -96,11 +102,14 @@
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.util.Text;
import org.apache.avro.LogicalType;
import org.apache.avro.LogicalTypes;
import org.apache.avro.Schema;
@@ -162,17 +171,29 @@ public class ArrowToAvroUtils {
* may be nullable. Record types must contain at least one child field and cannot contain multiple
* fields with the same name
*
+ * String fields that are dictionary-encoded will be represented as an Avro enum, so long as
+ * all the values meet the restrictions on Avro enums (non-null, valid identifiers). Other data
+ * types that are dictionary encoded, or string fields that do not meet the avro requirements,
+ * will be output as their decoded type.
+ *
* @param arrowFields The arrow fields used to generate the Avro schema
* @param typeName Name of the top level Avro record type
* @param namespace Namespace of the top level Avro record type
+ * @param dictionaries A dictionary provider is required if any fields use dictionary encoding
* @return An Avro record schema for the given list of fields, with the specified name and
* namespace
*/
public static Schema createAvroSchema(
- List arrowFields, String typeName, String namespace) {
+ List arrowFields, String typeName, String namespace, DictionaryProvider dictionaries) {
SchemaBuilder.RecordBuilder assembler =
SchemaBuilder.record(typeName).namespace(namespace);
- return buildRecordSchema(assembler, arrowFields, namespace);
+ return buildRecordSchema(assembler, arrowFields, namespace, dictionaries);
+ }
+
+ /** Overload provided for convenience, sets dictionaries = null. */
+ public static Schema createAvroSchema(
+ List arrowFields, String typeName, String namespace) {
+ return createAvroSchema(arrowFields, typeName, namespace, null);
}
/** Overload provided for convenience, sets namespace = null. */
@@ -185,61 +206,83 @@ public static Schema createAvroSchema(List arrowFields) {
return createAvroSchema(arrowFields, GENERIC_RECORD_TYPE_NAME);
}
+ /**
+ * Overload provided for convenience, sets name = GENERIC_RECORD_TYPE_NAME and namespace = null.
+ */
+ public static Schema createAvroSchema(List arrowFields, DictionaryProvider dictionaries) {
+ return createAvroSchema(arrowFields, GENERIC_RECORD_TYPE_NAME, null, dictionaries);
+ }
+
private static T buildRecordSchema(
- SchemaBuilder.RecordBuilder builder, List fields, String namespace) {
+ SchemaBuilder.RecordBuilder builder,
+ List fields,
+ String namespace,
+ DictionaryProvider dictionaries) {
if (fields.isEmpty()) {
throw new IllegalArgumentException("Record field must have at least one child field");
}
SchemaBuilder.FieldAssembler assembler = builder.namespace(namespace).fields();
for (Field field : fields) {
- assembler = buildFieldSchema(assembler, field, namespace);
+ assembler = buildFieldSchema(assembler, field, namespace, dictionaries);
}
return assembler.endRecord();
}
private static SchemaBuilder.FieldAssembler buildFieldSchema(
- SchemaBuilder.FieldAssembler assembler, Field field, String namespace) {
+ SchemaBuilder.FieldAssembler assembler,
+ Field field,
+ String namespace,
+ DictionaryProvider dictionaries) {
return assembler
.name(field.getName())
- .type(buildTypeSchema(SchemaBuilder.builder(), field, namespace))
+ .type(buildTypeSchema(SchemaBuilder.builder(), field, namespace, dictionaries))
.noDefault();
}
private static T buildTypeSchema(
- SchemaBuilder.TypeBuilder builder, Field field, String namespace) {
+ SchemaBuilder.TypeBuilder builder,
+ Field field,
+ String namespace,
+ DictionaryProvider dictionaries) {
// Nullable unions need special handling, since union types cannot be directly nested
if (field.getType().getTypeID() == ArrowType.ArrowTypeID.Union) {
boolean unionNullable = field.getChildren().stream().anyMatch(Field::isNullable);
if (unionNullable) {
SchemaBuilder.UnionAccumulator union = builder.unionOf().nullType();
- return addTypesToUnion(union, field.getChildren(), namespace);
+ return addTypesToUnion(union, field.getChildren(), namespace, dictionaries);
} else {
Field headType = field.getChildren().get(0);
List tailTypes = field.getChildren().subList(1, field.getChildren().size());
SchemaBuilder.UnionAccumulator union =
- buildBaseTypeSchema(builder.unionOf(), headType, namespace);
- return addTypesToUnion(union, tailTypes, namespace);
+ buildBaseTypeSchema(builder.unionOf(), headType, namespace, dictionaries);
+ return addTypesToUnion(union, tailTypes, namespace, dictionaries);
}
} else if (field.isNullable()) {
- return buildBaseTypeSchema(builder.nullable(), field, namespace);
+ return buildBaseTypeSchema(builder.nullable(), field, namespace, dictionaries);
} else {
- return buildBaseTypeSchema(builder, field, namespace);
+ return buildBaseTypeSchema(builder, field, namespace, dictionaries);
}
}
private static T buildArraySchema(
- SchemaBuilder.ArrayBuilder builder, Field listField, String namespace) {
+ SchemaBuilder.ArrayBuilder builder,
+ Field listField,
+ String namespace,
+ DictionaryProvider dictionaries) {
if (listField.getChildren().size() != 1) {
throw new IllegalArgumentException("List field must have exactly one child field");
}
Field itemField = listField.getChildren().get(0);
- return buildTypeSchema(builder.items(), itemField, namespace);
+ return buildTypeSchema(builder.items(), itemField, namespace, dictionaries);
}
private static T buildMapSchema(
- SchemaBuilder.MapBuilder builder, Field mapField, String namespace) {
+ SchemaBuilder.MapBuilder builder,
+ Field mapField,
+ String namespace,
+ DictionaryProvider dictionaries) {
if (mapField.getChildren().size() != 1) {
throw new IllegalArgumentException("Map field must have exactly one child field");
}
@@ -253,11 +296,14 @@ private static T buildMapSchema(
throw new IllegalArgumentException(
"Map keys must be of type string and cannot be nullable for conversion to Avro");
}
- return buildTypeSchema(builder.values(), valueField, namespace);
+ return buildTypeSchema(builder.values(), valueField, namespace, dictionaries);
}
private static T buildBaseTypeSchema(
- SchemaBuilder.BaseTypeBuilder builder, Field field, String namespace) {
+ SchemaBuilder.BaseTypeBuilder builder,
+ Field field,
+ String namespace,
+ DictionaryProvider dictionaries) {
ArrowType.ArrowTypeID typeID = field.getType().getTypeID();
@@ -269,6 +315,33 @@ private static T buildBaseTypeSchema(
return builder.booleanType();
case Int:
+ if (field.getDictionary() != null) {
+ if (dictionaries == null) {
+ throw new IllegalArgumentException(
+ "Field references a dictionary but no dictionaries were provided: "
+ + field.getName());
+ }
+ Dictionary dictionary = dictionaries.lookup(field.getDictionary().getId());
+ if (dictionary == null) {
+ throw new IllegalArgumentException(
+ "Field references a dictionary that does not exist: "
+ + field.getName()
+ + ", dictionary ID = "
+ + field.getDictionary().getId());
+ }
+ if (dictionaryIsValidEnum(dictionary)) {
+ String[] symbols = dictionarySymbols(dictionary);
+ return builder.enumeration(field.getName()).symbols(symbols);
+ } else {
+ Field decodedField =
+ new Field(
+ field.getName(),
+ dictionary.getVector().getField().getFieldType(),
+ dictionary.getVector().getField().getChildren());
+ return buildBaseTypeSchema(builder, decodedField, namespace, dictionaries);
+ }
+ }
+
ArrowType.Int intType = (ArrowType.Int) field.getType();
if (intType.getBitWidth() > 32 || (intType.getBitWidth() == 32 && !intType.getIsSigned())) {
return builder.longType();
@@ -328,7 +401,7 @@ private static T buildBaseTypeSchema(
String childNamespace =
namespace == null ? field.getName() : namespace + "." + field.getName();
return buildRecordSchema(
- builder.record(field.getName()), field.getChildren(), childNamespace);
+ builder.record(field.getName()), field.getChildren(), childNamespace, dictionaries);
case List:
case FixedSizeList:
@@ -339,13 +412,13 @@ private static T buildBaseTypeSchema(
new Field("item", itemField.getFieldType(), itemField.getChildren());
Field safeListField =
new Field(field.getName(), field.getFieldType(), List.of(safeItemField));
- return buildArraySchema(builder.array(), safeListField, namespace);
+ return buildArraySchema(builder.array(), safeListField, namespace, dictionaries);
} else {
- return buildArraySchema(builder.array(), field, namespace);
+ return buildArraySchema(builder.array(), field, namespace, dictionaries);
}
case Map:
- return buildMapSchema(builder.map(), field, namespace);
+ return buildMapSchema(builder.map(), field, namespace, dictionaries);
default:
throw new IllegalArgumentException(
@@ -354,9 +427,12 @@ private static T buildBaseTypeSchema(
}
private static T addTypesToUnion(
- SchemaBuilder.UnionAccumulator accumulator, List unionFields, String namespace) {
+ SchemaBuilder.UnionAccumulator accumulator,
+ List unionFields,
+ String namespace,
+ DictionaryProvider dictionaries) {
for (var field : unionFields) {
- accumulator = buildBaseTypeSchema(accumulator.and(), field, namespace);
+ accumulator = buildBaseTypeSchema(accumulator.and(), field, namespace, dictionaries);
}
return accumulator.endUnion();
}
@@ -373,30 +449,88 @@ private static LogicalType timestampLogicalType(ArrowType.Timestamp timestampTyp
}
}
+ private static boolean dictionaryIsValidEnum(Dictionary dictionary) {
+
+ if (dictionary.getVectorType().getTypeID() != ArrowType.ArrowTypeID.Utf8) {
+ return false;
+ }
+
+ VarCharVector vector = (VarCharVector) dictionary.getVector();
+ Set symbols = new HashSet<>();
+
+ for (int i = 0; i < vector.getValueCount(); i++) {
+ if (vector.isNull(i)) {
+ return false;
+ }
+ Text text = vector.getObject(i);
+ if (text == null) {
+ return false;
+ }
+ String symbol = text.toString();
+ if (!ENUM_REGEX.matcher(symbol).matches()) {
+ return false;
+ }
+ if (symbols.contains(symbol)) {
+ return false;
+ }
+ symbols.add(symbol);
+ }
+
+ return true;
+ }
+
+ private static String[] dictionarySymbols(Dictionary dictionary) {
+
+ VarCharVector vector = (VarCharVector) dictionary.getVector();
+ String[] symbols = new String[vector.getValueCount()];
+
+ for (int i = 0; i < vector.getValueCount(); i++) {
+ Text text = vector.getObject(i);
+ // This should never happen if dictionaryIsValidEnum() succeeded
+ if (text == null) {
+ throw new IllegalArgumentException("Illegal null value in enum");
+ }
+ symbols[i] = text.toString();
+ }
+
+ return symbols;
+ }
+
+ private static final Pattern ENUM_REGEX = Pattern.compile("^[A-Za-z_][A-Za-z0-9_]*$");
+
/**
* Create a composite Avro producer for a set of field vectors (typically the root set of a VSR).
*
* @param vectors The vectors that will be used to produce Avro data
* @return The resulting composite Avro producer
*/
- public static CompositeAvroProducer createCompositeProducer(List vectors) {
+ public static CompositeAvroProducer createCompositeProducer(
+ List vectors, DictionaryProvider dictionaries) {
List> producers = new ArrayList<>(vectors.size());
for (FieldVector vector : vectors) {
- BaseAvroProducer producer = createProducer(vector);
+ BaseAvroProducer producer = createProducer(vector, dictionaries);
producers.add(producer);
}
return new CompositeAvroProducer(producers);
}
- private static BaseAvroProducer createProducer(FieldVector vector) {
+ /** Overload provided for convenience, sets dictionaries = null. */
+ public static CompositeAvroProducer createCompositeProducer(List vectors) {
+
+ return createCompositeProducer(vectors, null);
+ }
+
+ private static BaseAvroProducer createProducer(
+ FieldVector vector, DictionaryProvider dictionaries) {
boolean nullable = vector.getField().isNullable();
- return createProducer(vector, nullable);
+ return createProducer(vector, nullable, dictionaries);
}
- private static BaseAvroProducer createProducer(FieldVector vector, boolean nullable) {
+ private static BaseAvroProducer createProducer(
+ FieldVector vector, boolean nullable, DictionaryProvider dictionaries) {
Preconditions.checkNotNull(vector, "Arrow vector object can't be null");
@@ -405,10 +539,34 @@ private static BaseAvroProducer createProducer(FieldVector vector, boolean nu
// Avro understands nullable types as a union of type | null
// Most nullable fields in a VSR will not be unions, so provide a special wrapper
if (nullable && minorType != Types.MinorType.UNION) {
- final BaseAvroProducer innerProducer = createProducer(vector, false);
+ final BaseAvroProducer innerProducer = createProducer(vector, false, dictionaries);
return new AvroNullableProducer<>(innerProducer);
}
+ if (vector.getField().getDictionary() != null) {
+ if (dictionaries == null) {
+ throw new IllegalArgumentException(
+ "Field references a dictionary but no dictionaries were provided: "
+ + vector.getField().getName());
+ }
+ Dictionary dictionary = dictionaries.lookup(vector.getField().getDictionary().getId());
+ if (dictionary == null) {
+ throw new IllegalArgumentException(
+ "Field references a dictionary that does not exist: "
+ + vector.getField().getName()
+ + ", dictionary ID = "
+ + vector.getField().getDictionary().getId());
+ }
+ // If a field is dictionary-encoded but cannot be represented as an Avro enum,
+ // then decode it before writing
+ if (dictionaryIsValidEnum(dictionary)) {
+ return new AvroEnumProducer((BaseIntVector) vector);
+ } else {
+ BaseAvroProducer dictProducer = createProducer(dictionary.getVector(), false, null);
+ return new DictionaryDecodingProducer<>((BaseIntVector) vector, dictProducer);
+ }
+ }
+
switch (minorType) {
case NULL:
return new AvroNullProducer((NullVector) vector);
@@ -486,21 +644,23 @@ private static BaseAvroProducer createProducer(FieldVector vector, boolean nu
Producer[] childProducers = new Producer[childVectors.size()];
for (int i = 0; i < childVectors.size(); i++) {
FieldVector childVector = childVectors.get(i);
- childProducers[i] = createProducer(childVector, childVector.getField().isNullable());
+ childProducers[i] =
+ createProducer(childVector, childVector.getField().isNullable(), dictionaries);
}
return new AvroStructProducer(structVector, childProducers);
case LIST:
ListVector listVector = (ListVector) vector;
FieldVector itemVector = listVector.getDataVector();
- Producer itemProducer = createProducer(itemVector, itemVector.getField().isNullable());
+ Producer itemProducer =
+ createProducer(itemVector, itemVector.getField().isNullable(), dictionaries);
return new AvroListProducer(listVector, itemProducer);
case FIXED_SIZE_LIST:
FixedSizeListVector fixedListVector = (FixedSizeListVector) vector;
FieldVector fixedItemVector = fixedListVector.getDataVector();
Producer fixedItemProducer =
- createProducer(fixedItemVector, fixedItemVector.getField().isNullable());
+ createProducer(fixedItemVector, fixedItemVector.getField().isNullable(), dictionaries);
return new AvroFixedSizeListProducer(fixedListVector, fixedItemProducer);
case MAP:
@@ -514,7 +674,7 @@ private static BaseAvroProducer createProducer(FieldVector vector, boolean nu
FieldVector valueVector = entryVector.getChildrenFromFields().get(1);
Producer keyProducer = new AvroStringProducer(keyVector);
Producer valueProducer =
- createProducer(valueVector, valueVector.getField().isNullable());
+ createProducer(valueVector, valueVector.getField().isNullable(), dictionaries);
Producer entryProducer =
new AvroStructProducer(entryVector, new Producer[] {keyProducer, valueProducer});
return new AvroMapProducer(mapVector, entryProducer);
diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java
index aedef7732e..a6e77e4050 100644
--- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java
+++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java
@@ -1071,8 +1071,8 @@ private static FieldType createFieldType(
}
private static String convertAliases(Set aliases) {
- JsonStringArrayList jsonList = new JsonStringArrayList();
- aliases.stream().forEach(a -> jsonList.add(a));
+ JsonStringArrayList jsonList = new JsonStringArrayList(aliases.size());
+ jsonList.addAll(aliases);
return jsonList.toString();
}
}
diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java
index 4123370061..e82fdc36fb 100644
--- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java
+++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java
@@ -17,13 +17,14 @@
package org.apache.arrow.adapter.avro;
import java.io.EOFException;
-import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.arrow.adapter.avro.consumers.CompositeAvroConsumer;
+import org.apache.arrow.adapter.avro.consumers.Consumer;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.util.ValueVectorUtility;
@@ -75,9 +76,11 @@ public static AvroToArrowVectorIterator create(
private void initialize() {
// create consumers
compositeConsumer = AvroToArrowUtils.createCompositeConsumer(schema, config);
- List vectors = new ArrayList<>();
- compositeConsumer.getConsumers().forEach(c -> vectors.add(c.getVector()));
- List fields = vectors.stream().map(t -> t.getField()).collect(Collectors.toList());
+ List vectors =
+ compositeConsumer.getConsumers().stream()
+ .map(Consumer::getVector)
+ .collect(Collectors.toList());
+ List fields = vectors.stream().map(ValueVector::getField).collect(Collectors.toList());
VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 0);
rootSchema = root.getSchema();
diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java
index 068566493e..eebfb7d241 100644
--- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java
+++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java
@@ -17,22 +17,22 @@
package org.apache.arrow.adapter.avro.producers;
import java.io.IOException;
-import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.BaseIntVector;
import org.apache.avro.io.Encoder;
/**
- * Producer that produces enum values from a dictionary-encoded {@link IntVector}, writes data to an
- * Avro encoder.
+ * Producer that produces enum values from a dictionary-encoded {@link BaseIntVector}, writes data
+ * to an Avro encoder.
*/
-public class AvroEnumProducer extends BaseAvroProducer {
+public class AvroEnumProducer extends BaseAvroProducer {
/** Instantiate an AvroEnumProducer. */
- public AvroEnumProducer(IntVector vector) {
+ public AvroEnumProducer(BaseIntVector vector) {
super(vector);
}
@Override
public void produce(Encoder encoder) throws IOException {
- encoder.writeEnum(vector.get(currentIndex++));
+ encoder.writeEnum((int) vector.getValueAsLong(currentIndex++));
}
}
diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java
new file mode 100644
index 0000000000..afeba08511
--- /dev/null
+++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java
@@ -0,0 +1,47 @@
+/*
+ * 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.adapter.avro.producers;
+
+import java.io.IOException;
+import org.apache.arrow.vector.BaseIntVector;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.avro.io.Encoder;
+
+/**
+ * Producer that decodes values from a dictionary-encoded {@link FieldVector}, writes the resulting
+ * values to an Avro encoder.
+ *
+ * @param Type of the underlying dictionary vector
+ */
+public class DictionaryDecodingProducer
+ extends BaseAvroProducer {
+
+ private final Producer dictProducer;
+
+ /** Instantiate a DictionaryDecodingProducer. */
+ public DictionaryDecodingProducer(BaseIntVector indexVector, Producer dictProducer) {
+ super(indexVector);
+ this.dictProducer = dictProducer;
+ }
+
+ @Override
+ public void produce(Encoder encoder) throws IOException {
+ int dicIndex = (int) vector.getValueAsLong(currentIndex++);
+ dictProducer.setPosition(dicIndex);
+ dictProducer.produce(encoder);
+ }
+}
diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java
index 2d70b45021..6d66ee9d45 100644
--- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java
+++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java
@@ -76,10 +76,14 @@
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.writer.BaseWriter;
import org.apache.arrow.vector.complex.writer.FieldWriter;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryEncoder;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.JsonStringArrayList;
@@ -2817,4 +2821,81 @@ record = datumReader.read(record, decoder);
}
}
}
+
+ @Test
+ public void testWriteDictEnumEncoded() throws Exception {
+
+ BufferAllocator allocator = new RootAllocator();
+
+ // Create a dictionary
+ FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector dictionaryVector =
+ new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+ dictionaryVector.allocateNew(3);
+ dictionaryVector.set(0, "apple".getBytes());
+ dictionaryVector.set(1, "banana".getBytes());
+ dictionaryVector.set(2, "cherry".getBytes());
+ dictionaryVector.setValueCount(3);
+
+ Dictionary dictionary =
+ new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null));
+ DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+ // Field definition
+ FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector stringVector =
+ new VarCharVector(new Field("enumField", stringField, null), allocator);
+ stringVector.allocateNew(10);
+ stringVector.setSafe(0, "apple".getBytes());
+ stringVector.setSafe(1, "banana".getBytes());
+ stringVector.setSafe(2, "cherry".getBytes());
+ stringVector.setSafe(3, "cherry".getBytes());
+ stringVector.setSafe(4, "apple".getBytes());
+ stringVector.setSafe(5, "banana".getBytes());
+ stringVector.setSafe(6, "apple".getBytes());
+ stringVector.setSafe(7, "cherry".getBytes());
+ stringVector.setSafe(8, "banana".getBytes());
+ stringVector.setSafe(9, "apple".getBytes());
+ stringVector.setValueCount(10);
+
+ IntVector encodedVector = (IntVector) DictionaryEncoder.encode(stringVector, dictionary);
+
+ // Set up VSR
+ List vectors = Arrays.asList(encodedVector);
+ int rowCount = 10;
+
+ try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
+
+ File dataFile = new File(TMP, "testWriteEnumEncoded.avro");
+
+ // Write an AVRO block using the producer classes
+ try (FileOutputStream fos = new FileOutputStream(dataFile)) {
+ BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null);
+ CompositeAvroProducer producer =
+ ArrowToAvroUtils.createCompositeProducer(vectors, dictionaries);
+ for (int row = 0; row < rowCount; row++) {
+ producer.produce(encoder);
+ }
+ encoder.flush();
+ }
+
+ // Set up reading the AVRO block as a GenericRecord
+ Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields(), dictionaries);
+ GenericDatumReader datumReader = new GenericDatumReader<>(schema);
+
+ try (InputStream inputStream = new FileInputStream(dataFile)) {
+
+ BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null);
+ GenericRecord record = null;
+
+ // Read and check values
+ for (int row = 0; row < rowCount; row++) {
+ record = datumReader.read(record, decoder);
+ // Values read from Avro should be the decoded enum values
+ assertEquals(stringVector.getObject(row).toString(), record.get("enumField").toString());
+ }
+ }
+ }
+ }
}
diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java
index d3e12e763a..d5e0357a8c 100644
--- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java
+++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java
@@ -20,11 +20,18 @@
import java.util.Arrays;
import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.UnionMode;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.avro.LogicalTypes;
@@ -1389,4 +1396,126 @@ public void testConvertUnionTypes() {
Schema.Type.STRING,
schema.getField("nullableDenseUnionField").schema().getTypes().get(3).getType());
}
+
+ @Test
+ public void testWriteDictEnumEncoded() {
+
+ BufferAllocator allocator = new RootAllocator();
+
+ // Create a dictionary
+ FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector dictionaryVector =
+ new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+ dictionaryVector.allocateNew(3);
+ dictionaryVector.set(0, "apple".getBytes());
+ dictionaryVector.set(1, "banana".getBytes());
+ dictionaryVector.set(2, "cherry".getBytes());
+ dictionaryVector.setValueCount(3);
+
+ Dictionary dictionary =
+ new Dictionary(
+ dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+ DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+ List fields =
+ Arrays.asList(
+ new Field(
+ "enumField",
+ new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null),
+ null));
+
+ Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
+
+ assertEquals(Schema.Type.RECORD, schema.getType());
+ assertEquals(1, schema.getFields().size());
+
+ Schema.Field enumField = schema.getField("enumField");
+
+ assertEquals(Schema.Type.ENUM, enumField.schema().getType());
+ assertEquals(3, enumField.schema().getEnumSymbols().size());
+ assertEquals("apple", enumField.schema().getEnumSymbols().get(0));
+ assertEquals("banana", enumField.schema().getEnumSymbols().get(1));
+ assertEquals("cherry", enumField.schema().getEnumSymbols().get(2));
+ }
+
+ @Test
+ public void testWriteDictEnumInvalid() {
+
+ BufferAllocator allocator = new RootAllocator();
+
+ // Create a dictionary
+ FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector dictionaryVector =
+ new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+ dictionaryVector.allocateNew(3);
+ dictionaryVector.set(0, "passion fruit".getBytes());
+ dictionaryVector.set(1, "banana".getBytes());
+ dictionaryVector.set(2, "cherry".getBytes());
+ dictionaryVector.setValueCount(3);
+
+ Dictionary dictionary =
+ new Dictionary(
+ dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+ DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+ List fields =
+ Arrays.asList(
+ new Field(
+ "enumField",
+ new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null),
+ null));
+
+ // Dictionary field contains values that are not valid enums
+ // Should be decoded and output as a string field
+
+ Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
+
+ assertEquals(Schema.Type.RECORD, schema.getType());
+ assertEquals(1, schema.getFields().size());
+
+ Schema.Field enumField = schema.getField("enumField");
+ assertEquals(Schema.Type.STRING, enumField.schema().getType());
+ }
+
+ @Test
+ public void testWriteDictEnumInvalid2() {
+
+ BufferAllocator allocator = new RootAllocator();
+
+ // Create a dictionary
+ FieldType dictionaryField = new FieldType(false, new ArrowType.Int(64, true), null);
+ BigIntVector dictionaryVector =
+ new BigIntVector(new Field("dictionary", dictionaryField, null), allocator);
+
+ dictionaryVector.allocateNew(3);
+ dictionaryVector.set(0, 123L);
+ dictionaryVector.set(1, 456L);
+ dictionaryVector.set(2, 789L);
+ dictionaryVector.setValueCount(3);
+
+ Dictionary dictionary =
+ new Dictionary(
+ dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+ DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+ List fields =
+ Arrays.asList(
+ new Field(
+ "enumField",
+ new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null),
+ null));
+
+ // Dictionary field encodes LONG values rather than STRING
+ // Should be doecded and output as a LONG field
+
+ Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
+
+ assertEquals(Schema.Type.RECORD, schema.getType());
+ assertEquals(1, schema.getFields().size());
+
+ Schema.Field enumField = schema.getField("enumField");
+ assertEquals(Schema.Type.LONG, enumField.schema().getType());
+ }
}
diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java
index 85e6a960b0..ceaf59aa72 100644
--- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java
+++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java
@@ -52,6 +52,7 @@
import org.apache.arrow.vector.TimeStampMilliVector;
import org.apache.arrow.vector.TimeStampNanoTZVector;
import org.apache.arrow.vector.TimeStampNanoVector;
+import org.apache.arrow.vector.TinyIntVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
@@ -60,10 +61,14 @@
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.writer.BaseWriter;
import org.apache.arrow.vector.complex.writer.FieldWriter;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryEncoder;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.avro.Schema;
@@ -78,16 +83,21 @@ public class RoundTripDataTest {
@TempDir public static File TMP;
- private static AvroToArrowConfig basicConfig(BufferAllocator allocator) {
- return new AvroToArrowConfig(allocator, 1000, null, Collections.emptySet(), false);
+ private static AvroToArrowConfig basicConfig(
+ BufferAllocator allocator, DictionaryProvider.MapDictionaryProvider dictionaries) {
+ return new AvroToArrowConfig(allocator, 1000, dictionaries, Collections.emptySet(), false);
}
private static VectorSchemaRoot readDataFile(
- Schema schema, File dataFile, BufferAllocator allocator) throws Exception {
+ Schema schema,
+ File dataFile,
+ BufferAllocator allocator,
+ DictionaryProvider.MapDictionaryProvider dictionaries)
+ throws Exception {
try (FileInputStream fis = new FileInputStream(dataFile)) {
BinaryDecoder decoder = new DecoderFactory().directBinaryDecoder(fis, null);
- return AvroToArrow.avroToArrow(schema, decoder, basicConfig(allocator));
+ return AvroToArrow.avroToArrow(schema, decoder, basicConfig(allocator, dictionaries));
}
}
@@ -95,11 +105,22 @@ private static void roundTripTest(
VectorSchemaRoot root, BufferAllocator allocator, File dataFile, int rowCount)
throws Exception {
+ roundTripTest(root, allocator, dataFile, rowCount, null);
+ }
+
+ private static void roundTripTest(
+ VectorSchemaRoot root,
+ BufferAllocator allocator,
+ File dataFile,
+ int rowCount,
+ DictionaryProvider dictionaries)
+ throws Exception {
+
// Write an AVRO block using the producer classes
try (FileOutputStream fos = new FileOutputStream(dataFile)) {
BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null);
CompositeAvroProducer producer =
- ArrowToAvroUtils.createCompositeProducer(root.getFieldVectors());
+ ArrowToAvroUtils.createCompositeProducer(root.getFieldVectors(), dictionaries);
for (int row = 0; row < rowCount; row++) {
producer.produce(encoder);
}
@@ -107,10 +128,14 @@ private static void roundTripTest(
}
// Generate AVRO schema
- Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields());
+ Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields(), dictionaries);
+
+ DictionaryProvider.MapDictionaryProvider roundTripDictionaries =
+ new DictionaryProvider.MapDictionaryProvider();
// Read back in and compare
- try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator)) {
+ try (VectorSchemaRoot roundTrip =
+ readDataFile(schema, dataFile, allocator, roundTripDictionaries)) {
assertEquals(root.getSchema(), roundTrip.getSchema());
assertEquals(rowCount, roundTrip.getRowCount());
@@ -119,6 +144,21 @@ private static void roundTripTest(
for (int row = 0; row < rowCount; row++) {
assertEquals(root.getVector(0).getObject(row), roundTrip.getVector(0).getObject(row));
}
+
+ if (dictionaries != null) {
+ for (long id : dictionaries.getDictionaryIds()) {
+ Dictionary originalDictionary = dictionaries.lookup(id);
+ Dictionary roundTripDictionary = roundTripDictionaries.lookup(id);
+ assertEquals(
+ originalDictionary.getVector().getValueCount(),
+ roundTripDictionary.getVector().getValueCount());
+ for (int j = 0; j < originalDictionary.getVector().getValueCount(); j++) {
+ assertEquals(
+ originalDictionary.getVector().getObject(j),
+ roundTripDictionary.getVector().getObject(j));
+ }
+ }
+ }
}
}
@@ -141,7 +181,7 @@ private static void roundTripByteArrayTest(
Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields());
// Read back in and compare
- try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator)) {
+ try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator, null)) {
assertEquals(root.getSchema(), roundTrip.getSchema());
assertEquals(rowCount, roundTrip.getRowCount());
@@ -1603,4 +1643,58 @@ public void testRoundTripNullableStructs() throws Exception {
roundTripTest(root, allocator, dataFile, rowCount);
}
}
+
+ @Test
+ public void testRoundTripEnum() throws Exception {
+
+ BufferAllocator allocator = new RootAllocator();
+
+ // Create a dictionary
+ FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector dictionaryVector =
+ new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+ dictionaryVector.allocateNew(3);
+ dictionaryVector.set(0, "apple".getBytes());
+ dictionaryVector.set(1, "banana".getBytes());
+ dictionaryVector.set(2, "cherry".getBytes());
+ dictionaryVector.setValueCount(3);
+
+ // For simplicity, ensure the index type matches what will be decoded during Avro enum decoding
+ Dictionary dictionary =
+ new Dictionary(
+ dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+ DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+ // Field definition
+ FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector stringVector =
+ new VarCharVector(new Field("enumField", stringField, null), allocator);
+ stringVector.allocateNew(10);
+ stringVector.setSafe(0, "apple".getBytes());
+ stringVector.setSafe(1, "banana".getBytes());
+ stringVector.setSafe(2, "cherry".getBytes());
+ stringVector.setSafe(3, "cherry".getBytes());
+ stringVector.setSafe(4, "apple".getBytes());
+ stringVector.setSafe(5, "banana".getBytes());
+ stringVector.setSafe(6, "apple".getBytes());
+ stringVector.setSafe(7, "cherry".getBytes());
+ stringVector.setSafe(8, "banana".getBytes());
+ stringVector.setSafe(9, "apple".getBytes());
+ stringVector.setValueCount(10);
+
+ TinyIntVector encodedVector =
+ (TinyIntVector) DictionaryEncoder.encode(stringVector, dictionary);
+
+ // Set up VSR
+ List vectors = Arrays.asList(encodedVector);
+ int rowCount = 10;
+
+ try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
+
+ File dataFile = new File(TMP, "testRoundTripEnums.avro");
+
+ roundTripTest(root, allocator, dataFile, rowCount, dictionaries);
+ }
+ }
}
diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java
index 864e2c8b59..37c0b4d9fe 100644
--- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java
+++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java
@@ -21,27 +21,50 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.avro.Schema;
+import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
public class RoundTripSchemaTest {
private void doRoundTripTest(List fields) {
+ doRoundTripTest(fields, null);
+ }
- AvroToArrowConfig config = new AvroToArrowConfig(null, 1, null, Collections.emptySet(), false);
+ private void doRoundTripTest(List fields, DictionaryProvider dictionaries) {
- Schema avroSchema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord");
+ DictionaryProvider.MapDictionaryProvider decodeDictionaries =
+ new DictionaryProvider.MapDictionaryProvider();
+ AvroToArrowConfig decodeConfig =
+ new AvroToArrowConfig(null, 1, decodeDictionaries, Collections.emptySet(), false);
+
+ Schema avroSchema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries);
org.apache.arrow.vector.types.pojo.Schema arrowSchema =
- AvroToArrowUtils.createArrowSchema(avroSchema, config);
+ AvroToArrowUtils.createArrowSchema(avroSchema, decodeConfig);
// Compare string representations - equality not defined for logical types
assertEquals(fields, arrowSchema.getFields());
+
+ for (int i = 0; i < fields.size(); i++) {
+ Field field = fields.get(i);
+ Field rtField = arrowSchema.getFields().get(i);
+ if (field.getDictionary() != null) {
+ // Dictionary content is not decoded until the data is consumed
+ Assertions.assertNotNull(rtField.getDictionary());
+ }
+ }
}
// Schema round trip for primitive types, nullable and non-nullable
@@ -440,4 +463,38 @@ public void testRoundTripStructType() {
doRoundTripTest(fields);
}
+
+ @Test
+ public void testRoundTripEnumType() {
+
+ BufferAllocator allocator = new RootAllocator();
+
+ FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null);
+ VarCharVector dictionaryVector =
+ new VarCharVector(new Field("dictionary", dictionaryField, null), allocator);
+
+ dictionaryVector.allocateNew(3);
+ dictionaryVector.set(0, "apple".getBytes());
+ dictionaryVector.set(1, "banana".getBytes());
+ dictionaryVector.set(2, "cherry".getBytes());
+ dictionaryVector.setValueCount(3);
+
+ // For simplicity, ensure the index type matches what will be decoded during Avro enum decoding
+ Dictionary dictionary =
+ new Dictionary(
+ dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true)));
+ DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary);
+
+ List fields =
+ Arrays.asList(
+ new Field(
+ "enumField",
+ new FieldType(
+ true,
+ new ArrowType.Int(8, true),
+ new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))),
+ null));
+
+ doRoundTripTest(fields, dictionaries);
+ }
}
diff --git a/adapter/jdbc/pom.xml b/adapter/jdbc/pom.xml
index f92863fbd6..a0819d7aef 100644
--- a/adapter/jdbc/pom.xml
+++ b/adapter/jdbc/pom.xml
@@ -23,7 +23,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
../../pom.xml
diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java
index 2366116fd0..b8389ee27c 100644
--- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java
+++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java
@@ -24,7 +24,6 @@
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VectorSchemaRoot;
-import org.apache.arrow.vector.types.pojo.ArrowType;
/** Composite consumer which hold all consumers. It manages the consume and cleanup process. */
public class CompositeJdbcConsumer implements JdbcConsumer {
@@ -46,9 +45,9 @@ public void consume(ResultSet rs) throws SQLException, IOException {
BaseConsumer consumer = (BaseConsumer) consumers[i];
JdbcFieldInfo fieldInfo =
new JdbcFieldInfo(rs.getMetaData(), consumer.columnIndexInResultSet);
- ArrowType arrowType = consumer.vector.getMinorType().getType();
+
throw new JdbcConsumerException(
- "Exception while consuming JDBC value", e, fieldInfo, arrowType);
+ "Exception while consuming JDBC value", e, fieldInfo, consumer.vector.getField());
} else {
throw e;
}
diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java
index 04e26d640c..98927f416c 100644
--- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java
+++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java
@@ -17,7 +17,7 @@
package org.apache.arrow.adapter.jdbc.consumer.exceptions;
import org.apache.arrow.adapter.jdbc.JdbcFieldInfo;
-import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
/**
* Exception while consuming JDBC data. This exception stores the JdbcFieldInfo for the column and
@@ -25,7 +25,7 @@
*/
public class JdbcConsumerException extends RuntimeException {
final JdbcFieldInfo fieldInfo;
- final ArrowType arrowType;
+ final Field field;
/**
* Construct JdbcConsumerException with all fields.
@@ -33,17 +33,17 @@ public class JdbcConsumerException extends RuntimeException {
* @param message error message
* @param cause original exception
* @param fieldInfo JdbcFieldInfo for the column
- * @param arrowType ArrowType for the corresponding vector
+ * @param field ArrowType for the corresponding vector
*/
public JdbcConsumerException(
- String message, Throwable cause, JdbcFieldInfo fieldInfo, ArrowType arrowType) {
+ String message, Throwable cause, JdbcFieldInfo fieldInfo, Field field) {
super(message, cause);
this.fieldInfo = fieldInfo;
- this.arrowType = arrowType;
+ this.field = field;
}
- public ArrowType getArrowType() {
- return this.arrowType;
+ public Field getField() {
+ return this.field;
}
public JdbcFieldInfo getFieldInfo() {
diff --git a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java
index c7dc9b2791..e5039ccf59 100644
--- a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java
+++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java
@@ -43,15 +43,19 @@ public void testZeroRowResultSet() throws Exception {
.setReuseVectorSchemaRoot(reuseVectorSchemaRoot)
.build();
- ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config);
- assertTrue(iter.hasNext(), "Iterator on zero row ResultSet should haveNext() before use");
- VectorSchemaRoot root = iter.next();
- assertNotNull(root, "VectorSchemaRoot from first next() result should never be null");
- assertEquals(
- 0, root.getRowCount(), "VectorSchemaRoot from empty ResultSet should have zero rows");
- assertFalse(
- iter.hasNext(),
- "hasNext() should return false on empty ResultSets after initial next() call");
+ try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) {
+ assertTrue(iter.hasNext(), "Iterator on zero row ResultSet should haveNext() before use");
+ VectorSchemaRoot root = iter.next();
+ assertNotNull(root, "VectorSchemaRoot from first next() result should never be null");
+ assertEquals(
+ 0, root.getRowCount(), "VectorSchemaRoot from empty ResultSet should have zero rows");
+ assertFalse(
+ iter.hasNext(),
+ "hasNext() should return false on empty ResultSets after initial next() call");
+ if (!reuseVectorSchemaRoot) {
+ root.close();
+ }
+ }
}
}
}
diff --git a/adapter/orc/pom.xml b/adapter/orc/pom.xml
index e60a7ceb3c..fbb72d6c19 100644
--- a/adapter/orc/pom.xml
+++ b/adapter/orc/pom.xml
@@ -23,7 +23,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
../../pom.xml
@@ -61,7 +61,7 @@ under the License.
org.apache.orc
orc-core
- 2.1.1
+ 2.3.0
test
diff --git a/algorithm/pom.xml b/algorithm/pom.xml
index e934eb7b22..116be9aebf 100644
--- a/algorithm/pom.xml
+++ b/algorithm/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-algorithm
Arrow Algorithms
diff --git a/arrow-variant/pom.xml b/arrow-variant/pom.xml
new file mode 100644
index 0000000000..fea724824b
--- /dev/null
+++ b/arrow-variant/pom.xml
@@ -0,0 +1,51 @@
+
+
+
+ 4.0.0
+
+ org.apache.arrow
+ arrow-java-root
+ 19.0.0
+
+ arrow-variant
+ Arrow Variant
+ Arrow Variant type support.
+
+
+
+ org.apache.arrow
+ arrow-memory-core
+
+
+ org.apache.arrow
+ arrow-vector
+
+
+ org.apache.parquet
+ parquet-variant
+ ${dep.parquet.version}
+
+
+ org.apache.arrow
+ arrow-memory-unsafe
+ test
+
+
+
diff --git a/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java b/arrow-variant/src/main/java/module-info.java
similarity index 69%
rename from vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java
rename to arrow-variant/src/main/java/module-info.java
index 207b0951a7..da94173969 100644
--- a/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java
+++ b/arrow-variant/src/main/java/module-info.java
@@ -14,10 +14,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.arrow.vector.holder;
-import org.apache.arrow.vector.holders.ExtensionHolder;
+@SuppressWarnings("requires-automatic")
+module org.apache.arrow.variant {
+ exports org.apache.arrow.variant;
+ exports org.apache.arrow.variant.extension;
+ exports org.apache.arrow.variant.impl;
+ exports org.apache.arrow.variant.holders;
-public class UuidHolder extends ExtensionHolder {
- public byte[] value;
+ requires org.apache.arrow.memory.core;
+ requires org.apache.arrow.vector;
+ requires parquet.variant;
}
diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java b/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java
new file mode 100644
index 0000000000..fa05cdd93f
--- /dev/null
+++ b/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java
@@ -0,0 +1,217 @@
+/*
+ * 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 java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.util.Objects;
+import java.util.UUID;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+
+/**
+ * Wrapper around parquet-variant's Variant implementation.
+ *
+ * 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 {
+
+ public static final String METADATA_VECTOR_NAME = "metadata";
+ public static final String VALUE_VECTOR_NAME = "value";
+
+ private final Field rootField;
+
+ /**
+ * Constructs a new VariantVector with the given name and allocator.
+ *
+ * @param name the name of the vector
+ * @param allocator the buffer allocator for memory management
+ */
+ public VariantVector(String name, BufferAllocator allocator) {
+ super(
+ name,
+ allocator,
+ new StructVector(
+ name,
+ allocator,
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ null,
+ AbstractStructVector.ConflictPolicy.CONFLICT_ERROR,
+ false));
+ rootField = createVariantField(name);
+ ((FieldVector) this.getUnderlyingVector())
+ .initializeChildrenFromFields(rootField.getChildren());
+ }
+
+ /**
+ * Creates a new VariantVector with the given name. The Variant Field schema has to be the same
+ * everywhere, otherwise ArrowBuffer loading might fail during serialization/deserialization and
+ * schema mismatches can occur. This includes CompleteType's VARIANT and VARIANT_REQUIRED types.
+ */
+ public static Field createVariantField(String name) {
+ return new Field(
+ name, new FieldType(true, VariantType.INSTANCE, null), createVariantChildFields());
+ }
+
+ /**
+ * Creates the child fields for the VariantVector. Metadata vector will be index 0 and value
+ * vector will be index 1.
+ */
+ public static List createVariantChildFields() {
+ return List.of(
+ new Field(METADATA_VECTOR_NAME, new FieldType(false, Binary.INSTANCE, null), null),
+ new Field(VALUE_VECTOR_NAME, new FieldType(false, Binary.INSTANCE, null), null));
+ }
+
+ @Override
+ public void initializeChildrenFromFields(List children) {
+ // No-op, as children are initialized in the constructor
+ }
+
+ @Override
+ public Field getField() {
+ return rootField;
+ }
+
+ public VarBinaryVector getMetadataVector() {
+ return getUnderlyingVector().getChild(METADATA_VECTOR_NAME, VarBinaryVector.class);
+ }
+
+ public VarBinaryVector getValueVector() {
+ return getUnderlyingVector().getChild(VALUE_VECTOR_NAME, VarBinaryVector.class);
+ }
+
+ @Override
+ public TransferPair makeTransferPair(ValueVector target) {
+ return new VariantTransferPair(this, (VariantVector) target);
+ }
+
+ @Override
+ public TransferPair getTransferPair(Field field, BufferAllocator allocator) {
+ return new VariantTransferPair(this, new VariantVector(field.getName(), 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 VariantTransferPair(this, new VariantVector(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 void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
+ getUnderlyingVector()
+ .copyFrom(fromIndex, thisIndex, ((VariantVector) from).getUnderlyingVector());
+ }
+
+ @Override
+ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
+ getUnderlyingVector()
+ .copyFromSafe(fromIndex, thisIndex, ((VariantVector) from).getUnderlyingVector());
+ }
+
+ @Override
+ public Object getObject(int index) {
+ if (isNull(index)) {
+ return null;
+ }
+ VarBinaryVector metadataVector = getMetadataVector();
+ VarBinaryVector valueVector = getValueVector();
+
+ int metadataStart = metadataVector.getStartOffset(index);
+ int metadataEnd = metadataVector.getEndOffset(index);
+ int valueStart = valueVector.getStartOffset(index);
+ int valueEnd = valueVector.getEndOffset(index);
+
+ return new Variant(
+ metadataVector.getDataBuffer(),
+ metadataStart,
+ metadataEnd,
+ valueVector.getDataBuffer(),
+ valueStart,
+ valueEnd);
+ }
+
+ /**
+ * Retrieves the variant value at the specified index into the provided holder.
+ *
+ * @param index the index of the value to retrieve
+ * @param holder the holder to populate with the variant data
+ */
+ public void get(int index, NullableVariantHolder holder) {
+ if (isNull(index)) {
+ holder.isSet = 0;
+ } else {
+ holder.isSet = 1;
+ VarBinaryVector metadataVector = getMetadataVector();
+ VarBinaryVector valueVector = getValueVector();
+ assert !metadataVector.isNull(index) && !valueVector.isNull(index);
+
+ holder.metadataStart = metadataVector.getStartOffset(index);
+ holder.metadataEnd = metadataVector.getEndOffset(index);
+ holder.metadataBuffer = metadataVector.getDataBuffer();
+ holder.valueStart = valueVector.getStartOffset(index);
+ holder.valueEnd = valueVector.getEndOffset(index);
+ holder.valueBuffer = valueVector.getDataBuffer();
+ }
+ }
+
+ /**
+ * Retrieves the variant value at the specified index into the provided non-nullable holder.
+ *
+ * @param index the index of the value to retrieve
+ * @param holder the holder to populate with the variant data
+ */
+ public void get(int index, VariantHolder holder) {
+ VarBinaryVector metadataVector = getMetadataVector();
+ VarBinaryVector valueVector = getValueVector();
+ assert !metadataVector.isNull(index) && !valueVector.isNull(index);
+
+ holder.metadataStart = metadataVector.getStartOffset(index);
+ holder.metadataEnd = metadataVector.getEndOffset(index);
+ holder.metadataBuffer = metadataVector.getDataBuffer();
+ holder.valueStart = valueVector.getStartOffset(index);
+ holder.valueEnd = valueVector.getEndOffset(index);
+ holder.valueBuffer = valueVector.getDataBuffer();
+ }
+
+ /**
+ * Sets the variant value at the specified index from the provided holder.
+ *
+ * @param index the index at which to set the value
+ * @param holder the holder containing the variant data to set
+ */
+ public void set(int index, VariantHolder holder) {
+ BitVectorHelper.setBit(getUnderlyingVector().getValidityBuffer(), index);
+ getMetadataVector()
+ .set(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer);
+ getValueVector().set(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer);
+ }
+
+ /**
+ * Sets the variant value at the specified index from the provided nullable holder.
+ *
+ * @param index the index at which to set the value
+ * @param holder the nullable holder containing the variant data to set
+ */
+ public void set(int index, NullableVariantHolder holder) {
+ BitVectorHelper.setValidityBit(getUnderlyingVector().getValidityBuffer(), index, holder.isSet);
+ if (holder.isSet == 0) {
+ return;
+ }
+ getMetadataVector()
+ .set(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer);
+ getValueVector().set(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer);
+ }
+
+ /**
+ * Sets the variant value at the specified index from the provided holder, with bounds checking.
+ *
+ * @param index the index at which to set the value
+ * @param holder the holder containing the variant data to set
+ */
+ public void setSafe(int index, VariantHolder holder) {
+ getUnderlyingVector().setIndexDefined(index);
+ getMetadataVector()
+ .setSafe(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer);
+ getValueVector().setSafe(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer);
+ }
+
+ /**
+ * Sets the variant value at the specified index from the provided nullable holder, with bounds
+ * checking.
+ *
+ * @param index the index at which to set the value
+ * @param holder the nullable holder containing the variant data to set
+ */
+ public void setSafe(int index, NullableVariantHolder holder) {
+ if (holder.isSet == 0) {
+ getUnderlyingVector().setNull(index);
+ return;
+ }
+ getUnderlyingVector().setIndexDefined(index);
+ getMetadataVector()
+ .setSafe(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer);
+ getValueVector().setSafe(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer);
+ }
+
+ /** Sets the value at the given index from the provided Variant. */
+ public void setSafe(int index, Variant variant) {
+ ByteBuffer metadataBuffer = variant.getMetadataBuffer();
+ ByteBuffer valueBuffer = variant.getValueBuffer();
+ int metadataLength = metadataBuffer.remaining();
+ int valueLength = valueBuffer.remaining();
+ try (ArrowBuf metaBuf = getAllocator().buffer(metadataLength);
+ ArrowBuf valBuf = getAllocator().buffer(valueLength)) {
+ metaBuf.setBytes(0, metadataBuffer.duplicate());
+ valBuf.setBytes(0, valueBuffer.duplicate());
+ getUnderlyingVector().setIndexDefined(index);
+ getMetadataVector().setSafe(index, 1, 0, metadataLength, metaBuf);
+ getValueVector().setSafe(index, 1, 0, valueLength, valBuf);
+ }
+ }
+
+ @Override
+ protected FieldReader getReaderImpl() {
+ return new org.apache.arrow.variant.impl.VariantReaderImpl(this);
+ }
+
+ @Override
+ public int hashCode(int index) {
+ return hashCode(index, null);
+ }
+
+ @Override
+ public int hashCode(int index, ArrowBufHasher hasher) {
+ return getUnderlyingVector().hashCode(index, hasher);
+ }
+
+ /**
+ * VariantTransferPair is a transfer pair for VariantVector. It transfers the metadata and value
+ * together using the underlyingVector's transfer pair.
+ */
+ protected static class VariantTransferPair implements TransferPair {
+ private final TransferPair pair;
+ private final VariantVector from;
+ private final VariantVector to;
+
+ public VariantTransferPair(VariantVector from, VariantVector to) {
+ this.from = from;
+ this.to = to;
+ this.pair = from.getUnderlyingVector().makeTransferPair((to).getUnderlyingVector());
+ }
+
+ @Override
+ public void transfer() {
+ pair.transfer();
+ }
+
+ @Override
+ public void splitAndTransfer(int startIndex, int length) {
+ pair.splitAndTransfer(startIndex, length);
+ }
+
+ @Override
+ public ValueVector getTo() {
+ return to;
+ }
+
+ @Override
+ public void copyValueSafe(int from, int to) {
+ pair.copyValueSafe(from, to);
+ }
+ }
+}
diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/holders/NullableVariantHolder.java b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/NullableVariantHolder.java
new file mode 100644
index 0000000000..b78d4a2013
--- /dev/null
+++ b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/NullableVariantHolder.java
@@ -0,0 +1,56 @@
+/*
+ * 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.holders;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.variant.extension.VariantType;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
+@SuppressWarnings("checkstyle:VisibilityModifier")
+public final class NullableVariantHolder extends ExtensionHolder {
+
+ public int isSet;
+ public int metadataStart;
+ public int metadataEnd;
+ public ArrowBuf metadataBuffer;
+ public int valueStart;
+ public int valueEnd;
+ public ArrowBuf valueBuffer;
+
+ public NullableVariantHolder() {}
+
+ @Override
+ public boolean equals(Object obj) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int hashCode() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public String toString() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public ArrowType type() {
+ return VariantType.INSTANCE;
+ }
+}
diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java
similarity index 51%
rename from vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java
rename to arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java
index 68029b1df5..e3947ac439 100644
--- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java
+++ b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java
@@ -14,34 +14,43 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.arrow.vector.complex.impl;
+package org.apache.arrow.variant.holders;
-import java.nio.ByteBuffer;
-import java.util.UUID;
-import org.apache.arrow.vector.UuidVector;
-import org.apache.arrow.vector.holder.UuidHolder;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.variant.extension.VariantType;
import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.pojo.ArrowType;
-public class UuidWriterImpl extends AbstractExtensionTypeWriter {
+@SuppressWarnings("checkstyle:VisibilityModifier")
+public final class VariantHolder extends ExtensionHolder {
- public UuidWriterImpl(UuidVector vector) {
- super(vector);
+ public final int isSet = 1;
+ public int metadataStart;
+ public int metadataEnd;
+ public ArrowBuf metadataBuffer;
+ public int valueStart;
+ public int valueEnd;
+ public ArrowBuf valueBuffer;
+
+ public VariantHolder() {}
+
+ @Override
+ public boolean equals(Object obj) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int hashCode() {
+ throw new UnsupportedOperationException();
}
@Override
- public void writeExtension(Object value) {
- UUID uuid = (UUID) value;
- ByteBuffer bb = ByteBuffer.allocate(16);
- bb.putLong(uuid.getMostSignificantBits());
- bb.putLong(uuid.getLeastSignificantBits());
- vector.setSafe(getPosition(), bb.array());
- vector.setValueCount(getPosition() + 1);
+ public String toString() {
+ throw new UnsupportedOperationException();
}
@Override
- public void write(ExtensionHolder holder) {
- UuidHolder uuidHolder = (UuidHolder) holder;
- vector.setSafe(getPosition(), uuidHolder.value);
- vector.setValueCount(getPosition() + 1);
+ public ArrowType type() {
+ return VariantType.INSTANCE;
}
}
diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/impl/NullableVariantHolderReaderImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/NullableVariantHolderReaderImpl.java
new file mode 100644
index 0000000000..1645529c0c
--- /dev/null
+++ b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/NullableVariantHolderReaderImpl.java
@@ -0,0 +1,69 @@
+/*
+ * 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.impl;
+
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+import org.apache.arrow.vector.complex.impl.AbstractFieldReader;
+import org.apache.arrow.vector.types.Types;
+
+public class NullableVariantHolderReaderImpl extends AbstractFieldReader {
+ private final NullableVariantHolder holder;
+
+ public NullableVariantHolderReaderImpl(NullableVariantHolder holder) {
+ this.holder = holder;
+ }
+
+ @Override
+ public int size() {
+ throw new UnsupportedOperationException("You can't call size on a Holder value reader.");
+ }
+
+ @Override
+ public boolean next() {
+ throw new UnsupportedOperationException("You can't call next on a single value reader.");
+ }
+
+ @Override
+ public void setPosition(int index) {
+ throw new UnsupportedOperationException("You can't call setPosition on a single value reader.");
+ }
+
+ @Override
+ public Types.MinorType getMinorType() {
+ return Types.MinorType.EXTENSIONTYPE;
+ }
+
+ @Override
+ public boolean isSet() {
+ return holder.isSet == 1;
+ }
+
+ /**
+ * Reads the variant holder data into the provided holder.
+ *
+ * @param h the holder to read into
+ */
+ public void read(NullableVariantHolder h) {
+ h.metadataStart = this.holder.metadataStart;
+ h.metadataEnd = this.holder.metadataEnd;
+ h.metadataBuffer = this.holder.metadataBuffer;
+ h.valueStart = this.holder.valueStart;
+ h.valueEnd = this.holder.valueEnd;
+ h.valueBuffer = this.holder.valueBuffer;
+ h.isSet = this.isSet() ? 1 : 0;
+ }
+}
diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantReaderImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantReaderImpl.java
new file mode 100644
index 0000000000..670104b7d1
--- /dev/null
+++ b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantReaderImpl.java
@@ -0,0 +1,73 @@
+/*
+ * 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.impl;
+
+import org.apache.arrow.variant.extension.VariantVector;
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+import org.apache.arrow.variant.holders.VariantHolder;
+import org.apache.arrow.vector.complex.impl.AbstractFieldReader;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.Field;
+
+public class VariantReaderImpl extends AbstractFieldReader {
+ private final VariantVector vector;
+
+ public VariantReaderImpl(VariantVector vector) {
+ this.vector = vector;
+ }
+
+ @Override
+ public Types.MinorType getMinorType() {
+ return this.vector.getMinorType();
+ }
+
+ @Override
+ public Field getField() {
+ return this.vector.getField();
+ }
+
+ @Override
+ public boolean isSet() {
+ return !this.vector.isNull(this.idx());
+ }
+
+ @Override
+ public void read(ExtensionHolder holder) {
+ if (holder instanceof VariantHolder) {
+ vector.get(idx(), (VariantHolder) holder);
+ } else if (holder instanceof NullableVariantHolder) {
+ vector.get(idx(), (NullableVariantHolder) holder);
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported holder type for VariantReader: " + holder.getClass());
+ }
+ }
+
+ public void read(VariantHolder h) {
+ this.vector.get(this.idx(), h);
+ }
+
+ public void read(NullableVariantHolder h) {
+ this.vector.get(this.idx(), h);
+ }
+
+ @Override
+ public Object readObject() {
+ return this.vector.getObject(this.idx());
+ }
+}
diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantWriterImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantWriterImpl.java
new file mode 100644
index 0000000000..266ddb75d2
--- /dev/null
+++ b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantWriterImpl.java
@@ -0,0 +1,121 @@
+/*
+ * 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.impl;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.variant.Variant;
+import org.apache.arrow.variant.extension.VariantVector;
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+import org.apache.arrow.variant.holders.VariantHolder;
+import org.apache.arrow.vector.complex.impl.AbstractExtensionTypeWriter;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
+/**
+ * Writer implementation for VARIANT extension type vectors.
+ *
+ * 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 {
+
+ private static final String UNSUPPORTED_TYPE_TEMPLATE = "Unsupported type for Variant: %s";
+
+ /**
+ * Constructs a new VariantWriterImpl for the given vector.
+ *
+ * @param vector the variant vector to write to
+ */
+ public VariantWriterImpl(VariantVector vector) {
+ super(vector);
+ }
+
+ /**
+ * Writes an extension type or variant value to the vector.
+ *
+ * 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 resultSet = (ArrayList) inVector.getObject(0);
+ assertEquals(2, resultSet.size());
+ assertEquals(variant1, resultSet.get(0));
+ assertEquals(variant2, resultSet.get(1));
+ }
+ }
+
+ @Test
+ public void testListVectorReaderForVariantExtensionType() {
+ try (ListVector inVector = ListVector.empty("input", allocator)) {
+ 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);
+ writer.endList();
+
+ writer.setPosition(1);
+ writer.startList();
+ extensionWriter.writeExtension(variant2);
+ extensionWriter.writeExtension(variant2);
+ writer.endList();
+
+ inVector.setValueCount(2);
+
+ UnionListReader reader = inVector.getReader();
+ reader.setPosition(0);
+ assertTrue(reader.next());
+ FieldReader variantReader = reader.reader();
+ NullableVariantHolder resultHolder = new NullableVariantHolder();
+ variantReader.read(resultHolder);
+ assertEquals(variant1, new Variant(resultHolder));
+
+ reader.setPosition(1);
+ assertTrue(reader.next());
+ variantReader = reader.reader();
+ variantReader.read(resultHolder);
+ assertEquals(variant2, new Variant(resultHolder));
+
+ assertTrue(reader.next());
+ variantReader = reader.reader();
+ variantReader.read(resultHolder);
+ assertEquals(variant2, new Variant(resultHolder));
+ }
+ }
+
+ @Test
+ public void testCopyFromForVariantExtensionType() {
+ try (ListVector inVector = ListVector.empty("input", allocator);
+ ListVector outVector = ListVector.empty("output", allocator)) {
+ 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);
+ writer.endList();
+
+ writer.setPosition(1);
+ writer.startList();
+ extensionWriter.writeExtension(variant2);
+ extensionWriter.writeExtension(variant2);
+ writer.endList();
+
+ inVector.setValueCount(2);
+
+ outVector.allocateNew();
+ outVector.copyFrom(0, 0, inVector);
+ outVector.copyFrom(1, 1, inVector);
+ outVector.setValueCount(2);
+
+ ArrayList resultSet0 = (ArrayList) outVector.getObject(0);
+ assertEquals(1, resultSet0.size());
+ assertEquals(variant1, resultSet0.get(0));
+
+ ArrayList resultSet1 = (ArrayList) outVector.getObject(1);
+ assertEquals(2, resultSet1.size());
+ assertEquals(variant2, resultSet1.get(0));
+ assertEquals(variant2, resultSet1.get(1));
+ }
+ }
+
+ @Test
+ public void testCopyValueSafeForVariantExtensionType() {
+ try (ListVector inVector = ListVector.empty("input", allocator)) {
+ 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);
+ writer.endList();
+
+ writer.setPosition(1);
+ writer.startList();
+ extensionWriter.writeExtension(variant2);
+ extensionWriter.writeExtension(variant2);
+ writer.endList();
+
+ inVector.setValueCount(2);
+
+ try (ListVector outVector = (ListVector) inVector.getTransferPair(allocator).getTo()) {
+ TransferPair tp = inVector.makeTransferPair(outVector);
+ tp.copyValueSafe(0, 0);
+ tp.copyValueSafe(1, 1);
+ outVector.setValueCount(2);
+
+ ArrayList resultSet0 = (ArrayList) outVector.getObject(0);
+ assertEquals(1, resultSet0.size());
+ assertEquals(variant1, resultSet0.get(0));
+
+ ArrayList resultSet1 = (ArrayList) outVector.getObject(1);
+ assertEquals(2, resultSet1.size());
+ assertEquals(variant2, resultSet1.get(0));
+ assertEquals(variant2, resultSet1.get(1));
+ }
+ }
+ }
+}
diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInMapVector.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInMapVector.java
new file mode 100644
index 0000000000..dd925810de
--- /dev/null
+++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInMapVector.java
@@ -0,0 +1,125 @@
+/*
+ * 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 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.MapVector;
+import org.apache.arrow.vector.complex.impl.UnionMapReader;
+import org.apache.arrow.vector.complex.impl.UnionMapWriter;
+import org.apache.arrow.vector.complex.reader.FieldReader;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class TestVariantInMapVector {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ public void init() {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void terminate() {
+ allocator.close();
+ }
+
+ @Test
+ public void testMapVectorWithVariantExtensionType() {
+ Variant variant1 = TestVariant.variantString("hello");
+ Variant variant2 = TestVariant.variantString("world");
+ try (final MapVector inVector = MapVector.empty("map", allocator, false)) {
+ inVector.allocateNew();
+ UnionMapWriter writer = inVector.getWriter();
+ writer.setPosition(0);
+
+ writer.startMap();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(0);
+ writer.value().extension(VariantType.INSTANCE).writeExtension(variant1, VariantType.INSTANCE);
+ writer.endEntry();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(1);
+ writer.value().extension(VariantType.INSTANCE).writeExtension(variant2, VariantType.INSTANCE);
+ writer.endEntry();
+ writer.endMap();
+
+ writer.setValueCount(1);
+
+ UnionMapReader mapReader = inVector.getReader();
+ mapReader.setPosition(0);
+ mapReader.next();
+ FieldReader variantReader = mapReader.value();
+ NullableVariantHolder holder = new NullableVariantHolder();
+ variantReader.read(holder);
+ assertEquals(variant1, new Variant(holder));
+
+ mapReader.next();
+ variantReader = mapReader.value();
+ variantReader.read(holder);
+ assertEquals(variant2, new Variant(holder));
+ }
+ }
+
+ @Test
+ public void testCopyFromForVariantExtensionType() {
+ Variant variant1 = TestVariant.variantString("hello");
+ Variant variant2 = TestVariant.variantString("world");
+ try (final MapVector inVector = MapVector.empty("in", allocator, false);
+ final MapVector outVector = MapVector.empty("out", allocator, false)) {
+ inVector.allocateNew();
+ UnionMapWriter writer = inVector.getWriter();
+ writer.setPosition(0);
+
+ writer.startMap();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(0);
+ writer.value().extension(VariantType.INSTANCE).writeExtension(variant1, VariantType.INSTANCE);
+ writer.endEntry();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(1);
+ writer.value().extension(VariantType.INSTANCE).writeExtension(variant2, VariantType.INSTANCE);
+ writer.endEntry();
+ writer.endMap();
+
+ writer.setValueCount(1);
+ outVector.allocateNew();
+ outVector.copyFrom(0, 0, inVector);
+ outVector.setValueCount(1);
+
+ UnionMapReader mapReader = outVector.getReader();
+ mapReader.setPosition(0);
+ mapReader.next();
+ FieldReader variantReader = mapReader.value();
+ NullableVariantHolder holder = new NullableVariantHolder();
+ variantReader.read(holder);
+ assertEquals(variant1, new Variant(holder));
+
+ mapReader.next();
+ variantReader = mapReader.value();
+ variantReader.read(holder);
+ assertEquals(variant2, new Variant(holder));
+ }
+ }
+}
diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantType.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantType.java
new file mode 100644
index 0000000000..017e71224b
--- /dev/null
+++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantType.java
@@ -0,0 +1,308 @@
+/*
+ * 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.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.ipc.ArrowStreamReader;
+import org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+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.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class TestVariantType {
+ BufferAllocator allocator;
+
+ @BeforeEach
+ void beforeEach() {
+ allocator = new RootAllocator();
+ }
+
+ @AfterEach
+ void afterEach() {
+ allocator.close();
+ }
+
+ @Test
+ void testConstants() {
+ assertNotNull(VariantType.INSTANCE);
+ }
+
+ @Test
+ void testStorageType() {
+ VariantType type = VariantType.INSTANCE;
+ assertEquals(ArrowType.Struct.INSTANCE, type.storageType());
+ assertInstanceOf(ArrowType.Struct.class, type.storageType());
+ }
+
+ @Test
+ void testExtensionName() {
+ VariantType type = VariantType.INSTANCE;
+ assertEquals("parquet.variant", type.extensionName());
+ }
+
+ @Test
+ void testExtensionEquals() {
+ VariantType type1 = VariantType.INSTANCE;
+ VariantType type2 = VariantType.INSTANCE;
+
+ assertTrue(type1.extensionEquals(type2));
+ }
+
+ @Test
+ void testIsComplex() {
+ VariantType type = VariantType.INSTANCE;
+ assertFalse(type.isComplex());
+ }
+
+ @Test
+ void testSerialize() {
+ VariantType type = VariantType.INSTANCE;
+ String serialized = type.serialize();
+ assertEquals("", serialized);
+ }
+
+ @Test
+ void testDeserializeValid() {
+ VariantType type = VariantType.INSTANCE;
+ ArrowType storageType = ArrowType.Struct.INSTANCE;
+
+ ArrowType deserialized = assertDoesNotThrow(() -> type.deserialize(storageType, ""));
+ assertInstanceOf(VariantType.class, deserialized);
+ assertEquals(VariantType.INSTANCE, deserialized);
+ }
+
+ @Test
+ void testDeserializeInvalidStorageType() {
+ VariantType type = VariantType.INSTANCE;
+ ArrowType wrongStorageType = ArrowType.Utf8.INSTANCE;
+
+ assertThrows(UnsupportedOperationException.class, () -> type.deserialize(wrongStorageType, ""));
+ }
+
+ @Test
+ void testGetNewVector() {
+ VariantType type = VariantType.INSTANCE;
+ try (FieldVector vector =
+ type.getNewVector("variant_field", FieldType.nullable(type), allocator)) {
+ assertInstanceOf(VariantVector.class, vector);
+ assertEquals("variant_field", vector.getField().getName());
+ assertEquals(type, vector.getField().getType());
+ }
+ }
+
+ @Test
+ void testGetNewVectorWithNullableFieldType() {
+ VariantType type = VariantType.INSTANCE;
+ FieldType nullableFieldType = FieldType.nullable(type);
+
+ try (FieldVector vector = type.getNewVector("nullable_variant", nullableFieldType, allocator)) {
+ assertInstanceOf(VariantVector.class, vector);
+ assertEquals("nullable_variant", vector.getField().getName());
+ assertTrue(vector.getField().isNullable());
+ }
+ }
+
+ @Test
+ void testGetNewVectorWithNonNullableFieldType() {
+ VariantType type = VariantType.INSTANCE;
+ FieldType nonNullableFieldType = FieldType.notNullable(type);
+
+ try (FieldVector vector =
+ type.getNewVector("non_nullable_variant", nonNullableFieldType, allocator)) {
+ assertInstanceOf(VariantVector.class, vector);
+ assertEquals("non_nullable_variant", vector.getField().getName());
+ }
+ }
+
+ @Test
+ void testIpcRoundTrip() {
+ VariantType type = VariantType.INSTANCE;
+
+ Schema schema = new Schema(Collections.singletonList(Field.nullable("variant", type)));
+ byte[] serialized = schema.serializeAsMessage();
+ Schema deserialized = Schema.deserializeMessage(ByteBuffer.wrap(serialized));
+ assertEquals(schema, deserialized);
+ }
+
+ @Test
+ void testVectorIpcRoundTrip() throws IOException {
+ VariantType type = VariantType.INSTANCE;
+
+ try (FieldVector vector = type.getNewVector("field", FieldType.nullable(type), allocator);
+ ArrowBuf metadataBuf1 = allocator.buffer(10);
+ ArrowBuf valueBuf1 = allocator.buffer(10);
+ ArrowBuf metadataBuf2 = allocator.buffer(10);
+ ArrowBuf valueBuf2 = allocator.buffer(10)) {
+ VariantVector variantVector = (VariantVector) vector;
+
+ byte[] metadata1 = new byte[] {1, 2, 3};
+ byte[] value1 = new byte[] {4, 5, 6, 7};
+ metadataBuf1.setBytes(0, metadata1);
+ valueBuf1.setBytes(0, value1);
+
+ byte[] metadata2 = new byte[] {8, 9};
+ byte[] value2 = new byte[] {10, 11, 12};
+ metadataBuf2.setBytes(0, metadata2);
+ valueBuf2.setBytes(0, value2);
+
+ NullableVariantHolder holder1 = new NullableVariantHolder();
+ holder1.isSet = 1;
+ holder1.metadataStart = 0;
+ holder1.metadataEnd = metadata1.length;
+ holder1.metadataBuffer = metadataBuf1;
+ holder1.valueStart = 0;
+ holder1.valueEnd = value1.length;
+ holder1.valueBuffer = valueBuf1;
+
+ NullableVariantHolder holder2 = new NullableVariantHolder();
+ holder2.isSet = 1;
+ holder2.metadataStart = 0;
+ holder2.metadataEnd = metadata2.length;
+ holder2.metadataBuffer = metadataBuf2;
+ holder2.valueStart = 0;
+ holder2.valueEnd = value2.length;
+ holder2.valueBuffer = valueBuf2;
+
+ variantVector.setSafe(0, holder1);
+ variantVector.setNull(1);
+ variantVector.setSafe(2, holder2);
+ variantVector.setValueCount(3);
+
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (VectorSchemaRoot root = new VectorSchemaRoot(Collections.singletonList(variantVector));
+ ArrowStreamWriter writer =
+ new ArrowStreamWriter(root, new DictionaryProvider.MapDictionaryProvider(), baos)) {
+ writer.start();
+ writer.writeBatch();
+ }
+
+ try (ArrowStreamReader reader =
+ new ArrowStreamReader(new ByteArrayInputStream(baos.toByteArray()), allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ assertEquals(3, root.getRowCount());
+ assertEquals(
+ new Schema(Collections.singletonList(variantVector.getField())), root.getSchema());
+
+ VariantVector actual = assertInstanceOf(VariantVector.class, root.getVector("field"));
+ assertFalse(actual.isNull(0));
+ assertTrue(actual.isNull(1));
+ assertFalse(actual.isNull(2));
+
+ NullableVariantHolder result1 = new NullableVariantHolder();
+ actual.get(0, result1);
+ assertEquals(1, result1.isSet);
+ assertEquals(metadata1.length, result1.metadataEnd - result1.metadataStart);
+ assertEquals(value1.length, result1.valueEnd - result1.valueStart);
+
+ assertNull(actual.getObject(1));
+
+ NullableVariantHolder result2 = new NullableVariantHolder();
+ actual.get(2, result2);
+ assertEquals(1, result2.isSet);
+ assertEquals(metadata2.length, result2.metadataEnd - result2.metadataStart);
+ assertEquals(value2.length, result2.valueEnd - result2.valueStart);
+ }
+ }
+ }
+
+ @Test
+ void testSingleton() {
+ VariantType type1 = VariantType.INSTANCE;
+ VariantType type2 = VariantType.INSTANCE;
+
+ // Same instance
+ assertSame(type1, type2);
+ assertTrue(type1.extensionEquals(type2));
+ }
+
+ @Test
+ void testExtensionTypeRegistry() {
+ // VariantType should be automatically registered via static initializer
+ ArrowType.ExtensionType registeredType =
+ ExtensionTypeRegistry.lookup(VariantType.EXTENSION_NAME);
+ assertNotNull(registeredType);
+ assertInstanceOf(VariantType.class, registeredType);
+ assertEquals(VariantType.INSTANCE, registeredType);
+ }
+
+ @Test
+ void testFieldMetadata() {
+ Map metadata = new HashMap<>();
+ metadata.put("key1", "value1");
+ metadata.put("key2", "value2");
+
+ FieldType fieldType = new FieldType(true, VariantType.INSTANCE, null, metadata);
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ Field field = new Field("test", fieldType, VariantVector.createVariantChildFields());
+
+ // Field metadata includes both custom metadata and extension type metadata
+ Map fieldMetadata = field.getMetadata();
+ assertEquals("value1", fieldMetadata.get("key1"));
+ assertEquals("value2", fieldMetadata.get("key2"));
+ // Extension type metadata is also present
+ assertTrue(fieldMetadata.containsKey("ARROW:extension:name"));
+ assertTrue(fieldMetadata.containsKey("ARROW:extension:metadata"));
+ }
+ }
+
+ @Test
+ void testFieldChildren() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ Field field = vector.getField();
+
+ assertNotNull(field.getChildren());
+ assertEquals(2, field.getChildren().size());
+
+ Field metadataField = field.getChildren().get(0);
+ assertEquals(VariantVector.METADATA_VECTOR_NAME, metadataField.getName());
+ assertEquals(ArrowType.Binary.INSTANCE, metadataField.getType());
+
+ Field valueField = field.getChildren().get(1);
+ assertEquals(VariantVector.VALUE_VECTOR_NAME, valueField.getName());
+ assertEquals(ArrowType.Binary.INSTANCE, valueField.getType());
+ }
+ }
+}
diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantVector.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantVector.java
new file mode 100644
index 0000000000..1c172e304f
--- /dev/null
+++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantVector.java
@@ -0,0 +1,844 @@
+/*
+ * 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.assertArrayEquals;
+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.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.variant.Variant;
+import org.apache.arrow.variant.holders.NullableVariantHolder;
+import org.apache.arrow.variant.holders.VariantHolder;
+import org.apache.arrow.variant.impl.VariantReaderImpl;
+import org.apache.arrow.variant.impl.VariantWriterImpl;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for VariantVector, VariantWriterImpl, and VariantReaderImpl. */
+class TestVariantVector {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ void beforeEach() {
+ allocator = new RootAllocator();
+ }
+
+ @AfterEach
+ void afterEach() {
+ allocator.close();
+ }
+
+ private VariantHolder createHolder(
+ ArrowBuf metadataBuf, byte[] metadata, ArrowBuf valueBuf, byte[] value) {
+ VariantHolder holder = new VariantHolder();
+ holder.metadataStart = 0;
+ holder.metadataEnd = metadata.length;
+ holder.metadataBuffer = metadataBuf;
+ holder.valueStart = 0;
+ holder.valueEnd = value.length;
+ holder.valueBuffer = valueBuf;
+ return holder;
+ }
+
+ private NullableVariantHolder createNullableHolder(
+ ArrowBuf metadataBuf, byte[] metadata, ArrowBuf valueBuf, byte[] value) {
+ NullableVariantHolder holder = new NullableVariantHolder();
+ holder.isSet = 1;
+ holder.metadataStart = 0;
+ holder.metadataEnd = metadata.length;
+ holder.metadataBuffer = metadataBuf;
+ holder.valueStart = 0;
+ holder.valueEnd = value.length;
+ holder.valueBuffer = valueBuf;
+ return holder;
+ }
+
+ private NullableVariantHolder createNullHolder() {
+ NullableVariantHolder holder = new NullableVariantHolder();
+ holder.isSet = 0;
+ return holder;
+ }
+
+ // ========== Basic Vector Tests ==========
+
+ @Test
+ void testVectorCreation() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ assertNotNull(vector);
+ assertEquals("test", vector.getField().getName());
+ assertNotNull(vector.getMetadataVector());
+ assertNotNull(vector.getValueVector());
+ }
+ }
+
+ @Test
+ void testSetAndGet() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2, 3};
+ byte[] value = new byte[] {4, 5, 6, 7};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ // Retrieve and verify
+ NullableVariantHolder result = new NullableVariantHolder();
+ vector.get(0, result);
+
+ assertEquals(1, result.isSet);
+ assertEquals(metadata.length, result.metadataEnd - result.metadataStart);
+ assertEquals(value.length, result.valueEnd - result.valueStart);
+
+ byte[] actualMetadata = new byte[metadata.length];
+ byte[] actualValue = new byte[value.length];
+ result.metadataBuffer.getBytes(result.metadataStart, actualMetadata);
+ result.valueBuffer.getBytes(result.valueStart, actualValue);
+
+ assertArrayEquals(metadata, actualMetadata);
+ assertArrayEquals(value, actualValue);
+ }
+ }
+
+ @Test
+ void testSetNull() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ NullableVariantHolder holder = createNullHolder();
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ assertTrue(vector.isNull(0));
+
+ NullableVariantHolder result = new NullableVariantHolder();
+ vector.get(0, result);
+ assertEquals(0, result.isSet);
+ }
+ }
+
+ @Test
+ void testMultipleValues() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf1 = allocator.buffer(10);
+ ArrowBuf valueBuf1 = allocator.buffer(10);
+ ArrowBuf metadataBuf2 = allocator.buffer(10);
+ ArrowBuf valueBuf2 = allocator.buffer(10)) {
+
+ byte[] metadata1 = new byte[] {1, 2};
+ byte[] value1 = new byte[] {3, 4, 5};
+ metadataBuf1.setBytes(0, metadata1);
+ valueBuf1.setBytes(0, value1);
+
+ NullableVariantHolder holder1 =
+ createNullableHolder(metadataBuf1, metadata1, valueBuf1, value1);
+
+ byte[] metadata2 = new byte[] {6, 7, 8};
+ byte[] value2 = new byte[] {9, 10};
+ metadataBuf2.setBytes(0, metadata2);
+ valueBuf2.setBytes(0, value2);
+
+ NullableVariantHolder holder2 =
+ createNullableHolder(metadataBuf2, metadata2, valueBuf2, value2);
+
+ vector.setSafe(0, holder1);
+ vector.setSafe(1, holder2);
+ vector.setValueCount(2);
+
+ // Verify first value
+ NullableVariantHolder result1 = new NullableVariantHolder();
+ vector.get(0, result1);
+ assertEquals(1, result1.isSet);
+
+ byte[] actualMetadata1 = new byte[metadata1.length];
+ byte[] actualValue1 = new byte[value1.length];
+ result1.metadataBuffer.getBytes(result1.metadataStart, actualMetadata1);
+ result1.valueBuffer.getBytes(result1.valueStart, actualValue1);
+ assertArrayEquals(metadata1, actualMetadata1);
+ assertArrayEquals(value1, actualValue1);
+
+ // Verify second value
+ NullableVariantHolder result2 = new NullableVariantHolder();
+ vector.get(1, result2);
+ assertEquals(1, result2.isSet);
+
+ byte[] actualMetadata2 = new byte[metadata2.length];
+ byte[] actualValue2 = new byte[value2.length];
+ result2.metadataBuffer.getBytes(result2.metadataStart, actualMetadata2);
+ result2.valueBuffer.getBytes(result2.valueStart, actualValue2);
+ assertArrayEquals(metadata2, actualMetadata2);
+ assertArrayEquals(value2, actualValue2);
+ }
+ }
+
+ @Test
+ void testNonNullableHolder() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2, 3};
+ byte[] value = new byte[] {4, 5, 6};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ VariantHolder holder = createHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ assertFalse(vector.isNull(0));
+
+ NullableVariantHolder result = new NullableVariantHolder();
+ vector.get(0, result);
+ assertEquals(1, result.isSet);
+ }
+ }
+
+ // ========== Writer Tests ==========
+
+ @Test
+ void testWriteWithVariantHolder() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ VariantWriterImpl writer = new VariantWriterImpl(vector);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2};
+ byte[] value = new byte[] {3, 4, 5};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ VariantHolder holder = createHolder(metadataBuf, metadata, valueBuf, value);
+
+ writer.setPosition(0);
+ writer.write(holder);
+
+ assertEquals(1, vector.getValueCount());
+ assertFalse(vector.isNull(0));
+ }
+ }
+
+ @Test
+ void testWriteWithNullableVariantHolder() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ VariantWriterImpl writer = new VariantWriterImpl(vector);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2};
+ byte[] value = new byte[] {3, 4, 5};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ writer.setPosition(0);
+ writer.write(holder);
+
+ assertEquals(1, vector.getValueCount());
+ assertFalse(vector.isNull(0));
+ }
+ }
+
+ @Test
+ void testWriteWithNullableVariantHolderNull() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ VariantWriterImpl writer = new VariantWriterImpl(vector)) {
+
+ NullableVariantHolder holder = createNullHolder();
+
+ writer.setPosition(0);
+ writer.write(holder);
+
+ assertEquals(1, vector.getValueCount());
+ assertTrue(vector.isNull(0));
+ }
+ }
+
+ @Test
+ void testWriteExtensionWithUnsupportedType() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ VariantWriterImpl writer = new VariantWriterImpl(vector)) {
+
+ writer.setPosition(0);
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () -> writer.writeExtension("invalid-type"));
+
+ assertTrue(exception.getMessage().contains("Unsupported type for Variant"));
+ }
+ }
+
+ @Test
+ void testWriteWithUnsupportedHolder() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ VariantWriterImpl writer = new VariantWriterImpl(vector)) {
+
+ ExtensionHolder unsupportedHolder =
+ new ExtensionHolder() {
+ @Override
+ public ArrowType type() {
+ return VariantType.INSTANCE;
+ }
+ };
+
+ writer.setPosition(0);
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () -> writer.write(unsupportedHolder));
+
+ assertTrue(exception.getMessage().contains("Unsupported type for Variant"));
+ }
+ }
+
+ // ========== Reader Tests ==========
+
+ @Test
+ void testReaderReadWithNullableVariantHolder() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2, 3};
+ byte[] value = new byte[] {4, 5, 6};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ VariantReaderImpl reader = (VariantReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ NullableVariantHolder result = new NullableVariantHolder();
+ reader.read(result);
+
+ assertEquals(1, result.isSet);
+ assertEquals(metadata.length, result.metadataEnd - result.metadataStart);
+ assertEquals(value.length, result.valueEnd - result.valueStart);
+ }
+ }
+
+ @Test
+ void testReaderReadWithNullableVariantHolderNull() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ VariantReaderImpl reader = (VariantReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ NullableVariantHolder holder = new NullableVariantHolder();
+ reader.read(holder);
+
+ assertEquals(0, holder.isSet);
+ }
+ }
+
+ @Test
+ void testReaderIsSet() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1};
+ byte[] value = new byte[] {2};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setNull(1);
+ vector.setValueCount(2);
+
+ VariantReaderImpl reader = (VariantReaderImpl) vector.getReader();
+
+ reader.setPosition(0);
+ assertTrue(reader.isSet());
+
+ reader.setPosition(1);
+ assertFalse(reader.isSet());
+ }
+ }
+
+ @Test
+ void testReaderGetMinorType() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ VariantReaderImpl reader = (VariantReaderImpl) vector.getReader();
+ assertEquals(vector.getMinorType(), reader.getMinorType());
+ }
+ }
+
+ @Test
+ void testReaderGetField() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ VariantReaderImpl reader = (VariantReaderImpl) vector.getReader();
+ assertEquals(vector.getField(), reader.getField());
+ assertEquals("test", reader.getField().getName());
+ }
+ }
+
+ @Test
+ void testReaderReadWithNonNullableVariantHolder() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2, 3};
+ byte[] value = new byte[] {4, 5, 6};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ VariantReaderImpl reader = (VariantReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ VariantHolder result = new VariantHolder();
+ reader.read(result);
+
+ // Verify the data was read correctly
+ byte[] actualMetadata = new byte[metadata.length];
+ byte[] actualValue = new byte[value.length];
+ result.metadataBuffer.getBytes(result.metadataStart, actualMetadata);
+ result.valueBuffer.getBytes(result.valueStart, actualValue);
+
+ assertArrayEquals(metadata, actualMetadata);
+ assertArrayEquals(value, actualValue);
+ assertEquals(1, result.isSet);
+ }
+ }
+
+ // ========== Transfer Pair Tests ==========
+
+ @Test
+ void testTransferPair() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2, 3};
+ byte[] value = new byte[] {4, 5, 6, 7};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ fromVector.setSafe(0, holder);
+ fromVector.setValueCount(1);
+
+ org.apache.arrow.vector.util.TransferPair transferPair =
+ fromVector.getTransferPair(allocator);
+ VariantVector toVector = (VariantVector) transferPair.getTo();
+
+ transferPair.transfer();
+
+ assertEquals(0, fromVector.getValueCount());
+ assertEquals(1, toVector.getValueCount());
+
+ NullableVariantHolder result = new NullableVariantHolder();
+ toVector.get(0, result);
+ assertEquals(1, result.isSet);
+
+ byte[] actualMetadata = new byte[metadata.length];
+ byte[] actualValue = new byte[value.length];
+ result.metadataBuffer.getBytes(result.metadataStart, actualMetadata);
+ result.valueBuffer.getBytes(result.valueStart, actualValue);
+
+ assertArrayEquals(metadata, actualMetadata);
+ assertArrayEquals(value, actualValue);
+
+ toVector.close();
+ }
+ }
+
+ @Test
+ void testSplitAndTransfer() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ ArrowBuf metadataBuf1 = allocator.buffer(10);
+ ArrowBuf valueBuf1 = allocator.buffer(10);
+ ArrowBuf metadataBuf2 = allocator.buffer(10);
+ ArrowBuf valueBuf2 = allocator.buffer(10);
+ ArrowBuf metadataBuf3 = allocator.buffer(10);
+ ArrowBuf valueBuf3 = allocator.buffer(10)) {
+
+ byte[] metadata1 = new byte[] {1};
+ byte[] value1 = new byte[] {2, 3};
+ metadataBuf1.setBytes(0, metadata1);
+ valueBuf1.setBytes(0, value1);
+
+ byte[] metadata2 = new byte[] {4, 5};
+ byte[] value2 = new byte[] {6};
+ metadataBuf2.setBytes(0, metadata2);
+ valueBuf2.setBytes(0, value2);
+
+ byte[] metadata3 = new byte[] {7, 8, 9};
+ byte[] value3 = new byte[] {10, 11, 12};
+ metadataBuf3.setBytes(0, metadata3);
+ valueBuf3.setBytes(0, value3);
+
+ NullableVariantHolder holder1 =
+ createNullableHolder(metadataBuf1, metadata1, valueBuf1, value1);
+ NullableVariantHolder holder2 =
+ createNullableHolder(metadataBuf2, metadata2, valueBuf2, value2);
+ NullableVariantHolder holder3 =
+ createNullableHolder(metadataBuf3, metadata3, valueBuf3, value3);
+
+ fromVector.setSafe(0, holder1);
+ fromVector.setSafe(1, holder2);
+ fromVector.setSafe(2, holder3);
+ fromVector.setValueCount(3);
+
+ org.apache.arrow.vector.util.TransferPair transferPair =
+ fromVector.getTransferPair(allocator);
+ VariantVector toVector = (VariantVector) transferPair.getTo();
+
+ // Split and transfer indices 1-2 (middle and last)
+ transferPair.splitAndTransfer(1, 2);
+
+ assertEquals(2, toVector.getValueCount());
+
+ // Verify transferred values
+ NullableVariantHolder result1 = new NullableVariantHolder();
+ toVector.get(0, result1);
+ assertEquals(1, result1.isSet);
+
+ byte[] actualMetadata1 = new byte[metadata2.length];
+ byte[] actualValue1 = new byte[value2.length];
+ result1.metadataBuffer.getBytes(result1.metadataStart, actualMetadata1);
+ result1.valueBuffer.getBytes(result1.valueStart, actualValue1);
+ assertArrayEquals(metadata2, actualMetadata1);
+ assertArrayEquals(value2, actualValue1);
+
+ NullableVariantHolder result2 = new NullableVariantHolder();
+ toVector.get(1, result2);
+ assertEquals(1, result2.isSet);
+
+ byte[] actualMetadata2 = new byte[metadata3.length];
+ byte[] actualValue2 = new byte[value3.length];
+ result2.metadataBuffer.getBytes(result2.metadataStart, actualMetadata2);
+ result2.valueBuffer.getBytes(result2.valueStart, actualValue2);
+ assertArrayEquals(metadata3, actualMetadata2);
+ assertArrayEquals(value3, actualValue2);
+
+ toVector.close();
+ }
+ }
+
+ @Test
+ void testCopyValueSafe() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ VariantVector toVector = new VariantVector("to", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2};
+ byte[] value = new byte[] {3, 4, 5};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ fromVector.setSafe(0, holder);
+ fromVector.setValueCount(1);
+
+ org.apache.arrow.vector.util.TransferPair transferPair =
+ fromVector.makeTransferPair(toVector);
+
+ transferPair.copyValueSafe(0, 0);
+ toVector.setValueCount(1);
+
+ // Verify the value was copied
+ NullableVariantHolder result = new NullableVariantHolder();
+ toVector.get(0, result);
+ assertEquals(1, result.isSet);
+
+ byte[] actualMetadata = new byte[metadata.length];
+ byte[] actualValue = new byte[value.length];
+ result.metadataBuffer.getBytes(result.metadataStart, actualMetadata);
+ result.valueBuffer.getBytes(result.valueStart, actualValue);
+
+ assertArrayEquals(metadata, actualMetadata);
+ assertArrayEquals(value, actualValue);
+
+ // Original vector should still have the value
+ NullableVariantHolder originalResult = new NullableVariantHolder();
+ fromVector.get(0, originalResult);
+ assertEquals(1, originalResult.isSet);
+ }
+ }
+
+ @Test
+ void testGetTransferPairWithField() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1};
+ byte[] value = new byte[] {2};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ fromVector.setSafe(0, holder);
+ fromVector.setValueCount(1);
+
+ org.apache.arrow.vector.util.TransferPair transferPair =
+ fromVector.getTransferPair(fromVector.getField(), allocator);
+ VariantVector toVector = (VariantVector) transferPair.getTo();
+
+ transferPair.transfer();
+
+ assertEquals(1, toVector.getValueCount());
+ assertEquals(fromVector.getField().getName(), toVector.getField().getName());
+
+ toVector.close();
+ }
+ }
+
+ // ========== Copy Operations Tests ==========
+
+ @Test
+ void testCopyFrom() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ VariantVector toVector = new VariantVector("to", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2, 3};
+ byte[] value = new byte[] {4, 5};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ fromVector.setSafe(0, holder);
+ fromVector.setValueCount(1);
+
+ toVector.allocateNew();
+ toVector.copyFrom(0, 0, fromVector);
+ toVector.setValueCount(1);
+
+ NullableVariantHolder result = new NullableVariantHolder();
+ toVector.get(0, result);
+ assertEquals(1, result.isSet);
+
+ byte[] actualMetadata = new byte[metadata.length];
+ byte[] actualValue = new byte[value.length];
+ result.metadataBuffer.getBytes(result.metadataStart, actualMetadata);
+ result.valueBuffer.getBytes(result.valueStart, actualValue);
+
+ assertArrayEquals(metadata, actualMetadata);
+ assertArrayEquals(value, actualValue);
+ }
+ }
+
+ @Test
+ void testCopyFromSafe() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ VariantVector toVector = new VariantVector("to", allocator);
+ ArrowBuf metadataBuf1 = allocator.buffer(10);
+ ArrowBuf valueBuf1 = allocator.buffer(10);
+ ArrowBuf metadataBuf2 = allocator.buffer(10);
+ ArrowBuf valueBuf2 = allocator.buffer(10)) {
+
+ byte[] metadata1 = new byte[] {1};
+ byte[] value1 = new byte[] {2, 3};
+ metadataBuf1.setBytes(0, metadata1);
+ valueBuf1.setBytes(0, value1);
+
+ NullableVariantHolder holder1 =
+ createNullableHolder(metadataBuf1, metadata1, valueBuf1, value1);
+
+ byte[] metadata2 = new byte[] {4, 5};
+ byte[] value2 = new byte[] {6};
+ metadataBuf2.setBytes(0, metadata2);
+ valueBuf2.setBytes(0, value2);
+
+ NullableVariantHolder holder2 =
+ createNullableHolder(metadataBuf2, metadata2, valueBuf2, value2);
+
+ fromVector.setSafe(0, holder1);
+ fromVector.setSafe(1, holder2);
+ fromVector.setValueCount(2);
+
+ // Copy without pre-allocating toVector
+ for (int i = 0; i < 2; i++) {
+ toVector.copyFromSafe(i, i, fromVector);
+ }
+ toVector.setValueCount(2);
+
+ // Verify both values
+ NullableVariantHolder result1 = new NullableVariantHolder();
+ toVector.get(0, result1);
+ assertEquals(1, result1.isSet);
+
+ byte[] actualMetadata1 = new byte[metadata1.length];
+ byte[] actualValue1 = new byte[value1.length];
+ result1.metadataBuffer.getBytes(result1.metadataStart, actualMetadata1);
+ result1.valueBuffer.getBytes(result1.valueStart, actualValue1);
+ assertArrayEquals(metadata1, actualMetadata1);
+ assertArrayEquals(value1, actualValue1);
+
+ NullableVariantHolder result2 = new NullableVariantHolder();
+ toVector.get(1, result2);
+ assertEquals(1, result2.isSet);
+
+ byte[] actualMetadata2 = new byte[metadata2.length];
+ byte[] actualValue2 = new byte[value2.length];
+ result2.metadataBuffer.getBytes(result2.metadataStart, actualMetadata2);
+ result2.valueBuffer.getBytes(result2.valueStart, actualValue2);
+ assertArrayEquals(metadata2, actualMetadata2);
+ assertArrayEquals(value2, actualValue2);
+ }
+ }
+
+ @Test
+ void testCopyFromWithNulls() {
+ try (VariantVector fromVector = new VariantVector("from", allocator);
+ VariantVector toVector = new VariantVector("to", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1};
+ byte[] value = new byte[] {2};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ fromVector.setSafe(0, holder);
+ fromVector.setNull(1);
+ fromVector.setSafe(2, holder);
+ fromVector.setValueCount(3);
+
+ toVector.allocateNew();
+ for (int i = 0; i < 3; i++) {
+ toVector.copyFromSafe(i, i, fromVector);
+ }
+ toVector.setValueCount(3);
+
+ assertFalse(toVector.isNull(0));
+ assertTrue(toVector.isNull(1));
+ assertFalse(toVector.isNull(2));
+ }
+ }
+
+ // ========== GetObject Tests ==========
+
+ @Test
+ void testGetObject() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1, 2};
+ byte[] value = new byte[] {3, 4, 5};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ Object obj = vector.getObject(0);
+ assertNotNull(obj);
+ assertTrue(obj instanceof Variant);
+ assertEquals(new Variant(metadata, value), obj);
+ }
+ }
+
+ @Test
+ void testGetObjectNull() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ Object obj = vector.getObject(0);
+ assertNull(obj);
+ }
+ }
+
+ // ========== Allocate and Capacity Tests ==========
+
+ @Test
+ void testAllocateNew() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ vector.allocateNew();
+ assertTrue(vector.getValueCapacity() > 0);
+ }
+ }
+
+ @Test
+ void testSetInitialCapacity() {
+ try (VariantVector vector = new VariantVector("test", allocator)) {
+ vector.setInitialCapacity(100);
+ vector.allocateNew();
+ assertTrue(vector.getValueCapacity() >= 100);
+ }
+ }
+
+ @Test
+ void testClearAndReuse() {
+ try (VariantVector vector = new VariantVector("test", allocator);
+ ArrowBuf metadataBuf = allocator.buffer(10);
+ ArrowBuf valueBuf = allocator.buffer(10)) {
+
+ byte[] metadata = new byte[] {1};
+ byte[] value = new byte[] {2};
+ metadataBuf.setBytes(0, metadata);
+ valueBuf.setBytes(0, value);
+
+ NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value);
+
+ vector.setSafe(0, holder);
+ vector.setValueCount(1);
+
+ assertFalse(vector.isNull(0));
+
+ vector.clear();
+ vector.allocateNew();
+
+ // After clear, vector should be empty
+ assertEquals(0, vector.getValueCount());
+ }
+ }
+}
diff --git a/bom/pom.xml b/bom/pom.xml
index 80f03d1205..6a4f741fca 100644
--- a/bom/pom.xml
+++ b/bom/pom.xml
@@ -29,7 +29,7 @@ under the License.
org.apache.arrow
arrow-bom
- 18.3.0
+ 19.0.0
pom
Arrow Bill of Materials
@@ -68,7 +68,7 @@ under the License.
scm:git:https://github.com/apache/arrow-java.git
scm:git:https://github.com/apache/arrow-java.git
- v18.3.0
+ v19.0.0
https://github.com/apache/arrow-java/tree/${project.scm.tag}
@@ -165,7 +165,7 @@ under the License.
${project.version}
- org.apache.arrow
+ org.apache.arrow.gandiva
arrow-gandiva
${project.version}
@@ -194,6 +194,11 @@ under the License.
arrow-tools
${project.version}
+
+ org.apache.arrow
+ arrow-variant
+ ${project.version}
+
@@ -208,7 +213,7 @@ under the License.
org.codehaus.mojo
versions-maven-plugin
- 2.18.0
+ 2.21.0
diff --git a/c/pom.xml b/c/pom.xml
index 290cb561c1..b0a7ffe41d 100644
--- a/c/pom.xml
+++ b/c/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-c-data
diff --git a/c/src/main/cpp/jni_wrapper.cc b/c/src/main/cpp/jni_wrapper.cc
index 35c2b7787e..3d7a194563 100644
--- a/c/src/main/cpp/jni_wrapper.cc
+++ b/c/src/main/cpp/jni_wrapper.cc
@@ -205,8 +205,9 @@ void TryCopyLastError(JNIEnv* env, InnerPrivateData* private_data) {
return;
}
+ jsize error_bytes_len = env->GetArrayLength(arr);
char* error_str = reinterpret_cast(error_bytes);
- private_data->last_error_ = std::string(error_str, std::strlen(error_str));
+ private_data->last_error_ = std::string(error_str, error_bytes_len);
env->ReleaseByteArrayElements(arr, error_bytes, JNI_ABORT);
}
@@ -326,19 +327,20 @@ void ArrowArrayStreamRelease(ArrowArrayStream* stream) {
jint JNI_OnLoad(JavaVM* vm, void* reserved) {
JNIEnv* env;
- if (vm->GetEnv(reinterpret_cast(&env), JNI_VERSION) != JNI_OK) {
- return JNI_ERR;
+ const int err_code = vm->GetEnv(reinterpret_cast(&env), JNI_VERSION);
+ if (err_code != JNI_OK) {
+ return err_code;
}
JNI_METHOD_START
- kObjectClass = CreateGlobalClassReference(env, "Ljava/lang/Object;");
+ kObjectClass = CreateGlobalClassReference(env, "java/lang/Object");
kRuntimeExceptionClass =
- CreateGlobalClassReference(env, "Ljava/lang/RuntimeException;");
+ CreateGlobalClassReference(env, "java/lang/RuntimeException");
kPrivateDataClass =
- CreateGlobalClassReference(env, "Lorg/apache/arrow/c/jni/PrivateData;");
+ CreateGlobalClassReference(env, "org/apache/arrow/c/jni/PrivateData");
kCDataExceptionClass =
- CreateGlobalClassReference(env, "Lorg/apache/arrow/c/jni/CDataJniException;");
+ CreateGlobalClassReference(env, "org/apache/arrow/c/jni/CDataJniException");
kStreamPrivateDataClass = CreateGlobalClassReference(
- env, "Lorg/apache/arrow/c/ArrayStreamExporter$ExportedArrayStreamPrivateData;");
+ env, "org/apache/arrow/c/ArrayStreamExporter$ExportedArrayStreamPrivateData");
kPrivateDataLastErrorField =
GetFieldID(env, kStreamPrivateDataClass, "lastError", "[B");
diff --git a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java
index b74fb1b473..f31a8a1faa 100644
--- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java
+++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java
@@ -58,7 +58,6 @@ void importArray(ArrowArray src) {
ArrowArray ownedArray = ArrowArray.allocateNew(allocator);
ownedArray.save(snapshot);
src.markReleased();
- src.close();
recursionLevel = 0;
diff --git a/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java b/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java
index 07a88cd8d7..34a9c4ec03 100644
--- a/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java
+++ b/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java
@@ -44,7 +44,6 @@ final class ArrowArrayStreamReader extends ArrowReader {
this.ownedStream = ArrowArrayStream.allocateNew(allocator);
this.ownedStream.save(snapshot);
stream.markReleased();
- stream.close();
}
@Override
diff --git a/c/src/main/java/org/apache/arrow/c/Data.java b/c/src/main/java/org/apache/arrow/c/Data.java
index 0b4da33b4e..f9d2ee4542 100644
--- a/c/src/main/java/org/apache/arrow/c/Data.java
+++ b/c/src/main/java/org/apache/arrow/c/Data.java
@@ -231,6 +231,22 @@ public static void exportArrayStream(
new ArrayStreamExporter(allocator).export(out, reader);
}
+ /**
+ * Equivalent to calling {@link #importField(BufferAllocator, ArrowSchema,
+ * CDataDictionaryProvider, boolean) importField(allocator, schema, provider, true)}.
+ *
+ * @param allocator Buffer allocator for allocating dictionary vectors
+ * @param schema C data interface struct representing the field [inout]
+ * @param provider A dictionary provider will be initialized with empty dictionary vectors
+ * (optional)
+ * @return Imported field object
+ * @see #importField(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean)
+ */
+ public static Field importField(
+ BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {
+ return importField(allocator, schema, provider, true);
+ }
+
/**
* Import Java Field from the C data interface.
*
@@ -241,19 +257,42 @@ public static void exportArrayStream(
* @param schema C data interface struct representing the field [inout]
* @param provider A dictionary provider will be initialized with empty dictionary vectors
* (optional)
+ * @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method
+ * completes.
* @return Imported field object
*/
public static Field importField(
- BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {
+ BufferAllocator allocator,
+ ArrowSchema schema,
+ CDataDictionaryProvider provider,
+ boolean closeImportedStructs) {
try {
SchemaImporter importer = new SchemaImporter(allocator);
return importer.importField(schema, provider);
} finally {
schema.release();
- schema.close();
+ if (closeImportedStructs) {
+ schema.close();
+ }
}
}
+ /**
+ * Equivalent to calling {@link #importSchema(BufferAllocator, ArrowSchema,
+ * CDataDictionaryProvider, boolean) importSchema(allocator, schema, provider, true)}.
+ *
+ * @param allocator Buffer allocator for allocating dictionary vectors
+ * @param schema C data interface struct representing the field
+ * @param provider A dictionary provider will be initialized with empty dictionary vectors
+ * (optional)
+ * @return Imported schema object
+ * @see #importSchema(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean)
+ */
+ public static Schema importSchema(
+ BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {
+ return importSchema(allocator, schema, provider, true);
+ }
+
/**
* Import Java Schema from the C data interface.
*
@@ -264,11 +303,16 @@ public static Field importField(
* @param schema C data interface struct representing the field
* @param provider A dictionary provider will be initialized with empty dictionary vectors
* (optional)
+ * @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method
+ * completes.
* @return Imported schema object
*/
public static Schema importSchema(
- BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {
- Field structField = importField(allocator, schema, provider);
+ BufferAllocator allocator,
+ ArrowSchema schema,
+ CDataDictionaryProvider provider,
+ boolean closeImportedStructs) {
+ Field structField = importField(allocator, schema, provider, closeImportedStructs);
if (structField.getType().getTypeID() != ArrowTypeID.Struct) {
throw new IllegalArgumentException(
"Cannot import schema: ArrowSchema describes non-struct type");
@@ -276,24 +320,67 @@ public static Schema importSchema(
return new Schema(structField.getChildren(), structField.getMetadata());
}
+ /**
+ * Equivalent to calling {@link #importIntoVector(BufferAllocator, ArrowArray, FieldVector,
+ * DictionaryProvider, boolean)} importIntoVector(allocator, array, vector, provider, true)}.
+ *
+ * @param allocator Buffer allocator
+ * @param array C data interface struct holding the array data
+ * @param vector Imported vector object [out]
+ * @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @see #importIntoVector(BufferAllocator, ArrowArray, FieldVector, DictionaryProvider, boolean)
+ */
+ public static void importIntoVector(
+ BufferAllocator allocator,
+ ArrowArray array,
+ FieldVector vector,
+ DictionaryProvider provider) {
+ importIntoVector(allocator, array, vector, provider, true);
+ }
+
/**
* Import Java vector from the C data interface.
*
- * The ArrowArray struct has its contents moved (as per the C data interface specification) to
- * a private object held alive by the resulting array.
+ *
On successful completion, the ArrowArray struct will have been moved (as per the C data
+ * interface specification) to a private object held alive by the resulting array.
*
* @param allocator Buffer allocator
* @param array C data interface struct holding the array data
* @param vector Imported vector object [out]
* @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method
+ * completes successfully.
*/
public static void importIntoVector(
BufferAllocator allocator,
ArrowArray array,
FieldVector vector,
- DictionaryProvider provider) {
+ DictionaryProvider provider,
+ boolean closeImportedStructs) {
ArrayImporter importer = new ArrayImporter(allocator, vector, provider);
importer.importArray(array);
+ if (closeImportedStructs) {
+ array.close();
+ }
+ }
+
+ /**
+ * Equivalent to calling {@link #importVector(BufferAllocator, ArrowArray, ArrowSchema,
+ * CDataDictionaryProvider, boolean) importVector(allocator, array, schema, provider, true)}.
+ *
+ * @param allocator Buffer allocator for allocating the output FieldVector
+ * @param array C data interface struct holding the array data
+ * @param schema C data interface struct holding the array type
+ * @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @return Imported vector object
+ * @see #importVector(BufferAllocator, ArrowArray, ArrowSchema, CDataDictionaryProvider, boolean)
+ */
+ public static FieldVector importVector(
+ BufferAllocator allocator,
+ ArrowArray array,
+ ArrowSchema schema,
+ CDataDictionaryProvider provider) {
+ return importVector(allocator, array, schema, provider, true);
}
/**
@@ -307,19 +394,42 @@ public static void importIntoVector(
* @param array C data interface struct holding the array data
* @param schema C data interface struct holding the array type
* @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method
+ * completes successfully and the ArrowSchema struct will be always be closed.
* @return Imported vector object
*/
public static FieldVector importVector(
BufferAllocator allocator,
ArrowArray array,
ArrowSchema schema,
- CDataDictionaryProvider provider) {
- Field field = importField(allocator, schema, provider);
+ CDataDictionaryProvider provider,
+ boolean closeImportedStructs) {
+ Field field = importField(allocator, schema, provider, closeImportedStructs);
FieldVector vector = field.createVector(allocator);
- importIntoVector(allocator, array, vector, provider);
+ importIntoVector(allocator, array, vector, provider, closeImportedStructs);
return vector;
}
+ /**
+ * Equivalent to calling {@link #importIntoVectorSchemaRoot(BufferAllocator, ArrowArray,
+ * VectorSchemaRoot, DictionaryProvider, boolean) importIntoVectorSchemaRoot(allocator, array,
+ * root, provider, true)}.
+ *
+ * @param allocator Buffer allocator
+ * @param array C data interface struct holding the record batch data
+ * @param root vector schema root to load into
+ * @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @see #importIntoVectorSchemaRoot(BufferAllocator, ArrowArray, VectorSchemaRoot,
+ * DictionaryProvider, boolean)
+ */
+ public static void importIntoVectorSchemaRoot(
+ BufferAllocator allocator,
+ ArrowArray array,
+ VectorSchemaRoot root,
+ DictionaryProvider provider) {
+ importIntoVectorSchemaRoot(allocator, array, root, provider, true);
+ }
+
/**
* Import record batch from the C data interface into vector schema root.
*
@@ -333,15 +443,18 @@ public static FieldVector importVector(
* @param array C data interface struct holding the record batch data
* @param root vector schema root to load into
* @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method
+ * completes successfully
*/
public static void importIntoVectorSchemaRoot(
BufferAllocator allocator,
ArrowArray array,
VectorSchemaRoot root,
- DictionaryProvider provider) {
+ DictionaryProvider provider,
+ boolean closeImportedStructs) {
try (StructVector structVector = StructVector.emptyWithDuplicates("", allocator)) {
structVector.initializeChildrenFromFields(root.getSchema().getFields());
- importIntoVector(allocator, array, structVector, provider);
+ importIntoVector(allocator, array, structVector, provider, closeImportedStructs);
StructVectorUnloader unloader = new StructVectorUnloader(structVector);
VectorLoader loader = new VectorLoader(root);
try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) {
@@ -350,6 +463,21 @@ public static void importIntoVectorSchemaRoot(
}
}
+ /**
+ * Equivalent to calling {@link #importVectorSchemaRoot(BufferAllocator, ArrowSchema,
+ * CDataDictionaryProvider, boolean) importVectorSchemaRoot(allocator, schema, provider, true)}.
+ *
+ * @param allocator Buffer allocator for allocating the output VectorSchemaRoot
+ * @param schema C data interface struct holding the record batch schema
+ * @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @return Imported vector schema root
+ * @see #importVectorSchemaRoot(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean)
+ */
+ public static VectorSchemaRoot importVectorSchemaRoot(
+ BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {
+ return importVectorSchemaRoot(allocator, schema, provider, true);
+ }
+
/**
* Import Java vector schema root from a C data interface Schema.
*
@@ -360,11 +488,37 @@ public static void importIntoVectorSchemaRoot(
* @param allocator Buffer allocator for allocating the output VectorSchemaRoot
* @param schema C data interface struct holding the record batch schema
* @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method
+ * completes
* @return Imported vector schema root
*/
public static VectorSchemaRoot importVectorSchemaRoot(
- BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {
- return importVectorSchemaRoot(allocator, null, schema, provider);
+ BufferAllocator allocator,
+ ArrowSchema schema,
+ CDataDictionaryProvider provider,
+ boolean closeImportedStructs) {
+ return importVectorSchemaRoot(allocator, null, schema, provider, closeImportedStructs);
+ }
+
+ /**
+ * Equivalent to calling {@link #importVectorSchemaRoot(BufferAllocator, ArrowArray, ArrowSchema,
+ * CDataDictionaryProvider, boolean) importVectorSchemaRoot(allocator, array, schema, provider,
+ * true)}.
+ *
+ * @param allocator Buffer allocator for allocating the output VectorSchemaRoot
+ * @param array C data interface struct holding the record batch data (optional)
+ * @param schema C data interface struct holding the record batch schema
+ * @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @return Imported vector schema root
+ * @see #importVectorSchemaRoot(BufferAllocator, ArrowArray, ArrowSchema, CDataDictionaryProvider,
+ * boolean)
+ */
+ public static VectorSchemaRoot importVectorSchemaRoot(
+ BufferAllocator allocator,
+ ArrowArray array,
+ ArrowSchema schema,
+ CDataDictionaryProvider provider) {
+ return importVectorSchemaRoot(allocator, array, schema, provider, true);
}
/**
@@ -383,29 +537,56 @@ public static VectorSchemaRoot importVectorSchemaRoot(
* @param array C data interface struct holding the record batch data (optional)
* @param schema C data interface struct holding the record batch schema
* @param provider Dictionary provider to load dictionary vectors to (optional)
+ * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method
+ * completes successfully and the ArrowSchema struct will be always be closed.
* @return Imported vector schema root
*/
public static VectorSchemaRoot importVectorSchemaRoot(
BufferAllocator allocator,
ArrowArray array,
ArrowSchema schema,
- CDataDictionaryProvider provider) {
+ CDataDictionaryProvider provider,
+ boolean closeImportedStructs) {
VectorSchemaRoot vsr =
- VectorSchemaRoot.create(importSchema(allocator, schema, provider), allocator);
+ VectorSchemaRoot.create(
+ importSchema(allocator, schema, provider, closeImportedStructs), allocator);
if (array != null) {
- importIntoVectorSchemaRoot(allocator, array, vsr, provider);
+ importIntoVectorSchemaRoot(allocator, array, vsr, provider, closeImportedStructs);
}
return vsr;
}
/**
- * Import an ArrowArrayStream as an {@link ArrowReader}.
+ * Equivalent to calling {@link #importArrayStream(BufferAllocator, ArrowArrayStream, boolean)
+ * importArrayStream(allocator, stream, true)}.
*
* @param allocator Buffer allocator for allocating the output data.
* @param stream C stream interface struct to import.
* @return Imported reader
+ * @see #importArrayStream(BufferAllocator, ArrowArrayStream, boolean)
*/
public static ArrowReader importArrayStream(BufferAllocator allocator, ArrowArrayStream stream) {
- return new ArrowArrayStreamReader(allocator, stream);
+ return importArrayStream(allocator, stream, true);
+ }
+
+ /**
+ * Import an ArrowArrayStream as an {@link ArrowReader}.
+ *
+ *
On successful completion, the ArrowArrayStream struct will have been moved (as per the C
+ * data interface specification) to a private object held alive by the resulting ArrowReader.
+ *
+ * @param allocator Buffer allocator for allocating the output data.
+ * @param stream C stream interface struct to import.
+ * @param closeImportedStructs if true, the ArrowArrayStream struct will be closed when this
+ * method completes successfully
+ * @return Imported reader
+ */
+ public static ArrowReader importArrayStream(
+ BufferAllocator allocator, ArrowArrayStream stream, boolean closeImportedStructs) {
+ ArrowArrayStreamReader reader = new ArrowArrayStreamReader(allocator, stream);
+ if (closeImportedStructs) {
+ stream.close();
+ }
+ return reader;
}
}
diff --git a/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java b/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java
index f712b400bf..46c93f5541 100644
--- a/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java
+++ b/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java
@@ -75,8 +75,23 @@ private synchronized void loadRemaining() {
}
private void load(String name) {
- final String libraryToLoad =
- name + "/" + getNormalizedArch() + "/" + System.mapLibraryName(name);
+ String libraryName = System.mapLibraryName(name);
+
+ // If 'arrow.cdata.library.path' is defined, try to load the native library from there
+ String libraryPath = System.getProperty("arrow.cdata.library.path");
+ if (libraryPath != null) {
+ try {
+ File libraryFile = new File(libraryPath, libraryName);
+ if (libraryFile.isFile()) {
+ System.load(libraryFile.getAbsolutePath());
+ return;
+ }
+ } catch (UnsatisfiedLinkError e) {
+ // Ignore this error and fall back to extracting from the JAR file
+ }
+ }
+
+ final String libraryToLoad = name + "/" + getNormalizedArch() + "/" + libraryName;
try {
File temp =
File.createTempFile("jnilib-", ".tmp", new File(System.getProperty("java.io.tmpdir")));
diff --git a/c/src/test/java/org/apache/arrow/c/ExceptionTest.java b/c/src/test/java/org/apache/arrow/c/ExceptionTest.java
new file mode 100644
index 0000000000..5bc96a8f99
--- /dev/null
+++ b/c/src/test/java/org/apache/arrow/c/ExceptionTest.java
@@ -0,0 +1,150 @@
+/*
+ * 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.c;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.catchThrowableOfType;
+
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.io.StringWriter;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.VectorLoader;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+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.junit.jupiter.api.Test;
+
+// Regression test for https://github.com/apache/arrow-java/issues/759
+final class ExceptionTest {
+ @Test
+ public void testException() throws IOException {
+ final Schema schema =
+ new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true))));
+ final List batches = new ArrayList<>();
+
+ try (BufferAllocator allocator = new RootAllocator();
+ VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+
+ final String exceptionMessage = "This is a message for testing exception.";
+
+ RuntimeException exToThrow = new RuntimeException(exceptionMessage);
+ batches.add(exToThrow);
+
+ StringWriter sw = new StringWriter();
+ PrintWriter pw = new PrintWriter(sw);
+ exToThrow.printStackTrace(pw);
+ final String expectExceptionMessage = sw.toString();
+
+ ArrowReader source = new ExceptionMemoryArrowReader(allocator, schema, batches);
+
+ try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator);
+ final VectorSchemaRoot importRoot = VectorSchemaRoot.create(schema, allocator)) {
+ final VectorLoader loader = new VectorLoader(importRoot);
+ Data.exportArrayStream(allocator, source, stream);
+
+ try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) {
+ IOException jniException = catchThrowableOfType(IOException.class, reader::loadNextBatch);
+ final String jniMessage = jniException.getMessage();
+ assertThat(jniMessage.endsWith(expectExceptionMessage + "}"));
+ }
+ }
+ }
+ }
+
+ static class ExceptionMemoryArrowReader extends ArrowReader {
+ private final Schema schema;
+ private final List batches; // set ArrowRecordBatch or Exception
+ private final DictionaryProvider provider;
+ private int nextBatch;
+
+ ExceptionMemoryArrowReader(BufferAllocator allocator, Schema schema, List batches) {
+ super(allocator);
+ this.schema = schema;
+ this.batches = batches;
+ this.provider = new CDataDictionaryProvider();
+ this.nextBatch = 0;
+ }
+
+ @Override
+ public Dictionary lookup(long id) {
+ return provider.lookup(id);
+ }
+
+ @Override
+ public Set getDictionaryIds() {
+ return provider.getDictionaryIds();
+ }
+
+ @Override
+ public Map getDictionaryVectors() {
+ return getDictionaryIds().stream()
+ .collect(Collectors.toMap(Function.identity(), this::lookup));
+ }
+
+ @Override
+ public boolean loadNextBatch() throws IOException {
+ if (nextBatch < batches.size()) {
+ Object object = batches.get(nextBatch++);
+ if (object instanceof RuntimeException) {
+ throw (RuntimeException) object;
+ }
+ VectorLoader loader = new VectorLoader(getVectorSchemaRoot());
+ loader.load((ArrowRecordBatch) object);
+ return true;
+ }
+ return false;
+ }
+
+ @Override
+ public long bytesRead() {
+ return 0;
+ }
+
+ @Override
+ protected void closeReadSource() throws IOException {
+ try {
+ for (Object object : batches) {
+ if (object instanceof ArrowRecordBatch) {
+ ArrowRecordBatch batch = (ArrowRecordBatch) object;
+ batch.close();
+ }
+ }
+ } catch (Exception e) {
+ throw new IOException(e);
+ }
+ }
+
+ @Override
+ protected Schema readSchema() {
+ return schema;
+ }
+ }
+}
diff --git a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java
index 6d68449c0b..f6ff88571e 100644
--- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java
+++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java
@@ -17,9 +17,7 @@
package org.apache.arrow.c;
import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
@@ -37,7 +35,6 @@
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
-import org.apache.arrow.memory.util.hash.ArrowBufHasher;
import org.apache.arrow.vector.BaseLargeVariableWidthVector;
import org.apache.arrow.vector.BaseVariableWidthVector;
import org.apache.arrow.vector.BigIntVector;
@@ -46,7 +43,6 @@
import org.apache.arrow.vector.DateMilliVector;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.DurationVector;
-import org.apache.arrow.vector.ExtensionTypeVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.FixedSizeBinaryVector;
import org.apache.arrow.vector.Float2Vector;
@@ -76,6 +72,7 @@
import org.apache.arrow.vector.UInt2Vector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.UInt8Vector;
+import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
@@ -94,6 +91,7 @@
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
import org.apache.arrow.vector.complex.impl.UnionMapWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.holders.IntervalDayHolder;
import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder;
import org.apache.arrow.vector.holders.NullableUInt4Holder;
@@ -102,7 +100,6 @@
import org.apache.arrow.vector.types.Types.MinorType;
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;
@@ -812,9 +809,8 @@ public void testEmptyRunEndEncodedVector() {
@Test
public void testExtensionTypeVector() {
- ExtensionTypeRegistry.register(new UuidType());
final Schema schema =
- new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+ new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
// Fill with data
UUID u1 = UUID.randomUUID();
@@ -832,13 +828,12 @@ public void testExtensionTypeVector() {
assertEquals(root.getSchema(), importedRoot.getSchema());
final Field field = importedRoot.getSchema().getFields().get(0);
- final UuidType expectedType = new UuidType();
assertEquals(
field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
- expectedType.extensionName());
+ UuidType.INSTANCE.extensionName());
assertEquals(
field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
- expectedType.serialize());
+ UuidType.INSTANCE.serialize());
final UuidVector deserialized = (UuidVector) importedRoot.getFieldVectors().get(0);
assertEquals(vector.getValueCount(), deserialized.getValueCount());
@@ -958,6 +953,50 @@ public void testVectorSchemaRootWithDuplicatedFieldNames() {
@Test
public void testSchema() {
+ Schema schema = createSchema();
+ // Consumer allocates empty ArrowSchema
+ try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) {
+ // Producer fills the schema with data
+ exportSchema(schema, consumerArrowSchema);
+
+ // Consumer imports schema
+ Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null);
+ assertEquals(schema.toJson(), importedSchema.toJson());
+ }
+ }
+
+ @Test
+ public void testSchemaStructReuse() {
+ Schema schema = createSchema();
+ // Consumer allocates empty ArrowSchema
+ try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) {
+ // Producer fills the schema with data
+ exportSchema(schema, consumerArrowSchema);
+
+ // Consumer imports schema
+ Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null, false);
+ assertEquals(schema.toJson(), importedSchema.toJson());
+
+ // Imported struct should be released but not closed
+ assertEquals(0, consumerArrowSchema.snapshot().release);
+ assertNotEquals(0, consumerArrowSchema.memoryAddress());
+
+ // Export and import again
+ exportSchema(schema, consumerArrowSchema);
+ importedSchema = Data.importSchema(allocator, consumerArrowSchema, null, false);
+ assertEquals(schema.toJson(), importedSchema.toJson());
+ assertEquals(0, consumerArrowSchema.snapshot().release);
+ assertNotEquals(0, consumerArrowSchema.memoryAddress());
+ }
+ }
+
+ private void exportSchema(Schema schema, ArrowSchema targetArrowSchema) {
+ try (ArrowSchema arrowSchema = ArrowSchema.wrap(targetArrowSchema.memoryAddress())) {
+ Data.exportSchema(allocator, schema, null, arrowSchema);
+ }
+ }
+
+ private static Schema createSchema() {
Field decimalField =
new Field("inner1", FieldType.nullable(new ArrowType.Decimal(19, 4, 128)), null);
Field strField = new Field("inner2", FieldType.nullable(new ArrowType.Utf8()), null);
@@ -968,16 +1007,7 @@ public void testSchema() {
Arrays.asList(decimalField, strField));
Field intField = new Field("col2", FieldType.nullable(new ArrowType.Int(32, true)), null);
Schema schema = new Schema(Arrays.asList(itemField, intField));
- // Consumer allocates empty ArrowSchema
- try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) {
- // Producer fills the schema with data
- try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress())) {
- Data.exportSchema(allocator, schema, null, arrowSchema);
- }
- // Consumer imports schema
- Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null);
- assertEquals(schema.toJson(), importedSchema.toJson());
- }
+ return schema;
}
@Test
@@ -1002,12 +1032,8 @@ public void testImportReleasedArray() {
try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator);
ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) {
// Producer creates structures from existing memory pointers
- try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress());
- ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) {
- // Producer exports vector into the C Data Interface structures
- try (final NullVector vector = new NullVector()) {
- Data.exportVector(allocator, vector, null, arrowArray, arrowSchema);
- }
+ try (final NullVector vector = new NullVector()) {
+ exportFieldVector(vector, consumerArrowSchema, consumerArrowArray);
}
// Release array structure
@@ -1025,6 +1051,45 @@ public void testImportReleasedArray() {
}
}
+ @Test
+ public void testArrayStructReuse() {
+ // Consumer allocates empty structures
+ try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator);
+ ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) {
+ // Producer creates structures from existing memory pointers
+ try (final NullVector vector = new NullVector()) {
+ exportFieldVector(vector, consumerArrowSchema, consumerArrowArray);
+ }
+ Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null, false);
+
+ // Imported structs should be released but not closed
+ assertEquals(0, consumerArrowSchema.snapshot().release);
+ assertNotEquals(0, consumerArrowSchema.memoryAddress());
+ assertEquals(0, consumerArrowArray.snapshot().release);
+ assertNotEquals(0, consumerArrowArray.memoryAddress());
+
+ try (final NullVector vector = new NullVector()) {
+ exportFieldVector(vector, consumerArrowSchema, consumerArrowArray);
+ }
+ Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null, false);
+
+ // Imported structs should be released but not closed
+ assertEquals(0, consumerArrowSchema.snapshot().release);
+ assertNotEquals(0, consumerArrowSchema.memoryAddress());
+ assertEquals(0, consumerArrowArray.snapshot().release);
+ assertNotEquals(0, consumerArrowArray.memoryAddress());
+ }
+ }
+
+ private void exportFieldVector(
+ FieldVector vector, ArrowSchema consumerArrowSchema, ArrowArray consumerArrowArray) {
+ try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress());
+ ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) {
+ // Producer exports vector into the C Data Interface structures
+ Data.exportVector(allocator, vector, null, arrowArray, arrowSchema);
+ }
+ }
+
private VectorSchemaRoot createTestVSR() {
BitVector bitVector = new BitVector("boolean", allocator);
@@ -1047,72 +1112,4 @@ private VectorSchemaRoot createTestVSR() {
return new VectorSchemaRoot(fields, vectors);
}
-
- static class UuidType extends ExtensionType {
-
- @Override
- public ArrowType storageType() {
- return new ArrowType.FixedSizeBinary(16);
- }
-
- @Override
- public String extensionName() {
- return "uuid";
- }
-
- @Override
- public boolean extensionEquals(ExtensionType other) {
- return other instanceof UuidType;
- }
-
- @Override
- public ArrowType deserialize(ArrowType storageType, String serializedData) {
- if (!storageType.equals(storageType())) {
- throw new UnsupportedOperationException(
- "Cannot construct UuidType from underlying type " + storageType);
- }
- return new UuidType();
- }
-
- @Override
- public String serialize() {
- return "";
- }
-
- @Override
- public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
- return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));
- }
- }
-
- static class UuidVector extends ExtensionTypeVector {
-
- public UuidVector(
- String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) {
- super(name, allocator, underlyingVector);
- }
-
- @Override
- public UUID getObject(int index) {
- 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) {
- return getUnderlyingVector().hashCode(index, hasher);
- }
-
- public void set(int index, UUID uuid) {
- ByteBuffer bb = ByteBuffer.allocate(16);
- bb.putLong(uuid.getMostSignificantBits());
- bb.putLong(uuid.getLeastSignificantBits());
- getUnderlyingVector().set(index, bb.array());
- }
- }
}
diff --git a/ci/docker/vcpkg-jni.dockerfile b/ci/docker/vcpkg-jni.dockerfile
index 55fa35e0d1..d6bd322a39 100644
--- a/ci/docker/vcpkg-jni.dockerfile
+++ b/ci/docker/vcpkg-jni.dockerfile
@@ -18,24 +18,10 @@
ARG base
FROM ${base}
-# Install the libraries required by Gandiva to run
-# Use enable llvm[enable-rtti] in the vcpkg.json to avoid link problems in Gandiva
-RUN vcpkg install \
- --clean-after-build \
- --x-install-root=${VCPKG_ROOT}/installed \
- --x-manifest-root=/arrow/ci/vcpkg \
- --x-feature=dev \
- --x-feature=flight \
- --x-feature=gcs \
- --x-feature=json \
- --x-feature=parquet \
- --x-feature=gandiva \
- --x-feature=s3
-
# Install Java
# We need Java for JNI headers, but we don't invoke Maven in this build.
ARG java=11
-RUN yum install -y java-$java-openjdk-devel && yum clean all
+RUN dnf install -y java-$java-openjdk-devel && dnf clean all
# For ci/scripts/{cpp,java}_*.sh
ENV ARROW_HOME=/tmp/local \
diff --git a/ci/scripts/jni_build.sh b/ci/scripts/jni_build.sh
index aec6fc325c..c000837987 100755
--- a/ci/scripts/jni_build.sh
+++ b/ci/scripts/jni_build.sh
@@ -66,7 +66,7 @@ cmake \
-DProtobuf_USE_STATIC_LIBS=ON \
-GNinja \
"${EXTRA_CMAKE_OPTIONS[@]}"
-cmake --build "${build_dir}"
+cmake --build "${build_dir}" --verbose
if [ "${ARROW_JAVA_BUILD_TESTS}" = "ON" ]; then
ctest \
--output-on-failure \
diff --git a/ci/scripts/jni_full_build.sh b/ci/scripts/jni_full_build.sh
index e9ad0ddbda..5d0aee0555 100755
--- a/ci/scripts/jni_full_build.sh
+++ b/ci/scripts/jni_full_build.sh
@@ -97,8 +97,10 @@ find ~/.m2/repository/org/apache/arrow \
-exec echo "{}" ";" \
-exec cp "{}" "${dist_dir}" ";"
-for artifact in "${dist_dir}"/*; do
+pushd "${dist_dir}"
+for artifact in *; do
sha256sum "${artifact}" >"${artifact}.sha256"
sha512sum "${artifact}" >"${artifact}.sha512"
done
+popd
github_actions_group_end
diff --git a/ci/scripts/jni_macos_build.sh b/ci/scripts/jni_macos_build.sh
index f7543b6f7a..65ab450666 100755
--- a/ci/scripts/jni_macos_build.sh
+++ b/ci/scripts/jni_macos_build.sh
@@ -59,73 +59,27 @@ fi
github_actions_group_begin "Building Arrow C++ libraries"
install_dir="${build_dir}/cpp-install"
-: "${ARROW_ACERO:=ON}"
-export ARROW_ACERO
-: "${ARROW_BUILD_TESTS:=OFF}"
-export ARROW_BUILD_TESTS
-: "${ARROW_DATASET:=ON}"
-export ARROW_DATASET
-: "${ARROW_GANDIVA:=ON}"
-export ARROW_GANDIVA
-: "${ARROW_ORC:=ON}"
-export ARROW_ORC
-: "${ARROW_PARQUET:=ON}"
-: "${ARROW_S3:=ON}"
-: "${CMAKE_BUILD_TYPE:=Release}"
-: "${CMAKE_UNITY_BUILD:=ON}"
-export ARROW_TEST_DATA="${arrow_dir}/testing/data"
-export PARQUET_TEST_DATA="${arrow_dir}/cpp/submodules/parquet-testing/data"
+export ARROW_BUILD_TESTS=OFF
+
+export ARROW_DATASET=ON
+export ARROW_GANDIVA=ON
+export ARROW_ORC=ON
+export ARROW_PARQUET=ON
+
export AWS_EC2_METADATA_DISABLED=TRUE
cmake \
-S "${arrow_dir}/cpp" \
-B "${build_dir}/cpp" \
- -DARROW_ACERO="${ARROW_ACERO}" \
- -DARROW_BUILD_SHARED=OFF \
- -DARROW_BUILD_TESTS="${ARROW_BUILD_TESTS}" \
- -DARROW_CSV="${ARROW_DATASET}" \
- -DARROW_DATASET="${ARROW_DATASET}" \
- -DARROW_SUBSTRAIT="${ARROW_DATASET}" \
- -DARROW_DEPENDENCY_USE_SHARED=OFF \
- -DARROW_GANDIVA="${ARROW_GANDIVA}" \
- -DARROW_GANDIVA_STATIC_LIBSTDCPP=ON \
- -DARROW_JSON="${ARROW_DATASET}" \
- -DARROW_ORC="${ARROW_ORC}" \
- -DARROW_PARQUET="${ARROW_PARQUET}" \
- -DARROW_S3="${ARROW_S3}" \
- -DARROW_USE_CCACHE="${ARROW_USE_CCACHE}" \
- -DAWSSDK_SOURCE=BUNDLED \
- -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \
- -DCMAKE_INSTALL_PREFIX="${install_dir}" \
- -DCMAKE_UNITY_BUILD="${CMAKE_UNITY_BUILD}" \
- -DGTest_SOURCE=BUNDLED \
- -DPARQUET_BUILD_EXAMPLES=OFF \
- -DPARQUET_BUILD_EXECUTABLES=OFF \
- -DPARQUET_REQUIRE_ENCRYPTION=OFF \
- -Dre2_SOURCE=BUNDLED \
- -GNinja
+ --preset=ninja-release-jni-macos \
+ -DCMAKE_INSTALL_PREFIX="${install_dir}"
cmake --build "${build_dir}/cpp" --target install
github_actions_group_end
-if [ "${ARROW_RUN_TESTS:-}" == "ON" ]; then
- github_actions_group_begin "Running Arrow C++ libraries tests"
- # MinIO is required
- exclude_tests="arrow-s3fs-test"
- # unstable
- exclude_tests="${exclude_tests}|arrow-acero-asof-join-node-test"
- exclude_tests="${exclude_tests}|arrow-acero-hash-join-node-test"
- ctest \
- --exclude-regex "${exclude_tests}" \
- --label-regex unittest \
- --output-on-failure \
- --parallel "$(sysctl -n hw.ncpu)" \
- --test-dir "${build_dir}/cpp" \
- --timeout 300
- github_actions_group_end
-fi
-
-export JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/protobuf_ep-install"
+JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/_deps/protobuf-build"
+JAVA_JNI_CMAKE_ARGS+=" -DProtobuf_SRC_ROOT_FOLDER=${build_dir}/cpp/_deps/protobuf-src"
+export JAVA_JNI_CMAKE_ARGS
"${source_dir}/ci/scripts/jni_build.sh" \
"${source_dir}" \
"${install_dir}" \
@@ -142,6 +96,7 @@ github_actions_group_begin "Checking shared dependencies for libraries"
pushd "${dist_dir}"
archery linking check-dependencies \
--allow CoreFoundation \
+ --allow Network \
--allow Security \
--allow libSystem \
--allow libarrow_cdata_jni \
diff --git a/ci/scripts/jni_manylinux_build.sh b/ci/scripts/jni_manylinux_build.sh
index a34ec0f420..3577c37ab3 100755
--- a/ci/scripts/jni_manylinux_build.sh
+++ b/ci/scripts/jni_manylinux_build.sh
@@ -53,33 +53,19 @@ if [ "${ARROW_USE_CCACHE}" == "ON" ]; then
fi
github_actions_group_begin "Building Arrow C++ libraries"
-devtoolset_version="$(rpm -qa "devtoolset-*-gcc" --queryformat '%{VERSION}' | grep -o "^[0-9]*")"
-devtoolset_include_cpp="/opt/rh/devtoolset-${devtoolset_version}/root/usr/include/c++/${devtoolset_version}"
-: "${ARROW_ACERO:=ON}"
-export ARROW_ACERO
-: "${ARROW_BUILD_TESTS:=OFF}"
-export ARROW_BUILD_TESTS
-: "${ARROW_DATASET:=ON}"
-export ARROW_DATASET
-: "${ARROW_GANDIVA:=ON}"
-export ARROW_GANDIVA
-: "${ARROW_GCS:=ON}"
-: "${ARROW_JEMALLOC:=OFF}"
-: "${ARROW_MIMALLOC:=ON}"
-: "${ARROW_RPATH_ORIGIN:=ON}"
-: "${ARROW_ORC:=ON}"
-export ARROW_ORC
-: "${ARROW_PARQUET:=ON}"
-: "${ARROW_S3:=ON}"
-: "${CMAKE_BUILD_TYPE:=release}"
-: "${CMAKE_UNITY_BUILD:=ON}"
+
: "${VCPKG_ROOT:=/opt/vcpkg}"
: "${VCPKG_FEATURE_FLAGS:=-manifests}"
-: "${VCPKG_TARGET_TRIPLET:=${VCPKG_DEFAULT_TRIPLET:-x64-linux-static-${CMAKE_BUILD_TYPE}}}"
-: "${GANDIVA_CXX_FLAGS:=-isystem;${devtoolset_include_cpp};-isystem;${devtoolset_include_cpp}/x86_64-redhat-linux;-lpthread}"
+: "${VCPKG_TARGET_TRIPLET:=${VCPKG_DEFAULT_TRIPLET:-x64-linux-static-release}}"
+export VCPKG_TARGET_TRIPLET
+
+export ARROW_BUILD_TESTS=OFF
+
+export ARROW_DATASET=ON
+export ARROW_GANDIVA=ON
+export ARROW_ORC=ON
+export ARROW_PARQUET=ON
-export ARROW_TEST_DATA="${arrow_dir}/testing/data"
-export PARQUET_TEST_DATA="${arrow_dir}/cpp/submodules/parquet-testing/data"
export AWS_EC2_METADATA_DISABLED=TRUE
install_dir="${build_dir}/cpp-install"
@@ -87,71 +73,12 @@ install_dir="${build_dir}/cpp-install"
cmake \
-S "${arrow_dir}/cpp" \
-B "${build_dir}/cpp" \
- -DARROW_ACERO="${ARROW_ACERO}" \
- -DARROW_BUILD_SHARED=OFF \
- -DARROW_BUILD_TESTS="${ARROW_BUILD_TESTS}" \
- -DARROW_CSV="${ARROW_DATASET}" \
- -DARROW_DATASET="${ARROW_DATASET}" \
- -DARROW_SUBSTRAIT="${ARROW_DATASET}" \
- -DARROW_DEPENDENCY_SOURCE="VCPKG" \
- -DARROW_DEPENDENCY_USE_SHARED=OFF \
- -DARROW_GANDIVA_PC_CXX_FLAGS="${GANDIVA_CXX_FLAGS}" \
- -DARROW_GANDIVA="${ARROW_GANDIVA}" \
- -DARROW_GCS="${ARROW_GCS}" \
- -DARROW_JEMALLOC="${ARROW_JEMALLOC}" \
- -DARROW_JSON="${ARROW_DATASET}" \
- -DARROW_MIMALLOC="${ARROW_MIMALLOC}" \
- -DARROW_ORC="${ARROW_ORC}" \
- -DARROW_PARQUET="${ARROW_PARQUET}" \
- -DARROW_RPATH_ORIGIN="${ARROW_RPATH_ORIGIN}" \
- -DARROW_S3="${ARROW_S3}" \
- -DARROW_USE_CCACHE="${ARROW_USE_CCACHE}" \
- -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \
- -DCMAKE_INSTALL_PREFIX="${install_dir}" \
- -DCMAKE_UNITY_BUILD="${CMAKE_UNITY_BUILD}" \
- -DGTest_SOURCE=BUNDLED \
- -DORC_SOURCE=BUNDLED \
- -DORC_PROTOBUF_EXECUTABLE="${VCPKG_ROOT}/installed/${VCPKG_TARGET_TRIPLET}/tools/protobuf/protoc" \
- -DPARQUET_BUILD_EXAMPLES=OFF \
- -DPARQUET_BUILD_EXECUTABLES=OFF \
- -DPARQUET_REQUIRE_ENCRYPTION=OFF \
- -DVCPKG_MANIFEST_MODE=OFF \
- -DVCPKG_TARGET_TRIPLET="${VCPKG_TARGET_TRIPLET}" \
- -GNinja
+ --preset=ninja-release-jni-linux \
+ -DCMAKE_INSTALL_PREFIX="${install_dir}"
cmake --build "${build_dir}/cpp"
cmake --install "${build_dir}/cpp"
github_actions_group_end
-if [ "${ARROW_RUN_TESTS:-OFF}" = "ON" ]; then
- github_actions_group_begin "Running Arrow C++ libraries tests"
- # MinIO is required
- exclude_tests="arrow-s3fs-test"
- case $(arch) in
- aarch64)
- # GCS testbench is crashed on aarch64:
- # ImportError: ../grpc/_cython/cygrpc.cpython-38-aarch64-linux-gnu.so:
- # undefined symbol: vtable for std::__cxx11::basic_ostringstream<
- # char, std::char_traits, std::allocator >
- exclude_tests="${exclude_tests}|arrow-gcsfs-test"
- ;;
- esac
- # unstable
- exclude_tests="${exclude_tests}|arrow-acero-asof-join-node-test"
- exclude_tests="${exclude_tests}|arrow-acero-hash-join-node-test"
- # external dependency
- exclude_tests="${exclude_tests}|arrow-gcsfs-test"
- # strptime
- exclude_tests="${exclude_tests}|arrow-utility-test"
- ctest \
- --exclude-regex "${exclude_tests}" \
- --label-regex unittest \
- --output-on-failure \
- --parallel "$(nproc)" \
- --test-dir "${build_dir}/cpp" \
- --timeout 300
- github_actions_group_end
-fi
-
JAVA_JNI_CMAKE_ARGS="-DCMAKE_TOOLCHAIN_FILE=${VCPKG_ROOT}/scripts/buildsystems/vcpkg.cmake"
JAVA_JNI_CMAKE_ARGS="${JAVA_JNI_CMAKE_ARGS} -DVCPKG_TARGET_TRIPLET=${VCPKG_TARGET_TRIPLET}"
export JAVA_JNI_CMAKE_ARGS
diff --git a/ci/scripts/jni_windows_build.sh b/ci/scripts/jni_windows_build.sh
index d01ef45f5a..6503ac63e5 100755
--- a/ci/scripts/jni_windows_build.sh
+++ b/ci/scripts/jni_windows_build.sh
@@ -68,7 +68,7 @@ cmake \
-B "${build_dir}/cpp" \
-DARROW_ACERO="${ARROW_ACERO}" \
-DARROW_BUILD_SHARED=OFF \
- -DARROW_BUILD_TESTS=ON \
+ -DARROW_BUILD_TESTS="${ARROW_BUILD_TESTS}" \
-DARROW_CSV="${ARROW_DATASET}" \
-DARROW_DATASET="${ARROW_DATASET}" \
-DARROW_SUBSTRAIT="${ARROW_DATASET}" \
diff --git a/compose.yaml b/compose.yaml
index b125c3c983..f5082a22aa 100644
--- a/compose.yaml
+++ b/compose.yaml
@@ -99,7 +99,7 @@ services:
cache_from:
- ${REPO}:${ARCH}-vcpkg-jni-${VCPKG}
args:
- base: ${ARROW_REPO}:${ARCH}-python-${PYTHON}-wheel-manylinux-2014-vcpkg-${VCPKG}
+ base: ${ARROW_REPO}:${ARCH}-cpp-jni-${VCPKG}
volumes:
- .:/arrow-java:delegated
- ${ARROW_REPO_ROOT}:/arrow:delegated
diff --git a/compression/pom.xml b/compression/pom.xml
index 3443f11478..92144addc4 100644
--- a/compression/pom.xml
+++ b/compression/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-compression
Arrow Compression
@@ -50,12 +50,12 @@ under the License.
org.apache.commons
commons-compress
- 1.27.1
+ 1.28.0
com.github.luben
zstd-jni
- 1.5.7-2
+ 1.5.7-6
diff --git a/dataset/pom.xml b/dataset/pom.xml
index efbe310ea2..df5620c641 100644
--- a/dataset/pom.xml
+++ b/dataset/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-dataset
@@ -32,8 +32,8 @@ under the License.
../../../cpp/release-build/
- 1.15.2
- 1.12.0
+ 1.17.0
+ 1.12.1
@@ -130,7 +130,7 @@ under the License.
org.apache.orc
orc-core
- 2.1.1
+ 2.3.0
test
@@ -156,7 +156,7 @@ under the License.
commons-io
commons-io
- 2.19.0
+ 2.21.0
test
diff --git a/dataset/src/main/cpp/jni_wrapper.cc b/dataset/src/main/cpp/jni_wrapper.cc
index 49cc85251c..e8087648eb 100644
--- a/dataset/src/main/cpp/jni_wrapper.cc
+++ b/dataset/src/main/cpp/jni_wrapper.cc
@@ -23,6 +23,7 @@
#include "arrow/array/concatenate.h"
#include "arrow/c/bridge.h"
#include "arrow/c/helpers.h"
+#include "arrow/compute/initialize.h"
#include "arrow/dataset/api.h"
#include "arrow/dataset/file_base.h"
#ifdef ARROW_CSV
@@ -807,6 +808,13 @@ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_ensureS3Fina
JNI_METHOD_END()
}
+JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_initialize(
+ JNIEnv* env, jobject) {
+ JNI_METHOD_START
+ JniAssertOkOrThrow(arrow::compute::Initialize());
+ JNI_METHOD_END()
+}
+
/*
* Class: org_apache_arrow_dataset_file_JniWrapper
* Method: makeFileSystemDatasetFactory
diff --git a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java
index 631b8b1bbe..5fb4816488 100644
--- a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java
+++ b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java
@@ -56,6 +56,7 @@ public void ensureLoaded() {
}
loadRemaining();
ensureS3FinalizedOnShutdown();
+ JniWrapper.get().initialize();
}
private synchronized void loadRemaining() {
diff --git a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java
index 6637c113d9..cfef098ec4 100644
--- a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java
+++ b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java
@@ -124,4 +124,7 @@ public native long createScanner(
* uninitialized, then this is a noop.
*/
public native void ensureS3Finalized();
+
+ /** Initialize Arrow Compute. */
+ public native void initialize();
}
diff --git a/dev/release/release.sh b/dev/release/release.sh
index 70e1f96454..f08a618c4f 100755
--- a/dev/release/release.sh
+++ b/dev/release/release.sh
@@ -19,6 +19,8 @@
set -eu
+SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
if [ "$#" -ne 2 ]; then
echo "Usage: $0 "
echo " e.g.: $0 19.0.1 1"
diff --git a/docs/source/flight.rst b/docs/source/flight.rst
index fabced8094..fd0fdf07bc 100644
--- a/docs/source/flight.rst
+++ b/docs/source/flight.rst
@@ -232,8 +232,8 @@ Servers can add other gRPC services. For example, to add the `Health Check servi
See the :external+arrow:ref:`best practices for C++ `.
-.. _`FlightClient`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightClient.html
-.. _`FlightProducer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightProducer.html
-.. _`FlightServer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightServer.html
-.. _`NoOpFlightProducer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/NoOpFlightProducer.html
-.. _`Location`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/Location.html
+.. _`FlightClient`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightClient.html
+.. _`FlightProducer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightProducer.html
+.. _`FlightServer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightServer.html
+.. _`NoOpFlightProducer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/NoOpFlightProducer.html
+.. _`Location`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/Location.html
diff --git a/docs/source/flight_sql.rst b/docs/source/flight_sql.rst
index 169a0e24bf..09ce1dda0d 100644
--- a/docs/source/flight_sql.rst
+++ b/docs/source/flight_sql.rst
@@ -29,4 +29,4 @@ over the network.
For usage information, see the `API documentation`_.
-.. _API documentation: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/sql/package-summary.html
+.. _API documentation: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.sql/org/apache/arrow/flight/sql/package-summary.html
diff --git a/docs/source/flight_sql_jdbc_driver.rst b/docs/source/flight_sql_jdbc_driver.rst
index 1806930943..4deb726b33 100644
--- a/docs/source/flight_sql_jdbc_driver.rst
+++ b/docs/source/flight_sql_jdbc_driver.rst
@@ -173,3 +173,126 @@ DriverManager#getConnection()
`_,
the username and password supplied on the URI supercede the username and
password arguments to the function call.
+
+OAuth 2.0 Authentication
+========================
+
+The driver supports OAuth 2.0 authentication for obtaining access tokens
+from an authorization server. Two OAuth flows are currently supported:
+
+* **Client Credentials** - For service-to-service authentication where no
+ user interaction is required. The application authenticates using its own
+ credentials (client ID and client secret).
+
+* **Token Exchange** (RFC 8693) - For exchanging one token for another,
+ commonly used for federated authentication, delegation, or impersonation
+ scenarios.
+
+OAuth Connection Properties
+---------------------------
+
+The following properties configure OAuth authentication. These properties
+should be provided via the ``Properties`` object when connecting, as they
+may contain special characters that are difficult to encode in a URI.
+
+**Common OAuth Properties**
+
+.. list-table::
+ :header-rows: 1
+
+ * - Parameter
+ - Type
+ - Required
+ - Default
+ - Description
+
+ * - oauth.flow
+ - String
+ - Yes (to enable OAuth)
+ - null
+ - The OAuth grant type. Supported values: ``client_credentials``,
+ ``token_exchange``
+
+ * - oauth.tokenUri
+ - String
+ - Yes
+ - null
+ - The OAuth 2.0 token endpoint URL (e.g.,
+ ``https://auth.example.com/oauth/token``)
+
+ * - oauth.clientId
+ - String
+ - Conditional
+ - null
+ - The OAuth 2.0 client ID. Required for ``client_credentials`` flow,
+ optional for ``token_exchange``
+
+ * - oauth.clientSecret
+ - String
+ - Conditional
+ - null
+ - The OAuth 2.0 client secret. Required for ``client_credentials`` flow,
+ optional for ``token_exchange``
+
+ * - oauth.scope
+ - String
+ - No
+ - null
+ - Space-separated list of OAuth scopes to request
+
+ * - oauth.resource
+ - String
+ - No
+ - null
+ - The resource indicator for the token request (RFC 8707)
+
+**Token Exchange Properties**
+
+These properties are specific to the ``token_exchange`` flow:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Parameter
+ - Type
+ - Required
+ - Default
+ - Description
+
+ * - oauth.exchange.subjectToken
+ - String
+ - Yes
+ - null
+ - The subject token to exchange (e.g., a JWT from an identity provider)
+
+ * - oauth.exchange.subjectTokenType
+ - String
+ - Yes
+ - null
+ - The token type URI of the subject token. Common values:
+ ``urn:ietf:params:oauth:token-type:access_token``,
+ ``urn:ietf:params:oauth:token-type:jwt``
+
+ * - oauth.exchange.actorToken
+ - String
+ - No
+ - null
+ - The actor token for delegation/impersonation scenarios
+
+ * - oauth.exchange.actorTokenType
+ - String
+ - No
+ - null
+ - The token type URI of the actor token
+
+ * - oauth.exchange.aud
+ - String
+ - No
+ - null
+ - The target audience for the exchanged token
+
+ * - oauth.exchange.requestedTokenType
+ - String
+ - No
+ - null
+ - The desired token type for the exchanged token
diff --git a/docs/source/jdbc.rst b/docs/source/jdbc.rst
index c0477cb06d..2f57c34bf8 100644
--- a/docs/source/jdbc.rst
+++ b/docs/source/jdbc.rst
@@ -95,7 +95,7 @@ Type Mapping
The JDBC to Arrow type mapping can be obtained at runtime from
`JdbcToArrowUtils.getArrowTypeFromJdbcType`_.
-.. _JdbcToArrowUtils.getArrowTypeFromJdbcType: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.html#getArrowTypeFromJdbcType-org.apache.arrow.adapter.jdbc.JdbcFieldInfo-java.util.Calendar-
+.. _JdbcToArrowUtils.getArrowTypeFromJdbcType: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.html#getArrowTypeFromJdbcType-org.apache.arrow.adapter.jdbc.JdbcFieldInfo-java.util.Calendar-
+--------------------+--------------------+-------+
| JDBC Type | Arrow Type | Notes |
@@ -171,8 +171,8 @@ The JDBC to Arrow type mapping can be obtained at runtime from
timezone of the calendar, else it will be a timestamp without
timezone.
-.. _setArraySubTypeByColumnIndexMap: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnIndexMap-java.util.Map-
-.. _setArraySubTypeByColumnNameMap: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnNameMap-java.util.Map-
+.. _setArraySubTypeByColumnIndexMap: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnIndexMap-java.util.Map-
+.. _setArraySubTypeByColumnNameMap: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnNameMap-java.util.Map-
.. _ARROW-17006: https://issues.apache.org/jira/browse/ARROW-17006
VectorSchemaRoot to PreparedStatement Parameter Conversion
@@ -213,7 +213,8 @@ Type Mapping
------------
The Arrow to JDBC type mapping can be obtained at runtime via
-a method on ColumnBinder.
+a method on ColumnBinder. The Flight SQL JDBC driver follows the same
+mapping, with additional support for the UUID extension type noted below.
+----------------------------+----------------------------+-------+
| Arrow Type | JDBC Type | Notes |
@@ -232,6 +233,8 @@ a method on ColumnBinder.
+----------------------------+----------------------------+-------+
| FixedSizeBinary | BINARY (setBytes) | |
+----------------------------+----------------------------+-------+
+| Uuid (extension) | OTHER (setObject) | \(3) |
++----------------------------+----------------------------+-------+
| Float32 | REAL (setFloat) | |
+----------------------------+----------------------------+-------+
| Int8 | TINYINT (setByte) | |
@@ -276,3 +279,6 @@ a method on ColumnBinder.
`_,
which will lead to the driver using the "default timezone" (that of
the Java VM).
+* \(3) For the Flight SQL JDBC driver, the Arrow UUID extension type
+ (``arrow.uuid``) maps to JDBC ``OTHER`` and is surfaced as
+ ``java.util.UUID`` values.
diff --git a/docs/source/memory.rst b/docs/source/memory.rst
index 58ef382dc9..4a71ed846a 100644
--- a/docs/source/memory.rst
+++ b/docs/source/memory.rst
@@ -333,18 +333,18 @@ How this works:
}
}
-.. _`ArrowBuf`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html
-.. _`ArrowBuf.print()`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html#print-java.lang.StringBuilder-int-org.apache.arrow.memory.BaseAllocator.Verbosity-
-.. _`BufferAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/BufferAllocator.html
-.. _`BufferLedger`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/BufferLedger.html
-.. _`RootAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/RootAllocator.html
-.. _`newChildAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/RootAllocator.html#newChildAllocator-java.lang.String-org.apache.arrow.memory.AllocationListener-long-long-
+.. _`ArrowBuf`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html
+.. _`ArrowBuf.print()`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html#print-java.lang.StringBuilder-int-org.apache.arrow.memory.BaseAllocator.Verbosity-
+.. _`BufferAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/BufferAllocator.html
+.. _`BufferLedger`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/BufferLedger.html
+.. _`RootAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html
+.. _`newChildAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html#newChildAllocator-java.lang.String-org.apache.arrow.memory.AllocationListener-long-long-
.. _`Netty`: https://netty.io/wiki/
.. _`sun.misc.unsafe`: https://web.archive.org/web/20210929024401/http://www.docjar.com/html/api/sun/misc/Unsafe.java.html
.. _`Direct Memory`: https://docs.oracle.com/en/java/javase/11/docs/api/java.base/java/nio/ByteBuffer.html
-.. _`ReferenceManager`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html
-.. _`ReferenceManager.release`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html#release--
-.. _`ReferenceManager.retain`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html#retain--
+.. _`ReferenceManager`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html
+.. _`ReferenceManager.release`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html#release--
+.. _`ReferenceManager.retain`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html#retain--
Arrow Memory In-Depth
=====================
diff --git a/docs/source/overview.rst b/docs/source/overview.rst
index be579c1495..1188054114 100644
--- a/docs/source/overview.rst
+++ b/docs/source/overview.rst
@@ -45,6 +45,9 @@ but some modules are JNI bindings to the C++ library.
* - arrow-vector
- An off-heap reference implementation for Arrow columnar data format.
- Native
+ * - arrow-vector-codegen
+ - Template files for Arrow datatypes suitable for code generation.
+ - Native
* - arrow-tools
- Java applications for working with Arrow ValueVectors.
- Native
diff --git a/docs/source/substrait.rst b/docs/source/substrait.rst
index b3678ac815..5ec07f1658 100644
--- a/docs/source/substrait.rst
+++ b/docs/source/substrait.rst
@@ -19,7 +19,7 @@
Substrait
=========
-The ``arrow-dataset`` module can execute Substrait_ plans via the :external+arrow:doc:`Acero `
+The ``arrow-dataset`` module can execute Substrait_ plans via the :external+arrow:doc:`Acero `
query engine.
Executing Queries Using Substrait Plans
diff --git a/docs/source/table.rst b/docs/source/table.rst
index 5aa95e153c..880ef84d29 100644
--- a/docs/source/table.rst
+++ b/docs/source/table.rst
@@ -364,15 +364,15 @@ If the table contains dictionary-encoded vectors and was constructed with a ``Di
Data.exportTable(bufferAllocator, table, outArrowArray);
-.. _`ArrowBuf`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html
-.. _`Data`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/c/Data.html
-.. _`DictionaryProvider`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/dictionary/DictionaryProvider.html
-.. _`Field`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Field.html
-.. _`FieldReader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/complex/reader/FieldReader.html
-.. _`FieldVector`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/FieldVector.html
-.. _`Row`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Row.html
-.. _`Schema`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Schema.html
-.. _`Table`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Table.html
-.. _`ValueHolder`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/holders/ValueHolder.html
-.. _`ValueVector`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/ValueVector.html
-.. _`VectorSchemaRoot`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorSchemaRoot.html
+.. _`ArrowBuf`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html
+.. _`Data`: https://arrow.apache.org/java/current/reference/org.apache.arrow.c/org/apache/arrow/c/Data.html
+.. _`DictionaryProvider`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/dictionary/DictionaryProvider.html
+.. _`Field`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Field.html
+.. _`FieldReader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/complex/reader/FieldReader.html
+.. _`FieldVector`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/FieldVector.html
+.. _`Row`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Row.html
+.. _`Schema`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Schema.html
+.. _`Table`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Table.html
+.. _`ValueHolder`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/holders/ValueHolder.html
+.. _`ValueVector`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/ValueVector.html
+.. _`VectorSchemaRoot`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorSchemaRoot.html
diff --git a/docs/source/vector_schema_root.rst b/docs/source/vector_schema_root.rst
index 3119122d9a..f4a497c4e5 100644
--- a/docs/source/vector_schema_root.rst
+++ b/docs/source/vector_schema_root.rst
@@ -153,11 +153,11 @@ A `Table`_ is an immutable tabular data structure, very similar to VectorSchemaR
See the :doc:`table` documentation for more information.
-.. _`ArrowRecordBatch`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.html
-.. _`Field`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Field.html
-.. _`Flight`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/package-summary.html
-.. _`Schema`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Schema.html
-.. _`Table`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Table.html
-.. _`VectorLoader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorLoader.html
-.. _`VectorSchemaRoot`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorSchemaRoot.html
-.. _`VectorUnloader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorUnloader.html
+.. _`ArrowRecordBatch`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.html
+.. _`Field`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Field.html
+.. _`Flight`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/package-summary.html
+.. _`Schema`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Schema.html
+.. _`Table`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Table.html
+.. _`VectorLoader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorLoader.html
+.. _`VectorSchemaRoot`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorSchemaRoot.html
+.. _`VectorUnloader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorUnloader.html
diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml
index 757de85769..fbed544a1b 100644
--- a/flight/flight-core/pom.xml
+++ b/flight/flight-core/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 18.3.0
+ 19.0.0
flight-core
@@ -134,7 +134,7 @@ under the License.
com.google.api.grpc
proto-google-common-protos
- 2.54.1
+ 2.66.0
test
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java
index 9cefccb3fe..ab4eab3048 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java
@@ -287,7 +287,11 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s
ArrowBuf body = null;
ArrowBuf appMetadata = null;
while (stream.available() > 0) {
- int tag = readRawVarint32(stream);
+ final int tagFirstByte = stream.read();
+ if (tagFirstByte == -1) {
+ break;
+ }
+ int tag = readRawVarint32(tagFirstByte, stream);
switch (tag) {
case DESCRIPTOR_TAG:
{
@@ -366,6 +370,10 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s
private static int readRawVarint32(InputStream is) throws IOException {
int firstByte = is.read();
+ return readRawVarint32(firstByte, is);
+ }
+
+ private static int readRawVarint32(int firstByte, InputStream is) throws IOException {
return CodedInputStream.readRawVarint32(firstByte, is);
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java
index f4f6486a3c..0939d232cf 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java
@@ -26,10 +26,20 @@ public interface CallHeaders {
/** Get the value of a metadata key. If multiple values are present, then get the last one. */
byte[] getByte(String key);
- /** Get all values present for the given metadata key. */
+ /**
+ * Get all values present for the given metadata key.
+ *
+ * @param key the metadata key
+ * @return an iterable of all values for the key. Returns an empty iterable if no value to return.
+ */
Iterable getAll(String key);
- /** Get all values present for the given metadata key. */
+ /**
+ * Get all values present for the given metadata key.
+ *
+ * @param key the metadata key
+ * @return an iterable of all values for the key. Returns an empty iterable if no value to return.
+ */
Iterable getAllByte(String key);
/**
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java
index 47fd6f1366..5ec01b9c83 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java
@@ -80,20 +80,18 @@ public ServerSessionMiddleware onCallStarted(
String sessionId = null;
final Iterable it = incomingHeaders.getAll("cookie");
- if (it != null) {
- findIdCookie:
- for (final String headerValue : it) {
- for (final String cookie : headerValue.split(" ;")) {
- final String[] cookiePair = cookie.split("=");
- if (cookiePair.length != 2) {
- // Soft failure: Ignore invalid cookie list field
- break;
- }
-
- if (sessionCookieName.equals(cookiePair[0]) && cookiePair[1].length() > 0) {
- sessionId = cookiePair[1];
- break findIdCookie;
- }
+ findIdCookie:
+ for (final String headerValue : it) {
+ for (final String cookie : headerValue.split(" ;")) {
+ final String[] cookiePair = cookie.split("=");
+ if (cookiePair.length != 2) {
+ // Soft failure: Ignore invalid cookie list field
+ break;
+ }
+
+ if (sessionCookieName.equals(cookiePair[0]) && cookiePair[1].length() > 0) {
+ sessionId = cookiePair[1];
+ break findIdCookie;
}
}
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java
index e5eb934001..b33e6b7ecc 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java
@@ -100,10 +100,7 @@ public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
@Override
public void onHeadersReceived(CallHeaders incomingHeaders) {
- final Iterable setCookieHeaders = incomingHeaders.getAll(SET_COOKIE_HEADER);
- if (setCookieHeaders != null) {
- factory.updateCookies(setCookieHeaders);
- }
+ factory.updateCookies(incomingHeaders.getAll(SET_COOKIE_HEADER));
}
@Override
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java
index f33e9b2f94..fe81f3fb23 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java
@@ -18,6 +18,7 @@
import io.grpc.CallCredentials;
import io.grpc.Metadata;
+import io.grpc.Status;
import java.util.concurrent.Executor;
import java.util.function.Consumer;
import org.apache.arrow.flight.CallHeaders;
@@ -36,9 +37,14 @@ public void applyRequestMetadata(
RequestInfo requestInfo, Executor executor, MetadataApplier metadataApplier) {
executor.execute(
() -> {
- final Metadata headers = new Metadata();
- credentialWriter.accept(new MetadataAdapter(headers));
- metadataApplier.apply(headers);
+ try {
+ final Metadata headers = new Metadata();
+ credentialWriter.accept(new MetadataAdapter(headers));
+ metadataApplier.apply(headers);
+ } catch (Throwable t) {
+ metadataApplier.fail(
+ Status.UNAUTHENTICATED.withCause(t).withDescription(t.getMessage()));
+ }
});
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java
index 45c32a86c6..fcba88d212 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java
@@ -87,13 +87,13 @@ public static void readIntoBuffer(
final InputStream stream, final ArrowBuf buf, final int size, final boolean fastPath)
throws IOException {
ReadableBuffer readableBuffer = fastPath ? getReadableBuffer(stream) : null;
+ byte[] heapBytes = new byte[size];
if (readableBuffer != null) {
- readableBuffer.readBytes(buf.nioBuffer(0, size));
+ readableBuffer.readBytes(heapBytes, 0, size);
} else {
- byte[] heapBytes = new byte[size];
ByteStreams.readFully(stream, heapBytes);
- buf.writeBytes(heapBytes);
}
+ buf.writeBytes(heapBytes);
buf.writerIndex(size);
}
}
diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java
index a1de16ede6..64a0769d63 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java
@@ -18,6 +18,7 @@
import io.grpc.Metadata;
import java.nio.charset.StandardCharsets;
+import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;
@@ -53,13 +54,17 @@ public byte[] getByte(String key) {
@Override
public Iterable getAll(String key) {
- return this.metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER));
+ final Iterable all =
+ this.metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER));
+ return all != null ? all : Collections.emptyList();
}
@Override
public Iterable getAllByte(String key) {
if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
- return this.metadata.getAll(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER));
+ final Iterable all =
+ this.metadata.getAll(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER));
+ return all != null ? all : Collections.emptyList();
}
return StreamSupport.stream(getAll(key).spliterator(), false)
.map(String::getBytes)
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java
index a54ce69812..8aef9c69a1 100644
--- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java
@@ -21,6 +21,7 @@
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.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
@@ -110,6 +111,16 @@ public void mixedProperties() {
testHeaders(headers);
}
+ @Test
+ public void getAllReturnsEmptyIterableForMissingKey() {
+ FlightCallHeaders headers = new FlightCallHeaders();
+
+ assertNotNull(headers.getAll("missing"));
+ assertFalse(headers.getAll("missing").iterator().hasNext());
+ assertNotNull(headers.getAllByte("missing-bin"));
+ assertFalse(headers.getAllByte("missing-bin").iterator().hasNext());
+ }
+
private void testHeaders(CallHeaders headers) {
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
HeaderProducer producer = new HeaderProducer();
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java
index a9a3e355bc..214614defd 100644
--- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java
@@ -20,6 +20,7 @@
import static org.apache.arrow.flight.Location.forGrpcInsecure;
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.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
@@ -119,6 +120,16 @@ public void testFlightMetadata() throws Exception {
}
}
+ @Test
+ public void getAllReturnsEmptyIterableForMissingKey() {
+ ErrorFlightMetadata metadata = new ErrorFlightMetadata();
+
+ assertNotNull(metadata.getAll("missing"));
+ assertFalse(metadata.getAll("missing").iterator().hasNext());
+ assertNotNull(metadata.getAllByte("missing-bin"));
+ assertFalse(metadata.getAllByte("missing-bin").iterator().hasNext());
+ }
+
private static class StatusRuntimeExceptionProducer extends NoOpFlightProducer {
private final PerfOuterClass.Perf perf;
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
index 0c63785c88..0f202ba2d9 100644
--- a/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
@@ -178,6 +178,12 @@ public static void shutdown() throws Exception {
AutoCloseables.close(server);
allocator.getChildAllocators().forEach(BufferAllocator::close);
+
+ // gRPC/Netty may still be releasing Arrow buffers asynchronously after server shutdown.
+ // Poll briefly to allow in-flight buffer releases to complete before closing the allocator.
+ for (int i = 0; i < 20 && allocator.getAllocatedMemory() > 0; i++) {
+ Thread.sleep(100);
+ }
AutoCloseables.close(allocator);
}
}
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java
new file mode 100644
index 0000000000..b0f5dcfcfc
--- /dev/null
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java
@@ -0,0 +1,36 @@
+/*
+ * 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.flight.grpc;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import io.grpc.Metadata;
+import org.junit.jupiter.api.Test;
+
+public class TestMetadataAdapter {
+
+ @Test
+ public void getAllReturnsEmptyIterableForMissingKey() {
+ MetadataAdapter headers = new MetadataAdapter(new Metadata());
+
+ assertNotNull(headers.getAll("missing"));
+ assertFalse(headers.getAll("missing").iterator().hasNext());
+ assertNotNull(headers.getAllByte("missing-bin"));
+ assertFalse(headers.getAllByte("missing-bin").iterator().hasNext());
+ }
+}
diff --git a/flight/flight-integration-tests/pom.xml b/flight/flight-integration-tests/pom.xml
index e7fb999149..5ee7c6fc14 100644
--- a/flight/flight-integration-tests/pom.xml
+++ b/flight/flight-integration-tests/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 18.3.0
+ 19.0.0
flight-integration-tests
@@ -58,7 +58,7 @@ under the License.
commons-cli
commons-cli
- 1.9.0
+ 1.11.0
org.slf4j
diff --git a/flight/flight-sql-jdbc-core/pom.xml b/flight/flight-sql-jdbc-core/pom.xml
index a95fbaca30..d6fa11688d 100644
--- a/flight/flight-sql-jdbc-core/pom.xml
+++ b/flight/flight-sql-jdbc-core/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 18.3.0
+ 19.0.0
flight-sql-jdbc-core
@@ -105,7 +105,7 @@ under the License.
commons-io
commons-io
- 2.19.0
+ 2.21.0
test
@@ -120,6 +120,31 @@ under the License.
test
+
+ com.squareup.okhttp3
+ mockwebserver3
+ 5.3.2
+ test
+
+
+ com.squareup.okhttp3
+ mockwebserver3-junit5
+ 5.3.2
+ test
+
+
+ com.squareup.okhttp3
+ okhttp-jvm
+ 5.3.2
+ test
+
+
+ com.squareup.okio
+ okio-jvm
+ 3.16.4
+ test
+
+
io.netty
netty-common
@@ -140,7 +165,7 @@ under the License.
org.bouncycastle
bcpkix-jdk18on
- 1.80
+ 1.83
@@ -151,8 +176,15 @@ under the License.
com.github.ben-manes.caffeine
caffeine
- 3.1.8
+ 3.2.3
+
+
+ com.nimbusds
+ oauth2-oidc-sdk
+ 11.20.1
+
+
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java
index 7185ddfe01..0110525fea 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java
@@ -45,6 +45,7 @@
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Arrays;
+import java.util.Collections;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.List;
@@ -75,18 +76,23 @@
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.Text;
import org.apache.calcite.avatica.AvaticaConnection;
import org.apache.calcite.avatica.AvaticaDatabaseMetaData;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/** Arrow Flight JDBC's implementation of {@link DatabaseMetaData}. */
public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
+ private static final Logger LOGGER = LoggerFactory.getLogger(ArrowDatabaseMetadata.class);
private static final String JAVA_REGEX_SPECIALS = "[]()|^-+*?{}$\\.";
private static final Charset CHARSET = StandardCharsets.UTF_8;
private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
@@ -164,6 +170,9 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData {
LONGNVARCHAR, SqlSupportsConvert.SQL_CONVERT_LONGVARCHAR_VALUE);
sqlTypesToFlightEnumConvertTypes.put(DATE, SqlSupportsConvert.SQL_CONVERT_DATE_VALUE);
sqlTypesToFlightEnumConvertTypes.put(TIMESTAMP, SqlSupportsConvert.SQL_CONVERT_TIMESTAMP_VALUE);
+
+ // Register the UUID extension type so it is always available for the driver
+ ExtensionTypeRegistry.register(UuidType.INSTANCE);
}
ArrowDatabaseMetadata(final AvaticaConnection connection) {
@@ -769,7 +778,34 @@ private T getSqlInfoAndCacheIfCacheIsEmpty(
}
}
}
- return desiredType.cast(cachedSqlInfo.get(sqlInfoCommand));
+ T value = desiredType.cast(cachedSqlInfo.get(sqlInfoCommand));
+ if (value != null) {
+ return value;
+ }
+ LOGGER.debug(
+ "SqlInfo {} not provided by server, returning default for type {}",
+ sqlInfoCommand.name(),
+ desiredType.getSimpleName());
+
+ // Return sensible defaults when SqlInfo is unavailable
+ if (desiredType == Long.class) {
+ return desiredType.cast(0L);
+ } else if (desiredType == Integer.class) {
+ return desiredType.cast(0);
+ } else if (desiredType == Boolean.class) {
+ return desiredType.cast(false);
+ } else if (desiredType == String.class) {
+ return desiredType.cast("");
+ } else if (desiredType == Map.class) {
+ return desiredType.cast(Collections.emptyMap());
+ } else if (desiredType == List.class) {
+ return desiredType.cast(Collections.emptyList());
+ }
+
+ throw new SQLException(
+ String.format(
+ "The value of the SqlInfo %s is null and it could not be cast to %s.",
+ sqlInfoCommand.name(), desiredType.getName()));
}
private Optional convertListSqlInfoToString(final List sqlInfoList) {
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
index 747287ed13..623c2b81be 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java
@@ -20,6 +20,9 @@
import io.netty.util.concurrent.DefaultThreadFactory;
import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -32,6 +35,7 @@
import org.apache.arrow.util.Preconditions;
import org.apache.calcite.avatica.AvaticaConnection;
import org.apache.calcite.avatica.AvaticaFactory;
+import org.apache.calcite.avatica.DriverVersion;
/** Connection to the Arrow Flight server. */
public final class ArrowFlightConnection extends AvaticaConnection {
@@ -40,6 +44,8 @@ public final class ArrowFlightConnection extends AvaticaConnection {
private final ArrowFlightSqlClientHandler clientHandler;
private final ArrowFlightConnectionConfigImpl config;
private ExecutorService executorService;
+ private int metadataResultSetCount;
+ private Map metadataResultSetMap = new HashMap<>();
/**
* Creates a new {@link ArrowFlightConnection}.
@@ -64,6 +70,7 @@ private ArrowFlightConnection(
this.config = Preconditions.checkNotNull(config, "Config cannot be null.");
this.allocator = Preconditions.checkNotNull(allocator, "Allocator cannot be null.");
this.clientHandler = Preconditions.checkNotNull(clientHandler, "Handler cannot be null.");
+ this.metadataResultSetCount = 0;
}
/**
@@ -86,13 +93,16 @@ static ArrowFlightConnection createNewConnection(
throws SQLException {
url = replaceSemiColons(url);
final ArrowFlightConnectionConfigImpl config = new ArrowFlightConnectionConfigImpl(properties);
- final ArrowFlightSqlClientHandler clientHandler = createNewClientHandler(config, allocator);
+ final ArrowFlightSqlClientHandler clientHandler =
+ createNewClientHandler(config, allocator, driver.getDriverVersion());
return new ArrowFlightConnection(
driver, factory, url, properties, config, allocator, clientHandler);
}
private static ArrowFlightSqlClientHandler createNewClientHandler(
- final ArrowFlightConnectionConfigImpl config, final BufferAllocator allocator)
+ final ArrowFlightConnectionConfigImpl config,
+ final BufferAllocator allocator,
+ final DriverVersion driverVersion)
throws SQLException {
try {
return new ArrowFlightSqlClientHandler.Builder()
@@ -116,6 +126,8 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
.withCatalog(config.getCatalog())
.withClientCache(config.useClientCache() ? new FlightClientCache() : null)
.withConnectTimeout(config.getConnectTimeout())
+ .withDriverVersion(driverVersion)
+ .withOAuthConfiguration(config.getOauthConfiguration())
.build();
} catch (final SQLException e) {
try {
@@ -166,6 +178,31 @@ synchronized ExecutorService getExecutorService() {
: executorService;
}
+ /**
+ * Registers a new metadata ResultSet and assigns it a unique ID. Metadata ResultSets are those
+ * created without an associated Statement.
+ *
+ * @param resultSet the ResultSet to register
+ * @return the assigned ID
+ */
+ int getNewMetadataResultSetId(ArrowFlightJdbcFlightStreamResultSet resultSet) {
+ metadataResultSetMap.put(metadataResultSetCount, resultSet);
+ return metadataResultSetCount++;
+ }
+
+ /**
+ * Unregisters a metadata ResultSet when it is closed. This method is called by metadata
+ * ResultSets during their close operation to remove themselves from the tracking map.
+ *
+ * @param id the ID of the ResultSet to unregister, or null if not a metadata ResultSet
+ */
+ void onResultSetClose(Integer id) {
+ if (id == null) {
+ return;
+ }
+ metadataResultSetMap.remove(id);
+ }
+
@Override
public Properties getClientInfo() {
final Properties copy = new Properties();
@@ -175,19 +212,41 @@ public Properties getClientInfo() {
@Override
public void close() throws SQLException {
- clientHandler.close();
- if (executorService != null) {
- executorService.shutdown();
+ Exception topLevelException = null;
+ try {
+ if (executorService != null) {
+ executorService.shutdown();
+ }
+ } catch (final Exception e) {
+ topLevelException = e;
+ }
+ // copies of the collections are used to avoid concurrent modification problems
+ ArrayList closeables = new ArrayList<>(statementMap.values());
+ closeables.addAll(new ArrayList<>(metadataResultSetMap.values()));
+ closeables.add(clientHandler);
+ closeables.addAll(allocator.getChildAllocators());
+ closeables.add(allocator);
+ try {
+ AutoCloseables.close(closeables);
+ } catch (final Exception e) {
+ if (topLevelException == null) {
+ topLevelException = e;
+ } else {
+ topLevelException.addSuppressed(e);
+ }
}
-
try {
- AutoCloseables.close(clientHandler);
- allocator.getChildAllocators().forEach(AutoCloseables::closeNoChecked);
- AutoCloseables.close(allocator);
-
super.close();
} catch (final Exception e) {
- throw AvaticaConnection.HELPER.createException(e.getMessage(), e);
+ if (topLevelException == null) {
+ topLevelException = e;
+ } else {
+ topLevelException.addSuppressed(e);
+ }
+ }
+ if (topLevelException != null) {
+ throw AvaticaConnection.HELPER.createException(
+ topLevelException.getMessage(), topLevelException);
}
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java
index 9b9eba51e5..f3d76ace92 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java
@@ -26,6 +26,7 @@
import org.apache.arrow.driver.jdbc.utils.SqlTypes;
import org.apache.arrow.memory.util.LargeMemoryUtil;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -135,12 +136,22 @@ public ResultSet getResultSet(long index, int count) throws SQLException {
private static ResultSet getResultSetNoBoundariesCheck(
ValueVector dataVector, long start, long count) throws SQLException {
+ int intStart = LargeMemoryUtil.checkedCastToInt(start);
+ int intCount = LargeMemoryUtil.checkedCastToInt(count);
+
+ // Create an index vector with 1-based indices (per JDBC spec) to return with value vector
+ IntVector indexVector = new IntVector("INDEX", dataVector.getAllocator());
+ indexVector.allocateNew(intCount);
+ for (int i = 0; i < intCount; i++) {
+ indexVector.set(i, i + 1);
+ }
+ indexVector.setValueCount(intCount);
+
TransferPair transferPair = dataVector.getTransferPair(dataVector.getAllocator());
- transferPair.splitAndTransfer(
- LargeMemoryUtil.checkedCastToInt(start), LargeMemoryUtil.checkedCastToInt(count));
- FieldVector vectorSlice = (FieldVector) transferPair.getTo();
+ transferPair.splitAndTransfer(intStart, intCount);
+ FieldVector valueVector = (FieldVector) transferPair.getTo();
- VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(vectorSlice);
+ VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(indexVector, valueVector);
return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(vectorSchemaRoot);
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
index 53e6120f62..12ef8030d7 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java
@@ -75,7 +75,9 @@ public Logger getParentLogger() {
public ArrowFlightConnection connect(final String url, final Properties info)
throws SQLException {
final Properties properties = new Properties(info);
- properties.putAll(info);
+ if (info != null) {
+ properties.putAll(info);
+ }
if (url != null) {
final Optional> maybeProperties = getUrlsArgs(url);
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
index aabaf01e63..2885f7895b 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
@@ -54,6 +54,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
private VectorSchemaRoot currentVectorSchemaRoot;
private Schema schema;
+ private Integer id = null; // used for metadata result sets only
/** Public constructor used by ArrowFlightJdbcFactory. */
ArrowFlightJdbcFlightStreamResultSet(
@@ -82,6 +83,7 @@ private ArrowFlightJdbcFlightStreamResultSet(
super(null, state, signature, resultSetMetaData, timeZone, firstFrame);
this.connection = connection;
this.flightInfo = flightInfo;
+ this.id = connection.getNewMetadataResultSetId(this);
}
/**
@@ -234,7 +236,12 @@ protected void cancel() {
@Override
public synchronized void close() {
+
try {
+ if (isClosed()) {
+ return;
+ }
+ this.connection.onResultSetClose(id);
if (flightEndpointDataQueue != null) {
// flightStreamQueue should close currentFlightStream internally
flightEndpointDataQueue.close();
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
index 0dc2b07c97..49334951de 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
@@ -19,23 +19,26 @@
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
+import java.sql.Types;
import java.util.HashSet;
import java.util.List;
-import java.util.Objects;
import java.util.Set;
import java.util.TimeZone;
import org.apache.arrow.driver.jdbc.utils.ConvertUtils;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.calcite.avatica.AvaticaConnection;
import org.apache.calcite.avatica.AvaticaResultSet;
import org.apache.calcite.avatica.AvaticaResultSetMetaData;
+import org.apache.calcite.avatica.AvaticaSite;
import org.apache.calcite.avatica.AvaticaStatement;
import org.apache.calcite.avatica.ColumnMetaData;
import org.apache.calcite.avatica.Meta;
import org.apache.calcite.avatica.Meta.Frame;
import org.apache.calcite.avatica.Meta.Signature;
import org.apache.calcite.avatica.QueryState;
+import org.apache.calcite.avatica.util.Cursor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -102,6 +105,33 @@ void populateData(final VectorSchemaRoot vectorSchemaRoot, final Schema schema)
execute2(new ArrowFlightJdbcCursor(vectorSchemaRoot), this.signature.columns);
}
+ /**
+ * The default method in AvaticaResultSet does not properly handle TIMESTASMP_WITH_TIMEZONE, so we
+ * override here to add support.
+ *
+ * @param columnIndex the first column is 1, the second is 2, ...
+ * @return Object
+ * @throws SQLException if there is an underlying exception
+ */
+ @Override
+ public Object getObject(int columnIndex) throws SQLException {
+ this.checkOpen();
+
+ Cursor.Accessor accessor;
+ try {
+ accessor = accessorList.get(columnIndex - 1);
+ } catch (IndexOutOfBoundsException e) {
+ throw AvaticaConnection.HELPER.createException("invalid column ordinal: " + columnIndex);
+ }
+
+ ColumnMetaData metaData = columnMetaDataList.get(columnIndex - 1);
+ if (metaData.type.id == Types.TIMESTAMP_WITH_TIMEZONE) {
+ return accessor.getTimestamp(localCalendar);
+ } else {
+ return AvaticaSite.get(accessor, metaData.type.id, localCalendar);
+ }
+ }
+
@Override
protected void cancel() {
signature.columns.clear();
@@ -128,12 +158,10 @@ public void close() {
} catch (final Exception e) {
exceptions.add(e);
}
- if (!Objects.isNull(statement)) {
- try {
- super.close();
- } catch (final Exception e) {
- exceptions.add(e);
- }
+ try {
+ super.close();
+ } catch (final Exception e) {
+ exceptions.add(e);
}
exceptions.parallelStream().forEach(e -> LOGGER.error(e.getMessage(), e));
exceptions.stream()
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
index 9c7112f1c3..64529b50c8 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
@@ -62,21 +62,25 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
parameterSchema == null
? new ArrayList<>()
: ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields());
-
+ StatementType statementType =
+ resultSetSchema == null || resultSetSchema.getFields().isEmpty()
+ ? StatementType.IS_DML
+ : StatementType.SELECT;
return new Signature(
columnMetaData,
sql,
parameters,
Collections.emptyMap(),
null, // unnecessary, as SQL requests use ArrowFlightJdbcCursor
- StatementType.SELECT);
+ statementType);
}
@Override
public void closeStatement(final StatementHandle statementHandle) {
PreparedStatement preparedStatement =
statementHandlePreparedStatementMap.remove(new StatementHandleKey(statementHandle));
- // Testing if the prepared statement was created because the statement can be not created until
+ // Testing if the prepared statement was created because the statement can be
+ // not created until
// this moment
if (preparedStatement != null) {
preparedStatement.close();
@@ -105,7 +109,8 @@ public ExecuteResult execute(
preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator())
.bind(typedValues);
- if (statementHandle.signature == null) {
+ if (statementHandle.signature == null
+ || statementHandle.signature.statementType == StatementType.IS_DML) {
// Update query
long updatedCount = preparedStatement.executeUpdate();
return new ExecuteResult(
@@ -220,7 +225,8 @@ public ExecuteResult prepareAndExecute(
MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null);
return new ExecuteResult(Collections.singletonList(metaResultSet));
} catch (SQLTimeoutException e) {
- // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and Runtime
+ // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and
+ // Runtime
// Exceptions.
throw new RuntimeException(e);
} catch (SQLException e) {
@@ -249,6 +255,20 @@ public boolean syncResults(
return false;
}
+ @Override
+ public ConnectionProperties connectionSync(ConnectionHandle ch, ConnectionProperties connProps) {
+ final ConnectionProperties result = super.connectionSync(ch, connProps);
+ final String newCatalog = this.connProps.getCatalog();
+ if (newCatalog != null) {
+ try {
+ ((ArrowFlightConnection) connection).getClientHandler().setCatalog(newCatalog);
+ } catch (SQLException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ return result;
+ }
+
void setDefaultConnectionProperties() {
// TODO Double-check this.
connProps
@@ -264,7 +284,8 @@ PreparedStatement getPreparedStatement(StatementHandle statementHandle) {
return statementHandlePreparedStatementMap.get(new StatementHandleKey(statementHandle));
}
- // Helper used to look up prepared statement instances later. Avatica doesn't give us the
+ // Helper used to look up prepared statement instances later. Avatica doesn't
+ // give us the
// signature in
// an UPDATE code path so we can't directly use StatementHandle as a map key.
private static final class StatementHandleKey {
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java
index dad1fa5f73..8362eb7627 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java
@@ -19,6 +19,7 @@
import java.util.function.IntSupplier;
import org.apache.arrow.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor;
+import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor;
@@ -65,9 +66,12 @@
import org.apache.arrow.vector.UInt2Vector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.UInt8Vector;
+import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarBinaryVector;
+import org.apache.arrow.vector.ViewVarCharVector;
import org.apache.arrow.vector.complex.DenseUnionVector;
import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.LargeListVector;
@@ -130,9 +134,15 @@ public static ArrowFlightJdbcAccessor createAccessor(
} else if (vector instanceof VarBinaryVector) {
return new ArrowFlightJdbcBinaryVectorAccessor(
(VarBinaryVector) vector, getCurrentRow, setCursorWasNull);
+ } else if (vector instanceof ViewVarBinaryVector) {
+ return new ArrowFlightJdbcBinaryVectorAccessor(
+ (ViewVarBinaryVector) vector, getCurrentRow, setCursorWasNull);
} else if (vector instanceof LargeVarBinaryVector) {
return new ArrowFlightJdbcBinaryVectorAccessor(
(LargeVarBinaryVector) vector, getCurrentRow, setCursorWasNull);
+ } else if (vector instanceof UuidVector) {
+ return new ArrowFlightJdbcUuidVectorAccessor(
+ (UuidVector) vector, getCurrentRow, setCursorWasNull);
} else if (vector instanceof FixedSizeBinaryVector) {
return new ArrowFlightJdbcBinaryVectorAccessor(
(FixedSizeBinaryVector) vector, getCurrentRow, setCursorWasNull);
@@ -163,6 +173,9 @@ public static ArrowFlightJdbcAccessor createAccessor(
} else if (vector instanceof LargeVarCharVector) {
return new ArrowFlightJdbcVarCharVectorAccessor(
(LargeVarCharVector) vector, getCurrentRow, setCursorWasNull);
+ } else if (vector instanceof ViewVarCharVector) {
+ return new ArrowFlightJdbcVarCharVectorAccessor(
+ (ViewVarCharVector) vector, getCurrentRow, setCursorWasNull);
} else if (vector instanceof DurationVector) {
return new ArrowFlightJdbcDurationVectorAccessor(
(DurationVector) vector, getCurrentRow, setCursorWasNull);
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java
index 30dfffce64..e71b6380a9 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java
@@ -27,6 +27,7 @@
import org.apache.arrow.vector.FixedSizeBinaryVector;
import org.apache.arrow.vector.LargeVarBinaryVector;
import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.ViewVarBinaryVector;
/**
* Accessor for the Arrow types: {@link FixedSizeBinaryVector}, {@link VarBinaryVector} and {@link
@@ -61,6 +62,13 @@ public ArrowFlightJdbcBinaryVectorAccessor(
this(vector::get, currentRowSupplier, setCursorWasNull);
}
+ public ArrowFlightJdbcBinaryVectorAccessor(
+ ViewVarBinaryVector vector,
+ IntSupplier currentRowSupplier,
+ ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
+ this(vector::get, currentRowSupplier, setCursorWasNull);
+ }
+
private ArrowFlightJdbcBinaryVectorAccessor(
ByteArrayGetter getter,
IntSupplier currentRowSupplier,
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java
new file mode 100644
index 0000000000..4bdbcbb63d
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.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.accessor.impl.binary;
+
+import java.util.UUID;
+import java.util.function.IntSupplier;
+import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor;
+import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.util.UuidUtility;
+
+/**
+ * Accessor for the Arrow UUID extension type ({@link UuidVector}).
+ *
+ * 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:
+ *
+ *
+ * {@link #getObject()} returns {@link java.util.UUID}
+ * {@link #getString()} returns the hyphenated string format (e.g.,
+ * "550e8400-e29b-41d4-a716-446655440000")
+ * {@link #getBytes()} returns the 16-byte binary representation
+ *
+ */
+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 17c2c16ebf..f0ea284239 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;
@@ -66,6 +68,7 @@
import org.apache.arrow.util.VisibleForTesting;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.calcite.avatica.DriverVersion;
import org.apache.calcite.avatica.Meta.StatementType;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
@@ -146,20 +149,26 @@ public List getStreams(final FlightInfo flightInfo)
try {
for (FlightEndpoint endpoint : flightInfo.getEndpoints()) {
if (endpoint.getLocations().isEmpty()) {
- // Create a stream using the current client only and do not close the client at the end.
+ // Create a stream using the current client only and do not close the client at
+ // the end.
endpoints.add(
new CloseableEndpointStreamPair(
sqlClient.getStream(endpoint.getTicket(), getOptions()), null));
} else {
// Clone the builder and then set the new endpoint on it.
- // GH-38574: Currently a new FlightClient will be made for each partition that returns a
- // non-empty Location then disposed of. It may be better to cache clients because a server
- // may report the same Locations. It would also be good to identify when the reported
+ // GH-38574: Currently a new FlightClient will be made for each partition that
+ // returns a
+ // non-empty Location then disposed of. It may be better to cache clients
+ // because a server
+ // may report the same Locations. It would also be good to identify when the
+ // reported
// location
- // is the same as the original connection's Location and skip creating a FlightClient in
+ // is the same as the original connection's Location and skip creating a
+ // FlightClient in
// that scenario.
- // Also copy the cache to the client so we can share a cache. Cache needs to cache
+ // Also copy the cache to the client so we can share a cache. Cache needs to
+ // cache
// negative attempts too.
List exceptions = new ArrayList<>();
CloseableEndpointStreamPair stream = null;
@@ -261,15 +270,86 @@ public FlightInfo getInfo(final String query) {
@Override
public void close() throws SQLException {
if (catalog.isPresent()) {
- sqlClient.closeSession(new CloseSessionRequest(), getOptions());
+ try {
+ sqlClient.closeSession(new CloseSessionRequest(), getOptions());
+ } catch (FlightRuntimeException fre) {
+ handleBenignCloseException(
+ fre, "Failed to close Flight SQL session.", "closing Flight SQL session");
+ }
}
try {
AutoCloseables.close(sqlClient);
+ } catch (FlightRuntimeException fre) {
+ handleBenignCloseException(
+ fre, "Failed to clean up client resources.", "closing Flight SQL client");
} catch (final Exception e) {
throw new SQLException("Failed to clean up client resources.", e);
}
}
+ /**
+ * Handles FlightRuntimeException during close operations, suppressing benign gRPC shutdown errors
+ * while re-throwing genuine failures.
+ *
+ * @param fre the FlightRuntimeException to handle
+ * @param sqlErrorMessage the SQLException message to use for genuine failures
+ * @param operationDescription description of the operation for logging
+ * @throws SQLException if the exception represents a genuine failure
+ */
+ private void handleBenignCloseException(
+ FlightRuntimeException fre, String sqlErrorMessage, String operationDescription)
+ throws SQLException {
+ if (isBenignCloseException(fre)) {
+ logSuppressedCloseException(fre, operationDescription);
+ } else {
+ throw new SQLException(sqlErrorMessage, fre);
+ }
+ }
+
+ /**
+ * Handles FlightRuntimeException during close operations, suppressing benign gRPC shutdown errors
+ * while re-throwing genuine failures as FlightRuntimeException.
+ *
+ * @param fre the FlightRuntimeException to handle
+ * @param operationDescription description of the operation for logging
+ * @throws FlightRuntimeException if the exception represents a genuine failure
+ */
+ private void handleBenignCloseException(FlightRuntimeException fre, String operationDescription)
+ throws FlightRuntimeException {
+ if (isBenignCloseException(fre)) {
+ logSuppressedCloseException(fre, operationDescription);
+ } else {
+ throw fre;
+ }
+ }
+
+ /**
+ * Determines if a FlightRuntimeException represents a benign close operation error that should be
+ * suppressed.
+ *
+ * @param fre the FlightRuntimeException to check
+ * @return true if the exception should be suppressed, false otherwise
+ */
+ private boolean isBenignCloseException(FlightRuntimeException fre) {
+ return fre.status().code().equals(FlightStatusCode.UNAVAILABLE)
+ || (fre.status().code().equals(FlightStatusCode.INTERNAL)
+ && fre.getMessage() != null
+ && fre.getMessage().contains("Connection closed after GOAWAY"));
+ }
+
+ /**
+ * Logs a suppressed close exception with appropriate level based on debug settings.
+ *
+ * @param fre the FlightRuntimeException being suppressed
+ * @param operationDescription description of the operation for logging
+ */
+ private void logSuppressedCloseException(
+ FlightRuntimeException fre, String operationDescription) {
+ // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during
+ // shutdown
+ LOGGER.debug("Suppressed error {}", operationDescription, fre);
+ }
+
/** A prepared statement handler. */
public interface PreparedStatement extends AutoCloseable {
/**
@@ -317,25 +397,40 @@ public interface PreparedStatement extends AutoCloseable {
/** A connection is created with catalog set as a session option. */
private void setSetCatalogInSessionIfPresent() {
if (catalog.isPresent()) {
- final SetSessionOptionsRequest setSessionOptionRequest =
- new SetSessionOptionsRequest(
- ImmutableMap.builder()
- .put(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog.get()))
- .build());
- final SetSessionOptionsResult result =
- sqlClient.setSessionOptions(setSessionOptionRequest, getOptions());
+ try {
+ setCatalog(catalog.get());
+ } catch (SQLException e) {
+ throw CallStatus.INVALID_ARGUMENT
+ .withDescription(e.getMessage())
+ .withCause(e)
+ .toRuntimeException();
+ }
+ }
+ }
+ /**
+ * Sets the catalog for the current session.
+ *
+ * @param catalog the catalog to set.
+ * @throws SQLException if an error occurs while setting the catalog.
+ */
+ public void setCatalog(final String catalog) throws SQLException {
+ final SetSessionOptionsRequest request =
+ new SetSessionOptionsRequest(
+ ImmutableMap.of(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog)));
+ try {
+ final SetSessionOptionsResult result = sqlClient.setSessionOptions(request, getOptions());
if (result.hasErrors()) {
- Map errors = result.getErrors();
- for (Map.Entry error : errors.entrySet()) {
+ final Map errors = result.getErrors();
+ for (final Map.Entry error : errors.entrySet()) {
LOGGER.warn(error.toString());
}
- throw CallStatus.INVALID_ARGUMENT
- .withDescription(
- String.format(
- "Cannot set session option for catalog = %s. Check log for details.", catalog))
- .toRuntimeException();
+ throw new SQLException(
+ String.format(
+ "Cannot set session option for catalog = %s. Check log for details.", catalog));
}
+ } catch (final FlightRuntimeException e) {
+ throw new SQLException(e);
}
}
@@ -385,14 +480,7 @@ public void close() {
try {
preparedStatement.close(getOptions());
} catch (FlightRuntimeException fre) {
- // ARROW-17785: suppress exceptions caused by flaky gRPC layer
- if (fre.status().code().equals(FlightStatusCode.UNAVAILABLE)
- || (fre.status().code().equals(FlightStatusCode.INTERNAL)
- && fre.getMessage().contains("Connection closed after GOAWAY"))) {
- LOGGER.warn("Supressed error closing PreparedStatement", fre);
- return;
- }
- throw fre;
+ handleBenignCloseException(fre, "closing PreparedStatement");
}
}
};
@@ -548,6 +636,9 @@ public FlightInfo getCrossReference(
/** Builder for {@link ArrowFlightSqlClientHandler}. */
public static final class Builder {
+ static final String USER_AGENT_TEMPLATE = "JDBC Flight SQL Driver %s";
+ static final String DEFAULT_VERSION = "(unknown or development build)";
+
private final Set middlewareFactories = new HashSet<>();
private final Set options = new HashSet<>();
private String host;
@@ -587,7 +678,10 @@ public static final class Builder {
@VisibleForTesting @Nullable Duration connectTimeout;
- // These two middleware are for internal use within build() and should not be exposed by builder
+ @VisibleForTesting @Nullable OAuthConfiguration oauthConfig;
+
+ // These two middleware are for internal use within build() and should not be
+ // exposed by builder
// APIs.
// Note that these middleware may not necessarily be registered.
@VisibleForTesting
@@ -597,6 +691,8 @@ public static final class Builder {
@VisibleForTesting
ClientCookieMiddleware.Factory cookieFactory = new ClientCookieMiddleware.Factory();
+ DriverVersion driverVersion;
+
public Builder() {}
/**
@@ -623,6 +719,7 @@ public Builder() {}
this.clientKeyPath = original.clientKeyPath;
this.allocator = original.allocator;
this.catalog = original.catalog;
+ this.oauthConfig = original.oauthConfig;
if (original.retainCookies) {
this.cookieFactory = original.cookieFactory;
@@ -631,6 +728,8 @@ public Builder() {}
if (original.retainAuth) {
this.authFactory = original.authFactory;
}
+
+ this.driverVersion = original.driverVersion;
}
/**
@@ -879,6 +978,28 @@ public Builder withConnectTimeout(Duration connectTimeout) {
return this;
}
+ /**
+ * Sets the driver version for this handler.
+ *
+ * @param driverVersion the driver version to set
+ * @return this builder instance
+ */
+ public Builder withDriverVersion(DriverVersion driverVersion) {
+ this.driverVersion = driverVersion;
+ return this;
+ }
+
+ /**
+ * Sets the OAuth configuration for this handler.
+ *
+ * @param oauthConfig the OAuth configuration
+ * @return this builder instance
+ */
+ public Builder withOAuthConfiguration(final OAuthConfiguration oauthConfig) {
+ this.oauthConfig = oauthConfig;
+ return this;
+ }
+
public String getCacheKey() {
return getLocation().toString();
}
@@ -898,7 +1019,8 @@ public Location getLocation() {
* @throws SQLException on error.
*/
public ArrowFlightSqlClientHandler build() throws SQLException {
- // Copy middleware so that the build method doesn't change the state of the builder fields
+ // Copy middleware so that the build method doesn't change the state of the
+ // builder fields
// itself.
Set buildTimeMiddlewareFactories =
new HashSet<>(this.middlewareFactories);
@@ -906,7 +1028,8 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
boolean isUsingUserPasswordAuth = username != null && token == null;
try {
- // Token should take priority since some apps pass in a username/password even when a token
+ // Token should take priority since some apps pass in a username/password even
+ // when a token
// is provided
if (isUsingUserPasswordAuth) {
buildTimeMiddlewareFactories.add(authFactory);
@@ -914,6 +1037,11 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
final NettyClientBuilder clientBuilder = new NettyClientBuilder();
clientBuilder.allocator(allocator);
+ String userAgent = String.format(USER_AGENT_TEMPLATE, DEFAULT_VERSION);
+ if (driverVersion != null && driverVersion.versionString != null) {
+ userAgent = String.format(USER_AGENT_TEMPLATE, driverVersion.versionString);
+ }
+
buildTimeMiddlewareFactories.add(new ClientCookieMiddleware.Factory());
buildTimeMiddlewareFactories.forEach(clientBuilder::intercept);
if (useEncryption) {
@@ -948,6 +1076,9 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
}
NettyChannelBuilder channelBuilder = clientBuilder.build();
+
+ channelBuilder.userAgent(userAgent);
+
if (connectTimeout != null) {
channelBuilder.withOption(
ChannelOption.CONNECT_TIMEOUT_MILLIS, (int) connectTimeout.toMillis());
@@ -956,9 +1087,15 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
FlightGrpcUtils.createFlightClient(
allocator, channelBuilder.build(), clientBuilder.middleware());
final ArrayList credentialOptions = new ArrayList<>();
- if (isUsingUserPasswordAuth) {
- // If the authFactory has already been used for a handshake, use the existing token.
- // This can occur if the authFactory is being re-used for a new connection spawned for
+ // Authentication priority: OAuth > token > username/password
+ if (oauthConfig != null) {
+ OAuthTokenProvider tokenProvider = oauthConfig.createTokenProvider();
+ credentialOptions.add(new CredentialCallOption(new OAuthCredentialWriter(tokenProvider)));
+ } else if (isUsingUserPasswordAuth) {
+ // If the authFactory has already been used for a handshake, use the existing
+ // token.
+ // This can occur if the authFactory is being re-used for a new connection
+ // spawned for
// getStream().
if (authFactory.getCredentialCallOption() != null) {
credentialOptions.add(authFactory.getCredentialCallOption());
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java
new file mode 100644
index 0000000000..9c377a5850
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java
@@ -0,0 +1,108 @@
+/*
+ * 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.ParseException;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.TokenErrorResponse;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.TokenResponse;
+import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
+import com.nimbusds.oauth2.sdk.token.AccessToken;
+import java.io.IOException;
+import java.net.URI;
+import java.sql.SQLException;
+import java.time.Instant;
+import org.apache.arrow.util.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Abstract base class for OAuth token providers that handles token caching, refresh logic, and
+ * common request/response handling.
+ */
+public abstract class AbstractOAuthTokenProvider implements OAuthTokenProvider {
+ protected static final int EXPIRATION_BUFFER_SECONDS = 30;
+ protected static final int DEFAULT_EXPIRATION_SECONDS = 3600;
+
+ private final Object tokenLock = new Object();
+ private volatile @Nullable TokenInfo cachedToken;
+
+ @VisibleForTesting URI tokenUri;
+
+ @VisibleForTesting @Nullable ClientAuthentication clientAuth;
+
+ @VisibleForTesting @Nullable Scope scope;
+
+ @Override
+ public String getValidToken() throws SQLException {
+ TokenInfo token = cachedToken;
+ if (token != null && !token.isExpired(EXPIRATION_BUFFER_SECONDS)) {
+ return token.getAccessToken();
+ }
+
+ synchronized (tokenLock) {
+ token = cachedToken;
+ if (token != null && !token.isExpired(EXPIRATION_BUFFER_SECONDS)) {
+ return token.getAccessToken();
+ }
+ cachedToken = fetchNewToken();
+ return cachedToken.getAccessToken();
+ }
+ }
+
+ /**
+ * Fetches a new token from the authorization server. This method handles the common
+ * request/response logic while delegating flow-specific request building to subclasses.
+ *
+ * @return the new token information
+ * @throws SQLException if token cannot be obtained
+ */
+ protected TokenInfo fetchNewToken() throws SQLException {
+ try {
+ TokenRequest request = buildTokenRequest();
+ TokenResponse response = TokenResponse.parse(request.toHTTPRequest().send());
+
+ if (!response.indicatesSuccess()) {
+ TokenErrorResponse errorResponse = response.toErrorResponse();
+ String errorMsg =
+ String.format(
+ "OAuth request failed: %s - %s",
+ errorResponse.getErrorObject().getCode(),
+ errorResponse.getErrorObject().getDescription());
+ throw new SQLException(errorMsg);
+ }
+
+ AccessToken accessToken = response.toSuccessResponse().getTokens().getAccessToken();
+ long expiresIn =
+ accessToken.getLifetime() > 0 ? accessToken.getLifetime() : DEFAULT_EXPIRATION_SECONDS;
+ Instant expiresAt = Instant.now().plusSeconds(expiresIn);
+
+ return new TokenInfo(accessToken.getValue(), expiresAt);
+ } catch (ParseException e) {
+ throw new SQLException("Failed to parse OAuth token response", e);
+ } catch (IOException e) {
+ throw new SQLException("Failed to send OAuth token request", e);
+ }
+ }
+
+ /**
+ * Builds the flow-specific token request.
+ *
+ * @return the token request to send to the authorization server
+ */
+ protected abstract TokenRequest buildTokenRequest();
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java
new file mode 100644
index 0000000000..7e6289819c
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java
@@ -0,0 +1,58 @@
+/*
+ * 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.ClientCredentialsGrant;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic;
+import com.nimbusds.oauth2.sdk.auth.Secret;
+import com.nimbusds.oauth2.sdk.id.ClientID;
+import java.net.URI;
+import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * OAuth 2.0 Client Credentials flow token provider (RFC 6749 Section 4.4).
+ *
+ * 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 {
+ private final OAuthTokenProvider tokenProvider;
+
+ public OAuthCredentialWriter(OAuthTokenProvider tokenProvider) {
+ this.tokenProvider = Objects.requireNonNull(tokenProvider, "tokenProvider cannot be null");
+ }
+
+ @Override
+ public void accept(CallHeaders headers) {
+ try {
+ String token = tokenProvider.getValidToken();
+ headers.insert(Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + token);
+ } catch (SQLException e) {
+ throw new OAuthTokenException("Failed to obtain OAuth token", e);
+ }
+ }
+}
diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java
similarity index 64%
rename from vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java
rename to flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java
index 1b1bf4e6e4..aceadb327b 100644
--- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java
@@ -14,18 +14,18 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.arrow.vector.complex.impl;
+package org.apache.arrow.driver.jdbc.client.oauth;
-import org.apache.arrow.vector.ExtensionTypeVector;
-import org.apache.arrow.vector.UuidVector;
-
-public class UuidWriterFactory implements ExtensionTypeWriterFactory {
+/**
+ * Runtime exception thrown when OAuth token operations fail. Used to wrap checked exceptions in
+ * contexts that don't allow them.
+ */
+public class OAuthTokenException extends RuntimeException {
+ public OAuthTokenException(String message) {
+ super(message);
+ }
- @Override
- public AbstractFieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) {
- if (extensionTypeVector instanceof UuidVector) {
- return new UuidWriterImpl((UuidVector) extensionTypeVector);
- }
- return null;
+ public OAuthTokenException(String message, Throwable cause) {
+ super(message, cause);
}
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java
new file mode 100644
index 0000000000..241611e432
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java
@@ -0,0 +1,33 @@
+/*
+ * 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;
+
+/**
+ * Interface for OAuth token providers that handle token acquisition and refresh. Implementations
+ * should cache tokens and automatically refresh them before expiration.
+ */
+public interface OAuthTokenProvider {
+ /**
+ * Gets a valid OAuth access token, refreshing if necessary.
+ *
+ * @return a valid access token string
+ * @throws SQLException if token cannot be obtained
+ */
+ String getValidToken() throws SQLException;
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java
new file mode 100644
index 0000000000..bbf7072d39
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java
@@ -0,0 +1,419 @@
+/*
+ * 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.ParseException;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
+import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic;
+import com.nimbusds.oauth2.sdk.auth.Secret;
+import com.nimbusds.oauth2.sdk.id.Audience;
+import com.nimbusds.oauth2.sdk.id.ClientID;
+import com.nimbusds.oauth2.sdk.token.TokenTypeURI;
+import com.nimbusds.oauth2.sdk.token.TypelessAccessToken;
+import com.nimbusds.oauth2.sdk.tokenexchange.TokenExchangeGrant;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Unified factory for creating OAuth token providers.
+ *
+ * This class provides a single entry point for creating all OAuth token providers with a
+ * consistent builder API. It supports:
+ *
+ *
+ * Client Credentials flow (RFC 6749 Section 4.4)
+ * Token Exchange flow (RFC 8693)
+ *
+ *
+ * Example usage:
+ *
+ *
{@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 List resources;
+ private @Nullable ClientAuthentication clientAuth;
+
+ TokenExchangeBuilder() {}
+
+ /**
+ * Sets the OAuth token endpoint URI (required).
+ *
+ * @param tokenUri the token endpoint URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder 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 TokenExchangeBuilder 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 subject token to exchange (required).
+ *
+ * @param subjectToken the subject token value
+ * @return this builder
+ */
+ public TokenExchangeBuilder subjectToken(String subjectToken) {
+ this.subjectToken = Objects.requireNonNull(subjectToken, "subjectToken cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the type of the subject token (required).
+ *
+ * @param subjectTokenType the subject token type URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder subjectTokenType(String subjectTokenType) {
+ this.subjectTokenType =
+ Objects.requireNonNull(subjectTokenType, "subjectTokenType cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets the optional actor token for delegation scenarios.
+ *
+ * @param actorToken the actor token value
+ * @return this builder
+ */
+ public TokenExchangeBuilder actorToken(@Nullable String actorToken) {
+ this.actorToken = actorToken;
+ return this;
+ }
+
+ /**
+ * Sets the type of the actor token.
+ *
+ * @param actorTokenType the actor token type URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder actorTokenType(@Nullable String actorTokenType) {
+ this.actorTokenType = actorTokenType;
+ return this;
+ }
+
+ /**
+ * Sets the target audience for the exchanged token.
+ *
+ * @param audience the target audience
+ * @return this builder
+ */
+ public TokenExchangeBuilder audience(@Nullable String audience) {
+ this.audience = audience;
+ return this;
+ }
+
+ /**
+ * Sets the requested token type for the exchanged token.
+ *
+ * @param requestedTokenType the requested token type URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder requestedTokenType(@Nullable String requestedTokenType) {
+ this.requestedTokenType = requestedTokenType;
+ return this;
+ }
+
+ /**
+ * Sets the OAuth scopes for the token request.
+ *
+ * @param scope the OAuth scope object
+ * @return this builder
+ */
+ public TokenExchangeBuilder scope(@Nullable Scope scope) {
+ this.scope = scope;
+ return this;
+ }
+
+ /**
+ * Sets the OAuth scopes from a space-separated string.
+ *
+ * @param scope the space-separated scope string
+ * @return this builder
+ */
+ public TokenExchangeBuilder scope(@Nullable String scope) {
+ this.scope = (scope != null && !scope.isEmpty()) ? Scope.parse(scope) : null;
+ return this;
+ }
+
+ /**
+ * Sets the target resource URIs (RFC 8707).
+ *
+ * @param resources the list of resource URIs
+ * @return this builder
+ */
+ public TokenExchangeBuilder resources(@Nullable List resources) {
+ this.resources = resources;
+ return this;
+ }
+
+ /**
+ * Sets a single target resource URI (RFC 8707).
+ *
+ * @param resource the resource URI
+ * @return this builder
+ */
+ public TokenExchangeBuilder resource(@Nullable URI resource) {
+ this.resources = resource != null ? Collections.singletonList(resource) : null;
+ return this;
+ }
+
+ /**
+ * Sets a single target resource URI from a string (RFC 8707).
+ *
+ * @param resource the resource URI string
+ * @return this builder
+ */
+ public TokenExchangeBuilder resource(@Nullable String resource) {
+ if (resource != null && !resource.isEmpty()) {
+ this.resources = Collections.singletonList(URI.create(resource));
+ } else {
+ this.resources = null;
+ }
+ return this;
+ }
+
+ /**
+ * Sets the client authentication.
+ *
+ * @param clientAuth the client authentication object
+ * @return this builder
+ */
+ public TokenExchangeBuilder clientAuthentication(@Nullable ClientAuthentication clientAuth) {
+ this.clientAuth = clientAuth;
+ return this;
+ }
+
+ /**
+ * Sets client authentication using client ID and secret.
+ *
+ * @param clientId the client ID
+ * @param clientSecret the client secret
+ * @return this builder
+ */
+ public TokenExchangeBuilder clientCredentials(String clientId, String clientSecret) {
+ Objects.requireNonNull(clientId, "clientId cannot be null");
+ Objects.requireNonNull(clientSecret, "clientSecret cannot be null");
+ this.clientAuth = new ClientSecretBasic(new ClientID(clientId), new Secret(clientSecret));
+ return this;
+ }
+
+ /**
+ * Builds a new TokenExchangeTokenProvider instance.
+ *
+ * @return the configured TokenExchangeTokenProvider
+ * @throws IllegalStateException if required parameters are missing
+ */
+ public TokenExchangeTokenProvider build() {
+ if (tokenUri == null) {
+ throw new IllegalStateException("tokenUri is required");
+ }
+ if (subjectToken == null) {
+ throw new IllegalStateException("subjectToken is required");
+ }
+ if (subjectTokenType == null) {
+ throw new IllegalStateException("subjectTokenType is required");
+ }
+
+ TokenExchangeGrant grant = createGrant();
+ return new TokenExchangeTokenProvider(tokenUri, grant, clientAuth, scope, resources);
+ }
+
+ private TokenExchangeGrant createGrant() {
+ try {
+ TypelessAccessToken subjectAccessToken = new TypelessAccessToken(subjectToken);
+ TokenTypeURI subjectTypeUri = TokenTypeURI.parse(subjectTokenType);
+
+ TypelessAccessToken actorAccessToken =
+ actorToken != null ? new TypelessAccessToken(actorToken) : null;
+ TokenTypeURI actorTypeUri =
+ actorTokenType != null ? TokenTypeURI.parse(actorTokenType) : null;
+ TokenTypeURI requestedTypeUri =
+ requestedTokenType != null ? TokenTypeURI.parse(requestedTokenType) : null;
+ List audienceList =
+ audience != null ? Collections.singletonList(new Audience(audience)) : null;
+
+ return new TokenExchangeGrant(
+ subjectAccessToken,
+ subjectTypeUri,
+ actorAccessToken,
+ actorTypeUri,
+ requestedTypeUri,
+ audienceList);
+ } catch (ParseException e) {
+ throw new IllegalStateException("Failed to create TokenExchangeGrant", e);
+ }
+ }
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java
new file mode 100644
index 0000000000..af433a2712
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java
@@ -0,0 +1,81 @@
+/*
+ * 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.Scope;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
+import com.nimbusds.oauth2.sdk.tokenexchange.TokenExchangeGrant;
+import java.net.URI;
+import java.util.List;
+import java.util.Objects;
+import org.apache.arrow.util.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * OAuth 2.0 Token Exchange flow token provider (RFC 8693).
+ *
+ * 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 resources;
+
+ /**
+ * Creates a new TokenExchangeTokenProvider with full configuration.
+ *
+ * @param tokenUri the OAuth token endpoint URI
+ * @param grant the token exchange grant containing subject/actor token information
+ * @param clientAuth optional client authentication
+ * @param scope optional OAuth scopes
+ * @param resource optional target resource URI (RFC 8707)
+ */
+ TokenExchangeTokenProvider(
+ URI tokenUri,
+ TokenExchangeGrant grant,
+ @Nullable ClientAuthentication clientAuth,
+ @Nullable Scope scope,
+ @Nullable List resource) {
+ this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null");
+ this.grant = Objects.requireNonNull(grant, "grant cannot be null");
+ this.scope = scope;
+ this.resources = resource;
+ this.clientAuth = clientAuth;
+ }
+
+ @Override
+ protected TokenRequest buildTokenRequest() {
+ TokenRequest.Builder builder;
+ if (clientAuth != null) {
+ builder = new TokenRequest.Builder(tokenUri, clientAuth, grant);
+ } else {
+ builder = new TokenRequest.Builder(tokenUri, grant);
+ }
+
+ if (scope != null) {
+ builder.scope(scope);
+ }
+ if (resources != null) {
+ builder.resources(resources.toArray(new URI[0]));
+ }
+
+ return builder.build();
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java
new file mode 100644
index 0000000000..f47cc8b053
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java
@@ -0,0 +1,45 @@
+/*
+ * 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.time.Instant;
+import java.util.Objects;
+
+/** Holds OAuth token information including the access token and expiration time. */
+public class TokenInfo {
+ private final String accessToken;
+ private final Instant expiresAt;
+
+ public TokenInfo(String accessToken, Instant expiresAt) {
+ this.accessToken = Objects.requireNonNull(accessToken, "accessToken cannot be null");
+ this.expiresAt = Objects.requireNonNull(expiresAt, "expiresAt cannot be null");
+ }
+
+ public String getAccessToken() {
+ return accessToken;
+ }
+
+ /**
+ * Checks if the token is expired or will expire within the buffer period.
+ *
+ * @param bufferSeconds seconds before actual expiration to consider token expired
+ * @return true if token should be refreshed
+ */
+ public boolean isExpired(int bufferSeconds) {
+ return Instant.now().plusSeconds(bufferSeconds).isAfter(expiresAt);
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java
index a035bbba49..d692f39372 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java
@@ -17,6 +17,7 @@
package org.apache.arrow.driver.jdbc.converter.impl;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ViewVarBinaryVector;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.calcite.avatica.AvaticaParameter;
@@ -29,7 +30,12 @@ public BinaryViewAvaticaParameterConverter(ArrowType.BinaryView type) {}
@Override
public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
- throw new UnsupportedOperationException("Not implemented");
+ byte[] value = (byte[]) typedValue.toJdbc(null);
+ if (vector instanceof ViewVarBinaryVector) {
+ ((ViewVarBinaryVector) vector).setSafe(index, value);
+ return true;
+ }
+ return false;
}
@Override
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java
index 076fefc42a..c9d9f2926b 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java
@@ -17,8 +17,10 @@
package org.apache.arrow.driver.jdbc.converter.impl;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ViewVarCharVector;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.util.Text;
import org.apache.calcite.avatica.AvaticaParameter;
import org.apache.calcite.avatica.remote.TypedValue;
@@ -29,7 +31,12 @@ public Utf8ViewAvaticaParameterConverter(ArrowType.Utf8View type) {}
@Override
public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) {
- throw new UnsupportedOperationException("Utf8View not supported");
+ String value = (String) typedValue.toLocal();
+ if (vector instanceof ViewVarCharVector) {
+ ((ViewVarCharVector) vector).setSafe(index, new Text(value));
+ return true;
+ }
+ return false;
}
@Override
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java
new file mode 100644
index 0000000000..b2157890cf
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java
@@ -0,0 +1,103 @@
+/*
+ * 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.converter.impl;
+
+import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeIdFromArrowType;
+import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeNameFromArrowType;
+
+import java.nio.ByteBuffer;
+import java.util.UUID;
+import org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.util.UuidUtility;
+import org.apache.calcite.avatica.AvaticaParameter;
+import org.apache.calcite.avatica.remote.TypedValue;
+import org.apache.calcite.avatica.util.ByteString;
+
+/**
+ * AvaticaParameterConverter for UUID Arrow extension type.
+ *
+ * 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 getHeaderAttributes() {
return headers;
}
+ /**
+ * Returns OAuth configuration if oauth.flow is specified, null otherwise.
+ *
+ * @return the OAuth configuration or null
+ * @throws SQLException if the OAuth configuration is invalid
+ */
+ public @Nullable OAuthConfiguration getOauthConfiguration() throws SQLException {
+ String flow = ArrowFlightConnectionProperty.OAUTH_FLOW.getString(properties);
+ if (flow == null) {
+ return null;
+ }
+
+ return new OAuthConfiguration.Builder()
+ .flow(flow)
+ .clientId(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.getString(properties))
+ .clientSecret(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.getString(properties))
+ .tokenUri(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.getString(properties))
+ .scope(ArrowFlightConnectionProperty.OAUTH_SCOPE.getString(properties))
+ .resource(ArrowFlightConnectionProperty.OAUTH_RESOURCE.getString(properties))
+ .subjectToken(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.getString(properties))
+ .subjectTokenType(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.getString(properties))
+ .actorToken(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN.getString(properties))
+ .actorTokenType(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE.getString(properties))
+ .audience(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_AUDIENCE.getString(properties))
+ .requestedTokenType(
+ ArrowFlightConnectionProperty.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.getString(properties))
+ .build();
+ }
+
/** Custom {@link ConnectionProperty} for the {@link ArrowFlightConnectionConfigImpl}. */
public enum ArrowFlightConnectionProperty implements ConnectionProperty {
HOST("host", null, Type.STRING, true),
@@ -232,6 +267,23 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty {
CATALOG("catalog", null, Type.STRING, false),
CONNECT_TIMEOUT_MILLIS("connectTimeoutMs", 10000, Type.NUMBER, false),
USE_CLIENT_CACHE("useClientCache", true, Type.BOOLEAN, false),
+
+ // OAuth configuration properties
+ OAUTH_FLOW("oauth.flow", null, Type.STRING, false),
+ OAUTH_CLIENT_ID("oauth.clientId", null, Type.STRING, false),
+ OAUTH_CLIENT_SECRET("oauth.clientSecret", null, Type.STRING, false),
+ OAUTH_TOKEN_URI("oauth.tokenUri", null, Type.STRING, false),
+ OAUTH_SCOPE("oauth.scope", null, Type.STRING, false),
+ OAUTH_RESOURCE("oauth.resource", null, Type.STRING, false),
+
+ // Token exchange specific properties
+ OAUTH_EXCHANGE_SUBJECT_TOKEN("oauth.exchange.subjectToken", null, Type.STRING, false),
+ OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE("oauth.exchange.subjectTokenType", null, Type.STRING, false),
+ OAUTH_EXCHANGE_ACTOR_TOKEN("oauth.exchange.actorToken", null, Type.STRING, false),
+ OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE("oauth.exchange.actorTokenType", null, Type.STRING, false),
+ OAUTH_EXCHANGE_AUDIENCE("oauth.exchange.aud", null, Type.STRING, false),
+ OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE(
+ "oauth.exchange.requestedTokenType", null, Type.STRING, false),
;
private final String camelName;
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java
index 4c2a9b865f..8f40d6698e 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java
@@ -19,6 +19,7 @@
import java.util.List;
import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement;
import org.apache.arrow.driver.jdbc.converter.impl.BinaryAvaticaParameterConverter;
+import org.apache.arrow.driver.jdbc.converter.impl.BinaryViewAvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.BoolAvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.DateAvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.DecimalAvaticaParameterConverter;
@@ -39,11 +40,17 @@
import org.apache.arrow.driver.jdbc.converter.impl.TimestampAvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter;
+import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter;
+import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor;
+import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
import org.apache.calcite.avatica.remote.TypedValue;
+import org.checkerframework.checker.nullness.qual.Nullable;
/**
* Convert Avatica PreparedStatement parameters from a list of TypedValue to Arrow and bind them to
@@ -108,9 +115,9 @@ public void bind(List typedValues, int index) {
* @param typedValue TypedValue to bind to the vector.
* @param index Vector index to bind the value at.
*/
- private void bind(FieldVector vector, TypedValue typedValue, int index) {
+ private void bind(FieldVector vector, @Nullable TypedValue typedValue, int index) {
try {
- if (typedValue.value == null) {
+ if (typedValue == null || typedValue.value == null) {
if (vector.getField().isNullable()) {
vector.setNull(index);
} else {
@@ -127,7 +134,7 @@ private void bind(FieldVector vector, TypedValue typedValue, int index) {
throw new UnsupportedOperationException(
String.format(
"Binding value of type %s is not yet supported for expected Arrow type %s",
- typedValue.type, vector.getField().getType()));
+ typedValue == null ? "null" : typedValue.type, vector.getField().getType()));
}
}
@@ -207,7 +214,7 @@ public Boolean visit(ArrowType.Utf8 type) {
@Override
public Boolean visit(ArrowType.Utf8View type) {
- throw new UnsupportedOperationException("Utf8View is unsupported");
+ return new Utf8ViewAvaticaParameterConverter(type).bindParameter(vector, typedValue, index);
}
@Override
@@ -222,7 +229,7 @@ public Boolean visit(ArrowType.Binary type) {
@Override
public Boolean visit(ArrowType.BinaryView type) {
- throw new UnsupportedOperationException("BinaryView is unsupported");
+ return new BinaryViewAvaticaParameterConverter(type).bindParameter(vector, typedValue, index);
}
@Override
@@ -287,5 +294,15 @@ public Boolean visit(ArrowType.RunEndEncoded type) {
throw new UnsupportedOperationException(
"No Avatica parameter binder implemented for type " + type);
}
+
+ @Override
+ public Boolean visit(ExtensionType type) {
+ if (type instanceof UuidType) {
+ return new UuidAvaticaParameterConverter().bindParameter(vector, typedValue, index);
+ }
+
+ // fallback to default implementation
+ return ArrowTypeVisitor.super.visit(type);
+ }
}
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java
index 5dd4c69c73..dd51ee5361 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java
@@ -43,8 +43,12 @@
import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter;
import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter;
+import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter;
import org.apache.arrow.flight.sql.FlightSqlColumnMetadata;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor;
+import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.calcite.avatica.AvaticaParameter;
import org.apache.calcite.avatica.ColumnMetaData;
@@ -294,5 +298,15 @@ public AvaticaParameter visit(ArrowType.RunEndEncoded type) {
throw new UnsupportedOperationException(
"No Avatica parameter binder implemented for type " + type);
}
+
+ @Override
+ public AvaticaParameter visit(ExtensionType type) {
+ if (type instanceof UuidType) {
+ return new UuidAvaticaParameterConverter().createParameter(field);
+ }
+
+ // fallback to default implementation
+ return ArrowTypeVisitor.super.visit(type);
+ }
}
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java
index 1b76ca0c95..7982d5bc73 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java
@@ -20,11 +20,13 @@
import java.sql.Types;
import java.util.HashMap;
import java.util.Map;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.pojo.ArrowType;
/** SQL Types utility functions. */
public class SqlTypes {
+
private static final Map typeIdToName = new HashMap<>();
static {
@@ -107,12 +109,17 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) {
}
break;
case Binary:
+ case BinaryView:
return Types.VARBINARY;
case FixedSizeBinary:
+ if (arrowType instanceof UuidType) {
+ return Types.OTHER;
+ }
return Types.BINARY;
case LargeBinary:
return Types.LONGVARBINARY;
case Utf8:
+ case Utf8View:
return Types.VARCHAR;
case LargeUtf8:
return Types.LONGVARCHAR;
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java
index 81579cc387..3ab1460b27 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java
@@ -1543,11 +1543,83 @@ public void testEmptySqlInfo() throws Exception {
try (final Connection testConnection =
FLIGHT_SERVER_EMPTY_SQLINFO_TEST_RULE.getConnection(false)) {
final DatabaseMetaData metaData = testConnection.getMetaData();
+
assertThat(metaData.getSQLKeywords(), is(""));
assertThat(metaData.getNumericFunctions(), is(""));
assertThat(metaData.getStringFunctions(), is(""));
assertThat(metaData.getSystemFunctions(), is(""));
assertThat(metaData.getTimeDateFunctions(), is(""));
+
+ assertThat(metaData.getMaxBinaryLiteralLength(), is(0));
+ assertThat(metaData.getMaxCharLiteralLength(), is(0));
+ assertThat(metaData.getMaxColumnNameLength(), is(0));
+ assertThat(metaData.getMaxColumnsInGroupBy(), is(0));
+ assertThat(metaData.getMaxColumnsInIndex(), is(0));
+ assertThat(metaData.getMaxColumnsInOrderBy(), is(0));
+ assertThat(metaData.getMaxColumnsInSelect(), is(0));
+ assertThat(metaData.getMaxColumnsInTable(), is(0));
+ assertThat(metaData.getMaxConnections(), is(0));
+ assertThat(metaData.getMaxCursorNameLength(), is(0));
+ assertThat(metaData.getMaxIndexLength(), is(0));
+ assertThat(metaData.getMaxSchemaNameLength(), is(0));
+ assertThat(metaData.getMaxProcedureNameLength(), is(0));
+ assertThat(metaData.getMaxCatalogNameLength(), is(0));
+ assertThat(metaData.getMaxRowSize(), is(0));
+ assertThat(metaData.getMaxStatementLength(), is(0));
+ assertThat(metaData.getMaxStatements(), is(0));
+ assertThat(metaData.getMaxTableNameLength(), is(0));
+ assertThat(metaData.getMaxTablesInSelect(), is(0));
+ assertThat(metaData.getMaxUserNameLength(), is(0));
+
+ assertThat(metaData.supportsColumnAliasing(), is(false));
+ assertThat(metaData.nullPlusNonNullIsNull(), is(false));
+ assertThat(metaData.supportsTableCorrelationNames(), is(false));
+ assertThat(metaData.supportsDifferentTableCorrelationNames(), is(false));
+ assertThat(metaData.supportsExpressionsInOrderBy(), is(false));
+ assertThat(metaData.supportsOrderByUnrelated(), is(false));
+ assertThat(metaData.supportsLikeEscapeClause(), is(false));
+ assertThat(metaData.supportsNonNullableColumns(), is(false));
+ assertThat(metaData.supportsIntegrityEnhancementFacility(), is(false));
+ assertThat(metaData.isCatalogAtStart(), is(false));
+ assertThat(metaData.supportsSelectForUpdate(), is(false));
+ assertThat(metaData.supportsStoredProcedures(), is(false));
+ assertThat(metaData.supportsCorrelatedSubqueries(), is(false));
+ assertThat(metaData.doesMaxRowSizeIncludeBlobs(), is(false));
+ assertThat(metaData.supportsTransactions(), is(false));
+ assertThat(metaData.dataDefinitionCausesTransactionCommit(), is(false));
+ assertThat(metaData.dataDefinitionIgnoredInTransactions(), is(false));
+ assertThat(metaData.supportsBatchUpdates(), is(false));
+ assertThat(metaData.supportsSavepoints(), is(false));
+ assertThat(metaData.supportsNamedParameters(), is(false));
+ assertThat(metaData.locatorsUpdateCopy(), is(false));
+ assertThat(metaData.supportsStoredFunctionsUsingCallSyntax(), is(false));
+ assertThat(metaData.supportsGroupBy(), is(false));
+ assertThat(metaData.supportsGroupByUnrelated(), is(false));
+ assertThat(metaData.supportsMinimumSQLGrammar(), is(false));
+ assertThat(metaData.supportsCoreSQLGrammar(), is(false));
+ assertThat(metaData.supportsExtendedSQLGrammar(), is(false));
+ assertThat(metaData.supportsANSI92EntryLevelSQL(), is(false));
+ assertThat(metaData.supportsANSI92IntermediateSQL(), is(false));
+ assertThat(metaData.supportsANSI92FullSQL(), is(false));
+ assertThat(metaData.supportsOuterJoins(), is(false));
+ assertThat(metaData.supportsFullOuterJoins(), is(false));
+ assertThat(metaData.supportsLimitedOuterJoins(), is(false));
+ assertThat(metaData.supportsSchemasInProcedureCalls(), is(false));
+ assertThat(metaData.supportsSchemasInIndexDefinitions(), is(false));
+ assertThat(metaData.supportsSchemasInPrivilegeDefinitions(), is(false));
+ assertThat(metaData.supportsCatalogsInIndexDefinitions(), is(false));
+ assertThat(metaData.supportsCatalogsInPrivilegeDefinitions(), is(false));
+ assertThat(metaData.supportsPositionedDelete(), is(false));
+ assertThat(metaData.supportsPositionedUpdate(), is(false));
+ assertThat(metaData.supportsSubqueriesInComparisons(), is(false));
+ assertThat(metaData.supportsSubqueriesInExists(), is(false));
+ assertThat(metaData.supportsSubqueriesInIns(), is(false));
+ assertThat(metaData.supportsSubqueriesInQuantifieds(), is(false));
+ assertThat(metaData.supportsUnion(), is(false));
+ assertThat(metaData.supportsUnionAll(), is(false));
+ assertThat(metaData.supportsConvert(), is(false));
+
+ assertThat(metaData.getDefaultTransactionIsolation(), is(Connection.TRANSACTION_NONE));
}
}
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java
index 06d101724c..cb6abacb2f 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java
@@ -129,7 +129,7 @@ public void testShouldGetResultSetReturnValidResultSet() throws SQLException {
try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet()) {
int count = 0;
while (resultSet.next()) {
- assertEquals((Object) resultSet.getInt(1), dataVector.getObject(count));
+ assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count));
count++;
}
}
@@ -142,7 +142,7 @@ public void testShouldGetResultSetReturnValidResultSetWithOffsets() throws SQLEx
try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet(3, 5)) {
int count = 0;
while (resultSet.next()) {
- assertEquals((Object) resultSet.getInt(1), dataVector.getObject(count + 3));
+ assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count + 3));
count++;
}
assertEquals(5, count);
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java
index 1977b61392..7127c7fc32 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java
@@ -39,11 +39,11 @@ public void testCookies() throws SQLException {
Statement statement = connection.createStatement()) {
// Expect client didn't receive cookies before any operation
- assertNull(FLIGHT_SERVER_TEST_EXTENSION.getMiddlewareCookieFactory().getCookie());
+ assertNull(FLIGHT_SERVER_TEST_EXTENSION.getInterceptorFactory().getCookie());
// Run another action for check if the cookies was sent by the server.
statement.execute(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD);
- assertEquals("k=v", FLIGHT_SERVER_TEST_EXTENSION.getMiddlewareCookieFactory().getCookie());
+ assertEquals("k=v", FLIGHT_SERVER_TEST_EXTENSION.getInterceptorFactory().getCookie());
}
}
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java
index ae355829d7..88fb9889b6 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java
@@ -201,6 +201,30 @@ public void testConnectWithInsensitiveCasePropertyKeys2() throws Exception {
}
}
+ /**
+ * Tests whether the {@link ArrowFlightJdbcDriver} can establish a successful connection to the
+ * Arrow Flight client when provided with null properties.
+ */
+ @Test
+ public void testConnectWithNullProperties() throws Exception {
+ final Driver driver = new ArrowFlightJdbcDriver();
+ try (Connection connection =
+ driver.connect(
+ "jdbc:arrow-flight://"
+ + dataSource.getConfig().getHost()
+ + ":"
+ + dataSource.getConfig().getPort()
+ + "?"
+ + "useEncryption=false"
+ + "&user="
+ + dataSource.getConfig().getUser()
+ + "&password="
+ + dataSource.getConfig().getPassword(),
+ null)) {
+ assertTrue(connection.isValid(300));
+ }
+ }
+
/**
* Tests whether an exception is thrown upon attempting to connect to a malformed URI.
*
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
index 774ad0081e..0369c3a162 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
@@ -20,6 +20,8 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertAll;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
@@ -83,6 +85,19 @@ public void testSimpleQueryNoParameterBinding() throws SQLException {
}
}
+ @Test
+ public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException {
+ final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD;
+ try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
+ boolean isResultSet = preparedStatement.execute();
+ assertTrue(isResultSet);
+ final ResultSet resultSet = preparedStatement.getResultSet();
+ CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet);
+ assertFalse(preparedStatement.getMoreResults());
+ assertEquals(-1, preparedStatement.getUpdateCount());
+ }
+ }
+
@Test
public void testQueryWithParameterBinding() throws SQLException {
final String query = "Fake query with parameters";
@@ -174,6 +189,20 @@ public void testUpdateQuery() throws SQLException {
}
}
+ @Test
+ public void testUpdateQueryWithExecute() throws SQLException {
+ String query = "Fake update with execute";
+ PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42);
+ try (final PreparedStatement stmt = connection.prepareStatement(query)) {
+ boolean isResultSet = stmt.execute();
+ assertFalse(isResultSet);
+ int updated = stmt.getUpdateCount();
+ assertEquals(42, updated);
+ assertFalse(stmt.getMoreResults());
+ assertEquals(-1, stmt.getUpdateCount());
+ }
+ }
+
@Test
public void testUpdateQueryWithParameters() throws SQLException {
String query = "Fake update with parameters";
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
index 8e872a1167..55722f60fb 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java
@@ -16,24 +16,42 @@
*/
package org.apache.arrow.driver.jdbc;
+import static java.lang.String.format;
+import static java.util.stream.IntStream.range;
+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 static org.junit.jupiter.api.Assertions.fail;
+import com.google.protobuf.Message;
import java.net.URISyntaxException;
import java.sql.Connection;
import java.sql.Driver;
import java.sql.DriverManager;
+import java.sql.ResultSet;
import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Map;
import java.util.Properties;
+import java.util.function.Consumer;
import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication;
import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
+import org.apache.arrow.flight.FlightMethod;
+import org.apache.arrow.flight.FlightProducer.ServerStreamListener;
+import org.apache.arrow.flight.NoOpSessionOptionValueVisitor;
+import org.apache.arrow.flight.SessionOptionValue;
+import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas;
+import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.util.Text;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -576,4 +594,178 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager
assertTrue(connection.isValid(0));
}
}
+
+ /**
+ * Test that the JDBC driver properly integrates driver version into client handler.
+ *
+ * @throws Exception on error.
+ */
+ @Test
+ public void testJdbcDriverVersionIntegration() throws Exception {
+ final Properties properties = new Properties();
+ properties.put(
+ ArrowFlightConnectionProperty.HOST.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getHost());
+ properties.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false);
+
+ // Create a driver instance and connect
+ ArrowFlightJdbcDriver driverVersion = new ArrowFlightJdbcDriver();
+
+ try (Connection connection =
+ ArrowFlightConnection.createNewConnection(
+ driverVersion,
+ new ArrowFlightJdbcFactory(),
+ "jdbc:arrow-flight-sql://localhost:" + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties,
+ allocator)) {
+
+ assertTrue(connection.isValid(0));
+
+ var actualUserAgent =
+ FLIGHT_SERVER_TEST_EXTENSION
+ .getInterceptorFactory()
+ .getHeader(FlightMethod.HANDSHAKE, "user-agent");
+
+ var expectedUserAgent =
+ "JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString;
+ // Driver appends version to grpc user-agent header. Assert the header starts
+ // with the
+ // expected
+ // value and ignored grpc version.
+ assertTrue(
+ actualUserAgent.startsWith(expectedUserAgent),
+ "Expected: " + expectedUserAgent + " but found: " + actualUserAgent);
+ }
+ }
+
+ @Test
+ public void testSetCatalogShouldUpdateSessionOptions() throws Exception {
+ final Properties properties = new Properties();
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put("useEncryption", false);
+
+ try (Connection connection =
+ DriverManager.getConnection(
+ "jdbc:arrow-flight-sql://"
+ + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+ + ":"
+ + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties)) {
+ final String catalog = "new_catalog";
+ connection.setCatalog(catalog);
+
+ final Map options = PRODUCER.getSessionOptions();
+ assertTrue(options.containsKey("catalog"));
+ String actualCatalog =
+ options
+ .get("catalog")
+ .acceptVisitor(
+ new NoOpSessionOptionValueVisitor() {
+ @Override
+ public String visit(String value) {
+ return value;
+ }
+ });
+ assertEquals(catalog, actualCatalog);
+ }
+ }
+
+ @Test
+ public void testStatementsClosedOnConnectionClose() throws Exception {
+ // create a connection
+ final Properties properties = new Properties();
+ properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost");
+ properties.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put("useEncryption", false);
+
+ Connection connection =
+ DriverManager.getConnection(
+ "jdbc:arrow-flight-sql://"
+ + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+ + ":"
+ + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties);
+
+ // create some statements
+ int numStatements = 3;
+ Statement[] statements = new Statement[numStatements];
+ for (int i = 0; i < numStatements; i++) {
+ statements[i] = connection.createStatement();
+ assertFalse(statements[i].isClosed());
+ }
+
+ // close the connection
+ connection.close();
+
+ // assert the statements are closed
+ for (int i = 0; i < numStatements; i++) {
+ assertTrue(statements[i].isClosed());
+ }
+ }
+
+ @Test
+ public void testResultSetsFromDatabaseMetadataClosedOnConnectionClose() throws Exception {
+ // set up the FlightProducer to respond to metadata queries
+ // getTableTypes() is being used, but any other method would work
+ int rowCount = 3;
+ final Message commandGetTableTypes = CommandGetTableTypes.getDefaultInstance();
+ final Consumer commandGetTableTypesResultProducer =
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator();
+ final VectorSchemaRoot root =
+ VectorSchemaRoot.create(Schemas.GET_TABLE_TYPES_SCHEMA, allocator)) {
+ final VarCharVector tableType = (VarCharVector) root.getVector("table_type");
+ range(0, rowCount)
+ .forEach(i -> tableType.setSafe(i, new Text(format("table_type #%d", i))));
+ root.setRowCount(rowCount);
+ listener.start(root);
+ listener.putNext();
+ } catch (final Throwable throwable) {
+ listener.error(throwable);
+ } finally {
+ listener.completed();
+ }
+ };
+ PRODUCER.addCatalogQuery(commandGetTableTypes, commandGetTableTypesResultProducer);
+
+ // create a connection
+ final Properties properties = new Properties();
+ properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost");
+ properties.put(
+ ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort());
+ properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest);
+ properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest);
+ properties.put("useEncryption", false);
+
+ Connection connection =
+ DriverManager.getConnection(
+ "jdbc:arrow-flight-sql://"
+ + FLIGHT_SERVER_TEST_EXTENSION.getHost()
+ + ":"
+ + FLIGHT_SERVER_TEST_EXTENSION.getPort(),
+ properties);
+
+ // create ResultSets from DatabaseMetadata
+ int numResultSets = 3;
+ ResultSet[] resultSets = new ResultSet[numResultSets];
+ for (int i = 0; i < numResultSets; i++) {
+ resultSets[i] = connection.getMetaData().getTableTypes();
+ assertFalse(resultSets[i].isClosed());
+ }
+
+ // close the connection
+ connection.close();
+
+ // assert the ResultSets are closed
+ for (int i = 0; i < numResultSets; i++) {
+ assertTrue(resultSets[i].isClosed());
+ }
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java
index aa586651f5..f71114e1b5 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java
@@ -25,6 +25,8 @@
import java.sql.SQLException;
import java.util.ArrayDeque;
import java.util.Deque;
+import java.util.HashMap;
+import java.util.Map;
import java.util.Properties;
import org.apache.arrow.driver.jdbc.authentication.Authentication;
import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication;
@@ -33,6 +35,7 @@
import org.apache.arrow.flight.CallHeaders;
import org.apache.arrow.flight.CallInfo;
import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.FlightMethod;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightServerMiddleware;
import org.apache.arrow.flight.Location;
@@ -67,7 +70,8 @@ public class FlightServerTestExtension
private final CertKeyPair certKeyPair;
private final File mTlsCACert;
- private final MiddlewareCookie.Factory middlewareCookieFactory = new MiddlewareCookie.Factory();
+ private final InterceptorMiddleware.Factory interceptorFactory =
+ new InterceptorMiddleware.Factory();
private FlightServerTestExtension(
final Properties properties,
@@ -126,12 +130,18 @@ public Connection getConnection(boolean useEncryption) throws SQLException {
return this.createDataSource().getConnection();
}
+ public Connection getConnection(String timezone) throws SQLException {
+ setUseEncryption(false);
+ properties.put("timezone", timezone);
+ return this.createDataSource().getConnection();
+ }
+
private void setUseEncryption(boolean useEncryption) {
properties.put("useEncryption", useEncryption);
}
- public MiddlewareCookie.Factory getMiddlewareCookieFactory() {
- return middlewareCookieFactory;
+ public InterceptorMiddleware.Factory getInterceptorFactory() {
+ return interceptorFactory;
}
@FunctionalInterface
@@ -143,7 +153,7 @@ private FlightServer initiateServer(Location location) throws IOException {
FlightServer.Builder builder =
FlightServer.builder(allocator, location, producer)
.headerAuthenticator(authentication.authenticate())
- .middleware(FlightServerMiddleware.Key.of("KEY"), middlewareCookieFactory);
+ .middleware(FlightServerMiddleware.Key.of("KEY"), interceptorFactory);
if (certKeyPair != null) {
builder.useTls(certKeyPair.cert, certKeyPair.key);
}
@@ -301,11 +311,11 @@ public FlightServerTestExtension build() {
* A middleware to handle with the cookies in the server. It is used to test if cookies are being
* sent properly.
*/
- static class MiddlewareCookie implements FlightServerMiddleware {
+ static class InterceptorMiddleware implements FlightServerMiddleware {
private final Factory factory;
- public MiddlewareCookie(Factory factory) {
+ public InterceptorMiddleware(Factory factory) {
this.factory = factory;
}
@@ -323,22 +333,33 @@ public void onCallCompleted(CallStatus callStatus) {}
public void onCallErrored(Throwable throwable) {}
/** A factory for the MiddlewareCookie. */
- static class Factory implements FlightServerMiddleware.Factory {
+ static class Factory implements FlightServerMiddleware.Factory {
+ private final Map receivedCallHeaders = new HashMap<>();
private boolean receivedCookieHeader = false;
private String cookie;
@Override
- public MiddlewareCookie onCallStarted(
+ public InterceptorMiddleware onCallStarted(
CallInfo callInfo, CallHeaders callHeaders, RequestContext requestContext) {
cookie = callHeaders.get("Cookie");
receivedCookieHeader = null != cookie;
- return new MiddlewareCookie(this);
+
+ receivedCallHeaders.put(callInfo.method(), callHeaders);
+ return new InterceptorMiddleware(this);
}
public String getCookie() {
return cookie;
}
+
+ public String getHeader(FlightMethod method, String key) {
+ CallHeaders headers = receivedCallHeaders.get(method);
+ if (headers == null) {
+ return null;
+ }
+ return headers.get(key);
+ }
}
}
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java
new file mode 100644
index 0000000000..5e782db031
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java
@@ -0,0 +1,474 @@
+/*
+ * 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;
+
+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.net.URI;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.util.Properties;
+import java.util.concurrent.TimeUnit;
+import mockwebserver3.MockResponse;
+import mockwebserver3.MockWebServer;
+import mockwebserver3.RecordedRequest;
+import mockwebserver3.junit5.StartStop;
+import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication;
+import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
+import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
+import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas;
+import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs;
+import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * Integration tests for OAuth authentication flows in the JDBC driver.
+ *
+ * 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 actualUuids = new ArrayList<>(expectedRowCount);
+ final List actualIds = new ArrayList<>(expectedRowCount);
+
+ // Query
+ int actualRowCount = 0;
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ for (; resultSet.next(); actualRowCount++) {
+ actualIds.add((Integer) resultSet.getObject("id"));
+ actualUuids.add((UUID) resultSet.getObject("uuid_col"));
+ }
+ }
+
+ // Assertions
+ int finalActualRowCount = actualRowCount;
+ assertAll(
+ "UUID ResultSet values are as expected",
+ () -> assertThat(finalActualRowCount, is(equalTo(expectedRowCount))),
+ () -> assertThat(actualIds.toArray(new Integer[0]), is(expectedIds)),
+ () -> assertThat(actualUuids.toArray(new UUID[0]), is(expectedUuids)));
+ }
+
+ @Test
+ public void testGetObjectReturnsUuid() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ resultSet.next();
+ Object result = resultSet.getObject("uuid_col");
+ assertThat(result, instanceOf(UUID.class));
+ assertThat(result, is(CoreMockedSqlProducers.UUID_1));
+ }
+ }
+
+ @Test
+ public void testGetObjectByIndexReturnsUuid() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ resultSet.next();
+ Object result = resultSet.getObject(2);
+ assertThat(result, instanceOf(UUID.class));
+ assertThat(result, is(CoreMockedSqlProducers.UUID_1));
+ }
+ }
+
+ @Test
+ public void testGetStringReturnsHyphenatedFormat() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ resultSet.next();
+ String result = resultSet.getString("uuid_col");
+ assertThat(result, is(CoreMockedSqlProducers.UUID_1.toString()));
+ }
+ }
+
+ @Test
+ public void testGetBytesReturns16ByteArray() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ resultSet.next();
+ byte[] result = resultSet.getBytes("uuid_col");
+ assertThat(result.length, is(16));
+ assertThat(result, is(UuidUtility.getBytesFromUUID(CoreMockedSqlProducers.UUID_1)));
+ }
+ }
+
+ @Test
+ public void testNullUuidHandling() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ // Skip to row 4 which has NULL UUID
+ resultSet.next(); // row 1
+ resultSet.next(); // row 2
+ resultSet.next(); // row 3
+ resultSet.next(); // row 4 (NULL UUID)
+
+ Object objResult = resultSet.getObject("uuid_col");
+ assertThat(objResult, nullValue());
+ assertThat(resultSet.wasNull(), is(true));
+
+ String strResult = resultSet.getString("uuid_col");
+ assertThat(strResult, nullValue());
+ assertThat(resultSet.wasNull(), is(true));
+
+ byte[] bytesResult = resultSet.getBytes("uuid_col");
+ assertThat(bytesResult, nullValue());
+ assertThat(resultSet.wasNull(), is(true));
+ }
+ }
+
+ @Test
+ public void testMultipleUuidRows() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ resultSet.next();
+ assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1));
+
+ resultSet.next();
+ assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_2));
+
+ resultSet.next();
+ assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_3));
+
+ resultSet.next();
+ assertThat(resultSet.getObject("uuid_col"), nullValue());
+ }
+ }
+
+ @Test
+ public void testUuidExtensionTypeInSchema() throws SQLException {
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) {
+ ResultSetMetaData metaData = resultSet.getMetaData();
+
+ assertThat(metaData.getColumnCount(), is(2));
+ assertThat(metaData.getColumnName(1), is("id"));
+ assertThat(metaData.getColumnName(2), is("uuid_col"));
+
+ assertThat(metaData.getColumnType(2), is(java.sql.Types.OTHER));
+ }
+ }
+
+ @Test
+ public void testPreparedStatementWithUuidParameter() throws SQLException {
+ try (PreparedStatement pstmt =
+ connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) {
+ pstmt.setObject(1, CoreMockedSqlProducers.UUID_1);
+ try (ResultSet rs = pstmt.executeQuery()) {
+ rs.next();
+ assertThat(rs.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1));
+ }
+ }
+ }
+
+ @Test
+ public void testPreparedStatementWithUuidStringParameter() throws SQLException {
+ try (PreparedStatement pstmt =
+ connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) {
+ pstmt.setString(1, CoreMockedSqlProducers.UUID_1.toString());
+ try (ResultSet rs = pstmt.executeQuery()) {
+ rs.next();
+ assertThat(rs.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1));
+ }
+ }
+ }
+
+ @Test
+ public void testPreparedStatementUpdateWithUuid() throws SQLException {
+ try (PreparedStatement pstmt =
+ connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_UPDATE_SQL_CMD)) {
+ pstmt.setObject(1, CoreMockedSqlProducers.UUID_3);
+ pstmt.setInt(2, 1);
+ int updated = pstmt.executeUpdate();
+ assertThat(updated, is(1));
+ }
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java
new file mode 100644
index 0000000000..0921ae2d38
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java
@@ -0,0 +1,164 @@
+/*
+ * 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;
+
+import com.google.common.collect.ImmutableList;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.sql.Types;
+import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.util.Calendar;
+import java.util.Collections;
+import java.util.TimeZone;
+import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.TimeStampVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.TimeUnit;
+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.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * Timestamps have a lot of nuances in JDBC. This class is here to test that timestamp behavior is
+ * correct for different types of Timestamp vectors as well as different methods of retrieving the
+ * timestamps in JDBC.
+ */
+public class TimestampResultSetTest {
+ private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer();
+
+ @RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION;
+
+ static {
+ FLIGHT_SERVER_TEST_EXTENSION =
+ FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER);
+ }
+
+ private static final String QUERY_STRING = "SELECT * FROM TIMESTAMPS";
+ private static final Schema QUERY_SCHEMA =
+ new Schema(
+ ImmutableList.of(
+ Field.nullable("no_tz", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)),
+ Field.nullable("utc", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")),
+ Field.nullable("utc+1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT+1")),
+ Field.nullable("utc-1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT-1"))));
+
+ @BeforeAll
+ public static void setup() throws SQLException {
+ Instant firstDay2025 = OffsetDateTime.of(2025, 1, 1, 0, 0, 0, 0, ZoneOffset.UTC).toInstant();
+
+ FLIGHT_SQL_PRODUCER.addSelectQuery(
+ QUERY_STRING,
+ QUERY_SCHEMA,
+ Collections.singletonList(
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) {
+ listener.start(root);
+ root.getFieldVectors()
+ .forEach(v -> ((TimeStampVector) v).setSafe(0, firstDay2025.toEpochMilli()));
+ root.setRowCount(1);
+ listener.putNext();
+ } catch (final Throwable throwable) {
+ listener.error(throwable);
+ } finally {
+ listener.completed();
+ }
+ }));
+ }
+
+ /**
+ * This test doesn't yet test anything other than ensuring all ResultSet methods to retrieve a
+ * timestamp succeed.
+ *
+ * This is a good starting point to add more tests to ensure the values are correct when we
+ * change the "local calendar" either through changing the JVM default or through the connection
+ * property.
+ */
+ @Test
+ public void test() {
+ TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
+ try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection("UTC")) {
+ try (PreparedStatement s = connection.prepareStatement(QUERY_STRING)) {
+ try (ResultSet rs = s.executeQuery()) {
+ int numCols = rs.getMetaData().getColumnCount();
+ try {
+ rs.next();
+ for (int i = 1; i <= numCols; i++) {
+ int type = rs.getMetaData().getColumnType(i);
+ String name = rs.getMetaData().getColumnName(i);
+ System.out.println(name);
+ System.out.print("- getDate:\t\t\t\t\t\t\t");
+ System.out.print(rs.getDate(i));
+ System.out.println();
+ System.out.print("- getTimestamp:\t\t\t\t\t\t");
+ System.out.print(rs.getTimestamp(i));
+ System.out.println();
+ System.out.print("- getString:\t\t\t\t\t\t");
+ System.out.print(rs.getString(i));
+ System.out.println();
+ System.out.print("- getObject:\t\t\t\t\t\t");
+ System.out.print(rs.getObject(i));
+ System.out.println();
+ System.out.print("- getObject(Timestamp.class):\t\t");
+ System.out.print(rs.getObject(i, Timestamp.class));
+ System.out.println();
+ System.out.print("- getTimestamp(default Calendar):\t");
+ System.out.print(rs.getTimestamp(i, Calendar.getInstance()));
+ System.out.println();
+ System.out.print("- getTimestamp(UTC Calendar):\t\t");
+ System.out.print(
+ rs.getTimestamp(i, Calendar.getInstance(TimeZone.getTimeZone("UTC"))));
+ System.out.println();
+ System.out.print("- getObject(LocalDateTime.class):\t");
+ System.out.print(rs.getObject(i, LocalDateTime.class));
+ System.out.println();
+ if (type == Types.TIMESTAMP_WITH_TIMEZONE) {
+ System.out.print("- getObject(Instant.class):\t\t\t");
+ System.out.print(rs.getObject(i, Instant.class));
+ System.out.println();
+ System.out.print("- getObject(OffsetDateTime.class):\t");
+ System.out.print(rs.getObject(i, OffsetDateTime.class));
+ System.out.println();
+ System.out.print("- getObject(ZonedDateTime.class):\t");
+ System.out.print(rs.getObject(i, ZonedDateTime.class));
+ System.out.println();
+ }
+ System.out.println();
+ }
+ System.out.println();
+ } catch (SQLException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+ } catch (SQLException e) {
+ throw new RuntimeException(e);
+ }
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java
index b56bf3c63d..1fbd2f86a9 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java
@@ -16,10 +16,12 @@
*/
package org.apache.arrow.driver.jdbc.accessor;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.function.IntSupplier;
import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor;
+import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor;
import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor;
@@ -46,6 +48,8 @@
import org.apache.arrow.vector.LargeVarCharVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarBinaryVector;
+import org.apache.arrow.vector.ViewVarCharVector;
import org.apache.arrow.vector.complex.DenseUnionVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.StructVector;
@@ -239,6 +243,18 @@ public void createAccessorForFixedSizeBinaryVector() {
}
}
+ @Test
+ public void createAccessorForViewVarBinaryVector() {
+ try (ValueVector valueVector =
+ new ViewVarBinaryVector("", rootAllocatorTestExtension.getRootAllocator())) {
+ ArrowFlightJdbcAccessor accessor =
+ ArrowFlightJdbcAccessorFactory.createAccessor(
+ valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {});
+
+ assertTrue(accessor instanceof ArrowFlightJdbcBinaryVectorAccessor);
+ }
+ }
+
@Test
public void createAccessorForTimeStampVector() {
try (ValueVector valueVector = rootAllocatorTestExtension.createTimeStampMilliVector()) {
@@ -340,6 +356,18 @@ public void createAccessorForLargeVarCharVector() {
}
}
+ @Test
+ public void createAccessorForViewVarCharVector() {
+ try (ValueVector valueVector =
+ new ViewVarCharVector("", rootAllocatorTestExtension.getRootAllocator())) {
+ ArrowFlightJdbcAccessor accessor =
+ ArrowFlightJdbcAccessorFactory.createAccessor(
+ valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {});
+
+ assertTrue(accessor instanceof ArrowFlightJdbcVarCharVectorAccessor);
+ }
+ }
+
@Test
public void createAccessorForDurationVector() {
try (ValueVector valueVector =
@@ -471,4 +499,15 @@ public void createAccessorForMapVector() {
assertTrue(accessor instanceof ArrowFlightJdbcMapVectorAccessor);
}
}
+
+ @Test
+ public void createAccessorForUuidVector() {
+ try (ValueVector valueVector = rootAllocatorTestExtension.createUuidVector()) {
+ ArrowFlightJdbcAccessor accessor =
+ ArrowFlightJdbcAccessorFactory.createAccessor(
+ valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {});
+
+ assertInstanceOf(ArrowFlightJdbcUuidVectorAccessor.class, accessor);
+ }
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java
new file mode 100644
index 0000000000..b7f341240c
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java
@@ -0,0 +1,188 @@
+/*
+ * 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.accessor.impl.binary;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.util.UUID;
+import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
+import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension;
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.util.UuidUtility;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * Tests for {@link ArrowFlightJdbcUuidVectorAccessor}.
+ *
+ *
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/ArrowFlightSqlClientHandlerBuilderTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java
index 6524eaf39a..a60a71f23d 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java
@@ -149,6 +149,7 @@ public void testDefaults() {
assertEquals(Optional.empty(), builder.catalog);
assertNull(builder.flightClientCache);
assertNull(builder.connectTimeout);
+ assertNull(builder.driverVersion);
}
@Test
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 catalog =
+ throwFromCloseSession ? Optional.of("test_catalog") : Optional.empty();
+ final Collection credentialOptions = new ArrayList<>();
+ ArrowFlightSqlClientHandler.Builder builder = new ArrowFlightSqlClientHandler.Builder();
+
+ if (throwFromCloseSession) {
+ doThrow(callStatus.toRuntimeException())
+ .when(sqlClient)
+ .closeSession(any(CloseSessionRequest.class), any(CallOption[].class));
+ } else {
+ doThrow(callStatus.toRuntimeException()).when(sqlClient).close();
+ }
+
+ ArrowFlightSqlClientHandler sqlClientHandler =
+ new ArrowFlightSqlClientHandler(
+ cacheKey, sqlClient, builder, credentialOptions, catalog, null);
+
+ if (shouldSuppress) {
+ assertDoesNotThrow(sqlClientHandler::close);
+ } else {
+ assertThrows(SQLException.class, sqlClientHandler::close);
+ }
+ }
+
+ private static Object[] testCloseHandlesFlightRuntimeException() {
+ CallStatus benignInternalError =
+ new CallStatus(FlightStatusCode.INTERNAL, null, "Connection closed after GOAWAY", null);
+ CallStatus notBenignInternalError =
+ new CallStatus(FlightStatusCode.INTERNAL, null, "Not a benign internal error", null);
+ CallStatus unavailableError = new CallStatus(FlightStatusCode.UNAVAILABLE, null, null, null);
+ CallStatus unknownError = new CallStatus(FlightStatusCode.UNKNOWN, null, null, null);
+ return new Object[] {
+ new Object[] {true, benignInternalError, true},
+ new Object[] {false, benignInternalError, true},
+ new Object[] {true, notBenignInternalError, false},
+ new Object[] {false, notBenignInternalError, false},
+ new Object[] {true, unavailableError, true},
+ new Object[] {false, unavailableError, true},
+ new Object[] {true, unknownError, false},
+ new Object[] {false, unknownError, false},
+ };
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java
new file mode 100644
index 0000000000..c258a7c652
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java
@@ -0,0 +1,296 @@
+/*
+ * 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 static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.nimbusds.oauth2.sdk.Scope;
+import java.net.URI;
+import java.sql.SQLException;
+import java.util.Collections;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Named;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+/** Tests for {@link OAuthConfiguration}. */
+public class OAuthConfigurationTest {
+
+ private static final String TOKEN_URI = "https://auth.example.com/token";
+ private static final String CLIENT_ID = "test-client-id";
+ private static final String CLIENT_SECRET = "test-client-secret";
+ private static final String SCOPE = "read write";
+ private static final String SUBJECT_TOKEN = "subject-token-value";
+ public static final String RESOURCE = "https://api.example.com/resource";
+
+ @FunctionalInterface
+ interface BuilderConfigurer {
+ void configure(OAuthConfiguration.Builder builder) throws SQLException;
+ }
+
+ static Stream createFlowCases() {
+ return Stream.of(
+ Arguments.of(
+ Named.of(
+ "string flow", (BuilderConfigurer) builder -> builder.flow("client_credentials"))),
+ Arguments.of(
+ Named.of(
+ "uppercase string flow",
+ (BuilderConfigurer) builder -> builder.flow("CLIENT_CREDENTIALS"))));
+ }
+
+ @ParameterizedTest
+ @MethodSource("createFlowCases")
+ public void testCreateFlowConfiguration(BuilderConfigurer flowConfigurer) throws SQLException {
+ OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder();
+ flowConfigurer.configure(builder);
+ OAuthConfiguration config =
+ builder.tokenUri(TOKEN_URI).clientId(CLIENT_ID).clientSecret(CLIENT_SECRET).build();
+
+ // Verify configuration creates correct provider type
+ OAuthTokenProvider provider = config.createTokenProvider();
+ assertInstanceOf(ClientCredentialsTokenProvider.class, provider);
+ }
+
+ @Test
+ public void testCreateClientCredentialsTokenProvider() throws SQLException {
+ OAuthConfiguration config =
+ new OAuthConfiguration.Builder()
+ .flow("client_credentials")
+ .tokenUri(TOKEN_URI)
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)
+ .scope(SCOPE)
+ .build();
+
+ OAuthTokenProvider provider = config.createTokenProvider();
+
+ assertNotNull(provider);
+ assertInstanceOf(ClientCredentialsTokenProvider.class, provider);
+
+ ClientCredentialsTokenProvider ccProvider = (ClientCredentialsTokenProvider) provider;
+ assertEquals(URI.create(TOKEN_URI), ccProvider.tokenUri);
+ assertEquals(CLIENT_ID, ccProvider.clientAuth.getClientID().getValue());
+ assertEquals(Scope.parse(SCOPE), ccProvider.scope);
+ }
+
+ @Test
+ public void testCreateTokenExchangeTokenProviderWithAllOptions() throws SQLException {
+ String subjectTokenType = "urn:ietf:params:oauth:token-type:access_token";
+ String actorToken = "actor-token-value";
+ String actorTokenType = "urn:ietf:params:oauth:token-type:jwt";
+ String audience = "https://api.example.com";
+ String requestedTokenType = "urn:ietf:params:oauth:token-type:access_token";
+
+ OAuthConfiguration config =
+ new OAuthConfiguration.Builder()
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .scope(SCOPE)
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)
+ .resource(RESOURCE)
+ .subjectToken(SUBJECT_TOKEN)
+ .subjectTokenType(subjectTokenType)
+ .actorToken(actorToken)
+ .actorTokenType(actorTokenType)
+ .audience(audience)
+ .requestedTokenType(requestedTokenType)
+ .build();
+
+ OAuthTokenProvider provider = config.createTokenProvider();
+
+ assertNotNull(provider);
+ assertInstanceOf(TokenExchangeTokenProvider.class, provider);
+
+ TokenExchangeTokenProvider teProvider = (TokenExchangeTokenProvider) provider;
+ assertEquals(URI.create(TOKEN_URI), teProvider.tokenUri);
+ assertNotNull(teProvider.grant);
+ assertEquals(SUBJECT_TOKEN, teProvider.grant.getSubjectToken().getValue());
+ assertEquals(subjectTokenType, teProvider.grant.getSubjectTokenType().getURI().toString());
+ assertEquals(actorToken, teProvider.grant.getActorToken().getValue());
+ assertEquals(actorTokenType, teProvider.grant.getActorTokenType().getURI().toString());
+ assertNotNull(teProvider.grant.getAudience());
+ assertEquals(1, teProvider.grant.getAudience().size());
+ assertEquals(audience, teProvider.grant.getAudience().get(0).getValue());
+ assertEquals(requestedTokenType, teProvider.grant.getRequestedTokenType().getURI().toString());
+ assertEquals(Scope.parse(SCOPE), teProvider.scope);
+ assertEquals(Collections.singletonList(URI.create(RESOURCE)), teProvider.resources);
+
+ assertEquals(CLIENT_ID, teProvider.clientAuth.getClientID().getValue());
+ }
+
+ static Stream generalValidationErrorCases() {
+ return Stream.of(
+ Arguments.of(
+ Named.of(
+ "null flow",
+ (BuilderConfigurer) builder -> builder.flow((String) null).tokenUri(TOKEN_URI)),
+ "OAuth flow cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "empty flow", (BuilderConfigurer) builder -> builder.flow("").tokenUri(TOKEN_URI)),
+ "OAuth flow cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "invalid flow",
+ (BuilderConfigurer) builder -> builder.flow("invalid_flow").tokenUri(TOKEN_URI)),
+ "Unsupported OAuth flow: invalid_flow"),
+ Arguments.of(
+ Named.of(
+ "null tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri((String) null)
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)),
+ "Token URI cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "empty tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri("")
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)),
+ "Token URI cannot be null or empty"),
+ Arguments.of(
+ Named.of(
+ "invalid tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri("not a valid uri ://")
+ .clientId(CLIENT_ID)
+ .clientSecret(CLIENT_SECRET)),
+ null),
+ Arguments.of(
+ Named.of(
+ "invalid tokenUri",
+ (BuilderConfigurer)
+ builder ->
+ builder.flow("client_credentials").tokenUri(TOKEN_URI).clientId(CLIENT_ID)),
+ // null means verify exception has message and cause
+ "clientSecret is required for client_credentials flow"));
+ }
+
+ @ParameterizedTest
+ @MethodSource("generalValidationErrorCases")
+ public void testGeneralValidationErrors(BuilderConfigurer configurer, String expectedMessage) {
+ SQLException exception =
+ assertThrows(
+ SQLException.class,
+ () -> {
+ OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder();
+ configurer.configure(builder);
+ builder.build();
+ });
+
+ if (expectedMessage != null) {
+ assertEquals(expectedMessage, exception.getMessage());
+ } else {
+ assertNotNull(exception.getMessage());
+ assertNotNull(exception.getCause());
+ }
+ }
+
+ static Stream flowSpecificValidationErrorCases() {
+ return Stream.of(
+ // client_credentials flow validation
+ Arguments.of(
+ Named.of(
+ "client_credentials: missing clientId",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("client_credentials")
+ .tokenUri(TOKEN_URI)
+ .clientSecret(CLIENT_SECRET)),
+ "clientId is required for client_credentials flow"),
+ Arguments.of(
+ Named.of(
+ "client_credentials: missing clientSecret",
+ (BuilderConfigurer)
+ builder ->
+ builder.flow("client_credentials").tokenUri(TOKEN_URI).clientId(CLIENT_ID)),
+ "clientSecret is required for client_credentials flow"),
+ // token_exchange flow validation
+ Arguments.of(
+ Named.of(
+ "token_exchange: missing subjectToken",
+ (BuilderConfigurer) builder -> builder.flow("token_exchange").tokenUri(TOKEN_URI)),
+ "subjectToken is required for token_exchange flow"),
+ Arguments.of(
+ Named.of(
+ "token_exchange: empty subjectToken",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .subjectToken("")
+ .subjectTokenType("urn:ietf:params:oauth:token-type:access_token")),
+ "subjectToken is required for token_exchange flow"),
+ Arguments.of(
+ Named.of(
+ "token_exchange: missing subjectTokenType",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .subjectToken(SUBJECT_TOKEN)),
+ "subjectTokenType is required for token_exchange flow"),
+ Arguments.of(
+ Named.of(
+ "token_exchange: empty subjectTokenType",
+ (BuilderConfigurer)
+ builder ->
+ builder
+ .flow("token_exchange")
+ .tokenUri(TOKEN_URI)
+ .subjectToken(SUBJECT_TOKEN)
+ .subjectTokenType("")),
+ "subjectTokenType is required for token_exchange flow"));
+ }
+
+ @ParameterizedTest
+ @MethodSource("flowSpecificValidationErrorCases")
+ public void testFlowSpecificValidationErrors(
+ BuilderConfigurer configurer, String expectedMessage) {
+ SQLException exception =
+ assertThrows(
+ SQLException.class,
+ () -> {
+ OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder();
+ configurer.configure(builder);
+ builder.build();
+ });
+
+ assertEquals(expectedMessage, exception.getMessage());
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java
new file mode 100644
index 0000000000..1a33f7f0ae
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java
@@ -0,0 +1,95 @@
+/*
+ * 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 static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.sql.SQLException;
+import org.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.FlightCallHeaders;
+import org.apache.arrow.flight.auth2.Auth2Constants;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+/** Tests for {@link OAuthCredentialWriter}. */
+@ExtendWith(MockitoExtension.class)
+public class OAuthCredentialWriterTest {
+
+ @Mock private OAuthTokenProvider mockTokenProvider;
+
+ @Test
+ public void testConstructorRejectsNullTokenProvider() {
+ assertThrows(NullPointerException.class, () -> new OAuthCredentialWriter(null));
+ }
+
+ @Test
+ public void testAcceptWritesBearerTokenToHeaders() throws SQLException {
+ String testToken = "test-access-token-12345";
+ when(mockTokenProvider.getValidToken()).thenReturn(testToken);
+
+ OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider);
+ CallHeaders headers = new FlightCallHeaders();
+
+ writer.accept(headers);
+
+ verify(mockTokenProvider).getValidToken();
+ assertEquals(
+ Auth2Constants.BEARER_PREFIX + testToken, headers.get(Auth2Constants.AUTHORIZATION_HEADER));
+ }
+
+ @Test
+ public void testAcceptThrowsOAuthTokenExceptionOnSQLException() throws SQLException {
+ SQLException sqlException = new SQLException("Token fetch failed");
+ when(mockTokenProvider.getValidToken()).thenThrow(sqlException);
+
+ OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider);
+ CallHeaders headers = new FlightCallHeaders();
+
+ OAuthTokenException exception =
+ assertThrows(OAuthTokenException.class, () -> writer.accept(headers));
+
+ assertEquals("Failed to obtain OAuth token", exception.getMessage());
+ assertEquals(sqlException, exception.getCause());
+ }
+
+ @Test
+ public void testAcceptCallsTokenProviderEachTime() throws SQLException {
+ when(mockTokenProvider.getValidToken())
+ .thenReturn("token1")
+ .thenReturn("token2")
+ .thenReturn("token3");
+
+ OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider);
+
+ CallHeaders headers1 = new FlightCallHeaders();
+ writer.accept(headers1);
+ assertEquals("Bearer token1", headers1.get(Auth2Constants.AUTHORIZATION_HEADER));
+
+ CallHeaders headers2 = new FlightCallHeaders();
+ writer.accept(headers2);
+ assertEquals("Bearer token2", headers2.get(Auth2Constants.AUTHORIZATION_HEADER));
+
+ CallHeaders headers3 = new FlightCallHeaders();
+ writer.accept(headers3);
+ assertEquals("Bearer token3", headers3.get(Auth2Constants.AUTHORIZATION_HEADER));
+ }
+}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java
new file mode 100644
index 0000000000..07751f0abc
--- /dev/null
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java
@@ -0,0 +1,160 @@
+/*
+ * 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.converter.impl;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+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.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.sql.Types;
+import java.util.UUID;
+import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension;
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.util.UuidUtility;
+import org.apache.calcite.avatica.AvaticaParameter;
+import org.apache.calcite.avatica.ColumnMetaData;
+import org.apache.calcite.avatica.remote.TypedValue;
+import org.apache.calcite.avatica.util.ByteString;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * Tests for {@link UuidAvaticaParameterConverter}.
+ *
+ * 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 expectedColumnMetaData =
@@ -78,6 +91,25 @@ public void testShouldConvertArrowFieldsToColumnMetaDataList() {
.setCatalogName("catalog1")
.setSchemaName("schema1")
.setTableName("table1")
+ .setColumnName("col1")
+ .setType(
+ Common.AvaticaType.newBuilder()
+ .setId(SqlTypes.getSqlTypeIdFromArrowType(ArrowType.Utf8.INSTANCE))
+ .setName(SqlTypes.getSqlTypeNameFromArrowType(ArrowType.Utf8.INSTANCE))
+ .build())
+ .build()),
+ ColumnMetaData.fromProto(
+ Common.ColumnMetaData.newBuilder()
+ .setCatalogName("catalog1")
+ .setSchemaName("schema1")
+ .setTableName("table1")
+ .setColumnName("col2")
+ .setType(
+ Common.AvaticaType.newBuilder()
+ .setId(SqlTypes.getSqlTypeIdFromArrowType(ArrowType.Utf8View.INSTANCE))
+ .setName(
+ SqlTypes.getSqlTypeNameFromArrowType(ArrowType.Utf8View.INSTANCE))
+ .build())
.build()));
final List actualColumnMetaData =
@@ -95,6 +127,8 @@ private void assertColumnMetaData(
assertThat(expectedColumnMetaData.catalogName, equalTo(actualColumnMetaData.catalogName));
assertThat(expectedColumnMetaData.schemaName, equalTo(actualColumnMetaData.schemaName));
assertThat(expectedColumnMetaData.tableName, equalTo(actualColumnMetaData.tableName));
+ assertThat(expectedColumnMetaData.columnName, equalTo(actualColumnMetaData.columnName));
+ assertThat(expectedColumnMetaData.type, equalTo(actualColumnMetaData.type));
assertThat(expectedColumnMetaData.readOnly, equalTo(actualColumnMetaData.readOnly));
assertThat(expectedColumnMetaData.autoIncrement, equalTo(actualColumnMetaData.autoIncrement));
assertThat(expectedColumnMetaData.precision, equalTo(actualColumnMetaData.precision));
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java
index 8197d7d95f..7c17755693 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java
@@ -28,8 +28,10 @@
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
+import java.util.UUID;
import java.util.function.Consumer;
import java.util.stream.IntStream;
import org.apache.arrow.flight.FlightProducer.ServerStreamListener;
@@ -40,10 +42,13 @@
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.TimeStampMilliVector;
import org.apache.arrow.vector.UInt4Vector;
+import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.TimeUnit;
@@ -52,6 +57,7 @@
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.Text;
+import org.apache.arrow.vector.util.UuidUtility;
/** Standard {@link MockFlightSqlProducer} instances for tests. */
// TODO Remove this once all tests are refactor to use only the queries they need.
@@ -62,6 +68,22 @@ public final class CoreMockedSqlProducers {
public static final String LEGACY_CANCELLATION_SQL_CMD = "SELECT * FROM TAKES_FOREVER";
public static final String LEGACY_REGULAR_WITH_EMPTY_SQL_CMD = "SELECT * FROM TEST_EMPTIES";
+ public static final String UUID_SQL_CMD = "SELECT * FROM UUID_TABLE";
+ public static final String UUID_PREPARED_SELECT_SQL_CMD =
+ "SELECT * FROM UUID_TABLE WHERE uuid_col = ?";
+ public static final String UUID_PREPARED_UPDATE_SQL_CMD =
+ "UPDATE UUID_TABLE SET uuid_col = ? WHERE id = ?";
+
+ public static final UUID UUID_1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000");
+ public static final UUID UUID_2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8");
+ public static final UUID UUID_3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479");
+
+ public static final Schema UUID_SCHEMA =
+ new Schema(
+ ImmutableList.of(
+ new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null),
+ new Field("uuid_col", new FieldType(true, UuidType.INSTANCE, null), null)));
+
private CoreMockedSqlProducers() {
// Prevent instantiation.
}
@@ -78,9 +100,109 @@ public static MockFlightSqlProducer getLegacyProducer() {
addLegacyMetadataSqlCmdSupport(producer);
addLegacyCancellationSqlCmdSupport(producer);
addQueryWithEmbeddedEmptyRoot(producer);
+ addUuidSqlCmdSupport(producer);
+ addUuidPreparedSelectSqlCmdSupport(producer);
+ addUuidPreparedUpdateSqlCmdSupport(producer);
return producer;
}
+ /**
+ * Gets a {@link MockFlightSqlProducer} configured with UUID test data.
+ *
+ * @return a new producer with UUID support.
+ */
+ public static MockFlightSqlProducer getUuidProducer() {
+ final MockFlightSqlProducer producer = new MockFlightSqlProducer();
+ addUuidSqlCmdSupport(producer);
+ return producer;
+ }
+
+ private static void addUuidPreparedUpdateSqlCmdSupport(final MockFlightSqlProducer producer) {
+ final String query = "UPDATE UUID_TABLE SET uuid_col = ? WHERE id = ?";
+ final Schema parameterSchema =
+ new Schema(
+ Arrays.asList(
+ new Field("", new FieldType(true, UuidType.INSTANCE, null), null),
+ Field.nullable("", new ArrowType.Int(32, true))));
+
+ producer.addUpdateQuery(query, 1);
+ producer.addExpectedParameters(
+ UUID_PREPARED_UPDATE_SQL_CMD,
+ parameterSchema,
+ Collections.singletonList(Arrays.asList(CoreMockedSqlProducers.UUID_3, 1)));
+ }
+
+ private static void addUuidPreparedSelectSqlCmdSupport(final MockFlightSqlProducer producer) {
+ final Schema parameterSchema =
+ new Schema(
+ Collections.singletonList(
+ new Field("", new FieldType(true, UuidType.INSTANCE, null), null)));
+
+ final Consumer uuidResultProvider =
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(UUID_SCHEMA, allocator)) {
+ root.allocateNew();
+ IntVector idVector = (IntVector) root.getVector("id");
+ UuidVector uuidVector = (UuidVector) root.getVector("uuid_col");
+ idVector.setSafe(0, 1);
+ uuidVector.setSafe(0, UuidUtility.getBytesFromUUID(CoreMockedSqlProducers.UUID_1));
+ root.setRowCount(1);
+ listener.start(root);
+ listener.putNext();
+ } catch (final Throwable throwable) {
+ listener.error(throwable);
+ } finally {
+ listener.completed();
+ }
+ };
+
+ producer.addSelectQuery(
+ UUID_PREPARED_SELECT_SQL_CMD, UUID_SCHEMA, Collections.singletonList(uuidResultProvider));
+ producer.addExpectedParameters(
+ UUID_PREPARED_SELECT_SQL_CMD,
+ parameterSchema,
+ Collections.singletonList(Collections.singletonList(CoreMockedSqlProducers.UUID_1)));
+ }
+
+ private static void addUuidSqlCmdSupport(final MockFlightSqlProducer producer) {
+ final Consumer uuidResultProvider =
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(UUID_SCHEMA, allocator)) {
+ root.allocateNew();
+
+ IntVector idVector = (IntVector) root.getVector("id");
+ UuidVector uuidVector = (UuidVector) root.getVector("uuid_col");
+
+ // Row 0: id=1, uuid=UUID_1
+ idVector.setSafe(0, 1);
+ uuidVector.setSafe(0, UuidUtility.getBytesFromUUID(UUID_1));
+
+ // Row 1: id=2, uuid=UUID_2
+ idVector.setSafe(1, 2);
+ uuidVector.setSafe(1, UuidUtility.getBytesFromUUID(UUID_2));
+
+ // Row 2: id=3, uuid=UUID_3
+ idVector.setSafe(2, 3);
+ uuidVector.setSafe(2, UuidUtility.getBytesFromUUID(UUID_3));
+
+ // Row 3: id=4, uuid=NULL
+ idVector.setSafe(3, 4);
+ uuidVector.setNull(3);
+
+ root.setRowCount(4);
+ listener.start(root);
+ listener.putNext();
+ } finally {
+ listener.completed();
+ }
+ };
+
+ producer.addSelectQuery(
+ UUID_SQL_CMD, UUID_SCHEMA, Collections.singletonList(uuidResultProvider));
+ }
+
private static void addQueryWithEmbeddedEmptyRoot(final MockFlightSqlProducer producer) {
final Schema querySchema =
new Schema(
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
index a8874c4869..45c2a96404 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
@@ -52,6 +52,9 @@
import org.apache.arrow.flight.PutResult;
import org.apache.arrow.flight.Result;
import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.SessionOptionValue;
+import org.apache.arrow.flight.SetSessionOptionsRequest;
+import org.apache.arrow.flight.SetSessionOptionsResult;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.flight.sql.FlightSqlProducer;
import org.apache.arrow.flight.sql.SqlInfoBuilder;
@@ -664,6 +667,22 @@ public SqlInfoBuilder getSqlInfoBuilder() {
return sqlInfoBuilder;
}
+ private final Map sessionOptions = new HashMap<>();
+
+ @Override
+ public void setSessionOptions(
+ final SetSessionOptionsRequest request,
+ final CallContext context,
+ final StreamListener listener) {
+ sessionOptions.putAll(request.getSessionOptions());
+ listener.onNext(new SetSessionOptionsResult(Collections.emptyMap()));
+ listener.onCompleted();
+ }
+
+ public Map getSessionOptions() {
+ return sessionOptions;
+ }
+
private static final class TicketConversionUtils {
private TicketConversionUtils() {
// Prevent instantiation.
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java
index 347e92a16c..4b299d63e0 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java
@@ -19,6 +19,7 @@
import java.math.BigDecimal;
import java.nio.charset.StandardCharsets;
import java.util.Random;
+import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;
import org.apache.arrow.memory.BufferAllocator;
@@ -53,6 +54,7 @@
import org.apache.arrow.vector.UInt2Vector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.UInt8Vector;
+import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.LargeListVector;
@@ -60,6 +62,7 @@
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter;
import org.apache.arrow.vector.complex.impl.UnionLargeListWriter;
import org.apache.arrow.vector.complex.impl.UnionListWriter;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.extension.AfterAllCallback;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
@@ -811,4 +814,23 @@ public FixedSizeListVector createFixedSizeListVector() {
return valueVector;
}
+
+ /**
+ * Create a UuidVector to be used in the accessor tests.
+ *
+ * @return UuidVector
+ */
+ public UuidVector createUuidVector() {
+ UuidVector valueVector = new UuidVector("", this.getRootAllocator());
+ valueVector.allocateNew(3);
+ valueVector.setSafe(
+ 0, UuidUtility.getBytesFromUUID(UUID.fromString("550e8400-e29b-41d4-a716-446655440000")));
+ valueVector.setSafe(
+ 1, UuidUtility.getBytesFromUUID(UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8")));
+ valueVector.setSafe(
+ 2, UuidUtility.getBytesFromUUID(UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479")));
+ valueVector.setValueCount(3);
+
+ return valueVector;
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java
index a6dd6b3275..c4858d787d 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java
@@ -21,6 +21,7 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import java.sql.Types;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.IntervalUnit;
@@ -40,9 +41,11 @@ public void testGetSqlTypeIdFromArrowType() {
assertEquals(Types.BINARY, getSqlTypeIdFromArrowType(new ArrowType.FixedSizeBinary(1024)));
assertEquals(Types.VARBINARY, getSqlTypeIdFromArrowType(new ArrowType.Binary()));
+ assertEquals(Types.VARBINARY, getSqlTypeIdFromArrowType(new ArrowType.BinaryView()));
assertEquals(Types.LONGVARBINARY, getSqlTypeIdFromArrowType(new ArrowType.LargeBinary()));
assertEquals(Types.VARCHAR, getSqlTypeIdFromArrowType(new ArrowType.Utf8()));
+ assertEquals(Types.VARCHAR, getSqlTypeIdFromArrowType(new ArrowType.Utf8View()));
assertEquals(Types.LONGVARCHAR, getSqlTypeIdFromArrowType(new ArrowType.LargeUtf8()));
assertEquals(Types.DATE, getSqlTypeIdFromArrowType(new ArrowType.Date(DateUnit.MILLISECOND)));
@@ -83,6 +86,8 @@ public void testGetSqlTypeIdFromArrowType() {
assertEquals(Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Map(true)));
assertEquals(Types.NULL, getSqlTypeIdFromArrowType(new ArrowType.Null()));
+
+ assertEquals(Types.OTHER, getSqlTypeIdFromArrowType(UuidType.INSTANCE));
}
@Test
@@ -94,9 +99,11 @@ public void testGetSqlTypeNameFromArrowType() {
assertEquals("BINARY", getSqlTypeNameFromArrowType(new ArrowType.FixedSizeBinary(1024)));
assertEquals("VARBINARY", getSqlTypeNameFromArrowType(new ArrowType.Binary()));
+ assertEquals("VARBINARY", getSqlTypeNameFromArrowType(new ArrowType.BinaryView()));
assertEquals("LONGVARBINARY", getSqlTypeNameFromArrowType(new ArrowType.LargeBinary()));
assertEquals("VARCHAR", getSqlTypeNameFromArrowType(new ArrowType.Utf8()));
+ assertEquals("VARCHAR", getSqlTypeNameFromArrowType(new ArrowType.Utf8View()));
assertEquals("LONGVARCHAR", getSqlTypeNameFromArrowType(new ArrowType.LargeUtf8()));
assertEquals("DATE", getSqlTypeNameFromArrowType(new ArrowType.Date(DateUnit.MILLISECOND)));
@@ -136,5 +143,7 @@ public void testGetSqlTypeNameFromArrowType() {
assertEquals("JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Map(true)));
assertEquals("NULL", getSqlTypeNameFromArrowType(new ArrowType.Null()));
+
+ assertEquals("OTHER", getSqlTypeNameFromArrowType(UuidType.INSTANCE));
}
}
diff --git a/flight/flight-sql-jdbc-driver/pom.xml b/flight/flight-sql-jdbc-driver/pom.xml
index 3776e97f3f..e6f23bcb08 100644
--- a/flight/flight-sql-jdbc-driver/pom.xml
+++ b/flight/flight-sql-jdbc-driver/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 18.3.0
+ 19.0.0
flight-sql-jdbc-driver
diff --git a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt
index 8bc43cbe0f..8476bd9995 100644
--- a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt
+++ b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt
@@ -345,6 +345,14 @@ License: https://www.apache.org/licenses/LICENSE-2.0
--------------------------------------------------------------------------------
+This binary artifact contains Nimbus OAuth 2.0 SDK with OpenID Connect extensions 11.20.1.
+
+Copyright: Copyright 2012-2024 Connect2id Ltd.
+Home page: https://connect2id.com/products/nimbus-oauth-openid-connect-sdk
+License: https://www.apache.org/licenses/LICENSE-2.0
+
+--------------------------------------------------------------------------------
+
This binary artifact contains Bouncycastle 1.80.
Copyright: Copyright (c) 2000-2024 The Legion of the Bouncy Castle Inc. (https://www.bouncycastle.org).
diff --git a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java
index c1bd111fb9..145744ad38 100644
--- a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java
+++ b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java
@@ -70,6 +70,10 @@ public class ITDriverJarValidation {
"LICENSE.txt",
"NOTICE.txt",
"arrow-git.properties",
+ "iso3166_1alpha2-codes.properties",
+ "iso3166_1alpha3-codes.properties",
+ "iso3166_1alpha-2-3-map.properties",
+ "iso3166_3-codes.properties",
"properties/flight.properties",
"META-INF/io.netty.versions.properties",
"META-INF/MANIFEST.MF",
diff --git a/flight/flight-sql/pom.xml b/flight/flight-sql/pom.xml
index 66ade30306..a58b76acda 100644
--- a/flight/flight-sql/pom.xml
+++ b/flight/flight-sql/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-flight
- 18.3.0
+ 19.0.0
flight-sql
@@ -95,7 +95,7 @@ under the License.
org.apache.commons
commons-dbcp2
- 2.13.0
+ 2.14.0
test
@@ -107,25 +107,25 @@ under the License.
org.apache.commons
commons-pool2
- 2.12.1
+ 2.13.1
test
org.apache.commons
commons-text
- 1.13.1
- test
-
-
- org.hamcrest
- hamcrest
+ 1.15.0
test
commons-cli
commons-cli
- 1.9.0
+ 1.11.0
true
+
+ org.assertj
+ assertj-core
+ test
+
diff --git a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
index 3f769363fb..e2934ab1e9 100644
--- a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
+++ b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
@@ -21,10 +21,7 @@
import static java.util.Collections.singletonList;
import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults;
import static org.apache.arrow.util.AutoCloseables.close;
-import static org.hamcrest.CoreMatchers.containsString;
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.CoreMatchers.notNullValue;
-import static org.hamcrest.CoreMatchers.nullValue;
+import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertAll;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -40,6 +37,7 @@
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.stream.IntStream;
import org.apache.arrow.flight.CancelFlightInfoRequest;
@@ -76,8 +74,7 @@
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.Text;
import org.apache.arrow.vector.util.VectorBatchAppender;
-import org.hamcrest.Matcher;
-import org.hamcrest.MatcherAssert;
+import org.assertj.core.api.Condition;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -247,16 +244,15 @@ private static List> getNonConformingResultsForGetSqlInfo(
@Test
public void testGetTablesSchema() {
final FlightInfo info = sqlClient.getTables(null, null, null, null, true);
- MatcherAssert.assertThat(
- info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA)));
+ assertThat(info.getSchemaOptional())
+ .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA));
}
@Test
public void testGetTablesSchemaExcludeSchema() {
final FlightInfo info = sqlClient.getTables(null, null, null, null, false);
- MatcherAssert.assertThat(
- info.getSchemaOptional(),
- is(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)));
+ assertThat(info.getSchemaOptional())
+ .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA));
}
@Test
@@ -266,8 +262,8 @@ public void testGetTablesResultNoSchema() throws Exception {
sqlClient.getTables(null, null, null, null, false).getEndpoints().get(0).getTicket())) {
assertAll(
() -> {
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA));
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA);
},
() -> {
final List> results = getResults(stream);
@@ -301,7 +297,7 @@ public void testGetTablesResultNoSchema() throws Exception {
asList(null /* TODO No catalog yet */, "SYSIBM", "SYSDUMMY1", "SYSTEM TABLE"),
asList(null /* TODO No catalog yet */, "APP", "FOREIGNTABLE", "TABLE"),
asList(null /* TODO No catalog yet */, "APP", "INTTABLE", "TABLE"));
- MatcherAssert.assertThat(results, is(expectedResults));
+ assertThat(results).isEqualTo(expectedResults);
});
}
}
@@ -318,8 +314,8 @@ public void testGetTablesResultFilteredNoSchema() throws Exception {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA),
() -> {
final List> results = getResults(stream);
final List> expectedResults =
@@ -327,7 +323,7 @@ public void testGetTablesResultFilteredNoSchema() throws Exception {
// catalog_name | schema_name | table_name | table_type | table_schema
asList(null /* TODO No catalog yet */, "APP", "FOREIGNTABLE", "TABLE"),
asList(null /* TODO No catalog yet */, "APP", "INTTABLE", "TABLE"));
- MatcherAssert.assertThat(results, is(expectedResults));
+ assertThat(results).isEqualTo(expectedResults);
});
}
}
@@ -343,11 +339,9 @@ public void testGetTablesResultFilteredWithSchema() throws Exception {
.getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA)),
+ assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA),
() -> {
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA));
+ assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA);
final List> results = getResults(stream);
final List> expectedResults =
ImmutableList.of(
@@ -487,7 +481,7 @@ public void testGetTablesResultFilteredWithSchema() throws Exception {
.getMetadataMap()),
null)))
.toJson()));
- MatcherAssert.assertThat(results, is(expectedResults));
+ assertThat(results).isEqualTo(expectedResults);
});
}
}
@@ -498,11 +492,11 @@ public void testSimplePreparedStatementSchema() throws Exception {
assertAll(
() -> {
final Schema actualSchema = preparedStatement.getResultSetSchema();
- MatcherAssert.assertThat(actualSchema, is(SCHEMA_INT_TABLE));
+ assertThat(actualSchema).isEqualTo(SCHEMA_INT_TABLE);
},
() -> {
final FlightInfo info = preparedStatement.execute();
- MatcherAssert.assertThat(info.getSchemaOptional(), is(Optional.of(SCHEMA_INT_TABLE)));
+ assertThat(info.getSchemaOptional()).isEqualTo(Optional.of(SCHEMA_INT_TABLE));
});
}
}
@@ -513,10 +507,8 @@ public void testSimplePreparedStatementResults() throws Exception {
final FlightStream stream =
sqlClient.getStream(preparedStatement.execute().getEndpoints().get(0).getTicket())) {
assertAll(
- () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)),
- () ->
- MatcherAssert.assertThat(
- getResults(stream), is(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY)));
+ () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE),
+ () -> assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY));
}
}
@@ -538,10 +530,8 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce
FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket());
assertAll(
- () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)),
- () ->
- MatcherAssert.assertThat(
- getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)));
+ () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE),
+ () -> assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_PARAMETER_BINDING));
}
}
}
@@ -579,8 +569,8 @@ public void testSimplePreparedStatementUpdateResults() throws SQLException {
deletedRows = deletePrepare.executeUpdate();
}
assertAll(
- () -> MatcherAssert.assertThat(updatedRows, is(10L)),
- () -> MatcherAssert.assertThat(deletedRows, is(10L)));
+ () -> assertThat(updatedRows).isEqualTo(10L),
+ () -> assertThat(deletedRows).isEqualTo(10L));
}
}
}
@@ -647,7 +637,7 @@ public void testBulkIngest() throws IOException {
null,
null));
- MatcherAssert.assertThat(updatedRows, is(-1L));
+ assertThat(updatedRows).isEqualTo(-1L);
// Ingest directly using VectorSchemaRoot
populateNext10RowsInIngestRootBatch(
@@ -672,7 +662,7 @@ public void testBulkIngest() throws IOException {
deletedRows = deletePrepare.executeUpdate();
}
- MatcherAssert.assertThat(deletedRows, is(30L));
+ assertThat(deletedRows).isEqualTo(30L);
}
}
}
@@ -709,8 +699,7 @@ public void testSimplePreparedStatementUpdateResultsWithoutParameters() throws S
final long deletedRows = deletePrepare.executeUpdate();
assertAll(
- () -> MatcherAssert.assertThat(updatedRows, is(1L)),
- () -> MatcherAssert.assertThat(deletedRows, is(1L)));
+ () -> assertThat(updatedRows).isEqualTo(1L), () -> assertThat(deletedRows).isEqualTo(1L));
}
}
@@ -719,19 +708,19 @@ public void testSimplePreparedStatementClosesProperly() {
final PreparedStatement preparedStatement = sqlClient.prepare("SELECT * FROM intTable");
assertAll(
() -> {
- MatcherAssert.assertThat(preparedStatement.isClosed(), is(false));
+ assertThat(preparedStatement.isClosed()).isEqualTo(false);
},
() -> {
preparedStatement.close();
- MatcherAssert.assertThat(preparedStatement.isClosed(), is(true));
+ assertThat(preparedStatement.isClosed()).isEqualTo(true);
});
}
@Test
public void testGetCatalogsSchema() {
final FlightInfo info = sqlClient.getCatalogs();
- MatcherAssert.assertThat(
- info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA)));
+ assertThat(info.getSchemaOptional())
+ .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA));
}
@Test
@@ -740,11 +729,11 @@ public void testGetCatalogsResults() throws Exception {
sqlClient.getStream(sqlClient.getCatalogs().getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA),
() -> {
List> catalogs = getResults(stream);
- MatcherAssert.assertThat(catalogs, is(emptyList()));
+ assertThat(catalogs).isEqualTo(emptyList());
});
}
}
@@ -752,9 +741,8 @@ public void testGetCatalogsResults() throws Exception {
@Test
public void testGetTableTypesSchema() {
final FlightInfo info = sqlClient.getTableTypes();
- MatcherAssert.assertThat(
- info.getSchemaOptional(),
- is(Optional.of(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA)));
+ assertThat(info.getSchemaOptional())
+ .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA));
}
@Test
@@ -763,8 +751,8 @@ public void testGetTableTypesResult() throws Exception {
sqlClient.getStream(sqlClient.getTableTypes().getEndpoints().get(0).getTicket())) {
assertAll(
() -> {
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA));
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA);
},
() -> {
final List> tableTypes = getResults(stream);
@@ -775,7 +763,7 @@ public void testGetTableTypesResult() throws Exception {
singletonList("SYSTEM TABLE"),
singletonList("TABLE"),
singletonList("VIEW"));
- MatcherAssert.assertThat(tableTypes, is(expectedTableTypes));
+ assertThat(tableTypes).isEqualTo(expectedTableTypes);
});
}
}
@@ -783,8 +771,8 @@ public void testGetTableTypesResult() throws Exception {
@Test
public void testGetSchemasSchema() {
final FlightInfo info = sqlClient.getSchemas(null, null);
- MatcherAssert.assertThat(
- info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA)));
+ assertThat(info.getSchemaOptional())
+ .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA));
}
@Test
@@ -793,8 +781,7 @@ public void testGetSchemasResult() throws Exception {
sqlClient.getStream(sqlClient.getSchemas(null, null).getEndpoints().get(0).getTicket())) {
assertAll(
() -> {
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA));
+ assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA);
},
() -> {
final List> schemas = getResults(stream);
@@ -812,7 +799,7 @@ public void testGetSchemasResult() throws Exception {
asList(null /* TODO Add catalog. */, "SYSIBM"),
asList(null /* TODO Add catalog. */, "SYSPROC"),
asList(null /* TODO Add catalog. */, "SYSSTAT"));
- MatcherAssert.assertThat(schemas, is(expectedSchemas));
+ assertThat(schemas).isEqualTo(expectedSchemas);
});
}
}
@@ -825,24 +812,24 @@ public void testGetPrimaryKey() {
final List> results = getResults(stream);
assertAll(
- () -> MatcherAssert.assertThat(results.size(), is(1)),
+ () -> assertThat(results.size()).isEqualTo(1),
() -> {
final List result = results.get(0);
assertAll(
- () -> MatcherAssert.assertThat(result.get(0), is("")),
- () -> MatcherAssert.assertThat(result.get(1), is("APP")),
- () -> MatcherAssert.assertThat(result.get(2), is("INTTABLE")),
- () -> MatcherAssert.assertThat(result.get(3), is("ID")),
- () -> MatcherAssert.assertThat(result.get(4), is("1")),
- () -> MatcherAssert.assertThat(result.get(5), notNullValue()));
+ () -> assertThat(result.get(0)).isEqualTo(""),
+ () -> assertThat(result.get(1)).isEqualTo("APP"),
+ () -> assertThat(result.get(2)).isEqualTo("INTTABLE"),
+ () -> assertThat(result.get(3)).isEqualTo("ID"),
+ () -> assertThat(result.get(4)).isEqualTo("1"),
+ () -> assertThat(result.get(5)).isNotNull());
});
}
@Test
public void testGetSqlInfoSchema() {
final FlightInfo info = sqlClient.getSqlInfo();
- MatcherAssert.assertThat(
- info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)));
+ assertThat(info.getSchemaOptional())
+ .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA));
}
@Test
@@ -851,11 +838,11 @@ public void testGetSqlInfoResults() throws Exception {
try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA),
() ->
- MatcherAssert.assertThat(
- getNonConformingResultsForGetSqlInfo(getResults(stream)), is(emptyList())));
+ assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream)))
+ .isEqualTo(emptyList()));
}
}
@@ -866,11 +853,11 @@ public void testGetSqlInfoResultsWithSingleArg() throws Exception {
try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA),
() ->
- MatcherAssert.assertThat(
- getNonConformingResultsForGetSqlInfo(getResults(stream), arg), is(emptyList())));
+ assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream), arg))
+ .isEqualTo(emptyList()));
}
}
@@ -895,11 +882,11 @@ public void testGetSqlInfoResultsWithManyArgs() throws Exception {
try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA),
() ->
- MatcherAssert.assertThat(
- getNonConformingResultsForGetSqlInfo(getResults(stream), args), is(emptyList())));
+ assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream), args))
+ .isEqualTo(emptyList()));
}
}
@@ -915,28 +902,30 @@ public void testGetCommandExportedKeys() throws Exception {
final List> results = getResults(stream);
- final List> matchers =
+ final List> matchers =
asList(
- nullValue(String.class), // pk_catalog_name
- is("APP"), // pk_schema_name
- is("FOREIGNTABLE"), // pk_table_name
- is("ID"), // pk_column_name
- nullValue(String.class), // fk_catalog_name
- is("APP"), // fk_schema_name
- is("INTTABLE"), // fk_table_name
- is("FOREIGNID"), // fk_column_name
- is("1"), // key_sequence
- containsString("SQL"), // fk_key_name
- containsString("SQL"), // pk_key_name
- is("3"), // update_rule
- is("3")); // delete_rule
+ new Condition<>(Objects::isNull, "pk_catalog_name expected to be null"),
+ new Condition<>(c -> c.equals("APP"), "pk_schema_name expected to equal APP"),
+ new Condition<>(
+ c -> c.equals("FOREIGNTABLE"), "pk_table_name should equal FOREIGNTABLE"),
+ new Condition<>(c -> c.equals("ID"), "pk_column_name should equal ID"),
+ new Condition<>(Objects::isNull, "fk_catalog_name expected to be null"),
+ new Condition<>(c -> c.equals("APP"), "fk_schema_name expected to be APP"),
+ new Condition<>(c -> c.equals("INTTABLE"), "fk_table_name expeced to be INTTABLE"),
+ new Condition<>(
+ c -> c.equals("FOREIGNID"), "fk_column_name expected to equal FOREIGNID"),
+ new Condition<>(c -> c.equals("1"), "key_sequence expected to equal 1"),
+ new Condition<>(c -> c.contains("SQL"), "fk_key_name expected to contain SQL"),
+ new Condition<>(c -> c.contains("SQL"), "pk_key_name expected to contain SQL"),
+ new Condition<>(c -> c.equals("3"), "update_rule expected to equal 3"),
+ new Condition<>(c -> c.equals("3"), "delete_rule expected to equal 3"));
final List assertions = new ArrayList<>();
assertEquals(1, results.size());
for (int i = 0; i < matchers.size(); i++) {
final String actual = results.get(0).get(i);
- final Matcher expected = matchers.get(i);
- assertions.add(() -> MatcherAssert.assertThat(actual, expected));
+ final Condition expected = matchers.get(i);
+ assertions.add(() -> assertThat(actual).satisfies(expected));
}
assertAll(assertions);
}
@@ -954,28 +943,30 @@ public void testGetCommandImportedKeys() throws Exception {
final List> results = getResults(stream);
- final List> matchers =
+ final List> matchers =
asList(
- nullValue(String.class), // pk_catalog_name
- is("APP"), // pk_schema_name
- is("FOREIGNTABLE"), // pk_table_name
- is("ID"), // pk_column_name
- nullValue(String.class), // fk_catalog_name
- is("APP"), // fk_schema_name
- is("INTTABLE"), // fk_table_name
- is("FOREIGNID"), // fk_column_name
- is("1"), // key_sequence
- containsString("SQL"), // fk_key_name
- containsString("SQL"), // pk_key_name
- is("3"), // update_rule
- is("3")); // delete_rule
+ new Condition<>(Objects::isNull, "pk_catalog_name expected to be null"),
+ new Condition<>(c -> c.equals("APP"), "pk_schema_name expected to equal APP"),
+ new Condition<>(
+ c -> c.equals("FOREIGNTABLE"), "pk_table_name should equal FOREIGNTABLE"),
+ new Condition<>(c -> c.equals("ID"), "pk_column_name should equal ID"),
+ new Condition<>(Objects::isNull, "fk_catalog_name expected to be null"),
+ new Condition<>(c -> c.equals("APP"), "fk_schema_name expected to be APP"),
+ new Condition<>(c -> c.equals("INTTABLE"), "fk_table_name expeced to be INTTABLE"),
+ new Condition<>(
+ c -> c.equals("FOREIGNID"), "fk_column_name expected to equal FOREIGNID"),
+ new Condition<>(c -> c.equals("1"), "key_sequence expected to equal 1"),
+ new Condition<>(c -> c.contains("SQL"), "fk_key_name expected to contain SQL"),
+ new Condition<>(c -> c.contains("SQL"), "pk_key_name expected to contain SQL"),
+ new Condition<>(c -> c.equals("3"), "update_rule expected to equal 3"),
+ new Condition<>(c -> c.equals("3"), "delete_rule expected to equal 3"));
assertEquals(1, results.size());
final List assertions = new ArrayList<>();
for (int i = 0; i < matchers.size(); i++) {
final String actual = results.get(0).get(i);
- final Matcher expected = matchers.get(i);
- assertions.add(() -> MatcherAssert.assertThat(actual, expected));
+ final Condition expected = matchers.get(i);
+ assertions.add(() -> assertThat(actual).satisfies(expected));
}
assertAll(assertions);
}
@@ -1431,7 +1422,7 @@ public void testGetTypeInfo() throws Exception {
null,
null,
null));
- MatcherAssert.assertThat(results, is(matchers));
+ assertThat(results).isEqualTo(matchers);
}
}
@@ -1465,7 +1456,7 @@ public void testGetTypeInfoWithFiltering() throws Exception {
null,
"10",
null));
- MatcherAssert.assertThat(results, is(matchers));
+ assertThat(results).isEqualTo(matchers);
}
}
@@ -1479,28 +1470,30 @@ public void testGetCommandCrossReference() throws Exception {
final List> results = getResults(stream);
- final List> matchers =
+ final List> matchers =
asList(
- nullValue(String.class), // pk_catalog_name
- is("APP"), // pk_schema_name
- is("FOREIGNTABLE"), // pk_table_name
- is("ID"), // pk_column_name
- nullValue(String.class), // fk_catalog_name
- is("APP"), // fk_schema_name
- is("INTTABLE"), // fk_table_name
- is("FOREIGNID"), // fk_column_name
- is("1"), // key_sequence
- containsString("SQL"), // fk_key_name
- containsString("SQL"), // pk_key_name
- is("3"), // update_rule
- is("3")); // delete_rule
+ new Condition<>(Objects::isNull, "pk_catalog_name expected to be null"),
+ new Condition<>(c -> c.equals("APP"), "pk_schema_name expected to equal APP"),
+ new Condition<>(
+ c -> c.equals("FOREIGNTABLE"), "pk_table_name should equal FOREIGNTABLE"),
+ new Condition<>(c -> c.equals("ID"), "pk_column_name should equal ID"),
+ new Condition<>(Objects::isNull, "fk_catalog_name expected to be null"),
+ new Condition<>(c -> c.equals("APP"), "fk_schema_name expected to be APP"),
+ new Condition<>(c -> c.equals("INTTABLE"), "fk_table_name expeced to be INTTABLE"),
+ new Condition<>(
+ c -> c.equals("FOREIGNID"), "fk_column_name expected to equal FOREIGNID"),
+ new Condition<>(c -> c.equals("1"), "key_sequence expected to equal 1"),
+ new Condition<>(c -> c.contains("SQL"), "fk_key_name expected to contain SQL"),
+ new Condition<>(c -> c.contains("SQL"), "pk_key_name expected to contain SQL"),
+ new Condition<>(c -> c.equals("3"), "update_rule expected to equal 3"),
+ new Condition<>(c -> c.equals("3"), "delete_rule expected to equal 3"));
assertEquals(1, results.size());
final List assertions = new ArrayList<>();
for (int i = 0; i < matchers.size(); i++) {
final String actual = results.get(0).get(i);
- final Matcher expected = matchers.get(i);
- assertions.add(() -> MatcherAssert.assertThat(actual, expected));
+ final Condition expected = matchers.get(i);
+ assertions.add(() -> assertThat(actual).satisfies(expected));
}
assertAll(assertions);
}
@@ -1509,7 +1502,7 @@ public void testGetCommandCrossReference() throws Exception {
@Test
public void testCreateStatementSchema() throws Exception {
final FlightInfo info = sqlClient.execute("SELECT * FROM intTable");
- MatcherAssert.assertThat(info.getSchemaOptional(), is(Optional.of(SCHEMA_INT_TABLE)));
+ assertThat(info.getSchemaOptional()).isEqualTo(Optional.of(SCHEMA_INT_TABLE));
// Consume statement to close connection before cache eviction
try (FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) {
@@ -1526,11 +1519,10 @@ public void testCreateStatementResults() throws Exception {
sqlClient.execute("SELECT * FROM intTable").getEndpoints().get(0).getTicket())) {
assertAll(
() -> {
- MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE));
+ assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE);
},
() -> {
- MatcherAssert.assertThat(
- getResults(stream), is(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY));
+ assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY);
});
}
}
@@ -1543,19 +1535,19 @@ public void testExecuteUpdate() {
sqlClient.executeUpdate(
"INSERT INTO INTTABLE (keyName, value) VALUES "
+ "('KEYNAME1', 1001), ('KEYNAME2', 1002), ('KEYNAME3', 1003)");
- MatcherAssert.assertThat(insertedCount, is(3L));
+ assertThat(insertedCount).isEqualTo(3L);
},
() -> {
long updatedCount =
sqlClient.executeUpdate(
"UPDATE INTTABLE SET keyName = 'KEYNAME1' "
+ "WHERE keyName = 'KEYNAME2' OR keyName = 'KEYNAME3'");
- MatcherAssert.assertThat(updatedCount, is(2L));
+ assertThat(updatedCount).isEqualTo(2L);
},
() -> {
long deletedCount =
sqlClient.executeUpdate("DELETE FROM INTTABLE WHERE keyName = 'KEYNAME1'");
- MatcherAssert.assertThat(deletedCount, is(3L));
+ assertThat(deletedCount).isEqualTo(3L);
});
}
@@ -1566,10 +1558,10 @@ public void testQueryWithNoResultsShouldNotHang() throws Exception {
final FlightStream stream =
sqlClient.getStream(preparedStatement.execute().getEndpoints().get(0).getTicket())) {
assertAll(
- () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)),
+ () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE),
() -> {
final List> result = getResults(stream);
- MatcherAssert.assertThat(result, is(emptyList()));
+ assertThat(result).isEqualTo(emptyList());
});
}
}
diff --git a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java
index 36d621ad64..ee1507b6af 100644
--- a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java
+++ b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java
@@ -18,7 +18,7 @@
import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults;
import static org.apache.arrow.util.AutoCloseables.close;
-import static org.hamcrest.CoreMatchers.is;
+import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertAll;
import org.apache.arrow.flight.FlightClient;
@@ -34,7 +34,6 @@
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Schema;
-import org.hamcrest.MatcherAssert;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -89,10 +88,10 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce
for (FlightEndpoint endpoint : flightInfo.getEndpoints()) {
try (FlightStream stream = sqlClient.getStream(endpoint.getTicket())) {
assertAll(
- () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)),
+ () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE),
() ->
- MatcherAssert.assertThat(
- getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)));
+ assertThat(getResults(stream))
+ .isEqualTo(EXPECTED_RESULTS_FOR_PARAMETER_BINDING));
}
}
}
diff --git a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java
index 71c0dc88e4..3f527f961e 100644
--- a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java
+++ b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java
@@ -22,7 +22,7 @@
import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults;
import static org.apache.arrow.util.AutoCloseables.close;
import static org.apache.arrow.vector.types.Types.MinorType.INT;
-import static org.hamcrest.CoreMatchers.is;
+import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertAll;
import com.google.common.collect.ImmutableList;
@@ -53,7 +53,6 @@
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.Text;
-import org.hamcrest.MatcherAssert;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -245,15 +244,15 @@ public void testGetTablesResultNoSchema() throws Exception {
sqlClient.getTables(null, null, null, null, false).getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA),
() -> {
final List> results = getResults(stream);
final List> expectedResults =
ImmutableList.of(
// catalog_name | schema_name | table_name | table_type | table_schema
asList(null, null, "test_table", "TABLE"));
- MatcherAssert.assertThat(results, is(expectedResults));
+ assertThat(results).isEqualTo(expectedResults);
});
}
}
@@ -264,15 +263,15 @@ public void testGetTableTypesResult() throws Exception {
sqlClient.getStream(sqlClient.getTableTypes().getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA)),
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA),
() -> {
final List> tableTypes = getResults(stream);
final List> expectedTableTypes =
ImmutableList.of(
// table_type
singletonList("TABLE"));
- MatcherAssert.assertThat(tableTypes, is(expectedTableTypes));
+ assertThat(tableTypes).isEqualTo(expectedTableTypes);
});
}
}
@@ -283,9 +282,9 @@ public void testGetSqlInfoResults() throws Exception {
try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) {
assertAll(
() ->
- MatcherAssert.assertThat(
- stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)),
- () -> MatcherAssert.assertThat(getResults(stream), is(emptyList())));
+ assertThat(stream.getSchema())
+ .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA),
+ () -> assertThat(getResults(stream)).isEqualTo(emptyList()));
}
}
@@ -303,7 +302,7 @@ public void testGetTypeInfo() throws Exception {
"Integer", "4", "400", null, null, "3", "true", null, "true", null, "true",
"Integer", null, null, "4", null, "10", null));
- MatcherAssert.assertThat(results, is(matchers));
+ assertThat(results).isEqualTo(matchers);
}
}
@@ -317,10 +316,8 @@ public void testExecuteQuery() throws Exception {
.get(0)
.getTicket())) {
assertAll(
- () ->
- MatcherAssert.assertThat(stream.getSchema(), is(FlightSqlTestProducer.FIXED_SCHEMA)),
- () ->
- MatcherAssert.assertThat(getResults(stream), is(singletonList(singletonList("1")))));
+ () -> assertThat(stream.getSchema()).isEqualTo(FlightSqlTestProducer.FIXED_SCHEMA),
+ () -> assertThat(getResults(stream)).isEqualTo(singletonList(singletonList("1"))));
}
}
}
diff --git a/flight/pom.xml b/flight/pom.xml
index 7b31e8ce91..a5a40a834a 100644
--- a/flight/pom.xml
+++ b/flight/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-flight
diff --git a/format/pom.xml b/format/pom.xml
index 9b4eebfe3c..c09fad32fb 100644
--- a/format/pom.xml
+++ b/format/pom.xml
@@ -23,7 +23,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-format
diff --git a/gandiva/pom.xml b/gandiva/pom.xml
index 167bf39cb9..d26edeb9d6 100644
--- a/gandiva/pom.xml
+++ b/gandiva/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
org.apache.arrow.gandiva
diff --git a/memory/memory-core/pom.xml b/memory/memory-core/pom.xml
index 840d3464ba..586fddae1a 100644
--- a/memory/memory-core/pom.xml
+++ b/memory/memory-core/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 18.3.0
+ 19.0.0
arrow-memory-core
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
index 5d052c2cde..d4d76f57f4 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
@@ -16,7 +16,7 @@
*/
package org.apache.arrow.memory;
-import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicLongFieldUpdater;
import org.apache.arrow.util.Preconditions;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -37,16 +37,24 @@ class Accountant implements AutoCloseable {
*/
protected final long reservation;
- private final AtomicLong peakAllocation = new AtomicLong();
+ // AtomicLongFieldUpdaters for memory accounting fields to reduce memory overhead
+ private static final AtomicLongFieldUpdater PEAK_ALLOCATION_UPDATER =
+ AtomicLongFieldUpdater.newUpdater(Accountant.class, "peakAllocation");
+ private static final AtomicLongFieldUpdater ALLOCATION_LIMIT_UPDATER =
+ AtomicLongFieldUpdater.newUpdater(Accountant.class, "allocationLimit");
+ private static final AtomicLongFieldUpdater LOCALLY_HELD_MEMORY_UPDATER =
+ AtomicLongFieldUpdater.newUpdater(Accountant.class, "locallyHeldMemory");
+
+ private volatile long peakAllocation = 0;
/**
* Maximum local memory that can be held. This can be externally updated. Changing it won't cause
* past memory to change but will change responses to future allocation efforts
*/
- private final AtomicLong allocationLimit = new AtomicLong();
+ private volatile long allocationLimit = 0;
/** Currently allocated amount of memory. */
- private final AtomicLong locallyHeldMemory = new AtomicLong();
+ private volatile long locallyHeldMemory = 0;
public Accountant(
@Nullable Accountant parent, String name, long reservation, long maxAllocation) {
@@ -64,7 +72,7 @@ public Accountant(
this.parent = parent;
this.name = name;
this.reservation = reservation;
- this.allocationLimit.set(maxAllocation);
+ ALLOCATION_LIMIT_UPDATER.set(this, maxAllocation);
if (reservation != 0) {
Preconditions.checkArgument(parent != null, "parent must not be null");
@@ -117,12 +125,12 @@ private AllocationOutcome.Status allocateBytesInternal(long size) {
}
private void updatePeak() {
- final long currentMemory = locallyHeldMemory.get();
+ final long currentMemory = locallyHeldMemory;
while (true) {
- final long previousPeak = peakAllocation.get();
+ final long previousPeak = peakAllocation;
if (currentMemory > previousPeak) {
- if (!peakAllocation.compareAndSet(previousPeak, currentMemory)) {
+ if (!PEAK_ALLOCATION_UPDATER.compareAndSet(this, previousPeak, currentMemory)) {
// peak allocation changed underneath us. try again.
continue;
}
@@ -166,7 +174,7 @@ private AllocationOutcome.Status allocate(
final boolean incomingUpdatePeak,
final boolean forceAllocation,
@Nullable AllocationOutcomeDetails details) {
- final long oldLocal = locallyHeldMemory.getAndAdd(size);
+ final long oldLocal = LOCALLY_HELD_MEMORY_UPDATER.getAndAdd(this, size);
final long newLocal = oldLocal + size;
// Borrowed from Math.addExact (but avoid exception here)
// Overflow if result has opposite sign of both arguments
@@ -174,7 +182,7 @@ private AllocationOutcome.Status allocate(
// failure
final boolean overflow = ((oldLocal ^ newLocal) & (size ^ newLocal)) < 0;
final long beyondReservation = newLocal - reservation;
- final boolean beyondLimit = overflow || newLocal > allocationLimit.get();
+ final boolean beyondLimit = overflow || newLocal > allocationLimit;
final boolean updatePeak = forceAllocation || (incomingUpdatePeak && !beyondLimit);
if (details != null) {
@@ -214,7 +222,7 @@ private AllocationOutcome.Status allocate(
public void releaseBytes(long size) {
// reduce local memory. all memory released above reservation should be released up the tree.
- final long newSize = locallyHeldMemory.addAndGet(-size);
+ final long newSize = LOCALLY_HELD_MEMORY_UPDATER.addAndGet(this, -size);
Preconditions.checkArgument(newSize >= 0, "Accounted size went negative.");
@@ -255,7 +263,7 @@ public String getName() {
* @return Limit in bytes.
*/
public long getLimit() {
- return allocationLimit.get();
+ return allocationLimit;
}
/**
@@ -274,7 +282,7 @@ public long getInitReservation() {
* @param newLimit The limit in bytes.
*/
public void setLimit(long newLimit) {
- allocationLimit.set(newLimit);
+ ALLOCATION_LIMIT_UPDATER.set(this, newLimit);
}
/**
@@ -284,7 +292,7 @@ public void setLimit(long newLimit) {
* @return Currently allocate memory in bytes.
*/
public long getAllocatedMemory() {
- return locallyHeldMemory.get();
+ return locallyHeldMemory;
}
/**
@@ -293,17 +301,17 @@ public long getAllocatedMemory() {
* @return The peak allocated memory in bytes.
*/
public long getPeakMemoryAllocation() {
- return peakAllocation.get();
+ return peakAllocation;
}
public long getHeadroom() {
- long localHeadroom = allocationLimit.get() - locallyHeldMemory.get();
+ long localHeadroom = allocationLimit - locallyHeldMemory;
if (parent == null) {
return localHeadroom;
}
// Amount of reserved memory left on top of what parent has
- long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory.get());
+ long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory);
return Math.min(localHeadroom, parent.getHeadroom() + reservedHeadroom);
}
}
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java
index 775a8925ad..9712be34d7 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java
@@ -24,7 +24,6 @@
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.ReadOnlyBufferException;
-import java.util.concurrent.atomic.AtomicLong;
import org.apache.arrow.memory.BaseAllocator.Verbosity;
import org.apache.arrow.memory.util.CommonUtil;
import org.apache.arrow.memory.util.HistoricalLog;
@@ -57,9 +56,8 @@ public final class ArrowBuf implements AutoCloseable {
private static final int DOUBLE_SIZE = Double.BYTES;
private static final int LONG_SIZE = Long.BYTES;
- private static final AtomicLong idGenerator = new AtomicLong(0);
private static final int LOG_BYTES_PER_ROW = 10;
- private final long id = idGenerator.incrementAndGet();
+
private final ReferenceManager referenceManager;
private final @Nullable BufferManager bufferManager;
private final long addr;
@@ -67,7 +65,8 @@ public final class ArrowBuf implements AutoCloseable {
private long writerIndex;
private final @Nullable HistoricalLog historicalLog =
BaseAllocator.DEBUG
- ? new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", id)
+ ? new HistoricalLog(
+ BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", System.identityHashCode(this))
: null;
private volatile long capacity;
@@ -136,7 +135,7 @@ public long capacity() {
/**
* Adjusts the capacity of this buffer. Size increases are NOT supported.
*
- * @param newCapacity Must be in in the range [0, length).
+ * @param newCapacity Must be in the range [0, length).
*/
public synchronized ArrowBuf capacity(long newCapacity) {
@@ -218,7 +217,8 @@ public long memoryAddress() {
@Override
public String toString() {
- return String.format("ArrowBuf[%d], address:%d, capacity:%d", id, memoryAddress(), capacity);
+ return String.format(
+ "ArrowBuf[%d], address:%d, capacity:%d", getId(), memoryAddress(), capacity);
}
@Override
@@ -1080,12 +1080,15 @@ public String toHexString(final long start, final int length) {
}
/**
- * Get the integer id assigned to this ArrowBuf for debugging purposes.
+ * Get the id assigned to this ArrowBuf for debugging purposes.
+ *
+ * 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, ReferenceManager {
private final @Nullable IdentityHashMap buffers =
BaseAllocator.DEBUG ? new IdentityHashMap<>() : null;
- private static final AtomicLong LEDGER_ID_GENERATOR = new AtomicLong(0);
- // unique ID assigned to each ledger
- private final long ledgerId = LEDGER_ID_GENERATOR.incrementAndGet();
- private final AtomicInteger bufRefCnt = new AtomicInteger(0); // start at zero so we can
- // manage request for retain
- // correctly
+
+ // AtomicIntegerFieldUpdater for bufRefCnt to reduce memory overhead
+ private static final AtomicIntegerFieldUpdater BUF_REF_CNT_UPDATER =
+ AtomicIntegerFieldUpdater.newUpdater(BufferLedger.class, "bufRefCnt");
+ // start at zero so we can manage request for retain correctly
+ private volatile int bufRefCnt = 0;
+
private final long lCreationTime = System.nanoTime();
private final BufferAllocator allocator;
private final AllocationManager allocationManager;
@@ -78,7 +78,7 @@ public BufferAllocator getAllocator() {
*/
@Override
public int getRefCount() {
- return bufRefCnt.get();
+ return bufRefCnt;
}
/**
@@ -86,7 +86,7 @@ public int getRefCount() {
* ArrowBufs managed by this ledger will share the ref count.
*/
void increment() {
- bufRefCnt.incrementAndGet();
+ BUF_REF_CNT_UPDATER.incrementAndGet(this);
}
/**
@@ -144,7 +144,7 @@ private int decrement(int decrement) {
allocator.assertOpen();
final int outcome;
synchronized (allocationManager) {
- outcome = bufRefCnt.addAndGet(-decrement);
+ outcome = BUF_REF_CNT_UPDATER.addAndGet(this, -decrement);
if (outcome == 0) {
lDestructionTime = System.nanoTime();
// refcount of this reference manager has dropped to 0
@@ -174,7 +174,7 @@ public void retain(int increment) {
if (historicalLog != null) {
historicalLog.recordEvent("retain(%d)", increment);
}
- final int originalReferenceCount = bufRefCnt.getAndAdd(increment);
+ final int originalReferenceCount = BUF_REF_CNT_UPDATER.getAndAdd(this, increment);
Preconditions.checkArgument(originalReferenceCount > 0);
}
@@ -472,13 +472,13 @@ public long getAccountedSize() {
void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) {
CommonUtil.indent(sb, indent)
.append("ledger[")
- .append(ledgerId)
+ .append(System.identityHashCode(this))
.append("] allocator: ")
.append(allocator.getName())
.append("), isOwning: ")
.append(", size: ")
.append(", references: ")
- .append(bufRefCnt.get())
+ .append(bufRefCnt)
.append(", life: ")
.append(lCreationTime)
.append("..")
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
index 91bd7cd905..be0749a215 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
@@ -18,6 +18,7 @@
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
+import java.lang.reflect.InaccessibleObjectException;
import java.lang.reflect.InvocationTargetException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
@@ -81,9 +82,18 @@ public Object run() {
BYTE_ARRAY_BASE_OFFSET = UNSAFE.arrayBaseOffset(byte[].class);
// get the offset of the address field in a java.nio.Buffer object
+ long maybeOffset;
Field addressField = java.nio.Buffer.class.getDeclaredField("address");
- addressField.setAccessible(true);
- BYTE_BUFFER_ADDRESS_OFFSET = UNSAFE.objectFieldOffset(addressField);
+ try {
+ addressField.setAccessible(true);
+ maybeOffset = UNSAFE.objectFieldOffset(addressField);
+ } catch (InaccessibleObjectException e) {
+ maybeOffset = -1;
+ logger.debug(
+ "Cannot access the address field of java.nio.Buffer. DirectBuffer operations wont be available",
+ e);
+ }
+ BYTE_BUFFER_ADDRESS_OFFSET = maybeOffset;
Constructor directBufferConstructor;
long address = -1;
@@ -109,6 +119,9 @@ public Object run() {
} catch (SecurityException e) {
logger.debug("Cannot get constructor for direct buffer allocation", e);
return e;
+ } catch (InaccessibleObjectException e) {
+ logger.debug("Cannot get constructor for direct buffer allocation", e);
+ return e;
}
}
});
@@ -156,7 +169,11 @@ public Object run() {
* @return address of the underlying memory.
*/
public static long getByteBufferAddress(ByteBuffer buf) {
- return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET);
+ if (BYTE_BUFFER_ADDRESS_OFFSET != -1) {
+ return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET);
+ }
+ throw new UnsupportedOperationException(
+ "Byte buffer address cannot be obtained because sun.misc.Unsafe or java.nio.DirectByteBuffer.(long, int) is not available");
}
private MemoryUtil() {}
diff --git a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
index a39004a9d0..ba5a539a87 100644
--- a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
+++ b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java
@@ -22,7 +22,9 @@
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
+import java.util.stream.Stream;
import java.util.stream.StreamSupport;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Utilities for AutoCloseable classes. */
public final class AutoCloseables {
@@ -33,7 +35,8 @@ private AutoCloseables() {}
* Returns a new {@link AutoCloseable} that calls {@link #close(Iterable)} on autoCloseables
* 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 Iterable
// like for example VectorContainer make sure we handle that properly
if (ac == null) {
@@ -111,12 +122,17 @@ public static void close(Iterable ac) throws Exception
/** Calls {@link #close(Iterable)} on the flattened list of closeables. */
@SafeVarargs
- public static void close(Iterable... closeables) throws Exception {
+ public static void close(@Nullable Iterable... closeables)
+ throws Exception {
+ if (closeables == null) {
+ return;
+ }
close(flatten(closeables));
}
@SafeVarargs
- private static Iterable flatten(Iterable... closeables) {
+ private static Iterable flatten(
+ Iterable... closeables) {
return new Iterable() {
// Cast from Iterable to Iterable is safe in this
// context
@@ -127,16 +143,18 @@ public Iterator iterator() {
return Arrays.stream(closeables)
.flatMap(
(Iterable i) ->
- StreamSupport.stream(
- ((Iterable) i).spliterator(), /* parallel= */ false))
+ i == null
+ ? Stream.empty()
+ : StreamSupport.stream(
+ ((Iterable) i).spliterator(), /* parallel= */ false))
.iterator();
}
};
}
/** Converts ac to a {@link Iterable} filtering out any null values. */
- public static Iterable iter(AutoCloseable... ac) {
- if (ac.length == 0) {
+ public static Iterable iter(@Nullable AutoCloseable... ac) {
+ if (ac == null || ac.length == 0) {
return Collections.emptyList();
} else {
final List nonNullAc = new ArrayList<>();
@@ -153,10 +171,11 @@ public static Iterable iter(AutoCloseable... ac) {
public static class RollbackCloseable implements AutoCloseable {
private boolean commit = false;
- private List closeables;
+ private final List closeables;
- public RollbackCloseable(AutoCloseable... closeables) {
- this.closeables = new ArrayList<>(Arrays.asList(closeables));
+ public RollbackCloseable(@Nullable AutoCloseable... closeables) {
+ this.closeables =
+ closeables == null ? new ArrayList<>() : new ArrayList<>(Arrays.asList(closeables));
}
public T add(T t) {
@@ -165,12 +184,18 @@ public T add(T t) {
}
/** Add all of list 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.(long, int) not available")) {
found = true;
break;
}
+ e = e.getCause();
}
assertTrue(found, "Expected exception was not thrown");
}
diff --git a/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java
new file mode 100644
index 0000000000..ba5b78178a
--- /dev/null
+++ b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java
@@ -0,0 +1,268 @@
+/*
+ * 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.util;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+public class TestAutoCloseables {
+
+ /** Closeable that records that it was closed and can optionally throw. */
+ private static final class TrackCloseable implements AutoCloseable {
+ private boolean closed;
+ private final Exception toThrow;
+
+ TrackCloseable() {
+ this.toThrow = null;
+ }
+
+ TrackCloseable(Exception toThrow) {
+ this.toThrow = toThrow;
+ }
+
+ @Override
+ public void close() throws Exception {
+ closed = true;
+ if (toThrow != null) {
+ throw toThrow;
+ }
+ }
+
+ boolean isClosed() {
+ return closed;
+ }
+ }
+
+ @Test
+ public void testCloseVarargsIgnoresNulls() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.close(a, null, b);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testCloseVarargsThrowsFirstExceptionAndSuppressesRest() throws Exception {
+ Exception e1 = new Exception("first");
+ Exception e2 = new Exception("second");
+ TrackCloseable c1 = new TrackCloseable(e1);
+ TrackCloseable c2 = new TrackCloseable(e2);
+ Exception thrown = assertThrows(Exception.class, () -> AutoCloseables.close(c1, c2));
+ assertEquals("first", thrown.getMessage());
+ assertTrue(Arrays.asList(thrown.getSuppressed()).contains(e2));
+ }
+
+ @Test
+ public void testCloseIterableNullIterableReturns() throws Exception {
+ AutoCloseables.close((List) null); // no exception
+ }
+
+ @Test
+ public void testCloseIterableIgnoresNullElements() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ List list = Arrays.asList(a, null, b);
+ AutoCloseables.close(list);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testCloseIterableWhenIterableIsAlsoAutoCloseable() throws Exception {
+ TrackCloseable iter = new TrackCloseable();
+ TrackCloseable inner = new TrackCloseable();
+ // When the Iterable itself implements AutoCloseable (e.g. VectorContainer),
+ // close(Iterable) calls close() on it and does not iterate over elements
+ class IterableCloseable implements Iterable, AutoCloseable {
+ @Override
+ @SuppressWarnings("unchecked")
+ public Iterator iterator() {
+ return (Iterator) Collections.singletonList(inner);
+ }
+
+ @Override
+ public void close() throws Exception {
+ iter.close();
+ }
+ }
+ AutoCloseables.close(new IterableCloseable());
+ assertTrue(iter.isClosed());
+ assertFalse(inner.isClosed());
+ }
+
+ @Test
+ public void testCloseIterableVarargsWithNullIterables() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ TrackCloseable c = new TrackCloseable();
+ List list1 = Arrays.asList(null, a, b);
+ List list2 = Collections.singletonList(c);
+ AutoCloseables.close(list1, null, list2);
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ assertTrue(c.isClosed());
+ }
+
+ @Test
+ public void testCloseThrowableSuppressesException() {
+ Exception e = new Exception("from close");
+ TrackCloseable c = new TrackCloseable(e);
+ Exception main = new Exception("main");
+ AutoCloseables.close(main, c);
+ assertTrue(c.isClosed());
+ assertEquals(1, main.getSuppressed().length);
+ assertEquals(e, main.getSuppressed()[0]);
+ }
+
+ @Test
+ public void testCloseThrowableWithNullCloseables() {
+ Exception main = new Exception("main");
+ AutoCloseables.close(main, (AutoCloseable) null);
+ assertEquals(0, main.getSuppressed().length);
+
+ AutoCloseables.close(main, (AutoCloseable[]) null); // no exception
+ }
+
+ @Test
+ public void testIterFiltersNulls() {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ Iterable it = AutoCloseables.iter(a, null, b);
+ List list = new ArrayList<>();
+ it.forEach(list::add);
+ assertEquals(2, list.size());
+ assertTrue(list.contains(a));
+ assertTrue(list.contains(b));
+ }
+
+ @Test
+ public void testIterEmptyVarargs() {
+ Iterable it = AutoCloseables.iter();
+ List list = new ArrayList<>();
+ it.forEach(list::add);
+ assertTrue(list.isEmpty());
+ }
+
+ @Test
+ public void testIterWithNull() {
+ AutoCloseables.iter((AutoCloseable) null); // no exception
+ }
+
+ @Test
+ public void testCloseNoCheckedWithNull() {
+ AutoCloseables.closeNoChecked(null); // no exception
+ }
+
+ @Test
+ public void testCloseNoCheckedWrapsException() {
+ Exception e = new Exception("close failed");
+ TrackCloseable c = new TrackCloseable(e);
+ RuntimeException re =
+ assertThrows(RuntimeException.class, () -> AutoCloseables.closeNoChecked(c));
+ assertSame(re.getCause(), e);
+ assertTrue(re.getMessage().contains("close failed"));
+ }
+
+ @Test
+ public void testNoop() throws Exception {
+ AutoCloseable noop = AutoCloseables.noop();
+ assertSame(noop, AutoCloseables.noop());
+ noop.close(); // no exception
+ }
+
+ @Test
+ public void testAllClosesCollectionOnClose() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ List list = Arrays.asList(a, b);
+ AutoCloseable all = AutoCloseables.all(list);
+ assertFalse(a.isClosed());
+ assertFalse(b.isClosed());
+ all.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testAllWithNullCollection() throws Exception {
+ AutoCloseable all = AutoCloseables.all(null);
+ all.close(); // no exception
+ }
+
+ @Test
+ public void testRollbackCloseableClosesWhenNotCommitted() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b);
+ rb.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableDoesNotCloseWhenCommitted() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b);
+ rb.commit();
+ rb.close();
+ assertFalse(a.isClosed());
+ assertFalse(b.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableAddAndAddAll() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ TrackCloseable b = new TrackCloseable();
+ TrackCloseable c = new TrackCloseable();
+ TrackCloseable d = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a);
+ rb.add(b);
+ rb.addAll(c, d);
+ rb.addAll((AutoCloseable[]) null); // null varargs shouldn't fail
+ rb.addAll((List) null); // null Iterable shouldn't fail
+ rb.close();
+ assertTrue(a.isClosed());
+ assertTrue(b.isClosed());
+ assertTrue(c.isClosed());
+ assertTrue(d.isClosed());
+ }
+
+ @Test
+ public void testRollbackCloseableWithNull() throws Exception {
+ AutoCloseables.rollbackable((AutoCloseable) null); // no exception
+ }
+
+ @Test
+ public void testRollbackCloseableWithNulls() throws Exception {
+ TrackCloseable a = new TrackCloseable();
+ AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, null);
+ rb.close();
+ assertTrue(a.isClosed());
+ }
+}
diff --git a/memory/memory-netty-buffer-patch/pom.xml b/memory/memory-netty-buffer-patch/pom.xml
index e9a63b2122..cb38efa345 100644
--- a/memory/memory-netty-buffer-patch/pom.xml
+++ b/memory/memory-netty-buffer-patch/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 18.3.0
+ 19.0.0
arrow-memory-netty-buffer-patch
diff --git a/memory/memory-netty/pom.xml b/memory/memory-netty/pom.xml
index 42f35efb33..f33eb95e44 100644
--- a/memory/memory-netty/pom.xml
+++ b/memory/memory-netty/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 18.3.0
+ 19.0.0
arrow-memory-netty
diff --git a/memory/memory-unsafe/pom.xml b/memory/memory-unsafe/pom.xml
index 0af306cbfc..d941fee645 100644
--- a/memory/memory-unsafe/pom.xml
+++ b/memory/memory-unsafe/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-memory
- 18.3.0
+ 19.0.0
arrow-memory-unsafe
diff --git a/memory/pom.xml b/memory/pom.xml
index 09a5bc2924..af953ccd21 100644
--- a/memory/pom.xml
+++ b/memory/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-memory
pom
diff --git a/performance/pom.xml b/performance/pom.xml
index 02bdf46a11..685f433f05 100644
--- a/performance/pom.xml
+++ b/performance/pom.xml
@@ -22,7 +22,7 @@ under the License.
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
arrow-performance
jar
diff --git a/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java b/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java
new file mode 100644
index 0000000000..395ba13b9d
--- /dev/null
+++ b/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java
@@ -0,0 +1,213 @@
+/*
+ * 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.memory;
+
+import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryMXBean;
+import java.lang.management.MemoryUsage;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+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.annotations.Warmup;
+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 memory footprint of Arrow memory objects.
+ *
+ * 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:
+ *
+ *
{@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 de24b12e73..0b0aa58232 100644
--- a/pom.xml
+++ b/pom.xml
@@ -23,12 +23,12 @@ under the License.
org.apache
apache
- 34
+ 35
org.apache.arrow
arrow-java-root
- 18.3.0
+ 19.0.0
pom
Apache Arrow Java Root POM
@@ -68,6 +68,7 @@ under the License.
bom
format
memory
+ arrow-variant
vector
tools
adapter/jdbc
@@ -81,7 +82,7 @@ under the License.
scm:git:https://github.com/apache/arrow-java.git
scm:git:https://github.com/apache/arrow-java.git
- v18.3.0
+ v19.0.0
https://github.com/apache/arrow-java/tree/${project.scm.tag}
@@ -91,26 +92,28 @@ under the License.
+ 1773307790
${project.build.directory}/generated-sources
1.9.0
5.12.2
2.0.17
33.4.8-jre
- 4.1.119.Final
- 1.71.0
- 4.30.2
- 2.18.3
- 3.4.1
+ 4.2.9.Final
+ 1.79.0
+ 4.33.4
+ 2.21.0
+ 3.4.3
25.2.10
- 1.12.0
- 5.17.0
+ 1.12.1
+ 1.17.0
+ 5.21.0
2
10.23.0
true
- 2.37.0
- 3.49.2
- 1.5.18
+ 2.42.0
+ 3.53.1
+ 1.5.32
none
-Xdoclint:none
@@ -123,6 +126,8 @@ under the License.
3.2.2
@@ -172,13 +177,13 @@ under the License.
org.assertj
assertj-core
- 3.27.3
+ 3.27.7
test
org.immutables
value-annotations
- 2.10.1
+ 2.12.1
provided
@@ -309,7 +314,7 @@ under the License.
org.immutables
value
- 2.10.1
+ 2.12.1
@@ -347,7 +352,7 @@ under the License.
org.jacoco
jacoco-maven-plugin
- 0.8.13
+ 0.8.14
+ codegen
+ ${basedir}/src/main/codegen
+
+ **/*.tdd
+ **/*.fmpp
+ **/*.ftl
+
+
+
+
+
diff --git a/vector/src/main/codegen/includes/vv_imports.ftl b/vector/src/main/codegen/includes/vv_imports.ftl
index 7f216a7b43..2bbcecc856 100644
--- a/vector/src/main/codegen/includes/vv_imports.ftl
+++ b/vector/src/main/codegen/includes/vv_imports.ftl
@@ -34,6 +34,7 @@ import org.apache.arrow.vector.complex.*;
import org.apache.arrow.vector.complex.reader.*;
import org.apache.arrow.vector.complex.impl.*;
import org.apache.arrow.vector.complex.writer.*;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
diff --git a/vector/src/main/codegen/templates/AbstractFieldReader.java b/vector/src/main/codegen/templates/AbstractFieldReader.java
index 25b071fab7..789295e959 100644
--- a/vector/src/main/codegen/templates/AbstractFieldReader.java
+++ b/vector/src/main/codegen/templates/AbstractFieldReader.java
@@ -29,9 +29,9 @@
* Source code generated using FreeMarker template ${.template_name}
*/
@SuppressWarnings("unused")
-abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{
+public abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{
- AbstractFieldReader(){
+ protected AbstractFieldReader(){
super();
}
@@ -108,6 +108,23 @@ public void copyAsField(String name, ${name}Writer writer) {
}
+
+ public void read(ExtensionHolder holder) {
+ fail("Extension");
+ }
+
+ public void read(int arrayIndex, ExtensionHolder holder) {
+ fail("RepeatedExtension");
+ }
+
+ public void copyAsValue(AbstractExtensionTypeWriter writer) {
+ fail("CopyAsValueExtension");
+ }
+
+ public void copyAsField(String name, AbstractExtensionTypeWriter writer) {
+ fail("CopyAsFieldExtension");
+ }
+
public FieldReader reader(String name) {
fail("reader(String name)");
return null;
@@ -126,4 +143,5 @@ public int size() {
private void fail(String name) {
throw new IllegalArgumentException(String.format("You tried to read a [%s] type when you are using a field reader of type [%s].", name, this.getClass().getSimpleName()));
}
+
}
diff --git a/vector/src/main/codegen/templates/AbstractFieldWriter.java b/vector/src/main/codegen/templates/AbstractFieldWriter.java
index ae5b97faef..4b4a17d932 100644
--- a/vector/src/main/codegen/templates/AbstractFieldWriter.java
+++ b/vector/src/main/codegen/templates/AbstractFieldWriter.java
@@ -107,14 +107,17 @@ public void endEntry() {
throw new IllegalStateException(String.format("You tried to end a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
}
+ @Override
public void write(ExtensionHolder var1) {
- this.fail("ExtensionType");
+ this.fail("Cannot write ExtensionHolder");
}
+ @Override
public void writeExtension(Object var1) {
- this.fail("ExtensionType");
+ this.fail("Cannot write extension object");
}
- public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
- this.fail("ExtensionType");
+ @Override
+ public void writeExtension(Object var1, ArrowType type) {
+ this.fail("Cannot write extension with type " + type);
}
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
diff --git a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java
index 951edd5eee..2e7792fcfe 100644
--- a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java
+++ b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java
@@ -295,7 +295,7 @@ public MapWriter map(boolean keysSorted) {
@Override
public ExtensionWriter extension(ArrowType arrowType) {
- return getWriter(MinorType.EXTENSIONTYPE).extension(arrowType);
+ return getWriter(MinorType.LIST).extension(arrowType);
}
@Override
@@ -325,7 +325,7 @@ public MapWriter map(String name, boolean keysSorted) {
@Override
public ExtensionWriter extension(String name, ArrowType arrowType) {
- return getWriter(MinorType.EXTENSIONTYPE).extension(name, arrowType);
+ return getWriter(MinorType.STRUCT).extension(name, arrowType);
}
<#list vv.types as type><#list type.minor as minor>
diff --git a/vector/src/main/codegen/templates/ArrowType.java b/vector/src/main/codegen/templates/ArrowType.java
index fd35c1cd2b..b428f09155 100644
--- a/vector/src/main/codegen/templates/ArrowType.java
+++ b/vector/src/main/codegen/templates/ArrowType.java
@@ -27,8 +27,10 @@
import org.apache.arrow.flatbuf.Type;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.complex.writer.FieldWriter;
import org.apache.arrow.vector.types.*;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ValueVector;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -331,6 +333,10 @@ public boolean equals(Object obj) {
public T accept(ArrowTypeVisitor visitor) {
return visitor.visit(this);
}
+
+ public FieldWriter getNewFieldWriter(ValueVector vector) {
+ throw new UnsupportedOperationException("WriterImpl not yet implemented.");
+ }
}
private static final int defaultDecimalBitWidth = 128;
diff --git a/vector/src/main/codegen/templates/BaseReader.java b/vector/src/main/codegen/templates/BaseReader.java
index e75e8a2974..c52345af21 100644
--- a/vector/src/main/codegen/templates/BaseReader.java
+++ b/vector/src/main/codegen/templates/BaseReader.java
@@ -73,7 +73,7 @@ public interface RepeatedMapReader extends MapReader{
public interface ScalarReader extends
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Reader,
- BaseReader {}
+ ExtensionReader, BaseReader {}
interface ComplexReader{
StructReader rootAsStruct();
diff --git a/vector/src/main/codegen/templates/BaseWriter.java b/vector/src/main/codegen/templates/BaseWriter.java
index 78da7fddc3..a4c98d7089 100644
--- a/vector/src/main/codegen/templates/BaseWriter.java
+++ b/vector/src/main/codegen/templates/BaseWriter.java
@@ -125,11 +125,12 @@ public interface ExtensionWriter extends BaseWriter {
void writeExtension(Object value);
/**
- * Adds the given extension type factory. This factory allows configuring writer implementations for specific ExtensionTypeVector.
+ * Writes the given extension type value.
*
- * @param factory the extension type factory to add
+ * @param value the extension type value to write
+ * @param type of the extension
*/
- void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory);
+ void writeExtension(Object value, ArrowType type);
}
public interface ScalarWriter extends
diff --git a/vector/src/main/codegen/templates/ComplexCopier.java b/vector/src/main/codegen/templates/ComplexCopier.java
index 4fff7059a7..6655f6c2a7 100644
--- a/vector/src/main/codegen/templates/ComplexCopier.java
+++ b/vector/src/main/codegen/templates/ComplexCopier.java
@@ -41,11 +41,8 @@ public class ComplexCopier {
* @param input field to read from
* @param output field to write to
*/
- public static void copy(FieldReader input, FieldWriter output) {
- writeValue(input, output);
- }
+ public static void copy(FieldReader reader, FieldWriter writer) {
- private static void writeValue(FieldReader reader, FieldWriter writer) {
final MinorType mt = reader.getMinorType();
switch (mt) {
@@ -61,7 +58,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
FieldReader childReader = reader.reader();
FieldWriter childWriter = getListWriterForReader(childReader, writer);
if (childReader.isSet()) {
- writeValue(childReader, childWriter);
+ copy(childReader, childWriter);
} else {
childWriter.writeNull();
}
@@ -79,8 +76,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
FieldReader structReader = reader.reader();
if (structReader.isSet()) {
writer.startEntry();
- writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()));
- writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()));
+ copy(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()));
+ copy(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()));
writer.endEntry();
} else {
writer.writeNull();
@@ -99,7 +96,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
if (childReader.getMinorType() != Types.MinorType.NULL) {
FieldWriter childWriter = getStructWriterForReader(childReader, writer, name);
if (childReader.isSet()) {
- writeValue(childReader, childWriter);
+ copy(childReader, childWriter);
} else {
childWriter.writeNull();
}
@@ -110,6 +107,16 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
writer.writeNull();
}
break;
+ case EXTENSIONTYPE:
+ if (reader.isSet()) {
+ Object value = reader.readObject();
+ if (value != null) {
+ writer.writeExtension(value, reader.getField().getType());
+ }
+ } else {
+ writer.writeNull();
+ }
+ break;
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
<#assign fields = minor.fields!type.fields />
<#assign uncappedName = name?uncap_first/>
@@ -162,6 +169,9 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr
return (FieldWriter) writer.map(name);
case LISTVIEW:
return (FieldWriter) writer.listView(name);
+ case EXTENSIONTYPE:
+ ExtensionWriter extensionWriter = writer.extension(name, reader.getField().getType());
+ return (FieldWriter) extensionWriter;
default:
throw new UnsupportedOperationException(reader.getMinorType().toString());
}
@@ -186,6 +196,9 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
return (FieldWriter) writer.list();
case LISTVIEW:
return (FieldWriter) writer.listView();
+ case EXTENSIONTYPE:
+ ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+ return (FieldWriter) extensionWriter;
default:
throw new UnsupportedOperationException(reader.getMinorType().toString());
}
@@ -211,6 +224,9 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w
return (FieldWriter) writer.listView();
case MAP:
return (FieldWriter) writer.map(false);
+ case EXTENSIONTYPE:
+ ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+ return (FieldWriter) extensionWriter;
default:
throw new UnsupportedOperationException(reader.getMinorType().toString());
}
diff --git a/vector/src/main/codegen/templates/DenseUnionWriter.java b/vector/src/main/codegen/templates/DenseUnionWriter.java
index 8515b759e6..9aeea5b054 100644
--- a/vector/src/main/codegen/templates/DenseUnionWriter.java
+++ b/vector/src/main/codegen/templates/DenseUnionWriter.java
@@ -55,7 +55,9 @@ public DenseUnionWriter(DenseUnionVector vector, NullableStructWriterFactory nul
public void setPosition(int index) {
super.setPosition(index);
for (BaseWriter writer : writers) {
- writer.setPosition(index);
+ if (writer != null) {
+ writer.setPosition(index);
+ }
}
}
diff --git a/vector/src/main/codegen/templates/NullReader.java b/vector/src/main/codegen/templates/NullReader.java
index 1d77248e96..88e6ea98ea 100644
--- a/vector/src/main/codegen/templates/NullReader.java
+++ b/vector/src/main/codegen/templates/NullReader.java
@@ -86,6 +86,10 @@ public void read(int arrayIndex, Nullable${name}Holder holder){
}
+ public void read(ExtensionHolder holder) {
+ holder.isSet = 0;
+ }
+
public int size(){
return 0;
}
diff --git a/vector/src/main/codegen/templates/PromotableWriter.java b/vector/src/main/codegen/templates/PromotableWriter.java
index 8d7d57bb9d..11d34f72c9 100644
--- a/vector/src/main/codegen/templates/PromotableWriter.java
+++ b/vector/src/main/codegen/templates/PromotableWriter.java
@@ -286,7 +286,7 @@ protected void setWriter(ValueVector v) {
writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory);
break;
case EXTENSIONTYPE:
- writer = new UnionExtensionWriter((ExtensionTypeVector) vector);
+ writer = ((ExtensionType) vector.getField().getType()).getNewFieldWriter(vector);
break;
default:
writer = type.getNewFieldWriter(vector);
@@ -541,13 +541,13 @@ public void writeLargeVarChar(String value) {
}
@Override
- public void writeExtension(Object value) {
- getWriter(MinorType.EXTENSIONTYPE).writeExtension(value);
+ public void writeExtension(Object value, ArrowType arrowType) {
+ getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value, arrowType);
}
@Override
- public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) {
- getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeWriterFactory(factory);
+ public void write(ExtensionHolder holder) {
+ getWriter(MinorType.EXTENSIONTYPE, holder.type()).write(holder);
}
@Override
diff --git a/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java b/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java
index f6e3f63caf..484199ab2a 100644
--- a/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java
+++ b/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java
@@ -35,6 +35,10 @@
<#include "/@includes/vv_imports.ftl" />
+<#function is_timestamp_tz type>
+ <#return type?starts_with("TimeStamp") && type?ends_with("TZ")>
+
+
/*
* This class is generated using freemarker and the ${.template_name} template.
*/
@@ -96,55 +100,30 @@ public void close() throws Exception {
public void setPosition(int index) {
super.setPosition(index);
}
- <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
- <#assign fields = minor.fields!type.fields />
- <#assign uncappedName = name?uncap_first/>
- <#if uncappedName == "int" ><#assign uncappedName = "integer" />
- <#if !minor.typeParams?? >
+ <#list vv.types as type><#list type.minor as minor>
+ <#assign lowerName = minor.class?uncap_first />
+ <#if lowerName == "int" ><#assign lowerName = "integer" />
+ <#assign upperName = minor.class?upper_case />
@Override
- public ${name}Writer ${uncappedName}() {
+ public ${minor.class}Writer ${lowerName}() {
return this;
}
+ <#if minor.typeParams?? >
@Override
- public ${name}Writer ${uncappedName}(String name) {
- structName = name;
- return writer.${uncappedName}(name);
+ public ${minor.class}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) {
+ return writer.${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name});
}
-
-
- @Override
- public DecimalWriter decimal() {
- return this;
- }
-
- @Override
- public DecimalWriter decimal(String name, int scale, int precision) {
- return writer.decimal(name, scale, precision);
- }
-
- @Override
- public DecimalWriter decimal(String name) {
- return writer.decimal(name);
- }
-
@Override
- public Decimal256Writer decimal256() {
- return this;
- }
-
- @Override
- public Decimal256Writer decimal256(String name, int scale, int precision) {
- return writer.decimal256(name, scale, precision);
+ public ${minor.class}Writer ${lowerName}(String name) {
+ structName = name;
+ return writer.${lowerName}(name);
}
- @Override
- public Decimal256Writer decimal256(String name) {
- return writer.decimal256(name);
- }
+
@Override
public StructWriter struct() {
@@ -215,87 +194,86 @@ public void end() {
}
@Override
- public void write(DecimalHolder holder) {
- if (writer.idx() >= (idx() + 1) * listSize) {
- throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
- }
- writer.write(holder);
- writer.setPosition(writer.idx() + 1);
- }
-
- @Override
- public void write(Decimal256Holder holder) {
+ public void writeNull() {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.write(holder);
- writer.setPosition(writer.idx() + 1);
+ writer.writeNull();
}
+ <#list vv.types as type>
+ <#list type.minor as minor>
+ <#assign name = minor.class?cap_first />
+ <#assign fields = minor.fields!type.fields />
+ <#assign uncappedName = name?uncap_first/>
@Override
- public void writeNull() {
+ public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeNull();
+ writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, );
+ writer.setPosition(writer.idx()+1);
}
- public void writeDecimal(long start, ArrowBuf buffer, ArrowType arrowType) {
+ <#if is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary">
+ @Override
+ public void write(${name}Holder holder) {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeDecimal(start, buffer, arrowType);
- writer.setPosition(writer.idx() + 1);
+ writer.write(holder);
+ writer.setPosition(writer.idx()+1);
}
- public void writeDecimal(BigDecimal value) {
+ <#elseif minor.class?starts_with("Decimal")>
+ @Override
+ public void write${name}(long start, ArrowBuf buffer, ArrowType arrowType) {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeDecimal(value);
- writer.setPosition(writer.idx() + 1);
+ writer.write${name}(start, buffer, arrowType);
+ writer.setPosition(writer.idx()+1);
}
- public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) {
+ @Override
+ public void write(${name}Holder holder) {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeBigEndianBytesToDecimal(value, arrowType);
- writer.setPosition(writer.idx() + 1);
+ writer.write(holder);
+ writer.setPosition(writer.idx()+1);
}
- public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) {
+ @Override
+ public void write${name}(BigDecimal value) {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeDecimal256(start, buffer, arrowType);
- writer.setPosition(writer.idx() + 1);
+ writer.write${name}(value);
+ writer.setPosition(writer.idx()+1);
}
- public void writeDecimal256(BigDecimal value) {
+ @Override
+ public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType){
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeDecimal256(value);
+ writer.writeBigEndianBytesTo${name}(value, arrowType);
writer.setPosition(writer.idx() + 1);
}
-
- public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) {
+ <#else>
+ @Override
+ public void write(${name}Holder holder) {
if (writer.idx() >= (idx() + 1) * listSize) {
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
}
- writer.writeBigEndianBytesToDecimal256(value, arrowType);
- writer.setPosition(writer.idx() + 1);
+ writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, );
+ writer.setPosition(writer.idx()+1);
}
+
-
- <#list vv.types as type>
- <#list type.minor as minor>
- <#assign name = minor.class?cap_first />
- <#assign fields = minor.fields!type.fields />
- <#assign uncappedName = name?uncap_first/>
- <#if minor.class?ends_with("VarBinary")>
+ <#if minor.class?ends_with("VarBinary")>
@Override
public void write${minor.class}(byte[] value) {
if (writer.idx() >= (idx() + 1) * listSize) {
@@ -349,27 +327,8 @@ public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) {
writer.write${minor.class}(value);
writer.setPosition(writer.idx() + 1);
}
-
-
- <#if !minor.typeParams?? >
- @Override
- public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) {
- if (writer.idx() >= (idx() + 1) * listSize) {
- throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
- }
- writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, );
- writer.setPosition(writer.idx() + 1);
- }
-
- public void write(${name}Holder holder) {
- if (writer.idx() >= (idx() + 1) * listSize) {
- throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
- }
- writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, );
- writer.setPosition(writer.idx() + 1);
- }
+
-
}
diff --git a/vector/src/main/codegen/templates/UnionListWriter.java b/vector/src/main/codegen/templates/UnionListWriter.java
index 9424533f29..394348f029 100644
--- a/vector/src/main/codegen/templates/UnionListWriter.java
+++ b/vector/src/main/codegen/templates/UnionListWriter.java
@@ -53,6 +53,7 @@ public class Union${listName}Writer extends AbstractFieldWriter {
private boolean inStruct = false;
private boolean listStarted = false;
private String structName;
+ private ArrowType extensionType;
<#if listName == "LargeList" || listName == "LargeListView">
private static final long OFFSET_WIDTH = 8;
<#else>
@@ -122,8 +123,6 @@ public void setPosition(int index) {
<#assign lowerName = minor.class?uncap_first />
<#if lowerName == "int" ><#assign lowerName = "integer" />
<#assign upperName = minor.class?upper_case />
- <#assign capName = minor.class?cap_first />
- <#assign vectName = capName />
@Override
public ${minor.class}Writer ${lowerName}() {
return this;
@@ -203,13 +202,13 @@ public MapWriter map(String name, boolean keysSorted) {
@Override
public ExtensionWriter extension(ArrowType arrowType) {
- writer.extension(arrowType);
- return writer;
+ extensionType = arrowType;
+ return this;
}
+
@Override
public ExtensionWriter extension(String name, ArrowType arrowType) {
- ExtensionWriter extensionWriter = writer.extension(name, arrowType);
- return extensionWriter;
+ return writer.extension(name, arrowType);
}
<#if listName == "LargeList">
@@ -336,14 +335,18 @@ public void writeNull() {
@Override
public void writeExtension(Object value) {
- writer.writeExtension(value);
+ writer.writeExtension(value, extensionType);
+ writer.setPosition(writer.idx() + 1);
}
+
@Override
- public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
- writer.addExtensionTypeWriterFactory(var1);
+ public void writeExtension(Object value, ArrowType type) {
+ writeExtension(value);
}
+
public void write(ExtensionHolder var1) {
writer.write(var1);
+ writer.setPosition(writer.idx() + 1);
}
<#list vv.types as type>
@@ -365,6 +368,7 @@ public void write(${name}Holder holder) {
}
<#elseif minor.class?starts_with("Decimal")>
+ @Override
public void write${name}(long start, ArrowBuf buffer, ArrowType arrowType) {
writer.write${name}(start, buffer, arrowType);
writer.setPosition(writer.idx()+1);
@@ -376,11 +380,13 @@ public void write(${name}Holder holder) {
writer.setPosition(writer.idx()+1);
}
+ @Override
public void write${name}(BigDecimal value) {
writer.write${name}(value);
writer.setPosition(writer.idx()+1);
}
+ @Override
public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType){
writer.writeBigEndianBytesTo${name}(value, arrowType);
writer.setPosition(writer.idx() + 1);
@@ -424,6 +430,7 @@ public void write(${name}Holder holder) {
writer.setPosition(writer.idx() + 1);
}
+ @Override
public void write${minor.class}(String value) {
writer.write${minor.class}(value);
writer.setPosition(writer.idx() + 1);
diff --git a/vector/src/main/codegen/templates/UnionMapWriter.java b/vector/src/main/codegen/templates/UnionMapWriter.java
index 8b2f091215..8bbf6ae0a4 100644
--- a/vector/src/main/codegen/templates/UnionMapWriter.java
+++ b/vector/src/main/codegen/templates/UnionMapWriter.java
@@ -243,4 +243,27 @@ public ExtensionWriter extension(ArrowType type) {
return super.extension(type);
}
}
+
+ public FixedSizeBinaryWriter fixedSizeBinary(int byteWidth) {
+ switch (mode) {
+ case KEY:
+ return entryWriter.fixedSizeBinary(MapVector.KEY_NAME, byteWidth);
+ case VALUE:
+ return entryWriter.fixedSizeBinary(MapVector.VALUE_NAME, byteWidth);
+ default:
+ return this;
+ }
+ }
+
+ @Override
+ public FixedSizeBinaryWriter fixedSizeBinary() {
+ switch (mode) {
+ case KEY:
+ return entryWriter.fixedSizeBinary(MapVector.KEY_NAME);
+ case VALUE:
+ return entryWriter.fixedSizeBinary(MapVector.VALUE_NAME);
+ default:
+ return this;
+ }
+ }
}
diff --git a/vector/src/main/codegen/templates/UnionReader.java b/vector/src/main/codegen/templates/UnionReader.java
index 96ad3e1b9b..0edae7ade0 100644
--- a/vector/src/main/codegen/templates/UnionReader.java
+++ b/vector/src/main/codegen/templates/UnionReader.java
@@ -79,6 +79,10 @@ public void read(int index, UnionHolder holder) {
}
private FieldReader getReaderForIndex(int index) {
+ return getReaderForIndex(index, null);
+ }
+
+ private FieldReader getReaderForIndex(int index, ArrowType type) {
int typeValue = data.getTypeValue(index);
FieldReader reader = (FieldReader) readers[typeValue];
if (reader != null) {
@@ -105,11 +109,26 @@ private FieldReader getReaderForIndex(int index) {
+ case EXTENSIONTYPE:
+ if(type == null) {
+ throw new RuntimeException("Cannot get Extension reader without an ArrowType");
+ }
+ return (FieldReader) getExtension(type);
default:
throw new UnsupportedOperationException("Unsupported type: " + MinorType.values()[typeValue]);
}
}
+ private ExtensionReader extensionReader;
+
+ private ExtensionReader getExtension(ArrowType type) {
+ if (extensionReader == null) {
+ extensionReader = data.getExtension(type).getReader();
+ extensionReader.setPosition(idx());
+ }
+ return extensionReader;
+ }
+
private SingleStructReaderImpl structReader;
private StructReader getStruct() {
@@ -240,4 +259,8 @@ public FieldReader reader() {
public boolean next() {
return getReaderForIndex(idx()).next();
}
+
+ public void read(ExtensionHolder holder){
+ getReaderForIndex(idx(), holder.type()).read(holder);
+ }
}
diff --git a/vector/src/main/codegen/templates/UnionVector.java b/vector/src/main/codegen/templates/UnionVector.java
index 67efdf60f7..c706591966 100644
--- a/vector/src/main/codegen/templates/UnionVector.java
+++ b/vector/src/main/codegen/templates/UnionVector.java
@@ -379,6 +379,22 @@ public MapVector getMap(String name, ArrowType arrowType) {
return mapVector;
}
+ private ExtensionTypeVector extensionVector;
+
+ public ExtensionTypeVector getExtension(ArrowType arrowType) {
+ if (extensionVector == null) {
+ int vectorCount = internalStruct.size();
+ extensionVector = addOrGet(null, MinorType.EXTENSIONTYPE, arrowType, ExtensionTypeVector.class);
+ if (internalStruct.size() > vectorCount) {
+ extensionVector.allocateNew();
+ if (callBack != null) {
+ callBack.doWork();
+ }
+ }
+ }
+ return extensionVector;
+ }
+
public int getTypeValue(int index) {
return typeBuffer.getByte(index * TYPE_WIDTH);
}
@@ -725,6 +741,8 @@ public ValueVector getVectorByType(int typeId, ArrowType arrowType) {
return getListView();
case MAP:
return getMap(name, arrowType);
+ case EXTENSIONTYPE:
+ return getExtension(arrowType);
default:
throw new UnsupportedOperationException("Cannot support type: " + MinorType.values()[typeId]);
}
diff --git a/vector/src/main/codegen/templates/UnionWriter.java b/vector/src/main/codegen/templates/UnionWriter.java
index 272edab17c..0db699fd8c 100644
--- a/vector/src/main/codegen/templates/UnionWriter.java
+++ b/vector/src/main/codegen/templates/UnionWriter.java
@@ -28,6 +28,8 @@
package org.apache.arrow.vector.complex.impl;
<#include "/@includes/vv_imports.ftl" />
+import java.util.HashMap;
+
import org.apache.arrow.vector.complex.writer.BaseWriter;
import org.apache.arrow.vector.types.Types.MinorType;
@@ -213,8 +215,31 @@ public MapWriter asMap(ArrowType arrowType) {
return getMapWriter(arrowType);
}
+ private java.util.Map extensionWriters = new HashMap<>();
+
private ExtensionWriter getExtensionWriter(ArrowType arrowType) {
- throw new UnsupportedOperationException("ExtensionTypes are not supported yet.");
+ ExtensionWriter w = extensionWriters.get(arrowType);
+ if (w == null) {
+ w = ((ExtensionType) arrowType).getNewFieldWriter(data.getExtension(arrowType));
+ w.setPosition(idx());
+ extensionWriters.put(arrowType, w);
+ }
+ return w;
+ }
+
+ public void writeExtension(Object value, ArrowType type) {
+ data.setType(idx(), MinorType.EXTENSIONTYPE);
+ ExtensionWriter w = getExtensionWriter(type);
+ w.setPosition(idx());
+ w.writeExtension(value);
+ }
+
+ @Override
+ public void write(ExtensionHolder holder) {
+ data.setType(idx(), MinorType.EXTENSIONTYPE);
+ ExtensionWriter w = getExtensionWriter(holder.type());
+ w.setPosition(idx());
+ w.write(holder);
}
BaseWriter getWriter(MinorType minorType) {
diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java
index 4be55396b7..df1ac74f9b 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java
@@ -49,9 +49,7 @@ public abstract class BaseFixedWidthVector extends BaseValueVector
protected final Field field;
private int allocationMonitor;
- protected ArrowBuf validityBuffer;
protected ArrowBuf valueBuffer;
- protected int valueCount;
/**
* Constructs a new instance.
@@ -72,6 +70,7 @@ public BaseFixedWidthVector(Field field, final BufferAllocator allocator, final
refreshValueCapacity();
}
+ @Override
public int getTypeWidth() {
return typeWidth;
}
@@ -87,7 +86,7 @@ public String getName() {
/* TODO:
* Once the entire hierarchy has been refactored, move common functions
- * like getNullCount(), splitAndTransferValidityBuffer to top level
+ * like getNullCount() to top level
* base class BaseValueVector.
*
* Along with this, some class members (validityBuffer) can also be
@@ -342,9 +341,9 @@ private void allocateBytes(int valueCount) {
* slice the source buffer so we have to explicitly allocate the validityBuffer of the target
* vector. This is unlike the databuffer which we can always slice for the target vector.
*/
- private void allocateValidityBuffer(final int validityBufferSize) {
- validityBuffer = allocator.buffer(validityBufferSize);
- validityBuffer.readerIndex(0);
+ @Override
+ protected void allocateValidityBuffer(final long validityBufferSize) {
+ super.allocateValidityBuffer(validityBufferSize);
refreshValueCapacity();
}
@@ -359,7 +358,7 @@ public int getBufferSizeFor(final int count) {
if (count == 0) {
return 0;
}
- return (count * typeWidth) + getValidityBufferSizeFromCount(count);
+ return (count * typeWidth) + BitVectorHelper.getValidityBufferSizeFromCount(count);
}
/**
@@ -372,7 +371,7 @@ public int getBufferSize() {
if (valueCount == 0) {
return 0;
}
- return (valueCount * typeWidth) + getValidityBufferSizeFromCount(valueCount);
+ return (valueCount * typeWidth) + BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
}
/**
@@ -536,10 +535,10 @@ private void setReaderAndWriterIndex() {
validityBuffer.writerIndex(0);
valueBuffer.writerIndex(0);
} else {
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
if (typeWidth == 0) {
/* specialized handling for BitVector */
- valueBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
+ valueBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
} else {
valueBuffer.writerIndex((long) valueCount * typeWidth);
}
@@ -656,72 +655,18 @@ private void splitAndTransferValueBuffer(
target.refreshValueCapacity();
}
- /**
- * Validity buffer has multiple cases of split and transfer depending on the starting position of
- * the source index.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, BaseFixedWidthVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- /* slice */
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator);
- target.refreshValueCapacity();
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- this.validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- this.validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
+ @Override
+ protected void sliceAndTransferValidityBuffer(
+ int startIndex, int length, BaseValueVector target) {
+ final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
+ final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length);
+
+ if (target.validityBuffer != null) {
+ target.validityBuffer.getReferenceManager().release();
}
+ ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
+ target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator);
+ ((BaseFixedWidthVector) target).refreshValueCapacity();
}
/*----------------------------------------------------------------*
diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java
index 7e0d0affc6..3fac195786 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java
@@ -52,10 +52,8 @@ public abstract class BaseLargeVariableWidthVector extends BaseValueVector
/* protected members */
public static final int OFFSET_WIDTH = 8; /* 8 byte unsigned int to track offsets */
protected static final byte[] emptyByteArray = new byte[] {};
- protected ArrowBuf validityBuffer;
protected ArrowBuf valueBuffer;
protected ArrowBuf offsetBuffer;
- protected int valueCount;
protected int lastSet;
protected final Field field;
@@ -375,14 +373,26 @@ private void setReaderAndWriterIndex() {
valueBuffer.readerIndex(0);
if (valueCount == 0) {
validityBuffer.writerIndex(0);
- offsetBuffer.writerIndex(0);
valueBuffer.writerIndex(0);
} else {
final long lastDataOffset = getStartOffset(valueCount);
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
- offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
valueBuffer.writerIndex(lastDataOffset);
}
+ // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.
+ // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers
+ // in other libraries. According to Arrow spec, we should still output the offset buffer which
+ // is [0].
+ final long requiredOffsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH;
+ if (offsetBuffer.capacity() < requiredOffsetBufferSize) {
+ ArrowBuf newOffsetBuffer = allocateOffsetBuffer(requiredOffsetBufferSize);
+ if (offsetBuffer.capacity() > 0) {
+ newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity());
+ }
+ offsetBuffer.getReferenceManager().release();
+ offsetBuffer = newOffsetBuffer;
+ }
+ offsetBuffer.writerIndex(requiredOffsetBufferSize);
}
/** Same as {@link #allocateNewSafe()}. */
@@ -501,10 +511,9 @@ private ArrowBuf allocateOffsetBuffer(final long size) {
}
/* allocate validity buffer */
- private void allocateValidityBuffer(final long size) {
- validityBuffer = allocator.buffer(size);
- validityBuffer.readerIndex(0);
- initValidityBuffer();
+ @Override
+ protected void allocateValidityBuffer(final long size) {
+ super.allocateValidityBuffer(size);
}
/**
@@ -633,7 +642,7 @@ public int getBufferSizeFor(final int valueCount) {
return 0;
}
- final long validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final long validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
final long offsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH;
/* get the end offset for this valueCount */
final long dataBufferSize = getStartOffset(valueCount);
@@ -809,69 +818,17 @@ private void splitAndTransferOffsetBuffer(
target.valueBuffer = transferBuffer(slicedBuffer, target.allocator);
}
- /*
- * Transfer the validity.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, BaseLargeVariableWidthVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
+ @Override
+ protected void sliceAndTransferValidityBuffer(
+ int startIndex, int length, BaseValueVector target) {
+ final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
+ final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length);
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain();
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- this.validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- this.validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
+ if (target.validityBuffer != null) {
+ target.validityBuffer.getReferenceManager().release();
}
+ target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
+ target.validityBuffer.getReferenceManager().retain();
}
/*----------------------------------------------------------------*
diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java
index 9befcb890f..37dfa20616 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java
@@ -48,6 +48,10 @@ public abstract class BaseValueVector implements ValueVector {
protected volatile FieldReader fieldReader;
+ protected ArrowBuf validityBuffer;
+
+ protected int valueCount;
+
protected BaseValueVector(BufferAllocator allocator) {
this.allocator = Preconditions.checkNotNull(allocator, "allocator cannot be null");
}
@@ -110,7 +114,14 @@ protected ArrowBuf releaseBuffer(ArrowBuf buffer) {
return buffer;
}
- /* number of bytes for the validity buffer for the given valueCount */
+ /**
+ * Compute the size of validity buffer required to manage a given number of elements in a vector.
+ *
+ * @param valueCount number of elements in the vector
+ * @return buffer size
+ * @deprecated -- use {@link BitVectorHelper#getValidityBufferSizeFromCount} instead.
+ */
+ @Deprecated(forRemoval = true, since = "18.4.0")
protected static int getValidityBufferSizeFromCount(final int valueCount) {
return DataSizeRoundingUtil.divideBy8Ceil(valueCount);
}
@@ -248,4 +259,116 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
throw new UnsupportedOperationException();
}
+
+ /**
+ * Transfer the validity buffer from `validityBuffer` to the target vector's `validityBuffer`.
+ * Start at `startIndex` and copy `length` number of elements. If the starting index is 8 byte
+ * aligned, then the buffer is sliced from that index and ownership is transferred. If not,
+ * individual bytes are copied.
+ *
+ * @param startIndex starting index
+ * @param length number of elements to be copied
+ * @param target target vector
+ */
+ protected void splitAndTransferValidityBuffer(
+ int startIndex, int length, BaseValueVector target) {
+ int offset = startIndex % 8;
+
+ if (length <= 0) {
+ return;
+ }
+ if (offset == 0) {
+ sliceAndTransferValidityBuffer(startIndex, length, target);
+ } else {
+ copyValidityBuffer(startIndex, length, target);
+ }
+ }
+
+ /**
+ * If the start index is 8 byte aligned, slice `validityBuffer` and transfer ownership to
+ * `target`'s `validityBuffer`.
+ *
+ * @param startIndex starting index
+ * @param length number of elements to be copied
+ * @param target target vector
+ */
+ protected void sliceAndTransferValidityBuffer(
+ int startIndex, int length, BaseValueVector target) {
+ final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
+ final int byteSizeTarget = getValidityBufferSizeFromCount(length);
+
+ if (target.validityBuffer != null) {
+ target.validityBuffer.getReferenceManager().release();
+ }
+ target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
+ target.validityBuffer.getReferenceManager().retain(1);
+ }
+
+ /**
+ * Allocate new validity buffer for `target` and copy bytes from `validityBuffer`. Precise details
+ * in the comments below.
+ *
+ * @param startIndex starting index
+ * @param length number of elements to be copied
+ * @param target target vector
+ */
+ protected void copyValidityBuffer(int startIndex, int length, BaseValueVector target) {
+ final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
+ final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
+ final int byteSizeTarget = getValidityBufferSizeFromCount(length);
+ final int offset = startIndex % 8;
+
+ /* Copy data
+ * When the first bit starts from the middle of a byte (offset != 0),
+ * copy data from src BitVector.
+ * Each byte in the target is composed by a part in i-th byte,
+ * another part in (i+1)-th byte.
+ */
+ target.allocateValidityBuffer(byteSizeTarget);
+
+ for (int i = 0; i < byteSizeTarget - 1; i++) {
+ byte b1 =
+ BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset);
+ byte b2 =
+ BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset);
+
+ target.validityBuffer.setByte(i, (b1 + b2));
+ }
+
+ /* Copying the last piece is done in the following manner:
+ * if the source vector has 1 or more bytes remaining, we copy
+ * the last piece as a byte formed by shifting data
+ * from the current byte and the next byte.
+ *
+ * if the source vector has no more bytes remaining
+ * (we are at the last byte), we copy the last piece as a byte
+ * by shifting data from the current byte.
+ */
+ if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
+ byte b1 =
+ BitVectorHelper.getBitsFromCurrentByte(
+ this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
+ byte b2 =
+ BitVectorHelper.getBitsFromNextByte(
+ this.validityBuffer, firstByteSource + byteSizeTarget, offset);
+
+ target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
+ } else {
+ byte b1 =
+ BitVectorHelper.getBitsFromCurrentByte(
+ this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
+ target.validityBuffer.setByte(byteSizeTarget - 1, b1);
+ }
+ }
+
+ /**
+ * Allocate new validity buffer for when the bytes need to be copied over.
+ *
+ * @param byteSizeTarget desired size of the buffer
+ */
+ protected void allocateValidityBuffer(long byteSizeTarget) {
+ validityBuffer = allocator.buffer(byteSizeTarget);
+ validityBuffer.readerIndex(0);
+ validityBuffer.setZero(0, validityBuffer.capacity());
+ }
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java
index 1609e64ca5..d5bd167256 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java
@@ -50,10 +50,8 @@ public abstract class BaseVariableWidthVector extends BaseValueVector
/* protected members */
public static final int OFFSET_WIDTH = 4; /* 4 byte unsigned int to track offsets */
protected static final byte[] emptyByteArray = new byte[] {};
- protected ArrowBuf validityBuffer;
protected ArrowBuf valueBuffer;
protected ArrowBuf offsetBuffer;
- protected int valueCount;
protected int lastSet;
protected final Field field;
@@ -87,7 +85,7 @@ public String getName() {
/* TODO:
* Once the entire hierarchy has been refactored, move common functions
- * like getNullCount(), splitAndTransferValidityBuffer to top level
+ * like getNullCount() to top level
* base class BaseValueVector.
*
* Along with this, some class members (validityBuffer) can also be
@@ -391,14 +389,26 @@ private void setReaderAndWriterIndex() {
valueBuffer.readerIndex(0);
if (valueCount == 0) {
validityBuffer.writerIndex(0);
- offsetBuffer.writerIndex(0);
valueBuffer.writerIndex(0);
} else {
final int lastDataOffset = getStartOffset(valueCount);
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
- offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
valueBuffer.writerIndex(lastDataOffset);
}
+ // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.
+ // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers
+ // in other libraries. According to Arrow spec, we should still output the offset buffer which
+ // is [0].
+ final long requiredOffsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH;
+ if (offsetBuffer.capacity() < requiredOffsetBufferSize) {
+ ArrowBuf newOffsetBuffer = allocateOffsetBuffer(requiredOffsetBufferSize);
+ if (offsetBuffer.capacity() > 0) {
+ newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity());
+ }
+ offsetBuffer.getReferenceManager().release();
+ offsetBuffer = newOffsetBuffer;
+ }
+ offsetBuffer.writerIndex(requiredOffsetBufferSize);
}
/** Same as {@link #allocateNewSafe()}. */
@@ -519,11 +529,9 @@ private ArrowBuf allocateOffsetBuffer(final long size) {
}
/* allocate validity buffer */
- private void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- initValidityBuffer();
+ @Override
+ protected void allocateValidityBuffer(final long size) {
+ super.allocateValidityBuffer(size);
}
/**
@@ -673,7 +681,7 @@ public int getBufferSizeFor(final int valueCount) {
return 0;
}
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH;
/* get the end offset for this valueCount */
final int dataBufferSize = offsetBuffer.getInt((long) valueCount * OFFSET_WIDTH);
@@ -856,70 +864,17 @@ private void splitAndTransferOffsetBuffer(
target.valueBuffer = transferBuffer(slicedBuffer, target.allocator);
}
- /*
- * Transfer the validity.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, BaseVariableWidthVector target) {
- if (length <= 0) {
- return;
- }
-
+ @Override
+ protected void sliceAndTransferValidityBuffer(
+ int startIndex, int length, BaseValueVector target) {
final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- final int byteSizeTarget = getValidityBufferSizeFromCount(length);
- final int offset = startIndex % 8;
-
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator);
- return;
- }
-
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
+ final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length);
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- this.validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
+ if (target.validityBuffer != null) {
+ target.validityBuffer.getReferenceManager().release();
}
+ final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
+ target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator);
}
/*----------------------------------------------------------------*
diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java
index beda91dc3f..ea9de8320e 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java
@@ -78,13 +78,11 @@ public abstract class BaseVariableWidthViewVector extends BaseValueVector
// The third 4 bytes of view are allocated for buffer index
public static final int BUF_INDEX_WIDTH = 4;
public static final byte[] EMPTY_BYTE_ARRAY = new byte[] {};
- protected ArrowBuf validityBuffer;
// The view buffer is used to store the variable width view elements
protected ArrowBuf viewBuffer;
// The external buffer which stores the long strings
protected List dataBuffers;
protected int initialDataBufferSize;
- protected int valueCount;
protected int lastSet;
protected final Field field;
@@ -117,7 +115,7 @@ public String getName() {
/* TODO:
* Once the entire hierarchy has been refactored, move common functions
- * like getNullCount(), splitAndTransferValidityBuffer to top level
+ * like getNullCount() to top level
* base class BaseValueVector.
*
* Along with this, some class members (validityBuffer) can also be
@@ -129,12 +127,6 @@ public String getName() {
* the top class as of now is not a good idea.
*/
- /* TODO:
- * Implement TransferPair functionality
- * https://github.com/apache/arrow/issues/40932
- *
- */
-
/**
* Get buffer that manages the validity (NULL or NON-NULL nature) of elements in the vector.
* Consider it as a buffer for internal bit vector data structure.
@@ -400,7 +392,7 @@ private void setReaderAndWriterIndex() {
validityBuffer.writerIndex(0);
viewBuffer.writerIndex(0);
} else {
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
viewBuffer.writerIndex(valueCount * ELEMENT_SIZE);
}
}
@@ -683,7 +675,7 @@ public int getBufferSizeFor(final int valueCount) {
return 0;
}
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
final int viewBufferSize = valueCount * ELEMENT_SIZE;
final int dataBufferSize = getDataBufferSize();
return validityBufferSize + viewBufferSize + dataBufferSize;
@@ -854,77 +846,22 @@ public void splitAndTransferTo(int startIndex, int length, BaseVariableWidthView
}
/* allocate validity buffer */
- private void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- initValidityBuffer();
+ @Override
+ protected void allocateValidityBuffer(final long size) {
+ super.allocateValidityBuffer(size);
}
- /*
- * Transfer the validity.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, BaseVariableWidthViewVector target) {
- if (length <= 0) {
- return;
- }
-
+ @Override
+ protected void sliceAndTransferValidityBuffer(
+ int startIndex, int length, BaseValueVector target) {
final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- final int byteSizeTarget = getValidityBufferSizeFromCount(length);
- final int offset = startIndex % 8;
-
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator);
- return;
- }
+ final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length);
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- this.validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
+ if (target.validityBuffer != null) {
+ target.validityBuffer.getReferenceManager().release();
}
+ final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
+ target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator);
}
/**
diff --git a/vector/src/main/java/org/apache/arrow/vector/BitVector.java b/vector/src/main/java/org/apache/arrow/vector/BitVector.java
index f8e3342625..ecee02f665 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BitVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BitVector.java
@@ -98,7 +98,7 @@ public MinorType getMinorType() {
*/
@Override
public void setInitialCapacity(int valueCount) {
- final int size = getValidityBufferSizeFromCount(valueCount);
+ final int size = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
if (size * 2L > MAX_ALLOCATION_SIZE) {
throw new OversizedAllocationException("Requested amount of memory is more than max allowed");
}
@@ -121,7 +121,7 @@ public int getBufferSizeFor(final int count) {
if (count == 0) {
return 0;
}
- return 2 * getValidityBufferSizeFromCount(count);
+ return 2 * BitVectorHelper.getValidityBufferSizeFromCount(count);
}
/**
@@ -165,7 +165,7 @@ private ArrowBuf splitAndTransferBuffer(
int startIndex, int length, ArrowBuf sourceBuffer, ArrowBuf destBuffer) {
int firstByteSource = BitVectorHelper.byteIndex(startIndex);
int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
+ int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length);
int offset = startIndex % 8;
if (length > 0) {
diff --git a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java
index 0ac56691a6..bc2c3da98f 100644
--- a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java
+++ b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java
@@ -135,11 +135,11 @@ public static void setValidityBit(ArrowBuf validityBuffer, int index, int value)
public static ArrowBuf setValidityBit(
ArrowBuf validityBuffer, BufferAllocator allocator, int valueCount, int index, int value) {
if (validityBuffer == null) {
- validityBuffer = allocator.buffer(getValidityBufferSize(valueCount));
+ validityBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount));
}
setValidityBit(validityBuffer, index, value);
if (index == (valueCount - 1)) {
- validityBuffer.writerIndex(getValidityBufferSize(valueCount));
+ validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
}
return validityBuffer;
@@ -165,7 +165,7 @@ public static int get(final ArrowBuf buffer, int index) {
* @param valueCount number of elements in the vector
* @return buffer size
*/
- public static int getValidityBufferSize(int valueCount) {
+ public static int getValidityBufferSizeFromCount(int valueCount) {
return DataSizeRoundingUtil.divideBy8Ceil(valueCount);
}
@@ -182,7 +182,7 @@ public static int getNullCount(final ArrowBuf validityBuffer, final int valueCou
return 0;
}
int count = 0;
- final int sizeInBytes = getValidityBufferSize(valueCount);
+ final int sizeInBytes = getValidityBufferSizeFromCount(valueCount);
// If value count is not a multiple of 8, then calculate number of used bits in the last byte
final int remainder = valueCount % 8;
final int fullBytesCount = remainder == 0 ? sizeInBytes : sizeInBytes - 1;
@@ -233,7 +233,7 @@ public static boolean checkAllBitsEqualTo(
if (valueCount == 0) {
return true;
}
- final int sizeInBytes = getValidityBufferSize(valueCount);
+ final int sizeInBytes = getValidityBufferSizeFromCount(valueCount);
// boundary check
validityBuffer.checkBytes(0, sizeInBytes);
@@ -325,7 +325,7 @@ public static ArrowBuf loadValidityBuffer(
sourceValidityBuffer == null || sourceValidityBuffer.capacity() == 0;
if (isValidityBufferNull
&& (fieldNode.getNullCount() == 0 || fieldNode.getNullCount() == valueCount)) {
- newBuffer = allocator.buffer(getValidityBufferSize(valueCount));
+ newBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount));
newBuffer.setZero(0, newBuffer.capacity());
if (fieldNode.getNullCount() != 0) {
/* all NULLs */
diff --git a/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java
index e22a973f3b..61a5574898 100644
--- a/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java
@@ -31,4 +31,7 @@ public interface FixedWidthVector extends ElementAddressableVector {
/** Zero out the underlying buffer backing this vector. */
void zeroVector();
+
+ /** Get the width of the type in bytes. */
+ int getTypeWidth();
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java
index abaefcfc12..50f2f066cc 100644
--- a/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java
@@ -155,12 +155,13 @@ public void set(int index, NullableTimeStampMicroTZHolder holder)
throws IllegalArgumentException {
if (holder.isSet < 0) {
throw new IllegalArgumentException();
- } else if (!this.timeZone.equals(holder.timezone)) {
- throw new IllegalArgumentException(
- String.format(
- "holder.timezone: %s not equal to vector timezone: %s",
- holder.timezone, this.timeZone));
} else if (holder.isSet > 0) {
+ if (!this.timeZone.equals(holder.timezone)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "holder.timezone: %s not equal to vector timezone: %s",
+ holder.timezone, this.timeZone));
+ }
BitVectorHelper.setBit(validityBuffer, index);
setValue(index, holder.value);
} else {
diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java
index b5e5fb1be1..9e4998396c 100644
--- a/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java
@@ -155,12 +155,13 @@ public void set(int index, NullableTimeStampMilliTZHolder holder)
throws IllegalArgumentException {
if (holder.isSet < 0) {
throw new IllegalArgumentException();
- } else if (!this.timeZone.equals(holder.timezone)) {
- throw new IllegalArgumentException(
- String.format(
- "holder.timezone: %s not equal to vector timezone: %s",
- holder.timezone, this.timeZone));
} else if (holder.isSet > 0) {
+ if (!this.timeZone.equals(holder.timezone)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "holder.timezone: %s not equal to vector timezone: %s",
+ holder.timezone, this.timeZone));
+ }
BitVectorHelper.setBit(validityBuffer, index);
setValue(index, holder.value);
} else {
diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java
index 2386b3a859..b44b3da8d3 100644
--- a/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java
@@ -154,12 +154,13 @@ public Long getObject(int index) {
public void set(int index, NullableTimeStampNanoTZHolder holder) throws IllegalArgumentException {
if (holder.isSet < 0) {
throw new IllegalArgumentException();
- } else if (!this.timeZone.equals(holder.timezone)) {
- throw new IllegalArgumentException(
- String.format(
- "holder.timezone: %s not equal to vector timezone: %s",
- holder.timezone, this.timeZone));
} else if (holder.isSet > 0) {
+ if (!this.timeZone.equals(holder.timezone)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "holder.timezone: %s not equal to vector timezone: %s",
+ holder.timezone, this.timeZone));
+ }
BitVectorHelper.setBit(validityBuffer, index);
setValue(index, holder.value);
} else {
diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java
index f1774f2703..a64a87f699 100644
--- a/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java
@@ -150,12 +150,13 @@ public Long getObject(int index) {
public void set(int index, NullableTimeStampSecTZHolder holder) throws IllegalArgumentException {
if (holder.isSet < 0) {
throw new IllegalArgumentException();
- } else if (!this.timeZone.equals(holder.timezone)) {
- throw new IllegalArgumentException(
- String.format(
- "holder.timezone: %s not equal to vector timezone: %s",
- holder.timezone, this.timeZone));
} else if (holder.isSet > 0) {
+ if (!this.timeZone.equals(holder.timezone)) {
+ throw new IllegalArgumentException(
+ String.format(
+ "holder.timezone: %s not equal to vector timezone: %s",
+ holder.timezone, this.timeZone));
+ }
BitVectorHelper.setBit(validityBuffer, index);
setValue(index, holder.value);
} else {
diff --git a/vector/src/main/java/org/apache/arrow/vector/UuidVector.java b/vector/src/main/java/org/apache/arrow/vector/UuidVector.java
new file mode 100644
index 0000000000..e1e61a5a2e
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/UuidVector.java
@@ -0,0 +1,458 @@
+/*
+ * 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 static org.apache.arrow.vector.extension.UuidType.UUID_BYTE_WIDTH;
+
+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.util.ArrowBufPointer;
+import org.apache.arrow.memory.util.ByteFunctionHelpers;
+import org.apache.arrow.memory.util.hash.ArrowBufHasher;
+import org.apache.arrow.util.Preconditions;
+import org.apache.arrow.vector.complex.impl.UuidReaderImpl;
+import org.apache.arrow.vector.complex.reader.FieldReader;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
+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;
+import org.apache.arrow.vector.util.UuidUtility;
+
+/**
+ * Vector implementation for UUID values using {@link UuidType}.
+ *
+ * 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 ExtensionTypeVector
+ implements ValueIterableVector, FixedWidthVector {
+ private final Field field;
+
+ /** The fixed byte width of UUID values (16 bytes). */
+ public static final int TYPE_WIDTH = UUID_BYTE_WIDTH;
+
+ /**
+ * Constructs a UUID vector with the given name, allocator, and underlying vector.
+ *
+ * @param name the name of the vector
+ * @param allocator the buffer allocator
+ * @param underlyingVector the underlying FixedSizeBinaryVector for storage
+ */
+ public UuidVector(
+ String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) {
+ super(name, allocator, underlyingVector);
+ this.field = new Field(name, FieldType.nullable(UuidType.INSTANCE), null);
+ }
+
+ /**
+ * Constructs a UUID vector with the given name, field type, allocator, and underlying vector.
+ *
+ * @param name the name of the vector
+ * @param fieldType the field type (should contain UuidType)
+ * @param allocator the buffer allocator
+ * @param underlyingVector the underlying FixedSizeBinaryVector for storage
+ */
+ public UuidVector(
+ String name,
+ FieldType fieldType,
+ BufferAllocator allocator,
+ FixedSizeBinaryVector underlyingVector) {
+ super(name, allocator, underlyingVector);
+ this.field = new Field(name, fieldType, null);
+ }
+
+ /**
+ * Constructs a UUID vector with the given name and allocator.
+ *
+ * Creates 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/VectorSchemaRoot.java b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java
index a7cb9ced72..4c1fbf761a 100644
--- a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java
+++ b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java
@@ -199,13 +199,18 @@ public FieldVector getVector(int index) {
*/
public VectorSchemaRoot addVector(int index, FieldVector vector) {
Preconditions.checkNotNull(vector);
- Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
+ Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size());
List newVectors = new ArrayList<>();
- for (int i = 0; i < fieldVectors.size(); i++) {
- if (i == index) {
- newVectors.add(vector);
+ if (index == fieldVectors.size()) {
+ newVectors.addAll(fieldVectors);
+ newVectors.add(vector);
+ } else {
+ for (int i = 0; i < fieldVectors.size(); i++) {
+ if (i == index) {
+ newVectors.add(vector);
+ }
+ newVectors.add(fieldVectors.get(i));
}
- newVectors.add(fieldVectors.get(i));
}
return new VectorSchemaRoot(newVectors);
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java b/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java
index abcf312c5e..bc2e3a6aab 100644
--- a/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java
+++ b/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java
@@ -43,6 +43,7 @@
import org.apache.arrow.vector.complex.ListViewVector;
import org.apache.arrow.vector.complex.NonNullableStructVector;
import org.apache.arrow.vector.complex.RunEndEncodedVector;
+import org.apache.arrow.vector.complex.RunEndEncodedVector.RangeIterator;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
@@ -270,42 +271,35 @@ protected boolean compareRunEndEncodedVectors(Range range) {
RunEndEncodedVector leftVector = (RunEndEncodedVector) left;
RunEndEncodedVector rightVector = (RunEndEncodedVector) right;
- final int leftRangeEnd = range.getLeftStart() + range.getLength();
- final int rightRangeEnd = range.getRightStart() + range.getLength();
+ final RunEndEncodedVector.RangeIterator leftIterator =
+ new RunEndEncodedVector.RangeIterator(leftVector, range.getLeftStart(), range.getLength());
+ final RunEndEncodedVector.RangeIterator rightIterator =
+ new RunEndEncodedVector.RangeIterator(
+ rightVector, range.getRightStart(), range.getLength());
FieldVector leftValuesVector = leftVector.getValuesVector();
FieldVector rightValuesVector = rightVector.getValuesVector();
RangeEqualsVisitor innerVisitor = createInnerVisitor(leftValuesVector, rightValuesVector, null);
- int leftLogicalIndex = range.getLeftStart();
- int rightLogicalIndex = range.getRightStart();
+ while (nextRun(leftIterator, rightIterator)) {
+ int leftPhysicalIndex = leftIterator.getRunIndex();
+ int rightPhysicalIndex = rightIterator.getRunIndex();
- while (leftLogicalIndex < leftRangeEnd) {
- // TODO: implement it more efficient
- // https://github.com/apache/arrow/issues/44157
- int leftPhysicalIndex = leftVector.getPhysicalIndex(leftLogicalIndex);
- int rightPhysicalIndex = rightVector.getPhysicalIndex(rightLogicalIndex);
- if (leftValuesVector.accept(
- innerVisitor, new Range(leftPhysicalIndex, rightPhysicalIndex, 1))) {
- int leftRunEnd = leftVector.getRunEnd(leftLogicalIndex);
- int rightRunEnd = rightVector.getRunEnd(rightLogicalIndex);
-
- int leftRunLength = Math.min(leftRunEnd, leftRangeEnd) - leftLogicalIndex;
- int rightRunLength = Math.min(rightRunEnd, rightRangeEnd) - rightLogicalIndex;
-
- if (leftRunLength != rightRunLength) {
- return false;
- } else {
- leftLogicalIndex = leftRunEnd;
- rightLogicalIndex = rightRunEnd;
- }
- } else {
+ if (leftIterator.getRunLength() != rightIterator.getRunLength()
+ || !leftValuesVector.accept(
+ innerVisitor, new Range(leftPhysicalIndex, rightPhysicalIndex, 1))) {
return false;
}
}
- return true;
+ return leftIterator.isEnd() && rightIterator.isEnd();
+ }
+
+ private static boolean nextRun(RangeIterator leftIterator, RangeIterator rightIterator) {
+ boolean left = leftIterator.nextRun();
+ boolean right = rightIterator.nextRun();
+ return left && right;
}
protected RangeEqualsVisitor createInnerVisitor(
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java
index 2921e43cb6..a57fbe473f 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java
@@ -46,11 +46,13 @@ public abstract class AbstractStructVector extends AbstractContainerVector {
private ConflictPolicy conflictPolicy;
static {
- String conflictPolicyStr =
- System.getProperty(STRUCT_CONFLICT_POLICY_JVM, ConflictPolicy.CONFLICT_REPLACE.toString());
+ String conflictPolicyStr = System.getProperty(STRUCT_CONFLICT_POLICY_JVM);
if (conflictPolicyStr == null) {
conflictPolicyStr = System.getenv(STRUCT_CONFLICT_POLICY_ENV);
}
+ if (conflictPolicyStr == null) {
+ conflictPolicyStr = ConflictPolicy.CONFLICT_REPLACE.toString();
+ }
ConflictPolicy conflictPolicy;
try {
conflictPolicy = ConflictPolicy.valueOf(conflictPolicyStr.toUpperCase(Locale.ROOT));
@@ -62,11 +64,11 @@ public abstract class AbstractStructVector extends AbstractContainerVector {
/** Policy to determine how to react when duplicate columns are encountered. */
public enum ConflictPolicy {
- // Ignore the conflict and append the field. This is the default behaviour
+ // Ignore the conflict and append the field.
CONFLICT_APPEND,
// Keep the existing field and ignore the newer one.
CONFLICT_IGNORE,
- // Replace the existing field with the newer one.
+ // Replace the existing field with the newer one. This is the default behaviour
CONFLICT_REPLACE,
// Refuse the new field and error out.
CONFLICT_ERROR
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java
index 12edd6557b..fac3f86bba 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java
@@ -52,7 +52,6 @@ public abstract class BaseLargeRepeatedValueViewVector extends BaseValueVector
protected ArrowBuf sizeBuffer;
protected FieldVector vector;
protected final CallBack repeatedCallBack;
- protected int valueCount;
protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH;
protected long sizeAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH;
private final String name;
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java
index fbe83bad52..ee1d65d3e3 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java
@@ -54,7 +54,6 @@ public abstract class BaseRepeatedValueVector extends BaseValueVector
protected ArrowBuf offsetBuffer;
protected FieldVector vector;
protected final CallBack repeatedCallBack;
- protected int valueCount;
protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH;
private final String name;
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java
index e6213316b5..fd7a4ff2c6 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java
@@ -52,7 +52,6 @@ public abstract class BaseRepeatedValueViewVector extends BaseValueVector
protected ArrowBuf sizeBuffer;
protected FieldVector vector;
protected final CallBack repeatedCallBack;
- protected int valueCount;
protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH;
protected long sizeAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH;
private final String name;
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java
index c762eb5172..e3b4ab477f 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java
@@ -69,12 +69,10 @@ public static FixedSizeListVector empty(String name, int size, BufferAllocator a
}
private FieldVector vector;
- private ArrowBuf validityBuffer;
private final int listSize;
private Field field;
private UnionFixedSizeListReader reader;
- private int valueCount;
private int validityAllocationSizeInBytes;
/**
@@ -110,7 +108,8 @@ public FixedSizeListVector(
this.listSize = ((ArrowType.FixedSizeList) field.getFieldType().getType()).getListSize();
Preconditions.checkArgument(listSize >= 0, "list size must be non-negative");
this.valueCount = 0;
- this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
+ this.validityAllocationSizeInBytes =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
}
@Override
@@ -189,7 +188,7 @@ public List getFieldBuffers() {
private void setReaderAndWriterIndex() {
validityBuffer.readerIndex(0);
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
}
/**
@@ -247,12 +246,10 @@ public boolean allocateNewSafe() {
return success;
}
- private void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- validityAllocationSizeInBytes = curSize;
- validityBuffer.setZero(0, validityBuffer.capacity());
+ @Override
+ protected void allocateValidityBuffer(final long size) {
+ super.allocateValidityBuffer(size);
+ validityAllocationSizeInBytes = (int) size;
}
@Override
@@ -268,7 +265,8 @@ private void reallocValidityBuffer() {
if (validityAllocationSizeInBytes > 0) {
newAllocationSize = validityAllocationSizeInBytes;
} else {
- newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+ newAllocationSize =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
}
}
@@ -311,7 +309,7 @@ public UnionFixedSizeListWriter getWriter() {
@Override
public void setInitialCapacity(int numRecords) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
vector.setInitialCapacity(numRecords * listSize);
}
@@ -328,7 +326,7 @@ public int getBufferSize() {
if (getValueCount() == 0) {
return 0;
}
- return getValidityBufferSizeFromCount(valueCount) + vector.getBufferSize();
+ return BitVectorHelper.getValidityBufferSizeFromCount(valueCount) + vector.getBufferSize();
}
@Override
@@ -336,7 +334,7 @@ public int getBufferSizeFor(int valueCount) {
if (valueCount == 0) {
return 0;
}
- return getValidityBufferSizeFromCount(valueCount)
+ return BitVectorHelper.getValidityBufferSizeFromCount(valueCount)
+ vector.getBufferSizeFor(valueCount * listSize);
}
@@ -647,71 +645,6 @@ public void splitAndTransfer(int startIndex, int length) {
to.setValueCount(length);
}
- /*
- * transfer the validity.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, FixedSizeListVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain(1);
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
- }
- }
-
@Override
public ValueVector getTo() {
return to;
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
index ed075352c9..92dd3eaef7 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
@@ -31,6 +31,7 @@
import org.apache.arrow.memory.util.ArrowBufPointer;
import org.apache.arrow.memory.util.ByteFunctionHelpers;
import org.apache.arrow.memory.util.CommonUtil;
+import org.apache.arrow.memory.util.LargeMemoryUtil;
import org.apache.arrow.memory.util.hash.ArrowBufHasher;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.AddOrGetResult;
@@ -94,11 +95,9 @@ public static LargeListVector empty(String name, BufferAllocator allocator) {
protected ArrowBuf offsetBuffer;
protected FieldVector vector;
protected final CallBack callBack;
- protected int valueCount;
protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH;
protected String defaultDataVectorName = DATA_VECTOR_NAME;
- protected ArrowBuf validityBuffer;
protected UnionLargeListReader reader;
private Field field;
private int validityAllocationSizeInBytes;
@@ -131,7 +130,8 @@ public LargeListVector(Field field, BufferAllocator allocator, CallBack callBack
this.field = field;
this.validityBuffer = allocator.getEmpty();
this.callBack = callBack;
- this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
+ this.validityAllocationSizeInBytes =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
this.lastSet = -1;
this.offsetBuffer = allocator.getEmpty();
this.vector = vector == null ? DEFAULT_DATA_VECTOR : vector;
@@ -156,7 +156,7 @@ public void initializeChildrenFromFields(List children) {
@Override
public void setInitialCapacity(int numRecords) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
offsetAllocationSizeInBytes = (long) (numRecords + 1) * OFFSET_WIDTH;
if (vector instanceof BaseFixedWidthVector || vector instanceof BaseVariableWidthVector) {
vector.setInitialCapacity(numRecords * RepeatedValueVector.DEFAULT_REPEAT_PER_RECORD);
@@ -184,7 +184,7 @@ public void setInitialCapacity(int numRecords) {
*/
@Override
public void setInitialCapacity(int numRecords, double density) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
if ((numRecords * density) >= Integer.MAX_VALUE) {
throw new OversizedAllocationException("Requested amount of memory is more than max allowed");
}
@@ -309,11 +309,14 @@ private void setReaderAndWriterIndex() {
offsetBuffer.readerIndex(0);
if (valueCount == 0) {
validityBuffer.writerIndex(0);
- offsetBuffer.writerIndex(0);
} else {
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
- offsetBuffer.writerIndex((valueCount + 1) * OFFSET_WIDTH);
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
}
+ // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.
+ // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers
+ // in other libraries. According to Arrow spec, we should still output the offset buffer which
+ // is [0].
+ offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
}
/**
@@ -374,12 +377,10 @@ public boolean allocateNewSafe() {
return success;
}
- private void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- validityAllocationSizeInBytes = curSize;
- validityBuffer.setZero(0, validityBuffer.capacity());
+ @Override
+ protected void allocateValidityBuffer(final long size) {
+ super.allocateValidityBuffer(size);
+ validityAllocationSizeInBytes = (int) size;
}
protected ArrowBuf allocateOffsetBuffer(final long size) {
@@ -442,7 +443,8 @@ private void reallocValidityBuffer() {
if (validityAllocationSizeInBytes > 0) {
newAllocationSize = validityAllocationSizeInBytes;
} else {
- newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+ newAllocationSize =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
}
}
newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
@@ -692,71 +694,6 @@ public void splitAndTransfer(int startIndex, int length) {
to.setValueCount(length);
}
- /*
- * transfer the validity.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, LargeListVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain(1);
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
- }
- }
-
@Override
public ValueVector getTo() {
return to;
@@ -821,7 +758,7 @@ public int getBufferSize() {
return 0;
}
final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH;
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return offsetBufferSize + validityBufferSize + vector.getBufferSize();
}
@@ -830,7 +767,7 @@ public int getBufferSizeFor(int valueCount) {
if (valueCount == 0) {
return 0;
}
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
long innerVectorValueCount = offsetBuffer.getLong((long) valueCount * OFFSET_WIDTH);
return ((valueCount + 1) * OFFSET_WIDTH)
@@ -928,10 +865,11 @@ public List getObject(int index) {
if (isSet(index) == 0) {
return null;
}
- final List vals = new JsonStringArrayList<>();
final long start = offsetBuffer.getLong((long) index * OFFSET_WIDTH);
final long end = offsetBuffer.getLong(((long) index + 1L) * OFFSET_WIDTH);
final ValueVector vv = getDataVector();
+ final List vals =
+ new JsonStringArrayList<>(LargeMemoryUtil.checkedCastToInt(end - start));
for (long i = start; i < end; i++) {
vals.add(vv.getObject(checkedCastToInt(i)));
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java
index 84c6f03edb..2da7eb057e 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java
@@ -77,7 +77,6 @@
public class LargeListViewVector extends BaseLargeRepeatedValueViewVector
implements PromotableVector, ValueIterableVector> {
- protected ArrowBuf validityBuffer;
protected UnionLargeListViewReader reader;
private CallBack callBack;
protected Field field;
@@ -113,7 +112,8 @@ public LargeListViewVector(Field field, BufferAllocator allocator, CallBack call
this.validityBuffer = allocator.getEmpty();
this.field = field;
this.callBack = callBack;
- this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
+ this.validityAllocationSizeInBytes =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
}
@Override
@@ -134,7 +134,7 @@ public void initializeChildrenFromFields(List children) {
@Override
public void setInitialCapacity(int numRecords) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords);
}
@@ -157,7 +157,7 @@ public void setInitialCapacity(int numRecords) {
*/
@Override
public void setInitialCapacity(int numRecords, double density) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords, density);
}
@@ -176,7 +176,7 @@ public void setInitialCapacity(int numRecords, double density) {
*/
@Override
public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialTotalCapacity(numRecords, totalNumberOfElements);
}
@@ -226,7 +226,7 @@ private void setReaderAndWriterIndex() {
offsetBuffer.writerIndex(0);
sizeBuffer.writerIndex(0);
} else {
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH);
sizeBuffer.writerIndex((long) valueCount * SIZE_WIDTH);
}
@@ -284,12 +284,10 @@ public boolean allocateNewSafe() {
return success;
}
+ @Override
protected void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- validityAllocationSizeInBytes = curSize;
- validityBuffer.setZero(0, validityBuffer.capacity());
+ super.allocateValidityBuffer(size);
+ validityAllocationSizeInBytes = (int) size;
}
@Override
@@ -323,7 +321,8 @@ private long getNewAllocationSize(int currentBufferCapacity) {
if (validityAllocationSizeInBytes > 0) {
newAllocationSize = validityAllocationSizeInBytes;
} else {
- newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+ newAllocationSize =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
}
}
newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
@@ -529,71 +528,6 @@ public void splitAndTransfer(int startIndex, int length) {
}
}
- /*
- * transfer the validity.
- */
- private void splitAndTransferValidityBuffer(
- int startIndex, int length, LargeListViewVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain(1);
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
- }
- }
-
@Override
public ValueVector getTo() {
return to;
@@ -629,7 +563,7 @@ public int getBufferSize() {
}
final int offsetBufferSize = valueCount * OFFSET_WIDTH;
final int sizeBufferSize = valueCount * SIZE_WIDTH;
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return offsetBufferSize + sizeBufferSize + validityBufferSize + vector.getBufferSize();
}
@@ -644,7 +578,7 @@ public int getBufferSizeFor(int valueCount) {
if (valueCount == 0) {
return 0;
}
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return super.getBufferSizeFor(valueCount) + validityBufferSize;
}
@@ -738,10 +672,10 @@ public List getObject(int index) {
if (isSet(index) == 0) {
return null;
}
- final List vals = new JsonStringArrayList<>();
final int start = offsetBuffer.getInt(index * OFFSET_WIDTH);
final int end = start + sizeBuffer.getInt((index) * SIZE_WIDTH);
final ValueVector vv = getDataVector();
+ final List vals = new JsonStringArrayList<>(end - start);
for (int i = start; i < end; i++) {
vals.add(vv.getObject(checkedCastToInt(i)));
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
index 3daeb6d77b..6c3993df63 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
@@ -74,7 +74,6 @@ public static ListVector empty(String name, BufferAllocator allocator) {
return new ListVector(name, allocator, FieldType.nullable(ArrowType.List.INSTANCE), null);
}
- protected ArrowBuf validityBuffer;
protected UnionListReader reader;
private CallBack callBack;
protected Field field;
@@ -108,7 +107,8 @@ public ListVector(Field field, BufferAllocator allocator, CallBack callBack) {
this.validityBuffer = allocator.getEmpty();
this.field = field;
this.callBack = callBack;
- this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
+ this.validityAllocationSizeInBytes =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
this.lastSet = -1;
}
@@ -130,7 +130,7 @@ public void initializeChildrenFromFields(List children) {
@Override
public void setInitialCapacity(int numRecords) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords);
}
@@ -153,7 +153,7 @@ public void setInitialCapacity(int numRecords) {
*/
@Override
public void setInitialCapacity(int numRecords, double density) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords, density);
}
@@ -172,7 +172,7 @@ public void setInitialCapacity(int numRecords, double density) {
*/
@Override
public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialTotalCapacity(numRecords, totalNumberOfElements);
}
@@ -267,11 +267,14 @@ private void setReaderAndWriterIndex() {
offsetBuffer.readerIndex(0);
if (valueCount == 0) {
validityBuffer.writerIndex(0);
- offsetBuffer.writerIndex(0);
} else {
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
- offsetBuffer.writerIndex((valueCount + 1) * OFFSET_WIDTH);
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
}
+ // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.
+ // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers
+ // in other libraries. According to Arrow spec, we should still output the offset buffer which
+ // is [0].
+ offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
}
/**
@@ -323,12 +326,10 @@ public boolean allocateNewSafe() {
return success;
}
+ @Override
protected void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- validityAllocationSizeInBytes = curSize;
- validityBuffer.setZero(0, validityBuffer.capacity());
+ super.allocateValidityBuffer(size);
+ validityAllocationSizeInBytes = (int) size;
}
/**
@@ -366,7 +367,8 @@ private long getNewAllocationSize(int currentBufferCapacity) {
if (validityAllocationSizeInBytes > 0) {
newAllocationSize = validityAllocationSizeInBytes;
} else {
- newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+ newAllocationSize =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
}
}
newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
@@ -573,70 +575,6 @@ public void splitAndTransfer(int startIndex, int length) {
}
}
- /*
- * transfer the validity.
- */
- private void splitAndTransferValidityBuffer(int startIndex, int length, ListVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain(1);
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
- }
- }
-
@Override
public ValueVector getTo() {
return to;
@@ -678,7 +616,7 @@ public int getBufferSize() {
return 0;
}
final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH;
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return offsetBufferSize + validityBufferSize + vector.getBufferSize();
}
@@ -687,7 +625,7 @@ public int getBufferSizeFor(int valueCount) {
if (valueCount == 0) {
return 0;
}
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return super.getBufferSizeFor(valueCount) + validityBufferSize;
}
@@ -784,10 +722,10 @@ public List getObject(int index) {
if (isSet(index) == 0) {
return null;
}
- final List vals = new JsonStringArrayList<>();
final int start = offsetBuffer.getInt(index * OFFSET_WIDTH);
final int end = offsetBuffer.getInt((index + 1) * OFFSET_WIDTH);
final ValueVector vv = getDataVector();
+ final List vals = new JsonStringArrayList<>(end - start);
for (int i = start; i < end; i++) {
vals.add(vv.getObject(i));
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java
index 9b4e6b4c0c..d41f61e291 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java
@@ -76,7 +76,6 @@
public class ListViewVector extends BaseRepeatedValueViewVector
implements PromotableVector, ValueIterableVector> {
- protected ArrowBuf validityBuffer;
protected UnionListViewReader reader;
private CallBack callBack;
protected Field field;
@@ -112,7 +111,8 @@ public ListViewVector(Field field, BufferAllocator allocator, CallBack callBack)
this.validityBuffer = allocator.getEmpty();
this.field = field;
this.callBack = callBack;
- this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
+ this.validityAllocationSizeInBytes =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION);
}
@Override
@@ -133,7 +133,7 @@ public void initializeChildrenFromFields(List children) {
@Override
public void setInitialCapacity(int numRecords) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords);
}
@@ -156,7 +156,7 @@ public void setInitialCapacity(int numRecords) {
*/
@Override
public void setInitialCapacity(int numRecords, double density) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords, density);
}
@@ -175,7 +175,7 @@ public void setInitialCapacity(int numRecords, double density) {
*/
@Override
public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) {
- validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
+ validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords);
super.setInitialTotalCapacity(numRecords, totalNumberOfElements);
}
@@ -225,9 +225,9 @@ private void setReaderAndWriterIndex() {
offsetBuffer.writerIndex(0);
sizeBuffer.writerIndex(0);
} else {
- validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
- offsetBuffer.writerIndex(valueCount * OFFSET_WIDTH);
- sizeBuffer.writerIndex(valueCount * SIZE_WIDTH);
+ validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
+ offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH);
+ sizeBuffer.writerIndex((long) valueCount * SIZE_WIDTH);
}
}
@@ -283,12 +283,10 @@ public boolean allocateNewSafe() {
return success;
}
+ @Override
protected void allocateValidityBuffer(final long size) {
- final int curSize = (int) size;
- validityBuffer = allocator.buffer(curSize);
- validityBuffer.readerIndex(0);
- validityAllocationSizeInBytes = curSize;
- validityBuffer.setZero(0, validityBuffer.capacity());
+ super.allocateValidityBuffer(size);
+ validityAllocationSizeInBytes = (int) size;
}
@Override
@@ -322,7 +320,8 @@ private long getNewAllocationSize(int currentBufferCapacity) {
if (validityAllocationSizeInBytes > 0) {
newAllocationSize = validityAllocationSizeInBytes;
} else {
- newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+ newAllocationSize =
+ BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
}
}
newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
@@ -446,14 +445,22 @@ public int hashCode(int index, ArrowBufHasher hasher) {
return ArrowBufPointer.NULL_HASH_CODE;
}
int hash = 0;
- final int start = offsetBuffer.getInt(index * OFFSET_WIDTH);
- final int end = sizeBuffer.getInt(index * OFFSET_WIDTH);
+ final int start = getElementStartIndex(index);
+ final int end = getElementEndIndex(index);
for (int i = start; i < end; i++) {
hash = ByteFunctionHelpers.combineHash(hash, vector.hashCode(i, hasher));
}
return hash;
}
+ private void setElementOffsetBuffer(int index, int value) {
+ offsetBuffer.setInt((long) index * OFFSET_WIDTH, value);
+ }
+
+ private void setElementSizeBuffer(int index, int value) {
+ sizeBuffer.setInt((long) index * SIZE_WIDTH, value);
+ }
+
private class TransferImpl implements TransferPair {
ListViewVector to;
@@ -499,7 +506,6 @@ public void splitAndTransfer(int startIndex, int length) {
valueCount);
to.clear();
if (length > 0) {
- final int startPoint = offsetBuffer.getInt((long) startIndex * OFFSET_WIDTH);
// we have to scan by index since there are out-of-order offsets
to.offsetBuffer = to.allocateBuffers((long) length * OFFSET_WIDTH);
to.sizeBuffer = to.allocateBuffers((long) length * SIZE_WIDTH);
@@ -508,9 +514,9 @@ public void splitAndTransfer(int startIndex, int length) {
int maxOffsetAndSizeSum = -1;
int minOffsetValue = -1;
for (int i = 0; i < length; i++) {
- final int offsetValue = offsetBuffer.getInt((long) (startIndex + i) * OFFSET_WIDTH);
- final int sizeValue = sizeBuffer.getInt((long) (startIndex + i) * SIZE_WIDTH);
- to.sizeBuffer.setInt((long) i * SIZE_WIDTH, sizeValue);
+ final int offsetValue = getElementStartIndex(startIndex + i);
+ final int sizeValue = getElementSize(startIndex + i);
+ to.setElementSizeBuffer(i, sizeValue);
if (maxOffsetAndSizeSum < offsetValue + sizeValue) {
maxOffsetAndSizeSum = offsetValue + sizeValue;
}
@@ -521,9 +527,9 @@ public void splitAndTransfer(int startIndex, int length) {
/* splitAndTransfer the offset buffer */
for (int i = 0; i < length; i++) {
- final int offsetValue = offsetBuffer.getInt((long) (startIndex + i) * OFFSET_WIDTH);
+ final int offsetValue = getElementStartIndex(startIndex + i);
final int relativeOffset = offsetValue - minOffsetValue;
- to.offsetBuffer.setInt((long) i * OFFSET_WIDTH, relativeOffset);
+ to.setElementOffsetBuffer(i, relativeOffset);
}
/* splitAndTransfer the validity buffer */
@@ -536,70 +542,6 @@ public void splitAndTransfer(int startIndex, int length) {
}
}
- /*
- * transfer the validity.
- */
- private void splitAndTransferValidityBuffer(int startIndex, int length, ListViewVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain(1);
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
- }
- }
-
@Override
public ValueVector getTo() {
return to;
@@ -634,7 +576,7 @@ public int getBufferSize() {
}
final int offsetBufferSize = valueCount * OFFSET_WIDTH;
final int sizeBufferSize = valueCount * SIZE_WIDTH;
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return offsetBufferSize + sizeBufferSize + validityBufferSize + vector.getBufferSize();
}
@@ -649,7 +591,7 @@ public int getBufferSizeFor(int valueCount) {
if (valueCount == 0) {
return 0;
}
- final int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
+ final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount);
return super.getBufferSizeFor(valueCount) + validityBufferSize;
}
@@ -743,10 +685,10 @@ public List getObject(int index) {
if (isSet(index) == 0) {
return null;
}
- final List vals = new JsonStringArrayList<>();
- final int start = offsetBuffer.getInt(index * OFFSET_WIDTH);
- final int end = start + sizeBuffer.getInt((index) * SIZE_WIDTH);
+ final int start = getElementStartIndex(index);
+ final int end = getElementEndIndex(index);
final ValueVector vv = getDataVector();
+ final List vals = new JsonStringArrayList<>(end - start);
for (int i = start; i < end; i++) {
vals.add(vv.getObject(i));
}
@@ -776,7 +718,7 @@ public boolean isEmpty(int index) {
if (isNull(index)) {
return true;
} else {
- return sizeBuffer.getInt(index * SIZE_WIDTH) == 0;
+ return getElementSize(index) == 0;
}
}
@@ -787,10 +729,7 @@ public boolean isEmpty(int index) {
* @return 1 if element at given index is not null, 0 otherwise
*/
public int isSet(int index) {
- final int byteIndex = index >> 3;
- final byte b = validityBuffer.getByte(byteIndex);
- final int bitIndex = index & 7;
- return (b >> bitIndex) & 0x01;
+ return BitVectorHelper.get(validityBuffer, index);
}
/**
@@ -840,8 +779,8 @@ public void setNull(int index) {
reallocValidityAndSizeAndOffsetBuffers();
}
- offsetBuffer.setInt(index * OFFSET_WIDTH, 0);
- sizeBuffer.setInt(index * SIZE_WIDTH, 0);
+ setElementOffsetBuffer(index, 0);
+ setElementSizeBuffer(index, 0);
BitVectorHelper.unsetBit(validityBuffer, index);
}
@@ -859,11 +798,11 @@ public int startNewValue(int index) {
if (index > 0) {
final int prevOffset = getMaxViewEndChildVectorByIndex(index);
- offsetBuffer.setInt(index * OFFSET_WIDTH, prevOffset);
+ setElementOffsetBuffer(index, prevOffset);
}
BitVectorHelper.setBit(validityBuffer, index);
- return offsetBuffer.getInt(index * OFFSET_WIDTH);
+ return getElementStartIndex(index);
}
/**
@@ -901,9 +840,9 @@ private void validateInvariants(int offset, int size) {
* @param value value to set
*/
public void setOffset(int index, int value) {
- validateInvariants(value, sizeBuffer.getInt(index * SIZE_WIDTH));
+ validateInvariants(value, getElementSize(index));
- offsetBuffer.setInt(index * OFFSET_WIDTH, value);
+ setElementOffsetBuffer(index, value);
}
/**
@@ -913,9 +852,9 @@ public void setOffset(int index, int value) {
* @param value value to set
*/
public void setSize(int index, int value) {
- validateInvariants(offsetBuffer.getInt(index * SIZE_WIDTH), value);
+ validateInvariants(getElementStartIndex(index), value);
- sizeBuffer.setInt(index * SIZE_WIDTH, value);
+ setElementSizeBuffer(index, value);
}
/**
@@ -951,12 +890,16 @@ public void setValueCount(int valueCount) {
@Override
public int getElementStartIndex(int index) {
- return offsetBuffer.getInt(index * OFFSET_WIDTH);
+ return offsetBuffer.getInt((long) index * OFFSET_WIDTH);
+ }
+
+ private int getElementSize(int index) {
+ return sizeBuffer.getInt((long) index * SIZE_WIDTH);
}
@Override
public int getElementEndIndex(int index) {
- return sizeBuffer.getInt(index * OFFSET_WIDTH);
+ return getElementStartIndex(index) + getElementSize(index);
}
@Override
@@ -1013,8 +956,8 @@ public double getDensity() {
@Override
public void validate() {
for (int i = 0; i < valueCount; i++) {
- final int offset = offsetBuffer.getInt(i * OFFSET_WIDTH);
- final int size = sizeBuffer.getInt(i * SIZE_WIDTH);
+ final int offset = getElementStartIndex(i);
+ final int size = getElementSize(i);
validateInvariants(offset, size);
}
}
@@ -1026,6 +969,6 @@ public void validate() {
* @param size number of elements in the list that was written
*/
public void endValue(int index, int size) {
- sizeBuffer.setInt(index * SIZE_WIDTH, size);
+ setElementSizeBuffer(index, size);
}
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java
index 23cda8401b..3f98322ba9 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java
@@ -22,7 +22,6 @@
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.AddOrGetResult;
-import org.apache.arrow.vector.BitVectorHelper;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.ZeroVector;
@@ -232,70 +231,6 @@ public void splitAndTransfer(int startIndex, int length) {
}
}
- /*
- * transfer the validity.
- */
- private void splitAndTransferValidityBuffer(int startIndex, int length, MapVector target) {
- int firstByteSource = BitVectorHelper.byteIndex(startIndex);
- int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = getValidityBufferSizeFromCount(length);
- int offset = startIndex % 8;
-
- if (length > 0) {
- if (offset == 0) {
- // slice
- if (target.validityBuffer != null) {
- target.validityBuffer.getReferenceManager().release();
- }
- target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
- target.validityBuffer.getReferenceManager().retain(1);
- } else {
- /* Copy data
- * When the first bit starts from the middle of a byte (offset != 0),
- * copy data from src BitVector.
- * Each byte in the target is composed by a part in i-th byte,
- * another part in (i+1)-th byte.
- */
- target.allocateValidityBuffer(byteSizeTarget);
-
- for (int i = 0; i < byteSizeTarget - 1; i++) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + i + 1, offset);
-
- target.validityBuffer.setByte(i, (b1 + b2));
- }
-
- /* Copying the last piece is done in the following manner:
- * if the source vector has 1 or more bytes remaining, we copy
- * the last piece as a byte formed by shifting data
- * from the current byte and the next byte.
- *
- * if the source vector has no more bytes remaining
- * (we are at the last byte), we copy the last piece as a byte
- * by shifting data from the current byte.
- */
- if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- byte b2 =
- BitVectorHelper.getBitsFromNextByte(
- validityBuffer, firstByteSource + byteSizeTarget, offset);
-
- target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
- } else {
- byte b1 =
- BitVectorHelper.getBitsFromCurrentByte(
- validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
- target.validityBuffer.setByte(byteSizeTarget - 1, b1);
- }
- }
- }
- }
-
@Override
public ValueVector getTo() {
return to;
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java
index 1bb9a3d6c0..b83e13449a 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java
@@ -28,6 +28,7 @@
import org.apache.arrow.memory.OutOfMemoryException;
import org.apache.arrow.memory.util.ByteFunctionHelpers;
import org.apache.arrow.memory.util.hash.ArrowBufHasher;
+import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.BaseIntVector;
import org.apache.arrow.vector.BaseValueVector;
import org.apache.arrow.vector.BigIntVector;
@@ -820,4 +821,101 @@ static int getPhysicalIndex(FieldVector runEndVector, int logicalIndex) {
return result;
}
+
+ public static class RangeIterator {
+
+ private final RunEndEncodedVector runEndEncodedVector;
+ private final int rangeEnd;
+ private int runIndex;
+ private int runEnd;
+ private int logicalPos;
+
+ /**
+ * Constructs a new RangeIterator for iterating over a range of values in a RunEndEncodedVector.
+ *
+ * @param runEndEncodedVector The vector to iterate over
+ * @param startIndex The logical start index of the range (inclusive)
+ * @param length The number of values to include in the range
+ * @throws IllegalArgumentException if startIndex is negative or (startIndex + length) exceeds
+ * vector bounds
+ */
+ public RangeIterator(RunEndEncodedVector runEndEncodedVector, int startIndex, int length) {
+ int rangeEnd = startIndex + length;
+ Preconditions.checkArgument(
+ startIndex >= 0, "startIndex %s must be non negative.", startIndex);
+ Preconditions.checkArgument(
+ rangeEnd <= runEndEncodedVector.getValueCount(),
+ "(startIndex + length) %s out of range[0, %s].",
+ rangeEnd,
+ runEndEncodedVector.getValueCount());
+
+ this.rangeEnd = rangeEnd;
+ this.runEndEncodedVector = runEndEncodedVector;
+ this.runIndex = runEndEncodedVector.getPhysicalIndex(startIndex) - 1;
+ this.runEnd = startIndex;
+ this.logicalPos = -1;
+ }
+
+ /**
+ * Advances to the next run in the range.
+ *
+ * @return true if there is another run available, false if iteration has completed
+ */
+ public boolean nextRun() {
+ logicalPos = runEnd;
+ if (logicalPos >= rangeEnd) {
+ return false;
+ }
+ updateRun();
+ return true;
+ }
+
+ private void updateRun() {
+ runIndex++;
+ runEnd = (int) ((BaseIntVector) runEndEncodedVector.runEndsVector).getValueAsLong(runIndex);
+ }
+
+ /**
+ * Advances to the next value in the range.
+ *
+ * @return true if there is another value available, false if iteration has completed
+ */
+ public boolean nextValue() {
+ logicalPos++;
+ if (logicalPos >= rangeEnd) {
+ return false;
+ }
+ if (logicalPos == runEnd) {
+ updateRun();
+ }
+ return true;
+ }
+
+ /**
+ * Gets the current run index (physical position in the run-ends vector).
+ *
+ * @return the current run index
+ */
+ public int getRunIndex() {
+ return runIndex;
+ }
+
+ /**
+ * Gets the length of the current run within the iterator's range.
+ *
+ * @return the number of remaining values in current run within the iterator's range
+ */
+ public int getRunLength() {
+ return Math.min(runEnd, rangeEnd) - logicalPos;
+ }
+
+ /**
+ * Checks if iteration has completed.
+ *
+ * @return true if all values in the range have been processed, false otherwise
+ */
+ public boolean isEnd() {
+ return logicalPos >= rangeEnd;
+ }
+ }
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java
index ca5f572034..5e5bb7fc21 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java
@@ -18,6 +18,7 @@
import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt;
import static org.apache.arrow.util.Preconditions.checkNotNull;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import java.util.ArrayList;
import java.util.Arrays;
@@ -89,7 +90,7 @@ public StructVector(
super(name, checkNotNull(allocator), fieldType, callBack);
this.validityBuffer = allocator.getEmpty();
this.validityAllocationSizeInBytes =
- BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION);
+ getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION);
}
/**
@@ -118,7 +119,7 @@ public StructVector(
allowConflictPolicyChanges);
this.validityBuffer = allocator.getEmpty();
this.validityAllocationSizeInBytes =
- BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION);
+ getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION);
}
/**
@@ -132,7 +133,7 @@ public StructVector(Field field, BufferAllocator allocator, CallBack callBack) {
super(field, checkNotNull(allocator), callBack);
this.validityBuffer = allocator.getEmpty();
this.validityAllocationSizeInBytes =
- BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION);
+ getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION);
}
/**
@@ -153,7 +154,7 @@ public StructVector(
super(field, checkNotNull(allocator), callBack, conflictPolicy, allowConflictPolicyChanges);
this.validityBuffer = allocator.getEmpty();
this.validityAllocationSizeInBytes =
- BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION);
+ getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION);
}
@Override
@@ -182,7 +183,7 @@ public List getFieldBuffers() {
private void setReaderAndWriterIndex() {
validityBuffer.readerIndex(0);
- validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSize(valueCount));
+ validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount));
}
/**
@@ -318,7 +319,7 @@ public void splitAndTransfer(int startIndex, int length) {
private void splitAndTransferValidityBuffer(int startIndex, int length, StructVector target) {
int firstByteSource = BitVectorHelper.byteIndex(startIndex);
int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
- int byteSizeTarget = BitVectorHelper.getValidityBufferSize(length);
+ int byteSizeTarget = getValidityBufferSizeFromCount(length);
int offset = startIndex % 8;
if (length > 0) {
@@ -464,7 +465,7 @@ public int getBufferSize() {
if (valueCount == 0) {
return 0;
}
- return super.getBufferSize() + BitVectorHelper.getValidityBufferSize(valueCount);
+ return super.getBufferSize() + getValidityBufferSizeFromCount(valueCount);
}
/**
@@ -478,18 +479,18 @@ public int getBufferSizeFor(final int valueCount) {
if (valueCount == 0) {
return 0;
}
- return super.getBufferSizeFor(valueCount) + BitVectorHelper.getValidityBufferSize(valueCount);
+ return super.getBufferSizeFor(valueCount) + getValidityBufferSizeFromCount(valueCount);
}
@Override
public void setInitialCapacity(int numRecords) {
- validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSize(numRecords);
+ validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords);
}
@Override
public void setInitialCapacity(int numRecords, double density) {
- validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSize(numRecords);
+ validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);
super.setInitialCapacity(numRecords, density);
}
@@ -547,7 +548,7 @@ private long getNewAllocationSize(int currentBufferCapacity) {
newAllocationSize = validityAllocationSizeInBytes;
} else {
newAllocationSize =
- BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION) * 2L;
+ getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION) * 2L;
}
}
newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java
deleted file mode 100644
index 09f0314c5f..0000000000
--- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java
+++ /dev/null
@@ -1,38 +0,0 @@
-/*
- * 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.complex.impl;
-
-import org.apache.arrow.vector.ExtensionTypeVector;
-import org.apache.arrow.vector.complex.writer.FieldWriter;
-
-/**
- * A factory interface for creating instances of {@link ExtensionTypeWriter}. This factory allows
- * configuring writer implementations for specific {@link ExtensionTypeVector}.
- *
- * @param the type of writer implementation for a specific {@link ExtensionTypeVector}.
- */
-public interface ExtensionTypeWriterFactory {
-
- /**
- * Returns an instance of the writer implementation for the given {@link ExtensionTypeVector}.
- *
- * @param vector the {@link ExtensionTypeVector} for which the writer implementation is to be
- * returned.
- * @return an instance of the writer implementation for the given {@link ExtensionTypeVector}.
- */
- T getWriterImpl(ExtensionTypeVector vector);
-}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java
new file mode 100644
index 0000000000..7a5312f6ed
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java
@@ -0,0 +1,123 @@
+/*
+ * 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.complex.impl;
+
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.util.UuidUtility;
+
+/**
+ * Reader implementation for reading UUID values from a {@link NullableUuidHolder}.
+ *
+ * This reader wraps a single UUID holder value and provides methods to read from it. Unlike
+ * {@link UuidReaderImpl} which reads from a vector, this reader operates on a holder instance.
+ *
+ * @see NullableUuidHolder
+ * @see UuidReaderImpl
+ */
+public class NullableUuidHolderReaderImpl extends AbstractFieldReader {
+ private final NullableUuidHolder holder;
+
+ /**
+ * Constructs a reader for the given UUID holder.
+ *
+ * @param holder the UUID holder to read from
+ */
+ public NullableUuidHolderReaderImpl(NullableUuidHolder holder) {
+ this.holder = holder;
+ }
+
+ @Override
+ public int size() {
+ throw new UnsupportedOperationException(
+ "size() is not supported on NullableUuidHolderReaderImpl. "
+ + "This reader wraps a single UUID holder value, not a collection. "
+ + "Use UuidReaderImpl for vector-based UUID reading.");
+ }
+
+ @Override
+ public boolean next() {
+ throw new UnsupportedOperationException(
+ "next() is not supported on NullableUuidHolderReaderImpl. "
+ + "This reader wraps a single UUID holder value, not an iterator. "
+ + "Use UuidReaderImpl for vector-based UUID reading.");
+ }
+
+ @Override
+ public void setPosition(int index) {
+ throw new UnsupportedOperationException(
+ "setPosition() is not supported on NullableUuidHolderReaderImpl. "
+ + "This reader wraps a single UUID holder value, not a vector. "
+ + "Use UuidReaderImpl for vector-based UUID reading.");
+ }
+
+ @Override
+ public Types.MinorType getMinorType() {
+ return Types.MinorType.EXTENSIONTYPE;
+ }
+
+ @Override
+ public boolean isSet() {
+ return holder.isSet == 1;
+ }
+
+ @Override
+ public void read(ExtensionHolder h) {
+ if (h instanceof NullableUuidHolder) {
+ NullableUuidHolder nullableHolder = (NullableUuidHolder) h;
+ nullableHolder.buffer = this.holder.buffer;
+ nullableHolder.isSet = this.holder.isSet;
+ nullableHolder.start = this.holder.start;
+ } else if (h instanceof UuidHolder) {
+ UuidHolder uuidHolder = (UuidHolder) h;
+ uuidHolder.buffer = this.holder.buffer;
+ uuidHolder.start = this.holder.start;
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported holder type: "
+ + h.getClass().getName()
+ + ". "
+ + "Only NullableUuidHolder and UuidHolder are supported for UUID values. "
+ + "Provided holder type cannot be used to read UUID data.");
+ }
+ }
+
+ @Override
+ public Object readObject() {
+ if (!isSet()) {
+ return null;
+ }
+ // Convert UUID bytes to Java UUID object
+ try {
+ return UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ } catch (Exception e) {
+ throw new RuntimeException(
+ String.format(
+ "Failed to read UUID from buffer. Invalid Arrow buffer state: "
+ + "capacity=%d, readableBytes=%d, readerIndex=%d, writerIndex=%d, refCnt=%d. "
+ + "The buffer must contain exactly 16 bytes of valid UUID data.",
+ holder.buffer.capacity(),
+ holder.buffer.readableBytes(),
+ holder.buffer.readerIndex(),
+ holder.buffer.writerIndex(),
+ holder.buffer.refCnt()),
+ e);
+ }
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java
index d341384bd9..93796aa77e 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java
@@ -60,11 +60,6 @@ public void writeExtension(Object var1) {
}
@Override
- public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) {
- this.writer = factory.getWriterImpl(vector);
- this.writer.setPosition(idx());
- }
-
public void write(ExtensionHolder holder) {
this.writer.write(holder);
}
@@ -76,4 +71,10 @@ public void setPosition(int index) {
this.writer.setPosition(index);
}
}
+
+ @Override
+ public void writeNull() {
+ this.vector.setNull(getPosition());
+ this.vector.setValueCount(getPosition() + 1);
+ }
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java
new file mode 100644
index 0000000000..bb7ae13e5b
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java
@@ -0,0 +1,94 @@
+/*
+ * 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.complex.impl;
+
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
+import org.apache.arrow.vector.types.Types.MinorType;
+import org.apache.arrow.vector.types.pojo.Field;
+
+/**
+ * Reader implementation for {@link UuidVector}.
+ *
+ *
Provides methods to read UUID values from a vector, including support for reading into {@link
+ * UuidHolder} and retrieving values as {@link java.util.UUID} objects.
+ *
+ * @see UuidVector
+ * @see org.apache.arrow.vector.extension.UuidType
+ */
+public class UuidReaderImpl extends AbstractFieldReader {
+
+ private final UuidVector vector;
+
+ /**
+ * Constructs a reader for the given UUID vector.
+ *
+ * @param vector the UUID vector to read from
+ */
+ public UuidReaderImpl(UuidVector vector) {
+ super();
+ this.vector = vector;
+ }
+
+ @Override
+ public MinorType getMinorType() {
+ return vector.getMinorType();
+ }
+
+ @Override
+ public Field getField() {
+ return vector.getField();
+ }
+
+ @Override
+ public boolean isSet() {
+ return !vector.isNull(idx());
+ }
+
+ @Override
+ public void read(ExtensionHolder holder) {
+ if (holder instanceof NullableUuidHolder) {
+ vector.get(idx(), (NullableUuidHolder) holder);
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported holder type for UuidReader: " + holder.getClass());
+ }
+ }
+
+ @Override
+ public void read(int arrayIndex, ExtensionHolder holder) {
+ if (holder instanceof NullableUuidHolder) {
+ vector.get(arrayIndex, (NullableUuidHolder) holder);
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported holder type for UuidReader: " + holder.getClass());
+ }
+ }
+
+ @Override
+ public void copyAsValue(AbstractExtensionTypeWriter writer) {
+ UuidWriterImpl impl = (UuidWriterImpl) writer;
+ impl.vector.copyFromSafe(idx(), impl.idx(), vector);
+ }
+
+ @Override
+ public Object readObject() {
+ return vector.getObject(idx());
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java
new file mode 100644
index 0000000000..944b7e2e62
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java
@@ -0,0 +1,81 @@
+/*
+ * 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.complex.impl;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
+/**
+ * Writer implementation for {@link UuidVector}.
+ *
+ *
Supports writing UUID values in multiple formats: {@link java.util.UUID}, byte arrays, and
+ * {@link ArrowBuf}. Also handles {@link UuidHolder} and {@link NullableUuidHolder}.
+ *
+ * @see UuidVector
+ * @see org.apache.arrow.vector.extension.UuidType
+ */
+public class UuidWriterImpl extends AbstractExtensionTypeWriter {
+
+ /**
+ * Constructs a writer for the given UUID vector.
+ *
+ * @param vector the UUID vector to write to
+ */
+ public UuidWriterImpl(UuidVector vector) {
+ super(vector);
+ }
+
+ @Override
+ public void writeExtension(Object value) {
+ if (value instanceof byte[]) {
+ vector.setSafe(getPosition(), (byte[]) value);
+ } else if (value instanceof ArrowBuf) {
+ vector.setSafe(getPosition(), (ArrowBuf) value);
+ } else if (value instanceof java.util.UUID) {
+ vector.setSafe(getPosition(), (java.util.UUID) value);
+ } else if (value instanceof ExtensionHolder) {
+ write((ExtensionHolder) value);
+ } else {
+ throw new IllegalArgumentException(
+ "Unsupported value type for UUID: "
+ + value.getClass().getName()
+ + ". "
+ + "Supported types are: byte[] (16 bytes), ArrowBuf (16 bytes), or java.util.UUID. "
+ + "Convert your value to one of these types before writing.");
+ }
+ vector.setValueCount(getPosition() + 1);
+ }
+
+ @Override
+ public void writeExtension(Object value, ArrowType type) {
+ writeExtension(value);
+ }
+
+ @Override
+ public void write(ExtensionHolder holder) {
+ if (holder instanceof UuidHolder) {
+ vector.setSafe(getPosition(), (UuidHolder) holder);
+ } else if (holder instanceof NullableUuidHolder) {
+ vector.setSafe(getPosition(), (NullableUuidHolder) holder);
+ }
+ vector.setValueCount(getPosition() + 1);
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/reader/ExtensionReader.java b/vector/src/main/java/org/apache/arrow/vector/complex/reader/ExtensionReader.java
new file mode 100644
index 0000000000..1ba7b27156
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/reader/ExtensionReader.java
@@ -0,0 +1,44 @@
+/*
+ * 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.complex.reader;
+
+import org.apache.arrow.vector.holders.ExtensionHolder;
+
+/** Interface for reading extension types. Extends the functionality of {@link BaseReader}. */
+public interface ExtensionReader extends BaseReader {
+
+ /**
+ * Reads to the given extension holder.
+ *
+ * @param holder the {@link ExtensionHolder} to read
+ */
+ void read(ExtensionHolder holder);
+
+ /**
+ * Reads and returns an object representation of the extension type.
+ *
+ * @return the object representation of the extension type
+ */
+ Object readObject();
+
+ /**
+ * Checks if the current value is set.
+ *
+ * @return true if the value is set, false otherwise
+ */
+ boolean isSet();
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java b/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java
index ca56214fda..780a4ee659 100644
--- a/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java
+++ b/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java
@@ -54,10 +54,12 @@
import org.apache.arrow.vector.TimeStampNanoVector;
import org.apache.arrow.vector.TimeStampSecTZVector;
import org.apache.arrow.vector.TimeStampSecVector;
+import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.ViewVarBinaryVector;
import org.apache.arrow.vector.ViewVarCharVector;
+import org.apache.arrow.vector.complex.writer.FieldWriter;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
@@ -177,6 +179,11 @@ public int hashCode() {
return Objects.hash(super.hashCode(), storageType, typeName, vendorName);
}
+ @Override
+ public FieldWriter getNewFieldWriter(ValueVector vector) {
+ throw new UnsupportedOperationException("WriterImpl not yet implemented.");
+ }
+
@Override
public String toString() {
return "OpaqueType("
diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java b/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java
new file mode 100644
index 0000000000..c249c6eda9
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java
@@ -0,0 +1,119 @@
+/*
+ * 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.extension;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.FixedSizeBinaryVector;
+import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.complex.impl.UuidWriterImpl;
+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;
+
+/**
+ * Extension type for UUID (Universally Unique Identifier) values.
+ *
+ * UUIDs are stored as 16-byte fixed-size binary values. This extension type provides a
+ * standardized way to represent UUIDs in Arrow, making them interoperable across different systems
+ * and languages.Ï€
+ *
+ *
The extension name is "arrow.uuid" and it uses {@link ArrowType.FixedSizeBinary} with 16 bytes
+ * as the storage type.
+ *
+ *
Usage:
+ *
+ *
{@code
+ * UuidVector vector = new UuidVector("uuid_col", allocator);
+ * vector.set(0, UUID.randomUUID());
+ * UUID value = vector.getObject(0);
+ * }
+ *
+ * @see UuidVector
+ * @see org.apache.arrow.vector.holders.UuidHolder
+ * @see org.apache.arrow.vector.holders.NullableUuidHolder
+ */
+public class UuidType extends ExtensionType {
+ /** Singleton instance of UuidType. */
+ public static final UuidType INSTANCE = new UuidType();
+
+ /** Extension name registered in the Arrow extension type registry. */
+ public static final String EXTENSION_NAME = "arrow.uuid";
+
+ /** Number of bytes used to store a UUID (128 bits = 16 bytes). */
+ public static final int UUID_BYTE_WIDTH = 16;
+
+ /** Number of characters in the standard UUID string representation (with hyphens). */
+ public static final int UUID_STRING_WIDTH = 36;
+
+ /** Storage type for UUID: FixedSizeBinary(16). */
+ public static final ArrowType STORAGE_TYPE = new ArrowType.FixedSizeBinary(UUID_BYTE_WIDTH);
+
+ private UuidType() {}
+
+ static {
+ ExtensionTypeRegistry.register(INSTANCE);
+ }
+
+ @Override
+ public ArrowType storageType() {
+ return STORAGE_TYPE;
+ }
+
+ @Override
+ public String extensionName() {
+ return EXTENSION_NAME;
+ }
+
+ @Override
+ public boolean extensionEquals(ExtensionType other) {
+ return other instanceof UuidType;
+ }
+
+ @Override
+ public ArrowType deserialize(ArrowType storageType, String serializedData) {
+ if (!storageType.equals(storageType())) {
+ throw new UnsupportedOperationException(
+ "Cannot construct UuidType from underlying type " + storageType);
+ }
+ return INSTANCE;
+ }
+
+ @Override
+ public String serialize() {
+ return "";
+ }
+
+ @Override
+ public boolean isComplex() {
+ return false;
+ }
+
+ @Override
+ public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
+ return new UuidVector(
+ name, fieldType, allocator, new FixedSizeBinaryVector(name, allocator, UUID_BYTE_WIDTH));
+ }
+
+ @Override
+ public FieldWriter getNewFieldWriter(ValueVector vector) {
+ return new UuidWriterImpl((UuidVector) vector);
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java b/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java
index fc7ed85878..4d3f767aef 100644
--- a/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java
+++ b/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java
@@ -16,7 +16,11 @@
*/
package org.apache.arrow.vector.holders;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
/** Base {@link ValueHolder} class for a {@link org.apache.arrow.vector.ExtensionTypeVector}. */
public abstract class ExtensionHolder implements ValueHolder {
public int isSet;
+
+ public abstract ArrowType type();
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java b/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java
new file mode 100644
index 0000000000..6a2b4ff604
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java
@@ -0,0 +1,45 @@
+/*
+ * 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.holders;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
+/**
+ * Value holder for nullable UUID values.
+ *
+ * The {@code isSet} field controls nullability: when {@code isSet = 1}, the holder contains a
+ * valid UUID in {@code buffer}; when {@code isSet = 0}, the holder represents a null value and
+ * {@code buffer} should not be accessed.
+ *
+ * @see UuidHolder
+ * @see org.apache.arrow.vector.UuidVector
+ * @see org.apache.arrow.vector.extension.UuidType
+ */
+public class NullableUuidHolder extends ExtensionHolder {
+ /** Buffer containing 16-byte UUID data. */
+ public ArrowBuf buffer;
+
+ /** Offset in the buffer where the UUID data starts. */
+ public int start = 0;
+
+ @Override
+ public ArrowType type() {
+ return UuidType.INSTANCE;
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java b/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java
new file mode 100644
index 0000000000..9ec0305f30
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java
@@ -0,0 +1,48 @@
+/*
+ * 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.holders;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+
+/**
+ * Value holder for non-nullable UUID values.
+ *
+ *
Contains a 16-byte UUID in {@code buffer} with {@code isSet} always 1.
+ *
+ * @see NullableUuidHolder
+ * @see org.apache.arrow.vector.UuidVector
+ * @see org.apache.arrow.vector.extension.UuidType
+ */
+public class UuidHolder extends ExtensionHolder {
+ /** Buffer containing 16-byte UUID data. */
+ public ArrowBuf buffer;
+
+ /** Offset in the buffer where the UUID data starts. */
+ public int start = 0;
+
+ /** Constructs a UuidHolder with isSet = 1. */
+ public UuidHolder() {
+ this.isSet = 1;
+ }
+
+ @Override
+ public ArrowType type() {
+ return UuidType.INSTANCE;
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java
index fe0803d298..e4bab7eb80 100644
--- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java
+++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java
@@ -20,6 +20,7 @@
import static com.fasterxml.jackson.core.JsonToken.END_OBJECT;
import static com.fasterxml.jackson.core.JsonToken.START_ARRAY;
import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.apache.arrow.vector.BufferLayout.BufferType.DATA;
import static org.apache.arrow.vector.BufferLayout.BufferType.OFFSET;
import static org.apache.arrow.vector.BufferLayout.BufferType.SIZE;
@@ -381,7 +382,7 @@ private class BufferHelper {
new BufferReader() {
@Override
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
- final int bufferSize = BitVectorHelper.getValidityBufferSize(count);
+ final int bufferSize = getValidityBufferSizeFromCount(count);
ArrowBuf buf = allocator.buffer(bufferSize);
// C++ integration test fails without this.
diff --git a/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java b/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java
new file mode 100644
index 0000000000..a1b0b54579
--- /dev/null
+++ b/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java
@@ -0,0 +1,77 @@
+/*
+ * 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.util;
+
+import static org.apache.arrow.vector.extension.UuidType.UUID_BYTE_WIDTH;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.UUID;
+import org.apache.arrow.memory.ArrowBuf;
+
+/**
+ * Utility class for UUID conversions and operations.
+ *
+ *
Provides methods to convert between {@link UUID} objects and byte representations used in
+ * Arrow vectors.
+ *
+ * @see org.apache.arrow.vector.UuidVector
+ * @see org.apache.arrow.vector.extension.UuidType
+ */
+public class UuidUtility {
+ /**
+ * Converts a UUID to a 16-byte array.
+ *
+ *
The UUID is stored in big-endian byte order, with the most significant bits first.
+ *
+ * @param uuid the UUID to convert
+ * @return a 16-byte array representing the UUID
+ */
+ public static byte[] getBytesFromUUID(UUID uuid) {
+ byte[] result = new byte[16];
+ long msb = uuid.getMostSignificantBits();
+ long lsb = uuid.getLeastSignificantBits();
+ for (int i = 15; i >= 8; i--) {
+ result[i] = (byte) (lsb & 0xFF);
+ lsb >>= 8;
+ }
+ for (int i = 7; i >= 0; i--) {
+ result[i] = (byte) (msb & 0xFF);
+ msb >>= 8;
+ }
+ return result;
+ }
+
+ /**
+ * Constructs a UUID from bytes stored in an ArrowBuf at the specified index.
+ *
+ *
Reads 16 bytes from the buffer starting at the given index and interprets them as a UUID in
+ * big-endian byte order.
+ *
+ * @param buffer the buffer containing UUID data
+ * @param index the byte offset in the buffer where the UUID starts
+ * @return the UUID constructed from the buffer data
+ */
+ public static UUID uuidFromArrowBuf(ArrowBuf buffer, long index) {
+ ByteBuffer buf = buffer.nioBuffer(index, UUID_BYTE_WIDTH);
+
+ buf.order(ByteOrder.BIG_ENDIAN);
+ long mostSigBits = buf.getLong(0);
+ long leastSigBits = buf.getLong(Long.BYTES);
+ return new UUID(mostSigBits, leastSigBits);
+ }
+}
diff --git a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java
index 0dc96a4d4b..e7c0d11cb9 100644
--- a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java
+++ b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java
@@ -24,13 +24,17 @@
import org.apache.arrow.memory.util.MemoryUtil;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.BaseFixedWidthVector;
+import org.apache.arrow.vector.BaseIntVector;
import org.apache.arrow.vector.BaseLargeVariableWidthVector;
import org.apache.arrow.vector.BaseVariableWidthVector;
import org.apache.arrow.vector.BaseVariableWidthViewVector;
+import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.BitVectorHelper;
import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.NullVector;
+import org.apache.arrow.vector.SmallIntVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.compare.TypeEqualsVisitor;
import org.apache.arrow.vector.compare.VectorVisitor;
@@ -39,6 +43,7 @@
import org.apache.arrow.vector.complex.LargeListVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.NonNullableStructVector;
+import org.apache.arrow.vector.complex.RunEndEncodedVector;
import org.apache.arrow.vector.complex.UnionVector;
/** Utility to append two vectors together. */
@@ -698,4 +703,98 @@ public ValueVector visit(ExtensionTypeVector deltaVector, Void value) {
deltaVector.getUnderlyingVector().accept(underlyingAppender, null);
return targetVector;
}
+
+ @Override
+ public ValueVector visit(RunEndEncodedVector deltaVector, Void value) {
+ Preconditions.checkArgument(
+ typeVisitor.equals(deltaVector),
+ "The deltaVector to append must have the same type as the targetVector");
+
+ if (deltaVector.getValueCount() == 0) {
+ return targetVector; // optimization, nothing to append, return
+ }
+
+ RunEndEncodedVector targetEncodedVector = (RunEndEncodedVector) targetVector;
+
+ final int targetLogicalValueCount = targetEncodedVector.getValueCount();
+
+ // Append the values vector first.
+ VectorAppender valueAppender = new VectorAppender(targetEncodedVector.getValuesVector());
+ deltaVector.getValuesVector().accept(valueAppender, null);
+
+ // Then append the run-ends vector.
+ BaseIntVector targetRunEndsVector = (BaseIntVector) targetEncodedVector.getRunEndsVector();
+ BaseIntVector deltaRunEndsVector = (BaseIntVector) deltaVector.getRunEndsVector();
+ appendRunEndsVector(targetRunEndsVector, deltaRunEndsVector, targetLogicalValueCount);
+
+ targetEncodedVector.setValueCount(targetLogicalValueCount + deltaVector.getValueCount());
+ return targetVector;
+ }
+
+ private void appendRunEndsVector(
+ BaseIntVector targetRunEndsVector,
+ BaseIntVector deltaRunEndsVector,
+ int targetLogicalValueCount) {
+ int targetPhysicalValueCount = targetRunEndsVector.getValueCount();
+ int newPhysicalValueCount = targetPhysicalValueCount + deltaRunEndsVector.getValueCount();
+
+ // make sure there is enough capacity
+ while (targetVector.getValueCapacity() < newPhysicalValueCount) {
+ targetVector.reAlloc();
+ }
+
+ // append validity buffer
+ BitVectorHelper.concatBits(
+ targetRunEndsVector.getValidityBuffer(),
+ targetRunEndsVector.getValueCount(),
+ deltaRunEndsVector.getValidityBuffer(),
+ deltaRunEndsVector.getValueCount(),
+ targetRunEndsVector.getValidityBuffer());
+
+ // shift and append data buffer
+ shiftAndAppendRunEndsDataBuffer(
+ targetRunEndsVector,
+ targetPhysicalValueCount,
+ deltaRunEndsVector.getDataBuffer(),
+ targetLogicalValueCount,
+ deltaRunEndsVector.getValueCount());
+
+ targetRunEndsVector.setValueCount(newPhysicalValueCount);
+ }
+
+ private void shiftAndAppendRunEndsDataBuffer(
+ BaseIntVector toRunEndVector,
+ int toIndex,
+ ArrowBuf fromRunEndBuffer,
+ int offset,
+ int physicalLength) {
+ ArrowBuf toRunEndBuffer = toRunEndVector.getDataBuffer();
+ if (toRunEndVector instanceof SmallIntVector) {
+ byte typeWidth = SmallIntVector.TYPE_WIDTH;
+ for (int i = 0; i < physicalLength; i++) {
+ toRunEndBuffer.setShort(
+ (long) (i + toIndex) * typeWidth,
+ fromRunEndBuffer.getShort((long) (i) * typeWidth) + offset);
+ }
+
+ } else if (toRunEndVector instanceof IntVector) {
+ byte typeWidth = IntVector.TYPE_WIDTH;
+ for (int i = 0; i < physicalLength; i++) {
+ toRunEndBuffer.setInt(
+ (long) (i + toIndex) * typeWidth,
+ fromRunEndBuffer.getInt((long) (i) * typeWidth) + offset);
+ }
+
+ } else if (toRunEndVector instanceof BigIntVector) {
+ byte typeWidth = BigIntVector.TYPE_WIDTH;
+ for (int i = 0; i < physicalLength; i++) {
+ toRunEndBuffer.setLong(
+ (long) (i + toIndex) * typeWidth,
+ fromRunEndBuffer.getLong((long) (i) * typeWidth) + offset);
+ }
+ } else {
+ throw new IllegalArgumentException(
+ "Run-end vector and must be of type int with size 16, 32, or 64 bits.");
+ }
+ }
}
diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java
index 5c7215437f..5cfe64b14e 100644
--- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java
+++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java
@@ -52,14 +52,22 @@ private void validateVectorCommon(ValueVector vector) {
if (vector instanceof FieldVector) {
FieldVector fieldVector = (FieldVector) vector;
- // TODO: https://github.com/apache/arrow/issues/41734
int typeBufferCount = TypeLayout.getTypeBufferCount(arrowType);
- validateOrThrow(
- fieldVector.getFieldBuffers().size() == typeBufferCount,
- "Expected %s buffers in vector of type %s, got %s.",
- typeBufferCount,
- vector.getField().getType().toString(),
- fieldVector.getFieldBuffers().size());
+ if (TypeLayout.getTypeLayout(arrowType).isFixedBufferCount()) {
+ validateOrThrow(
+ fieldVector.getFieldBuffers().size() == typeBufferCount,
+ "Expected %s buffers in vector of type %s, got %s.",
+ typeBufferCount,
+ vector.getField().getType().toString(),
+ fieldVector.getFieldBuffers().size());
+ } else {
+ validateOrThrow(
+ fieldVector.getFieldBuffers().size() >= typeBufferCount,
+ "Expected at least %s buffers in vector of type %s, got %s.",
+ typeBufferCount,
+ vector.getField().getType().toString(),
+ fieldVector.getFieldBuffers().size());
+ }
}
}
@@ -158,7 +166,12 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) {
@Override
public Void visit(BaseVariableWidthViewVector vector, Void value) {
- throw new UnsupportedOperationException("View vectors are not supported.");
+ final int valueCount = vector.getValueCount();
+ validateVectorCommon(vector);
+ validateOrThrow(vector.getFieldBuffers().size() >= 2, "Expected at least 2 buffers.");
+ validateValidityBuffer(vector, valueCount);
+ validateDataBuffer(vector, (long) valueCount * BaseVariableWidthViewVector.ELEMENT_SIZE);
+ return null;
}
@Override
diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java
index c62bff79f7..9da8cc813e 100644
--- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java
+++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java
@@ -121,7 +121,8 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) {
@Override
public Void visit(BaseVariableWidthViewVector vector, Void value) {
- throw new UnsupportedOperationException("View vectors are not supported.");
+ vector.validateScalars();
+ return null;
}
@Override
diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java
index daad41dbdc..395852ef79 100644
--- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java
+++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java
@@ -61,6 +61,8 @@
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarBinaryVector;
+import org.apache.arrow.vector.ViewVarCharVector;
import org.apache.arrow.vector.compare.VectorVisitor;
import org.apache.arrow.vector.complex.DenseUnionVector;
import org.apache.arrow.vector.complex.FixedSizeListVector;
@@ -380,7 +382,12 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) {
@Override
public Void visit(BaseVariableWidthViewVector vector, Void value) {
- throw new UnsupportedOperationException("View vectors are not supported.");
+ if (vector instanceof ViewVarCharVector) {
+ validateVectorCommon(vector, ArrowType.Utf8View.class);
+ } else if (vector instanceof ViewVarBinaryVector) {
+ validateVectorCommon(vector, ArrowType.BinaryView.class);
+ }
+ return null;
}
@Override
diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java
index 5004ba488c..2111410016 100644
--- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java
+++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java
@@ -107,8 +107,13 @@ public Void visit(BaseLargeVariableWidthVector left, Void value) {
}
@Override
- public Void visit(BaseVariableWidthViewVector left, Void value) {
- throw new UnsupportedOperationException("View vectors are not supported.");
+ public Void visit(BaseVariableWidthViewVector vector, Void value) {
+ if (vector.getValueCount() > 0) {
+ if (vector.getDataBuffer() == null || vector.getDataBuffer().capacity() == 0) {
+ throw new IllegalArgumentException("valueBuffer is null or capacity is 0");
+ }
+ }
+ return null;
}
@Override
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java
index 73a88b3a1e..b3455fe52c 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java
@@ -30,14 +30,21 @@
import java.util.Arrays;
import java.util.List;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListReader;
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter;
import org.apache.arrow.vector.complex.impl.UnionListReader;
import org.apache.arrow.vector.complex.reader.FieldReader;
+import org.apache.arrow.vector.holders.DurationHolder;
+import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
+import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
+import org.apache.arrow.vector.holders.TimeStampNanoTZHolder;
+import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.Text;
import org.apache.arrow.vector.util.TransferPair;
@@ -628,6 +635,206 @@ public void testWriteLargeVarBinaryHelpers() throws Exception {
}
}
+ @Test
+ public void testWriterTimeStampNanoTZField() {
+ try (final FixedSizeListVector vector =
+ FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) {
+ UnionFixedSizeListWriter writer = vector.getWriter();
+ writer.allocate();
+
+ final int valueCount = 10;
+
+ for (int i = 0; i < valueCount; i++) {
+ writer.startList();
+ writer.timeStampNanoTZ().writeTimeStampNanoTZ(i * 1000L);
+ writer.timeStampNanoTZ().writeTimeStampNanoTZ((i + 1) * 1000L);
+ writer.timeStampNanoTZ().writeTimeStampNanoTZ((i + 2) * 1000L);
+ writer.endList();
+ }
+ vector.setValueCount(valueCount);
+
+ UnionFixedSizeListReader reader = vector.getReader();
+ for (int i = 0; i < valueCount; i++) {
+ reader.setPosition(i);
+ assertTrue(reader.isSet());
+ assertTrue(reader.next());
+ assertEquals(i * 1000L, reader.reader().readLong().longValue());
+ assertTrue(reader.next());
+ assertEquals((i + 1) * 1000L, reader.reader().readLong().longValue());
+ assertTrue(reader.next());
+ assertEquals((i + 2) * 1000L, reader.reader().readLong().longValue());
+ assertFalse(reader.next());
+ }
+ }
+ }
+
+ @Test
+ public void testWriterUsingHolderTimeStampNanoTZField() {
+ try (final FixedSizeListVector vector =
+ FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) {
+ UnionFixedSizeListWriter writer = vector.getWriter();
+ writer.allocate();
+
+ TimeStampNanoTZHolder holder = new TimeStampNanoTZHolder();
+ holder.timezone = "SomeFakeTimeZone";
+ writer.startList();
+ holder.value = 12341234L;
+ writer.timeStampNanoTZ().write(holder);
+ holder.value = 55555L;
+ writer.timeStampNanoTZ().write(holder);
+
+ // Writing with a different timezone should throw
+ holder.timezone = "AsdfTimeZone";
+ holder.value = 77777;
+ IllegalArgumentException ex =
+ assertThrows(
+ IllegalArgumentException.class, () -> writer.timeStampNanoTZ().write(holder));
+ assertEquals(
+ "holder.timezone: AsdfTimeZone not equal to vector timezone: SomeFakeTimeZone",
+ ex.getMessage());
+
+ writer.endList();
+ vector.setValueCount(1);
+
+ Field expectedDataField =
+ new Field(
+ BaseRepeatedValueVector.DATA_VECTOR_NAME,
+ FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "SomeFakeTimeZone")),
+ null);
+ Field expectedField =
+ new Field(
+ vector.getName(),
+ FieldType.nullable(new ArrowType.FixedSizeList(3)),
+ List.of(expectedDataField));
+
+ assertEquals(expectedField, writer.getField());
+ }
+ }
+
+ @Test
+ public void testWriterUsingHolderTimestampMilliTZField() {
+ try (final FixedSizeListVector vector =
+ FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) {
+ UnionFixedSizeListWriter writer = vector.getWriter();
+ writer.allocate();
+
+ TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder();
+ holder.timezone = "SomeFakeTimeZone";
+ writer.startList();
+ holder.value = 12341234L;
+ writer.timeStampMilliTZ().write(holder);
+ holder.value = 55555L;
+ writer.timeStampMilliTZ().write(holder);
+
+ // Writing with a different timezone should throw
+ holder.timezone = "AsdfTimeZone";
+ holder.value = 77777;
+ IllegalArgumentException ex =
+ assertThrows(
+ IllegalArgumentException.class, () -> writer.timeStampMilliTZ().write(holder));
+ assertEquals(
+ "holder.timezone: AsdfTimeZone not equal to vector timezone: SomeFakeTimeZone",
+ ex.getMessage());
+
+ writer.endList();
+ vector.setValueCount(1);
+
+ Field expectedDataField =
+ new Field(
+ BaseRepeatedValueVector.DATA_VECTOR_NAME,
+ FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "SomeFakeTimeZone")),
+ null);
+ Field expectedField =
+ new Field(
+ vector.getName(),
+ FieldType.nullable(new ArrowType.FixedSizeList(3)),
+ List.of(expectedDataField));
+
+ assertEquals(expectedField, writer.getField());
+ }
+ }
+
+ @Test
+ public void testWriterUsingHolderDurationField() {
+ try (final FixedSizeListVector vector =
+ FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) {
+ UnionFixedSizeListWriter writer = vector.getWriter();
+ writer.allocate();
+
+ DurationHolder durationHolder = new DurationHolder();
+ durationHolder.unit = TimeUnit.MILLISECOND;
+
+ writer.startList();
+ durationHolder.value = 812374L;
+ writer.duration().write(durationHolder);
+ durationHolder.value = 143451L;
+ writer.duration().write(durationHolder);
+
+ // Writing with a different unit should throw
+ durationHolder.unit = TimeUnit.SECOND;
+ durationHolder.value = 8888888;
+ IllegalArgumentException ex =
+ assertThrows(
+ IllegalArgumentException.class, () -> writer.duration().write(durationHolder));
+ assertEquals("holder.unit: SECOND not equal to vector unit: MILLISECOND", ex.getMessage());
+
+ writer.endList();
+ vector.setValueCount(1);
+
+ Field expectedDataField =
+ new Field(
+ BaseRepeatedValueVector.DATA_VECTOR_NAME,
+ FieldType.nullable(new ArrowType.Duration(TimeUnit.MILLISECOND)),
+ null);
+ Field expectedField =
+ new Field(
+ vector.getName(),
+ FieldType.nullable(new ArrowType.FixedSizeList(3)),
+ List.of(expectedDataField));
+
+ assertEquals(expectedField, writer.getField());
+ }
+ }
+
+ @Test
+ public void testWriterUsingHolderFixedSizeBinaryField() {
+ try (final FixedSizeListVector vector =
+ FixedSizeListVector.empty("vector", /* size= */ 2, allocator)) {
+ UnionFixedSizeListWriter writer = vector.getWriter();
+ writer.allocate();
+
+ FixedSizeBinaryHolder holder1 =
+ TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22});
+ FixedSizeBinaryHolder holder2 =
+ TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21});
+
+ writer.startList();
+ writer.fixedSizeBinary().write(holder1);
+ holder1.buffer.close();
+ writer.fixedSizeBinary().write(holder2);
+ holder2.buffer.close();
+
+ writer.endList();
+ vector.setValueCount(1);
+
+ FieldReader reader = vector.getReader();
+ assertTrue(reader.isSet(), "shouldn't be null");
+
+ Field expectedDataField =
+ new Field(
+ BaseRepeatedValueVector.DATA_VECTOR_NAME,
+ FieldType.nullable(new ArrowType.FixedSizeBinary(2)),
+ null);
+ Field expectedField =
+ new Field(
+ vector.getName(),
+ FieldType.nullable(new ArrowType.FixedSizeList(2)),
+ List.of(expectedDataField));
+
+ assertEquals(expectedField, writer.getField());
+ }
+ }
+
private int[] convertListToIntArray(List list) {
int[] values = new int[list.size()];
for (int i = 0; i < list.size(); i++) {
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
index 101d942d2a..bf9bba9c78 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
@@ -25,18 +26,24 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.UUID;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
import org.apache.arrow.vector.complex.LargeListVector;
import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.impl.UnionLargeListReader;
import org.apache.arrow.vector.complex.impl.UnionLargeListWriter;
import org.apache.arrow.vector.complex.reader.FieldReader;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -943,7 +950,7 @@ public void testGetBufferSizeFor() {
int[] indices = new int[] {0, 2, 4, 6, 10, 14};
for (int valueCount = 1; valueCount <= 5; valueCount++) {
- int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount);
+ int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
int offsetBufferSize = (valueCount + 1) * LargeListVector.OFFSET_WIDTH;
int expectedSize =
@@ -1020,6 +1027,99 @@ public void testGetTransferPairWithField() throws Exception {
}
}
+ @Test
+ public void testCopyValueSafeForExtensionType() throws Exception {
+ try (LargeListVector inVector = LargeListVector.empty("input", allocator);
+ LargeListVector outVector = LargeListVector.empty("output", allocator)) {
+ UnionLargeListWriter writer = inVector.getWriter();
+ writer.allocate();
+
+ // Create first list with UUIDs
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startList();
+ ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u1);
+ extensionWriter.writeExtension(u2);
+ writer.endList();
+
+ // Create second list with UUIDs
+ writer.setPosition(1);
+ UUID u3 = UUID.randomUUID();
+ UUID u4 = UUID.randomUUID();
+ writer.startList();
+ extensionWriter = writer.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u3);
+ extensionWriter.writeExtension(u4);
+ extensionWriter.writeNull();
+
+ writer.endList();
+ writer.setValueCount(2);
+
+ // Use copyFromSafe with ExtensionTypeWriterFactory
+ // This internally calls TransferImpl.copyValueSafe with ExtensionTypeWriterFactory
+ outVector.allocateNew();
+ TransferPair tp = inVector.makeTransferPair(outVector);
+ tp.copyValueSafe(0, 0);
+ tp.copyValueSafe(1, 1);
+ outVector.setValueCount(2);
+
+ // Verify first list
+ UnionLargeListReader reader = outVector.getReader();
+ reader.setPosition(0);
+ assertTrue(reader.isSet(), "first list shouldn't be null");
+ reader.next();
+ FieldReader uuidReader = reader.reader();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ uuidReader.read(holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u1, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u2, actualUuid);
+
+ // Verify second list
+ reader.setPosition(1);
+ assertTrue(reader.isSet(), "second list shouldn't be null");
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u3, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u4, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ assertFalse(uuidReader.isSet(), "third element should be null");
+ }
+ }
+
+ @Test
+ public void testEmptyLargeListOffsetBuffer() {
+ // Test that LargeListVector has correct readableBytes after allocation.
+ // According to Arrow spec, offset buffer must have N+1 entries.
+ // Even when N=0, it should contain [0].
+ try (LargeListVector list = LargeListVector.empty("list", allocator)) {
+ list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ list.allocateNew();
+ list.setValueCount(0);
+
+ List buffers = list.getFieldBuffers();
+ assertTrue(
+ buffers.get(1).readableBytes() >= LargeListVector.OFFSET_WIDTH,
+ "Offset buffer should have at least "
+ + LargeListVector.OFFSET_WIDTH
+ + " bytes for offset[0]");
+ assertEquals(0L, list.getOffsetBuffer().getLong(0));
+ }
+ }
+
private void writeIntValues(UnionLargeListWriter writer, int[] values) {
writer.startList();
for (int v : values) {
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java
index 26e7bb4a0d..256aa99687 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertSame;
@@ -1062,7 +1063,7 @@ public void testGetBufferSizeFor() {
int[] indices = new int[] {0, 2, 4, 6, 10, 14};
for (int valueCount = 1; valueCount <= 5; valueCount++) {
- int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount);
+ int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
int offsetBufferSize = valueCount * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH;
int sizeBufferSize = valueCount * BaseLargeRepeatedValueViewVector.SIZE_WIDTH;
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
index 1d6fa39f9e..0c90b32abc 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
@@ -26,15 +27,20 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.UUID;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
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.extension.UuidType;
import org.apache.arrow.vector.holders.DurationHolder;
import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.Types.MinorType;
@@ -42,6 +48,7 @@
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -1123,7 +1130,7 @@ public void testGetBufferSizeFor() {
int[] indices = new int[] {0, 2, 4, 6, 10, 14};
for (int valueCount = 1; valueCount <= 5; valueCount++) {
- int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount);
+ int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
int offsetBufferSize = (valueCount + 1) * BaseRepeatedValueVector.OFFSET_WIDTH;
int expectedSize =
@@ -1198,6 +1205,200 @@ public void testGetTransferPairWithField() {
}
}
+ @Test
+ public void testListVectorWithExtensionType() throws Exception {
+ final FieldType type = FieldType.nullable(UuidType.INSTANCE);
+ try (final ListVector inVector = new ListVector("list", allocator, type, null)) {
+ UnionListWriter writer = inVector.getWriter();
+ writer.allocate();
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startList();
+ ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u1);
+ extensionWriter.writeExtension(u2);
+ writer.endList();
+
+ writer.setValueCount(1);
+
+ FieldReader reader = inVector.getReader();
+ assertTrue(reader.isSet(), "shouldn't be null");
+ Object result = inVector.getObject(0);
+ ArrayList resultSet = (ArrayList) result;
+ assertEquals(2, resultSet.size());
+ assertEquals(u1, resultSet.get(0));
+ assertEquals(u2, resultSet.get(1));
+ }
+ }
+
+ @Test
+ public void testListVectorReaderForExtensionType() throws Exception {
+ final FieldType type = FieldType.nullable(UuidType.INSTANCE);
+ try (final ListVector inVector = new ListVector("list", allocator, type, null)) {
+ UnionListWriter writer = inVector.getWriter();
+ writer.allocate();
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startList();
+ ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u1);
+ extensionWriter.writeExtension(u2);
+ writer.endList();
+
+ writer.setValueCount(1);
+
+ UnionListReader reader = inVector.getReader();
+ assertTrue(reader.isSet(), "shouldn't be null");
+ reader.setPosition(0);
+ reader.next();
+ FieldReader uuidReader = reader.reader();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ uuidReader.read(holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u1, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u2, actualUuid);
+ }
+ }
+
+ @Test
+ public void testCopyFromForExtensionType() throws Exception {
+ try (ListVector inVector = ListVector.empty("input", allocator);
+ ListVector outVector = ListVector.empty("output", allocator)) {
+ UnionListWriter writer = inVector.getWriter();
+ writer.allocate();
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startList();
+
+ writer.extension(UuidType.INSTANCE).writeExtension(u1);
+ writer.writeExtension(u2);
+ writer.writeNull();
+ writer.endList();
+
+ writer.setValueCount(3);
+
+ // copy values from input to output
+ outVector.allocateNew();
+ outVector.copyFrom(0, 0, inVector);
+ outVector.setValueCount(3);
+
+ UnionListReader reader = outVector.getReader();
+ assertTrue(reader.isSet(), "shouldn't be null");
+ reader.setPosition(0);
+ reader.next();
+ FieldReader uuidReader = reader.reader();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ uuidReader.read(holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u1, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u2, actualUuid);
+ }
+ }
+
+ @Test
+ public void testCopyValueSafeForExtensionType() throws Exception {
+ try (ListVector inVector = ListVector.empty("input", allocator);
+ ListVector outVector = ListVector.empty("output", allocator)) {
+ UnionListWriter writer = inVector.getWriter();
+ writer.allocate();
+
+ // Create first list with UUIDs
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startList();
+ ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u1);
+ extensionWriter.writeExtension(u2);
+ writer.endList();
+
+ // Create second list with UUIDs
+ writer.setPosition(1);
+ UUID u3 = UUID.randomUUID();
+ UUID u4 = UUID.randomUUID();
+ writer.startList();
+ extensionWriter = writer.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u3);
+ extensionWriter.writeExtension(u4);
+ extensionWriter.writeNull();
+
+ writer.endList();
+ writer.setValueCount(2);
+
+ // Use TransferPair with ExtensionTypeWriterFactory
+ // This tests the new makeTransferPair API with writerFactory parameter
+ outVector.allocateNew();
+ TransferPair transferPair = inVector.makeTransferPair(outVector);
+ transferPair.copyValueSafe(0, 0);
+ transferPair.copyValueSafe(1, 1);
+ outVector.setValueCount(2);
+
+ // Verify first list
+ UnionListReader reader = outVector.getReader();
+ reader.setPosition(0);
+ assertTrue(reader.isSet(), "first list shouldn't be null");
+ reader.next();
+ FieldReader uuidReader = reader.reader();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ uuidReader.read(holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u1, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u2, actualUuid);
+
+ // Verify second list
+ reader.setPosition(1);
+ assertTrue(reader.isSet(), "second list shouldn't be null");
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u3, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u4, actualUuid);
+ reader.next();
+ uuidReader = reader.reader();
+ assertFalse(uuidReader.isSet(), "third element should be null");
+ }
+ }
+
+ @Test
+ public void testEmptyListOffsetBuffer() {
+ // Test that ListVector has correct readableBytes after allocation.
+ // According to Arrow spec, offset buffer must have N+1 entries.
+ // Even when N=0, it should contain [0].
+ try (ListVector list = ListVector.empty("list", allocator)) {
+ list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ list.allocateNew();
+ list.setValueCount(0);
+
+ List buffers = list.getFieldBuffers();
+ assertTrue(
+ buffers.get(1).readableBytes() >= BaseRepeatedValueVector.OFFSET_WIDTH,
+ "Offset buffer should have at least "
+ + BaseRepeatedValueVector.OFFSET_WIDTH
+ + " bytes for offset[0]");
+ assertEquals(0, list.getOffsetBuffer().getInt(0));
+ }
+ }
+
private void writeIntValues(UnionListWriter writer, int[] values) {
writer.startList();
for (int v : values) {
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java
index 639585fc48..8ab0edb145 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -1075,7 +1076,7 @@ public void testGetBufferSizeFor() {
int[] indices = new int[] {0, 2, 4, 6, 10, 14};
for (int valueCount = 1; valueCount <= 5; valueCount++) {
- int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount);
+ int validityBufferSize = getValidityBufferSizeFromCount(valueCount);
int offsetBufferSize = valueCount * BaseRepeatedValueViewVector.OFFSET_WIDTH;
int sizeBufferSize = valueCount * BaseRepeatedValueViewVector.SIZE_WIDTH;
@@ -1549,55 +1550,7 @@ public void testOverwriteWithNull() {
public void testOutOfOrderOffset1() {
// [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]]
try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) {
- // Allocate buffers in listViewVector by calling `allocateNew` method.
- listViewVector.allocateNew();
-
- // Initialize the child vector using `initializeChildrenFromFields` method.
-
- FieldType fieldType = new FieldType(true, new ArrowType.Int(16, true), null, null);
- Field field = new Field("child-vector", fieldType, null);
- listViewVector.initializeChildrenFromFields(Collections.singletonList(field));
-
- // Set values in the child vector.
- FieldVector fieldVector = listViewVector.getDataVector();
- fieldVector.clear();
-
- SmallIntVector childVector = (SmallIntVector) fieldVector;
-
- childVector.allocateNew(7);
-
- childVector.set(0, 0);
- childVector.set(1, -127);
- childVector.set(2, 127);
- childVector.set(3, 50);
- childVector.set(4, 12);
- childVector.set(5, -7);
- childVector.set(6, 25);
-
- childVector.setValueCount(7);
-
- // Set validity, offset and size buffers using `setValidity`,
- // `setOffset` and `setSize` methods.
- listViewVector.setValidity(0, 1);
- listViewVector.setValidity(1, 0);
- listViewVector.setValidity(2, 1);
- listViewVector.setValidity(3, 1);
- listViewVector.setValidity(4, 1);
-
- listViewVector.setOffset(0, 4);
- listViewVector.setOffset(1, 7);
- listViewVector.setOffset(2, 0);
- listViewVector.setOffset(3, 0);
- listViewVector.setOffset(4, 3);
-
- listViewVector.setSize(0, 3);
- listViewVector.setSize(1, 0);
- listViewVector.setSize(2, 4);
- listViewVector.setSize(3, 0);
- listViewVector.setSize(4, 2);
-
- // Set value count using `setValueCount` method.
- listViewVector.setValueCount(5);
+ initializeListViewVectorAsInSpecification(listViewVector);
final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer();
final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer();
@@ -2216,6 +2169,105 @@ public void testRangeChildVector2() {
}
}
+ @Test
+ public void testGetElementStartIndexAndEndIndexOrderedOffsetsNoIntersection() {
+ /*
+ values = [10, 20, 30, 40, 50]
+ offsets = [0, 3]
+ sizes = [3, 2]
+ vector: [[10, 20, 30], [40, 50]]
+ */
+ try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) {
+ initializeListViewVector(
+ listViewVector, List.of(10, 20, 30, 40, 50), List.of(1, 1), List.of(0, 3), List.of(3, 2));
+
+ assertEquals(0, listViewVector.getElementStartIndex(0));
+ assertEquals(3, listViewVector.getElementEndIndex(0));
+ assertEquals(3, listViewVector.getElementStartIndex(1));
+ assertEquals(5, listViewVector.getElementEndIndex(1));
+
+ final FieldVector dataVec = listViewVector.getDataVector();
+ int elemIndex = 0;
+ int start = listViewVector.getElementStartIndex(elemIndex);
+ int end = listViewVector.getElementEndIndex(elemIndex);
+ List list = listViewVector.getObject(elemIndex);
+ assertEquals(end - start, list.size());
+ for (int j = 0; j < list.size(); j++) {
+ assertEquals(((SmallIntVector) dataVec).get(start + j), list.get(j));
+ }
+ }
+ }
+
+ @Test
+ public void testGetElementStartIndexAndEndIndexNotOrderedOffsetsNoIntersection() {
+ /*
+ values = [1, 2, 3, 4, 5, 6]
+ validity = [1, 1, 1]
+ offsets = [4, 2, 0]
+ sizes = [2, 2, 2]
+ vector: [[5, 6], [3, 4], [1, 2]]
+ */
+ try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) {
+ initializeListViewVector(
+ listViewVector,
+ List.of(1, 2, 3, 4, 5, 6),
+ List.of(1, 1, 1),
+ List.of(4, 2, 0),
+ List.of(2, 2, 2));
+
+ assertEquals(4, listViewVector.getElementStartIndex(0));
+ assertEquals(6, listViewVector.getElementEndIndex(0));
+ assertEquals(2, listViewVector.getElementStartIndex(1));
+ assertEquals(4, listViewVector.getElementEndIndex(1));
+ assertEquals(0, listViewVector.getElementStartIndex(2));
+ assertEquals(2, listViewVector.getElementEndIndex(2));
+ }
+ }
+
+ @Test
+ public void testGetElementStartIndexAndEndIndexOrderedOffsetsWithIntersection() {
+ /*
+ values = [1, 2, 3, 4, 5]
+ validity = [1, 1, 1]
+ offsets = [0, 1, 4]
+ sizes = [2, 3, 1]
+ vector: [[1, 2], [2, 3, 4], [5]]
+ */
+ try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) {
+ initializeListViewVector(
+ listViewVector,
+ List.of(1, 2, 3, 4, 5),
+ List.of(1, 1, 1),
+ List.of(0, 1, 4),
+ List.of(2, 3, 1));
+
+ assertEquals(0, listViewVector.getElementStartIndex(0));
+ assertEquals(2, listViewVector.getElementEndIndex(0));
+ assertEquals(1, listViewVector.getElementStartIndex(1));
+ assertEquals(4, listViewVector.getElementEndIndex(1));
+ assertEquals(4, listViewVector.getElementStartIndex(2));
+ assertEquals(5, listViewVector.getElementEndIndex(2));
+ }
+ }
+
+ @Test
+ public void testGetElementStartIndexAndEndIndexOrderedOffsetsAsInSpecification() {
+ try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) {
+ initializeListViewVectorAsInSpecification(listViewVector);
+
+ assertEquals(4, listViewVector.getElementStartIndex(0));
+ assertEquals(7, listViewVector.getElementEndIndex(0));
+ assertEquals(7, listViewVector.getElementStartIndex(1));
+ assertEquals(7, listViewVector.getElementEndIndex(1));
+ assertEquals(0, listViewVector.getElementStartIndex(2));
+ assertEquals(4, listViewVector.getElementEndIndex(2));
+ assertEquals(0, listViewVector.getElementStartIndex(3));
+ assertEquals(0, listViewVector.getElementEndIndex(3));
+ assertEquals(3, listViewVector.getElementStartIndex(4));
+ assertEquals(5, listViewVector.getElementEndIndex(4));
+ }
+ }
+
private void writeIntValues(UnionListViewWriter writer, int[] values) {
writer.startListView();
for (int v : values) {
@@ -2223,4 +2275,70 @@ private void writeIntValues(UnionListViewWriter writer, int[] values) {
}
writer.endListView();
}
+
+ /**
+ * ListViewVector from the specification .
+ */
+ private void initializeListViewVectorAsInSpecification(ListViewVector listViewVector) {
+ /*
+ values = [0, -127, 127, 50, 12, -7, 25]
+ validity = [1, 1, 1, 0, 1] (reversed)
+ offsets = [4, 7, 0, 0, 3]
+ sizes = [3, 0, 4, 0, 2]
+ vector: [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]]
+ */
+ initializeListViewVector(
+ listViewVector,
+ List.of(0, -127, 127, 50, 12, -7, 25),
+ List.of(1, 1, 1, 0, 1),
+ List.of(4, 7, 0, 0, 3),
+ List.of(3, 0, 4, 0, 2));
+ }
+
+ private void initializeListViewVector(
+ ListViewVector listViewVector,
+ List values,
+ List validity,
+ List offsets,
+ List sizes) {
+ // Allocate buffers in listViewVector by calling `allocateNew` method.
+ assert offsets.size() == sizes.size();
+ listViewVector.allocateNew();
+
+ // Initialize the child vector using `initializeChildrenFromFields` method.
+ FieldType fieldType = new FieldType(true, new ArrowType.Int(16, true), null, null);
+ Field field = new Field("child-vector", fieldType, null);
+ listViewVector.initializeChildrenFromFields(Collections.singletonList(field));
+
+ // Set values in the child vector.
+ FieldVector fieldVector = listViewVector.getDataVector();
+ fieldVector.clear();
+
+ SmallIntVector childVector = (SmallIntVector) fieldVector;
+ childVector.allocateNew(values.size());
+ for (int i = 0; i < values.size(); i++) {
+ childVector.set(i, values.get(i));
+ }
+ childVector.setValueCount(values.size());
+
+ // Set validity, offset and size buffers using `setValidity`,
+ // `setOffset` and `setSize` methods.
+ List reversedValidity = new ArrayList<>(validity);
+ Collections.reverse(reversedValidity);
+ for (int i = 0; i < reversedValidity.size(); i++) {
+ listViewVector.setValidity(i, reversedValidity.get(i));
+ }
+
+ for (int i = 0; i < offsets.size(); i++) {
+ listViewVector.setOffset(i, offsets.get(i));
+ }
+
+ for (int i = 0; i < sizes.size(); i++) {
+ listViewVector.setSize(i, sizes.get(i));
+ }
+
+ // Set value count using `setValueCount` method.
+ listViewVector.setValueCount(offsets.size());
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
index 313d83ec91..2f520f3882 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
@@ -16,16 +16,19 @@
*/
package org.apache.arrow.vector;
+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.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import java.util.UUID;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.complex.MapVector;
@@ -33,15 +36,20 @@
import org.apache.arrow.vector.complex.impl.UnionMapReader;
import org.apache.arrow.vector.complex.impl.UnionMapWriter;
import org.apache.arrow.vector.complex.reader.FieldReader;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
import org.apache.arrow.vector.complex.writer.FieldWriter;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.JsonStringArrayList;
import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -1263,4 +1271,395 @@ public void testMapTypeReturnsSupportedMapWriter() {
assertEquals(11, getResultValue(resultStruct));
}
}
+
+ @Test
+ public void testMapVectorWithExtensionType() throws Exception {
+ try (final MapVector inVector = MapVector.empty("map", allocator, false)) {
+ inVector.allocateNew();
+ UnionMapWriter writer = inVector.getWriter();
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startMap();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(0);
+ ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u1, UuidType.INSTANCE);
+ writer.endEntry();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(1);
+ extensionWriter = writer.value().extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u2, UuidType.INSTANCE);
+ writer.endEntry();
+ writer.endMap();
+
+ writer.setValueCount(1);
+
+ UnionMapReader mapReader = inVector.getReader();
+ mapReader.setPosition(0);
+ mapReader.next();
+ FieldReader uuidReader = mapReader.value();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ uuidReader.read(holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u1, actualUuid);
+ mapReader.next();
+ uuidReader = mapReader.value();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u2, actualUuid);
+ }
+ }
+
+ @Test
+ public void testCopyFromForExtensionType() throws Exception {
+ try (final MapVector inVector = MapVector.empty("in", allocator, false);
+ final MapVector outVector = MapVector.empty("out", allocator, false)) {
+ inVector.allocateNew();
+ UnionMapWriter writer = inVector.getWriter();
+ writer.setPosition(0);
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ writer.startMap();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(0);
+ ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(u1, UuidType.INSTANCE);
+ writer.endEntry();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(1);
+ extensionWriter.writeExtension(u2, UuidType.INSTANCE);
+ writer.endEntry();
+ writer.endMap();
+
+ writer.setValueCount(1);
+ outVector.allocateNew();
+ outVector.copyFrom(0, 0, inVector);
+ outVector.setValueCount(1);
+
+ UnionMapReader mapReader = outVector.getReader();
+ mapReader.setPosition(0);
+ mapReader.next();
+ FieldReader uuidReader = mapReader.value();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ uuidReader.read(holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u1, actualUuid);
+ mapReader.next();
+ uuidReader = mapReader.value();
+ uuidReader.read(holder);
+ actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(u2, actualUuid);
+ }
+ }
+
+ /**
+ * Regression test for GH-586: UnionMapWriter.fixedSizeBinary() should properly delegate to the
+ * entry writer for both key and value paths.
+ */
+ @Test
+ public void testFixedSizeBinaryWriter() {
+ try (MapVector mapVector = MapVector.empty("map_vector", allocator, false)) {
+ UnionMapWriter writer = mapVector.getWriter();
+ writer.allocate();
+
+ // populate input vector with the following records
+ // {[11, 22] -> [32, 21]}
+ // {1 -> [11, 22], 2 -> [32, 21]}
+ // null
+ // {[11, 22] -> 1, [32, 21] -> 2}
+ // {[11, 22] -> null}
+ // {null -> [32, 21]} - wrong "for a given entry, the "key" is non-nullable" - todo: it
+ // shouldn't work. Should it?
+ FixedSizeBinaryHolder holder1 =
+ TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22});
+ FixedSizeBinaryHolder holder2 =
+ TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21});
+
+ writer.setPosition(0); // optional
+ writer.startMap();
+ writer.startEntry();
+ writer
+ .key()
+ .fixedSizeBinary(holder1.byteWidth)
+ .write(holder1); // need to initialize with byteWidth - NPE otherwise
+ writer.value().fixedSizeBinary(holder2.byteWidth).write(holder2);
+ writer.endEntry();
+ holder1.buffer.close();
+ holder2.buffer.close();
+ writer.endMap();
+
+ // {1 -> [11, 22], 2 -> [32, 21]}
+ holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22});
+ holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21});
+ writer.setPosition(1);
+ writer.startMap();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(1);
+ writer.value().fixedSizeBinary().write(holder1);
+ writer.endEntry();
+ holder1.buffer.close();
+ writer.startEntry();
+ writer.key().bigInt().writeBigInt(2);
+ writer.value().fixedSizeBinary().write(holder2);
+ writer.endEntry();
+ writer.endMap();
+ holder2.buffer.close();
+
+ // {[11, 22] -> 1, [32, 21] -> 2}
+ holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22});
+ holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21});
+ writer.setPosition(3);
+ writer.startMap();
+ writer.startEntry();
+ writer.key().fixedSizeBinary().write(holder1);
+ writer.value().bigInt().writeBigInt(1);
+ writer.endEntry();
+ holder1.buffer.close();
+ writer.startEntry();
+ writer.key().fixedSizeBinary().write(holder2);
+ writer.value().bigInt().writeBigInt(2);
+ writer.endEntry();
+ writer.endMap();
+ holder2.buffer.close();
+
+ // {[11, 22] -> null}
+ holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22});
+ writer.setPosition(4);
+ writer.startMap();
+ writer.startEntry();
+ writer.key().fixedSizeBinary().write(holder1);
+ writer.endEntry();
+ writer.endMap();
+ holder1.buffer.close();
+
+ // {null -> [32, 21]}
+ holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21});
+ writer.setPosition(5);
+ writer.startMap();
+ writer.startEntry();
+ writer.value().fixedSizeBinary().write(holder2);
+ writer.endEntry();
+ writer.endMap();
+ holder2.buffer.close();
+
+ writer.setValueCount(6);
+
+ // assert the output vector is correct
+ FieldReader reader = mapVector.getReader();
+ assertTrue(reader.isSet(), "shouldn't be null");
+ reader.setPosition(1);
+ assertTrue(reader.isSet(), "shouldn't be null");
+ reader.setPosition(2);
+ assertFalse(reader.isSet(), "should be null");
+ reader.setPosition(3);
+ assertTrue(reader.isSet(), "shouldn't be null");
+ reader.setPosition(4);
+ assertTrue(reader.isSet(), "shouldn't be null");
+ reader.setPosition(5);
+ assertTrue(reader.isSet(), "shouldn't be null");
+
+ /* index 0 */
+ Object result = mapVector.getObject(0);
+ ArrayList resultSet = (ArrayList) result;
+ assertEquals(1, resultSet.size());
+ Map resultStruct = (Map) resultSet.get(0);
+ assertTrue(resultStruct.containsKey(MapVector.KEY_NAME));
+ assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME));
+ assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME));
+ assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME));
+
+ /* index 1 */
+ result = mapVector.getObject(1);
+ resultSet = (ArrayList) result;
+ assertEquals(2, resultSet.size());
+ resultStruct = (Map) resultSet.get(0);
+ assertEquals(1L, getResultKey(resultStruct));
+ assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME));
+ assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.VALUE_NAME));
+ resultStruct = (Map) resultSet.get(1);
+ assertEquals(2L, getResultKey(resultStruct));
+ assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME));
+ assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME));
+
+ /* index 2 */
+ result = mapVector.getObject(2);
+ assertNull(result);
+
+ /* index 3 */
+ result = mapVector.getObject(3);
+ resultSet = (ArrayList) result;
+ assertEquals(2, resultSet.size());
+ resultStruct = (Map) resultSet.get(0);
+ assertTrue(resultStruct.containsKey(MapVector.KEY_NAME));
+ assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME));
+ assertEquals(1L, getResultValue(resultStruct));
+ resultStruct = (Map) resultSet.get(1);
+ assertTrue(resultStruct.containsKey(MapVector.KEY_NAME));
+ assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.KEY_NAME));
+ assertEquals(2L, getResultValue(resultStruct));
+
+ /* index 4 */
+ result = mapVector.getObject(4);
+ resultSet = (ArrayList) result;
+ assertEquals(1, resultSet.size());
+ resultStruct = (Map) resultSet.get(0);
+ assertTrue(resultStruct.containsKey(MapVector.KEY_NAME));
+ assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME));
+ assertFalse(resultStruct.containsKey(MapVector.VALUE_NAME));
+
+ /* index 5 */
+ result = mapVector.getObject(5);
+ resultSet = (ArrayList) result;
+ assertEquals(1, resultSet.size());
+ resultStruct = (Map) resultSet.get(0);
+ assertFalse(resultStruct.containsKey(MapVector.KEY_NAME));
+ assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME));
+ assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME));
+ }
+ }
+
+ @Test
+ public void testFixedSizeBinaryFirstInitialization() {
+ try (MapVector mapVector = MapVector.empty("map_vector", allocator, false)) {
+ UnionMapWriter writer = mapVector.getWriter();
+ writer.allocate();
+
+ // populate input vector with the following records
+ // {[11, 22] -> [32, 21]}
+ FixedSizeBinaryHolder holder1 =
+ TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22});
+ FixedSizeBinaryHolder holder2 =
+ TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21});
+
+ writer.setPosition(0); // optional
+ writer.startMap();
+ writer.startEntry();
+ // require byteWidth parameter for first-time initialization of `key` or `value` writers
+ assertThrows(NullPointerException.class, () -> writer.key().fixedSizeBinary().write(holder1));
+ assertThrows(
+ NullPointerException.class, () -> writer.value().fixedSizeBinary().write(holder2));
+ writer.key().fixedSizeBinary(holder1.byteWidth).write(holder1);
+ writer.value().fixedSizeBinary(holder2.byteWidth).write(holder2);
+ writer.endEntry();
+ holder1.buffer.close();
+ holder2.buffer.close();
+ writer.endMap();
+
+ writer.setValueCount(1);
+
+ // assert the output vector is correct
+ FieldReader reader = mapVector.getReader();
+ assertTrue(reader.isSet(), "shouldn't be null");
+
+ /* index 0 */
+ Object result = mapVector.getObject(0);
+ ArrayList resultSet = (ArrayList) result;
+ assertEquals(1, resultSet.size());
+ Map resultStruct = (Map) resultSet.get(0);
+ assertTrue(resultStruct.containsKey(MapVector.KEY_NAME));
+ assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME));
+ assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME));
+ assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME));
+ }
+ }
+
+ @Test
+ public void testMapWithUuidKeyAndListUuidValue() throws Exception {
+ try (final MapVector mapVector = MapVector.empty("map", allocator, false)) {
+ mapVector.allocateNew();
+ UnionMapWriter writer = mapVector.getWriter();
+
+ // Create test UUIDs
+ UUID key1 = UUID.randomUUID();
+ UUID key2 = UUID.randomUUID();
+ UUID value1a = UUID.randomUUID();
+ UUID value1b = UUID.randomUUID();
+ UUID value2a = UUID.randomUUID();
+ UUID value2b = UUID.randomUUID();
+ UUID value2c = UUID.randomUUID();
+
+ // Write first map entry: {key1 -> [value1a, value1b]}
+ writer.setPosition(0);
+ writer.startMap();
+
+ writer.startEntry();
+ ExtensionWriter keyWriter = writer.key().extension(UuidType.INSTANCE);
+ keyWriter.writeExtension(key1, UuidType.INSTANCE);
+ ListWriter valueWriter = writer.value().list();
+ valueWriter.startList();
+ ExtensionWriter listItemWriter = valueWriter.extension(UuidType.INSTANCE);
+ listItemWriter.writeExtension(value1a, UuidType.INSTANCE);
+ listItemWriter = valueWriter.extension(UuidType.INSTANCE);
+ listItemWriter.writeExtension(value1b, UuidType.INSTANCE);
+ valueWriter.endList();
+ writer.endEntry();
+
+ writer.startEntry();
+ keyWriter = writer.key().extension(UuidType.INSTANCE);
+ keyWriter.writeExtension(key2, UuidType.INSTANCE);
+ valueWriter = writer.value().list();
+ valueWriter.startList();
+ listItemWriter = valueWriter.extension(UuidType.INSTANCE);
+ listItemWriter.writeExtension(value2a, UuidType.INSTANCE);
+ listItemWriter = valueWriter.extension(UuidType.INSTANCE);
+ listItemWriter.writeExtension(value2b, UuidType.INSTANCE);
+ listItemWriter = valueWriter.extension(UuidType.INSTANCE);
+ listItemWriter.writeExtension(value2c, UuidType.INSTANCE);
+ valueWriter.endList();
+ writer.endEntry();
+
+ writer.endMap();
+ writer.setValueCount(1);
+
+ // Read and verify the data
+ UnionMapReader mapReader = mapVector.getReader();
+ mapReader.setPosition(0);
+
+ // Read first entry
+ mapReader.next();
+ FieldReader keyReader = mapReader.key();
+ NullableUuidHolder keyHolder = new NullableUuidHolder();
+ keyReader.read(keyHolder);
+ UUID actualKey = UuidUtility.uuidFromArrowBuf(keyHolder.buffer, keyHolder.start);
+ assertEquals(key1, actualKey);
+
+ FieldReader valueReader = mapReader.value();
+ assertTrue(valueReader.isSet());
+ List listValue = (List) valueReader.readObject();
+ assertEquals(2, listValue.size());
+
+ // Verify first list item - readObject() returns UUID objects for extension types
+ UUID actualValue1a = (UUID) listValue.get(0);
+ assertEquals(value1a, actualValue1a);
+
+ // Verify second list item
+ UUID actualValue1b = (UUID) listValue.get(1);
+ assertEquals(value1b, actualValue1b);
+
+ // Read second entry
+ mapReader.next();
+ keyReader = mapReader.key();
+ keyReader.read(keyHolder);
+ actualKey = UuidUtility.uuidFromArrowBuf(keyHolder.buffer, keyHolder.start);
+ assertEquals(key2, actualKey);
+
+ valueReader = mapReader.value();
+ assertTrue(valueReader.isSet());
+ listValue = (List) valueReader.readObject();
+ assertEquals(3, listValue.size());
+
+ // Verify first list item - readObject() returns UUID objects for extension types
+ UUID actualValue2a = (UUID) listValue.get(0);
+ assertEquals(value2a, actualValue2a);
+
+ // Verify second list item
+ UUID actualValue2b = (UUID) listValue.get(1);
+ assertEquals(value2b, actualValue2b);
+
+ // Verify third list item
+ UUID actualValue2c = (UUID) listValue.get(2);
+ assertEquals(value2c, actualValue2c);
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java b/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java
index adf51c0730..9fa153e928 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java
@@ -148,12 +148,18 @@ public void testRangeCompare() {
assertTrue(
constantVector.accept(
new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 2, 13)));
- assertFalse(
- constantVector.accept(
- new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 10, 10)));
- assertFalse(
- constantVector.accept(
- new RangeEqualsVisitor(constantVector, constantVector), new Range(10, 1, 10)));
+
+ // throws exception if the range end is out the bound of the vector
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ constantVector.accept(
+ new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 10, 10)));
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ constantVector.accept(
+ new RangeEqualsVisitor(constantVector, constantVector), new Range(10, 1, 10)));
// Create REE vector representing: [1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5].
RunEndEncodedVector reeVector =
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java
index d40af9ae89..8c8a45f588 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java
@@ -35,6 +35,7 @@
import org.apache.arrow.vector.complex.impl.NullableStructWriter;
import org.apache.arrow.vector.complex.writer.Float8Writer;
import org.apache.arrow.vector.complex.writer.IntWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.holders.ComplexHolder;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.Types.MinorType;
@@ -42,7 +43,6 @@
import org.apache.arrow.vector.types.pojo.ArrowType.Struct;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.UuidType;
import org.apache.arrow.vector.util.TransferPair;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
@@ -160,17 +160,23 @@ public void testGetPrimitiveVectors() {
UnionVector unionVector = vector.addOrGetUnion("union");
unionVector.addVector(new BigIntVector("bigInt", allocator));
unionVector.addVector(new SmallIntVector("smallInt", allocator));
+ unionVector.addVector(new UuidVector("uuid", allocator));
// add varchar vector
vector.addOrGet(
"varchar", FieldType.nullable(MinorType.VARCHAR.getType()), VarCharVector.class);
+ // add extension vector
+ vector.addOrGet("extension", FieldType.nullable(UuidType.INSTANCE), UuidVector.class);
+
List primitiveVectors = vector.getPrimitiveVectors();
- assertEquals(4, primitiveVectors.size());
+ assertEquals(6, primitiveVectors.size());
assertEquals(MinorType.INT, primitiveVectors.get(0).getMinorType());
assertEquals(MinorType.BIGINT, primitiveVectors.get(1).getMinorType());
assertEquals(MinorType.SMALLINT, primitiveVectors.get(2).getMinorType());
- assertEquals(MinorType.VARCHAR, primitiveVectors.get(3).getMinorType());
+ assertEquals(MinorType.EXTENSIONTYPE, primitiveVectors.get(3).getMinorType());
+ assertEquals(MinorType.VARCHAR, primitiveVectors.get(4).getMinorType());
+ assertEquals(MinorType.EXTENSIONTYPE, primitiveVectors.get(5).getMinorType());
}
}
@@ -341,7 +347,7 @@ public void testGetTransferPairWithFieldAndCallBack() {
@Test
public void testStructVectorWithExtensionTypes() {
- UuidType uuidType = new UuidType();
+ UuidType uuidType = UuidType.INSTANCE;
Field uuidField = new Field("struct_child", FieldType.nullable(uuidType), null);
Field structField =
new Field("struct", FieldType.nullable(new ArrowType.Struct()), List.of(uuidField));
@@ -353,7 +359,7 @@ public void testStructVectorWithExtensionTypes() {
@Test
public void testStructVectorTransferPairWithExtensionType() {
- UuidType uuidType = new UuidType();
+ UuidType uuidType = UuidType.INSTANCE;
Field uuidField = new Field("uuid_child", FieldType.nullable(uuidType), null);
Field structField =
new Field("struct", FieldType.nullable(new ArrowType.Struct()), List.of(uuidField));
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java
index 82295f8037..d91b2004c0 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java
@@ -18,8 +18,10 @@
import java.util.Random;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
import org.apache.arrow.vector.types.pojo.FieldType;
public class TestUtils {
@@ -62,4 +64,26 @@ public static String generateRandomString(int length) {
}
return sb.toString();
}
+
+ /*
+ * Ensure the extension type is registered, as there might other tests trying to unregister the
+ * type. ex.: TestExtensionType#readUnderlyingType
+ */
+ public static void ensureRegistered(ArrowType.ExtensionType type) {
+ if (ExtensionTypeRegistry.lookup(type.extensionName()) == null) {
+ ExtensionTypeRegistry.register(type);
+ }
+ }
+
+ public static FixedSizeBinaryHolder fixedSizeBinaryHolder(
+ BufferAllocator allocator, byte[] array) {
+ FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder();
+ holder.byteWidth = array.length;
+ holder.buffer = allocator.buffer(array.length);
+ for (int i = 0; i < array.length; i++) {
+ holder.buffer.setByte(i, array[i]);
+ }
+
+ return holder;
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java b/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java
new file mode 100644
index 0000000000..99045d1cba
--- /dev/null
+++ b/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java
@@ -0,0 +1,276 @@
+/*
+ * 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 static org.apache.arrow.vector.TestUtils.ensureRegistered;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.UUID;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.ipc.ArrowStreamReader;
+import org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+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.UuidUtility;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class TestUuidType {
+ BufferAllocator allocator;
+
+ @BeforeEach
+ void beforeEach() {
+ allocator = new RootAllocator();
+ }
+
+ @AfterEach
+ void afterEach() {
+ allocator.close();
+ }
+
+ @Test
+ void testConstants() {
+ assertEquals("arrow.uuid", UuidType.EXTENSION_NAME);
+ assertNotNull(UuidType.INSTANCE);
+ assertNotNull(UuidType.STORAGE_TYPE);
+ assertInstanceOf(ArrowType.FixedSizeBinary.class, UuidType.STORAGE_TYPE);
+ assertEquals(
+ UuidType.UUID_BYTE_WIDTH,
+ ((ArrowType.FixedSizeBinary) UuidType.STORAGE_TYPE).getByteWidth());
+ }
+
+ @Test
+ void testStorageType() {
+ UuidType type = UuidType.INSTANCE;
+ assertEquals(UuidType.STORAGE_TYPE, type.storageType());
+ assertInstanceOf(ArrowType.FixedSizeBinary.class, type.storageType());
+ }
+
+ @Test
+ void testExtensionName() {
+ UuidType type = UuidType.INSTANCE;
+ assertEquals("arrow.uuid", type.extensionName());
+ }
+
+ @Test
+ void testExtensionEquals() {
+ UuidType type1 = UuidType.INSTANCE;
+ UuidType type2 = UuidType.INSTANCE;
+ UuidType type3 = UuidType.INSTANCE;
+
+ assertTrue(type1.extensionEquals(type2));
+ assertTrue(type1.extensionEquals(type3));
+ assertTrue(type2.extensionEquals(type3));
+ }
+
+ @Test
+ void testIsComplex() {
+ UuidType type = UuidType.INSTANCE;
+ assertFalse(type.isComplex());
+ }
+
+ @Test
+ void testSerialize() {
+ UuidType type = UuidType.INSTANCE;
+ String serialized = type.serialize();
+ assertEquals("", serialized);
+ }
+
+ @Test
+ void testDeserializeValid() {
+ UuidType type = UuidType.INSTANCE;
+ ArrowType storageType = new ArrowType.FixedSizeBinary(UuidType.UUID_BYTE_WIDTH);
+
+ ArrowType deserialized = assertDoesNotThrow(() -> type.deserialize(storageType, ""));
+ assertInstanceOf(UuidType.class, deserialized);
+ assertEquals(UuidType.INSTANCE, deserialized);
+ }
+
+ @Test
+ void testDeserializeInvalidStorageType() {
+ UuidType type = UuidType.INSTANCE;
+ ArrowType wrongStorageType = new ArrowType.FixedSizeBinary(32);
+
+ assertThrows(UnsupportedOperationException.class, () -> type.deserialize(wrongStorageType, ""));
+ }
+
+ @Test
+ void testGetNewVector() {
+ UuidType type = UuidType.INSTANCE;
+ try (FieldVector vector =
+ type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) {
+ assertInstanceOf(UuidVector.class, vector);
+ assertEquals("uuid_field", vector.getField().getName());
+ assertEquals(type, vector.getField().getType());
+ }
+ }
+
+ @Test
+ void testVectorOperations() {
+ UuidType type = UuidType.INSTANCE;
+ try (FieldVector vector =
+ type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) {
+ UuidVector uuidVector = (UuidVector) vector;
+
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ uuidVector.setSafe(0, uuid1);
+ uuidVector.setSafe(1, uuid2);
+ uuidVector.setNull(2);
+ uuidVector.setValueCount(3);
+
+ assertEquals(uuid1, uuidVector.getObject(0));
+ assertEquals(uuid2, uuidVector.getObject(1));
+ assertNull(uuidVector.getObject(2));
+ assertFalse(uuidVector.isNull(0));
+ assertFalse(uuidVector.isNull(1));
+ assertTrue(uuidVector.isNull(2));
+ }
+ }
+
+ @Test
+ void testIpcRoundTrip() {
+ UuidType type = UuidType.INSTANCE;
+ ensureRegistered(type);
+
+ Schema schema = new Schema(Collections.singletonList(Field.nullable("uuid", type)));
+ byte[] serialized = schema.serializeAsMessage();
+ Schema deserialized = Schema.deserializeMessage(ByteBuffer.wrap(serialized));
+ assertEquals(schema, deserialized);
+ }
+
+ @Test
+ void testVectorIpcRoundTrip() throws IOException {
+ UuidType type = UuidType.INSTANCE;
+ ensureRegistered(type);
+
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ try (FieldVector vector = type.getNewVector("field", FieldType.nullable(type), allocator)) {
+ UuidVector uuidVector = (UuidVector) vector;
+ uuidVector.setSafe(0, uuid1);
+ uuidVector.setNull(1);
+ uuidVector.setSafe(2, uuid2);
+ uuidVector.setValueCount(3);
+
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (VectorSchemaRoot root = new VectorSchemaRoot(Collections.singletonList(uuidVector));
+ ArrowStreamWriter writer =
+ new ArrowStreamWriter(root, new DictionaryProvider.MapDictionaryProvider(), baos)) {
+ writer.start();
+ writer.writeBatch();
+ }
+
+ try (ArrowStreamReader reader =
+ new ArrowStreamReader(new ByteArrayInputStream(baos.toByteArray()), allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ assertEquals(3, root.getRowCount());
+ assertEquals(
+ new Schema(Collections.singletonList(uuidVector.getField())), root.getSchema());
+
+ UuidVector actual = assertInstanceOf(UuidVector.class, root.getVector("field"));
+ assertFalse(actual.isNull(0));
+ assertTrue(actual.isNull(1));
+ assertFalse(actual.isNull(2));
+ assertEquals(uuid1, actual.getObject(0));
+ assertNull(actual.getObject(1));
+ assertEquals(uuid2, actual.getObject(2));
+ }
+ }
+ }
+
+ @Test
+ void testVectorByteArrayOperations() {
+ UuidType type = UuidType.INSTANCE;
+ try (FieldVector vector =
+ type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) {
+ UuidVector uuidVector = (UuidVector) vector;
+
+ UUID uuid = UUID.randomUUID();
+ byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid);
+
+ uuidVector.setSafe(0, uuidBytes);
+ uuidVector.setValueCount(1);
+
+ assertEquals(uuid, uuidVector.getObject(0));
+
+ // Verify the bytes match
+ byte[] actualBytes = new byte[UuidType.UUID_BYTE_WIDTH];
+ int offset = uuidVector.getStartOffset(0);
+ uuidVector.getDataBuffer().getBytes(offset, actualBytes);
+ assertArrayEquals(uuidBytes, actualBytes);
+ }
+ }
+
+ @Test
+ void testGetNewVectorWithCustomFieldType() {
+ UuidType type = UuidType.INSTANCE;
+ FieldType fieldType = new FieldType(false, type, null);
+
+ try (FieldVector vector = type.getNewVector("non_nullable_uuid", fieldType, allocator)) {
+ assertInstanceOf(UuidVector.class, vector);
+ assertEquals("non_nullable_uuid", vector.getField().getName());
+ assertFalse(vector.getField().isNullable());
+ }
+ }
+
+ @Test
+ void testSingleton() {
+ UuidType type1 = UuidType.INSTANCE;
+ UuidType type2 = UuidType.INSTANCE;
+
+ // Same instance
+ assertSame(type1, type2);
+ assertTrue(type1.extensionEquals(type2));
+ }
+
+ @Test
+ void testUnderlyingVector() {
+ UuidType type = UuidType.INSTANCE;
+ try (FieldVector vector =
+ type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) {
+ UuidVector uuidVector = (UuidVector) vector;
+ FixedSizeBinaryVector underlying = uuidVector.getUnderlyingVector();
+
+ assertInstanceOf(FixedSizeBinaryVector.class, underlying);
+ assertEquals(UuidType.UUID_BYTE_WIDTH, underlying.getByteWidth());
+ }
+ }
+}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java b/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java
new file mode 100644
index 0000000000..b5dd12d89c
--- /dev/null
+++ b/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java
@@ -0,0 +1,726 @@
+/*
+ * 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 static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+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.arrow.vector.complex.impl.NullableUuidHolderReaderImpl;
+import org.apache.arrow.vector.complex.impl.UuidReaderImpl;
+import org.apache.arrow.vector.complex.impl.UuidWriterImpl;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.util.UuidUtility;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for UuidVector, UuidWriterImpl, and UuidReaderImpl. */
+class TestUuidVector {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ void beforeEach() {
+ allocator = new RootAllocator();
+ }
+
+ @AfterEach
+ void afterEach() {
+ allocator.close();
+ }
+
+ // ========== Writer Tests ==========
+
+ @Test
+ void testWriteToExtensionVector() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ UUID uuid = UUID.randomUUID();
+ ByteBuffer bb = ByteBuffer.allocate(UuidType.UUID_BYTE_WIDTH);
+ bb.putLong(uuid.getMostSignificantBits());
+ bb.putLong(uuid.getLeastSignificantBits());
+
+ // Allocate ArrowBuf for the holder
+ try (ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) {
+ buf.setBytes(0, bb.array());
+
+ UuidHolder holder = new UuidHolder();
+ holder.buffer = buf;
+
+ writer.write(holder);
+ UUID result = vector.getObject(0);
+ assertEquals(uuid, result);
+ }
+ }
+ }
+
+ @Test
+ void testWriteExtensionWithUUID() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ UUID uuid = UUID.randomUUID();
+ writer.setPosition(0);
+ writer.writeExtension(uuid);
+
+ UUID result = vector.getObject(0);
+ assertEquals(uuid, result);
+ assertEquals(1, vector.getValueCount());
+ }
+ }
+
+ @Test
+ void testWriteExtensionWithByteArray() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ UUID uuid = UUID.randomUUID();
+ byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid);
+
+ writer.setPosition(0);
+ writer.writeExtension(uuidBytes);
+
+ UUID result = vector.getObject(0);
+ assertEquals(uuid, result);
+ assertEquals(1, vector.getValueCount());
+ }
+ }
+
+ @Test
+ void testWriteExtensionWithArrowBuf() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector);
+ ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) {
+ UUID uuid = UUID.randomUUID();
+ byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid);
+ buf.setBytes(0, uuidBytes);
+
+ writer.setPosition(0);
+ writer.writeExtension(buf);
+
+ UUID result = vector.getObject(0);
+ assertEquals(uuid, result);
+ assertEquals(1, vector.getValueCount());
+ }
+ }
+
+ @Test
+ void testWriteExtensionWithUnsupportedType() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ writer.setPosition(0);
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () -> writer.writeExtension("invalid-type"));
+
+ assertTrue(
+ exception.getMessage().contains("Unsupported value type for UUID: java.lang.String"));
+ }
+ }
+
+ @Test
+ void testWriteExtensionMultipleValues() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+ UUID uuid3 = UUID.randomUUID();
+
+ writer.setPosition(0);
+ writer.writeExtension(uuid1);
+ writer.setPosition(1);
+ writer.writeExtension(uuid2);
+ writer.setPosition(2);
+ writer.writeExtension(uuid3);
+
+ assertEquals(uuid1, vector.getObject(0));
+ assertEquals(uuid2, vector.getObject(1));
+ assertEquals(uuid3, vector.getObject(2));
+ assertEquals(3, vector.getValueCount());
+ }
+ }
+
+ @Test
+ void testWriteWithUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector);
+ ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) {
+ UUID uuid = UUID.randomUUID();
+ byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid);
+ buf.setBytes(0, uuidBytes);
+
+ UuidHolder holder = new UuidHolder();
+ holder.buffer = buf;
+ holder.isSet = 1;
+
+ writer.setPosition(0);
+ writer.write(holder);
+
+ UUID result = vector.getObject(0);
+ assertEquals(uuid, result);
+ assertEquals(1, vector.getValueCount());
+ }
+ }
+
+ @Test
+ void testWriteWithNullableUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector);
+ ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) {
+ UUID uuid = UUID.randomUUID();
+ byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid);
+ buf.setBytes(0, uuidBytes);
+
+ NullableUuidHolder holder = new NullableUuidHolder();
+ holder.buffer = buf;
+ holder.isSet = 1;
+
+ writer.setPosition(0);
+ writer.write(holder);
+
+ UUID result = vector.getObject(0);
+ assertEquals(uuid, result);
+ assertEquals(1, vector.getValueCount());
+ }
+ }
+
+ @Test
+ void testWriteWithNullableUuidHolderNull() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ NullableUuidHolder holder = new NullableUuidHolder();
+ holder.isSet = 0;
+
+ writer.setPosition(0);
+ writer.write(holder);
+
+ assertTrue(vector.isNull(0));
+ assertEquals(1, vector.getValueCount());
+ }
+ }
+
+ // ========== Reader Tests ==========
+
+ @Test
+ void testReaderCopyAsValueExtensionVector() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator);
+ UuidVector vectorForRead = new UuidVector("test2", allocator);
+ UuidWriterImpl writer = new UuidWriterImpl(vector)) {
+ UUID uuid = UUID.randomUUID();
+ vectorForRead.setValueCount(1);
+ vectorForRead.set(0, uuid);
+ UuidReaderImpl reader = (UuidReaderImpl) vectorForRead.getReader();
+ reader.copyAsValue(writer);
+ UuidReaderImpl reader2 = (UuidReaderImpl) vector.getReader();
+ NullableUuidHolder holder = new NullableUuidHolder();
+ reader2.read(0, holder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(uuid, actualUuid);
+ }
+ }
+
+ @Test
+ void testReaderReadWithUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ vector.setSafe(0, uuid);
+ vector.setValueCount(1);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ NullableUuidHolder holder = new NullableUuidHolder();
+ reader.read(holder);
+
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(uuid, actualUuid);
+ assertEquals(1, holder.isSet);
+ }
+ }
+
+ @Test
+ void testReaderReadWithNullableUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ vector.setSafe(0, uuid);
+ vector.setValueCount(1);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ NullableUuidHolder holder = new NullableUuidHolder();
+ reader.read(holder);
+
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(uuid, actualUuid);
+ assertEquals(1, holder.isSet);
+ }
+ }
+
+ @Test
+ void testReaderReadWithNullableUuidHolderNull() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ NullableUuidHolder holder = new NullableUuidHolder();
+ reader.read(holder);
+
+ assertEquals(0, holder.isSet);
+ }
+ }
+
+ @Test
+ void testReaderReadWithArrayIndexUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+ UUID uuid3 = UUID.randomUUID();
+
+ vector.setSafe(0, uuid1);
+ vector.setSafe(1, uuid2);
+ vector.setSafe(2, uuid3);
+ vector.setValueCount(3);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+
+ NullableUuidHolder holder = new NullableUuidHolder();
+ reader.read(1, holder);
+
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+ assertEquals(uuid2, actualUuid);
+ assertEquals(1, holder.isSet);
+ }
+ }
+
+ @Test
+ void testReaderReadWithArrayIndexNullableUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ vector.setSafe(0, uuid1);
+ vector.setNull(1);
+ vector.setSafe(2, uuid2);
+ vector.setValueCount(3);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+
+ NullableUuidHolder holder1 = new NullableUuidHolder();
+ reader.read(0, holder1);
+ assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder1.buffer, holder1.start));
+ assertEquals(1, holder1.isSet);
+
+ NullableUuidHolder holder2 = new NullableUuidHolder();
+ reader.read(1, holder2);
+ assertEquals(0, holder2.isSet);
+
+ NullableUuidHolder holder3 = new NullableUuidHolder();
+ reader.read(2, holder3);
+ assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder3.buffer, holder3.start));
+ assertEquals(1, holder3.isSet);
+ }
+ }
+
+ @Test
+ void testReaderReadWithUnsupportedHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ vector.setSafe(0, uuid);
+ vector.setValueCount(1);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ reader.setPosition(0);
+
+ // Create a mock unsupported holder
+ ExtensionHolder unsupportedHolder =
+ new ExtensionHolder() {
+ @Override
+ public ArrowType type() {
+ return null;
+ }
+ };
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () -> reader.read(unsupportedHolder));
+
+ assertTrue(exception.getMessage().contains("Unsupported holder type for UuidReader"));
+ }
+ }
+
+ @Test
+ void testReaderIsSet() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ vector.setSafe(0, uuid);
+ vector.setNull(1);
+ vector.setSafe(2, uuid);
+ vector.setValueCount(3);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+
+ reader.setPosition(0);
+ assertTrue(reader.isSet());
+
+ reader.setPosition(1);
+ assertFalse(reader.isSet());
+
+ reader.setPosition(2);
+ assertTrue(reader.isSet());
+ }
+ }
+
+ @Test
+ void testReaderReadObject() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ vector.setSafe(0, uuid1);
+ vector.setNull(1);
+ vector.setSafe(2, uuid2);
+ vector.setValueCount(3);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+
+ reader.setPosition(0);
+ assertEquals(uuid1, reader.readObject());
+
+ reader.setPosition(1);
+ assertNull(reader.readObject());
+
+ reader.setPosition(2);
+ assertEquals(uuid2, reader.readObject());
+ }
+ }
+
+ @Test
+ void testReaderGetMinorType() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ assertEquals(vector.getMinorType(), reader.getMinorType());
+ }
+ }
+
+ @Test
+ void testReaderGetField() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ assertEquals(vector.getField(), reader.getField());
+ assertEquals("test", reader.getField().getName());
+ }
+ }
+
+ @Test
+ void testHolderStartOffsetWithMultipleValues() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+ UUID uuid3 = UUID.randomUUID();
+
+ vector.setSafe(0, uuid1);
+ vector.setSafe(1, uuid2);
+ vector.setSafe(2, uuid3);
+ vector.setValueCount(3);
+
+ // Test UuidHolder with different indices
+ NullableUuidHolder holder = new NullableUuidHolder();
+ vector.get(0, holder);
+ assertEquals(0, holder.start);
+ assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start));
+
+ vector.get(1, holder);
+ assertEquals(16, holder.start); // UUID_BYTE_WIDTH = 16
+ assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start));
+
+ vector.get(2, holder);
+ assertEquals(32, holder.start); // 2 * UUID_BYTE_WIDTH = 32
+ assertEquals(uuid3, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start));
+ }
+ }
+
+ @Test
+ void testNullableHolderStartOffsetWithMultipleValues() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ vector.setSafe(0, uuid1);
+ vector.setNull(1);
+ vector.setSafe(2, uuid2);
+ vector.setValueCount(3);
+
+ // Test NullableUuidHolder with different indices
+ NullableUuidHolder holder1 = new NullableUuidHolder();
+ vector.get(0, holder1);
+ assertEquals(0, holder1.start);
+ assertEquals(1, holder1.isSet);
+ assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder1.buffer, holder1.start));
+
+ NullableUuidHolder holder2 = new NullableUuidHolder();
+ vector.get(1, holder2);
+ assertEquals(0, holder2.isSet);
+
+ NullableUuidHolder holder3 = new NullableUuidHolder();
+ vector.get(2, holder3);
+ assertEquals(32, holder3.start); // 2 * UUID_BYTE_WIDTH = 32
+ assertEquals(1, holder3.isSet);
+ assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder3.buffer, holder3.start));
+
+ // Verify all holders share the same buffer
+ assertEquals(holder1.buffer, holder3.buffer);
+ }
+ }
+
+ @Test
+ void testSetFromHolderWithStartOffset() throws Exception {
+ try (UuidVector sourceVector = new UuidVector("source", allocator);
+ UuidVector targetVector = new UuidVector("target", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ sourceVector.setSafe(0, uuid1);
+ sourceVector.setSafe(1, uuid2);
+ sourceVector.setValueCount(3);
+
+ // Get holder from index 1 (should have start = 16)
+ NullableUuidHolder holder = new NullableUuidHolder();
+ sourceVector.get(1, holder);
+ assertEquals(16, holder.start);
+
+ // Set target vector using holder with non-zero start offset
+ targetVector.setSafe(0, holder);
+ targetVector.setValueCount(1);
+
+ // Verify the value was copied correctly
+ assertEquals(uuid2, targetVector.getObject(0));
+ }
+ }
+
+ @Test
+ void testSetFromNullableHolderWithStartOffset() throws Exception {
+ try (UuidVector sourceVector = new UuidVector("source", allocator);
+ UuidVector targetVector = new UuidVector("target", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+
+ sourceVector.setSafe(0, uuid1);
+ sourceVector.setNull(1);
+ sourceVector.setSafe(2, uuid2);
+ sourceVector.setValueCount(3);
+
+ // Get holder from index 2 (should have start = 32)
+ NullableUuidHolder holder = new NullableUuidHolder();
+ sourceVector.get(2, holder);
+ assertEquals(32, holder.start);
+ assertEquals(1, holder.isSet);
+
+ // Set target vector using holder with non-zero start offset
+ targetVector.setSafe(0, holder);
+ targetVector.setValueCount(1);
+
+ // Verify the value was copied correctly
+ assertEquals(uuid2, targetVector.getObject(0));
+
+ // Test with null holder
+ NullableUuidHolder nullHolder = new NullableUuidHolder();
+ sourceVector.get(1, nullHolder);
+ assertEquals(0, nullHolder.isSet);
+
+ targetVector.setSafe(1, nullHolder);
+ targetVector.setValueCount(2);
+ assertTrue(targetVector.isNull(1));
+ }
+ }
+
+ @Test
+ void testGetStartOffset() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ vector.allocateNew(10);
+
+ // Test getStartOffset for various indices
+ assertEquals(0, vector.getStartOffset(0));
+ assertEquals(16, vector.getStartOffset(1));
+ assertEquals(32, vector.getStartOffset(2));
+ assertEquals(48, vector.getStartOffset(3));
+ assertEquals(160, vector.getStartOffset(10));
+ }
+ }
+
+ @Test
+ void testReaderWithStartOffsetMultipleReads() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+ UUID uuid3 = UUID.randomUUID();
+
+ vector.setSafe(0, uuid1);
+ vector.setSafe(1, uuid2);
+ vector.setSafe(2, uuid3);
+ vector.setValueCount(3);
+
+ UuidReaderImpl reader = (UuidReaderImpl) vector.getReader();
+ NullableUuidHolder holder = new NullableUuidHolder();
+
+ // Read from different positions and verify start offset
+ reader.read(0, holder);
+ assertEquals(0, holder.start);
+ assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start));
+
+ reader.read(1, holder);
+ assertEquals(16, holder.start);
+ assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start));
+
+ reader.read(2, holder);
+ assertEquals(32, holder.start);
+ assertEquals(uuid3, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start));
+ }
+ }
+
+ @Test
+ void testWriterWithExtensionHolder() throws Exception {
+ try (UuidVector sourceVector = new UuidVector("source", allocator);
+ UuidVector targetVector = new UuidVector("target", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ sourceVector.setSafe(0, uuid);
+ sourceVector.setValueCount(1);
+
+ // Get holder from source
+ NullableUuidHolder holder = new NullableUuidHolder();
+ sourceVector.get(0, holder);
+
+ // Write using UuidWriterImpl with ExtensionHolder
+ UuidWriterImpl writer = new UuidWriterImpl(targetVector);
+ writer.setPosition(0);
+ writer.writeExtension(holder);
+
+ assertEquals(uuid, targetVector.getObject(0));
+ }
+ }
+
+ @Test
+ void testNullableUuidHolderReaderImpl() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ vector.setSafe(0, uuid);
+ vector.setValueCount(1);
+
+ // Get holder from vector
+ NullableUuidHolder sourceHolder = new NullableUuidHolder();
+ vector.get(0, sourceHolder);
+ assertEquals(1, sourceHolder.isSet);
+ assertEquals(0, sourceHolder.start);
+
+ // Create reader from holder
+ NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder);
+ assertTrue(reader.isSet());
+ assertEquals(uuid, reader.readObject());
+
+ // Read into another holder
+ NullableUuidHolder targetHolder = new NullableUuidHolder();
+ reader.read(targetHolder);
+ assertEquals(1, targetHolder.isSet);
+ assertEquals(0, targetHolder.start);
+ assertEquals(uuid, UuidUtility.uuidFromArrowBuf(targetHolder.buffer, targetHolder.start));
+ }
+ }
+
+ @Test
+ void testNullableUuidHolderReaderImplWithNull() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ vector.setNull(0);
+ vector.setValueCount(1);
+
+ // Get null holder from vector
+ NullableUuidHolder sourceHolder = new NullableUuidHolder();
+ vector.get(0, sourceHolder);
+ assertEquals(0, sourceHolder.isSet);
+
+ // Create reader from null holder
+ NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder);
+ assertFalse(reader.isSet());
+ assertNull(reader.readObject());
+
+ // Read into another holder
+ NullableUuidHolder targetHolder = new NullableUuidHolder();
+ reader.read(targetHolder);
+ assertEquals(0, targetHolder.isSet);
+ }
+ }
+
+ @Test
+ void testNullableUuidHolderReaderImplReadIntoUuidHolder() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid = UUID.randomUUID();
+ vector.setSafe(0, uuid);
+ vector.setValueCount(1);
+
+ // Get holder from vector
+ NullableUuidHolder sourceHolder = new NullableUuidHolder();
+ vector.get(0, sourceHolder);
+
+ // Create reader from holder
+ NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder);
+
+ // Read into UuidHolder (non-nullable)
+ UuidHolder targetHolder = new UuidHolder();
+ reader.read(targetHolder);
+ assertEquals(0, targetHolder.start);
+ assertEquals(uuid, UuidUtility.uuidFromArrowBuf(targetHolder.buffer, targetHolder.start));
+ }
+ }
+
+ @Test
+ void testNullableUuidHolderReaderImplWithNonZeroStart() throws Exception {
+ try (UuidVector vector = new UuidVector("test", allocator)) {
+ UUID uuid1 = UUID.randomUUID();
+ UUID uuid2 = UUID.randomUUID();
+ vector.setSafe(0, uuid1);
+ vector.setSafe(1, uuid2);
+ vector.setValueCount(2);
+
+ // Get holder from index 1 (start = 16)
+ NullableUuidHolder sourceHolder = new NullableUuidHolder();
+ vector.get(1, sourceHolder);
+ assertEquals(1, sourceHolder.isSet);
+ assertEquals(16, sourceHolder.start);
+
+ // Create reader from holder
+ NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder);
+ assertEquals(uuid2, reader.readObject());
+
+ // Read into another holder and verify start is preserved
+ NullableUuidHolder targetHolder = new NullableUuidHolder();
+ reader.read(targetHolder);
+ assertEquals(16, targetHolder.start);
+ assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(targetHolder.buffer, targetHolder.start));
+ }
+ }
+}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java b/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
index daec331831..22c93b0cbe 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.apache.arrow.vector.TestUtils.newVarBinaryVector;
import static org.apache.arrow.vector.TestUtils.newVarCharVector;
import static org.apache.arrow.vector.TestUtils.newVector;
@@ -56,6 +57,10 @@
import org.apache.arrow.vector.complex.impl.UnionListViewWriter;
import org.apache.arrow.vector.complex.impl.UnionListWriter;
import org.apache.arrow.vector.holders.NullableIntHolder;
+import org.apache.arrow.vector.holders.NullableTimeStampMicroTZHolder;
+import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder;
+import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder;
+import org.apache.arrow.vector.holders.NullableTimeStampSecTZHolder;
import org.apache.arrow.vector.holders.NullableUInt4Holder;
import org.apache.arrow.vector.holders.NullableVarBinaryHolder;
import org.apache.arrow.vector.holders.NullableVarCharHolder;
@@ -1233,7 +1238,7 @@ public void testSplitAndTransfer3() {
// the size needed for the validity buffer
final long validitySize =
DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize(
- BaseValueVector.getValidityBufferSizeFromCount(2));
+ getValidityBufferSizeFromCount(2));
assertEquals(allocatedMem + validitySize, allocator.getAllocatedMemory());
// The validity and offset buffers are sliced from a same buffer.See
// BaseFixedWidthVector#allocateBytes.
@@ -2464,7 +2469,7 @@ public void testDefaultAllocNewAll() {
assertTrue(intVector.getValueCapacity() >= defaultCapacity);
expectedSize =
(defaultCapacity * IntVector.TYPE_WIDTH)
- + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity);
+ + getValidityBufferSizeFromCount(defaultCapacity);
assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05);
// verify that the wastage is within bounds for BigIntVector.
@@ -2473,7 +2478,7 @@ public void testDefaultAllocNewAll() {
assertTrue(bigIntVector.getValueCapacity() >= defaultCapacity);
expectedSize =
(defaultCapacity * bigIntVector.TYPE_WIDTH)
- + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity);
+ + getValidityBufferSizeFromCount(defaultCapacity);
assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05);
// verify that the wastage is within bounds for DecimalVector.
@@ -2482,7 +2487,7 @@ public void testDefaultAllocNewAll() {
assertTrue(decimalVector.getValueCapacity() >= defaultCapacity);
expectedSize =
(defaultCapacity * decimalVector.TYPE_WIDTH)
- + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity);
+ + getValidityBufferSizeFromCount(defaultCapacity);
assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05);
// verify that the wastage is within bounds for VarCharVector.
@@ -2492,7 +2497,7 @@ public void testDefaultAllocNewAll() {
assertTrue(varCharVector.getValueCapacity() >= defaultCapacity - 1);
expectedSize =
(defaultCapacity * VarCharVector.OFFSET_WIDTH)
- + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity)
+ + getValidityBufferSizeFromCount(defaultCapacity)
+ defaultCapacity * 8;
// wastage should be less than 5%.
assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05);
@@ -2501,7 +2506,7 @@ public void testDefaultAllocNewAll() {
beforeSize = childAllocator.getAllocatedMemory();
bitVector.allocateNew();
assertTrue(bitVector.getValueCapacity() >= defaultCapacity);
- expectedSize = BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity) * 2;
+ expectedSize = getValidityBufferSizeFromCount(defaultCapacity) * 2;
assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05);
}
}
@@ -2566,6 +2571,195 @@ public void testSetNullableVarCharHolderSafe() {
}
}
+ @Test
+ public void testTimeStampTZVectorSetSafeUnset() {
+ // reproduction of https://github.com/apache/arrow/issues/45084
+ try (TimeStampMicroTZVector vector = new TimeStampMicroTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+ // Set a valid value
+ NullableTimeStampMicroTZHolder validHolder = new NullableTimeStampMicroTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ // Unset the value using a holder with default (null) timezone
+ // The bug used to throw IllegalArgumentException because holder.timezone (null) !=
+ // vector.timezone ("UTC")
+ // The correct behaviour is to not throw an exception and to unset the value.
+ NullableTimeStampMicroTZHolder unsetHolder = new NullableTimeStampMicroTZHolder();
+ unsetHolder.isSet = 0;
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampMilliTZVectorSetSafeUnset() {
+ // reproduction of https://github.com/apache/arrow/issues/45084
+ try (TimeStampMilliTZVector vector = new TimeStampMilliTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampMilliTZHolder validHolder = new NullableTimeStampMilliTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampMilliTZHolder unsetHolder = new NullableTimeStampMilliTZHolder();
+ unsetHolder.isSet = 0;
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampNanoTZVectorSetSafeUnset() {
+ // reproduction of https://github.com/apache/arrow/issues/45084
+ try (TimeStampNanoTZVector vector = new TimeStampNanoTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampNanoTZHolder validHolder = new NullableTimeStampNanoTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampNanoTZHolder unsetHolder = new NullableTimeStampNanoTZHolder();
+ unsetHolder.isSet = 0;
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampSecTZVectorSetSafeUnset() {
+ // reproduction of https://github.com/apache/arrow/issues/45084
+ try (TimeStampSecTZVector vector = new TimeStampSecTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampSecTZHolder validHolder = new NullableTimeStampSecTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampSecTZHolder unsetHolder = new NullableTimeStampSecTZHolder();
+ unsetHolder.isSet = 0;
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampMicroTZVectorSetSafeUnsetExplicitTimezone() {
+ // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break
+ // workaround.
+ try (TimeStampMicroTZVector vector = new TimeStampMicroTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampMicroTZHolder validHolder = new NullableTimeStampMicroTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampMicroTZHolder unsetHolder = new NullableTimeStampMicroTZHolder();
+ unsetHolder.isSet = 0;
+ unsetHolder.timezone = "UTC";
+
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampMilliTZVectorSetSafeUnsetExplicitTimezone() {
+ // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break
+ // workaround.
+ try (TimeStampMilliTZVector vector = new TimeStampMilliTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampMilliTZHolder validHolder = new NullableTimeStampMilliTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampMilliTZHolder unsetHolder = new NullableTimeStampMilliTZHolder();
+ unsetHolder.isSet = 0;
+ unsetHolder.timezone = "UTC";
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampNanoTZVectorSetSafeUnsetExplicitTimezone() {
+ // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break
+ // workaround.
+ try (TimeStampNanoTZVector vector = new TimeStampNanoTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampNanoTZHolder validHolder = new NullableTimeStampNanoTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampNanoTZHolder unsetHolder = new NullableTimeStampNanoTZHolder();
+ unsetHolder.isSet = 0;
+ unsetHolder.timezone = "UTC";
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
+ @Test
+ public void testTimeStampSecTZVectorSetSafeUnsetExplicitTimezone() {
+ // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break
+ // workaround.
+ try (TimeStampSecTZVector vector = new TimeStampSecTZVector("vector", allocator, "UTC")) {
+ vector.allocateNew();
+
+ NullableTimeStampSecTZHolder validHolder = new NullableTimeStampSecTZHolder();
+ validHolder.isSet = 1;
+ validHolder.value = 1000L;
+ validHolder.timezone = "UTC";
+ vector.setSafe(0, validHolder);
+
+ assertEquals(1000L, vector.get(0));
+
+ NullableTimeStampSecTZHolder unsetHolder = new NullableTimeStampSecTZHolder();
+ unsetHolder.isSet = 0;
+ unsetHolder.timezone = "UTC";
+ vector.setSafe(0, unsetHolder);
+
+ assertNull(vector.getObject(0));
+ }
+ }
+
@Test
public void testSetNullableVarBinaryHolder() {
try (VarBinaryVector vector = new VarBinaryVector("", allocator)) {
@@ -3746,4 +3940,42 @@ public void testVectorLoadUnloadOnNonVariadicVectors() {
}
}
}
+
+ @Test
+ public void testEmptyVarCharOffsetBuffer() {
+ // Validates that offset buffer has at least OFFSET_WIDTH bytes (for offset[0]=0)
+ // even when valueCount is 0, per Arrow specification.
+ try (VarCharVector vector = newVarCharVector("varchar", allocator)) {
+ vector.allocateNew();
+ vector.setValueCount(0);
+
+ List buffers = vector.getFieldBuffers();
+ // buffers: [validity, offset, data]
+ assertTrue(
+ buffers.get(1).readableBytes() >= BaseVariableWidthVector.OFFSET_WIDTH,
+ "Offset buffer should have at least "
+ + BaseVariableWidthVector.OFFSET_WIDTH
+ + " bytes for offset[0]");
+ assertEquals(0, vector.getOffsetBuffer().getInt(0));
+ }
+ }
+
+ @Test
+ public void testEmptyLargeVarCharOffsetBuffer() {
+ // Validates that offset buffer has at least OFFSET_WIDTH bytes (for offset[0]=0)
+ // even when valueCount is 0, per Arrow specification.
+ try (LargeVarCharVector vector = new LargeVarCharVector("largevarchar", allocator)) {
+ vector.allocateNew();
+ vector.setValueCount(0);
+
+ List buffers = vector.getFieldBuffers();
+ // buffers: [validity, offset, data]
+ assertTrue(
+ buffers.get(1).readableBytes() >= BaseLargeVariableWidthVector.OFFSET_WIDTH,
+ "Offset buffer should have at least "
+ + BaseLargeVariableWidthVector.OFFSET_WIDTH
+ + " bytes for offset[0]");
+ assertEquals(0, vector.getOffsetBuffer().getLong(0));
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java
index 7a3a1bae63..baf5e672c8 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.apache.arrow.vector.TestUtils.newVector;
import static org.apache.arrow.vector.TestUtils.newViewVarBinaryVector;
import static org.apache.arrow.vector.TestUtils.newViewVarCharVector;
@@ -60,6 +61,7 @@
import org.apache.arrow.vector.util.ReusableByteArray;
import org.apache.arrow.vector.util.Text;
import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.validate.ValidateUtil;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -2367,7 +2369,7 @@ private void testSplitAndTransferOnValiditySplitHelper(
// the allocation only consists in the size needed for the validity buffer
final long validitySize =
DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize(
- BaseValueVector.getValidityBufferSizeFromCount(2));
+ getValidityBufferSizeFromCount(2));
// we allocate view and data buffers for the target vector
assertTrue(allocatedMem + validitySize < allocator.getAllocatedMemory());
// The validity is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes.
@@ -2444,7 +2446,7 @@ public void testSplitAndTransferWithLongStringsOnValiditySplit() {
final ViewVarBinaryVector sourceVector =
newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) {
testSplitAndTransferOnValiditySplitHelper(
- targetVector, sourceVector, startIndex, length, data);
+ targetVector, sourceVector, startIndex, length, binaryData);
}
}
@@ -2851,4 +2853,18 @@ public void testVectorLoadUnloadOnMixedTypes() {
}
}
}
+
+ @Test
+ public void testValidate() {
+ try (final ViewVarCharVector vector = new ViewVarCharVector("v", allocator)) {
+ vector.validateFull();
+ setVector(vector, STR1, STR2, STR3);
+ vector.validateFull();
+
+ vector.getDataBuffer().capacity(0);
+ ValidateUtil.ValidateException e =
+ assertThrows(ValidateUtil.ValidateException.class, () -> vector.validate());
+ assertTrue(e.getMessage().contains("Not enough capacity for data buffer"));
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java
index c121d94892..bd3113f8bc 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java
@@ -171,6 +171,26 @@ public void testAddVector() {
}
}
+ @Test
+ public void testAddVectorAtEnd() {
+ try (final IntVector intVector1 = new IntVector("intVector1", allocator);
+ final IntVector intVector2 = new IntVector("intVector2", allocator);
+ final IntVector intVector3 = new IntVector("intVector3", allocator); ) {
+
+ VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2));
+ assertEquals(2, original.getFieldVectors().size());
+
+ VectorSchemaRoot newRecordBatch = original.addVector(2, intVector3);
+ assertEquals(3, newRecordBatch.getFieldVectors().size());
+ assertEquals(intVector1, newRecordBatch.getFieldVectors().get(0));
+ assertEquals(intVector2, newRecordBatch.getFieldVectors().get(1));
+ assertEquals(intVector3, newRecordBatch.getFieldVectors().get(2));
+
+ original.close();
+ newRecordBatch.close();
+ }
+ }
+
@Test
public void testRemoveVector() {
try (final IntVector intVector1 = new IntVector("intVector1", allocator);
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java
index 6121fb67fe..782535fccc 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java
@@ -17,6 +17,7 @@
package org.apache.arrow.vector;
import static java.util.Arrays.asList;
+import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -215,7 +216,7 @@ public void testLoadValidityBuffer() throws IOException {
int count = 10;
ArrowBuf[] values = new ArrowBuf[4];
for (int i = 0; i < 4; i += 2) {
- ArrowBuf buf1 = allocator.buffer(BitVectorHelper.getValidityBufferSize(count));
+ ArrowBuf buf1 = allocator.buffer(getValidityBufferSizeFromCount(count));
ArrowBuf buf2 = allocator.buffer(count * 4); // integers
buf1.setZero(0, buf1.capacity());
buf2.setZero(0, buf2.capacity());
diff --git a/vector/src/test/java/org/apache/arrow/vector/UuidVector.java b/vector/src/test/java/org/apache/arrow/vector/UuidVector.java
deleted file mode 100644
index 5c90d45f60..0000000000
--- a/vector/src/test/java/org/apache/arrow/vector/UuidVector.java
+++ /dev/null
@@ -1,114 +0,0 @@
-/*
- * 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.nio.ByteBuffer;
-import java.util.UUID;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.util.hash.ArrowBufHasher;
-import org.apache.arrow.vector.types.pojo.Field;
-import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.UuidType;
-import org.apache.arrow.vector.util.TransferPair;
-
-public class UuidVector extends ExtensionTypeVector
- implements ValueIterableVector {
- private final Field field;
-
- public UuidVector(
- String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) {
- super(name, allocator, underlyingVector);
- this.field = new Field(name, FieldType.nullable(new UuidType()), null);
- }
-
- public UuidVector(String name, BufferAllocator allocator) {
- super(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));
- this.field = new Field(name, FieldType.nullable(new UuidType()), null);
- }
-
- @Override
- public UUID getObject(int index) {
- 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) {
- return getUnderlyingVector().hashCode(index, hasher);
- }
-
- public void set(int index, UUID uuid) {
- ByteBuffer bb = ByteBuffer.allocate(16);
- bb.putLong(uuid.getMostSignificantBits());
- bb.putLong(uuid.getLeastSignificantBits());
- getUnderlyingVector().set(index, bb.array());
- }
-
- @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 TransferPair makeTransferPair(ValueVector to) {
- return new TransferImpl((UuidVector) to);
- }
-
- public void setSafe(int index, byte[] value) {
- getUnderlyingVector().setIndexDefined(index);
- getUnderlyingVector().setSafe(index, value);
- }
-
- public class TransferImpl implements TransferPair {
- UuidVector to;
- ValueVector targetUnderlyingVector;
- TransferPair tp;
-
- public TransferImpl(UuidVector to) {
- this.to = to;
- targetUnderlyingVector = this.to.getUnderlyingVector();
- tp = getUnderlyingVector().makeTransferPair(targetUnderlyingVector);
- }
-
- public UuidVector getTo() {
- return this.to;
- }
-
- public void transfer() {
- tp.transfer();
- }
-
- public void splitAndTransfer(int startIndex, int length) {
- tp.splitAndTransfer(startIndex, length);
- }
-
- public void copyValueSafe(int fromIndex, int toIndex) {
- tp.copyValueSafe(fromIndex, toIndex);
- }
- }
-}
diff --git a/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java b/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java
index 08da786eb2..9624734356 100644
--- a/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java
+++ b/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java
@@ -22,6 +22,7 @@
import java.nio.charset.Charset;
import java.util.Arrays;
+import java.util.List;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
@@ -39,6 +40,7 @@
import org.apache.arrow.vector.complex.LargeListViewVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.ListViewVector;
+import org.apache.arrow.vector.complex.RunEndEncodedVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
import org.apache.arrow.vector.complex.impl.NullableStructWriter;
@@ -53,7 +55,9 @@
import org.apache.arrow.vector.holders.NullableUInt4Holder;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ArrowType.RunEndEncoded;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.junit.jupiter.api.AfterEach;
@@ -1003,6 +1007,54 @@ public void testLargeListViewVectorApproxEquals() {
}
}
+ @Test
+ public void testRunEndEncodedFloat8ApproxEquals() {
+ try (final Float8Vector vector1 = new Float8Vector("float", allocator);
+ final Float8Vector vector2 = new Float8Vector("float", allocator);
+ final Float8Vector vector3 = new Float8Vector("float", allocator);
+ final IntVector reeVector = new IntVector("ree", allocator)) {
+
+ final float epsilon = 1.0E-6f;
+ setVector(vector1, 1.1, 2.2);
+ setVector(vector2, 1.1 + epsilon / 2, 2.2 + epsilon / 2);
+ setVector(vector3, 1.1 + epsilon * 2, 2.2 + epsilon * 2);
+ setVector(reeVector, 1, 3);
+
+ ArrowType type = MinorType.FLOAT8.getType();
+ final FieldType valueType = FieldType.notNullable(type);
+ final FieldType runEndType = FieldType.notNullable(MinorType.INT.getType());
+
+ final Field valueField = new Field("value", valueType, null);
+ final Field runEndField = new Field("ree", runEndType, null);
+
+ Field field =
+ new Field(
+ "ree_float",
+ FieldType.notNullable(RunEndEncoded.INSTANCE),
+ List.of(runEndField, valueField));
+
+ try (final RunEndEncodedVector encodedVector1 =
+ new RunEndEncodedVector(field, allocator, reeVector, vector1, null);
+ final RunEndEncodedVector encodedVector2 =
+ new RunEndEncodedVector(field, allocator, reeVector, vector2, null);
+ final RunEndEncodedVector encodedVector3 =
+ new RunEndEncodedVector(field, allocator, reeVector, vector3, null)) {
+
+ encodedVector1.setValueCount(3);
+ encodedVector2.setValueCount(3);
+ encodedVector3.setValueCount(3);
+
+ Range range = new Range(0, 0, encodedVector1.getValueCount());
+ assertTrue(
+ new ApproxEqualsVisitor(encodedVector1, encodedVector2, epsilon, epsilon)
+ .rangeEquals(range));
+ assertFalse(
+ new ApproxEqualsVisitor(encodedVector1, encodedVector3, epsilon, epsilon)
+ .rangeEquals(range));
+ }
+ }
+ }
+
private void writeStructVector(NullableStructWriter writer, int value1, long value2) {
writer.start();
writer.integer("f0").writeInt(value1);
diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java
index 3bc02c6029..b2a8cf9ba4 100644
--- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java
+++ b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java
@@ -20,6 +20,7 @@
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.math.BigDecimal;
+import java.util.UUID;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.DecimalVector;
@@ -30,8 +31,10 @@
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.reader.FieldReader;
import org.apache.arrow.vector.complex.writer.BaseWriter;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
import org.apache.arrow.vector.complex.writer.FieldWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.holders.DecimalHolder;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -845,4 +848,110 @@ public void testCopyMapVectorWithMapValue() {
assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
}
}
+
+ @Test
+ public void testCopyListVectorWithExtensionType() {
+ try (ListVector from = ListVector.empty("v", allocator);
+ ListVector to = ListVector.empty("v", allocator)) {
+
+ UnionListWriter listWriter = from.getWriter();
+ listWriter.allocate();
+
+ for (int i = 0; i < COUNT; i++) {
+ listWriter.setPosition(i);
+ listWriter.startList();
+ ExtensionWriter extensionWriter = listWriter.extension(UuidType.INSTANCE);
+ extensionWriter.writeExtension(UUID.randomUUID());
+ extensionWriter.writeExtension(UUID.randomUUID());
+ listWriter.endList();
+ }
+ from.setValueCount(COUNT);
+
+ // copy values
+ FieldReader in = from.getReader();
+ FieldWriter out = to.getWriter();
+ for (int i = 0; i < COUNT; i++) {
+ in.setPosition(i);
+ out.setPosition(i);
+ ComplexCopier.copy(in, out);
+ }
+
+ to.setValueCount(COUNT);
+
+ // validate equals
+ assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+ }
+ }
+
+ @Test
+ public void testCopyMapVectorWithExtensionType() {
+ try (final MapVector from = MapVector.empty("v", allocator, false);
+ final MapVector to = MapVector.empty("v", allocator, false)) {
+
+ from.allocateNew();
+
+ UnionMapWriter mapWriter = from.getWriter();
+ for (int i = 0; i < COUNT; i++) {
+ mapWriter.setPosition(i);
+ mapWriter.startMap();
+ mapWriter.startEntry();
+ ExtensionWriter extensionKeyWriter = mapWriter.key().extension(UuidType.INSTANCE);
+ extensionKeyWriter.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
+ ExtensionWriter extensionValueWriter = mapWriter.value().extension(UuidType.INSTANCE);
+ extensionValueWriter.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
+ mapWriter.endEntry();
+ mapWriter.endMap();
+ }
+
+ from.setValueCount(COUNT);
+
+ // copy values
+ FieldReader in = from.getReader();
+ FieldWriter out = to.getWriter();
+ for (int i = 0; i < COUNT; i++) {
+ in.setPosition(i);
+ out.setPosition(i);
+ ComplexCopier.copy(in, out);
+ }
+ to.setValueCount(COUNT);
+
+ // validate equals
+ assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+ }
+ }
+
+ @Test
+ public void testCopyStructVectorWithExtensionType() {
+ try (final StructVector from = StructVector.empty("v", allocator);
+ final StructVector to = StructVector.empty("v", allocator)) {
+
+ from.allocateNewSafe();
+
+ NullableStructWriter structWriter = from.getWriter();
+ for (int i = 0; i < COUNT; i++) {
+ structWriter.setPosition(i);
+ structWriter.start();
+ ExtensionWriter extensionWriter1 = structWriter.extension("uuid1", UuidType.INSTANCE);
+ extensionWriter1.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
+ ExtensionWriter extensionWriter2 = structWriter.extension("uuid2", UuidType.INSTANCE);
+ extensionWriter2.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
+ structWriter.end();
+ }
+
+ from.setValueCount(COUNT);
+
+ // copy values
+ FieldReader in = from.getReader();
+ FieldWriter out = to.getWriter();
+ for (int i = 0; i < COUNT; i++) {
+ in.setPosition(i);
+ out.setPosition(i);
+ ComplexCopier.copy(in, out);
+ }
+ to.setValueCount(COUNT);
+
+ // validate equals
+ assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java
index 1556852c5a..5b6d65d6ba 100644
--- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java
+++ b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java
@@ -31,6 +31,7 @@
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.DirtyRootAllocator;
+import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.LargeVarBinaryVector;
import org.apache.arrow.vector.LargeVarCharVector;
import org.apache.arrow.vector.UuidVector;
@@ -41,6 +42,7 @@
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.holders.DurationHolder;
import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
import org.apache.arrow.vector.holders.NullableDecimalHolder;
@@ -48,15 +50,16 @@
import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder;
import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
import org.apache.arrow.vector.holders.UnionHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.UuidType;
import org.apache.arrow.vector.util.DecimalUtility;
import org.apache.arrow.vector.util.Text;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -100,7 +103,6 @@ public void testPromoteToUnion() throws Exception {
writer.integer("A").writeInt(10);
// we don't write anything in 3
-
writer.setPosition(4);
writer.integer("A").writeInt(100);
@@ -130,9 +132,21 @@ public void testPromoteToUnion() throws Exception {
binHolder.buffer = buf;
writer.fixedSizeBinary("A", 4).write(binHolder);
+ writer.setPosition(9);
+ UUID uuid = UUID.randomUUID();
+ writer.extension("A", UuidType.INSTANCE).writeExtension(uuid, UuidType.INSTANCE);
+ writer.end();
+
+ writer.setPosition(10);
+ UUID uuid2 = UUID.randomUUID();
+ UuidHolder uuidHolder = new UuidHolder();
+ uuidHolder.buffer = allocator.buffer(UuidType.UUID_BYTE_WIDTH);
+ uuidHolder.buffer.setBytes(0, UuidUtility.getBytesFromUUID(uuid2));
+ writer.extension("A", UuidType.INSTANCE).write(uuidHolder);
writer.end();
+ allocator.releaseBytes(UuidType.UUID_BYTE_WIDTH);
- container.setValueCount(9);
+ container.setValueCount(11);
final UnionVector uv = v.getChild("A", UnionVector.class);
@@ -169,6 +183,12 @@ public void testPromoteToUnion() throws Exception {
.order(ByteOrder.nativeOrder())
.getInt());
+ assertFalse(uv.isNull(9), "9 shouldn't be null");
+ assertEquals(uuid, uv.getObject(9));
+
+ assertFalse(uv.isNull(10), "10 shouldn't be null");
+ assertEquals(uuid2, uv.getObject(10));
+
container.clear();
container.allocateNew();
@@ -785,18 +805,17 @@ public void testExtensionType() throws Exception {
try (final NonNullableStructVector container =
NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator);
final UuidVector v =
- container.addOrGet("uuid", FieldType.nullable(new UuidType()), UuidVector.class);
+ container.addOrGet("uuid", FieldType.nullable(UuidType.INSTANCE), UuidVector.class);
final PromotableWriter writer = new PromotableWriter(v, container)) {
UUID u1 = UUID.randomUUID();
UUID u2 = UUID.randomUUID();
container.allocateNew();
container.setValueCount(1);
- writer.addExtensionTypeWriterFactory(new UuidWriterFactory());
writer.setPosition(0);
- writer.writeExtension(u1);
+ writer.writeExtension(u1, UuidType.INSTANCE);
writer.setPosition(1);
- writer.writeExtension(u2);
+ writer.writeExtension(u2, UuidType.INSTANCE);
container.setValueCount(2);
@@ -805,4 +824,29 @@ public void testExtensionType() throws Exception {
assertEquals(u2, uuidVector.getObject(1));
}
}
+
+ @Test
+ public void testExtensionTypeForList() throws Exception {
+ try (final ListVector container = ListVector.empty(EMPTY_SCHEMA_PATH, allocator);
+ final UuidVector v =
+ (UuidVector)
+ container.addOrGetVector(FieldType.nullable(UuidType.INSTANCE)).getVector();
+ final PromotableWriter writer = new PromotableWriter(v, container)) {
+ UUID u1 = UUID.randomUUID();
+ UUID u2 = UUID.randomUUID();
+ container.allocateNew();
+ container.setValueCount(1);
+
+ writer.setPosition(0);
+ writer.writeExtension(u1, UuidType.INSTANCE);
+ writer.setPosition(1);
+ writer.writeExtension(u2, UuidType.INSTANCE);
+
+ container.setValueCount(2);
+
+ FieldVector uuidVector = container.getDataVector();
+ assertEquals(u1, uuidVector.getObject(0));
+ assertEquals(u2, uuidVector.getObject(1));
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java
index 2745386db4..80d03cae6d 100644
--- a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java
+++ b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java
@@ -19,6 +19,7 @@
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.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -31,6 +32,7 @@
import java.util.HashSet;
import java.util.List;
import java.util.Set;
+import java.util.UUID;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
@@ -71,9 +73,11 @@
import org.apache.arrow.vector.complex.reader.Float8Reader;
import org.apache.arrow.vector.complex.reader.IntReader;
import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.holders.DecimalHolder;
import org.apache.arrow.vector.holders.DurationHolder;
import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
@@ -82,8 +86,11 @@
import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder;
import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder;
import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
import org.apache.arrow.vector.types.TimeUnit;
+import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;
import org.apache.arrow.vector.types.pojo.ArrowType.Int;
@@ -99,6 +106,7 @@
import org.apache.arrow.vector.util.JsonStringHashMap;
import org.apache.arrow.vector.util.Text;
import org.apache.arrow.vector.util.TransferPair;
+import org.apache.arrow.vector.util.UuidUtility;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -1098,6 +1106,13 @@ public void simpleUnion() throws Exception {
new UnionVector("union", allocator, /* field type */ null, /* call-back */ null);
UnionWriter unionWriter = new UnionWriter(vector);
unionWriter.allocate();
+
+ UUID uuid = UUID.randomUUID();
+ ByteBuffer bb = ByteBuffer.allocate(16);
+ bb.putLong(uuid.getMostSignificantBits());
+ bb.putLong(uuid.getLeastSignificantBits());
+ byte[] uuidByte = bb.array();
+
for (int i = 0; i < COUNT; i++) {
unionWriter.setPosition(i);
if (i % 5 == 0) {
@@ -1120,6 +1135,12 @@ public void simpleUnion() throws Exception {
holder.buffer = buf;
unionWriter.write(holder);
bufs.add(buf);
+ } else if (i % 5 == 4) {
+ UuidHolder holder = new UuidHolder();
+ holder.buffer = allocator.buffer(UuidType.UUID_BYTE_WIDTH);
+ holder.buffer.setBytes(0, uuidByte);
+ unionWriter.write(holder);
+ allocator.releaseBytes(UuidType.UUID_BYTE_WIDTH);
} else {
unionWriter.writeFloat4((float) i);
}
@@ -1145,6 +1166,10 @@ public void simpleUnion() throws Exception {
unionReader.read(holder);
assertEquals(i, holder.buffer.getInt(0));
assertEquals(4, holder.byteWidth);
+ } else if (i % 5 == 4) {
+ NullableUuidHolder holder = new NullableUuidHolder();
+ unionReader.read(holder);
+ assertEquals(UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start), uuid);
} else {
assertEquals((float) i, unionReader.readFloat(), 1e-12);
}
@@ -2489,4 +2514,79 @@ public void unionWithVarCharAndBinaryHelpers() throws Exception {
"row12", new String(vector.getLargeVarBinaryVector().get(11), StandardCharsets.UTF_8));
}
}
+
+ @Test
+ public void extensionWriterReader() throws Exception {
+ // test values
+ UUID u1 = UUID.randomUUID();
+
+ try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
+ // write
+
+ ComplexWriter writer = new ComplexWriterImpl("root", parent);
+ StructWriter rootWriter = writer.rootAsStruct();
+
+ {
+ ExtensionWriter extensionWriter = rootWriter.extension("uuid1", UuidType.INSTANCE);
+ extensionWriter.setPosition(0);
+ extensionWriter.writeExtension(u1, UuidType.INSTANCE);
+ }
+ // read
+ StructReader rootReader = new SingleStructReaderImpl(parent).reader("root");
+ {
+ FieldReader uuidReader = rootReader.reader("uuid1");
+ uuidReader.setPosition(0);
+ NullableUuidHolder uuidHolder = new NullableUuidHolder();
+ uuidReader.read(uuidHolder);
+ UUID actualUuid = UuidUtility.uuidFromArrowBuf(uuidHolder.buffer, 0);
+ assertEquals(u1, actualUuid);
+ assertTrue(uuidReader.isSet());
+ assertEquals(uuidReader.getMinorType(), MinorType.EXTENSIONTYPE);
+ assertInstanceOf(UuidType.class, uuidReader.getField().getFieldType().getType());
+ }
+ }
+ }
+
+ @Test
+ void testListOfDenseUnionWriterNPE() {
+ // Regression test for https://github.com/apache/arrow-java/issues/399
+ try (ListVector listVector = ListVector.empty("list", allocator)) {
+ listVector.addOrGetVector(FieldType.nullable(MinorType.DENSEUNION.getType()));
+ UnionListWriter listWriter = listVector.getWriter();
+
+ listWriter.startList();
+ listWriter.endList();
+ }
+ }
+
+ @Test
+ void testListOfDenseUnionWriterWithData() {
+ try (ListVector listVector = ListVector.empty("list", allocator)) {
+ listVector.addOrGetVector(FieldType.nullable(MinorType.DENSEUNION.getType()));
+
+ UnionListWriter listWriter = listVector.getWriter();
+ listWriter.startList();
+ listWriter.writeInt(100);
+ listWriter.writeBigInt(200L);
+ listWriter.endList();
+
+ listWriter.startList();
+ listWriter.writeFloat4(3.14f);
+ listWriter.endList();
+
+ listVector.setValueCount(2);
+
+ assertEquals(2, listVector.getValueCount());
+
+ List value0 = (List) listVector.getObject(0);
+ List value1 = (List) listVector.getObject(1);
+
+ assertEquals(2, value0.size());
+ assertEquals(100, value0.get(0));
+ assertEquals(200L, value0.get(1));
+
+ assertEquals(1, value1.size());
+ assertEquals(3.14f, value1.get(0));
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java
index bf1b9b0dfa..5bb5962704 100644
--- a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java
+++ b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java
@@ -20,20 +20,16 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import java.nio.ByteBuffer;
-import java.util.UUID;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.LargeVarBinaryVector;
import org.apache.arrow.vector.LargeVarCharVector;
-import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.complex.impl.LargeVarBinaryWriterImpl;
import org.apache.arrow.vector.complex.impl.LargeVarCharWriterImpl;
-import org.apache.arrow.vector.complex.impl.UuidWriterImpl;
import org.apache.arrow.vector.complex.impl.VarBinaryWriterImpl;
import org.apache.arrow.vector.complex.impl.VarCharWriterImpl;
-import org.apache.arrow.vector.holder.UuidHolder;
import org.apache.arrow.vector.util.Text;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
@@ -188,20 +184,4 @@ public void testWriteTextToLargeVarChar() throws Exception {
assertEquals(input, result);
}
}
-
- @Test
- public void testWriteToExtensionVector() throws Exception {
- try (UuidVector vector = new UuidVector("test", allocator);
- UuidWriterImpl writer = new UuidWriterImpl(vector)) {
- UUID uuid = UUID.randomUUID();
- ByteBuffer bb = ByteBuffer.allocate(16);
- bb.putLong(uuid.getMostSignificantBits());
- bb.putLong(uuid.getLeastSignificantBits());
- UuidHolder holder = new UuidHolder();
- holder.value = bb.array();
- writer.write(holder);
- UUID result = vector.getObject(0);
- assertEquals(uuid, result);
- }
- }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java b/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java
index d24708d66c..ae5ac0726c 100644
--- a/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java
+++ b/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java
@@ -16,6 +16,7 @@
*/
package org.apache.arrow.vector.types.pojo;
+import static org.apache.arrow.vector.TestUtils.ensureRegistered;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -43,10 +44,13 @@
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.ValueIterableVector;
+import org.apache.arrow.vector.ValueVector;
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.FieldWriter;
+import org.apache.arrow.vector.extension.UuidType;
import org.apache.arrow.vector.ipc.ArrowFileReader;
import org.apache.arrow.vector.ipc.ArrowFileWriter;
import org.apache.arrow.vector.types.FloatingPointPrecision;
@@ -59,9 +63,9 @@ public class TestExtensionType {
/** Test that a custom UUID type can be round-tripped through a temporary file. */
@Test
public void roundtripUuid() throws IOException {
- ExtensionTypeRegistry.register(new UuidType());
+ ensureRegistered(UuidType.INSTANCE);
final Schema schema =
- new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+ new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
UUID u1 = UUID.randomUUID();
@@ -89,7 +93,7 @@ public void roundtripUuid() throws IOException {
assertEquals(root.getSchema(), readerRoot.getSchema());
final Field field = readerRoot.getSchema().getFields().get(0);
- final UuidType expectedType = new UuidType();
+ final UuidType expectedType = UuidType.INSTANCE;
assertEquals(
field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
expectedType.extensionName());
@@ -113,9 +117,9 @@ public void roundtripUuid() throws IOException {
/** Test that a custom UUID type can be read as its underlying type. */
@Test
public void readUnderlyingType() throws IOException {
- ExtensionTypeRegistry.register(new UuidType());
+ ensureRegistered(UuidType.INSTANCE);
final Schema schema =
- new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+ new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
UUID u1 = UUID.randomUUID();
@@ -135,7 +139,7 @@ public void readUnderlyingType() throws IOException {
writer.end();
}
- ExtensionTypeRegistry.unregister(new UuidType());
+ ExtensionTypeRegistry.unregister(UuidType.INSTANCE);
try (final SeekableByteChannel channel =
Files.newByteChannel(Paths.get(file.getAbsolutePath()));
@@ -153,7 +157,7 @@ public void readUnderlyingType() throws IOException {
.getByteWidth());
final Field field = readerRoot.getSchema().getFields().get(0);
- final UuidType expectedType = new UuidType();
+ final UuidType expectedType = UuidType.INSTANCE;
assertEquals(
field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
expectedType.extensionName());
@@ -254,7 +258,7 @@ public void roundtripLocation() throws IOException {
@Test
public void testVectorCompare() {
- UuidType uuidType = new UuidType();
+ UuidType uuidType = UuidType.INSTANCE;
ExtensionTypeRegistry.register(uuidType);
try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
UuidVector a1 =
@@ -331,6 +335,11 @@ public String serialize() {
public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
return new LocationVector(name, allocator);
}
+
+ @Override
+ public FieldWriter getNewFieldWriter(ValueVector vector) {
+ throw new UnsupportedOperationException("Not yet implemented.");
+ }
}
public static class LocationVector extends ExtensionTypeVector
diff --git a/vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java b/vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java
deleted file mode 100644
index 5e2bd8881b..0000000000
--- a/vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * 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.types.pojo;
-
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.vector.FieldVector;
-import org.apache.arrow.vector.FixedSizeBinaryVector;
-import org.apache.arrow.vector.UuidVector;
-import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
-
-public class UuidType extends ExtensionType {
-
- @Override
- public ArrowType storageType() {
- return new ArrowType.FixedSizeBinary(16);
- }
-
- @Override
- public String extensionName() {
- return "uuid";
- }
-
- @Override
- public boolean extensionEquals(ExtensionType other) {
- return other instanceof UuidType;
- }
-
- @Override
- public ArrowType deserialize(ArrowType storageType, String serializedData) {
- if (!storageType.equals(storageType())) {
- throw new UnsupportedOperationException(
- "Cannot construct UuidType from underlying type " + storageType);
- }
- return new UuidType();
- }
-
- @Override
- public String serialize() {
- return "";
- }
-
- @Override
- public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
- return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));
- }
-}
diff --git a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java
index 4ee9630a4d..df5521a1ad 100644
--- a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java
+++ b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java
@@ -47,6 +47,7 @@
import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.LargeListVector;
import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.RunEndEncodedVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.UnionVector;
import org.apache.arrow.vector.holders.NullableBigIntHolder;
@@ -1025,6 +1026,72 @@ public void testAppendDenseUnionVectorMismatch() {
}
}
+ @Test
+ public void testAppendRunEndEncodedVector() {
+ final FieldType reeFieldType = FieldType.notNullable(ArrowType.RunEndEncoded.INSTANCE);
+ final Field runEndsField =
+ new Field("runEnds", FieldType.notNullable(Types.MinorType.INT.getType()), null);
+ final Field valuesField = Field.nullable("values", Types.MinorType.INT.getType());
+ final List children = Arrays.asList(runEndsField, valuesField);
+
+ final Field targetField = new Field("target", reeFieldType, children);
+ final Field deltaField = new Field("delta", reeFieldType, children);
+ try (RunEndEncodedVector target = new RunEndEncodedVector(targetField, allocator, null);
+ RunEndEncodedVector delta = new RunEndEncodedVector(deltaField, allocator, null)) {
+
+ // populate target
+ target.allocateNew();
+ // data: [1, 1, 2, null, 3, 3, 3] (7 values)
+ // values: [1, 2, null, 3]
+ // runEnds: [2, 3, 4, 7]
+ ValueVectorDataPopulator.setVector((IntVector) target.getValuesVector(), 1, 2, null, 3);
+ ValueVectorDataPopulator.setVector((IntVector) target.getRunEndsVector(), 2, 3, 4, 7);
+ target.setValueCount(7);
+
+ // populate delta
+ delta.allocateNew();
+ // data: [3, 4, 4, 5, null, null] (6 values)
+ // values: [3, 4, 5, null]
+ // runEnds: [1, 3, 4, 6]
+ ValueVectorDataPopulator.setVector((IntVector) delta.getValuesVector(), 3, 4, 5, null);
+ ValueVectorDataPopulator.setVector((IntVector) delta.getRunEndsVector(), 1, 3, 4, 6);
+ delta.setValueCount(6);
+
+ VectorAppender appender = new VectorAppender(target);
+ delta.accept(appender, null);
+
+ assertEquals(13, target.getValueCount());
+
+ final Field expectedField = new Field("expected", reeFieldType, children);
+ try (RunEndEncodedVector expected = new RunEndEncodedVector(expectedField, allocator, null)) {
+ expected.allocateNew();
+ // expected data: [1, 1, 2, null, 3, 3, 3, 3, 4, 4, 5, null, null] (13 values)
+ // expected values: [1, 2, null, 3, 3, 4, 5, null]
+ // expected runEnds: [2, 3, 4, 7, 8, 10, 11, 13]
+ ValueVectorDataPopulator.setVector(
+ (IntVector) expected.getValuesVector(), 1, 2, null, 3, 3, 4, 5, null);
+ ValueVectorDataPopulator.setVector(
+ (IntVector) expected.getRunEndsVector(), 2, 3, 4, 7, 8, 10, 11, 13);
+ expected.setValueCount(13);
+
+ assertVectorsEqual(expected, target);
+ }
+
+ // Check that delta is unchanged.
+ final Field expectedDeltaField = new Field("expectedDelta", reeFieldType, children);
+ try (RunEndEncodedVector expectedDelta =
+ new RunEndEncodedVector(expectedDeltaField, allocator, null)) {
+ expectedDelta.allocateNew();
+ ValueVectorDataPopulator.setVector(
+ (IntVector) expectedDelta.getValuesVector(), 3, 4, 5, null);
+ ValueVectorDataPopulator.setVector(
+ (IntVector) expectedDelta.getRunEndsVector(), 1, 3, 4, 6);
+ expectedDelta.setValueCount(6);
+ assertVectorsEqual(expectedDelta, delta);
+ }
+ }
+ }
+
@Test
public void testAppendVectorNegative() {
final int vectorLength = 10;