Skip to content

Commit 1ae7e00

Browse files
author
adeneche
committed
updated TestComplexWriter to ensure position is set properly by the various writers
1 parent 7d5aefc commit 1ae7e00

File tree

1 file changed

+108
-52
lines changed

1 file changed

+108
-52
lines changed

‎java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java‎

Lines changed: 108 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,10 @@ public void simpleNestedTypes() {
6565
IntWriter intWriter = rootWriter.integer("int");
6666
BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt");
6767
for (int i = 0; i < COUNT; i++) {
68-
intWriter.setPosition(i);
68+
rootWriter.start();
6969
intWriter.writeInt(i);
70-
bigIntWriter.setPosition(i);
7170
bigIntWriter.writeBigInt(i);
71+
rootWriter.end();
7272
}
7373
writer.setValueCount(COUNT);
7474
MapReader rootReader = new SingleMapReaderImpl(parent).reader("root");
@@ -83,23 +83,52 @@ public void simpleNestedTypes() {
8383

8484
@Test
8585
public void nullableMap() {
86-
MapVector parent = new MapVector("parent", allocator, null);
87-
ComplexWriter writer = new ComplexWriterImpl("root", parent);
88-
MapWriter rootWriter = writer.rootAsMap();
89-
for (int i = 0; i < COUNT; i++) {
90-
rootWriter.setPosition(i);
91-
rootWriter.start();
92-
if (i % 2 == 0) {
93-
MapWriter mapWriter = rootWriter.map("map");
94-
mapWriter.setPosition(i);
95-
mapWriter.start();
96-
mapWriter.bigInt("nested").writeBigInt(i);
97-
mapWriter.end();
86+
try (MapVector mapVector = new MapVector("parent", allocator, null)) {
87+
ComplexWriter writer = new ComplexWriterImpl("root", mapVector);
88+
MapWriter rootWriter = writer.rootAsMap();
89+
for (int i = 0; i < COUNT; i++) {
90+
rootWriter.start();
91+
if (i % 2 == 0) {
92+
MapWriter mapWriter = rootWriter.map("map");
93+
mapWriter.setPosition(i);
94+
mapWriter.start();
95+
mapWriter.bigInt("nested").writeBigInt(i);
96+
mapWriter.end();
97+
}
98+
rootWriter.end();
9899
}
99-
rootWriter.end();
100+
writer.setValueCount(COUNT);
101+
checkNullableMap(mapVector);
100102
}
101-
writer.setValueCount(COUNT);
102-
MapReader rootReader = new SingleMapReaderImpl(parent).reader("root");
103+
}
104+
105+
/**
106+
* This test is similar to {@link #nullableMap()} ()} but we get the inner map writer once at the beginning
107+
*/
108+
@Test
109+
public void nullableMap2() {
110+
try (MapVector mapVector = new MapVector("parent", allocator, null)) {
111+
ComplexWriter writer = new ComplexWriterImpl("root", mapVector);
112+
MapWriter rootWriter = writer.rootAsMap();
113+
MapWriter mapWriter = rootWriter.map("map");
114+
115+
for (int i = 0; i < COUNT; i++) {
116+
rootWriter.start();
117+
if (i % 2 == 0) {
118+
mapWriter.setPosition(i);
119+
mapWriter.start();
120+
mapWriter.bigInt("nested").writeBigInt(i);
121+
mapWriter.end();
122+
}
123+
rootWriter.end();
124+
}
125+
writer.setValueCount(COUNT);
126+
checkNullableMap(mapVector);
127+
}
128+
}
129+
130+
private void checkNullableMap(MapVector mapVector) {
131+
MapReader rootReader = new SingleMapReaderImpl(mapVector).reader("root");
103132
for (int i = 0; i < COUNT; i++) {
104133
rootReader.setPosition(i);
105134
assertTrue("index is set: " + i, rootReader.isSet());
@@ -113,7 +142,6 @@ public void nullableMap() {
113142
assertNull("index is not set: " + i, map.readObject());
114143
}
115144
}
116-
parent.close();
117145
}
118146

119147
@Test
@@ -129,7 +157,6 @@ public void testList() {
129157
rootWriter.list("list").endList();
130158
rootWriter.end();
131159

132-
rootWriter.setPosition(1);
133160
rootWriter.start();
134161
rootWriter.bigInt("int").writeBigInt(1);
135162
rootWriter.end();
@@ -152,7 +179,6 @@ public void listScalarType() {
152179
listVector.allocateNew();
153180
UnionListWriter listWriter = new UnionListWriter(listVector);
154181
for (int i = 0; i < COUNT; i++) {
155-
listWriter.setPosition(i);
156182
listWriter.startList();
157183
for (int j = 0; j < i % 7; j++) {
158184
listWriter.writeInt(j);
@@ -206,7 +232,6 @@ public void listMapType() {
206232
UnionListWriter listWriter = new UnionListWriter(listVector);
207233
MapWriter mapWriter = listWriter.map();
208234
for (int i = 0; i < COUNT; i++) {
209-
listWriter.setPosition(i);
210235
listWriter.startList();
211236
for (int j = 0; j < i % 7; j++) {
212237
mapWriter.start();
@@ -234,7 +259,6 @@ public void listListType() {
234259
listVector.allocateNew();
235260
UnionListWriter listWriter = new UnionListWriter(listVector);
236261
for (int i = 0; i < COUNT; i++) {
237-
listWriter.setPosition(i);
238262
listWriter.startList();
239263
for (int j = 0; j < i % 7; j++) {
240264
ListWriter innerListWriter = listWriter.list();
@@ -251,6 +275,32 @@ public void listListType() {
251275
}
252276
}
253277

278+
/**
279+
* This test is similar to {@link #listListType()} but we get the inner list writer once at the beginning
280+
*/
281+
@Test
282+
public void listListType2() {
283+
try (ListVector listVector = new ListVector("list", allocator, null)) {
284+
listVector.allocateNew();
285+
UnionListWriter listWriter = new UnionListWriter(listVector);
286+
ListWriter innerListWriter = listWriter.list();
287+
288+
for (int i = 0; i < COUNT; i++) {
289+
listWriter.startList();
290+
for (int j = 0; j < i % 7; j++) {
291+
innerListWriter.startList();
292+
for (int k = 0; k < i % 13; k++) {
293+
innerListWriter.integer().writeInt(k);
294+
}
295+
innerListWriter.endList();
296+
}
297+
listWriter.endList();
298+
}
299+
listWriter.setValueCount(COUNT);
300+
checkListOfLists(listVector);
301+
}
302+
}
303+
254304
private void checkListOfLists(final ListVector listVector) {
255305
UnionListReader listReader = new UnionListReader(listVector);
256306
for (int i = 0; i < COUNT; i++) {
@@ -266,56 +316,63 @@ private void checkListOfLists(final ListVector listVector) {
266316
}
267317
}
268318

269-
/**
270-
* This test is similar to {@link #listListType()} but we get the inner list writer once at the beginning
271-
*/
272319
@Test
273-
public void listListType2() {
320+
public void unionListListType() {
274321
try (ListVector listVector = new ListVector("list", allocator, null)) {
275322
listVector.allocateNew();
276323
UnionListWriter listWriter = new UnionListWriter(listVector);
277-
ListWriter innerListWriter = listWriter.list();
278-
279324
for (int i = 0; i < COUNT; i++) {
280-
listWriter.setPosition(i);
281325
listWriter.startList();
282326
for (int j = 0; j < i % 7; j++) {
327+
ListWriter innerListWriter = listWriter.list();
283328
innerListWriter.startList();
284329
for (int k = 0; k < i % 13; k++) {
285-
innerListWriter.integer().writeInt(k);
330+
if (k % 2 == 0) {
331+
innerListWriter.integer().writeInt(k);
332+
} else {
333+
innerListWriter.bigInt().writeBigInt(k);
334+
}
286335
}
287336
innerListWriter.endList();
288337
}
289338
listWriter.endList();
290339
}
291340
listWriter.setValueCount(COUNT);
292-
checkListOfLists(listVector);
341+
checkUnionList(listVector);
293342
}
294343
}
295344

345+
/**
346+
* This test is similar to {@link #unionListListType()} but we get the inner list writer once at the beginning
347+
*/
296348
@Test
297-
public void unionListListType() {
298-
ListVector listVector = new ListVector("list", allocator, null);
299-
listVector.allocateNew();
300-
UnionListWriter listWriter = new UnionListWriter(listVector);
301-
for (int i = 0; i < COUNT; i++) {
302-
listWriter.setPosition(i);
303-
listWriter.startList();
304-
for (int j = 0; j < i % 7; j++) {
305-
ListWriter innerListWriter = listWriter.list();
306-
innerListWriter.startList();
307-
for (int k = 0; k < i % 13; k++) {
308-
if (k % 2 == 0) {
309-
innerListWriter.integer().writeInt(k);
310-
} else {
311-
innerListWriter.bigInt().writeBigInt(k);
349+
public void unionListListType2() {
350+
try (ListVector listVector = new ListVector("list", allocator, null)) {
351+
listVector.allocateNew();
352+
UnionListWriter listWriter = new UnionListWriter(listVector);
353+
ListWriter innerListWriter = listWriter.list();
354+
355+
for (int i = 0; i < COUNT; i++) {
356+
listWriter.startList();
357+
for (int j = 0; j < i % 7; j++) {
358+
innerListWriter.startList();
359+
for (int k = 0; k < i % 13; k++) {
360+
if (k % 2 == 0) {
361+
innerListWriter.integer().writeInt(k);
362+
} else {
363+
innerListWriter.bigInt().writeBigInt(k);
364+
}
312365
}
366+
innerListWriter.endList();
313367
}
314-
innerListWriter.endList();
368+
listWriter.endList();
315369
}
316-
listWriter.endList();
370+
listWriter.setValueCount(COUNT);
371+
checkUnionList(listVector);
317372
}
318-
listWriter.setValueCount(COUNT);
373+
}
374+
375+
private void checkUnionList(ListVector listVector) {
319376
UnionListReader listReader = new UnionListReader(listVector);
320377
for (int i = 0; i < COUNT; i++) {
321378
listReader.setPosition(i);
@@ -332,7 +389,6 @@ public void unionListListType() {
332389
}
333390
}
334391
}
335-
listVector.clear();
336392
}
337393

338394
@Test
@@ -415,8 +471,8 @@ public void promotableWriterSchema() {
415471
MapVector parent = new MapVector("parent", allocator, null);
416472
ComplexWriter writer = new ComplexWriterImpl("root", parent);
417473
MapWriter rootWriter = writer.rootAsMap();
418-
BigIntWriter bigIntWriter = rootWriter.bigInt("a");
419-
VarCharWriter varCharWriter = rootWriter.varChar("a");
474+
rootWriter.bigInt("a");
475+
rootWriter.varChar("a");
420476

421477
Field field = parent.getField().getChildren().get(0).getChildren().get(0);
422478
Assert.assertEquals("a", field.getName());

0 commit comments

Comments
 (0)