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];