diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java index 83b985b701..83077409ae 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java @@ -20,11 +20,16 @@ import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.complex.LargeListVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.compression.NoCompressionCodec; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.util.TransferPair; /** * Helper class that handles converting a {@link VectorSchemaRoot} to a {@link ArrowRecordBatch}. @@ -74,6 +79,12 @@ public VectorUnloader( /** * Performs the depth first traversal of the Vectors to create an {@link ArrowRecordBatch} * suitable for serialization. + * + *

A variable-width, list or map vector whose offsets start past zero, or whose child vector + * holds elements past its last offset, as one imported through the C Data interface from a sliced + * array does, is written as the values between its first and last offset, with the offsets + * re-based to start at zero. Vectors with a union, view or extension vector at or below them are + * written as they are. */ public ArrowRecordBatch getRecordBatch() { List nodes = new ArrayList<>(); @@ -111,6 +122,10 @@ private void appendNodes( List buffers, List variadicBufferCounts) { FieldVector storageVector = ExtensionTypeVector.getStorageVector(vector); + if (firstOffset(storageVector) != 0 && splitKeepsValues(vector)) { + appendRebased(vector, vector.getValueCount(), nodes, buffers, variadicBufferCounts); + return; + } nodes.add( new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1)); List fieldBuffers = vector.getFieldBuffers(); @@ -137,8 +152,92 @@ private void appendNodes( buf.getReferenceManager().retain(); buffers.add(codec.compress(vector.getAllocator(), buf)); } + long coveredChildLength = coveredChildLength(storageVector); for (FieldVector child : vector.getChildrenFromFields()) { - appendNodes(child, nodes, buffers, variadicBufferCounts); + if (coveredChildLength >= 0 + && child.getValueCount() > coveredChildLength + && splitKeepsValues(child)) { + appendRebased(child, (int) coveredChildLength, nodes, buffers, variadicBufferCounts); + } else { + appendNodes(child, nodes, buffers, variadicBufferCounts); + } + } + } + + /** Appends the first {@code length} values of the vector, re-based to start at offset zero. */ + private void appendRebased( + FieldVector vector, + int length, + List nodes, + List buffers, + List variadicBufferCounts) { + TransferPair pair = vector.getTransferPair(vector.getField(), vector.getAllocator()); + try (FieldVector rebased = (FieldVector) pair.getTo()) { + pair.splitAndTransfer(0, length); + appendNodes(rebased, nodes, buffers, variadicBufferCounts); + } + } + + /** + * Whether splitting the vector keeps every value in its slot. A dense union's split does not when + * a type's offsets have gaps, so only the vector types listed here are re-based. + */ + private static boolean splitKeepsValues(FieldVector vector) { + if (vector instanceof BaseFixedWidthVector + || vector instanceof BaseVariableWidthVector + || vector instanceof BaseLargeVariableWidthVector + || vector instanceof NullVector) { + return true; + } + if (vector instanceof ListVector + || vector instanceof LargeListVector + || vector instanceof FixedSizeListVector + || vector instanceof NonNullableStructVector) { + for (FieldVector child : vector.getChildrenFromFields()) { + if (!splitKeepsValues(child)) { + return false; + } + } + return true; + } + return false; + } + + /** The first offset of a variable-width, list or map vector, or zero for any other vector. */ + private static long firstOffset(FieldVector vector) { + if (vector.getValueCount() == 0) { + return 0; + } + if (vector instanceof BaseVariableWidthVector) { + return ((BaseVariableWidthVector) vector).getStartOffset(0); + } + if (vector instanceof BaseLargeVariableWidthVector) { + return ((BaseLargeVariableWidthVector) vector).getStartOffset(0); + } + if (vector instanceof ListVector) { + return ((ListVector) vector).getElementStartIndex(0); + } + if (vector instanceof LargeListVector) { + return ((LargeListVector) vector).getElementStartIndex(0); + } + return 0; + } + + /** + * The number of child elements a list or map vector's offsets cover, which is its last offset, or + * -1 for a vector whose child is not addressed by offsets. + */ + private static long coveredChildLength(FieldVector vector) { + int valueCount = vector.getValueCount(); + if (valueCount == 0) { + return -1; + } + if (vector instanceof ListVector) { + return ((ListVector) vector).getElementEndIndex(valueCount - 1); + } + if (vector instanceof LargeListVector) { + return ((LargeListVector) vector).getElementEndIndex(valueCount - 1); } + return -1; } } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java index 9368b5ce3c..f3fc6faafb 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java @@ -316,7 +316,18 @@ private void setReaderAndWriterIndex() { // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers // in other libraries. According to Arrow spec, we should still output the offset buffer which // is [0]. - offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); + final long requiredOffsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH; + if (offsetBuffer.capacity() < requiredOffsetBufferSize) { + // An empty vector may still hold the allocator's shared empty buffer + ArrowBuf newOffsetBuffer = allocator.buffer(requiredOffsetBufferSize); + newOffsetBuffer.setZero(0, newOffsetBuffer.capacity()); + if (offsetBuffer.capacity() > 0) { + newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity()); + } + offsetBuffer.getReferenceManager().release(); + offsetBuffer = newOffsetBuffer; + } + offsetBuffer.writerIndex(requiredOffsetBufferSize); } /** diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java index 6c3993df63..f2dee09154 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java @@ -274,7 +274,18 @@ private void setReaderAndWriterIndex() { // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers // in other libraries. According to Arrow spec, we should still output the offset buffer which // is [0]. - offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); + final long requiredOffsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH; + if (offsetBuffer.capacity() < requiredOffsetBufferSize) { + // An empty vector may still hold the allocator's shared empty buffer + ArrowBuf newOffsetBuffer = allocator.buffer(requiredOffsetBufferSize); + newOffsetBuffer.setZero(0, newOffsetBuffer.capacity()); + if (offsetBuffer.capacity() > 0) { + newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity()); + } + offsetBuffer.getReferenceManager().release(); + offsetBuffer = newOffsetBuffer; + } + offsetBuffer.writerIndex(requiredOffsetBufferSize); } /** diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java index b6823fbf3c..9f92a23a76 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java @@ -1216,4 +1216,17 @@ public void testSetValueCountAcceptsMaxIntChildValueCount() { } } } + + @Test + public void testGetFieldBuffersOfEmptyVector() { + // An empty vector holds the allocator's shared empty buffer as its offsets until it allocates + try (LargeListVector vector = LargeListVector.empty("list", allocator)) { + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + ArrowBuf offsets = vector.getFieldBuffers().get(1); + assertTrue(offsets.capacity() >= LargeListVector.OFFSET_WIDTH); + assertEquals(LargeListVector.OFFSET_WIDTH, offsets.readableBytes()); + assertEquals(0L, offsets.getLong(0)); + assertEquals(0, allocator.getEmpty().writerIndex()); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java index 0c90b32abc..548e05ace5 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java @@ -1406,4 +1406,27 @@ private void writeIntValues(UnionListWriter writer, int[] values) { } writer.endList(); } + + @Test + public void testSplitAndTransferZeroLength() { + try (ListVector listVector = ListVector.empty("sourceVector", allocator)) { + listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + UnionListWriter listWriter = listVector.getWriter(); + listWriter.allocate(); + listWriter.setPosition(0); + listWriter.startList(); + listWriter.integer().writeInt(1); + listWriter.endList(); + listVector.setValueCount(1); + + TransferPair pair = listVector.getTransferPair(allocator); + pair.splitAndTransfer(0, 0); + try (ListVector to = (ListVector) pair.getTo()) { + ArrowBuf offsets = to.getFieldBuffers().get(1); + assertTrue(offsets.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH); + assertEquals(0, offsets.getInt(0)); + assertEquals(0, allocator.getEmpty().writerIndex()); + } + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 2f520f3882..99bb178dd4 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -31,6 +31,7 @@ import java.util.UUID; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.impl.UnionMapReader; @@ -1662,4 +1663,35 @@ public void testMapWithUuidKeyAndListUuidValue() throws Exception { assertEquals(value2c, actualValue2c); } } + + @Test + public void testSplitAndTransferZeroLength() { + try (MapVector mapVector = MapVector.empty("sourceVector", allocator, false)) { + FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null, null); + AddOrGetResult addResult = mapVector.addOrGetVector(type); + FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(), null, null); + FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType()); + addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType, BigIntVector.class); + addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType, Float8Vector.class); + UnionMapWriter mapWriter = mapVector.getWriter(); + mapWriter.allocate(); + mapWriter.setPosition(0); + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().bigInt().writeBigInt(10); + mapWriter.value().float8().writeFloat8(1.0); + mapWriter.endEntry(); + mapWriter.endMap(); + mapVector.setValueCount(1); + + TransferPair pair = mapVector.getTransferPair(allocator); + pair.splitAndTransfer(0, 0); + try (MapVector to = (MapVector) pair.getTo()) { + ArrowBuf offsets = to.getFieldBuffers().get(1); + assertTrue(offsets.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH); + assertEquals(0, offsets.getInt(0)); + assertEquals(0, allocator.getEmpty().writerIndex()); + } + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java new file mode 100644 index 0000000000..e4ba77594d --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java @@ -0,0 +1,433 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.complex.DenseUnionVector; +import org.apache.arrow.vector.complex.LargeListVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.ipc.message.ArrowBuffer; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +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.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * A vector imported through the C Data interface from a sliced array keeps the offsets of the + * slice, which start past zero, and the bytes or child elements before them. {@link VectorUnloader} + * must write only what the offsets cover. + */ +public class TestVectorUnloadNonZeroStartOffset { + + private BufferAllocator allocator; + + @BeforeEach + public void init() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void terminate() { + allocator.close(); + } + + private ArrowBuf ints(int... values) { + ArrowBuf buf = allocator.buffer(4L * values.length); + for (int i = 0; i < values.length; i++) { + buf.setInt(4L * i, values[i]); + } + buf.writerIndex(4L * values.length); + return buf; + } + + private ArrowBuf longs(long... values) { + ArrowBuf buf = allocator.buffer(8L * values.length); + for (int i = 0; i < values.length; i++) { + buf.setLong(8L * i, values[i]); + } + buf.writerIndex(8L * values.length); + return buf; + } + + private ArrowBuf bytes(String values) { + byte[] data = values.getBytes(StandardCharsets.UTF_8); + ArrowBuf buf = allocator.buffer(data.length); + buf.setBytes(0, data); + buf.writerIndex(data.length); + return buf; + } + + private ArrowBuf allValid() { + return longs(-1L); + } + + private ArrowBuf typeIds(int... values) { + ArrowBuf buf = allocator.buffer(values.length); + for (int i = 0; i < values.length; i++) { + buf.setByte(i, values[i]); + } + buf.writerIndex(values.length); + return buf; + } + + private static void load(FieldVector vector, int length, ArrowBuf... buffers) { + vector.loadFieldBuffers(new ArrowFieldNode(length, 0), Arrays.asList(buffers)); + for (ArrowBuf buf : buffers) { + buf.close(); + } + } + + /** Rows 2..3 of ["aa", "bb", "cc", "dd"]: the data buffer still holds "aabb" in front. */ + private VarCharVector varCharSlice() { + VarCharVector vector = new VarCharVector("s", allocator); + load(vector, 2, allValid(), ints(4, 6, 8), bytes("aabbccdd")); + return vector; + } + + private LargeVarCharVector largeVarCharSlice() { + LargeVarCharVector vector = new LargeVarCharVector("s", allocator); + load(vector, 2, allValid(), longs(4, 6, 8), bytes("aabbccdd")); + return vector; + } + + private void loadEightInts(FieldVector child) { + load(child, 8, allValid(), ints(0, 1, 2, 3, 4, 5, 6, 7)); + } + + /** Two rows of [[0, 1], [2, 3], [4, 5], [6, 7]] whose child still holds all eight elements. */ + private ListVector listSlice(int... offsets) { + ListVector vector = ListVector.empty("l", allocator); + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + load(vector, 2, allValid(), ints(offsets)); + loadEightInts(vector.getDataVector()); + return vector; + } + + private LargeListVector largeListSlice(long... offsets) { + LargeListVector vector = LargeListVector.empty("l", allocator); + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + load(vector, 2, allValid(), longs(offsets)); + loadEightInts(vector.getDataVector()); + return vector; + } + + /** Rows 1..2 of [{0: 10, 1: 11}, {2: 12, 3: 13}, {4: 14, 5: 15}, {6: 16, 7: 17}]. */ + private MapVector mapSlice() { + MapVector vector = MapVector.empty("m", allocator, false); + Field key = new Field("key", FieldType.notNullable(MinorType.INT.getType()), null); + Field value = new Field("value", FieldType.nullable(MinorType.INT.getType()), null); + Field entries = + new Field( + MapVector.DATA_VECTOR_NAME, + FieldType.notNullable(ArrowType.Struct.INSTANCE), + Arrays.asList(key, value)); + vector.initializeChildrenFromFields(Collections.singletonList(entries)); + load(vector, 2, allValid(), ints(2, 4, 6)); + StructVector struct = (StructVector) vector.getDataVector(); + load(struct, 8, allValid()); + loadEightInts(struct.getChild("key")); + load(struct.getChild("value"), 8, allValid(), ints(10, 11, 12, 13, 14, 15, 16, 17)); + return vector; + } + + /** + * The first slice of [["aa"], ["bb"], ["cc"]]: list offsets start at zero, its child's do not. + */ + private ListVector listOfVarCharSlice() { + ListVector vector = ListVector.empty("l", allocator); + vector.addOrGetVector(FieldType.nullable(MinorType.VARCHAR.getType())); + load(vector, 2, allValid(), ints(0, 1, 2)); + load(vector.getDataVector(), 2, allValid(), ints(2, 4, 6), bytes("aabbcc")); + return vector; + } + + /** Two empty rows over an inner list that still holds [[1], [2]]. */ + private ListVector emptyRowsOverNestedList() { + ListVector vector = ListVector.empty("l", allocator); + vector.addOrGetVector(FieldType.nullable(ArrowType.List.INSTANCE)); + ListVector inner = (ListVector) vector.getDataVector(); + inner.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + load(vector, 2, allValid(), ints(0, 0, 0)); + load(inner, 2, allValid(), ints(0, 1, 2)); + load(inner.getDataVector(), 2, allValid(), ints(1, 2)); + return vector; + } + + /** Two empty rows over an inner list that still holds two large lists, [[1]] and [[2]]. */ + private ListVector emptyRowsOverNestedLargeList() { + ListVector vector = ListVector.empty("l", allocator); + vector.addOrGetVector(FieldType.nullable(ArrowType.List.INSTANCE)); + ListVector inner = (ListVector) vector.getDataVector(); + inner.addOrGetVector(FieldType.nullable(ArrowType.LargeList.INSTANCE)); + LargeListVector largeList = (LargeListVector) inner.getDataVector(); + largeList.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + load(vector, 2, allValid(), ints(0, 0, 0)); + load(inner, 2, allValid(), ints(0, 1, 2)); + load(largeList, 2, allValid(), longs(0, 1, 2)); + load(largeList.getDataVector(), 2, allValid(), ints(1, 2)); + return vector; + } + + /** + * One row over a dense union whose int child skips an element: union slots 0..3 hold child + * elements 0, 1, 3 and 4, that is 10, 11, 13 and 14. + */ + private ListVector listOfDenseUnion(int... offsets) { + Field union = + new Field( + "u", + FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {0})), + Collections.singletonList(Field.nullable("i", MinorType.INT.getType()))); + Field list = + new Field( + "l", FieldType.nullable(ArrowType.List.INSTANCE), Collections.singletonList(union)); + ListVector vector = (ListVector) list.createVector(allocator); + load(vector, 1, allValid(), ints(offsets)); + DenseUnionVector unionVector = (DenseUnionVector) vector.getDataVector(); + load(unionVector, 4, typeIds(0, 0, 0, 0), ints(0, 1, 3, 4)); + load(unionVector.getChildrenFromFields().get(0), 5, allValid(), ints(10, 11, 12, 13, 14)); + return vector; + } + + private static List sizes(ArrowRecordBatch batch) { + return batch.getBuffersLayout().stream().map(ArrowBuffer::getSize).collect(Collectors.toList()); + } + + private static List lengths(ArrowRecordBatch batch) { + return batch.getNodes().stream().map(ArrowFieldNode::getLength).collect(Collectors.toList()); + } + + private VectorSchemaRoot reload(VectorSchemaRoot root, ArrowRecordBatch batch) { + VectorSchemaRoot loaded = VectorSchemaRoot.create(root.getSchema(), allocator); + new VectorLoader(loaded).load(batch); + return loaded; + } + + @Test + public void varCharIsWrittenFromItsFirstOffset() { + try (VarCharVector vector = varCharSlice(); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(1L, 12L, 4L), sizes(batch)); + assertEquals(0, batch.getBuffers().get(1).getInt(0)); + assertEquals(4, batch.getBuffers().get(1).getInt(8)); + VarCharVector read = (VarCharVector) loaded.getVector(0); + assertEquals("cc", read.getObject(0).toString()); + assertEquals("dd", read.getObject(1).toString()); + } + } + + @Test + public void largeVarCharIsWrittenFromItsFirstOffset() { + try (LargeVarCharVector vector = largeVarCharSlice(); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(1L, 24L, 4L), sizes(batch)); + assertEquals(0L, batch.getBuffers().get(1).getLong(0)); + LargeVarCharVector read = (LargeVarCharVector) loaded.getVector(0); + assertEquals("cc", read.getObject(0).toString()); + assertEquals("dd", read.getObject(1).toString()); + } + } + + @Test + public void listChildIsWrittenBetweenItsOffsets() { + try (ListVector vector = listSlice(2, 4, 6); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 4), lengths(batch)); + assertEquals(Arrays.asList(1L, 12L, 1L, 16L), sizes(batch)); + assertEquals(0, batch.getBuffers().get(1).getInt(0)); + ListVector read = (ListVector) loaded.getVector(0); + assertEquals(Arrays.asList(2, 3), read.getObject(0)); + assertEquals(Arrays.asList(4, 5), read.getObject(1)); + } + } + + @Test + public void largeListChildIsWrittenBetweenItsOffsets() { + try (LargeListVector vector = largeListSlice(2, 4, 6); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 4), lengths(batch)); + assertEquals(Arrays.asList(1L, 24L, 1L, 16L), sizes(batch)); + assertEquals(0L, batch.getBuffers().get(1).getLong(0)); + LargeListVector read = (LargeListVector) loaded.getVector(0); + assertEquals(Arrays.asList(2, 3), read.getObject(0)); + assertEquals(Arrays.asList(4, 5), read.getObject(1)); + } + } + + @Test + public void mapEntriesAreWrittenBetweenItsOffsets() { + try (MapVector vector = mapSlice(); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 4, 4, 4), lengths(batch)); + assertEquals(Arrays.asList(1L, 12L, 1L, 1L, 16L, 1L, 16L), sizes(batch)); + MapVector read = (MapVector) loaded.getVector(0); + assertEquals(vector.getObject(0), read.getObject(0)); + assertEquals(vector.getObject(1), read.getObject(1)); + StructVector entries = (StructVector) read.getDataVector(); + assertEquals(2, ((IntVector) entries.getChild("key")).get(0)); + assertEquals(15, ((IntVector) entries.getChild("value")).get(3)); + } + } + + @Test + public void listChildIsWrittenUpToItsLastOffset() { + try (ListVector vector = listSlice(0, 2, 4); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 4), lengths(batch)); + assertEquals(Arrays.asList(1L, 12L, 1L, 16L), sizes(batch)); + // the list's own offsets start at zero, so they are written as they are + assertEquals( + vector.getOffsetBuffer().memoryAddress(), batch.getBuffers().get(1).memoryAddress()); + ListVector read = (ListVector) loaded.getVector(0); + assertEquals(Arrays.asList(0, 1), read.getObject(0)); + assertEquals(Arrays.asList(2, 3), read.getObject(1)); + } + } + + @Test + public void nestedVarCharIsWrittenFromItsFirstOffset() { + try (ListVector vector = listOfVarCharSlice(); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 2), lengths(batch)); + assertEquals(Arrays.asList(1L, 12L, 1L, 12L, 4L), sizes(batch)); + assertEquals(0, batch.getBuffers().get(3).getInt(0)); + ListVector read = (ListVector) loaded.getVector(0); + assertEquals(vector.getObject(0), read.getObject(0)); + assertEquals(vector.getObject(1), read.getObject(1)); + assertEquals("bb", ((VarCharVector) read.getDataVector()).getObject(0).toString()); + } + } + + @Test + public void vectorWithZeroStartOffsetIsWrittenWithItsOwnBuffers() { + try (VarCharVector vector = new VarCharVector("s", allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + load(vector, 2, allValid(), ints(0, 2, 4), bytes("aabb")); + try (ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch()) { + assertEquals(Arrays.asList(1L, 12L, 4L), sizes(batch)); + assertEquals( + vector.getOffsetBuffer().memoryAddress(), batch.getBuffers().get(1).memoryAddress()); + assertEquals( + vector.getDataBuffer().memoryAddress(), batch.getBuffers().get(2).memoryAddress()); + } + } + } + + @Test + public void emptyCoveredRangeOverNestedListKeepsItsFirstOffset() { + try (ListVector vector = emptyRowsOverNestedList(); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 0, 0), lengths(batch)); + assertEquals(Arrays.asList(1L, 12L, 0L, 4L, 0L, 0L), sizes(batch)); + assertEquals(0, batch.getBuffers().get(3).getInt(0)); + ListVector read = (ListVector) loaded.getVector(0); + assertEquals(Collections.emptyList(), read.getObject(0)); + assertEquals(Collections.emptyList(), read.getObject(1)); + } + } + + @Test + public void emptyCoveredRangeOverNestedLargeListKeepsItsFirstOffset() { + try (ListVector vector = emptyRowsOverNestedLargeList(); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(2, 0, 0, 0), lengths(batch)); + assertEquals(Arrays.asList(1L, 12L, 0L, 4L, 0L, 8L, 0L, 0L), sizes(batch)); + assertEquals(0L, batch.getBuffers().get(5).getLong(0)); + ListVector read = (ListVector) loaded.getVector(0); + assertEquals(Collections.emptyList(), read.getObject(0)); + assertEquals(Collections.emptyList(), read.getObject(1)); + } + } + + @Test + public void denseUnionChildIsWrittenUntrimmed() { + try (ListVector vector = listOfDenseUnion(0, 3); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(1, 4, 5), lengths(batch)); + assertEquals(Arrays.asList(10, 11, 13), ((ListVector) loaded.getVector(0)).getObject(0)); + } + } + + @Test + public void listOverDenseUnionIsWrittenWithItsOwnOffsets() { + try (ListVector vector = listOfDenseUnion(1, 3); + VectorSchemaRoot root = VectorSchemaRoot.of(vector); + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + VectorSchemaRoot loaded = reload(root, batch)) { + assertEquals(Arrays.asList(1, 4, 5), lengths(batch)); + assertEquals(Arrays.asList(11, 13), ((ListVector) loaded.getVector(0)).getObject(0)); + } + } + + @Test + public void failedSplitDoesNotLeak() { + try (ListVector vector = listSlice(2, 4, 9); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + assertThrows(IllegalArgumentException.class, () -> new VectorUnloader(root).getRecordBatch()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void failedSplitReleasesTheBuffersOfEarlierVectors() { + try (IntVector first = new IntVector("i", allocator); + ListVector malformed = listSlice(2, 4, 9); + VectorSchemaRoot root = VectorSchemaRoot.of(first, malformed)) { + load(first, 2, allValid(), ints(1, 2)); + assertThrows(IllegalArgumentException.class, () -> new VectorUnloader(root).getRecordBatch()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } +}