DRILL-7341: Vector reAlloc may fail after exchange
closes #1838
diff --git a/exec/vector/src/main/codegen/templates/FixedValueVectors.java b/exec/vector/src/main/codegen/templates/FixedValueVectors.java
index 508e484..5741555 100644
--- a/exec/vector/src/main/codegen/templates/FixedValueVectors.java
+++ b/exec/vector/src/main/codegen/templates/FixedValueVectors.java
@@ -210,10 +210,17 @@
// a zero-length buffer. Instead, just allocate a 256 byte
// buffer if we start at 0.
- final long newAllocationSize = allocationSizeInBytes == 0
+ long newAllocationSize = allocationSizeInBytes == 0
? 256
: allocationSizeInBytes * 2L;
+ final int currentCapacity = data.capacity();
+ // Some operations, such as Value Vector#exchange, can be change DrillBuf data field without corresponding allocation size changes.
+ // Check that the size of the allocation is sufficient to copy the old buffer.
+ while (newAllocationSize < currentCapacity) {
+ newAllocationSize *= 2L;
+ }
+
// TODO: Replace this with MAX_BUFFER_SIZE once all
// code is aware of the maximum vector size.
@@ -222,8 +229,7 @@
}
reallocRaw((int) newAllocationSize);
- final int halfNewCapacity = data.capacity() / 2;
- data.setZero(halfNewCapacity, halfNewCapacity);
+ data.setZero(currentCapacity, data.capacity() - currentCapacity);
}
@Override
diff --git a/exec/vector/src/main/codegen/templates/VariableLengthVectors.java b/exec/vector/src/main/codegen/templates/VariableLengthVectors.java
index 8dd8eb1..1bb758f 100644
--- a/exec/vector/src/main/codegen/templates/VariableLengthVectors.java
+++ b/exec/vector/src/main/codegen/templates/VariableLengthVectors.java
@@ -390,7 +390,14 @@
}
public void reAlloc() {
- final long newAllocationSize = allocationSizeInBytes*2L;
+ long newAllocationSize = allocationSizeInBytes*2L;
+
+ // Some operations, such as Value Vector#exchange, can be change DrillBuf data field without corresponding allocation size changes.
+ // Check that the size of the allocation is sufficient to copy the old buffer.
+ while (newAllocationSize < data.capacity()) {
+ newAllocationSize *= 2L;
+ }
+
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
throw new OversizedAllocationException("Unable to expand the buffer. Max allowed buffer size is reached.");
}
diff --git a/exec/vector/src/main/java/org/apache/drill/exec/vector/BitVector.java b/exec/vector/src/main/java/org/apache/drill/exec/vector/BitVector.java
index 634e4d6..7b2cd28 100644
--- a/exec/vector/src/main/java/org/apache/drill/exec/vector/BitVector.java
+++ b/exec/vector/src/main/java/org/apache/drill/exec/vector/BitVector.java
@@ -179,7 +179,14 @@
* Allocate new buffer with double capacity, and copy data into the new buffer. Replace vector's buffer with new buffer, and release old one
*/
public void reAlloc() {
- final long newAllocationSize = allocationSizeInBytes * 2L;
+ long newAllocationSize = allocationSizeInBytes * 2L;
+
+ // Some operations, such as Value Vector#exchange, can be change DrillBuf data field without corresponding allocation size changes.
+ // Check that the size of the allocation is sufficient to copy the old buffer.
+ while (newAllocationSize < data.capacity()) {
+ newAllocationSize *= 2L;
+ }
+
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
throw new OversizedAllocationException("Requested amount of memory is more than max allowed allocation size");
}
diff --git a/exec/vector/src/test/java/org/apache/drill/exec/vector/VariableLengthVectorTest.java b/exec/vector/src/test/java/org/apache/drill/exec/vector/VariableLengthVectorTest.java
index 4c6aeed..c77854f 100644
--- a/exec/vector/src/test/java/org/apache/drill/exec/vector/VariableLengthVectorTest.java
+++ b/exec/vector/src/test/java/org/apache/drill/exec/vector/VariableLengthVectorTest.java
@@ -89,6 +89,30 @@
}
}
+ @Test
+ public void testDRILL7341() {
+ try (RootAllocator allocator = new RootAllocator(10_000_000)) {
+ final MaterializedField field = MaterializedField.create("stringCol", Types.optional(TypeProtos.MinorType.VARCHAR));
+ final NullableVarCharVector sourceVector = new NullableVarCharVector(field, allocator);
+ final NullableVarCharVector targetVector = new NullableVarCharVector(field, allocator);
+
+ sourceVector.allocateNew();
+ targetVector.allocateNew();
+
+ try {
+ final NullableVarCharVector.Mutator sourceMutator = sourceVector.getMutator();
+ sourceMutator.setValueCount(sourceVector.getValueCapacity() * 4);
+
+ targetVector.exchange(sourceVector);
+ final NullableVarCharVector.Mutator targetMutator = targetVector.getMutator();
+ targetMutator.setValueCount(targetVector.getValueCapacity() * 2);
+ } finally {
+ sourceVector.clear();
+ targetVector.clear();
+ }
+ }
+ }
+
/**
* Set 10000 values. Then go back and set new values starting at the 1001 the record.
*/
