diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java b/java/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java index 393df96b2969e..0a043b51067ef 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java @@ -205,14 +205,12 @@ public void setSafe(int index, NullableViewVarBinaryHolder holder) { */ @Override public TransferPair getTransferPair(String ref, BufferAllocator allocator) { - // TODO: https://github.com/apache/arrow/issues/40932 - throw new UnsupportedOperationException("Unsupported operation"); + return new TransferImpl(ref, allocator); } @Override public TransferPair getTransferPair(Field field, BufferAllocator allocator) { - // TODO: https://github.com/apache/arrow/issues/40932 - throw new UnsupportedOperationException("Unsupported operation"); + return new TransferImpl(field, allocator); } /** @@ -223,7 +221,42 @@ public TransferPair getTransferPair(Field field, BufferAllocator allocator) { */ @Override public TransferPair makeTransferPair(ValueVector to) { - // TODO: https://github.com/apache/arrow/issues/40932 - throw new UnsupportedOperationException("Unsupported operation"); + return new TransferImpl((ViewVarBinaryVector) to); + } + + private class TransferImpl implements TransferPair { + ViewVarBinaryVector to; + + public TransferImpl(String ref, BufferAllocator allocator) { + to = new ViewVarBinaryVector(ref, field.getFieldType(), allocator); + } + + public TransferImpl(Field field, BufferAllocator allocator) { + to = new ViewVarBinaryVector(field, allocator); + } + + public TransferImpl(ViewVarBinaryVector to) { + this.to = to; + } + + @Override + public ViewVarBinaryVector getTo() { + return to; + } + + @Override + public void transfer() { + transferTo(to); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + splitAndTransferTo(startIndex, length, to); + } + + @Override + public void copyValueSafe(int fromIndex, int toIndex) { + to.copyFromSafe(fromIndex, toIndex, ViewVarBinaryVector.this); + } } } diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java b/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java index d2c03930ca37a..fece93de9bf14 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java @@ -67,7 +67,8 @@ private void populateVarcharVector(final VarCharVector vector, int valueCount, S vector.setValueCount(valueCount); } - private void populateViewVarcharVector(final ViewVarCharVector vector, int valueCount, String[] compareArray) { + private void populateBaseVariableWidthViewVector(final BaseVariableWidthViewVector vector, int valueCount, + String[] compareArray) { for (int i = 0; i < valueCount; i += 3) { final String s = String.format("%010d", i); vector.set(i, s.getBytes(StandardCharsets.UTF_8)); @@ -120,11 +121,16 @@ public void testWithEmptyVector() { transferPair = varCharVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); - // BaseVariableWidthViewVector + // BaseVariableWidthViewVector: ViewVarCharVector ViewVarCharVector viewVarCharVector = new ViewVarCharVector("", allocator); transferPair = viewVarCharVector.getTransferPair(allocator); transferPair.splitAndTransfer(0, 0); assertEquals(0, transferPair.getTo().getValueCount()); + // BaseVariableWidthVector: ViewVarBinaryVector + ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("", allocator); + transferPair = viewVarBinaryVector.getTransferPair(allocator); + transferPair.splitAndTransfer(0, 0); + assertEquals(0, transferPair.getTo().getValueCount()); // BaseLargeVariableWidthVector LargeVarCharVector largeVarCharVector = new LargeVarCharVector("", allocator); transferPair = largeVarCharVector.getTransferPair(allocator); @@ -225,36 +231,46 @@ public void test() throws Exception { } } - @Test - public void testView() throws Exception { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { - viewVarCharVector.allocateNew(10000, 1000); - - final int valueCount = 500; - final String[] compareArray = new String[valueCount]; + private void testView(BaseVariableWidthViewVector vector) { + vector.allocateNew(10000, 1000); + final int valueCount = 500; + final String[] compareArray = new String[valueCount]; - populateViewVarcharVector(viewVarCharVector, valueCount, compareArray); + populateBaseVariableWidthViewVector(vector, valueCount, compareArray); - final TransferPair tp = viewVarCharVector.getTransferPair(allocator); - final ViewVarCharVector newViewVarCharVector = (ViewVarCharVector) tp.getTo(); - final int[][] startLengths = {{0, 201}, {201, 0}, {201, 200}, {401, 99}}; + final TransferPair tp = vector.getTransferPair(allocator); + final BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo();; + final int[][] startLengths = {{0, 201}, {201, 0}, {201, 200}, {401, 99}}; - for (final int[] startLength : startLengths) { - final int start = startLength[0]; - final int length = startLength[1]; - tp.splitAndTransfer(start, length); - for (int i = 0; i < length; i++) { - final boolean expectedSet = ((start + i) % 3) == 0; - if (expectedSet) { - final byte[] expectedValue = compareArray[start + i].getBytes(StandardCharsets.UTF_8); - assertFalse(newViewVarCharVector.isNull(i)); - assertArrayEquals(expectedValue, newViewVarCharVector.get(i)); - } else { - assertTrue(newViewVarCharVector.isNull(i)); - } + for (final int[] startLength : startLengths) { + final int start = startLength[0]; + final int length = startLength[1]; + tp.splitAndTransfer(start, length); + for (int i = 0; i < length; i++) { + final boolean expectedSet = ((start + i) % 3) == 0; + if (expectedSet) { + final byte[] expectedValue = compareArray[start + i].getBytes(StandardCharsets.UTF_8); + assertFalse(newVector.isNull(i)); + assertArrayEquals(expectedValue, newVector.get(i)); + } else { + assertTrue(newVector.isNull(i)); } - newViewVarCharVector.clear(); } + newVector.clear(); + } + } + + @Test + public void testUtf8View() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + testView(viewVarCharVector); + } + } + + @Test + public void testBinaryView() throws Exception { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { + testView(viewVarBinaryVector); } } @@ -282,35 +298,47 @@ public void testMemoryConstrainedTransfer() { } } - @Test - public void testMemoryConstrainedTransferInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { - // Here we have the target vector being transferred with a long string - // hence, the data buffer will be allocated. - // The default data buffer allocation takes - // BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * BaseVariableWidthViewVector.ELEMENT_SIZE - // set limit = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * - // BaseVariableWidthViewVector.ELEMENT_SIZE - final int setLimit = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * - BaseVariableWidthViewVector.ELEMENT_SIZE; - allocator.setLimit(setLimit); - - viewVarCharVector.allocateNew(16000, 1000); + private void testMemoryConstrainedTransferInViews(BaseVariableWidthViewVector vector) { + // Here we have the target vector being transferred with a long string + // hence, the data buffer will be allocated. + // The default data buffer allocation takes + // BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * BaseVariableWidthViewVector.ELEMENT_SIZE + // set limit = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * + // BaseVariableWidthViewVector.ELEMENT_SIZE + final int setLimit = BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * + BaseVariableWidthViewVector.ELEMENT_SIZE; + allocator.setLimit(setLimit); - final int valueCount = 1000; + vector.allocateNew(16000, 1000); - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final int valueCount = 1000; - final TransferPair tp = viewVarCharVector.getTransferPair(allocator); - final ViewVarCharVector newViewVarCharVector = (ViewVarCharVector) tp.getTo(); - final int[][] startLengths = {{0, 700}, {700, 299}}; + populateBaseVariableWidthViewVector(vector, valueCount, null); - for (final int[] startLength : startLengths) { - final int start = startLength[0]; - final int length = startLength[1]; - tp.splitAndTransfer(start, length); - newViewVarCharVector.clear(); - } + final TransferPair tp = vector.getTransferPair(allocator); + final BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); + + final int[][] startLengths = {{0, 700}, {700, 299}}; + + for (final int[] startLength : startLengths) { + final int start = startLength[0]; + final int length = startLength[1]; + tp.splitAndTransfer(start, length); + newVector.clear(); + } + } + + @Test + public void testMemoryConstrainedTransferInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + testMemoryConstrainedTransferInViews(viewVarCharVector); + } + } + + @Test + public void testMemoryConstrainedTransferInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { + testMemoryConstrainedTransferInViews(viewVarBinaryVector); } } @@ -345,34 +373,45 @@ public void testTransfer() { } } - @Test - public void testTransferInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { - viewVarCharVector.allocateNew(16000, 1000); + private void testTransferInViews(BaseVariableWidthViewVector vector) { + vector.allocateNew(16000, 1000); - final int valueCount = 500; - final String[] compareArray = new String[valueCount]; - populateViewVarcharVector(viewVarCharVector, valueCount, compareArray); + final int valueCount = 500; + final String[] compareArray = new String[valueCount]; + populateBaseVariableWidthViewVector(vector, valueCount, compareArray); - final TransferPair tp = viewVarCharVector.getTransferPair(allocator); - final ViewVarCharVector newViewVarCharVector = (ViewVarCharVector) tp.getTo(); - tp.transfer(); + final TransferPair tp = vector.getTransferPair(allocator); + final BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); + tp.transfer(); - assertEquals(0, viewVarCharVector.valueCount); - assertEquals(valueCount, newViewVarCharVector.valueCount); + assertEquals(0, vector.valueCount); + assertEquals(valueCount, newVector.valueCount); - for (int i = 0; i < valueCount; i++) { - final boolean expectedSet = (i % 3) == 0; - if (expectedSet) { - final byte[] expectedValue = compareArray[i].getBytes(StandardCharsets.UTF_8); - assertFalse(newViewVarCharVector.isNull(i)); - assertArrayEquals(expectedValue, newViewVarCharVector.get(i)); - } else { - assertTrue(newViewVarCharVector.isNull(i)); - } + for (int i = 0; i < valueCount; i++) { + final boolean expectedSet = (i % 3) == 0; + if (expectedSet) { + final byte[] expectedValue = compareArray[i].getBytes(StandardCharsets.UTF_8); + assertFalse(newVector.isNull(i)); + assertArrayEquals(expectedValue, newVector.get(i)); + } else { + assertTrue(newVector.isNull(i)); } + } + + newVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testTransferInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + testTransferInViews(viewVarCharVector); + } + } + + @Test + public void testTransferInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { + testTransferInViews(viewVarBinaryVector); } } @@ -424,21 +463,31 @@ public void testSplitAndTransferNon() { } } - @Test - public void testSplitAndTransferNonInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + private void testSplitAndTransferNonInViews(BaseVariableWidthViewVector vector) { + vector.allocateNew(16000, 1000); + final int valueCount = 500; + populateBaseVariableWidthViewVector(vector, valueCount, null); - viewVarCharVector.allocateNew(16000, 1000); - final int valueCount = 500; - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final TransferPair tp = vector.getTransferPair(allocator); + BaseVariableWidthViewVector newVector = (BaseVariableWidthViewVector) tp.getTo(); - final TransferPair tp = viewVarCharVector.getTransferPair(allocator); - ViewVarCharVector newViewVarCharVector = (ViewVarCharVector) tp.getTo(); + tp.splitAndTransfer(0, 0); + assertEquals(0, newVector.getValueCount()); - tp.splitAndTransfer(0, 0); - assertEquals(0, newViewVarCharVector.getValueCount()); + newVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testSplitAndTransferNonInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + testSplitAndTransferNonInViews(viewVarCharVector); + } + } + + @Test + public void testSplitAndTransferNonInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { + testSplitAndTransferNonInViews(viewVarBinaryVector); } } @@ -460,21 +509,31 @@ public void testSplitAndTransferAll() { } } - @Test - public void testSplitAndTransferAllInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + private void testSplitAndTransferAllInViews(BaseVariableWidthViewVector vector) { + vector.allocateNew(16000, 1000); + final int valueCount = 500; + populateBaseVariableWidthViewVector(vector, valueCount, null); - viewVarCharVector.allocateNew(16000, 1000); - final int valueCount = 500; - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final TransferPair tp = vector.getTransferPair(allocator); + BaseVariableWidthViewVector newViewVarCharVector = (BaseVariableWidthViewVector) tp.getTo(); - final TransferPair tp = viewVarCharVector.getTransferPair(allocator); - ViewVarCharVector newViewVarCharVector = (ViewVarCharVector) tp.getTo(); + tp.splitAndTransfer(0, valueCount); + assertEquals(valueCount, newViewVarCharVector.getValueCount()); - tp.splitAndTransfer(0, valueCount); - assertEquals(valueCount, newViewVarCharVector.getValueCount()); + newViewVarCharVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testSplitAndTransferAllInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + testSplitAndTransferAllInViews(viewVarCharVector); + } + } + + @Test + public void testSplitAndTransferAllInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator)) { + testSplitAndTransferAllInViews(viewVarBinaryVector); } } @@ -499,24 +558,35 @@ public void testInvalidStartIndex() { } } - @Test - public void testInvalidStartIndexInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); - final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + private void testInvalidStartIndexInViews(BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { + vector.allocateNew(16000, 1000); + final int valueCount = 500; + populateBaseVariableWidthViewVector(vector, valueCount, null); - viewVarCharVector.allocateNew(16000, 1000); - final int valueCount = 500; - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final TransferPair tp = vector.makeTransferPair(newVector); - final TransferPair tp = viewVarCharVector.makeTransferPair(newViewVarCharVector); + IllegalArgumentException e = assertThrows( + IllegalArgumentException.class, + () -> tp.splitAndTransfer(valueCount, 10)); - IllegalArgumentException e = assertThrows( - IllegalArgumentException.class, - () -> tp.splitAndTransfer(valueCount, 10)); + assertEquals("Invalid parameters startIndex: 500, length: 10 for valueCount: 500", e.getMessage()); - assertEquals("Invalid parameters startIndex: 500, length: 10 for valueCount: 500", e.getMessage()); + newVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testInvalidStartIndexInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); + final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + testInvalidStartIndexInViews(viewVarCharVector, newViewVarCharVector); + } + } + + @Test + public void testInvalidStartIndexInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); + final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { + testInvalidStartIndexInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @@ -541,24 +611,35 @@ public void testInvalidLength() { } } - @Test - public void testInvalidLengthInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); - final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + private void testInvalidLengthInViews(BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { + vector.allocateNew(16000, 1000); + final int valueCount = 500; + populateBaseVariableWidthViewVector(vector, valueCount, null); - viewVarCharVector.allocateNew(16000, 1000); - final int valueCount = 500; - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final TransferPair tp = vector.makeTransferPair(newVector); - final TransferPair tp = viewVarCharVector.makeTransferPair(newViewVarCharVector); + IllegalArgumentException e = assertThrows( + IllegalArgumentException.class, + () -> tp.splitAndTransfer(0, valueCount * 2)); - IllegalArgumentException e = assertThrows( - IllegalArgumentException.class, - () -> tp.splitAndTransfer(0, valueCount * 2)); + assertEquals("Invalid parameters startIndex: 0, length: 1000 for valueCount: 500", e.getMessage()); - assertEquals("Invalid parameters startIndex: 0, length: 1000 for valueCount: 500", e.getMessage()); + newVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testInvalidLengthInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); + final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + testInvalidLengthInViews(viewVarCharVector, newViewVarCharVector); + } + } + + @Test + public void testInvalidLengthInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); + final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { + testInvalidLengthInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @@ -580,21 +661,33 @@ public void testZeroStartIndexAndLength() { } } - @Test - public void testZeroStartIndexAndLengthInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); - final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + private void testZeroStartIndexAndLengthInViews(BaseVariableWidthViewVector vector, + BaseVariableWidthViewVector newVector) { + vector.allocateNew(0, 0); + final int valueCount = 0; + populateBaseVariableWidthViewVector(vector, valueCount, null); - viewVarCharVector.allocateNew(0, 0); - final int valueCount = 0; - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final TransferPair tp = vector.makeTransferPair(newVector); - final TransferPair tp = viewVarCharVector.makeTransferPair(newViewVarCharVector); + tp.splitAndTransfer(0, 0); + assertEquals(valueCount, newVector.getValueCount()); - tp.splitAndTransfer(0, 0); - assertEquals(valueCount, newViewVarCharVector.getValueCount()); + newVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testZeroStartIndexAndLengthInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); + final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + testZeroStartIndexAndLengthInViews(viewVarCharVector, newViewVarCharVector); + } + } + + @Test + public void testZeroStartIndexAndLengthInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); + final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { + testZeroStartIndexAndLengthInViews(viewVarBinaryVector, newViewVarBinaryVector); } } @@ -616,21 +709,32 @@ public void testZeroLength() { } } - @Test - public void testZeroLengthInViews() { - try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); - final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + private void testZeroLengthInViews(BaseVariableWidthViewVector vector, BaseVariableWidthViewVector newVector) { + vector.allocateNew(16000, 1000); + final int valueCount = 500; + populateBaseVariableWidthViewVector(vector, valueCount, null); - viewVarCharVector.allocateNew(16000, 1000); - final int valueCount = 500; - populateViewVarcharVector(viewVarCharVector, valueCount, null); + final TransferPair tp = vector.makeTransferPair(newVector); - final TransferPair tp = viewVarCharVector.makeTransferPair(newViewVarCharVector); + tp.splitAndTransfer(500, 0); + assertEquals(0, newVector.getValueCount()); - tp.splitAndTransfer(500, 0); - assertEquals(0, newViewVarCharVector.getValueCount()); + newVector.clear(); + } - newViewVarCharVector.clear(); + @Test + public void testZeroLengthInUtf8Views() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator); + final ViewVarCharVector newViewVarCharVector = new ViewVarCharVector("newvector", allocator)) { + testZeroLengthInViews(viewVarCharVector, newViewVarCharVector); + } + } + + @Test + public void testZeroLengthInBinaryViews() { + try (final ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("myvector", allocator); + final ViewVarBinaryVector newViewVarBinaryVector = new ViewVarBinaryVector("newvector", allocator)) { + testZeroLengthInViews(viewVarBinaryVector, newViewVarBinaryVector); } } diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java index 817941ecb46d6..ebf9b58da7b40 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java @@ -37,7 +37,9 @@ import java.util.List; import java.util.Objects; import java.util.Random; +import java.util.function.BiConsumer; import java.util.function.Function; +import java.util.stream.IntStream; import java.util.stream.Stream; import org.apache.arrow.memory.ArrowBuf; @@ -86,9 +88,12 @@ public class TestVarCharViewVector { private BufferAllocator allocator; + private Random random; + @BeforeEach public void prepare() { allocator = new RootAllocator(Integer.MAX_VALUE); + random = new Random(); } @AfterEach @@ -1717,6 +1722,44 @@ public void testCopyFromSafeWithNulls(Function validateVector = (targetVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); + }; + try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { - sourceVector.allocateNew(1024 * 10, 1024); - - sourceVector.set(0, STR4); - sourceVector.set(1, STR5); - sourceVector.set(2, STR6); - sourceVector.setValueCount(3); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(0, 2, targetVector); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem < allocator.getAllocatedMemory()); - - // The validity buffer is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. - // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. - assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); + testSplitAndTransferOnSlicedBufferHelper(targetVector, sourceVector, + startIndex, length, data); } - assertArrayEquals(STR4, targetVector.get(0)); - assertArrayEquals(STR5, targetVector.get(1)); + validateVector.accept(targetVector, data); + } + + final byte [][] binaryData = generateBinaryDataArray(3, 10); + + try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { + try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + testSplitAndTransferOnSlicedBufferHelper(targetVector, sourceVector, + startIndex, length, binaryData); + } + validateVector.accept(targetVector, binaryData); } } @@ -1760,36 +1803,58 @@ public void testSplitAndTransfer1() { * With a long string included. */ @Test - public void testSplitAndTransfer2() { + public void testSplitAndTransferWithLongStringsOnSlicedBuffer() { + final byte [][] data = new byte[][]{STR2, STR5, STR6}; + final int startIndex = 0; + final int length = 2; + + BiConsumer validateVector = (targetVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); + }; + try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { - sourceVector.allocateNew(1024 * 10, 1024); - - sourceVector.set(0, STR2); - sourceVector.set(1, STR5); - sourceVector.set(2, STR6); - sourceVector.setValueCount(3); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(0, 2, targetVector); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem < allocator.getAllocatedMemory()); - - // The validity buffer is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. - // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. - assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR2, sourceVector.get(0)); - assertArrayEquals(STR5, sourceVector.get(1)); - assertArrayEquals(STR6, sourceVector.get(2)); + testSplitAndTransferOnSlicedBufferHelper(targetVector, sourceVector, + startIndex, length, data); } - assertArrayEquals(STR2, targetVector.get(0)); - assertArrayEquals(STR5, targetVector.get(1)); + validateVector.accept(targetVector, data); + } + + final byte [][] binaryData = generateBinaryDataArray(3, 18); + try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { + try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + testSplitAndTransferOnSlicedBufferHelper(targetVector, sourceVector, + startIndex, length, binaryData); + } + validateVector.accept(targetVector, binaryData); + } + } + + private void testSplitAndTransferOnSlicedVectorHelper(BaseVariableWidthViewVector sourceVector, + BaseVariableWidthViewVector targetVector, int startIndex, int length, byte[][] data) { + sourceVector.allocateNew(1024 * 10, 1024); + + for (int i = 0; i < data.length; i++) { + sourceVector.set(i, data[i]); + } + sourceVector.setValueCount(data.length); + + final long allocatedMem = allocator.getAllocatedMemory(); + final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); + final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); + + sourceVector.splitAndTransferTo(startIndex, length, targetVector); + // we allocate view and data buffers for the target vector + assertTrue(allocatedMem < allocator.getAllocatedMemory()); + // The validity buffer is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. + // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. + assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); + // since the new view buffer is allocated, the refcnt is the same as the source vector. + assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); + + for (int i = startIndex; i < length ; i++) { + assertArrayEquals(data[i], targetVector.get(i - startIndex)); } } @@ -1800,35 +1865,31 @@ public void testSplitAndTransfer2() { * With short strings. */ @Test - public void testSplitAndTransfer3() { + public void testSplitAndTransferWithShortStringsOnSlicedVector() { + byte [][] data = new byte[][]{STR4, STR5, STR6}; + final int startIndex = 0; + final int length = 2; + + BiConsumer validateVector = (sourceVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], sourceVector.get(i))); + }; + try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { - sourceVector.allocateNew(1024 * 10, 1024); - - sourceVector.set(0, STR4); - sourceVector.set(1, STR5); - sourceVector.set(2, STR6); - sourceVector.setValueCount(3); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(0, 2, targetVector); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem < allocator.getAllocatedMemory()); - // The validity buffer is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. - // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. - assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR4, targetVector.get(0)); - assertArrayEquals(STR5, targetVector.get(1)); + testSplitAndTransferOnSlicedVectorHelper(sourceVector, targetVector, + startIndex, length, data); + } + validateVector.accept(sourceVector, data); + } + + byte [][] binaryData = generateBinaryDataArray(3, 10); + try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { + testSplitAndTransferOnSlicedVectorHelper(sourceVector, targetVector, + startIndex, length, binaryData); } - assertArrayEquals(STR4, sourceVector.get(0)); - assertArrayEquals(STR5, sourceVector.get(1)); - assertArrayEquals(STR6, sourceVector.get(2)); + validateVector.accept(sourceVector, binaryData); } } @@ -1839,35 +1900,77 @@ public void testSplitAndTransfer3() { * With a long string included. */ @Test - public void testSplitAndTransfer4() { + public void testSplitAndTransferWithLongStringsOnSlicedVector() { + final byte [][] data = new byte[][]{STR2, STR5, STR6}; + final int startIndex = 0; + final int length = 2; + + BiConsumer validateVector = (sourceVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], sourceVector.get(i))); + }; + try (final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator)) { - sourceVector.allocateNew(1024 * 10, 1024); - - sourceVector.set(0, STR2); - sourceVector.set(1, STR5); - sourceVector.set(2, STR6); - sourceVector.setValueCount(3); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(0, 2, targetVector); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem < allocator.getAllocatedMemory()); - // The validity buffer is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. - // Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0. - assertEquals(validityRefCnt + 1, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR2, targetVector.get(0)); - assertArrayEquals(STR5, targetVector.get(1)); + testSplitAndTransferOnSlicedVectorHelper(sourceVector, targetVector, + startIndex, length, data); + } + validateVector.accept(sourceVector, data); + } + + final byte [][] binaryData = generateBinaryDataArray(3, 20); + try (final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator)) { + testSplitAndTransferOnSlicedVectorHelper(sourceVector, targetVector, + startIndex, length, binaryData); + } + validateVector.accept(sourceVector, binaryData); + } + } + + private void testSplitAndTransferOnValiditySplitHelper( + BaseVariableWidthViewVector targetVector, BaseVariableWidthViewVector sourceVector, + int startIndex, int length, byte[][] data) { + sourceVector.allocateNew(1024 * 10, 1024); + + sourceVector.set(0, new byte[0]); + sourceVector.setNull(1); + for (int i = 0; i < data.length; i++) { + if (data[i] == null) { + sourceVector.setNull(i); + } else { + sourceVector.set(i, data[i]); + } + } + sourceVector.setValueCount(data.length); + + final long allocatedMem = allocator.getAllocatedMemory(); + final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); + final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); + + sourceVector.splitAndTransferTo(startIndex, length, targetVector); + // the allocation only consists in the size needed for the validity buffer + final long validitySize = + DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( + BaseValueVector.getValidityBufferSizeFromCount(2)); + // we allocate view and data buffers for the target vector + assertTrue(allocatedMem + validitySize < allocator.getAllocatedMemory()); + // The validity is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. + // Since values up to the startIndex are empty/null validity refcnt should not change. + assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); + // since the new view buffer is allocated, the refcnt is the same as the source vector. + assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); + + for (int i = startIndex; i < startIndex + length; i++) { + assertArrayEquals(data[i], targetVector.get(i - startIndex)); + } + + for (int i = 0; i < data.length; i++) { + if (data[i] == null) { + assertTrue(sourceVector.isNull(i)); + } else { + assertArrayEquals(data[i], sourceVector.get(i)); } - assertArrayEquals(STR2, sourceVector.get(0)); - assertArrayEquals(STR5, sourceVector.get(1)); - assertArrayEquals(STR6, sourceVector.get(2)); } } @@ -1878,43 +1981,24 @@ public void testSplitAndTransfer4() { * With short strings. */ @Test - public void testSplitAndTransfer5() { + public void testSplitAndTransferWithShortStringsOnValiditySplit() { + final byte [][] data = new byte[][]{new byte[0], null, STR4, STR5, STR6}; + final int startIndex = 2; + final int length = 2; + try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { - sourceVector.allocateNew(1024 * 10, 1024); - - sourceVector.set(0, new byte[0]); - sourceVector.setNull(1); - sourceVector.set(2, STR4); - sourceVector.set(3, STR5); - sourceVector.set(4, STR6); - sourceVector.setValueCount(5); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(2, 2, targetVector); - // the allocation only consists in the size needed for the validity buffer - final long validitySize = - DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( - BaseValueVector.getValidityBufferSizeFromCount(2)); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem + validitySize < allocator.getAllocatedMemory()); - // The validity is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. - // Since values up to the startIndex are empty/null validity refcnt should not change. - assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR4, targetVector.get(0)); - assertArrayEquals(STR5, targetVector.get(1)); - - assertArrayEquals(new byte[0], sourceVector.get(0)); - assertTrue(sourceVector.isNull(1)); - assertArrayEquals(STR4, sourceVector.get(2)); - assertArrayEquals(STR5, sourceVector.get(3)); - assertArrayEquals(STR6, sourceVector.get(4)); + testSplitAndTransferOnValiditySplitHelper(targetVector, sourceVector, + startIndex, length, data); + } + + final byte [][] binaryData = generateBinaryDataArray(5, 10); + binaryData[0] = new byte[0]; + binaryData[1] = null; + try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator); + final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + testSplitAndTransferOnValiditySplitHelper(targetVector, sourceVector, + startIndex, length, binaryData); } } @@ -1925,44 +2009,59 @@ public void testSplitAndTransfer5() { * With long strings. */ @Test - public void testSplitAndTransfer6() { + public void testSplitAndTransferWithLongStringsOnValiditySplit() { + final byte [][] data = new byte[][]{new byte[0], null, STR1, STR2, STR3}; + final int startIndex = 2; + final int length = 2; + try (final ViewVarCharVector targetVector = newViewVarCharVector("split-target", allocator); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { - sourceVector.allocateNew(1024 * 10, 1024); - - sourceVector.set(0, new byte[0]); - sourceVector.setNull(1); - sourceVector.set(2, STR1); - sourceVector.set(3, STR2); - sourceVector.set(4, STR3); - sourceVector.setValueCount(5); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(2, 2, targetVector); - // the allocation consists in the size needed for the validity buffer and the long string - // allocation - final long validitySize = - DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( - BaseValueVector.getValidityBufferSizeFromCount(2)); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem + validitySize < allocator.getAllocatedMemory()); - // The validity is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. - // Since values up to the startIndex are empty/null validity refcnt should not change. - assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR1, targetVector.get(0)); - assertArrayEquals(STR2, targetVector.get(1)); - - assertArrayEquals(new byte[0], sourceVector.get(0)); - assertTrue(sourceVector.isNull(1)); - assertArrayEquals(STR1, sourceVector.get(2)); - assertArrayEquals(STR2, sourceVector.get(3)); - assertArrayEquals(STR3, sourceVector.get(4)); + testSplitAndTransferOnValiditySplitHelper(targetVector, sourceVector, + startIndex, length, data); + } + + final byte [][] binaryData = generateBinaryDataArray(5, 18); + binaryData[0] = new byte[0]; + binaryData[1] = null; + + try (final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", allocator); + final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + testSplitAndTransferOnValiditySplitHelper(targetVector, sourceVector, + startIndex, length, data); + } + } + + private void testSplitAndTransferOnAllocatorToAllocator(BaseVariableWidthViewVector targetVector, + BaseVariableWidthViewVector sourceVector, int startIndex, int length, byte[][] data) { + sourceVector.allocateNew(50, data.length); + + for (int i = 0; i < data.length; i++) { + sourceVector.set(i, data[i]); + } + sourceVector.setValueCount(data.length); + + final long allocatedMem = allocator.getAllocatedMemory(); + final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); + final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); + + sourceVector.splitAndTransferTo(startIndex, length, targetVector); + + if (sourceVector.getDataBuffers().isEmpty()) { + // no extra allocation as strings are all inline + assertEquals(allocatedMem, allocator.getAllocatedMemory()); + } else { + // extra allocation as some strings are not inline + assertTrue(allocatedMem < allocator.getAllocatedMemory()); + } + + // the refcnts of each buffer for this test should be the same as what + // the source allocator ended up with. + assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); + // since the new view buffer is allocated, the refcnt is the same as the source vector. + assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); + + for (int i = 0; i < data.length; i++) { + assertArrayEquals(data[i], sourceVector.get(i)); } } @@ -1973,39 +2072,36 @@ public void testSplitAndTransfer6() { * With short strings. */ @Test - public void testSplitAndTransfer7() { + public void testSplitAndTransferWithShortStringsOnAllocatorToAllocator() { final int maxAllocation = 512; + final byte [][] data = new byte[][]{STR4, STR5, STR6}; + final int startIndex = 0; + final int length = 2; + + BiConsumer validateVector = (targetVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); + }; + try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", 256, maxAllocation); final ViewVarCharVector targetVector = newViewVarCharVector("split-target", targetAllocator)) { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", 256, maxAllocation); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { - sourceVector.allocateNew(50, 3); - - sourceVector.set(0, STR4); - sourceVector.set(1, STR5); - sourceVector.set(2, STR6); - sourceVector.setValueCount(3); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(0, 2, targetVector); - // no extra allocation as strings are all inline - assertEquals(allocatedMem, allocator.getAllocatedMemory()); - - // the refcnts of each buffer for this test should be the same as what - // the source allocator ended up with. - assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR4, sourceVector.get(0)); - assertArrayEquals(STR5, sourceVector.get(1)); - assertArrayEquals(STR6, sourceVector.get(2)); + testSplitAndTransferOnAllocatorToAllocator(targetVector, sourceVector, + startIndex, length, data); } - assertArrayEquals(STR4, targetVector.get(0)); - assertArrayEquals(STR5, targetVector.get(1)); + validateVector.accept(targetVector, data); + } + + final byte [][] binaryData = generateBinaryDataArray(3, 10); + try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", 256, maxAllocation); + final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", targetAllocator)) { + try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", 256, maxAllocation); + final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { + testSplitAndTransferOnAllocatorToAllocator(targetVector, sourceVector, + startIndex, length, binaryData); + } + validateVector.accept(targetVector, binaryData); } } @@ -2016,12 +2112,21 @@ public void testSplitAndTransfer7() { * With long strings. */ @Test - public void testSplitAndTransfer8() { + public void testSplitAndTransferWithLongStringsOnAllocatorToAllocator() { final int initialReservation = 1024; // Here we have the target vector being transferred with a long string // hence, the data buffer will be allocated. // The default data buffer allocation takes // BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * BaseVariableWidthViewVector.ELEMENT_SIZE + final byte [][] data = new byte[][]{STR1, STR2, STR3}; + final int startIndex = 0; + final int length = 2; + + BiConsumer validateVector = (targetVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); + }; + final int maxAllocation = initialReservation + BaseVariableWidthViewVector.INITIAL_VIEW_VALUE_ALLOCATION * BaseVariableWidthViewVector.ELEMENT_SIZE; try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", @@ -2030,136 +2135,169 @@ public void testSplitAndTransfer8() { try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", initialReservation, maxAllocation); final ViewVarCharVector sourceVector = newViewVarCharVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { - sourceVector.allocateNew(48, 3); - - sourceVector.set(0, STR1); - sourceVector.set(1, STR2); - sourceVector.set(2, STR3); - sourceVector.setValueCount(3); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - sourceVector.splitAndTransferTo(0, 2, targetVector); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem < allocator.getAllocatedMemory()); - - // the refcnts of each buffer for this test should be the same as what - // the source allocator ended up with. - assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR1, sourceVector.get(0)); - assertArrayEquals(STR2, sourceVector.get(1)); - assertArrayEquals(STR3, sourceVector.get(2)); + testSplitAndTransferOnAllocatorToAllocator(targetVector, sourceVector, + startIndex, length, data); } - assertArrayEquals(STR1, targetVector.get(0)); - assertArrayEquals(STR2, targetVector.get(1)); + validateVector.accept(targetVector, data); } - } - @Test - public void testReallocAfterVectorTransfer1() { - try (final ViewVarCharVector vector = new ViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { - /* 4096 values with 16 bytes per record */ - final int bytesPerRecord = 32; - vector.allocateNew(4096 * bytesPerRecord, 4096); - int valueCapacity = vector.getValueCapacity(); - assertTrue(valueCapacity >= 4096); - - /* populate the vector */ - for (int i = 0; i < valueCapacity; i++) { - if ((i & 1) == 1) { - vector.set(i, STR1); - } else { - vector.set(i, STR2); - } + final byte [][] binaryData = generateBinaryDataArray(3, 18); + + try (final BufferAllocator targetAllocator = allocator.newChildAllocator("target-alloc", + initialReservation, maxAllocation); + final ViewVarBinaryVector targetVector = newViewVarBinaryVector("split-target", targetAllocator)) { + try (final BufferAllocator sourceAllocator = allocator.newChildAllocator("source-alloc", + initialReservation, maxAllocation); + final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, sourceAllocator)) { + testSplitAndTransferOnAllocatorToAllocator(targetVector, sourceVector, + startIndex, length, binaryData); } + validateVector.accept(targetVector, binaryData); + } + } - /* Check the vector output */ - for (int i = 0; i < valueCapacity; i++) { - if ((i & 1) == 1) { - assertArrayEquals(STR1, vector.get(i)); - } else { - assertArrayEquals(STR2, vector.get(i)); - } + private void testReallocAfterVectorTransferHelper(BaseVariableWidthViewVector vector, + byte[] str1, byte[] str2) { + /* 4096 values with 16 bytes per record */ + final int bytesPerRecord = 32; + vector.allocateNew(4096 * bytesPerRecord, 4096); + int valueCapacity = vector.getValueCapacity(); + assertTrue(valueCapacity >= 4096); + + /* populate the vector */ + for (int i = 0; i < valueCapacity; i++) { + if ((i & 1) == 1) { + vector.set(i, str1); + } else { + vector.set(i, str2); } + } - /* trigger first realloc */ - vector.setSafe(valueCapacity, STR2, 0, STR2.length); - assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); - while (vector.getByteCapacity() < bytesPerRecord * vector.getValueCapacity()) { - vector.reallocViewBuffer(); - vector.reallocViewDataBuffer(); + /* Check the vector output */ + for (int i = 0; i < valueCapacity; i++) { + if ((i & 1) == 1) { + assertArrayEquals(str1, vector.get(i)); + } else { + assertArrayEquals(str2, vector.get(i)); } + } - /* populate the remaining vector */ - for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { - if ((i & 1) == 1) { - vector.set(i, STR1); - } else { - vector.set(i, STR2); - } + /* trigger first realloc */ + vector.setSafe(valueCapacity, str2, 0, str2.length); + assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); + while (vector.getByteCapacity() < bytesPerRecord * vector.getValueCapacity()) { + vector.reallocViewBuffer(); + vector.reallocViewDataBuffer(); + } + + /* populate the remaining vector */ + for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { + if ((i & 1) == 1) { + vector.set(i, str1); + } else { + vector.set(i, str2); } + } - /* Check the vector output */ - valueCapacity = vector.getValueCapacity(); - for (int i = 0; i < valueCapacity; i++) { - if ((i & 1) == 1) { - assertArrayEquals(STR1, vector.get(i)); - } else { - assertArrayEquals(STR2, vector.get(i)); - } + /* Check the vector output */ + valueCapacity = vector.getValueCapacity(); + for (int i = 0; i < valueCapacity; i++) { + if ((i & 1) == 1) { + assertArrayEquals(str1, vector.get(i)); + } else { + assertArrayEquals(str2, vector.get(i)); } + } + + /* trigger second realloc */ + vector.setSafe(valueCapacity + bytesPerRecord, str2, 0, str2.length); + assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); + while (vector.getByteCapacity() < bytesPerRecord * vector.getValueCapacity()) { + vector.reallocViewBuffer(); + vector.reallocViewDataBuffer(); + } - /* trigger second realloc */ - vector.setSafe(valueCapacity + bytesPerRecord, STR2, 0, STR2.length); - assertTrue(vector.getValueCapacity() >= 2 * valueCapacity); - while (vector.getByteCapacity() < bytesPerRecord * vector.getValueCapacity()) { - vector.reallocViewBuffer(); - vector.reallocViewDataBuffer(); + /* populate the remaining vector */ + for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { + if ((i & 1) == 1) { + vector.set(i, str1); + } else { + vector.set(i, str2); } + } - /* populate the remaining vector */ - for (int i = valueCapacity; i < vector.getValueCapacity(); i++) { - if ((i & 1) == 1) { - vector.set(i, STR1); - } else { - vector.set(i, STR2); - } + /* Check the vector output */ + valueCapacity = vector.getValueCapacity(); + for (int i = 0; i < valueCapacity; i++) { + if ((i & 1) == 1) { + assertArrayEquals(str1, vector.get(i)); + } else { + assertArrayEquals(str2, vector.get(i)); } + } - /* Check the vector output */ - valueCapacity = vector.getValueCapacity(); - for (int i = 0; i < valueCapacity; i++) { - if ((i & 1) == 1) { - assertArrayEquals(STR1, vector.get(i)); - } else { - assertArrayEquals(STR2, vector.get(i)); - } + /* We are potentially working with 4x the size of vector buffer + * that we initially started with. + * Now let's transfer the vector. + */ + + TransferPair transferPair = vector.getTransferPair(allocator); + transferPair.transfer(); + BaseVariableWidthViewVector toVector = (BaseVariableWidthViewVector) transferPair.getTo(); + valueCapacity = toVector.getValueCapacity(); + + for (int i = 0; i < valueCapacity; i++) { + if ((i & 1) == 1) { + assertArrayEquals(str1, toVector.get(i)); + } else { + assertArrayEquals(str2, toVector.get(i)); } + } + toVector.close(); + } - /* We are potentially working with 4x the size of vector buffer - * that we initially started with. - * Now let's transfer the vector. - */ + @Test + public void testReallocAfterVectorTransfer() { + try (final ViewVarCharVector vector = new ViewVarCharVector(EMPTY_SCHEMA_PATH, allocator)) { + testReallocAfterVectorTransferHelper(vector, STR1, STR2); + } - TransferPair transferPair = vector.getTransferPair(allocator); - transferPair.transfer(); - ViewVarCharVector toVector = (ViewVarCharVector) transferPair.getTo(); - valueCapacity = toVector.getValueCapacity(); + try (final ViewVarBinaryVector vector = new ViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { + testReallocAfterVectorTransferHelper(vector, generateRandomBinaryData(12), + generateRandomBinaryData(13)); + } + } - for (int i = 0; i < valueCapacity; i++) { - if ((i & 1) == 1) { - assertArrayEquals(STR1, toVector.get(i)); - } else { - assertArrayEquals(STR2, toVector.get(i)); - } - } + private void testSplitAndTransferWithMultipleDataBuffersHelper(BaseVariableWidthViewVector sourceVector, + BaseVariableWidthViewVector targetVector, int startIndex, int length, byte[][] data) { + sourceVector.allocateNew(48, 4); + + for (int i = 0; i < data.length; i++) { + sourceVector.set(i, data[i]); + } + sourceVector.setValueCount(data.length); + + // we should have multiple data buffers + assertTrue(sourceVector.getDataBuffers().size() > 1); + + final long allocatedMem = allocator.getAllocatedMemory(); + final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); + final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - toVector.close(); + // split and transfer with slice starting at the beginning: + // this should not allocate anything new + sourceVector.splitAndTransferTo(startIndex, length, targetVector); + // we allocate view and data buffers for the target vector + assertTrue(allocatedMem < allocator.getAllocatedMemory()); + + // the refcnts of each buffer for this test should be the same as what + // the source allocator ended up with. + assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); + // since the new view buffer is allocated, the refcnt is the same as the source vector. + assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); + + for (int i = 0; i < data.length; i++) { + assertArrayEquals(data[i], sourceVector.get(i)); } } @@ -2171,45 +2309,31 @@ public void testReallocAfterVectorTransfer1() { * Check multi-data buffer source copying */ @Test - public void testSplitAndTransfer9() { + public void testSplitAndTransferWithMultipleDataBuffers() { + final String str4 = generateRandomString(35); + final byte[][] data = new byte[][]{STR1, STR2, STR3, str4.getBytes(StandardCharsets.UTF_8)}; + final int startIndex = 1; + final int length = 3; + + BiConsumer validateVector = (targetVector, expectedData) -> { + IntStream.range(startIndex, length).forEach(i -> + assertArrayEquals(expectedData[i], targetVector.get(i - startIndex))); + }; + try (final ViewVarCharVector targetVector = new ViewVarCharVector("target", allocator)) { - String str4 = generateRandomString(35); try (final ViewVarCharVector sourceVector = new ViewVarCharVector("source", allocator)) { - sourceVector.allocateNew(48, 4); - - sourceVector.set(0, STR1); - sourceVector.set(1, STR2); - sourceVector.set(2, STR3); - sourceVector.set(3, str4.getBytes(StandardCharsets.UTF_8)); - sourceVector.setValueCount(4); - - // we should have multiple data buffers - assertTrue(sourceVector.getDataBuffers().size() > 1); - - final long allocatedMem = allocator.getAllocatedMemory(); - final int validityRefCnt = sourceVector.getValidityBuffer().refCnt(); - final int dataRefCnt = sourceVector.getDataBuffer().refCnt(); - - // split and transfer with slice starting at the beginning: - // this should not allocate anything new - sourceVector.splitAndTransferTo(1, 3, targetVector); - // we allocate view and data buffers for the target vector - assertTrue(allocatedMem < allocator.getAllocatedMemory()); - - // the refcnts of each buffer for this test should be the same as what - // the source allocator ended up with. - assertEquals(validityRefCnt, sourceVector.getValidityBuffer().refCnt()); - // since the new view buffer is allocated, the refcnt is the same as the source vector. - assertEquals(dataRefCnt, sourceVector.getDataBuffer().refCnt()); - - assertArrayEquals(STR1, sourceVector.get(0)); - assertArrayEquals(STR2, sourceVector.get(1)); - assertArrayEquals(STR3, sourceVector.get(2)); - assertArrayEquals(str4.getBytes(StandardCharsets.UTF_8), sourceVector.get(3)); + testSplitAndTransferWithMultipleDataBuffersHelper(sourceVector, targetVector, + startIndex, length, data); + } + validateVector.accept(targetVector, data); + } + + try (final ViewVarBinaryVector targetVector = new ViewVarBinaryVector("target", allocator)) { + try (final ViewVarBinaryVector sourceVector = new ViewVarBinaryVector("source", allocator)) { + testSplitAndTransferWithMultipleDataBuffersHelper(sourceVector, targetVector, + startIndex, length, data); } - assertArrayEquals(STR2, targetVector.get(0)); - assertArrayEquals(STR3, targetVector.get(1)); - assertArrayEquals(str4.getBytes(StandardCharsets.UTF_8), targetVector.get(2)); + validateVector.accept(targetVector, data); } }