From b89cbe6751eaa0aec4b00d8fecf477d07eb20977 Mon Sep 17 00:00:00 2001 From: Vibhatha Abeykoon Date: Fri, 7 Jun 2024 16:23:57 +0530 Subject: [PATCH] feat: adding rest of the test cases: snapshot 1: requires transferPair to proceed --- .../templates/AbstractFieldWriter.java | 12 + .../main/codegen/templates/BaseWriter.java | 2 + .../main/codegen/templates/StructWriters.java | 22 ++ .../vector/complex/impl/PromotableWriter.java | 3 + .../complex/writer/TestComplexWriter.java | 259 ++++++++++++------ 5 files changed, 217 insertions(+), 81 deletions(-) diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java index 6c2368117f7c2..05071e908e666 100644 --- a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java +++ b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java @@ -184,6 +184,12 @@ public ListWriter list() { return null; } + @Override + public ListWriter listView() { + fail("ListView"); + return null; + } + @Override public MapWriter map() { fail("Map"); @@ -202,6 +208,12 @@ public ListWriter list(String name) { return null; } + @Override + public ListWriter listView(String name) { + fail("ListView"); + return null; + } + @Override public MapWriter map(String name) { fail("Map"); diff --git a/java/vector/src/main/codegen/templates/BaseWriter.java b/java/vector/src/main/codegen/templates/BaseWriter.java index 35df256b324b5..91680c599241d 100644 --- a/java/vector/src/main/codegen/templates/BaseWriter.java +++ b/java/vector/src/main/codegen/templates/BaseWriter.java @@ -62,6 +62,7 @@ public interface StructWriter extends BaseWriter { void copyReaderToField(String name, FieldReader reader); StructWriter struct(String name); ListWriter list(String name); + ListWriter listView(String name); MapWriter map(String name); MapWriter map(String name, boolean keysSorted); void start(); @@ -73,6 +74,7 @@ public interface ListWriter extends BaseWriter { void endList(); StructWriter struct(); ListWriter list(); + ListWriter listView(); MapWriter map(); MapWriter map(boolean keysSorted); void copyReader(FieldReader reader); diff --git a/java/vector/src/main/codegen/templates/StructWriters.java b/java/vector/src/main/codegen/templates/StructWriters.java index b676173ac39d9..dca027a49547c 100644 --- a/java/vector/src/main/codegen/templates/StructWriters.java +++ b/java/vector/src/main/codegen/templates/StructWriters.java @@ -200,6 +200,28 @@ public ListWriter list(String name) { return writer; } + @Override + public ListWriter listView(String name) { + String finalName = handleCase(name); + FieldWriter writer = fields.get(finalName); + int vectorCount = container.size(); + if(writer == null) { + FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.LISTVIEW.getType(), null, null); + writer = new PromotableWriter(container.addOrGet(name, fieldType, ListViewVector.class), container, getNullableStructWriterFactory()); + if (container.size() > vectorCount) { + writer.allocate(); + } + writer.setPosition(idx()); + fields.put(finalName, writer); + } else { + if (writer instanceof PromotableWriter) { + // ensure writers are initialized + ((PromotableWriter)writer).getWriter(MinorType.LISTVIEW); + } + } + return writer; + } + @Override public MapWriter map(String name) { return map(name, false); diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java index c59b997286d2d..e93b233b265db 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java @@ -244,6 +244,9 @@ private void setWriter(ValueVector v) { case LIST: writer = new UnionListWriter((ListVector) vector, nullableStructWriterFactory); break; + case LISTVIEW: + writer = new UnionListViewWriter((ListViewVector) vector, nullableStructWriterFactory); + break; case MAP: writer = new UnionMapWriter((MapVector) vector); break; 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 ad68cbba01fe3..50cdc117f4862 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 @@ -551,9 +551,9 @@ public void listDurationType() { public void listViewDurationType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); - UnionListViewWriter listWriter = new UnionListViewWriter(listViewVector); - createListTypeVectorWithDurationType(listWriter); - listWriter.setValueCount(COUNT); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createListTypeVectorWithDurationType(listViewWriter); + listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithDurationType(listReader); } @@ -615,9 +615,9 @@ public void listViewFixedSizeBinaryType() throws Exception { List buffers = new ArrayList<>(); try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); - UnionListViewWriter listWriter = new UnionListViewWriter(listViewVector); - createListTypeVectorWithFixedSizeBinaryType(listWriter, buffers); - listWriter.setValueCount(COUNT); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createListTypeVectorWithFixedSizeBinaryType(listViewWriter, buffers); + listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithFixedSizeBinaryType(listReader); } @@ -666,9 +666,9 @@ public void listScalarTypeNullable() { public void listViewScalarTypeNullable() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); - UnionListViewWriter listWriter = new UnionListViewWriter(listViewVector); - createScalarTypeVectorWithNullableType(listWriter); - listWriter.setValueCount(COUNT); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createScalarTypeVectorWithNullableType(listViewWriter); + listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkScalarTypeVectorWithNullableType(listReader); } @@ -702,10 +702,10 @@ private void checkListTypeVectorWithStructType(FieldReader reader) { public void listStructType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); - UnionListWriter listWriter = new UnionListWriter(listVector); - StructWriter structWriter = listWriter.struct(); - createListTypeVectorWithStructType(listWriter, structWriter); - listWriter.setValueCount(COUNT); + UnionListWriter listViewWriter = new UnionListWriter(listVector); + StructWriter structWriter = listViewWriter.struct(); + createListTypeVectorWithStructType(listViewWriter, structWriter); + listViewWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); checkListTypeVectorWithStructType(listReader); } @@ -715,10 +715,10 @@ public void listStructType() { public void listViewStructType() { try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { listViewVector.allocateNew(); - UnionListViewWriter listWriter = new UnionListViewWriter(listViewVector); - StructWriter structWriter = listWriter.struct(); - createListTypeVectorWithStructType(listWriter, structWriter); - listWriter.setValueCount(COUNT); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + StructWriter structWriter = listViewWriter.struct(); + createListTypeVectorWithStructType(listViewWriter, structWriter); + listViewWriter.setValueCount(COUNT); UnionListViewReader listReader = new UnionListViewReader(listViewVector); checkListTypeVectorWithStructType(listReader); } @@ -747,7 +747,19 @@ public void listListType() { createNestedListTypeVector(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); - checkListOfLists(listReader); + checkListOfListTypes(listReader); + } + } + + @Test + public void listViewListType() { + try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { + listViewVector.allocateNew(); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createNestedListTypeVector(listViewWriter); + listViewWriter.setValueCount(COUNT); + UnionListViewReader listReader = new UnionListViewReader(listViewVector); + checkListOfListTypes(listReader); } } @@ -777,11 +789,23 @@ public void listListType2() { createNestedListTypeVector2(listWriter); listWriter.setValueCount(COUNT); UnionListReader listReader = new UnionListReader(listVector); - checkListOfLists(listReader); + checkListOfListTypes(listReader); + } + } + + @Test + public void listViewListType2() { + try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { + listViewVector.allocateNew(); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createNestedListTypeVector2(listViewWriter); + listViewWriter.setValueCount(COUNT); + UnionListViewReader listReader = new UnionListViewReader(listViewVector); + checkListOfListTypes(listReader); } } - private void checkListOfLists(final FieldReader reader) { + private void checkListOfListTypes(final FieldReader reader) { for (int i = 0; i < COUNT; i++) { reader.setPosition(i); for (int j = 0; j < i % 7; j++) { @@ -795,29 +819,46 @@ private void checkListOfLists(final FieldReader reader) { } } + private void createUnionListOfListTypeVector(FieldWriter writer) { + for (int i = 0; i < COUNT; i++) { + writer.startList(); + for (int j = 0; j < i % 7; j++) { + FieldWriter innerListWriter = (FieldWriter) writer.list(); + innerListWriter.startList(); + for (int k = 0; k < i % 13; k++) { + if (k % 2 == 0) { + innerListWriter.integer().writeInt(k); + } else { + innerListWriter.bigInt().writeBigInt(k); + } + } + innerListWriter.endList(); + } + writer.endList(); + } + } + @Test public void unionListListType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); - 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++) { - if (k % 2 == 0) { - innerListWriter.integer().writeInt(k); - } else { - innerListWriter.bigInt().writeBigInt(k); - } - } - innerListWriter.endList(); - } - listWriter.endList(); - } + createUnionListOfListTypeVector(listWriter); listWriter.setValueCount(COUNT); - checkUnionList(listVector); + UnionListReader listReader = new UnionListReader(listVector); + checkUnionListType(listReader); + } + } + + @Test + public void unionListViewListType() { + try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { + listViewVector.allocateNew(); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createUnionListOfListTypeVector(listViewWriter); + listViewWriter.setValueCount(COUNT); + UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); + checkUnionListType(listViewReader); } } @@ -829,35 +870,54 @@ public void unionListListType2() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); - ListWriter innerListWriter = listWriter.list(); + createUnionListTypeVector2(listWriter); + listWriter.setValueCount(COUNT); + UnionListReader listReader = new UnionListReader(listVector); + checkUnionListType(listReader); + } + } - for (int i = 0; i < COUNT; i++) { - listWriter.startList(); - for (int j = 0; j < i % 7; j++) { - innerListWriter.startList(); - for (int k = 0; k < i % 13; k++) { - if (k % 2 == 0) { - innerListWriter.integer().writeInt(k); - } else { - innerListWriter.bigInt().writeBigInt(k); - } + /** + * This test is similar to {@link #unionListViewListType()} + * but we get the inner list writer once at the beginning. + */ + @Test + public void unionListViewListType2() { + try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { + listViewVector.allocateNew(); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + createUnionListTypeVector2(listViewWriter); + listViewWriter.setValueCount(COUNT); + UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); + checkUnionListType(listViewReader); + } + } + + private void createUnionListTypeVector2(FieldWriter writer) { + FieldWriter innerListWriter = (FieldWriter) writer.list(); + for (int i = 0; i < COUNT; i++) { + writer.startList(); + for (int j = 0; j < i % 7; j++) { + innerListWriter.startList(); + for (int k = 0; k < i % 13; k++) { + if (k % 2 == 0) { + innerListWriter.integer().writeInt(k); + } else { + innerListWriter.bigInt().writeBigInt(k); } - innerListWriter.endList(); } - listWriter.endList(); + innerListWriter.endList(); } - listWriter.setValueCount(COUNT); - checkUnionList(listVector); + writer.endList(); } } - private void checkUnionList(ListVector listVector) { - UnionListReader listReader = new UnionListReader(listVector); + private void checkUnionListType(FieldReader reader) { for (int i = 0; i < COUNT; i++) { - listReader.setPosition(i); + reader.setPosition(i); for (int j = 0; j < i % 7; j++) { - listReader.next(); - FieldReader innerListReader = listReader.reader(); + reader.next(); + FieldReader innerListReader = reader.reader(); for (int k = 0; k < i % 13; k++) { innerListReader.next(); if (k % 2 == 0) { @@ -875,27 +935,11 @@ public void testListMapType() { try (ListVector listVector = ListVector.empty("list", allocator)) { listVector.allocateNew(); UnionListWriter listWriter = new UnionListWriter(listVector); - MapWriter innerMapWriter = listWriter.map(true); - for (int i = 0; i < COUNT; i++) { - listWriter.startList(); - for (int j = 0; j < i % 7; j++) { - innerMapWriter.startMap(); - for (int k = 0; k < i % 13; k++) { - innerMapWriter.startEntry(); - innerMapWriter.key().integer().writeInt(k); - if (k % 2 == 0) { - innerMapWriter.value().bigInt().writeBigInt(k); - } - innerMapWriter.endEntry(); - } - innerMapWriter.endMap(); - } - listWriter.endList(); - } + createListTypeVectorWithMapType(listWriter); listWriter.setValueCount(COUNT); - checkListMap(listVector); - + UnionListReader listReader = new UnionListReader(listVector); + checkListTypeMap(listReader); // Verify that the map vector has keysSorted = true MapVector mapVector = (MapVector) listVector.getDataVector(); ArrowType arrowType = mapVector.getField().getFieldType().getType(); @@ -903,13 +947,49 @@ public void testListMapType() { } } - private void checkListMap(ListVector listVector) { - UnionListReader listReader = new UnionListReader(listVector); + @Test + public void testListViewMapType() { + try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { + listViewVector.allocateNew(); + UnionListViewWriter listViewWriter = new UnionListViewWriter(listViewVector); + + createListTypeVectorWithMapType(listViewWriter); + listViewWriter.setValueCount(COUNT); + UnionListViewReader listViewReader = new UnionListViewReader(listViewVector); + checkListTypeMap(listViewReader); + // Verify that the map vector has keysSorted = true + MapVector mapVector = (MapVector) listViewVector.getDataVector(); + ArrowType arrowType = mapVector.getField().getFieldType().getType(); + assertTrue(((ArrowType.Map) arrowType).getKeysSorted()); + } + } + + private static void createListTypeVectorWithMapType(FieldWriter writer) { + MapWriter innerMapWriter = writer.map(true); for (int i = 0; i < COUNT; i++) { - listReader.setPosition(i); + writer.startList(); for (int j = 0; j < i % 7; j++) { - listReader.next(); - UnionMapReader mapReader = (UnionMapReader) listReader.reader(); + innerMapWriter.startMap(); + for (int k = 0; k < i % 13; k++) { + innerMapWriter.startEntry(); + innerMapWriter.key().integer().writeInt(k); + if (k % 2 == 0) { + innerMapWriter.value().bigInt().writeBigInt(k); + } + innerMapWriter.endEntry(); + } + innerMapWriter.endMap(); + } + writer.endList(); + } + } + + private void checkListTypeMap(FieldReader reader) { + for (int i = 0; i < COUNT; i++) { + reader.setPosition(i); + for (int j = 0; j < i % 7; j++) { + reader.next(); + UnionMapReader mapReader = (UnionMapReader) reader.reader(); for (int k = 0; k < i % 13; k++) { mapReader.next(); assertEquals(k, mapReader.key().readInteger().intValue(), "record key: " + i); @@ -1433,6 +1513,8 @@ public void testSingleStructWriter1() { Float4Writer float4Writer = singleStructWriter.float4("float4Field"); Float8Writer float8Writer = singleStructWriter.float8("float8Field"); ListWriter listWriter = singleStructWriter.list("listField"); + // TODO: we need to implement transferPair functionality here + ListWriter listViewWriter = singleStructWriter.listView("listViewField"); MapWriter mapWriter = singleStructWriter.map("mapField", false); int intValue = 100; @@ -1456,6 +1538,14 @@ public void testSingleStructWriter1() { listWriter.integer().writeInt(intValue + i + 3); listWriter.endList(); + listViewWriter.setPosition(i); + listViewWriter.startList(); + listViewWriter.integer().writeInt(intValue + i); + listViewWriter.integer().writeInt(intValue + i + 1); + listViewWriter.integer().writeInt(intValue + i + 2); + listViewWriter.integer().writeInt(intValue + i + 3); + listViewWriter.endList(); + mapWriter.setPosition(i); mapWriter.startMap(); mapWriter.startEntry(); @@ -1494,6 +1584,7 @@ public void testSingleStructWriter1() { Float4Reader float4Reader = singleStructReader.reader("float4Field"); Float8Reader float8Reader = singleStructReader.reader("float8Field"); UnionListReader listReader = (UnionListReader) singleStructReader.reader("listField"); + UnionListViewReader listViewReader = (UnionListViewReader) singleStructReader.reader("listViewField"); UnionMapReader mapReader = (UnionMapReader) singleStructReader.reader("mapField"); for (int i = 0; i < initialCapacity; i++) { @@ -1502,6 +1593,7 @@ public void testSingleStructWriter1() { float4Reader.setPosition(i); float8Reader.setPosition(i); listReader.setPosition(i); + listViewReader.setPosition(i); mapReader.setPosition(i); assertEquals(intValue + i, intReader.readInteger().intValue()); @@ -1514,6 +1606,11 @@ public void testSingleStructWriter1() { assertEquals(intValue + i + j, listReader.reader().readInteger().intValue()); } + for (int j = 0; j < 4; j++) { + listViewReader.next(); + assertEquals(intValue + i + j, listViewReader.reader().readInteger().intValue()); + } + for (int k = 0; k < 4; k += 2) { mapReader.next(); assertEquals(intValue + k + i, mapReader.key().readInteger().intValue());