diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListViewReader.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListViewReader.java index 4bcd028de3..cd4a0c17dc 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListViewReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListViewReader.java @@ -33,6 +33,7 @@ public class UnionLargeListViewReader extends AbstractFieldReader { private final ValueVector data; private int currentOffset; private int size; + private int remaining; /** * Constructor for UnionLargeListViewReader. @@ -60,6 +61,7 @@ public void setPosition(int index) { if (vector.getOffsetBuffer().capacity() == 0) { currentOffset = 0; size = 0; + remaining = 0; } else { currentOffset = vector @@ -67,6 +69,7 @@ public void setPosition(int index) { .getInt(index * (long) BaseLargeRepeatedValueViewVector.OFFSET_WIDTH); size = vector.getSizeBuffer().getInt(index * (long) BaseLargeRepeatedValueViewVector.SIZE_WIDTH); + remaining = size; } } @@ -102,12 +105,10 @@ public int size() { @Override public boolean next() { - // Here, the currentOffSet keeps track of the current position in the vector inside the list at - // set position. - // And, size keeps track of the elements count in the list, so to make sure we traverse - // the full list, we need to check if the currentOffset is less than the currentOffset + size - if (currentOffset < currentOffset + size) { + // Yield exactly the element count stored with this list view, beginning at its stored offset. + if (remaining > 0) { data.getReader().setPosition(checkedCastToInt(currentOffset++)); + remaining--; return true; } else { return false; diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionListViewReader.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionListViewReader.java index 17ac1150fd..c4bc2b5015 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionListViewReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionListViewReader.java @@ -31,6 +31,7 @@ public class UnionListViewReader extends AbstractFieldReader { private final ValueVector data; private int currentOffset; private int size; + private int remaining; /** * Constructor for UnionListViewReader. @@ -58,10 +59,12 @@ public void setPosition(int index) { if (vector.getOffsetBuffer().capacity() == 0) { currentOffset = 0; size = 0; + remaining = 0; } else { currentOffset = vector.getOffsetBuffer().getInt(index * (long) BaseRepeatedValueViewVector.OFFSET_WIDTH); size = vector.getSizeBuffer().getInt(index * (long) BaseRepeatedValueViewVector.SIZE_WIDTH); + remaining = size; } } @@ -97,12 +100,10 @@ public int size() { @Override public boolean next() { - // Here, the currentOffSet keeps track of the current position in the vector inside the list at - // set position. - // And, size keeps track of the elements count in the list, so to make sure we traverse - // the full list, we need to check if the currentOffset is less than the currentOffset + size - if (currentOffset < currentOffset + size) { + // Yield exactly the element count stored with this list view, beginning at its stored offset. + if (remaining > 0) { data.getReader().setPosition(currentOffset++); + remaining--; return true; } else { return false; diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java index 256aa99687..1cdf10308a 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java @@ -29,6 +29,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector; import org.apache.arrow.vector.complex.LargeListViewVector; +import org.apache.arrow.vector.complex.impl.UnionLargeListViewReader; import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -2230,6 +2231,38 @@ public void testRangeChildVector2() { } } + @Test + public void testDirectReaderIteratesLargeListViewRange() { + try (LargeListViewVector largeListViewVector = + LargeListViewVector.empty("largelistview", allocator)) { + largeListViewVector.allocateNew(); + FieldType fieldType = new FieldType(true, new ArrowType.Int(32, true), null, null); + largeListViewVector.initializeChildrenFromFields( + Collections.singletonList(new Field("child-vector", fieldType, null))); + IntVector childVector = (IntVector) largeListViewVector.getDataVector(); + childVector.allocateNew(5); + for (int i = 0; i < 5; i++) { + childVector.set(i, 10 + i); + } + childVector.setValueCount(5); + largeListViewVector.setValidity(0, 1); + largeListViewVector.setOffset(0, 2); + largeListViewVector.setSize(0, 3); + largeListViewVector.setValueCount(1); + + UnionLargeListViewReader reader = new UnionLargeListViewReader(largeListViewVector); + reader.setPosition(0); + assertTrue(reader.next()); + assertEquals(12, ((Number) reader.reader().readObject()).intValue()); + assertTrue(reader.next()); + assertEquals(13, ((Number) reader.reader().readObject()).intValue()); + assertTrue(reader.next()); + assertEquals(14, ((Number) reader.reader().readObject()).intValue()); + assertFalse(reader.next()); + assertFalse(reader.next()); + } + } + private void writeIntValues(UnionLargeListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java index 8ab0edb145..96403f7eef 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java @@ -32,9 +32,11 @@ import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; +import org.apache.arrow.vector.complex.impl.UnionListViewReader; import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; +import org.apache.arrow.vector.holders.UnionHolder; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -142,6 +144,69 @@ public void testBasicListViewVector() { } } + @Test + public void testCopyFromNonEmptyListView() { + try (ListViewVector inVector = ListViewVector.empty("input", allocator); + ListViewVector outVector = ListViewVector.empty("output", allocator)) { + UnionListViewWriter writer = inVector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writeIntValues(writer, new int[] {10, 20}); + writer.setValueCount(1); + + outVector.allocateNew(); + outVector.copyFrom(0, 0, inVector); + outVector.setValueCount(1); + + assertEquals(Arrays.asList(10, 20), outVector.getObject(0)); + } + } + + @Test + public void testReaderIteratesListViewRangeAndResets() { + try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { + initializeListViewVector( + listViewVector, + List.of(10, 11, 20, 21, 22), + List.of(1, 1, 1), + List.of(0, 2, 5), + List.of(2, 3, 0)); + UnionListViewReader reader = listViewVector.getReader(); + + reader.setPosition(0); + assertTrue(reader.next()); + assertEquals(10, ((Number) reader.reader().readObject()).intValue()); + assertTrue(reader.next()); + assertEquals(11, ((Number) reader.reader().readObject()).intValue()); + assertFalse(reader.next()); + assertFalse(reader.next()); + + reader.setPosition(1); + assertTrue(reader.next()); + assertEquals(20, ((Number) reader.reader().readObject()).intValue()); + assertTrue(reader.next()); + assertEquals(21, ((Number) reader.reader().readObject()).intValue()); + assertTrue(reader.next()); + assertEquals(22, ((Number) reader.reader().readObject()).intValue()); + assertFalse(reader.next()); + assertFalse(reader.next()); + + reader.setPosition(1); + assertTrue(reader.next()); + assertEquals(20, ((Number) reader.reader().readObject()).intValue()); + + reader.setPosition(2); + assertFalse(reader.next()); + assertFalse(reader.next()); + + reader.setPosition(1); + UnionHolder holder = new UnionHolder(); + reader.read(2, holder); + assertEquals(22, ((Number) holder.reader.readObject()).intValue()); + assertFalse(reader.next()); + } + } + @Test public void testImplicitNullVectors() { try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) {