Skip to content

GH-1194: Preserve empty list offset buffers #1229

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
prashanthbdremio wants to merge 1 commit into
base: main
Choose a base branch
from
Open

GH-1194: Preserve empty list offset buffers #1229

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 38 additions & 17 deletions vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,18 +305,33 @@ public void exportCDataBuffers(List<ArrowBuf> buffers, ArrowBuf buffersPtr, long

/** Set the reader and writer indexes for the inner buffers. */
private void setReaderAndWriterIndex() {
final long requiredOffsetBufferCapacity = (long) (valueCount + 1) * OFFSET_WIDTH;
validityBuffer.readerIndex(0);
offsetBuffer.readerIndex(0);
if (valueCount == 0) {
validityBuffer.writerIndex(0);
ensureEmptyOffsetBufferCapacity(requiredOffsetBufferCapacity);
} else {
validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
}
// IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.
// 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);
offsetBuffer.writerIndex(requiredOffsetBufferCapacity);
}

private void ensureEmptyOffsetBufferCapacity(long requiredCapacity) {
if (offsetBuffer.capacity() >= requiredCapacity) {
return;
}
long previousOffsetAllocationSizeInBytes = offsetAllocationSizeInBytes;
ArrowBuf oldOffsetBuffer = offsetBuffer;
offsetBuffer = allocateOffsetBuffer(requiredCapacity);
offsetBuffer.setBytes(
0, oldOffsetBuffer, 0, Math.min(oldOffsetBuffer.capacity(), requiredCapacity));
offsetAllocationSizeInBytes = previousOffsetAllocationSizeInBytes;
oldOffsetBuffer.getReferenceManager().release();
}

/**
Expand Down Expand Up @@ -674,24 +689,30 @@ public void splitAndTransfer(int startIndex, int length) {
startIndex,
length,
valueCount);
final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH);
final long sliceLength =
offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint;
to.clear();
to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH);
/* splitAndTransfer offset buffer */
for (int i = 0; i < length + 1; i++) {
final long relativeOffset =
offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint;
to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset);
if (length > 0) {
final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH);
final long sliceLength =
offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint;
to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH);
/* splitAndTransfer offset buffer */
for (int i = 0; i < length + 1; i++) {
final long relativeOffset =
offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint;
to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset);
}
/* splitAndTransfer validity buffer */
splitAndTransferValidityBuffer(startIndex, length, to);
/* splitAndTransfer data buffer */
dataTransferPair.splitAndTransfer(
checkedCastToInt(startPoint), checkedCastToInt(sliceLength));
to.lastSet = length - 1;
to.setValueCount(length);
} else {
to.ensureEmptyOffsetBufferCapacity(OFFSET_WIDTH);
dataTransferPair.splitAndTransfer(0, 0);
to.setValueCount(0);
}
/* splitAndTransfer validity buffer */
splitAndTransferValidityBuffer(startIndex, length, to);
/* splitAndTransfer data buffer */
dataTransferPair.splitAndTransfer(
checkedCastToInt(startPoint), checkedCastToInt(sliceLength));
to.lastSet = length - 1;
to.setValueCount(length);
}

