From 93ca23bc1e3c802bdb534735347957f73edbed95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Fri, 13 Nov 2026 11:03:58 +0100 Subject: [PATCH] GH-1471: Pair dense union children with their type ids DenseUnionVector.getField() lists the children in the order their type ids were registered, while the child vectors are created when first used. getChildrenFromFields() and the transfer pair paired them by creation order, which wrote the buffers of one child under the field of another in exports and mixed up children in transfers. Both now follow the registered type ids. registerNewTypeId() on a vector without declared type ids, such as a transfer target, also skips the type ids in use. --- .../codegen/templates/DenseUnionVector.java | 63 ++++++++--- .../arrow/vector/TestDenseUnionVector.java | 102 ++++++++++++++++++ 2 files changed, 150 insertions(+), 15 deletions(-) diff --git a/vector/src/main/codegen/templates/DenseUnionVector.java b/vector/src/main/codegen/templates/DenseUnionVector.java index e979738ffc..71930bce11 100644 --- a/vector/src/main/codegen/templates/DenseUnionVector.java +++ b/vector/src/main/codegen/templates/DenseUnionVector.java @@ -166,9 +166,26 @@ public void initializeChildrenFromFields(List 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 getChildrenFromFields() { - return internalStruct.getChildrenFromFields(); + List 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 addOrGet(byte typeId, MinorType minorType, Class 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)" diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java index 0fc02a77a2..b92708265d 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java +++ b/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 fields = vector.getField().getChildren(); + List 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];