Skip to content
Open
263 changes: 236 additions & 27 deletions c/src/main/java/org/apache/arrow/c/ArrayImporter.java
Original file line number Diff line number Diff line change
Expand Up @@ -25,29 +25,48 @@
import java.util.List;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.util.MemoryUtil;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.ExtensionTypeVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.types.UnionMode;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;
import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.util.TransferPair;

/** Importer for {@link ArrowArray}. */
final class ArrayImporter {
private static final int MAX_IMPORT_RECURSION_LEVEL = 64;

private final BufferAllocator allocator;
private final FieldVector vector;
// The field the vector was created from, or null to use the vector's own field. Unlike the
// field of an ExtensionTypeVector that does not override getField(), it has the extension type.
private final Field field;
private final DictionaryProvider dictionaryProvider;

private ReferenceCountedArrowArray underlyingAllocation;
private int recursionLevel;

ArrayImporter(
BufferAllocator allocator, FieldVector vector, DictionaryProvider dictionaryProvider) {
this(allocator, vector, /* field= */ null, dictionaryProvider);
}

ArrayImporter(
BufferAllocator allocator,
FieldVector vector,
Field field,
DictionaryProvider dictionaryProvider) {
this.allocator = Preconditions.checkNotNull(allocator);
this.vector = Preconditions.checkNotNull(vector);
this.field = field;
this.dictionaryProvider = dictionaryProvider;
}

Expand All @@ -71,8 +90,7 @@ void importArray(ArrowArray src) {
}
}

