Repository navigation
GH-88: Support importing C Data Interface arrays with a non-zero offset #1419
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
25296e0
21df835
b997324
0321931
112a0ff
90e7afe
0dcebc1
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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; | ||
| } | ||
|
|
||
|
|
@@ -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( | ||
|
|
@@ -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); | ||
| 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(); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. When With type ids The same array imports correctly with offset 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
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed in 90e7afe with the second option:
|
||
| try (FieldVector unslicedVector = unslicedField.createVector(allocator)) { | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Keeping the duplicates in the temporary vector is not enough on its own. 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.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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
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 |
||
| // 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"); | ||
|
|
@@ -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 " | ||
|
|
@@ -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; | ||
| } | ||
| } | ||
| } | ||
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added in 0dcebc1.
testImportArrayWithNegativeOffsetFailscovers 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.
testImportSliceWithOffsetBeyondIntRangeimports 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 anArithmeticExceptionfromMath.addExact, and with bounds checking offcheckedCastToIntcould 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).