Refactor extension type factory

Author: jhrotko

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

@@ -140,11 +140,6 @@ public int size() {
     return -1;
   }
 
-  @Override
-  public ExtensionTypeWriterFactory getExtensionTypeWriterFactory() {
-    throw new IllegalStateException("The current reader doesn't support reading extension type");
-  }
-
   private void fail(String name) {
     throw new IllegalArgumentException(String.format("You tried to read a [%s] type when you are using a field reader of type [%s].", name, this.getClass().getSimpleName()));
   }

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

@@ -107,15 +107,14 @@ 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()));
   }
 
+  @Override
   public void write(ExtensionHolder var1)  {
     this.fail("ExtensionType");
   }
+  @Override
   public void writeExtension(Object var1)  {
     this.fail("ExtensionType");
   }
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory 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 />

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

@@ -27,8 +27,10 @@
 
 import org.apache.arrow.flatbuf.Type;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.complex.writer.FieldWriter;
 import org.apache.arrow.vector.types.*;
 import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.ValueVector;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -331,6 +333,8 @@ public boolean equals(Object obj) {
     public <T> T accept(ArrowTypeVisitor<T> visitor) {
       return visitor.visit(this);
     }
+
+    public abstract FieldWriter getNewFieldWriter(ValueVector vector);
   }
 
   private static final int defaultDecimalBitWidth = 128;

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

@@ -123,13 +123,6 @@ public interface ExtensionWriter extends BaseWriter {
      * @param value the extension type value to write
      */
     void writeExtension(Object value);
-
-    /**
-     * Adds the given extension type factory. This factory allows configuring writer implementations for specific ExtensionTypeVector.
-     *
-     * @param factory the extension type factory to add
-     */
-    void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory);
   }
 
   public interface ScalarWriter extends

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

