diff --git a/java/vector/src/main/codegen/templates/MapWriters.java b/java/vector/src/main/codegen/templates/MapWriters.java index 696bbf655ca..51327b43af0 100644 --- a/java/vector/src/main/codegen/templates/MapWriters.java +++ b/java/vector/src/main/codegen/templates/MapWriters.java @@ -112,6 +112,11 @@ public MapWriter map(String name) { } writer.setPosition(idx()); fields.put(name.toLowerCase(), writer); + } else { + if (writer instanceof PromotableWriter) { + // ensure writers are initialized + ((PromotableWriter)writer).getWriter(MinorType.MAP); + } } return writer; } @@ -149,6 +154,11 @@ public ListWriter list(String name) { } writer.setPosition(idx()); fields.put(name.toLowerCase(), writer); + } else { + if (writer instanceof PromotableWriter) { + // ensure writers are initialized + ((PromotableWriter)writer).getWriter(MinorType.LIST); + } } return writer; } @@ -210,6 +220,11 @@ public void end() { } writer.setPosition(idx()); fields.put(name.toLowerCase(), writer); + } else { + if (writer instanceof PromotableWriter) { + // ensure writers are initialized + ((PromotableWriter)writer).getWriter(MinorType.${upperName}); + } } return writer; } diff --git a/java/vector/src/main/codegen/templates/UnionWriter.java b/java/vector/src/main/codegen/templates/UnionWriter.java index 460ec1c0d95..efb66f168f5 100644 --- a/java/vector/src/main/codegen/templates/UnionWriter.java +++ b/java/vector/src/main/codegen/templates/UnionWriter.java @@ -25,6 +25,8 @@ package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> +import org.apache.arrow.vector.complex.writer.BaseWriter; +import org.apache.arrow.vector.types.Types.MinorType; /* * This class is generated using freemarker and the ${.template_name} template. @@ -100,6 +102,28 @@ public ListWriter asList() { return getListWriter(); } + BaseWriter getWriter(MinorType minorType) { + switch (minorType) { + case MAP: + return getMapWriter(); + case LIST: + return getListWriter(); + <#list vv.types as type> + <#list type.minor as minor> + <#assign name = minor.class?cap_first /> + <#assign fields = minor.fields!type.fields /> + <#assign uncappedName = name?uncap_first/> + <#if !minor.class?starts_with("Decimal")> + case ${name?upper_case}: + return get${name}Writer(); + + + + default: + throw new UnsupportedOperationException("Unknown type: " + minorType); + } + } + <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> <#assign fields = minor.fields!type.fields /> <#assign uncappedName = name?uncap_first/> 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 c282688530b..94ff82c04bd 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 @@ -94,19 +94,19 @@ public void setPosition(int index) { protected FieldWriter getWriter(MinorType type) { if (state == State.UNION) { - return writer; - } - if (state == State.UNTYPED) { + ((UnionWriter)writer).getWriter(type); + } else if (state == State.UNTYPED) { if (type == null) { + // ??? return null; } ValueVector v = listVector.addOrGetVector(type).getVector(); v.allocateNew(); setWriter(v); writer.setPosition(position); - } - if (type != this.type) { - return promoteToUnion(); + } else if (type != this.type) { + promoteToUnion(); + ((UnionWriter)writer).getWriter(type); } return writer; } @@ -133,7 +133,7 @@ private FieldWriter promoteToUnion() { unionVector.addVector((FieldVector)tp.getTo()); writer = new UnionWriter(unionVector); writer.setPosition(idx()); - for (int i = 0; i < idx(); i++) { + for (int i = 0; i <= idx(); i++) { unionVector.getMutator().setType(i, vector.getMinorType()); } vector = null; 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 398aea915b3..9419f88de5b 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 @@ -45,6 +45,7 @@ import org.apache.arrow.vector.types.pojo.ArrowType.Union; import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.Text; import org.junit.Assert; import org.junit.Test; @@ -362,11 +363,38 @@ public void promotableWriter() { MapReader rootReader = new SingleMapReaderImpl(parent).reader("root"); for (int i = 0; i < 100; i++) { rootReader.setPosition(i); - Assert.assertEquals(i, rootReader.reader("a").readLong().intValue()); + FieldReader reader = rootReader.reader("a"); + Long value = reader.readLong(); + Assert.assertNotNull("index: " + i, value); + Assert.assertEquals(i, value.intValue()); } for (int i = 100; i < 200; i++) { rootReader.setPosition(i); - Assert.assertEquals(Integer.toString(i), rootReader.reader("a").readText().toString()); + FieldReader reader = rootReader.reader("a"); + Text value = reader.readText(); + Assert.assertEquals(Integer.toString(i), value.toString()); } } + + /** + * Even without writing to the writer, the union schema is created correctly + */ + @Test + public void promotableWriterSchema() { + MapVector parent = new MapVector("parent", allocator, null); + ComplexWriter writer = new ComplexWriterImpl("root", parent); + MapWriter rootWriter = writer.rootAsMap(); + BigIntWriter bigIntWriter = rootWriter.bigInt("a"); + VarCharWriter varCharWriter = rootWriter.varChar("a"); + + Field field = parent.getField().getChildren().get(0).getChildren().get(0); + Assert.assertEquals("a", field.getName()); + Assert.assertEquals(Union.TYPE_TYPE, field.getType().getTypeType()); + + Assert.assertEquals(Int.TYPE_TYPE, field.getChildren().get(0).getType().getTypeType()); + Int intType = (Int) field.getChildren().get(0).getType(); + Assert.assertEquals(64, intType.getBitWidth()); + Assert.assertTrue(intType.getIsSigned()); + Assert.assertEquals(Utf8.TYPE_TYPE, field.getChildren().get(1).getType().getTypeType()); + } } \ No newline at end of file