ARROW-403: [Java] Create transfer pairs for internal vectors in UnionVector transfer impl

Author: vkorukanti

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

@@ -246,12 +246,6 @@ public TransferPair makeTransferPair(ValueVector target) {
     return new TransferImpl((UnionVector) target);
   }
 
-  public void transferTo(org.apache.arrow.vector.complex.UnionVector target) {
-    typeVector.makeTransferPair(target.typeVector).transfer();
-    internalMap.makeTransferPair(target.internalMap).transfer();
-    target.valueCount = valueCount;
-  }
-
   public void copyFrom(int inIndex, int outIndex, UnionVector from) {
     from.getReader().setPosition(inIndex);
     getWriter().setPosition(outIndex);
@@ -275,20 +269,27 @@ public FieldVector addVector(FieldVector v) {
   }
 
   private class TransferImpl implements TransferPair {
-
-    UnionVector to;
+    private final TransferPair internalMapVectorTransferPair;
+    private final TransferPair typeVectorTransferPair;
+    private final UnionVector to;
 
     public TransferImpl(String name, BufferAllocator allocator) {
       to = new UnionVector(name, allocator, null);
+      internalMapVectorTransferPair = internalMap.makeTransferPair(to.internalMap);
+      typeVectorTransferPair = typeVector.makeTransferPair(to.typeVector);
     }
 
     public TransferImpl(UnionVector to) {
       this.to = to;
+      internalMapVectorTransferPair = internalMap.makeTransferPair(to.internalMap);
+      typeVectorTransferPair = typeVector.makeTransferPair(to.typeVector);
     }
 
     @Override
     public void transfer() {
-      transferTo(to);
+      internalMapVectorTransferPair.transfer();
+      typeVectorTransferPair.transfer();
+      to.valueCount = valueCount;
     }
 
     @Override

java/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java

@@ -21,8 +21,12 @@
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.complex.UnionVector;
+import org.apache.arrow.vector.holders.NullableBitHolder;
+import org.apache.arrow.vector.holders.NullableIntHolder;
 import org.apache.arrow.vector.holders.NullableUInt4Holder;
 import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.Types.MinorType;
+import org.apache.arrow.vector.util.TransferPair;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -76,4 +80,54 @@ public void testUnionVector() throws Exception {
     }
   }
 
+  @Test
+  public void testTransfer() throws Exception {
+    try (UnionVector srcVector = new UnionVector(EMPTY_SCHEMA_PATH, allocator, null)) {
+      srcVector.allocateNew();
+
+      // write some data
+      final UnionVector.Mutator mutator = srcVector.getMutator();
+      mutator.setType(0, MinorType.INT);
+      mutator.setSafe(0, newIntHolder(5));
+      mutator.setType(1, MinorType.BIT);
+      mutator.setSafe(1, newBitHolder(false));
+      mutator.setType(3, MinorType.INT);
+      mutator.setSafe(3, newIntHolder(10));
+      mutator.setType(5, MinorType.BIT);
+      mutator.setSafe(5, newBitHolder(false));
+      mutator.setValueCount(6);
+
+      try(UnionVector destVector = new UnionVector(EMPTY_SCHEMA_PATH, allocator, null)) {
+        TransferPair pair = srcVector.makeTransferPair(destVector);
+
+        // Creating the transfer should transfer the type of the field at least.
+        assertEquals(srcVector.getField(), destVector.getField());
+
+        // transfer
+        pair.transfer();
+
+        assertEquals(srcVector.getField(), destVector.getField());
+
+        // now check the values are transferred
+        assertEquals(srcVector.getAccessor().getValueCount(), destVector.getAccessor().getValueCount());
+        for(int i=0; i<srcVector.getAccessor().getValueCount(); i++) {
+          assertEquals("Different values at index " + i, srcVector.getAccessor().get(i), destVector.getAccessor().get(i));
+        }
+      }
+    }
+  }
+
+  private static NullableIntHolder newIntHolder(int value) {
+    final NullableIntHolder holder = new NullableIntHolder();
+    holder.isSet = 1;
+    holder.value = value;
+    return holder;
+  }
+
+  private static NullableBitHolder newBitHolder(boolean value) {
+    final NullableBitHolder holder = new NullableBitHolder();
+    holder.isSet = 1;
+    holder.value = value ? 1 : 0;
+    return holder;
+  }
 }