apacheGH-836: Added support of ExtensionType for ComplexCopier (apache#837)

Updated ComplexCopier to support ExtensionType - it contains two
**copy** methods
```
public static void copy(FieldReader input, FieldWriter output) //for not breaking  existing logic

public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory)
```
Also updated ComplexCopier tests.
Closes apache#836.

Author: xxlaykxx

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

@@ -34,6 +34,7 @@ import org.apache.arrow.vector.complex.*;
 import org.apache.arrow.vector.complex.reader.*;
 import org.apache.arrow.vector.complex.impl.*;
 import org.apache.arrow.vector.complex.writer.*;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;

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

@@ -109,6 +109,10 @@ 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");
   }

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

@@ -49,6 +49,7 @@ public interface RepeatedStructReader extends StructReader{
     boolean next();
     int size();
     void copyAsValue(StructWriter writer);
+    void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory);
   }
 
   public interface ListReader extends BaseReader{
@@ -59,6 +60,7 @@ public interface RepeatedListReader extends ListReader{
     boolean next();
     int size();
     void copyAsValue(ListWriter writer);
+    void copyAsValue(ListWriter writer, ExtensionTypeWriterFactory writerFactory);
   }
 
   public interface MapReader extends BaseReader{
@@ -69,6 +71,7 @@ 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/ComplexCopier.java

@@ -42,10 +42,14 @@ public class ComplexCopier {
    * @param output field to write to
    */
   public static void copy(FieldReader input, FieldWriter output) {
-    writeValue(input, output);
+    writeValue(input, output, null);
   }
 
-  private static void writeValue(FieldReader reader, FieldWriter writer) {
+  public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
+    writeValue(input, output, extensionTypeWriterFactory);
+  }
+
+  private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
     final MinorType mt = reader.getMinorType();
 
       switch (mt) {
@@ -61,7 +65,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
             FieldReader childReader = reader.reader();
             FieldWriter childWriter = getListWriterForReader(childReader, writer);
             if (childReader.isSet()) {
-              writeValue(childReader, childWriter);
+              writeValue(childReader, childWriter, extensionTypeWriterFactory);
             } else {
               childWriter.writeNull();
             }
@@ -79,8 +83,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
             FieldReader structReader = reader.reader();
             if (structReader.isSet()) {
               writer.startEntry();
-              writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()));
-              writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()));
+              writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key()), extensionTypeWriterFactory);
+              writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value()), extensionTypeWriterFactory);
               writer.endEntry();
             } else {
               writer.writeNull();
@@ -99,7 +103,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
             if (childReader.getMinorType() != Types.MinorType.NULL) {
               FieldWriter childWriter = getStructWriterForReader(childReader, writer, name);
               if (childReader.isSet()) {
-                writeValue(childReader, childWriter);
+                writeValue(childReader, childWriter, extensionTypeWriterFactory);
               } else {
                 childWriter.writeNull();
               }
@@ -110,6 +114,20 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
           writer.writeNull();
         }
         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);
+          }
+        } else {
+          writer.writeNull();
+        }
+        break;
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
   <#assign fields = minor.fields!type.fields />
   <#assign uncappedName = name?uncap_first/>
@@ -162,6 +180,9 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr
       return (FieldWriter) writer.map(name);
     case LISTVIEW:
       return (FieldWriter) writer.listView(name);
+    case EXTENSIONTYPE:
+      ExtensionWriter extensionWriter = writer.extension(name, reader.getField().getType());
+      return (FieldWriter) extensionWriter;
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());
     }
@@ -186,6 +207,9 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
       return (FieldWriter) writer.list();
     case LISTVIEW:
       return (FieldWriter) writer.listView();
+    case EXTENSIONTYPE:
+      ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+      return (FieldWriter) extensionWriter;
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());
     }
@@ -211,6 +235,9 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w
         return (FieldWriter) writer.listView();
       case MAP:
         return (FieldWriter) writer.map(false);
+      case EXTENSIONTYPE:
+        ExtensionWriter extensionWriter = writer.extension(reader.getField().getType());
+        return (FieldWriter) extensionWriter;
       default:
         throw new UnsupportedOperationException(reader.getMinorType().toString());
     }

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

@@ -86,6 +86,7 @@ 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/java/org/apache/arrow/vector/BaseValueVector.java

@@ -22,6 +22,7 @@
 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;
@@ -248,4 +249,128 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
   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