@Override
Expand Down
21 changes: 20 additions & 1 deletion vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,18 +263,33 @@ public void exportCDataBuffers(List<ArrowBuf> buffers, ArrowBuf buffersPtr, long

/** Set the reader and writer indexes for the inner buffers. */
private void setReaderAndWriterIndex() {
final long requiredOffsetBufferCapacity = (long) (valueCount + 1) * OFFSET_WIDTH;
validityBuffer.readerIndex(0);
offsetBuffer.readerIndex(0);
if (valueCount == 0) {
validityBuffer.writerIndex(0);
ensureEmptyOffsetBufferCapacity(requiredOffsetBufferCapacity);
} else {
validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
}
// IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.
// 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);
offsetBuffer.writerIndex(requiredOffsetBufferCapacity);
}

private void ensureEmptyOffsetBufferCapacity(long requiredCapacity) {
if (offsetBuffer.capacity() >= requiredCapacity) {
return;
}
long previousOffsetAllocationSizeInBytes = offsetAllocationSizeInBytes;
ArrowBuf oldOffsetBuffer = offsetBuffer;
offsetBuffer = allocateOffsetBuffer(requiredCapacity);
offsetBuffer.setBytes(
0, oldOffsetBuffer, 0, Math.min(oldOffsetBuffer.capacity(), requiredCapacity));
offsetAllocationSizeInBytes = previousOffsetAllocationSizeInBytes;
oldOffsetBuffer.getReferenceManager().release();
}

/**
Expand Down Expand Up @@ -572,6 +587,10 @@ public void splitAndTransfer(int startIndex, int length) {
dataTransferPair.splitAndTransfer(startPoint, sliceLength);
to.lastSet = length - 1;
to.setValueCount(length);
} else {
to.ensureEmptyOffsetBufferCapacity(OFFSET_WIDTH);
dataTransferPair.splitAndTransfer(0, 0);
to.setValueCount(0);
}
}

Expand Down
85 changes: 75 additions & 10 deletions vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1102,24 +1102,89 @@ public void testCopyValueSafeForExtensionType() throws Exception {

@Test
public void testEmptyLargeListOffsetBuffer() {
// Test that LargeListVector has correct readableBytes after allocation.
// According to Arrow spec, offset buffer must have N+1 entries.
// Even when N=0, it should contain [0].
try (LargeListVector list = LargeListVector.empty("list", allocator)) {
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
list.allocateNew();
list.setValueCount(0);

List<ArrowBuf> buffers = list.getFieldBuffers();
assertTrue(
buffers.get(1).readableBytes() >= LargeListVector.OFFSET_WIDTH,
"Offset buffer should have at least "
+ LargeListVector.OFFSET_WIDTH
+ " bytes for offset[0]");
assertEquals(0L, list.getOffsetBuffer().getLong(0));
assertEmptyLargeListOffsetBuffer(list);
}
}

@Test
public void testUnallocatedEmptyLargeListOffsetBuffer() {
try (LargeListVector list = LargeListVector.empty("list", allocator)) {
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
list.setValueCount(0);

assertEmptyLargeListOffsetBuffer(list);
}
}

@Test
public void testSplitAndTransferEmptyLargeListOffsetBuffer() {
try (LargeListVector source = LargeListVector.empty("source", allocator);
LargeListVector target = LargeListVector.empty("target", allocator)) {
source.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
target.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
source.allocateNew();
source.setValueCount(0);

TransferPair transferPair = source.makeTransferPair(target);
transferPair.splitAndTransfer(0, 0);

assertEmptyLargeListOffsetBuffer(target);
}
}

@Test
public void testSplitAndTransferEmptyLargeListAllocatesOffsetBuffer() {
try (LargeListVector fromVector = LargeListVector.empty("fromVector", allocator);
LargeListVector toVector = LargeListVector.empty("toVector", allocator)) {
fromVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
fromVector.allocateNew();
fromVector.setValueCount(0);

TransferPair transferPair = fromVector.makeTransferPair(toVector);
transferPair.splitAndTransfer(0, 0);

assertAllocatedEmptyLargeListOffsetBuffer(toVector);
}
}

@Test
public void testSplitAndTransferEmptyNestedLargeListAllocatesOffsetBuffers() {
try (LargeListVector fromVector = LargeListVector.empty("fromVector", allocator);
LargeListVector toVector = LargeListVector.empty("toVector", allocator)) {
fromVector.addOrGetVector(FieldType.nullable(MinorType.LARGELIST.getType()));
LargeListVector childVector = (LargeListVector) fromVector.getDataVector();
childVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
fromVector.allocateNew();
fromVector.setValueCount(0);

TransferPair transferPair = fromVector.makeTransferPair(toVector);
transferPair.splitAndTransfer(0, 0);

assertAllocatedEmptyLargeListOffsetBuffer(toVector);
assertAllocatedEmptyLargeListOffsetBuffer((LargeListVector) toVector.getDataVector());
}
}

private ArrowBuf assertEmptyLargeListOffsetBuffer(LargeListVector list) {
List<ArrowBuf> buffers = list.getFieldBuffers();
ArrowBuf offsetBuffer = buffers.get(1);
assertEquals(LargeListVector.OFFSET_WIDTH, offsetBuffer.readableBytes());
assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH);
assertEquals(0L, offsetBuffer.getLong(0));
return offsetBuffer;
}

private void assertAllocatedEmptyLargeListOffsetBuffer(LargeListVector list) {
ArrowBuf offsetBuffer = list.getOffsetBuffer();
assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH);
assertEquals(0L, offsetBuffer.getLong(0));
}

private void writeIntValues(UnionLargeListWriter writer, int[] values) {
writer.startList();
for (int v : values) {
Expand Down
85 changes: 75 additions & 10 deletions vector/src/test/java/org/apache/arrow/vector/TestListVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1381,24 +1381,89 @@ public void testCopyValueSafeForExtensionType() throws Exception {

@Test
public void testEmptyListOffsetBuffer() {
// Test that ListVector has correct readableBytes after allocation.
// According to Arrow spec, offset buffer must have N+1 entries.
// Even when N=0, it should contain [0].
try (ListVector list = ListVector.empty("list", allocator)) {
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
list.allocateNew();
list.setValueCount(0);

List<ArrowBuf> buffers = list.getFieldBuffers();
assertTrue(
buffers.get(1).readableBytes() >= BaseRepeatedValueVector.OFFSET_WIDTH,
"Offset buffer should have at least "
+ BaseRepeatedValueVector.OFFSET_WIDTH
+ " bytes for offset[0]");
assertEquals(0, list.getOffsetBuffer().getInt(0));
assertEmptyListOffsetBuffer(list);
}
}

@Test
public void testUnallocatedEmptyListOffsetBuffer() {
try (ListVector list = ListVector.empty("list", allocator)) {
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
list.setValueCount(0);

assertEmptyListOffsetBuffer(list);
}
}

@Test
public void testSplitAndTransferEmptyListOffsetBuffer() {
try (ListVector source = ListVector.empty("source", allocator);
ListVector target = ListVector.empty("target", allocator)) {
source.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
target.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
source.allocateNew();
source.setValueCount(0);

TransferPair transferPair = source.makeTransferPair(target);
transferPair.splitAndTransfer(0, 0);

assertEmptyListOffsetBuffer(target);
}
}

@Test
public void testSplitAndTransferEmptyListAllocatesOffsetBuffer() {
try (ListVector fromVector = ListVector.empty("fromVector", allocator);
ListVector toVector = ListVector.empty("toVector", allocator)) {
fromVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
fromVector.allocateNew();
fromVector.setValueCount(0);

TransferPair transferPair = fromVector.makeTransferPair(toVector);
transferPair.splitAndTransfer(0, 0);

assertAllocatedEmptyListOffsetBuffer(toVector);
}
}

@Test
public void testSplitAndTransferEmptyNestedListAllocatesOffsetBuffers() {
try (ListVector fromVector = ListVector.empty("fromVector", allocator);
ListVector toVector = ListVector.empty("toVector", allocator)) {
fromVector.addOrGetVector(FieldType.nullable(MinorType.LIST.getType()));
ListVector childVector = (ListVector) fromVector.getDataVector();
childVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
fromVector.allocateNew();
fromVector.setValueCount(0);

TransferPair transferPair = fromVector.makeTransferPair(toVector);
transferPair.splitAndTransfer(0, 0);

assertAllocatedEmptyListOffsetBuffer(toVector);
assertAllocatedEmptyListOffsetBuffer((ListVector) toVector.getDataVector());
}
}

private ArrowBuf assertEmptyListOffsetBuffer(ListVector list) {
List<ArrowBuf> buffers = list.getFieldBuffers();
ArrowBuf offsetBuffer = buffers.get(1);
assertEquals(BaseRepeatedValueVector.OFFSET_WIDTH, offsetBuffer.readableBytes());
assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH);
assertEquals(0, offsetBuffer.getInt(0));
return offsetBuffer;
}

private void assertAllocatedEmptyListOffsetBuffer(ListVector list) {
ArrowBuf offsetBuffer = list.getOffsetBuffer();
assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH);
assertEquals(0, offsetBuffer.getInt(0));
}

private void writeIntValues(UnionListWriter writer, int[] values) {
writer.startList();
for (int v : values) {
Expand Down
Loading
Loading