private void importChild(ArrayImporter parent, ArrowArray src) {
ArrowArray.Snapshot snapshot = src.snapshot();
private void importChild(ArrayImporter parent, ArrowArray.Snapshot snapshot) {
checkState(snapshot.release != NULL, "Cannot import released ArrowArray");
recursionLevel = parent.recursionLevel + 1;
checkState(
Expand All @@ -84,34 +102,181 @@ private void importChild(ArrayImporter parent, ArrowArray src) {
}

private void doImport(ArrowArray.Snapshot snapshot) {
// A non-zero offset is only meaningful for arrays that have buffers to offset into.
// Null arrays carry no buffers, so their offset is inert and safe to ignore.
checkState(snapshot.offset >= 0, "ArrowArray struct has negative offset (%s)", snapshot.offset);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

The new error for a negative offset has no test. There is also no test for offset plus length being larger than an int can hold.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added in 0dcebc1. testImportArrayWithNegativeOffsetFails covers the negative offset, and the importer now rejects a negative length the same way (testImportArrayWithNegativeLengthFails).

Offset plus length beyond an int is fine for most types: the importer skips the whole bytes of rows before the slice by advancing the buffer pointers, so the temporary vector holds fewer than 8 rows before the slice. testImportSliceWithOffsetBeyondIntRange imports 5 rows at offset 2^31 + 3, with the buffer pointers moved back by the rows that are skipped. Run-end encoded arrays are the exception, because their run ends count from the start of the array, so their rows cannot be skipped. Such an import used to fail with an ArithmeticException from Math.addExact, and with bounds checking off checkedCastToInt could truncate the offset. It now checks the row count first and fails with a message that gives the offset, the length and the rows needed (testImportRunEndEncodedSliceBeyondIntRangeFails).

checkState(snapshot.length >= 0, "ArrowArray struct has negative length (%s)", snapshot.length);
if (snapshot.offset != 0 && snapshot.length != 0) {
if (storageType().getTypeID() == ArrowTypeID.Struct) {
importSlicedStruct(snapshot);
} else {
importSliced(snapshot);
}
} else {
importUnsliced(snapshot, /* skippedRows= */ 0, /* trimChildren= */ false);
}
}

/**
* Imports a struct with a non-zero offset into the target itself, because transferring a struct
* matches its children by name, which loses children that have the same name as another. The
* offset of a struct applies to its children, so it is pushed into them, and only the validity
* bitmap of the struct itself is sliced.
*/
private void importSlicedStruct(ArrowArray.Snapshot snapshot) {
long skippedRows = snapshot.offset & ~7L;
int unslicedLength = unslicedLength(snapshot, skippedRows);
int offset = (int) (snapshot.offset - skippedRows);
int length = (int) snapshot.length;
importChildren(snapshot, snapshot.offset, /* childRowsPerRow= */ 1);
importDictionary(snapshot);

snapshot.offset = 0;
snapshot.length = unslicedLength;
snapshot.null_count = unknownNullCount(snapshot);
try (StructVector unsliced = StructVector.empty("", allocator)) {
importBuffers(unsliced, snapshot, skippedRows);
TransferPair pair = unsliced.getTransferPair(allocator);
pair.splitAndTransfer(offset, length);
try (StructVector sliced = (StructVector) pair.getTo()) {
vector.loadFieldBuffers(
new ArrowFieldNode(length, sliced.getNullCount()),
Collections.singletonList(sliced.getValidityBuffer()));
}
}
}

/**
* Java vectors cannot represent an offset, so the rows up to the end of the slice are imported
* into a temporary vector and the slice is then transferred into the target. Leading rows that
* fill whole bytes of a bitmap are skipped by advancing the buffer pointers, so that the
* temporary vector, and any validity bitmap synthesized for it, has fewer than 8 rows before the
* slice.
*/
private void importSliced(ArrowArray.Snapshot snapshot) {
// Run-end encoded arrays have no buffers, and their children are not indexed by row
long skippedRows =
storageType().getTypeID() == ArrowTypeID.RunEndEncoded ? 0 : snapshot.offset & ~7L;
int unslicedLength = unslicedLength(snapshot, skippedRows);
int offset = (int) (snapshot.offset - skippedRows);
int length = (int) snapshot.length;

snapshot.offset = 0;
snapshot.length = unslicedLength;
snapshot.null_count = unknownNullCount(snapshot);

Field unslicedField = field != null ? field : vector.getField();

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.

When field is null (the public Data.importIntoVector(allocator, array, vector, provider)), the temporary vector is built from vector.getField(). For a dense union, that doesn't reproduce the target. DenseUnionVector.getField() walks typeFields, which is indexed by type id, so it returns the children sorted by type id. The C array and the target keep the producer's order.

With type ids [10, 5] (children a, b), the temporary vector has children [b, a]. importUnsliced loads C child 0 (a) into the temporary type-5 vector, and createTransferPairs() then rewrites the target's typeMapFields in the temporary vector's order. If both children are IntVector, checkTransferable passes and rows read values from the wrong child, with no error. If the classes differ, it throws a misleading "its field create a ..." error.

The same array imports correctly with offset 0, and through Data.importVector, which passes the field from the C schema. So only the sliced import without a field is affected. All the dense union tests go through importVector, and testImportDenseUnionWithTypeIdsWithNonZeroOffset uses ascending ids with children of different types, so none of them catch this.

Field field = new Field("u",
    FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[] {10, 5})),
    Arrays.asList(
        Field.nullable("a", MinorType.INT.getType()),
        Field.nullable("b", MinorType.INT.getType())));
// source rows: a=100, b=201, a=102, b=203
Data.exportVector(allocator, source, null, array);
slice(array, 1, 3);
try (FieldVector target = field.createVector(allocator)) {
  Data.importIntoVector(allocator, array, target, null);
  // expected [201, 102, 203]
}

Could we build the temporary field in the target's child order (children from getChildrenFromFields(), type ids from typeMapFields), or make DenseUnionVector.getField() keep the child order?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Fixed in 90e7afe with the second option: DenseUnionVector.getField() now lists the children in their order, each with its type id, instead of sorting them by type id. Exports had the same mismatch: Data.exportVector with a schema takes the schema from getField() but the array's children from getChildrenFromFields(). createTransferPairs() now also sets the target's nextTypeId, which stayed at 0 for a target created by a transfer pair, so that target's getField() lists its children too.

testImportIntoDenseUnionWithDescendingTypeIdsWithNonZeroOffset (0dcebc1) is your example, imported with importIntoVector and no field. Before the fix its first row came back as 100. testImportDenseUnionWithDescendingTypeIdsWithNonZeroOffset imports every slice through a schema exported from getField(), and TestDenseUnionVector.testGetFieldKeepsChildOrder checks the field of the vector and of a transfer target.

try (FieldVector unslicedVector = unslicedField.createVector(allocator)) {

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.

Field.createVector builds a struct with the default CONFLICT_REPLACE policy, so duplicate child names collapse into one. importIntoVectorSchemaRoot uses StructVector.emptyWithDuplicates on purpose for this case (see testVectorSchemaRootWithDuplicatedFieldNames), but the sliced path loses that. With columns [a, a, b] and a non-zero offset, the temporary struct has 2 children, and importUnsliced fails with ArrowArray struct has 3 children (expected 2). The same array imports fine with offset 0.

Keeping the duplicates in the temporary vector is not enough on its own. StructTransferPair maps children by name (from.getChild(child)), so both a columns would be routed to the first one.

One option is to create the temporary struct with the target's conflict policy and transfer struct children by position. Another is to skip the temporary vector for structs: the importer already pushes the offset into struct children, so only the struct's own validity bitmap would need shifting.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Fixed in 0dcebc1 with the second option: a struct with a non-zero offset is now imported into the target itself. Its offset is pushed into its children, each of which is imported into the child at the same position, and only the struct's own validity bitmap is sliced, through a temporary struct without children. A transfer pair into the target would have called the target's setValueCount(), which allocates and zeroes the buffers of every child.

testImportVectorSchemaRootWithDuplicatedFieldNamesWithNonZeroOffset imports columns [a, a, b] at offsets 3, 8 and 13, with different values in the two a columns. Before the fix it failed in checkTransferable, because the collapsed temporary struct [a, b] put a VarCharVector against the target's second a.

This also removes the case my earlier reply left unsupported: an extension vector that does not report its type, as a child of the caller's struct, is now imported like a top-level vector, so testImportIntoStructWithPlainExtensionChildWithNonZeroOffset checks the values instead of the exception. Below a list, map or union it still fails with the same error.

// When the target is an extension vector whose getField() reports its storage type, the
// temporary vector is a storage vector, so transfer between the storage vectors
FieldVector from = ExtensionTypeVector.getStorageVector(unslicedVector);
FieldVector to = ExtensionTypeVector.getStorageVector(vector);
checkTransferable(from, to);
ArrayImporter unslicedImporter =
new ArrayImporter(allocator, unslicedVector, unslicedField, dictionaryProvider);
unslicedImporter.recursionLevel = recursionLevel;
unslicedImporter.underlyingAllocation = underlyingAllocation;
unslicedImporter.importUnsliced(snapshot, skippedRows, /* trimChildren= */ true);
from.makeTransferPair(to).splitAndTransfer(offset, length);
}
}

/**
* Returns how many rows a temporary vector needs for the rows of {@code snapshot} that follow its
* first {@code skippedRows} rows, which only run-end encoded arrays, whose rows cannot be
* skipped, can make more than a vector can hold.
*/
private static int unslicedLength(ArrowArray.Snapshot snapshot, long skippedRows) {
long length = Math.addExact(snapshot.offset - skippedRows, snapshot.length);
checkState(
snapshot.offset == 0
|| snapshot.length == 0
|| vector.getField().getType().getTypeID() == ArrowType.ArrowTypeID.Null,
"ArrowArray struct has non-zero offset (%s), which is not supported",
snapshot.offset);
length <= Integer.MAX_VALUE,
"Cannot import an ArrowArray with offset %s and length %s, because a vector cannot hold"
+ " the %s rows it has to import",
snapshot.offset,
snapshot.length,
length);
return (int) length;
}