+   * aligned, then the buffer is sliced from that index and ownership is transferred. If not,
+   * individual bytes are copied.
+   *
+   * @param startIndex starting index
+   * @param length number of elements to be copied
+   * @param target target vector
+   */
+  protected void splitAndTransferValidityBuffer(
+      int startIndex, int length, BaseValueVector target) {
+    int offset = startIndex % 8;
+
+    if (length <= 0) {
+      return;
+    }
+    if (offset == 0) {
+      sliceAndTransferValidityBuffer(startIndex, length, target);
+    } else {
+      copyValidityBuffer(startIndex, length, target);
+    }
+  }
+
+  /**
+   * If the start index is 8 byte aligned, slice `validityBuffer` and transfer ownership to
+   * `target`'s `validityBuffer`.
+   *
+   * @param startIndex starting index
+   * @param length number of elements to be copied
+   * @param target target vector
+   */
+  protected void sliceAndTransferValidityBuffer(
+      int startIndex, int length, BaseValueVector target) {
+    final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
+    final int byteSizeTarget = getValidityBufferSizeFromCount(length);
+
+    if (target.validityBuffer != null) {
+      target.validityBuffer.getReferenceManager().release();
+    }
+    target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget);
+    target.validityBuffer.getReferenceManager().retain(1);
+  }
+
+  /**
+   * Allocate new validity buffer for `target` and copy bytes from `validityBuffer`. Precise details
+   * in the comments below.
+   *
+   * @param startIndex starting index
+   * @param length number of elements to be copied
+   * @param target target vector
+   */
+  protected void copyValidityBuffer(int startIndex, int length, BaseValueVector target) {
+    final int firstByteSource = BitVectorHelper.byteIndex(startIndex);
+    final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
+    final int byteSizeTarget = getValidityBufferSizeFromCount(length);
+    final int offset = startIndex % 8;
+
+    /* Copy data
+     * When the first bit starts from the middle of a byte (offset != 0),
+     * copy data from src BitVector.
+     * Each byte in the target is composed by a part in i-th byte,
+     * another part in (i+1)-th byte.
+     */
+    target.allocateValidityBuffer(byteSizeTarget);
+
+    for (int i = 0; i < byteSizeTarget - 1; i++) {
+      byte b1 =
+          BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset);
+      byte b2 =
+          BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset);
+
+      target.validityBuffer.setByte(i, (b1 + b2));
+    }
+
+    /* Copying the last piece is done in the following manner:
+     * if the source vector has 1 or more bytes remaining, we copy
+     * the last piece as a byte formed by shifting data
+     * from the current byte and the next byte.
+     *
+     * if the source vector has no more bytes remaining
+     * (we are at the last byte), we copy the last piece as a byte
+     * by shifting data from the current byte.
+     */
+    if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) {
+      byte b1 =
+          BitVectorHelper.getBitsFromCurrentByte(
+              this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
+      byte b2 =
+          BitVectorHelper.getBitsFromNextByte(
+              this.validityBuffer, firstByteSource + byteSizeTarget, offset);
+
+      target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2);
+    } else {
+      byte b1 =
+          BitVectorHelper.getBitsFromCurrentByte(
+              this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset);
+      target.validityBuffer.setByte(byteSizeTarget - 1, b1);
+    }
+  }
+
+  /**
+   * Allocate new validity buffer for when the bytes need to be copied over.
+   *
+   * @param byteSizeTarget desired size of the buffer
+   */
+  protected void allocateValidityBuffer(long byteSizeTarget) {
+    validityBuffer = allocator.buffer(byteSizeTarget);
+    validityBuffer.readerIndex(0);
+    validityBuffer.setZero(0, validityBuffer.capacity());
+  }
 }

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

@@ -27,6 +27,7 @@
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.compare.VectorVisitor;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.impl.NullReader;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
@@ -329,6 +330,18 @@ 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();
+  }
+
   @Override
   public String getName() {
     return this.getField().getName();

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

@@ -22,6 +22,7 @@
 import org.apache.arrow.memory.OutOfMemoryException;
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
 import org.apache.arrow.vector.compare.VectorVisitor;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.Field;
@@ -309,6 +310,30 @@ public interface ValueVector extends Closeable, Iterable<ValueVector> {
    */
   void copyFromSafe(int fromIndex, int thisIndex, ValueVector from);
 
+  /**
+   * Copy a cell value from a particular index in source vector to a particular position in this
+   * vector.
+   *
+   * @param fromIndex position to copy from in source vector
+   * @param thisIndex position to copy to in this vector
+   * @param from source vector
+   * @param writerFactory the extension type writer factory to use for copying extension type values
+   */
+  void copyFrom(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory);
+
+  /**
+   * Same as {@link #copyFrom(int, int, ValueVector)} except that it handles the case when the
+   * capacity of the vector needs to be expanded before copy.
+   *
+   * @param fromIndex position to copy from in source vector
+   * @param thisIndex position to copy to in this vector
+   * @param from source vector
+   * @param writerFactory the extension type writer factory to use for copying extension type values
+   */
+  void copyFromSafe(
+      int fromIndex, int thisIndex, ValueVector from, ExtensionTypeWriterFactory writerFactory);
+
   /**
    * Accept a generic {@link VectorVisitor} and return the result.
    *

vector/src/main/java/org/apache/arrow/vector/complex/AbstractContainerVector.java

@@ -21,6 +21,7 @@
 import org.apache.arrow.vector.DensityAwareVector;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.complex.impl.ExtensionTypeWriterFactory;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList;
@@ -151,6 +152,18 @@ 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();
+  }
+
   @Override
   public String getName() {
     return name;