From f78cd309ee94b340ac6c81649863a92b20a30d9a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 10:42:58 -0600 Subject: [PATCH 1/3] fix(vector): write only the values a vector's offsets cover when unloading VectorUnloader wrote a variable-width vector's data buffer from byte 0 up to its last offset, and a list's child vector from element 0. A vector imported through the C Data interface from a sliced array has offsets that start past zero and keeps the values of the rows before the slice in front, so every such batch was serialized with bytes no row references, and the k-th slice of an array carried the k-1 slices before it. When a variable-width, list or map vector's offsets start past zero, re-base it through splitAndTransfer(0, valueCount) and unload the result, closing it once its buffers are retained by the record batch. The re-based vector only copies the offsets and shares the other buffers. When a list or map starts at zero but its child runs past the last offset, as the first slice of an array does, keep its own buffers and re-base only the child to the covered length. Vectors whose offsets start at zero are written with their own buffers as before. Closes #1424 --- .../apache/arrow/vector/VectorUnloader.java | 71 ++++- .../TestVectorUnloadNonZeroStartOffset.java | 301 ++++++++++++++++++ 2 files changed, 371 insertions(+), 1 deletion(-) create mode 100644 vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java 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..838b49112b 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,14 @@ import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.complex.LargeListVector; +import org.apache.arrow.vector.complex.ListVector; 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 +77,11 @@ 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. */ public ArrowRecordBatch getRecordBatch() { List nodes = new ArrayList<>(); @@ -111,6 +119,10 @@ private void appendNodes( List buffers, List variadicBufferCounts) { FieldVector storageVector = ExtensionTypeVector.getStorageVector(vector); + if (firstOffset(storageVector) != 0) { + appendRebased(vector, vector.getValueCount(), nodes, buffers, variadicBufferCounts); + return; + } nodes.add( new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1)); List fieldBuffers = vector.getFieldBuffers(); @@ -137,8 +149,65 @@ 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) { + 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()); + pair.splitAndTransfer(0, length); + try (FieldVector rebased = (FieldVector) pair.getTo()) { + appendNodes(rebased, nodes, buffers, variadicBufferCounts); + } + } + + /** 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/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..cc277bee86 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java @@ -0,0 +1,301 @@ +/* + * 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 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.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.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 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; + } + + 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()); + } + } + } +} From d811f92a100daac457b479b705def67be17e5cac Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 1 Nov 2026 08:34:15 -0700 Subject: [PATCH 2/3] fix(vector): handle unions, empty lists and split failures when re-basing - Re-base or trim only vectors whose split keeps every value in its slot: fixed-width, variable-width and null vectors, and lists, large lists, fixed-size lists, maps and structs whose children qualify. A dense union's split assumes each type's offsets are contiguous, so vectors with a union, view, run-end encoded or extension vector at or below them are written as before. - Reallocate a ListVector's offsets buffer in getFieldBuffers() when it cannot hold [0], as BaseVariableWidthVector does. A zero-length split leaves the allocator's shared empty buffer in place, so an empty covered range over a nested list wrote its offsets with no bytes, after setting the writer index of the shared buffer. - Run splitAndTransfer inside the try-with-resources that closes its target, so a split that fails on malformed offsets does not leak it. --- .../apache/arrow/vector/VectorUnloader.java | 38 ++++++- .../arrow/vector/complex/ListVector.java | 13 ++- .../apache/arrow/vector/TestListVector.java | 23 ++++ .../apache/arrow/vector/TestMapVector.java | 32 ++++++ .../TestVectorUnloadNonZeroStartOffset.java | 102 ++++++++++++++++++ 5 files changed, 203 insertions(+), 5 deletions(-) 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 838b49112b..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,8 +20,10 @@ 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; @@ -81,7 +83,8 @@ public VectorUnloader( *

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. + * 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<>(); @@ -119,7 +122,7 @@ private void appendNodes( List buffers, List variadicBufferCounts) { FieldVector storageVector = ExtensionTypeVector.getStorageVector(vector); - if (firstOffset(storageVector) != 0) { + if (firstOffset(storageVector) != 0 && splitKeepsValues(vector)) { appendRebased(vector, vector.getValueCount(), nodes, buffers, variadicBufferCounts); return; } @@ -151,7 +154,9 @@ private void appendNodes( } long coveredChildLength = coveredChildLength(storageVector); for (FieldVector child : vector.getChildrenFromFields()) { - if (coveredChildLength >= 0 && child.getValueCount() > coveredChildLength) { + if (coveredChildLength >= 0 + && child.getValueCount() > coveredChildLength + && splitKeepsValues(child)) { appendRebased(child, (int) coveredChildLength, nodes, buffers, variadicBufferCounts); } else { appendNodes(child, nodes, buffers, variadicBufferCounts); @@ -167,12 +172,37 @@ private void appendRebased( List buffers, List variadicBufferCounts) { TransferPair pair = vector.getTransferPair(vector.getField(), vector.getAllocator()); - pair.splitAndTransfer(0, length); 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) { 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/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 index cc277bee86..522af57db4 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java @@ -17,6 +17,7 @@ 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; @@ -26,6 +27,7 @@ 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; @@ -34,6 +36,7 @@ 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; @@ -90,6 +93,15 @@ 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) { @@ -161,6 +173,39 @@ private ListVector listOfVarCharSlice() { 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; + } + + /** + * 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()); } @@ -298,4 +343,61 @@ public void vectorWithZeroStartOffsetIsWrittenWithItsOwnBuffers() { } } } + + @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 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()); + } } From bc8bd9c73aa7c8da05415461fd115cb1f79fd623 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Mon, 2 Nov 2026 05:57:41 +0100 Subject: [PATCH 3/3] fix(vector): reallocate LargeListVector offsets that cannot hold [0] A zero-length split of a ListVector only clears its target, so a LargeListVector below it keeps the allocator's shared empty buffer as its offsets. getFieldBuffers() then set the writer index of the shared buffer to 8: the LargeList offsets were written with no bytes when a later vector reset it, and when none did, every other buffer still on the shared empty buffer was sized 8 bytes and writing the batch threw IndexOutOfBoundsException. Reallocate the offsets buffer in getFieldBuffers() when it cannot hold [0], as ListVector now does. --- .../arrow/vector/complex/LargeListVector.java | 13 +++++++- .../arrow/vector/TestLargeListVector.java | 13 ++++++++ .../TestVectorUnloadNonZeroStartOffset.java | 30 +++++++++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) 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/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/TestVectorUnloadNonZeroStartOffset.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java index 522af57db4..e4ba77594d 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java @@ -185,6 +185,21 @@ private ListVector emptyRowsOverNestedList() { 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. @@ -359,6 +374,21 @@ public void emptyCoveredRangeOverNestedListKeepsItsFirstOffset() { } } + @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);