ARROW-347: Add method to pass CallBack when creating a transfer pair

supersedes and closes apache#182

Author: Julien Le Dem <julien@dremio.com>

Closes apache#425 from julienledem/arrow_347 and squashes the following commits:

3c47b82 [Julien Le Dem] ARROW-347: Add method to pass CallBack when creating a transfer pair

Author: julienledem

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

@@ -236,12 +236,17 @@ public Field getField() {
 
   @Override
   public TransferPair getTransferPair(BufferAllocator allocator) {
-    return new TransferImpl(name, allocator);
+    return getTransferPair(name, allocator);
   }
 
   @Override
   public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
-    return new TransferImpl(ref, allocator);
+    return getTransferPair(ref, allocator, null);
+  }
+
+  @Override
+  public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
+    return new org.apache.arrow.vector.complex.UnionVector.TransferImpl(ref, allocator, callBack);
   }
 
   @Override
@@ -276,8 +281,8 @@ private class TransferImpl implements TransferPair {
     private final TransferPair typeVectorTransferPair;
     private final UnionVector to;
 
-    public TransferImpl(String name, BufferAllocator allocator) {
-      to = new UnionVector(name, allocator, null);
+    public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) {
+      to = new UnionVector(name, allocator, callBack);
       internalMapVectorTransferPair = internalMap.makeTransferPair(to.internalMap);
       typeVectorTransferPair = typeVector.makeTransferPair(to.typeVector);
     }

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

@@ -24,6 +24,8 @@
 import org.apache.arrow.vector.schema.ArrowFieldNode;
 
 import io.netty.buffer.ArrowBuf;
+import org.apache.arrow.vector.util.CallBack;
+import org.apache.arrow.vector.util.TransferPair;
 
 
 public abstract class BaseDataValueVector extends BaseValueVector implements BufferBacked {
@@ -87,6 +89,11 @@ public void close() {
     super.close();
   }
 
+  @Override
+  public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
+    return getTransferPair(ref, allocator);
+  }
+
   @Override
   public ArrowBuf[] getBuffers(boolean clear) {
     ArrowBuf[] out;

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

@@ -24,6 +24,7 @@
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.util.CallBack;
 import org.apache.arrow.vector.util.TransferPair;
 
 import io.netty.buffer.ArrowBuf;
@@ -106,6 +107,8 @@ public interface ValueVector extends Closeable, Iterable<ValueVector> {
 
   TransferPair getTransferPair(String ref, BufferAllocator allocator);
 
+  TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack);
+
   /**
    * Returns a new {@link org.apache.arrow.vector.util.TransferPair transfer pair} that is used to transfer underlying
    * buffers into the target vector.

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

@@ -29,6 +29,7 @@
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType.Null;
 import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.util.CallBack;
 import org.apache.arrow.vector.util.TransferPair;
 
 import io.netty.buffer.ArrowBuf;
@@ -159,6 +160,11 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
     return defaultPair;
   }
 
+  @Override
+  public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
+    return defaultPair;
+  }
+
   @Override
   public TransferPair makeTransferPair(ValueVector target) {
     return defaultPair;

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

@@ -29,6 +29,7 @@
 import org.apache.arrow.vector.ZeroVector;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
+import org.apache.arrow.vector.util.CallBack;
 import org.apache.arrow.vector.util.SchemaChangeRuntimeException;
 
 import com.google.common.base.Preconditions;
@@ -44,15 +45,17 @@ public abstract class BaseRepeatedValueVector extends BaseValueVector implements
 
   protected final UInt4Vector offsets;
   protected FieldVector vector;
+  protected final CallBack callBack;
 
-  protected BaseRepeatedValueVector(String name, BufferAllocator allocator) {
-    this(name, allocator, DEFAULT_DATA_VECTOR);
+  protected BaseRepeatedValueVector(String name, BufferAllocator allocator, CallBack callBack) {
+    this(name, allocator, DEFAULT_DATA_VECTOR, callBack);
   }
 
-  protected BaseRepeatedValueVector(String name, BufferAllocator allocator, FieldVector vector) {
+  protected BaseRepeatedValueVector(String name, BufferAllocator allocator, FieldVector vector, CallBack callBack) {
     super(name, allocator);
     this.offsets = new UInt4Vector(OFFSETS_VECTOR_NAME, allocator);
     this.vector = Preconditions.checkNotNull(vector, "data vector cannot be null");
+    this.callBack = callBack;
   }
 
   @Override
@@ -154,9 +157,12 @@ public int size() {
   public <T extends ValueVector> AddOrGetResult<T> addOrGetVector(MinorType minorType, DictionaryEncoding dictionary) {
     boolean created = false;
     if (vector instanceof ZeroVector) {
-      vector = minorType.getNewVector(DATA_VECTOR_NAME, allocator, dictionary, null);
+      vector = minorType.getNewVector(DATA_VECTOR_NAME, allocator, dictionary, callBack);
       // returned vector must have the same field
       created = true;
+      if (callBack != null) {
+        callBack.doWork();
+      }
     }
 
     if (vector.getField().getType().getTypeID() != minorType.getType().getTypeID()) {

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

@@ -24,10 +24,6 @@
 import java.util.Collections;
 import java.util.List;
 
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ObjectArrays;
-
-import io.netty.buffer.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.OutOfMemoryException;
 import org.apache.arrow.vector.AddOrGetResult;
@@ -52,24 +48,27 @@
 import org.apache.arrow.vector.util.JsonStringArrayList;
 import org.apache.arrow.vector.util.TransferPair;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ObjectArrays;
+
+import io.netty.buffer.ArrowBuf;
+
 public class ListVector extends BaseRepeatedValueVector implements FieldVector {
 
   final UInt4Vector offsets;
   final BitVector bits;
   private final List<BufferBacked> innerVectors;
   private Mutator mutator = new Mutator();
   private Accessor accessor = new Accessor();
-  private UnionListWriter writer;
   private UnionListReader reader;
   private CallBack callBack;
   private final DictionaryEncoding dictionary;
 
   public ListVector(String name, BufferAllocator allocator, DictionaryEncoding dictionary, CallBack callBack) {
-    super(name, allocator);
+    super(name, allocator, callBack);
     this.bits = new BitVector("$bits$", allocator);
     this.offsets = getOffsetVector();
     this.innerVectors = Collections.unmodifiableList(Arrays.<BufferBacked>asList(bits, offsets));
-    this.writer = new UnionListWriter(this);
     this.reader = new UnionListReader(this);
     this.dictionary = dictionary;
     this.callBack = callBack;
@@ -86,6 +85,8 @@ public void initializeChildrenFromFields(List<Field> children) {
     if (!addOrGetVector.isCreated()) {
       throw new IllegalArgumentException("Child vector already existed: " + addOrGetVector.getVector());
     }
+
+    addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren());
   }
 
   @Override
@@ -111,7 +112,7 @@ public List<BufferBacked> getFieldInnerVectors() {
   }
 
   public UnionListWriter getWriter() {
-    return writer;
+    return new UnionListWriter(this);
   }
 
   @Override
@@ -139,7 +140,12 @@ public FieldVector getDataVector() {
 
   @Override
   public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
-    return new TransferImpl(ref, allocator);
+    return getTransferPair(ref, allocator, null);
+  }
+
+  @Override
+  public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
+    return new TransferImpl(ref, allocator, callBack);
   }
 
   @Override
@@ -152,8 +158,8 @@ private class TransferImpl implements TransferPair {
     ListVector to;
     TransferPair pairs[] = new TransferPair[3];
 
-    public TransferImpl(String name, BufferAllocator allocator) {
-      this(new ListVector(name, allocator, dictionary, null));
+    public TransferImpl(String name, BufferAllocator allocator, CallBack callBack) {
+      this(new ListVector(name, allocator, dictionary, callBack));
     }
 
     public TransferImpl(ListVector to) {
@@ -172,6 +178,7 @@ public void transfer() {
       for (TransferPair pair : pairs) {
         pair.transfer();
       }
+      to.lastSet = lastSet;
     }
 
     @Override
@@ -282,9 +289,12 @@ public ArrowBuf[] getBuffers(boolean clear) {
   }
 
   public UnionVector promoteToUnion() {
-    UnionVector vector = new UnionVector(name, allocator, null);
+    UnionVector vector = new UnionVector(name, allocator, callBack);
     replaceDataVector(vector);
     reader = new UnionListReader(this);
+    if (callBack != null) {
+      callBack.doWork();
+    }
     return vector;
   }
 

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

@@ -115,6 +115,11 @@ public int getBufferSizeFor(final int valueCount) {
 
   @Override
   public TransferPair getTransferPair(BufferAllocator allocator) {
+    return getTransferPair(name, allocator, null);
+  }
+
+  @Override
+  public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
     return new MapTransferPair(this, new MapVector(name, allocator, callBack), false);
   }
 

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

@@ -86,7 +86,7 @@ public FieldReader getReader() {
 
   @Override
   public TransferPair getTransferPair(BufferAllocator allocator) {
-    return new NullableMapTransferPair(this, new NullableMapVector(name, allocator, dictionary, callBack), false);
+    return new NullableMapTransferPair(this, new NullableMapVector(name, allocator, dictionary, null), false);
   }
 
   @Override
@@ -96,6 +96,11 @@ public TransferPair makeTransferPair(ValueVector to) {
 
   @Override
   public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
+    return new NullableMapTransferPair(this, new NullableMapVector(ref, allocator, dictionary, null), false);
+  }
+
+  @Override
+  public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) {
     return new NullableMapTransferPair(this, new NullableMapVector(ref, allocator, dictionary, callBack), false);
   }
 

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

@@ -142,7 +142,7 @@ public boolean isEmptyMap() {
   }
 
   protected FieldWriter getWriter() {
-    return getWriter(type);
+    return writer;
   }
 
   private FieldWriter promoteToUnion() {