GH-471: Fix ListView reader iteration bounds

Author: PNHD

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,13 +61,15 @@ public void setPosition(int index) {
     if (vector.getOffsetBuffer().capacity() == 0) {
       currentOffset = 0;
       size = 0;
+      remaining = 0;
     } else {
       currentOffset =
           vector
               .getOffsetBuffer()
               .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;

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;

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) {

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)) {