diff --git a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java index e92e832991..d7b4ecd6a5 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -25,13 +25,18 @@ import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; +import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; +import org.apache.arrow.vector.types.pojo.Field; /** Importer for {@link ArrowArray}. */ final class ArrayImporter { @@ -39,6 +44,9 @@ final class ArrayImporter { private final BufferAllocator allocator; private final FieldVector vector; + // The field the vector was created from, or null to use the vector's own field. Unlike the + // field of an ExtensionTypeVector that does not override getField(), it has the extension type. + private final Field field; private final DictionaryProvider dictionaryProvider; private ReferenceCountedArrowArray underlyingAllocation; @@ -46,8 +54,17 @@ final class ArrayImporter { ArrayImporter( BufferAllocator allocator, FieldVector vector, DictionaryProvider dictionaryProvider) { + this(allocator, vector, /* field= */ null, dictionaryProvider); + } + + ArrayImporter( + BufferAllocator allocator, + FieldVector vector, + Field field, + DictionaryProvider dictionaryProvider) { this.allocator = Preconditions.checkNotNull(allocator); this.vector = Preconditions.checkNotNull(vector); + this.field = field; this.dictionaryProvider = dictionaryProvider; } @@ -71,8 +88,7 @@ void importArray(ArrowArray src) { } } - private void importChild(ArrayImporter parent, ArrowArray src) { - ArrowArray.Snapshot snapshot = src.snapshot(); + private void importChild(ArrayImporter parent, ArrowArray.Snapshot snapshot) { checkState(snapshot.release != NULL, "Cannot import released ArrowArray"); recursionLevel = parent.recursionLevel + 1; checkState( @@ -84,15 +100,78 @@ private void importChild(ArrayImporter parent, ArrowArray src) { } private void doImport(ArrowArray.Snapshot snapshot) { - // A non-zero offset is only meaningful for arrays that have buffers to offset into. - // Null arrays carry no buffers, so their offset is inert and safe to ignore. - checkState( - snapshot.offset == 0 - || snapshot.length == 0 - || vector.getField().getType().getTypeID() == ArrowType.ArrowTypeID.Null, - "ArrowArray struct has non-zero offset (%s), which is not supported", - snapshot.offset); + checkState(snapshot.offset >= 0, "ArrowArray struct has negative offset (%s)", snapshot.offset); + if (snapshot.offset != 0 && snapshot.length != 0) { + importSliced(snapshot); + } else { + importUnsliced(snapshot, /* skippedRows= */ 0, /* trimChildren= */ false); + } + } + + /** + * Java vectors cannot represent an offset, so the rows up to the end of the slice are imported + * into a temporary vector and the slice is then transferred into the target. Leading rows that + * fill whole bytes of a bitmap are skipped by advancing the buffer pointers, so that the + * temporary vector, and any validity bitmap synthesized for it, has fewer than 8 rows before the + * slice. + */ + private void importSliced(ArrowArray.Snapshot snapshot) { + // Run-end encoded arrays have no buffers, and their children are not indexed by row + long skippedRows = + storageType().getTypeID() == ArrowTypeID.RunEndEncoded ? 0 : snapshot.offset & ~7L; + int offset = checkedCastToInt(snapshot.offset - skippedRows); + int length = checkedCastToInt(snapshot.length); + + snapshot.offset = 0; + snapshot.length = Math.addExact(offset, length); + snapshot.null_count = unknownNullCount(snapshot); + + Field unslicedField = field != null ? field : vector.getField(); + try (FieldVector unslicedVector = unslicedField.createVector(allocator)) { + // When the target is an extension vector whose getField() reports its storage type, the + // temporary vector is a storage vector, so transfer between the storage vectors + FieldVector from = ExtensionTypeVector.getStorageVector(unslicedVector); + FieldVector to = ExtensionTypeVector.getStorageVector(vector); + checkTransferable(from, to); + ArrayImporter unslicedImporter = + new ArrayImporter(allocator, unslicedVector, unslicedField, dictionaryProvider); + unslicedImporter.recursionLevel = recursionLevel; + unslicedImporter.underlyingAllocation = underlyingAllocation; + unslicedImporter.importUnsliced(snapshot, skippedRows, /* trimChildren= */ true); + from.makeTransferPair(to).splitAndTransfer(offset, length); + } + } + /** + * Transfer pairs of nested vectors create or reuse the children of the target based on the + * classes of the source's children, so every vector of the target must be an instance of the + * class of its counterpart in the source. + */ + private static void checkTransferable(FieldVector from, FieldVector to) { + if (!from.getClass().isInstance(to)) { + String reason = + to instanceof ExtensionTypeVector && !(from instanceof ExtensionTypeVector) + ? "its getField() does not report its extension type" + : "its field creates a " + from.getClass().getSimpleName(); + throw new UnsupportedOperationException( + String.format( + "Cannot import an ArrowArray with a non-zero offset into %s %s, because %s", + to.getClass().getSimpleName(), to.getName(), reason)); + } + List fromChildren = from.getChildrenFromFields(); + List toChildren = to.getChildrenFromFields(); + for (int i = 0; i < Math.min(fromChildren.size(), toChildren.size()); i++) { + checkTransferable(fromChildren.get(i), toChildren.get(i)); + } + } + + /** + * Imports an array whose offset is 0, skipping the first {@code skippedRows} rows of its buffers, + * which must be a multiple of 8. With {@code trimChildren}, children whose rows line up with the + * rows of the array are imported only for the rows of the array, rather than in full. + */ + private void importUnsliced( + ArrowArray.Snapshot snapshot, long skippedRows, boolean trimChildren) { // First import children (required for reconstituting parent array data) long[] children = NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children)); @@ -103,11 +182,31 @@ private void doImport(ArrowArray.Snapshot snapshot) { "ArrowArray struct has %s children (expected %s)", children.length, childVectors.size()); + List childFields = field != null ? field.getChildren() : null; + long childRowsPerRow = trimChildren ? childRowsPerRow(storageType()) : 0; for (int i = 0; i < children.length; i++) { checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i); + ArrowArray.Snapshot child = ArrowArray.wrap(children[i]).snapshot(); + if (childRowsPerRow != 0) { + long firstRow = Math.multiplyExact(skippedRows, childRowsPerRow); + long rowCount = Math.multiplyExact(snapshot.length, childRowsPerRow); + checkState( + child.length - firstRow >= rowCount, + "ArrowArray struct child at position %s has length %s, but its parent needs %s rows", + i, + child.length, + firstRow + rowCount); + child.offset = Math.addExact(child.offset, firstRow); + child.length = rowCount; + child.null_count = unknownNullCount(child); + } + Field childField = + childFields != null && childFields.size() == children.length + ? childFields.get(i) + : null; ArrayImporter childImporter = - new ArrayImporter(allocator, childVectors.get(i), dictionaryProvider); - childImporter.importChild(this, ArrowArray.wrap(children[i])); + new ArrayImporter(allocator, childVectors.get(i), childField, dictionaryProvider); + childImporter.importChild(this, child); } } @@ -124,7 +223,7 @@ private void doImport(ArrowArray.Snapshot snapshot) { ArrayImporter dictionaryImporter = new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider); - dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary)); + dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary).snapshot()); } // Import main data @@ -133,7 +232,8 @@ private void doImport(ArrowArray.Snapshot snapshot) { NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers)); try (final BufferImportTypeVisitor visitor = - new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) { + new BufferImportTypeVisitor( + allocator, underlyingAllocation, fieldNode, bufferPointers, skippedRows)) { final List buffers; if (bufferPointers == null || bufferPointers.length == 0) { buffers = Collections.emptyList(); @@ -150,4 +250,42 @@ private void doImport(ArrowArray.Snapshot snapshot) { e); } } + + /** + * Returns the null count to import a different range of the rows of {@code snapshot} with than + * the range its null count covers. Without a validity bitmap it must be 0 for the buffers to + * load, otherwise it is unknown. + */ + private static long unknownNullCount(ArrowArray.Snapshot snapshot) { + boolean hasBitmap = + snapshot.n_buffers > 0 + && snapshot.buffers != NULL + && MemoryUtil.getLong(snapshot.buffers) != NULL; + return hasBitmap ? -1 : 0; + } + + /** The type of the vector, or the storage type of an extension type. */ + private ArrowType storageType() { + ArrowType type = vector.getField().getType(); + return type instanceof ArrowType.ExtensionType + ? ((ArrowType.ExtensionType) type).storageType() + : type; + } + + /** + * Returns how many rows of each child a row of an array of {@code type} spans, or 0 if its + * children are not indexed by its rows. + */ + private static long childRowsPerRow(ArrowType type) { + switch (type.getTypeID()) { + case Struct: + return 1; + case Union: + return ((ArrowType.Union) type).getMode() == UnionMode.Sparse ? 1 : 0; + case FixedSizeList: + return ((ArrowType.FixedSizeList) type).getListSize(); + default: + return 0; + } + } } diff --git a/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java b/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java index 10f690fc87..fae944c6fb 100644 --- a/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java +++ b/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java @@ -62,6 +62,7 @@ class BufferImportTypeVisitor implements ArrowType.ArrowTypeVisitor imported; BufferImportTypeVisitor( @@ -69,10 +70,26 @@ class BufferImportTypeVisitor implements ArrowType.ArrowTypeVisitor(); } @@ -83,6 +100,10 @@ public void close() throws Exception { @VisibleForTesting ArrowBuf importBuffer(ArrowType type, int index, long capacity) { + return importBuffer(type, index, /* skippedBytes= */ 0, capacity); + } + + private ArrowBuf importBuffer(ArrowType type, int index, long skippedBytes, long capacity) { checkState( buffers.length > index, "Expected at least %s buffers for type %s, but found %s", @@ -102,24 +123,28 @@ ArrowBuf importBuffer(ArrowType type, int index, long capacity) { } } - ArrowBuf buf = underlyingAllocation.unsafeAssociateAllocation(allocator, capacity, bufferPtr); + ArrowBuf buf = + underlyingAllocation.unsafeAssociateAllocation( + allocator, capacity, Math.addExact(bufferPtr, skippedBytes)); imported.add(buf); return buf; } + /** Imports a buffer that holds {@code bitsPerSlot} bits for each row. */ private ArrowBuf importFixedBits(ArrowType type, int index, long bitsPerSlot) { final long capacity = DataSizeRoundingUtil.divideBy8Ceil(bitsPerSlot * fieldNode.getLength()); - return importBuffer(type, index, capacity); + return importBuffer(type, index, Math.multiplyExact(skippedRows, bitsPerSlot) / 8, capacity); } + /** Imports a buffer that holds {@code bytesPerSlot} bytes for each row. */ private ArrowBuf importFixedBytes(ArrowType type, int index, long bytesPerSlot) { final long capacity = bytesPerSlot * fieldNode.getLength(); - return importBuffer(type, index, capacity); + return importBuffer(type, index, Math.multiplyExact(skippedRows, bytesPerSlot), capacity); } private ArrowBuf importOffsets(ArrowType type, long bytesPerSlot) { final long capacity = bytesPerSlot * (fieldNode.getLength() + 1); - return importBuffer(type, 1, capacity); + return importBuffer(type, 1, Math.multiplyExact(skippedRows, bytesPerSlot), capacity); } private ArrowBuf importData(ArrowType type, long capacity) { diff --git a/c/src/main/java/org/apache/arrow/c/Data.java b/c/src/main/java/org/apache/arrow/c/Data.java index d870b5884d..dfe4d00166 100644 --- a/c/src/main/java/org/apache/arrow/c/Data.java +++ b/c/src/main/java/org/apache/arrow/c/Data.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.c; +import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; @@ -35,8 +36,12 @@ /** * Functions for working with the C data interface. * - *

This API is EXPERIMENTAL. Note that currently only 64bit systems are supported. Importing - * {@link ArrowArray ArrowArrays} with a non-zero offset is not supported. + *

This API is EXPERIMENTAL. Note that currently only 64bit systems are supported. + * + *

Importing {@link ArrowArray ArrowArrays} with a non-zero offset is not always zero-copy: + * buffers that cannot be sliced at the offset, such as offset buffers or view data, are copied. + * Importing such an array into an existing vector fails if an extension vector nested in it does + * not report its extension type in {@link FieldVector#getField()}. */ public final class Data { @@ -358,7 +363,22 @@ public static void importIntoVector( FieldVector vector, DictionaryProvider provider, boolean closeImportedStructs) { - ArrayImporter importer = new ArrayImporter(allocator, vector, provider); + importIntoVector(allocator, array, vector, /* field= */ null, provider, closeImportedStructs); + } + + /** + * Import Java vector from the C data interface, given the field that the vector was created from, + * if any. The field reports the type of extension vectors whose getField() only reports their + * storage type, which an import with a non-zero offset needs to create a temporary vector. + */ + private static void importIntoVector( + BufferAllocator allocator, + ArrowArray array, + FieldVector vector, + Field field, + DictionaryProvider provider, + boolean closeImportedStructs) { + ArrayImporter importer = new ArrayImporter(allocator, vector, field, provider); importer.importArray(array); if (closeImportedStructs) { array.close(); @@ -407,7 +427,7 @@ public static FieldVector importVector( boolean closeImportedStructs) { Field field = importField(allocator, schema, provider, closeImportedStructs); FieldVector vector = field.createVector(allocator); - importIntoVector(allocator, array, vector, provider, closeImportedStructs); + importIntoVector(allocator, array, vector, field, provider, closeImportedStructs); return vector; } @@ -454,8 +474,10 @@ public static void importIntoVectorSchemaRoot( DictionaryProvider provider, boolean closeImportedStructs) { try (StructVector structVector = StructVector.emptyWithDuplicates("", allocator)) { - structVector.initializeChildrenFromFields(root.getSchema().getFields()); - importIntoVector(allocator, array, structVector, provider, closeImportedStructs); + List fields = root.getSchema().getFields(); + structVector.initializeChildrenFromFields(fields); + Field field = new Field("", structVector.getField().getFieldType(), fields); + importIntoVector(allocator, array, structVector, field, provider, closeImportedStructs); StructVectorUnloader unloader = new StructVectorUnloader(structVector); VectorLoader loader = new VectorLoader(root); try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) { diff --git a/c/src/test/java/org/apache/arrow/c/DictionaryTest.java b/c/src/test/java/org/apache/arrow/c/DictionaryTest.java index 8cd4913f22..9f02a3967d 100644 --- a/c/src/test/java/org/apache/arrow/c/DictionaryTest.java +++ b/c/src/test/java/org/apache/arrow/c/DictionaryTest.java @@ -28,7 +28,6 @@ import java.util.Collections; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; @@ -100,37 +99,113 @@ void roundtrip(FieldVector vector, DictionaryProvider provider, Class clazz) @Test public void testWithDictionary() throws Exception { - DictionaryProvider.MapDictionaryProvider provider = - new DictionaryProvider.MapDictionaryProvider(); - // create dictionary and provider + try (DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(); + IntVector encodedVector = createEncodedVector(provider, "bb", "bb", "cc", "aa")) { + roundtrip(encodedVector, provider, IntVector.class); + } + } + + @Test + public void testWithDictionaryAndNonZeroOffset() throws Exception { + try (DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(); + IntVector encodedVector = createEncodedVector(provider, "bb", "bb", "cc", "aa"); + ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { + Data.exportVector( + allocator, encodedVector, provider, consumerArrowArray, consumerArrowSchema); + + // Slicing a dictionary encoded array slices the indices, not the dictionary + ArrowArray.Snapshot snapshot = consumerArrowArray.snapshot(); + snapshot.offset = 1; + snapshot.length = 2; + consumerArrowArray.save(snapshot); + + try (CDataDictionaryProvider cDictionaryProvider = new CDataDictionaryProvider(); + FieldVector imported = + Data.importVector( + allocator, consumerArrowArray, consumerArrowSchema, cDictionaryProvider)) { + IntVector indices = (IntVector) imported; + assertEquals(2, indices.getValueCount()); + assertEquals(1, indices.get(0)); + assertEquals(2, indices.get(1)); + + Dictionary importedDictionary = cDictionaryProvider.lookup(0L); + assertEquals(3, importedDictionary.getVector().getValueCount()); + try (VarCharVector decoded = + (VarCharVector) DictionaryEncoder.decode(imported, importedDictionary)) { + assertEquals("bb", new String(decoded.get(0), StandardCharsets.UTF_8)); + assertEquals("cc", new String(decoded.get(1), StandardCharsets.UTF_8)); + } + } + } + } + + @Test + public void testWithDictionaryAndOffsetsPastTheFirstByte() throws Exception { + String[] values = new String[20]; + for (int i = 0; i < values.length; i++) { + values[i] = i % 5 == 4 ? null : new String[] {"aa", "bb", "cc"}[i % 3]; + } + try (DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(); + IntVector encodedVector = createEncodedVector(provider, values)) { + for (int offset : new int[] {8, 9, 16}) { + try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { + Data.exportVector( + allocator, encodedVector, provider, consumerArrowArray, consumerArrowSchema); + ArrowArray.Snapshot snapshot = consumerArrowArray.snapshot(); + snapshot.offset = offset; + snapshot.length = values.length - offset; + consumerArrowArray.save(snapshot); + + try (CDataDictionaryProvider cDictionaryProvider = new CDataDictionaryProvider(); + FieldVector imported = + Data.importVector( + allocator, consumerArrowArray, consumerArrowSchema, cDictionaryProvider); + VarCharVector decoded = + (VarCharVector) + DictionaryEncoder.decode(imported, cDictionaryProvider.lookup(0L))) { + assertEquals(values.length - offset, decoded.getValueCount()); + for (int i = 0; i < decoded.getValueCount(); i++) { + String expected = values[offset + i]; + String actual = + decoded.isNull(i) ? null : new String(decoded.get(i), StandardCharsets.UTF_8); + assertEquals(expected, actual, "offset " + offset + ", index " + i); + } + } + } + } + } + } + + /** Registers an aa/bb/cc dictionary with {@code provider} and encodes values with it. */ + private IntVector createEncodedVector( + DictionaryProvider.MapDictionaryProvider provider, String... values) { final VarCharVector dictVector = new VarCharVector("dict", allocator); dictVector.allocateNewSafe(); dictVector.setSafe(0, "aa".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8)); dictVector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8)); dictVector.setValueCount(3); - Dictionary dictionary = new Dictionary(dictVector, new DictionaryEncoding(0L, false, /* indexType= */ null)); provider.put(dictionary); - // create vector and encode it - final VarCharVector vector = new VarCharVector("vector", allocator); - vector.allocateNewSafe(); - vector.setSafe(0, "bb".getBytes(StandardCharsets.UTF_8)); - vector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8)); - vector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8)); - vector.setSafe(3, "aa".getBytes(StandardCharsets.UTF_8)); - vector.setValueCount(4); - - // get the encoded vector - IntVector encodedVector = (IntVector) DictionaryEncoder.encode(vector, dictionary); - - // Perform roundtrip using C Data Interface - roundtrip(encodedVector, provider, IntVector.class); - - // Close all - AutoCloseables.close((AutoCloseable) vector, encodedVector, dictVector); + try (VarCharVector vector = new VarCharVector("vector", allocator)) { + vector.allocateNewSafe(); + for (int i = 0; i < values.length; i++) { + if (values[i] == null) { + vector.setNull(i); + } else { + vector.setSafe(i, values[i].getBytes(StandardCharsets.UTF_8)); + } + } + vector.setValueCount(values.length); + return (IntVector) DictionaryEncoder.encode(vector, dictionary); + } } @Test diff --git a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java index ed1690a9e7..fdd3193c61 100644 --- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java +++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java @@ -29,12 +29,15 @@ import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.function.IntFunction; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BigIntVector; @@ -83,6 +86,7 @@ import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.ZeroVector; import org.apache.arrow.vector.compare.VectorEqualsVisitor; +import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.LargeListViewVector; @@ -95,13 +99,16 @@ import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.extension.JsonType; import org.apache.arrow.vector.extension.OpaqueType; +import org.apache.arrow.vector.extension.OpaqueVector; import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.IntervalDayHolder; +import org.apache.arrow.vector.holders.NullableIntHolder; import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.table.Table; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.UnionMode; 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; @@ -109,12 +116,14 @@ import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.ValueVectorUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; public class RoundtripTest { private static final String EMPTY_SCHEMA_PATH = ""; @@ -748,42 +757,16 @@ public void testEmptyLargeListViewVector() { @Test public void testMapVector() { - int count = 5; try (final MapVector vector = MapVector.empty("v", allocator, false)) { - vector.allocateNew(); - UnionMapWriter mapWriter = vector.getWriter(); - for (int i = 0; i < count; i++) { - mapWriter.startMap(); - for (int j = 0; j < i + 1; j++) { - mapWriter.startEntry(); - mapWriter.key().bigInt().writeBigInt(j); - mapWriter.value().integer().writeInt(j); - mapWriter.endEntry(); - } - mapWriter.endMap(); - } - mapWriter.setValueCount(count); - + writeTestMaps(vector); assertTrue(roundtrip(vector, MapVector.class)); } } @Test public void testUnionVector() { - final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); - uInt4Holder.value = 100; - uInt4Holder.isSet = 1; - try (UnionVector vector = UnionVector.empty("v", allocator)) { - vector.allocateNew(); - - // write some data - vector.setType(0, MinorType.UINT4); - vector.setSafe(0, uInt4Holder); - vector.setType(2, MinorType.UINT4); - vector.setSafe(2, uInt4Holder); - vector.setValueCount(4); - + writeTestUnion(vector); assertTrue(roundtrip(vector, UnionVector.class)); } } @@ -1126,22 +1109,918 @@ public void testImportArrayWithNonZeroOffset() { try (IntVector source = new IntVector("source", allocator); IntVector destination = new IntVector("destination", allocator); ArrowArray array = ArrowArray.allocateNew(allocator)) { - setVector(source, 1, 2, 3); + setVector(source, 1, 2, null, 4, 5); Data.exportVector(allocator, source, null, array); + slice(array, 1, 3); + + Data.importIntoVector(allocator, array, destination, null); + + assertEquals(3, destination.getValueCount()); + assertEquals(1, destination.getNullCount()); + assertEquals(2, destination.get(0)); + assertTrue(destination.isNull(1)); + assertEquals(4, destination.get(2)); + } + } + @Test + public void testImportIntVectorWithNonZeroOffset() { + try (IntVector source = new IntVector("source", allocator)) { + setVector(source, rows(i -> i % 3 == 0 ? null : i, Integer[]::new)); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportVarCharWithNonZeroOffset() { + try (VarCharVector source = new VarCharVector("source", allocator)) { + setVector(source, rows(i -> i % 4 == 1 ? null : "v".repeat(i % 5), String[]::new)); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportListWithNonZeroOffset() { + try (ListVector source = ListVector.empty("source", allocator)) { + setVector(source, testLists()); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportStructWithNonZeroOffset() { + // The offset of a struct applies to its children, which are exported at full length + try (StructVector source = StructVector.empty("source", allocator); + StructVector destination = StructVector.empty("destination", allocator); + ArrowArray array = ArrowArray.allocateNew(allocator)) { + Map> values = new HashMap<>(); + values.put("a", Arrays.asList(1, 2, 3, 4)); + values.put("b", Arrays.asList(10, null, 30, 40)); + setVector(source, values); + destination.initializeChildrenFromFields(source.getField().getChildren()); + Data.exportVector(allocator, source, null, array); + slice(array, 1, 2); + + Data.importIntoVector(allocator, array, destination, null); + + assertEquals(2, destination.getValueCount()); + IntVector a = (IntVector) destination.getChild("a"); + IntVector b = (IntVector) destination.getChild("b"); + assertEquals(2, a.getValueCount()); + assertEquals(2, a.get(0)); + assertEquals(3, a.get(1)); + assertTrue(b.isNull(0)); + assertEquals(30, b.get(1)); + } + } + + @Test + public void testImportStructWithNullsWithNonZeroOffset() { + try (StructVector source = StructVector.empty("source", allocator)) { + setVector(source, testStructValues()); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportStructWithChildOffsetsWithNonZeroOffset() { + // The offset of a child adds to the offset of the struct + try (StructVector source = StructVector.empty("source", allocator)) { + setVector(source, testStructValues()); + for (int childOffset : new int[] {3, 8, 11}) { + for (int offset = 0; offset < ROW_COUNT - childOffset; offset++) { + int length = ROW_COUNT - childOffset - offset; + try (ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + Data.exportVector(allocator, source, null, array, schema); + ArrowArray.Snapshot snapshot = array.snapshot(); + long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + for (long child : children) { + slice(ArrowArray.wrap(child), childOffset, ROW_COUNT - childOffset); + } + slice(array, offset, length); + + try (StructVector imported = + (StructVector) Data.importVector(allocator, array, schema, null)) { + assertEquals(length, imported.getValueCount()); + for (int i = 0; i < length; i++) { + String row = "child offset " + childOffset + ", offset " + offset + ", index " + i; + assertEquals(source.isNull(offset + i), imported.isNull(i), row); + for (String name : source.getChildFieldNames()) { + assertEquals( + source.getChild(name).getObject(childOffset + offset + i), + imported.getChild(name).getObject(i), + row + ", child " + name); + } + } + } + } + } + } + } + } + + @Test + public void testImportStructWithShortChildWithNonZeroOffset() { + try (StructVector source = StructVector.empty("source", allocator); + StructVector destination = StructVector.empty("destination", allocator); + ArrowArray array = ArrowArray.allocateNew(allocator)) { + setVector(source, testStructValues()); + destination.initializeChildrenFromFields(source.getField().getChildren()); + Data.exportVector(allocator, source, null, array); ArrowArray.Snapshot snapshot = array.snapshot(); - snapshot.offset = 1; - array.save(snapshot); + for (long child : NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children)) { + slice(ArrowArray.wrap(child), 0, 13); + } + // The slice needs 14 rows of each child + slice(array, 9, 5); Exception e = assertThrows( IllegalStateException.class, () -> Data.importIntoVector(allocator, array, destination, null)); assertEquals( - "ArrowArray struct has non-zero offset (1), which is not supported", e.getMessage()); + "ArrowArray struct child at position 0 has length 13, but its parent needs 14 rows", + e.getMessage()); + } + } + + @Test + public void testImportVectorSchemaRootWithSlicedChildren() { + // Arrow C++ exports record_batch->Slice(8, 2) as a struct whose children carry the offset + try (VectorSchemaRoot source = createTestVSR(); + VectorSchemaRoot destination = VectorSchemaRoot.create(source.getSchema(), allocator); + ArrowArray array = ArrowArray.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, source, null, array); + ArrowArray.Snapshot snapshot = array.snapshot(); + slice(array, 0, 2); + for (long child : NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children)) { + slice(ArrowArray.wrap(child), 8, 2); + } + + Data.importIntoVectorSchemaRoot(allocator, array, destination, null); + + assertEquals(2, destination.getRowCount()); + BitVector bits = (BitVector) destination.getVector("boolean"); + VarCharVector strings = (VarCharVector) destination.getVector("varchar"); + assertEquals(0, bits.get(0)); + assertEquals(1, bits.get(1)); + assertEquals("test8", new String(strings.get(0), StandardCharsets.UTF_8)); + assertEquals("test9", new String(strings.get(1), StandardCharsets.UTF_8)); + } + } + + @Test + public void testImportVectorSchemaRootWithSlicedStructOfPlainExtension() { + // The schema of the root has the extension type that the vector's getField() does not report + Field field = + new Field( + "s", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Collections.singletonList(Field.nullable("p", PlainType.INSTANCE))); + Schema schema = new Schema(Collections.singletonList(field)); + withExtensionType( + PlainType.INSTANCE, + () -> { + try (VectorSchemaRoot source = VectorSchemaRoot.create(schema, allocator); + VectorSchemaRoot destination = VectorSchemaRoot.create(schema, allocator); + ArrowArray array = ArrowArray.allocateNew(allocator)) { + StructVector struct = (StructVector) source.getVector("s"); + setTestStrings(((PlainVector) struct.getChild("p")).getUnderlyingVector()); + for (int i = 0; i < ROW_COUNT; i++) { + struct.setIndexDefined(i); + } + source.setRowCount(ROW_COUNT); + Data.exportVectorSchemaRoot(allocator, source, null, array); + ArrowArray.Snapshot snapshot = array.snapshot(); + slice(array, 0, 5); + long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); + for (long child : children) { + slice(ArrowArray.wrap(child), 9, 5); + } + + Data.importIntoVectorSchemaRoot(allocator, array, destination, null); + + assertEquals(5, destination.getRowCount()); + StructVector imported = (StructVector) destination.getVector("s"); + assertInstanceOf(PlainVector.class, imported.getChild("p")); + for (int i = 0; i < 5; i++) { + assertEquals(struct.getObject(9 + i), imported.getObject(i), "index " + i); + } + } + }); + } + + @Test + public void testImportExtensionTypeVectorWithNonZeroOffset() { + try (UuidVector source = new UuidVector("source", allocator)) { + setTestUuids(source); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportStructWithExtensionTypeChildWithNonZeroOffset() { + try (StructVector source = StructVector.empty("source", allocator)) { + UuidVector child = + source.addOrGet("u", FieldType.nullable(UuidType.INSTANCE), UuidVector.class); + setTestUuids(child); + setStructRowsDefined(source); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportBitVectorWithNonZeroOffset() { + try (BitVector source = new BitVector("source", allocator)) { + setVector(source, rows(i -> i % 5 == 0 ? null : i % 2, Integer[]::new)); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportLargeListWithNonZeroOffset() { + try (LargeListVector source = LargeListVector.empty("source", allocator)) { + setVector(source, testLists()); + assertSlicedImportsMatchSource(source); + } + } + + @Test + @SuppressWarnings("unchecked") + public void testImportFixedSizeListWithNonZeroOffset() { + try (FixedSizeListVector source = FixedSizeListVector.empty("source", 2, allocator)) { + setVector(source, rows(i -> Arrays.asList(2 * i, 2 * i + 1), List[]::new)); + for (int i = 1; i < ROW_COUNT; i += 6) { + source.setNull(i); + } + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportMapWithNonZeroOffset() { + try (MapVector source = MapVector.empty("source", allocator, false)) { + writeTestMaps(source); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportUnionWithNonZeroOffset() { + try (UnionVector source = UnionVector.empty("source", allocator)) { + writeTestUnion(source); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportDenseUnionWithNonZeroOffset() { + final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); + uInt4Holder.isSet = 1; + final NullableIntHolder intHolder = new NullableIntHolder(); + try (DenseUnionVector source = DenseUnionVector.empty("source", allocator)) { + source.allocateNew(); + byte uint4TypeId = source.registerNewTypeId(Field.nullable("u", MinorType.UINT4.getType())); + byte intTypeId = source.registerNewTypeId(Field.nullable("i", MinorType.INT.getType())); + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 3 == 0) { + uInt4Holder.value = 100 + i; + source.setTypeId(i, uint4TypeId); + source.setSafe(i, uInt4Holder); + } else { + intHolder.isSet = i % 5 == 0 ? 0 : 1; + intHolder.value = -i; + source.setTypeId(i, intTypeId); + source.setSafe(i, intHolder); + } + } + source.setValueCount(ROW_COUNT); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportDenseUnionWithTypeIdsWithNonZeroOffset() { + // The type ids differ from the positions of the children they select + Field field = + new Field( + "source", + FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {5, 10})), + Arrays.asList( + Field.nullable("i", MinorType.INT.getType()), + Field.nullable("s", MinorType.VARCHAR.getType()))); + try (DenseUnionVector source = (DenseUnionVector) field.createVector(allocator)) { + source.allocateNew(); + IntVector ints = (IntVector) source.getVectorByType((byte) 5); + VarCharVector strings = (VarCharVector) source.getVectorByType((byte) 10); + int intCount = 0; + int stringCount = 0; + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 3 == 0) { + source.setTypeId(i, (byte) 10); + source.setOffset(i, stringCount); + strings.setSafe(stringCount++, ("s" + i).getBytes(StandardCharsets.UTF_8)); + } else { + source.setTypeId(i, (byte) 5); + source.setOffset(i, intCount); + ints.setSafe(intCount++, i); + } + } + source.setValueCount(ROW_COUNT); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportDenseUnionWithOffsetGapsWithNonZeroOffset() { + // The offsets of a dense union only have to increase, so they may skip values of a child + try (DenseUnionVector source = + (DenseUnionVector) denseUnionWithIntChild("source").createVector(allocator)) { + setTestValuesWithOffsetGaps(source); + assertSlicedImportsMatchSource(source); + } + } + + @Test + public void testImportStructWithDenseUnionWithOffsetGapsWithNonZeroOffset() { + // The struct sets the value count of the dense union after transferring it + Field field = + new Field( + "source", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Collections.singletonList(denseUnionWithIntChild("u"))); + try (StructVector source = (StructVector) field.createVector(allocator)) { + for (int i = 0; i < OFFSETS_WITH_GAPS.length; i++) { + if (i % 6 != 5) { + source.setIndexDefined(i); + } + } + source.setValueCount(OFFSETS_WITH_GAPS.length); + setTestValuesWithOffsetGaps((DenseUnionVector) source.getChild("u")); + assertSlicedImportsMatchSource(source); + } + } + + private static final int[] OFFSETS_WITH_GAPS = {0, 1, 3, 4, 5, 8, 9, 12, 13, 14, 17, 19}; + + private static Field denseUnionWithIntChild(String name) { + return new Field( + name, + FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {0})), + Collections.singletonList(Field.nullable("i", MinorType.INT.getType()))); + } + + /** Selects the values at {@link #OFFSETS_WITH_GAPS} from the child of a dense union. */ + private static void setTestValuesWithOffsetGaps(DenseUnionVector vector) { + vector.allocateNew(); + IntVector child = (IntVector) vector.getVectorByType((byte) 0); + setVector(child, rows(i -> 100 + i, Integer[]::new)); + for (int i = 0; i < OFFSETS_WITH_GAPS.length; i++) { + vector.setTypeId(i, (byte) 0); + vector.setOffset(i, OFFSETS_WITH_GAPS[i]); + } + vector.setValueCount(OFFSETS_WITH_GAPS.length); + // setValueCount() sets the child's value count to the number of rows selecting it + child.setValueCount(ROW_COUNT); + } + + @Test + public void testImportRunEndEncodedWithNonZeroOffset() { + try (RunEndEncodedVector source = RunEndEncodedVector.empty("source", allocator)) { + setVector(source, List.of(2, 3, 6, 10, 11, 15, 20), List.of(1, 2, 3, 4, 5, 6, 7)); + assertSlicedImportsMatchSource(source); } } + @Test + public void testImportViewVarCharWithNonZeroOffset() { + try (ViewVarCharVector source = new ViewVarCharVector("source", allocator)) { + setVector( + source, + rows( + i -> + i % 4 == 1 + ? null + : i % 2 == 0 ? "short " + i : "a string longer than twelve bytes " + i, + String[]::new)); + assertSlicedImportsMatchSource(source); + } + } + + @ParameterizedTest + @ValueSource(ints = {1 << 16, (1 << 16) + 3}) + public void testImportSliceWithoutValidityBitmapAllocatesOnlyTheSlice(int offset) { + // Producers commonly omit the validity bitmap of a column without nulls. Neither the bitmap + // synthesized for the import nor the imported buffers may grow with the offset. + final int length = 10; + try (IntVector source = new IntVector("source", allocator); + ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + source.allocateNew(offset + length); + for (int i = 0; i < offset + length; i++) { + source.set(i, i); + } + source.setValueCount(offset + length); + Data.exportVector(allocator, source, null, array, schema); + slice(array, offset, length); + dropValidityBitmap(array); + + try (BufferAllocator importAllocator = + allocator.newChildAllocator("import", 0, Long.MAX_VALUE); + IntVector imported = + (IntVector) Data.importVector(importAllocator, array, schema, null)) { + assertEquals(length, imported.getValueCount()); + assertEquals(0, imported.getNullCount()); + for (int i = 0; i < length; i++) { + assertEquals(offset + i, imported.get(i)); + } + assertTrue( + importAllocator.getAllocatedMemory() < 1024, + "allocated " + importAllocator.getAllocatedMemory() + " bytes"); + } + } + } + + @ParameterizedTest + @ValueSource(ints = {10, 1 << 16, (1 << 16) + 3}) + public void testImportStructSliceWithoutValidityBitmapsAllocatesOnlyTheSlice(int offset) { + // The children of a sliced struct are exported at full length, and only the rows of the + // slice may be imported from them + final int rowCount = (1 << 16) + 20; + final int length = 10; + try (StructVector source = StructVector.empty("source", allocator); + ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + IntVector child = + source.addOrGet("a", FieldType.nullable(MinorType.INT.getType()), IntVector.class); + child.allocateNew(rowCount); + for (int i = 0; i < rowCount; i++) { + child.set(i, i); + source.setIndexDefined(i); + } + source.setValueCount(rowCount); + Data.exportVector(allocator, source, null, array, schema); + ArrowArray.Snapshot snapshot = array.snapshot(); + dropValidityBitmap(array); + dropValidityBitmap(ArrowArray.wrap(NativeUtil.toJavaArray(snapshot.children, 1)[0])); + slice(array, offset, length); + + try (BufferAllocator importAllocator = + allocator.newChildAllocator("import", 0, Long.MAX_VALUE); + StructVector imported = + (StructVector) Data.importVector(importAllocator, array, schema, null)) { + assertEquals(length, imported.getValueCount()); + IntVector importedChild = (IntVector) imported.getChild("a"); + for (int i = 0; i < length; i++) { + assertFalse(imported.isNull(i)); + assertEquals(offset + i, importedChild.get(i)); + } + assertTrue( + importAllocator.getAllocatedMemory() < 1024, + "allocated " + importAllocator.getAllocatedMemory() + " bytes"); + } + } + } + + @Test + public void testImportOpaqueVectorWithNonZeroOffset() { + // OpaqueVector overrides getField() but not makeTransferPair() + withExtensionType( + OPAQUE_TYPE, + () -> { + try (OpaqueVector source = + (OpaqueVector) Field.nullable("source", OPAQUE_TYPE).createVector(allocator)) { + setTestStrings((VarCharVector) source.getUnderlyingVector()); + assertSlicedImportsMatchSource(source); + } + }); + } + + @Test + public void testImportStructWithOpaqueChildWithNonZeroOffset() { + Field field = + new Field( + "source", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Collections.singletonList(Field.nullable("o", OPAQUE_TYPE))); + withExtensionType( + OPAQUE_TYPE, + () -> { + try (StructVector source = (StructVector) field.createVector(allocator)) { + OpaqueVector child = (OpaqueVector) source.getChild("o"); + setTestStrings((VarCharVector) child.getUnderlyingVector()); + setStructRowsDefined(source); + assertSlicedImportsMatchSource(source); + } + }); + } + + @Test + public void testImportListOfOpaqueWithNonZeroOffset() { + Field field = + new Field( + "source", + FieldType.nullable(ArrowType.List.INSTANCE), + Collections.singletonList(Field.nullable("item", OPAQUE_TYPE))); + withExtensionType( + OPAQUE_TYPE, + () -> { + try (ListVector source = (ListVector) field.createVector(allocator)) { + OpaqueVector items = (OpaqueVector) source.getDataVector(); + writeTestStringLists(source, (VarCharVector) items.getUnderlyingVector()); + assertSlicedImportsMatchSource(source); + } + }); + } + + @Test + public void testImportPlainExtensionVectorWithNonZeroOffset() { + // PlainVector overrides neither getField() nor makeTransferPair() + Field field = Field.nullable("source", PlainType.INSTANCE); + withExtensionType( + PlainType.INSTANCE, + () -> { + try (PlainVector source = (PlainVector) field.createVector(allocator)) { + setTestStrings(source.getUnderlyingVector()); + assertSlicedImportsMatchSource(field, source); + } + }); + } + + @Test + public void testImportIntoPlainExtensionVectorWithNonZeroOffset() { + // The target reports only its storage type, so the temporary vector is a plain VarCharVector + Field field = Field.nullable("source", PlainType.INSTANCE); + withExtensionType( + PlainType.INSTANCE, + () -> { + try (PlainVector source = (PlainVector) field.createVector(allocator); + PlainVector destination = (PlainVector) field.createVector(allocator); + ArrowArray array = ArrowArray.allocateNew(allocator)) { + setTestStrings(source.getUnderlyingVector()); + Data.exportVector(allocator, source, null, array); + slice(array, 9, 5); + + Data.importIntoVector(allocator, array, destination, null); + + assertEquals(5, destination.getValueCount()); + for (int i = 0; i < 5; i++) { + assertEquals(source.getObject(9 + i), destination.getObject(i), "index " + i); + } + } + }); + } + + @Test + public void testImportStructWithPlainExtensionChildWithNonZeroOffset() { + Field field = + new Field( + "source", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Collections.singletonList(Field.nullable("p", PlainType.INSTANCE))); + withExtensionType( + PlainType.INSTANCE, + () -> { + try (StructVector source = (StructVector) field.createVector(allocator)) { + setTestStrings(((PlainVector) source.getChild("p")).getUnderlyingVector()); + setStructRowsDefined(source); + assertSlicedImportsMatchSource(field, source); + } + }); + } + + @Test + public void testImportListOfPlainExtensionWithNonZeroOffset() { + Field field = + new Field( + "source", + FieldType.nullable(ArrowType.List.INSTANCE), + Collections.singletonList(Field.nullable("item", PlainType.INSTANCE))); + withExtensionType( + PlainType.INSTANCE, + () -> { + try (ListVector source = (ListVector) field.createVector(allocator)) { + PlainVector items = (PlainVector) source.getDataVector(); + writeTestStringLists(source, items.getUnderlyingVector()); + assertSlicedImportsMatchSource(field, source); + } + }); + } + + @Test + public void testImportIntoStructWithPlainExtensionChildWithNonZeroOffset() { + // Without the imported field the importer cannot create a PlainVector for the struct's child + Field field = + new Field( + "source", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Collections.singletonList(Field.nullable("p", PlainType.INSTANCE))); + withExtensionType( + PlainType.INSTANCE, + () -> { + try (StructVector source = (StructVector) field.createVector(allocator); + StructVector destination = (StructVector) field.createVector(allocator); + ArrowArray array = ArrowArray.allocateNew(allocator)) { + setTestStrings(((PlainVector) source.getChild("p")).getUnderlyingVector()); + setStructRowsDefined(source); + Data.exportVector(allocator, source, null, array); + slice(array, 9, 5); + + Exception e = + assertThrows( + UnsupportedOperationException.class, + () -> Data.importIntoVector(allocator, array, destination, null)); + assertEquals( + "Cannot import an ArrowArray with a non-zero offset into PlainVector p, because" + + " its getField() does not report its extension type", + e.getMessage()); + assertInstanceOf(PlainVector.class, destination.getChild("p")); + } + }); + } + + private static final int ROW_COUNT = 20; + + /** Returns the values of the {@link #ROW_COUNT} rows of a test vector. */ + private static T[] rows(IntFunction value, IntFunction newArray) { + return IntStream.range(0, ROW_COUNT).mapToObj(value).toArray(newArray); + } + + /** Lists of up to 3 values, and some nulls. */ + @SuppressWarnings({"unchecked", "rawtypes"}) + private static List[] testLists() { + List[] lists = + rows( + i -> + i % 5 == 2 + ? null + : IntStream.range(0, i % 4) + .mapToObj(j -> 10 * i + j) + .collect(Collectors.toList()), + List[]::new); + return lists; + } + + /** Two int children, which are both null in the rows where the struct is null. */ + private static Map> testStructValues() { + Map> values = new HashMap<>(); + values.put("a", Arrays.asList(rows(i -> i % 7 == 0 ? null : i, Integer[]::new))); + values.put( + "b", Arrays.asList(rows(i -> i % 7 == 0 || i % 3 == 0 ? null : 10 * i, Integer[]::new))); + return values; + } + + private static void setTestStrings(VarCharVector vector) { + setVector(vector, rows(i -> i % 4 == 1 ? null : "s" + i, String[]::new)); + } + + private static void setTestUuids(UuidVector vector) { + vector.setValueCount(ROW_COUNT); + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 4 == 0) { + vector.setNull(i); + } else { + vector.set(i, new UUID(i, -i)); + } + } + } + + /** Defines all but some rows of a struct whose children are set already. */ + private static void setStructRowsDefined(StructVector vector) { + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 6 != 5) { + vector.setIndexDefined(i); + } + } + vector.setValueCount(ROW_COUNT); + } + + /** + * Imports the slices of {@code source} that start at each offset and either have one row or end + * at the end of the source, which covers offsets before and after the first byte of a bitmap, + * aligned to a byte or not. + */ + private void assertSlicedImportsMatchSource(FieldVector source) { + assertSlicedImportsMatchSource(source.getField(), source); + } + + /** Like {@link #assertSlicedImportsMatchSource(FieldVector)}, exporting the type of field. */ + private void assertSlicedImportsMatchSource(Field field, FieldVector source) { + int valueCount = source.getValueCount(); + for (int offset = 1; offset < valueCount; offset++) { + assertSlicedImportMatchesSource(field, source, offset, 1); + assertSlicedImportMatchesSource(field, source, offset, valueCount - offset); + } + } + + /** + * Imports a slice of {@code source}, exported with the type of {@code field}, and compares it + * with the same range of the source. + */ + private void assertSlicedImportMatchesSource( + Field field, FieldVector source, int offset, int length) { + try (ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + Data.exportField(allocator, field, null, schema); + Data.exportVector(allocator, source, null, array); + slice(array, offset, length); + try (FieldVector imported = Data.importVector(allocator, array, schema, null)) { + if (canValidate(imported)) { + ValueVectorUtility.validateFull(imported); + } + assertSameVectorClasses(source, imported); + assertEquals(length, imported.getValueCount()); + for (int i = 0; i < length; i++) { + assertEquals( + source.getObject(offset + i), + imported.getObject(i), + "offset " + offset + ", length " + length + ", index " + i); + } + } + } + } + + /** + * Returns whether {@link ValueVectorUtility#validateFull} supports {@code vector}, which it does + * not for fixed-size lists, maps, run-end encoded vectors, or extension vectors that report their + * storage type. + */ + private static boolean canValidate(FieldVector vector) { + if (vector instanceof FixedSizeListVector + || vector instanceof MapVector + || vector instanceof RunEndEncodedVector + || (vector instanceof ExtensionTypeVector + && !(vector.getField().getType() instanceof ExtensionType))) { + return false; + } + return vector.getChildrenFromFields().stream().allMatch(RoundtripTest::canValidate); + } + + private static void assertSameVectorClasses(FieldVector expected, FieldVector actual) { + assertEquals(expected.getClass(), actual.getClass()); + List expectedChildren = expected.getChildrenFromFields(); + List actualChildren = actual.getChildrenFromFields(); + assertEquals(expectedChildren.size(), actualChildren.size()); + for (int i = 0; i < expectedChildren.size(); i++) { + assertSameVectorClasses(expectedChildren.get(i), actualChildren.get(i)); + } + } + + private static final OpaqueType OPAQUE_TYPE = + new OpaqueType(ArrowType.Utf8.INSTANCE, "test.type", "test.vendor"); + + /** Runs {@code test} with {@code type} registered, so that importing it creates its vector. */ + private static void withExtensionType(ExtensionType type, Runnable test) { + ExtensionType previous = ExtensionTypeRegistry.lookup(type.extensionName()); + ExtensionTypeRegistry.register(type); + try { + test.run(); + } finally { + ExtensionTypeRegistry.unregister(type); + if (previous != null) { + ExtensionTypeRegistry.register(previous); + } + } + } + + /** An extension type whose vector overrides neither getField() nor makeTransferPair(). */ + private static final class PlainType extends ExtensionType { + static final PlainType INSTANCE = new PlainType(); + + @Override + public ArrowType storageType() { + return ArrowType.Utf8.INSTANCE; + } + + @Override + public String extensionName() { + return "test.plain"; + } + + @Override + public boolean extensionEquals(ExtensionType other) { + return other instanceof PlainType; + } + + @Override + public ArrowType deserialize(ArrowType storageType, String serializedData) { + return INSTANCE; + } + + @Override + public String serialize() { + return ""; + } + + @Override + public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { + return new PlainVector(name, allocator); + } + } + + private static final class PlainVector extends ExtensionTypeVector { + PlainVector(String name, BufferAllocator allocator) { + super(name, allocator, new VarCharVector(name, allocator)); + } + + @Override + public Object getObject(int index) { + return getUnderlyingVector().getObject(index); + } + + @Override + public int hashCode(int index) { + return getUnderlyingVector().hashCode(index); + } + + @Override + public int hashCode(int index, ArrowBufHasher hasher) { + return getUnderlyingVector().hashCode(index, hasher); + } + } + + /** Writes lists of up to 3 strings and some nulls, whose items are stored in values. */ + private static void writeTestStringLists(ListVector vector, VarCharVector values) { + vector.allocateNew(); + values.allocateNew(); + int valueCount = 0; + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 5 == 2) { + vector.setNull(i); + continue; + } + vector.startNewValue(i); + for (int j = 0; j < i % 4; j++) { + if (j == 1) { + values.setNull(valueCount); + } else { + values.setSafe(valueCount, ("s" + i + j).getBytes(StandardCharsets.UTF_8)); + } + valueCount++; + } + vector.endValue(i, i % 4); + } + vector.setValueCount(ROW_COUNT); + } + + /** Rewrite an exported array in place as one whose producer omitted the validity bitmap. */ + private static void dropValidityBitmap(ArrowArray array) { + ArrowArray.Snapshot snapshot = array.snapshot(); + assertEquals(0, snapshot.null_count); + MemoryUtil.putLong(snapshot.buffers, NativeUtil.NULL); + } + + /** Writes maps of up to 3 entries and some nulls. */ + private static void writeTestMaps(MapVector vector) { + vector.allocateNew(); + UnionMapWriter mapWriter = vector.getWriter(); + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 5 == 3) { + continue; + } + mapWriter.setPosition(i); + mapWriter.startMap(); + for (int j = 0; j < i % 4; j++) { + mapWriter.startEntry(); + mapWriter.key().bigInt().writeBigInt(j); + mapWriter.value().integer().writeInt(10 * i + j); + mapWriter.endEntry(); + } + mapWriter.endMap(); + } + mapWriter.setValueCount(ROW_COUNT); + } + + /** Writes uint4 and int values, and some nulls. */ + private static void writeTestUnion(UnionVector vector) { + final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); + uInt4Holder.isSet = 1; + final NullableIntHolder intHolder = new NullableIntHolder(); + intHolder.isSet = 1; + vector.allocateNew(); + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 3 == 0) { + uInt4Holder.value = 100 + i; + vector.setType(i, MinorType.UINT4); + vector.setSafe(i, uInt4Holder); + } else if (i % 3 == 1) { + intHolder.value = -i; + vector.setType(i, MinorType.INT); + vector.setSafe(i, intHolder); + } + } + vector.setValueCount(ROW_COUNT); + } + + /** Rewrite an exported array in place as a slice, as a C producer would. */ + private static void slice(ArrowArray array, long offset, long length) { + ArrowArray.Snapshot snapshot = array.snapshot(); + snapshot.offset = offset; + snapshot.length = length; + array.save(snapshot); + } + @Test public void testImportEmptyArrayWithNonZeroOffset() { try (IntVector source = new IntVector("source", allocator); @@ -1149,11 +2028,7 @@ public void testImportEmptyArrayWithNonZeroOffset() { ArrowArray array = ArrowArray.allocateNew(allocator)) { setVector(source, 1, 2, 3); Data.exportVector(allocator, source, null, array); - - ArrowArray.Snapshot snapshot = array.snapshot(); - snapshot.offset = source.getValueCount(); - snapshot.length = 0; - array.save(snapshot); + slice(array, source.getValueCount(), 0); Data.importIntoVector(allocator, array, destination, null); assertEquals(0, destination.getValueCount()); @@ -1166,14 +2041,7 @@ public void testImportNullArrayWithNonZeroOffset() { NullVector destination = new NullVector("destination"); ArrowArray array = ArrowArray.allocateNew(allocator)) { Data.exportVector(allocator, source, null, array); - - // Mimic a sliced null array. Arrow C++ propagates the slice offset onto the struct - // even for null arrays, which export no buffers at all (n_buffers == 0), so there is - // nothing for the offset to apply to and the import is safe. - ArrowArray.Snapshot snapshot = array.snapshot(); - snapshot.offset = 2; - snapshot.length = 8; - array.save(snapshot); + slice(array, 2, 8); Data.importIntoVector(allocator, array, destination, null); assertEquals(8, destination.getValueCount()); diff --git a/docs/source/cdata.rst b/docs/source/cdata.rst index d2b5ab0cc3..588e2f1e8c 100644 --- a/docs/source/cdata.rst +++ b/docs/source/cdata.rst @@ -24,8 +24,9 @@ through :external+arrow:ref:`c-data-interface`, even between different language .. note:: - The Arrow Java C Data Interface implementation does not support importing arrays with - a non-zero offset. + Importing arrays with a non-zero offset (for example a sliced record batch) is not always + zero-copy: buffers that cannot be sliced at the offset, such as offset buffers or view data, + are copied. Java to Python -------------- diff --git a/vector/src/main/codegen/templates/DenseUnionVector.java b/vector/src/main/codegen/templates/DenseUnionVector.java index e9a9a2878b..4bcfba3941 100644 --- a/vector/src/main/codegen/templates/DenseUnionVector.java +++ b/vector/src/main/codegen/templates/DenseUnionVector.java @@ -622,13 +622,16 @@ public TransferImpl(DenseUnionVector to) { createTransferPairs(); } + // The transfer pairs follow the order of the children, while typeFields and childVectors are + // indexed by type id. typeMapFields maps the position of a child to its type id. private void createTransferPairs() { for (int i = 0; i < nextTypeId; i++) { + byte typeId = typeMapFields[i]; ValueVector srcVec = internalStruct.getVectorById(i); ValueVector dstVec = to.internalStruct.getVectorById(i); - to.typeFields[i] = typeFields[i]; - to.typeMapFields[i] = typeMapFields[i]; - to.childVectors[i] = dstVec; + to.typeFields[typeId] = typeFields[typeId]; + to.typeMapFields[i] = typeId; + to.childVectors[typeId] = dstVec; internalTransferPairs[i] = srcVec.makeTransferPair(dstVec); } } @@ -646,7 +649,7 @@ public void transfer() { for (int i = 0; i < nextTypeId; i++) { if (internalTransferPairs[i] != null) { internalTransferPairs[i].transfer(); - to.childVectors[i] = internalTransferPairs[i].getTo(); + to.childVectors[typeMapFields[i]] = internalTransferPairs[i].getTo(); } } to.valueCount = valueCount; @@ -669,29 +672,54 @@ public void splitAndTransfer(int startIndex, int length) { to.reallocOffsetBuffer(); } - int [] typeCounts = new int[nextTypeId]; - int [] typeStarts = new int[nextTypeId]; - for (int i = 0; i < typeCounts.length; i++) { - typeCounts[i] = 0; - typeStarts[i] = -1; + int[] childIndexes = new int[typeFields.length]; + Arrays.fill(childIndexes, -1); + for (int i = 0; i < nextTypeId; i++) { + childIndexes[typeMapFields[i]] = i; } + // The window selects childCounts[i] values of child i, which are renumbered from 0. The + // offsets of a child only have to increase, so the selected values may have gaps. + int[] childCounts = new int[nextTypeId]; + int[] childStarts = new int[nextTypeId]; + int[] childEnds = new int[nextTypeId]; + boolean[] contiguous = new boolean[nextTypeId]; + Arrays.fill(contiguous, true); for (int i = startIndex; i < startIndex + length; i++) { byte typeId = typeBuffer.getByte(i); if (typeId >= 0) { - to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH, typeCounts[typeId]); - typeCounts[typeId] += 1; - if (typeStarts[typeId] == -1) { - typeStarts[typeId] = offsetBuffer.getInt((long) i * OFFSET_WIDTH); + int child = childIndexes[typeId]; + int offset = offsetBuffer.getInt((long) i * OFFSET_WIDTH); + if (childCounts[child] == 0) { + childStarts[child] = offset; + } else if (offset != childEnds[child]) { + contiguous[child] = false; } + childEnds[child] = offset + 1; + to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH, childCounts[child]); + childCounts[child] += 1; } } - // transfer vector values + // transfer vector values, and copy those of children whose selected values have gaps + for (int i = 0; i < nextTypeId; i++) { + if (childCounts[i] > 0 && contiguous[i]) { + internalTransferPairs[i].splitAndTransfer(childStarts[i], childCounts[i]); + to.childVectors[typeMapFields[i]] = internalTransferPairs[i].getTo(); + } + } + for (int i = startIndex; i < startIndex + length; i++) { + byte typeId = typeBuffer.getByte(i); + if (typeId >= 0 && !contiguous[childIndexes[typeId]]) { + internalTransferPairs[childIndexes[typeId]].copyValueSafe( + offsetBuffer.getInt((long) i * OFFSET_WIDTH), + to.offsetBuffer.getInt((long) (i - startIndex) * OFFSET_WIDTH)); + } + } for (int i = 0; i < nextTypeId; i++) { - if (typeCounts[i] > 0 && typeStarts[i] != -1) { - internalTransferPairs[i].splitAndTransfer(typeStarts[i], typeCounts[i]); - to.childVectors[i] = internalTransferPairs[i].getTo(); + if (!contiguous[i]) { + internalTransferPairs[i].getTo().setValueCount(childCounts[i]); + to.childVectors[typeMapFields[i]] = internalTransferPairs[i].getTo(); } } diff --git a/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java b/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java index 8249d93460..04f6200bd8 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/ExtensionTypeVector.java @@ -154,8 +154,15 @@ public TransferPair getTransferPair(Field field, BufferAllocator allocator, Call return underlyingVector.getTransferPair(field, allocator, callBack); } + /** + * Make a transfer pair to {@code target}. An extension vector target receives the data in its + * underlying vector, any other target is passed to the transfer pair of the underlying vector. + */ @Override public TransferPair makeTransferPair(ValueVector target) { + if (target instanceof ExtensionTypeVector) { + return new TransferImpl((ExtensionTypeVector) target); + } return underlyingVector.makeTransferPair(target); } @@ -305,4 +312,35 @@ public BufferAllocator getAllocator() { public OUT accept(VectorVisitor visitor, IN value) { return visitor.visit(this, value); } + + /** Transfers the underlying vector into the underlying vector of another extension vector. */ + private class TransferImpl implements TransferPair { + private final ExtensionTypeVector to; + private final TransferPair underlyingPair; + + TransferImpl(ExtensionTypeVector to) { + this.to = to; + this.underlyingPair = underlyingVector.makeTransferPair(to.getUnderlyingVector()); + } + + @Override + public void transfer() { + underlyingPair.transfer(); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + underlyingPair.splitAndTransfer(startIndex, length); + } + + @Override + public ValueVector getTo() { + return to; + } + + @Override + public void copyValueSafe(int from, int to) { + underlyingPair.copyValueSafe(from, to); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java index 9ac30730c4..845720c188 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java @@ -22,6 +22,8 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -45,6 +47,7 @@ import org.apache.arrow.vector.util.JsonStringHashMap; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.ValueVectorUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -365,6 +368,137 @@ public void testSplitAndTransferWithMixedVectors() throws Exception { } } + @Test + public void testSplitAndTransferWithTypeIds() { + try (DenseUnionVector source = newVectorWithTypeIds(); + DenseUnionVector target = (DenseUnionVector) source.getField().createVector(allocator)) { + source.makeTransferPair(target).splitAndTransfer(1, 3); + + ValueVectorUtility.validateFull(target); + assertEquals(3, target.getValueCount()); + assertEquals(new Text("a"), target.getObject(0)); + assertEquals(2, target.getObject(1)); + assertEquals(new Text("b"), target.getObject(2)); + } + } + + @Test + public void testTransferPairWithTypeIdsCreatesTarget() { + try (DenseUnionVector source = newVectorWithTypeIds()) { + TransferPair pair = source.getTransferPair(allocator); + try (DenseUnionVector target = (DenseUnionVector) pair.getTo()) { + assertEquals(source.getField(), target.getField()); + + pair.splitAndTransfer(1, 2); + assertEquals(2, target.getValueCount()); + assertEquals(new Text("a"), target.getObject(0)); + assertEquals(2, target.getObject(1)); + + pair.transfer(); + assertEquals(4, target.getValueCount()); + assertEquals(1, target.getObject(0)); + assertEquals(new Text("b"), target.getObject(3)); + } + } + } + + /** Creates [1, "a", 2, "b"] in a dense union of ints with type id 5 and strings with 10. */ + private DenseUnionVector newVectorWithTypeIds() { + Field field = + new Field( + "duv", + FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {5, 10})), + Arrays.asList( + Field.nullable("i", MinorType.INT.getType()), + Field.nullable("s", MinorType.VARCHAR.getType()))); + DenseUnionVector vector = (DenseUnionVector) field.createVector(allocator); + vector.allocateNew(); + IntVector ints = (IntVector) vector.getVectorByType((byte) 5); + VarCharVector strings = (VarCharVector) vector.getVectorByType((byte) 10); + ints.setSafe(0, 1); + ints.setSafe(1, 2); + strings.setSafe(0, new Text("a")); + strings.setSafe(1, new Text("b")); + for (int i = 0; i < 4; i++) { + vector.setTypeId(i, (byte) (i % 2 == 0 ? 5 : 10)); + vector.setOffset(i, i / 2); + } + vector.setValueCount(4); + return vector; + } + + @Test + public void testSplitAndTransferWithOffsetGaps() { + Field field = newFieldWithIntChild(); + try (DenseUnionVector source = (DenseUnionVector) field.createVector(allocator); + DenseUnionVector target = (DenseUnionVector) field.createVector(allocator)) { + setValuesWithOffsetGaps(source); + + source.makeTransferPair(target).splitAndTransfer(1, 2); + + ValueVectorUtility.validateFull(target); + assertEquals(2, target.getValueCount()); + assertEquals(102, target.getObject(0)); + assertEquals(107, target.getObject(1)); + // The values the window selects are copied, as they are not contiguous in the child + assertEquals(2, target.getVectorByType((byte) 0).getValueCount()); + + source.makeTransferPair(target).splitAndTransfer(0, 2); + assertEquals(100, target.getObject(0)); + assertEquals(102, target.getObject(1)); + } + } + + @Test + public void testSplitAndTransferWithOffsetGapsInStruct() { + // The struct sets the value count of the dense union after the split + Field field = + new Field( + "struct", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Collections.singletonList(newFieldWithIntChild())); + try (StructVector source = (StructVector) field.createVector(allocator); + StructVector target = (StructVector) field.createVector(allocator)) { + DenseUnionVector duv = (DenseUnionVector) source.getChild("duv"); + for (int i = 0; i < 3; i++) { + source.setIndexDefined(i); + } + source.setValueCount(3); + setValuesWithOffsetGaps(duv); + + source.makeTransferPair(target).splitAndTransfer(1, 2); + + ValueVectorUtility.validateFull(target); + assertEquals(2, target.getValueCount()); + assertEquals(102, target.getObject(0).get("duv")); + assertEquals(107, target.getObject(1).get("duv")); + } + } + + private static Field newFieldWithIntChild() { + return new Field( + "duv", + FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {0})), + Collections.singletonList(Field.nullable("i", MinorType.INT.getType()))); + } + + /** Selects 100, 102 and 107 from a child that holds 100 to 107. */ + private static void setValuesWithOffsetGaps(DenseUnionVector vector) { + vector.allocateNew(); + IntVector child = (IntVector) vector.getVectorByType((byte) 0); + for (int i = 0; i < 8; i++) { + child.setSafe(i, 100 + i); + } + int[] offsets = {0, 2, 7}; + for (int i = 0; i < offsets.length; i++) { + vector.setTypeId(i, (byte) 0); + vector.setOffset(i, offsets[i]); + } + vector.setValueCount(offsets.length); + // setValueCount() sets the value count of the child to the number of rows selecting it + child.setValueCount(8); + } + @Test public void testSplitAndTransferDuvInStruct() { try (StructVector struct = StructVector.empty("struct", allocator)) { 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 ae5ac0726c..14c8dcc33a 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 @@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -55,6 +56,7 @@ import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; +import org.apache.arrow.vector.util.TransferPair; import org.apache.arrow.vector.util.VectorBatchAppender; import org.apache.arrow.vector.validate.ValidateVectorVisitor; import org.junit.jupiter.api.Test; @@ -197,6 +199,37 @@ public void testNullCheck() { assertTrue(e.getMessage().contains("underlyingVector cannot be null.")); } + @Test + public void testTransferPairBetweenExtensionVectors() { + // LocationVector overrides neither makeTransferPair() nor getField() + try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); + final LocationVector source = new LocationVector("source", allocator); + final LocationVector target = new LocationVector("target", allocator)) { + source.allocateNew(); + for (int i = 0; i < 4; i++) { + source.set(i, i, -i); + } + source.setValueCount(4); + + TransferPair pair = source.makeTransferPair(target); + assertSame(target, pair.getTo()); + + pair.splitAndTransfer(1, 2); + assertEquals(2, target.getValueCount()); + assertEquals(source.getObject(1), target.getObject(0)); + assertEquals(source.getObject(2), target.getObject(1)); + + pair.copyValueSafe(3, 1); + assertEquals(source.getObject(3), target.getObject(1)); + + java.util.Map last = source.getObject(3); + pair.transfer(); + assertEquals(4, target.getValueCount()); + assertEquals(last, target.getObject(3)); + assertEquals(0, source.getValueCount()); + } + } + /** Test that a custom Location type can be round-tripped through a temporary file. */ @Test public void roundtripLocation() throws IOException {