GH-891: Add ExtensionTypeWriterFactory to TransferPair

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

@@ -109,10 +109,6 @@ public void copyAsField(String name, ${name}Writer writer) {
 
  </#list></#list>
 
- public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory) {
- fail("CopyAsValue StructWriter");
- }
-
  public void read(ExtensionHolder holder) {
  fail("Extension");
  }
@@ -147,4 +143,5 @@ public int size() {
  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,14 +107,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()));
  }
 
+ @Override
  public void write(ExtensionHolder var1) {
- this.fail("ExtensionType");
+ this.fail("Cannot write ExtensionHolder");
  }
+ @Override
  public void writeExtension(Object var1) {
- this.fail("ExtensionType");
+ this.fail("Cannot write extension object");
  }
- public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
- this.fail("ExtensionType");
+ @Override
+ public void writeExtension(Object var1, ArrowType type) {
+ this.fail("Cannot write extension with type " + type);
  }
 
  <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />

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,10 @@ public boolean equals(Object obj) {
  public <T> T accept(ArrowTypeVisitor<T> visitor) {
  return visitor.visit(this);
  }
+
+ public FieldWriter getNewFieldWriter(ValueVector vector) {
+ throw new UnsupportedOperationException("WriterImpl not yet implemented.");
+ }
  }
 
  private static final int defaultDecimalBitWidth = 128;

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

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

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

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

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

@@ -41,15 +41,8 @@ public class ComplexCopier {
  * @param input field to read from
  * @param output field to write to
  */
- public static void copy(FieldReader input, FieldWriter output) {
- writeValue(input, output, null);
- }
-
- public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
- writeValue(input, output, extensionTypeWriterFactory);
- }
+ public static void copy(FieldReader reader, FieldWriter writer) {
 
- private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
  final MinorType mt = reader.getMinorType();
 
  switch (mt) {
@@ -65,7 +58,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
  FieldReader childReader = reader.reader();
  FieldWriter childWriter = getListWriterForReader(childReader, writer);
  if (childReader.isSet()) {
-writeValue(childReader, childWriter, extensionTypeWriterFactory);
+copy(childReader, childWriter);
  } else {
  childWriter.writeNull();
  }
@@ -83,8 +76,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
  FieldReader structReader = reader.reader();
  if (structReader.isSet()) {
  writer.startEntry();
-writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()), extensionTypeWriterFactory);
-writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()), extensionTypeWriterFactory);
+copy(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()));
+copy(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()));
  writer.endEntry();
  } else {
  writer.writeNull();
@@ -103,7 +96,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
  if (childReader.getMinorType() != Types.MinorType.NULL) {
  FieldWriter childWriter = getStructWriterForReader(childReader, writer, name);
  if (childReader.isSet()) {
-writeValue(childReader, childWriter, extensionTypeWriterFactory);
+copy(childReader, childWriter);
  } else {
  childWriter.writeNull();
  }
@@ -115,14 +108,10 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
  }
  break;
  case EXTENSIONTYPE:
- if (extensionTypeWriterFactory == null) {
- throw new IllegalArgumentException("Must provide ExtensionTypeWriterFactory");
- }
  if (reader.isSet()) {
  Object value = reader.readObject();
  if (value != null) {
- writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory);
- writer.writeExtension(value);
+ writer.writeExtension(value, reader.getField().getType());
  }
  } else {
  writer.writeNull();

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);
@@ -541,17 +541,13 @@ public void writeLargeVarChar(String value) {
  }
 
  @Override
- public void writeExtension(Object value) {
- getWriter(MinorType.EXTENSIONTYPE).writeExtension(value);
+ public void writeExtension(Object value, ArrowType arrowType) {
+ getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value, arrowType);
  }
 
  @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 write(ExtensionHolder holder) {
+ getWriter(MinorType.EXTENSIONTYPE, holder.type()).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,13 +337,13 @@ 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 writeExtension(Object value, ArrowType type) {
+writeExtension(value);
  }
 
  public void write(ExtensionHolder var1) {

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

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

@@ -28,6 +28,8 @@
 package org.apache.arrow.vector.complex.impl;
 
 <#include "/@includes/vv_imports.ftl" />
+import java.util.HashMap;
+
 import org.apache.arrow.vector.complex.writer.BaseWriter;
 import org.apache.arrow.vector.types.Types.MinorType;
 
@@ -213,8 +215,31 @@ public MapWriter asMap(ArrowType arrowType) {
  return getMapWriter(arrowType);
  }
 
+ private java.util.Map<ArrowType, ExtensionWriter> extensionWriters = new HashMap<>();
+
  private ExtensionWriter getExtensionWriter(ArrowType arrowType) {
- throw new UnsupportedOperationException("ExtensionTypes are not supported yet.");
+ ExtensionWriter w = extensionWriters.get(arrowType);
+ if (w == null) {
+ w = ((ExtensionType) arrowType).getNewFieldWriter(data.getExtension(arrowType));
+ w.setPosition(idx());
+ extensionWriters.put(arrowType, w);
+ }
+ return w;
+ }
+
+ public void writeExtension(Object value, ArrowType type) {
+ data.setType(idx(), MinorType.EXTENSIONTYPE);
+ ExtensionWriter w = getExtensionWriter(type);
+ w.setPosition(idx());
+ w.writeExtension(value);
+ }
+
+ @Override
+ public void write(ExtensionHolder holder) {
+ data.setType(idx(), MinorType.EXTENSIONTYPE);
+ ExtensionWriter w = getExtensionWriter(holder.type());
+ w.setPosition(idx());
+ w.write(holder);
  }
 
  BaseWriter getWriter(MinorType minorType) {

vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java

@@ -22,7 +22,6 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.ReferenceManager;
 import org.apache.arrow.util.Preconditions;
-import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.util.DataSizeRoundingUtil;
 import org.apache.arrow.vector.util.TransferPair;
@@ -261,18 +260,6 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
  throw new UnsupportedOperationException();
  }
 
- @Override
- public void copyFrom(
- int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public void copyFromSafe(
- int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory) {
- throw new UnsupportedOperationException();
- }
-
  /**
  * Transfer the validity buffer from `validityBuffer` to the target vector's `validityBuffer`.
  * Start at `startIndex` and copy `length` number of elements. If the starting index is 8 byte