diff --git a/vector/src/main/codegen/templates/AbstractFieldReader.java b/vector/src/main/codegen/templates/AbstractFieldReader.java index c7c5b4d78..556fb576c 100644 --- a/vector/src/main/codegen/templates/AbstractFieldReader.java +++ b/vector/src/main/codegen/templates/AbstractFieldReader.java @@ -109,10 +109,6 @@ public void copyAsField(String name, ${name}Writer writer) { - public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory) { - fail("CopyAsValue StructWriter"); - } - public void read(ExtensionHolder holder) { fail("Extension"); } @@ -147,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 ae5b97fae..4b4a17d93 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/ArrowType.java b/vector/src/main/codegen/templates/ArrowType.java index fd35c1cd2..b428f0915 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 4c6f49ab9..c52345af2 100644 --- a/vector/src/main/codegen/templates/BaseReader.java +++ b/vector/src/main/codegen/templates/BaseReader.java @@ -49,7 +49,6 @@ public interface RepeatedStructReader extends StructReader{ boolean next(); int size(); void copyAsValue(StructWriter writer); - void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory); } public interface ListReader extends BaseReader{ @@ -60,7 +59,6 @@ public interface RepeatedListReader extends ListReader{ boolean next(); int size(); void copyAsValue(ListWriter writer); - void copyAsValue(ListWriter writer, ExtensionTypeWriterFactory writerFactory); } public interface MapReader extends BaseReader{ @@ -71,7 +69,6 @@ public interface RepeatedMapReader extends MapReader{ boolean next(); int size(); void copyAsValue(MapWriter writer); - void copyAsValue(MapWriter writer, ExtensionTypeWriterFactory writerFactory); } public interface ScalarReader extends diff --git a/vector/src/main/codegen/templates/BaseWriter.java b/vector/src/main/codegen/templates/BaseWriter.java index 78da7fddc..a4c98d708 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 4df5478f4..6655f6c2a 100644 --- a/vector/src/main/codegen/templates/ComplexCopier.java +++ b/vector/src/main/codegen/templates/ComplexCopier.java @@ -41,15 +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, null); - } - - public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory) { - writeValue(input, output, extensionTypeWriterFactory); - } + public static void copy(FieldReader reader, FieldWriter writer) { - private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) { final MinorType mt = reader.getMinorType(); switch (mt) { @@ -65,7 +58,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension FieldReader childReader = reader.reader(); FieldWriter childWriter = getListWriterForReader(childReader, writer); if (childReader.isSet()) { - writeValue(childReader, childWriter, extensionTypeWriterFactory); + copy(childReader, childWriter); } else { childWriter.writeNull(); } @@ -83,8 +76,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension FieldReader structReader = reader.reader(); if (structReader.isSet()) { writer.startEntry(); - writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()), extensionTypeWriterFactory); - writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()), extensionTypeWriterFactory); + copy(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key())); + copy(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value())); writer.endEntry(); } else { writer.writeNull(); @@ -103,7 +96,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension if (childReader.getMinorType() != Types.MinorType.NULL) { FieldWriter childWriter = getStructWriterForReader(childReader, writer, name); if (childReader.isSet()) { - writeValue(childReader, childWriter, extensionTypeWriterFactory); + copy(childReader, childWriter); } else { childWriter.writeNull(); } @@ -115,14 +108,10 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension } break; case EXTENSIONTYPE: - if (extensionTypeWriterFactory == null) { - throw new IllegalArgumentException("Must provide ExtensionTypeWriterFactory"); - } if (reader.isSet()) { Object value = reader.readObject(); if (value != null) { - writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory); - writer.writeExtension(value); + writer.writeExtension(value, reader.getField().getType()); } } else { writer.writeNull(); diff --git a/vector/src/main/codegen/templates/NullReader.java b/vector/src/main/codegen/templates/NullReader.java index 052963347..88e6ea98e 100644 --- a/vector/src/main/codegen/templates/NullReader.java +++ b/vector/src/main/codegen/templates/NullReader.java @@ -86,7 +86,6 @@ public void read(int arrayIndex, Nullable${name}Holder holder){ } - public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory){} public void read(ExtensionHolder holder) { holder.isSet = 0; } diff --git a/vector/src/main/codegen/templates/PromotableWriter.java b/vector/src/main/codegen/templates/PromotableWriter.java index d22eb00b2..11d34f72c 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,17 +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 addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory, ArrowType arrowType) { - getWriter(MinorType.EXTENSIONTYPE, arrowType).addExtensionTypeWriterFactory(factory); + public void write(ExtensionHolder holder) { + getWriter(MinorType.EXTENSIONTYPE, holder.type()).write(holder); } @Override diff --git a/vector/src/main/codegen/templates/UnionListWriter.java b/vector/src/main/codegen/templates/UnionListWriter.java index 3c41ac72b..4b5473923 100644 --- a/vector/src/main/codegen/templates/UnionListWriter.java +++ b/vector/src/main/codegen/templates/UnionListWriter.java @@ -204,13 +204,13 @@ public MapWriter map(String name, boolean keysSorted) { @Override public ExtensionWriter extension(ArrowType arrowType) { - this.extensionType = arrowType; + 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"> @@ -337,13 +337,13 @@ 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, extensionType); + public void writeExtension(Object value, ArrowType type) { + writeExtension(value); } public void write(ExtensionHolder var1) { diff --git a/vector/src/main/codegen/templates/UnionReader.java b/vector/src/main/codegen/templates/UnionReader.java index 96ad3e1b9..0edae7ade 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 67efdf60f..c70659196 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 272edab17..0db699fd8 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/BaseValueVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java index cc57cde29..37dfa2061 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java @@ -22,7 +22,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.util.Preconditions; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.util.DataSizeRoundingUtil; import org.apache.arrow.vector.util.TransferPair; @@ -261,18 +260,6 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } - @Override - public void copyFrom( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException(); - } - - @Override - public void copyFromSafe( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - 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 diff --git a/vector/src/main/java/org/apache/arrow/vector/NullVector.java b/vector/src/main/java/org/apache/arrow/vector/NullVector.java index 0d6dab283..6bfe540d2 100644 --- a/vector/src/main/java/org/apache/arrow/vector/NullVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/NullVector.java @@ -27,7 +27,6 @@ import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.impl.NullReader; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; @@ -330,18 +329,6 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } - @Override - public void copyFrom( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException(); - } - - @Override - public void copyFromSafe( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException(); - } - @Override public String getName() { return this.getField().getName(); diff --git a/vector/src/main/java/org/apache/arrow/vector/ValueVector.java b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java index e0628c2ee..3a5058256 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ValueVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java @@ -22,7 +22,6 @@ import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.compare.VectorVisitor; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; @@ -310,30 +309,6 @@ public interface ValueVector extends Closeable, Iterable { */ void copyFromSafe(int fromIndex, int thisIndex, ValueVector from); - /** - * Copy a cell value from a particular index in source vector to a particular position in this - * vector. - * - * @param fromIndex position to copy from in source vector - * @param thisIndex position to copy to in this vector - * @param from source vector - * @param writerFactory the extension type writer factory to use for copying extension type values - */ - void copyFrom( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory); - - /** - * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the - * capacity of the vector needs to be expanded before copy. - * - * @param fromIndex position to copy from in source vector - * @param thisIndex position to copy to in this vector - * @param from source vector - * @param writerFactory the extension type writer factory to use for copying extension type values - */ - void copyFromSafe( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory); - /** * Accept a generic {@link VectorVisitor} and return the result. * diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/AbstractContainerVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/AbstractContainerVector.java index 429f9884b..a6a71cf1a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/AbstractContainerVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/AbstractContainerVector.java @@ -21,7 +21,6 @@ import org.apache.arrow.vector.DensityAwareVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList; @@ -152,18 +151,6 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } - @Override - public void copyFrom( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException(); - } - - @Override - public void copyFromSafe( - int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException(); - } - @Override public String getName() { return name; 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 48c8127e2..997b5a8b7 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 @@ -49,7 +49,6 @@ import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.impl.UnionLargeListReader; import org.apache.arrow.vector.complex.impl.UnionLargeListWriter; import org.apache.arrow.vector.complex.reader.FieldReader; @@ -483,42 +482,12 @@ public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { */ @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { - copyFrom(inIndex, outIndex, from, null); - } - - /** - * Copy a cell value from a particular index in source vector to a particular position in this - * vector. - * - * @param inIndex position to copy from in source vector - * @param outIndex position to copy to in this vector - * @param from source vector - * @param writerFactory the extension type writer factory to use for copying extension type values - */ - @Override - public void copyFrom( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); FieldReader in = from.getReader(); in.setPosition(inIndex); UnionLargeListWriter out = getWriter(); out.setPosition(outIndex); - ComplexCopier.copy(in, out, writerFactory); - } - - /** - * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the - * capacity of the vector needs to be expanded before copy. - * - * @param inIndex position to copy from in source vector - * @param outIndex position to copy to in this vector - * @param from source vector - * @param writerFactory the extension type writer factory to use for copying extension type values - */ - @Override - public void copyFromSafe( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - copyFrom(inIndex, outIndex, from, writerFactory); + ComplexCopier.copy(in, out); } /** 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 992a66444..2da7eb057 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 @@ -41,7 +41,6 @@ import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.impl.UnionLargeListViewReader; import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListReader; @@ -347,20 +346,6 @@ public void copyFrom(int inIndex, int outIndex, ValueVector from) { "LargeListViewVector does not support copyFrom operation yet."); } - @Override - public void copyFromSafe( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException( - "LargeListViewVector does not support copyFromSafe operation yet."); - } - - @Override - public void copyFrom( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - throw new UnsupportedOperationException( - "LargeListViewVector does not support copyFrom operation yet."); - } - @Override public FieldVector getDataVector() { return vector; 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 89549257c..93a313ef4 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 @@ -42,7 +42,6 @@ import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.complex.reader.FieldReader; @@ -401,42 +400,12 @@ public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { */ @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { - copyFrom(inIndex, outIndex, from, null); - } - - /** - * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the - * capacity of the vector needs to be expanded before copy. - * - * @param inIndex position to copy from in source vector - * @param outIndex position to copy to in this vector - * @param from source vector - * @param writerFactory the extension type writer factory to use for copying extension type values - */ - @Override - public void copyFromSafe( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - copyFrom(inIndex, outIndex, from, writerFactory); - } - - /** - * Copy a cell value from a particular index in source vector to a particular position in this - * vector. - * - * @param inIndex position to copy from in source vector - * @param outIndex position to copy to in this vector - * @param from source vector - * @param writerFactory the extension type writer factory to use for copying extension type values - */ - @Override - public void copyFrom( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); FieldReader in = from.getReader(); in.setPosition(inIndex); FieldWriter out = getWriter(); out.setPosition(outIndex); - ComplexCopier.copy(in, out, writerFactory); + ComplexCopier.copy(in, out); } /** 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 278424042..8711db5e0 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 @@ -42,7 +42,6 @@ import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.impl.ComplexCopier; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.impl.UnionListViewReader; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.complex.reader.FieldReader; @@ -339,12 +338,6 @@ public void copyFromSafe(int inIndex, int outIndex, ValueVector from) { copyFrom(inIndex, outIndex, from); } - @Override - public void copyFromSafe( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { - copyFrom(inIndex, outIndex, from, writerFactory); - } - @Override public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); @@ -352,18 +345,12 @@ public OUT accept(VectorVisitor visitor, IN value) { @Override public void copyFrom(int inIndex, int outIndex, ValueVector from) { - copyFrom(inIndex, outIndex, from, null); - } - - @Override - public void copyFrom( - int inIndex, int outIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) { Preconditions.checkArgument(this.getMinorType() == from.getMinorType()); FieldReader in = from.getReader(); in.setPosition(inIndex); FieldWriter out = getWriter(); out.setPosition(outIndex); - ComplexCopier.copy(in, out, writerFactory); + ComplexCopier.copy(in, out); } @Override diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractBaseReader.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractBaseReader.java index bf074ecb9..b2e95663f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractBaseReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractBaseReader.java @@ -115,14 +115,4 @@ public void copyAsValue(ListWriter writer) { public void copyAsValue(MapWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } - - @Override - public void copyAsValue(ListWriter writer, ExtensionTypeWriterFactory writerFactory) { - ComplexCopier.copy(this, (FieldWriter) writer, writerFactory); - } - - @Override - public void copyAsValue(MapWriter writer, ExtensionTypeWriterFactory writerFactory) { - ComplexCopier.copy(this, (FieldWriter) writer, writerFactory); - } } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java index fccff6c21..bc98ab262 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java @@ -17,6 +17,8 @@ package org.apache.arrow.vector.complex.impl; import org.apache.arrow.vector.ExtensionTypeVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.types.pojo.Field; /** 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 a01d59155..000000000 --- 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 AbstractExtensionTypeWriter}. 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/UnionExtensionWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java index 4219069cb..93796aa77 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); } @@ -79,6 +74,7 @@ public void setPosition(int index) { @Override public void writeNull() { - this.writer.writeNull(); + this.vector.setNull(getPosition()); + this.vector.setValueCount(getPosition() + 1); } } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListReader.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListReader.java index a9104cb0d..be236c316 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListReader.java @@ -105,8 +105,4 @@ public boolean next() { public void copyAsValue(UnionLargeListWriter writer) { ComplexCopier.copy(this, (FieldWriter) writer); } - - public void copyAsValue(UnionLargeListWriter writer, ExtensionTypeWriterFactory writerFactory) { - ComplexCopier.copy(this, (FieldWriter) writer, writerFactory); - } } 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 index bb35b960d..0550b015a 100644 --- 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 @@ -61,7 +61,6 @@ public boolean isSet() { return !vector.isNull(idx()); } - @Override public void read(ExtensionHolder holder) { if (holder instanceof UuidHolder) { vector.get(idx(), (UuidHolder) holder); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java deleted file mode 100644 index 35988129c..000000000 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java +++ /dev/null @@ -1,45 +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.UuidVector; - -/** - * Factory for creating {@link UuidWriterImpl} instances. - * - *

This factory is used to create writers for UUID extension type vectors. - * - * @see UuidWriterImpl - * @see org.apache.arrow.vector.extension.UuidType - */ -public class UuidWriterFactory implements ExtensionTypeWriterFactory { - - /** - * Creates a writer implementation for the given extension type vector. - * - * @param extensionTypeVector the vector to create a writer for - * @return a {@link UuidWriterImpl} if the vector is a {@link UuidVector}, null otherwise - */ - @Override - public AbstractFieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) { - if (extensionTypeVector instanceof UuidVector) { - return new UuidWriterImpl((UuidVector) extensionTypeVector); - } - return null; - } -} 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 index 8a78add11..2199c0dca 100644 --- 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 @@ -18,9 +18,12 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.complex.writer.FieldWriter; 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}. @@ -56,6 +59,11 @@ public void writeExtension(Object value) { vector.setValueCount(getPosition() + 1); } + @Override + public void writeExtension(Object value, ArrowType type) { + writeExtension(value); + } + @Override public void write(ExtensionHolder holder) { if (holder instanceof UuidHolder) { 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 ca56214fd..780a4ee65 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 index f0f2636c8..7a7af298e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java +++ b/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java @@ -20,6 +20,9 @@ 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; @@ -106,4 +109,9 @@ public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocato 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 fc7ed8587..4d3f767ae 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 index e5398d82c..7fa50ca76 100644 --- a/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java +++ b/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java @@ -17,6 +17,8 @@ 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. @@ -32,4 +34,9 @@ public class NullableUuidHolder extends ExtensionHolder { /** Buffer containing 16-byte UUID data. */ public ArrowBuf buffer; + + @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 index 484e05c24..8a0a66e43 100644 --- a/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java +++ b/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java @@ -17,6 +17,8 @@ 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. @@ -35,4 +37,9 @@ public class UuidHolder extends ExtensionHolder { public UuidHolder() { this.isSet = 1; } + + @Override + public ArrowType type() { + return UuidType.INSTANCE; + } } 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 d5cbf925b..759c84651 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java @@ -26,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.UuidHolder; 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; @@ -1021,6 +1027,79 @@ 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(); + UuidHolder holder = new UuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0); + 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, 0); + assertEquals(u3, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0); + assertEquals(u4, actualUuid); + reader.next(); + uuidReader = reader.reader(); + assertFalse(uuidReader.isSet(), "third element should be null"); + } + } + private void writeIntValues(UnionLargeListWriter writer, int[] values) { writer.startList(); for (int v : values) { 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 41a95a8d1..51f8dade6 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java @@ -35,7 +35,6 @@ 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.impl.UuidWriterFactory; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; import org.apache.arrow.vector.extension.UuidType; @@ -1216,8 +1215,7 @@ public void testListVectorWithExtensionType() throws Exception { UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); writer.startList(); - ExtensionWriter extensionWriter = writer.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); + ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE); extensionWriter.writeExtension(u1); extensionWriter.writeExtension(u2); writer.endList(); @@ -1244,8 +1242,7 @@ public void testListVectorReaderForExtensionType() throws Exception { UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); writer.startList(); - ExtensionWriter extensionWriter = writer.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); + ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE); extensionWriter.writeExtension(u1); extensionWriter.writeExtension(u2); writer.endList(); @@ -1279,19 +1276,18 @@ public void testCopyFromForExtensionType() throws Exception { UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); writer.startList(); - ExtensionWriter extensionWriter = writer.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); - extensionWriter.writeExtension(u2); - extensionWriter.writeNull(); + + writer.extension(UuidType.INSTANCE).writeExtension(u1); + writer.writeExtension(u2); + writer.writeNull(); writer.endList(); - writer.setValueCount(1); + writer.setValueCount(3); // copy values from input to output outVector.allocateNew(); - outVector.copyFrom(0, 0, inVector, new UuidWriterFactory()); - outVector.setValueCount(1); + outVector.copyFrom(0, 0, inVector); + outVector.setValueCount(3); UnionListReader reader = outVector.getReader(); assertTrue(reader.isSet(), "shouldn't be null"); @@ -1310,6 +1306,79 @@ public void testCopyFromForExtensionType() throws Exception { } } + @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(); + UuidHolder holder = new UuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0); + 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, 0); + assertEquals(u3, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0); + assertEquals(u4, actualUuid); + reader.next(); + uuidReader = reader.reader(); + assertFalse(uuidReader.isSet(), "third element should be null"); + } + } + private void writeIntValues(UnionListWriter writer, int[] values) { writer.startList(); for (int v : values) { 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 df8f338f4..bfac1237a 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -35,7 +35,6 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.impl.UnionMapWriter; -import org.apache.arrow.vector.complex.impl.UuidWriterFactory; 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; @@ -1284,15 +1283,13 @@ public void testMapVectorWithExtensionType() throws Exception { writer.startMap(); writer.startEntry(); writer.key().bigInt().writeBigInt(0); - ExtensionWriter extensionWriter = writer.value().extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); + 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(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u2); + extensionWriter = writer.value().extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u2, UuidType.INSTANCE); writer.endEntry(); writer.endMap(); @@ -1326,21 +1323,18 @@ public void testCopyFromForExtensionType() throws Exception { writer.startMap(); writer.startEntry(); writer.key().bigInt().writeBigInt(0); - ExtensionWriter extensionWriter = writer.value().extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); + 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(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u2); + extensionWriter.writeExtension(u2, UuidType.INSTANCE); writer.endEntry(); writer.endMap(); writer.setValueCount(1); outVector.allocateNew(); - outVector.copyFrom(0, 0, inVector, new UuidWriterFactory()); + outVector.copyFrom(0, 0, inVector); outVector.setValueCount(1); UnionMapReader mapReader = outVector.getReader(); @@ -1576,4 +1570,103 @@ public void testFixedSizeBinaryFirstInitialization() { 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(); + UuidHolder keyHolder = new UuidHolder(); + keyReader.read(keyHolder); + UUID actualKey = UuidUtility.uuidFromArrowBuf(keyHolder.buffer, 0); + 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, 0); + 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/TestStructVector.java b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java index b8abfe1ef..c5e23c8e8 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java @@ -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()); } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java b/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java index 3d70238ec..a3690461c 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java @@ -33,6 +33,7 @@ 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; @@ -358,7 +359,13 @@ void testReaderReadWithUnsupportedHolder() throws Exception { reader.setPosition(0); // Create a mock unsupported holder - ExtensionHolder unsupportedHolder = new ExtensionHolder() {}; + ExtensionHolder unsupportedHolder = + new ExtensionHolder() { + @Override + public ArrowType type() { + return null; + } + }; IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> reader.read(unsupportedHolder)); @@ -377,7 +384,13 @@ void testReaderReadWithArrayIndexUnsupportedHolder() throws Exception { UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); // Create a mock unsupported holder - ExtensionHolder unsupportedHolder = new ExtensionHolder() {}; + ExtensionHolder unsupportedHolder = + new ExtensionHolder() { + @Override + public ArrowType type() { + return null; + } + }; IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> reader.read(0, unsupportedHolder)); 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 493a4b26a..b2a8cf9ba 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 @@ -860,8 +860,7 @@ public void testCopyListVectorWithExtensionType() { for (int i = 0; i < COUNT; i++) { listWriter.setPosition(i); listWriter.startList(); - ExtensionWriter extensionWriter = listWriter.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); + ExtensionWriter extensionWriter = listWriter.extension(UuidType.INSTANCE); extensionWriter.writeExtension(UUID.randomUUID()); extensionWriter.writeExtension(UUID.randomUUID()); listWriter.endList(); @@ -874,7 +873,7 @@ public void testCopyListVectorWithExtensionType() { for (int i = 0; i < COUNT; i++) { in.setPosition(i); out.setPosition(i); - ComplexCopier.copy(in, out, new UuidWriterFactory()); + ComplexCopier.copy(in, out); } to.setValueCount(COUNT); @@ -896,12 +895,10 @@ public void testCopyMapVectorWithExtensionType() { mapWriter.setPosition(i); mapWriter.startMap(); mapWriter.startEntry(); - ExtensionWriter extensionKeyWriter = mapWriter.key().extension(new UuidType()); - extensionKeyWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionKeyWriter.writeExtension(UUID.randomUUID()); - ExtensionWriter extensionValueWriter = mapWriter.value().extension(new UuidType()); - extensionValueWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionValueWriter.writeExtension(UUID.randomUUID()); + 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(); } @@ -914,7 +911,7 @@ public void testCopyMapVectorWithExtensionType() { for (int i = 0; i < COUNT; i++) { in.setPosition(i); out.setPosition(i); - ComplexCopier.copy(in, out, new UuidWriterFactory()); + ComplexCopier.copy(in, out); } to.setValueCount(COUNT); @@ -934,12 +931,10 @@ public void testCopyStructVectorWithExtensionType() { for (int i = 0; i < COUNT; i++) { structWriter.setPosition(i); structWriter.start(); - ExtensionWriter extensionWriter1 = structWriter.extension("timestamp1", new UuidType()); - extensionWriter1.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter1.writeExtension(UUID.randomUUID()); - ExtensionWriter extensionWriter2 = structWriter.extension("timestamp2", new UuidType()); - extensionWriter2.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter2.writeExtension(UUID.randomUUID()); + 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(); } @@ -951,7 +946,7 @@ public void testCopyStructVectorWithExtensionType() { for (int i = 0; i < COUNT; i++) { in.setPosition(i); out.setPosition(i); - ComplexCopier.copy(in, out, new UuidWriterFactory()); + ComplexCopier.copy(in, out); } to.setValueCount(COUNT); 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 a4594024f..5b6d65d6b 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; @@ -49,6 +50,7 @@ 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; @@ -57,6 +59,7 @@ import org.apache.arrow.vector.types.pojo.FieldType; 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); @@ -810,22 +829,22 @@ public void testExtensionType() throws Exception { public void testExtensionTypeForList() throws Exception { try (final ListVector container = ListVector.empty(EMPTY_SCHEMA_PATH, allocator); final UuidVector v = - (UuidVector) container.addOrGetVector(FieldType.nullable(new UuidType())).getVector(); + (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.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); - UuidVector uuidVector = (UuidVector) container.getDataVector(); + 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 46c259bda..34425c342 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 @@ -66,7 +66,6 @@ import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.impl.UnionReader; import org.apache.arrow.vector.complex.impl.UnionWriter; -import org.apache.arrow.vector.complex.impl.UuidWriterFactory; import org.apache.arrow.vector.complex.reader.BaseReader.StructReader; import org.apache.arrow.vector.complex.reader.BigIntReader; import org.apache.arrow.vector.complex.reader.FieldReader; @@ -87,6 +86,7 @@ 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; @@ -1106,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) { @@ -1128,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); } @@ -1153,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, 0), uuid); } else { assertEquals((float) i, unionReader.readFloat(), 1e-12); } @@ -2510,10 +2527,9 @@ public void extensionWriterReader() throws Exception { StructWriter rootWriter = writer.rootAsStruct(); { - ExtensionWriter extensionWriter = rootWriter.extension("uuid1", new UuidType()); + ExtensionWriter extensionWriter = rootWriter.extension("uuid1", UuidType.INSTANCE); extensionWriter.setPosition(0); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u1, UuidType.INSTANCE); } // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); 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 2ac4045aa..ae5ac0726 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 @@ -44,10 +44,12 @@ 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; @@ -333,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