Skip to content

Commit

Permalink
fix: updating test cases for new interface
Browse files Browse the repository at this point in the history
  • Loading branch information
@vibhatha
vibhatha committed Jul 22, 2024
1 parent 8c60a39 commit 682dd6d
Show file tree
Hide file tree
Showing 4 changed files with 122 additions and 61 deletions.
5 changes: 5 additions & 0 deletions java/vector/src/main/codegen/templates/PromotableViewWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,11 @@ protected FieldWriter getWriter(MinorType type, ArrowType arrowType) {
return writer;
}

@Override
public StructWriter struct() {
return getWriter(MinorType.LISTVIEW).struct();
}

<#list vv.types as type><#list type.minor as minor>
<#assign lowerName = minor.class?uncap_first />
<#if lowerName == "int" ><#assign lowerName = "integer" /></#if>
Expand Down
13 changes: 13 additions & 0 deletions java/vector/src/main/codegen/templates/UnionReader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,8 @@ private FieldReader getReaderForIndex(int index) {
return (FieldReader) getStruct();
case LIST:
return (FieldReader) getList();
case LISTVIEW:
return (FieldReader) getListView();
case MAP:
return (FieldReader) getMap();
<#list vv.types as type>
Expand Down Expand Up @@ -130,6 +132,17 @@ private FieldReader getList() {
return listReader;
}

private UnionListViewReader listViewReader;

private FieldReader getListView() {
if (listViewReader == null) {
listViewReader = new UnionListViewReader(data.getListView());
listViewReader.setPosition(idx());
readers[MinorType.LISTVIEW.ordinal()] = listViewReader;
}
return listViewReader;
}

private UnionMapReader mapReader;

private FieldReader getMap() {
Expand Down
2 changes: 2 additions & 0 deletions java/vector/src/main/codegen/templates/UnionVector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -717,6 +717,8 @@ public ValueVector getVectorByType(int typeId, ArrowType arrowType) {
return getStruct();
case LIST:
return getList();
case LISTVIEW:
return getListView();
case MAP:
return getMap(name, arrowType);
default:
Expand Down
163 changes: 102 additions & 61 deletions java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,6 @@
import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder;
import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;
import org.apache.arrow.vector.types.pojo.ArrowType.Int;
Expand Down Expand Up @@ -270,58 +269,6 @@ private void checkListTypeVectorWithTimeStampMilliTZType(FieldReader reader) {
}
}

private void complexCopierHelper(MinorType type) {
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
StructWriter rootWriter = writer.rootAsStruct();
ListWriter listWriter;
if (type == MinorType.LIST) {
listWriter = rootWriter.list("list");
} else if (type == MinorType.LISTVIEW) {
listWriter = rootWriter.listView("listView");
} else {
throw new UnsupportedOperationException(
"Unsupported type" + type + "for complexCopierHelper util.");
}
StructWriter innerStructWriter = listWriter.struct();
IntWriter outerIntWriter = listWriter.integer();
rootWriter.start();
listWriter.startList();
outerIntWriter.writeInt(1);
outerIntWriter.writeInt(2);
innerStructWriter.start();
IntWriter intWriter = innerStructWriter.integer("a");
intWriter.writeInt(1);
innerStructWriter.end();
innerStructWriter.start();
intWriter = innerStructWriter.integer("a");
intWriter.writeInt(2);
innerStructWriter.end();
listWriter.endList();
rootWriter.end();
writer.setValueCount(1);

StructVector structVector = (StructVector) parent.getChild("root");
TransferPair tp = structVector.getTransferPair(allocator);
tp.splitAndTransfer(0, 1);
NonNullableStructVector toStructVector = (NonNullableStructVector) tp.getTo();
JsonStringHashMap<?, ?> toMapValue = (JsonStringHashMap<?, ?>) toStructVector.getObject(0);
JsonStringArrayList<?> object;
if (type == MinorType.LIST) {
object = (JsonStringArrayList<?>) toMapValue.get("list");
} else {
object = (JsonStringArrayList<?>) toMapValue.get("listView");
}
assertEquals(1, object.get(0));
assertEquals(2, object.get(1));
JsonStringHashMap<?, ?> innerStruct = (JsonStringHashMap<?, ?>) object.get(2);
assertEquals(1, innerStruct.get("a"));
innerStruct = (JsonStringHashMap<?, ?>) object.get(3);
assertEquals(2, innerStruct.get("a"));
toStructVector.close();
}
}

private void createNullsWithListWriters(FieldWriter writer) {
for (int i = 0; i < COUNT; i++) {
writer.setPosition(i);
Expand Down Expand Up @@ -971,7 +918,18 @@ public void listListType() {
try (ListVector listVector = ListVector.empty("list", allocator)) {
listVector.allocateNew();
UnionListWriter listWriter = new UnionListWriter(listVector);
createNestedListTypeVector(listWriter);
for (int i = 0; i < COUNT; i++) {
listWriter.startList();
for (int j = 0; j < i % 7; j++) {
ListWriter innerListWriter = listWriter.list();
innerListWriter.startList();
for (int k = 0; k < i % 13; k++) {
innerListWriter.integer().writeInt(k);
}
innerListWriter.endList();
}
listWriter.endList();
}
listWriter.setValueCount(COUNT);
UnionListReader listReader = new UnionListReader(listVector);
checkListOfListTypes(listReader);
Expand All @@ -983,7 +941,18 @@ public void listViewListType() {
try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) {
listViewVector.allocateNew();
UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector);
createNestedListTypeVector(listViewWriter);
for (int i = 0; i < COUNT; i++) {
listViewWriter.startListView();
for (int j = 0; j < i % 7; j++) {
ListWriter innerListWriter = listViewWriter.listView();
innerListWriter.startListView();
for (int k = 0; k < i % 13; k++) {
innerListWriter.integer().writeInt(k);
}
innerListWriter.endListView();
}
listViewWriter.endListView();
}
listViewWriter.setValueCount(COUNT);
UnionListViewReader listReader = new UnionListViewReader(listViewVector);
checkListOfListTypes(listReader);
Expand Down Expand Up @@ -1568,12 +1537,84 @@ public void fixedSizeBinaryWriters() throws Exception {

@Test
public void complexCopierWithList() {
complexCopierHelper(MinorType.LIST);
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
StructWriter rootWriter = writer.rootAsStruct();
ListWriter listWriter = rootWriter.list("list");

StructWriter innerStructWriter = listWriter.struct();
IntWriter outerIntWriter = listWriter.integer();
rootWriter.start();
listWriter.startList();
outerIntWriter.writeInt(1);
outerIntWriter.writeInt(2);
innerStructWriter.start();
IntWriter intWriter = innerStructWriter.integer("a");
intWriter.writeInt(1);
innerStructWriter.end();
innerStructWriter.start();
intWriter = innerStructWriter.integer("a");
intWriter.writeInt(2);
innerStructWriter.end();
listWriter.endList();
rootWriter.end();
writer.setValueCount(1);

StructVector structVector = (StructVector) parent.getChild("root");
TransferPair tp = structVector.getTransferPair(allocator);
tp.splitAndTransfer(0, 1);
NonNullableStructVector toStructVector = (NonNullableStructVector) tp.getTo();
JsonStringHashMap<?, ?> toMapValue = (JsonStringHashMap<?, ?>) toStructVector.getObject(0);
JsonStringArrayList<?> object = (JsonStringArrayList<?>) toMapValue.get("list");
assertEquals(1, object.get(0));
assertEquals(2, object.get(1));
JsonStringHashMap<?, ?> innerStruct = (JsonStringHashMap<?, ?>) object.get(2);
assertEquals(1, innerStruct.get("a"));
innerStruct = (JsonStringHashMap<?, ?>) object.get(3);
assertEquals(2, innerStruct.get("a"));
toStructVector.close();
}
}

@Test
public void complexCopierWithListView() {
complexCopierHelper(MinorType.LISTVIEW);
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
StructWriter rootWriter = writer.rootAsStruct();
ListWriter listViewWriter = rootWriter.listView("listView");

StructWriter innerStructWriter = listViewWriter.struct();
IntWriter outerIntWriter = listViewWriter.integer();
rootWriter.start();
listViewWriter.startListView();
outerIntWriter.writeInt(1);
outerIntWriter.writeInt(2);
innerStructWriter.start();
IntWriter intWriter = innerStructWriter.integer("a");
intWriter.writeInt(1);
innerStructWriter.end();
innerStructWriter.start();
intWriter = innerStructWriter.integer("a");
intWriter.writeInt(2);
innerStructWriter.end();
listViewWriter.endListView();
rootWriter.end();
writer.setValueCount(1);

StructVector structVector = (StructVector) parent.getChild("root");
TransferPair tp = structVector.getTransferPair(allocator);
tp.splitAndTransfer(0, 1);
NonNullableStructVector toStructVector = (NonNullableStructVector) tp.getTo();
JsonStringHashMap<?, ?> toMapValue = (JsonStringHashMap<?, ?>) toStructVector.getObject(0);
JsonStringArrayList<?> object = (JsonStringArrayList<?>) toMapValue.get("listView");
assertEquals(1, object.get(0));
assertEquals(2, object.get(1));
JsonStringHashMap<?, ?> innerStruct = (JsonStringHashMap<?, ?>) object.get(2);
assertEquals(1, innerStruct.get("a"));
innerStruct = (JsonStringHashMap<?, ?>) object.get(3);
assertEquals(2, innerStruct.get("a"));
toStructVector.close();
}
}

@Test
Expand Down Expand Up @@ -1805,17 +1846,17 @@ public void testListViewOfListViewWriterWithNulls() {
for (int i = 0; i < COUNT; i++) {
listViewWriter.setPosition(i);
if (i % 2 == 0) {
listViewWriter.startList();
listViewWriter.startListView();
if (i % 4 == 0) {
listViewWriter.listView().writeNull();
} else {
listViewWriter.listView().startList();
listViewWriter.listView().startListView();
listViewWriter.listView().integer().writeNull();
listViewWriter.listView().integer().writeInt(i);
listViewWriter.listView().integer().writeInt(i * 2);
listViewWriter.listView().endList();
listViewWriter.listView().endListView();
}
listViewWriter.endList();
listViewWriter.endListView();
} else {
listViewWriter.writeNull();
}
Expand Down

0 comments on commit 682dd6d

Please sign in to comment.