ARROW-333: Make writers update their internal schema even when no data is written

Make PromotableWriter predefine writers when asked

Author: Julien Le Dem <julien@dremio.com>

Closes #170 from julienledem/promotable_writer_preset and squashes the following commits:

972eb9c [Julien Le Dem] ARROW-333: Make writers update their internal schema even when no data is written Make PromotableWriter predefine writers when asked

Author: julienledem

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;
   }

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();
+      </#if>
+    </#list>
+    </#list>
+    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/>

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;

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());
+  }
 }