Skip to content

Commit 4e7e887

Browse files
author
Yicong Huang
committed
fix: ensure offset buffer has at least one entry for empty variable-width vectors
1 parent 855f7ad commit 4e7e887

File tree

8 files changed

+126
-88
lines changed

8 files changed

+126
-88
lines changed

‎adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -43,15 +43,19 @@ public void testZeroRowResultSet() throws Exception {
4343
.setReuseVectorSchemaRoot(reuseVectorSchemaRoot)
4444
.build();
4545

46-
ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config);
47-
assertTrue(iter.hasNext(), "Iterator on zero row ResultSet should haveNext() before use");
48-
VectorSchemaRoot root = iter.next();
49-
assertNotNull(root, "VectorSchemaRoot from first next() result should never be null");
50-
assertEquals(
51-
0, root.getRowCount(), "VectorSchemaRoot from empty ResultSet should have zero rows");
52-
assertFalse(
53-
iter.hasNext(),
54-
"hasNext() should return false on empty ResultSets after initial next() call");
46+
try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) {
47+
assertTrue(iter.hasNext(), "Iterator on zero row ResultSet should haveNext() before use");
48+
VectorSchemaRoot root = iter.next();
49+
assertNotNull(root, "VectorSchemaRoot from first next() result should never be null");
50+
assertEquals(
51+
0, root.getRowCount(), "VectorSchemaRoot from empty ResultSet should have zero rows");
52+
assertFalse(
53+
iter.hasNext(),
54+
"hasNext() should return false on empty ResultSets after initial next() call");
55+
if (!reuseVectorSchemaRoot) {
56+
root.close();
57+
}
58+
}
5559
}
5660
}
5761
}

