GH-200: Map transferred children by type id

DenseUnionVector keeps typeFields and childVectors by type id and maps
the i-th registered type to its id through typeMapFields, but the
transfer pair indexed them by i. With declared type ids such as [5, 10],
the target's setValueCount then failed on a null child, and
splitAndTransfer indexed per-type arrays of length nextTypeId by type
id.

Author: sharanggupta

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

@@ -103,7 +103,8 @@ public class DenseUnionVector extends AbstractContainerVector implements FieldVe
    */
   private Field[] typeFields = new Field[Byte.MAX_VALUE + 1];
   /**
-   * The index is the index into the typeFields array, and the value is the logical field id.
+   * The index is the order in which the type was registered, and the value is its type id, i.e.
+   * the index into the typeFields and childVectors arrays.
    */
   private byte[] typeMapFields = new byte[Byte.MAX_VALUE + 1];
 
@@ -623,12 +624,16 @@ public TransferImpl(DenseUnionVector to) {
     }
 
     private void createTransferPairs() {
+      // Expects the i-th registered type to be the i-th child of the internal struct. Type
+      // fields and child vectors are indexed by its type id, which differs from i when the
+      // union declares type ids.
       for (int i = 0; i < nextTypeId; i++) {
+        byte typeId = typeMapFields[i];
         ValueVector srcVec = internalStruct.getVectorById(i);
         ValueVector dstVec = to.internalStruct.getVectorById(i);
-        to.typeFields[i] = typeFields[i];
-        to.typeMapFields[i] = typeMapFields[i];
-        to.childVectors[i] = dstVec;
+        to.typeFields[typeId] = typeFields[typeId];
+        to.typeMapFields[i] = typeId;
+        to.childVectors[typeId] = dstVec;
         internalTransferPairs[i] = srcVec.makeTransferPair(dstVec);
       }
       to.nextTypeId = nextTypeId;
@@ -647,7 +652,7 @@ public void transfer() {
       for (int i = 0; i < nextTypeId; i++) {
         if (internalTransferPairs[i] != null) {
           internalTransferPairs[i].transfer();
-          to.childVectors[i] = internalTransferPairs[i].getTo();
+          to.childVectors[typeMapFields[i]] = internalTransferPairs[i].getTo();
         }
       }
       to.valueCount = valueCount;
@@ -670,8 +675,8 @@ public void splitAndTransfer(int startIndex, int length) {
         to.reallocOffsetBuffer();
       }
 
-      int [] typeCounts = new int[nextTypeId];
-      int [] typeStarts = new int[nextTypeId];
+      int [] typeCounts = new int[Byte.MAX_VALUE + 1];
+      int [] typeStarts = new int[Byte.MAX_VALUE + 1];
       for (int i = 0; i < typeCounts.length; i++) {
         typeCounts[i] = 0;
         typeStarts[i] = -1;
@@ -690,9 +695,10 @@ public void splitAndTransfer(int startIndex, int length) {
 
       // transfer vector values
       for (int i = 0; i < nextTypeId; i++) {
-        if (typeCounts[i] > 0 && typeStarts[i] != -1) {
-          internalTransferPairs[i].splitAndTransfer(typeStarts[i], typeCounts[i]);
-          to.childVectors[i] = internalTransferPairs[i].getTo();
+        byte typeId = typeMapFields[i];
+        if (typeCounts[typeId] > 0 && typeStarts[typeId] != -1) {
+          internalTransferPairs[i].splitAndTransfer(typeStarts[typeId], typeCounts[typeId]);
+          to.childVectors[typeId] = internalTransferPairs[i].getTo();
         }
       }
 

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

@@ -191,6 +191,48 @@ public void testTransfer() throws Exception {
     }
   }
 
+  @Test
+  public void testTransferWithDeclaredTypeIds() {
+    try (DenseUnionVector srcVector = newVectorWithTypeIds5And10()) {
+      Field srcField = srcVector.getField();
+      TransferPair pair = srcVector.getTransferPair(allocator);
+      try (DenseUnionVector destVector = (DenseUnionVector) pair.getTo()) {
+        pair.transfer();
+        destVector.setValueCount(3);
+
+        assertEquals(srcField, destVector.getField());
+        assertEquals(5, destVector.getObject(0));
+        assertEquals(true, destVector.getObject(1));
+        assertEquals(7, destVector.getObject(2));
+
+        TransferPair secondPair = destVector.getTransferPair(allocator);
+        try (DenseUnionVector secondDestVector = (DenseUnionVector) secondPair.getTo()) {
+          secondPair.transfer();
+
+          assertEquals(srcField, secondDestVector.getField());
+          assertEquals(3, secondDestVector.getValueCount());
+          assertEquals(5, secondDestVector.getObject(0));
+          assertEquals(true, secondDestVector.getObject(1));
+          assertEquals(7, secondDestVector.getObject(2));
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testSplitAndTransferWithDeclaredTypeIds() {
+    try (DenseUnionVector srcVector = newVectorWithTypeIds5And10()) {
+      TransferPair pair = srcVector.getTransferPair(allocator);
+      try (DenseUnionVector destVector = (DenseUnionVector) pair.getTo()) {
+        pair.splitAndTransfer(1, 2);
+
+        assertEquals(2, destVector.getValueCount());
+        assertEquals(true, destVector.getObject(0));
+        assertEquals(7, destVector.getObject(1));
+      }
+    }
+  }
+
   @Test
   public void testSplitAndTransfer() throws Exception {
     try (DenseUnionVector sourceVector =
@@ -706,4 +748,26 @@ private static NullableFloat4Holder newFloat4Holder(float value) {
     holder.value = value;
     return holder;
   }
+
+  /** Returns a dense union with declared type ids 5 (int) and 10 (bit), as IPC can produce. */
+  private DenseUnionVector newVectorWithTypeIds5And10() {
+    Field field =
+        new Field(
+            "u",
+            FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {5, 10})),
+            List.of(
+                Field.nullable("i", MinorType.INT.getType()),
+                Field.nullable("b", MinorType.BIT.getType())));
+    DenseUnionVector vector = (DenseUnionVector) field.createVector(allocator);
+    vector.allocateNew();
+    vector.setTypeId(0, (byte) 5);
+    vector.setSafe(0, newIntHolder(5));
+    vector.setTypeId(1, (byte) 10);
+    vector.setSafe(1, newBitHolder(true));
+    vector.setTypeId(2, (byte) 5);
+    vector.setSafe(2, newIntHolder(7));
+    vector.setValueCount(3);
+    return vector;
+  }
+
 }