/**
* Transfer pairs of nested vectors create or reuse the children of the target based on the
* classes of the source's children, so every vector of the target must be an instance of the
* class of its counterpart in the source.
*/
private static void checkTransferable(FieldVector from, FieldVector to) {
if (!from.getClass().isInstance(to)) {
String reason =
to instanceof ExtensionTypeVector && !(from instanceof ExtensionTypeVector)
? "its getField() does not report its extension type"
: "its field creates a " + from.getClass().getSimpleName();
throw new UnsupportedOperationException(
String.format(
"Cannot import an ArrowArray with a non-zero offset into %s %s, because %s",
to.getClass().getSimpleName(), to.getName(), reason));
}
List<FieldVector> fromChildren = from.getChildrenFromFields();
List<FieldVector> toChildren = to.getChildrenFromFields();
for (int i = 0; i < Math.min(fromChildren.size(), toChildren.size()); i++) {
checkTransferable(fromChildren.get(i), toChildren.get(i));
}
}

/**
* Imports an array whose offset is 0, skipping the first {@code skippedRows} rows of its buffers,
* which must be a multiple of 8. With {@code trimChildren}, children whose rows line up with the
* rows of the array are imported only for the rows of the array, rather than in full.
*/
private void importUnsliced(
ArrowArray.Snapshot snapshot, long skippedRows, boolean trimChildren) {
// First import children (required for reconstituting parent array data)
importChildren(snapshot, skippedRows, trimChildren ? childRowsPerRow(storageType()) : 0);
importDictionary(snapshot);
importBuffers(vector, snapshot, skippedRows);
}

