GH-721: Allow using 1GB+ data buffers in variable width vectors (#722)

## What's Changed

Allow actually reaching MAX_BUFFER_SIZE at reallocating variable width
vectors instead of exceeding it calculating the next power of 2.
For unit testing the maximum allocation size has be increased to 2MB -
1byte to simulate the default maximum behavior. Due to this change
needed some updates in existing unit tests because of the round ups used
at calculating the required buffer sizes.

Closes #721.

Author: gszadovszky

pom.xml

@@ -327,8 +327,8 @@ under the License.
               <io.netty.tryReflectionSetAccessible>true</io.netty.tryReflectionSetAccessible>
               <user.timezone>UTC</user.timezone>
               <!-- Note: changing the below configuration might increase the max allocation size for a vector
-              which in turn can cause OOM. -->
-              <arrow.vector.max_allocation_bytes>1048576</arrow.vector.max_allocation_bytes>
+              which in turn can cause OOM. Using 2MB - 1byte to simulate the defaul limit of 2^31 - 1 bytes. -->
+              <arrow.vector.max_allocation_bytes>2097151</arrow.vector.max_allocation_bytes>
             </systemPropertyVariables>
             <useModulePath>false</useModulePath>
           </configuration>

vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java

@@ -571,10 +571,13 @@ public void reallocDataBuffer(long desiredAllocSize) {
       return;
     }
 
-    final long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize);
+    final long newAllocationSize =
+        Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE);
     assert newAllocationSize >= 1;
 
-    checkDataBufferSize(newAllocationSize);
+    if (newAllocationSize < desiredAllocSize) {
+      checkDataBufferSize(desiredAllocSize);
+    }
 
     final ArrowBuf newBuf = allocator.buffer(newAllocationSize);
     newBuf.setBytes(0, valueBuffer, 0, valueBuffer.capacity());

vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java

@@ -550,15 +550,18 @@ public void reallocViewBuffer(long desiredAllocSize) {
     if (desiredAllocSize == 0) {
       return;
     }
-    long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize);
+    long newAllocationSize = Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE);
     assert newAllocationSize >= 1;
 
-    checkDataBufferSize(newAllocationSize);
     // for each set operation, we have to allocate 16 bytes
     // here we are adjusting the desired allocation-based allocation size
     // to align with the 16bytes requirement.
     newAllocationSize = roundUpToMultipleOf16(newAllocationSize);
 
+    if (newAllocationSize < desiredAllocSize) {
+      checkDataBufferSize(desiredAllocSize);
+    }
+
     final ArrowBuf newBuf = allocator.buffer(newAllocationSize);
     newBuf.setBytes(0, viewBuffer, 0, viewBuffer.capacity());
 
@@ -587,10 +590,13 @@ public void reallocViewDataBuffer(long desiredAllocSize) {
       return;
     }
 
-    final long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize);
+    final long newAllocationSize =
+        Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE);
     assert newAllocationSize >= 1;
 
-    checkDataBufferSize(newAllocationSize);
+    if (newAllocationSize < desiredAllocSize) {
+      checkDataBufferSize(desiredAllocSize);
+    }
 
     final ArrowBuf newBuf = allocator.buffer(newAllocationSize);
     dataBuffers.add(newBuf);

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

@@ -95,7 +95,7 @@ public void init() {
   private static final byte[] STR5 = "EEE5".getBytes(utf8Charset);
   private static final byte[] STR6 = "FFFFF6".getBytes(utf8Charset);
   private static final int MAX_VALUE_COUNT =
-      (int) (Integer.getInteger("arrow.vector.max_allocation_bytes", Integer.MAX_VALUE) / 7);
+      (int) (Integer.getInteger("arrow.vector.max_allocation_bytes", Integer.MAX_VALUE) / 9);
   private static final int MAX_VALUE_COUNT_8BYTE = (int) (MAX_VALUE_COUNT / 2);
 
   @AfterEach

vector/src/test/java/org/apache/arrow/vector/TestVectorReAlloc.java

@@ -24,6 +24,7 @@
 import java.nio.charset.StandardCharsets;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.memory.util.CommonUtil;
 import org.apache.arrow.vector.complex.DenseUnionVector;
 import org.apache.arrow.vector.complex.FixedSizeListVector;
 import org.apache.arrow.vector.complex.ListVector;
@@ -222,6 +223,17 @@ public void testVariableAllocateAfterReAlloc() throws Exception {
     }
   }
 
+  @Test
+  public void testVariableReAllocAbove1GB() throws Exception {
+    try (final VarCharVector vector = new VarCharVector("", allocator)) {
+      long desiredSizeAboveLastPowerOf2 =
+          CommonUtil.nextPowerOfTwo(BaseVariableWidthVector.MAX_ALLOCATION_SIZE) / 2 + 1;
+      vector.reallocDataBuffer(desiredSizeAboveLastPowerOf2);
+
+      assertTrue(vector.getDataBuffer().capacity() >= desiredSizeAboveLastPowerOf2);
+    }
+  }
+
   @Test
   public void testLargeVariableAllocateAfterReAlloc() throws Exception {
     try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) {

vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java

@@ -28,6 +28,7 @@
 import java.util.stream.Stream;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.memory.util.CommonUtil;
 import org.apache.arrow.vector.BaseValueVector;
 import org.apache.arrow.vector.BaseVariableWidthViewVector;
 import org.apache.arrow.vector.BigIntVector;
@@ -309,7 +310,15 @@ public void testAppendEmptyVariableWidthVector() {
 
   @Test
   public void testAppendLargeAndSmallVariableVectorsWithinLimit() {
-    int sixteenthOfMaxAllocation = Math.toIntExact(BaseValueVector.MAX_ALLOCATION_SIZE / 16);
+    // Using the max power of 2 allocation size to avoid hitting the max limit at round ups
+    long maxPowerOfTwoAllocationSize =
+        CommonUtil.nextPowerOfTwo(BaseValueVector.MAX_ALLOCATION_SIZE);
+    if (maxPowerOfTwoAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) {
+      maxPowerOfTwoAllocationSize =
+          CommonUtil.nextPowerOfTwo(BaseValueVector.MAX_ALLOCATION_SIZE / 2);
+    }
+
+    int sixteenthOfMaxAllocation = Math.toIntExact(maxPowerOfTwoAllocationSize / 16);
     try (VarCharVector target = makeVarCharVec(1, sixteenthOfMaxAllocation);
         VarCharVector delta = makeVarCharVec(sixteenthOfMaxAllocation, 1)) {
       new VectorAppender(delta).visit(target, null);