andygrove commented on code in PR #1419:
URL: https://github.com/apache/arrow-java/pull/1419#discussion_r4460641283
##########
c/src/main/java/org/apache/arrow/c/ArrayImporter.java:
##########
@@ -84,15 +100,78 @@ 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
- || snapshot.length == 0
- || vector.getField().getType().getTypeID() ==
ArrowType.ArrowTypeID.Null,
- "ArrowArray struct has non-zero offset (%s), which is not supported",
- snapshot.offset);
+ checkState(snapshot.offset >= 0, "ArrowArray struct has negative offset
(%s)", snapshot.offset);
+ if (snapshot.offset != 0 && snapshot.length != 0) {
+ importSliced(snapshot);
+ } else {
+ importUnsliced(snapshot, /* skippedRows= */ 0, /* trimChildren= */
false);
+ }
+ }
+
+ /**
+ * 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 offset = checkedCastToInt(snapshot.offset - skippedRows);
+ int length = checkedCastToInt(snapshot.length);
+
+ snapshot.offset = 0;
+ snapshot.length = Math.addExact(offset, length);
+ snapshot.null_count = unknownNullCount(snapshot);
+
+ Field unslicedField = field != null ? field : vector.getField();
Review Comment:
Fixed in 90e7afefa 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` (0dcebc16e)
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.
##########
c/src/main/java/org/apache/arrow/c/ArrayImporter.java:
##########
@@ -84,15 +100,78 @@ 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
- || snapshot.length == 0
- || vector.getField().getType().getTypeID() ==
ArrowType.ArrowTypeID.Null,
- "ArrowArray struct has non-zero offset (%s), which is not supported",
- snapshot.offset);
+ checkState(snapshot.offset >= 0, "ArrowArray struct has negative offset
(%s)", snapshot.offset);
+ if (snapshot.offset != 0 && snapshot.length != 0) {
+ importSliced(snapshot);
+ } else {
+ importUnsliced(snapshot, /* skippedRows= */ 0, /* trimChildren= */
false);
+ }
+ }
+
+ /**
+ * 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 offset = checkedCastToInt(snapshot.offset - skippedRows);
+ int length = checkedCastToInt(snapshot.length);
+
+ snapshot.offset = 0;
+ snapshot.length = Math.addExact(offset, length);
+ snapshot.null_count = unknownNullCount(snapshot);
+
+ Field unslicedField = field != null ? field : vector.getField();
+ try (FieldVector unslicedVector = unslicedField.createVector(allocator)) {
Review Comment:
Fixed in 0dcebc16e 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.
##########
vector/src/main/codegen/templates/DenseUnionVector.java:
##########
@@ -669,29 +672,54 @@ public void splitAndTransfer(int startIndex, int length) {
to.reallocOffsetBuffer();
}
- int [] typeCounts = new int[nextTypeId];
- int [] typeStarts = new int[nextTypeId];
- for (int i = 0; i < typeCounts.length; i++) {
- typeCounts[i] = 0;
- typeStarts[i] = -1;
+ int[] childIndexes = new int[typeFields.length];
+ Arrays.fill(childIndexes, -1);
+ for (int i = 0; i < nextTypeId; i++) {
+ childIndexes[typeMapFields[i]] = i;
}
+ // The window selects childCounts[i] values of child i, which are
renumbered from 0. The
+ // offsets of a child only have to increase, so the selected values may
have gaps.
+ int[] childCounts = new int[nextTypeId];
+ int[] childStarts = new int[nextTypeId];
+ int[] childEnds = new int[nextTypeId];
+ boolean[] contiguous = new boolean[nextTypeId];
+ Arrays.fill(contiguous, true);
for (int i = startIndex; i < startIndex + length; i++) {
byte typeId = typeBuffer.getByte(i);
if (typeId >= 0) {
- to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH,
typeCounts[typeId]);
- typeCounts[typeId] += 1;
- if (typeStarts[typeId] == -1) {
- typeStarts[typeId] = offsetBuffer.getInt((long) i * OFFSET_WIDTH);
+ int child = childIndexes[typeId];
+ int offset = offsetBuffer.getInt((long) i * OFFSET_WIDTH);
+ if (childCounts[child] == 0) {
+ childStarts[child] = offset;
+ } else if (offset != childEnds[child]) {
+ contiguous[child] = false;
}
+ childEnds[child] = offset + 1;
+ to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH,
childCounts[child]);
+ childCounts[child] += 1;
}
}
- // transfer vector values
+ // transfer vector values, and copy those of children whose selected
values have gaps
+ for (int i = 0; i < nextTypeId; i++) {
+ if (childCounts[i] > 0 && contiguous[i]) {
+ internalTransferPairs[i].splitAndTransfer(childStarts[i],
childCounts[i]);
+ to.childVectors[typeMapFields[i]] = internalTransferPairs[i].getTo();
+ }
+ }
+ for (int i = startIndex; i < startIndex + length; i++) {
+ byte typeId = typeBuffer.getByte(i);
+ if (typeId >= 0 && !contiguous[childIndexes[typeId]]) {
+ internalTransferPairs[childIndexes[typeId]].copyValueSafe(
Review Comment:
Fixed in 90e7afefa. The values of children with gaps are now copied before
any child is transferred. If a copy throws `UnsupportedOperationException`, the
target is cleared and the split throws an `UnsupportedOperationException` that
names the union, the child and its class, with the original as the cause. The
copy itself is the check, rather than a list of classes: a check on the child's
class would miss a struct with a run-end encoded field, whose copy fails the
same way.
`testSplitAndTransferWithOffsetGapsInRunEndEncodedChild` and
`testSplitAndTransferWithOffsetGapsInLargeListViewChild` check the message, and
that the target has no rows and its int child no values. Before the fix the int
child had already been transferred.
##########
c/src/test/java/org/apache/arrow/c/RoundtripTest.java:
##########
@@ -1126,34 +1109,926 @@ public void testImportArrayWithNonZeroOffset() {
try (IntVector source = new IntVector("source", allocator);
IntVector destination = new IntVector("destination", allocator);
ArrowArray array = ArrowArray.allocateNew(allocator)) {
- setVector(source, 1, 2, 3);
+ setVector(source, 1, 2, null, 4, 5);
Data.exportVector(allocator, source, null, array);
+ slice(array, 1, 3);
+
+ Data.importIntoVector(allocator, array, destination, null);
+
+ assertEquals(3, destination.getValueCount());
+ assertEquals(1, destination.getNullCount());
+ assertEquals(2, destination.get(0));
+ assertTrue(destination.isNull(1));
+ assertEquals(4, destination.get(2));
+ }
+ }
+ @Test
+ public void testImportIntVectorWithNonZeroOffset() {
+ try (IntVector source = new IntVector("source", allocator)) {
+ setVector(source, rows(i -> i % 3 == 0 ? null : i, Integer[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportVarCharWithNonZeroOffset() {
+ try (VarCharVector source = new VarCharVector("source", allocator)) {
+ setVector(source, rows(i -> i % 4 == 1 ? null : "v".repeat(i % 5),
String[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportListWithNonZeroOffset() {
+ try (ListVector source = ListVector.empty("source", allocator)) {
+ setVector(source, testLists());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithNonZeroOffset() {
+ // The offset of a struct applies to its children, which are exported at
full length
+ try (StructVector source = StructVector.empty("source", allocator);
+ StructVector destination = StructVector.empty("destination",
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ Map<String, List<Integer>> values = new HashMap<>();
+ values.put("a", Arrays.asList(1, 2, 3, 4));
+ values.put("b", Arrays.asList(10, null, 30, 40));
+ setVector(source, values);
+
destination.initializeChildrenFromFields(source.getField().getChildren());
+ Data.exportVector(allocator, source, null, array);
+ slice(array, 1, 2);
+
+ Data.importIntoVector(allocator, array, destination, null);
+
+ assertEquals(2, destination.getValueCount());
+ IntVector a = (IntVector) destination.getChild("a");
+ IntVector b = (IntVector) destination.getChild("b");
+ assertEquals(2, a.getValueCount());
+ assertEquals(2, a.get(0));
+ assertEquals(3, a.get(1));
+ assertTrue(b.isNull(0));
+ assertEquals(30, b.get(1));
+ }
+ }
+
+ @Test
+ public void testImportStructWithNullsWithNonZeroOffset() {
+ try (StructVector source = StructVector.empty("source", allocator)) {
+ setVector(source, testStructValues());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithChildOffsetsWithNonZeroOffset() {
+ // The offset of a child adds to the offset of the struct
+ try (StructVector source = StructVector.empty("source", allocator)) {
+ setVector(source, testStructValues());
+ for (int childOffset : new int[] {3, 8, 11}) {
+ for (int offset = 0; offset < ROW_COUNT - childOffset; offset++) {
+ int length = ROW_COUNT - childOffset - offset;
+ try (ArrowArray array = ArrowArray.allocateNew(allocator);
+ ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
+ Data.exportVector(allocator, source, null, array, schema);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ long[] children = NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children);
+ for (long child : children) {
+ slice(ArrowArray.wrap(child), childOffset, ROW_COUNT -
childOffset);
+ }
+ slice(array, offset, length);
+
+ try (StructVector imported =
+ (StructVector) Data.importVector(allocator, array, schema,
null)) {
+ assertEquals(length, imported.getValueCount());
+ for (int i = 0; i < length; i++) {
+ String row = "child offset " + childOffset + ", offset " +
offset + ", index " + i;
+ assertEquals(source.isNull(offset + i), imported.isNull(i),
row);
+ for (String name : source.getChildFieldNames()) {
+ assertEquals(
+ source.getChild(name).getObject(childOffset + offset +
i),
+ imported.getChild(name).getObject(i),
+ row + ", child " + name);
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testImportStructWithShortChildWithNonZeroOffset() {
+ try (StructVector source = StructVector.empty("source", allocator);
+ StructVector destination = StructVector.empty("destination",
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ setVector(source, testStructValues());
+
destination.initializeChildrenFromFields(source.getField().getChildren());
+ Data.exportVector(allocator, source, null, array);
ArrowArray.Snapshot snapshot = array.snapshot();
- snapshot.offset = 1;
- array.save(snapshot);
+ for (long child : NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children)) {
+ slice(ArrowArray.wrap(child), 0, 13);
+ }
+ // The slice needs 14 rows of each child
+ slice(array, 9, 5);
Exception e =
assertThrows(
IllegalStateException.class,
() -> Data.importIntoVector(allocator, array, destination,
null));
assertEquals(
- "ArrowArray struct has non-zero offset (1), which is not supported",
e.getMessage());
+ "ArrowArray struct child at position 0 has length 13, but its parent
needs 14 rows",
+ e.getMessage());
+ }
+ }
+
+ @Test
+ public void testImportVectorSchemaRootWithSlicedChildren() {
+ // Arrow C++ exports record_batch->Slice(8, 2) as a struct whose children
carry the offset
+ try (VectorSchemaRoot source = createTestVSR();
+ VectorSchemaRoot destination =
VectorSchemaRoot.create(source.getSchema(), allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ Data.exportVectorSchemaRoot(allocator, source, null, array);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ slice(array, 0, 2);
+ for (long child : NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children)) {
+ slice(ArrowArray.wrap(child), 8, 2);
+ }
+
+ Data.importIntoVectorSchemaRoot(allocator, array, destination, null);
+
+ assertEquals(2, destination.getRowCount());
+ BitVector bits = (BitVector) destination.getVector("boolean");
+ VarCharVector strings = (VarCharVector) destination.getVector("varchar");
+ assertEquals(0, bits.get(0));
+ assertEquals(1, bits.get(1));
+ assertEquals("test8", new String(strings.get(0),
StandardCharsets.UTF_8));
+ assertEquals("test9", new String(strings.get(1),
StandardCharsets.UTF_8));
+ }
+ }
+
+ @Test
+ public void testImportVectorSchemaRootWithSlicedStructOfPlainExtension() {
+ // The schema of the root has the extension type that the vector's
getField() does not report
+ Field field =
+ new Field(
+ "s",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(Field.nullable("p",
PlainType.INSTANCE)));
+ Schema schema = new Schema(Collections.singletonList(field));
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (VectorSchemaRoot source = VectorSchemaRoot.create(schema,
allocator);
+ VectorSchemaRoot destination = VectorSchemaRoot.create(schema,
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ StructVector struct = (StructVector) source.getVector("s");
+ setTestStrings(((PlainVector)
struct.getChild("p")).getUnderlyingVector());
+ for (int i = 0; i < ROW_COUNT; i++) {
+ struct.setIndexDefined(i);
+ }
+ source.setRowCount(ROW_COUNT);
+ Data.exportVectorSchemaRoot(allocator, source, null, array);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ slice(array, 0, 5);
+ long[] children = NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children);
+ for (long child : children) {
+ slice(ArrowArray.wrap(child), 9, 5);
+ }
+
+ Data.importIntoVectorSchemaRoot(allocator, array, destination,
null);
+
+ assertEquals(5, destination.getRowCount());
+ StructVector imported = (StructVector) destination.getVector("s");
+ assertInstanceOf(PlainVector.class, imported.getChild("p"));
+ for (int i = 0; i < 5; i++) {
+ assertEquals(struct.getObject(9 + i), imported.getObject(i),
"index " + i);
+ }
+ }
+ });
+ }
+
+ @Test
+ public void testImportExtensionTypeVectorWithNonZeroOffset() {
+ try (UuidVector source = new UuidVector("source", allocator)) {
+ setTestUuids(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithExtensionTypeChildWithNonZeroOffset() {
+ try (StructVector source = StructVector.empty("source", allocator)) {
+ UuidVector child =
+ source.addOrGet("u", FieldType.nullable(UuidType.INSTANCE),
UuidVector.class);
+ setTestUuids(child);
+ setStructRowsDefined(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportBitVectorWithNonZeroOffset() {
+ try (BitVector source = new BitVector("source", allocator)) {
+ setVector(source, rows(i -> i % 5 == 0 ? null : i % 2, Integer[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportLargeListWithNonZeroOffset() {
+ try (LargeListVector source = LargeListVector.empty("source", allocator)) {
+ setVector(source, testLists());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testImportFixedSizeListWithNonZeroOffset() {
+ try (FixedSizeListVector source = FixedSizeListVector.empty("source", 2,
allocator)) {
+ setVector(source, rows(i -> Arrays.asList(2 * i, 2 * i + 1),
List[]::new));
+ for (int i = 1; i < ROW_COUNT; i += 6) {
+ source.setNull(i);
+ }
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportMapWithNonZeroOffset() {
+ try (MapVector source = MapVector.empty("source", allocator, false)) {
+ writeTestMaps(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportUnionWithNonZeroOffset() {
+ try (UnionVector source = UnionVector.empty("source", allocator)) {
+ writeTestUnion(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportDenseUnionWithNonZeroOffset() {
+ final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
+ uInt4Holder.isSet = 1;
+ final NullableIntHolder intHolder = new NullableIntHolder();
+ try (DenseUnionVector source = DenseUnionVector.empty("source",
allocator)) {
+ source.allocateNew();
+ byte uint4TypeId = source.registerNewTypeId(Field.nullable("u",
MinorType.UINT4.getType()));
+ byte intTypeId = source.registerNewTypeId(Field.nullable("i",
MinorType.INT.getType()));
+ for (int i = 0; i < ROW_COUNT; i++) {
+ if (i % 3 == 0) {
+ uInt4Holder.value = 100 + i;
+ source.setTypeId(i, uint4TypeId);
+ source.setSafe(i, uInt4Holder);
+ } else {
+ intHolder.isSet = i % 5 == 0 ? 0 : 1;
+ intHolder.value = -i;
+ source.setTypeId(i, intTypeId);
+ source.setSafe(i, intHolder);
+ }
+ }
+ source.setValueCount(ROW_COUNT);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportDenseUnionWithTypeIdsWithNonZeroOffset() {
+ // The type ids differ from the positions of the children they select
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[]
{5, 10})),
+ Arrays.asList(
+ Field.nullable("i", MinorType.INT.getType()),
+ Field.nullable("s", MinorType.VARCHAR.getType())));
+ try (DenseUnionVector source = (DenseUnionVector)
field.createVector(allocator)) {
+ source.allocateNew();
+ IntVector ints = (IntVector) source.getVectorByType((byte) 5);
+ VarCharVector strings = (VarCharVector) source.getVectorByType((byte)
10);
+ int intCount = 0;
+ int stringCount = 0;
+ for (int i = 0; i < ROW_COUNT; i++) {
+ if (i % 3 == 0) {
+ source.setTypeId(i, (byte) 10);
+ source.setOffset(i, stringCount);
+ strings.setSafe(stringCount++, ("s" +
i).getBytes(StandardCharsets.UTF_8));
+ } else {
+ source.setTypeId(i, (byte) 5);
+ source.setOffset(i, intCount);
+ ints.setSafe(intCount++, i);
+ }
+ }
+ source.setValueCount(ROW_COUNT);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportDenseUnionWithOffsetGapsWithNonZeroOffset() {
+ // The offsets of a dense union only have to increase, so they may skip
values of a child
+ try (DenseUnionVector source =
+ (DenseUnionVector)
denseUnionWithIntChild("source").createVector(allocator)) {
+ setTestValuesWithOffsetGaps(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithDenseUnionWithOffsetGapsWithNonZeroOffset() {
+ // The struct sets the value count of the dense union after transferring it
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(denseUnionWithIntChild("u")));
+ try (StructVector source = (StructVector) field.createVector(allocator)) {
+ for (int i = 0; i < OFFSETS_WITH_GAPS.length; i++) {
+ if (i % 6 != 5) {
+ source.setIndexDefined(i);
+ }
+ }
+ source.setValueCount(OFFSETS_WITH_GAPS.length);
+ setTestValuesWithOffsetGaps((DenseUnionVector) source.getChild("u"));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ private static final int[] OFFSETS_WITH_GAPS = {0, 1, 3, 4, 5, 8, 9, 12, 13,
14, 17, 19};
+
+ private static Field denseUnionWithIntChild(String name) {
+ return new Field(
+ name,
+ FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[]
{0})),
+ Collections.singletonList(Field.nullable("i",
MinorType.INT.getType())));
+ }
+
+ /** Selects the values at {@link #OFFSETS_WITH_GAPS} from the child of a
dense union. */
+ private static void setTestValuesWithOffsetGaps(DenseUnionVector vector) {
+ vector.allocateNew();
+ IntVector child = (IntVector) vector.getVectorByType((byte) 0);
+ setVector(child, rows(i -> 100 + i, Integer[]::new));
+ for (int i = 0; i < OFFSETS_WITH_GAPS.length; i++) {
+ vector.setTypeId(i, (byte) 0);
+ vector.setOffset(i, OFFSETS_WITH_GAPS[i]);
+ }
+ vector.setValueCount(OFFSETS_WITH_GAPS.length);
+ // setValueCount() sets the child's value count to the number of rows
selecting it
+ child.setValueCount(ROW_COUNT);
+ }
+
+ @Test
+ public void testImportRunEndEncodedWithNonZeroOffset() {
+ try (RunEndEncodedVector source = RunEndEncodedVector.empty("source",
allocator)) {
+ setVector(source, List.of(2, 3, 6, 10, 11, 15, 20), List.of(1, 2, 3, 4,
5, 6, 7));
+ assertSlicedImportsMatchSource(source);
}
}
+ @Test
+ public void testImportViewVarCharWithNonZeroOffset() {
Review Comment:
Added in 0dcebc16e, each through `assertSlicedImportsMatchSource`:
`ListView`, `LargeListView`, `ViewVarBinary`, `VarBinary`, `LargeVarChar`,
`LargeVarBinary`, `Decimal`, `Decimal256`, `FixedSizeBinary`, `Float16`, and
timestamp with time zone (microseconds, `Asia/Tokyo`). They all passed without
changes to the importer. The helper now compares binary values with
`assertArrayEquals`, since `getObject()` returns a `byte[]`. `validateFull`
supports neither list views nor float16, so the helper skips it for them, as it
already did for maps and run-end encoded vectors.
##########
c/src/test/java/org/apache/arrow/c/RoundtripTest.java:
##########
@@ -1126,34 +1109,926 @@ public void testImportArrayWithNonZeroOffset() {
try (IntVector source = new IntVector("source", allocator);
IntVector destination = new IntVector("destination", allocator);
ArrowArray array = ArrowArray.allocateNew(allocator)) {
- setVector(source, 1, 2, 3);
+ setVector(source, 1, 2, null, 4, 5);
Data.exportVector(allocator, source, null, array);
+ slice(array, 1, 3);
+
+ Data.importIntoVector(allocator, array, destination, null);
+
+ assertEquals(3, destination.getValueCount());
+ assertEquals(1, destination.getNullCount());
+ assertEquals(2, destination.get(0));
+ assertTrue(destination.isNull(1));
+ assertEquals(4, destination.get(2));
+ }
+ }
+ @Test
+ public void testImportIntVectorWithNonZeroOffset() {
+ try (IntVector source = new IntVector("source", allocator)) {
+ setVector(source, rows(i -> i % 3 == 0 ? null : i, Integer[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportVarCharWithNonZeroOffset() {
+ try (VarCharVector source = new VarCharVector("source", allocator)) {
+ setVector(source, rows(i -> i % 4 == 1 ? null : "v".repeat(i % 5),
String[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportListWithNonZeroOffset() {
+ try (ListVector source = ListVector.empty("source", allocator)) {
+ setVector(source, testLists());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithNonZeroOffset() {
Review Comment:
Added in 0dcebc16e. `testImportStructOfMapAndListOfStructsWithNonZeroOffset`
imports, at every offset, a struct with a nullable map and a non-nullable list
of structs whose `x` is nullable and `y` is not.
`testImportListWithChildOffsetWithNonZeroOffset` slices a list at offset 1 and
its child at offset 2, and checks the lists against hand-computed values.
##########
c/src/test/java/org/apache/arrow/c/RoundtripTest.java:
##########
@@ -1126,34 +1109,926 @@ public void testImportArrayWithNonZeroOffset() {
try (IntVector source = new IntVector("source", allocator);
IntVector destination = new IntVector("destination", allocator);
ArrowArray array = ArrowArray.allocateNew(allocator)) {
- setVector(source, 1, 2, 3);
+ setVector(source, 1, 2, null, 4, 5);
Data.exportVector(allocator, source, null, array);
+ slice(array, 1, 3);
+
+ Data.importIntoVector(allocator, array, destination, null);
+
+ assertEquals(3, destination.getValueCount());
+ assertEquals(1, destination.getNullCount());
+ assertEquals(2, destination.get(0));
+ assertTrue(destination.isNull(1));
+ assertEquals(4, destination.get(2));
+ }
+ }
+ @Test
+ public void testImportIntVectorWithNonZeroOffset() {
+ try (IntVector source = new IntVector("source", allocator)) {
+ setVector(source, rows(i -> i % 3 == 0 ? null : i, Integer[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportVarCharWithNonZeroOffset() {
+ try (VarCharVector source = new VarCharVector("source", allocator)) {
+ setVector(source, rows(i -> i % 4 == 1 ? null : "v".repeat(i % 5),
String[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportListWithNonZeroOffset() {
+ try (ListVector source = ListVector.empty("source", allocator)) {
+ setVector(source, testLists());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithNonZeroOffset() {
+ // The offset of a struct applies to its children, which are exported at
full length
+ try (StructVector source = StructVector.empty("source", allocator);
+ StructVector destination = StructVector.empty("destination",
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ Map<String, List<Integer>> values = new HashMap<>();
+ values.put("a", Arrays.asList(1, 2, 3, 4));
+ values.put("b", Arrays.asList(10, null, 30, 40));
+ setVector(source, values);
+
destination.initializeChildrenFromFields(source.getField().getChildren());
+ Data.exportVector(allocator, source, null, array);
+ slice(array, 1, 2);
+
+ Data.importIntoVector(allocator, array, destination, null);
+
+ assertEquals(2, destination.getValueCount());
+ IntVector a = (IntVector) destination.getChild("a");
+ IntVector b = (IntVector) destination.getChild("b");
+ assertEquals(2, a.getValueCount());
+ assertEquals(2, a.get(0));
+ assertEquals(3, a.get(1));
+ assertTrue(b.isNull(0));
+ assertEquals(30, b.get(1));
+ }
+ }
+
+ @Test
+ public void testImportStructWithNullsWithNonZeroOffset() {
+ try (StructVector source = StructVector.empty("source", allocator)) {
+ setVector(source, testStructValues());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithChildOffsetsWithNonZeroOffset() {
+ // The offset of a child adds to the offset of the struct
+ try (StructVector source = StructVector.empty("source", allocator)) {
+ setVector(source, testStructValues());
+ for (int childOffset : new int[] {3, 8, 11}) {
+ for (int offset = 0; offset < ROW_COUNT - childOffset; offset++) {
+ int length = ROW_COUNT - childOffset - offset;
+ try (ArrowArray array = ArrowArray.allocateNew(allocator);
+ ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
+ Data.exportVector(allocator, source, null, array, schema);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ long[] children = NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children);
+ for (long child : children) {
+ slice(ArrowArray.wrap(child), childOffset, ROW_COUNT -
childOffset);
+ }
+ slice(array, offset, length);
+
+ try (StructVector imported =
+ (StructVector) Data.importVector(allocator, array, schema,
null)) {
+ assertEquals(length, imported.getValueCount());
+ for (int i = 0; i < length; i++) {
+ String row = "child offset " + childOffset + ", offset " +
offset + ", index " + i;
+ assertEquals(source.isNull(offset + i), imported.isNull(i),
row);
+ for (String name : source.getChildFieldNames()) {
+ assertEquals(
+ source.getChild(name).getObject(childOffset + offset +
i),
+ imported.getChild(name).getObject(i),
+ row + ", child " + name);
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testImportStructWithShortChildWithNonZeroOffset() {
+ try (StructVector source = StructVector.empty("source", allocator);
+ StructVector destination = StructVector.empty("destination",
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ setVector(source, testStructValues());
+
destination.initializeChildrenFromFields(source.getField().getChildren());
+ Data.exportVector(allocator, source, null, array);
ArrowArray.Snapshot snapshot = array.snapshot();
- snapshot.offset = 1;
- array.save(snapshot);
+ for (long child : NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children)) {
+ slice(ArrowArray.wrap(child), 0, 13);
+ }
+ // The slice needs 14 rows of each child
+ slice(array, 9, 5);
Exception e =
assertThrows(
IllegalStateException.class,
() -> Data.importIntoVector(allocator, array, destination,
null));
assertEquals(
- "ArrowArray struct has non-zero offset (1), which is not supported",
e.getMessage());
+ "ArrowArray struct child at position 0 has length 13, but its parent
needs 14 rows",
+ e.getMessage());
+ }
+ }
+
+ @Test
+ public void testImportVectorSchemaRootWithSlicedChildren() {
+ // Arrow C++ exports record_batch->Slice(8, 2) as a struct whose children
carry the offset
+ try (VectorSchemaRoot source = createTestVSR();
+ VectorSchemaRoot destination =
VectorSchemaRoot.create(source.getSchema(), allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ Data.exportVectorSchemaRoot(allocator, source, null, array);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ slice(array, 0, 2);
+ for (long child : NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children)) {
+ slice(ArrowArray.wrap(child), 8, 2);
+ }
+
+ Data.importIntoVectorSchemaRoot(allocator, array, destination, null);
+
+ assertEquals(2, destination.getRowCount());
+ BitVector bits = (BitVector) destination.getVector("boolean");
+ VarCharVector strings = (VarCharVector) destination.getVector("varchar");
+ assertEquals(0, bits.get(0));
+ assertEquals(1, bits.get(1));
+ assertEquals("test8", new String(strings.get(0),
StandardCharsets.UTF_8));
+ assertEquals("test9", new String(strings.get(1),
StandardCharsets.UTF_8));
+ }
+ }
+
+ @Test
+ public void testImportVectorSchemaRootWithSlicedStructOfPlainExtension() {
+ // The schema of the root has the extension type that the vector's
getField() does not report
+ Field field =
+ new Field(
+ "s",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(Field.nullable("p",
PlainType.INSTANCE)));
+ Schema schema = new Schema(Collections.singletonList(field));
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (VectorSchemaRoot source = VectorSchemaRoot.create(schema,
allocator);
+ VectorSchemaRoot destination = VectorSchemaRoot.create(schema,
allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ StructVector struct = (StructVector) source.getVector("s");
+ setTestStrings(((PlainVector)
struct.getChild("p")).getUnderlyingVector());
+ for (int i = 0; i < ROW_COUNT; i++) {
+ struct.setIndexDefined(i);
+ }
+ source.setRowCount(ROW_COUNT);
+ Data.exportVectorSchemaRoot(allocator, source, null, array);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ slice(array, 0, 5);
+ long[] children = NativeUtil.toJavaArray(snapshot.children, (int)
snapshot.n_children);
+ for (long child : children) {
+ slice(ArrowArray.wrap(child), 9, 5);
+ }
+
+ Data.importIntoVectorSchemaRoot(allocator, array, destination,
null);
+
+ assertEquals(5, destination.getRowCount());
+ StructVector imported = (StructVector) destination.getVector("s");
+ assertInstanceOf(PlainVector.class, imported.getChild("p"));
+ for (int i = 0; i < 5; i++) {
+ assertEquals(struct.getObject(9 + i), imported.getObject(i),
"index " + i);
+ }
+ }
+ });
+ }
+
+ @Test
+ public void testImportExtensionTypeVectorWithNonZeroOffset() {
+ try (UuidVector source = new UuidVector("source", allocator)) {
+ setTestUuids(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithExtensionTypeChildWithNonZeroOffset() {
+ try (StructVector source = StructVector.empty("source", allocator)) {
+ UuidVector child =
+ source.addOrGet("u", FieldType.nullable(UuidType.INSTANCE),
UuidVector.class);
+ setTestUuids(child);
+ setStructRowsDefined(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportBitVectorWithNonZeroOffset() {
+ try (BitVector source = new BitVector("source", allocator)) {
+ setVector(source, rows(i -> i % 5 == 0 ? null : i % 2, Integer[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportLargeListWithNonZeroOffset() {
+ try (LargeListVector source = LargeListVector.empty("source", allocator)) {
+ setVector(source, testLists());
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testImportFixedSizeListWithNonZeroOffset() {
+ try (FixedSizeListVector source = FixedSizeListVector.empty("source", 2,
allocator)) {
+ setVector(source, rows(i -> Arrays.asList(2 * i, 2 * i + 1),
List[]::new));
+ for (int i = 1; i < ROW_COUNT; i += 6) {
+ source.setNull(i);
+ }
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportMapWithNonZeroOffset() {
+ try (MapVector source = MapVector.empty("source", allocator, false)) {
+ writeTestMaps(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportUnionWithNonZeroOffset() {
+ try (UnionVector source = UnionVector.empty("source", allocator)) {
+ writeTestUnion(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportDenseUnionWithNonZeroOffset() {
+ final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
+ uInt4Holder.isSet = 1;
+ final NullableIntHolder intHolder = new NullableIntHolder();
+ try (DenseUnionVector source = DenseUnionVector.empty("source",
allocator)) {
+ source.allocateNew();
+ byte uint4TypeId = source.registerNewTypeId(Field.nullable("u",
MinorType.UINT4.getType()));
+ byte intTypeId = source.registerNewTypeId(Field.nullable("i",
MinorType.INT.getType()));
+ for (int i = 0; i < ROW_COUNT; i++) {
+ if (i % 3 == 0) {
+ uInt4Holder.value = 100 + i;
+ source.setTypeId(i, uint4TypeId);
+ source.setSafe(i, uInt4Holder);
+ } else {
+ intHolder.isSet = i % 5 == 0 ? 0 : 1;
+ intHolder.value = -i;
+ source.setTypeId(i, intTypeId);
+ source.setSafe(i, intHolder);
+ }
+ }
+ source.setValueCount(ROW_COUNT);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportDenseUnionWithTypeIdsWithNonZeroOffset() {
+ // The type ids differ from the positions of the children they select
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[]
{5, 10})),
+ Arrays.asList(
+ Field.nullable("i", MinorType.INT.getType()),
+ Field.nullable("s", MinorType.VARCHAR.getType())));
+ try (DenseUnionVector source = (DenseUnionVector)
field.createVector(allocator)) {
+ source.allocateNew();
+ IntVector ints = (IntVector) source.getVectorByType((byte) 5);
+ VarCharVector strings = (VarCharVector) source.getVectorByType((byte)
10);
+ int intCount = 0;
+ int stringCount = 0;
+ for (int i = 0; i < ROW_COUNT; i++) {
+ if (i % 3 == 0) {
+ source.setTypeId(i, (byte) 10);
+ source.setOffset(i, stringCount);
+ strings.setSafe(stringCount++, ("s" +
i).getBytes(StandardCharsets.UTF_8));
+ } else {
+ source.setTypeId(i, (byte) 5);
+ source.setOffset(i, intCount);
+ ints.setSafe(intCount++, i);
+ }
+ }
+ source.setValueCount(ROW_COUNT);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportDenseUnionWithOffsetGapsWithNonZeroOffset() {
+ // The offsets of a dense union only have to increase, so they may skip
values of a child
+ try (DenseUnionVector source =
+ (DenseUnionVector)
denseUnionWithIntChild("source").createVector(allocator)) {
+ setTestValuesWithOffsetGaps(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @Test
+ public void testImportStructWithDenseUnionWithOffsetGapsWithNonZeroOffset() {
+ // The struct sets the value count of the dense union after transferring it
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(denseUnionWithIntChild("u")));
+ try (StructVector source = (StructVector) field.createVector(allocator)) {
+ for (int i = 0; i < OFFSETS_WITH_GAPS.length; i++) {
+ if (i % 6 != 5) {
+ source.setIndexDefined(i);
+ }
+ }
+ source.setValueCount(OFFSETS_WITH_GAPS.length);
+ setTestValuesWithOffsetGaps((DenseUnionVector) source.getChild("u"));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ private static final int[] OFFSETS_WITH_GAPS = {0, 1, 3, 4, 5, 8, 9, 12, 13,
14, 17, 19};
+
+ private static Field denseUnionWithIntChild(String name) {
+ return new Field(
+ name,
+ FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[]
{0})),
+ Collections.singletonList(Field.nullable("i",
MinorType.INT.getType())));
+ }
+
+ /** Selects the values at {@link #OFFSETS_WITH_GAPS} from the child of a
dense union. */
+ private static void setTestValuesWithOffsetGaps(DenseUnionVector vector) {
+ vector.allocateNew();
+ IntVector child = (IntVector) vector.getVectorByType((byte) 0);
+ setVector(child, rows(i -> 100 + i, Integer[]::new));
+ for (int i = 0; i < OFFSETS_WITH_GAPS.length; i++) {
+ vector.setTypeId(i, (byte) 0);
+ vector.setOffset(i, OFFSETS_WITH_GAPS[i]);
+ }
+ vector.setValueCount(OFFSETS_WITH_GAPS.length);
+ // setValueCount() sets the child's value count to the number of rows
selecting it
+ child.setValueCount(ROW_COUNT);
+ }
+
+ @Test
+ public void testImportRunEndEncodedWithNonZeroOffset() {
+ try (RunEndEncodedVector source = RunEndEncodedVector.empty("source",
allocator)) {
+ setVector(source, List.of(2, 3, 6, 10, 11, 15, 20), List.of(1, 2, 3, 4,
5, 6, 7));
+ assertSlicedImportsMatchSource(source);
}
}
+ @Test
+ public void testImportViewVarCharWithNonZeroOffset() {
+ try (ViewVarCharVector source = new ViewVarCharVector("source",
allocator)) {
+ setVector(
+ source,
+ rows(
+ i ->
+ i % 4 == 1
+ ? null
+ : i % 2 == 0 ? "short " + i : "a string longer than
twelve bytes " + i,
+ String[]::new));
+ assertSlicedImportsMatchSource(source);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {1 << 16, (1 << 16) + 3})
+ public void testImportSliceWithoutValidityBitmapAllocatesOnlyTheSlice(int
offset) {
+ // Producers commonly omit the validity bitmap of a column without nulls.
Neither the bitmap
+ // synthesized for the import nor the imported buffers may grow with the
offset.
+ final int length = 10;
+ try (IntVector source = new IntVector("source", allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator);
+ ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
+ source.allocateNew(offset + length);
+ for (int i = 0; i < offset + length; i++) {
+ source.set(i, i);
+ }
+ source.setValueCount(offset + length);
+ Data.exportVector(allocator, source, null, array, schema);
+ slice(array, offset, length);
+ dropValidityBitmap(array);
+
+ try (BufferAllocator importAllocator =
+ allocator.newChildAllocator("import", 0, Long.MAX_VALUE);
+ IntVector imported =
+ (IntVector) Data.importVector(importAllocator, array, schema,
null)) {
+ assertEquals(length, imported.getValueCount());
+ assertEquals(0, imported.getNullCount());
+ for (int i = 0; i < length; i++) {
+ assertEquals(offset + i, imported.get(i));
+ }
+ assertTrue(
+ importAllocator.getAllocatedMemory() < 1024,
+ "allocated " + importAllocator.getAllocatedMemory() + " bytes");
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {10, 1 << 16, (1 << 16) + 3})
+ public void
testImportStructSliceWithoutValidityBitmapsAllocatesOnlyTheSlice(int offset) {
+ // The children of a sliced struct are exported at full length, and only
the rows of the
+ // slice may be imported from them
+ final int rowCount = (1 << 16) + 20;
+ final int length = 10;
+ try (StructVector source = StructVector.empty("source", allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator);
+ ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
+ IntVector child =
+ source.addOrGet("a", FieldType.nullable(MinorType.INT.getType()),
IntVector.class);
+ child.allocateNew(rowCount);
+ for (int i = 0; i < rowCount; i++) {
+ child.set(i, i);
+ source.setIndexDefined(i);
+ }
+ source.setValueCount(rowCount);
+ Data.exportVector(allocator, source, null, array, schema);
+ ArrowArray.Snapshot snapshot = array.snapshot();
+ dropValidityBitmap(array);
+
dropValidityBitmap(ArrowArray.wrap(NativeUtil.toJavaArray(snapshot.children,
1)[0]));
+ slice(array, offset, length);
+
+ try (BufferAllocator importAllocator =
+ allocator.newChildAllocator("import", 0, Long.MAX_VALUE);
+ StructVector imported =
+ (StructVector) Data.importVector(importAllocator, array, schema,
null)) {
+ assertEquals(length, imported.getValueCount());
+ IntVector importedChild = (IntVector) imported.getChild("a");
+ for (int i = 0; i < length; i++) {
+ assertFalse(imported.isNull(i));
+ assertEquals(offset + i, importedChild.get(i));
+ }
+ assertTrue(
+ importAllocator.getAllocatedMemory() < 1024,
+ "allocated " + importAllocator.getAllocatedMemory() + " bytes");
+ }
+ }
+ }
+
+ @Test
+ public void testImportOpaqueVectorWithNonZeroOffset() {
+ // OpaqueVector overrides getField() but not makeTransferPair()
+ withExtensionType(
+ OPAQUE_TYPE,
+ () -> {
+ try (OpaqueVector source =
+ (OpaqueVector) Field.nullable("source",
OPAQUE_TYPE).createVector(allocator)) {
+ setTestStrings((VarCharVector) source.getUnderlyingVector());
+ assertSlicedImportsMatchSource(source);
+ }
+ });
+ }
+
+ @Test
+ public void testImportStructWithOpaqueChildWithNonZeroOffset() {
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(Field.nullable("o", OPAQUE_TYPE)));
+ withExtensionType(
+ OPAQUE_TYPE,
+ () -> {
+ try (StructVector source = (StructVector)
field.createVector(allocator)) {
+ OpaqueVector child = (OpaqueVector) source.getChild("o");
+ setTestStrings((VarCharVector) child.getUnderlyingVector());
+ setStructRowsDefined(source);
+ assertSlicedImportsMatchSource(source);
+ }
+ });
+ }
+
+ @Test
+ public void testImportListOfOpaqueWithNonZeroOffset() {
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.List.INSTANCE),
+ Collections.singletonList(Field.nullable("item", OPAQUE_TYPE)));
+ withExtensionType(
+ OPAQUE_TYPE,
+ () -> {
+ try (ListVector source = (ListVector) field.createVector(allocator))
{
+ OpaqueVector items = (OpaqueVector) source.getDataVector();
+ writeTestStringLists(source, (VarCharVector)
items.getUnderlyingVector());
+ assertSlicedImportsMatchSource(source);
+ }
+ });
+ }
+
+ @Test
+ public void testImportPlainExtensionVectorWithNonZeroOffset() {
+ // PlainVector overrides neither getField() nor makeTransferPair()
+ Field field = Field.nullable("source", PlainType.INSTANCE);
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (PlainVector source = (PlainVector)
field.createVector(allocator)) {
+ setTestStrings(source.getUnderlyingVector());
+ assertSlicedImportsMatchSource(field, source);
+ }
+ });
+ }
+
+ @Test
+ public void testImportIntoPlainExtensionVectorWithNonZeroOffset() {
+ // The target reports only its storage type, so the temporary vector is a
plain VarCharVector
+ Field field = Field.nullable("source", PlainType.INSTANCE);
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (PlainVector source = (PlainVector)
field.createVector(allocator);
+ PlainVector destination = (PlainVector)
field.createVector(allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ setTestStrings(source.getUnderlyingVector());
+ Data.exportVector(allocator, source, null, array);
+ slice(array, 9, 5);
+
+ Data.importIntoVector(allocator, array, destination, null);
+
+ assertEquals(5, destination.getValueCount());
+ for (int i = 0; i < 5; i++) {
+ assertEquals(source.getObject(9 + i), destination.getObject(i),
"index " + i);
+ }
+ }
+ });
+ }
+
+ @Test
+ public void testImportStructWithPlainExtensionChildWithNonZeroOffset() {
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(Field.nullable("p",
PlainType.INSTANCE)));
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (StructVector source = (StructVector)
field.createVector(allocator)) {
+ setTestStrings(((PlainVector)
source.getChild("p")).getUnderlyingVector());
+ setStructRowsDefined(source);
+ assertSlicedImportsMatchSource(field, source);
+ }
+ });
+ }
+
+ @Test
+ public void testImportListOfPlainExtensionWithNonZeroOffset() {
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.List.INSTANCE),
+ Collections.singletonList(Field.nullable("item",
PlainType.INSTANCE)));
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (ListVector source = (ListVector) field.createVector(allocator))
{
+ PlainVector items = (PlainVector) source.getDataVector();
+ writeTestStringLists(source, items.getUnderlyingVector());
+ assertSlicedImportsMatchSource(field, source);
+ }
+ });
+ }
+
+ @Test
+ public void testImportIntoStructWithPlainExtensionChildWithNonZeroOffset() {
+ // Without the imported field the importer cannot create a PlainVector for
the struct's child
+ Field field =
+ new Field(
+ "source",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Collections.singletonList(Field.nullable("p",
PlainType.INSTANCE)));
+ withExtensionType(
+ PlainType.INSTANCE,
+ () -> {
+ try (StructVector source = (StructVector)
field.createVector(allocator);
+ StructVector destination = (StructVector)
field.createVector(allocator);
+ ArrowArray array = ArrowArray.allocateNew(allocator)) {
+ setTestStrings(((PlainVector)
source.getChild("p")).getUnderlyingVector());
+ setStructRowsDefined(source);
+ Data.exportVector(allocator, source, null, array);
+ slice(array, 9, 5);
+
+ Exception e =
+ assertThrows(
+ UnsupportedOperationException.class,
+ () -> Data.importIntoVector(allocator, array, destination,
null));
+ assertEquals(
+ "Cannot import an ArrowArray with a non-zero offset into
PlainVector p, because"
+ + " its getField() does not report its extension type",
+ e.getMessage());
+ assertInstanceOf(PlainVector.class, destination.getChild("p"));
+ }
+ });
+ }
+
+ private static final int ROW_COUNT = 20;
+
+ /** Returns the values of the {@link #ROW_COUNT} rows of a test vector. */
+ private static <T> T[] rows(IntFunction<T> value, IntFunction<T[]> newArray)
{
+ return IntStream.range(0, ROW_COUNT).mapToObj(value).toArray(newArray);
+ }
+
+ /** Lists of up to 3 values, and some nulls. */
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ private static List<Integer>[] testLists() {
+ List[] lists =
+ rows(
+ i ->
+ i % 5 == 2
+ ? null
+ : IntStream.range(0, i % 4)
+ .mapToObj(j -> 10 * i + j)
+ .collect(Collectors.toList()),
+ List[]::new);
+ return lists;
+ }
+
+ /** Two int children, which are both null in the rows where the struct is
null. */
+ private static Map<String, List<Integer>> testStructValues() {
+ Map<String, List<Integer>> values = new HashMap<>();
+ values.put("a", Arrays.asList(rows(i -> i % 7 == 0 ? null : i,
Integer[]::new)));
+ values.put(
+ "b", Arrays.asList(rows(i -> i % 7 == 0 || i % 3 == 0 ? null : 10 * i,
Integer[]::new)));
+ return values;
+ }
+
+ private static void setTestStrings(VarCharVector vector) {
+ setVector(vector, rows(i -> i % 4 == 1 ? null : "s" + i, String[]::new));
+ }
+
+ private static void setTestUuids(UuidVector vector) {
+ vector.setValueCount(ROW_COUNT);
+ for (int i = 0; i < ROW_COUNT; i++) {
+ if (i % 4 == 0) {
+ vector.setNull(i);
+ } else {
+ vector.set(i, new UUID(i, -i));
+ }
+ }
+ }
+
+ /** Defines all but some rows of a struct whose children are set already. */
+ private static void setStructRowsDefined(StructVector vector) {
+ for (int i = 0; i < ROW_COUNT; i++) {
+ if (i % 6 != 5) {
+ vector.setIndexDefined(i);
+ }
+ }
+ vector.setValueCount(ROW_COUNT);
+ }
+
+ /**
+ * Imports the slices of {@code source} that start at each offset and either
have one row or end
+ * at the end of the source, which covers offsets before and after the first
byte of a bitmap,
+ * aligned to a byte or not.
+ */
+ private void assertSlicedImportsMatchSource(FieldVector source) {
+ assertSlicedImportsMatchSource(source.getField(), source);
+ }
+
+ /** Like {@link #assertSlicedImportsMatchSource(FieldVector)}, exporting the
type of field. */
+ private void assertSlicedImportsMatchSource(Field field, FieldVector source)
{
+ int valueCount = source.getValueCount();
+ for (int offset = 1; offset < valueCount; offset++) {
+ assertSlicedImportMatchesSource(field, source, offset, 1);
+ assertSlicedImportMatchesSource(field, source, offset, valueCount -
offset);
+ }
+ }
+
+ /**
+ * Imports a slice of {@code source}, exported with the type of {@code
field}, and compares it
+ * with the same range of the source.
+ */
+ private void assertSlicedImportMatchesSource(
Review Comment:
Added in 0dcebc16e: the helper now checks the imported vector's null count
against the number of null rows in the slice of the source, for every test that
uses it. Unions and run-end encoded vectors have no validity bitmap of their
own and always report a null count of 0, so for them it expects 0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
