GH-87: [Vector] Add ExtensionWriter

vector/src/main/codegen/templates/AbstractFieldWriter.java

@@ -107,6 +107,16 @@ public void endEntry() {
     throw new IllegalStateException(String.format("You tried to end a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
   }
 
+  public <T extends ExtensionHolder> void write(T var1)  {
+    this.fail("ExtensionType");
+  }
+  public void writeExtensionType(Object var1)  {
+    this.fail("ExtensionType");
+  }
+  public <T extends ExtensionTypeWriterFactory> void addExtensionTypeFactory(T var1) {
+    this.fail("ExtensionType");
+  }
+
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
   <#assign fields = minor.fields!type.fields />
   <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) />
@@ -241,6 +251,18 @@ public MapWriter map(String name, boolean keysSorted) {
     fail("Map");
     return null;
   }
+
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    fail("Extension");
+    return null;
+  }
+
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    fail("Extension");
+    return null;
+  }
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java

@@ -293,6 +293,11 @@ public MapWriter map(boolean keysSorted) {
     return getWriter(MinorType.MAP, new ArrowType.Map(keysSorted));
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    return getWriter(MinorType.EXTENSIONTYPE).extension(arrowType);
+  }
+
   @Override
   public StructWriter struct(String name) {
     return getWriter(MinorType.STRUCT).struct(name);
@@ -318,6 +323,11 @@ public MapWriter map(String name, boolean keysSorted) {
     return getWriter(MinorType.STRUCT).map(name, keysSorted);
   }
 
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    return getWriter(MinorType.EXTENSIONTYPE).extension(name, arrowType);
+  }
+
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

vector/src/main/codegen/templates/BaseWriter.java

@@ -61,6 +61,7 @@ public interface StructWriter extends BaseWriter {
 
     void copyReaderToField(String name, FieldReader reader);
     StructWriter struct(String name);
+    ExtensionWriter extension(String name, ArrowType arrowType);
     ListWriter list(String name);
     ListWriter listView(String name);
     MapWriter map(String name);
@@ -79,6 +80,7 @@ public interface ListWriter extends BaseWriter {
     ListWriter listView();
     MapWriter map();
     MapWriter map(boolean keysSorted);
+    ExtensionWriter extension(ArrowType arrowType);
     void copyReader(FieldReader reader);
 
     <#list vv.types as type><#list type.minor as minor>
@@ -101,6 +103,13 @@ public interface MapWriter extends ListWriter {
     MapWriter value();
   }
 
+  public interface ExtensionWriter extends BaseWriter {
+    void writeNull();
+    <T extends ExtensionHolder> void write(T var1);
+    void writeExtensionType(Object var1);
+    <T extends ExtensionTypeWriterFactory> void addExtensionTypeFactory(T var1);
+  }
+
   public interface ScalarWriter extends
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Writer, </#list></#list> BaseWriter {}
 

vector/src/main/codegen/templates/PromotableWriter.java

@@ -285,6 +285,9 @@ protected void setWriter(ValueVector v) {
       case UNION:
         writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory);
         break;
+      case EXTENSIONTYPE:
+        writer = new UnionExtensionWriter((ExtensionTypeVector) vector);
+        break;
       default:
         writer = type.getNewFieldWriter(vector);
         break;
@@ -316,6 +319,7 @@ protected boolean requiresArrowType(MinorType type) {
         || type == MinorType.MAP
         || type == MinorType.DURATION
         || type == MinorType.FIXEDSIZEBINARY
+        || type == MinorType.EXTENSIONTYPE
         || (type.name().startsWith("TIMESTAMP") && type.name().endsWith("TZ"));
   }
 
@@ -536,6 +540,16 @@ public void writeLargeVarChar(String value) {
     getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value);
   }
 
+  @Override
+  public void writeExtensionType(Object value) {
+    getWriter(MinorType.EXTENSIONTYPE).writeExtensionType(value);
+  }
+
+  @Override
+  public <T extends ExtensionTypeWriterFactory> void addExtensionTypeFactory(T var1) {
+    getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeFactory(var1);
+  }
+
   @Override
   public void allocate() {
     getWriter().allocate();

vector/src/main/codegen/templates/StructWriters.java

@@ -83,6 +83,9 @@ public class ${mode}StructWriter extends AbstractFieldWriter {
         fields.put(handleCase(child.getName()), writer);
         break;
       }
+      case EXTENSIONTYPE:
+        extension(child.getName(), child.getType());
+        break;
       case UNION:
         FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.UNION.getType(), null, null);
         UnionWriter writer = new UnionWriter(container.addOrGet(child.getName(), fieldType, UnionVector.class), getNullableStructWriterFactory());
@@ -159,6 +162,29 @@ public StructWriter struct(String name) {
     return writer;
   }
 
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    String finalName = handleCase(name);
+    FieldWriter writer = fields.get(finalName);
+    if(writer == null){
+      int vectorCount=container.size();
+      FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null);
+      ExtensionTypeVector vector = container.addOrGet(name, fieldType, ExtensionTypeVector.class);
+      writer = new PromotableWriter(vector, container, getNullableStructWriterFactory());
+      if(vectorCount != container.size()) {
+        writer.allocate();
+      }
+      writer.setPosition(idx());
+      fields.put(finalName, writer);
+    } else {
+      if (writer instanceof PromotableWriter) {
+        // ensure writers are initialized
+        ((PromotableWriter)writer).getWriter(MinorType.EXTENSIONTYPE, arrowType);
+      }
+    }
+    return (ExtensionWriter) writer;
+  }
+
   @Override
   public void close() throws Exception {
     clear();

vector/src/main/codegen/templates/UnionListWriter.java

@@ -201,6 +201,17 @@ public MapWriter map(String name, boolean keysSorted) {
     return mapWriter;
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    writer.extension(arrowType);
+    return writer;
+  }
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    ExtensionWriter extensionWriter = writer.extension(name, arrowType);
+    return extensionWriter;
+  }
+
   <#if listName == "LargeList">
   @Override
   public void startList() {
@@ -323,6 +334,18 @@ public void writeNull() {
     }
   }
 
