Add makeTransferPair w/ extended factory

Author: jhrotko

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

@@ -30,6 +30,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.ipc.message.ArrowFieldNode;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.util.CallBack;
@@ -901,4 +902,10 @@ public int hashCode(int index, ArrowBufHasher hasher) {
   public <OUT, IN> OUT accept(VectorVisitor<OUT, IN> visitor, IN value) {
     return visitor.visit(this, value);
   }
+
+  @Override
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
 }

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

@@ -384,4 +384,10 @@ protected void allocateValidityBuffer(long byteSizeTarget) {
     validityBuffer.readerIndex(0);
     validityBuffer.setZero(0, validityBuffer.capacity());
   }
+
+  @Override
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
 }

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

@@ -31,6 +31,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.ipc.message.ArrowFieldNode;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.util.CallBack;
@@ -1443,4 +1444,10 @@ public <OUT, IN> OUT accept(VectorVisitor<OUT, IN> visitor, IN value) {
   public final int getEndOffset(int index) {
     return offsetBuffer.getInt((long) (index + 1) * OFFSET_WIDTH);
   }
+
+  @Override
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
 }

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

@@ -201,6 +201,12 @@ public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((NullVector) target);
   }
 
+  @Override
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    return makeTransferPair(target);
+  }
+
   @Override
   public FieldReader getReader() {
     return NullReader.INSTANCE;

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

@@ -190,6 +190,8 @@ public interface ValueVector extends Closeable, Iterable<ValueVector> {
    */
   TransferPair makeTransferPair(ValueVector target);
 
+  TransferPair makeTransferPair(ValueVector target, ExtensionTypeWriterFactory writerFactory);
+
   /**
    * Get a reader for this vector.
    *

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

@@ -30,6 +30,7 @@
 import org.apache.arrow.vector.types.pojo.ArrowType.Struct;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.util.CallBack;
+import org.apache.arrow.vector.util.TransferPair;
 import org.apache.arrow.vector.util.ValueVectorUtility;
 
 /**
@@ -168,4 +169,10 @@ public void copyFromSafe(
   public String getName() {
     return name;
   }
+
+  @Override
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    throw new UnsupportedOperationException();
+  }
 }

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

@@ -63,7 +63,6 @@
 import org.apache.arrow.vector.util.OversizedAllocationException;
 import org.apache.arrow.vector.util.SchemaChangeRuntimeException;
 import org.apache.arrow.vector.util.TransferPair;
-import org.apache.arrow.vector.util.TransferPairWithExtendedType;
 
 /**
  * A list vector contains lists of a specific type of elements. Its structure contains 3 elements.
@@ -442,8 +441,7 @@ private void reallocValidityBuffer() {
       if (validityAllocationSizeInBytes > 0) {
         newAllocationSize = validityAllocationSizeInBytes;
       } else {
-        newAllocationSize =
-            BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+        newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
       }
     }
     newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
@@ -570,6 +568,12 @@ public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((LargeListVector) target);
   }
 
+  @Override
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    return new TransferImpl((LargeListVector) target, writerFactory);
+  }
+
   @Override
   public long getValidityBufferAddress() {
     return validityBuffer.memoryAddress();
@@ -649,10 +653,11 @@ public UnionVector promoteToUnion() {
     return vector;
   }
 
-  private class TransferImpl implements TransferPairWithExtendedType {
+  private class TransferImpl implements TransferPair {
 
     LargeListVector to;
     TransferPair dataTransferPair;
+    ExtensionTypeWriterFactory writerFactory;
 
     public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) {
       this(new LargeListVector(name, allocator, field.getFieldType(), callBack));
@@ -662,6 +667,11 @@ public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) {
       this(new LargeListVector(field, allocator, callBack));
     }
 
+    public TransferImpl(LargeListVector to, ExtensionTypeWriterFactory writerFactory) {
+      this(to);
+      this.writerFactory = writerFactory;
+    }
+
     public TransferImpl(LargeListVector to) {
       this.to = to;
       to.addOrGetVector(vector.getField().getFieldType());
@@ -730,13 +740,7 @@ public ValueVector getTo() {
 
     @Override
     public void copyValueSafe(int from, int to) {
-      this.to.copyFrom(from, to, LargeListVector.this);
-    }
-
-    @Override
-    public void copyValueSafe(
-        int from, int to, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
-      this.to.copyFrom(from, to, LargeListVector.this, extensionTypeWriterFactory);
+      this.to.copyFrom(from, to, LargeListVector.this, writerFactory);
     }
   }
 

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

@@ -56,7 +56,6 @@
 import org.apache.arrow.vector.util.JsonStringArrayList;
 import org.apache.arrow.vector.util.OversizedAllocationException;
 import org.apache.arrow.vector.util.TransferPair;
-import org.apache.arrow.vector.util.TransferPairWithExtendedType;
 
 /**
  * A list vector contains lists of a specific type of elements. Its structure contains 3 elements.
@@ -366,8 +365,7 @@ private long getNewAllocationSize(int currentBufferCapacity) {
       if (validityAllocationSizeInBytes > 0) {
         newAllocationSize = validityAllocationSizeInBytes;
       } else {
-        newAllocationSize =
-            BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
+        newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L;
       }
     }
     newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize);
@@ -475,6 +473,11 @@ public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((ListVector) target);
   }
 
+  public TransferPair makeTransferPair(
+      ValueVector target, ExtensionTypeWriterFactory writerFactory) {
+    return new TransferImpl((ListVector) target, writerFactory);
+  }
+
   @Override
   public long getValidityBufferAddress() {
     return validityBuffer.memoryAddress();
@@ -529,10 +532,11 @@ public <OUT, IN> OUT accept(VectorVisitor<OUT, IN> visitor, IN value) {
     return visitor.visit(this, value);
   }
 
-  private class TransferImpl implements TransferPairWithExtendedType {
+  private class TransferImpl implements TransferPair {
 
     ListVector to;
     TransferPair dataTransferPair;
+    ExtensionTypeWriterFactory writerFactory;
 
     public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) {
       this(new ListVector(name, allocator, field.getFieldType(), callBack));
@@ -542,6 +546,11 @@ public TransferImpl(Field field, BufferAllocator allocator, CallBack callBack) {
       this(new ListVector(field, allocator, callBack));
     }
 
+    public TransferImpl(ListVector to, ExtensionTypeWriterFactory writerFactory) {
+      this(to);
+      this.writerFactory = writerFactory;
+    }
+
     public TransferImpl(ListVector to) {
       this.to = to;
       to.addOrGetVector(vector.getField().getFieldType());
@@ -611,13 +620,7 @@ public ValueVector getTo() {
 
     @Override
     public void copyValueSafe(int from, int to) {
-      this.to.copyFrom(from, to, ListVector.this);
-    }
-
-    @Override
-    public void copyValueSafe(
-        int from, int to, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
-      this.to.copyFrom(from, to, ListVector.this, extensionTypeWriterFactory);
+      this.to.copyFrom(from, to, ListVector.this, writerFactory);
     }
   }
 

vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java

@@ -79,6 +79,11 @@ public void setPosition(int index) {
 
   @Override
   public void writeNull() {
-    this.writer.writeNull();
+    if (this.writer != null) {
+      this.writer.writeNull();
+    } else {
+      this.vector.setNull(getPosition());
+      this.vector.setValueCount(getPosition() + 1);
+    }
   }
 }