From 25296e04ccbe1fbfe1c9dc1c6a82a1bebf6755e5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 08:09:43 -0600 Subject: [PATCH 1/5] fix(c): support importing arrays with a non-zero offset The C Data Interface importer ignored, and later rejected, the offset of an ArrowArray struct, so a sliced record batch exported from C++ could not be imported. Java vectors cannot represent an offset, so the importer now loads the leading offset + length values into a temporary vector and transfers the requested window into the target with splitAndTransfer. Buffers that can be sliced are shared; offset buffers and unaligned validity bitmaps are copied. Closes #88 --- .../org/apache/arrow/c/ArrayImporter.java | 43 ++- c/src/main/java/org/apache/arrow/c/Data.java | 6 +- .../org/apache/arrow/c/DictionaryTest.java | 76 +++-- .../org/apache/arrow/c/RoundtripTest.java | 305 +++++++++++++++--- docs/source/cdata.rst | 5 +- 5 files changed, 350 insertions(+), 85 deletions(-) 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..4e4518b13a 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -25,12 +25,12 @@ 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.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.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.DictionaryEncoding; /** Importer for {@link ArrowArray}. */ @@ -84,15 +84,40 @@ 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); + } + } + + /** + * Java vectors cannot represent an offset, so the leading {@code offset + length} values are + * imported into a temporary vector and the requested window is then transferred into the target. + */ + private void importSliced(ArrowArray.Snapshot snapshot) { + int offset = checkedCastToInt(snapshot.offset); + int length = checkedCastToInt(snapshot.length); + + snapshot.length = Math.addExact(snapshot.offset, snapshot.length); + snapshot.offset = 0; + // The producer's null count only covers the slice. Without a validity bitmap the null count + // must be 0 for the buffers to load, otherwise it is unknown. + boolean hasBitmap = snapshot.n_buffers > 0 && MemoryUtil.getLong(snapshot.buffers) != NULL; + snapshot.null_count = hasBitmap ? -1 : 0; + + try (FieldVector unslicedVector = vector.getField().createVector(allocator)) { + ArrayImporter unslicedImporter = + new ArrayImporter(allocator, unslicedVector, dictionaryProvider); + unslicedImporter.recursionLevel = recursionLevel; + unslicedImporter.underlyingAllocation = underlyingAllocation; + unslicedImporter.importUnsliced(snapshot); + unslicedVector.makeTransferPair(vector).splitAndTransfer(offset, length); + } + } + private void importUnsliced(ArrowArray.Snapshot snapshot) { // First import children (required for reconstituting parent array data) long[] children = NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children)); 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..76521c8c7c 100644 --- a/c/src/main/java/org/apache/arrow/c/Data.java +++ b/c/src/main/java/org/apache/arrow/c/Data.java @@ -35,8 +35,10 @@ /** * 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. */ public final class Data { 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..68fb11cb8d 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,70 @@ 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)) { + roundtrip(encodedVector, provider, IntVector.class); + } + } + + @Test + public void testWithDictionaryAndNonZeroOffset() throws Exception { + try (DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(); + IntVector encodedVector = createEncodedVector(provider); + 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)); + } + } + } + } + + /** Registers an aa/bb/cc dictionary with {@code provider} and encodes bb/bb/cc/aa with it. */ + private IntVector createEncodedVector(DictionaryProvider.MapDictionaryProvider provider) { 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(); + 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); + 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..447178299e 100644 --- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java +++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java @@ -83,6 +83,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; @@ -97,6 +98,7 @@ import org.apache.arrow.vector.extension.OpaqueType; 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; @@ -748,42 +750,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 +1102,262 @@ 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 testImportArrayWithUnalignedOffset() { + try (IntVector source = new IntVector("source", allocator)) { + Integer[] values = new Integer[20]; + for (int i = 0; i < values.length; i++) { + values[i] = i % 3 == 0 ? null : i; + } + setVector(source, values); + assertSlicedImportMatchesSource(source, 5, 11); + } + } + + @Test + public void testImportVarCharWithNonZeroOffset() { + try (VarCharVector source = new VarCharVector("source", allocator)) { + setVector(source, "a", "bb", null, "dddd", "e"); + assertSlicedImportMatchesSource(source, 2, 2); + } + } + + @Test + public void testImportListWithNonZeroOffset() { + try (ListVector source = ListVector.empty("source", allocator)) { + setVector(source, Arrays.asList(1, 2), null, Arrays.asList(3), Arrays.asList(4, 5, 6)); + assertSlicedImportMatchesSource(source, 1, 3); + } + } + + @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 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(); - snapshot.offset = 1; - array.save(snapshot); + slice(array, 0, 2); + for (long child : NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children)) { + slice(ArrowArray.wrap(child), 8, 2); + } - 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()); + 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 testImportExtensionTypeVectorWithNonZeroOffset() { + try (UuidVector source = new UuidVector("source", allocator)) { + source.setValueCount(3); + for (int i = 0; i < 3; i++) { + source.set(i, UUID.randomUUID()); + } + assertSlicedImportMatchesSource(source, 1, 2); } } + @Test + public void testImportStructWithExtensionTypeChildWithNonZeroOffset() { + try (StructVector source = StructVector.empty("source", allocator)) { + UuidVector child = + source.addOrGet("u", FieldType.nullable(UuidType.INSTANCE), UuidVector.class); + child.setValueCount(3); + for (int i = 0; i < 3; i++) { + child.set(i, UUID.randomUUID()); + source.setIndexDefined(i); + } + source.setValueCount(3); + assertSlicedImportMatchesSource(source, 1, 2); + } + } + + @Test + public void testImportBitVectorWithUnalignedOffset() { + try (BitVector source = new BitVector("source", allocator)) { + Integer[] values = new Integer[20]; + for (int i = 0; i < values.length; i++) { + values[i] = i % 5 == 0 ? null : i % 2; + } + setVector(source, values); + assertSlicedImportMatchesSource(source, 3, 13); + } + } + + @Test + public void testImportLargeListWithNonZeroOffset() { + try (LargeListVector source = LargeListVector.empty("source", allocator)) { + setVector(source, Arrays.asList(1, 2), null, Arrays.asList(3), Arrays.asList(4, 5, 6)); + assertSlicedImportMatchesSource(source, 1, 3); + } + } + + @Test + public void testImportFixedSizeListWithNonZeroOffset() { + try (FixedSizeListVector source = FixedSizeListVector.empty("source", 2, allocator)) { + setVector(source, Arrays.asList(1, 2), Arrays.asList(3, 4), Arrays.asList(5, 6)); + assertSlicedImportMatchesSource(source, 1, 2); + } + } + + @Test + public void testImportMapWithNonZeroOffset() { + try (MapVector source = MapVector.empty("source", allocator, false)) { + writeTestMaps(source); + assertSlicedImportMatchesSource(source, 2, 2); + } + } + + @Test + public void testImportUnionWithNonZeroOffset() { + try (UnionVector source = UnionVector.empty("source", allocator)) { + writeTestUnion(source); + assertSlicedImportMatchesSource(source, 1, 3); + } + } + + @Test + public void testImportDenseUnionWithNonZeroOffset() { + final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); + uInt4Holder.value = 100; + uInt4Holder.isSet = 1; + final NullableIntHolder intHolder = new NullableIntHolder(); + intHolder.value = -7; + intHolder.isSet = 1; + 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())); + source.setTypeId(0, uint4TypeId); + source.setSafe(0, uInt4Holder); + source.setTypeId(1, intTypeId); + source.setSafe(1, intHolder); + source.setTypeId(2, uint4TypeId); + source.setSafe(2, uInt4Holder); + source.setTypeId(3, intTypeId); + source.setSafe(3, intHolder); + source.setValueCount(4); + assertSlicedImportMatchesSource(source, 1, 3); + } + } + + @Test + public void testImportRunEndEncodedWithNonZeroOffset() { + try (RunEndEncodedVector source = RunEndEncodedVector.empty("source", allocator)) { + setVector(source, List.of(2, 3, 6), List.of(1, 2, 3)); + assertSlicedImportMatchesSource(source, 1, 4); + } + } + + @Test + public void testImportViewVarCharWithNonZeroOffset() { + try (ViewVarCharVector source = new ViewVarCharVector("source", allocator)) { + setVector(source, "short", "a string longer than twelve bytes", null, "another long string"); + assertSlicedImportMatchesSource(source, 1, 3); + } + } + + /** Imports a slice of {@code source} and compares it with the same range of the source. */ + private void assertSlicedImportMatchesSource(FieldVector source, int offset, int length) { + try (ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { + Data.exportVector(allocator, source, null, array, schema); + slice(array, offset, length); + try (FieldVector imported = Data.importVector(allocator, array, schema, null)) { + assertEquals(length, imported.getValueCount()); + for (int i = 0; i < length; i++) { + assertEquals(source.getObject(offset + i), imported.getObject(i), "index " + i); + } + } + } + } + + private static void writeTestMaps(MapVector vector) { + vector.allocateNew(); + UnionMapWriter mapWriter = vector.getWriter(); + for (int i = 0; i < 5; 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(5); + } + + private static void writeTestUnion(UnionVector vector) { + final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); + uInt4Holder.value = 100; + uInt4Holder.isSet = 1; + vector.allocateNew(); + vector.setType(0, MinorType.UINT4); + vector.setSafe(0, uInt4Holder); + vector.setType(2, MinorType.UINT4); + vector.setSafe(2, uInt4Holder); + vector.setValueCount(4); + } + + /** 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 +1365,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 +1378,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 -------------- From 21df8354c801fa4176e3244a54f8bd4ce88295d1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 12:29:00 -0600 Subject: [PATCH 2/5] Update c/src/main/java/org/apache/arrow/c/ArrayImporter.java MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: JB Onofré --- c/src/main/java/org/apache/arrow/c/ArrayImporter.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 4e4518b13a..68b2c74db3 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -104,7 +104,10 @@ private void importSliced(ArrowArray.Snapshot snapshot) { snapshot.offset = 0; // The producer's null count only covers the slice. Without a validity bitmap the null count // must be 0 for the buffers to load, otherwise it is unknown. - boolean hasBitmap = snapshot.n_buffers > 0 && MemoryUtil.getLong(snapshot.buffers) != NULL; + boolean hasBitmap = + snapshot.n_buffers > 0 + && snapshot.buffers != NULL + && MemoryUtil.getLong(snapshot.buffers) != NULL; snapshot.null_count = hasBitmap ? -1 : 0; try (FieldVector unslicedVector = vector.getField().createVector(allocator)) { From b997324fe1acc36c3b313ff3aefff90065f097f5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 13:15:31 -0600 Subject: [PATCH 3/5] fix(vector): let ExtensionTypeVector transfer into extension vectors ExtensionTypeVector.makeTransferPair passed an extension vector target to the transfer pair of its underlying vector, which casts the target to its own class. Transfers between extension vectors that do not override makeTransferPair, such as OpaqueVector, failed with a ClassCastException. Transfer into the underlying vector of an extension vector target instead. --- .../arrow/vector/ExtensionTypeVector.java | 38 +++++++++++++++++++ .../vector/types/pojo/TestExtensionType.java | 33 ++++++++++++++++ 2 files changed, 71 insertions(+) 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/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 { From 0321931e1108cf875378bf32d58e98c998edc5c8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 13:15:31 -0600 Subject: [PATCH 4/5] fix(vector): split dense unions with type ids or offset gaps DenseUnionVector's transfer pair indexed its per-child state by type id, although it is sized and ordered by child. Splitting a dense union whose type ids are not 0 to n - 1 failed with an ArrayIndexOutOfBoundsException, and transfers placed the target's children under the wrong type ids. The split also assumed that the values a window selects from a child are contiguous. The offsets of a child only have to increase, so for offsets with gaps the split returned the wrong values. Values with gaps are now copied, while contiguous values are still split. --- .../codegen/templates/DenseUnionVector.java | 62 +++++--- .../arrow/vector/TestDenseUnionVector.java | 134 ++++++++++++++++++ 2 files changed, 179 insertions(+), 17 deletions(-) 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/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)) { From 112a0ff3810acfdb32139f3f5ddfa56620aa60b9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 13:15:31 -0600 Subject: [PATCH 5/5] fix(c): bound sliced imports by the slice and keep extension types A sliced import loaded every row before the slice into a temporary vector, so the validity bitmap synthesized for an array without one, and the imported buffers, grew with the offset and stayed alive with the target. Skip the leading rows that fill whole bytes of a bitmap by advancing the buffer pointers, and import only the rows of the slice from the children of structs, sparse unions and fixed-size lists. Create the temporary vector from the imported field, which has the extension types of vectors whose getField() reports their storage type, and transfer between storage vectors at the top level. Importing into an existing vector whose nested extension vectors do not report their type fails with a clear error instead of replacing them. --- .../org/apache/arrow/c/ArrayImporter.java | 156 +++- .../arrow/c/BufferImportTypeVisitor.java | 33 +- c/src/main/java/org/apache/arrow/c/Data.java | 28 +- .../org/apache/arrow/c/DictionaryTest.java | 61 +- .../org/apache/arrow/c/RoundtripTest.java | 797 ++++++++++++++++-- 5 files changed, 968 insertions(+), 107 deletions(-) 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 68b2c74db3..d7b4ecd6a5 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -27,11 +27,16 @@ 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( @@ -88,39 +104,74 @@ private void doImport(ArrowArray.Snapshot snapshot) { if (snapshot.offset != 0 && snapshot.length != 0) { importSliced(snapshot); } else { - importUnsliced(snapshot); + importUnsliced(snapshot, /* skippedRows= */ 0, /* trimChildren= */ false); } } /** - * Java vectors cannot represent an offset, so the leading {@code offset + length} values are - * imported into a temporary vector and the requested window is then transferred into the target. + * 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) { - int offset = checkedCastToInt(snapshot.offset); + // 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.length = Math.addExact(snapshot.offset, snapshot.length); snapshot.offset = 0; - // The producer's null count only covers the slice. Without a validity bitmap the null count - // must be 0 for the buffers to load, otherwise it is unknown. - boolean hasBitmap = - snapshot.n_buffers > 0 - && snapshot.buffers != NULL - && MemoryUtil.getLong(snapshot.buffers) != NULL; - snapshot.null_count = hasBitmap ? -1 : 0; + snapshot.length = Math.addExact(offset, length); + snapshot.null_count = unknownNullCount(snapshot); - try (FieldVector unslicedVector = vector.getField().createVector(allocator)) { + 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, dictionaryProvider); + new ArrayImporter(allocator, unslicedVector, unslicedField, dictionaryProvider); unslicedImporter.recursionLevel = recursionLevel; unslicedImporter.underlyingAllocation = underlyingAllocation; - unslicedImporter.importUnsliced(snapshot); - unslicedVector.makeTransferPair(vector).splitAndTransfer(offset, length); + unslicedImporter.importUnsliced(snapshot, skippedRows, /* trimChildren= */ true); + from.makeTransferPair(to).splitAndTransfer(offset, length); } } - private void importUnsliced(ArrowArray.Snapshot snapshot) { + /** + * 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)); @@ -131,11 +182,31 @@ private void importUnsliced(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); } } @@ -152,7 +223,7 @@ private void importUnsliced(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 @@ -161,7 +232,8 @@ private void importUnsliced(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(); @@ -178,4 +250,42 @@ private void importUnsliced(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 76521c8c7c..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; @@ -39,6 +40,8 @@ * *

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 { @@ -360,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(); @@ -409,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; } @@ -456,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 68fb11cb8d..9f02a3967d 100644 --- a/c/src/test/java/org/apache/arrow/c/DictionaryTest.java +++ b/c/src/test/java/org/apache/arrow/c/DictionaryTest.java @@ -101,7 +101,7 @@ void roundtrip(FieldVector vector, DictionaryProvider provider, Class clazz) public void testWithDictionary() throws Exception { try (DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); - IntVector encodedVector = createEncodedVector(provider)) { + IntVector encodedVector = createEncodedVector(provider, "bb", "bb", "cc", "aa")) { roundtrip(encodedVector, provider, IntVector.class); } } @@ -110,7 +110,7 @@ public void testWithDictionary() throws Exception { public void testWithDictionaryAndNonZeroOffset() throws Exception { try (DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); - IntVector encodedVector = createEncodedVector(provider); + IntVector encodedVector = createEncodedVector(provider, "bb", "bb", "cc", "aa"); ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { Data.exportVector( @@ -142,8 +142,48 @@ public void testWithDictionaryAndNonZeroOffset() throws Exception { } } - /** Registers an aa/bb/cc dictionary with {@code provider} and encodes bb/bb/cc/aa with it. */ - private IntVector createEncodedVector(DictionaryProvider.MapDictionaryProvider provider) { + @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)); @@ -156,11 +196,14 @@ private IntVector createEncodedVector(DictionaryProvider.MapDictionaryProvider p try (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); + 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); } } 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 447178299e..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; @@ -96,6 +99,7 @@ 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; @@ -104,6 +108,7 @@ 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; @@ -111,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 = ""; @@ -1117,30 +1124,26 @@ public void testImportArrayWithNonZeroOffset() { } @Test - public void testImportArrayWithUnalignedOffset() { + public void testImportIntVectorWithNonZeroOffset() { try (IntVector source = new IntVector("source", allocator)) { - Integer[] values = new Integer[20]; - for (int i = 0; i < values.length; i++) { - values[i] = i % 3 == 0 ? null : i; - } - setVector(source, values); - assertSlicedImportMatchesSource(source, 5, 11); + 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, "a", "bb", null, "dddd", "e"); - assertSlicedImportMatchesSource(source, 2, 2); + 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, Arrays.asList(1, 2), null, Arrays.asList(3), Arrays.asList(4, 5, 6)); - assertSlicedImportMatchesSource(source, 1, 3); + setVector(source, testLists()); + assertSlicedImportsMatchSource(source); } } @@ -1171,6 +1174,77 @@ public void testImportStructWithNonZeroOffset() { } } + @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(); + 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 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 @@ -1196,14 +1270,52 @@ public void testImportVectorSchemaRootWithSlicedChildren() { } } + @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)) { - source.setValueCount(3); - for (int i = 0; i < 3; i++) { - source.set(i, UUID.randomUUID()); - } - assertSlicedImportMatchesSource(source, 1, 2); + setTestUuids(source); + assertSlicedImportsMatchSource(source); } } @@ -1212,41 +1324,37 @@ public void testImportStructWithExtensionTypeChildWithNonZeroOffset() { try (StructVector source = StructVector.empty("source", allocator)) { UuidVector child = source.addOrGet("u", FieldType.nullable(UuidType.INSTANCE), UuidVector.class); - child.setValueCount(3); - for (int i = 0; i < 3; i++) { - child.set(i, UUID.randomUUID()); - source.setIndexDefined(i); - } - source.setValueCount(3); - assertSlicedImportMatchesSource(source, 1, 2); + setTestUuids(child); + setStructRowsDefined(source); + assertSlicedImportsMatchSource(source); } } @Test - public void testImportBitVectorWithUnalignedOffset() { + public void testImportBitVectorWithNonZeroOffset() { try (BitVector source = new BitVector("source", allocator)) { - Integer[] values = new Integer[20]; - for (int i = 0; i < values.length; i++) { - values[i] = i % 5 == 0 ? null : i % 2; - } - setVector(source, values); - assertSlicedImportMatchesSource(source, 3, 13); + 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, Arrays.asList(1, 2), null, Arrays.asList(3), Arrays.asList(4, 5, 6)); - assertSlicedImportMatchesSource(source, 1, 3); + setVector(source, testLists()); + assertSlicedImportsMatchSource(source); } } @Test + @SuppressWarnings("unchecked") public void testImportFixedSizeListWithNonZeroOffset() { try (FixedSizeListVector source = FixedSizeListVector.empty("source", 2, allocator)) { - setVector(source, Arrays.asList(1, 2), Arrays.asList(3, 4), Arrays.asList(5, 6)); - assertSlicedImportMatchesSource(source, 1, 2); + 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); } } @@ -1254,7 +1362,7 @@ public void testImportFixedSizeListWithNonZeroOffset() { public void testImportMapWithNonZeroOffset() { try (MapVector source = MapVector.empty("source", allocator, false)) { writeTestMaps(source); - assertSlicedImportMatchesSource(source, 2, 2); + assertSlicedImportsMatchSource(source); } } @@ -1262,92 +1370,647 @@ public void testImportMapWithNonZeroOffset() { public void testImportUnionWithNonZeroOffset() { try (UnionVector source = UnionVector.empty("source", allocator)) { writeTestUnion(source); - assertSlicedImportMatchesSource(source, 1, 3); + assertSlicedImportsMatchSource(source); } } @Test public void testImportDenseUnionWithNonZeroOffset() { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); - uInt4Holder.value = 100; uInt4Holder.isSet = 1; final NullableIntHolder intHolder = new NullableIntHolder(); - intHolder.value = -7; - intHolder.isSet = 1; 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())); - source.setTypeId(0, uint4TypeId); - source.setSafe(0, uInt4Holder); - source.setTypeId(1, intTypeId); - source.setSafe(1, intHolder); - source.setTypeId(2, uint4TypeId); - source.setSafe(2, uInt4Holder); - source.setTypeId(3, intTypeId); - source.setSafe(3, intHolder); - source.setValueCount(4); - assertSlicedImportMatchesSource(source, 1, 3); + 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), List.of(1, 2, 3)); - assertSlicedImportMatchesSource(source, 1, 4); + 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, "short", "a string longer than twelve bytes", null, "another long string"); - assertSlicedImportMatchesSource(source, 1, 3); + setVector( + source, + rows( + i -> + i % 4 == 1 + ? null + : i % 2 == 0 ? "short " + i : "a string longer than twelve bytes " + i, + String[]::new)); + assertSlicedImportsMatchSource(source); } } - /** Imports a slice of {@code source} and compares it with the same range of the source. */ - private void assertSlicedImportMatchesSource(FieldVector source, int offset, int length) { - try (ArrowArray array = ArrowArray.allocateNew(allocator); + @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), "index " + 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 < 5; i++) { + for (int i = 0; i < ROW_COUNT; i++) { + if (i % 5 == 3) { + continue; + } + mapWriter.setPosition(i); mapWriter.startMap(); - for (int j = 0; j < i + 1; j++) { + for (int j = 0; j < i % 4; j++) { mapWriter.startEntry(); mapWriter.key().bigInt().writeBigInt(j); - mapWriter.value().integer().writeInt(j); + mapWriter.value().integer().writeInt(10 * i + j); mapWriter.endEntry(); } mapWriter.endMap(); } - mapWriter.setValueCount(5); + mapWriter.setValueCount(ROW_COUNT); } + /** Writes uint4 and int values, and some nulls. */ private static void writeTestUnion(UnionVector vector) { final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder(); - uInt4Holder.value = 100; uInt4Holder.isSet = 1; + final NullableIntHolder intHolder = new NullableIntHolder(); + intHolder.isSet = 1; vector.allocateNew(); - vector.setType(0, MinorType.UINT4); - vector.setSafe(0, uInt4Holder); - vector.setType(2, MinorType.UINT4); - vector.setSafe(2, uInt4Holder); - vector.setValueCount(4); + 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. */