diff --git a/.github/workflows/jarbuild.yml b/.github/workflows/jarbuild.yml index 0e9908b151..760561eb60 100644 --- a/.github/workflows/jarbuild.yml +++ b/.github/workflows/jarbuild.yml @@ -16,7 +16,7 @@ # under the License. name: JarBuild -on: +on: workflow_dispatch: inputs: arrow_branch: @@ -165,10 +165,9 @@ jobs: fail-fast: false matrix: platform: - - { runs_on: macos-15-intel, arch: "x86_64"} - { runs_on: macos-15, arch: "aarch_64" } env: - MACOSX_DEPLOYMENT_TARGET: "14.0" + MACOSX_DEPLOYMENT_TARGET: "15.0" steps: - name: Download source archive uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8 @@ -216,8 +215,30 @@ jobs: ./bootstrap-vcpkg.sh echo "VCPKG_ROOT=${PWD}/arrow/vcpkg" >> ${GITHUB_ENV} echo "${PWD}/arrow/vcpkg" >> ${GITHUB_PATH} + - name: Clean up disk space + run: | + echo "=== Free disk space before cleanup ===" + df -h / + + echo "" + echo "=== Removing Xcode simulators ===" + sudo rm -rf /Library/Developer/CoreSimulator/Caches || : + echo "Removed /Library/Developer/CoreSimulator/Caches" + + echo "" + echo "=== Removing user simulator data ===" + rm -rf ~/Library/Developer/CoreSimulator || : + echo "Removed ~/Library/Developer/CoreSimulator" + + echo "" + echo "=== Free disk space after cleanup ===" + df -h / - name: Install dependencies run: | + echo "=== Free disk space at start of dependency installation ===" + df -h / + + echo "" # Ensure updating python@XXX with the "--overwrite" option. # If python@XXX is updated without "--overwrite", it causes # a conflict error. Because Python 3 installed not by @@ -264,6 +285,11 @@ jobs: # bundled Protobuf. brew uninstall protobuf + echo "" + echo "=== Free disk space before LLVM build ===" + df -h / + + echo "" # Use vcpkg to install LLVM. vcpkg install \ --clean-after-build \ @@ -271,7 +297,12 @@ jobs: --x-manifest-root=arrow/ci/vcpkg \ --overlay-ports=arrow/ci/vcpkg/overlay/llvm/ \ --x-feature=gandiva-llvm - + + echo "" + echo "=== Free disk space after LLVM build ===" + df -h / + + echo "" brew bundle --file=Brewfile - name: Prepare ccache run: | @@ -284,10 +315,18 @@ jobs: restore-keys: jni-macos-${{ matrix.platform.arch }}- - name: Build run: | + echo "=== Free disk space at start of build ===" + df -h / + + echo "" set -e # make brew Java available to CMake export JAVA_HOME=$(brew --prefix openjdk@11)/libexec/openjdk.jdk/Contents/Home ci/scripts/jni_macos_build.sh . arrow build jni + + echo "" + echo "=== Free disk space at end of build ===" + df -h / - name: Compress into single artifact to keep directory structure run: tar -cvzf jni-macos-${{ matrix.platform.arch }}.tar.gz jni/ - name: Upload artifacts @@ -312,7 +351,6 @@ jobs: tar -xf apache-arrow-java-*.tar.gz --strip-components=1 tar -xvzf jni-linux-x86_64.tar.gz tar -xvzf jni-linux-aarch_64.tar.gz - tar -xvzf jni-macos-x86_64.tar.gz tar -xvzf jni-macos-aarch_64.tar.gz - name: Test that shared libraries exist run: | @@ -328,11 +366,6 @@ jobs: test -f jni/arrow_orc_jni/aarch_64/libarrow_orc_jni.so test -f jni/gandiva_jni/aarch_64/libgandiva_jni.so - test -f jni/arrow_cdata_jni/x86_64/libarrow_cdata_jni.dylib - test -f jni/arrow_dataset_jni/x86_64/libarrow_dataset_jni.dylib - test -f jni/arrow_orc_jni/x86_64/libarrow_orc_jni.dylib - test -f jni/gandiva_jni/x86_64/libgandiva_jni.dylib - test -f jni/arrow_cdata_jni/aarch_64/libarrow_cdata_jni.dylib test -f jni/arrow_dataset_jni/aarch_64/libarrow_dataset_jni.dylib test -f jni/arrow_orc_jni/aarch_64/libarrow_orc_jni.dylib @@ -478,4 +511,4 @@ jobs: $artifact done env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} \ No newline at end of file diff --git a/.github/workflows/rc.yml b/.github/workflows/rc.yml index 7e3cf5f6f2..9319a18ecf 100644 --- a/.github/workflows/rc.yml +++ b/.github/workflows/rc.yml @@ -380,7 +380,6 @@ jobs: tar -xf apache-arrow-java-*.tar.gz --strip-components=1 tar -xvzf jni-linux-x86_64.tar.gz tar -xvzf jni-linux-aarch_64.tar.gz - tar -xvzf jni-macos-x86_64.tar.gz tar -xvzf jni-macos-aarch_64.tar.gz tar -xvzf jni-windows-x86_64.tar.gz - name: Test that shared libraries exist diff --git a/vector/src/main/codegen/templates/AbstractFieldReader.java b/vector/src/main/codegen/templates/AbstractFieldReader.java index c7c5b4d78d..556fb576ce 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 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/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 4c6f49ab9b..c52345af21 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 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 4df5478f48..6655f6c2a7 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 0529633478..88e6ea98ea 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 d22eb00b2c..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,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 94723e6c9d..80383254f0 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 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..2374f6de13 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java @@ -72,6 +72,7 @@ public BaseFixedWidthVector(Field field, final BufferAllocator allocator, final refreshValueCapacity(); } + @Override public int getTypeWidth() { return typeWidth; } 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 6abf030dd2..9befcb890f 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; @@ -249,16 +248,4 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { 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(); - } } 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/NullVector.java b/vector/src/main/java/org/apache/arrow/vector/NullVector.java index 0d6dab2837..6bfe540d23 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 e0628c2ee1..3a5058256c 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ValueVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java @@ -22,7 +22,6 @@ import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.compare.VectorVisitor; -import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; @@ -310,30 +309,6 @@ public interface ValueVector extends Closeable, Iterable { */ 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 429f9884bb..a6a71cf1a4 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 3a900d0159..ed075352c9 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 @@ -48,7 +48,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; @@ -484,42 +483,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 6bfdea3a0f..84c6f03edb 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; @@ -348,20 +347,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 1e82fa22f2..3daeb6d77b 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; @@ -402,42 +401,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 fb703a6fb3..9b4e6b4c0c 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; @@ -340,12 +339,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); @@ -353,18 +346,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 bf074ecb90..b2e95663f7 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/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/UnionExtensionWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java index 4219069cba..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); } @@ -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 a9104cb0d2..be236c3166 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/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/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/test/java/org/apache/arrow/vector/TestLargeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java index 101d942d2a..c811114365 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java @@ -22,20 +22,26 @@ import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.nio.ByteBuffer; 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.holder.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.types.pojo.UuidType; import org.apache.arrow.vector.util.TransferPair; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -1020,6 +1026,83 @@ 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(new UuidType()); + extensionWriter.writeExtension(u1, new UuidType()); + extensionWriter.writeExtension(u2, new UuidType()); + writer.endList(); + + // Create second list with UUIDs + writer.setPosition(1); + UUID u3 = UUID.randomUUID(); + UUID u4 = UUID.randomUUID(); + writer.startList(); + extensionWriter = writer.extension(new UuidType()); + 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); + ByteBuffer bb = ByteBuffer.wrap(holder.value); + UUID actualUuid = new UUID(bb.getLong(), bb.getLong()); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + bb = ByteBuffer.wrap(holder.value); + actualUuid = new UUID(bb.getLong(), bb.getLong()); + 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); + bb = ByteBuffer.wrap(holder.value); + actualUuid = new UUID(bb.getLong(), bb.getLong()); + assertEquals(u3, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + bb = ByteBuffer.wrap(holder.value); + actualUuid = new UUID(bb.getLong(), bb.getLong()); + 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 ace36334ee..5817564a78 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.holder.UuidHolder; @@ -1216,7 +1215,6 @@ public void testListVectorWithExtensionType() throws Exception { UUID u2 = UUID.randomUUID(); writer.startList(); ExtensionWriter extensionWriter = writer.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); extensionWriter.writeExtension(u1); extensionWriter.writeExtension(u2); writer.endList(); @@ -1244,9 +1242,8 @@ public void testListVectorReaderForExtensionType() throws Exception { UUID u2 = UUID.randomUUID(); writer.startList(); ExtensionWriter extensionWriter = writer.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); - extensionWriter.writeExtension(u2); + extensionWriter.writeExtension(u1, new UuidType()); + extensionWriter.writeExtension(u2, new UuidType()); writer.endList(); writer.setValueCount(1); @@ -1280,23 +1277,80 @@ public void testCopyFromForExtensionType() throws Exception { UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); writer.startList(); + + writer.extension(new UuidType()).writeExtension(u1, new UuidType()); + writer.writeExtension(u2, new UuidType()); + 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(); + UuidHolder holder = new UuidHolder(); + uuidReader.read(holder); + ByteBuffer bb = ByteBuffer.wrap(holder.value); + UUID actualUuid = new UUID(bb.getLong(), bb.getLong()); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + bb = ByteBuffer.wrap(holder.value); + actualUuid = new UUID(bb.getLong(), bb.getLong()); + 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(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); extensionWriter.writeExtension(u1); extensionWriter.writeExtension(u2); - extensionWriter.writeNull(); writer.endList(); - writer.setValueCount(1); + // Create second list with UUIDs + writer.setPosition(1); + UUID u3 = UUID.randomUUID(); + UUID u4 = UUID.randomUUID(); + writer.startList(); + extensionWriter = writer.extension(new UuidType()); + extensionWriter.writeExtension(u3); + extensionWriter.writeExtension(u4); + extensionWriter.writeNull(); - // copy values from input to output + writer.endList(); + writer.setValueCount(2); + + // Use TransferPair with ExtensionTypeWriterFactory + // This tests the new makeTransferPair API with writerFactory parameter outVector.allocateNew(); - outVector.copyFrom(0, 0, inVector, new UuidWriterFactory()); - outVector.setValueCount(1); + TransferPair transferPair = inVector.makeTransferPair(outVector); + transferPair.copyValueSafe(0, 0); + transferPair.copyValueSafe(1, 1); + outVector.setValueCount(2); + // Verify first list UnionListReader reader = outVector.getReader(); - assertTrue(reader.isSet(), "shouldn't be null"); reader.setPosition(0); + assertTrue(reader.isSet(), "first list shouldn't be null"); reader.next(); FieldReader uuidReader = reader.reader(); UuidHolder holder = new UuidHolder(); @@ -1310,6 +1364,25 @@ public void testCopyFromForExtensionType() throws Exception { bb = ByteBuffer.wrap(holder.value); actualUuid = new UUID(bb.getLong(), bb.getLong()); 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); + bb = ByteBuffer.wrap(holder.value); + actualUuid = new UUID(bb.getLong(), bb.getLong()); + assertEquals(u3, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + bb = ByteBuffer.wrap(holder.value); + actualUuid = new UUID(bb.getLong(), bb.getLong()); + assertEquals(u4, actualUuid); + reader.next(); + uuidReader = reader.reader(); + assertFalse(uuidReader.isSet(), "third element should be null"); } } 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 1a1810d0f7..929470b7fa 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -34,7 +34,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; @@ -1282,14 +1281,11 @@ public void testMapVectorWithExtensionType() throws Exception { writer.startEntry(); writer.key().bigInt().writeBigInt(0); ExtensionWriter extensionWriter = writer.value().extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u1, new UuidType()); 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, new UuidType()); writer.endEntry(); writer.endMap(); @@ -1326,20 +1322,18 @@ public void testCopyFromForExtensionType() throws Exception { writer.startEntry(); writer.key().bigInt().writeBigInt(0); ExtensionWriter extensionWriter = writer.value().extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u1, new UuidType()); 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, new UuidType()); 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(); 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..2898e38abb 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(new UuidType()), 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/TestUtils.java b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java index 82295f8037..c28751aa58 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java @@ -20,6 +20,7 @@ import org.apache.arrow.memory.BufferAllocator; 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 +63,14 @@ 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); + } + } } 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 738e8905e3..ebc16e90c0 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 @@ -861,9 +861,8 @@ public void testCopyListVectorWithExtensionType() { listWriter.setPosition(i); listWriter.startList(); ExtensionWriter extensionWriter = listWriter.extension(new UuidType()); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(UUID.randomUUID()); - extensionWriter.writeExtension(UUID.randomUUID()); + extensionWriter.writeExtension(UUID.randomUUID(), new UuidType()); + extensionWriter.writeExtension(UUID.randomUUID(), new UuidType()); listWriter.endList(); } from.setValueCount(COUNT); @@ -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); @@ -897,11 +896,9 @@ public void testCopyMapVectorWithExtensionType() { mapWriter.startMap(); mapWriter.startEntry(); ExtensionWriter extensionKeyWriter = mapWriter.key().extension(new UuidType()); - extensionKeyWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionKeyWriter.writeExtension(UUID.randomUUID()); + extensionKeyWriter.writeExtension(UUID.randomUUID(), new UuidType()); ExtensionWriter extensionValueWriter = mapWriter.value().extension(new UuidType()); - extensionValueWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionValueWriter.writeExtension(UUID.randomUUID()); + extensionValueWriter.writeExtension(UUID.randomUUID(), new UuidType()); 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", new UuidType()); + extensionWriter1.writeExtension(UUID.randomUUID(), new UuidType()); + ExtensionWriter extensionWriter2 = structWriter.extension("uuid2", new UuidType()); + extensionWriter2.writeExtension(UUID.randomUUID(), new UuidType()); 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 7b8b1f9ef9..c9fd216c6d 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.holder.UuidHolder; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.holders.NullableDecimalHolder; @@ -100,7 +102,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 +131,22 @@ public void testPromoteToUnion() throws Exception { binHolder.buffer = buf; writer.fixedSizeBinary("A", 4).write(binHolder); + writer.setPosition(9); + UUID uuid = UUID.randomUUID(); + writer.extension("A", new UuidType()).writeExtension(uuid, new UuidType()); + writer.end(); + + writer.setPosition(10); + UUID uuid2 = UUID.randomUUID(); + UuidHolder uuidHolder = new UuidHolder(); + ByteBuffer bb = ByteBuffer.allocate(16); + bb.putLong(uuid2.getMostSignificantBits()); + bb.putLong(uuid2.getLeastSignificantBits()); + uuidHolder.value = bb.array(); + writer.extension("A", new UuidType()).write(uuidHolder); writer.end(); - 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(); @@ -791,12 +811,11 @@ public void testExtensionType() throws Exception { UUID u2 = UUID.randomUUID(); container.allocateNew(); container.setValueCount(1); - writer.addExtensionTypeWriterFactory(new UuidWriterFactory()); writer.setPosition(0); - writer.writeExtension(u1); + writer.writeExtension(u1, new UuidType()); writer.setPosition(1); - writer.writeExtension(u2); + writer.writeExtension(u2, new UuidType()); container.setValueCount(2); @@ -816,16 +835,15 @@ public void testExtensionTypeForList() throws Exception { UUID u2 = UUID.randomUUID(); container.allocateNew(); container.setValueCount(1); - writer.addExtensionTypeWriterFactory(new UuidWriterFactory()); writer.setPosition(0); - writer.writeExtension(u1); + writer.writeExtension(u1, new UuidType()); writer.setPosition(1); - writer.writeExtension(u2); + writer.writeExtension(u2, new UuidType()); 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/impl/UuidWriterFactory.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java deleted file mode 100644 index 1b1bf4e6e4..0000000000 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java +++ /dev/null @@ -1,31 +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; - -public class UuidWriterFactory implements ExtensionTypeWriterFactory { - - @Override - public AbstractFieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) { - if (extensionTypeVector instanceof UuidVector) { - return new UuidWriterImpl((UuidVector) extensionTypeVector); - } - return null; - } -} diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java index 68029b1df5..d0fb008a47 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java +++ b/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java @@ -21,6 +21,7 @@ import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.holder.UuidHolder; import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; public class UuidWriterImpl extends AbstractExtensionTypeWriter { @@ -38,6 +39,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) { UuidHolder uuidHolder = (UuidHolder) holder; 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 f374eb41e4..5e516c14b4 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; @@ -1105,6 +1104,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) { @@ -1127,6 +1133,10 @@ public void simpleUnion() throws Exception { holder.buffer = buf; unionWriter.write(holder); bufs.add(buf); + } else if (i % 5 == 4) { + UuidHolder holder = new UuidHolder(); + holder.value = uuidByte; + unionWriter.write(holder); } else { unionWriter.writeFloat4((float) i); } @@ -1152,6 +1162,12 @@ public void simpleUnion() throws Exception { unionReader.read(holder); assertEquals(i, holder.buffer.getInt(0)); assertEquals(4, holder.byteWidth); + } else if (i % 5 == 4) { + UuidHolder holder = new UuidHolder(); + unionReader.read(holder); + ByteBuffer uuidBb = ByteBuffer.wrap(holder.value); + UUID actualUuid = new UUID(uuidBb.getLong(), uuidBb.getLong()); + assertEquals(actualUuid, uuid); } else { assertEquals((float) i, unionReader.readFloat(), 1e-12); } @@ -2511,8 +2527,7 @@ public void extensionWriterReader() throws Exception { { ExtensionWriter extensionWriter = rootWriter.extension("uuid1", new UuidType()); extensionWriter.setPosition(0); - extensionWriter.addExtensionTypeWriterFactory(new UuidWriterFactory()); - extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u1, new UuidType()); } // read StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); @@ -2521,7 +2536,7 @@ public void extensionWriterReader() throws Exception { uuidReader.setPosition(0); UuidHolder uuidHolder = new UuidHolder(); uuidReader.read(uuidHolder); - final ByteBuffer bb = ByteBuffer.wrap(uuidHolder.value); + ByteBuffer bb = ByteBuffer.wrap(uuidHolder.value); UUID actualUuid = new UUID(bb.getLong(), bb.getLong()); assertEquals(u1, actualUuid); assertTrue(uuidReader.isSet()); 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 269cff0670..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,21 +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.UuidReaderImpl; -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; @@ -189,39 +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); - } - } - - @Test - public 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(); - UuidHolder holder = new UuidHolder(); - reader2.read(0, holder); - final ByteBuffer bb = ByteBuffer.wrap(holder.value); - UUID actualUuid = new UUID(bb.getLong(), bb.getLong()); - assertEquals(uuid, actualUuid); - } - } } diff --git a/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java b/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java index 207b0951a7..a206bc265e 100644 --- a/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java +++ b/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java @@ -17,7 +17,14 @@ package org.apache.arrow.vector.holder; import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.UuidType; public class UuidHolder extends ExtensionHolder { public byte[] value; + + @Override + public ArrowType type() { + return new UuidType(); + } } 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..ff8e45afc3 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,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.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.types.FloatingPointPrecision; @@ -59,7 +62,7 @@ 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(new UuidType()); final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new UuidType()))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); @@ -113,7 +116,7 @@ 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(new UuidType()); final Schema schema = new Schema(Collections.singletonList(Field.nullable("a", new UuidType()))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); @@ -331,6 +334,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 index 5e2bd8881b..eb2666eb25 100644 --- 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 @@ -20,10 +20,12 @@ 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.ExtensionType; public class UuidType extends ExtensionType { - @Override public ArrowType storageType() { return new ArrowType.FixedSizeBinary(16); @@ -57,4 +59,9 @@ public String serialize() { public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16)); } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + return new UuidWriterImpl((UuidVector) vector); + } }