GH-1471: Pair dense union children with their type ids

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

@@ -166,9 +166,26 @@ public void initializeChildrenFromFields(List<Field> children) {
     }
   }
 
+  /**
+   * Returns the children in the order of {@link #getField()}, which is the order their type ids
+   * were registered in, followed by any child that was added without registering a type id. The
+   * vectors of the children may have been created in another order.
+   */
   @Override
   public List<FieldVector> getChildrenFromFields() {
-    return internalStruct.getChildrenFromFields();
+    List<FieldVector> children = new ArrayList<>(internalStruct.size());
+    for (int i = 0; i < typeCount(); i++) {
+      ValueVector child = childVectors[typeMapFields[i]];
+      if (child != null) {
+        children.add((FieldVector) child);
+      }
+    }
+    for (FieldVector child : internalStruct.getChildrenFromFields()) {
+      if (children.stream().noneMatch(registered -> registered == child)) {
+        children.add(child);
+      }
+    }
+    return children;
   }
 
   @Override
@@ -241,14 +258,17 @@ public synchronized byte registerNewTypeId(Field field) {
               typeFields.length + " relative types. Please use union of union instead");
     }
     byte typeId = nextTypeId;
-    if (this.fieldType != null) {
-      int[] typeIds = ((ArrowType.Union) this.fieldType.getType()).getTypeIds();
-      if (typeIds != null) {
-        int thisTypeId = typeIds[nextTypeId];
-        if (thisTypeId > Byte.MAX_VALUE) {
-          throw new IllegalStateException("Dense union vector types must be bytes. " + thisTypeId + " is too large");
-        }
-        typeId = (byte) thisTypeId;
+    int[] typeIds = this.fieldType == null ? null : ((ArrowType.Union) this.fieldType.getType()).getTypeIds();
+    if (typeIds != null) {
+      int thisTypeId = typeIds[nextTypeId];
+      if (thisTypeId > Byte.MAX_VALUE) {
+        throw new IllegalStateException("Dense union vector types must be bytes. " + thisTypeId + " is too large");
+      }
+      typeId = (byte) thisTypeId;
+    } else {
+      // A transfer pair can give this vector type ids other than 0 to n - 1, so skip those in use
+      while (isRegistered(typeId)) {
+        typeId = (byte) ((typeId + 1) % typeFields.length);
       }
     }
     typeFields[typeId] = field;
@@ -257,6 +277,15 @@ public synchronized byte registerNewTypeId(Field field) {
     return typeId;
   }
 
+  private boolean isRegistered(byte typeId) {
+    for (int i = 0; i < typeCount(); i++) {
+      if (typeMapFields[i] == typeId) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private <T extends FieldVector> T addOrGet(byte typeId, MinorType minorType, Class<T> c) {
     return internalStruct.addOrGet(fieldName(typeId, minorType), fieldType(minorType), c);
   }
@@ -625,16 +654,20 @@ public TransferImpl(DenseUnionVector to) {
     }
 
     // The transfer pairs follow the order of the children, while typeFields and childVectors are
-    // indexed by type id. typeMapFields maps the position of a child to its type id.
+    // indexed by type id. typeMapFields maps the position of a child to its type id. The vectors
+    // of the children may have been created in another order, so internalStruct, whose transfer
+    // pair created the children of the target under the same names, is not indexed by position.
     private void createTransferPairs() {
       for (int i = 0; i < typeCount(); i++) {
         byte typeId = typeMapFields[i];
-        ValueVector srcVec = internalStruct.getVectorById(i);
-        ValueVector dstVec = to.internalStruct.getVectorById(i);
         to.typeFields[typeId] = typeFields[typeId];
         to.typeMapFields[i] = typeId;
-        to.childVectors[typeId] = dstVec;
-        internalTransferPairs[i] = srcVec.makeTransferPair(dstVec);
+        ValueVector srcVec = childVectors[typeId];
+        if (srcVec != null) {
+          ValueVector dstVec = to.internalStruct.getChild(srcVec.getName());
+          to.childVectors[typeId] = dstVec;
+          internalTransferPairs[i] = srcVec.makeTransferPair(dstVec);
+        }
       }
       to.nextTypeId = (byte) Math.max(to.typeCount(), typeCount());
     }
@@ -724,7 +757,7 @@ public void splitAndTransfer(int startIndex, int length) {
                 to.offsetBuffer.getInt((long) (i - startIndex) * OFFSET_WIDTH));
           } catch (UnsupportedOperationException e) {
             to.clear();
-            ValueVector childVector = internalStruct.getVectorById(child);
+            ValueVector childVector = childVectors[typeId];
             throw new UnsupportedOperationException(
                 String.format(
                     "Cannot split dense union %s: the values its rows select from child %s (%s)"

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

@@ -39,8 +39,10 @@
 import org.apache.arrow.vector.holders.NullableBigIntHolder;
 import org.apache.arrow.vector.holders.NullableBitHolder;
 import org.apache.arrow.vector.holders.NullableFloat4Holder;
+import org.apache.arrow.vector.holders.NullableFloat8Holder;
 import org.apache.arrow.vector.holders.NullableIntHolder;
 import org.apache.arrow.vector.holders.NullableUInt4Holder;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.testing.ValueVectorDataPopulator;
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.Types.MinorType;
@@ -424,6 +426,106 @@ public void testGetFieldKeepsChildOrder() {
     }
   }
 
+  /** Registers an int child and then a double child, but creates the double vector first. */
+  private DenseUnionVector newVectorWithChildrenCreatedOutOfOrder() {
+    DenseUnionVector vector = DenseUnionVector.empty("duv", allocator);
+    vector.allocateNew();
+    byte intTypeId = vector.registerNewTypeId(Field.nullable("i", MinorType.INT.getType()));
+    byte doubleTypeId = vector.registerNewTypeId(Field.nullable("d", MinorType.FLOAT8.getType()));
+
+    NullableFloat8Holder doubleHolder = new NullableFloat8Holder();
+    doubleHolder.isSet = 1;
+    NullableIntHolder intHolder = new NullableIntHolder();
+    intHolder.isSet = 1;
+
+    vector.setTypeId(0, doubleTypeId);
+    doubleHolder.value = 1.5;
+    vector.setSafe(0, doubleHolder);
+    vector.setTypeId(1, intTypeId);
+    intHolder.value = 42;
+    vector.setSafe(1, intHolder);
+    vector.setTypeId(2, doubleTypeId);
+    doubleHolder.value = 2.5;
+    vector.setSafe(2, doubleHolder);
+    vector.setValueCount(3);
+    return vector;
+  }
+
+  @Test
+  public void testChildrenCreatedOutOfOrderFollowGetField() {
+    try (DenseUnionVector vector = newVectorWithChildrenCreatedOutOfOrder()) {
+      List<Field> fields = vector.getField().getChildren();
+      List<FieldVector> children = vector.getChildrenFromFields();
+      assertEquals(fields.size(), children.size());
+      for (int i = 0; i < fields.size(); i++) {
+        assertEquals(fields.get(i).getType(), children.get(i).getField().getType());
+      }
+    }
+  }
+
+  @Test
+  public void testTransferWithChildrenCreatedOutOfOrder() {
+    try (DenseUnionVector source = newVectorWithChildrenCreatedOutOfOrder()) {
+      TransferPair pair = source.getTransferPair(allocator);
+      try (DenseUnionVector target = (DenseUnionVector) pair.getTo()) {
+        assertEquals(source.getField().getChildren(), target.getField().getChildren());
+
+        pair.splitAndTransfer(1, 2);
+        assertEquals(2, target.getValueCount());
+        assertEquals(42, target.getObject(0));
+        assertEquals(2.5, target.getObject(1));
+
+        pair.transfer();
+        assertEquals(3, target.getValueCount());
+        assertEquals(1.5, target.getObject(0));
+        assertEquals(42, target.getObject(1));
+        assertEquals(2.5, target.getObject(2));
+      }
+    }
+  }
+
+  @Test
+  public void testUnloadWithChildrenCreatedOutOfOrder() {
+    // The unloader takes the schema from getField() and the buffers from getChildrenFromFields()
+    try (VectorSchemaRoot root = VectorSchemaRoot.of(newVectorWithChildrenCreatedOutOfOrder());
+        ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+        VectorSchemaRoot loaded = VectorSchemaRoot.create(root.getSchema(), allocator)) {
+      new VectorLoader(loaded).load(batch);
+
+      FieldVector vector = loaded.getVector(0);
+      assertEquals(3, vector.getValueCount());
+      assertEquals(1.5, vector.getObject(0));
+      assertEquals(42, vector.getObject(1));
+      assertEquals(2.5, vector.getObject(2));
+    }
+  }
+
+  @Test
+  public void testRegisterNewTypeIdSkipsTypeIdsOfTransferTarget() {
+    Field field =
+        new Field(
+            "duv",
+            FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {2, 0})),
+            Arrays.asList(
+                Field.nullable("i", MinorType.INT.getType()),
+                Field.nullable("s", MinorType.VARCHAR.getType())));
+    try (DenseUnionVector source = (DenseUnionVector) field.createVector(allocator);
+        DenseUnionVector target = (DenseUnionVector) source.getTransferPair(allocator).getTo()) {
+      byte typeId = target.registerNewTypeId(Field.nullable("d", MinorType.FLOAT8.getType()));
+
+      assertEquals(3, typeId);
+      Field expected =
+          new Field(
+              "duv",
+              FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {2, 0, 3})),
+              Arrays.asList(
+                  Field.nullable("i", MinorType.INT.getType()),
+                  Field.nullable("s", MinorType.VARCHAR.getType()),
+                  Field.nullable("d", MinorType.FLOAT8.getType())));
+      assertEquals(expected, target.getField());
+    }
+  }
+
   @Test
   public void testMaximumNumberOfChildren() {
     int[] typeIds = new int[Byte.MAX_VALUE + 1];