Skip to content

Commit 5b1c63a

Browse files
PNHDPNHD
committed
GH-471: Fix ListView reader iteration bounds
1 parent 53a9ccd commit 5b1c63a

4 files changed

Lines changed: 110 additions & 10 deletions

File tree

‎vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionLargeListViewReader.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ public class UnionLargeListViewReader extends AbstractFieldReader {
3333
private final ValueVector data;
3434
private int currentOffset;
3535
private int size;
36+
private int remaining;
3637

3738
/**
3839
* Constructor for UnionLargeListViewReader.
@@ -60,13 +61,15 @@ public void setPosition(int index) {
6061
if (vector.getOffsetBuffer().capacity() == 0) {
6162
currentOffset = 0;
6263
size = 0;
64+
remaining = 0;
6365
} else {
6466
currentOffset =
6567
vector
6668
.getOffsetBuffer()
6769
.getInt(index * (long) BaseLargeRepeatedValueViewVector.OFFSET_WIDTH);
6870
size =
6971
vector.getSizeBuffer().getInt(index * (long) BaseLargeRepeatedValueViewVector.SIZE_WIDTH);
72+
remaining = size;
7073
}
7174
}
7275

@@ -102,12 +105,10 @@ public int size() {
102105

103106
@Override
104107
public boolean next() {
105-
// Here, the currentOffSet keeps track of the current position in the vector inside the list at
106-
// set position.
107-
// And, size keeps track of the elements count in the list, so to make sure we traverse
108-
// the full list, we need to check if the currentOffset is less than the currentOffset + size
109-
if (currentOffset < currentOffset + size) {
108+
// Yield exactly the element count stored with this list view, beginning at its stored offset.
109+
if (remaining > 0) {
110110
data.getReader().setPosition(checkedCastToInt(currentOffset++));
111+
remaining--;
111112
return true;
112113
} else {
113114
return false;

‎vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionListViewReader.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ public class UnionListViewReader extends AbstractFieldReader {
3131
private final ValueVector data;
3232
private int currentOffset;
3333
private int size;
34+
private int remaining;
3435

3536
/**
3637
* Constructor for UnionListViewReader.
@@ -58,10 +59,12 @@ public void setPosition(int index) {
5859
if (vector.getOffsetBuffer().capacity() == 0) {
5960
currentOffset = 0;
6061
size = 0;
62+
remaining = 0;
6163
} else {
6264
currentOffset =
6365
vector.getOffsetBuffer().getInt(index * (long) BaseRepeatedValueViewVector.OFFSET_WIDTH);
6466
size = vector.getSizeBuffer().getInt(index * (long) BaseRepeatedValueViewVector.SIZE_WIDTH);
67+
remaining = size;
6568
}
6669
}
6770

@@ -97,12 +100,10 @@ public int size() {
97100

98101
@Override
99102
public boolean next() {
100-
// Here, the currentOffSet keeps track of the current position in the vector inside the list at
101-
// set position.
102-
// And, size keeps track of the elements count in the list, so to make sure we traverse
103-
// the full list, we need to check if the currentOffset is less than the currentOffset + size
104-
if (currentOffset < currentOffset + size) {
103+
// Yield exactly the element count stored with this list view, beginning at its stored offset.
104+
if (remaining > 0) {
105105
data.getReader().setPosition(currentOffset++);
106+
remaining--;
106107
return true;
107108
} else {
108109
return false;

‎vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.apache.arrow.memory.BufferAllocator;
3030
import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector;
3131
import org.apache.arrow.vector.complex.LargeListViewVector;
32+
import org.apache.arrow.vector.complex.impl.UnionLargeListViewReader;
3233
import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter;
3334
import org.apache.arrow.vector.types.Types.MinorType;
3435
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -2230,6 +2231,38 @@ public void testRangeChildVector2() {
22302231
}
22312232
}
22322233

2234+
@Test
2235+
public void testDirectReaderIteratesLargeListViewRange() {
2236+
try (LargeListViewVector largeListViewVector =
2237+
LargeListViewVector.empty("largelistview", allocator)) {
2238+
largeListViewVector.allocateNew();
2239+
FieldType fieldType = new FieldType(true, new ArrowType.Int(32, true), null, null);
2240+
largeListViewVector.initializeChildrenFromFields(
2241+
Collections.singletonList(new Field("child-vector", fieldType, null)));
2242+
IntVector childVector = (IntVector) largeListViewVector.getDataVector();
2243+
childVector.allocateNew(5);
2244+
for (int i = 0; i < 5; i++) {
2245+
childVector.set(i, 10 + i);
2246+
}
2247+
childVector.setValueCount(5);
2248+
largeListViewVector.setValidity(0, 1);
2249+
largeListViewVector.setOffset(0, 2);
2250+
largeListViewVector.setSize(0, 3);
2251+
largeListViewVector.setValueCount(1);
2252+
2253+
UnionLargeListViewReader reader = new UnionLargeListViewReader(largeListViewVector);
2254+
reader.setPosition(0);
2255+
assertTrue(reader.next());
2256+
assertEquals(12, ((Number) reader.reader().readObject()).intValue());
2257+
assertTrue(reader.next());
2258+
assertEquals(13, ((Number) reader.reader().readObject()).intValue());
2259+
assertTrue(reader.next());
2260+
assertEquals(14, ((Number) reader.reader().readObject()).intValue());
2261+
assertFalse(reader.next());
2262+
assertFalse(reader.next());
2263+
}
2264+
}
2265+
22332266
private void writeIntValues(UnionLargeListViewWriter writer, int[] values) {
22342267
writer.startListView();
22352268
for (int v : values) {

‎vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,11 @@
3232
import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector;
3333
import org.apache.arrow.vector.complex.ListVector;
3434
import org.apache.arrow.vector.complex.ListViewVector;
35+
import org.apache.arrow.vector.complex.impl.UnionListViewReader;
3536
import org.apache.arrow.vector.complex.impl.UnionListViewWriter;
3637
import org.apache.arrow.vector.holders.DurationHolder;
3738
import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
39+
import org.apache.arrow.vector.holders.UnionHolder;
3840
import org.apache.arrow.vector.types.TimeUnit;
3941
import org.apache.arrow.vector.types.Types.MinorType;
4042
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -142,6 +144,69 @@ public void testBasicListViewVector() {
142144
}
143145
}
144146

147+
@Test
148+
public void testCopyFromNonEmptyListView() {
149+
try (ListViewVector inVector = ListViewVector.empty("input", allocator);
150+
ListViewVector outVector = ListViewVector.empty("output", allocator)) {
151+
UnionListViewWriter writer = inVector.getWriter();
152+
writer.allocate();
153+
writer.setPosition(0);
154+
writeIntValues(writer, new int[] {10, 20});
155+
writer.setValueCount(1);
156+
157+
outVector.allocateNew();
158+
outVector.copyFrom(0, 0, inVector);
159+
outVector.setValueCount(1);
160+
161+
assertEquals(Arrays.asList(10, 20), outVector.getObject(0));
162+
}
163+
}
164+
165+
@Test
166+
public void testReaderIteratesListViewRangeAndResets() {
167+
try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) {
168+
initializeListViewVector(
169+
listViewVector,
170+
List.of(10, 11, 20, 21, 22),
171+
List.of(1, 1, 1),
172+
List.of(0, 2, 5),
173+
List.of(2, 3, 0));
174+
UnionListViewReader reader = listViewVector.getReader();
175+
176+
reader.setPosition(0);
177+
assertTrue(reader.next());
178+
assertEquals(10, ((Number) reader.reader().readObject()).intValue());
179+
assertTrue(reader.next());
180+
assertEquals(11, ((Number) reader.reader().readObject()).intValue());
181+
assertFalse(reader.next());
182+
assertFalse(reader.next());
183+
184+
reader.setPosition(1);
185+
assertTrue(reader.next());
186+
assertEquals(20, ((Number) reader.reader().readObject()).intValue());
187+
assertTrue(reader.next());
188+
assertEquals(21, ((Number) reader.reader().readObject()).intValue());
189+
assertTrue(reader.next());
190+
assertEquals(22, ((Number) reader.reader().readObject()).intValue());
191+
assertFalse(reader.next());
192+
assertFalse(reader.next());
193+
194+
reader.setPosition(1);
195+
assertTrue(reader.next());
196+
assertEquals(20, ((Number) reader.reader().readObject()).intValue());
197+
198+
reader.setPosition(2);
199+
assertFalse(reader.next());
200+
assertFalse(reader.next());
201+
202+
reader.setPosition(1);
203+
UnionHolder holder = new UnionHolder();
204+
reader.read(2, holder);
205+
assertEquals(22, ((Number) holder.reader.readObject()).intValue());
206+
assertFalse(reader.next());
207+
}
208+
}
209+
145210
@Test
146211
public void testImplicitNullVectors() {
147212
try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) {

0 commit comments

Comments
 (0)