/**
* Imports the children of {@code snapshot}. With a non-zero {@code childRowsPerRow}, which is how
* many rows of each child a row of the array spans, each child is imported only for the rows of
* the array that follow its first {@code skippedRows} rows, rather than in full.
*/
private void importChildren(
ArrowArray.Snapshot snapshot, long skippedRows, long childRowsPerRow) {
long[] children =
NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children));
if (children != null && children.length > 0) {
List<FieldVector> childVectors = vector.getChildrenFromFields();
checkState(
children.length == childVectors.size(),
"ArrowArray struct has %s children (expected %s)",
children.length,
childVectors.size());
for (int i = 0; i < children.length; i++) {
checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i);
ArrayImporter childImporter =
new ArrayImporter(allocator, childVectors.get(i), dictionaryProvider);
childImporter.importChild(this, ArrowArray.wrap(children[i]));
if (children == null || children.length == 0) {
return;
}
List<FieldVector> childVectors = vector.getChildrenFromFields();
checkState(
children.length == childVectors.size(),
"ArrowArray struct has %s children (expected %s)",
children.length,
childVectors.size());
List<Field> childFields = field != null ? field.getChildren() : null;
for (int i = 0; i < children.length; i++) {
checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i);
ArrowArray.Snapshot child = ArrowArray.wrap(children[i]).snapshot();
if (childRowsPerRow != 0) {
long firstRow = Math.multiplyExact(skippedRows, childRowsPerRow);
long rowCount = Math.multiplyExact(snapshot.length, childRowsPerRow);
checkState(
child.length - firstRow >= rowCount,
"ArrowArray struct child at position %s has length %s, but its parent needs %s rows",
i,
child.length,
firstRow + rowCount);
child.offset = Math.addExact(child.offset, firstRow);
child.length = rowCount;
child.null_count = unknownNullCount(child);
}
Field childField =
childFields != null && childFields.size() == children.length ? childFields.get(i) : null;
ArrayImporter childImporter =
new ArrayImporter(allocator, childVectors.get(i), childField, dictionaryProvider);
childImporter.importChild(this, child);
}
}

// Handle import of a dictionary encoded vector
/** Imports the dictionary of a dictionary encoded array. */
private void importDictionary(ArrowArray.Snapshot snapshot) {
if (snapshot.dictionary != NULL) {
DictionaryEncoding encoding = vector.getField().getDictionary();
checkNotNull(encoding, "Missing encoding on import of ArrowArray with dictionary");
Expand All @@ -124,23 +289,29 @@ private void doImport(ArrowArray.Snapshot snapshot) {

ArrayImporter dictionaryImporter =
new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider);
dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary));
dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary).snapshot());
}
}

// Import main data
/**
* Imports the buffers of {@code snapshot} into {@code target}, skipping the first {@code
* skippedRows} rows, which must be a multiple of 8.
*/
private void importBuffers(FieldVector target, ArrowArray.Snapshot snapshot, long skippedRows) {
ArrowFieldNode fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count);
long[] bufferPointers =
NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers));

try (final BufferImportTypeVisitor visitor =
new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) {
new BufferImportTypeVisitor(
allocator, underlyingAllocation, fieldNode, bufferPointers, skippedRows)) {
final List<ArrowBuf> buffers;
if (bufferPointers == null || bufferPointers.length == 0) {
buffers = Collections.emptyList();
} else {
buffers = vector.getField().getType().accept(visitor);
buffers = target.getField().getType().accept(visitor);
}
vector.loadFieldBuffers(fieldNode, buffers);
target.loadFieldBuffers(fieldNode, buffers);
} catch (Exception e) {
throw new IllegalArgumentException(
"Could not load buffers for field "
Expand All @@ -150,4 +321,42 @@ private void doImport(ArrowArray.Snapshot snapshot) {
e);
}
}

/**
* Returns the null count to import a different range of the rows of {@code snapshot} with than
* the range its null count covers. Without a validity bitmap it must be 0 for the buffers to
* load, otherwise it is unknown.
*/
private static long unknownNullCount(ArrowArray.Snapshot snapshot) {
boolean hasBitmap =
snapshot.n_buffers > 0
&& snapshot.buffers != NULL
&& MemoryUtil.getLong(snapshot.buffers) != NULL;
return hasBitmap ? -1 : 0;
}

/** The type of the vector, or the storage type of an extension type. */
private ArrowType storageType() {
ArrowType type = vector.getField().getType();
return type instanceof ArrowType.ExtensionType
? ((ArrowType.ExtensionType) type).storageType()
: type;
}

/**
* Returns how many rows of each child a row of an array of {@code type} spans, or 0 if its
* children are not indexed by its rows.
*/
private static long childRowsPerRow(ArrowType type) {
switch (type.getTypeID()) {
case Struct:
return 1;
case Union:
return ((ArrowType.Union) type).getMode() == UnionMode.Sparse ? 1 : 0;
case FixedSizeList:
return ((ArrowType.FixedSizeList) type).getListSize();
default:
return 0;
}
}
}
Loading
Loading