GH-810: added support of ExtensionType for UnionVector

Author: xxlaykxx

vector/src/main/codegen/includes/vv_imports.ftl

@@ -43,6 +43,8 @@ import org.apache.arrow.vector.util.JsonStringArrayList;
 import java.util.Arrays;
 import java.util.Random;
 import java.util.List;
+import java.util.HashMap;
+import java.util.Map;
 
 import java.io.Closeable;
 import java.io.InputStream;

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

@@ -109,7 +109,7 @@ public void copyAsField(String name, ${name}Writer writer) {
 
   </#list></#list>
 
-  public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory) {
+  public void copyAsValue(StructWriter writer, ExtensionTypeFactory writerFactory) {
     fail("CopyAsValue StructWriter");
   }
 

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

@@ -110,10 +110,10 @@ public void endEntry() {
   public void write(ExtensionHolder var1)  {
     this.fail("ExtensionType");
   }
-  public void writeExtension(Object var1)  {
+  public void writeExtension(Object var1, ExtensionType var2)  {
     this.fail("ExtensionType");
   }
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
+  public void addExtensionTypeWriterFactory(ExtensionTypeFactory var1, ExtensionType var2) {
     this.fail("ExtensionType");
   }
 

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

@@ -49,7 +49,7 @@ public interface RepeatedStructReader extends StructReader{
     boolean next();
     int size();
     void copyAsValue(StructWriter writer);
-    void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory);
+    void copyAsValue(StructWriter writer, ExtensionTypeFactory writerFactory);
   }
 
   public interface ListReader extends BaseReader{
@@ -60,7 +60,7 @@ public interface RepeatedListReader extends ListReader{
     boolean next();
     int size();
     void copyAsValue(ListWriter writer);
-    void copyAsValue(ListWriter writer, ExtensionTypeWriterFactory writerFactory);
+    void copyAsValue(ListWriter writer, ExtensionTypeFactory writerFactory);
   }
 
   public interface MapReader extends BaseReader{
@@ -71,7 +71,7 @@ public interface RepeatedMapReader extends MapReader{
     boolean next();
     int size();
     void copyAsValue(MapWriter writer);
-    void copyAsValue(MapWriter writer, ExtensionTypeWriterFactory writerFactory);
+    void copyAsValue(MapWriter writer, ExtensionTypeFactory writerFactory);
   }
 
   public interface ScalarReader extends

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

@@ -122,14 +122,14 @@ public interface ExtensionWriter extends BaseWriter {
      *
      * @param value the extension type value to write
      */
-    void writeExtension(Object value);
+    void writeExtension(Object value, ExtensionType extensionType);
 
     /**
      * 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);
+    void addExtensionTypeWriterFactory(ExtensionTypeFactory factory, ExtensionType extensionType);
   }
 
   public interface ScalarWriter extends

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

@@ -19,6 +19,7 @@
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.complex.writer.FieldWriter;
 import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
 
 <@pp.dropOutputFile />
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/ComplexCopier.java" />
@@ -45,11 +46,11 @@ public static void copy(FieldReader input, FieldWriter output) {
     writeValue(input, output, null);
   }
 
-  public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
+  public static void copy(FieldReader input, FieldWriter output, ExtensionTypeFactory extensionTypeWriterFactory) {
     writeValue(input, output, extensionTypeWriterFactory);
   }
 
-  private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
+  private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeFactory extensionTypeWriterFactory) {
     final MinorType mt = reader.getMinorType();
 
       switch (mt) {
@@ -120,9 +121,10 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
         }
         if (reader.isSet()) {
           Object value = reader.readObject();
+          ExtensionType extensionType = (ExtensionType) reader.getField().getType();
           if (value != null) {
-            writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory);
-            writer.writeExtension(value);
+            writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory, extensionType);
+            writer.writeExtension(value, extensionType);
           }
         } else {
           writer.writeNull();

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

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

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

@@ -15,7 +15,8 @@
  * limitations under the License.
  */
 
-<@pp.dropOutputFile />
+import org.apache.arrow.vector.complex.impl.ExtensionTypeFactory;
+import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;<@pp.dropOutputFile />
 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/PromotableWriter.java" />
 
 <#include "/@includes/license.ftl" />
@@ -541,17 +542,13 @@ public void writeLargeVarChar(String value) {
   }
 
   @Override
-  public void writeExtension(Object value) {
-    getWriter(MinorType.EXTENSIONTYPE).writeExtension(value);
+  public void writeExtension(Object value, ExtensionType extensionType) {
+    getWriter(MinorType.EXTENSIONTYPE, extensionType).writeExtension(value, extensionType);
   }
 
   @Override
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) {
-    getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeWriterFactory(factory);
-  }
-
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory, ArrowType arrowType) {
-    getWriter(MinorType.EXTENSIONTYPE, arrowType).addExtensionTypeWriterFactory(factory);
+  public void addExtensionTypeWriterFactory(ExtensionTypeFactory var1, ExtensionType var2) {
+    getWriter(MinorType.EXTENSIONTYPE, var2).addExtensionTypeWriterFactory(var1, var2);
   }
 
   @Override

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

@@ -16,6 +16,7 @@
  */
 
 import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeFactory;
 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,7 @@
 
 import java.lang.UnsupportedOperationException;
 import java.math.BigDecimal;
+import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
 
 <@pp.dropOutputFile />
 <#list ["List", "ListView", "LargeList", "LargeListView"] as listName>
@@ -336,14 +338,14 @@ public void writeNull() {
   }
 
   @Override
-  public void writeExtension(Object value) {
-    writer.writeExtension(value);
+  public void writeExtension(Object value, ExtensionType extensionType) {
+    writer.writeExtension(value, extensionType);
     writer.setPosition(writer.idx() + 1);
   }
 
   @Override
-  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
-    writer.addExtensionTypeWriterFactory(var1, extensionType);
+  public void addExtensionTypeWriterFactory(ExtensionTypeFactory var1, ExtensionType var2) {
+    writer.addExtensionTypeWriterFactory(var1, var2);
   }
 
   public void write(ExtensionHolder var1) {

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

@@ -16,6 +16,8 @@
  */
 
 
+import java.util.HashMap;
+import java.util.Map;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.Field;
 
@@ -43,9 +45,12 @@ public class UnionReader extends AbstractFieldReader {
 
   private BaseReader[] readers = new BaseReader[NUM_SUPPORTED_TYPES];
   public UnionVector data;
-
+  private ExtensionTypeFactory extensionTypeFactory;
+  private Map<ExtensionType, ExtensionReader> extensionTypeReaders = new HashMap<>();
+  
   public UnionReader(UnionVector data) {
     this.data = data;
+    this.extensionTypeFactory = data.getExtensionTypeFactory();
   }
 
   public MinorType getMinorType() {
@@ -79,6 +84,10 @@ public void read(int index, UnionHolder holder) {
   }
 
   private FieldReader getReaderForIndex(int index) {
+    return getReaderForIndex(index, null);
+  }
+  
+    private FieldReader getReaderForIndex(int index, ExtensionType extensionType) {
     int typeValue = data.getTypeValue(index);
     FieldReader reader = (FieldReader) readers[typeValue];
     if (reader != null) {
@@ -95,6 +104,11 @@ private FieldReader getReaderForIndex(int index) {
       return (FieldReader) getListView();
     case MAP:
       return (FieldReader) getMap();
+    case EXTENSIONTYPE:
+      if(extensionType == null) {
+        throw new IllegalStateException("Cannot read extension type without extensionType");
+      }
+      return (FieldReader) getExtension(extensionType);
     <#list vv.types as type>
       <#list type.minor as minor>
         <#assign name = minor.class?cap_first />
@@ -214,6 +228,20 @@ public void copyAsValue(${name}Writer writer){
     </#list>
   </#list>
 
+  public void read(ExtensionHolder holder){
+    getReaderForIndex(idx(), extensionTypeFactory.getExtensionTypeByHolder(holder)).read(holder);
+  }
+
+  private ExtensionReader getExtension(ExtensionType arrowType) {
+    ExtensionReader extensionReader = extensionTypeReaders.get(arrowType);
+    if (extensionReader == null) {
+      extensionReader = extensionTypeFactory.getReaderImpl(data.getExtensionTypeVector(arrowType));
+      extensionReader.setPosition(idx());
+      extensionTypeReaders.put(arrowType, extensionReader);
+    }
+    return extensionReader;
+  }
+
   @Override
   public void copyAsValue(ListWriter writer) {
     ComplexCopier.copy(this, (FieldWriter) writer);