‎flight/flight-core/src/test/java/org/apache/arrow/flight/TestDictionaryUtils.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,10 @@ public void testReuseSchema() {
5555

5656
@Test
5757
public void testCreateSchema() {
58-
try (BufferAllocator allocator = new RootAllocator(1024)) {
58+
try (BufferAllocator allocator = new RootAllocator(1024);
59+
VarCharVector dictVec = new VarCharVector("dict vector", allocator)) {
5960
DictionaryEncoding dictionaryEncoding =
6061
new DictionaryEncoding(0, true, new ArrowType.Int(8, true));
61-
VarCharVector dictVec = new VarCharVector("dict vector", allocator);
6262
Dictionary dictionary = new Dictionary(dictVec, dictionaryEncoding);
6363
DictionaryProvider dictProvider = new DictionaryProvider.MapDictionaryProvider(dictionary);
6464
TreeSet<Long> dictionaryUsed = new TreeSet<>();

‎vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,7 @@ private void setReaderAndWriterIndex() {
379379
} else {
380380
final long lastDataOffset = getStartOffset(valueCount);
381381
validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount));
382+
offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH);
382383
valueBuffer.writerIndex(lastDataOffset);
383384
}
384385
// IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`.

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -509,7 +509,11 @@ public void close() {
509509
/** Initializes the struct's members from the given Fields. */
510510
public void initializeChildrenFromFields(List<Field> children) {
511511
for (Field field : children) {
512+
FieldVector oldVector = getChild(field.getName());
512513
FieldVector vector = (FieldVector) this.add(field.getName(), field.getFieldType());
514+
if (oldVector != null && oldVector != vector) {
515+
oldVector.close();
516+
}
513517
vector.initializeChildrenFromFields(field.getChildren());
514518
}
515519
}

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

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -412,23 +412,25 @@ public void testGetFieldTypeInfo() throws Exception {
412412
final Field field = new Field("union", fieldType, children);
413413

414414
MinorType minorType = MinorType.DENSEUNION;
415-
DenseUnionVector vector = (DenseUnionVector) minorType.getNewVector(field, allocator, null);
416-
vector.initializeChildrenFromFields(children);
415+
try (DenseUnionVector vector =
416+
(DenseUnionVector) minorType.getNewVector(field, allocator, null)) {
417+
vector.initializeChildrenFromFields(children);
417418

418-
assertEquals(vector.getField(), field);
419+
assertEquals(vector.getField(), field);
419420

420-
// Union has 2 child vectors
421-
assertEquals(2, vector.size());
421+
// Union has 2 child vectors
422+
assertEquals(2, vector.size());
422423

423-
// Check child field 0
424-
VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int");
425-
assertEquals(0, intChild.ordinal);
426-
assertEquals(intChild.vector.getField(), children.get(0));
424+
// Check child field 0
425+
VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int");
426+
assertEquals(0, intChild.ordinal);
427+
assertEquals(intChild.vector.getField(), children.get(0));
427428

428-
// Check child field 1
429-
VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar");
430-
assertEquals(1, varcharChild.ordinal);
431-
assertEquals(varcharChild.vector.getField(), children.get(1));
429+
// Check child field 1
430+
VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar");
431+
assertEquals(1, varcharChild.ordinal);
432+
assertEquals(varcharChild.vector.getField(), children.get(1));
433+
}
432434
}
433435

434436
@Test

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

Lines changed: 68 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -111,68 +111,91 @@ private void populateDenseUnionVector(final DenseUnionVector vector, int valueCo
111111
@Test
112112
public void testWithEmptyVector() {
113113
// MapVector use TransferImpl from ListVector
114-
ListVector listVector = ListVector.empty("", allocator);
115-
TransferPair transferPair = listVector.getTransferPair(allocator);
116-
transferPair.splitAndTransfer(0, 0);
117-
assertEquals(0, transferPair.getTo().getValueCount());
114+
try (ListVector listVector = ListVector.empty("", allocator)) {
115+
TransferPair transferPair = listVector.getTransferPair(allocator);
116+
transferPair.splitAndTransfer(0, 0);
117+
assertEquals(0, transferPair.getTo().getValueCount());
118+
transferPair.getTo().close();
119+
}
118120
// BaseFixedWidthVector
119-
IntVector intVector = new IntVector("", allocator);
120-
transferPair = intVector.getTransferPair(allocator);
121-
transferPair.splitAndTransfer(0, 0);
122-
assertEquals(0, transferPair.getTo().getValueCount());
121+
try (IntVector intVector = new IntVector("", allocator)) {
122+
TransferPair transferPair = intVector.getTransferPair(allocator);
123+
transferPair.splitAndTransfer(0, 0);
124+
assertEquals(0, transferPair.getTo().getValueCount());
125+
transferPair.getTo().close();
126+
}
123127
// BaseVariableWidthVector
124-
VarCharVector varCharVector = new VarCharVector("", allocator);
125-
transferPair = varCharVector.getTransferPair(allocator);
126-
transferPair.splitAndTransfer(0, 0);
127-
assertEquals(0, transferPair.getTo().getValueCount());
128+
try (VarCharVector varCharVector = new VarCharVector("", allocator)) {
129+
TransferPair transferPair = varCharVector.getTransferPair(allocator);
130+
transferPair.splitAndTransfer(0, 0);
131+
assertEquals(0, transferPair.getTo().getValueCount());
132+
transferPair.getTo().close();
133+
}
128134
// BaseVariableWidthViewVector: ViewVarCharVector
129-
ViewVarCharVector viewVarCharVector = new ViewVarCharVector("", allocator);
130-
transferPair = viewVarCharVector.getTransferPair(allocator);
131-
transferPair.splitAndTransfer(0, 0);
132-
assertEquals(0, transferPair.getTo().getValueCount());
135+
try (ViewVarCharVector viewVarCharVector = new ViewVarCharVector("", allocator)) {
136+
TransferPair transferPair = viewVarCharVector.getTransferPair(allocator);
137+
transferPair.splitAndTransfer(0, 0);
138+
assertEquals(0, transferPair.getTo().getValueCount());
139+
transferPair.getTo().close();
140+
}
133141
// BaseVariableWidthVector: ViewVarBinaryVector
134-
ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("", allocator);
135-
transferPair = viewVarBinaryVector.getTransferPair(allocator);
136-
transferPair.splitAndTransfer(0, 0);
137-
assertEquals(0, transferPair.getTo().getValueCount());
142+
try (ViewVarBinaryVector viewVarBinaryVector = new ViewVarBinaryVector("", allocator)) {
143+
TransferPair transferPair = viewVarBinaryVector.getTransferPair(allocator);
144+
transferPair.splitAndTransfer(0, 0);
145+
assertEquals(0, transferPair.getTo().getValueCount());
146+
transferPair.getTo().close();
147+
}
138148
// BaseLargeVariableWidthVector
139-
LargeVarCharVector largeVarCharVector = new LargeVarCharVector("", allocator);
140-
transferPair = largeVarCharVector.getTransferPair(allocator);
141-
transferPair.splitAndTransfer(0, 0);
142-
assertEquals(0, transferPair.getTo().getValueCount());
149+
try (LargeVarCharVector largeVarCharVector = new LargeVarCharVector("", allocator)) {
150+
TransferPair transferPair = largeVarCharVector.getTransferPair(allocator);
151+
transferPair.splitAndTransfer(0, 0);
152+
assertEquals(0, transferPair.getTo().getValueCount());
153+
transferPair.getTo().close();
154+
}
143155
// StructVector
144-
StructVector structVector = StructVector.empty("", allocator);
145-
transferPair = structVector.getTransferPair(allocator);
146-
transferPair.splitAndTransfer(0, 0);
147-
assertEquals(0, transferPair.getTo().getValueCount());
156+
try (StructVector structVector = StructVector.empty("", allocator)) {
157+
TransferPair transferPair = structVector.getTransferPair(allocator);
158+
transferPair.splitAndTransfer(0, 0);
159+
assertEquals(0, transferPair.getTo().getValueCount());
160+
transferPair.getTo().close();
161+
}
148162
// FixedSizeListVector
149-
FixedSizeListVector fixedSizeListVector = FixedSizeListVector.empty("", 0, allocator);
150-
transferPair = fixedSizeListVector.getTransferPair(allocator);
151-
transferPair.splitAndTransfer(0, 0);
152-
assertEquals(0, transferPair.getTo().getValueCount());
163+
try (FixedSizeListVector fixedSizeListVector = FixedSizeListVector.empty("", 0, allocator)) {
164+
TransferPair transferPair = fixedSizeListVector.getTransferPair(allocator);
165+
transferPair.splitAndTransfer(0, 0);
166+
assertEquals(0, transferPair.getTo().getValueCount());
167+
transferPair.getTo().close();
168+
}
153169
// FixedSizeBinaryVector
154-
FixedSizeBinaryVector fixedSizeBinaryVector = new FixedSizeBinaryVector("", allocator, 4);
155-
transferPair = fixedSizeBinaryVector.getTransferPair(allocator);
156-
transferPair.splitAndTransfer(0, 0);
157-
assertEquals(0, transferPair.getTo().getValueCount());
170+
try (FixedSizeBinaryVector fixedSizeBinaryVector =
171+
new FixedSizeBinaryVector("", allocator, 4)) {
172+
TransferPair transferPair = fixedSizeBinaryVector.getTransferPair(allocator);
173+
transferPair.splitAndTransfer(0, 0);
174+
assertEquals(0, transferPair.getTo().getValueCount());
175+
transferPair.getTo().close();
176+
}
158177
// UnionVector
159-
UnionVector unionVector = UnionVector.empty("", allocator);
160-
transferPair = unionVector.getTransferPair(allocator);
161-
transferPair.splitAndTransfer(0, 0);
162-
assertEquals(0, transferPair.getTo().getValueCount());
178+
try (UnionVector unionVector = UnionVector.empty("", allocator)) {
179+
TransferPair transferPair = unionVector.getTransferPair(allocator);
180+
transferPair.splitAndTransfer(0, 0);
181+
assertEquals(0, transferPair.getTo().getValueCount());
182+
transferPair.getTo().close();
183+
}
163184
// DenseUnionVector
164-
DenseUnionVector duv = DenseUnionVector.empty("", allocator);
165-
transferPair = duv.getTransferPair(allocator);
166-
transferPair.splitAndTransfer(0, 0);
167-
assertEquals(0, transferPair.getTo().getValueCount());
185+
try (DenseUnionVector duv = DenseUnionVector.empty("", allocator)) {
186+
TransferPair transferPair = duv.getTransferPair(allocator);
187+
transferPair.splitAndTransfer(0, 0);
188+
assertEquals(0, transferPair.getTo().getValueCount());
189+
transferPair.getTo().close();
190+
}
168191

169192
// non empty from vector
170193

171194
// BaseFixedWidthVector
172195
IntVector fromIntVector = new IntVector("", allocator);
173196
fromIntVector.allocateNew(100);
174197
populateIntVector(fromIntVector, 100);
175-
transferPair = fromIntVector.getTransferPair(allocator);
198+
TransferPair transferPair = fromIntVector.getTransferPair(allocator);
176199
IntVector toIntVector = (IntVector) transferPair.getTo();
177200
transferPair.splitAndTransfer(0, 0);
178201
assertEquals(0, toIntVector.getValueCount());

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

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -402,23 +402,24 @@ public void testGetFieldTypeInfo() throws Exception {
402402
final Field field = new Field("union", fieldType, children);
403403

404404
MinorType minorType = MinorType.UNION;
405-
UnionVector vector = (UnionVector) minorType.getNewVector(field, allocator, null);
406-
vector.initializeChildrenFromFields(children);
405+
try (UnionVector vector = (UnionVector) minorType.getNewVector(field, allocator, null)) {
406+
vector.initializeChildrenFromFields(children);
407407

408-
assertTrue(vector.getField().equals(field));
408+
assertTrue(vector.getField().equals(field));
409409

410-
// Union has 2 child vectors
411-
assertEquals(2, vector.size());
410+
// Union has 2 child vectors
411+
assertEquals(2, vector.size());
412412

413-
// Check child field 0
414-
VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int");
415-
assertEquals(0, intChild.ordinal);
416-
assertEquals(intChild.vector.getField(), children.get(0));
413+
// Check child field 0
414+
VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int");
415+
assertEquals(0, intChild.ordinal);
416+
assertEquals(intChild.vector.getField(), children.get(0));
417417

418-
// Check child field 1
419-
VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar");
420-
assertEquals(1, varcharChild.ordinal);
421-
assertEquals(varcharChild.vector.getField(), children.get(1));
418+
// Check child field 1
419+
VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar");
420+
assertEquals(1, varcharChild.ordinal);
421+
assertEquals(varcharChild.vector.getField(), children.get(1));
422+
}
422423
}
423424

424425
@Test

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1154,9 +1154,10 @@ public void testSplitAndTransfer1() {
11541154
final int dataRefCnt = sourceVector.getDataBuffer().refCnt();
11551155

11561156
// split and transfer with slice starting at the beginning: this should not allocate
1157-
// anything new
1157+
// anything new. The target vector's initial offset buffer (4 bytes) is freed during
1158+
// transfer.
11581159
sourceVector.splitAndTransferTo(0, 2, targetVector);
1159-
assertEquals(allocatedMem, allocator.getAllocatedMemory());
1160+
assertEquals(allocatedMem - 4, allocator.getAllocatedMemory());
11601161
// The validity and offset buffers are sliced from a same buffer.See
11611162
// BaseFixedWidthVector#allocateBytes.
11621163
// Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0.
@@ -1192,9 +1193,10 @@ public void testSplitAndTransfer2() {
11921193
final int dataRefCnt = sourceVector.getDataBuffer().refCnt();
11931194

11941195
// split and transfer with slice starting at the beginning: this should not allocate
1195-
// anything new
1196+
// anything new. The target vector's initial offset buffer (4 bytes) is freed during
1197+
// transfer.
11961198
sourceVector.splitAndTransferTo(0, 2, targetVector);
1197-
assertEquals(allocatedMem, allocator.getAllocatedMemory());
1199+
assertEquals(allocatedMem - 4, allocator.getAllocatedMemory());
11981200
// The validity and offset buffers are sliced from a same buffer.See
11991201
// BaseFixedWidthVector#allocateBytes.
12001202
// Therefore, the refcnt of the validity buffer is increased once since the startIndex is 0.
@@ -1239,7 +1241,8 @@ public void testSplitAndTransfer3() {
12391241
final long validitySize =
12401242
DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize(
12411243
getValidityBufferSizeFromCount(2));
1242-
assertEquals(allocatedMem + validitySize, allocator.getAllocatedMemory());
1244+
// The target vector's initial offset buffer (4 bytes) is freed during transfer.
1245+
assertEquals(allocatedMem + validitySize - 4, allocator.getAllocatedMemory());
12431246
// The validity and offset buffers are sliced from a same buffer.See
12441247
// BaseFixedWidthVector#allocateBytes.
12451248
// Since values up to the startIndex are empty/null, the offset buffer doesn't need to be
@@ -3701,7 +3704,7 @@ public void testEmptyBufBehavior() {
37013704
assertEquals(1, vector.getOffsetBuffer().refCnt());
37023705
assertEquals(0, vector.getDataBuffer().capacity());
37033706
assertEquals(0, vector.getValidityBuffer().capacity());
3704-
assertEquals(0, vector.getOffsetBuffer().capacity());
3707+
assertEquals(4, vector.getOffsetBuffer().capacity());
37053708

37063709
vector.allocateNew(valueCount);
37073710
assertEquals(1, vector.getDataBuffer().refCnt());

0 commit comments

Comments
 (0)