Skip to content

Commit ee21e3b

Browse files
prashanthbdremioprashanthbdremio
authored
MINOR: Preserve empty list offset buffer (#30)
Preserve the Arrow-required empty list offset buffer entry while avoiding malformed Netty buffer state. ListVector and LargeListVector now materialize a one-entry offset buffer when an empty vector still needs to expose offset[0], then set the writer index from (valueCount + 1) * OFFSET_WIDTH.
1 parent 2334954 commit ee21e3b

5 files changed

Lines changed: 267 additions & 23 deletions

File tree

‎vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java‎

Lines changed: 37 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -305,10 +305,12 @@ public void exportCDataBuffers(List<ArrowBuf> buffers, ArrowBuf buffersPtr, long
305305

306306
/** Set the reader and writer indexes for the inner buffers. */
307307
private void setReaderAndWriterIndex() {
308+
final long requiredOffsetBufferCapacity = (long) (valueCount + 1) * OFFSET_WIDTH;
308309
validityBuffer.readerIndex(0);
309310
offsetBuffer.readerIndex(0);
310311
if (valueCount == 0) {
311312
validityBuffer.writerIndex(0);
313+
ensureEmptyOffsetBufferCapacity(requiredOffsetBufferCapacity);
312314
} else {
313315
validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
314316
}
@@ -319,6 +321,19 @@ private void setReaderAndWriterIndex() {
319321
offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
320322
}
321323

324+
private void ensureEmptyOffsetBufferCapacity(long requiredCapacity) {
325+
if (offsetBuffer.capacity() >= requiredCapacity) {
326+
return;
327+
}
328+
long previousOffsetAllocationSizeInBytes = offsetAllocationSizeInBytes;
329+
ArrowBuf oldOffsetBuffer = offsetBuffer;
330+
offsetBuffer = allocateOffsetBuffer(requiredCapacity);
331+
offsetBuffer.setBytes(
332+
0, oldOffsetBuffer, 0, Math.min(oldOffsetBuffer.capacity(), requiredCapacity));
333+
offsetAllocationSizeInBytes = previousOffsetAllocationSizeInBytes;
334+
oldOffsetBuffer.getReferenceManager().release();
335+
}
336+
322337
/**
323338
* Get the inner vectors.
324339
*
@@ -674,24 +689,30 @@ public void splitAndTransfer(int startIndex, int length) {
674689
startIndex,
675690
length,
676691
valueCount);
677-
final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH);
678-
final long sliceLength =
679-
offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint;
680692
to.clear();
681-
to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH);
682-
/* splitAndTransfer offset buffer */
683-
for (int i = 0; i < length + 1; i++) {
684-
final long relativeOffset =
685-
offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint;
686-
to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset);
693+
if (length > 0) {
694+
final long startPoint = offsetBuffer.getLong((long) startIndex * OFFSET_WIDTH);
695+
final long sliceLength =
696+
offsetBuffer.getLong((long) (startIndex + length) * OFFSET_WIDTH) - startPoint;
697+
to.offsetBuffer = to.allocateOffsetBuffer((length + 1) * OFFSET_WIDTH);
698+
/* splitAndTransfer offset buffer */
699+
for (int i = 0; i < length + 1; i++) {
700+
final long relativeOffset =
701+
offsetBuffer.getLong((long) (startIndex + i) * OFFSET_WIDTH) - startPoint;
702+
to.offsetBuffer.setLong((long) i * OFFSET_WIDTH, relativeOffset);
703+
}
704+
/* splitAndTransfer validity buffer */
705+
splitAndTransferValidityBuffer(startIndex, length, to);
706+
/* splitAndTransfer data buffer */
707+
dataTransferPair.splitAndTransfer(
708+
checkedCastToInt(startPoint), checkedCastToInt(sliceLength));
709+
to.lastSet = length - 1;
710+
to.setValueCount(length);
711+
} else {
712+
to.ensureEmptyOffsetBufferCapacity(OFFSET_WIDTH);
713+
dataTransferPair.splitAndTransfer(0, 0);
714+
to.setValueCount(0);
687715
}
688-
/* splitAndTransfer validity buffer */
689-
splitAndTransferValidityBuffer(startIndex, length, to);
690-
/* splitAndTransfer data buffer */
691-
dataTransferPair.splitAndTransfer(
692-
checkedCastToInt(startPoint), checkedCastToInt(sliceLength));
693-
to.lastSet = length - 1;
694-
to.setValueCount(length);
695716
}
696717

697718
@Override

‎vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,10 +263,12 @@ public void exportCDataBuffers(List<ArrowBuf> buffers, ArrowBuf buffersPtr, long
263263

264264
/** Set the reader and writer indexes for the inner buffers. */
265265
private void setReaderAndWriterIndex() {
266+
final long requiredOffsetBufferCapacity = (long) (valueCount + 1) * OFFSET_WIDTH;
266267
validityBuffer.readerIndex(0);
267268
offsetBuffer.readerIndex(0);
268269
if (valueCount == 0) {
269270
validityBuffer.writerIndex(0);
271+
ensureEmptyOffsetBufferCapacity(requiredOffsetBufferCapacity);
270272
} else {
271273
validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
272274
}
@@ -277,6 +279,19 @@ private void setReaderAndWriterIndex() {
277279
offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
278280
}
279281

282+
private void ensureEmptyOffsetBufferCapacity(long requiredCapacity) {
283+
if (offsetBuffer.capacity() >= requiredCapacity) {
284+
return;
285+
}
286+
long previousOffsetAllocationSizeInBytes = offsetAllocationSizeInBytes;
287+
ArrowBuf oldOffsetBuffer = offsetBuffer;
288+
offsetBuffer = allocateOffsetBuffer(requiredCapacity);
289+
offsetBuffer.setBytes(
290+
0, oldOffsetBuffer, 0, Math.min(oldOffsetBuffer.capacity(), requiredCapacity));
291+
offsetAllocationSizeInBytes = previousOffsetAllocationSizeInBytes;
292+
oldOffsetBuffer.getReferenceManager().release();
293+
}
294+
280295
/**
281296
* Get the inner vectors.
282297
*
@@ -572,6 +587,10 @@ public void splitAndTransfer(int startIndex, int length) {
572587
dataTransferPair.splitAndTransfer(startPoint, sliceLength);
573588
to.lastSet = length - 1;
574589
to.setValueCount(length);
590+
} else {
591+
to.ensureEmptyOffsetBufferCapacity(OFFSET_WIDTH);
592+
dataTransferPair.splitAndTransfer(0, 0);
593+
to.setValueCount(0);
575594
}
576595
}
577596

‎vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -962,6 +962,91 @@ public void testGetBufferSizeFor() {
962962
}
963963
}
964964

965+
@Test
966+
public void testEmptyLargeListOffsetBuffer() {
967+
try (LargeListVector list = LargeListVector.empty("list", allocator)) {
968+
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
969+
list.allocateNew();
970+
list.setValueCount(0);
971+
972+
assertEmptyLargeListOffsetBuffer(list);
973+
}
974+
}
975+
976+
@Test
977+
public void testUnallocatedEmptyLargeListOffsetBuffer() {
978+
try (LargeListVector list = LargeListVector.empty("list", allocator)) {
979+
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
980+
list.setValueCount(0);
981+
982+
assertEmptyLargeListOffsetBuffer(list);
983+
}
984+
}
985+
986+
@Test
987+
public void testSplitAndTransferEmptyLargeListOffsetBuffer() {
988+
try (LargeListVector source = LargeListVector.empty("source", allocator);
989+
LargeListVector target = LargeListVector.empty("target", allocator)) {
990+
source.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
991+
target.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
992+
source.allocateNew();
993+
source.setValueCount(0);
994+
995+
TransferPair transferPair = source.makeTransferPair(target);
996+
transferPair.splitAndTransfer(0, 0);
997+
998+
assertEmptyLargeListOffsetBuffer(target);
999+
}
1000+
}
1001+
1002+
@Test
1003+
public void testSplitAndTransferEmptyLargeListAllocatesOffsetBuffer() {
1004+
try (LargeListVector fromVector = LargeListVector.empty("fromVector", allocator);
1005+
LargeListVector toVector = LargeListVector.empty("toVector", allocator)) {
1006+
fromVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1007+
fromVector.allocateNew();
1008+
fromVector.setValueCount(0);
1009+
1010+
TransferPair transferPair = fromVector.makeTransferPair(toVector);
1011+
transferPair.splitAndTransfer(0, 0);
1012+
1013+
assertAllocatedEmptyLargeListOffsetBuffer(toVector);
1014+
}
1015+
}
1016+
1017+
@Test
1018+
public void testSplitAndTransferEmptyNestedLargeListAllocatesOffsetBuffers() {
1019+
try (LargeListVector fromVector = LargeListVector.empty("fromVector", allocator);
1020+
LargeListVector toVector = LargeListVector.empty("toVector", allocator)) {
1021+
fromVector.addOrGetVector(FieldType.nullable(MinorType.LARGELIST.getType()));
1022+
LargeListVector childVector = (LargeListVector) fromVector.getDataVector();
1023+
childVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1024+
fromVector.allocateNew();
1025+
fromVector.setValueCount(0);
1026+
1027+
TransferPair transferPair = fromVector.makeTransferPair(toVector);
1028+
transferPair.splitAndTransfer(0, 0);
1029+
1030+
assertAllocatedEmptyLargeListOffsetBuffer(toVector);
1031+
assertAllocatedEmptyLargeListOffsetBuffer((LargeListVector) toVector.getDataVector());
1032+
}
1033+
}
1034+
1035+
private ArrowBuf assertEmptyLargeListOffsetBuffer(LargeListVector list) {
1036+
List<ArrowBuf> buffers = list.getFieldBuffers();
1037+
ArrowBuf offsetBuffer = buffers.get(1);
1038+
assertEquals(LargeListVector.OFFSET_WIDTH, offsetBuffer.readableBytes());
1039+
assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH);
1040+
assertEquals(0L, offsetBuffer.getLong(0));
1041+
return offsetBuffer;
1042+
}
1043+
1044+
private void assertAllocatedEmptyLargeListOffsetBuffer(LargeListVector list) {
1045+
ArrowBuf offsetBuffer = list.getOffsetBuffer();
1046+
assertTrue(offsetBuffer.capacity() >= LargeListVector.OFFSET_WIDTH);
1047+
assertEquals(0L, offsetBuffer.getLong(0));
1048+
}
1049+
9651050
@Test
9661051
public void testIsEmpty() {
9671052
try (final LargeListVector vector = LargeListVector.empty("list", allocator)) {
@@ -1101,7 +1186,7 @@ public void testCopyValueSafeForExtensionType() throws Exception {
11011186
}
11021187

11031188
@Test
1104-
public void testEmptyLargeListOffsetBuffer() {
1189+
public void testEmptyLargeListOffsetBufferSize() {
11051190
// Test that LargeListVector has correct readableBytes after allocation.
11061191
// According to Arrow spec, offset buffer must have N+1 entries.
11071192
// Even when N=0, it should contain [0].

‎vector/src/test/java/org/apache/arrow/vector/TestListVector.java‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1142,6 +1142,91 @@ public void testGetBufferSizeFor() {
11421142
}
11431143
}
11441144

1145+
@Test
1146+
public void testEmptyListOffsetBuffer() {
1147+
try (ListVector list = ListVector.empty("list", allocator)) {
1148+
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1149+
list.allocateNew();
1150+
list.setValueCount(0);
1151+
1152+
assertEmptyListOffsetBuffer(list);
1153+
}
1154+
}
1155+
1156+
@Test
1157+
public void testUnallocatedEmptyListOffsetBuffer() {
1158+
try (ListVector list = ListVector.empty("list", allocator)) {
1159+
list.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1160+
list.setValueCount(0);
1161+
1162+
assertEmptyListOffsetBuffer(list);
1163+
}
1164+
}
1165+
1166+
@Test
1167+
public void testSplitAndTransferEmptyListOffsetBuffer() {
1168+
try (ListVector source = ListVector.empty("source", allocator);
1169+
ListVector target = ListVector.empty("target", allocator)) {
1170+
source.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1171+
target.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1172+
source.allocateNew();
1173+
source.setValueCount(0);
1174+
1175+
TransferPair transferPair = source.makeTransferPair(target);
1176+
transferPair.splitAndTransfer(0, 0);
1177+
1178+
assertEmptyListOffsetBuffer(target);
1179+
}
1180+
}
1181+
1182+
@Test
1183+
public void testSplitAndTransferEmptyListAllocatesOffsetBuffer() {
1184+
try (ListVector fromVector = ListVector.empty("fromVector", allocator);
1185+
ListVector toVector = ListVector.empty("toVector", allocator)) {
1186+
fromVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1187+
fromVector.allocateNew();
1188+
fromVector.setValueCount(0);
1189+
1190+
TransferPair transferPair = fromVector.makeTransferPair(toVector);
1191+
transferPair.splitAndTransfer(0, 0);
1192+
1193+
assertAllocatedEmptyListOffsetBuffer(toVector);
1194+
}
1195+
}
1196+
1197+
@Test
1198+
public void testSplitAndTransferEmptyNestedListAllocatesOffsetBuffers() {
1199+
try (ListVector fromVector = ListVector.empty("fromVector", allocator);
1200+
ListVector toVector = ListVector.empty("toVector", allocator)) {
1201+
fromVector.addOrGetVector(FieldType.nullable(MinorType.LIST.getType()));
1202+
ListVector childVector = (ListVector) fromVector.getDataVector();
1203+
childVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
1204+
fromVector.allocateNew();
1205+
fromVector.setValueCount(0);
1206+
1207+
TransferPair transferPair = fromVector.makeTransferPair(toVector);
1208+
transferPair.splitAndTransfer(0, 0);
1209+
1210+
assertAllocatedEmptyListOffsetBuffer(toVector);
1211+
assertAllocatedEmptyListOffsetBuffer((ListVector) toVector.getDataVector());
1212+
}
1213+
}
1214+
1215+
private ArrowBuf assertEmptyListOffsetBuffer(ListVector list) {
1216+
List<ArrowBuf> buffers = list.getFieldBuffers();
1217+
ArrowBuf offsetBuffer = buffers.get(1);
1218+
assertEquals(BaseRepeatedValueVector.OFFSET_WIDTH, offsetBuffer.readableBytes());
1219+
assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH);
1220+
assertEquals(0, offsetBuffer.getInt(0));
1221+
return offsetBuffer;
1222+
}
1223+
1224+
private void assertAllocatedEmptyListOffsetBuffer(ListVector list) {
1225+
ArrowBuf offsetBuffer = list.getOffsetBuffer();
1226+
assertTrue(offsetBuffer.capacity() >= BaseRepeatedValueVector.OFFSET_WIDTH);
1227+
assertEquals(0, offsetBuffer.getInt(0));
1228+
}
1229+
11451230
@Test
11461231
public void testIsEmpty() {
11471232
try (final ListVector vector = ListVector.empty("list", allocator)) {
@@ -1380,7 +1465,7 @@ public void testCopyValueSafeForExtensionType() throws Exception {
13801465
}
13811466

13821467
@Test
1383-
public void testEmptyListOffsetBuffer() {
1468+
public void testEmptyListOffsetBufferSize() {
13841469
// Test that ListVector has correct readableBytes after allocation.
13851470
// According to Arrow spec, offset buffer must have N+1 entries.
13861471
// Even when N=0, it should contain [0].

0 commit comments

Comments
 (0)