This is an automated email from the ASF dual-hosted git repository.
jbonofre pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new 24d53195d GH-1424: Write only the values a vector's offsets cover in
VectorUnloader (#1425)
24d53195d is described below
commit 24d53195dc059300438da62ee6f05f5aae87e082
Author: Andy Grove <[email protected]>
AuthorDate: Sun Nov 1 22:40:36 2026 -0700
GH-1424: Write only the values a vector's offsets cover in VectorUnloader
(#1425)
## What's Changed
`VectorUnloader` wrote a variable-width vector's data buffer from byte 0
up to its last offset, and a list's child vector from element 0. A
vector imported through the C Data interface from a sliced array keeps
the offsets of the slice, which start past zero, and the bytes or child
elements before them, so every such batch was serialized with values no
row references, and the k-th slice of an array carried the k-1 slices
before it. The Rust and C++ IPC writers re-base offsets in this case;
Java's `VectorAppender` was taught to in #1180.
`VectorUnloader.appendNodes` now re-bases a variable-width, list or map
vector whose first offset is not zero through `splitAndTransfer(0,
valueCount)` and writes the result, closing the temporary vector once
the record batch has retained its buffers. The copy is limited to the
offsets, and to the validity bits when the start is not byte-aligned;
the data buffers are shared. A list or map whose offsets start at zero
but whose child runs past the last offset, as the first slice of a list
array does, keeps its own buffers and only has its child re-based to the
covered length. Vectors whose offsets already start at zero are written
with their own buffers as before, which a test pins by buffer address.
Only vectors whose split keeps every value in its slot are re-based or
trimmed: fixed-width, variable-width and null vectors, and list, large
list, fixed-size list, map and struct vectors whose children all
qualify. A dense union's split assumes each type's offsets are
contiguous, which the format does not require, so a vector with a union,
view, run-end encoded or extension vector at or below it is written as
on `main`.
Three things to be aware of:
- A list whose offsets point past the end of its child, which is
malformed, now fails inside `splitAndTransfer` at unload time instead of
producing a stream that fails to load.
- `ListVector.getFieldBuffers()` now reallocates an offsets buffer that
cannot hold `[0]`, as `BaseVariableWidthVector` has since #989. A
zero-length split leaves a list with the allocator's shared empty buffer
as its offsets, which was written with no bytes after
`getFieldBuffers()` set the writer index of the shared buffer.
- The C Data module's `StructVectorUnloader` has a near-copy of
`appendNodes` that is left as it is. It only feeds
`Data.importIntoVectorSchemaRoot`, where the batch goes straight into a
`VectorLoader`, and re-basing there would add an offsets copy to every
sliced import whether or not it is ever serialized.
Tested with a new `TestVectorUnloadNonZeroStartOffset`, which builds the
vectors a C Data import of a slice produces from raw buffers for
VarChar, LargeVarChar, List, LargeList, Map, a list whose child runs
past its last offset, and a VarChar with a non-zero start nested inside
a list that starts at zero. Each case asserts the node lengths and
buffer sizes written, that the written offsets start at zero, and that
the batch loads back to the right values. It also covers an empty
covered range over a nested list, a list over a dense union whose child
skips an element, and a split that fails on malformed offsets, which
must not leak. `TestListVector` and `TestMapVector` check that the
target of a zero-length split writes one offset.
Closes #1424.
---------
Co-authored-by: JB Onofré <[email protected]>
---
.../org/apache/arrow/vector/VectorUnloader.java | 101 ++++-
.../arrow/vector/complex/LargeListVector.java | 13 +-
.../apache/arrow/vector/complex/ListVector.java | 13 +-
.../apache/arrow/vector/TestLargeListVector.java | 13 +
.../org/apache/arrow/vector/TestListVector.java | 23 ++
.../org/apache/arrow/vector/TestMapVector.java | 32 ++
.../vector/TestVectorUnloadNonZeroStartOffset.java | 433 +++++++++++++++++++++
7 files changed, 625 insertions(+), 3 deletions(-)
diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
index 83b985b70..83077409a 100644
--- a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
+++ b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
@@ -20,11 +20,16 @@ import java.util.ArrayList;
import java.util.List;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.complex.FixedSizeListVector;
+import org.apache.arrow.vector.complex.LargeListVector;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.NonNullableStructVector;
import org.apache.arrow.vector.compression.CompressionCodec;
import org.apache.arrow.vector.compression.CompressionUtil;
import org.apache.arrow.vector.compression.NoCompressionCodec;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.util.TransferPair;
/**
* Helper class that handles converting a {@link VectorSchemaRoot} to a {@link
ArrowRecordBatch}.
@@ -74,6 +79,12 @@ public class VectorUnloader {
/**
* Performs the depth first traversal of the Vectors to create an {@link
ArrowRecordBatch}
* suitable for serialization.
+ *
+ * <p>A variable-width, list or map vector whose offsets start past zero, or
whose child vector
+ * holds elements past its last offset, as one imported through the C Data
interface from a sliced
+ * array does, is written as the values between its first and last offset,
with the offsets
+ * re-based to start at zero. Vectors with a union, view or extension vector
at or below them are
+ * written as they are.
*/
public ArrowRecordBatch getRecordBatch() {
List<ArrowFieldNode> nodes = new ArrayList<>();
@@ -111,6 +122,10 @@ public class VectorUnloader {
List<ArrowBuf> buffers,
List<Long> variadicBufferCounts) {
FieldVector storageVector = ExtensionTypeVector.getStorageVector(vector);
+ if (firstOffset(storageVector) != 0 && splitKeepsValues(vector)) {
+ appendRebased(vector, vector.getValueCount(), nodes, buffers,
variadicBufferCounts);
+ return;
+ }
nodes.add(
new ArrowFieldNode(vector.getValueCount(), includeNullCount ?
vector.getNullCount() : -1));
List<ArrowBuf> fieldBuffers = vector.getFieldBuffers();
@@ -137,8 +152,92 @@ public class VectorUnloader {
buf.getReferenceManager().retain();
buffers.add(codec.compress(vector.getAllocator(), buf));
}
+ long coveredChildLength = coveredChildLength(storageVector);
for (FieldVector child : vector.getChildrenFromFields()) {
- appendNodes(child, nodes, buffers, variadicBufferCounts);
+ if (coveredChildLength >= 0
+ && child.getValueCount() > coveredChildLength
+ && splitKeepsValues(child)) {
+ appendRebased(child, (int) coveredChildLength, nodes, buffers,
variadicBufferCounts);
+ } else {
+ appendNodes(child, nodes, buffers, variadicBufferCounts);
+ }
+ }
+ }
+
+ /** Appends the first {@code length} values of the vector, re-based to start
at offset zero. */
+ private void appendRebased(
+ FieldVector vector,
+ int length,
+ List<ArrowFieldNode> nodes,
+ List<ArrowBuf> buffers,
+ List<Long> variadicBufferCounts) {
+ TransferPair pair = vector.getTransferPair(vector.getField(),
vector.getAllocator());
+ try (FieldVector rebased = (FieldVector) pair.getTo()) {
+ pair.splitAndTransfer(0, length);
+ appendNodes(rebased, nodes, buffers, variadicBufferCounts);
+ }
+ }
+
+ /**
+ * Whether splitting the vector keeps every value in its slot. A dense
union's split does not when
+ * a type's offsets have gaps, so only the vector types listed here are
re-based.
+ */
+ private static boolean splitKeepsValues(FieldVector vector) {
+ if (vector instanceof BaseFixedWidthVector
+ || vector instanceof BaseVariableWidthVector
+ || vector instanceof BaseLargeVariableWidthVector
+ || vector instanceof NullVector) {
+ return true;
+ }
+ if (vector instanceof ListVector
+ || vector instanceof LargeListVector
+ || vector instanceof FixedSizeListVector
+ || vector instanceof NonNullableStructVector) {
+ for (FieldVector child : vector.getChildrenFromFields()) {
+ if (!splitKeepsValues(child)) {
+ return false;
+ }
+ }
+ return true;
+ }
+ return false;
+ }
+
+ /** The first offset of a variable-width, list or map vector, or zero for
any other vector. */
+ private static long firstOffset(FieldVector vector) {
+ if (vector.getValueCount() == 0) {
+ return 0;
+ }
+ if (vector instanceof BaseVariableWidthVector) {
+ return ((BaseVariableWidthVector) vector).getStartOffset(0);
+ }
+ if (vector instanceof BaseLargeVariableWidthVector) {
+ return ((BaseLargeVariableWidthVector) vector).getStartOffset(0);
+ }
+ if (vector instanceof ListVector) {
+ return ((ListVector) vector).getElementStartIndex(0);
+ }
+ if (vector instanceof LargeListVector) {
+ return ((LargeListVector) vector).getElementStartIndex(0);
+ }
+ return 0;
+ }
+
+ /**
+ * The number of child elements a list or map vector's offsets cover, which
is its last offset, or
+ * -1 for a vector whose child is not addressed by offsets.
+ */
+ private static long coveredChildLength(FieldVector vector) {
+ int valueCount = vector.getValueCount();
+ if (valueCount == 0) {
+ return -1;
+ }
+ if (vector instanceof ListVector) {
+ return ((ListVector) vector).getElementEndIndex(valueCount - 1);
+ }
+ if (vector instanceof LargeListVector) {
+ return ((LargeListVector) vector).getElementEndIndex(valueCount - 1);
}
+ return -1;
}
}
diff --git
a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
index 9368b5ce3..f3fc6faaf 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
@@ -316,7 +316,18 @@ public class LargeListVector extends BaseValueVector
// Both are set to 0 means 0 bytes are written to the IPC stream which
will crash IPC readers
// in other libraries. According to Arrow spec, we should still output the
offset buffer which
// is [0].
- offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
+ final long requiredOffsetBufferSize = (long) (valueCount + 1) *
OFFSET_WIDTH;
+ if (offsetBuffer.capacity() < requiredOffsetBufferSize) {
+ // An empty vector may still hold the allocator's shared empty buffer
+ ArrowBuf newOffsetBuffer = allocator.buffer(requiredOffsetBufferSize);
+ newOffsetBuffer.setZero(0, newOffsetBuffer.capacity());
+ if (offsetBuffer.capacity() > 0) {
+ newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity());
+ }
+ offsetBuffer.getReferenceManager().release();
+ offsetBuffer = newOffsetBuffer;
+ }
+ offsetBuffer.writerIndex(requiredOffsetBufferSize);
}
/**
diff --git
a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
index 6c3993df6..f2dee0915 100644
--- a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
+++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
@@ -274,7 +274,18 @@ public class ListVector extends BaseRepeatedValueVector
// Both are set to 0 means 0 bytes are written to the IPC stream which
will crash IPC readers
// in other libraries. According to Arrow spec, we should still output the
offset buffer which
// is [0].
- offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
+ final long requiredOffsetBufferSize = (long) (valueCount + 1) *
OFFSET_WIDTH;
+ if (offsetBuffer.capacity() < requiredOffsetBufferSize) {
+ // An empty vector may still hold the allocator's shared empty buffer
+ ArrowBuf newOffsetBuffer = allocator.buffer(requiredOffsetBufferSize);
+ newOffsetBuffer.setZero(0, newOffsetBuffer.capacity());
+ if (offsetBuffer.capacity() > 0) {
+ newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity());
+ }
+ offsetBuffer.getReferenceManager().release();
+ offsetBuffer = newOffsetBuffer;
+ }
+ offsetBuffer.writerIndex(requiredOffsetBufferSize);
}
/**
diff --git
a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
index b6823fbf3..9f92a23a7 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
@@ -1216,4 +1216,17 @@ public class TestLargeListVector {
}
}
}
+
+ @Test
+ public void testGetFieldBuffersOfEmptyVector() {
+ // An empty vector holds the allocator's shared empty buffer as its
offsets until it allocates
+ try (LargeListVector vector = LargeListVector.empty("list", allocator)) {
+ vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ ArrowBuf offsets = vector.getFieldBuffers().get(1);
+ assertTrue(offsets.capacity() >= LargeListVector.OFFSET_WIDTH);
+ assertEquals(LargeListVector.OFFSET_WIDTH, offsets.readableBytes());
+ assertEquals(0L, offsets.getLong(0));
+ assertEquals(0, allocator.getEmpty().writerIndex());
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
index 0c90b32ab..548e05ace 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
@@ -1406,4 +1406,27 @@ public class TestListVector {
}
writer.endList();
}
+
+ @Test
+ public void testSplitAndTransferZeroLength() {
+ try (ListVector listVector = ListVector.empty("sourceVector", allocator)) {
+ listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ UnionListWriter listWriter = listVector.getWriter();
+ listWriter.allocate();
+ listWriter.setPosition(0);
+ listWriter.startList();
+ listWriter.integer().writeInt(1);
+ listWriter.endList();
+ listVector.setValueCount(1);
+
+ TransferPair pair = listVector.getTransferPair(allocator);
+ pair.splitAndTransfer(0, 0);
+ try (ListVector to = (ListVector) pair.getTo()) {
+ ArrowBuf offsets = to.getFieldBuffers().get(1);
+ assertTrue(offsets.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH);
+ assertEquals(0, offsets.getInt(0));
+ assertEquals(0, allocator.getEmpty().writerIndex());
+ }
+ }
+ }
}
diff --git a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
index 2f520f388..99bb178dd 100644
--- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
+++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
@@ -31,6 +31,7 @@ import java.util.Map;
import java.util.UUID;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.complex.impl.UnionMapReader;
@@ -1662,4 +1663,35 @@ public class TestMapVector {
assertEquals(value2c, actualValue2c);
}
}
+
+ @Test
+ public void testSplitAndTransferZeroLength() {
+ try (MapVector mapVector = MapVector.empty("sourceVector", allocator,
false)) {
+ FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null,
null);
+ AddOrGetResult<StructVector> addResult = mapVector.addOrGetVector(type);
+ FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(),
null, null);
+ FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType());
+ addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType,
BigIntVector.class);
+ addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType,
Float8Vector.class);
+ UnionMapWriter mapWriter = mapVector.getWriter();
+ mapWriter.allocate();
+ mapWriter.setPosition(0);
+ mapWriter.startMap();
+ mapWriter.startEntry();
+ mapWriter.key().bigInt().writeBigInt(10);
+ mapWriter.value().float8().writeFloat8(1.0);
+ mapWriter.endEntry();
+ mapWriter.endMap();
+ mapVector.setValueCount(1);
+
+ TransferPair pair = mapVector.getTransferPair(allocator);
+ pair.splitAndTransfer(0, 0);
+ try (MapVector to = (MapVector) pair.getTo()) {
+ ArrowBuf offsets = to.getFieldBuffers().get(1);
+ assertTrue(offsets.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH);
+ assertEquals(0, offsets.getInt(0));
+ assertEquals(0, allocator.getEmpty().writerIndex());
+ }
+ }
+ }
}
diff --git
a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java
b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java
new file mode 100644
index 000000000..e4ba77594
--- /dev/null
+++
b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadNonZeroStartOffset.java
@@ -0,0 +1,433 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.complex.DenseUnionVector;
+import org.apache.arrow.vector.complex.LargeListVector;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.ipc.message.ArrowBuffer;
+import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.types.Types.MinorType;
+import org.apache.arrow.vector.types.UnionMode;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * A vector imported through the C Data interface from a sliced array keeps
the offsets of the
+ * slice, which start past zero, and the bytes or child elements before them.
{@link VectorUnloader}
+ * must write only what the offsets cover.
+ */
+public class TestVectorUnloadNonZeroStartOffset {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ public void init() {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void terminate() {
+ allocator.close();
+ }
+
+ private ArrowBuf ints(int... values) {
+ ArrowBuf buf = allocator.buffer(4L * values.length);
+ for (int i = 0; i < values.length; i++) {
+ buf.setInt(4L * i, values[i]);
+ }
+ buf.writerIndex(4L * values.length);
+ return buf;
+ }
+
+ private ArrowBuf longs(long... values) {
+ ArrowBuf buf = allocator.buffer(8L * values.length);
+ for (int i = 0; i < values.length; i++) {
+ buf.setLong(8L * i, values[i]);
+ }
+ buf.writerIndex(8L * values.length);
+ return buf;
+ }
+
+ private ArrowBuf bytes(String values) {
+ byte[] data = values.getBytes(StandardCharsets.UTF_8);
+ ArrowBuf buf = allocator.buffer(data.length);
+ buf.setBytes(0, data);
+ buf.writerIndex(data.length);
+ return buf;
+ }
+
+ private ArrowBuf allValid() {
+ return longs(-1L);
+ }
+
+ private ArrowBuf typeIds(int... values) {
+ ArrowBuf buf = allocator.buffer(values.length);
+ for (int i = 0; i < values.length; i++) {
+ buf.setByte(i, values[i]);
+ }
+ buf.writerIndex(values.length);
+ return buf;
+ }
+
+ private static void load(FieldVector vector, int length, ArrowBuf...
buffers) {
+ vector.loadFieldBuffers(new ArrowFieldNode(length, 0),
Arrays.asList(buffers));
+ for (ArrowBuf buf : buffers) {
+ buf.close();
+ }
+ }
+
+ /** Rows 2..3 of ["aa", "bb", "cc", "dd"]: the data buffer still holds
"aabb" in front. */
+ private VarCharVector varCharSlice() {
+ VarCharVector vector = new VarCharVector("s", allocator);
+ load(vector, 2, allValid(), ints(4, 6, 8), bytes("aabbccdd"));
+ return vector;
+ }
+
+ private LargeVarCharVector largeVarCharSlice() {
+ LargeVarCharVector vector = new LargeVarCharVector("s", allocator);
+ load(vector, 2, allValid(), longs(4, 6, 8), bytes("aabbccdd"));
+ return vector;
+ }
+
+ private void loadEightInts(FieldVector child) {
+ load(child, 8, allValid(), ints(0, 1, 2, 3, 4, 5, 6, 7));
+ }
+
+ /** Two rows of [[0, 1], [2, 3], [4, 5], [6, 7]] whose child still holds all
eight elements. */
+ private ListVector listSlice(int... offsets) {
+ ListVector vector = ListVector.empty("l", allocator);
+ vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ load(vector, 2, allValid(), ints(offsets));
+ loadEightInts(vector.getDataVector());
+ return vector;
+ }
+
+ private LargeListVector largeListSlice(long... offsets) {
+ LargeListVector vector = LargeListVector.empty("l", allocator);
+ vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ load(vector, 2, allValid(), longs(offsets));
+ loadEightInts(vector.getDataVector());
+ return vector;
+ }
+
+ /** Rows 1..2 of [{0: 10, 1: 11}, {2: 12, 3: 13}, {4: 14, 5: 15}, {6: 16, 7:
17}]. */
+ private MapVector mapSlice() {
+ MapVector vector = MapVector.empty("m", allocator, false);
+ Field key = new Field("key",
FieldType.notNullable(MinorType.INT.getType()), null);
+ Field value = new Field("value",
FieldType.nullable(MinorType.INT.getType()), null);
+ Field entries =
+ new Field(
+ MapVector.DATA_VECTOR_NAME,
+ FieldType.notNullable(ArrowType.Struct.INSTANCE),
+ Arrays.asList(key, value));
+ vector.initializeChildrenFromFields(Collections.singletonList(entries));
+ load(vector, 2, allValid(), ints(2, 4, 6));
+ StructVector struct = (StructVector) vector.getDataVector();
+ load(struct, 8, allValid());
+ loadEightInts(struct.getChild("key"));
+ load(struct.getChild("value"), 8, allValid(), ints(10, 11, 12, 13, 14, 15,
16, 17));
+ return vector;
+ }
+
+ /**
+ * The first slice of [["aa"], ["bb"], ["cc"]]: list offsets start at zero,
its child's do not.
+ */
+ private ListVector listOfVarCharSlice() {
+ ListVector vector = ListVector.empty("l", allocator);
+ vector.addOrGetVector(FieldType.nullable(MinorType.VARCHAR.getType()));
+ load(vector, 2, allValid(), ints(0, 1, 2));
+ load(vector.getDataVector(), 2, allValid(), ints(2, 4, 6),
bytes("aabbcc"));
+ return vector;
+ }
+
+ /** Two empty rows over an inner list that still holds [[1], [2]]. */
+ private ListVector emptyRowsOverNestedList() {
+ ListVector vector = ListVector.empty("l", allocator);
+ vector.addOrGetVector(FieldType.nullable(ArrowType.List.INSTANCE));
+ ListVector inner = (ListVector) vector.getDataVector();
+ inner.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ load(vector, 2, allValid(), ints(0, 0, 0));
+ load(inner, 2, allValid(), ints(0, 1, 2));
+ load(inner.getDataVector(), 2, allValid(), ints(1, 2));
+ return vector;
+ }
+
+ /** Two empty rows over an inner list that still holds two large lists,
[[1]] and [[2]]. */
+ private ListVector emptyRowsOverNestedLargeList() {
+ ListVector vector = ListVector.empty("l", allocator);
+ vector.addOrGetVector(FieldType.nullable(ArrowType.List.INSTANCE));
+ ListVector inner = (ListVector) vector.getDataVector();
+ inner.addOrGetVector(FieldType.nullable(ArrowType.LargeList.INSTANCE));
+ LargeListVector largeList = (LargeListVector) inner.getDataVector();
+ largeList.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
+ load(vector, 2, allValid(), ints(0, 0, 0));
+ load(inner, 2, allValid(), ints(0, 1, 2));
+ load(largeList, 2, allValid(), longs(0, 1, 2));
+ load(largeList.getDataVector(), 2, allValid(), ints(1, 2));
+ return vector;
+ }
+
+ /**
+ * One row over a dense union whose int child skips an element: union slots
0..3 hold child
+ * elements 0, 1, 3 and 4, that is 10, 11, 13 and 14.
+ */
+ private ListVector listOfDenseUnion(int... offsets) {
+ Field union =
+ new Field(
+ "u",
+ FieldType.nullable(new ArrowType.Union(UnionMode.Dense, new int[]
{0})),
+ Collections.singletonList(Field.nullable("i",
MinorType.INT.getType())));
+ Field list =
+ new Field(
+ "l", FieldType.nullable(ArrowType.List.INSTANCE),
Collections.singletonList(union));
+ ListVector vector = (ListVector) list.createVector(allocator);
+ load(vector, 1, allValid(), ints(offsets));
+ DenseUnionVector unionVector = (DenseUnionVector) vector.getDataVector();
+ load(unionVector, 4, typeIds(0, 0, 0, 0), ints(0, 1, 3, 4));
+ load(unionVector.getChildrenFromFields().get(0), 5, allValid(), ints(10,
11, 12, 13, 14));
+ return vector;
+ }
+
+ private static List<Long> sizes(ArrowRecordBatch batch) {
+ return
batch.getBuffersLayout().stream().map(ArrowBuffer::getSize).collect(Collectors.toList());
+ }
+
+ private static List<Integer> lengths(ArrowRecordBatch batch) {
+ return
batch.getNodes().stream().map(ArrowFieldNode::getLength).collect(Collectors.toList());
+ }
+
+ private VectorSchemaRoot reload(VectorSchemaRoot root, ArrowRecordBatch
batch) {
+ VectorSchemaRoot loaded = VectorSchemaRoot.create(root.getSchema(),
allocator);
+ new VectorLoader(loaded).load(batch);
+ return loaded;
+ }
+
+ @Test
+ public void varCharIsWrittenFromItsFirstOffset() {
+ try (VarCharVector vector = varCharSlice();
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(1L, 12L, 4L), sizes(batch));
+ assertEquals(0, batch.getBuffers().get(1).getInt(0));
+ assertEquals(4, batch.getBuffers().get(1).getInt(8));
+ VarCharVector read = (VarCharVector) loaded.getVector(0);
+ assertEquals("cc", read.getObject(0).toString());
+ assertEquals("dd", read.getObject(1).toString());
+ }
+ }
+
+ @Test
+ public void largeVarCharIsWrittenFromItsFirstOffset() {
+ try (LargeVarCharVector vector = largeVarCharSlice();
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(1L, 24L, 4L), sizes(batch));
+ assertEquals(0L, batch.getBuffers().get(1).getLong(0));
+ LargeVarCharVector read = (LargeVarCharVector) loaded.getVector(0);
+ assertEquals("cc", read.getObject(0).toString());
+ assertEquals("dd", read.getObject(1).toString());
+ }
+ }
+
+ @Test
+ public void listChildIsWrittenBetweenItsOffsets() {
+ try (ListVector vector = listSlice(2, 4, 6);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 4), lengths(batch));
+ assertEquals(Arrays.asList(1L, 12L, 1L, 16L), sizes(batch));
+ assertEquals(0, batch.getBuffers().get(1).getInt(0));
+ ListVector read = (ListVector) loaded.getVector(0);
+ assertEquals(Arrays.asList(2, 3), read.getObject(0));
+ assertEquals(Arrays.asList(4, 5), read.getObject(1));
+ }
+ }
+
+ @Test
+ public void largeListChildIsWrittenBetweenItsOffsets() {
+ try (LargeListVector vector = largeListSlice(2, 4, 6);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 4), lengths(batch));
+ assertEquals(Arrays.asList(1L, 24L, 1L, 16L), sizes(batch));
+ assertEquals(0L, batch.getBuffers().get(1).getLong(0));
+ LargeListVector read = (LargeListVector) loaded.getVector(0);
+ assertEquals(Arrays.asList(2, 3), read.getObject(0));
+ assertEquals(Arrays.asList(4, 5), read.getObject(1));
+ }
+ }
+
+ @Test
+ public void mapEntriesAreWrittenBetweenItsOffsets() {
+ try (MapVector vector = mapSlice();
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 4, 4, 4), lengths(batch));
+ assertEquals(Arrays.asList(1L, 12L, 1L, 1L, 16L, 1L, 16L), sizes(batch));
+ MapVector read = (MapVector) loaded.getVector(0);
+ assertEquals(vector.getObject(0), read.getObject(0));
+ assertEquals(vector.getObject(1), read.getObject(1));
+ StructVector entries = (StructVector) read.getDataVector();
+ assertEquals(2, ((IntVector) entries.getChild("key")).get(0));
+ assertEquals(15, ((IntVector) entries.getChild("value")).get(3));
+ }
+ }
+
+ @Test
+ public void listChildIsWrittenUpToItsLastOffset() {
+ try (ListVector vector = listSlice(0, 2, 4);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 4), lengths(batch));
+ assertEquals(Arrays.asList(1L, 12L, 1L, 16L), sizes(batch));
+ // the list's own offsets start at zero, so they are written as they are
+ assertEquals(
+ vector.getOffsetBuffer().memoryAddress(),
batch.getBuffers().get(1).memoryAddress());
+ ListVector read = (ListVector) loaded.getVector(0);
+ assertEquals(Arrays.asList(0, 1), read.getObject(0));
+ assertEquals(Arrays.asList(2, 3), read.getObject(1));
+ }
+ }
+
+ @Test
+ public void nestedVarCharIsWrittenFromItsFirstOffset() {
+ try (ListVector vector = listOfVarCharSlice();
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 2), lengths(batch));
+ assertEquals(Arrays.asList(1L, 12L, 1L, 12L, 4L), sizes(batch));
+ assertEquals(0, batch.getBuffers().get(3).getInt(0));
+ ListVector read = (ListVector) loaded.getVector(0);
+ assertEquals(vector.getObject(0), read.getObject(0));
+ assertEquals(vector.getObject(1), read.getObject(1));
+ assertEquals("bb", ((VarCharVector)
read.getDataVector()).getObject(0).toString());
+ }
+ }
+
+ @Test
+ public void vectorWithZeroStartOffsetIsWrittenWithItsOwnBuffers() {
+ try (VarCharVector vector = new VarCharVector("s", allocator);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+ load(vector, 2, allValid(), ints(0, 2, 4), bytes("aabb"));
+ try (ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch())
{
+ assertEquals(Arrays.asList(1L, 12L, 4L), sizes(batch));
+ assertEquals(
+ vector.getOffsetBuffer().memoryAddress(),
batch.getBuffers().get(1).memoryAddress());
+ assertEquals(
+ vector.getDataBuffer().memoryAddress(),
batch.getBuffers().get(2).memoryAddress());
+ }
+ }
+ }
+
+ @Test
+ public void emptyCoveredRangeOverNestedListKeepsItsFirstOffset() {
+ try (ListVector vector = emptyRowsOverNestedList();
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 0, 0), lengths(batch));
+ assertEquals(Arrays.asList(1L, 12L, 0L, 4L, 0L, 0L), sizes(batch));
+ assertEquals(0, batch.getBuffers().get(3).getInt(0));
+ ListVector read = (ListVector) loaded.getVector(0);
+ assertEquals(Collections.emptyList(), read.getObject(0));
+ assertEquals(Collections.emptyList(), read.getObject(1));
+ }
+ }
+
+ @Test
+ public void emptyCoveredRangeOverNestedLargeListKeepsItsFirstOffset() {
+ try (ListVector vector = emptyRowsOverNestedLargeList();
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(2, 0, 0, 0), lengths(batch));
+ assertEquals(Arrays.asList(1L, 12L, 0L, 4L, 0L, 8L, 0L, 0L),
sizes(batch));
+ assertEquals(0L, batch.getBuffers().get(5).getLong(0));
+ ListVector read = (ListVector) loaded.getVector(0);
+ assertEquals(Collections.emptyList(), read.getObject(0));
+ assertEquals(Collections.emptyList(), read.getObject(1));
+ }
+ }
+
+ @Test
+ public void denseUnionChildIsWrittenUntrimmed() {
+ try (ListVector vector = listOfDenseUnion(0, 3);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(1, 4, 5), lengths(batch));
+ assertEquals(Arrays.asList(10, 11, 13), ((ListVector)
loaded.getVector(0)).getObject(0));
+ }
+ }
+
+ @Test
+ public void listOverDenseUnionIsWrittenWithItsOwnOffsets() {
+ try (ListVector vector = listOfDenseUnion(1, 3);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector);
+ ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch();
+ VectorSchemaRoot loaded = reload(root, batch)) {
+ assertEquals(Arrays.asList(1, 4, 5), lengths(batch));
+ assertEquals(Arrays.asList(11, 13), ((ListVector)
loaded.getVector(0)).getObject(0));
+ }
+ }
+
+ @Test
+ public void failedSplitDoesNotLeak() {
+ try (ListVector vector = listSlice(2, 4, 9);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+ assertThrows(IllegalArgumentException.class, () -> new
VectorUnloader(root).getRecordBatch());
+ }
+ assertEquals(0, allocator.getAllocatedMemory());
+ }
+
+ @Test
+ public void failedSplitReleasesTheBuffersOfEarlierVectors() {
+ try (IntVector first = new IntVector("i", allocator);
+ ListVector malformed = listSlice(2, 4, 9);
+ VectorSchemaRoot root = VectorSchemaRoot.of(first, malformed)) {
+ load(first, 2, allValid(), ints(1, 2));
+ assertThrows(IllegalArgumentException.class, () -> new
VectorUnloader(root).getRecordBatch());
+ }
+ assertEquals(0, allocator.getAllocatedMemory());
+ }
+}
