Skip to content

GH-200: Keep registered type count when transferring a DenseUnionVector #1381

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sharanggupta wants to merge 3 commits into
base: main
Choose a base branch
from
+318 −10
Open

GH-200: Keep registered type count when transferring a DenseUnionVector #1381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3367abd
GH-200: Keep registered type count when transferring a DenseUnionVector
sharanggupta Oct 21, 2026
f41e9a8
GH-200: Test a second transfer for every vector kind
sharanggupta Oct 22, 2026
7f8447e
GH-200: Map transferred children by type id
sharanggupta Nov 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 17 additions & 10 deletions vector/src/main/codegen/templates/DenseUnionVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];

Expand Down Expand Up @@ -623,14 +624,19 @@ 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;

@jbonofre jbonofre Nov 6, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This cause a regression for unions whose type ids aren't 0..n-1, for example [5, 10], which is what IPC or the C Data interface gives us when the schema declares type ids.

createTransferPairs fills to.childVectors at the loop index (0 and 1), but setChildVectorValueCounts reads childVectors[typeMapFields[i]], i.e. childVectors[5] and childVectors[10], which are null on the target. Before this change to.nextTypeId was 0, so that loop never ran. Now any setValueCount on the target after transfer() (including through VectorSchemaRoot.setRowCount) throws:

java.lang.NullPointerException: Cannot invoke "org.apache.arrow.vector.ValueVector.setValueCount(int)" because "this.childVectors[...]" is null
    at org.apache.arrow.vector.complex.DenseUnionVector.setChildVectorValueCounts(DenseUnionVector.java)
    at org.apache.arrow.vector.complex.DenseUnionVector.setValueCount(DenseUnionVector.java)

Here's an example to reproduce the regression:

Field field =
    new Field(
        "u",
        FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {5, 10})),
        asList(
            Field.nullable("i", MinorType.INT.getType()),
            Field.nullable("b", MinorType.BIT.getType())));
try (DenseUnionVector vector = (DenseUnionVector) field.createVector(allocator)) {
  vector.allocateNew();
  vector.setTypeId(0, (byte) 5);
  vector.setSafe(0, intHolder(5));
  vector.setTypeId(1, (byte) 10);
  vector.setSafe(1, bitHolder(true));
  vector.setTypeId(2, (byte) 5);
  vector.setSafe(2, intHolder(7));
  vector.setValueCount(3);

  TransferPair pair = vector.getTransferPair(allocator);
  try (DenseUnionVector target = (DenseUnionVector) pair.getTo()) {
    pair.transfer();
    target.setValueCount(3); // NPE
  }
}

@sharanggupta sharanggupta Nov 7, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the reproducer. The transfer pair indexed typeFields and childVectors by registration order instead of by type id, which stayed hidden while the target's nextTypeId was 0. 7f8447e maps them through typeMapFields in createTransferPairs, transfer() and splitAndTransfer(); the last one also sized its per-type arrays by the type count, so it overflowed for id 10. Your example is now testTransferWithDeclaredTypeIds, with a second transfer and a splitAndTransfer variant.

}

@Override
Expand All @@ -646,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;
Expand All @@ -669,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;
Expand All @@ -689,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();
}
}

Expand Down
64 changes: 64 additions & 0 deletions vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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;
}

}
Loading
Loading