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 92dd3eaef7..fdd917a920 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 @@ -305,10 +305,12 @@ public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { + final long requiredOffsetBufferCapacity = (long) (valueCount + 1) * OFFSET_WIDTH; validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); + ensureEmptyOffsetBufferCapacity(requiredOffsetBufferCapacity); } else { validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } @@ -319,6 +321,19 @@ private void setReaderAndWriterIndex() { offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); } + private void ensureEmptyOffsetBufferCapacity(long requiredCapacity) { + if (offsetBuffer.capacity() >= requiredCapacity) { + return; + } + long previousOffsetAllocationSizeInBytes = offsetAllocationSizeInBytes; + ArrowBuf oldOffsetBuffer = offsetBuffer; + offsetBuffer = allocateOffsetBuffer(requiredCapacity); + offsetBuffer.setBytes( + 0, oldOffsetBuffer, 0, Math.min(oldOffsetBuffer.capacity(), requiredCapacity)); + offsetAllocationSizeInBytes = previousOffsetAllocationSizeInBytes; + oldOffsetBuffer.getReferenceManager().release(); + } + /** * Get the inner vectors. * @@ -674,24 +689,30 @@ public void splitAndTransfer(int startIndex, int length) { startIndex, length, valueCount); - final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH); - final long sliceLength = - offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint; to.clear(); - to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH); - /* splitAndTransfer offset buffer */ - for (int i = 0; i < length + 1; i++) { - final long relativeOffset = - offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint; - to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset); + if (length > 0) { + final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH); + final long sliceLength = + offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint; + to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH); + /* splitAndTransfer offset buffer */ + for (int i = 0; i < length + 1; i++) { + final long relativeOffset = + offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint; + to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset); + } + /* splitAndTransfer validity buffer */ + splitAndTransferValidityBuffer(startIndex, length, to); + /* splitAndTransfer data buffer */ + dataTransferPair.splitAndTransfer( + checkedCastToInt(startPoint), checkedCastToInt(sliceLength)); + to.lastSet = length - 1; + to.setValueCount(length); + } else { + to.ensureEmptyOffsetBufferCapacity(OFFSET_WIDTH); + dataTransferPair.splitAndTransfer(0, 0); + to.setValueCount(0); } - /* splitAndTransfer validity buffer */ - splitAndTransferValidityBuffer(startIndex, length, to); - /* splitAndTransfer data buffer */ - dataTransferPair.splitAndTransfer( - checkedCastToInt(startPoint), checkedCastToInt(sliceLength)); - to.lastSet = length - 1; - to.setValueCount(length); } @Override 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..2802cd7577 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 @@ -263,10 +263,12 @@ public void exportCDataBuffers(List buffers, ArrowBuf buffersPtr, long /** Set the reader and writer indexes for the inner buffers. */ private void setReaderAndWriterIndex() { + final long requiredOffsetBufferCapacity = (long) (valueCount + 1) * OFFSET_WIDTH; validityBuffer.readerIndex(0); offsetBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); + ensureEmptyOffsetBufferCapacity(requiredOffsetBufferCapacity); } else { validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } @@ -277,6 +279,19 @@ private void setReaderAndWriterIndex() { offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); } + private void ensureEmptyOffsetBufferCapacity(long requiredCapacity) { + if (offsetBuffer.capacity() >= requiredCapacity) { + return; + } + long previousOffsetAllocationSizeInBytes = offsetAllocationSizeInBytes; + ArrowBuf oldOffsetBuffer = offsetBuffer; + offsetBuffer = allocateOffsetBuffer(requiredCapacity); + offsetBuffer.setBytes( + 0, oldOffsetBuffer, 0, Math.min(oldOffsetBuffer.capacity(), requiredCapacity)); + offsetAllocationSizeInBytes = previousOffsetAllocationSizeInBytes; + oldOffsetBuffer.getReferenceManager().release(); + } + /** * Get the inner vectors. * @@ -572,6 +587,10 @@ public void splitAndTransfer(int startIndex, int length) { dataTransferPair.splitAndTransfer(startPoint, sliceLength); to.lastSet = length - 1; to.setValueCount(length); + } else { + to.ensureEmptyOffsetBufferCapacity(OFFSET_WIDTH); + dataTransferPair.splitAndTransfer(0, 0); + to.setValueCount(0); } } 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 bf9bba9c78..e91dbabbfb 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java @@ -962,6 +962,91 @@ public void testGetBufferSizeFor() { } } + @Test + public void testEmptyLargeListOffsetBuffer() { + try (LargeListVector list = LargeListVector.empty("list", allocator)) { + list.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + list.allocateNew(); + list.setValueCount(0); + + assertEmptyLargeListOffsetBuffer(list); + } + } + + @Test + public void testUnallocatedEmptyLargeListOffsetBuffer() { + try (LargeListVector list = LargeListVector.empty("list", allocator)) { + list.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + list.setValueCount(0); + + assertEmptyLargeListOffsetBuffer(list); + } + } + + @Test + public void testSplitAndTransferEmptyLargeListOffsetBuffer() { + try (LargeListVector source = LargeListVector.empty("source", allocator); + LargeListVector target = LargeListVector.empty("target", allocator)) { + source.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + target.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + source.allocateNew(); + source.setValueCount(0); + + TransferPair transferPair = source.makeTransferPair(target); + transferPair.splitAndTransfer(0, 0); + + assertEmptyLargeListOffsetBuffer(target); + } + } + + @Test + public void testSplitAndTransferEmptyLargeListAllocatesOffsetBuffer() { + try (LargeListVector fromVector = LargeListVector.empty("fromVector", allocator); + LargeListVector toVector = LargeListVector.empty("toVector", allocator)) { + fromVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + fromVector.allocateNew(); + fromVector.setValueCount(0); + + TransferPair transferPair = fromVector.makeTransferPair(toVector); + transferPair.splitAndTransfer(0, 0); + + assertAllocatedEmptyLargeListOffsetBuffer(toVector); + } + } + + @Test + public void testSplitAndTransferEmptyNestedLargeListAllocatesOffsetBuffers() { + try (LargeListVector fromVector = LargeListVector.empty("fromVector", allocator); + LargeListVector toVector = LargeListVector.empty("toVector", allocator)) { + fromVector.addOrGetVector(FieldType.nullable(MinorType.LARGELIST.getType())); + LargeListVector childVector = (LargeListVector) fromVector.getDataVector(); + childVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + fromVector.allocateNew(); + fromVector.setValueCount(0); + + TransferPair transferPair = fromVector.makeTransferPair(toVector); + transferPair.splitAndTransfer(0, 0); + + assertAllocatedEmptyLargeListOffsetBuffer(toVector); + assertAllocatedEmptyLargeListOffsetBuffer((LargeListVector) toVector.getDataVector()); + } + } + + private ArrowBuf assertEmptyLargeListOffsetBuffer(LargeListVector list) { + List buffers = list.getFieldBuffers(); + ArrowBuf offsetBuffer = buffers.get(1); + assertEquals(LargeListVector.OFFSET_WIDTH, offsetBuffer.readableBytes()); + assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH); + assertEquals(0L, offsetBuffer.getLong(0)); + return offsetBuffer; + } + + private void assertAllocatedEmptyLargeListOffsetBuffer(LargeListVector list) { + ArrowBuf offsetBuffer = list.getOffsetBuffer(); + assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH); + assertEquals(0L, offsetBuffer.getLong(0)); + } + @Test public void testIsEmpty() { try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { @@ -1101,7 +1186,7 @@ public void testCopyValueSafeForExtensionType() throws Exception { } @Test - public void testEmptyLargeListOffsetBuffer() { + public void testEmptyLargeListOffsetBufferSize() { // Test that LargeListVector has correct readableBytes after allocation. // According to Arrow spec, offset buffer must have N+1 entries. // Even when N=0, it should contain [0]. 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..60b675a1b5 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java @@ -1142,6 +1142,91 @@ public void testGetBufferSizeFor() { } } + @Test + public void testEmptyListOffsetBuffer() { + try (ListVector list = ListVector.empty("list", allocator)) { + list.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + list.allocateNew(); + list.setValueCount(0); + + assertEmptyListOffsetBuffer(list); + } + } + + @Test + public void testUnallocatedEmptyListOffsetBuffer() { + try (ListVector list = ListVector.empty("list", allocator)) { + list.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + list.setValueCount(0); + + assertEmptyListOffsetBuffer(list); + } + } + + @Test + public void testSplitAndTransferEmptyListOffsetBuffer() { + try (ListVector source = ListVector.empty("source", allocator); + ListVector target = ListVector.empty("target", allocator)) { + source.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + target.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + source.allocateNew(); + source.setValueCount(0); + + TransferPair transferPair = source.makeTransferPair(target); + transferPair.splitAndTransfer(0, 0); + + assertEmptyListOffsetBuffer(target); + } + } + + @Test + public void testSplitAndTransferEmptyListAllocatesOffsetBuffer() { + try (ListVector fromVector = ListVector.empty("fromVector", allocator); + ListVector toVector = ListVector.empty("toVector", allocator)) { + fromVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + fromVector.allocateNew(); + fromVector.setValueCount(0); + + TransferPair transferPair = fromVector.makeTransferPair(toVector); + transferPair.splitAndTransfer(0, 0); + + assertAllocatedEmptyListOffsetBuffer(toVector); + } + } + + @Test + public void testSplitAndTransferEmptyNestedListAllocatesOffsetBuffers() { + try (ListVector fromVector = ListVector.empty("fromVector", allocator); + ListVector toVector = ListVector.empty("toVector", allocator)) { + fromVector.addOrGetVector(FieldType.nullable(MinorType.LIST.getType())); + ListVector childVector = (ListVector) fromVector.getDataVector(); + childVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + fromVector.allocateNew(); + fromVector.setValueCount(0); + + TransferPair transferPair = fromVector.makeTransferPair(toVector); + transferPair.splitAndTransfer(0, 0); + + assertAllocatedEmptyListOffsetBuffer(toVector); + assertAllocatedEmptyListOffsetBuffer((ListVector) toVector.getDataVector()); + } + } + + private ArrowBuf assertEmptyListOffsetBuffer(ListVector list) { + List buffers = list.getFieldBuffers(); + ArrowBuf offsetBuffer = buffers.get(1); + assertEquals(BaseRepeatedValueVector.OFFSET_WIDTH, offsetBuffer.readableBytes()); + assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH); + assertEquals(0, offsetBuffer.getInt(0)); + return offsetBuffer; + } + + private void assertAllocatedEmptyListOffsetBuffer(ListVector list) { + ArrowBuf offsetBuffer = list.getOffsetBuffer(); + assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH); + assertEquals(0, offsetBuffer.getInt(0)); + } + @Test public void testIsEmpty() { try (final ListVector vector = ListVector.empty("list", allocator)) { @@ -1380,7 +1465,7 @@ public void testCopyValueSafeForExtensionType() throws Exception { } @Test - public void testEmptyListOffsetBuffer() { + public void testEmptyListOffsetBufferSize() { // Test that ListVector has correct readableBytes after allocation. // According to Arrow spec, offset buffer must have N+1 entries. // Even when N=0, it should contain [0]. diff --git a/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java b/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java index adf4eba10c..a51c3e74fa 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java @@ -28,10 +28,13 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.complex.BaseRepeatedValueVector; 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; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; @@ -111,13 +114,16 @@ private void populateDenseUnionVector(final DenseUnionVector vector, int valueCo @Test public void testWithEmptyVector() { // MapVector use TransferImpl from ListVector - ListVector listVector = ListVector.empty("", allocator); - TransferPair transferPair = listVector.getTransferPair(allocator); - transferPair.splitAndTransfer(0, 0); - assertEquals(0, transferPair.getTo().getValueCount()); + try (ListVector listVector = ListVector.empty("", allocator)) { + TransferPair transferPair = listVector.getTransferPair(allocator); + try (ValueVector toVector = transferPair.getTo()) { + transferPair.splitAndTransfer(0, 0); + assertEquals(0, toVector.getValueCount()); + } + } // BaseFixedWidthVector IntVector intVector = new IntVector("", allocator); - transferPair = intVector.getTransferPair(allocator); + TransferPair transferPair = intVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); // BaseVariableWidthVector @@ -911,6 +917,11 @@ public void testListVectorZeroStartIndexAndLength() { tp.splitAndTransfer(0, 0); assertEquals(valueCount, newListVector.getValueCount()); + List buffers = newListVector.getFieldBuffers(); + ArrowBuf offsetBuffer = buffers.get(1); + assertEquals(BaseRepeatedValueVector.OFFSET_WIDTH, offsetBuffer.readableBytes()); + assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH); + assertEquals(0, offsetBuffer.getInt(0)); newListVector.clear(); } @@ -935,6 +946,29 @@ public void testLargeListViewVectorZeroStartIndexAndLength() { } } + @Test + public void testLargeListVectorZeroStartIndexAndLength() { + try (final LargeListVector listVector = LargeListVector.empty("largelist", allocator); + final LargeListVector newListVector = LargeListVector.empty("newList", allocator)) { + + listVector.allocateNew(); + final int valueCount = 0; + listVector.setValueCount(valueCount); + + final TransferPair tp = listVector.makeTransferPair(newListVector); + + tp.splitAndTransfer(0, 0); + assertEquals(valueCount, newListVector.getValueCount()); + List buffers = newListVector.getFieldBuffers(); + ArrowBuf offsetBuffer = buffers.get(1); + assertEquals(LargeListVector.OFFSET_WIDTH, offsetBuffer.readableBytes()); + assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH); + assertEquals(0L, offsetBuffer.getLong(0)); + + newListVector.clear(); + } + } + @Test public void testStructVectorZeroStartIndexAndLength() { Map metadata = new HashMap<>();