Make ComplexVectors support writing of ExtensionType

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

@@ -15,7 +15,8 @@
  * limitations under the License.
  */
 
-<@pp.dropOutputFile />
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
+import org.apache.arrow.vector.holders.Float2Holder;<@pp.dropOutputFile />
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/AbstractFieldWriter.java" />
 
 
@@ -107,6 +108,17 @@ 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 +253,19 @@ 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>

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

@@ -283,6 +283,11 @@ public ListWriter listView() {
     return getWriter(MinorType.LISTVIEW).listView();
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    return getWriter(MinorType.EXTENSIONTYPE).extension(arrowType);
+  }
+
   @Override
   public MapWriter map() {
     return getWriter(MinorType.LIST).map();
@@ -313,6 +318,11 @@ public MapWriter map(String name) {
     return getWriter(MinorType.STRUCT).map(name);
   }
 
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    return getWriter(MinorType.EXTENSIONTYPE).extension(name, arrowType);
+  }
+
   @Override
   public MapWriter map(String name, boolean keysSorted) {
     return getWriter(MinorType.STRUCT).map(name, keysSorted);

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

@@ -15,7 +15,8 @@
  * limitations under the License.
  */
 
-<@pp.dropOutputFile />
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
+import org.apache.arrow.vector.holders.ValueHolder;<@pp.dropOutputFile />
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/writer/BaseWriter.java" />
 
 
@@ -61,6 +62,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);
@@ -69,6 +71,13 @@ public interface StructWriter extends BaseWriter {
     void end();
   }
 
+  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 ListWriter extends BaseWriter {
     void startList();
     void endList();
@@ -79,6 +88,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>

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

@@ -15,14 +15,16 @@
  * limitations under the License.
  */
 
-<@pp.dropOutputFile />
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
+import org.apache.arrow.vector.types.Types.MinorType;<@pp.dropOutputFile />
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/PromotableWriter.java" />
 
 <#include "/@includes/license.ftl" />
 
 package org.apache.arrow.vector.complex.impl;
 
 import java.util.Locale;
+import org.apache.arrow.vector.ExtensionTypeVector;
 <#include "/@includes/vv_imports.ftl" />
 
 /**
@@ -285,6 +287,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 +321,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 +542,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();

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

@@ -77,6 +77,9 @@ public class ${mode}StructWriter extends AbstractFieldWriter {
         map(child.getName(), arrowType.getKeysSorted());
         break;
       }
+      case EXTENSIONTYPE:
+        extension(child.getName(), child.getType());
+        break;
       case DENSEUNION: {
         FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.DENSEUNION.getType(), null, null);
         DenseUnionWriter writer = new DenseUnionWriter(container.addOrGet(child.getName(), fieldType, DenseUnionVector.class), getNullableStructWriterFactory());
@@ -136,6 +139,29 @@ public Field getField() {
       return container.getField();
   }
 
+  @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 StructWriter struct(String name) {
     String finalName = handleCase(name);

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

@@ -192,6 +192,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;
+  }
+
   @Override
   public void startList() {
     int start = vector.startNewValue(idx());

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

@@ -16,6 +16,7 @@
  */
 
 import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.writer.Decimal256Writer;
 import org.apache.arrow.vector.complex.writer.DecimalWriter;
 import org.apache.arrow.vector.holders.Decimal256Holder;
@@ -24,6 +25,8 @@
 
 import java.lang.UnsupportedOperationException;
 import java.math.BigDecimal;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.types.Types.MinorType;
 
 <@pp.dropOutputFile />
 <#list ["List", "ListView", "LargeList", "LargeListView"] as listName>
@@ -201,6 +204,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 +337,20 @@ 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 />

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

@@ -16,13 +16,16 @@
  */
 
 import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
 import org.apache.arrow.vector.complex.writer.Decimal256Writer;
 import org.apache.arrow.vector.complex.writer.DecimalWriter;
 import org.apache.arrow.vector.holders.Decimal256Holder;
 import org.apache.arrow.vector.holders.DecimalHolder;
 
 import java.lang.UnsupportedOperationException;
 import java.math.BigDecimal;
+import org.apache.arrow.vector.types.pojo.ArrowType;
 
 <@pp.dropOutputFile />
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionMapWriter.java" />
@@ -231,4 +234,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);
+    }
+  }
 }

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

@@ -208,6 +208,10 @@ private MapWriter getMapWriter(ArrowType arrowType) {
     return mapWriter;
   }
 
+  private ExtensionWriter getExtensionWriter(ArrowType arrowType) {
+    throw new UnsupportedOperationException("ExtensionTypes are not supported yet.");
+  }
+
   public MapWriter asMap(ArrowType arrowType) {
     data.setType(idx(), MinorType.MAP);
     return getMapWriter(arrowType);
@@ -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>

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