diff --git a/java/vector/src/main/codegen/templates/PromotableViewWriter.java b/java/vector/src/main/codegen/templates/PromotableViewWriter.java index 73adb6bfc7334..618de8a026d66 100644 --- a/java/vector/src/main/codegen/templates/PromotableViewWriter.java +++ b/java/vector/src/main/codegen/templates/PromotableViewWriter.java @@ -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" /> diff --git a/java/vector/src/main/codegen/templates/UnionReader.java b/java/vector/src/main/codegen/templates/UnionReader.java index 243bd832255c2..615ea3a536a15 100644 --- a/java/vector/src/main/codegen/templates/UnionReader.java +++ b/java/vector/src/main/codegen/templates/UnionReader.java @@ -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> @@ -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() { diff --git a/java/vector/src/main/codegen/templates/UnionVector.java b/java/vector/src/main/codegen/templates/UnionVector.java index e0d7faf4c23ce..e0fd0e4644313 100644 --- a/java/vector/src/main/codegen/templates/UnionVector.java +++ b/java/vector/src/main/codegen/templates/UnionVector.java @@ -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: diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java index 19eac31b173dc..80febfe507229 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java @@ -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; @@ -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); @@ -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); @@ -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); @@ -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 @@ -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(); }