@@ -111,8 +111,6 @@ public static void copy(FieldReader reader, FieldWriter writer) {
         if (reader.isSet()) {
           Object value = reader.readObject();
           if (value != null) {
-            ExtensionTypeWriterFactory writerFactory = reader.getExtensionTypeWriterFactory();
-            writer.addExtensionTypeWriterFactory(writerFactory);
             writer.writeExtension(value);
           }
         } else {

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

@@ -86,7 +86,6 @@ public void read(int arrayIndex, Nullable${name}Holder holder){
   }
   </#list></#list>
 
-  public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory){}
   public void read(ExtensionHolder holder) {
     holder.isSet = 0;
   }

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

@@ -286,7 +286,7 @@ protected void setWriter(ValueVector v) {
         writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory);
         break;
       case EXTENSIONTYPE:
-        writer = new UnionExtensionWriter((ExtensionTypeVector) vector);
+        writer = ((ExtensionType) vector.getField().getType()).getNewFieldWriter(vector);
         break;
       default:
         writer = type.getNewFieldWriter(vector);
@@ -325,6 +325,9 @@ protected boolean requiresArrowType(MinorType type) {
 
   @Override
   protected FieldWriter getWriter(MinorType type, ArrowType arrowType) {
+    if(type == MinorType.EXTENSIONTYPE) {
+      lastExtensionType = arrowType;
+    }
     if (state == State.UNION) {
       if (requiresArrowType(type)) {
         ((UnionWriter) writer).getWriter(type, arrowType);
@@ -540,18 +543,25 @@ public void writeLargeVarChar(String value) {
     getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value);
   }
 
+  protected ArrowType lastExtensionType;
+
   @Override
   public void writeExtension(Object value) {
-    getWriter(MinorType.EXTENSIONTYPE).writeExtension(value);
+    FieldWriter writer = getWriter(MinorType.EXTENSIONTYPE, lastExtensionType);
+    if(writer instanceof UnionWriter) {
+      ((UnionWriter) writer).writeExtension(value, lastExtensionType);
+    } else {
+      writer.writeExtension(value);
+    }
   }
 
-  @Override
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) {
-    getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeWriterFactory(factory);
+  public void writeExtension(Object value, ArrowType arrowType) {
+    getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value);
   }
 
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory, ArrowType arrowType) {
-    getWriter(MinorType.EXTENSIONTYPE, arrowType).addExtensionTypeWriterFactory(factory);
+  @Override
+  public void write(ExtensionHolder holder) {
+    getWriter(MinorType.EXTENSIONTYPE, lastExtensionType).write(holder);
   }
 
   @Override

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

@@ -204,13 +204,13 @@ public MapWriter map(String name, boolean keysSorted) {
 
   @Override
   public ExtensionWriter extension(ArrowType arrowType) {
-    this.extensionType = arrowType;
+    extensionType = arrowType;
     return this;
   }
+
   @Override
   public ExtensionWriter extension(String name, ArrowType arrowType) {
-    ExtensionWriter extensionWriter = writer.extension(name, arrowType);
-    return extensionWriter;
+    return writer.extension(name, arrowType);
   }
 
   <#if listName == "LargeList">
@@ -337,15 +337,10 @@ public void writeNull() {
 
   @Override
   public void writeExtension(Object value) {
-    writer.writeExtension(value);
+    writer.writeExtension(value, extensionType);
     writer.setPosition(writer.idx() + 1);
   }
 
-  @Override
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
-    writer.addExtensionTypeWriterFactory(var1, extensionType);
-  }
-
   public void write(ExtensionHolder var1) {
     writer.write(var1);
     writer.setPosition(writer.idx() + 1);

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

@@ -79,6 +79,10 @@ public void read(int index, UnionHolder holder) {
   }
 
   private FieldReader getReaderForIndex(int index) {
+    return getReaderForIndex(index, null);
+  }
+
+  private FieldReader getReaderForIndex(int index, ArrowType type) {
     int typeValue = data.getTypeValue(index);
     FieldReader reader = (FieldReader) readers[typeValue];
     if (reader != null) {
@@ -105,11 +109,26 @@ private FieldReader getReaderForIndex(int index) {
         </#if>
       </#list>
     </#list>
+    case EXTENSIONTYPE:
+      if(type == null) {
+        throw new RuntimeException("Cannot get Extension reader without an ArrowType");
+      }
+      return (FieldReader) getExtension(type);
     default:
       throw new UnsupportedOperationException("Unsupported type: " + MinorType.values()[typeValue]);
     }
   }
 
+  private ExtensionReader extensionReader;
+
+  private ExtensionReader getExtension(ArrowType type) {
+    if (extensionReader == null) {
+      extensionReader = data.getExtension(type).getReader();
+      extensionReader.setPosition(idx());
+    }
+    return extensionReader;
+  }
+
   private SingleStructReaderImpl structReader;
 
   private StructReader getStruct() {
@@ -240,4 +259,8 @@ public FieldReader reader() {
   public boolean next() {
     return getReaderForIndex(idx()).next();
   }
+
+  public void read(ExtensionHolder holder){
+    getReaderForIndex(idx(), holder.type()).read(holder);
+  }
 }

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

@@ -379,6 +379,22 @@ public MapVector getMap(String name, ArrowType arrowType) {
     return mapVector;
   }
 
+  private ExtensionTypeVector extensionVector;
+
+  public ExtensionTypeVector getExtension(ArrowType arrowType) {
+    if (extensionVector == null) {
+      int vectorCount = internalStruct.size();
+      extensionVector = addOrGet(null, MinorType.EXTENSIONTYPE, arrowType, ExtensionTypeVector.class);
+      if (internalStruct.size() > vectorCount) {
+        extensionVector.allocateNew();
+        if (callBack != null) {
+          callBack.doWork();
+        }
+      }
+    }
+    return extensionVector;
+  }
+
   public int getTypeValue(int index) {
     return typeBuffer.getByte(index * TYPE_WIDTH);
   }
@@ -725,6 +741,8 @@ public ValueVector getVectorByType(int typeId, ArrowType arrowType) {
           return getListView();
         case MAP:
           return getMap(name, arrowType);
+        case EXTENSIONTYPE:
+          return getExtension(arrowType);
         default:
           throw new UnsupportedOperationException("Cannot support type: " + MinorType.values()[typeId]);
       }