+  @Override
+  public void writeExtensionType(Object value) {
+    writer.writeExtensionType(value);
+  }
+  @Override
+  public <T extends ExtensionTypeWriterFactory> void addExtensionTypeFactory(T var1) {
+    writer.addExtensionTypeFactory(var1);
+  }
+  public <T extends ExtensionHolder> void write(T var1) {
+    writer.write(var1);
+  }
+
   <#list vv.types as type>
     <#list type.minor as minor>
       <#assign name = minor.class?cap_first />

vector/src/main/codegen/templates/UnionMapWriter.java

@@ -231,4 +231,16 @@ public MapWriter map() {
         return super.map();
     }
   }
+
+  @Override
+  public ExtensionWriter extension(ArrowType type) {
+    switch (mode) {
+      case KEY:
+        return entryWriter.extension(MapVector.KEY_NAME, type);
+      case VALUE:
+        return entryWriter.extension(MapVector.VALUE_NAME, type);
+      default:
+        return super.extension(type);
+    }
+  }
 }

vector/src/main/codegen/templates/UnionWriter.java

@@ -213,6 +213,10 @@ public MapWriter asMap(ArrowType arrowType) {
     return getMapWriter(arrowType);
   }
 
+  private ExtensionWriter getExtensionWriter(ArrowType arrowType) {
+    throw new UnsupportedOperationException("ExtensionTypes are not supported yet.");
+  }
+
   BaseWriter getWriter(MinorType minorType) {
     return getWriter(minorType, null);
   }
@@ -227,6 +231,8 @@ BaseWriter getWriter(MinorType minorType, ArrowType arrowType) {
       return getListViewWriter();
     case MAP:
       return getMapWriter(arrowType);
+    case EXTENSIONTYPE:
+      return getExtensionWriter(arrowType);
     <#list vv.types as type>
       <#list type.minor as minor>
         <#assign name = minor.class?cap_first />
@@ -460,6 +466,20 @@ public MapWriter map(String name, boolean keysSorted) {
     return getStructWriter().map(name, keysSorted);
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    data.setType(idx(), MinorType.EXTENSIONTYPE);
+    getListWriter().setPosition(idx());
+    return getListWriter().extension(arrowType);
+  }
+
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    data.setType(idx(), MinorType.EXTENSIONTYPE);
+    getStructWriter().setPosition(idx());
+    return getStructWriter().extension(name, arrowType);
+  }
+
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java

@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector.complex.impl;
+
+import org.apache.arrow.vector.ExtensionTypeVector;
+
+public interface ExtensionTypeWriterFactory<T extends AbstractFieldWriter> {
+  T getWriterImpl(ExtensionTypeVector vector);
+}

vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java

@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector.complex.impl;
+
+import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.pojo.Field;
+
+public class UnionExtensionWriter extends AbstractFieldWriter {
+  protected ExtensionTypeVector vector;
+  protected AbstractFieldWriter writer;
+
+  public UnionExtensionWriter(ExtensionTypeVector vector) {
+    this.vector = vector;
+  }
+
+  @Override
+  public void allocate() {
+    vector.allocateNew();
+  }
+
+  @Override
+  public void clear() {
+    vector.clear();
+  }
+
+  @Override
+  public int getValueCapacity() {
+    return vector.getValueCapacity();
+  }
+
+  @Override
+  public Field getField() {
+    return vector.getField();
+  }
+
+  @Override
+  public void close() throws Exception {
+    vector.close();
+  }
+
+  @Override
+  public void writeExtensionType(Object var1) {
+    this.writer.writeExtensionType(var1);
+  }
+
+  @Override
+  public <S extends ExtensionTypeWriterFactory> void addExtensionTypeFactory(S var1) {
+    this.writer = var1.getWriterImpl(vector);
+    this.writer.setPosition(idx());
+  }
+
+  public <T extends ExtensionHolder> void write(T var1) {
+    this.writer.write(var1);
+  }
+
+  @Override
+  public void setPosition(int index) {
+    super.setPosition(index);
+    if (this.writer != null) {
+      this.writer.setPosition(index);
+    }
+  }
+}

vector/src/main/java/org/apache/arrow/vector/complex/writer/FieldWriter.java

@@ -25,7 +25,7 @@
  * Composite of all writer types. Writers are convenience classes for incrementally adding values to
  * {@linkplain org.apache.arrow.vector.ValueVector}s.
  */
-public interface FieldWriter extends StructWriter, ListWriter, MapWriter, ScalarWriter {
+public interface FieldWriter extends StructWriter, ListWriter, MapWriter, ScalarWriter, ExtensionWriter {
   void allocate();
 
   void clear();

vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java

@@ -0,0 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector.holders;
+
+public abstract class ExtensionHolder implements ValueHolder {
+  public int isSet;
+}